Update api/utils.py

api/utils.py

@@ -1,185 +1,264 @@
-# Blackbox.py
+# utils.py
 
-from __future__ import annotations
-
-from aiohttp import ClientSession
-import uuid
+from datetime import datetime
 import json
-import re
-
-from ..typing import AsyncResult, Messages, ImageType
-from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
-from ..image import ImageResponse, to_data_uri
-
-class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
-    label = "Blackbox AI"
-    url = "https://www.blackbox.ai"
-    api_endpoint = "https://www.blackbox.ai/api/chat"
-    working = True
-    supports_stream = True
-    supports_system_message = True
-    supports_message_history = True
-    
-    default_model = 'blackboxai'
+import uuid
+import asyncio
+import random
+import string
+from typing import Any, Dict, Optional
+
+import httpx
+from fastapi import HTTPException
+from api.config import (
+    MODEL_MAPPING,
+    BASE_URL,
+    AGENT_MODE,
+    TRENDING_AGENT_MODE,
+)
+from api.models import ChatRequest
+from api.logger import setup_logger
+
+logger = setup_logger(__name__)
+
+# Helper function to create a UUID-based chat ID
+def generate_chat_id() -> str:
+    return str(uuid.uuid4())
+
+# Helper function to create chat completion data
+def create_chat_completion_data(
+    content: str, model: str, timestamp: int, finish_reason: Optional[str] = None
+) -> Dict[str, Any]:
+    return {
+        "id": f"chatcmpl-{uuid.uuid4()}",
+        "object": "chat.completion.chunk",
+        "created": timestamp,
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "delta": {"content": content, "role": "assistant"},
+                "finish_reason": finish_reason,
+            }
+        ],
+        "usage": None,
+    }
+
+# Function to convert message to dictionary format, ensuring base64 data
+def message_to_dict(message):
+    if isinstance(message.content, str):
+        content = message.content
+    elif isinstance(message.content, list) and len(message.content) > 0:
+        content = message.content[0].get("text", "")
+    else:
+        content = ""
     
-    image_models = ['Image Generation', 'repomap']
-   
-    userSelectedModel = ['gpt-4o', 'gemini-pro', 'claude-sonnet-3.5', 'blackboxai-pro']
+    if isinstance(message.content, list) and len(message.content) == 2 and "image_url" in message.content[1]:
+        # Ensure base64 images are always included for all models
+        return {
+            "role": message.role,
+            "content": content,
+            "data": {
+                "imageBase64": message.content[1]["image_url"]["url"],
+                "fileText": "",
+                "title": "snapshot",
+            },
+        }
+    return {"role": message.role, "content": content}
+
+# Function to strip model prefix from content if present
+def strip_model_prefix(content: str) -> str:
+    """Function retained but no model_prefix is used."""
+    # Since model_prefix is removed, this function can simply return content
+    return content
+
+# Function to get the correct referer URL for logging
+def get_referer_url(chat_id: str, model: str) -> str:
+    """Generate the referer URL based on specific models listed in MODEL_REFERERS."""
+    # MODEL_REFERERS has been removed; referer_url defaults to BASE_URL
+    return BASE_URL
+
+# Process streaming response without model prefixes and referers
+async def process_streaming_response(request: ChatRequest):
+    chat_id = generate_chat_id()
+    referer_url = get_referer_url(chat_id, request.model)
+    logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}")
+
+    agent_mode = AGENT_MODE.get(request.model, {})
+    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
     
