feat: record reranker call outcome into ModelUsageTracker

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
2026-07-02 15:42:20 +08:00
co-authored by Copilot
parent 41096369d3
commit 66fc388bfb
2 changed files with 67 additions and 0 deletions
@@ -9,6 +9,7 @@ from loguru import logger
from app.config.settings import settings from app.config.settings import settings
from app.domain.retrieval import Reranker, RetrievedChunk from app.domain.retrieval import Reranker, RetrievedChunk
from app.shared.model_usage_tracker import get_model_usage_tracker
class OpenAICompatibleReranker(Reranker): class OpenAICompatibleReranker(Reranker):
@@ -37,10 +38,26 @@ class OpenAICompatibleReranker(Reranker):
scores = self._call_reranker(query, texts) scores = self._call_reranker(query, texts)
except Exception as exc: except Exception as exc:
logger.warning("Reranker call failed ({}), falling back to original order: {}", type(exc).__name__, exc) logger.warning("Reranker call failed ({}), falling back to original order: {}", type(exc).__name__, exc)
# Record the failure so the Status page reflects real reranker health.
get_model_usage_tracker().record(
provider="reranker",
model=self._model,
success=False,
latency_ms=int((time.time() - start) * 1000),
error=str(exc),
)
return chunks[:top_k] return chunks[:top_k]
elapsed_ms = int((time.time() - start) * 1000) elapsed_ms = int((time.time() - start) * 1000)
logger.debug("Reranker scored {} chunks in {}ms", len(chunks), elapsed_ms) logger.debug("Reranker scored {} chunks in {}ms", len(chunks), elapsed_ms)
# TEI/Cohere-style rerank responses carry no token usage field —
# only call success/latency is meaningful for this role.
get_model_usage_tracker().record(
provider="reranker",
model=self._model,
success=True,
latency_ms=elapsed_ms,
)
ranked = sorted( ranked = sorted(
[(score, chunk) for score, chunk in zip(scores, chunks)], [(score, chunk) for score, chunk in zip(scores, chunks)],
@@ -0,0 +1,50 @@
"""Verifies OpenAICompatibleReranker records call outcome (no tokens) into ModelUsageTracker."""
from __future__ import annotations
from unittest.mock import patch
import pytest
from app.domain.retrieval import RetrievedChunk
from app.infrastructure.vectorstore.cross_encoder_reranker import OpenAICompatibleReranker
from app.shared.model_usage_tracker import get_model_usage_tracker
@pytest.fixture(autouse=True)
def _reset_tracker():
"""Clear the process-wide tracker before and after each test in this file."""
get_model_usage_tracker()._entries.clear()
yield
get_model_usage_tracker()._entries.clear()
def _chunk(chunk_id: str, text: str) -> RetrievedChunk:
"""Build a minimal RetrievedChunk for reranker tests."""
return RetrievedChunk(chunk_id=chunk_id, doc_id="doc-1", doc_title="Doc", text=text, score=0.0)
def test_successful_rerank_records_call_without_tokens():
"""A successful rerank() call is recorded with call_count_ok but zero tokens."""
reranker = OpenAICompatibleReranker(base_url="http://example.test", model="bge-reranker-v2-m3")
with patch.object(reranker, "_call_reranker", return_value=[0.9, 0.1]):
result = reranker.rerank("query", [_chunk("c1", "a"), _chunk("c2", "b")], top_k=2)
assert len(result) == 2
entry = get_model_usage_tracker().get("reranker", "bge-reranker-v2-m3")
assert entry is not None
assert entry.call_count_ok == 1
assert entry.total_tokens == 0
def test_failed_rerank_records_error_and_falls_back():
"""A rerank() call that raises internally is recorded as an error but still returns a fallback list."""
reranker = OpenAICompatibleReranker(base_url="http://example.test", model="bge-reranker-v2-m3")
with patch.object(reranker, "_call_reranker", side_effect=RuntimeError("gateway down")):
result = reranker.rerank("query", [_chunk("c1", "a")], top_k=1)
assert len(result) == 1 # existing fallback behavior: original order, unscored
entry = get_model_usage_tracker().get("reranker", "bge-reranker-v2-m3")
assert entry is not None
assert entry.call_count_error == 1
assert entry.status == "error"