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
+12
-1
@@ -23,7 +23,7 @@ GUARDRAIL_BLOCK_THRESHOLD = float(os.environ.get("GUARDRAIL_BLOCK_THRESHOLD", "0
|
||||
GUARDRAIL_WARN_THRESHOLD = float(os.environ.get("GUARDRAIL_WARN_THRESHOLD", "0.70"))
|
||||
GUARDRAIL_VERSION = "similarity-v1"
|
||||
GIT_COMMIT = os.environ.get("GIT_COMMIT", "unknown").strip()
|
||||
GATEWAY_VERSION = os.environ.get("GATEWAY_VERSION", "2.8.0").strip()
|
||||
GATEWAY_VERSION = os.environ.get("GATEWAY_VERSION", "2.9.0").strip()
|
||||
|
||||
VM_PUSH_URL = os.environ.get("VM_PUSH_URL", "http://host.docker.internal:8428/api/v1/import/prometheus")
|
||||
VM_PUSH_INTERVAL = int(os.environ.get("VM_PUSH_INTERVAL", "30"))
|
||||
@@ -33,6 +33,15 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(mess
|
||||
log = logging.getLogger("memory-gateway")
|
||||
|
||||
SPARSE_VECTOR_NAME = "bm25"
|
||||
|
||||
# Re-ranking: catena di fallback resiliente (frigate → brain locale).
|
||||
# Il default nel codice è OFF; il deploy imposta RERANK_ENABLED=true e la catena.
|
||||
RERANK_ENABLED = os.environ.get("RERANK_ENABLED", "false").lower() == "true"
|
||||
RERANK_MODEL = os.environ.get("RERANK_MODEL", "bge-reranker-v2-m3")
|
||||
RERANK_CANDIDATES = int(os.environ.get("RERANK_CANDIDATES", "16"))
|
||||
RERANK_TIMEOUT_MS = int(os.environ.get("RERANK_TIMEOUT_MS", "10000"))
|
||||
RERANK_RETRY_COOLDOWN_S = int(os.environ.get("RERANK_RETRY_COOLDOWN_S", "60"))
|
||||
RERANK_CHAIN = os.environ.get("RERANK_CHAIN", "")
|
||||
_metrics: dict[str, Any] = {
|
||||
"requests": Counter(),
|
||||
"duration_sum": Counter(),
|
||||
@@ -40,4 +49,6 @@ _metrics: dict[str, Any] = {
|
||||
"errors": Counter(),
|
||||
"search_queries": 0,
|
||||
"search_hits": 0,
|
||||
"rerank_calls": Counter(),
|
||||
"rerank_duration_sum": Counter(),
|
||||
}
|
||||
|
||||
@@ -18,6 +18,7 @@ from qdrant_client.http import models as qm
|
||||
import cleanup
|
||||
import embed as embedding
|
||||
import metrics
|
||||
import rerank
|
||||
import state
|
||||
from config import (
|
||||
COLLECTION,
|
||||
@@ -80,6 +81,7 @@ async def _lifespan(_app: FastAPI):
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
await embedding.close_http()
|
||||
await rerank.close_http()
|
||||
|
||||
|
||||
app = FastAPI(title="Memory Gateway", version=GATEWAY_VERSION, lifespan=_lifespan)
|
||||
|
||||
@@ -20,6 +20,11 @@ def record_search(hits: int) -> None:
|
||||
_metrics["search_hits"] += hits
|
||||
|
||||
|
||||
def record_rerank(backend: str, ok: bool, took_ms: int) -> None:
|
||||
_metrics["rerank_calls"][(backend, "ok" if ok else "fail")] += 1
|
||||
_metrics["rerank_duration_sum"][backend] += took_ms
|
||||
|
||||
|
||||
def snapshot(qdrant: Any, collection: str) -> dict:
|
||||
try:
|
||||
points = qdrant.get_collection(collection).points_count
|
||||
@@ -34,6 +39,8 @@ def snapshot(qdrant: Any, collection: str) -> dict:
|
||||
"errors": {f"{endpoint}:{status}": count for (endpoint, status), count in _metrics["errors"].items()},
|
||||
"search_queries": _metrics["search_queries"],
|
||||
"search_hits": _metrics["search_hits"],
|
||||
"rerank_calls": {f"{backend}:{outcome}": count for (backend, outcome), count in _metrics["rerank_calls"].items()},
|
||||
"rerank_avg_ms": {backend: round(total / _metrics["rerank_calls"][(backend, "ok")], 2) for backend, total in _metrics["rerank_duration_sum"].items() if _metrics["rerank_calls"][(backend, "ok")]},
|
||||
"points": points,
|
||||
}
|
||||
|
||||
@@ -50,6 +57,10 @@ def prometheus_lines(qdrant: Any, collection: str) -> list[str]:
|
||||
lines.append(f'qmem_errors_total{{endpoint="{endpoint}",status="{status}"}} {count}')
|
||||
lines.append(f"qmem_search_queries_total {_metrics['search_queries']}")
|
||||
lines.append(f"qmem_search_hits_total {_metrics['search_hits']}")
|
||||
for (backend, outcome), count in _metrics["rerank_calls"].items():
|
||||
lines.append(f'qmem_rerank_calls_total{{backend="{backend}",outcome="{outcome}"}} {count}')
|
||||
for backend, s in _metrics["rerank_duration_sum"].items():
|
||||
lines.append(f'qmem_rerank_duration_seconds_sum{{backend="{backend}"}} {s / 1000:.6f}')
|
||||
try:
|
||||
lines.append(f"qmem_points {qdrant.get_collection(collection).points_count}")
|
||||
except Exception: # noqa: BLE001
|
||||
|
||||
@@ -57,3 +57,4 @@ class SearchIn(BaseModel):
|
||||
level: Optional[Literal["L1_ROOT", "L2_SUBTOPIC", "L3_DETAIL"]] = None
|
||||
topic: Optional[str] = None
|
||||
include_private: bool = Field(default=False, description="Includi i record privati (solo ricerche esplicite)")
|
||||
rerank: Optional[bool] = Field(default=None, description="Override per-query dello stadio rerank (None = default server)")
|
||||
|
||||
@@ -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)
|
||||
+33
-4
@@ -13,6 +13,7 @@ from qdrant_client.http import models as qm
|
||||
import config
|
||||
import guardrail
|
||||
import metrics
|
||||
import rerank
|
||||
import state
|
||||
import store
|
||||
from audit import audit, now_iso, require_auth
|
||||
@@ -27,6 +28,7 @@ from config import (
|
||||
GUARDRAIL_VERSION,
|
||||
GUARDRAIL_WARN_THRESHOLD,
|
||||
MAX_TEXT_LEN,
|
||||
RERANK_CANDIDATES,
|
||||
)
|
||||
from models import MemoryIn, SearchIn
|
||||
|
||||
@@ -131,13 +133,40 @@ async def add_memory(
|
||||
|
||||
@router.post("/v1/memories:search")
|
||||
async def search_memories(body: SearchIn, key: str = Depends(require_auth)) -> dict:
|
||||
use_rerank = rerank.enabled() and body.rerank is not False
|
||||
# Con reranking attivo recuperiamo più candidati di top_k per dare margine allo stadio di rerank
|
||||
limit = max(body.top_k, RERANK_CANDIDATES) if use_rerank else body.top_k
|
||||
vector = await state.embed(body.query)
|
||||
sparse = state.sparse_encode(body.query) if body.hybrid else None
|
||||
hits = store.search(state.qdrant, COLLECTION, body, vector, sparse)
|
||||
hits = store.search(state.qdrant, COLLECTION, body, vector, sparse, limit=limit)
|
||||
results = store.format_results(hits)
|
||||
audit(key, "search", query_hash=hashlib.sha256(body.query.encode()).hexdigest()[:16], top_k=body.top_k, min_score=body.min_score, hits=len(results))
|
||||
|
||||
rerank_info: dict = {"enabled": use_rerank, "used": False}
|
||||
if use_rerank and len(results) >= 2:
|
||||
rr = await rerank.rerank(body.query, [r["text"] or "" for r in results])
|
||||
if rr:
|
||||
scores, backend, took_ms = rr
|
||||
for r, s in zip(results, scores):
|
||||
r["rerank_score"] = round(rerank.normalize_score(s), 4)
|
||||
results.sort(key=lambda r: r["rerank_score"], reverse=True)
|
||||
rerank_info.update(used=True, backend=backend, took_ms=took_ms, candidates=len(results))
|
||||
else:
|
||||
rerank_info["reason"] = "tutti i nodi rerank non raggiungibili (ordine di fusione preservato)"
|
||||
elif use_rerank:
|
||||
rerank_info["reason"] = "candidati insufficienti"
|
||||
|
||||
results = results[: body.top_k]
|
||||
audit(
|
||||
key,
|
||||
"search",
|
||||
query_hash=hashlib.sha256(body.query.encode()).hexdigest()[:16],
|
||||
top_k=body.top_k,
|
||||
min_score=body.min_score,
|
||||
hits=len(results),
|
||||
rerank_backend=rerank_info.get("backend"),
|
||||
)
|
||||
metrics.record_search(len(results))
|
||||
return {"results": results, "min_score": body.min_score, "total_hits": len(results)}
|
||||
return {"results": results, "min_score": body.min_score, "total_hits": len(results), "rerank": rerank_info}
|
||||
|
||||
|
||||
@router.get("/v1/memories/{memory_id}")
|
||||
@@ -218,7 +247,7 @@ async def status(request: Request) -> dict:
|
||||
|
||||
@router.get("/v1/version")
|
||||
async def version() -> dict:
|
||||
return {"version": GATEWAY_VERSION, "git_commit": __import__("config").GIT_COMMIT, "guardrail_version": GUARDRAIL_VERSION, "guardrail_enabled": GUARDRAIL_ENABLED, "guardrail_block_threshold": GUARDRAIL_BLOCK_THRESHOLD, "guardrail_warn_threshold": GUARDRAIL_WARN_THRESHOLD, "embedding_model": EMBED_MODEL, "collection": COLLECTION}
|
||||
return {"version": GATEWAY_VERSION, "git_commit": __import__("config").GIT_COMMIT, "guardrail_version": GUARDRAIL_VERSION, "guardrail_enabled": GUARDRAIL_ENABLED, "guardrail_block_threshold": GUARDRAIL_BLOCK_THRESHOLD, "guardrail_warn_threshold": GUARDRAIL_WARN_THRESHOLD, "embedding_model": EMBED_MODEL, "collection": COLLECTION, "rerank_enabled": rerank.enabled(), "rerank_model": config.RERANK_MODEL, "rerank_nodes": [n.name for n in rerank._get_chain()]}
|
||||
|
||||
|
||||
@router.get("/v1/metrics")
|
||||
|
||||
+8
-6
@@ -2,7 +2,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from typing import Any
|
||||
from typing import Any, Optional
|
||||
|
||||
from qdrant_client.http import models as qm
|
||||
|
||||
@@ -27,25 +27,27 @@ def search_filter(body: SearchIn) -> qm.Filter | None:
|
||||
return None
|
||||
|
||||
|
||||
def search(qdrant: Any, collection: str, body: SearchIn, vector: list[float], sparse: Any) -> list[Any]:
|
||||
def search(qdrant: Any, collection: str, body: SearchIn, vector: list[float], sparse: Any, limit: Optional[int] = None) -> list[Any]:
|
||||
"""Ricerca ibrida o densa. Con reranking attivo limit > top_k per dare candidati extra allo stadio di rerank."""
|
||||
eff_limit = limit if limit is not None else body.top_k
|
||||
qfilter = search_filter(body)
|
||||
if body.hybrid and sparse is not None:
|
||||
return qdrant.query_points(
|
||||
collection_name=collection,
|
||||
prefetch=[
|
||||
qm.Prefetch(query=vector, using="", limit=body.top_k * 4, score_threshold=body.min_score),
|
||||
qm.Prefetch(query=sparse, using=SPARSE_VECTOR_NAME, limit=body.top_k * 4),
|
||||
qm.Prefetch(query=vector, using="", limit=max(body.top_k * 4, eff_limit), score_threshold=body.min_score),
|
||||
qm.Prefetch(query=sparse, using=SPARSE_VECTOR_NAME, limit=max(body.top_k * 4, eff_limit)),
|
||||
],
|
||||
query=qm.FusionQuery(fusion=qm.Fusion.RRF),
|
||||
query_filter=qfilter,
|
||||
limit=body.top_k,
|
||||
limit=eff_limit,
|
||||
with_payload=True,
|
||||
).points
|
||||
return qdrant.query_points(
|
||||
collection_name=collection,
|
||||
query=vector,
|
||||
query_filter=qfilter,
|
||||
limit=body.top_k,
|
||||
limit=eff_limit,
|
||||
score_threshold=body.min_score,
|
||||
with_payload=True,
|
||||
).points
|
||||
|
||||
@@ -0,0 +1,205 @@
|
||||
"""Test dello stadio rerank: catena di fallback, cooldown, degrada con grazia."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import rerank
|
||||
from test_api import auth_headers, make_record
|
||||
|
||||
CHAIN = json.dumps(
|
||||
[
|
||||
{"name": "primario", "url": "http://primario:9002", "key": "k1", "timeout_ms": 500},
|
||||
{"name": "fallback", "url": "http://fallback:9003", "key": "k2", "timeout_ms": 5000},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_chain(monkeypatch):
|
||||
rerank.reset_chain_cache()
|
||||
rerank._http = None
|
||||
yield
|
||||
rerank.reset_chain_cache()
|
||||
rerank._http = None
|
||||
|
||||
|
||||
def _use_chain(monkeypatch, raw=CHAIN):
|
||||
# rerank.py importa i valori di config con `from config import`: si patchano
|
||||
# gli attributi del modulo rerank, non config.
|
||||
monkeypatch.setattr(rerank, "RERANK_CHAIN", raw)
|
||||
rerank.reset_chain_cache()
|
||||
|
||||
|
||||
def _mock_client(handler) -> list[str]:
|
||||
"""Client con transport mockato; ritorna la lista in cui registrare le chiamate."""
|
||||
calls: list[str] = []
|
||||
|
||||
def tracking_handler(request):
|
||||
calls.append(f"{request.url.host}{request.url.path}")
|
||||
return handler(request)
|
||||
|
||||
rerank._http = httpx.AsyncClient(transport=httpx.MockTransport(tracking_handler))
|
||||
return calls
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# parse della catena
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_parse_chain_valida(monkeypatch):
|
||||
_use_chain(monkeypatch)
|
||||
nodes = rerank._get_chain()
|
||||
assert [n.name for n in nodes] == ["primario", "fallback"]
|
||||
assert nodes[0].key == "k1"
|
||||
assert nodes[0].timeout_ms == 500
|
||||
assert nodes[1].timeout_ms == 5000
|
||||
|
||||
|
||||
def test_parse_chain_json_invalido(monkeypatch):
|
||||
_use_chain(monkeypatch, raw="non-json")
|
||||
assert rerank._get_chain() == []
|
||||
assert not rerank.enabled()
|
||||
|
||||
|
||||
def test_parse_chain_scarta_url_senza_schema(monkeypatch):
|
||||
_use_chain(monkeypatch, raw=json.dumps([{"name": "x", "url": "primario:9002"}]))
|
||||
assert rerank._get_chain() == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# fallback e cooldown
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_fallback_primario_500(monkeypatch):
|
||||
_use_chain(monkeypatch)
|
||||
|
||||
def handler(request):
|
||||
if request.url.host == "primario":
|
||||
return httpx.Response(500)
|
||||
return httpx.Response(200, json={"results": [{"index": 1, "relevance_score": 2.0}, {"index": 0, "relevance_score": -1.0}]})
|
||||
|
||||
_mock_client(handler)
|
||||
scores, backend, _took = asyncio.run(rerank.rerank("q", ["docA", "docB"]))
|
||||
assert backend == "fallback"
|
||||
# gli score tornano allineati alla posizione originaria dei documenti
|
||||
assert scores == pytest.approx([-1.0, 2.0])
|
||||
|
||||
|
||||
def test_cooldown_salta_nodo_fallito(monkeypatch):
|
||||
_use_chain(monkeypatch)
|
||||
|
||||
def handler(request):
|
||||
if request.url.host == "primario":
|
||||
return httpx.Response(500)
|
||||
return httpx.Response(200, json={"results": [{"index": 0, "relevance_score": 1.0}]})
|
||||
|
||||
_mock_client(handler)
|
||||
asyncio.run(rerank.rerank("q", ["a", "b"]))
|
||||
scores, backend, _ = asyncio.run(rerank.rerank("q", ["a", "b"]))
|
||||
assert backend == "fallback" # il primario è in cooldown e non viene richiamato
|
||||
|
||||
|
||||
def test_cooldown_non_blocca_per_sempre(monkeypatch):
|
||||
_use_chain(monkeypatch)
|
||||
calls: list[str] = []
|
||||
|
||||
def handler(request):
|
||||
calls.append(request.url.host)
|
||||
if request.url.host == "primario":
|
||||
return httpx.Response(500)
|
||||
return httpx.Response(200, json={"results": [{"index": 0, "relevance_score": 1.0}]})
|
||||
|
||||
_mock_client(handler)
|
||||
asyncio.run(rerank.rerank("q", ["a"]))
|
||||
# svuota il cooldown: il primario torna eleggibile
|
||||
rerank._down_until.clear()
|
||||
asyncio.run(rerank.rerank("q", ["a"]))
|
||||
assert calls == ["primario", "fallback", "primario", "fallback"]
|
||||
|
||||
|
||||
def test_tutti_nodi_falliti_restifica_none(monkeypatch):
|
||||
_use_chain(monkeypatch)
|
||||
_mock_client(lambda request: httpx.Response(500))
|
||||
assert asyncio.run(rerank.rerank("q", ["a", "b"])) is None
|
||||
|
||||
|
||||
def test_enabled_richiede_catena(monkeypatch):
|
||||
_use_chain(monkeypatch, raw="")
|
||||
assert not rerank.enabled()
|
||||
_use_chain(monkeypatch)
|
||||
monkeypatch.setattr(rerank, "RERANK_ENABLED", True)
|
||||
assert rerank.enabled()
|
||||
monkeypatch.setattr(rerank, "RERANK_ENABLED", False)
|
||||
assert not rerank.enabled()
|
||||
|
||||
|
||||
def test_normalize_score():
|
||||
assert rerank.normalize_score(0.0) == pytest.approx(0.5)
|
||||
assert rerank.normalize_score(10.0) > 0.99
|
||||
assert rerank.normalize_score(-10.0) < 0.01
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# integrazione endpoint search
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_search_rerank_riordina(client, monkeypatch):
|
||||
for text in ["alpha", "beta", "gamma"]:
|
||||
r = client.post("/v1/memories", json=make_record(text=text), headers=auth_headers())
|
||||
assert r.status_code == 200
|
||||
|
||||
async def fake_rerank(query, docs):
|
||||
# inverte: beta (index 1) primo, poi gamma/alpha
|
||||
return [0.1, 0.9, 0.5], "finto", 12
|
||||
|
||||
monkeypatch.setattr(rerank, "enabled", lambda: True)
|
||||
monkeypatch.setattr(rerank, "rerank", fake_rerank)
|
||||
|
||||
resp = client.post("/v1/memories:search", json={"query": "q", "top_k": 3}, headers=auth_headers())
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["rerank"]["used"] is True
|
||||
assert data["rerank"]["backend"] == "finto"
|
||||
texts = [r["text"] for r in data["results"]]
|
||||
assert texts == ["beta", "gamma", "alpha"]
|
||||
assert data["results"][0]["rerank_score"] == pytest.approx(rerank.normalize_score(0.9), abs=0.01)
|
||||
|
||||
|
||||
def test_search_rerank_disattivato_per_query(client, monkeypatch):
|
||||
for text in ["alpha", "beta"]:
|
||||
client.post("/v1/memories", json=make_record(text=text), headers=auth_headers())
|
||||
|
||||
async def fail_rerank(query, docs):
|
||||
raise AssertionError("rerank non deve essere chiamato con rerank=false")
|
||||
|
||||
monkeypatch.setattr(rerank, "enabled", lambda: True)
|
||||
monkeypatch.setattr(rerank, "rerank", fail_rerank)
|
||||
|
||||
resp = client.post("/v1/memories:search", json={"query": "q", "top_k": 2, "rerank": False}, headers=auth_headers())
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["rerank"]["enabled"] is False
|
||||
|
||||
|
||||
def test_search_rerank_fallito_degrada_con_grazia(client, monkeypatch):
|
||||
for text in ["alpha", "beta"]:
|
||||
client.post("/v1/memories", json=make_record(text=text), headers=auth_headers())
|
||||
|
||||
async def fail_rerank(query, docs):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(rerank, "enabled", lambda: True)
|
||||
monkeypatch.setattr(rerank, "rerank", fail_rerank)
|
||||
|
||||
resp = client.post("/v1/memories:search", json={"query": "q", "top_k": 2}, headers=auth_headers())
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["rerank"]["used"] is False
|
||||
assert "non raggiungibili" in data["rerank"]["reason"]
|
||||
assert len(data["results"]) == 2 # ordine di fusione preservato
|
||||
Reference in New Issue
Block a user