I have the following code which works

I have the following code which works perfectly fine, and I can then use the SQLTableRetrieverQueryEngine to query the database and get answers. But I want to be able to create the index once so I can reuse it without recreating it every time. Is it possibile to achieve that @Logan M ?

Here is my code:

def table_index_builder(sql_database, table_context_dict: dict[str, Literal]) -> ObjectIndex:
    table_node_mapping = SQLTableNodeMapping(sql_database)
    table_schema_objs: list[SQLTableSchema] = []

    # Build our table schema
    for table_name in table_context_dict:
        # one SQLTableSchema for each table
        table_schema = (SQLTableSchema(
            table_name=table_name, context_str=table_context_dict[table_name]))
        table_schema_objs.append(table_schema)

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

    return obj_index


def build_query_engine(sql_database: SQLDatabase, service_context: ServiceContext, table_context_dict: dict[str, Literal]) -> SQLTableRetrieverQueryEngine:
    # build object index
    obj_index = table_index_builder(
        sql_database,
        table_context_dict=table_context_dict)

    query_engine = SQLTableRetrieverQueryEngine(
        sql_database=sql_database,
        table_retriever=obj_index.as_retriever(similarity_top_k=2),
        service_context=service_context
    )

    return query_engine

Find answers from the community

I have the following code which works