Update main.py

main.py

@@ -90,105 +90,104 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             return cls.default_model
 
     @classmethod
-async def create_async_generator(
-    cls,
-    model: str,
-    messages: List[Dict[str, str]],
-    proxy: Optional[str] = None,
-    image: Optional[Any] = None,
-    image_name: Optional[str] = None,
-    **kwargs
-) -> Any:
-    model = cls.get_model(model)
+    async def create_async_generator(
+        cls,
+        model: str,
+        messages: List[Dict[str, str]],
+        proxy: Optional[str] = None,
+        image: Optional[Any] = None,
+        image_name: Optional[str] = None,
+        **kwargs
+    ) -> Any:
+        model = cls.get_model(model)
 
-    # Check if the model is working
-    if not cls.working or model not in cls.models:
-        raise ModelNotWorkingException(model)
+        # Check if the model is working
+        if not cls.working or model not in cls.models:
+            raise ModelNotWorkingException(model)
 
-    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"
-    }
+        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"
+        }
 
-    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']}"
+        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:
-            messages[-1]["data"] = {
-                "fileText": image_name,
-                "imageBase64": to_data_uri(image)
-            }
-        
-        random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
+        async with ClientSession(headers=headers) as session:
+            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,
-            "previewToken": None,
-            "userId": None,
-            "codeModelMode": True,
-            "agentMode": {},
-            "trendingAgentMode": {},
-            "userSelectedModel": None,
-            "userSystemPrompt": None,
-            "isMicMode": False,
-            "maxTokens": 99999999,
-            "playgroundTopP": 0.9,
-            "playgroundTemperature": 0.5,
-            "isChromeExt": False,
-            "githubToken": None,
-            "clickedAnswer2": False,
-            "clickedAnswer3": False,
-            "clickedForceWebSearch": False,
-            "visitFromDelta": False,
-            "mobileClient": False,
-            "webSearchMode": False,
-        }
+            data = {
+                "messages": messages,
+                "id": random_id,
+                "previewToken": None,
+                "userId": None,
+                "codeModelMode": True,
+                "agentMode": {},
+                "trendingAgentMode": {},
+                "userSelectedModel": None,
+                "userSystemPrompt": None,
+                "isMicMode": False,
+                "maxTokens": 99999999,
+                "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()
-            response_text = await response.text()
+            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]
             
-            if not response_text:  # Check for empty response
-                raise ModelNotWorkingException(model)
+            async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
+                response.raise_for_status()
+                response_text = await response.text()
+                
+                if not response_text:  # Check for empty response
+                    raise ModelNotWorkingException(model)
 
-            if model == 'ImageGenerationLV45LJp':
-                url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
-                if url_match:
-                    image_url = url_match.group(0)
-                    yield ImageResponse(image_url, alt=messages[-1]['content'])
+                if model == 'ImageGenerationLV45LJp':
+                    url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
+                    if url_match:
+                        image_url = url_match.group(0)
+                        yield ImageResponse(image_url, alt=messages[-1]['content'])
+                    else:
+                        raise Exception("Image URL not found in the response")
                 else:
-                    raise Exception("Image URL not found in the response")
-            else:
-                async for chunk in response.content.iter_any():
-                    if chunk:
-                        decoded_chunk = chunk.decode(errors='ignore')  # Handle decoding errors
-                        decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
-                        if decoded_chunk.strip():
-                            yield decoded_chunk
-
+                    async for chunk in response.content.iter_any():
+                        if chunk:
+                            decoded_chunk = chunk.decode(errors='ignore')  # Handle decoding errors
+                            decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
+                            if decoded_chunk.strip():
+                                yield decoded_chunk
 
 # FastAPI app setup
 app = FastAPI()
@@ -276,4 +275,4 @@ async def chat_completions(request: ChatRequest):
 
 @app.get("/niansuhai/v1/models")
 async def get_models():
-    return {"models": Blackbox.models}
+    return {"models": Blackbox.models}