Upload server.py

server.py

@@ -12,6 +12,8 @@ from get_embedding_function import get_embedding_function
 from langchain_community.vectorstores import Chroma  
 from langchain.prompts import ChatPromptTemplate
 from langchain_community.llms.ollama import Ollama
+from pydantic import BaseModel
+
 
 PROMPT_TEMPLATE = """
 Answer the question based only on the following context:
@@ -41,6 +43,11 @@ def get_embedding_function():
 def greet_json():
     return {"Hello": "World!"}
 
+
+class QueryRequest(BaseModel):
+    query_text: str
+
+
 @app.get("/train")
 def train():
 
@@ -139,8 +146,9 @@ def clear_database():
 
 
 @app.get("/query")
-def query(query_text: str):
-    # Prepare the DB.
+def query(request: QueryRequest):
+    query_text = request.query_text
+        # Prepare the DB.
     embedding_function = get_embedding_function()
     db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)