Update main.py

main.py

@@ -5,13 +5,18 @@ import random
 import string
 import uuid
 import json
-from aiohttp import ClientSession
+import logging
+from aiohttp import ClientSession, ClientError
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import List, Dict, Any, Optional
 from datetime import datetime
 from fastapi.responses import StreamingResponse
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
     def __init__(self, model: str):
@@ -167,28 +172,32 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             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()
-                response_text = await response.text()
-                
-                # Check if the response is empty
-                if not response_text.strip():
-                    raise ModelNotWorkingException(model)
+            try:
+                async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
+                    response.raise_for_status()
+                    response_text = await response.text()
+                    
+                    # Check if the response is empty
+                    if not response_text.strip():
+                        raise ModelNotWorkingException(model)
 
-                if model == 'ImageGenerationLV45LJp':
-                    url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
-                    if url_match:
-                        image_url = url_match.group(0)
-                        yield ImageResponse(image_url, alt=messages[-1]['content'])
+                    if model == 'ImageGenerationLV45LJp':
+                        url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
+                        if url_match:
+                            image_url = url_match.group(0)
+                            yield ImageResponse(image_url, alt=messages[-1]['content'])
+                        else:
+                            raise Exception("Image URL not found in the response")
                     else:
-                        raise Exception("Image URL not found in the response")
-                else:
-                    async for chunk in response.content.iter_any():
-                        if chunk:
-                            decoded_chunk = chunk.decode(errors='ignore')  # Handle decoding errors
-                            decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
-                            if decoded_chunk.strip():
-                                yield decoded_chunk
+                        async for chunk in response.content.iter_any():
+                            if chunk:
+                                decoded_chunk = chunk.decode(errors='ignore')  # Handle decoding errors
+                                decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
+                                if decoded_chunk.strip():
+                                    yield decoded_chunk
+            except ClientError as e:
+                logger.error(f"HTTP request failed: {e}")
+                raise HTTPException(status_code=503, detail="Service is unavailable. Please try again later.")
 
 # FastAPI app setup
 app = FastAPI()