Update api/utils.py

api/utils.py

@@ -1,258 +1,185 @@
-from datetime import datetime
-import json
-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
+# Blackbox.py
 
-logger = setup_logger(__name__)
+from __future__ import annotations
 
-# Helper function to create a random 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 = ""
+from aiohttp import ClientSession
+import uuid
+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
     
-    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, {})
+    default_model = 'blackboxai'
     
-    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
-    }
-
-    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,
+    image_models = ['Image Generation', 'repomap']
+   
+    userSelectedModel = ['gpt-4o', 'gemini-pro', 'claude-sonnet-3.5', 'blackboxai-pro']
+    
+    agentMode = {
+        'Image Generation': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
     }
-
-    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, {})
     
-    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"},
     }
-    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': '/',
+    
+    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",
     }
 
-    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)
+    @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
+            }
 
-    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,
-    }
+        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"
+        }
 
-    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:
+        async with ClientSession(headers=headers) as session:
+            async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
                 response.raise_for_status()
-                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,
-    }
+                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()