updated

main.py

@@ -1,10 +1,10 @@
 from fastapi import FastAPI, Depends, HTTPException, status
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from jose import JWTError
 from schemas import UserRegister, TokenResponse, RefreshTokenRequest, QueryInput
 from auth import  register_user, get_db, authenticate_user, create_token, verify_token, verify_access_token, Session
-from utils import search_with_llm
+from utils import generate_stream
 from fastapi.security import OAuth2PasswordRequestForm
 from pathlib import Path
 from datetime import timedelta
@@ -21,6 +21,8 @@ REFRESH_TOKEN_EXPIRE_DAYS = 7
 
 app = FastAPI()
 
+
+
 # Entry Endpoint 
 @app.get('/')
 def index() -> FileResponse:
@@ -95,18 +97,12 @@ async def refresh(refresh_request: RefreshTokenRequest):
             headers={"WWW-Authenticate": "Bearer"},
         )
     
-
-
 @app.post("/search")
 async def search(
     query_input: QueryInput,
     username: str = Depends(verify_access_token),
 ):
-    try:
-        response = search_with_llm(query_input.query)
-        return {"response": response}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")
+    return StreamingResponse(generate_stream(query_input.query), media_type="text/plain")
 
 # WebSocket endpoint for streaming
 @app.on_event("startup")

utils.py

@@ -1,21 +1,23 @@
-import os
+import asyncio
 import ollama
-from typing import List, Dict
+from typing import List
 
 def cosine_similarity(embedding_0, embedding_1):
     pass
 
-
 def generate_embedding(model, text: str, model_type: str) -> List[float]:
     pass
 
-def search_with_llm(query: str, model: str = "llama3.2"):
-    try:
-        response = ollama.chat(
-            model=model,
-            messages=[{"role": "user", "content": query}]
-        )
-        return response["message"]["content"]
-    except Exception as e:
-        return f"❌ Error processing request: {str(e)}"
+async def generate_stream(query: str):
+    """Stream responses from Ollama in real-time."""
+    stream = ollama.chat(
+        model="llama3.2",  # Choose your model (mistral, llama2, gemma)
+        messages=[{"role": "user", "content": query}],
+        stream=True  # Enable streaming
+    )
+    
+    for chunk in stream:
+        if "message" in chunk and "content" in chunk["message"]:
+            yield chunk["message"]["content"]
+            await asyncio.sleep(0)  # Allow async executi