Update main.py

main.py

@@ -1,3 +1,5 @@
+# main.py
+
 import os
 import re
 import random
@@ -12,11 +14,13 @@ from typing import List, Dict, Any, Optional, AsyncGenerator, Union
 
 from datetime import datetime
 
-from aiohttp import ClientSession, ClientTimeout, ClientError, ClientResponseError
+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  # Import the new Blackbox class
+
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
@@ -37,6 +41,7 @@ if not API_KEYS or API_KEYS == ['']:
 # Process available models
 if AVAILABLE_MODELS:
     AVAILABLE_MODELS = [model.strip() for model in AVAILABLE_MODELS.split(',') if model.strip()]
+    Blackbox.models = [model for model in Blackbox.models if model in AVAILABLE_MODELS]
 else:
     AVAILABLE_MODELS = []  # If empty, all models are available
 
@@ -96,361 +101,33 @@ 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)
 
-# Define ImageResponseData class
-class ImageResponseData:
-    def __init__(self, images: str, alt: str):
-        self.images = images
-        self.alt = alt
-
-# Mock implementation of to_data_uri
-def to_data_uri(image: Any) -> str:
-    return "data:image/png;base64,..."  # Replace with actual base64 data
-
-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',
-    ]
-
-    # Filter models based on AVAILABLE_MODELS
-    if AVAILABLE_MODELS:
-        models = [model for model in models if model in AVAILABLE_MODELS]
-
-    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": "/?model=blackboxai",
-        "gpt-4o": "/?model=gpt-4o",
-        "gemini-pro": "/?model=gemini-pro",
-        "claude-sonnet-3.5": "/?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) -> str:
-        if model in cls.models:
-            return model
-        elif model in cls.model_aliases:
-            return cls.model_aliases[model]
-        else:
-            return cls.default_model
-
-    @staticmethod
-    def generate_random_string(length: int = 7) -> str:
-        characters = string.ascii_letters + string.digits
-        return ''.join(random.choices(characters, k=length))
-
-    @staticmethod
-    def generate_next_action() -> str:
-        return uuid.uuid4().hex
+def calculate_estimated_cost(prompt_tokens: int, completion_tokens: int) -> float:
+    """
+    Calculate the estimated cost based on the number of tokens.
+    Replace the pricing below with your actual pricing model.
+    """
+    # Example pricing: $0.00000268 per token
+    cost_per_token = 0.00000268
+    return round((prompt_tokens + completion_tokens) * cost_per_token, 8)
 
