Update main.py

main.py

@@ -1,4 +1,3 @@
-import os
 import re
 import random
 import string
@@ -6,21 +5,13 @@ import uuid
 import json
 import logging
 import asyncio
+import base64
 from aiohttp import ClientSession, ClientTimeout, ClientError
-from fastapi import FastAPI, HTTPException, Request, Depends, Header, status
-from fastapi.responses import StreamingResponse, JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, Field, validator
-from typing import List, Dict, Any, Optional, Union, AsyncGenerator, Literal
+from fastapi import FastAPI, HTTPException, Request
+from pydantic import BaseModel
+from typing import List, Dict, Any, Optional, AsyncGenerator
 from datetime import datetime
-from slowapi import Limiter, _rate_limit_exceeded_handler
-from slowapi.util import get_remote_address
-from slowapi.errors import RateLimitExceeded
-import tiktoken
-from dotenv import load_dotenv
-
-# Load environment variables from .env file
-load_dotenv()
+from fastapi.responses import StreamingResponse
 
 # Configure logging
 logging.basicConfig(
@@ -32,57 +23,6 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Initialize FastAPI app
-app = FastAPI(title="OpenAI-Compatible API")
-
-# Configure CORS (adjust origins as needed)
-origins = [
-    "*",  # Allow all origins; replace with specific origins in production
-]
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Initialize Rate Limiter from environment variable
-RATE_LIMIT = os.getenv("RATE_LIMIT", "60/minute")  # Default to 60 requests per minute
-limiter = Limiter(key_func=get_remote_address, default_limits=[RATE_LIMIT])
-app.state.limiter = limiter
-app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
-
-# API Key Authentication
-API_KEYS = set(api_key.strip() for api_key in os.getenv("API_KEYS", "").split(",") if api_key.strip())
-
-async def get_api_key(authorization: Optional[str] = Header(None)):
-    """
-    Dependency to validate API Key from the Authorization header.
-    """
-    if authorization is None:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Authorization header missing",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    parts = authorization.split()
-    if parts[0].lower() != "bearer" or len(parts) != 2:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Invalid authorization header format",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    token = parts[1]
-    if token not in API_KEYS:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Invalid API Key",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    return token
-
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
     def __init__(self, model: str):
@@ -90,162 +30,34 @@ 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)
 
-# Mock implementations for ImageResponse and to_data_uri (custom functionality)
+# Proper implementation for ImageResponse and to_data_uri
 class ImageResponse:
-    def __init__(self, url: str, alt: str):
-        self.url = url
+    def __init__(self, data_uri: str, alt: str):
+        self.data_uri = data_uri
         self.alt = alt
 
-def to_data_uri(image: Any) -> str:
-    return "data:image/png;base64,..."  # Replace with actual base64 data if needed
+def to_data_uri(image: bytes, mime_type: str = "image/png") -> str:
+    encoded = base64.b64encode(image).decode('utf-8')
+    return f"data:{mime_type};base64,{encoded}"
 
-# Token Counting using tiktoken
-def count_tokens(messages: List[Dict[str, Any]], model: str) -> int:
-    """
-    Counts the number of tokens in the messages using tiktoken.
-    Adjust the encoding based on the model.
-    """
+def decode_base64_image(data_uri: str) -> bytes:
     try:
-        encoding = tiktoken.get_encoding("cl100k_base")  # Adjust encoding as per model
-    except:
-        encoding = tiktoken.get_encoding("cl100k_base")  # Default encoding
-    tokens = 0
-    for message in messages:
-        if isinstance(message['content'], list):
-            for content_part in message['content']:
-                if isinstance(content_part, dict):
-                    if content_part.get('type') == 'text':
-                        tokens += len(encoding.encode(content_part['text']))
-                    elif content_part.get('type') == 'image_url':
-                        tokens += len(encoding.encode(content_part['image_url']['url']))
-        else:
-            tokens += len(encoding.encode(message['content']))
-    return tokens
+        header, encoded = data_uri.split(",", 1)
+        return base64.b64decode(encoded)
+    except Exception as e:
+        logger.error(f"Error decoding base64 image: {e}")
+        raise e
 
-# Blackbox Class: Handles interaction with the external AI service
 class Blackbox:
