code: """SQL Structured Store."""
  
  from collections import defaultdict
  from enum import Enum
  from typing import Any, Optional, Sequence, Union
  
  from sqlalchemy import Table
  
  from llama_index.legacy.core.base_query_engine import BaseQueryEngine
  from llama_index.legacy.core.base_retriever import BaseRetriever
  from llama_index.legacy.data_structs.table import SQLStructTable
  from llama_index.legacy.indices.common.struct_store.schema import SQLContextContainer
  from llama_index.legacy.indices.common.struct_store.sql import (
      SQLStructDatapointExtractor,
  )
  from llama_index.legacy.indices.struct_store.base import BaseStructStoreIndex
  from llama_index.legacy.indices.struct_store.container_builder import (
      SQLContextContainerBuilder,
  )
  from llama_index.legacy.schema import BaseNode
  from llama_index.legacy.service_context import ServiceContext
  from llama_index.legacy.utilities.sql_wrapper import SQLDatabase
  
  
  class SQLQueryMode(str, Enum):
      SQL = "sql"
      NL = "nl"
  
  
  class SQLStructStoreIndex(BaseStructStoreIndex[SQLStructTable]):
      """SQL Struct Store Index.
  
      The SQLStructStoreIndex is an index that uses a SQL database
      under the hood. During index construction, the data can be inferred
      from unstructured documents given a schema extract prompt,
      or it can be pre-loaded in the database.
  
      During query time, the user can either specify a raw SQL query
      or a natural language query to retrieve their data.
  
      NOTE: this is deprecated.
  
      Args:
          documents (Optional[Sequence[DOCUMENTS_INPUT]]): Documents to index.
              NOTE: in the SQL index, this is an optional field.
          sql_database (Optional[SQLDatabase]): SQL database to use,
              including table names to specify.
              See :ref:`Ref-Struct-Store` for more details.
          table_name (Optional[str]): Name of the table to use
              for extracting data.
              Either table_name or table must be specified.
          table (Optional[Table]): SQLAlchemy Table object to use.
              Specifying the Table object explicitly, instead of
              the table name, allows you to pass in a view.
              Either table_name or table must be specified.
          sql_context_container (Optional[SQLContextContainer]): SQL context container.
              an be generated from a SQLContextContainerBuilder.
              See :ref:`Ref-Struct-Store` for more details.
  
      """
  
      index_struct_cls = SQLStructTable
  
      def __init__(
          self,
          nodes: Optional[Sequence[BaseNode]] = None,
          index_struct: Optional[SQLStructTable] = None,
          service_context: Optional[ServiceContext] = None,
          sql_database: Optional[SQLDatabase] = None,
          table_name: Optional[str] = None,
          table: Optional[Table] = None,
          ref_doc_id_column: Optional[str] = None,
          sql_context_container: Optional[SQLContextContainer] = None,
          **kwargs: Any,
      ) -> None:
          """Initialize params."""
          if sql_database is None:
              raise ValueError("sql_database must be specified")
          self.sql_database = sql_database
          # needed here for data extractor
          self._ref_doc_id_column = ref_doc_id_column
          self._table_name = table_name
          self._table = table
  
          # if documents aren't specified, pass in a blank []
          if index_struct is None:
              nodes = nodes or []
  
          super().__init__(
              nodes=nodes,
              index_struct=index_struct,
              service_context=service_context,
              **kwargs,
          )
  
          # TODO: index_struct context_dict is deprecated,
          # we're migrating storage of information to here.
          if sql_context_container is None:
              container_builder = SQLContextContainerBuilder(sql_database)
              sql_context_container = container_builder.build_context_container()
          self.sql_context_container = sql_context_container
  
      @property
      def ref_doc_id_column(self) -> Optional[str]:
          return self._ref_doc_id_column
  
      def _build_index_from_nodes(self, nodes: Sequence[BaseNode]) -> SQLStructTable:
          """Build index from nodes."""
          index_struct = self.index_struct_cls()
          if len(nodes) == 0:
              return index_struct
          else:
              data_extractor = SQLStructDatapointExtractor(
                  self._service_context.llm,
                  self.schema_extract_prompt,
                  self.output_parser,
                  self.sql_database,
                  table_name=self._table_name,
                  table=self._table,
                  ref_doc_id_column=self._ref_doc_id_column,
              )
              # group nodes by ids
              source_to_node = defaultdict(list)
              for node in nodes:
                  source_to_node[node.ref_doc_id].append(node)
  
              for node_set in source_to_node.values():
                  data_extractor.insert_datapoint_from_nodes(node_set)
          return index_struct
  
      def _insert(self, nodes: Sequence[BaseNode], **insert_kwargs: Any) -> None:
          """Insert a document."""
          data_extractor = SQLStructDatapointExtractor(
              self._service_context.llm,
              self.schema_extract_prompt,
              self.output_parser,
              self.sql_database,
              table_name=self._table_name,
              table=self._table,
              ref_doc_id_column=self._ref_doc_id_column,
          )
          data_extractor.insert_datapoint_from_nodes(nodes)
  
      def as_retriever(self, **kwargs: Any) -> BaseRetriever:
          raise NotImplementedError("Not supported")
  
      def as_query_engine(
          self, query_mode: Union[str, SQLQueryMode] = SQLQueryMode.NL, **kwargs: Any
      ) -> BaseQueryEngine:
          # NOTE: lazy import
          from llama_index.legacy.indices.struct_store.sql_query import (
              NLStructStoreQueryEngine,
              SQLStructStoreQueryEngine,
          )
  
          if query_mode == SQLQueryMode.NL:
              return NLStructStoreQueryEngine(self, **kwargs)
          elif query_mode == SQLQueryMode.SQL:
              return SQLStructStoreQueryEngine(self, **kwargs)
          else:
              raise ValueError(f"Unknown query mode: {query_mode}")
  
  
  GPTSQLStructStoreIndex = SQLStructStoreIndex
  
