Update main.py

main.py

@@ -22,8 +22,7 @@ class ModelNotWorkingException(Exception):
 # Mock implementations for ImageResponse and to_data_uri
 class ImageResponse:
     def __init__(self, url: str, alt: str):
-        self.url = url
-        self.alt = alt
+        self.url = url self.alt = alt
 
 def to_data_uri(image: Any) -> str:
     return "data:image/png;base64,..."  # Replace with actual base64 data
@@ -38,8 +37,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     url = "https://www.blackbox.ai"
     api_endpoint = "https://www.blackbox.ai/api/chat"
     working = True
-    supports_stream = True
-    supports_system_message = True
+    supports_stream = True supports_system_message = True
     supports_message_history = True
     
     default_model = 'blackbox'
@@ -78,27 +76,20 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         "flux": "ImageGenerationLV45LJp",
     }
 
-    @classmethod
-    def get_model(cls, model: str) -> str:
+    @classmethod def get_model(cls, model: str) -> str:
         if model in cls.models:
-            return model
-        elif model in cls.userSelectedModel:
-            return model
-        elif model in cls.model_aliases:
+            return model elif model in cls.userSelectedModel:
+            return model elif model in cls.model_aliases:
             return cls.model_aliases[model]
         else:
-            return cls.default_model
-
-    @classmethod
-    async def create_async_generator(
+            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:
+        **kwargs ) -> Any:
         model = cls.get_model(model)
 
         # Check if the model is working
@@ -166,8 +157,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                 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:
+ async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
                 response.raise_for_status()
                 if model == 'ImageGenerationLV45LJp':
                     response_text = await response.text()
@@ -180,8 +170,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 = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_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
 
@@ -249,9 +238,7 @@ async def chat_completions(request: ChatRequest):
             if isinstance(chunk, ImageResponse):
                 response_content += f"![image]({chunk.url})\n"
             else:
-                response_content += chunk
-
-        return {
+                response_content += chunk return {
             "id": f"chatcmpl-{uuid.uuid4()}",
             "object": "chat.completion",
             "created": int(datetime.now().timestamp()),
@@ -271,4 +258,4 @@ async def chat_completions(request: ChatRequest):
 
 @app.get("/niansuhai/v1/models")
 async def get_models():
-    return {"models": Blackbox.models}
+    return {"models": Blackbox.models}