Update main.py

main.py

@@ -1,3 +1,4 @@
+import os
 import re
 import random
 import string
@@ -5,13 +6,14 @@ import uuid
 import json
 import logging
 import asyncio
-import base64
-from aiohttp import ClientSession, ClientTimeout, ClientError
-from fastapi import FastAPI, HTTPException, Request
-from pydantic import BaseModel
+import time
+from collections import defaultdict
 from typing import List, Dict, Any, Optional, AsyncGenerator
-from datetime import datetime
+
+from aiohttp import ClientSession, ClientTimeout, ClientError
+from fastapi import FastAPI, HTTPException, Request, Depends, Header
 from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
 
 # Configure logging
 logging.basicConfig(
@@ -23,6 +25,47 @@ logging.basicConfig(
 )
 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
+
+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.")
+
+# Simple in-memory rate limiter
+rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
+
+async def get_api_key(authorization: str = Header(...)) -> str:
+    """
+    Dependency to extract and validate the API key from the Authorization header.
+    Expects the header in the format: Authorization: Bearer <API_KEY>
+    """
+    if not authorization.startswith('Bearer '):
+        logger.warning("Invalid authorization header format.")
+        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}")
+        raise HTTPException(status_code=401, detail='Invalid API key')
+    return api_key
+
+async def rate_limiter(api_key: str = Depends(get_api_key)):
+    """
+    Dependency to enforce rate limiting per API key.
+    Raises HTTP 429 if the rate limit is exceeded.
+    """
+    current_time = time.time()
+    window_start = rate_limit_store[api_key]["timestamp"]
+    if current_time - window_start > 60:
+        # Reset the count and timestamp after the time window
+        rate_limit_store[api_key] = {"count": 1, "timestamp": current_time}
+    else:
+        if rate_limit_store[api_key]["count"] >= RATE_LIMIT:
+            logger.warning(f"Rate limit exceeded for API key: {api_key}")
+            raise HTTPException(status_code=429, detail='Rate limit exceeded')
+        rate_limit_store[api_key]["count"] += 1
+
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
     def __init__(self, model: str):
@@ -30,23 +73,14 @@ 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)
 
-# Implementation for ImageResponse and to_data_uri
+# Mock implementations for ImageResponse and to_data_uri
 class ImageResponse:
-    def __init__(self, data_uri: str, alt: str):
-        self.data_uri = data_uri
+    def __init__(self, url: str, alt: str):
+        self.url = url
         self.alt = alt
 
-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}"
-
-def decode_base64_image(data_uri: str) -> bytes:
-    try:
-        header, encoded = data_uri.split(",", 1)
-        return base64.b64decode(encoded)
-    except Exception as e:
-        logger.error(f"Error decoding base64 image: {e}")
-        raise e
+def to_data_uri(image: Any) -> str:
+    return "data:image/png;base64,..."  # Replace with actual base64 data
 
 class Blackbox:
     url = "https://www.blackbox.ai"
@@ -158,7 +192,7 @@ class Blackbox:
         if model in cls.models:
             return model
         elif model in cls.userSelectedModel:
-            return cls.userSelectedModel[model]
+            return model
         elif model in cls.model_aliases:
             return cls.model_aliases[model]
         else:
@@ -168,9 +202,9 @@ class Blackbox:
     async def create_async_generator(
         cls,
         model: str,
-        messages: List[Dict[str, Any]],
+        messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
-        image: Optional[str] = None,  # Expecting a base64 string
+        image: Any = None,
         image_name: Optional[str] = None,
         webSearchMode: bool = False,
         **kwargs
@@ -200,39 +234,24 @@ class Blackbox:
             "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
         }
 
-        if model in cls.model_prefixes and messages:
+        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))
-        user_message = {
-            "id": random_id,
-            "role": 'user',
-            "content": 'Hi'  # This should be dynamically set based on input
-        }
+        messages[-1]['id'] = random_id
+        messages[-1]['role'] = 'user'
         if image is not None:
