Update main.py

main.py

@@ -55,6 +55,7 @@ def get_env_vars():
         'secret_api_endpoint': os.getenv('SECRET_API_ENDPOINT'),
         'secret_api_endpoint_2': os.getenv('SECRET_API_ENDPOINT_2'),
         'secret_api_endpoint_3': os.getenv('SECRET_API_ENDPOINT_3'),
+        'secret_api_endpoint_4': "https://text.pollinations.ai/openai",
         'mistral_api': "https://api.mistral.ai",
         'mistral_key': os.getenv('MISTRAL_KEY'),
         'image_endpoint': os.getenv("IMAGE_ENDPOINT"),
@@ -74,6 +75,25 @@ mistral_models = {
     "codestral-latest"
 }
 
+pollinations_models = {
+    "openai",
+    "openai-large",
+    "openai-reasoning",
+    "qwen-coder",
+    "llama",
+    "mistral",
+    "searchgpt",
+    "deepseek",
+    "claude-hybridspace",
+    "deepseek-r1",
+    "deepseek-reasoner",
+    "llamalight",
+    "gemini",
+    "gemini-thinking",
+    "hormoz",
+    "llama-scaleway"
+}
+
 alternate_models = {
     "gpt-4o-mini",
     "deepseek-v3",
@@ -341,6 +361,9 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
         custom_headers = {
             "Authorization": f"Bearer {env_vars['mistral_key']}"
         }
+    elif model_to_use in pollinations_models:
+        endpoint = env_vars['secret_api_endpoint_4']
+        custom_headers = {}
     elif model_to_use in alternate_models:
         endpoint = env_vars['secret_api_endpoint_2']
         custom_headers = {}
@@ -348,7 +371,6 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
         endpoint = env_vars['secret_api_endpoint']
         custom_headers = {}
 
-    # Get a scraper from the pool
     scraper = get_scraper()
 
     async def stream_generator(payload_dict):
@@ -790,6 +812,11 @@ async def startup_event():
     available_model_ids = load_model_ids("models.json")
     print(f"Loaded {len(available_model_ids)} model IDs")
     
+    # Add all pollinations models to available_model_ids
+    available_model_ids.extend(list(pollinations_models))
+    available_model_ids = list(set(available_model_ids))  # Remove duplicates
+    print(f"Added Pollinations models. Total available models: {len(available_model_ids)}")
+    
     # Preload scrapers
     for _ in range(MAX_SCRAPERS):
         scraper_pool.append(cloudscraper.create_scraper())
@@ -804,22 +831,105 @@ async def startup_event():
         missing_vars.append('SECRET_API_ENDPOINT_2')
     if not env_vars['secret_api_endpoint_3']:
         missing_vars.append('SECRET_API_ENDPOINT_3')
+    if not env_vars['secret_api_endpoint_4']:
+        missing_vars.append('SECRET_API_ENDPOINT_4')
+    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 not env_vars['image_endpoint']:
+        missing_vars.append('IMAGE_ENDPOINT')
     
     if missing_vars:
-        print(f"WARNING: The following required environment variables are missing: {', '.join(missing_vars)}")
+        print(f"WARNING: The following environment variables are missing: {', '.join(missing_vars)}")
+        print("Some functionality may be limited.")
+
+    print("Server started successfully!")
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    # Close the httpx client
+    client = get_async_client()
+    await client.aclose()
+    
+    # Clear scraper pool
+    scraper_pool.clear()
     
-    print("API started successfully with high-performance optimizations")
+    # Persist usage data
+    usage_tracker.save_data()
+    
+    print("Server shutdown complete!")
+
+# Server maintenance endpoint
 
+# Health check endpoint
+@app.get("/health")
+async def health_check():
+    """Health check endpoint for monitoring"""
+    env_vars = get_env_vars()
+    missing_critical_vars = []
+    
+    # Check critical environment variables
+    if not env_vars['api_keys'] or env_vars['api_keys'] == ['']:
+        missing_critical_vars.append('API_KEYS')
+    if not env_vars['secret_api_endpoint']:
+        missing_critical_vars.append('SECRET_API_ENDPOINT')
+    
+    # Check if models are loaded
+    models_loaded = len(available_model_ids) > 0
+    
+    status = "healthy"
+    if missing_critical_vars or not models_loaded:
+        status = "degraded"
+    
+    return {
+        "status": status,
+        "timestamp": datetime.datetime.utcnow().isoformat(),
+        "uptime": time.time() - usage_tracker.start_time,
+        "models_loaded": models_loaded,
+        "model_count": len(available_model_ids),
+        "issues": {
+            "missing_env_vars": missing_critical_vars,
+            "models_available": models_loaded
+        }
+    }
+
+# Error handlers
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request, exc):
+    """Format HTTP exceptions in a consistent way"""
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={"error": exc.detail}
+    )
+
+@app.exception_handler(Exception)
+async def general_exception_handler(request, exc):
+    """Handle unexpected exceptions gracefully"""
+    # Log the error for debugging
+    print(f"Unexpected error: {str(exc)}")
+    
+    return JSONResponse(
+        status_code=500,
+        content={"error": "An unexpected error occurred. Please try again later."}
+    )
+
+# Static files endpoint for serving CSS, JS, etc.
+
+# Documentation
+
+# Run the server when executed directly
 if __name__ == "__main__":
     import uvicorn
+    
+    port = int(os.getenv("PORT", 7860))
+    
+    print(f"Starting Lokiai AI server on port {port}")
     uvicorn.run(
-        app, 
+        "main:app", 
         host="0.0.0.0", 
-        port=7860,
-        workers=4,  # Multiple workers for better CPU utilization
-        loop="uvloop",  # Use uvloop for faster async operations
-        http="httptools",  # Faster HTTP parsing
-        log_level="warning",  # Reduce logging overhead
-        limit_concurrency=100,  # Limit concurrent connections
-        timeout_keep_alive=5  # Reduce idle connection time
+        port=port,
+        workers=4,
+        reload=False,
+        log_level="info"
     )