diff --git a/.github/pr-assets/5711-cached-token-cost.png b/.github/pr-assets/5711-cached-token-cost.png new file mode 100644 index 00000000000..f6b9d1819b5 Binary files /dev/null and b/.github/pr-assets/5711-cached-token-cost.png differ diff --git a/api/oss/src/core/tracing/utils/trees.py b/api/oss/src/core/tracing/utils/trees.py index adf7181524a..f5b2db8c537 100644 --- a/api/oss/src/core/tracing/utils/trees.py +++ b/api/oss/src/core/tracing/utils/trees.py @@ -575,6 +575,13 @@ def _cumulate_tree_dfs( "rerank", ] +# Prompt tokens served from a provider's cache, which price at a much lower rate than fresh +# input (a tenth of it, for some models). The ingest adapters disagree on the field name: +# `logfire_adapter` writes `cache_read` (matching the `gen_ai.usage.cache_read.input_tokens` +# the runner emits), while `vercelai_adapter` writes `cached` (from `ai.usage.cachedInputTokens`). +# Read every alias, or the cost is right for one integration and overstated for the other. +CACHE_READ_TOKEN_KEYS = ("cache_read", "cached") + def calculate_costs(span_idx: Dict[str, OTelFlatSpan]): for span in span_idx.values(): @@ -588,27 +595,59 @@ def calculate_costs(span_idx: Dict[str, OTelFlatSpan]): "model" ) or attr.get("ag", {}).get("data", {}).get("parameters", {}).get("model") - prompt_tokens = ( + tokens: dict = ( attr.get("ag", {}) .get("metrics", {}) .get("tokens", {}) .get("incremental", {}) - .get("prompt", 0.0) ) - completion_tokens = ( - attr.get("ag", {}) - .get("metrics", {}) - .get("tokens", {}) - .get("incremental", {}) - .get("completion", 0.0) + prompt_tokens = tokens.get("prompt", 0.0) + + completion_tokens = tokens.get("completion", 0.0) + + # Only a real positive number qualifies. Non-numeric or negative garbage from a + # foreign OTLP source must degrade to "no cache kwarg", not reach the int() below + # and turn into a swallowed exception that drops the span's ENTIRE cost. + cache_read_tokens = next( + ( + value + for key in CACHE_READ_TOKEN_KEYS + if isinstance((value := tokens.get(key)), (int, float)) + and not isinstance(value, bool) + and value > 0 + ), + 0.0, ) try: + # litellm's convention is that `prompt_tokens` INCLUDES the cached tokens and + # that it prices the cached slice separately (it normalizes Anthropic-style + # usage, where the input count excludes them, on the way in). So the cached + # count is passed ALONGSIDE the prompt total and must not be subtracted from + # it first -- doing that would understate cost instead of overstating it. + # + # Passed only when non-zero so a span with no caching calls exactly the + # signature it always did: the SDK pins `litellm>=1,<2`, and on a 1.x old + # enough to lack the parameter an unconditional kwarg would raise TypeError, + # which the `except` below would swallow into "no costs at all" for EVERY span. + # + # int(), and not incidentally: litellm reads the cached slice back off + # `Usage.prompt_tokens_details.cached_tokens`, and its `Usage` model only + # derives that wrapper from an int. Hand it the float this metric is stored + # as and `prompt_tokens_details` comes back None, so the cached tokens are + # billed at the full input rate again -- silently, with no error to catch. + cache_kwargs = ( + {"cache_read_input_tokens": int(cache_read_tokens)} + if cache_read_tokens + else {} + ) + costs = cost_calculator.cost_per_token( model=model, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, + **cache_kwargs, ) if not costs: @@ -646,6 +685,7 @@ def calculate_costs(span_idx: Dict[str, OTelFlatSpan]): model=model, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, + cache_read_tokens=cache_read_tokens, ) diff --git a/api/oss/tests/pytest/unit/tracing/utils/test_trees.py b/api/oss/tests/pytest/unit/tracing/utils/test_trees.py index be94c669dae..ee05b184825 100644 --- a/api/oss/tests/pytest/unit/tracing/utils/test_trees.py +++ b/api/oss/tests/pytest/unit/tracing/utils/test_trees.py @@ -39,6 +39,8 @@ def _span( links=None, prompt_tokens: float = 0.0, completion_tokens: float = 0.0, + cache_read_tokens: float | None = None, + cache_token_key: str = "cache_read", prompt_cost: float = 0.0, completion_cost: float = 0.0, errors: int = 0, @@ -50,14 +52,18 @@ def _span( ) -> OTelFlatSpan: total_tokens = prompt_tokens + completion_tokens total_cost = prompt_cost + completion_cost + incremental_tokens = { + "prompt": prompt_tokens, + "completion": completion_tokens, + "total": total_tokens, + } + # Under `cache_token_key` because the ingest adapters disagree on the name: the logfire + # path writes `cache_read`, the Vercel AI path writes `cached`. Absent entirely when + # None, which is the shape of a span from a provider or model without prompt caching. + if cache_read_tokens is not None: + incremental_tokens[cache_token_key] = cache_read_tokens metrics = { - "tokens": { - "incremental": { - "prompt": prompt_tokens, - "completion": completion_tokens, - "total": total_tokens, - } - }, + "tokens": {"incremental": incremental_tokens}, "costs": { "incremental": { "prompt": prompt_cost, @@ -252,6 +258,298 @@ def _raise(**_): assert "incremental" in span_idx[ROOT_UUID].attributes["ag"]["metrics"]["costs"] +@pytest.mark.parametrize("cache_token_key", ["cache_read", "cached"]) +def test_calculate_costs_passes_cached_tokens_to_the_pricer( + monkeypatch, + cache_token_key, +): + """Cached prompt tokens must reach litellm, which prices them far below fresh input. + + Both spellings have to be honoured: the logfire ingest path writes `cache_read`, the + Vercel AI path writes `cached`. Reading only one silently overstates the other's cost. + """ + span = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=25978, + completion_tokens=100, + cache_read_tokens=24540, + cache_token_key=cache_token_key, + span_type=SpanType.CHAT, + ) + span_idx = {span.span_id: span} + + seen = {} + + def _capture(**kwargs): + seen.update(kwargs) + return (0.005, 0.001) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _capture, + ) + + calculate_costs(span_idx) + + assert seen["cache_read_input_tokens"] == 24540 + # The cached count is a SUBSET of the prompt total, not an addition to it. litellm + # re-prices that slice itself, so the prompt total is passed through untouched; + # subtracting the cached tokens here would understate cost instead of overstating it. + assert seen["prompt_tokens"] == 25978 + + +def test_calculate_costs_sends_the_cached_count_as_an_int(monkeypatch): + """The count must reach litellm as an int, not the float this metric is stored as. + + litellm reads the cached slice back off `Usage.prompt_tokens_details.cached_tokens`, + and its `Usage` model only builds that wrapper from an int -- given a float it leaves + `prompt_tokens_details` None and bills every token at the full input rate again. There + is no exception to catch: the whole fix degrades to a no-op. Verified against litellm + 1.92.0, where `cost_per_token(..., cache_read_input_tokens=24540)` prices a 25,978-token + Gemini prompt at $0.001168 and the same call with `24540.0` at $0.007793. + """ + span = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=25978.0, + completion_tokens=100.0, + cache_read_tokens=24540.0, + span_type=SpanType.CHAT, + ) + span_idx = {span.span_id: span} + + seen = {} + + def _capture(**kwargs): + seen.update(kwargs) + return (0.001, 0.002) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _capture, + ) + + calculate_costs(span_idx) + + assert isinstance(seen["cache_read_input_tokens"], int) + assert not isinstance(seen["cache_read_input_tokens"], bool) + assert seen["cache_read_input_tokens"] == 24540 + + +def test_calculate_costs_omits_cache_kwarg_when_nothing_was_cached(monkeypatch): + """A span with no caching must call exactly the signature it always did. + + The SDK pins `litellm>=1,<2`, and `calculate_costs` swallows every exception. On a 1.x + old enough to lack the parameter, passing it unconditionally would raise TypeError and + be swallowed into "no costs at all" for EVERY span, not just cached ones -- so the + pricer is called here with a signature that accepts nothing else. + """ + span = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=10, + completion_tokens=20, + span_type=SpanType.CHAT, + ) + span_idx = {span.span_id: span} + + def _legacy_signature(model, prompt_tokens, completion_tokens): + return (0.1, 0.2) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _legacy_signature, + ) + + calculate_costs(span_idx) + + costs = span_idx[ROOT_UUID].attributes["ag"]["metrics"]["costs"]["incremental"] + assert costs["prompt"] == pytest.approx(0.1) + assert costs["completion"] == pytest.approx(0.2) + assert costs["total"] == pytest.approx(0.3) + + +def test_calculate_costs_ignores_a_zero_cached_count(monkeypatch): + """An explicit zero is a cache miss, not a cached slice: still the legacy signature.""" + span = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=10, + completion_tokens=20, + cache_read_tokens=0, + span_type=SpanType.CHAT, + ) + span_idx = {span.span_id: span} + + def _legacy_signature(model, prompt_tokens, completion_tokens): + return (0.1, 0.2) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _legacy_signature, + ) + + calculate_costs(span_idx) + + costs = span_idx[ROOT_UUID].attributes["ag"]["metrics"]["costs"]["incremental"] + assert costs["total"] == pytest.approx(0.3) + + +def test_calculate_costs_ignores_a_non_numeric_cached_count(monkeypatch): + """Garbage in the cache field must not cost the span its ENTIRE cost. + + A foreign OTLP source can write anything under `tokens.incremental`. A truthy + non-numeric value that reached `int()` would raise inside the try, and the bare + `except` would swallow it into "no costs at all" for the span -- a regression + against the pre-fix behaviour, where a garbage cache field was simply ignored + and prompt/completion were still priced. + """ + span = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=10, + completion_tokens=20, + cache_read_tokens="24540", + span_type=SpanType.CHAT, + ) + span_idx = {span.span_id: span} + + def _legacy_signature(model, prompt_tokens, completion_tokens): + return (0.1, 0.2) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _legacy_signature, + ) + + calculate_costs(span_idx) + + costs = span_idx[ROOT_UUID].attributes["ag"]["metrics"]["costs"]["incremental"] + assert costs["prompt"] == pytest.approx(0.1) + assert costs["completion"] == pytest.approx(0.2) + assert costs["total"] == pytest.approx(0.3) + + +def test_calculate_costs_ignores_a_negative_cached_count(monkeypatch): + """A negative count is not a cached slice: still the legacy signature. + + Negative is truthy, so without the numeric guard it would reach the pricer, + where "fresh = prompt - cached" arithmetic turns it into an overstated cost. + """ + span = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=10, + completion_tokens=20, + cache_read_tokens=-5, + span_type=SpanType.CHAT, + ) + span_idx = {span.span_id: span} + + def _legacy_signature(model, prompt_tokens, completion_tokens): + return (0.1, 0.2) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _legacy_signature, + ) + + calculate_costs(span_idx) + + costs = span_idx[ROOT_UUID].attributes["ag"]["metrics"]["costs"]["incremental"] + assert costs["total"] == pytest.approx(0.3) + + +def test_calculate_costs_prefers_cache_read_over_cached_when_both_present(monkeypatch): + """`cache_read` wins when both aliases appear on one span. + + No ingest adapter writes both today, but the precedence should be pinned: + `cache_read` is the spelling the runner emits and the logfire path stores. + """ + span = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=25978, + completion_tokens=100, + cache_read_tokens=24540, + span_type=SpanType.CHAT, + ) + span.attributes["ag"]["metrics"]["tokens"]["incremental"]["cached"] = 999 + span_idx = {span.span_id: span} + + seen = {} + + def _capture(**kwargs): + seen.update(kwargs) + return (0.001, 0.002) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _capture, + ) + + calculate_costs(span_idx) + + assert seen["cache_read_input_tokens"] == 24540 + + +def test_calculate_costs_bills_cached_tokens_below_fresh_input(monkeypatch): + """End-to-end shape of #5711, with the pricer modelling litellm's published contract. + + The reported case: a 25,978-token prompt of which 24,540 came from cache, on a model + whose cached rate is a tenth of its input rate. Billing every token at the full input + rate is what produced the 6.6x overstatement. + """ + input_rate = 0.30 / 1_000_000 + cached_rate = input_rate / 10 + output_rate = 2.50 / 1_000_000 + + def _priced(model, prompt_tokens, completion_tokens, cache_read_input_tokens=0): + fresh = prompt_tokens - cache_read_input_tokens + return ( + fresh * input_rate + cache_read_input_tokens * cached_rate, + completion_tokens * output_rate, + ) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _priced, + ) + + cached = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=25978, + completion_tokens=100, + cache_read_tokens=24540, + span_type=SpanType.CHAT, + ) + uncached = _span( + span_id=CHILD_A_UUID, + span_name="root", + prompt_tokens=25978, + completion_tokens=100, + span_type=SpanType.CHAT, + ) + + calculate_costs({cached.span_id: cached}) + calculate_costs({uncached.span_id: uncached}) + + cached_costs = cached.attributes["ag"]["metrics"]["costs"]["incremental"] + uncached_costs = uncached.attributes["ag"]["metrics"]["costs"]["incremental"] + + assert cached_costs["total"] < uncached_costs["total"] + # Compared on the prompt component, which is the part caching changes -- the issue's + # 6.6x is a prompt-side ratio, and folding in the (identical) completion cost would + # dilute it to ~5.7x. Same call, same token counts: the only difference is whether the + # cached slice was priced as cached. Before the fix both paths produced the high number. + assert uncached_costs["prompt"] / cached_costs["prompt"] == pytest.approx( + 6.67, + abs=0.01, + ) + + def test_calculate_and_propagate_metrics_runs_full_pipeline(monkeypatch): root = _span( span_id=ROOT_UUID, diff --git a/sdks/python/agenta/sdk/litellm/litellm.py b/sdks/python/agenta/sdk/litellm/litellm.py index 06dd439afdc..87efea8b0c4 100644 --- a/sdks/python/agenta/sdk/litellm/litellm.py +++ b/sdks/python/agenta/sdk/litellm/litellm.py @@ -1,5 +1,5 @@ import warnings -from typing import Dict +from typing import Any, Dict, Optional from opentelemetry.trace import SpanKind import agenta as ag @@ -10,6 +10,47 @@ log = get_module_logger(__name__) +def _read(source: Any, key: str) -> Any: + """Read ``key`` off a usage payload that may arrive as a dict or as an object.""" + if source is None: + return None + if isinstance(source, dict): + return source.get(key) + return getattr(source, key, None) + + +def _extract_token_usage(response_obj: Any) -> Dict[str, Optional[float]]: + """The token counts recorded under ``metrics.unit.tokens`` for one LLM call. + + ``cache_read`` is the slice of the prompt the provider served from its cache, which + prices far below fresh input. It is a SUBSET of ``prompt_tokens``, not an addition to + it, so it is recorded alongside the prompt total and never deducted from it; the + cost calculation re-prices that slice at the cached rate. + """ + usage = _read(response_obj, "usage") + + prompt_tokens = _read(usage, "prompt_tokens") + completion_tokens = _read(usage, "completion_tokens") + total_tokens = _read(usage, "total_tokens") + + # OpenAI and Google both report the cached count at + # ``prompt_tokens_details.cached_tokens``; Anthropic-style usage surfaces it flat as + # ``cache_read_input_tokens``. Check the nested form first: on a provider that reports + # both, the nested one is the OpenAI-convention value matching ``prompt_tokens``. + cache_read_tokens = _read(_read(usage, "prompt_tokens_details"), "cached_tokens") + if cache_read_tokens is None: + cache_read_tokens = _read(usage, "cache_read_input_tokens") + + # Falsy -> None throughout, matching how prompt/completion/total have always been + # recorded: a zero carries no more information than an absent field here. + return { + "prompt": float(prompt_tokens) if prompt_tokens else None, + "completion": float(completion_tokens) if completion_tokens else None, + "total": float(total_tokens) if total_tokens else None, + "cache_read": float(cache_read_tokens) if cache_read_tokens else None, + } + + def litellm_handler(): if ag.tracing is None: warnings.warn( @@ -174,29 +215,8 @@ def log_success_event( namespace="metrics.unit.costs", ) - # Handle both dict and object attribute access for usage, and safely handle None - usage = getattr(response_obj, "usage", None) - if isinstance(usage, dict): - prompt_tokens = usage.get("prompt_tokens") - completion_tokens = usage.get("completion_tokens") - total_tokens = usage.get("total_tokens") - elif usage is not None: - prompt_tokens = getattr(usage, "prompt_tokens", None) - completion_tokens = getattr(usage, "completion_tokens", None) - total_tokens = getattr(usage, "total_tokens", None) - else: - prompt_tokens = completion_tokens = total_tokens = None - span.set_attributes( - attributes=( - { - "prompt": float(prompt_tokens) if prompt_tokens else None, - "completion": float(completion_tokens) - if completion_tokens - else None, - "total": float(total_tokens) if total_tokens else None, - } - ), + attributes=_extract_token_usage(response_obj), namespace="metrics.unit.tokens", ) @@ -311,31 +331,8 @@ async def async_log_success_event( namespace="metrics.unit.costs", ) - # Handle both dict and object attribute access for usage - usage = getattr(response_obj, "usage", None) - if usage is None: - prompt_tokens = None - completion_tokens = None - total_tokens = None - elif isinstance(usage, dict): - prompt_tokens = usage.get("prompt_tokens") - completion_tokens = usage.get("completion_tokens") - total_tokens = usage.get("total_tokens") - else: - prompt_tokens = getattr(usage, "prompt_tokens", None) - completion_tokens = getattr(usage, "completion_tokens", None) - total_tokens = getattr(usage, "total_tokens", None) - span.set_attributes( - attributes=( - { - "prompt": float(prompt_tokens) if prompt_tokens else None, - "completion": float(completion_tokens) - if completion_tokens - else None, - "total": float(total_tokens) if total_tokens else None, - } - ), + attributes=_extract_token_usage(response_obj), namespace="metrics.unit.tokens", ) diff --git a/sdks/python/oss/tests/pytest/unit/test_litellm_token_usage.py b/sdks/python/oss/tests/pytest/unit/test_litellm_token_usage.py new file mode 100644 index 00000000000..307b316e71c --- /dev/null +++ b/sdks/python/oss/tests/pytest/unit/test_litellm_token_usage.py @@ -0,0 +1,121 @@ +"""Token extraction from a litellm response's usage object. + +Covers the cached-prompt-token count the cost calculation needs (see #5711): without it +every cached token is billed at the full input rate, which overstates cost by up to 6.6x +on a workload that replays a long prefix. +""" + +from types import SimpleNamespace + +import pytest + +from agenta.sdk.litellm.litellm import _extract_token_usage + + +def _response(usage): + return SimpleNamespace(usage=usage) + + +def test_reads_openai_style_usage_objects(): + usage = SimpleNamespace( + prompt_tokens=25978, + completion_tokens=100, + total_tokens=26078, + prompt_tokens_details=SimpleNamespace(cached_tokens=24540), + ) + + assert _extract_token_usage(_response(usage)) == { + "prompt": 25978.0, + "completion": 100.0, + "total": 26078.0, + # A SUBSET of `prompt`, not an addition to it: the cost calculation re-prices this + # slice at the provider's cached rate rather than adding a fourth token bucket. + "cache_read": 24540.0, + } + + +def test_reads_usage_delivered_as_a_dict(): + """litellm hands the usage object through as a dict on some providers.""" + usage = { + "prompt_tokens": 300, + "completion_tokens": 50, + "total_tokens": 350, + "prompt_tokens_details": {"cached_tokens": 128}, + } + + assert _extract_token_usage(_response(usage)) == { + "prompt": 300.0, + "completion": 50.0, + "total": 350.0, + "cache_read": 128.0, + } + + +def test_falls_back_to_the_flat_anthropic_style_cache_field(): + """Anthropic-style usage reports the cached count flat, with no nested details.""" + usage = SimpleNamespace( + prompt_tokens=300, + completion_tokens=50, + total_tokens=350, + cache_read_input_tokens=128, + ) + + assert _extract_token_usage(_response(usage))["cache_read"] == 128.0 + + +def test_prefers_the_nested_count_when_a_provider_reports_both(): + """The nested value is the OpenAI-convention one that matches `prompt_tokens`.""" + usage = SimpleNamespace( + prompt_tokens=300, + completion_tokens=50, + total_tokens=350, + prompt_tokens_details=SimpleNamespace(cached_tokens=128), + cache_read_input_tokens=999, + ) + + assert _extract_token_usage(_response(usage))["cache_read"] == 128.0 + + +@pytest.mark.parametrize( + "usage", + [ + pytest.param( + SimpleNamespace(prompt_tokens=300, completion_tokens=50, total_tokens=350), + id="no-cache-fields", + ), + pytest.param( + SimpleNamespace( + prompt_tokens=300, + completion_tokens=50, + total_tokens=350, + prompt_tokens_details=SimpleNamespace(cached_tokens=0), + ), + id="explicit-cache-miss", + ), + pytest.param( + SimpleNamespace( + prompt_tokens=300, + completion_tokens=50, + total_tokens=350, + prompt_tokens_details=None, + ), + id="null-details", + ), + ], +) +def test_records_no_cached_count_when_nothing_was_cached(usage): + """Absent, zero, and null all mean the same thing: nothing to re-price.""" + extracted = _extract_token_usage(_response(usage)) + + assert extracted["cache_read"] is None + assert extracted["prompt"] == 300.0 + + +def test_tolerates_a_response_without_usage(): + """A failed or streaming-partial response can carry no usage object at all.""" + assert _extract_token_usage(_response(None)) == { + "prompt": None, + "completion": None, + "total": None, + "cache_read": None, + }