-            try:
-                image_bytes = decode_base64_image(image)
-                data_uri = to_data_uri(image_bytes)
-                user_message['data'] = {
-                    'fileText': '',
-                    'imageBase64': data_uri,
-                    'title': image_name or "Uploaded Image"
-                }
-                user_message['content'] = 'FILE:BB\n$#$\n\n$#$\n' + user_message['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.")
+            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.")
         
-        # Update the last message with user_message
-        if messages:
-            messages[-1] = user_message
-        else:
-            messages.append(user_message)
-
         data = {
             "messages": messages,
             "id": random_id,
@@ -280,15 +299,7 @@ class Blackbox:
                             if url_match:
                                 image_url = url_match.group(0)
                                 logger.info(f"Image URL found: {image_url}")
-                                
-                                # 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'])
+                                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")
@@ -349,7 +360,6 @@ class ChatRequest(BaseModel):
     messages: List[Message]
     stream: Optional[bool] = False
     webSearchMode: Optional[bool] = False
-    image: Optional[str] = None  # Add image field for base64 data
 
 def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
@@ -367,32 +377,25 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,
     }
 
-@app.post("/niansuhai/v1/chat/completions")
-async def chat_completions(request: ChatRequest, req: Request):
-    logger.info(f"Received chat completions request: model='{request.model}' messages={request.messages} stream={request.stream} webSearchMode={request.webSearchMode} image={request.image}")
+@app.post("/niansuhai/v1/chat/completions", dependencies=[Depends(rate_limiter)])
+async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
+    """
+    Endpoint to handle chat completions.
+    Protected by API key and rate limiter.
+    """
+    logger.info(f"Received chat completions request from API key: {api_key} | Request: {request}")
     try:
-        # Validate that all messages have string content
-        for idx, msg in enumerate(request.messages):
-            if not isinstance(msg.content, str):
-                logger.error(f"Message at index {idx} has invalid content type: {type(msg.content)}")
-                raise HTTPException(
-                    status_code=422,
-                    detail=[{
-                        "type": "string_type",
-                        "loc": ["body", "messages", idx, "content"],
-                        "msg": "Input should be a valid string",
-                        "input": msg.content
-                    }]
-                )
-
-        # Convert Pydantic messages to dicts
+        # Validate that the requested model is available
+        if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
+            logger.warning(f"Attempt to use unavailable model: {request.model}")
+            raise HTTPException(status_code=400, detail="Requested model is not available.")
+
         messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
         
         async_generator = Blackbox.create_async_generator(
             model=request.model,
             messages=messages,
-            proxy=None,  # Pass proxy if needed
-            image=request.image,  # Pass the base64 image
+            image=None,
             image_name=None,
             webSearchMode=request.webSearchMode
         )
@@ -402,8 +405,7 @@ async def chat_completions(request: ChatRequest, req: Request):
                 try:
                     async for chunk in async_generator:
                         if isinstance(chunk, ImageResponse):
-                            # Use the base64 data URI directly
-                            image_markdown = f"![{chunk.alt}]({chunk.data_uri})"
+                            image_markdown = f"![image]({chunk.url})"
                             response_chunk = create_response(image_markdown, request.model)
                         else:
                             response_chunk = create_response(chunk, request.model)
@@ -426,11 +428,11 @@ async def chat_completions(request: ChatRequest, req: Request):
             response_content = ""
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponse):
-                    response_content += f"![{chunk.alt}]({chunk.data_uri})\n"
+                    response_content += f"![image]({chunk.url})\n"
                 else:
                     response_content += chunk
 
-            logger.info("Completed non-streaming response generation.")
+            logger.info(f"Completed non-streaming response generation for API key: {api_key}")
             return {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
@@ -462,26 +464,33 @@ async def chat_completions(request: ChatRequest, req: Request):
         logger.exception("An unexpected error occurred while processing the chat completions request.")
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.get("/niansuhai/v1/models")
-async def get_models():
-    logger.info("Fetching available models.")
+@app.get("/niansuhai/v1/models", dependencies=[Depends(rate_limiter)])
+async def get_models(api_key: str = Depends(get_api_key)):
+    """
+    Endpoint to fetch available models.
+    Protected by API key and rate limiter.
+    """
+    logger.info(f"Fetching available models for API key: {api_key}")
     return {"data": [{"id": model} for model in Blackbox.models]}
 
 # Additional endpoints for better functionality
-@app.get("/niansuhai/v1/health")
-async def health_check():
+@app.get("/niansuhai/v1/health", dependencies=[Depends(rate_limiter)])
+async def health_check(api_key: str = Depends(get_api_key)):
     """Health check endpoint to verify the service is running."""
+    logger.info(f"Health check requested by API key: {api_key}")
     return {"status": "ok"}
 
-@app.get("/niansuhai/v1/models/{model}/status")
-async def model_status(model: str):
+@app.get("/niansuhai/v1/models/{model}/status", dependencies=[Depends(rate_limiter)])
+async def model_status(model: str, api_key: str = Depends(get_api_key)):
     """Check if a specific model is available."""
+    logger.info(f"Model status requested for '{model}' by API key: {api_key}")
     if model in Blackbox.models:
         return {"model": model, "status": "available"}
     elif model in Blackbox.model_aliases:
         actual_model = Blackbox.model_aliases[model]
         return {"model": actual_model, "status": "available via alias"}
     else:
+        logger.warning(f"Model not found: {model}")
         raise HTTPException(status_code=404, detail="Model not found")
 
 if __name__ == "__main__":