-    url = "https://www.blackbox.ai"
-    api_endpoint = os.getenv("EXTERNAL_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',
-        "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',
-        *image_models,
-        'Niansuh',
-    ]
-
-    agentMode = {
-        'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
-        'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"},
-    }
-    trendingAgentMode = {
-        "blackboxai": {},
-        "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
-        "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
-        'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
-        'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
-        'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
-        'PythonAgent': {'mode': True, 'id': "Python Agent"},
-        'JavaAgent': {'mode': True, 'id': "Java Agent"},
-        'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"},
-        'HTMLAgent': {'mode': True, 'id': "HTML Agent"},
-        'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"},
-        'AndroidDeveloper': {'mode': True, 'id': "Android Developer"},
-        'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"},
-        'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"},
-        'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"},
-        'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
-        'ReactAgent': {'mode': True, 'id': "React Agent"},
-        'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
-        'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
-    }
-
-    userSelectedModel = {
-        "gpt-4o": "gpt-4o",
-        "gemini-pro": "gemini-pro",
-        'claude-sonnet-3.5': "claude-sonnet-3.5",
-    }
-
-    model_prefixes = {
-        'gpt-4o': '@GPT-4o',
-        'gemini-pro': '@Gemini-PRO',
-        'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
-        'PythonAgent': '@Python Agent',
-        'JavaAgent': '@Java Agent',
-        'JavaScriptAgent': '@JavaScript Agent',
-        'HTMLAgent': '@HTML Agent',
-        'GoogleCloudAgent': '@Google Cloud Agent',
-        'AndroidDeveloper': '@Android Developer',
-        'SwiftDeveloper': '@Swift Developer',
-        'Next.jsAgent': '@Next.js Agent',
-        'MongoDBAgent': '@MongoDB Agent',
-        'PyTorchAgent': '@PyTorch Agent',
-        'ReactAgent': '@React Agent',
-        'XcodeAgent': '@Xcode Agent',
-        'AngularJSAgent': '@AngularJS Agent',
-        'blackboxai-pro': '@BLACKBOXAI-PRO',
-        'ImageGeneration': '@Image Generation',
-        'Niansuh': '@Niansuh',
-    }
-
-    model_referers = {
-        "blackboxai": f"{url}/?model=blackboxai",
-        "gpt-4o": f"{url}/?model=gpt-4o",
-        "gemini-pro": f"{url}/?model=gemini-pro",
-        "claude-sonnet-3.5": f"{url}/?model=claude-sonnet-3.5"
-    }
-
-    model_aliases = {
-        "gemini-flash": "gemini-1.5-flash",
-        "claude-3.5-sonnet": "claude-sonnet-3.5",
-        "flux": "ImageGeneration",
-        "niansuh": "Niansuh",
-    }
-
-    @classmethod
-    def get_model(cls, model: str) -> str:
-        if model in cls.models:
-            return model
-        elif model in cls.userSelectedModel:
-            return model
-        elif model in cls.model_aliases:
-            return cls.model_aliases[model]
-        else:
-            return cls.default_model
+    # ... [existing Blackbox class definition]
 
     @classmethod
     async def create_async_generator(
         cls,
         model: str,
-        messages: List[Dict[str, Any]],
+        messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
-        image: Any = None,
+        image: Optional[str] = None,  # Expecting a base64 string
         image_name: Optional[str] = None,
         webSearchMode: bool = False,
         **kwargs
@@ -258,52 +70,34 @@ class Blackbox:
             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",
+            # ... [existing headers]
         }
 
         if model in cls.model_prefixes:
             prefix = cls.model_prefixes[model]
-            if messages and isinstance(messages[0]['content'], list):
-                # Prepend prefix to the first text message
-                for content_part in messages[0]['content']:
-                    if isinstance(content_part, dict) and content_part.get('type') == 'text' and not content_part['text'].startswith(prefix):
-                        logger.debug(f"Adding prefix '{prefix}' to the first text message.")
-                        content_part['text'] = f"{prefix} {content_part['text']}"
-                        break
-            elif messages and isinstance(messages[0]['content'], str) and not messages[0]['content'].startswith(prefix):
+            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))
-        # Assuming the last message is from the user
-        if messages:
-            last_message = messages[-1]
-            if isinstance(last_message['content'], list):
-                for content_part in last_message['content']:
-                    if isinstance(content_part, dict) and content_part.get('type') == 'text':
-                        content_part['role'] = 'user'
-            else:
-                last_message['id'] = random_id
-                last_message['role'] = 'user'
+        messages[-1]['id'] = random_id
+        messages[-1]['role'] = 'user'
         
         if image is not None:
-            # Process image if required
-            # This implementation assumes that image URLs are handled by the external service
-            pass  # Implement as needed
-
+            try:
+                image_bytes = decode_base64_image(image)
+                data_uri = to_data_uri(image_bytes)
+                messages[-1]['data'] = {
+                    'fileText': '',
+                    'imageBase64': data_uri,
+                    'title': image_name
+                }
+                messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content']
+                logger.debug("Image data added to the message.")
+            except Exception as e:
+                logger.error(f"Failed to decode base64 image: {e}")
+                raise HTTPException(status_code=400, detail="Invalid image data provided.")
+        
         data = {
             "messages": messages,
             "id": random_id,
@@ -314,7 +108,7 @@ class Blackbox:
             "trendingAgentMode": {},
             "isMicMode": False,
             "userSystemPrompt": None,
-            "maxTokens": int(os.getenv("MAX_TOKENS", "4096")),
+            "maxTokens": 99999999,
             "playgroundTopP": 0.9,
             "playgroundTemperature": 0.5,
             "isChromeExt": False,
@@ -351,85 +145,77 @@ class Blackbox:
                             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'])
+                                
+                                # Fetch the image data
+                                async with session.get(image_url) as img_response:
+                                    img_response.raise_for_status()
+                                    image_bytes = await img_response.read()
+                                    data_uri = to_data_uri(image_bytes)
+                                    logger.info("Image converted to base64 data URI.")
+                                
+                                yield ImageResponse(data_uri, alt=messages[-1]['content'])
                             else:
                                 logger.error("Image URL not found in the response.")
                                 raise Exception("Image URL not found in the response")
                         else:
-                            async for chunk, info 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():
-                                        yield decoded_chunk
+                            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.")
+                    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.")
+                    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))
 
-# Pydantic Models
-class TextContent(BaseModel):
-    type: Literal["text"] = Field(..., description="Type of content, e.g., 'text'.")
-    text: str = Field(..., description="The text content.")
-
-class ImageURLContent(BaseModel):
-    type: Literal["image_url"] = Field(..., description="Type of content, e.g., 'image_url'.")
-    image_url: Dict[str, str] = Field(..., description="Dictionary containing the image URL.")
-
-Content = Union[TextContent, ImageURLContent]
+# FastAPI app setup
+app = FastAPI()
 
 class Message(BaseModel):
-    role: str = Field(..., description="The role of the message author.")
-    content: Union[str, List[Content]] = Field(..., description="The content of the message. Can be a string or a list of content parts.")
-
-    @validator('content', pre=True)
-    def validate_content(cls, v):
-        if isinstance(v, list):
-            processed_content = []
-            for item in v:
-                if 'type' not in item:
-                    raise ValueError("Each content part must have a 'type' field.")
-                if item['type'] == 'text':
-                    processed_content.append(TextContent(**item))
-                elif item['type'] == 'image_url':
-                    processed_content.append(ImageURLContent(**item))
-                else:
-                    raise ValueError(f"Unsupported content type: {item['type']}")
-            return processed_content
-        elif isinstance(v, str):
-            return v
-        else:
-            raise ValueError("Content must be either a string or a list of content parts.")
+    role: str
+    content: str
 
 class ChatRequest(BaseModel):
-    model: str = Field(..., description="ID of the model to use.")
-    messages: List[Message] = Field(..., description="A list of messages comprising the conversation.")
-    stream: Optional[bool] = Field(False, description="Whether to stream the response.")
-    webSearchMode: Optional[bool] = Field(False, description="Whether to enable web search mode.")
-
-class ChatCompletionChoice(BaseModel):
-    index: int
-    delta: Dict[str, Any]
-    finish_reason: Optional[str] = None
-
-class ChatCompletionResponse(BaseModel):
-    id: str
-    object: str
-    created: int
     model: str
-    choices: List[ChatCompletionChoice]
-    usage: Optional[Dict[str, int]] = None
+    messages: List[Message]
+    stream: Optional[bool] = False
+    webSearchMode: Optional[bool] = False
+    image: Optional[str] = None  # Add image field for base64 data
 
-# Utility Function to Create Response
 def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -443,58 +229,36 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
                 "finish_reason": finish_reason,
             }
         ],
-        "usage": None,  # To be populated if usage metrics are available
+        "usage": None,
     }
 
-# Endpoint: Chat Completions
-@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
-@limiter.limit("60/minute")  # Example: 60 requests per minute per IP
-async def chat_completions(
-    chat_request: ChatRequest,            # Renamed from 'request' to 'chat_request'
-    request: Request,                     # Added 'request: Request' parameter
-    api_key: str = Depends(get_api_key)
-):
-    logger.info(f"Received chat completions request: {chat_request}")
+@app.post("/niansuhai/v1/chat/completions")
+async def chat_completions(request: ChatRequest, req: Request):
+    logger.info(f"Received chat completions request: {request}")
     try:
