Upload groq_proxy_service.py

groq_proxy_service.py

@@ -6,8 +6,15 @@ Lightweight FastAPI proxy that exposes a single POST /estimate_audience
 endpoint.  Deployed on HuggingFace Spaces (Slot 1 and Slot 2 of the
 audience estimation pool in music_chart_server.py).
 
-Rate limits honoured (lightweight Groq model per account):
-  60 RPM  / 1K RPD / 6K TPM / 500K TPD
+Design: dumb executor.
+  music_chart_server owns ALL model-selection and prompt logic.  It sends
+  both `model` and `prompt` in every request body so that this service
+  simply fires whatever it receives at Groq and returns the raw result.
+  This keeps control centralised — updating model preference or prompts
+  requires changes in one place only (music_chart_server.py).
+
+Rate limits honoured (same Groq free-tier limits apply to all models):
+  30 RPM  / varies RPD / varies TPM / varies TPD  (see music_chart_server config)
 
 This service enforces a strict 1 req/s inter-request gate so that
 music_chart_server's own 1 req/s per-slot reservation is always met,
@@ -16,9 +23,6 @@ even if multiple upstream callers arrive simultaneously.
 Environment variables (injected as HF Secrets — never hardcoded):
   GROQ_API_KEY   — the Groq API key for this Space's account (required)
 
-Model:
-  llama-3.1-8b-instant  (same lightweight model as the native slot)
-
 Usage:
   uvicorn groq_proxy_service:app --host 0.0.0.0 --port 7860
 """
@@ -31,7 +35,7 @@ import threading
 import time
 from typing import Optional
 
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 from groq import Groq
 from pydantic import BaseModel
@@ -53,14 +57,12 @@ if not GROQ_API_KEY:
         "All /estimate_audience calls will fail until it is configured."
     )
 
-# Lightweight model — fits easily within 6K TPM per account
-GROQ_MODEL = "llama-3.1-8b-instant"
-
-# Minimum seconds between consecutive Groq calls on this instance
-# 1.0 s  ⟹  60 RPM ceiling, matching the per-slot reservation in the pool
+# Minimum seconds between consecutive Groq calls on this instance.
+# music_chart_server already enforces 1 req/s per slot, but this gate
+# protects against accidental concurrent upstream callers.
 SLOT_INTERVAL: float = float(os.environ.get("SLOT_INTERVAL", "1.0"))
 
-# Maximum tokens in Groq response — audience JSON is tiny
+# Maximum tokens in Groq response — audience JSON is tiny regardless of model
 MAX_TOKENS: int = int(os.environ.get("MAX_TOKENS", "256"))
 
 # ─────────────────────────── RATE-LIMIT GATE ────────────────────────────────
@@ -111,41 +113,6 @@ def _get_client() -> Groq:
     return _groq_client
 
 
-# ─────────────────────────── PROMPT ─────────────────────────────────────────
-
-_PROMPT_TEMPLATE = """\
-Estimate audience distribution by country for artist {ARTIST}, song "{SONG_TITLE}".
-
-Return ONLY JSON:
-{{
-  "AAA": 0.00,
-  "BBB": 0.00,
-  "CCC": 0.00,
-  "OTHERS": 0.00
-}}
-
-Rules:
-- ISO3 country codes
-- Values are decimals, sum = 1.00
-- Top 3 countries + OTHERS only
-- No country > 0.50
-- Highest share = primary song language country (skip if English)
-- If genre is globally distributed (e.g., amapiano, afrobeats, hip-hop, pop), enforce OTHERS ≥ 0.15
-- If not a global genre, OTHERS can be any value
-- Prefer realistic streaming distribution
-- If unsure, distribute evenly
-- No text, JSON only
-Non Negotiable: Return ONLY valid JSON. No text, no explanation, no reasoning.
-"""
-
-
-def _build_prompt(artist: str, song_title: str) -> str:
-    return _PROMPT_TEMPLATE.format(
-        ARTIST=artist,
-        SONG_TITLE=song_title or artist,
-    )
-
-
 # ─────────────────────────── JSON PARSING ───────────────────────────────────
 
 def _parse_audience_json(raw: str, artist: str) -> dict:
@@ -180,21 +147,27 @@ def _parse_audience_json(raw: str, artist: str) -> dict:
 app = FastAPI(
     title="Groq Audience Proxy",
     description=(
-        "Slot proxy for the music_chart_server audience estimation pool. "
-        "Enforces 1 req/s to stay within Groq lightweight-model limits."
+        "Dumb-executor slot proxy for the music_chart_server audience pool. "
+        "Model and prompt are supplied by the caller — this service only fires "
+        "the request at Groq and returns the result. "
+        "Enforces 1 req/s to stay within Groq rate limits."
     ),
-    version="1.0.0",
+    version="2.0.0",
 )
 
 
 class AudienceRequest(BaseModel):
-    artist: str
+    artist:     str
     song_title: str = ""
+    # Centralised control fields — sent by music_chart_server on every request.
+    # The proxy uses these directly; it never picks models or builds prompts itself.
+    model:  str = ""   # e.g. "llama-3.1-8b-instant", "llama-3.3-70b-versatile"
+    prompt: str = ""   # fully-rendered prompt string from music_chart_server
 
 
 class AudienceResponse(BaseModel):
-    result: dict
-    model: str
+    result:       dict
+    model:        str
     wait_seconds: float
 
 
@@ -202,80 +175,95 @@ class AudienceResponse(BaseModel):
 def health():
     """Liveness probe used by HuggingFace Spaces and the upstream pool."""
     return {
-        "status": "ok",
-        "model": GROQ_MODEL,
+        "status":                "ok",
         "slot_interval_seconds": SLOT_INTERVAL,
-        "groq_key_configured": bool(GROQ_API_KEY),
+        "groq_key_configured":   bool(GROQ_API_KEY),
+        "note":                  "model selected by music_chart_server per request",
     }
 
 
 @app.post("/estimate_audience", response_model=AudienceResponse)
 def estimate_audience(req: AudienceRequest):
     """
