feat(token-tracking): attach usage hook to advisor's self-created LLM client

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
2026-07-02 14:39:13 +08:00
co-authored by Copilot
parent 613d167e81
commit 29e76d021c
2 changed files with 35 additions and 1 deletions
+2 -1
View File
@@ -180,9 +180,10 @@ async def analyze(
if client is None:
from openai import AsyncOpenAI
from rag_eval.metrics.factory import resolve_openai_client_kwargs
from rag_eval.metrics.factory import attach_usage_hook, resolve_openai_client_kwargs
client = AsyncOpenAI(**resolve_openai_client_kwargs(judge_model, settings))
attach_usage_hook(client)
owns_client = True
try:
+33
View File
@@ -132,6 +132,39 @@ def test_analyze_does_not_close_injected_client() -> None:
assert fake.closed is False
def test_analyze_attaches_usage_hook_to_self_created_client(monkeypatch) -> None:
"""A self-created client gets the token-usage hook attached (not the injected-client path)."""
captured: dict = {}
fake = _FakeClient(captured)
hook_calls = []
import openai
import rag_eval.metrics.factory as factory_mod
monkeypatch.setattr(openai, "AsyncOpenAI", lambda **kwargs: fake)
monkeypatch.setattr(
factory_mod, "resolve_openai_client_kwargs", lambda *a, **k: {"api_key": "x"}
)
monkeypatch.setattr(factory_mod, "attach_usage_hook", lambda c: hook_calls.append(c))
asyncio.run(analyze([_diagnosis()], "scn", "gpt-4o", _Settings()))
assert hook_calls == [fake]
def test_analyze_does_not_attach_hook_for_injected_client() -> None:
"""An injected chat_client is assumed to already have the hook attached by its owner."""
hook_calls = []
import rag_eval.metrics.factory as factory_mod
import unittest.mock as mock
with mock.patch.object(factory_mod, "attach_usage_hook", lambda c: hook_calls.append(c)):
fake = _FakeClient({})
asyncio.run(analyze([_diagnosis()], "scn", "gpt-4o", _Settings(), chat_client=fake))
assert hook_calls == []
def test_is_reasoning_model_detection() -> None:
assert _is_reasoning_model("gpt-5")
assert _is_reasoning_model("gpt-5.5")