From 29e76d021cdf1f0c4c78f5b2ef7b7494ca432640 Mon Sep 17 00:00:00 2001 From: wangwei Date: Thu, 2 Jul 2026 14:39:13 +0800 Subject: [PATCH] feat(token-tracking): attach usage hook to advisor's self-created LLM client Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- rag_eval/advisor/llm_analyzer.py | 3 ++- tests/test_advisor_llm_analyzer.py | 33 ++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/rag_eval/advisor/llm_analyzer.py b/rag_eval/advisor/llm_analyzer.py index 749e008..9a59f39 100644 --- a/rag_eval/advisor/llm_analyzer.py +++ b/rag_eval/advisor/llm_analyzer.py @@ -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: diff --git a/tests/test_advisor_llm_analyzer.py b/tests/test_advisor_llm_analyzer.py index 55c47d4..65066b6 100644 --- a/tests/test_advisor_llm_analyzer.py +++ b/tests/test_advisor_llm_analyzer.py @@ -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")