Skip to content

test(integration): embeddings, rerank, completions and moderations cost cases - #42020

Merged
kerry-berri merged 5 commits into
mainfrom
litellm_cost_shard_embeddings_rerank
Sep 21, 2026
Merged

kerry-berri merged 5 commits into
mainfrom
litellm_cost_shard_embeddings_rerank

Conversation

@devin-ai-integration

@devin-ai-integration devin-ai-integration Bot commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • Cost shard has no embeddings, rerank, completions or moderations cases
  • Per query rerank pricing and input only pricing never asserted
  • Bedrock, Vertex, Gemini, Cohere embedding wires never exercised

How it solves it:

  • 22 scripted cases across the four endpoints with exact expectations
  • Prefix table learns cohere, bedrock invoke, vertex embeddings, text-completion-openai
  • Zero spend rows may omit x-litellm-response-cost (moderations)

User Flow

Before: a proxy admin has no shard proving what embeddings, rerank, legacy completions or moderations cost

  1. They send POST https://litellm-domain/v1/embeddings with a list of 3 inputs and get usage.prompt_tokens: 21 back
  2. They send POST https://litellm-domain/v1/rerank with 3 documents and get meta.billed_units.search_units: 3 back
  3. They open https://litellm-domain/ui/?page=logs and have to hand check the spend on each row, because tests/integration/run.py cost only covers chat, responses and messages

After: each of those calls has a scripted case with the exact spend, header and row pinned

  1. They run tests/integration/run.py cost
  2. Each new case posts the same request to /v1/embeddings, /v1/rerank, /v1/completions or /v1/moderations, receives a literal provider body, and asserts the response cost header, the persisted spend row and its token counts
  3. A pricing change on any of the four endpoints fails the shard before release

Relevant issues

Builds on #41999 (harness fields this data needs). Diff shows both until that one lands, then this rebases to main

Affected release

Linear ticket

Resolves LIT-8192

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • The handful of test files covering my change pass locally, e.g. uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*, make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
  • My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

Screenshots / Proof of Fix

Test data only, the proxy code path is unchanged. Proof is the shard against the scripted upstream that the shard itself starts (no provider spend by design)

Before (9bd648b)

  1. tests/integration/run.py cost -k "embeddings or rerank or completions or moderations" selects nothing

After (8137d87)

  1. tests/integration/run.py cost -k "embeddings or rerank or completions or moderations" against a fresh proxy and empty Redis: 22 passed, 366 deselected in 36.83s
  2. make lint passes

Type

✅ Test

Caveats (if any)

Three things the run surfaced, pinned as current behaviour rather than fixed here. Gemini embeddings ignore the body's usageMetadata.promptTokenCount and recount the input locally (the case sends 7 in the body and is billed for 3). Bedrock Cohere embeddings carry no token count on the wire, so that row is also a local recount. Rerank rows persist 0 prompt and 0 completion tokens even though the cost calculator sets completion tokens to the search units

Two cases were dropped: Together rerank has no per query rate path and always bills 0, and Azure /v1/completions against a well formed OpenAI completions body fails inside the proxy's response conversion ('NoneType' object has no attribute 'get'), which looks like a real bug worth its own ticket

Link to Devin session: https://app.devin.ai/sessions/5425681f498543a19de84adcf8daeed8
Open in Devin Desktop: https://app.devin.ai/desktop/session/5425681f498543a19de84adcf8daeed8?variant=devin


Note

Low Risk
Test and fixture data only; no production billing or proxy logic changes in this diff.

Overview
Expands the cost integration shard with 22 scripted cases that assert spend, x-litellm-response-cost, and ledger rows for /v1/embeddings, /v1/rerank, /v1/completions, and /v1/moderations—covering OpenAI/Azure embeddings, Cohere/Bedrock rerank (including per-query pricing), legacy text completions (including streaming), Together completions/embeddings, Vertex/Gemini/Fireworks embedding wires, and zero-spend moderations.

The harness gains cost-map fields (input_cost_per_query, etc.), provider routing for cohere, bedrock, text-completion-openai, and vertex_ai-embedding-models, broader OpenAI model-prefix rules for embedding/moderation endpoints, and Vertex credentials for any vertex_ai* provider. test_cost_tracking now treats a missing cost header as OK when expected spend is 0. contracts.json registers the new parametrized tests under quota_management.spend_tracking.cost_matrix.logs_cost.

Reviewed by Cursor Bugbot for commit 23f5df0. Bugbot is set up for automated code reviews on this repo. Configure here.

kerry and others added 2 commits September 19, 2026 20:58
…st cases

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

I'll fix CI failures and address comments from users with write access. I'll skip comments containing "(aside)".

  • Disable automatic comment, CI, and merge conflict monitoring

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@greptile-apps

greptile-apps Bot commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 5/5

The PR appears safe to merge, with no outstanding correctness, security, or repository-rule issue.

Summary

This PR extends the integration cost-tracking harness with scripted embeddings, rerank, legacy completions, and moderation scenarios.

  • Adds 22 exact cost, token-count, response-header, and persisted-row cases.
  • Extends provider-prefix, credentials, and pricing-field support needed by the new fixtures.
  • Tightens zero-cost header validation to accept only an absent header or a numeric zero.
  • Registers the added cases in the integration contract coverage map.

Reviews (2) · Last reviewed commit: "test(integration): require a zero cost h..."

Comment thread tests/integration/cost_calculation/test_cost_tracking.py Outdated
…ase bills nothing

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

Pushed 23f5df0 tightening the zero cost header assertion per Greptile. No findings rebutted, CI on the new head is pending.

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit 23f5df0. Configure here.

@codecov

codecov Bot commented Sep 19, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

kerry added 2 commits September 19, 2026 21:22
…sions' into litellm_cost_shard_embeddings_rerank

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	tests/integration/contracts.json
#	tests/integration/cost_calculation/cost_tracking_case.py
#	tests/integration/cost_calculation/cost_tracking_cases.json
@kerry-berri
kerry-berri merged commit dcb195e into main Sep 21, 2026
91 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant