Upload production/production_server.py

production/production_server.py

@@ -5,7 +5,7 @@ Features:
 - FastAPI with async endpoints
 - Distributed rate limiting (Redis-backed token bucket)
 - Circuit breaker for external API resilience
-- Automatic fallback: NIM (primary) -> Cloudflare Workers AI (fallback)
+- Automatic fallback: NIM (primary) -> Cloudflare Workers AI (secondary) -> Gemini (tertiary) -> MLX
 - Request/response caching with Redis TTL
 - Multi-tenant session isolation
 - Health checks and graceful shutdown
@@ -13,6 +13,8 @@ Features:
 - Cost tracking and budget enforcement
 - Connection pooling for all HTTP clients
 - Cloudflare Workers AI support via OpenAI-compatible API
+- Google Gemini support via OpenAI-compatible API
+- MLX local support for Apple Silicon
 """
 
 import asyncio
@@ -58,10 +60,15 @@ NIM_API_BASE = os.environ.get("NIM_API_BASE", "https://integrate.api.nvidia.com/
 CLOUDFLARE_API_KEY = os.environ.get("CLOUDFLARE_API_KEY", "")
 CLOUDFLARE_ACCOUNT_ID = os.environ.get("CLOUDFLARE_ACCOUNT_ID", "")
 
+# Google Gemini / AI Studio
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
+GEMINI_API_BASE = os.environ.get("GEMINI_API_BASE", "https://generativelanguage.googleapis.com/v1beta/openai")
+
 # Fallback configuration
 FALLBACK_ENABLED = os.environ.get("FALLBACK_ENABLED", "true").lower() == "true"
 FALLBACK_PRIMARY = os.environ.get("FALLBACK_PRIMARY", "nim")
 FALLBACK_SECONDARY = os.environ.get("FALLBACK_SECONDARY", "cloudflare")
+FALLBACK_TERTIARY = os.environ.get("FALLBACK_TERTIARY", "gemini")
 
 # MLX (local Apple Silicon)
 MLX_API_BASE = os.environ.get("MLX_API_BASE", "http://localhost:8000/v1")
@@ -304,6 +311,7 @@ class CircuitBreaker:
 class FallbackConfig:
     primary: str = "nim"
     secondary: str = "cloudflare"
+    tertiary: str = "gemini"
     enabled: bool = True
 
 class FallbackManager:
@@ -338,18 +346,33 @@ class FallbackManager:
                 to_provider=self.config.secondary,
                 reason="circuit_open",
             ).inc()
-            logger.warning(f"Fallback: {self.config.primary} unavailable, switching to {self.config.secondary}")
+            logger.warning(
+                f"Fallback: {self.config.primary} unavailable, switching to {self.config.secondary}"
+            )
             return self.config.secondary, self._get_provider_config(self.config.secondary)
         
+        tertiary_breaker = CircuitBreaker(self.redis, self.config.tertiary)
+        if await tertiary_breaker.can_execute():
+            FALLBACK_COUNT.labels(
+                from_provider=self.config.secondary,
+                to_provider=self.config.tertiary,
+                reason="secondary_down",
+            ).inc()
+            logger.warning(
+                f"Fallback: both {self.config.primary} and {self.config.secondary} down, "
+                f"switching to {self.config.tertiary}"
+            )
+            return self.config.tertiary, self._get_provider_config(self.config.tertiary)
+        
         if MLX_ENABLED:
             mlx_breaker = CircuitBreaker(self.redis, "mlx")
             if await mlx_breaker.can_execute():
                 FALLBACK_COUNT.labels(
-                    from_provider=self.config.primary,
+                    from_provider=self.config.tertiary,
                     to_provider="mlx",
-                    reason="both_down",
+                    reason="all_cloud_down",
                 ).inc()
-                logger.warning("Both cloud providers down — falling back to MLX local")
+                logger.warning("All cloud providers down — falling back to MLX local")
                 return "mlx", self._get_provider_config("mlx")
         
         raise HTTPException(status_code=503, detail="All LLM providers unavailable.")
@@ -370,6 +393,13 @@ class FallbackManager:
                 "cost_per_1m_input": 0.0,
                 "cost_per_1m_output": 0.0,
             },
+            "gemini": {
+                "api_base": GEMINI_API_BASE,
+                "api_key": GEMINI_API_KEY,
+                "rpm_limit": 60,
+                "cost_per_1m_input": 0.075,
+                "cost_per_1m_output": 0.30,
+            },
             "mlx": {
                 "api_base": MLX_API_BASE,
                 "api_key": "no-key",
@@ -419,7 +449,7 @@ class ConcurrencyLimiter:
 # ---------------------------------------------------------------------------
 
 class ChatRequest(BaseModel):
-    model: str = Field(..., description="Model ID (e.g., @cf/meta/llama-3.1-8b-instruct)")
+    model: str = Field(..., description="Model ID (e.g., gemma-4-31b-bf16)")
     messages: list[dict] = Field(..., description="OpenAI-compatible messages")
     temperature: Optional[float] = 0.7
     max_tokens: Optional[int] = 4096
@@ -540,7 +570,7 @@ async def _init_schema():
 
 app = FastAPI(
     title="ml-intern Production API",
-    description="Production-grade API with NIM/Cloudflare fallback and MLX local support",
+    description="Production-grade API with NIM/Cloudflare/Gemini fallback and MLX local support",
     version="1.0.0",
     lifespan=lifespan,
 )
@@ -715,7 +745,7 @@ async def health_check():
     
     circuits = {}
     fallback_status = {}
-    for provider in ["nim", "cloudflare", "mlx"]:
+    for provider in ["nim", "cloudflare", "gemini", "mlx"]:
         try:
             state = await redis_manager.get_circuit_state(provider)
             circuits[provider] = state["state"]
@@ -749,7 +779,6 @@ async def chat_completions(request: ChatRequest, background_tasks: BackgroundTas
     
     await concurrency_limiter.acquire()
     try:
-        # 1. Determine provider
         if request.provider_override:
             provider = request.provider_override
             provider_config = fallback_manager._get_provider_config(provider)
@@ -761,7 +790,6 @@ async def chat_completions(request: ChatRequest, background_tasks: BackgroundTas
         
         request.state.provider = provider
         
-        # 2. Rate limiting
         rpm = provider_config.get("rpm_limit", DEFAULT_RPM_LIMIT)
         rate_limit_key = f"{provider}:{session_id}"
         allowed, retry_after = await redis_manager.check_rate_limit(rate_limit_key, rpm)
@@ -769,7 +797,6 @@ async def chat_completions(request: ChatRequest, background_tasks: BackgroundTas
             logger.warning(f"Rate limit exceeded for {rate_limit_key}")
             raise HTTPException(status_code=429, detail=f"Rate limit exceeded. Retry after {retry_after:.1f}s", headers={"Retry-After": str(int(retry_after))})
         
-        # 3. Check cache
         if not request.stream:
             cache_key = generate_cache_key(request)
             cached = await redis_manager.get_cache(cache_key)
@@ -790,19 +817,15 @@ async def chat_completions(request: ChatRequest, background_tasks: BackgroundTas
                     fallback_used=False,
                 )
         
-        # 4. Budget
         cost_tracker = CostTracker(session_id, provider=provider, model=request.model)
         
-        # 5. Call LLM
         response = await call_llm(provider, provider_config, request, session_id)
         
-        # Record success
         breaker = CircuitBreaker(redis_manager, provider)
         await breaker.record_success()
         
         cost_tracker.record_spend(response.cost_usd)
         
-        # 6. Cache response
         if not request.stream:
             cache_key = generate_cache_key(request)
             await redis_manager.set_cache(
@@ -815,11 +838,9 @@ async def chat_completions(request: ChatRequest, background_tasks: BackgroundTas
                 }),
             )
         
-        # 7. Persist
         if DATABASE_URL:
             background_tasks.add_task(_persist_request, session_id, request, response, provider)
         
-        # Mark fallback
         if provider != FALLBACK_PRIMARY and FALLBACK_ENABLED:
             response.fallback_used = True
         
@@ -875,12 +896,24 @@ async def list_models():
             {"id": "cloudflare/@cf/meta/llama-3.1-70b-instruct", "object": "model", "owned_by": "cloudflare"},
             {"id": "cloudflare/@cf/mistral/mistral-7b-instruct", "object": "model", "owned_by": "cloudflare"},
             {"id": "cloudflare/@cf/qwen/qwen1.5-14b-chat-awq", "object": "model", "owned_by": "cloudflare"},
+            {"id": "cloudflare/@cf/google/gemma-4-26b-a4b-it", "object": "model", "owned_by": "cloudflare"},
+        ])
+    
+    if GEMINI_API_KEY:
+        models.extend([
+            {"id": "gemini/gemini-2.5-pro-preview", "object": "model", "owned_by": "google"},
+            {"id": "gemini/gemini-2.5-flash-preview", "object": "model", "owned_by": "google"},
+            {"id": "gemini/gemma-4-26b", "object": "model", "owned_by": "google"},
+            {"id": "gemini/gemma-4-9b", "object": "model", "owned_by": "google"},
         ])
     
     if MLX_ENABLED:
         models.extend([
             {"id": "mlx/llama-3.1-8b", "object": "model", "owned_by": "mlx"},
             {"id": "mlx/llama-3.1-70b", "object": "model", "owned_by": "mlx"},
+            {"id": "mlx/gemma-4-26b-a4b-it", "object": "model", "owned_by": "mlx"},
+            {"id": "mlx/gemma-4-31b-bf16", "object": "model", "owned_by": "mlx"},
+            {"id": "mlx/gemma-4-e4b-it", "object": "model", "owned_by": "mlx"},
         ])
     
     return {"object": "list", "data": models}
@@ -888,7 +921,7 @@ async def list_models():
 @app.get("/v1/fallback/status")
 async def fallback_status():
     status = {}
-    for provider in ["nim", "cloudflare", "mlx"]:
+    for provider in ["nim", "cloudflare", "gemini", "mlx"]:
         breaker = CircuitBreaker(redis_manager, provider)
         can_execute = await breaker.can_execute()
         state = await redis_manager.get_circuit_state(provider)
@@ -903,6 +936,7 @@ async def fallback_status():
         "fallback_enabled": FALLBACK_ENABLED,
         "primary": FALLBACK_PRIMARY,
         "secondary": FALLBACK_SECONDARY,
+        "tertiary": FALLBACK_TERTIARY,
         "providers": status,
         "active_provider": await _get_active_provider_name(),
     }