From f33bdaa52d382a82d4db848212caaf65bee108e7 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 13 Jul 2026 09:38:17 +0100 Subject: [PATCH] fix: minor usage tracking improvements --- application/api/v1/routes.py | 11 ++-- application/llm/base.py | 18 +++--- application/llm/openai.py | 46 +++++++++++++ application/usage.py | 14 +++- tests/api/v1/test_routes_extended.py | 25 +++++++ tests/llm/test_openai_responses.py | 97 ++++++++++++++++++++++++++++ 6 files changed, 197 insertions(+), 14 deletions(-) diff --git a/application/api/v1/routes.py b/application/api/v1/routes.py index 61389d5e..f6970471 100644 --- a/application/api/v1/routes.py +++ b/application/api/v1/routes.py @@ -121,10 +121,13 @@ def _conversation_belongs_to_agent( def _response_usage(agent: Any) -> Dict[str, Any]: - """Return the upstream provider's latest OpenAI-shaped usage object.""" - usage = getattr(getattr(agent, "llm", None), "_last_usage", None) - if isinstance(usage, dict): - return usage + """Return the turn's cumulative usage in Chat Completions shape. + + Reads the per-instance accumulator so multi-round tool turns report + the sum of every LLM call, matching what ``token_usage`` rows bill; + the accumulator carries provider-exact counts whenever the upstream + reported them (see ``_prefer_provider_usage``). + """ tokens = getattr(getattr(agent, "llm", None), "token_usage", {}) or {} prompt = int(tokens.get("prompt_tokens", 0) or 0) completion = int(tokens.get("generated_tokens", 0) or 0) diff --git a/application/llm/base.py b/application/llm/base.py index f608c0c8..79c8397b 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -317,10 +317,12 @@ class BaseLLM(ABC): ): # Non-streaming counterpart to ``_emit_stream_finished_log``. Paired # with ``llm_gen_start`` so cost dashboards can join start/finish for - # non-streaming calls just as they do for streams. Token counts are - # client-side estimates from ``gen_token_usage``; ``status`` is - # ``"error"`` when the call raised. A distinct event name keeps - # non-stream calls out of stream dashboards. + # non-streaming calls just as they do for streams. Token counts come + # from ``gen_token_usage``: provider-exact when the vendor reported + # usage (OpenAI-family chat + Responses), tiktoken estimates + # otherwise; ``status`` is ``"error"`` when the call raised. A + # distinct event name keeps non-stream calls out of stream + # dashboards. extra = { "model": model, "provider": self.provider_name, @@ -361,10 +363,10 @@ class BaseLLM(ABC): error=None, ): # Paired with ``llm_stream_start`` so cost dashboards can sum tokens - # by user/agent/provider. Token counts are client-side estimates - # from ``stream_token_usage``; vendor-reported counts (incl. - # ``cached_tokens`` for prompt caching) require per-provider - # extraction in each ``_raw_gen_stream`` and aren't wired yet. + # by user/agent/provider. Token counts come from + # ``stream_token_usage``: provider-exact when the vendor reported + # usage (OpenAI-family via ``stream_options.include_usage`` and the + # Responses API), tiktoken estimates for providers that don't. extra = { "model": model, "provider": self.provider_name, diff --git a/application/llm/openai.py b/application/llm/openai.py index b0b6b65d..d2d2ec9a 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -442,8 +442,10 @@ class OpenAILLM(BaseLLM): request_params["tools"] = tools if response_format: request_params["response_format"] = response_format + self._last_usage = None response = self.client.chat.completions.create(**request_params) logging.debug("OpenAI request completed") + self._record_chat_usage(getattr(response, "usage", None)) if tools: return response.choices[0] else: @@ -506,11 +508,19 @@ class OpenAILLM(BaseLLM): request_params["tools"] = tools if response_format: request_params["response_format"] = response_format + # Ask for the terminal usage-only chunk (choices=[]) so token rows + # get provider-exact counts; servers that ignore stream_options + # leave the tiktoken-estimate fallback in place. + stream_options = dict(request_params.get("stream_options") or {}) + stream_options.setdefault("include_usage", True) + request_params["stream_options"] = stream_options + self._last_usage = None response = self.client.chat.completions.create(**request_params) try: for line in response: logging.debug(f"OpenAI stream line: {line}") + self._record_chat_usage(getattr(line, "usage", None)) if not getattr(line, "choices", None): continue @@ -782,6 +792,42 @@ class OpenAILLM(BaseLLM): result["summary"] = serialized return result + def _record_chat_usage(self, usage) -> None: + """Capture provider-reported Chat Completions usage for this call. + + Counterpart to ``_record_responses_metadata`` for the chat path; + provider totals are stored as-is (see ``_prefer_provider_usage``). + """ + if usage is None: + return + try: + prompt = int(getattr(usage, "prompt_tokens", 0) or 0) + completion = int(getattr(usage, "completion_tokens", 0) or 0) + total = int(getattr(usage, "total_tokens", 0) or 0) + except (TypeError, ValueError): + return + if not prompt and not completion: + # A zeroed usage object (some proxies) must not clobber estimates. + return + result = { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": total or prompt + completion, + } + input_details = getattr(usage, "prompt_tokens_details", None) + output_details = getattr(usage, "completion_tokens_details", None) + try: + cached = int(getattr(input_details, "cached_tokens", 0) or 0) + reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0) + except (TypeError, ValueError): + cached = 0 + reasoning = 0 + if cached: + result["prompt_tokens_details"] = {"cached_tokens": cached} + if reasoning: + result["completion_tokens_details"] = {"reasoning_tokens": reasoning} + self._last_usage = result + def _record_responses_metadata(self, response): rid = getattr(response, "id", None) if rid: diff --git a/application/usage.py b/application/usage.py index 2e625031..5f7c55c5 100644 --- a/application/usage.py +++ b/application/usage.py @@ -1,5 +1,6 @@ import logging import time +from typing import Any, Dict from application.storage.db.repositories.token_usage import TokenUsageRepository from application.storage.db.session import db_session @@ -140,8 +141,17 @@ def _persist_call_usage(llm, call_usage): logger.exception("token_usage persist failed") -def _prefer_provider_usage(llm, call_usage): - """Replace estimates with upstream counts when a provider reported them.""" +def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, int]: + """Replace estimates with upstream counts when a provider reported them. + + Invariant: provider totals are billing-parity bins. Upstream + ``prompt_tokens`` already includes cached-read tokens and + ``completion_tokens`` already includes reasoning/refusal tokens, so + they map 1:1 onto our two columns. Never subtract the + ``*_tokens_details`` breakdowns (``cached_tokens``, + ``reasoning_tokens``) back out of these bins — that would break + parity with what providers bill. + """ reported = getattr(llm, "_last_usage", None) if not isinstance(reported, dict): return call_usage diff --git a/tests/api/v1/test_routes_extended.py b/tests/api/v1/test_routes_extended.py index 7b4ed34c..9c247ca5 100644 --- a/tests/api/v1/test_routes_extended.py +++ b/tests/api/v1/test_routes_extended.py @@ -194,6 +194,31 @@ def _patch_v1_db(conn): yield +@pytest.mark.unit +def test_response_usage_reports_cumulative_turn_totals(): + from types import SimpleNamespace + + from application.api.v1.routes import _response_usage + + # Multi-round tool turns must report the accumulator's turn total, + # not the final round's provider snapshot. + agent = SimpleNamespace( + llm=SimpleNamespace( + _last_usage={ + "prompt_tokens": 5, + "completion_tokens": 1, + "total_tokens": 6, + }, + token_usage={"prompt_tokens": 30, "generated_tokens": 12}, + ) + ) + assert _response_usage(agent) == { + "prompt_tokens": 30, + "completion_tokens": 12, + "total_tokens": 42, + } + + @pytest.mark.unit class TestLookupAgentHappy: def test_returns_agent_for_valid_key(self, pg_conn): diff --git a/tests/llm/test_openai_responses.py b/tests/llm/test_openai_responses.py index 2a5eb408..1a9e52d2 100644 --- a/tests/llm/test_openai_responses.py +++ b/tests/llm/test_openai_responses.py @@ -415,6 +415,103 @@ def test_chat_path_drops_tool_controls_when_tools_are_unavailable(monkeypatch): assert "parallel_tool_calls" not in request +@pytest.mark.unit +def test_chat_nonstream_records_provider_usage(monkeypatch): + llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions")) + llm._last_usage = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + llm.client.chat.completions.create = MagicMock( + return_value=_ns( + choices=[_ns(message=_ns(content="ok"))], + usage=_ns(prompt_tokens=11, completion_tokens=5, total_tokens=16), + ) + ) + + result = llm._raw_gen(llm, "model", [{"role": "user", "content": "hi"}]) + + assert result == "ok" + assert llm._last_usage == { + "prompt_tokens": 11, + "completion_tokens": 5, + "total_tokens": 16, + } + + +@pytest.mark.unit +def test_chat_stream_requests_and_records_terminal_usage(monkeypatch): + llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions")) + llm._last_usage = {"prompt_tokens": 9, "completion_tokens": 9, "total_tokens": 18} + llm.client.chat.completions.create = MagicMock( + return_value=[ + _ns(choices=[_ns(delta=_ns(content="Hi"), finish_reason=None)]), + _ns( + choices=[], + usage=_ns(prompt_tokens=7, completion_tokens=2, total_tokens=9), + ), + ] + ) + + out = list( + llm._raw_gen_stream(llm, "model", [{"role": "user", "content": "hi"}]) + ) + + assert out == ["Hi"] + request = llm.client.chat.completions.create.call_args.kwargs + assert request["stream_options"] == {"include_usage": True} + assert llm._last_usage == { + "prompt_tokens": 7, + "completion_tokens": 2, + "total_tokens": 9, + } + + +@pytest.mark.unit +def test_chat_stream_without_usage_clears_stale_counts(monkeypatch): + llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions")) + llm._last_usage = {"prompt_tokens": 9, "completion_tokens": 9, "total_tokens": 18} + llm.client.chat.completions.create = MagicMock( + return_value=[_ns(choices=[_ns(delta=_ns(content="Hi"), finish_reason=None)])] + ) + + assert list( + llm._raw_gen_stream(llm, "model", [{"role": "user", "content": "hi"}]) + ) == ["Hi"] + assert llm._last_usage is None + + +@pytest.mark.unit +def test_zeroed_chat_usage_does_not_clobber_estimates(monkeypatch): + llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions")) + llm.client.chat.completions.create = MagicMock( + return_value=_ns( + choices=[_ns(message=_ns(content="ok"))], + usage=_ns(prompt_tokens=0, completion_tokens=0, total_tokens=0), + ) + ) + + llm._raw_gen(llm, "model", [{"role": "user", "content": "hi"}]) + + assert llm._last_usage is None + + +@pytest.mark.unit +def test_record_chat_usage_captures_detail_bins(monkeypatch): + llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions")) + llm._record_chat_usage(_ns( + prompt_tokens=10, + completion_tokens=7, + total_tokens=17, + prompt_tokens_details=_ns(cached_tokens=4), + completion_tokens_details=_ns(reasoning_tokens=3), + )) + assert llm._last_usage == { + "prompt_tokens": 10, + "completion_tokens": 7, + "total_tokens": 17, + "prompt_tokens_details": {"cached_tokens": 4}, + "completion_tokens_details": {"reasoning_tokens": 3}, + } + + @pytest.mark.unit def test_build_responses_params_store_with_previous_id(monkeypatch): llm = _make_llm(monkeypatch, _responses_caps(), store_responses=True)