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,..."  # Replace with actual base64 data
+    return "data:image/png;base64,..."
 
 class AsyncGeneratorProvider:
     pass
@@ -98,65 +98,36 @@ 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,
-            "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"
+            "user-agent": "Mozilla/5.0 ..."
         }
 
         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 image is not None:
+            if model == 'ImageGenerationLV45LJp' and 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,
-                "previewToken": None,
-                "userId": None,
-                "codeModelMode": True,
-                "agentMode": {},
-                "trendingAgentMode": {},
-                "userSelectedModel": None,
-                "userSystemPrompt": None,
-                "isMicMode": False,
+                "agentMode": cls.agentMode.get(model, {}),
+                "trendingAgentMode": cls.trendingAgentMode.get(model, {}),
+                "userSelectedModel": cls.userSelectedModel.get(model, None),
                 "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':
@@ -170,7 +141,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                 else:
                     async for chunk in response.content.iter_any():
                         if chunk:
-                            decoded_chunk = chunk.decode(errors='ignore')  # Handle decoding errors
+                            decoded_chunk = chunk.decode(errors='ignore')
                             decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
                             if decoded_chunk.strip():
                                 yield decoded_chunk
@@ -203,26 +174,31 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,
     }
 
-@app.post("/v1/chat/completions")
+@app.post("/niansuhai/v1/chat/completions")
 async def chat_completions(request: ChatRequest):
     messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
 
     async_generator = Blackbox.create_async_generator(
         model=request.model,
-        messages=messages
+        messages=messages,
+        image=None,  # Pass the image if required
+        image_name=None  # Pass image name if required
     )
 
     if request.stream:
         async def generate():
             async for chunk in async_generator:
-                yield f"data: {json.dumps(create_response(chunk, request.model))}\n\n"
+                yield f"data: {json.dumps(create_response(chunk.url 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:
-            response_content += chunk if isinstance(chunk, str) else chunk.content  # Concatenate response
+            if isinstance(chunk, ImageResponse):
+                response_content += f"Image URL: {chunk.url}\n"
+            else:
+                response_content += chunk  # Concatenate text responses
 
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",