Update main.py

main.py

@@ -1,3 +1,4 @@
+import os
 import re
 import random
 import string
@@ -5,12 +6,22 @@ import uuid
 import json
 import logging
 import asyncio
+import time
 from aiohttp import ClientSession, ClientTimeout, ClientError
-from fastapi import FastAPI, HTTPException, Request
-from pydantic import BaseModel
+from fastapi import FastAPI, HTTPException, Request, Depends, Header, status
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
 from typing import List, Dict, Any, Optional, AsyncGenerator
 from datetime import datetime
-from fastapi.responses import StreamingResponse
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+import tiktoken
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
 
 # Configure logging
 logging.basicConfig(
@@ -22,6 +33,57 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Initialize FastAPI app
+app = FastAPI(title="OpenAI-Compatible API")
+
+# Configure CORS (adjust origins as needed)
+origins = [
+    "*",  # Allow all origins; replace with specific origins in production
+]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize Rate Limiter from environment variable
+RATE_LIMIT = os.getenv("RATE_LIMIT", "60/minute")  # Default to 60 requests per minute
+limiter = Limiter(key_func=get_remote_address, default_limits=[RATE_LIMIT])
+app.state.limiter = limiter
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+
+# API Key Authentication
+API_KEYS = set(os.getenv("API_KEYS", "").split(","))  # Load API keys from environment variable
+
+async def get_api_key(authorization: Optional[str] = Header(None)):
+    """
+    Dependency to validate API Key from the Authorization header.
+    """
+    if authorization is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Authorization header missing",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    parts = authorization.split()
+    if parts[0].lower() != "bearer" or len(parts) != 2:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid authorization header format",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    token = parts[1]
+    if token not in API_KEYS:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid API Key",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    return token
+
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
     def __init__(self, model: str):
@@ -29,18 +91,34 @@ 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)
 
-# Mock implementations for ImageResponse and to_data_uri
+# Mock implementations for ImageResponse and to_data_uri (custom functionality)
 class ImageResponse:
     def __init__(self, url: str, alt: str):
         self.url = url
         self.alt = alt
 
 def to_data_uri(image: Any) -> str:
-    return "data:image/png;base64,..."  # Replace with actual base64 data
+    return "data:image/png;base64,..."  # Replace with actual base64 data if needed
+
+# Token Counting using tiktoken
+def count_tokens(messages: List[Dict[str, str]], model: str) -> int:
+    """
+    Counts the number of tokens in the messages using tiktoken.
+    Adjust the encoding based on the model.
+    """
+    try:
+        encoding = tiktoken.get_encoding("cl100k_base")  # Adjust encoding as per model
+    except:
+        encoding = tiktoken.get_encoding("cl100k_base")  # Default encoding
+    tokens = 0
+    for message in messages:
+        tokens += len(encoding.encode(message['content']))
+    return tokens
 
+# Blackbox Class: Handles interaction with the external AI service
 class Blackbox:
     url = "https://www.blackbox.ai"
-    api_endpoint = "https://www.blackbox.ai/api/chat"
+    api_endpoint = os.getenv("EXTERNAL_API_ENDPOINT", "https://www.blackbox.ai/api/chat")
     working = True
     supports_stream = True
     supports_system_message = True
@@ -100,13 +178,13 @@ class Blackbox:
         'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
         'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
     }
-    
+
     userSelectedModel = {
         "gpt-4o": "gpt-4o",
         "gemini-pro": "gemini-pro",
         'claude-sonnet-3.5': "claude-sonnet-3.5",
     }
-    
+
     model_prefixes = {
         'gpt-4o': '@GPT-4o',
         'gemini-pro': '@Gemini-PRO',
@@ -128,14 +206,14 @@ class Blackbox:
         'ImageGeneration': '@Image Generation',
         'Niansuh': '@Niansuh',
     }
-    
+
     model_referers = {
         "blackboxai": f"{url}/?model=blackboxai",
         "gpt-4o": f"{url}/?model=gpt-4o",
         "gemini-pro": f"{url}/?model=gemini-pro",
         "claude-sonnet-3.5": f"{url}/?model=claude-sonnet-3.5"
     }
-    
+
     model_aliases = {
         "gemini-flash": "gemini-1.5-flash",
         "claude-3.5-sonnet": "claude-sonnet-3.5",
@@ -218,7 +296,7 @@ class Blackbox:
             "trendingAgentMode": {},
             "isMicMode": False,
             "userSystemPrompt": None,
-            "maxTokens": 99999999,
+            "maxTokens": int(os.getenv("MAX_TOKENS", "4096")),
             "playgroundTopP": 0.9,
             "playgroundTemperature": 0.5,
             "isChromeExt": False,
@@ -294,29 +372,41 @@ class Blackbox:
             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
-app = FastAPI()
-
+# Pydantic Models
 class Message(BaseModel):
-    role: str
-    content: str
+    role: str = Field(..., description="The role of the message author.")
+    content: str = Field(..., description="The content of the message.")
 
 class ChatRequest(BaseModel):
+    model: str = Field(..., description="ID of the model to use.")
+    messages: List[Message] = Field(..., description="A list of messages comprising the conversation.")
+    stream: Optional[bool] = Field(False, description="Whether to stream the response.")
+    webSearchMode: Optional[bool] = Field(False, description="Whether to enable web search mode.")
+
+class ChatCompletionChoice(BaseModel):
+    index: int
+    delta: Dict[str, Any]
+    finish_reason: Optional[str] = None
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str
+    created: int
     model: str
-    messages: List[Message]
-    stream: Optional[bool] = False
-    webSearchMode: Optional[bool] = False
+    choices: List[ChatCompletionChoice]
+    usage: Optional[Dict[str, int]] = None
 
+# Utility Function to Create Response
 def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
     return {
         "id": f"chatcmpl-{uuid.uuid4()}",
@@ -330,19 +420,22 @@ def create_response(content: str, model: str, finish_reason: Optional[str] = Non
                 "finish_reason": finish_reason,
             }
         ],
-        "usage": None,
+        "usage": None,  # To be populated if usage metrics are available
     }
 
