Update main.py

main.py

@@ -8,14 +8,15 @@ import logging
 import asyncio
 import time
 from collections import defaultdict
-from typing import List, Dict, Any, Optional, AsyncGenerator, Union
+from typing import List, Dict, Any, Optional, AsyncGenerator, Union, Callable, Type, Tuple
 
 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 tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type, RetryError
 
 # Configure logging
 logging.basicConfig(
@@ -29,6 +30,7 @@ logger = logging.getLogger(__name__)
 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
+RETRY_ATTEMPTS = int(os.getenv('RETRY_ATTEMPTS', '5'))  # Retry attempts
 
 if not API_KEYS or API_KEYS == ['']:
     logger.error("No API keys found. Please set the API_KEYS environment variable.")
@@ -53,10 +55,7 @@ async def cleanup_rate_limit_stores():
     """
     while True:
         current_time = time.time()
-        ips_to_delete = [
-            ip for ip, value in rate_limit_store.items()
-            if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2
-        ]
+        ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
         for ip in ips_to_delete:
             del rate_limit_store[ip]
             logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
@@ -99,25 +98,47 @@ 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)
 
-# ImageResponse class
+# Mock implementations for ImageResponse and to_data_uri
 class ImageResponse:
     def __init__(self, url: str, alt: str):
         self.url = url
         self.alt = alt
 
-# Placeholder classes for AsyncGeneratorProvider and ProviderModelMixin
-class AsyncGeneratorProvider:
-    pass  # Implement as per your actual provider's requirements
-
-class ProviderModelMixin:
-    pass  # Implement as per your actual provider's requirements
-
-class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
-    label = "Blackbox AI"
+def to_data_uri(image: Any) -> str:
+    return "data:image/png;base64,..."  # Replace with actual base64 data
+
+# Retry Decorator
+def async_retry(
+    retries: int = 5,
+    exceptions: Tuple[Type[BaseException], ...] = (ClientError, asyncio.TimeoutError),
+    initial_delay: float = 1.0,
+    max_delay: float = 10.0,
+    backoff_multiplier: float = 2.0,
+    jitter: float = 0.1,
+) -> Callable:
+    """
+    Asynchronous retry decorator with exponential backoff and jitter.
+    """
+    def decorator(func: Callable) -> Callable:
+        @retry(
+            stop=stop_after_attempt(retries),
+            wait=wait_exponential(multiplier=initial_delay, min=initial_delay, max=max_delay) + wait_exponential(multiplier=0, max=jitter),
+            retry=retry_if_exception_type(exceptions),
+            reraise=True,
+        )
+        async def wrapper(*args, **kwargs):
+            try:
+                return await func(*args, **kwargs)
+            except exceptions as e:
+                logger.warning(f"Function {func.__name__} failed with {e}. Retrying...")
+                raise
+        return wrapper
+    return decorator
+
+class Blackbox:
     url = "https://www.blackbox.ai"
     api_endpoint = "https://www.blackbox.ai/api/chat"
     working = True
-    supports_gpt_4 = True
     supports_stream = True
     supports_system_message = True
     supports_message_history = True
@@ -127,7 +148,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     models = [
         default_model,
         'blackboxai-pro',
-        *image_models,
         "llama-3.1-8b",
         'llama-3.1-70b',
         'llama-3.1-405b',
@@ -148,6 +168,8 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         'ReactAgent',
         'XcodeAgent',
         'AngularJSAgent',
+        *image_models,
+        'Niansuh',
     ]
 
     # Filter models based on AVAILABLE_MODELS
@@ -156,8 +178,8 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
 
     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'},
@@ -184,7 +206,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         "gpt-4o": "gpt-4o",
         "gemini-pro": "gemini-pro",
         'claude-sonnet-3.5': "claude-sonnet-3.5",
-        "niansuh": "Niansuh",  # Added based on model_aliases
     }
 
     model_prefixes = {
@@ -210,10 +231,10 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     }
 
     model_referers = {
-        "blackboxai": "/?model=blackboxai",
-        "gpt-4o": "/?model=gpt-4o",
-        "gemini-pro": "/?model=gemini-pro",
-        "claude-sonnet-3.5": "/?model=claude-sonnet-3.5"
+        "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 = {
@@ -234,132 +255,89 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         else:
             return cls.default_model if cls.default_model in cls.models else None
 
-    @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
-
-    @staticmethod
-    def generate_next_router_state_tree() -> str:
-        router_state = [
-            "",
-            {
-                "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_retry(
+        retries=RETRY_ATTEMPTS,
+        exceptions=(ClientError, asyncio.TimeoutError),
+        initial_delay=1.0,
+        max_delay=10.0,
+        backoff_multiplier=2.0,
+        jitter=0.1,
+    )
     async def create_async_generator(
         cls,
         model: str,
         messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
-        websearch: bool = False,
+        image: Any = None,
+        image_name: Optional[str] = None,
+        webSearchMode: bool = False,
         **kwargs
-    ) -> AsyncGenerator[Union[str, ImageResponse], None]:
+    ) -> AsyncGenerator[Any, 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, ImageResponse]: Segments of the generated response or ImageResponse objects.
+        Create an asynchronous generator to interact with the external API.
         """
         model = cls.get_model(model)
         if model is None:
             logger.error(f"Model {model} is not available.")
             raise ModelNotWorkingException(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'
-        }
+        logger.info(f"Selected model: {model}")
 
-        headers_api_chat = {
-            'Content-Type': 'application/json',
-            'Referer': referer_url
+        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",
         }
-        headers_api_chat_combined = {**common_headers, **headers_api_chat}
-
-        payload_api_chat = {
-            "messages": [
-                {
-                    "id": chat_id,
-                    "content": formatted_prompt,
-                    "role": "user"
-                }
-            ],
-            "id": chat_id,
+
+        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'
+
+        logger.debug(f"Generated message ID: {random_id} for model: {model}")
+
+        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": agent_mode,
-            "trendingAgentMode": trending_agent_mode,
+            "agentMode": {},
+            "trendingAgentMode": {},
             "isMicMode": False,
             "userSystemPrompt": None,
-            "maxTokens": 1024,
+            "maxTokens": 99999999,
             "playgroundTopP": 0.9,
             "playgroundTemperature": 0.5,
             "isChromeExt": False,
@@ -369,111 +347,101 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             "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': '/'
+            "userSelectedModel": None,
+            "webSearchMode": webSearchMode,
         }
-        headers_chat_combined = {**common_headers, **headers_chat}
-
-        data_chat = '[]'
 
-        async with ClientSession(headers=common_headers) as session:
-            try:
-                # Send initial chat request with streaming
-                async with session.post(
-                    cls.api_endpoint,
-                    headers=headers_api_chat_combined,
-                    json=payload_api_chat,
-                    proxy=proxy,
-                    timeout=ClientTimeout(total=600),  # Adjust timeout as needed
-                ) as response_api_chat:
-                    response_api_chat.raise_for_status()
-                    
-                    # Stream the response in chunks
-                    async for data in response_api_chat.content.iter_chunked(1024):
-                        decoded_data = data.decode('utf-8', errors='ignore')
-                        cleaned_data = cls.clean_response(decoded_data)
-                        
-                        # Check for image response
-                        image_match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_data)
-                        if image_match:
-                            image_url = image_match.group(1)
-                            image_response = ImageResponse(url=image_url, alt="Generated Image")
-                            yield image_response
-                            continue  # Continue to the next chunk
-
-                        # Check for web search sources
-                        if websearch:
-                            source_match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_data, re.DOTALL)
-                            if source_match:
-                                source_part = source_match.group(1).strip()
-                                answer_part = cleaned_data[source_match.end():].strip()
+        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
+
+        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.")
+                            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:
+                        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:
-                                    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"{position}. [{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."
-                                yield final_response
-                                continue  # Continue to the next chunk
-
-                        # Yield the cleaned data chunk
-                        if cleaned_data.strip():
-                            yield cleaned_data.strip()
-            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)}"
+                                    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
+        except RetryError as re:
+            logger.error(f"All retry attempts failed for {cls.api_endpoint}: {re}")
+            raise HTTPException(status_code=502, detail="Error communicating with the external API.")
+
+# FastAPI app setup
+app = FastAPI()
 
