Update main.py

main.py

@@ -1,4 +1,4 @@
-# app/main.py
+# main.py
 
 import os
 import re
@@ -18,9 +18,8 @@ from aiohttp import ClientSession, ClientTimeout, ClientError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
 from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
 from pydantic import BaseModel
-
-from .blackbox import Blackbox, ImageResponse
-from .image import to_data_uri, ImageType
+from io import BytesIO
+import base64
 
 # Configure logging
 logging.basicConfig(
@@ -31,6 +30,10 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 # Load environment variables
+from dotenv import load_dotenv
+
+load_dotenv()
+
 API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
 RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60'))  # Requests per minute
 AVAILABLE_MODELS = os.getenv('AVAILABLE_MODELS', '')  # Comma-separated available models
@@ -101,6 +104,322 @@ class ModelNotWorkingException(Exception):
         self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed."
         super().__init__(self.message)
 
+# Image Handling Functions
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp', 'svg'}
+
+def is_allowed_extension(filename: str) -> bool:
+    """
+    Checks if the given filename has an allowed extension.
+    """
+    return '.' in filename and \
+           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
+def is_data_uri_an_image(data_uri: str) -> bool:
+    """
+    Checks if the given data URI represents an image.
+    """
+    match = re.match(r'data:image/(\w+);base64,', data_uri)
+    if not match:
+        raise ValueError("Invalid data URI image.")
+    image_format = match.group(1).lower()
+    if image_format not in ALLOWED_EXTENSIONS and image_format != "svg+xml":
+        raise ValueError("Invalid image format (from MIME type).")
+    return True
+
+def extract_data_uri(data_uri: str) -> bytes:
+    """
+    Extracts the binary data from the given data URI.
+    """
+    return base64.b64decode(data_uri.split(",")[1])
+
+def to_data_uri(image: str) -> str:
+    """
+    Validates and returns the data URI for an image.
+    """
+    is_data_uri_an_image(image)
+    return image
+
+class ImageResponseCustom:
+    def __init__(self, url: str, alt: str):
+        self.url = url
+        self.alt = alt
+
+# Blackbox AI Integration
+class Blackbox:
+    url = "https://www.blackbox.ai"
+    api_endpoint = "https://www.blackbox.ai/api/chat"
+    working = True
+    supports_stream = True
+    supports_system_message = True
+    supports_message_history = True
+
+    default_model = 'blackboxai'
+    image_models = ['ImageGeneration']
+    models = [
+        default_model,
+        'blackboxai-pro',
+        *image_models,
+        "llama-3.1-8b",
+        'llama-3.1-70b',
+        'llama-3.1-405b',
+        'gpt-4o',
+        'gemini-pro',
+        'gemini-1.5-flash',
+        'claude-sonnet-3.5',
+        'PythonAgent',
+        'JavaAgent',
+        'JavaScriptAgent',
+        'HTMLAgent',
+        'GoogleCloudAgent',
+        'AndroidDeveloper',
+        'SwiftDeveloper',
+        'Next.jsAgent',
+        'MongoDBAgent',
+        'PyTorchAgent',
+        'ReactAgent',
+        'XcodeAgent',
+        'AngularJSAgent',
+    ]
+
+    agentMode = {
+        'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
+        'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"},
+    }
+
+    trendingAgentMode = {
+        "blackboxai": {},
+        "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
+        "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
+        'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
+        'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
+        'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
+        'PythonAgent': {'mode': True, 'id': "Python Agent"},
+        'JavaAgent': {'mode': True, 'id': "Java Agent"},
+        'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"},
+        'HTMLAgent': {'mode': True, 'id': "HTML Agent"},
+        'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"},
+        'AndroidDeveloper': {'mode': True, 'id': "Android Developer"},
+        'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"},
+        'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"},
+        'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"},
+        'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
+        'ReactAgent': {'mode': True, 'id': "React Agent"},
+        'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
+        'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
+    }
+
+    userSelectedModel = {
+        "gpt-4o": "gpt-4o",
+        "gemini-pro": "gemini-pro",
+        'claude-sonnet-3.5': "claude-sonnet-3.5",
+    }
+
+    model_prefixes = {
+        'gpt-4o': '@GPT-4o',
+        'gemini-pro': '@Gemini-PRO',
+        'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
+        'PythonAgent': '@Python Agent',
+        'JavaAgent': '@Java Agent',
+        'JavaScriptAgent': '@JavaScript Agent',
+        'HTMLAgent': '@HTML Agent',
+        'GoogleCloudAgent': '@Google Cloud Agent',
+        'AndroidDeveloper': '@Android Developer',
+        'SwiftDeveloper': '@Swift Developer',
+        'Next.jsAgent': '@Next.js Agent',
+        'MongoDBAgent': '@MongoDB Agent',
+        'PyTorchAgent': '@PyTorch Agent',
+        'ReactAgent': '@React Agent',
+        'XcodeAgent': '@Xcode Agent',
+        'AngularJSAgent': '@AngularJS Agent',
+        'blackboxai-pro': '@BLACKBOXAI-PRO',
+        'ImageGeneration': '@Image Generation',
+        'Niansuh': '@Niansuh',
+    }
+
+    model_referers = {
+        "blackboxai": f"{url}/?model=blackboxai",
+        "gpt-4o": f"{url}/?model=gpt-4o",
+        "gemini-pro": f"{url}/?model=gemini-pro",
+        "claude-sonnet-3.5": f"{url}/?model=claude-sonnet-3.5"
+    }
+
+    model_aliases = {
+        "gemini-flash": "gemini-1.5-flash",
+        "claude-3.5-sonnet": "claude-sonnet-3.5",
+        "flux": "ImageGeneration",
+        "niansuh": "Niansuh",
+    }
+
+    @classmethod
+    def get_model(cls, model: str) -> Optional[str]:
+        if model in cls.models:
+            return model
+        elif model in cls.userSelectedModel and cls.userSelectedModel[model] in cls.models:
+            return cls.userSelectedModel[model]
+        elif model in cls.model_aliases and cls.model_aliases[model] in cls.models:
+            return cls.model_aliases[model]
+        else:
+            return cls.default_model if cls.default_model in cls.models else None
+
+    @classmethod
+    async def create_async_generator(
+        cls,
+        model: str,
+        messages: List[Dict[str, str]],
+        proxy: Optional[str] = None,
+        image: Optional[str] = None,
+        image_name: Optional[str] = None,
+        webSearchMode: bool = False,
+        **kwargs
+    ) -> AsyncGenerator[Union[str, ImageResponseCustom], None]:
+        model = cls.get_model(model)
+        if model is None:
+            logger.error(f"Model {model} is not available.")
+            raise ModelNotWorkingException(model)
+
+        logger.info(f"Selected model: {model}")
+
+        if not cls.working or model not in cls.models:
+            logger.error(f"Model {model} is not working or not supported.")
+            raise ModelNotWorkingException(model)
+        
+        headers = {
+            "accept": "*/*",
+            "accept-language": "en-US,en;q=0.9",
+            "cache-control": "no-cache",
+            "content-type": "application/json",
+            "origin": cls.url,
+            "pragma": "no-cache",
+            "priority": "u=1, i",
+            "referer": cls.model_referers.get(model, cls.url),
+            "sec-ch-ua": '"Chromium";v="129", "Not=A?Brand";v="8"',
+            "sec-ch-ua-mobile": "?0",
+            "sec-ch-ua-platform": '"Linux"',
+            "sec-fetch-dest": "empty",
+            "sec-fetch-mode": "cors",
+            "sec-fetch-site": "same-origin",
+            "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
+        }
+
+        if model in cls.model_prefixes:
+            prefix = cls.model_prefixes[model]
+            if not messages[0]['content'].startswith(prefix):
+                logger.debug(f"Adding prefix '{prefix}' to the first message.")
+                messages[0]['content'] = f"{prefix} {messages[0]['content']}"
+        
+        random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
+        messages[-1]['id'] = random_id
+        messages[-1]['role'] = 'user'
+
+        # Don't log the full message content for privacy
+        logger.debug(f"Generated message ID: {random_id} for model: {model}")
+
+        if image is not None:
+            messages[-1]['data'] = {
+                'fileText': '',
+                'imageBase64': image,
+                'title': image_name
+            }
+            messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content']
+            logger.debug("Image data added to the message.")
+        
+        data = {
+            "messages": messages,
+            "id": random_id,
+            "previewToken": None,
+            "userId": None,
+            "codeModelMode": True,
+            "agentMode": {},
+            "trendingAgentMode": {},
+            "isMicMode": False,
+            "userSystemPrompt": None,
+            "maxTokens": 99999999,
+            "playgroundTopP": 0.9,
+            "playgroundTemperature": 0.5,
+            "isChromeExt": False,
+            "githubToken": None,
+            "clickedAnswer2": False,
+            "clickedAnswer3": False,
+            "clickedForceWebSearch": False,
+            "visitFromDelta": False,
+            "mobileClient": False,
+            "userSelectedModel": None,
+            "webSearchMode": webSearchMode,
+        }
+
+        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]
+        logger.info(f"Sending request to {cls.api_endpoint} with data (excluding messages).")
+
+        timeout = ClientTimeout(total=60)  # Set an appropriate timeout
+        retry_attempts = 10  # Set the number of retry attempts
+
+        for attempt in range(retry_attempts):
+            try:
+                async with ClientSession(headers=headers, timeout=timeout) as session:
+                    async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
+                        response.raise_for_status()
+                        logger.info(f"Received response with status {response.status}")
+                        if model in cls.image_models:
+                            response_text = await response.text()
+                            # Extract image URL from the response
+                            url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
+                            if url_match:
+                                image_url = url_match.group(0)
+                                logger.info(f"Image URL found: {image_url}")
+                                yield ImageResponseCustom(url=image_url, alt=messages[-1]['content'])
+                            else:
+                                logger.error("Image URL not found in the response.")
+                                raise Exception("Image URL not found in the response")
+                        else:
+                            full_response = ""
+                            search_results_json = ""
+                            try:
+                                async for chunk, _ in response.content.iter_chunks():
+                                    if chunk:
+                                        decoded_chunk = chunk.decode(errors='ignore')
+                                        decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
+                                        if decoded_chunk.strip():
+                                            if '$~~~$' in decoded_chunk:
+                                                search_results_json += decoded_chunk
+                                            else:
+                                                full_response += decoded_chunk
+                                                yield decoded_chunk
+                                logger.info("Finished streaming response chunks.")
+                            except Exception as e:
+                                logger.exception("Error while iterating over response chunks.")
+                                raise e
+                            if data["webSearchMode"] and search_results_json:
+                                match = re.search(r'\$~~~\$(.*?)\$~~~\$', search_results_json, re.DOTALL)
+                                if match:
+                                    try:
+                                        search_results = json.loads(match.group(1))
+                                        formatted_results = "\n\n**Sources:**\n"
+                                        for i, result in enumerate(search_results[:5], 1):
+                                            formatted_results += f"{i}. [{result['title']}]({result['link']})\n"
+                                        logger.info("Formatted search results.")
+                                        yield formatted_results
+                                    except json.JSONDecodeError as je:
+                                        logger.error("Failed to parse search results JSON.")
+                                        raise je
+                break  # Exit the retry loop if successful
+            except ClientError as ce:
+                logger.error(f"Client error occurred: {ce}. Retrying attempt {attempt + 1}/{retry_attempts}")
+                if attempt == retry_attempts - 1:
+                    raise HTTPException(status_code=502, detail="Error communicating with the external API.")
+            except asyncio.TimeoutError:
+                logger.error(f"Request timed out. Retrying attempt {attempt + 1}/{retry_attempts}")
+                if attempt == retry_attempts - 1:
+                    raise HTTPException(status_code=504, detail="External API request timed out.")
+            except Exception as e:
+                logger.error(f"Unexpected error: {e}. Retrying attempt {attempt + 1}/{retry_attempts}")
+                if attempt == retry_attempts - 1:
+                    raise HTTPException(status_code=500, detail=str(e))
+
 # FastAPI app setup
 app = FastAPI()
 
@@ -227,7 +546,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                 try:
                     assistant_content = ""
                     async for chunk in async_generator:
-                        if isinstance(chunk, ImageResponse):
+                        if isinstance(chunk, ImageResponseCustom):
                             # Handle image responses if necessary
                             image_markdown = f"![image]({chunk.url})\n"
                             assistant_content += image_markdown
@@ -293,7 +612,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         else:
             response_content = ""
             async for chunk in async_generator:
-                if isinstance(chunk, ImageResponse):
+                if isinstance(chunk, ImageResponseCustom):
                     response_content += f"![image]({chunk.url})\n"
                 else:
                     response_content += chunk
@@ -401,4 +720,4 @@ async def http_exception_handler(request: Request, exc: HTTPException):
 # Run the application
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True)
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)