Update main.py

main.py

@@ -6,68 +6,23 @@ import uuid
 import json
 import logging
 import asyncio
-import contextvars
+import time
+from collections import defaultdict
 from typing import List, Dict, Any, Optional, AsyncGenerator, Union
 from datetime import datetime
 
-from slowapi import Limiter, _rate_limit_exceeded_handler
-from slowapi.util import get_remote_address
-from slowapi.errors import RateLimitExceeded
-from slowapi.middleware import SlowAPIMiddleware
-from fastapi import FastAPI, HTTPException, Request, Depends, Security
-from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-from pydantic import BaseModel, Field
-
 from aiohttp import ClientSession, ClientTimeout, ClientError
-
-# Define a context variable for client_ip
-client_ip_var = contextvars.ContextVar("client_ip", default="N/A")
-
-# Custom logging filter to inject client_ip from context variable
-class ContextFilter(logging.Filter):
-    def filter(self, record):
-        record.client_ip = client_ip_var.get()
-        return True
-
-# Custom logging formatter to handle missing client_ip
-class SafeFormatter(logging.Formatter):
-    def format(self, record):
-        if not hasattr(record, 'client_ip'):
-            record.client_ip = 'N/A'
-        return super().format(record)
+from fastapi import FastAPI, HTTPException, Request, Depends, Header
+from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
+from pydantic import BaseModel
 
 # Configure logging
-logger = logging.getLogger("main")  # Use a specific logger name if needed
-logger.setLevel(logging.INFO)
-
-# Create handlers
-handler = logging.StreamHandler()
-handler.setLevel(logging.INFO)
-
-# Create and set the custom formatter
-formatter = SafeFormatter(
-    fmt="%(asctime)s [%(levelname)s] %(name)s [IP: %(client_ip)s]: %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S"
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    handlers=[logging.StreamHandler()]
 )
-handler.setFormatter(formatter)
-
-# Add the custom filter to the handler
-handler.addFilter(ContextFilter())
-
-# Add handlers to the logger
-logger.addHandler(handler)
-
-# Initialize the limiter with slowapi
-limiter = Limiter(key_func=get_remote_address, default_limits=["60/minute"])
-app = FastAPI()
-
-# Register the rate limit exceeded handler
-app.state.limiter = limiter
-app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
-
-# Add SlowAPI middleware
-app.add_middleware(SlowAPIMiddleware)
+logger = logging.getLogger(__name__)
 
 # Load environment variables
 API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
@@ -78,51 +33,49 @@ 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.")
 
-# Define API key security using HTTPBearer (Bearer token)
-security = HTTPBearer()
-
-async def get_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
-    """
-    Dependency to extract and validate the API key from the Authorization header.
-    Expects the header in the format: Authorization: Bearer YOUR_API_KEY
-    """
-    if not credentials:
-        logger.warning("Authorization header missing")
-        raise HTTPException(
-            status_code=401,
-            detail="Authorization header missing",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    if credentials.scheme.lower() != "bearer":
-        logger.warning(f"Invalid authentication scheme: {credentials.scheme}")
-        raise HTTPException(
-            status_code=401,
-            detail="Invalid authentication scheme. Expected 'Bearer'.",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    api_key = credentials.credentials
-    if not api_key:
-        logger.warning("API key missing in Authorization header")
-        raise HTTPException(
-            status_code=401,
-            detail="API key missing",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    if api_key not in API_KEYS:
-        logger.warning(f"Invalid API key: {api_key}")
-        raise HTTPException(
-            status_code=403,
-            detail="Invalid API key",
-            headers={"WWW-Authenticate": "Bearer"},
-        )
-    return api_key
-
 # 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
+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(None)) -> str:
+    if authorization is None or not authorization.startswith('Bearer '):
+        logger.warning("Invalid or missing authorization header.")
+        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(req: Request, api_key: str = Depends(get_api_key)):
+    current_time = time.time()
+    # Rate limiting per API key
+    window_start = rate_limit_store[api_key]["timestamp"]
+    if current_time - window_start > 60:
+        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 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):
     def __init__(self, model: str):
@@ -272,13 +225,13 @@ class Blackbox:
     ) -> AsyncGenerator[Any, None]:
         model = cls.get_model(model)
         if model is None:
-            logger.error(f"Model {model} is not available. | NiansuhAI")
+            logger.error(f"Model {model} is not available.")
             raise ModelNotWorkingException(model)
 
         logger.info(f"Selected model: {model}")
 
         if not cls.working or model not in cls.models:
-            logger.error(f"Model {model} is not working or not supported. | NiansuhAI")
+            logger.error(f"Model {model} is not working or not supported.")
             raise ModelNotWorkingException(model)
         
         headers = {
@@ -290,13 +243,13 @@ class Blackbox:
             "pragma": "no-cache",
             "priority": "u=1, i",
             "referer": cls.model_referers.get(model, cls.url),
-            "sec-ch-ua": '"Chromium";v="130", "Not=A?Brand";v="99"',
+            "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/130.0.0.0 Safari/537.36",
+            "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:
@@ -371,7 +324,7 @@ class Blackbox:
                                 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 | NiansuhAI")
+                                raise Exception("Image URL not found in the response")
                         else:
                             full_response = ""
                             search_results_json = ""
@@ -403,31 +356,30 @@ class Blackbox:
                                     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. | NiansuhAI")
+                    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. | NiansuhAI")
+                    raise HTTPException(status_code=504, detail="External API request timed out.")
             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))
 
-# FastAPI app setup with slowapi rate limiting
+# FastAPI app setup
+app = FastAPI()
+
+# Middleware to enhance security
 @app.middleware("http")
 async def security_middleware(request: Request, call_next):
