Update main.py

main.py

@@ -1,7 +1,7 @@
 import os
 import re
 from dotenv import load_dotenv
-from fastapi import FastAPI, HTTPException, Request, Depends, Security, APIRouter
+from fastapi import FastAPI, HTTPException, Request, Depends, Security
 from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse, FileResponse
 from fastapi.security import APIKeyHeader
 from pydantic import BaseModel
@@ -12,7 +12,7 @@ import json
 import datetime
 import time
 import threading
-from typing import Optional, Dict, List, Any, Generator, Set # Import Set
+from typing import Optional, Dict, List, Any, Generator
 import asyncio
 from starlette.status import HTTP_403_FORBIDDEN
 import cloudscraper
@@ -66,10 +66,7 @@ def get_env_vars():
     }
 
 # Configuration for models - use sets for faster lookups
-# IMPORTANT: These will be updated in memory. For persistence,
-# you would need to save these changes to a file (like models.json)
-# or a database and reload them on startup.
-mistral_models: Set[str] = {
+mistral_models = {
     "mistral-large-latest",
     "pixtral-large-latest",
     "mistral-moderation-latest",
@@ -81,7 +78,7 @@ mistral_models: Set[str] = {
     "codestral-latest"
 }
 
-pollinations_models: Set[str] = {
+pollinations_models = {
     "openai",
     "openai-large",
     "openai-xlarge",
@@ -103,7 +100,7 @@ pollinations_models: Set[str] = {
     "openai-audio",
     "llama-scaleway"
 }
-alternate_models: Set[str] = {
+alternate_models = {
     "o1",
     "llama-4-scout",
     "o4-mini",
@@ -117,7 +114,7 @@ alternate_models: Set[str] = {
     "o3"
 }
 
-claude_3_models: Set[str] = {  # Models for the new endpoint
+claude_3_models = {  # Models for the new endpoint
     "claude-3-7-sonnet",
     "claude-3-7-sonnet-thinking",
     "claude 3.5 haiku",
@@ -131,7 +128,7 @@ claude_3_models: Set[str] = {  # Models for the new endpoint
 }
 
 # Supported image generation models
-supported_image_models: Set[str] = {
+supported_image_models = {
     "Flux Pro Ultra",
     "grok-2-aurora",
     "Flux Pro",
@@ -161,14 +158,10 @@ class ImageGenerationPayload(BaseModel):
     number: int
 
 
-# Pydantic model for updating models via admin endpoint
-class UpdateModelsPayload(BaseModel):
-    provider: str # e.g., "mistral", "pollinations", "alternate", "claude_3", "image"
-    models: List[str] # The new list of model IDs for the provider
 
 # Server status global variable
 server_status = True
-available_model_ids: List[str] = [] # This will be updated based on the sets
+available_model_ids: List[str] = []
 
 # Create a reusable httpx client pool with connection pooling
 @lru_cache(maxsize=1)
@@ -197,20 +190,20 @@ async def verify_api_key(
 ) -> bool:
     # Allow bypass if the referer is from /playground or /image-playground
     referer = request.headers.get("referer", "")
-    if referer.startswith(("https://parthsadaria-lokiai.hf.space/playground",
+    if referer.startswith(("https://parthsadaria-lokiai.hf.space/playground", 
                            "https://parthsadaria-lokiai.hf.space/image-playground")):
         return True
-
+    
     if not api_key:
         raise HTTPException(
             status_code=HTTP_403_FORBIDDEN,
             detail="No API key provided"
         )
-
+    
     # Only clean if needed
     if api_key.startswith('Bearer '):
         api_key = api_key[7:]  # Remove 'Bearer ' prefix
-
+    
     # Get API keys from environment
     valid_api_keys = get_env_vars().get('api_keys', [])
     if not valid_api_keys or valid_api_keys == ['']:
@@ -218,14 +211,14 @@ async def verify_api_key(
             status_code=HTTP_403_FORBIDDEN,
             detail="API keys not configured on server"
         )
-
+    
     # Fast check with set operation
     if api_key not in set(valid_api_keys):
         raise HTTPException(
             status_code=HTTP_403_FORBIDDEN,
             detail="Invalid API key"
         )
-
+    
     return True
 
 # Pre-load and cache models.json
@@ -241,27 +234,10 @@ def load_models_data():
 
 # Async wrapper for models data
 async def get_models():
-    # Combine models from all active sets for the /models endpoint
-    all_models = list(mistral_models) + \
-                 list(pollinations_models) + \
-                 list(alternate_models) + \
-                 list(claude_3_models) + \
-                 list(supported_image_models) # Include image models
-
-    # Fetch additional models from models.json if it exists and add them
-    models_from_file = load_models_data()
-    if models_from_file:
-        # Assuming models.json contains a list of dicts with 'id'
-        all_models.extend([model.get('id') for model in models_from_file if model.get('id')])
-
-    # Remove duplicates and sort for a consistent list
-    unique_models = sorted(list(set(all_models)))
-
-    # Format as a list of dictionaries for compatibility with existing /models endpoint
-    formatted_models = [{"id": model_id, "name": model_id} for model_id in unique_models]
-
-    return formatted_models
-
+    models_data = load_models_data()
+    if not models_data:
+        raise HTTPException(status_code=500, detail="Error loading available models")
+    return models_data
 
 # Enhanced async streaming - now with real-time SSE support
 async def generate_search_async(query: str, systemprompt: Optional[str] = None, stream: bool = True):
@@ -391,47 +367,47 @@ async def dynamic_ai_page(request: Request):
     user_agent = request.headers.get('user-agent', 'Unknown User')
     client_ip = request.client.host
     location = f"IP: {client_ip}"
-
+    
     prompt = f"""
-    Generate a dynamic HTML page for a user with the following details: with name "LOKI.AI"
+    Generate a dynamic HTML page for a user with the following details: with name "LOKI.AI" 
     - User-Agent: {user_agent}
     - Location: {location}
     - Style: Cyberpunk, minimalist, or retro
-
+    
     Make sure the HTML is clean and includes a heading, also have cool animations a motivational message, and a cool background.
     Wrap the generated HTML in triple backticks (```).
     """
-
+    
     payload = {
         "model": "mistral-small-latest",
         "messages": [{"role": "user", "content": prompt}]
     }
-
+    
     headers = {
         "Authorization": "Bearer playground"
     }
-
-    response = requests.post("[https://parthsadaria-lokiai.hf.space/chat/completions](https://parthsadaria-lokiai.hf.space/chat/completions)", json=payload, headers=headers)
+    
+    response = requests.post("https://parthsadaria-lokiai.hf.space/chat/completions", json=payload, headers=headers)
     data = response.json()
-
+    
     # Extract HTML from ``` blocks
     html_content = re.search(r"```(.*?)```", data['choices'][0]['message']['content'], re.DOTALL)
     if html_content:
         html_content = html_content.group(1).strip()
-
+    
     # Remove the first word
     if html_content:
         html_content = ' '.join(html_content.split(' ')[1:])
-
+    
     return HTMLResponse(content=html_content)
-
+    
 @app.get("/playground", response_class=HTMLResponse)
 async def playground():
     html_content = read_html_file("playground.html")
     if html_content is None:
         return HTMLResponse(content="<h1>playground.html not found</h1>", status_code=404)
     return HTMLResponse(content=html_content)
-
+    
 @app.get("/image-playground", response_class=HTMLResponse)
 async def playground():
     html_content = read_html_file("image-playground.html")
@@ -548,18 +524,12 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
     model_to_use = payload.model or "gpt-4o-mini"
 
     # Validate model availability - fast lookup with set
-    # Check if the model is in any of the currently active sets
-    if model_to_use not in mistral_models and \
-       model_to_use not in pollinations_models and \
-       model_to_use not in alternate_models and \
-       model_to_use not in claude_3_models and \
-       model_to_use not in supported_image_models: # Also check image models
+    if available_model_ids and model_to_use not in set(available_model_ids):
         raise HTTPException(
             status_code=400,
             detail=f"Model '{model_to_use}' is not available. Check /models for the available model list."
         )
 
-
     # Log request without blocking
     asyncio.create_task(log_request(request, model_to_use))
     usage_tracker.record_request(model=model_to_use, endpoint="/chat/completions")
@@ -589,11 +559,7 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
     elif model_to_use in claude_3_models:  # Use the new endpoint
         endpoint = env_vars['secret_api_endpoint_5']
         custom_headers = {}
-    # Add check for image models here if they use a different endpoint than /images/generations
-    # elif model_to_use in supported_image_models:
-    #     endpoint = env_vars['YOUR_IMAGE_COMPLETIONS_ENDPOINT'] # Define a new env var if needed
-    #     custom_headers = {}
-    else: # Default endpoint
+    else:
         endpoint = env_vars['secret_api_endpoint']
         custom_headers = {
             "Origin": header_url,
@@ -607,36 +573,15 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
     async def real_time_stream_generator():
         try:
             async with httpx.AsyncClient(timeout=60.0) as client:
-                # Adjust the endpoint path based on the provider if necessary
-                # For example, Mistral uses /v1/chat/completions, Pollinations might use something else
-                # Based on your existing code, it seems most use /v1/chat/completions,
-                # but this is a point to verify with the actual provider APIs.
-                api_path = "/v1/chat/completions"
-                if model_to_use in mistral_models:
-                     api_path = "/v1/chat/completions" # Or the correct path for Mistral
-                elif model_to_use in pollinations_models:
-                     api_path = "/v1/chat/completions" # Or the correct path for Pollinations
-                elif model_to_use in alternate_models:
-                     api_path = "/v1/chat/completions" # Or the correct path for Alternate
-                elif model_to_use in claude_3_models:
-                     api_path = "/v1/chat/completions" # Or the correct path for Claude 3
-
-                async with client.stream("POST", f"{endpoint}{api_path}", json=payload_dict, headers=custom_headers) as response:
+                async with client.stream("POST", f"{endpoint}/v1/chat/completions", json=payload_dict, headers=custom_headers) as response:
                     if response.status_code >= 400:
                         error_messages = {
                             422: "Unprocessable entity. Check your payload.",
                             400: "Bad request. Verify input data.",
                             403: "Forbidden. You do not have access to this resource.",
                             404: "The requested resource was not found.",
-                            500: "Internal Server Error from upstream API."
                         }
-                        detail = error_messages.get(response.status_code, f"Error code: {response.status_code} from upstream API.")
-                        try:
-                            # Attempt to get more detail from the upstream response body
-                            error_body = await response.aread()
-                            detail += f" Upstream response: {error_body.decode()}"
-                        except Exception:
-                            pass # Ignore errors reading the body
+                        detail = error_messages.get(response.status_code, f"Error code: {response.status_code}")
                         raise HTTPException(status_code=response.status_code, detail=detail)
 
                     # Stream the response in real-time with minimal buffering
@@ -671,32 +616,8 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
         async for chunk in real_time_stream_generator():
             response_content.append(chunk)
 
-        # The collected content might be multiple SSE data chunks.
-        # For non-streaming, we expect a single JSON object.
-        # This part might need adjustment based on the *actual* non-streaming
-        # response format of the upstream APIs. Assuming it's a single JSON:
-        try:
-            # Attempt to parse the full collected content as JSON
-            # This assumes the upstream non-streaming response is a single JSON blob
-            # without the 'data: ' prefix and multiple lines.
-            full_response_text = "".join(response_content)
-            # Remove potential 'data: ' prefixes if they exist even in non-stream
-            full_response_text = re.sub(r'^data: ', '', full_response_text, flags=re.MULTILINE)
-            # Remove empty lines
-            full_response_text = "\n".join([line for line in full_response_text.splitlines() if line.strip()])
-
-            # If the upstream API sends multiple JSON objects even in non-stream,
-            # you might need to process them differently, e.g., concatenate content.
-            # For now, assume a single JSON object is expected.
-            json_response = json.loads(full_response_text)
-            return JSONResponse(content=json_response)
-        except json.JSONDecodeError:
-             # If parsing fails, return the raw collected content and a server error
-             print(f"Warning: Failed to decode JSON for non-streaming response. Raw content: {response_content}")
-             raise HTTPException(status_code=500, detail="Failed to parse upstream API JSON response.")
-        except Exception as e:
-            print(f"Warning: Unexpected error processing non-streaming response: {e}")
-            raise HTTPException(status_code=500, detail=f"An error occurred processing non-streaming response: {str(e)}")
+        return JSONResponse(content=json.loads(''.join(response_content)))
+
 
 
 # New image generation endpoint
@@ -716,7 +637,7 @@ async def create_image(payload: ImageGenerationPayload, authenticated: bool = De
     if payload.model not in supported_image_models:
         raise HTTPException(
             status_code=400,
-            detail=f"Model '{payload.model}' is not supported for image generation.  Supported models are: {list(supported_image_models)}"
+            detail=f"Model '{payload.model}' is not supported for image generation.  Supported models are: {supported_image_models}"
         )
 
     # Log the request
@@ -730,16 +651,12 @@ async def create_image(payload: ImageGenerationPayload, authenticated: bool = De
         "number": payload.number
     }
 
-    # Target API endpoint for image generation
-    target_api_url = os.getenv('NEW_IMG') # Ensure this env var is set
-
-    if not target_api_url:
-         raise HTTPException(status_code=500, detail="Image generation API endpoint (NEW_IMG) is not configured.")
-
+    # Target API endpoint
+    target_api_url = os.getenv('NEW_IMG')
 
     try:
         # Use a timeout for the image generation request
-        async with httpx.AsyncClient(timeout=120.0) as client: # Increased timeout for image generation
+        async with httpx.AsyncClient(timeout=60.0) as client:
             response = await client.post(target_api_url, json=api_payload)
 
         if response.status_code != 200:
@@ -757,6 +674,7 @@ async def create_image(payload: ImageGenerationPayload, authenticated: bool = De
         raise HTTPException(status_code=500, detail=f"An unexpected error occurred during image generation: {e}")
 
 
+
 # Asynchronous logging function
 async def log_request(request, model):
     # Get minimal data for logging
@@ -920,7 +838,7 @@ def generate_usage_html(usage_data):
     <body>
         <div class="container">
             <div class="logo">
-                <img src="data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iMjAwIiBoZWlnaHQ9IjIwMCIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj48cGF0aCBkPSJNMTAwIDM1TDUwIDkwaDEwMHoiIGZpbGw9IiMzYTZlZTAiLz4+PGNpcmNsZSBjeD0iMTAwIiBjeT0iMTQwIiByPSIzMCIgZmlsbD0iIzNhNmVlMCIvPjwvc3ZnPg==" alt="Lokai AI Logo">
+                <img src="data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iMjAwIiBoZWlnaHQ9IjIwMCIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj48cGF0aCBkPSJNMTAwIDM1TDUwIDkwaDEwMHoiIGZpbGw9IiMzYTZlZTAiLz48Y2lyY2xlIGN4PSIxMDAiIGN5PSIxNDAiIHI9IjMwIiBmaWxsPSIjM2E2ZWUwIi8+PC9zdmc+" alt="Lokai AI Logo">
                 <h1>Lokiai AI</h1>
             </div>
 
@@ -1018,34 +936,33 @@ async def get_meme():
         raise HTTPException(status_code=500, detail="Failed to retrieve meme")
 
 # Utility function for loading model IDs - optimized to run once at startup
-# This function is now less critical as model availability is checked against sets
-# However, it's still used to initially populate available_model_ids from models.json
-def load_model_ids_from_file(json_file_path):
+def load_model_ids(json_file_path):
     try:
         with open(json_file_path, 'r') as f:
             models_data = json.load(f)
-            # Extract 'id' from each model object
+            # Extract 'id' from each model object and use a set for fast lookups
             return [model['id'] for model in models_data if 'id' in model]
     except Exception as e:
-        print(f"Error loading model IDs from {json_file_path}: {str(e)}")
+        print(f"Error loading model IDs: {str(e)}")
         return []
 
 @app.on_event("startup")
 async def startup_event():
     global available_model_ids
-    # Load initial models from models.json
-    available_model_ids = load_model_ids_from_file("models.json")
+    available_model_ids = load_model_ids("models.json")
+    print(f"Loaded {len(available_model_ids)} model IDs")
 
-    # Add models from hardcoded sets to the available_model_ids list for the /models endpoint
-    # Note: The actual model availability check in /chat/completions uses the sets directly.
+    # Add all pollinations models to available_model_ids
     available_model_ids.extend(list(pollinations_models))
+    # Add alternate models to available_model_ids
     available_model_ids.extend(list(alternate_models))
+    # Add mistral models to available_model_ids
     available_model_ids.extend(list(mistral_models))
+     # Add claude models
     available_model_ids.extend(list(claude_3_models))
-    available_model_ids.extend(list(supported_image_models)) # Add image models
 
     available_model_ids = list(set(available_model_ids))  # Remove duplicates
-    print(f"Initial available models for /models endpoint: {len(available_model_ids)}")
+    print(f"Total available models: {len(available_model_ids)}")
 
     # Preload scrapers
     for _ in range(MAX_SCRAPERS):
@@ -1067,14 +984,10 @@ async def startup_event():
         missing_vars.append('SECRET_API_ENDPOINT_4')
     if not env_vars['secret_api_endpoint_5']: # Check the new endpoint
         missing_vars.append('SECRET_API_ENDPOINT_5')
-    # Check Mistral keys only if Mistral models are defined in the sets
-    if mistral_models and (not env_vars.get('mistral_api') or not env_vars.get('mistral_key')):
-         if not env_vars.get('mistral_api'): missing_vars.append('MISTRAL_API')
-         if not env_vars.get('mistral_key'): missing_vars.append('MISTRAL_KEY')
-    # Check image endpoint only if image models are defined in the sets
-    if supported_image_models and not os.getenv('NEW_IMG'):
-         missing_vars.append('NEW_IMG')
-
+    if not env_vars['mistral_api'] and any(model in mistral_models for model in available_model_ids):
+        missing_vars.append('MISTRAL_API')
+    if not env_vars['mistral_key'] and any(model in mistral_models for model in available_model_ids):
+        missing_vars.append('MISTRAL_KEY')
 
     if missing_vars:
         print(f"WARNING: The following environment variables are missing: {', '.join(missing_vars)}")
@@ -1096,6 +1009,7 @@ async def shutdown_event():
 
     print("Server shutdown complete!")
 
+# Health check endpoint
 # Health check endpoint
 @app.get("/health")
 async def health_check():
@@ -1104,25 +1018,22 @@ async def health_check():
     missing_critical_vars = []
 
     # Check critical environment variables
-    if not env_vars.get('api_keys') or env_vars['api_keys'] == ['']:
+    if not env_vars['api_keys'] or env_vars['api_keys'] == ['']:
         missing_critical_vars.append('API_KEYS')
-    if not env_vars.get('secret_api_endpoint'):
+    if not env_vars['secret_api_endpoint']:
         missing_critical_vars.append('SECRET_API_ENDPOINT')
-    if not env_vars.get('secret_api_endpoint_2'):
+    if not env_vars['secret_api_endpoint_2']:
         missing_critical_vars.append('SECRET_API_ENDPOINT_2')
-    if not env_vars.get('secret_api_endpoint_3'):
+    if not env_vars['secret_api_endpoint_3']:
         missing_critical_vars.append('SECRET_API_ENDPOINT_3')
-    if not env_vars.get('secret_api_endpoint_4'):
+    if not env_vars['secret_api_endpoint_4']:
         missing_critical_vars.append('SECRET_API_ENDPOINT_4')
-    if not env_vars.get('secret_api_endpoint_5'): # Check the new endpoint
+    if not env_vars['secret_api_endpoint_5']: # Check the new endpoint
         missing_critical_vars.append('SECRET_API_ENDPOINT_5')
-    if not env_vars.get('mistral_api'):
+    if not env_vars['mistral_api']:
         missing_critical_vars.append('MISTRAL_API')
-    if not env_vars.get('mistral_key'):
+    if not env_vars['mistral_key']:
         missing_critical_vars.append('MISTRAL_KEY')
-    if not os.getenv('NEW_IMG'):
-         missing_critical_vars.append('NEW_IMG')
-
 
     health_status = {
         "status": "healthy" if not missing_critical_vars else "unhealthy",
@@ -1132,79 +1043,6 @@ async def health_check():
     }
     return JSONResponse(content=health_status)
 
-
-# --- Admin Endpoints ---
-# Create a separate APIRouter for admin endpoints for better organization
-admin_router = APIRouter(prefix="/admin", tags=["Admin"])
-
-@admin_router.post("/update_models", dependencies=[Depends(verify_api_key)])
-async def update_provider_models(payload: UpdateModelsPayload):
-    """
-    Updates the list of available models for a specific provider.
-    Requires API key authentication.
-    """
-    global mistral_models, pollinations_models, alternate_models, claude_3_models, supported_image_models, available_model_ids
-
-    provider = payload.provider.lower()
-    new_models_list = payload.models
-
-    # Map provider names to the corresponding global sets
-    provider_model_sets = {
-        "mistral": mistral_models,
-        "pollinations": pollinations_models,
-        "alternate": alternate_models,
-        "claude_3": claude_3_models,
-        "image": supported_image_models # Use "image" as the provider name for image models
-    }
-
-    if provider not in provider_model_sets:
-        raise HTTPException(status_code=400, detail=f"Unknown provider: {provider}. Valid providers are: {list(provider_model_sets.keys())}")
-
-    # Update the models set for the specified provider
-    # Using set() ensures uniqueness and efficient lookup
-    provider_model_sets[provider].clear() # Clear existing models
-    provider_model_sets[provider].update(new_models_list) # Add new models
-
-    # Rebuild the overall available_model_ids list for the /models endpoint
-    # This is important so the /models endpoint reflects the changes
-    available_model_ids = list(mistral_models) + \
-                         list(pollinations_models) + \
-                         list(alternate_models) + \
-                         list(claude_3_models) + \
-                         list(supported_image_models)
-    available_model_ids = list(set(available_model_ids)) # Remove duplicates
-
-    print(f"Updated models for provider '{provider}'. New models: {list(provider_model_sets[provider])}")
-    print(f"Total available models for /models endpoint: {len(available_model_ids)}")
-
-
-    # TODO: Implement persistence (e.g., save to models.json)
-    # For example: save_models_to_file("models.json", provider_model_sets)
-    # And modify startup_event to load from this file.
-
-    return {"message": f"Models updated successfully for provider: {provider}", "new_models_count": len(new_models_list)}
-
-@admin_router.get("/view_models", dependencies=[Depends(verify_api_key)])
-async def view_current_models():
-    """
-    Returns the currently active model sets for all providers.
-    Requires API key authentication.
-    """
-    # Return the current state of the model sets
-    return {
-        "mistral": list(mistral_models),
-        "pollinations": list(pollinations_models),
-        "alternate": list(alternate_models),
-        "claude_3": list(claude_3_models),
-        "image": list(supported_image_models)
-    }
-
-# Add the admin router to the main app
-app.include_router(admin_router)
-
-
 if __name__ == "__main__":
     import uvicorn
-    # Note: For production, consider using a process manager like Gunicorn
-    # with multiple workers for better performance and reliability.
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)