Update main.py

main.py

@@ -1,21 +1,23 @@
-from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
-from pydantic import BaseModel
-from typing import List, Dict, Any, Optional, AsyncGenerator, Union
+from __future__ import annotations
+
+import os
+import re
+import random
+import string
 import uuid
+import json
 import logging
-import os
+import asyncio
 import time
 from collections import defaultdict
+from typing import List, Dict, Any, Optional, AsyncGenerator, Union
+
 from datetime import datetime
-import asyncio
-import json
-import re
-import random
-import string
-from aiohttp import ClientSession, ClientTimeout, ClientError
 
-# Assuming the updated Blackbox class is defined above or imported
+from aiohttp import ClientSession, ClientTimeout, ClientError, ClientResponseError
+from fastapi import FastAPI, HTTPException, Request, Depends, Header
+from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
+from pydantic import BaseModel
 
 # Configure logging
 logging.basicConfig(
@@ -53,7 +55,10 @@ async def cleanup_rate_limit_stores():
     """
     while True:
         current_time = time.time()
-        ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
+        ips_to_delete = [
+            ip for ip, value in rate_limit_store.items()
+            if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2
+        ]
         for ip in ips_to_delete:
             del rate_limit_store[ip]
             logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
@@ -89,7 +94,418 @@ async def get_api_key(request: Request, authorization: str = Header(None)) -> st
         raise HTTPException(status_code=401, detail='Invalid API key')
     return api_key
 
-# Define request models
+# Custom exception for model not working
+class ModelNotWorkingException(Exception):
+    def __init__(self, model: str):
+        self.model = model
+        self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed."
+        super().__init__(self.message)
+
+# Mock implementations for ImageResponse and to_data_uri
+class ImageResponse:
+    def __init__(self, images: str, alt: str):
+        self.images = images
+        self.alt = alt
+
+def to_data_uri(image: Any) -> str:
+    return "data:image/png;base64,..."  # Replace with actual base64 data
+
+# Placeholder classes for AsyncGeneratorProvider and ProviderModelMixin
+class AsyncGeneratorProvider:
+    pass  # Implement as per your actual provider's requirements
+
+class ProviderModelMixin:
+    pass  # Implement as per your actual provider's requirements
+
+class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
+    label = "Blackbox AI"
+    url = "https://www.blackbox.ai"
+    api_endpoint = "https://www.blackbox.ai/api/chat"
+    working = True
+    supports_gpt_4 = True
+    supports_stream = True
+    supports_system_message = True
+    supports_message_history = True
+
+    default_model = 'blackboxai'
+    image_models = ['ImageGeneration']
+    models = [
+        default_model,
+        'blackboxai-pro',
+        *image_models,
+        "llama-3.1-8b",
+        'llama-3.1-70b',
+        'llama-3.1-405b',
+        'gpt-4o',
+        'gemini-pro',
+        'gemini-1.5-flash',
+        'claude-sonnet-3.5',
+        'PythonAgent',
+        'JavaAgent',
+        'JavaScriptAgent',
+        'HTMLAgent',
+        'GoogleCloudAgent',
+        'AndroidDeveloper',
+        'SwiftDeveloper',
+        'Next.jsAgent',
+        'MongoDBAgent',
+        'PyTorchAgent',
+        'ReactAgent',
+        'XcodeAgent',
+        'AngularJSAgent',
+    ]
+
+    # Filter models based on AVAILABLE_MODELS
+    if AVAILABLE_MODELS:
+        models = [model for model in models if model in AVAILABLE_MODELS]
+
+    agentMode = {
+        'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
+    }
+
+    trendingAgentMode = {
+        "blackboxai": {},
+        "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-405b"},
+        'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
+        'PythonAgent': {'mode': True, 'id': "Python Agent"},
+        'JavaAgent': {'mode': True, 'id': "Java Agent"},
+        'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"},
+        'HTMLAgent': {'mode': True, 'id': "HTML Agent"},
+        'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"},
+        'AndroidDeveloper': {'mode': True, 'id': "Android Developer"},
+        'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"},
+        'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"},
+        'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"},
+        'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
+        'ReactAgent': {'mode': True, 'id': "React Agent"},
+        'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
+        'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
+    }
+
+    userSelectedModel = {
+        "gpt-4o": "gpt-4o",
+        "gemini-pro": "gemini-pro",
+        'claude-sonnet-3.5': "claude-sonnet-3.5",
+        "niansuh": "Niansuh",  # Added based on model_aliases
+    }
+
+    model_prefixes = {
+        'gpt-4o': '@GPT-4o',
+        'gemini-pro': '@Gemini-PRO',
+        'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
+        'PythonAgent': '@Python Agent',
+        'JavaAgent': '@Java Agent',
+        'JavaScriptAgent': '@JavaScript Agent',
+        'HTMLAgent': '@HTML Agent',
+        'GoogleCloudAgent': '@Google Cloud Agent',
+        'AndroidDeveloper': '@Android Developer',
+        'SwiftDeveloper': '@Swift Developer',
+        'Next.jsAgent': '@Next.js Agent',
+        'MongoDBAgent': '@MongoDB Agent',
+        'PyTorchAgent': '@PyTorch Agent',
+        'ReactAgent': '@React Agent',
+        'XcodeAgent': '@Xcode Agent',
+        'AngularJSAgent': '@AngularJS Agent',
+        'blackboxai-pro': '@BLACKBOXAI-PRO',
+        'ImageGeneration': '@Image Generation',
+        'Niansuh': '@Niansuh',
+    }
+
+    model_referers = {
+        "blackboxai": "/?model=blackboxai",
+        "gpt-4o": "/?model=gpt-4o",
+        "gemini-pro": "/?model=gemini-pro",
+        "claude-sonnet-3.5": "/?model=claude-sonnet-3.5"
+    }
+
+    model_aliases = {
+        "gemini-flash": "gemini-1.5-flash",
+        "claude-3.5-sonnet": "claude-sonnet-3.5",
+        "flux": "ImageGeneration",
+        "niansuh": "Niansuh",
+    }
+
+    @classmethod
+    def get_model(cls, model: str) -> Optional[str]:
+        if model in cls.models:
+            return model
+        elif model in cls.userSelectedModel and cls.userSelectedModel[model] in cls.models:
+            return cls.userSelectedModel[model]
+        elif model in cls.model_aliases and cls.model_aliases[model] in cls.models:
+            return cls.model_aliases[model]
+        else:
+            return cls.default_model if cls.default_model in cls.models else None
+
+    @staticmethod
+    def generate_random_string(length: int = 7) -> str:
+        characters = string.ascii_letters + string.digits
+        return ''.join(random.choices(characters, k=length))
+
+    @staticmethod
+    def generate_next_action() -> str:
+        return uuid.uuid4().hex
+
+    @staticmethod
+    def generate_next_router_state_tree() -> str:
+        router_state = [
+            "",
+            {
+                "children": [
+                    "(chat)",
+                    {
+                        "children": [
+                            "__PAGE__",
+                            {}
+                        ]
+                    }
+                ]
+            },
+            None,
+            None,
+            True
+        ]
+        return json.dumps(router_state)
+
+    @staticmethod
+    def clean_response(text: str) -> str:
+        pattern = r'^\$\@\$v=undefined-rv1\$\@\$'
+        cleaned_text = re.sub(pattern, '', text)
+        return cleaned_text
+
+    @classmethod
+    async def create_async_generator(
+        cls,
+        model: str,
+        messages: List[Dict[str, str]],
+        proxy: Optional[str] = None,
+        websearch: bool = False,
+        **kwargs
+    ) -> AsyncGenerator[Union[str, ImageResponse], None]:
+        """
+        Creates an asynchronous generator for streaming responses from Blackbox AI.
+
+        Parameters:
+            model (str): Model to use for generating responses.
+            messages (List[Dict[str, str]]): Message history.
+            proxy (Optional[str]): Proxy URL, if needed.
+            websearch (bool): Enables or disables web search mode.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            Union[str, ImageResponse]: Segments of the generated response or ImageResponse objects.
+        """
+        model = cls.get_model(model)
+        if model is None:
+            logger.error(f"Model {model} is not available.")
+            raise ModelNotWorkingException(model)
+
+        chat_id = cls.generate_random_string()
+        next_action = cls.generate_next_action()
+        next_router_state_tree = cls.generate_next_router_state_tree()
+
+        agent_mode = cls.agentMode.get(model, {})
+        trending_agent_mode = cls.trendingAgentMode.get(model, {})
+
+        prefix = cls.model_prefixes.get(model, "")
+
+        formatted_prompt = ""
+        for message in messages:
+            role = message.get('role', '').capitalize()
+            content = message.get('content', '')
+            if role and content:
+                formatted_prompt += f"{role}: {content}\n"
+
+        if prefix:
+            formatted_prompt = f"{prefix} {formatted_prompt}".strip()
+
+        referer_path = cls.model_referers.get(model, f"/?model={model}")
+        referer_url = f"{cls.url}{referer_path}"
+
+        common_headers = {
+            'accept': '*/*',
+            'accept-language': 'en-US,en;q=0.9',
+            'cache-control': 'no-cache',
+            'origin': cls.url,
+            'pragma': 'no-cache',
+            'priority': 'u=1, i',
+            'sec-ch-ua': '"Chromium";v="129", "Not=A?Brand";v="8"',
+            'sec-ch-ua-mobile': '?0',
+            'sec-ch-ua-platform': '"Linux"',
+            'sec-fetch-dest': 'empty',
+            'sec-fetch-mode': 'cors',
+            'sec-fetch-site': 'same-origin',
+            'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) '
+                          'AppleWebKit/537.36 (KHTML, like Gecko) '
+                          'Chrome/129.0.0.0 Safari/537.36'
+        }
+
+        headers_api_chat = {
+            'Content-Type': 'application/json',
+            'Referer': referer_url
+        }
+        headers_api_chat_combined = {**common_headers, **headers_api_chat}
+
+        payload_api_chat = {
+            "messages": [
+                {
+                    "id": chat_id,
+                    "content": formatted_prompt,
+                    "role": "user"
+                }
+            ],
+            "id": chat_id,
+            "previewToken": None,
+            "userId": None,
+            "codeModelMode": True,
+            "agentMode": agent_mode,
+            "trendingAgentMode": trending_agent_mode,
+            "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,
+            "webSearchMode": websearch,
+            "userSelectedModel": cls.userSelectedModel.get(model, model)
+        }
+
+        headers_chat = {
+            'Accept': 'text/x-component',
+            'Content-Type': 'text/plain;charset=UTF-8',
+            'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
+            'next-action': next_action,
+            'next-router-state-tree': next_router_state_tree,
+            'next-url': '/'
+        }
+        headers_chat_combined = {**common_headers, **headers_chat}
+
+        data_chat = '[]'
+
+        async with ClientSession(headers=common_headers) as session:
+            try:
+                # Send initial chat request
+                async with session.post(
+                    cls.api_endpoint,
+                    headers=headers_api_chat_combined,
+                    json=payload_api_chat,
+                    proxy=proxy
+                ) as response_api_chat:
+                    response_api_chat.raise_for_status()
+                    text = await response_api_chat.text()
+                    cleaned_response = cls.clean_response(text)
+
+                    if model in cls.image_models:
+                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
+                        if match:
+                            image_url = match.group(1)
+                            image_response = ImageResponse(images=image_url, alt="Generated Image")
+                            yield image_response
+                        else:
+                            yield cleaned_response
+                    else:
+                        if websearch:
+                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
+                            if match:
+                                source_part = match.group(1).strip()
+                                answer_part = cleaned_response[match.end():].strip()
+                                try:
+                                    sources = json.loads(source_part)
+                                    source_formatted = "**Sources:**\n"
+                                    for item in sources[:5]:
+                                        title = item.get('title', 'No Title')
+                                        link = item.get('link', '#')
+                                        position = item.get('position', '')
+                                        source_formatted += f"{position}. [{title}]({link})\n"
+                                    final_response = f"{answer_part}\n\n{source_formatted}"
+                                except json.JSONDecodeError:
+                                    final_response = f"{answer_part}\n\nSource information is unavailable."
+                            else:
+                                final_response = cleaned_response
+                        else:
+                            if '$~~~$' in cleaned_response:
+                                final_response = cleaned_response.split('$~~~$')[0].strip()
+                            else:
+                                final_response = cleaned_response
+
+                        yield final_response
+            except ClientResponseError as e:
+                error_text = f"Error {e.status}: {e.message}"
+                try:
+                    error_response = await e.response.text()
+                    cleaned_error = cls.clean_response(error_response)
+                    error_text += f" - {cleaned_error}"
+                except Exception:
+                    pass
+                yield error_text
+            except Exception as e:
+                yield f"Unexpected error during /api/chat request: {str(e)}"
+
+            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
+
+            try:
+                # Send follow-up chat request
+                async with session.post(
+                    chat_url,
+                    headers=headers_chat_combined,
+                    data=data_chat,
+                    proxy=proxy
+                ) as response_chat:
+                    response_chat.raise_for_status()
+                    pass
+            except ClientResponseError as e:
+                error_text = f"Error {e.status}: {e.message}"
+                try:
+                    error_response = await e.response.text()
+                    cleaned_error = cls.clean_response(error_response)
+                    error_text += f" - {cleaned_error}"
+                except Exception:
+                    pass
+                yield error_text
+            except Exception as e:
+                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
+
+# FastAPI app setup
+app = FastAPI()
+
+# Add the cleanup task when the app starts
+@app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(cleanup_rate_limit_stores())
+    logger.info("Started rate limit store cleanup task.")
+
+# Middleware to enhance security and enforce Content-Type for specific endpoints
+@app.middleware("http")
+async def security_middleware(request: Request, call_next):
+    client_ip = request.client.host
+    # Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
+    if request.method == "POST" and request.url.path == "/v1/chat/completions":
+        content_type = request.headers.get("Content-Type")
+        if content_type != "application/json":
+            logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": {
+                        "message": "Content-Type must be application/json",
+                        "type": "invalid_request_error",
+                        "param": None,
+                        "code": None
+                    }
+                },
+            )
+    response = await call_next(request)
+    return response
+
+# Request Models
 class Message(BaseModel):
     role: str
     content: str
@@ -140,45 +556,6 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,  # To be filled in non-streaming responses
     }
 
-# Custom exception for model not working
-class ModelNotWorkingException(Exception):
-    def __init__(self, model: str):
-        self.model = model
-        self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed."
-        super().__init__(self.message)
-
-# Initialize FastAPI app
-app = FastAPI()
-
-# Add the cleanup task when the app starts
-@app.on_event("startup")
-async def startup_event():
-    asyncio.create_task(cleanup_rate_limit_stores())
-    logger.info("Started rate limit store cleanup task.")
-
-# Middleware to enhance security and enforce Content-Type for specific endpoints
-@app.middleware("http")
-async def security_middleware(request: Request, call_next):
-    client_ip = request.client.host
-    # Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
-    if request.method == "POST" and request.url.path == "/v1/chat/completions":
-        content_type = request.headers.get("Content-Type")
-        if content_type != "application/json":
-            logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
-            return JSONResponse(
-                status_code=400,
-                content={
-                    "error": {
-                        "message": "Content-Type must be application/json",
-                        "type": "invalid_request_error",
-                        "param": None,
-                        "code": None
-                    }
-                },
-            )
-    response = await call_next(request)
-    return response
-
 # FastAPI Endpoints
 
 @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
@@ -199,8 +576,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         async_generator = Blackbox.create_async_generator(
             model=request.model,
             messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
-            image=None,
-            image_name=None,
+            proxy=None,  # Add proxy if needed
             websearch=request.webSearchMode
         )
 
@@ -234,7 +610,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                         yield f"data: {json.dumps(response_chunk)}\n\n"
                     
                     # After all chunks are sent, send the final message with finish_reason
-                    prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
+                    prompt_tokens = sum(len(msg['content'].split()) for msg in request.messages)
                     completion_tokens = len(assistant_content.split())
                     total_tokens = prompt_tokens + completion_tokens
                     estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)