Update app.py

app.py

@@ -7,7 +7,7 @@ from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 
 # Configure logging
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
@@ -27,60 +27,81 @@ app.add_middleware(
 
 # Load your API key from the environment (defaults to "change_me")
 API_KEY = os.environ.get("API_KEY", "change_me")
-
-print(f"API key={API_KEY}")
+logger.debug(f"Loaded API key: {API_KEY}")
 
 # URL of the running Ollama server (adjust as needed)
 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."""
-    # 1. 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", "")
-    if not prompt_text:
-        raise HTTPException(status_code=400, detail="No prompt provided")
-
-    # 2. Validate API key
-    auth_header = request.headers.get("Authorization")
-    logger.info(f"Received Authorization header: {auth_header}")
-
-    if not auth_header or not auth_header.startswith("Bearer "):
-        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
-
-    token = auth_header.split(" ")[1]
-    if token != API_KEY:
-        raise HTTPException(status_code=401, detail=f"Invalid API key: {token}")
-
-    # 3. Prepare request payload
-    payload = {"model": model, "prompt": prompt_text}
-
-    # 4. Stream response from Ollama
-    async def stream_response():
-        try:
-            async with httpx.AsyncClient() as client:
-                async with client.stream(
-                    "POST", OLLAMA_SERVER_URL, json=payload, headers={"Content-Type": "application/json"}
-                ) as response:
-                    response.raise_for_status()
-                    async for chunk in response.aiter_text():
-                        yield chunk
-        except httpx.RequestError as exc:
-            logger.error(f"Request error while communicating with Ollama: {str(exc)}")
-            yield json.dumps({"error": "Unable to communicate 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}"})
-
-    return StreamingResponse(stream_response(), media_type="application/json")
+    try:
+        # 1. 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", "")
+
+        if not prompt_text:
+            logger.error("No prompt provided in the request")
+            raise HTTPException(status_code=400, detail="No prompt provided")
+
+        logger.debug(f"Request body: {body}")
+
+        # 2. Validate API key
+        auth_header = request.headers.get("Authorization")
+        logger.debug(f"Received Authorization header: {auth_header}")
+
+        if not auth_header or not auth_header.startswith("Bearer "):
+            logger.error("Missing or invalid Authorization header")
+            raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
+
+        token = auth_header.split(" ")[1]
+        if token != API_KEY:
+            logger.error(f"Invalid API key provided: {token}")
+            raise HTTPException(status_code=401, detail="Invalid API key")
+
+        # 3. Prepare request payload
+        payload = {"model": model, "prompt": prompt_text}
+        logger.debug(f"Payload prepared for Ollama: {payload}")
+
+        # 4. Stream response from Ollama
+        async def stream_response():
+            try:
+                async with httpx.AsyncClient() as client:
+                    async with client.stream(
+                        "POST", OLLAMA_SERVER_URL, json=payload, headers={"Content-Type": "application/json"}
+                    ) as response:
+                        logger.info(f"Response status code from Ollama: {response.status_code}")
+
+                        if response.status_code != 200:
+                            logger.error(f"HTTP error: {response.status_code} - {await response.text()}")
+                            yield json.dumps({"error": f"HTTP error: {response.status_code}"})
+                            return
+
+                        async for chunk in response.aiter_text():
+                            logger.debug(f"Chunk received: {chunk}")
+                            yield chunk
+            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}"})
+
+        return StreamingResponse(stream_response(), media_type="application/json")
+
+    except Exception as e:
+        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")
     uvicorn.run(app, host="0.0.0.0", port=7860)