Update main.py

main.py

@@ -1,21 +1,21 @@
-import os
-import re
-import random
-import string
+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
 import uuid
-import json
 import logging
-import asyncio
+import os
 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
-from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
-from pydantic import BaseModel
+
+# Assuming the updated Blackbox class is defined above or imported
 
 # Configure logging
 logging.basicConfig(
@@ -89,333 +89,7 @@ async def get_api_key(request: Request, authorization: str = Header(None)) -> st
         raise HTTPException(status_code=401, detail='Invalid API key')
     return api_key
 
-# 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, url: str, alt: str):
-        self.url = url
-        self.alt = alt
-
-def to_data_uri(image: Any) -> str:
-    return "data:image/png;base64,..."  # Replace with actual base64 data
-
-class Blackbox:
-    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'
-    image_models = ['ImageGeneration']
-    models = [
-        default_model,
-        'blackboxai-pro',
-        "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',
-        *image_models,
-        'Niansuh',
-    ]
-
-    # 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"},
-        'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"},
-    }
-    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",
-    }
-
-    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": f"{url}/?model=blackboxai",
-        "gpt-4o": f"{url}/?model=gpt-4o",
-        "gemini-pro": f"{url}/?model=gemini-pro",
-        "claude-sonnet-3.5": f"{url}/?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
-
-    @classmethod
-    async def create_async_generator(
-        cls,
-        model: str,
-        messages: List[Dict[str, str]],
-        proxy: Optional[str] = None,
-        image: Any = None,
-        image_name: Optional[str] = None,
-        webSearchMode: bool = False,
-        **kwargs
-    ) -> AsyncGenerator[Any, None]:
-        model = cls.get_model(model)
-        if model is None:
-            logger.error(f"Model {model} is not available.")
-            raise ModelNotWorkingException(model)
-
-        logger.info(f"Selected model: {model}")
-
-        if not cls.working or model not in cls.models:
-            logger.error(f"Model {model} is not working or not supported.")
-            raise ModelNotWorkingException(model)
-        
-        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": cls.model_referers.get(model, cls.url),
-            "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",
-        }
-
-        if model in cls.model_prefixes:
-            prefix = cls.model_prefixes[model]
-            if not messages[0]['content'].startswith(prefix):
-                logger.debug(f"Adding prefix '{prefix}' to the first message.")
-                messages[0]['content'] = f"{prefix} {messages[0]['content']}"
-        
-        random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
-        messages[-1]['id'] = random_id
-        messages[-1]['role'] = 'user'
-
-        # Don't log the full message content for privacy
-        logger.debug(f"Generated message ID: {random_id} for model: {model}")
-
-        if image is not None:
-            messages[-1]['data'] = {
-                'fileText': '',
-                'imageBase64': to_data_uri(image),
-                'title': image_name
-            }
-            messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content']
-            logger.debug("Image data added to the message.")
-        
-        data = {
-            "messages": messages,
-            "id": random_id,
-            "previewToken": None,
-            "userId": None,
-            "codeModelMode": True,
-            "agentMode": {},
-            "trendingAgentMode": {},
-            "isMicMode": False,
-            "userSystemPrompt": None,
-            "maxTokens": 99999999,
-            "playgroundTopP": 0.9,
-            "playgroundTemperature": 0.5,
-            "isChromeExt": False,
-            "githubToken": None,
-            "clickedAnswer2": False,
-            "clickedAnswer3": False,
-            "clickedForceWebSearch": False,
-            "visitFromDelta": False,
-            "mobileClient": False,
-            "userSelectedModel": None,
-            "webSearchMode": webSearchMode,
-        }
-
-        if model in cls.agentMode:
-            data["agentMode"] = cls.agentMode[model]
-        elif model in cls.trendingAgentMode:
-            data["trendingAgentMode"] = cls.trendingAgentMode[model]
-        elif model in cls.userSelectedModel:
-            data["userSelectedModel"] = cls.userSelectedModel[model]
-        logger.info(f"Sending request to {cls.api_endpoint} with data (excluding messages).")
-
-        timeout = ClientTimeout(total=60)  # Set an appropriate timeout
-        retry_attempts = 10  # Set the number of retry attempts
-
-        for attempt in range(retry_attempts):
-            try:
-                async with ClientSession(headers=headers, timeout=timeout) as session:
-                    async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
-                        response.raise_for_status()
-                        logger.info(f"Received response with status {response.status}")
-                        if model == 'ImageGeneration':
-                            response_text = await response.text()
-                            url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
-                            if url_match:
-                                image_url = url_match.group(0)
-                                logger.info(f"Image URL found.")
-                                yield ImageResponse(image_url, alt=messages[-1]['content'])
-                            else:
-                                logger.error("Image URL not found in the response.")
-                                raise Exception("Image URL not found in the response")
-                        else:
-                            full_response = ""
-                            search_results_json = ""
-                            try:
-                                async for chunk, _ in response.content.iter_chunks():
-                                    if chunk:
-                                        decoded_chunk = chunk.decode(errors='ignore')
-                                        decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
-                                        if decoded_chunk.strip():
-                                            if '$~~~$' in decoded_chunk:
-                                                search_results_json += decoded_chunk
-                                            else:
-                                                full_response += decoded_chunk
-                                                yield decoded_chunk
-                                logger.info("Finished streaming response chunks.")
-                            except Exception as e:
-                                logger.exception("Error while iterating over response chunks.")
-                                raise e
-                            if data["webSearchMode"] and search_results_json:
-                                match = re.search(r'\$~~~\$(.*?)\$~~~\$', search_results_json, re.DOTALL)
-                                if match:
-                                    try:
-                                        search_results = json.loads(match.group(1))
-                                        formatted_results = "\n\n**Sources:**\n"
-                                        for i, result in enumerate(search_results[:5], 1):
-                                            formatted_results += f"{i}. [{result['title']}]({result['link']})\n"
-                                        logger.info("Formatted search results.")
-                                        yield formatted_results
-                                    except json.JSONDecodeError as je:
-                                        logger.error("Failed to parse search results JSON.")
-                                        raise je
-                break  # Exit the retry loop if successful
-            except ClientError as ce:
-                logger.error(f"Client error occurred: {ce}. Retrying attempt {attempt + 1}/{retry_attempts}")
-                if attempt == retry_attempts - 1:
-                    raise HTTPException(status_code=502, detail="Error communicating with the external API.")
-            except asyncio.TimeoutError:
-                logger.error(f"Request timed out. Retrying attempt {attempt + 1}/{retry_attempts}")
-                if attempt == retry_attempts - 1:
-                    raise HTTPException(status_code=504, detail="External API request timed out.")
-            except Exception as e:
-                logger.error(f"Unexpected error: {e}. Retrying attempt {attempt + 1}/{retry_attempts}")
-                if attempt == retry_attempts - 1:
-                    raise HTTPException(status_code=500, detail=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
+# Define request models
 class Message(BaseModel):
     role: str
     content: str
@@ -466,6 +140,47 @@ 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)])
 async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
     client_ip = req.client.host
