Create app.py

app.py

@@ -0,0 +1,294 @@
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import List, Dict, Optional, Union
+from fastapi.responses import StreamingResponse
+import logging
+import uuid
+import time
+import json
+import asyncio
+import random
+import httpx
+from fake_useragent import UserAgent
+
+app = FastAPI()
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+MODEL_MAPPING = {
+    "keyless-gpt-4o-mini": "gpt-4o-mini",
+    "keyless-claude-3-haiku": "claude-3-haiku-20240307",
+    "keyless-mixtral-8x7b": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "keyless-meta-Llama-3.1-70B-Instruct-Turbo": "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"
+}
+
+class ModelInfo(BaseModel):
+    id: str
+    object: str = "model"
+    created: int = int(time.time())
+    owned_by: str = "custom"
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
+    n: Optional[int] = 1
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, List[str]]] = None
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = 0.0
+    frequency_penalty: Optional[float] = 0.0
+    logit_bias: Optional[Dict[str, float]] = None
+    user: Optional[str] = None
+
+class ChatCompletionResponseChoice(BaseModel):
+    index: int
+    message: ChatMessage
+    finish_reason: Optional[str] = None
+
+class ChatCompletionResponseUsage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[ChatCompletionResponseChoice]
+    usage: ChatCompletionResponseUsage
+
+class DeltaMessage(BaseModel):
+    role: Optional[str] = None
+    content: Optional[str] = None
+
+class ChatCompletionStreamResponseChoice(BaseModel):
+    index: int
+    delta: DeltaMessage
+    finish_reason: Optional[str] = None
+
+class ChatCompletionStreamResponse(BaseModel):
+    id: str
+    object: str = "chat.completion.chunk"
+    created: int
+    model: str
+    choices: List[ChatCompletionStreamResponseChoice]
+
+# Store active conversations
+conversations: Dict[str, List[ChatMessage]] = {}
+
+ua = UserAgent()
+
+def get_next_user_agent():
+    return ua.random
+
+async def update_vqd_token(user_agent):
+    async with httpx.AsyncClient() as client:
+        try:
+            await client.get("https://duckduckgo.com/country.json", headers={"User-Agent": user_agent})
+            headers = {"x-vqd-accept": "1", "User-Agent": user_agent}
+            response = await client.get("https://duckduckgo.com/duckchat/v1/status", headers=headers)
+            if response.status_code == 200:
+                vqd_token = response.headers.get("x-vqd-4", "")
+                logging.info(f"Fetched new x-vqd-4 token: {vqd_token}")
+                return vqd_token
+            else:
+                logging.warning(f"Failed to fetch x-vqd-4 token. Status code: {response.status_code}")
+                return ""
+        except Exception as e:
+            logging.error(f"Error fetching x-vqd-4 token: {str(e)}")
+            return ""
+
+async def chat_with_duckduckgo(query: str, model: str, conversation_history: List[ChatMessage]):
+    original_model = MODEL_MAPPING.get(model, model)
+    user_agent = get_next_user_agent()
+    vqd_token = await update_vqd_token(user_agent)
+    if not vqd_token:
+        raise HTTPException(status_code=500, detail="Failed to obtain VQD token")
+
+    # If there is a system message, add it before the first user message (DDG AI doesnt let us send system messages, so this is a workaround -- fundamentally, it works the same way when setting a system prompt)
+    system_message = next((msg for msg in conversation_history if msg.role == "system"), None)
+    user_messages = [{"role": msg.role, "content": msg.content} for msg in conversation_history if msg.role == "user"]
+    
+    if system_message and user_messages:
+        user_messages[0]["content"] = f"{system_message.content}\n\n{user_messages[0]['content']}"
+
+    payload = {
+        "messages": user_messages,
+        "model": original_model
+    }
+
+    headers = {
+        "x-vqd-4": vqd_token,
+        "Content-Type": "application/json",
+        "User-Agent": user_agent
+    }
+
+    logging.info(f"Sending payload to DuckDuckGo with User-Agent: {user_agent}")
+
+    async with httpx.AsyncClient() as client:
+        try:
+            response = await client.post("https://duckduckgo.com/duckchat/v1/chat", json=payload, headers=headers)
+            if response.status_code == 200:
+                full_response = ""
+                async for line in response.aiter_lines():
+                    if line.startswith("data: "):
+                        data = line[6:].strip()
+                        if data == "[DONE]":
+                            break
+                        try:
+                            json_data = json.loads(data)
+                            message = json_data.get("message", "")
+                            full_response += message
+                            yield message
+                        except json.JSONDecodeError:
+                            logging.warning(f"Failed to parse JSON: {data}")
+            elif response.status_code == 429:
+                logging.warning("Rate limit exceeded. Changing User-Agent and retrying.")
+                for attempt in range(5):  # Try up to 5 times
+                    user_agent = get_next_user_agent()
+                    vqd_token = await update_vqd_token(user_agent)
+                    headers["User-Agent"] = user_agent
+                    headers["x-vqd-4"] = vqd_token
+                    logging.info(f"Retrying with new User-Agent: {user_agent}")
+                    response = await client.post("https://duckduckgo.com/duckchat/v1/chat", json=payload, headers=headers)
+                    if response.status_code == 200:
+                        async for line in response.aiter_lines():
+                            if line.startswith("data: "):
+                                data = line[6:].strip()
+                                if data == "[DONE]":
+                                    break
+                                try:
+                                    json_data = json.loads(data)
+                                    message = json_data.get("message", "")
+                                    yield message
+                                except json.JSONDecodeError:
+                                    logging.warning(f"Failed to parse JSON: {data}")
+                        break
+                else:
+                    raise HTTPException(status_code=429, detail="Rate limit exceeded. Please try again later.")
+            else:
+                logging.error(f"Error response from DuckDuckGo. Status code: {response.status_code}")
+                raise HTTPException(status_code=response.status_code, detail=f"Error communicating with DuckDuckGo: {response.text}")
+        except httpx.HTTPStatusError as e:
+            logging.error(f"HTTP error occurred: {str(e)}")
+            raise HTTPException(status_code=e.response.status_code, detail=str(e))
+        except httpx.RequestError as e:
+            logging.error(f"Request error occurred: {str(e)}")
+            raise HTTPException(status_code=500, detail=str(e))
+        except Exception as e:
+            logging.error(f"Unexpected error in chat_with_duckduckgo: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")
+
+@app.get("/v1/models")
+async def list_models():
+    logging.info("Listing available models")
+    models = [ModelInfo(id=model_id) for model_id in MODEL_MAPPING.keys()]
+    return {"data": models, "object": "list"}
+
+@app.post("/v1/chat/completions")
+async def chat_completion(request: ChatCompletionRequest):
+    conversation_id = str(uuid.uuid4())
+
+    logging.info(f"Received chat completion request for conversation {conversation_id}")
+    logging.info(f"Request: {request.model_dump()}")
+
+    conversation_history = conversations.get(conversation_id, [])
+    conversation_history.extend(request.messages)
+
+    async def generate():
+        try:
+            full_response = ""
+            async for chunk in chat_with_duckduckgo(" ".join([msg.content for msg in request.messages]), request.model, conversation_history):
+                full_response += chunk
+                
+                response = ChatCompletionStreamResponse(
+                    id=conversation_id,
+                    created=int(time.time()),
+                    model=request.model,
+                    choices=[
+                        ChatCompletionStreamResponseChoice(
+                            index=0,
+                            delta=DeltaMessage(content=chunk),
+                            finish_reason=None
+                        )
+                    ]
+                )
+                yield f"data: {response.model_dump_json()}\n\n"
+                await asyncio.sleep(random.uniform(0.05, 0.1))
+
+            final_response = ChatCompletionStreamResponse(
+                id=conversation_id,
+                created=int(time.time()),
+                model=request.model,
+                choices=[
+                    ChatCompletionStreamResponseChoice(
+                        index=0,
+                        delta=DeltaMessage(),
+                        finish_reason="stop"
+                    )
+                ]
+            )
+            yield f"data: {final_response.model_dump_json()}\n\n"
+            yield "data: [DONE]\n\n"
+        except Exception as e:
+            logging.error(f"Error during streaming: {str(e)}")
+            yield f"data: {json.dumps({'error': str(e)})}\n\n"
+
+    if request.stream:
+        return StreamingResponse(generate(), media_type="text/event-stream")
+    else:
+        full_response = ""
+        async for chunk in chat_with_duckduckgo(" ".join([msg.content for msg in request.messages]), request.model, conversation_history):
+            full_response += chunk
+
+        response = ChatCompletionResponse(
+            id=conversation_id,
+            created=int(time.time()),
+            model=request.model,
+            choices=[
+                ChatCompletionResponseChoice(
+                    index=0,
+                    message=ChatMessage(role="assistant", content=full_response),
+                    finish_reason="stop"
+                )
+            ],
+            usage=ChatCompletionResponseUsage(
+                prompt_tokens=sum(len(msg.content.split()) for msg in conversation_history),
+                completion_tokens=len(full_response.split()),
+                total_tokens=sum(len(msg.content.split()) for msg in conversation_history) + len(full_response.split())
+            )
+        )
+
+        conversation_history.append(ChatMessage(role="assistant", content=full_response))
+        conversations[conversation_id] = conversation_history
+        
+        return response
+
+@app.delete("/v1/conversations/{conversation_id}")
+async def end_conversation(conversation_id: str):
+    if conversation_id in conversations:
+        del conversations[conversation_id]
+        logging.info(f"Conversation {conversation_id} ended and context cleared")
+        return {"message": f"Conversation {conversation_id} ended and context cleared."}
+    else:
+        logging.warning(f"Attempt to end non-existent conversation {conversation_id}")
+        raise HTTPException(status_code=404, detail="Conversation not found")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=1337)