Update main.py

main.py

@@ -1,5 +1,68 @@
+import os
+import re
+import random
+import string
+import uuid
+import json
+import logging
+import asyncio
+import time
+from collections import defaultdict
+from typing import List, Dict, Any, Optional, AsyncGenerator, Union
+
+from datetime import datetime
+
+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
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    handlers=[logging.StreamHandler()]
+)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+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
+
+if not API_KEYS or API_KEYS == ['']:
+    logger.error("No API keys found. Please set the API_KEYS environment variable.")
+    raise Exception("API_KEYS environment variable not set.")
+
+# Process available models
+if AVAILABLE_MODELS:
+    AVAILABLE_MODELS = [model.strip() for model in AVAILABLE_MODELS.split(',') if model.strip()]
+else:
+    AVAILABLE_MODELS = []  # If empty, all models are available
+
 # Simple in-memory rate limiter based solely on IP addresses
+rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
+
+# Define cleanup interval and window
+CLEANUP_INTERVAL = 60  # seconds
+RATE_LIMIT_WINDOW = 60  # seconds
+
+async def cleanup_rate_limit_stores():
+    """
+    Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
+    """
+    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]
+        for ip in ips_to_delete:
+            del rate_limit_store[ip]
+            logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
+        await asyncio.sleep(CLEANUP_INTERVAL)
+
 async def rate_limiter_per_ip(request: Request):
+    """
+    Rate limiter that enforces a limit based on the client's IP address.
+    """
     client_ip = request.client.host
     current_time = time.time()
 
@@ -9,18 +72,43 @@ async def rate_limiter_per_ip(request: Request):
     else:
         if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
             logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
-            raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address')
+            raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
         rate_limit_store[client_ip]["count"] += 1
 
-CLEANUP_INTERVAL = 60  # seconds
-RATE_LIMIT_WINDOW = 60  # seconds
+async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
+    """
+    Dependency to extract and validate the API key from the Authorization header.
+    """
+    client_ip = request.client.host
+    if authorization is None or not authorization.startswith('Bearer '):
+        logger.warning(f"Invalid or missing authorization header from IP: {client_ip}")
+        raise HTTPException(status_code=401, detail='Invalid authorization header format')
+    api_key = authorization[7:]
+    if api_key not in API_KEYS:
+        logger.warning(f"Invalid API key attempted: {api_key} from IP: {client_ip}")
+        raise HTTPException(status_code=401, detail='Invalid API key')
+    return api_key
+
+# 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 try another model or wait for it to be fixed."
+        super().__init__(self.message)
+
+# Mock implementations for ImageResponse and to_data_uri
+class ImageResponse:
+    def __init__(self, url: str, alt: str):
+        self.url = url
+        self.alt = alt
+
+def to_data_uri(image: Any) -> str:
+    return "data:image/png;base64,..."  # Replace with actual base64 data
 
 class Blackbox:
-    label = "Blackbox AI"
     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
@@ -30,7 +118,6 @@ class Blackbox:
     models = [
         default_model,
         'blackboxai-pro',
-        *image_models,
         "llama-3.1-8b",
         'llama-3.1-70b',
         'llama-3.1-405b',
@@ -51,12 +138,18 @@ class Blackbox:
         'ReactAgent',
         'XcodeAgent',
         'AngularJSAgent',
+        *image_models,
+        'Niansuh',
     ]
 
+    # 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'},
@@ -104,65 +197,33 @@ class Blackbox:
         '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"
+        "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) -> str:
+    def get_model(cls, model: str) -> Optional[str]:
         if model in cls.models:
             return model
-        elif model in cls.model_aliases:
+        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
-
-    @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
+            return cls.default_model if cls.default_model in cls.models else None
 
     @classmethod
     async def create_async_generator(
@@ -170,77 +231,73 @@ class Blackbox:
         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, Dict[str, Any]], None]:
-        """
-        Creates an asynchronous generator for streaming responses from Blackbox AI.
-        """
+    ) -> AsyncGenerator[Any, None]:
         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, {})
+        logger.info(f"Selected model: {model}")
 
-        prefix = cls.model_prefixes.get(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)
         
-        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 = {
+            "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 = {
-            '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,
+        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': 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,
@@ -250,66 +307,81 @@ class Blackbox:
             "clickedForceWebSearch": False,
             "visitFromDelta": False,
             "mobileClient": False,
-            "webSearchMode": websearch,
-            "userSelectedModel": cls.userSelectedModel.get(model, model)
+            "userSelectedModel": None,
+            "webSearchMode": webSearchMode,
         }
 
-        async with ClientSession(headers=common_headers) as session:
+        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 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)
-                            yield {"type": "image", "url": image_url, "alt": "Generated Image"}
-                        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 = "**Source:**\n"
-                                    for item in sources:
-                                        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."
+                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:
-                                final_response = cleaned_response
+                                logger.error("Image URL not found in the response.")
+                                raise Exception("Image URL not found in the 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
+                            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:
-                yield f"Unexpected error during /api/chat request: {str(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()
@@ -354,13 +426,47 @@ class ChatRequest(BaseModel):
     temperature: Optional[float] = 1.0
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, List[str]]] = None
     max_tokens: Optional[int] = None
     presence_penalty: Optional[float] = 0.0
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None
     user: Optional[str] = None
