Update main.py

main.py

@@ -49,13 +49,126 @@ def decode_base64_image(data_uri: str) -> bytes:
         raise e
 
 class Blackbox:
-    # ... [existing Blackbox class definition]
+    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',
+    ]
+
+    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) -> str:
+        if model in cls.models:
+            return model
+        elif model in cls.userSelectedModel:
+            return cls.userSelectedModel[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: List[Dict[str, str]],
+        messages: List[Dict[str, Any]],
         proxy: Optional[str] = None,
         image: Optional[str] = None,  # Expecting a base64 string
         image_name: Optional[str] = None,
@@ -70,7 +183,21 @@ class Blackbox:
             raise ModelNotWorkingException(model)
         
         headers = {
-            # ... [existing 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:
@@ -80,24 +207,32 @@ class Blackbox:
                 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'
-        
+        user_message = {
+            "id": random_id,
+            "role": 'user',
+            "content": 'Hi'  # This should be dynamically set based on input
+        }
         if image is not None:
             try:
                 image_bytes = decode_base64_image(image)
                 data_uri = to_data_uri(image_bytes)
-                messages[-1]['data'] = {
+                user_message['data'] = {
                     'fileText': '',
                     'imageBase64': data_uri,
-                    'title': image_name
+                    'title': image_name or "Uploaded Image"
                 }
-                messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content']
+                user_message['content'] = 'FILE:BB\n$#$\n\n$#$\n' + user_message['content']
                 logger.debug("Image data added to the message.")
             except Exception as e:
                 logger.error(f"Failed to decode base64 image: {e}")
                 raise HTTPException(status_code=400, detail="Invalid image data provided.")
         
+        # Update the last message with user_message
+        if messages:
+            messages[-1] = user_message
+        else:
+            messages.append(user_message)
+
         data = {
             "messages": messages,
             "id": random_id,
@@ -234,8 +369,9 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
 
 @app.post("/niansuhai/v1/chat/completions")
 async def chat_completions(request: ChatRequest, req: Request):
-    logger.info(f"Received chat completions request: {request}")
+    logger.info(f"Received chat completions request: model='{request.model}' messages={request.messages} stream={request.stream} webSearchMode={request.webSearchMode} image={request.image}")
     try:
+        # Convert Pydantic messages to dicts
         messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
         
         async_generator = Blackbox.create_async_generator(
@@ -252,6 +388,7 @@ async def chat_completions(request: ChatRequest, req: Request):
                 try:
                     async for chunk in async_generator:
                         if isinstance(chunk, ImageResponse):
+                            # Use the base64 data URI directly
                             image_markdown = f"![{chunk.alt}]({chunk.data_uri})"
                             response_chunk = create_response(image_markdown, request.model)
                         else: