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, Prompt
+from models import SQL, Prompt, PromptResponse
 import json
 
 app = FastAPI()
@@ -47,6 +47,6 @@ async def query_database(query_data: SQL):
 async def llm(prompt_data: Prompt):
     global query_engine
     response = query_engine.query(prompt_data.prompt)
-    return JSONResponse(json.dumps({"promt":prompt_data.prompt,"response":response}))
+    return PromptResponse(prompt=prompt_data.prompt,response= response)
 
 uvicorn.run(app,host='0.0.0.0',port=7860)

models/__init__.py

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

models/prompt.py

@@ -2,4 +2,7 @@ from pydantic import BaseModel
 
 
 class Prompt(BaseModel):
-    prompt:str
+    prompt:str
+
+class PromptResponse(Prompt):
+    response:str