Update main.py

main.py

@@ -11,7 +11,7 @@ from collections import defaultdict
 from typing import List, Dict, Any, Optional, Union
 from datetime import datetime
 
-from aiohttp import ClientSession, ClientResponseError
+from aiohttp import ClientSession, ClientTimeout, ClientError, ClientResponseError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
@@ -27,11 +27,18 @@ logger = logging.getLogger(__name__)
 # Load environment variables
 API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
 RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60'))  # Requests per minute
+AVAILABLE_MODELS = os.getenv('AVAILABLE_MODELS', '')  # Comma-separated available models
 
 if not API_KEYS or API_KEYS == ['']:
     logger.error("No API keys found. Please set the API_KEYS environment variable.")
     raise Exception("API_KEYS environment variable not set.")
 
+# Process available models
+if AVAILABLE_MODELS:
+    AVAILABLE_MODELS = [model.strip() for model in AVAILABLE_MODELS.split(',') if model.strip()]
+else:
+    AVAILABLE_MODELS = []  # If empty, all models are available
+
 # Simple in-memory rate limiter based solely on IP addresses
 rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 
@@ -49,9 +56,11 @@ class Blackbox:
     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',
@@ -59,9 +68,25 @@ class Blackbox:
         'gemini-pro',
         'gemini-1.5-flash',
         'claude-sonnet-3.5',
+        'PythonAgent',
+        'JavaAgent',
+        'JavaScriptAgent',
+        'HTMLAgent',
+        'GoogleCloudAgent',
+        'AndroidDeveloper',
+        'SwiftDeveloper',
+        'Next.jsAgent',
+        'MongoDBAgent',
+        'PyTorchAgent',
+        'ReactAgent',
+        'XcodeAgent',
+        'AngularJSAgent',
     ]
 
-    agentMode = {}
+    agentMode = {
+        'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
+    }
+
     trendingAgentMode = {
         "blackboxai": {},
         "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
@@ -69,6 +94,19 @@ class Blackbox:
         '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 = {
@@ -77,23 +115,12 @@ class Blackbox:
         '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',
-        'blackboxai-pro': '@BLACKBOXAI-PRO',
-    }
-
-    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 = {
+        "gpt-3.5-turbo": "blackboxai",
+        "gpt-4": "gpt-4o",
         "gemini-flash": "gemini-1.5-flash",
         "claude-3.5-sonnet": "claude-sonnet-3.5",
+        "flux": "ImageGeneration",
     }
 
     @classmethod
@@ -110,31 +137,6 @@ class Blackbox:
         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\$\@\$'
@@ -142,61 +144,39 @@ class Blackbox:
         return cleaned_text
 
     @classmethod
-    async def generate_response(
+    async def create_completion(
         cls,
         model: str,
         messages: List[Dict[str, str]],
-        proxy: Optional[str] = None,
         **kwargs
     ) -> str:
+        """
+        Creates a completion using the Blackbox AI API.
+        """
         model = cls.get_model(model)
-        chat_id = cls.generate_random_string()
-        next_action = cls.generate_next_action()
-        next_router_state_tree = cls.generate_next_router_state_tree()
+        if model is None:
+            logger.error(f"Model {model} is not available.")
+            raise ModelNotWorkingException(model)
 
-        agent_mode = cls.agentMode.get(model, {})
-        trending_agent_mode = cls.trendingAgentMode.get(model, {})
+        chat_id = cls.generate_random_string()
 
-        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 = {
+        headers = {
+            'Content-Type': 'application/json',
             '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'
+            'referer': f"{cls.url}/?model={model}",
+            '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 = {
+        payload = {
             "messages": [
                 {
                     "id": chat_id,
@@ -208,8 +188,8 @@ class Blackbox:
             "previewToken": None,
             "userId": None,
             "codeModelMode": True,
-            "agentMode": agent_mode,
-            "trendingAgentMode": trending_agent_mode,
+            "agentMode": cls.agentMode.get(model, {}),
+            "trendingAgentMode": cls.trendingAgentMode.get(model, {}),
             "isMicMode": False,
             "userSystemPrompt": None,
             "maxTokens": 1024,
@@ -226,16 +206,15 @@ class Blackbox:
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        async with ClientSession(headers=common_headers) as session:
+        async with ClientSession() as session:
             try:
                 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()
+                    headers=headers,
+                    json=payload
+                ) as response:
+                    response.raise_for_status()
+                    text = await response.text()
                     cleaned_response = cls.clean_response(text)
                     return cleaned_response
             except ClientResponseError as e:
@@ -246,9 +225,9 @@ class Blackbox:
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                return error_text
+                raise HTTPException(status_code=e.status, detail=error_text)
             except Exception as e:
-                return f"Unexpected error during /api/chat request: {str(e)}"
+                raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")
 
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
@@ -282,7 +261,7 @@ async def rate_limiter_per_ip(request: Request):
     else:
         if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
             logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
-            raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
+            raise HTTPException(status_code=429, detail='Rate limit exceeded')
         rate_limit_store[client_ip]["count"] += 1
 
 async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
@@ -362,14 +341,16 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
             raise HTTPException(status_code=400, detail="Requested model is not available.")
 
-        # Process the request with actual message content, but don't log itresponse_content = await Blackbox.generate_response(
+        # Process the request with actual message content, but don't log it
+        response_content = await Blackbox.create_completion(
             model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-            temperature=request.temperature,
-            max_tokens=request.max_tokens
+            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
         )
 
         logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
+        return {
+            "id": f"chatcmpl-{uuid.uuid4()}",
+            "object": "chat.logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
             "object": "chat.completion",
@@ -377,12 +358,12 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             "model": request.model,
             "choices": [
                 {
-                    "index": 0,
                     "message": {
                         "role": "assistant",
                         "content": response_content
                     },
-                    "finish_reason": "stop"
+                    "finish_reason": "stop",
+                    "index": 0
                 }
             ],
             "usage": {