-@app.post("/v1/chat/completions")
-async def chat_completions(request: ChatRequest, req: Request):
+# Endpoint: Chat Completions
+@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+@limiter.limit("60/minute")  # Example: 60 requests per minute per IP
+async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
     logger.info(f"Received chat completions request: {request}")
     try:
         messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
-        
+        prompt_tokens = count_tokens(messages, request.model)
+
         async_generator = Blackbox.create_async_generator(
             model=request.model,
             messages=messages,
-            image=None,
+            image=None,  # Adjust if image handling is required
             image_name=None,
             webSearchMode=request.webSearchMode
         )
@@ -350,15 +443,17 @@ async def chat_completions(request: ChatRequest, req: Request):
         if request.stream:
             async def generate():
                 try:
+                    completion_tokens = 0
                     async for chunk in async_generator:
                         if isinstance(chunk, ImageResponse):
                             image_markdown = f"![image]({chunk.url})"
                             response_chunk = create_response(image_markdown, request.model)
+                            yield f"data: {json.dumps(response_chunk)}\n\n"
+                            completion_tokens += len(image_markdown.split())
                         else:
                             response_chunk = create_response(chunk, request.model)
-                        
-                        # Yield each chunk in SSE format
-                        yield f"data: {json.dumps(response_chunk)}\n\n"
+                            yield f"data: {json.dumps(response_chunk)}\n\n"
+                            completion_tokens += len(chunk.split())
                     
                     # Signal the end of the stream
                     yield "data: [DONE]\n\n"
@@ -373,34 +468,36 @@ async def chat_completions(request: ChatRequest, req: Request):
             return StreamingResponse(generate(), media_type="text/event-stream")
         else:
             response_content = ""
+            completion_tokens = 0
             async for chunk in async_generator:
                 if isinstance(chunk, ImageResponse):
                     response_content += f"![image]({chunk.url})\n"
+                    completion_tokens += len(f"![image]({chunk.url})\n".split())
                 else:
                     response_content += chunk
+                    completion_tokens += len(chunk.split())
+
+            total_tokens = prompt_tokens + completion_tokens
 
             logger.info("Completed non-streaming response generation.")
-            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
-                    }
+            return ChatCompletionResponse(
+                id=f"chatcmpl-{uuid.uuid4()}",
+                object="chat.completion",
+                created=int(datetime.now().timestamp()),
+                model=request.model,
+                choices=[
+                    ChatCompletionChoice(
+                        index=0,
+                        delta={"content": response_content, "role": "assistant"},
+                        finish_reason="stop"
+                    )
                 ],
-                "usage": {
-                    "prompt_tokens": sum(len(msg['content'].split()) for msg in messages),
-                    "completion_tokens": len(response_content.split()),
-                    "total_tokens": sum(len(msg['content'].split()) for msg in messages) + len(response_content.split())
-                },
-            }
+                usage={
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": total_tokens
+                }
+            )
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e}")
         raise HTTPException(status_code=503, detail=str(e))
@@ -411,19 +508,17 @@ 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))
 
+# Endpoint: List Models
 @app.get("/v1/models")
-async def get_models():
+@limiter.limit("60/minute")
+async def get_models(api_key: str = Depends(get_api_key)):
     logger.info("Fetching available models.")
     return {"data": [{"id": model} for model in Blackbox.models]}
 
-# Additional endpoints for better functionality
-@app.get("/v1/health")
-async def health_check():
-    """Health check endpoint to verify the service is running."""
-    return {"status": "ok"}
-
+# Endpoint: Model Status
 @app.get("/v1/models/{model}/status")
-async def model_status(model: str):
+@limiter.limit("60/minute")
+async def model_status(model: str, api_key: str = Depends(get_api_key)):
     """Check if a specific model is available."""
     if model in Blackbox.models:
         return {"model": model, "status": "available"}
@@ -433,6 +528,14 @@ async def model_status(model: str):
     else:
         raise HTTPException(status_code=404, detail="Model not found")
 
+# Endpoint: Health Check
+@app.get("/v1/health")
+@limiter.limit("60/minute")
+async def health_check():
+    """Health check endpoint to verify the service is running."""
+    return {"status": "ok"}
+
+# Run the application
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)