Update main.py

main.py

@@ -17,7 +17,44 @@ import time
 from usage_tracker import UsageTracker
 from starlette.middleware.base import BaseHTTPMiddleware
 from collections import defaultdict
+from fastapi import Security #new
+from fastapi.security import APIKeyHeader
+from starlette.exceptions import HTTPException
+from starlette.status import HTTP_403_FORBIDDEN
 
+# API key header scheme
+api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
+
+# Function to validate API key
+async def verify_api_key(api_key: str = Security(api_key_header)) -> bool:
+    if not api_key:
+        raise HTTPException(
+            status_code=HTTP_403_FORBIDDEN,
+            detail="No API key provided"
+        )
+        
+    # Clean the API key by removing 'Bearer ' if present
+    if api_key.startswith('Bearer '):
+        api_key = api_key[7:]  # Remove 'Bearer ' prefix
+        
+    # Get API keys from environment
+    api_keys_str = os.getenv('API_KEYS')
+    if not api_keys_str:
+        raise HTTPException(
+            status_code=HTTP_403_FORBIDDEN,
+            detail="API keys not configured on server"
+        )
+        
+    valid_api_keys = api_keys_str.split(',')
+    
+    # Check if the provided key is valid
+    if api_key not in valid_api_keys:
+        raise HTTPException(
+            status_code=HTTP_403_FORBIDDEN,
+            detail="Invalid API key"
+        )
+        
+    return True
 class RateLimitMiddleware(BaseHTTPMiddleware):
     def __init__(self, app, requests_per_second: int = 2):
         super().__init__(app)
@@ -62,7 +99,6 @@ app = FastAPI()
 app.add_middleware(RateLimitMiddleware, requests_per_second=2)
 
 # Get API keys and secret endpoint from environment variables
-api_keys_str = os.getenv('API_KEYS') #deprecated -_-
 valid_api_keys = api_keys_str.split(',') if api_keys_str else []
 secret_api_endpoint = os.getenv('SECRET_API_ENDPOINT')
 secret_api_endpoint_2 = os.getenv('SECRET_API_ENDPOINT_2')
@@ -75,7 +111,7 @@ if not secret_api_endpoint or not secret_api_endpoint_2 or not secret_api_endpoi
     raise HTTPException(status_code=500, detail="API endpoint(s) are not configured in environment variables.")
 
 # Define models that should use the secondary endpoint
-alternate_models = {"gpt-4o-mini", "claude-3-haiku", "llama-3.1-70b", "mixtral-8x7b"}
+# alternate_models = {"gpt-4o-mini", "claude-3-haiku", "llama-3.1-70b", "mixtral-8x7b"}
 
 available_model_ids = []
 class Payload(BaseModel):
@@ -154,7 +190,7 @@ async def ping():
     return {"message": "pong", "response_time": f"{response_time:.6f} seconds"}
     
 @app.get("/searchgpt")
-async def search_gpt(q: str, stream: Optional[bool] = False, systemprompt: Optional[str] = None):
+async def search_gpt(q: str, stream: Optional[bool] = False, systemprompt: Optional[str] = None,authenticated: bool = Depends(verify_api_key)):
     if not q:
         raise HTTPException(status_code=400, detail="Query parameter 'q' is required")
     usage_tracker.record_request(endpoint="/searchgpt")
@@ -191,12 +227,12 @@ async def get_models():
         raise HTTPException(status_code=500, detail="Error decoding models.json")
 @app.get("api/v1/models")
 @app.get("/models")
-async def fetch_models():
+async def fetch_models(authenticated: bool = Depends(verify_api_key)):
     return await get_models()
 server_status = True
 @app.post("/chat/completions")
 @app.post("api/v1/chat/completions")
-async def get_completion(payload: Payload, request: Request):
+async def get_completion(payload: Payload, request: Request,authenticated: bool = Depends(verify_api_key)):
     # Check server status
     
     
@@ -216,7 +252,7 @@ async def get_completion(payload: Payload, request: Request):
     payload_dict["model"] = model_to_use
     # payload_dict["stream"] = payload_dict.get("stream", False)
     # Select the appropriate endpoint
-    endpoint = secret_api_endpoint_2 if model_to_use in alternate_models else secret_api_endpoint
+    endpoint = secret_api_endpoint
     
     # Current time and IP logging
     current_time = (datetime.datetime.utcnow() + datetime.timedelta(hours=5, minutes=30)).strftime("%Y-%m-%d %I:%M:%S %p")
@@ -286,6 +322,7 @@ async def generate_image(
     private: Optional[bool] = None,
     enhance: Optional[bool] = None,
     request: Request = None,  # Access raw POST data
+    authenticated: bool = Depends(verify_api_key)
 ):
     """
     Generate an image using the Image Generation API.