Update app.py

app.py

@@ -33,11 +33,12 @@ logger.debug(f"Loaded API key: {API_KEY}")
 OLLAMA_SERVER_URL = "http://localhost:11434/api/generate"
 logger.debug(f"Ollama server URL: {OLLAMA_SERVER_URL}")
 
+
 @app.post("/api/generate")
 async def generate(request: Request):
     """Endpoint that generates text based on the prompt."""
     try:
-        # 1. Parse the incoming request
+        # Parse the incoming request
         body = await request.json()
         model = body.get("model", "hf.co/abanm/Dubs-Q8_0-GGUF:latest")  # Default model
         prompt_text = body.get("prompt", "")
@@ -48,7 +49,7 @@ async def generate(request: Request):
 
         logger.debug(f"Request body: {body}")
 
-        # 2. Validate API key
+        # Validate API key
         auth_header = request.headers.get("Authorization")
         logger.debug(f"Received Authorization header: {auth_header}")
 
@@ -61,14 +62,14 @@ async def generate(request: Request):
             logger.error(f"Invalid API key provided: {token}")
             raise HTTPException(status_code=401, detail="Invalid API key")
 
-        # 3. Prepare request payload
+        # Prepare request payload
         payload = {"model": model, "prompt": prompt_text}
         logger.debug(f"Payload prepared for Ollama: {payload}")
 
-        # 4. Stream response from Ollama
+        # Stream response from Ollama
         async def stream_response():
             try:
-                async with httpx.AsyncClient() as client:
+                async with httpx.AsyncClient(timeout=httpx.Timeout(60.0)) as client:
                     async with client.stream(
                         "POST", OLLAMA_SERVER_URL, json=payload, headers={"Content-Type": "application/json"}
                     ) as response:
@@ -82,12 +83,15 @@ async def generate(request: Request):
                         async for chunk in response.aiter_text():
                             logger.debug(f"Chunk received: {chunk}")
                             yield chunk
+            except httpx.ReadTimeout:
+                logger.error("ReadTimeout while waiting for response chunks")
+                yield json.dumps({"error": "Server response timeout. Try again later."})
             except httpx.RequestError as exc:
                 logger.error(f"Request error while communicating with Ollama: {str(exc)}")
                 yield json.dumps({"error": "Network error occurred while communicating with Ollama"})
-            except httpx.HTTPStatusError as exc:
-                logger.error(f"HTTP error from Ollama: {exc.response.status_code} - {exc.response.text}")
-                yield json.dumps({"error": f"HTTP error: {exc.response.text}"})
+            except Exception as exc:
+                logger.exception(f"Unexpected error during streaming: {str(exc)}")
+                yield json.dumps({"error": "An unexpected error occurred during streaming."})
 
         return StreamingResponse(stream_response(), media_type="application/json")
 
@@ -95,12 +99,14 @@ async def generate(request: Request):
         logger.exception(f"Unexpected error: {str(e)}")
         raise HTTPException(status_code=500, detail="An unexpected error occurred")
 
+
 @app.get("/health")
 async def health():
     """Health check endpoint."""
     logger.info("Health check endpoint accessed")
     return {"status": "OK"}
 
+
 if __name__ == "__main__":
     import uvicorn
     logger.info("Starting FastAPI application")