Update app/main.py

app/main.py

@@ -1,4 +1,5 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, Form
+from fastapi.responses import HTMLResponse
 from fastapi.templating import Jinja2Templates
 from fastapi.staticfiles import StaticFiles
 
@@ -9,24 +10,26 @@ app = FastAPI()
 templates = Jinja2Templates(directory="templates")
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-chunks, heads = load_data()
+# Load your RAG data once
+documents, embeddings, model = load_data()
 
-@app.get("/")
-def home(request: Request):
+@app.get("/", response_class=HTMLResponse)
+async def home(request: Request):
     return templates.TemplateResponse(
         "index.html",
-        {"request": request}
+        {"request": request, "query": "", "answer": None},
     )
 
-@app.post("/search")
-async def search(request: Request):
-    body = await request.json()
-    query = body["query"]
+@app.post("/", response_class=HTMLResponse)
+async def ask(request: Request, query: str = Form(...)):
+    results = retrieve_chunks(query, documents, embeddings, model)
+    answer = results[0] if results else "No relevant answer found."
 
-    retrieved = retrieve_chunks(query, chunks, heads)
-    answer = "\n\n".join(retrieved[:2])
-
-    return {
-        "answer": answer,
-        "sources": retrieved
-    }
+    return templates.TemplateResponse(
+        "index.html",
+        {
+            "request": request,
+            "query": query,
+            "answer": answer,
+        },
+    )