fix(api): price cached input tokens at the cached rate - #5965
Open
WhoamiI00 wants to merge 1 commit into
Open
Conversation
Traced cost billed every prompt token at the full input rate, including the slice a provider served from its cache. On the case in Agenta-AI#5711 -- a 25,978-token Gemini prompt with 24,540 cached -- that reports $0.007793 of prompt cost where $0.001168 is correct, 6.67x too high. It is worst on agent workloads, which replay a long prefix on every call, and it runs in OSS as well as cloud. `calculate_costs` read only `prompt` and `completion` out of `tokens.incremental` and called litellm's `cost_per_token` with those alone, even though litellm accepts `cache_read_input_tokens` and its price map carries a separate, much lower cached rate (a tenth of input, for Gemini Flash). Read the cached count and forward it. Three things this has to get right: - The count is a SUBSET of `prompt_tokens`, not an addition to it. litellm's convention is that `prompt_tokens` already includes the cached slice, and it normalizes Anthropic-style usage on the way in, so it is passed alongside the prompt total rather than deducted from it. Deducting would understate cost. - It must arrive as an int. litellm reads the slice back off `Usage.prompt_tokens_details.cached_tokens`, and its `Usage` model only derives that wrapper from an int -- given the float this metric is stored as, `prompt_tokens_details` is None and every token is billed at the full rate again, silently. Without the coercion the rest of this fix is a no-op. - The adapters disagree on the field name: `logfire_adapter` writes `cache_read` (matching what the runner emits), `vercelai_adapter` writes `cached`. Both are read, or the cost is right for one integration only. The kwarg is passed only when the count is non-zero, so an uncached span calls exactly the signature it always did. The SDK pins `litellm>=1,<2`, and on a 1.x without the parameter an unconditional kwarg would raise TypeError into the bare `except`, dropping costs for every span rather than just cached ones. On the SDK side the litellm handler never recorded the count at all, so `cache_read` could not reach the API for SDK-traced calls. It now reads `prompt_tokens_details.cached_tokens` (OpenAI and Google) with a fallback to a flat `cache_read_input_tokens` (Anthropic-style). The extraction was duplicated between the sync and async paths and is now one helper. Closes Agenta-AI#5711
dosubot
Bot
added
the
size:L
|
@WhoamiI00 is attempting to deploy a commit to the agenta projects Team on Vercel. A member of the Team first needs to authorize it. |
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe SDK now extracts cached prompt tokens from LiteLLM responses and records them in tracing metrics. API cost calculation recognizes ChangesCached prompt token billing
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant LiteLLM
participant LiteLLMTracing
participant APITracing
participant CostPricer
LiteLLM->>LiteLLMTracing: usage response with cached prompt tokens
LiteLLMTracing->>APITracing: cache_read token metric
APITracing->>CostPricer: cached input token count
CostPricer-->>APITracing: calculated cost
Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: c112c05a-d634-4dbd-9299-a6d7091cfca5
⛔ Files ignored due to path filters (1)
.github/pr-assets/5711-cached-token-cost.pngis excluded by!**/*.png
📒 Files selected for processing (4)
api/oss/src/core/tracing/utils/trees.pyapi/oss/tests/pytest/unit/tracing/utils/test_trees.pysdks/python/agenta/sdk/litellm/litellm.pysdks/python/oss/tests/pytest/unit/test_litellm_token_usage.py
Comment on lines
+609
to
+635
| cache_read_tokens = next( | ||
| (tokens[key] for key in CACHE_READ_TOKEN_KEYS if tokens.get(key)), | ||
| 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 {} | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Validate the cached-token subset before pricing.
cache_read_tokens comes from span metrics without validation. A negative value or a value greater than prompt_tokens reaches the pricer. A pricer that calculates fresh input as prompt_tokens - cache_read_input_tokens can then record a negative or otherwise invalid prompt cost.
Require a whole, non-negative cached count that does not exceed prompt_tokens. Reject or omit invalid values. Add coverage for both bounds.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Traced cost bills every prompt token at the full input rate, including the slice a provider served from its cache. On the case in #5711 — a 25,978-token Gemini prompt of which 24,540 were cached — that reports $0.007793 of prompt cost where $0.001168 is correct: 6.67x too high. It runs in OSS as well as cloud, and it is worst on agent workloads, which replay a long prefix on every call.
calculate_costsread onlypromptandcompletionout oftokens.incrementaland called litellm'scost_per_tokenwith those alone — even though that function acceptscache_read_input_tokensand litellm's price map carries a separate, much lower cached rate (a tenth of input, for Gemini Flash).One correction to the issue
The issue says the cached count is never recorded. That is true of the SDK path but not of the OTLP path:
logfire_adapter.pyalready mapsgen_ai.usage.cache_read.input_tokens→ag.metrics.unit.tokens.cache_read, andspan_data_builders.pyrewritesunit.tokens.*→tokens.incremental.*, so it does land inag.metrics.tokens.incremental.cache_read. The runner already emits that attribute. The count was being recorded and then dropped on the floor at the one place that prices tokens.Three things this has to get right
1. The cached count is a subset of
prompt_tokens, not an addition to it. litellm's convention is thatprompt_tokensalready includes the cached slice (it normalizes Anthropic-style usage, where the input count excludes them, on the way in). So the count is passed alongside the prompt total, never deducted from it — deducting would understate cost instead of overstating it.2. It must arrive as an
int. This one is not obvious and it is the whole fix. litellm reads the slice back offUsage.prompt_tokens_details.cached_tokens, and itsUsagemodel only derives that wrapper from an int:Agenta stores token metrics as floats, so passing the value through as-is leaves
prompt_tokens_detailsatNoneand every token is billed at the full input rate again — silently, with no exception to catch. I had this wrong at first; the end-to-end demo below is what caught it, since the mocked unit tests happily showed the kwarg being passed.3. The adapters disagree on the field name.
logfire_adapterwritescache_read(matching what the runner emits),vercelai_adapterwritescached(fromai.usage.cachedInputTokens). Both are read, or the cost is right for one integration and overstated for the other.Backward compatibility
The kwarg is passed only when the count is 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 raiseTypeErrorstraight into the existing bareexcept— dropping costs for every span, not just cached ones.SDK side
The litellm handler never recorded the cached count, so
cache_readcould not reach the API for SDK-traced calls. It now readsprompt_tokens_details.cached_tokens(OpenAI and Google both report it there) with a fallback to a flatcache_read_input_tokens(Anthropic-style). The extraction was duplicated between the sync and async paths; it is now one helper.Closes #5711
Testing
Verified locally
Run on Linux (
python:3.11-slimcontainer) because litellm 1.92.0 ships manylinux wheels only.calculate_costs, no mocks. TheBEFOREline is literally the call the old code made.pytest oss/tests/pytest/unit/tracing(api) — 92 passed.pytest oss/tests/pytest/unit(SDK) — 2022 passed, 4 skipped, 10 xfailed.python run-tests.py --layer unit(api) — 2210 passed, 4 failed. The 4 failures are intest_web_entrypoint_email_env.py(SMTP entrypoint config) and are pre-existing: they fail identically on a clean checkout ofmainwith this branch's changes stashed. Nothing this PR touches is involved.ruff format --checkandruff checkon all four changed files — clean.trees.pyalone fails exactly the three cached-behaviour tests, while the two "signature unchanged" tests keep passing in both directions.Added or updated tests
api/oss/tests/pytest/unit/tracing/utils/test_trees.py— the_spanhelper takes an optional cached count under a configurable key, plus:test_calculate_costs_passes_cached_tokens_to_the_pricer— parametrized overcache_readandcached, asserting the count is forwarded and the prompt total is passed through untouched.test_calculate_costs_sends_the_cached_count_as_an_int— the float trap in point 2. Without the coercion the fix silently does nothing, so this is the test that matters most.test_calculate_costs_omits_cache_kwarg_when_nothing_was_cached— calls a pricer whose signature accepts nothing else, so a regression to an unconditional kwarg fails loudly.test_calculate_costs_ignores_a_zero_cached_count— an explicit zero is a cache miss.test_calculate_costs_bills_cached_tokens_below_fresh_input— end-to-end ratio against a pricer modelling litellm's published contract.sdks/python/oss/tests/pytest/unit/test_litellm_token_usage.py(new) — 8 cases over_extract_token_usage: OpenAI-style objects, dict-shaped usage, the flat Anthropic-style field, precedence when both are present, absent/zero/null cache details, and a response with no usage object.QA follow-up
Worth a maintainer's eye on two things I deliberately left alone:
cumulate_tokenshas a hardcoded prompt/completion/total shape, socache_readdoes not roll up to parent spans. Costs are correct either way (they cumulate from the corrected per-span costs); this only means the cached token count is not visible intokens.cumulative. Extending it changes the rollup shape the frontend reads, which felt out of scope here.cache_readandcachedfixes the symptom. Normalizingvercelai_adapterto emitcache_readwould fix the cause, but it changes ingest behaviour for existing spans, so I left it.Also worth noting:
cache_creationis recorded bylogfire_adapterand still not priced. Cache writes bill above the normal input rate on some providers, so that is a separate, smaller understatement — happy to follow up if you want it in scope.Demo
Backend-only change, so the demo is captured terminal output rather than a UI recording: the real
calculate_costsagainst the real litellm price map, before and after, plus a check that an uncached span is unchanged.Checklist