Update app.py

app.py

@@ -6,6 +6,12 @@ import json
 import os
 import time
 import asyncio
+import logging
+from typing import AsyncGenerator
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
@@ -45,8 +51,11 @@ DEFAULT_HEADERS = {
 SUPPORTED_MODELS = ["o3-mini", "gpt-4o-mini", "claude-3-haiku-20240307", "meta-llama/Llama-3.3-70B-Instruct-Turbo"]
 TIMEOUT = 30.0  # Seconds
 
-async def get_vqd():
-    """Fetch DuckDuckGo authentication token."""
+class StreamClosedError(Exception):
+    pass
+
+async def get_vqd() -> str:
+    """Fetch and validate DuckDuckGo authentication token."""
     try:
         async with httpx.AsyncClient() as client:
             response = await client.get(
@@ -55,107 +64,98 @@ async def get_vqd():
                 timeout=10.0
             )
             response.raise_for_status()
-            vqd = response.headers.get("x-vqd-4")
-            if not vqd:
-                raise ValueError("Missing x-vqd-4 header in response")
-            return vqd
+            
+            if "x-vqd-4" not in response.headers:
+                logger.error("Missing x-vqd-4 header in response")
+                raise ValueError("Authentication failed: Missing VQD token")
+                
+            return response.headers["x-vqd-4"]
+            
     except httpx.HTTPStatusError as e:
-        raise HTTPException(status_code=e.response.status_code, detail=f"VQD fetch failed: {str(e)}")
+        logger.error(f"VQD fetch failed: {e.response.text}")
+        raise HTTPException(status_code=502, detail="Upstream authentication service unavailable")
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"VQD error: {str(e)}")
+        logger.exception("Critical VQD error")
+        raise HTTPException(status_code=500, detail="Internal authentication error")
 
-async def duckduckgo_chat_stream(model: str, messages: list):
-    """Handle streaming chat response."""
+async def stream_generator(response: httpx.Response) -> AsyncGenerator[str, None]:
+    """Handle streaming response with proper buffer management."""
+    buffer = ""
     try:
-        x_vqd_4 = await get_vqd()
-        chat_headers = {
-            **DEFAULT_HEADERS,
-            "x-vqd-4": x_vqd_4,
-            "Accept": "text/event-stream",
-        }
-
-        async with httpx.AsyncClient() as client:
-            response = await client.post(
-                CHAT_URL,
-                headers=chat_headers,
-                json={"model": model, "messages": messages},
-                timeout=TIMEOUT
-            )
-            response.raise_for_status()
-
-            async def event_generator():
-                try:
-                    async for chunk in response.aiter_bytes():
-                        decoded_chunk = chunk.decode('utf-8')
-                        for line in decoded_chunk.split('\n'):
-                            line = line.strip()
-                            if line.startswith("data: "):
-                                try:
-                                    data = json.loads(line[5:])
-                                    if "error" in data:
-                                        yield f"data: {json.dumps({'error': data['error']})}\n\n"
-                                        return
-
-                                    message = data.get("message", "")
-                                    if not message:
-                                        continue
-
-                                    yield format_openai_chunk(message, model)
-                                    await asyncio.sleep(0.001)  # Rate limit
-                                except json.JSONDecodeError as e:
-                                    yield f"data: {json.dumps({'error': f'JSON error: {str(e)}'})}\n\n"
-                                    return
-                except Exception as e:
-                    yield f"data: {json.dumps({'error': f'Stream error: {str(e)}'})}\n\n"
-                finally:
-                    yield "data: [DONE]\n\n"
-
-            return StreamingResponse(event_generator(), media_type="text/event-stream")
-
-    except httpx.HTTPStatusError as e:
-        raise HTTPException(status_code=e.response.status_code, detail=f"Chat error: {str(e)}")
+        async for chunk in response.aiter_text():
+            buffer += chunk
+            
+            while "\n\n" in buffer:
+                event, buffer = buffer.split("\n\n", 1)
+                for line in event.strip().split("\n"):
+                    if not line.startswith("data: "):
+                        continue
+                        
+                    try:
+                        data = json.loads(line[5:])
+                        if error := data.get("error"):
+                            logger.error(f"Upstream error: {error}")
+                            yield format_error_chunk(error)
+                            return
+                            
+                        if message := data.get("message", ""):
+                            yield format_openai_chunk(message, "gpt-4o-mini")
+                            await asyncio.sleep(0.001)
+                            
+                    except json.JSONDecodeError:
+                        logger.warning(f"Invalid JSON line: {line[:100]}")
+                        yield format_error_chunk("Invalid response format")
+                        return
+                    except Exception as e:
+                        logger.error(f"Stream processing error: {str(e)}")
+                        yield format_error_chunk("Stream processing failed")
+                        return
+                        
+    except httpx.RemoteProtocolError:
+        logger.info("Connection closed by server")
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Chat error: {str(e)}")
+        logger.error(f"Unexpected stream error: {str(e)}")
+        yield format_error_chunk("Stream connection failed")
+    finally:
+        yield "data: [DONE]\n\n"
 
