Multiple SQL Tables

I managed to figure out how to use retrieval with a database to avoid overflowing the context window, however, I don't know if this is the recommended method. Could someone let me know if this is the optimal approach if you're not passing in a set list of tables?

Plain Text

from llama_index.indices.struct_store.sql_query import (
    SQLTableRetrieverQueryEngine,
)
from llama_index.objects import (
    SQLTableNodeMapping,
    ObjectIndex,
    SQLTableSchema,
)
from llama_index import VectorStoreIndex

table_node_mapping = SQLTableNodeMapping(sql_database)

from llama_index.retrievers import NLSQLRetriever

# default retrieval (return_raw=True)
nl_sql_retriever = NLSQLRetriever(sql_database, return_raw=True)

# Get all tables wrapped as SQLTableSchema Objects
table_schema_objs = nl_sql_retriever._get_tables("")

obj_index = ObjectIndex.from_objects(
    table_schema_objs,
    table_node_mapping,
    VectorStoreIndex,
)

query_engine = SQLTableRetrieverQueryEngine(
    sql_database, obj_index.as_retriever(similarity_top_k=1)
)

response = query_engine.query("How many albums do the rolling stones have?")
print(response)

Specifically this line nl_sql_retriever._get_tables("") I'm using it to get the list of tables wrapped in the SQLTableSchema class, but it might exist somewhere else more conveniently.

I was referencing this tutorial: https://docs.llamaindex.ai/en/stable/examples/index_structs/struct_indices/SQLIndexDemo.html

Also, a side-question, I found this old tutorial that looks like its been deleted. Not sure if any part of this is relevant anymore either using the context builder: https://gpt-index.readthedocs.io/en/v0.6.22/examples/index_structs/struct_indices/SQLIndexDemo-ManyTables.html

Anyway, any help is appreciated! Thanks

Find answers from the community

Multiple SQL Tables