add llm endpoint

main.py

@@ -7,7 +7,7 @@ from query_engine import set_query_engine
 from llama_index.core.indices.struct_store import NLSQLTableQueryEngine
 import os
 from huggingface_hub import hf_hub_download
-from models import SQL
+from models import SQL, Prompt
 import json
 
 app = FastAPI()
@@ -43,9 +43,10 @@ async def query_database(query_data: SQL):
     df = conn.execute(query_data.query).fetch_df()
     return JSONResponse(json.loads(df.to_json(orient = "records")))
 
-# @app.post("/llm")
-# async def llm(prompt: str):
-#     response = query_engine.query(prompt)
-#     return JSONResponse({"promt":prompt,"response":response})
+@app.post("/llm")
+async def llm(prompt_data: Prompt):
+    global query_engine
+    response = query_engine.query(prompt_data.promt)
+    return JSONResponse({"promt":prompt_data.promt,"response":response})
 
 uvicorn.run(app,host='0.0.0.0',port=7860)

models/__init__.py

@@ -1,5 +1,7 @@
 from .sql_model import SQL
+from .prompt import Prompt
 
 __all__ = [
-    "SQL"
+    "SQL",
+    "Prompt"
 ]

models/prompt.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class Prompt(BaseModel):
+    promt:str