Update main.py

main.py

@@ -1,27 +1,32 @@
-from __future__ import annotations
-
 import re
 import random
 import string
 import uuid
 import json
 import logging
-from aiohttp import ClientSession, ClientError
-from fastapi import FastAPI, HTTPException
+import asyncio
+from aiohttp import ClientSession, ClientTimeout, ClientError
+from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any, Optional, AsyncGenerator
 from datetime import datetime
 from fastapi.responses import StreamingResponse
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    handlers=[
+        logging.StreamHandler()
+    ]
+)
 logger = logging.getLogger(__name__)
 
 # 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."
+        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)
 
 # Mock implementations for ImageResponse and to_data_uri
@@ -33,43 +38,67 @@ class ImageResponse:
 def to_data_uri(image: Any) -> str:
     return "data:image/png;base64,..."  # Replace with actual base64 data
 
-class AsyncGeneratorProvider:
-    pass
-
-class ProviderModelMixin:
-    pass
-
-class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
+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 = 'blackbox'
+
+    default_model = 'blackboxai'
+    image_models = ['ImageGeneration']
     models = [
-        'blackbox',
-        'gemini-1.5-flash',
+        default_model,
+        'blackboxai-pro',
         "llama-3.1-8b",
         'llama-3.1-70b',
         'llama-3.1-405b',
-        'ImageGenerationLV45LJp',
         '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',
+        *image_models,
+        'Niansuh',
     ]
 
     agentMode = {
-        'ImageGenerationLV45LJp': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
+        'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
+        'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"},
     }
-
     trendingAgentMode = {
-        "blackbox": {},
+        "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 = {
@@ -78,9 +107,40 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         '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",
-        "flux": "ImageGenerationLV45LJp",
+        "claude-3.5-sonnet": "claude-sonnet-3.5",
+        "flux": "ImageGeneration",
+        "niansuh": "Niansuh",
     }
 
     @classmethod
@@ -100,15 +160,18 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         model: str,
         messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
-        image: Optional[Any] = None,
+        image: Any = None,
         image_name: Optional[str] = None,
+        webSearchMode: bool = False,
         **kwargs
-    ) -> Any:
+    ) -> AsyncGenerator[Any, None]:
         model = cls.get_model(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",
@@ -116,89 +179,130 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             "content-type": "application/json",
             "origin": cls.url,
             "pragma": "no-cache",
-            "referer": f"{cls.url}/",
-            "sec-ch-ua": '"Not;A=Brand";v="24", "Chromium";v="128"',
+            "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/128.0.0.0 Safari/537.36"
+            "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.userSelectedModel:
-            prefix = f"@{cls.userSelectedModel[model]}"
+        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']}"
         
-        async with ClientSession(headers=headers) as session:
-            if image is not None:
-                messages[-1]["data"] = {
-                    "fileText": image_name,
-                    "imageBase64": to_data_uri(image)
-                }
-            
-            random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
-
-            data = {
-                "messages": messages,
-                "id": random_id,
-                "previewToken": None,
-                "userId": None,
-                "codeModelMode": True,
-                "agentMode": {},
-                "trendingAgentMode": {},
-                "userSelectedModel": None,
-                "userSystemPrompt": None,
-                "isMicMode": False,
-                "maxTokens": 8192,
-                "playgroundTopP": 0.9,
-                "playgroundTemperature": 0.5,
-                "isChromeExt": False,
-                "githubToken": None,
-                "clickedAnswer2": False,
-                "clickedAnswer3": False,
-                "clickedForceWebSearch": False,
-                "visitFromDelta": False,
-                "mobileClient": False,
-                "webSearchMode": False,
+        random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
+        messages[-1]['id'] = random_id
+        messages[-1]['role'] = 'user'
+        if image is not None:
+            messages[-1]['data'] = {
+                'fileText': '',
+                'imageBase64': to_data_uri(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]
-            
-            try:
-                async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
-                    response.raise_for_status()
-                    response_text = await response.text()
+        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: {data}")
 
-                    # Check if the response is empty or just whitespace
-                    if not response_text.strip():
-                        raise ModelNotWorkingException(model)
+        timeout = ClientTimeout(total=60)  # Set an appropriate timeout
+        retry_attempts = 10  # Set the number of retry attempts
 
-                    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'])
+        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 == 'ImageGeneration':
+                            response_text = await response.text()
+                            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 ImageResponse(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:
-                            raise Exception("Image URL not found in the response")
-                    
-                    # Handle normal responses
-                    async for chunk in response.content.iter_any():
-                        if chunk:
-                            decoded_chunk = chunk.decode(errors='ignore')
-                            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.")
+                            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. | NiansuhAI")
+            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. | NiansuhAI")
+            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()
@@ -210,7 +314,8 @@ class Message(BaseModel):
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
-    stream: Optional[bool] = False  # Add this for streaming
+    stream: Optional[bool] = False
+    webSearchMode: Optional[bool] = False
 
 def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
@@ -228,62 +333,106 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,
     }
 
-@app.post("/niansuhai/v1/chat/completions")
-async def chat_completions(request: ChatRequest):
-    # Validate the model
-    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]
-
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatRequest, req: Request):
+    logger.info(f"Received chat completions request: {request}")
     try:
+        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,
-            image_name=None
+            image_name=None,
+            webSearchMode=request.webSearchMode
         )
