Update main.py

main.py

@@ -20,7 +20,7 @@ class ImageResponse:
 
 def to_data_uri(image: Any) -> str:
     # Placeholder for actual image encoding
-    return "data:image/png;base64,..."
+    return "data:image/png;base64,..."  # Replace with actual base64 data
 
 class AsyncGeneratorProvider:
     pass
@@ -98,36 +98,65 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         headers = {
             "accept": "*/*",
             "accept-language": "en-US,en;q=0.9",
+            "cache-control": "no-cache",
             "content-type": "application/json",
             "origin": cls.url,
-            "user-agent": "Mozilla/5.0 ..."
+            "pragma": "no-cache",
+            "referer": f"{cls.url}/",
+            "sec-ch-ua": '"Not;A=Brand";v="24", "Chromium";v="128"',
+            "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/128.0.0.0 Safari/537.36"
         }
 
         if model in cls.userSelectedModel:
             prefix = f"@{cls.userSelectedModel[model]}"
             if not messages[0]['content'].startswith(prefix):
                 messages[0]['content'] = f"{prefix} {messages[0]['content']}"
-
+        
         async with ClientSession(headers=headers) as session:
-            if model == 'ImageGenerationLV45LJp' and image is not None:
+            if image is not None:
                 messages[-1]["data"] = {
                     "fileText": image_name,
                     "imageBase64": to_data_uri(image)
                 }
-
+            
             random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
 
             data = {
                 "messages": messages,
                 "id": random_id,
-                "agentMode": cls.agentMode.get(model, {}),
-                "trendingAgentMode": cls.trendingAgentMode.get(model, {}),
-                "userSelectedModel": cls.userSelectedModel.get(model, None),
+                "previewToken": None,
+                "userId": None,
+                "codeModelMode": True,
+                "agentMode": {},
+                "trendingAgentMode": {},
+                "userSelectedModel": None,
+                "userSystemPrompt": None,
+                "isMicMode": False,
                 "maxTokens": 1024,
                 "playgroundTopP": 0.9,
                 "playgroundTemperature": 0.5,
+                "isChromeExt": False,
+                "githubToken": None,
+                "clickedAnswer2": False,
+                "clickedAnswer3": False,
+                "clickedForceWebSearch": False,
+                "visitFromDelta": False,
+                "mobileClient": False,
+                "webSearchMode": False,
             }
 
+            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]
+            
             async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
                 response.raise_for_status()
                 if model == 'ImageGenerationLV45LJp':
@@ -141,7 +170,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                 else:
                     async for chunk in response.content.iter_any():
                         if chunk:
-                            decoded_chunk = chunk.decode(errors='ignore')
+                            decoded_chunk = chunk.decode(errors='ignore')  # Handle decoding errors
                             decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
                             if decoded_chunk.strip():
                                 yield decoded_chunk
@@ -180,25 +209,20 @@ async def chat_completions(request: ChatRequest):
 
     async_generator = Blackbox.create_async_generator(
         model=request.model,
-        messages=messages,
-        image=None,  # Pass the image if required
-        image_name=None  # Pass image name if required
+        messages=messages
     )
 
     if request.stream:
         async def generate():
             async for chunk in async_generator:
-                yield f"data: {json.dumps(create_response(chunk.url if isinstance(chunk, ImageResponse) else chunk, request.model))}\n\n"
+                yield f"data: {json.dumps(create_response(chunk.content if isinstance(chunk, ImageResponse) else chunk, request.model))}\n\n"
             yield "data: [DONE]\n\n"
 
         return StreamingResponse(generate(), media_type="text/event-stream")
     else:
         response_content = ""
         async for chunk in async_generator:
-            if isinstance(chunk, ImageResponse):
-                response_content += f"Image URL: {chunk.url}\n"
-            else:
-                response_content += chunk  # Concatenate text responses
+            response_content += chunk.content if isinstance(chunk, ImageResponse) else chunk  # Concatenate response
 
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",