-            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
+# Add the cleanup task when the app starts
+@app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(cleanup_rate_limit_stores())
+    logger.info("Started rate limit store cleanup task.")
 
-            try:
-                # Send follow-up chat request (if necessary)
-                async with session.post(
-                    chat_url,
-                    headers=headers_chat_combined,
-                    data=data_chat,
-                    proxy=proxy,
-                    timeout=ClientTimeout(total=600),
-                ) as response_chat:
-                    response_chat.raise_for_status()
-                    # If there's additional streaming data from the chat URL, handle it here
-                    async for data in response_chat.content.iter_chunked(1024):
-                        decoded_data = data.decode('utf-8', errors='ignore')
-                        cleaned_data = cls.clean_response(decoded_data)
-                        if cleaned_data.strip():
-                            yield cleaned_data.strip()
-            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)}"
+# Middleware to enhance security and enforce Content-Type for specific endpoints
+@app.middleware("http")
+async def security_middleware(request: Request, call_next):
+    client_ip = request.client.host
+    # Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
+    if request.method == "POST" and request.url.path == "/v1/chat/completions":
+        content_type = request.headers.get("Content-Type")
+        if content_type != "application/json":
+            logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": {
+                        "message": "Content-Type must be application/json",
+                        "type": "invalid_request_error",
+                        "param": None,
+                        "code": None
+                    }
+                },
+            )
+    response = await call_next(request)
+    return response
 
 # Request Models
 class Message(BaseModel):
