Update main.py

main.py

@@ -1,4 +1,3 @@
-# main.py
 import os
 import re
 import random
@@ -17,25 +16,13 @@ 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 import FastAPI, HTTPException, Request, Depends, Header, Security
 from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
+from fastapi.security.api_key import APIKeyHeader
 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")
 
@@ -57,21 +44,21 @@ 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)
+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"
 )
-console_handler.setFormatter(formatter)
+handler.setFormatter(formatter)
 
-# Add the custom filter to the console handler
-console_handler.addFilter(ContextFilter())
+# Add the custom filter to the handler
+handler.addFilter(ContextFilter())
 
 # Add handlers to the logger
-logger.addHandler(console_handler)
+logger.addHandler(handler)
 
 # Initialize the limiter with slowapi
 limiter = Limiter(key_func=get_remote_address, default_limits=["60/minute"])
@@ -84,106 +71,50 @@ 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)
+# 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.")
+
+# Define API key security
+API_KEY_NAME = "X-API-Key"
+api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
+
+async def get_api_key(api_key: str = Security(api_key_header)):
+    if not api_key:
+        logger.warning("API key missing")
+        raise HTTPException(
+            status_code=401,
+            detail="API key missing",
+            headers={"WWW-Authenticate": "API Key"},
         )
-        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,
-    }
-
+    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": "API Key"},
+        )
+    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
+
+# 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
@@ -230,9 +161,7 @@ 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 = {
@@ -472,8 +401,81 @@ class Blackbox:
                 if attempt == retry_attempts - 1:
                     raise HTTPException(status_code=500, detail=str(e))
 
-@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)):
+# FastAPI app setup with slowapi rate limiting
+@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(limiter.limit("60/minute"))])
+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]
 
@@ -501,14 +503,6 @@ 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)
                         
@@ -529,14 +523,6 @@ 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