feat(gateway): stadio rerank con catena di fallback resiliente (frigate→brain)
- gateway/rerank.py: catena da RERANK_CHAIN (JSON, per-nodo key+timeout),
cooldown 60s sui nodi falliti, score sigmoide [0,1], degrada con grazia
all'ordine di fusione se tutti i nodi sono giù
- routes: /v1/memories:search applica il rerank post-fusione (fetch esteso a
RERANK_CANDIDATES), risposta con rerank{used,backend,took_ms}, flag
per-query rerank=false; /v1/version espone lo stato rerank
- store: search() accetta limit esteso; models: SearchIn.rerank
- metrics: qmem_rerank_calls_total + durata per backend
- test: 10 nuovi (fallback, cooldown, degradazione, integrazione) — 46 pass
This commit is contained in:
enne2
@@ -0,0 +1,175 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user