Update main.py

main.py

@@ -1,17 +1,27 @@
+import os
+from dotenv import load_dotenv
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse, HTMLResponse
 from pydantic import BaseModel
 import requests
-import time
 from functools import lru_cache
 
+# Load environment variables from .env file
+load_dotenv()
+
 app = FastAPI()
 
 # Create a session for reusing the HTTP connection
 session = requests.Session()
 
-# Valid API keys
-valid_api_keys = ['PARTH-SADARIA-NI-API-CHAWI', 'HEET-NI-CHEESEWADI-API-KEY']
+# Get API keys and secret endpoint from environment variables
+api_keys_str = os.getenv('API_KEYS')
+valid_api_keys = api_keys_str.split(',') if api_keys_str else []
+secret_api_endpoint = os.getenv('SECRET_API_ENDPOINT')
+
+# Check if the endpoint is set in the environment
+if not secret_api_endpoint:
+    raise HTTPException(status_code=500, detail="Secret API endpoint is not configured in environment variables.")
 
 class Payload(BaseModel):
     model: str
@@ -27,9 +37,9 @@ async def root():
         <meta name="viewport" content="width=device-width, initial-scale=1.0">
         <title>Loki.AI API</title>
         <style>
-            body { font-family: Arial, sans-serif; text-align: center; margin-top: 50px; }
+            body { font-family: Arial, sans-serif; text-align: center; margin-top: 50px; background-color: #121212; color: white; }
             h1 { color: #4CAF50; }
-            a { color: #007BFF; text-decoration: none; }
+            a { color: #BB86FC; text-decoration: none; }
             a:hover { text-decoration: underline; }
         </style>
     </head>
@@ -47,7 +57,7 @@ async def root():
 @lru_cache(maxsize=1)
 def get_cached_models():
     try:
-        response = session.get("https://lokiai.vercel.app/api/v1/models", timeout=3)
+        response = session.get(f"{secret_api_endpoint}/api/v1/models", timeout=3)
         response.raise_for_status()
         return response.json()
     except requests.exceptions.RequestException as e:
@@ -65,16 +75,13 @@ async def get_completion(payload: Payload, request: Request):
     if api_key not in valid_api_keys:
         raise HTTPException(status_code=403, detail="Forbidden: Invalid API key. Join dsc.gg/chadgang and DM @mr_leaderyt on Discord for a free API key :)")
     
-    # Set the URL for all requests
-    url = "https://gpt.tiptopuni.com/api/openai/v1/chat/completions"
-    
     # Prepare the payload for streaming
     payload_dict = {**payload.dict(), "stream": True}
 
     # Define a generator to stream the response
     def stream_generator():
         try:
-            with session.post(url, json=payload_dict, stream=True, timeout=15) as response:
+            with session.post(secret_api_endpoint, json=payload_dict, stream=True, timeout=15) as response:
                 response.raise_for_status()
                 for chunk in response.iter_content(chunk_size=1024):
                     if chunk:  # Only yield non-empty chunks