Update main.py

main.py

@@ -8,17 +8,10 @@ import json
 from aiohttp import ClientSession
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any, Optional, AsyncGenerator
 from datetime import datetime
 from fastapi.responses import StreamingResponse
 
-# Custom exception for model not working
-class ModelNotWorkingException(Exception):
-    def __init__(self, model: str):
-        self.model = model
-        self.message = f"The model '{model}' is currently not working. Please wait for NiansuhAI to fix this. Thank you for your patience."
-        super().__init__(self.message)
-
 # Mock implementations for ImageResponse and to_data_uri
 class ImageResponse:
     def __init__(self, url: str, alt: str):
@@ -26,6 +19,7 @@ class ImageResponse:
         self.alt = alt
 
 def to_data_uri(image: Any) -> str:
+    # Placeholder for actual image encoding
     return "data:image/png;base64,..."  # Replace with actual base64 data
 
 class AsyncGeneratorProvider:
@@ -47,7 +41,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         'blackbox',
         'gemini-1.5-flash',
         "llama-3.1-8b",
-        'llama-3.1-70b',  # Example of a non-working model
+        'llama-3.1-70b',
         'llama-3.1-405b',
         'ImageGenerationLV45LJp',
         'gpt-4o',
@@ -55,19 +49,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         'claude-sonnet-3.5',
     ]
 
-    # Define the working status of models
-    model_status = {
-        'blackbox': True,
-        'gemini-1.5-flash': True,
-        'llama-3.1-8b': True,
-        'llama-3.1-70b': False,  # Non-working model
-        'llama-3.1-405b': True,
-        'ImageGenerationLV45LJp': True,
-        'gpt-4o': True,
-        'gemini-pro': True,
-        'claude-sonnet-3.5': True,
-    }
-
     agentMode = {
         'ImageGenerationLV45LJp': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
     }
@@ -111,12 +92,11 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         image: Optional[Any] = None,
         image_name: Optional[str] = None,
         **kwargs
-    ) -> Any:
-        model = cls.get_model(model)
+    ) -> AsyncGenerator:
+        if not messages:
+            raise ValueError("Messages cannot be empty")
 
-        # Check if the model is working
-        if not cls.model_status.get(model, False):
-            raise ModelNotWorkingException(model)
+        model = cls.get_model(model)
 
         headers = {
             "accept": "*/*",
@@ -139,14 +119,14 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             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))
 
             data = {
@@ -173,6 +153,13 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                 "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':
@@ -184,19 +171,12 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                     else:
                         raise Exception("Image URL not found in the response")
                 else:
-                    response_content = ""
                     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():
-                                response_content += decoded_chunk
-
-                    # Check if the response content is empty
-                    if not response_content.strip():
-                        raise ModelNotWorkingException(model)
-
-                    yield response_content
+                                yield decoded_chunk
 
 # FastAPI app setup
 app = FastAPI()
@@ -228,32 +208,24 @@ 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}")
-
     messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
 
-    try:
-        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
-        )
-    except ModelNotWorkingException as e:
-        raise HTTPException(status_code=503, detail=str(e))
+    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
+    )
 
     if request.stream:
         async def generate():
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponse):
+                    # Format the response as a Markdown image
                     image_markdown = f"![image]({chunk.url})"
                     yield f"data: {json.dumps(create_response(image_markdown, request.model))}\n\n"
                 else:
-                    response_with_footer = f"{chunk}\nNiansuhAI"  # Concatenating here
-                    yield f"data: {json.dumps(create_response(response_with_footer, request.model))}\n\n"
+                    yield f"data: {json.dumps(create_response(chunk, request.model))}\n\n"
             yield "data: [DONE]\n\n"
 
         return StreamingResponse(generate(), media_type="text/event-stream")
@@ -261,12 +233,10 @@ async def chat_completions(request: ChatRequest):
         response_content = ""
         async for chunk in async_generator:
             if isinstance(chunk, ImageResponse):
+                # Add Markdown image to the response
                 response_content += f"![image]({chunk.url})\n"
             else:
-                response_content += chunk
-        
-        # Append "\nNiansuhAI" to the final response content
-        response_content += "\nNiansuhAI"
+                response_content += chunk  # Concatenate text responses
 
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
@@ -285,10 +255,3 @@ async def chat_completions(request: ChatRequest):
             ],
             "usage": None,
         }
-
-
-
-@app.get("/niansuhai/v1/models")
-async def get_models():
-    return {"models": Blackbox.models}
-