feat(token-tracking): add HTTP response hook and attach_usage_hook, wire into build_models

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
2026-07-02 14:36:59 +08:00
co-authored by Copilot
parent 8b896e4e7f
commit 613d167e81
3 changed files with 2117 additions and 0 deletions
+47
View File
@@ -5,12 +5,15 @@ from __future__ import annotations
import logging
from typing import Any
import httpx
from openai import AsyncOpenAI
from rag_eval.compat import ensure_ragas_import_compat
from rag_eval.settings import EvaluationSettings
from rag_eval.shared.models import Scenario
from .token_tracker import get_current_tracker
ensure_ragas_import_compat()
from ragas.embeddings.base import embedding_factory
@@ -32,6 +35,48 @@ from .pipeline import MetricPipeline
logger = logging.getLogger("rag_eval.metrics.factory")
async def _usage_response_hook(response: httpx.Response) -> None:
"""Record token usage from an OpenAI-compatible HTTP response, if a tracker is active.
Applies to both chat-completions and embeddings responses since both
return top-level `model` and `usage` fields in OpenAI-compatible APIs.
Never raises — a broken/incompatible gateway response must not affect scoring.
"""
tracker = get_current_tracker()
if tracker is None:
return
try:
await response.aread()
data = response.json()
usage = data.get("usage")
if not usage:
# Gateway did not report usage at all — skip rather than record a
# misleading 0/0 call.
return
model = data.get("model") or "unknown"
tracker.record(
model,
int(usage.get("prompt_tokens", 0) or 0),
int(usage.get("completion_tokens", 0) or 0),
)
except Exception: # noqa: BLE001
logger.debug("[factory] usage hook failed to parse response", exc_info=True)
def attach_usage_hook(client: AsyncOpenAI) -> None:
"""Attach the token-usage response hook to an AsyncOpenAI client (idempotent).
Safe to call multiple times on the same client (e.g. when judge and
embedding models share one client) — the hook is only appended once.
"""
httpx_client = getattr(client, "_client", None)
if httpx_client is None or not hasattr(httpx_client, "event_hooks"):
return
hooks = httpx_client.event_hooks.setdefault("response", [])
if _usage_response_hook not in hooks:
hooks.append(_usage_response_hook)
def _resolve_openai_client_kwargs(
model: str,
settings: EvaluationSettings,
@@ -126,8 +171,10 @@ def build_models(
)
llm_client = AsyncOpenAI(**llm_kwargs)
attach_usage_hook(llm_client)
# Only allocate a second client when the embedding model needs different settings.
emb_client = AsyncOpenAI(**emb_kwargs) if emb_kwargs != llm_kwargs else llm_client
attach_usage_hook(emb_client)
# RAGAS structured-output judge calls can be truncated by the upstream default
# 1024 completion budget, especially for faithfulness and GPT-5 family models.