Update main.py

main.py

@@ -3,8 +3,8 @@ from __future__ import annotations
 import re
 import random
 import string
-import json
 import uuid
+import json
 from aiohttp import ClientSession
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
@@ -41,7 +41,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     supports_stream = True
     supports_system_message = True
     supports_message_history = True
-
+    
     default_model = 'blackbox'
     models = [
         'blackbox',
@@ -66,13 +66,13 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
         'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
     }
-
+    
     userSelectedModel = {
         "gpt-4o": "gpt-4o",
         "gemini-pro": "gemini-pro",
         'claude-sonnet-3.5': "claude-sonnet-3.5",
     }
-
+    
     model_aliases = {
         "gemini-flash": "gemini-1.5-flash",
         "flux": "ImageGenerationLV45LJp",
@@ -124,16 +124,16 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
 
         if model in cls.userSelectedModel:
             prefix = f"@{cls.userSelectedModel[model]}"
-            if messages and not messages[0]['content'].startswith(prefix):
+            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 and messages:
+            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 = {
@@ -147,7 +147,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                 "userSelectedModel": None,
                 "userSystemPrompt": None,
                 "isMicMode": False,
-                "maxTokens": 4096,
+                "maxTokens": 8192,
                 "playgroundTopP": 0.9,
                 "playgroundTemperature": 0.5,
                 "isChromeExt": False,
@@ -166,7 +166,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:
                 response.raise_for_status()
                 if model == 'ImageGenerationLV45LJp':
@@ -180,7 +180,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
@@ -195,7 +195,7 @@ class Message(BaseModel):
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
-    stream: Optional[bool] = False
+    stream: Optional[bool] = False  # Add this for streaming
 
 def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
@@ -215,6 +215,7 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
 
 @app.post("/niansuhai/v1/chat/completions")
 async def chat_completions(request: ChatRequest):
+    # Validate the model
     valid_models = Blackbox.models + list(Blackbox.userSelectedModel.keys()) + list(Blackbox.model_aliases.keys())
     if request.model not in valid_models:
         raise HTTPException(status_code=400, detail=f"Invalid model name: {request.model}. Valid models are: {valid_models}")
@@ -222,7 +223,7 @@ async def chat_completions(request: ChatRequest):
     messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
 
     try:
-        async_generator = await Blackbox.create_async_generator(
+        async_generator = Blackbox.create_async_generator(
             model=request.model,
             messages=messages,
             image=None,
@@ -270,4 +271,4 @@ async def chat_completions(request: ChatRequest):
 
 @app.get("/niansuhai/v1/models")
 async def get_models():
-    return {"models": Blackbox.models}
+    return {"models": Blackbox.models}