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, ClientTimeout, ClientError, ClientResponseError
+from aiohttp import ClientSession, ClientResponseError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
@@ -27,18 +27,11 @@ 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()})
 
@@ -56,11 +49,9 @@ 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',
@@ -68,25 +59,9 @@ 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 = {
-        'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
-    }
-
+    agentMode = {}
     trendingAgentMode = {
         "blackboxai": {},
         "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
@@ -94,19 +69,6 @@ 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 = {
@@ -115,12 +77,23 @@ 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
@@ -137,6 +110,31 @@ 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\$\@\$'
@@ -144,39 +142,61 @@ class Blackbox:
         return cleaned_text
 
     @classmethod
-    async def create_completion(
+    async def generate_response(
         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)
-        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()
 
-        headers = {
-            'Content-Type': 'application/json',
+        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,
-            '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'
+            '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'
         }
 
-        payload = {
+        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,
@@ -188,8 +208,8 @@ class Blackbox:
             "previewToken": None,
             "userId": None,
             "codeModelMode": True,
-            "agentMode": cls.agentMode.get(model, {}),
-            "trendingAgentMode": cls.trendingAgentMode.get(model, {}),
+            "agentMode": agent_mode,
+            "trendingAgentMode": trending_agent_mode,
             "isMicMode": False,
             "userSystemPrompt": None,
             "maxTokens": 1024,
@@ -206,15 +226,16 @@ class Blackbox:
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        async with ClientSession() as session:
+        async with ClientSession(headers=common_headers) as session:
             try:
                 async with session.post(
                     cls.api_endpoint,
-                    headers=headers,
-                    json=payload
-                ) as response:
-                    response.raise_for_status()
-                    text = await response.text()
+                    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)
                     return cleaned_response
             except ClientResponseError as e:
@@ -225,9 +246,9 @@ class Blackbox:
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                raise HTTPException(status_code=e.status, detail=error_text)
+                return error_text
             except Exception as e:
-                raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")
+                return f"Unexpected error during /api/chat request: {str(e)}"
 
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
@@ -261,7 +282,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')
+            raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
         rate_limit_store[client_ip]["count"] += 1
 
 async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
@@ -328,7 +349,7 @@ class ChatRequest(BaseModel):
     user: Optional[str] = None
 
 @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)):  # Make sure this function is async
+async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
     client_ip = req.client.host
     # Redact user messages only for logging purposes
     redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
@@ -341,10 +362,12 @@ 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 it
-        response_content = await Blackbox.create_completion(  # Correct usage of await inside async function
+        # Process the request with actual message content
+        response_content = await Blackbox.generate_response(
             model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
+            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
+            temperature=request.temperature,
+            max_tokens=request.max_tokens
         )
 
         logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
@@ -355,19 +378,19 @@ 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",
-                    "index": 0
+                    "finish_reason": "stop"
                 }
             ],
             "usage": {
                 "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
                 "completion_tokens": len(response_content.split()),
                 "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
-            }
+            },
         }
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
@@ -379,7 +402,6 @@ 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: GET /v1/models
 @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
 async def get_models(req: Request):
@@ -413,4 +435,4 @@ async def http_exception_handler(request: Request, exc: HTTPException):
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)