-    agentMode = {
-        'Image Generation': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
+    headers_api_chat = {
+        'Accept': '*/*',
+        'Accept-Language': 'en-US,en;q=0.9',
+        'Cache-Control': 'no-cache',
+        'Origin': BASE_URL,
+        'Pragma': 'no-cache',
+        'Priority': 'u=1, i',
+        'User-Agent': (
+            'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
+            'AppleWebKit/537.36 (KHTML, like Gecko) '
+            'Chrome/130.0.0.0 Safari/537.36'
+        ),
+        'Content-Type': 'application/json',
+        'Referer': referer_url,  # Retain referer if necessary
     }
-    
-    trendingAgentMode = {
-        "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
-        "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
-        'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
-        'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405"},
-        #
-        'Python Agent': {'mode': True, 'id': "Python Agent"},
-        'Java Agent': {'mode': True, 'id': "Java Agent"},
-        'JavaScript Agent': {'mode': True, 'id': "JavaScript Agent"},
-        'HTML Agent': {'mode': True, 'id': "HTML Agent"},
-        'Google Cloud Agent': {'mode': True, 'id': "Google Cloud Agent"},
-        'Android Developer': {'mode': True, 'id': "Android Developer"},
-        'Swift Developer': {'mode': True, 'id': "Swift Developer"},
-        'Next.js Agent': {'mode': True, 'id': "Next.js Agent"},
-        'MongoDB Agent': {'mode': True, 'id': "MongoDB Agent"},
-        'PyTorch Agent': {'mode': True, 'id': "PyTorch Agent"},
-        'React Agent': {'mode': True, 'id': "React Agent"},
-        'Xcode Agent': {'mode': True, 'id': "Xcode Agent"},
-        'AngularJS Agent': {'mode': True, 'id': "AngularJS Agent"},
-        'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
-        #
-        'repomap': {'mode': True, 'id': "repomap"},
-        #
-        'Heroku Agent': {'mode': True, 'id': "Heroku Agent"},
-        'Godot Agent': {'mode': True, 'id': "Godot Agent"},
-        'Go Agent': {'mode': True, 'id': "Go Agent"},
-        'Gitlab Agent': {'mode': True, 'id': "Gitlab Agent"},
-        'Git Agent': {'mode': True, 'id': "Git Agent"},
-        'Flask Agent': {'mode': True, 'id': "Flask Agent"},
-        'Firebase Agent': {'mode': True, 'id': "Firebase Agent"},
-        'FastAPI Agent': {'mode': True, 'id': "FastAPI Agent"},
-        'Erlang Agent': {'mode': True, 'id': "Erlang Agent"},
-        'Electron Agent': {'mode': True, 'id': "Electron Agent"},
-        'Docker Agent': {'mode': True, 'id': "Docker Agent"},
-        'DigitalOcean Agent': {'mode': True, 'id': "DigitalOcean Agent"},
-        'Bitbucket Agent': {'mode': True, 'id': "Bitbucket Agent"},
-        'Azure Agent': {'mode': True, 'id': "Azure Agent"},
-        'Flutter Agent': {'mode': True, 'id': "Flutter Agent"},
-        'Youtube Agent': {'mode': True, 'id': "Youtube Agent"},
-        'builder Agent': {'mode': True, 'id': "builder Agent"},
+
+    if request.model == 'o1-preview':
+        delay_seconds = random.randint(1, 60)
+        logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' (Chat ID: {chat_id})")
+        await asyncio.sleep(delay_seconds)
+
+    json_data = {
+        "agentMode": agent_mode,
+        "clickedAnswer2": False,
+        "clickedAnswer3": False,
+        "clickedForceWebSearch": False,
+        "codeModelMode": True,
+        "githubToken": None,
+        "id": chat_id,
+        "isChromeExt": False,
+        "isMicMode": False,
+        "maxTokens": request.max_tokens,
+        "messages": [message_to_dict(msg) for msg in request.messages],
+        "mobileClient": False,
+        "playgroundTemperature": request.temperature,
+        "playgroundTopP": request.top_p,
+        "previewToken": None,
+        "trendingAgentMode": trending_agent_mode,
+        "userId": None,
+        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
+        "userSystemPrompt": None,
+        "validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
+        "visitFromDelta": False,
     }
+
+    async with httpx.AsyncClient() as client:
+        try:
+            async with client.stream(
+                "POST",
+                f"{BASE_URL}/api/chat",
+                headers=headers_api_chat,
+                json=json_data,
+                timeout=100,
+            ) as response:
+                response.raise_for_status()
+                async for line in response.aiter_lines():
+                    timestamp = int(datetime.now().timestamp())
+                    if line:
+                        content = line
+                        if content.startswith("$@$v=undefined-rv1$@$"):
+                            content = content[21:]
+                        cleaned_content = strip_model_prefix(content)
+                        yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
+
+                yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
+                yield "data: [DONE]\n\n"
+        except httpx.HTTPStatusError as e:
+            logger.error(f"HTTP error occurred for Chat ID {chat_id}: {e}")
+            raise HTTPException(status_code=e.response.status_code, detail=str(e))
+        except httpx.RequestError as e:
+            logger.error(f"Error occurred during request for Chat ID {chat_id}: {e}")
+            raise HTTPException(status_code=500, detail=str(e))
+
+# Process non-streaming response without model prefixes and referers
+async def process_non_streaming_response(request: ChatRequest):
+    chat_id = generate_chat_id()
+    referer_url = get_referer_url(chat_id, request.model)
+    logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}")
+
+    agent_mode = AGENT_MODE.get(request.model, {})
+    trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
     
-    models = [default_model, *userSelectedModel, *list(agentMode.keys()), *list(trendingAgentMode.keys())]
-    
-    model_aliases = {
-        "gemini-flash": "gemini-1.5-flash",
-        "claude-3.5-sonnet": "claude-sonnet-3.5",
-        "flux": "Image Generation",
+    headers_api_chat = {
+        'Accept': '*/*',
+        'Accept-Language': 'en-US,en;q=0.9',
+        'Cache-Control': 'no-cache',
+        'Origin': BASE_URL,
+        'Pragma': 'no-cache',
+        'Priority': 'u=1, i',
+        'User-Agent': (
+            'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
+            'AppleWebKit/537.36 (KHTML, like Gecko) '
+            'Chrome/130.0.0.0 Safari/537.36'
+        ),
+        'Content-Type': 'application/json',
+        'Referer': referer_url,  # Retain referer if necessary
+    }
+    headers_chat = {
+        'Accept': '*/*',
+        'Accept-Language': 'en-US,en;q=0.9',
+        'Cache-Control': 'no-cache',
+        'Content-Type': 'text/plain;charset=UTF-8',
+        'Origin': BASE_URL,
+        'Pragma': 'no-cache',
+        'Priority': 'u=1, i',
+        'Referer': referer_url,
+        'Next-Action': str(uuid.uuid4()),
+        'Next-Router-State-Tree': json.dumps([""]),
+        'Next-URL': '/',
     }
 
-    @staticmethod
-    def generate_id() -> str:
-        """Generate a UUID-based message ID."""
-        return str(uuid.uuid4())
-        
-    @classmethod
-    def get_model(cls, model: str) -> str:
-        if model in cls.models:
-            return model
-        elif model in cls.model_aliases:
-            return cls.model_aliases[model]
-        else:
-            return cls.default_model
-
-    @classmethod
-    async def create_async_generator(
-        cls,
-        model: str,
-        messages: Messages,
-        proxy: str = None,
-        web_search: bool = False,
-        image: ImageType = None,
-        image_name: str = None,
-        **kwargs
-    ) -> AsyncResult:
-        model = cls.get_model(model)
-        message_id = cls.generate_id()
-        # Removed add_prefix_to_messages since prefixes are no longer used
-        messages_processed = messages.copy()
-
-        if image is not None:
-            messages_processed[-1]['data'] = {
-                'fileText': '',
-                'imageBase64': to_data_uri(image),
-                'title': image_name
-            }
+    if request.model == 'o1-preview':
+        delay_seconds = random.randint(20, 60)
+        logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' (Chat ID: {chat_id})")
+        await asyncio.sleep(delay_seconds)
 
-        headers = {
-            'Accept': '*/*',
-            'Accept-Language': 'en-US,en;q=0.9',
-            'Cache-Control': 'no-cache',
-            'Content-Type': 'application/json',
-            'Origin': cls.url,
-            'Pragma': 'no-cache',
-            'Priority': 'u=1, i',
-            'Referer': f'{cls.url}/',
-            'User-Agent': (
-                'Mozilla/5.0 (X11; Linux x86_64) '
-                'AppleWebKit/537.36 (KHTML, like Gecko) '
-                'Chrome/130.0.0.0 Safari/537.36'
-            )
-        }
-        
-        data = {
-            "messages": messages_processed,
-            "id": message_id,
-            "previewToken": None,
-            "userId": None,
-            "codeModelMode": True,
-            "agentMode": cls.agentMode.get(model, {}) if model in cls.agentMode else {},
-            "trendingAgentMode": cls.trendingAgentMode.get(model, {}) if model in cls.trendingAgentMode else {},
-            "isMicMode": False,
-            "userSystemPrompt": None,
-            "maxTokens": 1024,
-            "playgroundTopP": 0.9,
-            "playgroundTemperature": 0.5,
-            "isChromeExt": False,
-            "githubToken": None,
-            "clickedAnswer2": False,
-            "clickedAnswer3": False,
-            "clickedForceWebSearch": False,
-            "visitFromDelta": False,
-            "mobileClient": False,
-            "userSelectedModel": model if model in cls.userSelectedModel else None,
-            "webSearchMode": web_search,
-            "validated": "69783381-2ce4-4dbd-ac78-35e9063feabc"
-        }
+    json_data = {
+        "agentMode": agent_mode,
+        "clickedAnswer2": False,
+        "clickedAnswer3": False,
+        "clickedForceWebSearch": False,
+        "codeModelMode": True,
+        "githubToken": None,
+        "id": chat_id,
+        "isChromeExt": False,
+        "isMicMode": False,
+        "maxTokens": request.max_tokens,
+        "messages": [message_to_dict(msg) for msg in request.messages],
+        "mobileClient": False,
+        "playgroundTemperature": request.temperature,
+        "playgroundTopP": request.top_p,
+        "previewToken": None,
+        "trendingAgentMode": trending_agent_mode,
+        "userId": None,
+        "userSelectedModel": MODEL_MAPPING.get(request.model, request.model),
+        "userSystemPrompt": None,
+        "validated": "69783381-2ce4-4dbd-ac78-35e9063feabc",
+        "visitFromDelta": False,
+    }
 
-        async with ClientSession(headers=headers) as session:
-            async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
+    full_response = ""
+    async with httpx.AsyncClient() as client:
+        try:
+            async with client.stream(
+                method="POST", url=f"{BASE_URL}/api/chat", headers=headers_api_chat, json=json_data
+            ) as response:
                 response.raise_for_status()
-                response_text = await response.text()
-                
-                if model in cls.image_models:
-                    image_matches = re.findall(r'!\[.*?\]\((https?://[^\)]+)\)', response_text)
-                    if image_matches:
-                        image_url = image_matches[0]
-                        image_response = ImageResponse(images=[image_url], alt="Generated Image")
-                        yield image_response
-                        return
-
-                json_match = re.search(r'\$~~~\$(.*?)\$~~~\$', response_text, re.DOTALL)
-                if json_match:
-                    search_results = json.loads(json_match.group(1))
-                    answer = response_text.split('$~~~$')[-1].strip()
-                    
-                    formatted_response = f"{answer}\n\n**Source:**"
-                    for i, result in enumerate(search_results, 1):
-                        formatted_response += f"\n{i}. {result['title']}: {result['link']}"
-                    
-                    yield formatted_response
-                else:
-                    yield response_text.strip()
+                async for chunk in response.aiter_text():
+                    full_response += chunk
+        except httpx.HTTPStatusError as e:
+            logger.error(f"HTTP error occurred for Chat ID {chat_id}: {e}")
+            raise HTTPException(status_code=e.response.status_code, detail=str(e))
+        except httpx.RequestError as e:
+            logger.error(f"Error occurred during request for Chat ID {chat_id}: {e}")
+            raise HTTPException(status_code=500, detail=str(e))
+    
+    if full_response.startswith("$@$v=undefined-rv1$@$"):
+        full_response = full_response[21:]
+
+    cleaned_full_response = strip_model_prefix(full_response)
+
+    return {
+        "id": f"chatcmpl-{uuid.uuid4()}",
+        "object": "chat.completion",
+        "created": int(datetime.now().timestamp()),
+        "model": request.model,
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": cleaned_full_response},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": None,
+    }