-    @staticmethod
-    def generate_next_router_state_tree() -> str:
-        router_state = [
-            "",
+def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
+    return {
+        "id": f"chatcmpl-{uuid.uuid4()}",
+        "object": "chat.completion",
+        "created": int(datetime.now().timestamp()),
+        "model": model,
+        "choices": [
             {
-                "children": [
-                    "(chat)",
-                    {
-                        "children": [
-                            "__PAGE__",
-                            {}
-                        ]
-                    }
-                ]
-            },
-            None,
-            None,
-            True
-        ]
-        return json.dumps(router_state)
-
-    @staticmethod
-    def clean_response(text: str) -> str:
-        pattern = r'^\$\@\$v=undefined-rv1\$\@\$'
-        cleaned_text = re.sub(pattern, '', text)
-        return cleaned_text
-
-    @classmethod
-    async def create_async_generator(
-        cls,
-        model: str,
-        messages: List[Dict[str, str]],
-        proxy: Optional[str] = None,
-        websearch: bool = False,
-        **kwargs
-    ) -> AsyncGenerator[Union[str, ImageResponseData], None]:
-        """
-        Creates an asynchronous generator for streaming responses from Blackbox AI.
-
-        Parameters:
-            model (str): Model to use for generating responses.
-            messages (List[Dict[str, str]]): Message history.
-            proxy (Optional[str]): Proxy URL, if needed.
-            websearch (bool): Enables or disables web search mode.
-            **kwargs: Additional keyword arguments.
-
-        Yields:
-            Union[str, ImageResponseData]: Segments of the generated response or ImageResponseData objects.
-        """
-        model = cls.get_model(model)
-
-        chat_id = cls.generate_random_string()
-        next_action = cls.generate_next_action()
-        next_router_state_tree = cls.generate_next_router_state_tree()
-
-        agent_mode = cls.agentMode.get(model, {})
-        trending_agent_mode = cls.trendingAgentMode.get(model, {})
-
-        prefix = cls.model_prefixes.get(model, "")
-        
-        formatted_prompt = ""
-        for message in messages:
-            role = message.get('role', '').capitalize()
-            content = message.get('content', '')
-            if role and content:
-                formatted_prompt += f"{role}: {content}\n"
-
-        if prefix:
-            formatted_prompt = f"{prefix} {formatted_prompt}".strip()
-
-        referer_path = cls.model_referers.get(model, f"/?model={model}")
-        referer_url = f"{cls.url}{referer_path}"
-
-        common_headers = {
-            'accept': '*/*',
-            'accept-language': 'en-US,en;q=0.9',
-            'cache-control': 'no-cache',
-            'origin': cls.url,
-            'pragma': 'no-cache',
-            'priority': 'u=1, i',
-            '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'
-        }
-
-        headers_api_chat = {
-            'Content-Type': 'application/json',
-            'Referer': referer_url
-        }
-        headers_api_chat_combined = {**common_headers, **headers_api_chat}
-
-        payload_api_chat = {
-            "messages": [
-                {
-                    "id": chat_id,
-                    "content": formatted_prompt,
-                    "role": "user"
-                }
-            ],
-            "id": chat_id,
-            "previewToken": None,
-            "userId": None,
-            "codeModelMode": True,
-            "agentMode": agent_mode,
-            "trendingAgentMode": trending_agent_mode,
-            "isMicMode": False,
-            "userSystemPrompt": None,
-            "maxTokens": 1024,
-            "playgroundTopP": 0.9,
-            "playgroundTemperature": 0.5,
-            "isChromeExt": False,
-            "githubToken": None,
-            "clickedAnswer2": False,
-            "clickedAnswer3": False,
-            "clickedForceWebSearch": False,
-            "visitFromDelta": False,
-            "mobileClient": False,
-            "webSearchMode": websearch,
-            "userSelectedModel": cls.userSelectedModel.get(model, model)
-        }
-
-        headers_chat = {
-            'Accept': 'text/x-component',
-            'Content-Type': 'text/plain;charset=UTF-8',
-            'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
-            'next-action': next_action,
-            'next-router-state-tree': next_router_state_tree,
-            'next-url': '/'
-        }
-        headers_chat_combined = {**common_headers, **headers_chat}
-
-        data_chat = '[]'
-
-        async with ClientSession(headers=common_headers) as session:
-            try:
-                async with session.post(
-                    cls.api_endpoint,
-                    headers=headers_api_chat_combined,
-                    json=payload_api_chat,
-                    proxy=proxy
-                ) as response_api_chat:
-                    response_api_chat.raise_for_status()
-                    text = await response_api_chat.text()
-                    cleaned_response = cls.clean_response(text)
-
-                    if model in cls.image_models:
-                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
-                        if match:
-                            image_url = match.group(1)
-                            image_response = ImageResponseData(images=image_url, alt="Generated Image")
-                            yield image_response
-                        else:
-                            yield cleaned_response
-                    else:
-                        if websearch:
-                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
-                            if match:
-                                source_part = match.group(1).strip()
-                                answer_part = cleaned_response[match.end():].strip()
-                                try:
-                                    sources = json.loads(source_part)
-                                    source_formatted = "**Sources:**\n"
-                                    for item in sources[:5]:
-                                        title = item.get('title', 'No Title')
-                                        link = item.get('link', '#')
-                                        position = item.get('position', '')
-                                        source_formatted += f"- [{title}]({link})\n"
-                                    final_response = f"{answer_part}\n\n{source_formatted}"
-                                except json.JSONDecodeError:
-                                    final_response = f"{answer_part}\n\nSource information is unavailable."
-                            else:
-                                final_response = cleaned_response
-                        else:
-                            if '$~~~$' in cleaned_response:
-                                final_response = cleaned_response.split('$~~~$')[0].strip()
-                            else:
-                                final_response = cleaned_response
-
-                        yield final_response
-            except ClientResponseError as e:
-                error_text = f"Error {e.status}: {e.message}"
-                try:
-                    error_response = await e.response.text()
-                    cleaned_error = cls.clean_response(error_response)
-                    error_text += f" - {cleaned_error}"
-                except Exception:
-                    pass
-                yield error_text
-            except Exception as e:
-                yield f"Unexpected error during /api/chat request: {str(e)}"
-
-            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
-
-            try:
-                async with session.post(
-                    chat_url,
-                    headers=headers_chat_combined,
-                    data=data_chat,
-                    proxy=proxy
-                ) as response_chat:
-                    response_chat.raise_for_status()
-                    # No action needed based on the original code
-            except ClientResponseError as e:
-                error_text = f"Error {e.status}: {e.message}"
-                try:
-                    error_response = await e.response.text()
-                    cleaned_error = cls.clean_response(error_response)
-                    error_text += f" - {cleaned_error}"
-                except Exception:
-                    pass
-                yield error_text
-            except Exception as e:
-                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": content
+                },
+                "finish_reason": finish_reason
+            }
+        ],
+        "usage": None,  # To be filled in non-streaming responses
+    }
 
 # FastAPI app setup
 app = FastAPI()
@@ -502,39 +179,11 @@ class ChatRequest(BaseModel):
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
-    websearch: Optional[bool] = False  # Custom parameter
+    webSearchMode: Optional[bool] = False  # Custom parameter
 
 class TokenizerRequest(BaseModel):
     text: str
 
-def calculate_estimated_cost(prompt_tokens: int, completion_tokens: int) -> float:
-    """
-    Calculate the estimated cost based on the number of tokens.
-    Replace the pricing below with your actual pricing model.
-    """
-    # Example pricing: $0.00000268 per token
-    cost_per_token = 0.00000268
-    return round((prompt_tokens + completion_tokens) * cost_per_token, 8)
-
-def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
-    return {
-        "id": f"chatcmpl-{uuid.uuid4()}",
-        "object": "chat.completion",
-        "created": int(datetime.now().timestamp()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": content
-                },
-                "finish_reason": finish_reason
-            }
-        ],
-        "usage": None,  # To be filled in non-streaming responses
-    }
-
 @app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
 async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
     client_ip = req.client.host
@@ -552,55 +201,41 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         # Process the request with actual message content, but don't log it
         async_generator = Blackbox.create_async_generator(
             model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
-            websearch=request.websearch
+            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
+            websearch=request.webSearchMode
         )
 
         if request.stream:
             async def generate():
                 try:
                     assistant_content = ""
-                    for chunk in async_generator:
-                        if isinstance(chunk, ImageResponseData):
+                    async for chunk in async_generator:
+                        if isinstance(chunk, ImageResponse):
                             # Handle image responses if necessary
                             image_markdown = f"![image]({chunk.images})\n"
                             assistant_content += image_markdown
-                            response_chunk = {
-                                "id": f"chatcmpl-{uuid.uuid4()}",
-                                "object": "chat.completion.chunk",  # Change to 'chat.completion.chunk'
-                                "created": int(datetime.now().timestamp()),
-                                "model": request.model,
-                                "choices": [
-                                    {
-                                        "index": 0,
-                                        "delta": {"role": "assistant", "content": " "},  # Initial space or any starter
-                                        "finish_reason": None
-                                    }
-                                ]
-                            }
-                            yield f"data: {json.dumps(response_chunk)}\n\n"
-                            response_chunk["choices"][0]["delta"]["content"] = image_markdown.strip()
-                            yield f"data: {json.dumps(response_chunk)}\n\n"
+                            response_chunk = create_response(image_markdown, request.model, finish_reason=None)
                         else:
                             assistant_content += chunk
                             # Yield the chunk as a partial choice
                             response_chunk = {
                                 "id": f"chatcmpl-{uuid.uuid4()}",
-                                "object": "chat.completion.chunk",  # Change to 'chat.completion.chunk'
+                                "object": "chat.completion.chunk",
                                 "created": int(datetime.now().timestamp()),
                                 "model": request.model,
                                 "choices": [
                                     {
                                         "index": 0,
-                                        "delta": {"role": "assistant", "content": chunk},
-                                        "finish_reason": None
+                                        "delta": {"content": chunk, "role": "assistant"},
+                                        "finish_reason": None,
                                     }
-                                ]
+                                ],
+                                "usage": None,  # Usage can be updated if you track tokens in real-time
                             }
-                            yield f"data: {json.dumps(response_chunk)}\n\n"
+                        yield f"data: {json.dumps(response_chunk)}\n\n"
                     
                     # After all chunks are sent, send the final message with finish_reason
-                    prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
+                    prompt_tokens = sum(len(msg['content'].split()) for msg in request.messages)
                     completion_tokens = len(assistant_content.split())
                     total_tokens = prompt_tokens + completion_tokens
                     estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
@@ -637,11 +272,11 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                     error_response = {"error": str(e)}
                     yield f"data: {json.dumps(error_response)}\n\n"
 
-            return StreamingResponse(generate(), media_type="text/event-stream")  # Use SSE media type
+            return StreamingResponse(generate(), media_type="text/event-stream")
         else:
             response_content = ""
             async for chunk in async_generator:
-                if isinstance(chunk, ImageResponseData):
+                if isinstance(chunk, ImageResponse):
                     response_content += f"![image]({chunk.images})\n"
                 else:
                     response_content += chunk