-    except ModelNotWorkingException as e:
-        raise HTTPException(status_code=503, detail=str(e))
 
-    if request.stream:
-        async def generate():
+        if request.stream:
+            async def generate():
+                try:
+                    async for chunk in async_generator:
+                        if isinstance(chunk, ImageResponse):
+                            image_markdown = f"![image]({chunk.url})"
+                            response_chunk = create_response(image_markdown, request.model)
+                        else:
+                            response_chunk = create_response(chunk, request.model)
+                        
+                        # Yield each chunk in SSE format
+                        yield f"data: {json.dumps(response_chunk)}\n\n"
+                    
+                    # Signal the end of the stream
+                    yield "data: [DONE]\n\n"
+                except HTTPException as he:
+                    error_response = {"error": he.detail}
+                    yield f"data: {json.dumps(error_response)}\n\n"
+                except Exception as e:
+                    logger.exception("Error during streaming response generation.")
+                    error_response = {"error": str(e)}
+                    yield f"data: {json.dumps(error_response)}\n\n"
+
+            return StreamingResponse(generate(), media_type="text/event-stream")
+        else:
+            response_content = ""
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponse):
-                    image_markdown = f"![image]({chunk.url})"
-                    yield f"data: {json.dumps(create_response(image_markdown, request.model))}\n\n"
+                    response_content += f"![image]({chunk.url})\n"
                 else:
-                    yield f"data: {json.dumps(create_response(chunk, request.model))}\n\n"
-            yield "data: [DONE]\n\n"
+                    response_content += chunk
 
-        return StreamingResponse(generate(), media_type="text/event-stream")
-    else:
-        response_content = ""
-        async for chunk in async_generator:
-            if isinstance(chunk, ImageResponse):
-                response_content += f"![image]({chunk.url})\n"
-            else:
-                response_content += chunk
-
-        return {
-            "id": f"chatcmpl-{uuid.uuid4()}",
-            "object": "chat.completion",
-            "created": int(datetime.now().timestamp()),
-            "model": request.model,
-            "choices": [
-                {
-                    "message": {
-                        "role": "assistant",
-                        "content": response_content
-                    },
-                    "finish_reason": "stop",
-                    "index": 0
-                }
-            ],
-            "usage": None,
-        }
+            logger.info("Completed non-streaming response generation.")
+            return {
+                "id": f"chatcmpl-{uuid.uuid4()}",
+                "object": "chat.completion",
+                "created": int(datetime.now().timestamp()),
+                "model": request.model,
+                "choices": [
+                    {
+                        "message": {
+                            "role": "assistant",
+                            "content": response_content
+                        },
+                        "finish_reason": "stop",
+                        "index": 0
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": sum(len(msg['content'].split()) for msg in messages),
+                    "completion_tokens": len(response_content.split()),
+                    "total_tokens": sum(len(msg['content'].split()) for msg in messages) + len(response_content.split())
+                },
+            }
+    except ModelNotWorkingException as e:
+        logger.warning(f"Model not working: {e}")
+        raise HTTPException(status_code=503, detail=str(e))
+    except HTTPException as he:
+        logger.warning(f"HTTPException: {he.detail}")
+        raise he
+    except Exception as e:
+        logger.exception("An unexpected error occurred while processing the chat completions request.")
+        raise HTTPException(status_code=500, detail=str(e))
 
-@app.get("/niansuhai/v1/models")
+@app.get("/v1/models")
 async def get_models():
-    return {"models": Blackbox.models}
+    logger.info("Fetching available models.")
+    return {"data": [{"id": model} for model in Blackbox.models]}
+
+# Additional endpoints for better functionality
+@app.get("/v1/health")
+async def health_check():
+    """Health check endpoint to verify the service is running."""
+    return {"status": "ok"}
+
+@app.get("/v1/models/{model}/status")
+async def model_status(model: str):
+    """Check if a specific model is available."""
+    if model in Blackbox.models:
+        return {"model": model, "status": "available"}
+    elif model in Blackbox.model_aliases:
+        actual_model = Blackbox.model_aliases[model]
+        return {"model": actual_model, "status": "available via alias"}
+    else:
+        raise HTTPException(status_code=404, detail="Model not found")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)