Update app.py

app.py

@@ -1,36 +1,30 @@
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
-import httpx
+import requests
 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()
 
-# CORS Configuration
+# CORS settings
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # Allow all origins. In production, change this to specific domains.
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Environment Configuration
+# Environment variable configuration
 STATUS_URL = os.environ.get("STATUS_URL", "https://duckduckgo.com/duckchat/v1/status")
 CHAT_URL = os.environ.get("CHAT_URL", "https://duckduckgo.com/duckchat/v1/chat")
 REFERER = os.environ.get("REFERER", "https://duckduckgo.com/")
 ORIGIN = os.environ.get("ORIGIN", "https://duckduckgo.com")
 USER_AGENT = os.environ.get("USER_AGENT", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36")
-COOKIE = os.environ.get("COOKIE", "dcm=3; s=l; bf=1")
+COOKIE = os.environ.get("COOKIE", "dcm=3; s=l; bf=1")  # Get cookie from environment variable
 
 DEFAULT_HEADERS = {
     "User-Agent": USER_AGENT,
@@ -48,184 +42,220 @@ DEFAULT_HEADERS = {
     "TE": "trailers",
 }
 
-SUPPORTED_MODELS = ["o3-mini", "gpt-4o-mini", "claude-3-haiku-20240307", "meta-llama/Llama-3.3-70B-Instruct-Turbo"]
-TIMEOUT = 30.0  # Seconds
+SUPPORTED_MODELS = [
+    "o3-mini",
+    "gpt-4o-mini",
+    "claude-3-haiku-20240307",
+    "meta-llama/Llama-3.3-70B-Instruct-Turbo",
+]
 
-class StreamClosedError(Exception):
-    pass
+async def get_vqd():
+    """Get the VQD value for DuckDuckGo Chat."""
+    headers = {**DEFAULT_HEADERS, "x-vqd-accept": "1"}
+    try:
+        response = requests.get(STATUS_URL, headers=headers)
+        response.raise_for_status()  # Raise HTTPError if status code is not 200
+        vqd = response.headers.get("x-vqd-4")
+        if not vqd:
+            raise ValueError("x-vqd-4 header not found in the response.")
+        return vqd
+    except requests.exceptions.RequestException as e:
+        raise HTTPException(status_code=500, detail=f"HTTP request failed: {e}")
+    except ValueError as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-async def get_vqd() -> str:
-    """Fetch and validate DuckDuckGo authentication token."""
+async def duckduckgo_chat_stream(model, messages):
+    """Interact with DuckDuckGo Chat with streaming output."""
     try:
-        async with httpx.AsyncClient() as client:
-            response = await client.get(
-                STATUS_URL,
-                headers={**DEFAULT_HEADERS, "x-vqd-accept": "1"},
-                timeout=10.0
-            )
-            response.raise_for_status()
-            
-            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:
-        logger.error(f"VQD fetch failed: {e.response.text}")
-        raise HTTPException(status_code=502, detail="Upstream authentication service unavailable")
+        x_vqd_4 = await get_vqd()
+
+        chat_headers = {
+            **DEFAULT_HEADERS,
+            "x-vqd-4": x_vqd_4,
+            "Accept": "text/event-stream",  # Ensure we accept SSE
+        }
+
+        body = json.dumps({
+            "model": model,
+            "messages": messages,
+        })
+
+        response = requests.post(CHAT_URL, headers=chat_headers, data=body, stream=True)
+        response.raise_for_status()
+
+        async def event_stream():
+            try:
+                for line in response.iter_lines():
+                    if line:
+                        decoded_line = line.decode('utf-8')
+                        if decoded_line.startswith("data: "):
+                            content = decoded_line[5:].strip()
+                            # Check if this is the final marker
+                            if content == "[DONE]":
+                                yield f"data: [DONE]\n\n"
+                                break
+                            try:
+                                json_data = json.loads(content)
+                                message_content = json_data.get("message", "")
+                                if message_content:
+                                    # Build OpenAI style streaming response
+                                    openai_stream_response = {
+                                        "id": f"chatcmpl-{int(time.time() * 1000)}",
+                                        "object": "chat.completion.chunk",
+                                        "created": int(time.time()),
+                                        "model": model,
+                                        "choices": [
+                                            {
+                                                "delta": {"content": message_content},
+                                                "index": 0,
+                                                "finish_reason": None,
+                                            }
+                                        ],
+                                    }
+                                    yield f"data: {json.dumps(openai_stream_response)}\n\n"
+                                await asyncio.sleep(0.01)  # Prevent high CPU usage
+                            except json.JSONDecodeError as e:
+                                print(f"JSON decode error: {e}, line: {decoded_line}")
+                                yield f"data: {json.dumps({'error': 'JSON decode error'})}\n\n"
+                                break  # Stop the stream
+            except requests.exceptions.RequestException as e:
+                print(f"Request error: {e}")
+                yield f"data: {json.dumps({'error': 'Request error'})}\n\n"
+            except Exception as e:
+                print(f"An error occurred: {e}")
+                yield f"data: {json.dumps({'error': 'An error occurred'})}\n\n"
+            finally:
+                yield "data: [DONE]\n\n"  # End SSE stream
+
+        return StreamingResponse(event_stream(), media_type="text/event-stream")
+
+    except requests.exceptions.RequestException as e:
+        raise HTTPException(status_code=500, detail=f"HTTP request failed: {e}")
     except Exception as e:
-        logger.exception("Critical VQD error")
-        raise HTTPException(status_code=500, detail="Internal authentication error")
+        raise HTTPException(status_code=500, detail=f"Error during chat: {e}")
 
-async def stream_generator(response: httpx.Response) -> AsyncGenerator[str, None]:
-    """Handle streaming response with proper buffer management."""
-    buffer = ""
+async def duckduckgo_chat_non_stream(model, messages):
+    """Interact with DuckDuckGo Chat without streaming output."""
     try:
-        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
-                        
+        x_vqd_4 = await get_vqd()
+
+        chat_headers = {
+            **DEFAULT_HEADERS,
+            "x-vqd-4": x_vqd_4,
+        }
+
+        body = json.dumps({
+            "model": model,
+            "messages": messages,
+        })
+
+        response = requests.post(CHAT_URL, headers=chat_headers, data=body)
+        response.raise_for_status()
+
+        full_message = ""
+        for line in response.iter_lines():
+            if line:
+                decoded_line = line.decode('utf-8')
+                if decoded_line.startswith("data: "):
                     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:
-        logger.error(f"Unexpected stream error: {str(e)}")
-        yield format_error_chunk("Stream connection failed")
-    finally:
-        yield "data: [DONE]\n\n"
+                        json_data = json.loads(decoded_line[5:])
+                        full_message += json_data.get("message", "")
+                    except json.JSONDecodeError as e:
+                        print(f"JSON decode error: {e}, line: {decoded_line}")
+                        pass  # Ignore decoding errors
 
-async def duckduckgo_chat_stream(model: str, messages: list) -> StreamingResponse:
-    """Robust streaming handler with connection monitoring."""
-    try:
-        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={
-                    **DEFAULT_HEADERS,
-                    "x-vqd-4": vqd_token,
-                    "Accept": "text/event-stream",
-                },
-                json={"model": model, "messages": messages},
-                timeout=TIMEOUT
-            )
-            response.raise_for_status()
-            
-            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:
-        logger.error(f"Upstream API error: {e.response.text}")
-        raise HTTPException(status_code=502, detail="Upstream service unavailable")
+        return full_message
+
+    except requests.exceptions.RequestException as e:
+        raise HTTPException(status_code=500, detail=f"HTTP request failed: {e}")
     except Exception as e:
-        logger.exception("Stream setup failed")
-        raise HTTPException(status_code=500, detail="Stream initialization failed")
-
-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",
-        "created": int(time.time()),
-        "model": model,
-        "choices": [{
-            "delta": {"content": content},
-            "index": 0,
-            "finish_reason": None
-        }]
-    }) + "\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"
+        raise HTTPException(status_code=500, detail=f"Error during chat: {e}")
 
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request):
     try:
-        data = await request.json()
-        model = data.get("model", "o3-mini")
-        messages = data.get("messages", [])
-        stream = data.get("stream", False)
+        body = await request.json()
+        if not body:
+            raise HTTPException(status_code=400, detail="Invalid request body")
 
-        # Validate input
+        model = body.get("model", "o3-mini")
         if model not in SUPPORTED_MODELS:
-            raise HTTPException(400, f"Unsupported model: {model}")
-        if not messages or not isinstance(messages, list):
-            raise HTTPException(400, "Invalid messages format")
-            
-        # Process messages
-        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
+            raise HTTPException(
+                status_code=400,
+                detail=f"Model \"{model}\" is not supported. Supported models are: {', '.join(SUPPORTED_MODELS)}."
             )
-        }]
+
+        messages = body.get("messages")
+        if not messages:
+            raise HTTPException(status_code=400, detail="No message content provided")
+
+        stream = body.get("stream", False)  # Get the stream parameter, default is False
+
+        # Process system message
+        system_message = next((msg for msg in messages if msg.get("role") == "system"), None)
+        system_prompt = f"You will play the role of a {system_message['content']}.\n" if system_message else ""
+
+        # Extract and format the conversation history
+        history_messages = "\n".join(
+            f"{msg['role']}: {msg['content']}"
+            for msg in messages
+            if msg.get("role") != "system" and msg != messages[-1]
+        )
+
+        # Get the last user message
+        last_user_message = messages[-1]
+        current_question = last_user_message["content"] if last_user_message.get("role") == "user" else ""
+
+        # Build the combined message
+        combined_message_content = (
+            f"{system_prompt}Below is the conversation history:\n{history_messages}\n"
+            f"User's current question: {current_question}"
+        )
+        combined_message = {"role": "user", "content": combined_message_content}
 
         if stream:
-            return await duckduckgo_chat_stream(model, payload)
+            return await duckduckgo_chat_stream(model, [combined_message])
         else:
-            # Non-streaming implementation
-            raise HTTPException(501, "Non-streaming mode temporarily disabled")
+            response_text = await duckduckgo_chat_non_stream(model, [combined_message])
 
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.exception("API handler error")
-        raise HTTPException(500, "Internal server error")
+            # Build OpenAI style response
+            openai_response = {
+                "id": f"chatcmpl-{int(time.time() * 1000)}",  # Unique ID
+                "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
+                },
+            }
+
+            return JSONResponse(content=openai_response)
 
-@app.get("/health")
-async def health_check():
-    return {"status": "ok", "timestamp": int(time.time())}
+    except HTTPException as e:
+        raise e  # Re-raise HTTPException so FastAPI can handle it
+    except Exception as e:
+        print(f"API error: {e}")
+        raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
 
 @app.exception_handler(HTTPException)
-async def http_error_handler(request: Request, exc: HTTPException):
+async def http_exception_handler(request: Request, exc: HTTPException):
     return JSONResponse(
         status_code=exc.status_code,
-        content={"error": exc.detail},
+        content={"detail": exc.detail},
     )
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+@app.get("/")
+async def greet_json():
+    return {"Hello": "World!"}