-        # Process messages for token counting and sending to Blackbox
-        processed_messages = []
-        for msg in chat_request.messages:
-            if isinstance(msg.content, list):
-                # Convert list of content parts to a structured format
-                combined_content = []
-                for part in msg.content:
-                    if isinstance(part, TextContent):
-                        combined_content.append({"type": part.type, "text": part.text})
-                    elif isinstance(part, ImageURLContent):
-                        combined_content.append({"type": part.type, "image_url": part.image_url})
-                processed_messages.append({"role": msg.role, "content": combined_content})
-            else:
-                processed_messages.append({"role": msg.role, "content": msg.content})
-
-        prompt_tokens = count_tokens(processed_messages, chat_request.model)
-
+        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
+        
         async_generator = Blackbox.create_async_generator(
-            model=chat_request.model,
-            messages=processed_messages,
-            image=None,  # Adjust if image handling is required
+            model=request.model,
+            messages=messages,
+            proxy=None,  # Pass proxy if needed
+            image=request.image,  # Pass the base64 image
             image_name=None,
-            webSearchMode=chat_request.webSearchMode
+            webSearchMode=request.webSearchMode
         )
 
-        if chat_request.stream:
+        if request.stream:
             async def generate():
                 try:
-                    completion_tokens = 0
                     async for chunk in async_generator:
                         if isinstance(chunk, ImageResponse):
-                            image_markdown = f"![image]({chunk.url})"
-                            response_chunk = create_response(image_markdown, chat_request.model)
-                            yield f"data: {json.dumps(response_chunk)}\n\n"
-                            completion_tokens += len(image_markdown.split())
+                            image_markdown = f"![{chunk.alt}]({chunk.data_uri})"
+                            response_chunk = create_response(image_markdown, request.model)
                         else:
-                            response_chunk = create_response(chunk, chat_request.model)
-                            yield f"data: {json.dumps(response_chunk)}\n\n"
-                            completion_tokens += len(chunk.split())
+                            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"
@@ -509,36 +273,34 @@ async def chat_completions(
             return StreamingResponse(generate(), media_type="text/event-stream")
         else:
             response_content = ""
-            completion_tokens = 0
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponse):
-                    response_content += f"![image]({chunk.url})\n"
-                    completion_tokens += len(f"![image]({chunk.url})\n".split())
+                    response_content += f"![{chunk.alt}]({chunk.data_uri})\n"
                 else:
                     response_content += chunk
-                    completion_tokens += len(chunk.split())
-
-            total_tokens = prompt_tokens + completion_tokens
 
             logger.info("Completed non-streaming response generation.")
-            return ChatCompletionResponse(
-                id=f"chatcmpl-{uuid.uuid4()}",
-                object="chat.completion",
-                created=int(datetime.now().timestamp()),
-                model=chat_request.model,
-                choices=[
-                    ChatCompletionChoice(
-                        index=0,
-                        delta={"content": response_content, "role": "assistant"},
-                        finish_reason="stop"
-                    )
+            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
-                }
-            )
+                "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))
@@ -549,24 +311,19 @@ async def chat_completions(
         logger.exception("An unexpected error occurred while processing the chat completions request.")
         raise HTTPException(status_code=500, detail=str(e))
 
-# Endpoint: List Models
-@app.get("/v1/models", response_model=Dict[str, List[Dict[str, str]]])
-@limiter.limit("60/minute")
-async def get_models(
-    request: Request,                      # Ensure 'request: Request' parameter is present
-    api_key: str = Depends(get_api_key)
-):
+@app.get("/niansuhai/v1/models")
+async def get_models():
     logger.info("Fetching available models.")
     return {"data": [{"id": model} for model in Blackbox.models]}
 
-# Endpoint: Model Status
-@app.get("/v1/models/{model}/status", response_model=Dict[str, str])
-@limiter.limit("60/minute")
-async def model_status(
-    model: str,
-    request: Request,                      # Ensure 'request: Request' parameter is present
-    api_key: str = Depends(get_api_key)
-):
+# Additional endpoints for better functionality
+@app.get("/niansuhai/v1/health")
+async def health_check():
+    """Health check endpoint to verify the service is running."""
+    return {"status": "ok"}
+
+@app.get("/niansuhai/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"}
@@ -576,16 +333,6 @@ async def model_status(
     else:
         raise HTTPException(status_code=404, detail="Model not found")
 
-# Endpoint: Health Check
-@app.get("/v1/health", response_model=Dict[str, str])
-@limiter.limit("60/minute")
-async def health_check(
-    request: Request                       # Ensure 'request: Request' parameter is present
-):
-    """Health check endpoint to verify the service is running."""
-    return {"status": "ok"}
-
-# Run the application
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)