-    Estimate audience distribution for the given artist / song.
+    Execute an audience estimation Groq call using the model and prompt
+    supplied by music_chart_server in the request body.
 
     The gate serialises calls so that Groq receives at most 1 request per
-    SLOT_INTERVAL seconds from this Space — aligning with the 60 RPM limit
-    of the lightweight Groq model.
+    SLOT_INTERVAL seconds from this Space.
 
-    Returns:
+    Request body:
       {
-        "result": {"USA": 0.35, "GBR": 0.25, ...},
-        "model":  "llama-3.1-8b-instant",
+        "artist":     "Artist Name",
+        "song_title": "Track Title",
+        "model":      "llama-3.1-8b-instant",   <- chosen by music_chart_server
+        "prompt":     "<full rendered prompt>"   <- built by music_chart_server
+      }
+
+    Response:
+      {
+        "result":       {"USA": 0.35, "GBR": 0.25, ...},
+        "model":        "llama-3.1-8b-instant",
         "wait_seconds": 0.0
       }
 
     Error responses:
-      429 — gate is already holding a request (should not occur in normal use
-             because the upstream pool issues at most 1 req/s to this slot)
+      422 — missing required fields
       500 — Groq call failed or response could not be parsed
     """
-    artist = req.artist.strip()
+    artist     = req.artist.strip()
     song_title = req.song_title.strip()
+    model      = req.model.strip()
+    prompt     = req.prompt.strip()
 
     if not artist:
         raise HTTPException(status_code=422, detail="'artist' must not be empty.")
+    if not model:
+        raise HTTPException(status_code=422, detail="'model' must not be empty.")
+    if not prompt:
+        raise HTTPException(status_code=422, detail="'prompt' must not be empty.")
 
-    log.info(f"estimate_audience: artist='{artist}' song='{song_title}'")
+    log.info(
+        f"estimate_audience: artist='{artist}' song='{song_title}' model='{model}'"
+    )
 
     # ── Acquire slot (blocks until window is clear) ───────────────────────────
     wait_s = _gate.acquire()
     if wait_s > 0:
-        log.debug(f"Gate wait: {wait_s:.3f}s for '{artist}'")
+        log.debug(f"Gate wait: {wait_s:.3f}s for '{artist}' model='{model}'")
 
-    # ── Call Groq ─────────────────────────────────────────────────────────────
+    # ── Call Groq with the caller-supplied model and prompt ───────────────────
     try:
         client = _get_client()
         completion = client.chat.completions.create(
-            model=GROQ_MODEL,
-            messages=[
-                {"role": "user", "content": _build_prompt(artist, song_title)}
-            ],
+            model=model,
+            messages=[{"role": "user", "content": prompt}],
             temperature=0.4,
             max_tokens=MAX_TOKENS,
             stream=False,
         )
         raw: str = completion.choices[0].message.content or ""
-        log.debug(f"Groq raw response for '{artist}': {raw!r}")
+        log.debug(f"Groq raw response for '{artist}' model='{model}': {raw!r}")
     except Exception as exc:
-        log.error(f"Groq API error for '{artist}': {exc}")
+        log.error(f"Groq API error for '{artist}' model='{model}': {exc}")
         raise HTTPException(
             status_code=500,
-            detail=f"Groq API call failed: {exc}",
+            detail=f"Groq API call failed (model={model}): {exc}",
         )
 
     # ── Parse and return ──────────────────────────────────────────────────────
     try:
         result = _parse_audience_json(raw, artist)
     except Exception as exc:
-        log.error(f"JSON parse error for '{artist}': {exc}  raw={raw!r}")
+        log.error(
+            f"JSON parse error for '{artist}' model='{model}': {exc}  raw={raw!r}"
+        )
         raise HTTPException(
             status_code=500,
-            detail=f"Failed to parse Groq response: {exc}",
+            detail=f"Failed to parse Groq response (model={model}): {exc}",
         )
 
-    log.info(f"estimate_audience OK: '{artist}' → {result}")
-    return AudienceResponse(result=result, model=GROQ_MODEL, wait_seconds=wait_s)
+    log.info(f"estimate_audience OK: '{artist}' model='{model}' -> {result}")
+    return AudienceResponse(result=result, model=model, wait_seconds=wait_s)
 
 
 # ─────────────────────────── ENTRYPOINT ─────────────────────────────────────