Update main.py

main.py

@@ -1,18 +1,19 @@
 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
-
+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
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 
 # Configure logging
@@ -25,26 +26,48 @@ logger = logging.getLogger(__name__)
 
 # Load environment variables
 API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
-RATE_LIMIT_PER_MINUTE = int(os.getenv('RATE_LIMIT_PER_MINUTE', '60'))  # Requests per minute per IP
+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 per IP
-rate_limit_store_ip = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
+# Simple in-memory rate limiter
+rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
+ip_rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
+
+async def get_api_key(authorization: str = Header(...)) -> str:
+    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(request: Request):
-    client_host = request.client.host
+async def rate_limiter(req: Request, api_key: str = Depends(get_api_key)):
     current_time = time.time()
-    window_start = rate_limit_store_ip[client_host]["timestamp"]
+    # Rate limiting per API key
+    window_start = rate_limit_store[api_key]["timestamp"]
     if current_time - window_start > 60:
-        rate_limit_store_ip[client_host] = {"count": 1, "timestamp": current_time}
+        rate_limit_store[api_key] = {"count": 1, "timestamp": current_time}
     else:
-        if rate_limit_store_ip[client_host]["count"] >= RATE_LIMIT_PER_MINUTE:
-            logger.warning(f"Rate limit exceeded for IP: {client_host}")
-            raise HTTPException(status_code=429, detail='Rate limit exceeded.')
-        rate_limit_store_ip[client_host]["count"] += 1
+        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 for API key')
+        rate_limit_store[api_key]["count"] += 1
+
+    # Rate limiting per IP address
+    client_ip = req.client.host
+    window_start_ip = ip_rate_limit_store[client_ip]["timestamp"]
+    if current_time - window_start_ip > 60:
+        ip_rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
+    else:
+        if ip_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')
+        ip_rate_limit_store[client_ip]["count"] += 1
 
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
@@ -195,7 +218,7 @@ class Blackbox:
         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",
@@ -211,7 +234,7 @@ class Blackbox:
             "sec-fetch-dest": "empty",
             "sec-fetch-mode": "cors",
             "sec-fetch-site": "same-origin",
-            "user-agent": "Mozilla/5.0 (X11; Linux x86_64)",
+            "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:
@@ -219,7 +242,7 @@ class Blackbox:
             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'
@@ -235,7 +258,7 @@ class Blackbox:
             }
             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,
@@ -268,8 +291,8 @@ class Blackbox:
             data["userSelectedModel"] = cls.userSelectedModel[model]
         logger.info(f"Sending request to {cls.api_endpoint} with data (excluding messages).")
 
-        timeout = ClientTimeout(total=30)  # Reduced timeout for faster response
-        retry_attempts = 3  # Reduced retry attempts for faster failure handling
+        timeout = ClientTimeout(total=60)  # Set an appropriate timeout
+        retry_attempts = 10  # Set the number of retry attempts
 
         for attempt in range(retry_attempts):
             try:
@@ -335,204 +358,162 @@ class Blackbox:
 # FastAPI app setup
 app = FastAPI()
 
-# Implement per-IP rate limiting middleware
-@app.middleware("http")
-async def rate_limit_middleware(request: Request, call_next):
-    await rate_limiter(request)
-    response = await call_next(request)
-    return response
-
-# Pydantic models for OpenAI API
 class Message(BaseModel):
     role: str
     content: str
 
-class ChatCompletionRequest(BaseModel):
+class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
     temperature: Optional[float] = 1.0
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
     stream: Optional[bool] = False
-    stop: Optional[Any] = None  # Can be a string or list of strings
+    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
 
-def create_chat_completion_response(content: str, model: str, usage: Dict[str, int]) -> Dict[str, Any]:
+def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
-        "object": "chat.completion",
+        "object": "chat.completion.chunk",
         "created": int(datetime.now().timestamp()),
         "model": model,
         "choices": [
             {
                 "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": content
-                },
-                "finish_reason": "stop"
+                "delta": {"content": content, "role": "assistant"},
+                "finish_reason": finish_reason,
             }
         ],
-        "usage": usage
+        "usage": None,
     }
 
-def create_stream_response_chunk(content: str, role: Optional[str] = None, finish_reason: Optional[str] = None):
-    delta = {}
-    if role:
-        delta['role'] = role
-    if content:
-        delta['content'] = content
-    return {
-        "object": "chat.completion.chunk",
-        "created": int(datetime.now().timestamp()),
-        "model": "",  # Model name can be added if necessary
-        "choices": [
-            {
-                "delta": delta,
-                "index": 0,
-                "finish_reason": finish_reason
-            }
-        ]
-    }
+@app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter)])
+async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
+    # Redact user messages only for logging purposes
+    redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
 
-@app.post("/v1/chat/completions")
-async def chat_completions(request: ChatCompletionRequest, authorization: str = Header(None)):
-    # Verify API key
-    if not authorization or 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.')
-
-    logger.info(f"Received chat completion request for model: {request.model}")
+    logger.info(f"Received chat completions request from API key: {api_key} | Model: {request.model} | Messages: {redacted_messages}")
 
-    # Validate model
-    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="The model is not available.")
-
-    # Process the request
     try:
-        # Convert messages to dicts
-        messages = [msg.dict() for msg in request.messages]
-
-        # Check if the user is requesting image generation
-        image_generation_requested = any(
-            re.search(r'\b(generate|create|draw)\b.*\b(image|picture|art)\b', msg['content'], re.IGNORECASE)
-            for msg in messages if msg['role'] == 'user'
+        # 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.")
+
+        # Process the request with actual message content, but don't log it
+        async_generator = Blackbox.create_async_generator(
+            model=request.model,
+            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
+            image=None,
+            image_name=None,
+            webSearchMode=request.webSearchMode
         )
 
-        if image_generation_requested:
-            model = 'ImageGeneration'
-            # For image generation, use the last message as prompt
-            prompt = messages[-1]['content']
-            # Build messages for the Blackbox.create_async_generator
-            messages = [{"role": "user", "content": prompt}]
-            async_generator = Blackbox.create_async_generator(
-                model=model,
-                messages=messages,
-                image=None,
-                image_name=None,
-                webSearchMode=False
-            )
-
-            # Collect images
-            images = []
-            count = 0
-            async for response in async_generator:
-                if isinstance(response, ImageResponse):
-                    images.append(response.url)
-                    count += 1
-                    if count >= request.n:
-                        break
-
-            # Build response content with image URLs
-            response_content = "\n".join(f"![Generated Image]({url})" for url in images)
-            completion_tokens = len(response_content.split())
-        else:
-            # Use the requested model
-            async_generator = Blackbox.create_async_generator(
-                model=request.model,
-                messages=messages,
-                image=None,
-                image_name=None,
-                webSearchMode=False
-            )
-            # Usage tracking
-            completion_tokens = 0  # Will be updated as we process the response
-
-        prompt_tokens = sum(len(msg['content'].split()) for msg in messages)
-
         if request.stream:
             async def generate():
-                nonlocal completion_tokens
                 try:
-                    # Initial delta with role
-                    initial_chunk = create_stream_response_chunk(content=None, role="assistant")
-                    yield f"data: {json.dumps(initial_chunk)}\n\n"
-
                     async for chunk in async_generator:
-                        if isinstance(chunk, str):
-                            completion_tokens += len(chunk.split())
-                            response_chunk = create_stream_response_chunk(content=chunk)
-                            yield f"data: {json.dumps(response_chunk)}\n\n"
-                        elif isinstance(chunk, ImageResponse):
-                            content = f"![Generated Image]({chunk.url})"
-                            completion_tokens += len(content.split())
-                            response_chunk = create_stream_response_chunk(content=content)
-                            yield f"data: {json.dumps(response_chunk)}\n\n"
+                        if isinstance(chunk, ImageResponse):
+                            image_markdown = f"![image]({chunk.url})"
+                            response_chunk = create_response(image_markdown, request.model)
                         else:
-                            pass  # Handle other types if necessary
-                    # Finish reason
-                    final_chunk = create_stream_response_chunk(content=None, finish_reason="stop")
-                    yield f"data: {json.dumps(final_chunk)}\n\n"
+                            response_chunk = create_response(chunk, request.model)
+                        
+                        yield f"data: {json.dumps(response_chunk)}\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("Error during streaming response generation.")
-                    yield f"data: {json.dumps({'error': str(e)})}\n\n"
+                    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, str):
+                if isinstance(chunk, ImageResponse):
+                    response_content += f"![image]({chunk.url})\n"
+                else:
                     response_content += chunk
-                elif isinstance(chunk, ImageResponse):
-                    response_content += f"![Generated Image]({chunk.url})\n"
-            completion_tokens = len(response_content.split())
-            usage = {
-                "prompt_tokens": prompt_tokens,
-                "completion_tokens": completion_tokens,
-                "total_tokens": prompt_tokens + completion_tokens
+
+            logger.info(f"Completed non-streaming response generation for API key: {api_key}")
+            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": 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())
+                },
             }
-            return create_chat_completion_response(response_content, request.model, usage)
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e}")
         raise HTTPException(status_code=503, detail=str(e))
+    except HTTPException as he:
+        logger.warning(f"HTTPException: {he.detail}")
+        raise he
     except Exception as e:
         logger.exception("An unexpected error occurred while processing the chat completions request.")
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.get("/v1/models")
-async def get_models(authorization: str = Header(None)):
-    # Verify API key
-    if not authorization or 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.')
-
-    logger.info("Fetching available models.")
-    # Return models in OpenAI format
-    models_data = [{"id": model, "object": "model", "owned_by": "organization-owner", "permission": []} for model in Blackbox.models]
-    return {"data": models_data, "object": "list"}
+@app.get("/v1/models", dependencies=[Depends(rate_limiter)])
+async def get_models(api_key: str = Depends(get_api_key)):
+    logger.info(f"Fetching available models for API key: {api_key}")
+    return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
+
+# Additional endpoints for better functionality
+@app.get("/v1/health", dependencies=[Depends(rate_limiter)])
+async def health_check(api_key: str = Depends(get_api_key)):
+    logger.info(f"Health check requested by API key: {api_key}")
+    return {"status": "ok"}
+
+@app.get("/v1/models/{model}/status", dependencies=[Depends(rate_limiter)])
+async def model_status(model: str, api_key: str = Depends(get_api_key)):
+    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")
+
+# Custom exception handler to match OpenAI's error format
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request: Request, exc: HTTPException):
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={
+            "error": {
+                "message": exc.detail,
+                "type": "invalid_request_error",
+                "param": None,
+                "code": None
+            }
+        },
+    )
 
 if __name__ == "__main__":
     import uvicorn