Update main.py

main.py

@@ -3,15 +3,22 @@ from __future__ import annotations
 import re
 import random
 import string
-import uuid
 import json
+import uuid
 from aiohttp import ClientSession
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from typing import List, Dict, Any, Optional, AsyncGenerator
+from typing import List, Dict, Any, Optional
 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):
@@ -19,7 +26,6 @@ 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:
@@ -35,7 +41,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     supports_stream = True
     supports_system_message = True
     supports_message_history = True
-    
+
     default_model = 'blackbox'
     models = [
         'blackbox',
@@ -60,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",
@@ -92,12 +98,13 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         image: Optional[Any] = None,
         image_name: Optional[str] = None,
         **kwargs
-    ) -> AsyncGenerator:
-        if not messages:
-            raise ValueError("Messages cannot be empty")
-
+    ) -> 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)
+
         headers = {
             "accept": "*/*",
             "accept-language": "en-US,en;q=0.9",
@@ -117,11 +124,11 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
 
         if model in cls.userSelectedModel:
             prefix = f"@{cls.userSelectedModel[model]}"
-            if not messages[0]['content'].startswith(prefix):
+            if messages and 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 image is not None and messages:
                 messages[-1]["data"] = {
                     "fileText": image_name,
                     "imageBase64": to_data_uri(image)
@@ -173,7 +180,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
@@ -188,7 +195,7 @@ class Message(BaseModel):
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
-    stream: Optional[bool] = False  # Add this for streaming
+    stream: Optional[bool] = False
 
 def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
@@ -208,20 +215,26 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
 
 @app.post("/niansuhai/v1/chat/completions")
 async def chat_completions(request: ChatRequest):
+    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]
 
-    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
-    )
+    try:
+        async_generator = await Blackbox.create_async_generator(
+            model=request.model,
+            messages=messages,
+            image=None,
+            image_name=None
+        )
+    except ModelNotWorkingException as e:
+        raise HTTPException(status_code=503, detail=str(e))
 
     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:
@@ -233,10 +246,9 @@ 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  # Concatenate text responses
+                response_content += chunk
 
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",
@@ -255,3 +267,7 @@ async def chat_completions(request: ChatRequest):
             ],
             "usage": None,
         }
+
+@app.get("/niansuhai/v1/models")
+async def get_models():
+    return {"models": Blackbox.models}