groq_proxy_service.py

#!/usr/bin/env python3
"""
groq_proxy_service.py
─────────────────────
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).

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,
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)

Usage:
  uvicorn groq_proxy_service:app --host 0.0.0.0 --port 7860
"""

import json
import logging
import os
import re
import threading
import time
from typing import Optional

from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from groq import Groq
from pydantic import BaseModel

# ─────────────────────────── LOGGING ────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
log = logging.getLogger("groq_proxy")

# ─────────────────────────── CONFIGURATION ──────────────────────────────────

GROQ_API_KEY: str = os.environ.get("GROQ_API_KEY", "")
if not GROQ_API_KEY:
    log.warning(
        "GROQ_API_KEY environment variable is not set. "
        "All /estimate_audience calls will fail until it is configured."
    )

# 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 regardless of model
MAX_TOKENS: int = int(os.environ.get("MAX_TOKENS", "256"))

# ─────────────────────────── RATE-LIMIT GATE ────────────────────────────────

class _SlotGate:
    """
    Serialising gate that ensures at most 1 Groq call per SLOT_INTERVAL seconds.

    All concurrent HTTP requests queue here; each is released only when the
    previous call's slot window has elapsed.  Thread-safe via a single lock.
    """

    def __init__(self, interval: float):
        self._interval = interval
        self._lock = threading.Lock()
        self._next_allowed: float = 0.0   # monotonic timestamp

    def acquire(self) -> float:
        """
        Block until the current slot is available, then reserve it.
        Returns the number of seconds this call had to wait.
        """
        with self._lock:
            now = time.monotonic()
            wait = max(0.0, self._next_allowed - now)
            if wait > 0:
                time.sleep(wait)
            # Reserve the slot: next call may start only after SLOT_INTERVAL
            self._next_allowed = time.monotonic() + self._interval
            return wait


_gate = _SlotGate(interval=SLOT_INTERVAL)

# ─────────────────────────── GROQ CLIENT ────────────────────────────────────

_groq_client: Optional[Groq] = None
_groq_lock = threading.Lock()


def _get_client() -> Groq:
    global _groq_client
    with _groq_lock:
        if _groq_client is None:
            if not GROQ_API_KEY:
                raise RuntimeError("GROQ_API_KEY is not configured on this Space.")
            _groq_client = Groq(api_key=GROQ_API_KEY)
    return _groq_client


# ─────────────────────────── JSON PARSING ───────────────────────────────────

def _parse_audience_json(raw: str, artist: str) -> dict:
    """Strip markdown fences, extract JSON object, normalise proportions."""
    clean = re.sub(r"```(?:json)?|```", "", raw).strip()
    match = re.search(r"\{[^{}]+\}", clean, re.DOTALL)
    if not match:
        raise ValueError(f"No JSON object in model response: {clean!r}")

    parsed: dict = {
        k.upper(): float(v)
        for k, v in json.loads(match.group()).items()
    }
    if not parsed:
        raise ValueError("Parsed audience dict is empty.")

    total = sum(parsed.values())
    if total <= 0:
        raise ValueError(f"Non-positive proportion total ({total}) for '{artist}'.")

    if abs(total - 1.0) > 0.05:
        log.debug(
            f"Proportions sum to {total:.4f} for '{artist}'; re-normalising."
        )
        parsed = {k: round(v / total, 4) for k, v in parsed.items()}

    return parsed


# ─────────────────────────── FASTAPI APP ────────────────────────────────────

app = FastAPI(
    title="Groq Audience Proxy",
    description=(
        "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="2.0.0",
)


class AudienceRequest(BaseModel):
    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
    wait_seconds: float


@app.get("/health")
def health():
    """Liveness probe used by HuggingFace Spaces and the upstream pool."""
    return {
        "status":                "ok",
        "slot_interval_seconds": SLOT_INTERVAL,
        "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):
    """
    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.

    Request body:
      {
        "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:
      422 — missing required fields
      500 — Groq call failed or response could not be parsed
    """
    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}' 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}' model='{model}'")

    # ── Call Groq with the caller-supplied model and prompt ───────────────────
    try:
        client = _get_client()
        completion = client.chat.completions.create(
            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}' model='{model}': {raw!r}")
    except Exception as exc:
        log.error(f"Groq API error for '{artist}' model='{model}': {exc}")
        raise HTTPException(
            status_code=500,
            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}' model='{model}': {exc}  raw={raw!r}"
        )
        raise HTTPException(
            status_code=500,
            detail=f"Failed to parse Groq response (model={model}): {exc}",
        )

    log.info(f"estimate_audience OK: '{artist}' model='{model}' -> {result}")
    return AudienceResponse(result=result, model=model, wait_seconds=wait_s)


# ─────────────────────────── ENTRYPOINT ─────────────────────────────────────

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "groq_proxy_service:app",
        host="0.0.0.0",
        port=int(os.environ.get("PORT", "7860")),
        log_level="info",
    )