Update main.py

main.py

@@ -1,3 +1,4 @@
+# main.py
 import os
 import re
 import random
@@ -7,16 +8,34 @@ import json
 import logging
 import asyncio
 import time
+import contextvars
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, AsyncGenerator, Union
 from datetime import datetime
-import contextvars  # Added for context variables
 
-from aiohttp import ClientSession, ClientTimeout, ClientError
+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, Header
 from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
 from pydantic import BaseModel
 
+from sqlalchemy.orm import Session
+
+from aiohttp import ClientSession, ClientTimeout, ClientError
+
+from database import SessionLocal, engine, get_db
+from models import Base, Image, Log
+
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Create all tables
+Base.metadata.create_all(bind=engine)
+
 # Define a context variable for client_ip
 client_ip_var = contextvars.ContextVar("client_ip", default="N/A")
 
@@ -26,77 +45,145 @@ class ContextFilter(logging.Filter):
         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)
+
 # Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s [IP: %(client_ip)s]: %(message)s",
-    handlers=[logging.StreamHandler()]
+logger = logging.getLogger("main")  # Use a specific logger name if needed
+logger.setLevel(logging.INFO)
+
+# Create handlers
+console_handler = logging.StreamHandler()
+console_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"
 )
+console_handler.setFormatter(formatter)
+
+# Add the custom filter to the console handler
+console_handler.addFilter(ContextFilter())
+
+# Add handlers to the logger
+logger.addHandler(console_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)
+
+from logging import Handler
+
+class DBLogHandler(Handler):
+    def __init__(self, db: Session):
+        super().__init__()
+        self.db = db
+
+    def emit(self, record):
+        log_entry = Log(
+            level=record.levelname,
+            message=record.getMessage(),
+            client_ip=getattr(record, 'client_ip', None)
+        )
+        try:
+            self.db.add(log_entry)
+            self.db.commit()
+        except Exception as e:
+            # Handle exceptions (e.g., rollback)
+            self.db.rollback()
+            print(f"Failed to log to database: {e}")
+
+# Dependency to add DBLogHandler
+async def get_db_log_handler(request: Request):
+    db = next(get_db())
+    db_log_handler = DBLogHandler(db)
+    logger.addHandler(db_log_handler)
+    try:
+        yield
+    finally:
+        logger.removeHandler(db_log_handler)
+        db.close()
+
+@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":
+        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={
+                    "error": {
+                        "message": "Content-Type must be application/json",
+                        "type": "invalid_request_error",
+                        "param": None,
+                        "code": None
+                    }
+                },
+            )
+    
+    # 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
+
+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[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_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
+    return {
+        "id": f"chatcmpl-{uuid.uuid4()}",
+        "object": "chat.completion.chunk",
+        "created": int(datetime.now().timestamp()),
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "delta": {"content": content, "role": "assistant"},
+                "finish_reason": finish_reason,
+            }
+        ],
+        "usage": None,
+    }
 
-# Get the root logger and add the ContextFilter
-logger = logging.getLogger(__name__)
-logger.addFilter(ContextFilter())
-
-# 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
-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 | NiansuhAI')
-    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 | NiansuhAI')
-    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} | NiansuhAI")
-            raise HTTPException(status_code=429, detail='Rate limit exceeded for API key | NiansuhAI')
-        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 | NiansuhAI')
-        ip_rate_limit_store[client_ip]["count"] += 1
-
-# 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
@@ -143,7 +230,9 @@ class Blackbox:
     ]
 
     # Filter models based on AVAILABLE_MODELS
+    AVAILABLE_MODELS = os.getenv("AVAILABLE_MODELS", "")
     if AVAILABLE_MODELS:
+        AVAILABLE_MODELS = [model.strip() for model in AVAILABLE_MODELS.split(',') if model.strip()]
         models = [model for model in models if model in AVAILABLE_MODELS]
 
     agentMode = {
@@ -383,80 +472,8 @@ class Blackbox:
                 if attempt == retry_attempts - 1:
                     raise HTTPException(status_code=500, detail=str(e))
 
-# FastAPI app setup
-app = FastAPI()
-
-# Middleware to enhance security and log client IP
-@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":
-        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={
-                    "error": {
-                        "message": "Content-Type must be application/json",
-                        "type": "invalid_request_error",
-                        "param": None,
-                        "code": None
-                    }
-                },
-            )
-    
-    # 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
-
-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[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_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
-    return {
-        "id": f"chatcmpl-{uuid.uuid4()}",
-        "object": "chat.completion.chunk",
-        "created": int(datetime.now().timestamp()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "delta": {"content": content, "role": "assistant"},
-                "finish_reason": finish_reason,
-            }
-        ],
-        "usage": None,
-    }
-
-@app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter)])
-async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
+@app.post("/v1/chat/completions", dependencies=[Depends(limiter.limit("60/minute")), Depends(get_db_log_handler)])
+async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key), db: Session = Depends(get_db)):
     # Redact user messages only for logging purposes
     redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
 
@@ -484,6 +501,14 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
                         if isinstance(chunk, ImageResponse):
                             image_markdown = f"![image]({chunk.url})"
                             response_chunk = create_response(image_markdown, request.model)
+                            
+                            # Store image in the database
+                            image_entry = Image(
+                                image_url=chunk.url,
+                                description=request.messages[-1].get('content', '')
+                            )
+                            db.add(image_entry)
+                            db.commit()
                         else:
                             response_chunk = create_response(chunk, request.model)
                         
@@ -504,6 +529,14 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponse):
                     response_content += f"![image]({chunk.url})\n"
+                    
+                    # Store image in the database
+                    image_entry = Image(
+                        image_url=chunk.url,
+                        description=request.messages[-1].get('content', '')
+                    )
+                    db.add(image_entry)
+                    db.commit()
                 else:
                     response_content += chunk
 
@@ -549,20 +582,17 @@ async def chat_completions_get():
 
 @app.get("/v1/models")
 async def get_models(req: Request):
-    client_ip = req.client.host
     logger.info(f"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):
-    client_ip = req.client.host
     logger.info(f"Health check requested")
     return {"status": "ok"}
 
 @app.get("/v1/models/{model}/status")
 async def model_status(model: str, req: Request):
-    client_ip = req.client.host
     logger.info(f"Model status requested for '{model}'")
     if model in Blackbox.models:
         return {"model": model, "status": "available"}
@@ -576,8 +606,7 @@ 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):
-    client_ip = request.client.host
-    logger.error(f"HTTPException: {exc.detail} | Client IP: {client_ip}")
+    logger.error(f"HTTPException: {exc.detail}")
     return JSONResponse(
         status_code=exc.status_code,
         content={
@@ -596,7 +625,6 @@ class TokenizerRequest(BaseModel):
 
 @app.post("/v1/tokenizer")
 async def tokenizer(request: TokenizerRequest, req: Request):
-    client_ip = req.client.host
     text = request.text
     token_count = len(text.split())
     logger.info(f"Tokenizer called | Tokens: {token_count}")