app.py

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
import requests
import json
import os
import time
import asyncio

app = FastAPI()

# CORS settings
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins. In production, change this to specific domains.
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 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")  # Get cookie from environment variable

DEFAULT_HEADERS = {
    "User-Agent": USER_AGENT,
    "Accept": "text/event-stream",
    "Accept-Language": "en-US,en;q=0.5",
    "Referer": REFERER,
    "Content-Type": "application/json",
    "Origin": ORIGIN,
    "Connection": "keep-alive",
    "Cookie": COOKIE,
    "Sec-Fetch-Dest": "empty",
    "Sec-Fetch-Mode": "cors",
    "Sec-Fetch-Site": "same-origin",
    "Pragma": "no-cache",
    "TE": "trailers",
}

SUPPORTED_MODELS = [
    "o3-mini",
    "gpt-4o-mini",
    "claude-3-haiku-20240307",
    "meta-llama/Llama-3.3-70B-Instruct-Turbo",
]

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 duckduckgo_chat_stream(model, messages):
    """Interact with DuckDuckGo Chat with streaming output."""
    try:
        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:
        raise HTTPException(status_code=500, detail=f"Error during chat: {e}")

async def duckduckgo_chat_non_stream(model, messages):
    """Interact with DuckDuckGo Chat without streaming output."""
    try:
        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:
                        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

        return full_message

    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"HTTP request failed: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error during chat: {e}")

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    try:
        body = await request.json()
        if not body:
            raise HTTPException(status_code=400, detail="Invalid request body")

        model = body.get("model", "o3-mini")
        if model not in SUPPORTED_MODELS:
            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, [combined_message])
        else:
            response_text = await duckduckgo_chat_non_stream(model, [combined_message])

            # 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)

    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_exception_handler(request: Request, exc: HTTPException):
    return JSONResponse(
        status_code=exc.status_code,
        content={"detail": exc.detail},
    )

@app.get("/")
async def greet_json():
    return {"Hello": "World!"}