Update main.py

main.py

@@ -47,7 +47,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         'blackbox',
         'gemini-1.5-flash',
         "llama-3.1-8b",
-        'llama-3.1-70b',
+        'llama-3.1-70b',  # Example of a non-working model
         'llama-3.1-405b',
         'ImageGenerationLV45LJp',
         'gpt-4o',
@@ -55,6 +55,19 @@ 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"},
     }
@@ -102,7 +115,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         model = cls.get_model(model)
 
         # Check if the model is working
-        if not cls.working or model not in cls.models:
+        if not cls.model_status.get(model, False):
             raise ModelNotWorkingException(model)
 
         headers = {
@@ -226,8 +239,8 @@ async def chat_completions(request: ChatRequest):
         async_generator = Blackbox.create_async_generator(
             model=request.model,
             messages=messages,
-            image=None,
-            image_name=None
+            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))
@@ -236,6 +249,7 @@ async def chat_completions(request: ChatRequest):
         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:
@@ -249,7 +263,7 @@ async def chat_completions(request: ChatRequest):
             if isinstance(chunk, ImageResponse):
                 response_content += f"![image]({chunk.url})\n"
             else:
-                response_content += chunk
+                response_content += chunk  # Concatenate text responses
 
         return {
             "id": f"chatcmpl-{uuid.uuid4()}",