-    client_ip = request.client.host
-    # Set the client_ip in the context variable
-    client_ip_var.set(client_ip)
-    
     # Enforce that POST requests to sensitive endpoints must have a valid Content-Type
-    if request.method == "POST" and request.url.path == "/v1/chat/completions":
+    if request.method == "POST" and request.url.path in ["/v1/chat/completions", "/v1/completions"]:
         content_type = request.headers.get("Content-Type")
         if content_type != "application/json":
-            logger.warning("Invalid Content-Type for /v1/chat/completions")
             return JSONResponse(
                 status_code=400,
                 content={
@@ -439,33 +391,24 @@ async def security_middleware(request: Request, call_next):
                     }
                 },
             )
-    
-    # Log the incoming request
-    logger.info(f"Incoming request: {request.method} {request.url.path}")
-    
     response = await call_next(request)
-    
-    # Log the response status
-    logger.info(f"Response status: {response.status_code}")
-    
     return response
 
 class Message(BaseModel):
     role: str
     content: str
-    name: Optional[str] = None  # Optional field as per OpenAI's API
 
 class ChatRequest(BaseModel):
     model: str
     messages: List[Message]
-    temperature: Optional[float] = Field(1.0, ge=0.0, le=2.0)
-    top_p: Optional[float] = Field(1.0, ge=0.0, le=1.0)
-    n: Optional[int] = Field(1, ge=1, le=10)
+    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] = Field(None, ge=1, le=4096)
-    presence_penalty: Optional[float] = Field(0.0, ge=-2.0, le=2.0)
-    frequency_penalty: Optional[float] = Field(0.0, ge=-2.0, le=2.0)
+    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
@@ -486,15 +429,8 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
         "usage": None,
     }
 
-@app.post("/v1/chat/completions", dependencies=[Depends(limiter.limit("60/minute"))])
-async def chat_completions(
-    request: ChatRequest,
-    req: Request,
-    api_key: str = Depends(get_api_key)
-):
-    """
-    Handles the /v1/chat/completions endpoint, emulating OpenAI's API.
-    """
+@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]
 
@@ -504,7 +440,7 @@ async def chat_completions(
         # 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. | NiansuhAI")
+            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(
@@ -545,7 +481,7 @@ async def chat_completions(
                 else:
                     response_content += chunk
 
-            logger.info(f"Completed non-streaming response generation for API key: {api_key} | Model: {request.model}")
+            logger.info(f"Completed non-streaming response generation for API key: {api_key}")
             return {
                 "id": f"chatcmpl-{uuid.uuid4()}",
                 "object": "chat.completion",
@@ -584,18 +520,18 @@ async def chat_completions_get():
     return RedirectResponse(url='about:blank')
 
 @app.get("/v1/models")
-async def get_models(req: Request):
-    logger.info(f"Fetching available models")
+async def get_models():
+    logger.info("Fetching available models")
     return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
 
 # Additional endpoints for better functionality
-@app.get("/v1/health")
-async def health_check(req: Request):
-    logger.info(f"Health check requested")
+@app.get("/v1/health", dependencies=[Depends(rate_limiter)])
+async def health_check(req: Request, 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")
-async def model_status(model: str, req: Request):
+async def model_status(model: str):
     logger.info(f"Model status requested for '{model}'")
     if model in Blackbox.models:
         return {"model": model, "status": "available"}
@@ -609,7 +545,6 @@ async def model_status(model: str, req: Request):
 # Custom exception handler to match OpenAI's error format
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request: Request, exc: HTTPException):
-    logger.error(f"HTTPException: {exc.detail}")
     return JSONResponse(
         status_code=exc.status_code,
         content={
@@ -627,12 +562,68 @@ class TokenizerRequest(BaseModel):
     text: str
 
 @app.post("/v1/tokenizer")
-async def tokenizer(request: TokenizerRequest, req: Request):
+async def tokenizer(request: TokenizerRequest):
     text = request.text
     token_count = len(text.split())
-    logger.info(f"Tokenizer called | Tokens: {token_count}")
     return {"text": text, "tokens": token_count}
 
+# New endpoint: /v1/completions to support text completions
+class CompletionRequest(BaseModel):
+    model: str
+    prompt: str
+    max_tokens: Optional[int] = 16
+    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
+    logprobs: Optional[int] = None
+    echo: Optional[bool] = False
+    presence_penalty: Optional[float] = 0.0
+    frequency_penalty: Optional[float] = 0.0
+    best_of: Optional[int] = 1
+    logit_bias: Optional[Dict[str, float]] = None
+    user: Optional[str] = None
+
+@app.post("/v1/completions", dependencies=[Depends(rate_limiter)])
+async def completions(request: CompletionRequest, req: Request, api_key: str = Depends(get_api_key)):
+    logger.info(f"Received completion request from API key: {api_key} | Model: {request.model}")
+
+    try:
+        # 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.")
+
+        # Simulate a simple completion by echoing the prompt
+        completion_text = f"{request.prompt} [Completed by {request.model}]"
+
+        return {
+            "id": f"cmpl-{uuid.uuid4()}",
+            "object": "text_completion",
+            "created": int(datetime.now().timestamp()),
+            "model": request.model,
+            "choices": [
+                {
+                    "text": completion_text,
+                    "index": 0,
+                    "logprobs": None,
+                    "finish_reason": "length"
+                }
+            ],
+            "usage": {
+                "prompt_tokens": len(request.prompt.split()),
+                "completion_tokens": len(completion_text.split()),
+                "total_tokens": len(request.prompt.split()) + len(completion_text.split())
+            }
+        }
+    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 completions request.")
+        raise HTTPException(status_code=500, detail=str(e))
+
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)