Update rag.py

rag.py

@@ -1,32 +1,22 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
 from typing import List, Optional
 from llms import initialize_llms, answer_query_with_chunks, query_to_retrieve
 from chroma_operations.retrieve import search_similar_chunks
-
-# Initialize FastAPI
-app = FastAPI()
-
-
-# Define request model
-class RAGRequest(BaseModel):
-    query_text: str
-    file_names: List[str]
-    collection_name: str = "rag_collection"
-
+import asyncio
 
 # Load LLM once at startup
 llms = initialize_llms()
 llm = llms["llm"]
 
-
-@app.post("/ask")
-async def ask_question(
-    query_text: str, file_names: List[str], collection_name: str = "rag_collection"
-):
+def ask_question(
+    query_text: str, 
+    file_names: List[str], 
+    collection_name: str = "rag_collection"
+) -> dict:
     try:
+        # Await async LLM calls
         query_search = query_to_retrieve(query_text, llm)
-
+        
+        # Await ChromaDB search
         retrieved_docs = search_similar_chunks(
             query_search,
             document_names=file_names,
@@ -34,34 +24,11 @@ async def ask_question(
         )
 
         if not retrieved_docs:
-            raise HTTPException(status_code=404, detail="No matching documents found.")
-
-        answer = answer_query_with_chunks(request.query_text, retrieved_docs, llm)
-        return {"answer": answer, "chunks": retrieved_docs}
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-
-# For backward compatibility with older clients
-@app.post("/ask-single")
-async def ask_question_single(
-    query_text: str, file_name: str, collection_name: str = "rag_collection"
-):
-    try:
-        query_search = query_to_retrieve(query_text, llm)
-
-        retrieved_docs = search_similar_chunks(
-            query_search,
-            document_names=[file_name],  # Convert single filename to a list
-            collection_name=collection_name,
-        )
-
-        if not retrieved_docs:
-            raise HTTPException(status_code=404, detail="No matching documents found.")
+            return {"answer": "No matching documents found.", "chunks": []}
 
+        # Await answer generation
         answer = answer_query_with_chunks(query_text, retrieved_docs, llm)
         return {"answer": answer, "chunks": retrieved_docs}
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        return {"answer": f"Error processing request: {str(e)}", "chunks": []}