add llm endpoint

main.py

@@ -22,15 +22,19 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# query_engine:NLSQLTableQueryEngine  = set_query_engine()
+query_engine:NLSQLTableQueryEngine
 query_file_path = ""
+llm_file_path = ""
 
 @app.on_event("startup")
 def startup():
     dataset_name = "pdrMottaS/afabd-duckdb"
-    duckdb_filename = "afabd.db"
     global query_file_path
-    query_file_path = hf_hub_download(repo_id=dataset_name, filename=duckdb_filename, repo_type="dataset")
+    global llm_file_path
+    global query_engine
+    llm_file_path = hf_hub_download(repo_id=dataset_name, filename='llm_afabd.db', repo_type="dataset")
+    query_file_path = hf_hub_download(repo_id=dataset_name, filename='afabd.db', repo_type="dataset")
+    query_engine = set_query_engine(llm_file_path)
 
 @app.post("/sql")
 async def query_database(query_data: SQL):

query_engine.py

@@ -5,7 +5,7 @@ from sqlalchemy import create_engine
 from llama_index.core import SQLDatabase
 from llama_index.core.indices.struct_store import NLSQLTableQueryEngine
 
-def set_query_engine():
+def set_query_engine(path:str):
     llm = Groq(
         model="llama3-8b-8192",
         api_key="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"
@@ -13,10 +13,10 @@ def set_query_engine():
     embedding = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
     Settings.llm = llm
     Settings.embed_model = embedding
-    engine = create_engine("duckdb:///llm_view.db")
+    engine = create_engine(f"duckdb:///{path}")
     sql_database = SQLDatabase(engine)
     return NLSQLTableQueryEngine(
         sql_database=sql_database, # The SQL database instance to query
-        tables=["afabd"], # List of tables to include in the query engine
+        tables=["escola", "curso", "avaliacao"], # List of tables to include in the query engine
         llm=llm, # The language model used for processing natural language queries
     )