diff --git a/application/alembic/versions/0031_token_usage_cache_tokens.py b/application/alembic/versions/0031_token_usage_cache_tokens.py new file mode 100644 index 00000000..1749b9c2 --- /dev/null +++ b/application/alembic/versions/0031_token_usage_cache_tokens.py @@ -0,0 +1,35 @@ +"""0031 token_usage cache tokens — persist the prompt-cache breakdown. + +Providers report how much of each prompt was served from the prompt cache +(``cached_tokens``) and, on newer OpenAI-family models, how much was written +to it (``cache_write_tokens``). The LLM clients already parsed both and the +usage layer discarded them, so ``token_usage`` could not show a cache hit +rate. These two nullable columns carry the breakdown; NULL means the +provider reported nothing (distinct from 0). ``prompt_tokens`` keeps its +meaning as the provider's total. Idempotent both ways. + +Revision ID: 0031_token_usage_cache_tokens +Revises: 0030_superseded_messages +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0031_token_usage_cache_tokens" +down_revision: Union[str, None] = "0030_superseded_messages" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute("ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS cached_tokens integer;") + op.execute( + "ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS cache_write_tokens integer;" + ) + + +def downgrade() -> None: + op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS cache_write_tokens;") + op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS cached_tokens;") diff --git a/application/llm/base.py b/application/llm/base.py index 8e1aabc6..d331ca17 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -710,6 +710,7 @@ class BaseLLM(ABC): completion_tokens, latency_ms, cached_tokens=None, + cache_write_tokens=None, error=None, ): # Non-streaming counterpart to ``_emit_stream_finished_log``. Paired @@ -730,6 +731,8 @@ class BaseLLM(ABC): } if cached_tokens is not None: extra["cached_tokens"] = int(cached_tokens) + if cache_write_tokens is not None: + extra["cache_write_tokens"] = int(cache_write_tokens) if error is not None: extra["error_class"] = type(error).__name__ logging.info("llm_gen_finished", extra=extra) @@ -757,6 +760,7 @@ class BaseLLM(ABC): completion_tokens, latency_ms, cached_tokens=None, + cache_write_tokens=None, error=None, ): # Paired with ``llm_stream_start`` so cost dashboards can sum tokens @@ -774,6 +778,8 @@ class BaseLLM(ABC): } if cached_tokens is not None: extra["cached_tokens"] = int(cached_tokens) + if cache_write_tokens is not None: + extra["cache_write_tokens"] = int(cache_write_tokens) if error is not None: extra["error_class"] = type(error).__name__ logging.info("llm_stream_finished", extra=extra) diff --git a/application/llm/openai.py b/application/llm/openai.py index 1f60e568..65751206 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -1307,17 +1307,36 @@ class OpenAILLM(BaseLLM): output_details = getattr(usage, "completion_tokens_details", None) try: cached = int(getattr(input_details, "cached_tokens", 0) or 0) + written = int(getattr(input_details, "cache_write_tokens", 0) or 0) reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0) except (TypeError, ValueError): cached = 0 + written = 0 reasoning = 0 - if cached: - result["prompt_tokens_details"] = {"cached_tokens": cached} + details = self._prompt_cache_details(cached, written) + if details: + result["prompt_tokens_details"] = details if reasoning: result["completion_tokens_details"] = {"reasoning_tokens": reasoning} self._last_usage = result self._last_usage_claimed = False + @staticmethod + def _prompt_cache_details(cached: int, written: int) -> dict: + """Build the ``prompt_tokens_details`` breakdown from the two cache bins. + + Only non-zero bins are included so a provider that never reports + caching produces no details (and no misleading zero downstream). + ``cache_write_tokens`` is reported by newer OpenAI-family deployments + that charge for cache writes. + """ + details = {} + if cached: + details["cached_tokens"] = cached + if written: + details["cache_write_tokens"] = written + return details + @staticmethod def _function_call_ids(response): """call_ids of every ``function_call`` item in a Responses output. @@ -1357,9 +1376,11 @@ class OpenAILLM(BaseLLM): "total_tokens": int(getattr(usage, "total_tokens", 0) or prompt + completion), } cached = int(getattr(input_details, "cached_tokens", 0) or 0) + written = int(getattr(input_details, "cache_write_tokens", 0) or 0) reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0) - if cached: - result["prompt_tokens_details"] = {"cached_tokens": cached} + details = self._prompt_cache_details(cached, written) + if details: + result["prompt_tokens_details"] = details if reasoning: result["completion_tokens_details"] = {"reasoning_tokens": reasoning} self._last_usage = result diff --git a/application/storage/db/models.py b/application/storage/db/models.py index aa80b19b..a65ea3c1 100644 --- a/application/storage/db/models.py +++ b/application/storage/db/models.py @@ -235,6 +235,12 @@ token_usage_table = Table( # Added in ``0015_token_usage_model_id``. Canonical model id (catalog # name for built-ins, UUID for BYOM); NULL on un-backfilled rows. Column("model_id", Text), + # Added in ``0031_token_usage_cache_tokens``. Prompt-cache breakdowns of + # ``prompt_tokens`` as reported by the provider (reads; and writes on + # newer OpenAI-family models). NULL = not reported, distinct from 0 = no + # cache activity, so hit-rate queries stay honest across providers. + Column("cached_tokens", Integer), + Column("cache_write_tokens", Integer), ) user_logs_table = Table( diff --git a/application/storage/db/repositories/token_usage.py b/application/storage/db/repositories/token_usage.py index c23ec90b..7292dbd1 100644 --- a/application/storage/db/repositories/token_usage.py +++ b/application/storage/db/repositories/token_usage.py @@ -35,6 +35,8 @@ class TokenUsageRepository: request_id: Optional[str] = None, model_id: Optional[str] = None, timestamp: Optional[datetime] = None, + cached_tokens: Optional[int] = None, + cache_write_tokens: Optional[int] = None, ) -> None: # Attribution guard: the ``token_usage_attribution_chk`` CHECK # constraint requires at least one of ``user_id`` / ``api_key`` @@ -60,12 +62,14 @@ class TokenUsageRepository: INSERT INTO token_usage ( user_id, api_key, agent_id, prompt_tokens, generated_tokens, + cached_tokens, cache_write_tokens, source, request_id, model_id, timestamp ) VALUES ( :user_id, :api_key, CAST(:agent_id AS uuid), :prompt_tokens, :generated_tokens, + :cached_tokens, :cache_write_tokens, :source, :request_id, :model_id, COALESCE(:timestamp, now()) ) """ @@ -76,6 +80,8 @@ class TokenUsageRepository: "agent_id": agent_id_uuid, "prompt_tokens": prompt_tokens, "generated_tokens": generated_tokens, + "cached_tokens": cached_tokens, + "cache_write_tokens": cache_write_tokens, "source": source, "request_id": request_id, "model_id": model_id, diff --git a/application/usage.py b/application/usage.py index a1d31d08..2b10c7aa 100644 --- a/application/usage.py +++ b/application/usage.py @@ -131,6 +131,11 @@ def _persist_call_usage(llm, call_usage): agent_id=str(agent_id) if agent_id else None, prompt_tokens=call_usage["prompt_tokens"], generated_tokens=call_usage["generated_tokens"], + # Present only when the provider reported the breakdown; + # persisted as NULL otherwise so "unknown" never reads as + # "0% cache hits". + cached_tokens=call_usage.get("cached_tokens"), + cache_write_tokens=call_usage.get("cache_write_tokens"), source=( getattr(llm, "_token_usage_source", None) or "agent_stream" ), @@ -151,6 +156,12 @@ def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, in ``*_tokens_details`` breakdowns (``cached_tokens``, ``reasoning_tokens``) back out of these bins — that would break parity with what providers bill. + + The prompt-cache sub-bins ARE carried alongside (``cached_tokens``, + ``cache_write_tokens``; Anthropic's ``cache_creation_tokens`` maps to + the latter) so persistence and the finish events can chart them. They + are added only when the provider reported them. The rest of + ``call_usage`` (e.g. ``model``) is preserved rather than replaced. """ reported = getattr(llm, "_last_usage", None) if not isinstance(reported, dict): @@ -174,10 +185,22 @@ def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, in # Slotted/immutable LLM stand-ins can't record the claim; this # call still gets the provider counts, which is correct for them. pass - return { + merged = { + **call_usage, "prompt_tokens": int(prompt or 0), "generated_tokens": int(completion or 0), } + details = reported.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + written = details.get("cache_write_tokens") + if written is None: + written = details.get("cache_creation_tokens") + if cached is not None: + merged["cached_tokens"] = int(cached or 0) + if written is not None: + merged["cache_write_tokens"] = int(written or 0) + return merged def gen_token_usage(func): @@ -224,6 +247,8 @@ def gen_token_usage(func): prompt_tokens=call_usage["prompt_tokens"], completion_tokens=call_usage["generated_tokens"], latency_ms=int((time.monotonic() - started_at) * 1000), + cached_tokens=call_usage.get("cached_tokens"), + cache_write_tokens=call_usage.get("cache_write_tokens"), error=error, ) except Exception: @@ -272,6 +297,8 @@ def stream_token_usage(func): prompt_tokens=call_usage["prompt_tokens"], completion_tokens=call_usage["generated_tokens"], latency_ms=int((time.monotonic() - started_at) * 1000), + cached_tokens=call_usage.get("cached_tokens"), + cache_write_tokens=call_usage.get("cache_write_tokens"), error=error, ) except Exception: diff --git a/tests/llm/handlers/test_round_usage_persistence.py b/tests/llm/handlers/test_round_usage_persistence.py index 2ada0b9e..f45d418a 100644 --- a/tests/llm/handlers/test_round_usage_persistence.py +++ b/tests/llm/handlers/test_round_usage_persistence.py @@ -179,8 +179,12 @@ class TestPerRoundUsageRows: list(handler.handle_streaming(agent, first, {}, [])) assert len(rows) == 2 - assert rows[0] == {"prompt_tokens": 100, "generated_tokens": 10} - assert rows[1] == {"prompt_tokens": 200, "generated_tokens": 20} + # Provider-reported counts replace the estimates; the rest of the + # call record (e.g. ``model``) rides along, hence a subset check + # rather than dict equality. + counts = lambda row: {k: row[k] for k in ("prompt_tokens", "generated_tokens")} # noqa: E731 + assert counts(rows[0]) == {"prompt_tokens": 100, "generated_tokens": 10} + assert counts(rows[1]) == {"prompt_tokens": 200, "generated_tokens": 20} class TestPreferProviderUsageClaim: diff --git a/tests/llm/test_base.py b/tests/llm/test_base.py index 518dfa67..67ce9254 100644 --- a/tests/llm/test_base.py +++ b/tests/llm/test_base.py @@ -2347,3 +2347,38 @@ class TestBuildConversationFromMessagesEmpty: result = handler._build_conversation_from_messages(messages) assert result is not None assert len(result.get("queries", [])) >= 1 + + +class TestFinishedLogCacheBins: + """``cached_tokens`` / ``cache_write_tokens`` ride on the finish events + only when a provider reported them, so dashboards can compute a cache + hit rate without a sentinel zero polluting providers that never report.""" + + def test_stream_finished_carries_cache_bins_when_given(self, caplog): + import logging as _logging + + llm = StubLLM() + with caplog.at_level(_logging.INFO, logger="root"): + llm._emit_stream_finished_log( + "m1", + prompt_tokens=100, + completion_tokens=5, + latency_ms=10, + cached_tokens=80, + cache_write_tokens=7, + ) + evt = next(r for r in caplog.records if r.message == "llm_stream_finished") + assert evt.cached_tokens == 80 + assert evt.cache_write_tokens == 7 + + def test_gen_finished_omits_cache_bins_when_none(self, caplog): + import logging as _logging + + llm = StubLLM() + with caplog.at_level(_logging.INFO, logger="root"): + llm._emit_gen_finished_log( + "m1", prompt_tokens=100, completion_tokens=5, latency_ms=10 + ) + evt = next(r for r in caplog.records if r.message == "llm_gen_finished") + assert not hasattr(evt, "cached_tokens") + assert not hasattr(evt, "cache_write_tokens") diff --git a/tests/llm/test_openai_responses.py b/tests/llm/test_openai_responses.py index 9ae2e645..ab00a543 100644 --- a/tests/llm/test_openai_responses.py +++ b/tests/llm/test_openai_responses.py @@ -1550,3 +1550,38 @@ def test_responses_gen_stream_retries_without_tools(monkeypatch): assert len(create.calls) == 2 assert "tools" in create.calls[0] assert "tools" not in create.calls[1] + + +@pytest.mark.unit +def test_record_chat_usage_captures_cache_write_bin(monkeypatch): + """Newer OpenAI-family deployments report cache writes next to cache reads.""" + 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, cache_write_tokens=6), + completion_tokens_details=None, + )) + assert llm._last_usage["prompt_tokens_details"] == { + "cached_tokens": 4, + "cache_write_tokens": 6, + } + + +@pytest.mark.unit +def test_record_responses_metadata_captures_cache_write_bin(monkeypatch): + llm = _make_llm(monkeypatch, _responses_caps()) + llm._record_responses_metadata(_ns( + id="resp_1", + usage=_ns( + input_tokens=10, + output_tokens=7, + total_tokens=17, + input_tokens_details=_ns(cached_tokens=0, cache_write_tokens=9), + output_tokens_details=None, + ), + )) + # A pure cache-write turn (first request of a prefix) has reads=0 and + # writes>0; the details must still be recorded. + assert llm._last_usage["prompt_tokens_details"] == {"cache_write_tokens": 9} diff --git a/tests/storage/db/repositories/test_token_usage.py b/tests/storage/db/repositories/test_token_usage.py index c4781b4c..681548ef 100644 --- a/tests/storage/db/repositories/test_token_usage.py +++ b/tests/storage/db/repositories/test_token_usage.py @@ -290,3 +290,32 @@ class TestBucketedTotals: ) assert len(rows) == 1 assert rows[0]["prompt_tokens"] == 10 + + +class TestCacheBins: + def _row(self, pg_conn, user_id): + return pg_conn.execute( + text( + "SELECT cached_tokens, cache_write_tokens FROM token_usage " + "WHERE user_id = :u ORDER BY id DESC LIMIT 1" + ), + {"u": user_id}, + ).fetchone() + + def test_persists_cache_bins(self, pg_conn): + repo = TokenUsageRepository(pg_conn) + repo.insert( + user_id="cache-user", + prompt_tokens=1000, + generated_tokens=10, + cached_tokens=800, + cache_write_tokens=100, + ) + assert tuple(self._row(pg_conn, "cache-user")) == (800, 100) + + def test_cache_bins_default_to_null_not_zero(self, pg_conn): + """NULL means "provider did not report"; 0 means "reported no cache + activity". Keeping them distinct is what makes a hit-rate query honest.""" + repo = TokenUsageRepository(pg_conn) + repo.insert(user_id="cache-user-2", prompt_tokens=10, generated_tokens=1) + assert tuple(self._row(pg_conn, "cache-user-2")) == (None, None) diff --git a/tests/storage/db/test_migration_0031.py b/tests/storage/db/test_migration_0031.py new file mode 100644 index 00000000..1836ac3a --- /dev/null +++ b/tests/storage/db/test_migration_0031.py @@ -0,0 +1,100 @@ +"""Migration round-trip test for 0031_token_usage_cache_tokens.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from sqlalchemy import text + + +pytestmark = pytest.mark.integration + + +def _alembic_ini() -> Path: + return Path(__file__).resolve().parents[3] / "application" / "alembic.ini" + + +def _run_alembic(url: str, *args: str) -> None: + subprocess.check_call( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + ) + + +def _alembic_heads(url: str) -> list[str]: + out = subprocess.check_output( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "heads"], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + text=True, + ) + return [line for line in out.splitlines() if line.strip()] + + +def _alembic_version(conn) -> str: + return conn.execute(text("SELECT version_num FROM alembic_version")).scalar() + + +def _column_exists(conn, table: str, column: str) -> bool: + row = conn.execute( + text( + "SELECT 1 FROM information_schema.columns " + "WHERE table_name = :t AND column_name = :c AND table_schema = 'public'" + ), + {"t": table, "c": column}, + ).fetchone() + return row is not None + + +_0031 = "0031_token_usage_cache_tokens" +_0030 = "0030_superseded_messages" + + +class TestMigration0031RoundTrip: + def test_single_head(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + assert len(_alembic_heads(url)) == 1 + + def test_head_has_cache_columns(self, pg_engine): + with pg_engine.connect() as conn: + assert _alembic_version(conn) >= _0031 + assert _column_exists(conn, "token_usage", "cached_tokens") + assert _column_exists(conn, "token_usage", "cache_write_tokens") + + def test_downgrade_drops_then_upgrade_restores(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0030) + with pg_engine.connect() as conn: + assert _alembic_version(conn) == _0030 + assert not _column_exists(conn, "token_usage", "cached_tokens") + assert not _column_exists(conn, "token_usage", "cache_write_tokens") + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + assert _alembic_version(conn) >= _0031 + assert _column_exists(conn, "token_usage", "cached_tokens") + + def test_existing_rows_read_null_cache_bins(self, pg_engine): + """Rows written before 0031 must read NULL (unknown), not 0.""" + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0030) + with pg_engine.begin() as conn: + conn.execute( + text( + "INSERT INTO token_usage (user_id, prompt_tokens, generated_tokens) " + "VALUES ('u-mig31', 10, 1)" + ) + ) + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + row = conn.execute( + text( + "SELECT cached_tokens, cache_write_tokens FROM token_usage " + "WHERE user_id = 'u-mig31'" + ) + ).fetchone() + assert tuple(row) == (None, None) diff --git a/tests/test_usage.py b/tests/test_usage.py index 0b6053fc..dbc3cdc8 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -600,3 +600,115 @@ class TestCountPromptTokens: ] tokens = _count_prompt_tokens(messages, tools=None) assert tokens > 0 + + +# ── prompt-cache breakdown ──────────────────────────────────────────────────── +# +# Providers report ``cached_tokens`` (and, on newer OpenAI-family models, +# ``cache_write_tokens``) as breakdowns of ``prompt_tokens``. The prompt bin +# stays the provider total (never subtract the details back out); the two +# sub-bins ride alongside so persistence and the finish events can chart a +# cache hit rate. + + +class _ReportingLLM: + decoded_token = {"sub": "user_1"} + user_api_key = None + agent_id = None + token_usage = {"prompt_tokens": 0, "generated_tokens": 0} + + def __init__(self, details): + self._last_usage = { + "prompt_tokens": 1000, + "completion_tokens": 10, + "total_tokens": 1010, + } + if details is not None: + self._last_usage["prompt_tokens_details"] = details + self._last_usage_claimed = False + self.emitted = [] + + def _emit_gen_finished_log(self, model, **kwargs): + self.emitted.append(kwargs) + + +@pytest.mark.unit +def test_prefer_provider_usage_carries_cache_bins(): + from application.usage import _prefer_provider_usage + + llm = _ReportingLLM({"cached_tokens": 800, "cache_write_tokens": 100}) + usage = _prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1}) + assert usage["prompt_tokens"] == 1000 + assert usage["generated_tokens"] == 10 + assert usage["cached_tokens"] == 800 + assert usage["cache_write_tokens"] == 100 + + +@pytest.mark.unit +def test_prefer_provider_usage_maps_anthropic_cache_creation_to_writes(): + from application.usage import _prefer_provider_usage + + llm = _ReportingLLM({"cached_tokens": 3, "cache_creation_tokens": 9}) + usage = _prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1}) + assert usage["cached_tokens"] == 3 + assert usage["cache_write_tokens"] == 9 + + +@pytest.mark.unit +def test_prefer_provider_usage_without_details_has_no_cache_keys(): + from application.usage import _prefer_provider_usage + + llm = _ReportingLLM(None) + usage = _prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1}) + assert "cached_tokens" not in usage + assert "cache_write_tokens" not in usage + + +@pytest.mark.unit +def test_prefer_provider_usage_keeps_the_rest_of_the_call_record(): + from application.usage import _prefer_provider_usage + + llm = _ReportingLLM(None) + usage = _prefer_provider_usage( + llm, {"prompt_tokens": 1, "generated_tokens": 1, "model": "m"} + ) + assert usage["model"] == "m" + + +@pytest.mark.unit +def test_decorator_persists_and_emits_cache_bins(monkeypatch): + _install_fake_token_repo(monkeypatch) + llm = _ReportingLLM({"cached_tokens": 800, "cache_write_tokens": 100}) + + @gen_token_usage + def wrapped(self, model, messages, stream, tools, **kwargs): + _ = (model, messages, stream, tools, kwargs) + return "ok" + + wrapped(llm, "m", [{"role": "user", "content": "hi"}], False, None) + + row = _FakeTokenUsageRepo.last_instance.inserted[0] + assert row["prompt_tokens"] == 1000 + assert row["cached_tokens"] == 800 + assert row["cache_write_tokens"] == 100 + assert llm.emitted[0]["cached_tokens"] == 800 + assert llm.emitted[0]["cache_write_tokens"] == 100 + + +@pytest.mark.unit +def test_decorator_omits_cache_bins_when_provider_reports_none(monkeypatch): + _install_fake_token_repo(monkeypatch) + llm = _ReportingLLM(None) + + @gen_token_usage + def wrapped(self, model, messages, stream, tools, **kwargs): + _ = (model, messages, stream, tools, kwargs) + return "ok" + + wrapped(llm, "m", [{"role": "user", "content": "hi"}], False, None) + + row = _FakeTokenUsageRepo.last_instance.inserted[0] + assert row["cached_tokens"] is None + assert row["cache_write_tokens"] is None + assert llm.emitted[0]["cached_tokens"] is None + assert llm.emitted[0]["cache_write_tokens"] is None