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>

rag_eval/metrics/factory.py

@@ -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.

tests/test_token_usage_hook.py

@@ -0,0 +1,75 @@
+"""Tests for the token-usage HTTP response hook and attach_usage_hook wiring."""
+from __future__ import annotations
+
+import asyncio
+import json
+
+import httpx
+
+from rag_eval.metrics.factory import _usage_response_hook, attach_usage_hook
+from rag_eval.metrics.token_tracker import track_token_usage
+
+
+def _fake_response(payload: dict | None) -> httpx.Response:
+    """Build a real httpx.Response with a JSON (or broken) body for hook testing."""
+    content = b"not json" if payload is None else json.dumps(payload).encode("utf-8")
+    return httpx.Response(200, content=content, request=httpx.Request("POST", "http://test/x"))
+
+
+class TestUsageResponseHook:
+    def test_records_usage_when_tracker_active(self):
+        with track_token_usage() as tracker:
+            response = _fake_response({
+                "model": "gpt-5",
+                "usage": {"prompt_tokens": 120, "completion_tokens": 45},
+            })
+            asyncio.run(_usage_response_hook(response))
+        assert tracker.as_dict() == {
+            "gpt-5": {"input_tokens": 120, "output_tokens": 45, "calls": 1}
+        }
+
+    def test_noop_when_no_tracker_active(self):
+        response = _fake_response({"model": "gpt-5", "usage": {"prompt_tokens": 1, "completion_tokens": 1}})
+        # Must not raise even though no tracker is active.
+        asyncio.run(_usage_response_hook(response))
+
+    def test_noop_when_response_has_no_usage_field(self):
+        with track_token_usage() as tracker:
+            response = _fake_response({"model": "gpt-5"})
+            asyncio.run(_usage_response_hook(response))
+        assert tracker.as_dict() == {}
+
+    def test_noop_on_non_json_response(self):
+        with track_token_usage() as tracker:
+            response = _fake_response(None)
+            asyncio.run(_usage_response_hook(response))
+        assert tracker.as_dict() == {}
+
+    def test_embedding_response_without_completion_tokens_defaults_output_to_zero(self):
+        """Embeddings responses omit completion_tokens; output should default to 0."""
+        with track_token_usage() as tracker:
+            response = _fake_response({
+                "model": "Qwen/Qwen3-Embedding-4B",
+                "usage": {"prompt_tokens": 30, "total_tokens": 30},
+            })
+            asyncio.run(_usage_response_hook(response))
+        assert tracker.as_dict() == {
+            "Qwen/Qwen3-Embedding-4B": {"input_tokens": 30, "output_tokens": 0, "calls": 1}
+        }
+
+
+class TestAttachUsageHook:
+    def test_attaches_hook_to_client_event_hooks(self):
+        from openai import AsyncOpenAI
+
+        client = AsyncOpenAI(api_key="sk-test", base_url="http://localhost:1")
+        attach_usage_hook(client)
+        assert _usage_response_hook in client._client.event_hooks["response"]
+
+    def test_idempotent_when_called_twice_on_same_client(self):
+        from openai import AsyncOpenAI
+
+        client = AsyncOpenAI(api_key="sk-test", base_url="http://localhost:1")
+        attach_usage_hook(client)
+        attach_usage_hook(client)
+        assert client._client.event_hooks["response"].count(_usage_response_hook) == 1