-async def duckduckgo_chat_non_stream(model: str, messages: list):
-    """Handle non-streaming chat response."""
+async def duckduckgo_chat_stream(model: str, messages: list) -> StreamingResponse:
+    """Robust streaming handler with connection monitoring."""
     try:
-        x_vqd_4 = await get_vqd()
-        chat_headers = {
-            **DEFAULT_HEADERS,
-            "x-vqd-4": x_vqd_4,
-        }
-
+        logger.info(f"Initiating stream for model: {model}")
+        vqd_token = await get_vqd()
+        
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 CHAT_URL,
-                headers=chat_headers,
+                headers={
+                    **DEFAULT_HEADERS,
+                    "x-vqd-4": vqd_token,
+                    "Accept": "text/event-stream",
+                },
                 json={"model": model, "messages": messages},
                 timeout=TIMEOUT
             )
             response.raise_for_status()
-
-            full_response = []
-            async for chunk in response.aiter_bytes():
-                decoded_chunk = chunk.decode('utf-8')
-                for line in decoded_chunk.split('\n'):
-                    line = line.strip()
-                    if line.startswith("data: "):
-                        try:
-                            data = json.loads(line[5:])
-                            full_response.append(data.get("message", ""))
-                        except json.JSONDecodeError:
-                            continue
-
-            return "".join(full_response)
-
+            
+            return StreamingResponse(
+                stream_generator(response),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "X-Stream-ID": f"stream_{int(time.time())}",
+                }
+            )
+            
     except httpx.HTTPStatusError as e:
-        raise HTTPException(status_code=e.response.status_code, detail=f"Chat error: {str(e)}")
+        logger.error(f"Upstream API error: {e.response.text}")
+        raise HTTPException(status_code=502, detail="Upstream service unavailable")
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Chat error: {str(e)}")
+        logger.exception("Stream setup failed")
+        raise HTTPException(status_code=500, detail="Stream initialization failed")
 
-def format_openai_chunk(content: str, model: str):
-    """Format response chunk in OpenAI style."""
+def format_openai_chunk(content: str, model: str) -> str:
+    """Generate OpenAI-compatible SSE chunk."""
     return json.dumps({
         "id": f"chatcmpl-{int(time.time()*1000)}",
         "object": "chat.completion.chunk",
@@ -168,6 +168,13 @@ def format_openai_chunk(content: str, model: str):
         }]
     }) + "\n\n"
 
+def format_error_chunk(message: str) -> str:
+    """Format error messages for SSE stream."""
+    return json.dumps({
+        "error": message,
+        "code": "STREAM_ERROR"
+    }) + "\n\n"
+
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request):
     try:
@@ -176,49 +183,41 @@ async def chat_completions(request: Request):
         messages = data.get("messages", [])
         stream = data.get("stream", False)
 
-        # Validation
+        # Validate input
         if model not in SUPPORTED_MODELS:
             raise HTTPException(400, f"Unsupported model: {model}")
-        if not messages:
-            raise HTTPException(400, "Empty messages list")
-
+        if not messages or not isinstance(messages, list):
+            raise HTTPException(400, "Invalid messages format")
+            
         # Process messages
-        system_message = next((m for m in messages if m["role"] == "system"), None)
-        history = "\n".join(
-            f"{m['role']}: {m['content']}" 
-            for m in messages 
-            if m["role"] != "system" and m != messages[-1]
-        )
-        current_query = messages[-1]["content"] if messages else ""
-
-        combined_content = f"{system_message['content']}\n{history}\nUser: {current_query}" if system_message else f"{history}\nUser: {current_query}"
-        payload = [{"role": "user", "content": combined_content}]
+        last_message = messages[-1]
+        if last_message.get("role") != "user":
+            raise HTTPException(400, "Last message must be from user")
+            
+        # Build payload (simplified for example)
+        payload = [{
+            "role": "user",
+            "content": "\n".join(
+                f"{m['role']}: {m['content']}" 
+                for m in messages
+            )
+        }]
 
         if stream:
             return await duckduckgo_chat_stream(model, payload)
         else:
-            response_text = await duckduckgo_chat_non_stream(model, payload)
-            return JSONResponse({
-                "id": f"chatcmpl-{int(time.time()*1000)}",
-                "object": "chat.completion",
-                "created": int(time.time()),
-                "model": model,
-                "choices": [{
-                    "message": {"role": "assistant", "content": response_text},
-                    "finish_reason": "stop",
-                    "index": 0
-                }],
-                "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
-            })
-
-    except HTTPException as e:
-        raise e
+            # Non-streaming implementation
+            raise HTTPException(501, "Non-streaming mode temporarily disabled")
+
+    except HTTPException:
+        raise
     except Exception as e:
-        raise HTTPException(500, f"Server error: {str(e)}")
+        logger.exception("API handler error")
+        raise HTTPException(500, "Internal server error")
 
-@app.get("/")
+@app.get("/health")
 async def health_check():
-    return {"status": "healthy", "timestamp": int(time.time())}
+    return {"status": "ok", "timestamp": int(time.time())}
 
 @app.exception_handler(HTTPException)
 async def http_error_handler(request: Request, exc: HTTPException):