@@ -486,7 +201,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
             image=None,
             image_name=None,
-            webSearchMode=request.webSearchMode
+            websearch=request.webSearchMode
         )
 
         if request.stream:
@@ -496,7 +211,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                     async for chunk in async_generator:
                         if isinstance(chunk, ImageResponse):
                             # Handle image responses if necessary
-                            image_markdown = f"![image]({chunk.url})\n"
+                            image_markdown = f"![image]({chunk.images})\n"
                             assistant_content += image_markdown
                             response_chunk = create_response(image_markdown, request.model, finish_reason=None)
                         else:
@@ -519,9 +234,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
-                    # *** Key Correction Starts Here ***
                     prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
-                    # *** Key Correction Ends Here ***
                     completion_tokens = len(assistant_content.split())
                     total_tokens = prompt_tokens + completion_tokens
                     estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
@@ -563,7 +276,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             response_content = ""
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponse):
-                    response_content += f"![image]({chunk.url})\n"
+                    response_content += f"![image]({chunk.images})\n"
                 else:
                     response_content += chunk
 
@@ -606,7 +319,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         logger.exception(f"An unexpected error occurred while processing the chat completions request from IP: {client_ip}.")
         raise HTTPException(status_code=500, detail=str(e))
 
-# Endpoint: POST /v1/tokenizer
+# Tokenizer Endpoint
 @app.post("/v1/tokenizer", dependencies=[Depends(rate_limiter_per_ip)])
 async def tokenizer(request: TokenizerRequest, req: Request):
     client_ip = req.client.host
@@ -615,14 +328,14 @@ async def tokenizer(request: TokenizerRequest, req: Request):
     logger.info(f"Tokenizer requested from IP: {client_ip} | Text length: {len(text)}")
     return {"text": text, "tokens": token_count}
 
-# Endpoint: GET /v1/models
+# Get Models Endpoint
 @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
 async def get_models(req: Request):
     client_ip = req.client.host
     logger.info(f"Fetching available models from IP: {client_ip}")
     return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
 
-# Endpoint: GET /v1/models/{model}/status
+# Model Status Endpoint
 @app.get("/v1/models/{model}/status", dependencies=[Depends(rate_limiter_per_ip)])
 async def model_status(model: str, req: Request):
     client_ip = req.client.host
@@ -636,21 +349,21 @@ async def model_status(model: str, req: Request):
         logger.warning(f"Model not found: {model} from IP: {client_ip}")
         raise HTTPException(status_code=404, detail="Model not found")
 
-# Endpoint: GET /v1/health
+# Health Check Endpoint
 @app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
 async def health_check(req: Request):
     client_ip = req.client.host
     logger.info(f"Health check requested from IP: {client_ip}")
     return {"status": "ok"}
 
-# Endpoint: GET /v1/chat/completions (GET method)
+# Redirect GET requests to /v1/chat/completions to 'about:blank'
 @app.get("/v1/chat/completions")
 async def chat_completions_get(req: Request):
     client_ip = req.client.host
     logger.info(f"GET request made to /v1/chat/completions from IP: {client_ip}, redirecting to 'about:blank'")
     return RedirectResponse(url='about:blank')
 
-# Custom exception handler to match OpenAI's error format
+# Custom Exception Handler
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request: Request, exc: HTTPException):
     client_ip = request.client.host