Update main.py

main.py

@@ -12,6 +12,13 @@ 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:
@@ -94,7 +100,11 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         **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)
+
         headers = {
             "accept": "*/*",
             "accept-language": "en-US,en;q=0.9",
@@ -212,18 +222,20 @@ async def chat_completions(request: ChatRequest):
 
     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 = 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:
@@ -235,10 +247,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()}",
@@ -258,7 +269,6 @@ async def chat_completions(request: ChatRequest):
             "usage": None,
         }
 
-
 @app.get("/niansuhai/v1/models")
 async def get_models():
-    return {"models": Blackbox.models}
+    return {"models": Blackbox.models}