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:
+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_WARN_THRESHOLD = float(os.environ.get("GUARDRAIL_WARN_THRESHOLD", "0.70"))
|
||||||
GUARDRAIL_VERSION = "similarity-v1"
|
GUARDRAIL_VERSION = "similarity-v1"
|
||||||
GIT_COMMIT = os.environ.get("GIT_COMMIT", "unknown").strip()
|
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_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"))
|
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")
|
log = logging.getLogger("memory-gateway")
|
||||||
|
|
||||||
SPARSE_VECTOR_NAME = "bm25"
|
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] = {
|
_metrics: dict[str, Any] = {
|
||||||
"requests": Counter(),
|
"requests": Counter(),
|
||||||
"duration_sum": Counter(),
|
"duration_sum": Counter(),
|
||||||
@@ -40,4 +49,6 @@ _metrics: dict[str, Any] = {
|
|||||||
"errors": Counter(),
|
"errors": Counter(),
|
||||||
"search_queries": 0,
|
"search_queries": 0,
|
||||||
"search_hits": 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 cleanup
|
||||||
import embed as embedding
|
import embed as embedding
|
||||||
import metrics
|
import metrics
|
||||||
|
import rerank
|
||||||
import state
|
import state
|
||||||
from config import (
|
from config import (
|
||||||
COLLECTION,
|
COLLECTION,
|
||||||
@@ -80,6 +81,7 @@ async def _lifespan(_app: FastAPI):
|
|||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
pass
|
pass
|
||||||
await embedding.close_http()
|
await embedding.close_http()
|
||||||
|
await rerank.close_http()
|
||||||
|
|
||||||
|
|
||||||
app = FastAPI(title="Memory Gateway", version=GATEWAY_VERSION, lifespan=_lifespan)
|
app = FastAPI(title="Memory Gateway", version=GATEWAY_VERSION, lifespan=_lifespan)
|
||||||
|
|||||||
@@ -20,6 +20,11 @@ def record_search(hits: int) -> None:
|
|||||||
_metrics["search_hits"] += hits
|
_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:
|
def snapshot(qdrant: Any, collection: str) -> dict:
|
||||||
try:
|
try:
|
||||||
points = qdrant.get_collection(collection).points_count
|
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()},
|
"errors": {f"{endpoint}:{status}": count for (endpoint, status), count in _metrics["errors"].items()},
|
||||||
"search_queries": _metrics["search_queries"],
|
"search_queries": _metrics["search_queries"],
|
||||||
"search_hits": _metrics["search_hits"],
|
"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,
|
"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_errors_total{{endpoint="{endpoint}",status="{status}"}} {count}')
|
||||||
lines.append(f"qmem_search_queries_total {_metrics['search_queries']}")
|
lines.append(f"qmem_search_queries_total {_metrics['search_queries']}")
|
||||||
lines.append(f"qmem_search_hits_total {_metrics['search_hits']}")
|
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:
|
try:
|
||||||
lines.append(f"qmem_points {qdrant.get_collection(collection).points_count}")
|
lines.append(f"qmem_points {qdrant.get_collection(collection).points_count}")
|
||||||
except Exception: # noqa: BLE001
|
except Exception: # noqa: BLE001
|
||||||
|
|||||||
@@ -57,3 +57,4 @@ class SearchIn(BaseModel):
|
|||||||
level: Optional[Literal["L1_ROOT", "L2_SUBTOPIC", "L3_DETAIL"]] = None
|
level: Optional[Literal["L1_ROOT", "L2_SUBTOPIC", "L3_DETAIL"]] = None
|
||||||
topic: Optional[str] = None
|
topic: Optional[str] = None
|
||||||
include_private: bool = Field(default=False, description="Includi i record privati (solo ricerche esplicite)")
|
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 config
|
||||||
import guardrail
|
import guardrail
|
||||||
import metrics
|
import metrics
|
||||||
|
import rerank
|
||||||
import state
|
import state
|
||||||
import store
|
import store
|
||||||
from audit import audit, now_iso, require_auth
|
from audit import audit, now_iso, require_auth
|
||||||
@@ -27,6 +28,7 @@ from config import (
|
|||||||
GUARDRAIL_VERSION,
|
GUARDRAIL_VERSION,
|
||||||
GUARDRAIL_WARN_THRESHOLD,
|
GUARDRAIL_WARN_THRESHOLD,
|
||||||
MAX_TEXT_LEN,
|
MAX_TEXT_LEN,
|
||||||
|
RERANK_CANDIDATES,
|
||||||
)
|
)
|
||||||
from models import MemoryIn, SearchIn
|
from models import MemoryIn, SearchIn
|
||||||
|
|
||||||
@@ -131,13 +133,40 @@ async def add_memory(
|
|||||||
|
|
||||||
@router.post("/v1/memories:search")
|
@router.post("/v1/memories:search")
|
||||||
async def search_memories(body: SearchIn, key: str = Depends(require_auth)) -> dict:
|
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)
|
vector = await state.embed(body.query)
|
||||||
sparse = state.sparse_encode(body.query) if body.hybrid else None
|
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)
|
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))
|
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}")
|
@router.get("/v1/memories/{memory_id}")
|
||||||
@@ -218,7 +247,7 @@ async def status(request: Request) -> dict:
|
|||||||
|
|
||||||
@router.get("/v1/version")
|
@router.get("/v1/version")
|
||||||
async def version() -> dict:
|
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")
|
@router.get("/v1/metrics")
|
||||||
|
|||||||
+8
-6
@@ -2,7 +2,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import hashlib
|
import hashlib
|
||||||
from typing import Any
|
from typing import Any, Optional
|
||||||
|
|
||||||
from qdrant_client.http import models as qm
|
from qdrant_client.http import models as qm
|
||||||
|
|
||||||
@@ -27,25 +27,27 @@ def search_filter(body: SearchIn) -> qm.Filter | None:
|
|||||||
return 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)
|
qfilter = search_filter(body)
|
||||||
if body.hybrid and sparse is not None:
|
if body.hybrid and sparse is not None:
|
||||||
return qdrant.query_points(
|
return qdrant.query_points(
|
||||||
collection_name=collection,
|
collection_name=collection,
|
||||||
prefetch=[
|
prefetch=[
|
||||||
qm.Prefetch(query=vector, using="", limit=body.top_k * 4, score_threshold=body.min_score),
|
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=body.top_k * 4),
|
qm.Prefetch(query=sparse, using=SPARSE_VECTOR_NAME, limit=max(body.top_k * 4, eff_limit)),
|
||||||
],
|
],
|
||||||
query=qm.FusionQuery(fusion=qm.Fusion.RRF),
|
query=qm.FusionQuery(fusion=qm.Fusion.RRF),
|
||||||
query_filter=qfilter,
|
query_filter=qfilter,
|
||||||
limit=body.top_k,
|
limit=eff_limit,
|
||||||
with_payload=True,
|
with_payload=True,
|
||||||
).points
|
).points
|
||||||
return qdrant.query_points(
|
return qdrant.query_points(
|
||||||
collection_name=collection,
|
collection_name=collection,
|
||||||
query=vector,
|
query=vector,
|
||||||
query_filter=qfilter,
|
query_filter=qfilter,
|
||||||
limit=body.top_k,
|
limit=eff_limit,
|
||||||
score_threshold=body.min_score,
|
score_threshold=body.min_score,
|
||||||
with_payload=True,
|
with_payload=True,
|
||||||
).points
|
).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