Update inference proxy (Groq/Gemini/CPU priority chain)

app.py

@@ -1,19 +1,14 @@
 """
-deploy/tgi_space/app.py — Fallback inference server if TGI fails on CPU Space.
+deploy/tgi_space/app.py — Smart inference proxy for persona generation.
 
-Serves OpenAI-compatible /v1/chat/completions using HuggingFace transformers directly.
-Use this Dockerfile instead when TGI has RUNTIME_ERROR on a Space:
+Priority chain (first available wins):
+  1. GROQ_API_KEY → Groq Cloud (fast, free: 14,400 req/day with llama-3.1-8b-instant)
+  2. GEMINI_API_KEY → Gemini Flash (generous free: 1,500 req/day, 1M tok/day)
+  3. Local CPU → transformers (slow fallback, only for smoke-testing)
 
-    FROM python:3.11-slim
-    WORKDIR /app
-    COPY app.py .
-    RUN pip install --no-cache-dir fastapi uvicorn transformers accelerate torch --index-url https://download.pytorch.org/whl/cpu
-    ENV PORT=7860
-    ENV MODEL_ID=Qwen/Qwen2.5-0.5B-Instruct
-    ENV HF_HOME=/data
-    CMD ["python", "app.py"]
-
-Usage (from setup_spaces.py): set FALLBACK_DOCKERFILE=deploy/tgi_space/Dockerfile.fallback
+To activate a fast provider, set the env var in the HF Space settings:
+  - Groq:   GROQ_API_KEY  = gsk_...  (free at https://console.groq.com)
+  - Gemini: GEMINI_API_KEY = AIza...  (free at https://aistudio.google.com)
 """
 from __future__ import annotations
 
@@ -24,38 +19,48 @@ import time
 import uuid
 from typing import Any, Dict, List, Optional
 
-import torch
-from fastapi import FastAPI, HTTPException
+import httpx
+from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
 
-app = FastAPI(title="Persona Inference Server")
+app = FastAPI(title="Persona Inference Proxy")
+
+# ---------------------------------------------------------------------------
+# Config (env vars)
+# ---------------------------------------------------------------------------
 
 MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-0.5B-Instruct")
 PORT = int(os.environ.get("PORT", 7860))
 MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", 600))
 
-# Load model at startup (float32 for CPU, bfloat16 if available)
+GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
+GROQ_MODEL = os.environ.get("GROQ_MODEL", "llama-3.1-8b-instant")
+
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
+GEMINI_MODEL = os.environ.get("GEMINI_MODEL", "gemini-1.5-flash")
+
+# Local model (loaded lazily — only if no fast provider)
 _tokenizer = None
 _model = None
+_local_loaded = False
 
 
-def _load_model():
-    global _tokenizer, _model
-    print(f"Loading {MODEL_ID} ...")
-    _tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_ID, trust_remote_code=True
-    )
-    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-    _model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=dtype,
-        device_map="auto",
-        trust_remote_code=True,
-    )
-    _model.eval()
-    print(f"Model loaded: {MODEL_ID}  dtype={dtype}")
+def _infer_mode() -> str:
+    if GROQ_API_KEY:
+        return "groq"
+    if GEMINI_API_KEY:
+        return "gemini"
+    return "local-cpu"
+
+
+def _active_model() -> str:
+    mode = _infer_mode()
+    if mode == "groq":
+        return GROQ_MODEL
+    if mode == "gemini":
+        return GEMINI_MODEL
+    return MODEL_ID
 
 
 # ---------------------------------------------------------------------------
@@ -76,24 +81,77 @@ class ChatRequest(BaseModel):
 
 
 # ---------------------------------------------------------------------------
-# Endpoints
+# Provider implementations
 # ---------------------------------------------------------------------------
 
