pi-qmem/gateway/rerank.py

"""Stadio di re-ranking (cross-encoder) con catena di fallback resiliente.

La catena è definita da RERANK_CHAIN (JSON): il primo nodo raggiungibile vince.
Dopo un fallimento il nodo entra in cooldown (RERANK_RETRY_COOLDOWN_S) e la
richiesta passa al successivo; se tutti i nodi sono in cooldown si ritenta
comunque il primo (meglio di un fallimento immediato). Se nessun nodo risponde
la ricerca degrada con grazia all'ordine di fusione ibrida (nessun errore al
client): il reranking è un miglioramento, non una dipendenza.
"""
from __future__ import annotations

import json
import math
import time
from dataclasses import dataclass
from typing import Optional

import httpx

import metrics
from config import (
    RERANK_CHAIN,
    RERANK_ENABLED,
    RERANK_MODEL,
    RERANK_RETRY_COOLDOWN_S,
    RERANK_TIMEOUT_MS,
    log,
)


@dataclass(frozen=True)
class RerankNode:
    """Un endpoint reranker nella catena di fallback."""

    name: str
    url: str
    key: str
    timeout_ms: int


def parse_chain(raw: str) -> list[RerankNode]:
    """Parsa RERANK_CHAIN: JSON [{name, url, key, timeout_ms}]. URL senza schema → scartato."""
    try:
        entries = json.loads(raw) if raw else []
    except (json.JSONDecodeError, TypeError):
        log.error("RERANK_CHAIN non è JSON valido: reranking disattivato")
        return []
    if not isinstance(entries, list):
        log.error("RERANK_CHAIN non è una lista: reranking disattivato")
        return []
    nodes: list[RerankNode] = []
    for entry in entries:
        if not isinstance(entry, dict) or not entry.get("url"):
            continue
        url = str(entry["url"]).rstrip("/")
        if not url.startswith(("http://", "https://")):
            continue
        nodes.append(
            RerankNode(
                name=str(entry.get("name") or url),
                url=url,
                key=str(entry.get("key") or ""),
                timeout_ms=int(entry.get("timeout_ms", RERANK_TIMEOUT_MS)),
            )
        )
    return nodes


_chain: Optional[list[RerankNode]] = None
_down_until: dict[str, float] = {}  # url → monotonic deadline del cooldown
_http: Optional[httpx.AsyncClient] = None


def _get_chain() -> list[RerankNode]:
    global _chain
    if _chain is None:
        _chain = parse_chain(RERANK_CHAIN)
    return _chain


def reset_chain_cache() -> None:
    """Forza il re-parse della catena (usato dai test)."""
    global _chain
    _chain = None
    _down_until.clear()


def get_http() -> httpx.AsyncClient:
    global _http
    if _http is None:
        _http = httpx.AsyncClient(timeout=30)
    return _http


async def close_http() -> None:
    global _http
    if _http is not None:
        await _http.aclose()
        _http = None


def enabled() -> bool:
    """Reranking attivo: flag env + catena configurata non vuota."""
    return RERANK_ENABLED and bool(_get_chain())


def live_nodes() -> tuple[list[RerankNode], bool]:
    """Nodi fuori cooldown; all_down=True se nessun nodo è live (forza retry totale)."""
    chain = _get_chain()
    now = time.monotonic()
    live = [n for n in chain if _down_until.get(n.url, 0) <= now]
    return live, bool(chain) and len(live) < len(chain)


async def rerank(query: str, docs: list[str]) -> Optional[tuple[list[float], str, int]]:
    """Reranka i documenti rispetto alla query tramite la catena di fallback.

    Ritorna (scores, backend_name, took_ms) dove scores è allineato a docs
    (logit sigmoide in [0,1]), oppure None se tutti i nodi falliscono.
    """
    chain = _get_chain()
    if not chain or not docs:
        return None
    live, all_down = live_nodes()
    if not live:
        # tutti in cooldown: ritenta comunque il primo (meglio di un fallimento immediato)
        live = [chain[0]]
    payload = {"model": RERANK_MODEL, "query": query, "documents": docs, "top_n": len(docs)}
    started = time.monotonic()
    for node in live:
        headers = {"Content-Type": "application/json"}
        if node.key:
            headers["Authorization"] = f"Bearer {node.key}"
        try:
            t0 = time.monotonic()
            response = await get_http().post(
                f"{node.url}/v1/rerank",
                json=payload,
                headers=headers,
                timeout=httpx.Timeout(node.timeout_ms / 1000.0),
            )
            response.raise_for_status()
            data = response.json()
            # Il risultato è [{index, relevance_score}] ordinato per rilevanza:
            # riportiamo ogni score alla posizione originaria del documento.
            scores = [0.0] * len(docs)
            for item in data.get("results", []):
                idx = int(item["index"])
                if 0 <= idx < len(docs):
                    scores[idx] = float(item.get("relevance_score", 0.0))
            took = int((time.monotonic() - started) * 1000)
            metrics.record_rerank(node.name, True, took)
            return scores, node.name, took
        except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc:
            took = int((time.monotonic() - t0) * 1000)
            _down_until[node.url] = time.monotonic() + RERANK_RETRY_COOLDOWN_S
            metrics.record_rerank(node.name, False, took)
            log.warning(
                "rerank: nodo '%s' fallito dopo %dms (%s: %s) → cooldown %ds",
                node.name,
                took,
                exc.__class__.__name__,
                exc,
                RERANK_RETRY_COOLDOWN_S,
            )
    return None


def normalize_score(logit: float) -> float:
    """Sigmoide: logit di rilevanza → punteggio [0,1] leggibile nei risultati."""
    if logit >= 0:
        z = math.exp(-logit)
        return 1.0 / (1.0 + z)
    z = math.exp(logit)
    return z / (1.0 + z)