apis: llama_index.legacy.indices.struct_store.container_builder.SQLContextContainerBuilder llama_index.legacy.indices.common.struct_store.sql.SQLStructDatapointExtractor llama_index.legacy.indices.struct_store.sql_query.NLStructStoreQueryEngine llama_index.legacy.indices.struct_store.sql_query.SQLStructStoreQueryEngine
extract_api: [((5106, 5332), 'llama_index.legacy.indices.common.struct_store.sql.SQLStructDatapointExtractor', 'SQLStructDatapointExtractor', (['self._service_context.llm', 'self.schema_extract_prompt', 'self.output_parser', 'self.sql_database'], {'table_name': 'self._table_name', 'table': 'self._table', 'ref_doc_id_column': 'self._ref_doc_id_column'}), '(self._service_context.llm, self.\n    schema_extract_prompt, self.output_parser, self.sql_database,\n    table_name=self._table_name, table=self._table, ref_doc_id_column=self.\n    _ref_doc_id_column)\n', (5133, 5332), False, 'from llama_index.legacy.indices.common.struct_store.sql import SQLStructDatapointExtractor\n'), ((3747, 3787), 'llama_index.legacy.indices.struct_store.container_builder.SQLContextContainerBuilder', 'SQLContextContainerBuilder', (['sql_database'], {}), '(sql_database)\n', (3773, 3787), False, 'from llama_index.legacy.indices.struct_store.container_builder import SQLContextContainerBuilder\n'), ((4304, 4530), 'llama_index.legacy.indices.common.struct_store.sql.SQLStructDatapointExtractor', 'SQLStructDatapointExtractor', (['self._service_context.llm', 'self.schema_extract_prompt', 'self.output_parser', 'self.sql_database'], {'table_name': 'self._table_name', 'table': 'self._table', 'ref_doc_id_column': 'self._ref_doc_id_column'}), '(self._service_context.llm, self.\n    schema_extract_prompt, self.output_parser, self.sql_database,\n    table_name=self._table_name, table=self._table, ref_doc_id_column=self.\n    _ref_doc_id_column)\n', (4331, 4530), False, 'from llama_index.legacy.indices.common.struct_store.sql import SQLStructDatapointExtractor\n'), ((4706, 4723), 'collections.defaultdict', 'defaultdict', (['list'], {}), '(list)\n', (4717, 4723), False, 'from collections import defaultdict\n'), ((5969, 6009), 'llama_index.legacy.indices.struct_store.sql_query.NLStructStoreQueryEngine', 'NLStructStoreQueryEngine', (['self'], {}), '(self, **kwargs)\n', (5993, 6009), False, 'from llama_index.legacy.indices.struct_store.sql_query import NLStructStoreQueryEngine, SQLStructStoreQueryEngine\n'), ((6074, 6115), 'llama_index.legacy.indices.struct_store.sql_query.SQLStructStoreQueryEngine', 'SQLStructStoreQueryEngine', (['self'], {}), '(self, **kwargs)\n', (6099, 6115), False, 'from llama_index.legacy.indices.struct_store.sql_query import NLStructStoreQueryEngine, SQLStructStoreQueryEngine\n')]