-@app.on_event("startup")
-async def startup():
-    loop = asyncio.get_event_loop()
-    await loop.run_in_executor(None, _load_model)
+async def _call_groq(req: ChatRequest) -> dict:
+    payload = {
+        "model": GROQ_MODEL,
+        "messages": [{"role": m.role, "content": m.content} for m in req.messages],
+        "max_tokens": min(req.max_tokens, MAX_NEW_TOKENS),
+        "temperature": req.temperature,
+        "top_p": req.top_p,
+    }
+    async with httpx.AsyncClient(timeout=90.0) as client:
+        r = await client.post(
+            "https://api.groq.com/openai/v1/chat/completions",
+            json=payload,
+            headers={
+                "Authorization": f"Bearer {GROQ_API_KEY}",
+                "Content-Type": "application/json",
+            },
+        )
+    if r.status_code != 200:
+        raise HTTPException(status_code=r.status_code, detail=f"Groq error: {r.text[:200]}")
+    return r.json()
+
+
+async def _call_gemini(req: ChatRequest) -> dict:
+    """Call Gemini via its OpenAI-compatible endpoint."""
+    payload = {
+        "model": GEMINI_MODEL,
+        "messages": [{"role": m.role, "content": m.content} for m in req.messages],
+        "max_tokens": min(req.max_tokens, MAX_NEW_TOKENS),
+        "temperature": req.temperature,
+    }
+    async with httpx.AsyncClient(timeout=90.0) as client:
+        r = await client.post(
+            "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
+            json=payload,
+            headers={
+                "Authorization": f"Bearer {GEMINI_API_KEY}",
+                "Content-Type": "application/json",
+            },
+        )
+    if r.status_code != 200:
+        raise HTTPException(status_code=r.status_code, detail=f"Gemini error: {r.text[:200]}")
+    return r.json()
 
 
-@app.get("/health")
-async def health():
-    return {"status": "ok", "model": MODEL_ID, "loaded": _model is not None}
+def _load_local_model():
+    global _tokenizer, _model, _local_loaded
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    print(f"Loading local model {MODEL_ID} on CPU ...")
+    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    _model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float32,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    _model.eval()
+    _local_loaded = True
+    print(f"Local model loaded: {MODEL_ID}")
 
 
-@app.post("/v1/chat/completions")
-async def chat_completions(req: ChatRequest):
-    if _model is None or _tokenizer is None:
-        raise HTTPException(status_code=503, detail="Model not loaded yet")
+async def _call_local(req: ChatRequest) -> dict:
+    global _local_loaded
+    if not _local_loaded:
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(None, _load_local_model)
+
+    import torch
 
     msgs = [{"role": m.role, "content": m.content} for m in req.messages]
 
@@ -121,16 +179,51 @@ async def chat_completions(req: ChatRequest):
         "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
         "object": "chat.completion",
         "created": int(time.time()),
-        "model": req.model,
-        "choices": [{
-            "index": 0,
-            "message": {"role": "assistant", "content": content},
-            "finish_reason": "stop",
-        }],
+        "model": MODEL_ID,
+        "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
         "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
     }
 
 
+# ---------------------------------------------------------------------------
+# Startup: eagerly load local model only if no fast provider configured
+# ---------------------------------------------------------------------------
+
+@app.on_event("startup")
+async def startup():
+    mode = _infer_mode()
+    print(f"Inference mode: {mode}  active_model: {_active_model()}")
+    if mode == "local-cpu":
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(None, _load_local_model)
+
+
+# ---------------------------------------------------------------------------
+# Endpoints
+# ---------------------------------------------------------------------------
+
+@app.get("/health")
+async def health():
+    mode = _infer_mode()
+    loaded = True if mode != "local-cpu" else _local_loaded
+    return {
+        "status": "ok",
+        "mode": mode,
+        "model": _active_model(),
+        "loaded": loaded,
+    }
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(req: ChatRequest):
+    mode = _infer_mode()
+    if mode == "groq":
+        return await _call_groq(req)
+    if mode == "gemini":
+        return await _call_gemini(req)
+    return await _call_local(req)
+
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=PORT)