@@ -526,40 +494,6 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,  # To be filled in non-streaming responses
     }
 
-# Initialize FastAPI app
-app = FastAPI()
-
-# Add the cleanup task when the app starts
-@app.on_event("startup")
-async def startup_event():
-    asyncio.create_task(cleanup_rate_limit_stores())
-    logger.info("Started rate limit store cleanup task.")
-
-# Middleware to enhance security and enforce Content-Type for specific endpoints
-@app.middleware("http")
-async def security_middleware(request: Request, call_next):
-    client_ip = request.client.host
-    # Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
-    if request.method == "POST" and request.url.path == "/v1/chat/completions":
-        content_type = request.headers.get("Content-Type")
-        if content_type != "application/json":
-            logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
-            return JSONResponse(
-                status_code=400,
-                content={
-                    "error": {
-                        "message": "Content-Type must be application/json",
-                        "type": "invalid_request_error",
-                        "param": None,
-                        "code": None
-                    }
-                },
-            )
-    response = await call_next(request)
-    return response
-
-# FastAPI Endpoints
-
 @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
@@ -578,28 +512,24 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         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
-            proxy=None,  # Add proxy if needed
-            websearch=request.webSearchMode
+            image=None,
+            image_name=None,
+            webSearchMode=request.webSearchMode
         )
 
         if request.stream:
             async def generate():
                 try:
                     assistant_content = ""
-                    prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
-                    completion_tokens = 0
-
                     async for chunk in async_generator:
                         if isinstance(chunk, ImageResponse):
                             # Handle image responses if necessary
-                            image_markdown = f"![{chunk.alt}]({chunk.url})\n"
+                            image_markdown = f"![image]({chunk.url})\n"
                             assistant_content += image_markdown
                             response_chunk = create_response(image_markdown, request.model, finish_reason=None)
-                            yield f"data: {json.dumps(response_chunk)}\n\n"
                         else:
-                            # Assuming 'chunk' is a string of text
                             assistant_content += chunk
-                            completion_tokens += len(chunk.split())
+                            # Yield the chunk as a partial choice
                             response_chunk = {
                                 "id": f"chatcmpl-{uuid.uuid4()}",
                                 "object": "chat.completion.chunk",
@@ -614,9 +544,11 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                                 ],
                                 "usage": None,  # Usage can be updated if you track tokens in real-time
                             }
-                            yield f"data: {json.dumps(response_chunk)}\n\n"
-
-                    # After all chunks are sent, calculate tokens and estimated cost
+                        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)
+                    completion_tokens = len(assistant_content.split())
                     total_tokens = prompt_tokens + completion_tokens
                     estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
 
@@ -657,7 +589,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             response_content = ""
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponse):
-                    response_content += f"![{chunk.alt}]({chunk.url})\n"
+                    response_content += f"![image]({chunk.url})\n"
                 else:
                     response_content += chunk
 
@@ -705,9 +637,18 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
 async def tokenizer(request: TokenizerRequest, req: Request):
     client_ip = req.client.host
     text = request.text
-    token_count = len(text.split())
     logger.info(f"Tokenizer requested from IP: {client_ip} | Text length: {len(text)}")
-    return {"text": text, "tokens": token_count}
+    
+    try:
+        # Example integration: Assuming Blackbox has a tokenizer endpoint
+        result = await Blackbox.process_tokenizer_request(text)
+        token_count = result.get("tokens", len(text.split()))
+        return {"text": text, "tokens": token_count}
+    except HTTPException as he:
+        raise he
+    except Exception as e:
+        logger.exception(f"An unexpected error occurred during tokenization from IP: {client_ip}.")
+        raise HTTPException(status_code=500, detail=str(e))
 
 # Endpoint: GET /v1/models
 @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])