services/embedder_client.py

"""
Encode text to embeddings, either via a separate embedder HF Space
(EMBEDDER_URL set) or an in-process model (EMBEDDER_URL unset).

Why a circuit breaker: if the embedder Space is asleep, crashed, or just
slow, we do NOT want every /check-submission call to hang for the full
timeout one by one while participants wait. After a few consecutive
failures we "open" the breaker and skip calling the embedder entirely for
a cooldown window, going straight to the local-fallback path (or, if there
is no local model loaded either, letting the caller degrade to fuzzy-only
matching). This keeps failures cheap and bounded instead of compounding.
"""

import logging
import threading
import time
from typing import Optional

import httpx
import numpy as np

from config import (
    EMBEDDER_CIRCUIT_COOLDOWN_SECONDS,
    EMBEDDER_CIRCUIT_FAILURE_THRESHOLD,
    EMBEDDER_TIMEOUT_SECONDS,
    EMBEDDER_URL,
    MODEL_NAME,
    NER_TIMEOUT_SECONDS,
)

logger = logging.getLogger(__name__)

_EMBEDDER_API_KEY_HEADER = "X-API-Key"


class _CircuitBreaker:
    def __init__(self, failure_threshold: int, cooldown_seconds: float) -> None:
        self._failure_threshold = failure_threshold
        self._cooldown_seconds = cooldown_seconds
        self._lock = threading.Lock()
        self._consecutive_failures = 0
        self._opened_at: Optional[float] = None

    def is_open(self) -> bool:
        with self._lock:
            if self._opened_at is None:
                return False
            if time.monotonic() - self._opened_at >= self._cooldown_seconds:
                # Cooldown elapsed -- allow one probe attempt through.
                self._opened_at = None
                self._consecutive_failures = 0
                return False
            return True

    def record_success(self) -> None:
        with self._lock:
            self._consecutive_failures = 0
            self._opened_at = None

    def record_failure(self) -> None:
        with self._lock:
            self._consecutive_failures += 1
            if self._consecutive_failures >= self._failure_threshold and self._opened_at is None:
                self._opened_at = time.monotonic()
                logger.warning(
                    "Embedder circuit breaker OPEN after %d consecutive failures — "
                    "falling back to local/fuzzy matching for %.0fs",
                    self._consecutive_failures,
                    self._cooldown_seconds,
                )


_breaker = _CircuitBreaker(EMBEDDER_CIRCUIT_FAILURE_THRESHOLD, EMBEDDER_CIRCUIT_COOLDOWN_SECONDS)
_client: Optional[httpx.Client] = None


def _get_http_client() -> httpx.Client:
    global _client
    if _client is None:
        _client = httpx.Client(timeout=EMBEDDER_TIMEOUT_SECONDS)
    return _client


def is_remote_configured() -> bool:
    return bool(EMBEDDER_URL)


def is_remote_available() -> bool:
    """True if the remote embedder is configured and the breaker isn't open."""
    return is_remote_configured() and not _breaker.is_open()


def encode_remote(texts: list[str], api_key: str = "") -> Optional[np.ndarray]:
    """Try the remote embedder Space. Returns None (never raises) on any
    failure so callers can fall back cleanly -- a slow/dead embedder must
    never be able to break or hang a duplicate check."""
    if not is_remote_available():
        return None
    try:
        headers = {_EMBEDDER_API_KEY_HEADER: api_key} if api_key else {}
        resp = _get_http_client().post(
            f"{EMBEDDER_URL}/embed", json={"texts": texts}, headers=headers
        )
        resp.raise_for_status()
        data = resp.json()
        _breaker.record_success()
        return np.asarray(data["embeddings"])
    except Exception as exc:
        logger.warning("Embedder Space call failed, will fall back: %s", exc)
        _breaker.record_failure()
        return None


_ner_client: Optional[httpx.Client] = None


def _get_ner_http_client() -> httpx.Client:
    # Separate client (and much longer timeout) from the embed one above --
    # NER is a one-off batch call from the QA batch, possibly the Space's
    # very first request ever if the NER model hasn't been touched yet
    # (cold start + model download), not a per-submission call that needs
    # to fail fast.
    global _ner_client
    if _ner_client is None:
        _ner_client = httpx.Client(timeout=NER_TIMEOUT_SECONDS)
    return _ner_client


def ner_remote(texts: list[str], api_key: str = "") -> Optional[list[list[dict]]]:
    """Try the remote embedder Space's /ner endpoint. Returns None (never
    raises) on any failure -- including the Space being unconfigured, or
    responding with available=False (its own NER model failed to load) --
    so callers (pii_service.scan_pii_batch) can fall back to a local model
    or plain regex+list detection. No circuit breaker: this is called at
    most once per QA batch run, never per-request, so there's no pile-up
    risk to guard against."""
    if not is_remote_configured():
        return None
    try:
        headers = {_EMBEDDER_API_KEY_HEADER: api_key} if api_key else {}
        resp = _get_ner_http_client().post(
            f"{EMBEDDER_URL}/ner", json={"texts": texts}, headers=headers
        )
        resp.raise_for_status()
        data = resp.json()
        if not data.get("available", False):
            logger.warning("Embedder Space's NER model isn't loaded there — falling back")
            return None
        return data["entities"]
    except Exception as exc:
        logger.warning("Embedder Space /ner call failed, will fall back: %s", exc)
        return None