+    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: Request)])
+@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
     # Redact user messages only for logging purposes
@@ -375,35 +481,119 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             raise HTTPException(status_code=400, detail="Requested model is not available.")
 
         # Process the request with actual message content, but don't log it
-        response_content = await Blackbox.create_async_generator(
+        async_generator = Blackbox.create_async_generator(
             model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-            temperature=request.temperature,
-            max_tokens=request.max_tokens
+            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
+            image=None,
+            image_name=None,
+            webSearchMode=request.webSearchMode
         )
 
-        logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
-        return {
-            "id": f"chatcmpl-{uuid.uuid4()}",
-            "object": "chat.completion",
-            "created": int(datetime.now().timestamp()),
-            "model": request.model,
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": response_content
-                    },
-                    "finish_reason": "stop"
-                }
-            ],
-            "usage": {
-                "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
-                "completion_tokens": len(response_content.split()),
-                "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
-            },
-        }
+        if request.stream:
+            async def generate():
+                try:
+                    assistant_content = ""
+                    async for chunk in async_generator:
+                        if isinstance(chunk, ImageResponse):
+                            # Handle image responses if necessary
+                            image_markdown = f"![image]({chunk.url})\n"
+                            assistant_content += image_markdown
+                            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",
+                                "created": int(datetime.now().timestamp()),
+                                "model": request.model,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "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"
+                    
+                    # 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)
+
+                    final_response = {
+                        "id": f"chatcmpl-{uuid.uuid4()}",
+                        "object": "chat.completion",
+                        "created": int(datetime.now().timestamp()),
+                        "model": request.model,
+                        "choices": [
+                            {
+                                "message": {
+                                    "role": "assistant",
+                                    "content": assistant_content
+                                },
+                                "finish_reason": "stop",
+                                "index": 0
+                            }
+                        ],
+                        "usage": {
+                            "prompt_tokens": prompt_tokens,
+                            "completion_tokens": completion_tokens,
+                            "total_tokens": total_tokens,
+                            "estimated_cost": estimated_cost
+                        },
+                    }
+                    yield f"data: {json.dumps(final_response)}\n\n"
+                    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(f"Error during streaming response generation from IP: {client_ip}.")
+                    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):
+                    response_content += f"![image]({chunk.url})\n"
+                else:
+                    response_content += chunk
+
+            prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
+            completion_tokens = len(response_content.split())
+            total_tokens = prompt_tokens + completion_tokens
+            estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
+
+            logger.info(f"Completed non-streaming response generation for API key: {api_key} | IP: {client_ip}")
+
+            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": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": total_tokens,
+                    "estimated_cost": estimated_cost
+                },
+            }
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
         raise HTTPException(status_code=503, detail=str(e))
@@ -414,6 +604,15 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         logger.exception(f"An unexpected error occurred while processing the chat completions request from IP: {client_ip}.")
         raise HTTPException(status_code=500, detail=str(e))
 
+# Endpoint: POST /v1/tokenizer
+@app.post("/v1/tokenizer", dependencies=[Depends(rate_limiter_per_ip)])
+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}
+
 # Endpoint: GET /v1/models
 @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
 async def get_models(req: Request):
@@ -421,6 +620,20 @@ async def get_models(req: Request):
     logger.info(f"Fetching available models from IP: {client_ip}")
     return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
 
+# Endpoint: GET /v1/models/{model}/status
+@app.get("/v1/models/{model}/status", dependencies=[Depends(rate_limiter_per_ip)])
+async def model_status(model: str, req: Request):
+    client_ip = req.client.host
+    logger.info(f"Model status requested for '{model}' from IP: {client_ip}")
+    if model in Blackbox.models:
+        return {"model": model, "status": "available"}
+    elif model in Blackbox.model_aliases and Blackbox.model_aliases[model] in Blackbox.models:
+        actual_model = Blackbox.model_aliases[model]
+        return {"model": actual_model, "status": "available via alias"}
+    else:
+        logger.warning(f"Model not found: {model} from IP: {client_ip}")
+        raise HTTPException(status_code=404, detail="Model not found")
+
 # Endpoint: GET /v1/health
 @app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
 async def health_check(req: Request):
@@ -428,6 +641,13 @@ async def health_check(req: Request):
     logger.info(f"Health check requested from IP: {client_ip}")
     return {"status": "ok"}
 
+# Endpoint: GET /v1/chat/completions (GET method)
+@app.get("/v1/chat/completions")
+async def chat_completions_get(req: Request):
+    client_ip = req.client.host
+    logger.info(f"GET request made to /v1/chat/completions from IP: {client_ip}, redirecting to 'about:blank'")
+    return RedirectResponse(url='about:blank')
+
 # Custom exception handler to match OpenAI's error format
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request: Request, exc: HTTPException):
@@ -444,3 +664,8 @@ async def http_exception_handler(request: Request, exc: HTTPException):
             }
         },
     )
+
+# Run the application
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)