Skip to content

fix(databricks): carry service_tier through request, streaming and cost calc - #41837

Closed
devin-ai-integration[bot] wants to merge 15 commits into
mainfrom
litellm_databricks_service_tier_pricing
Closed

devin-ai-integration[bot] wants to merge 15 commits into
mainfrom
litellm_databricks_service_tier_pricing

Conversation

@devin-ai-integration

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

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • Databricks service_tier: priority is dropped from the outgoing request
  • Responses API bridge ignores service_tier in the request body
  • Streamed responses lose the returned service_tier in the stream wrapper and aggregation
  • Responses API output drops the served service_tier
  • Databricks priority traffic is priced at standard rates, so spend is under-counted
  • A returned default tier did not override a requested priority tier in cost calc

How it solves it:

  • Add service_tier to Databricks supported params and carry it on responses and stream chunks
  • Read service_tier from the Responses request body, then kwargs
  • Forward service_tier through the Databricks cost path to tiered pricing keys
  • Stream wrapper stamps the last seen service_tier on emitted and usage chunks; stream_chunk_builder keeps it
  • Responses API output carries the served service_tier from the chat completion
  • Cost calc resolves the tier from the response first, then the request
  • Integration test proves a priority request is forwarded and charged at priority rates (tests/integration/pricing/test_service_tier_pricing.py, red when _resolve_service_tier returns None)

User Flow

Before: a developer sends a Databricks priority request through the proxy and is billed the standard price

  1. They send POST https://litellm-domain/v1/chat/completions with "model": "databricks/<endpoint>" and "service_tier": "priority"
  2. The request reaches Databricks without service_tier, so it is served on the standard tier
  3. The response has no service_tier field and x-litellm-response-cost matches standard pricing
  4. https://litellm-domain/ui/?page=logs shows the request at standard spend even when Databricks did serve it as priority

After: the same request is served and billed as priority, and a fallback to default is billed as default

  1. They send the same POST https://litellm-domain/v1/chat/completions with "service_tier": "priority"
  2. The request reaches Databricks with service_tier: priority
  3. The response carries "service_tier": "priority" (also on streamed chunks) and x-litellm-response-cost is the priority price
  4. If Databricks falls back and returns "service_tier": "default", the cost header and the logs page show the default price
  5. The same works through POST https://litellm-domain/v1/responses with "service_tier": "priority" in the body

Relevant issues

Fixes #41834

Affected release

Linear ticket

Resolves LIT-8121

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

No Databricks workspace endpoint was reachable from this environment, so the live proxy run is still owed. The run below is the SDK against a registered Databricks model with priority pricing at 2x standard (100 prompt tokens, 20 cached, 50 completion), run with LITELLM_LOCAL_MODEL_COST_MAP=True

Shared setup: register databricks/offline-tier-test with input_cost_per_token=2e-6, output_cost_per_token=4e-6, cache_read_input_token_cost=1e-6 and the _priority keys at 2x, then call get_optional_params(custom_llm_provider="databricks", service_tier="priority", drop_params=True) and litellm.completion_cost(...)

Before (4e21178)

  1. get_optional_params output contains service_tier: False
  2. completion_cost with response service_tier="default": 0.00038
  3. completion_cost with response service_tier="priority": 0.00038 (priority ignored)
  4. completion_cost with requested priority and returned default: 0.00038

After (211a545)

  1. get_optional_params output contains service_tier: True
  2. completion_cost with response service_tier="default": 0.00038
  3. completion_cost with response service_tier="priority": 0.00076
  4. completion_cost with requested priority and returned default: 0.00038

Type

🐛 Bug Fix

Caveats (if any)

Medium

  • completion_cost tier precedence changed for every provider: response tier, then usage tier, then requested tier
    • Before, a requested tier (service_tier= or optional_params) won over the tier on the response
    • test_completion_cost_service_tier_priority was rewritten to the new order
  • Live proxy proof against a real Databricks endpoint not yet captured (no workspace host available here)

Low

  • Gemini trafficType mapping moved into the resolver as the last fallback, same effective order

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

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

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


Note

Medium Risk
Changes global completion_cost service-tier precedence and spend headers/logs for every provider; incorrect tier resolution would mis-price requests, especially priority Databricks traffic.

Overview
Fixes Databricks service_tier being dropped so priority (and other tiered) traffic is forwarded, surfaced on responses/streams, and billed with the right tier keys.

Databricks & Responses API: service_tier is now a supported Databricks param, copied onto non-streaming and streaming chunks, and passed into databricks_cost_per_token for _priority pricing. The Responses API bridge reads service_tier from the request body (then kwargs) and copies the served tier from the chat completion onto the Responses output.

Streaming: CustomStreamWrapper remembers the latest chunk service_tier and stamps it on emitted chunks (including usage-only frames); stream_chunk_builder picks up the last non-null tier from reversed chunks.

Cost calculation (all providers): Inline tier extraction is replaced by _resolve_service_tier, which bills using provider-reported tier first (response, then usage), then the requested/optional_params tier, with Gemini traffic_type as a final fallback. Logging and token breakdowns use this resolved tier—so a returned default tier is charged at default rates even when the client asked for priority.

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

@devin-ai-integration
devin-ai-integration Bot requested a review from a team September 18, 2026 17:04
@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

  • Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
  • Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

  • Disable automatic comment, CI, and merge conflict monitoring

@CLAassistant

CLAassistant commented Sep 18, 2026 •

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 all sign our Contributor License Agreement before we can accept your contribution.
0 out of 2 committers have signed the CLA.

❌ devin-ai-integration[bot]
❌ kerry-berri
You have signed the CLA already but the status is still pending? Let us recheck it.

@codspeed

codspeed Bot commented Sep 18, 2026 •

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_databricks_service_tier_pricing (7a37dc5) with main (2dccc0d)1

Open in CodSpeed

Footnotes

  1. No successful run was found on main (e26a645) during the generation of this report, so 2dccc0d was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

@codecov

codecov Bot commented Sep 18, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.72727% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
litellm/litellm_core_utils/streaming_handler.py 92.30% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented Sep 18, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 5/5

The PR appears safe to merge, with no new actionable findings or outstanding previous findings

Findings

  1. P2 Tests violate required coding conventions ▶

Summary

This PR carries Databricks service_tier through requests, responses, streaming aggregation, and tier-aware cost calculation

  • The Responses API bridge forwards the requested tier and exposes the returned tier
  • Cost calculation prefers the returned tier over the requested tier
  • Regression tests cover tier propagation, pricing, and recorded spend
  • The latest changes complete Final annotations in the added tests without changing their behavior

Reviews (4) · Last reviewed commit: "style(tests): apply Final bindings and d..."

Comment thread litellm/llms/databricks/chat/transformation.py
Comment thread litellm/cost_calculator.py
Comment on lines +1867 to +1868
usage_no_tier = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
priority_reference = completion_cost(

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.

P2 Tests violate required coding conventions

The added tests omit Final annotations and add comments that repeat their assertions. CLAUDE.md requires Final bindings and prohibits redundant source comments. Apply those conventions throughout the added tests before merging

Context Used: CLAUDE.md (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

@greptileai please re-review, the streaming wrapper, Responses output, and test convention findings are addressed in the latest commits

kvnloo commented Sep 20, 2026

Copy link
Copy Markdown

The remaining Greptile finding about global precedence points to a useful modeling simplification here: requested tier and served tier are two different facts and probably should not be collapsed into one provider-agnostic service_tier before pricing policy runs.

For Databricks the distinction matters materially because priority can fall back:

requested_service_tier = "priority"   # caller intent
served_service_tier    = "default"    # provider evidence

The user's billing/UX expectation is:

  • outbound request should preserve the requested tier;
  • response/logs should expose the tier actually served when the provider reports it;
  • Databricks pricing should use the served tier when present;
  • if Databricks omits a returned tier, requested tier can remain an estimate/fallback;
  • unrelated providers should keep their existing precedence unless their own contract says a returned tier is authoritative.

That avoids needing a global rule like:

response > usage > explicit request > optional params

which is exactly what creates the cross-provider compatibility change called out in review.

Conceptually, the lifecycle becomes:

Request intent
  requested_service_tier
          |
          v
provider request
          |
          v
Provider observation
  served_service_tier?   <-- do not manufacture this
          |
          v
provider-specific pricing policy
  Databricks: served ?? requested
  existing providers: preserve existing contract

This also gives observability a clearer contract. A log can distinguish “priority requested, default served” from “priority requested, provider never reported a tier”. Those are operationally different events even if both eventually use a fallback pricing rule.

I would keep the public response's existing service_tier field as the served value if that matches the API contract. The request value already exists in request/optional params, so this does not necessarily require a new public type. The important part is to keep the two facts separate internally until a provider-specific cost policy chooses one.

Regression shape

A compact matrix would protect both Databricks and compatibility:

provider contract requested returned price from
Databricks priority priority priority
Databricks priority default default
Databricks priority absent priority fallback/estimate
Databricks absent default default
unrelated provider with existing explicit-request precedence priority different tier existing behavior unchanged

Then separately assert the returned tier survives Chat, Responses, Messages, and stream aggregation unchanged.

This is the same useful pattern as the context-limit work elsewhere in the repo: carry facts with provenance first; apply policy once at the domain boundary that actually owns the decision. Here the domain boundary is provider-specific pricing, not generic response normalization.

That should let #41837 fix the full Databricks UX (request -> actual service -> response -> spend) without turning a Databricks fallback rule into a global LiteLLM behavior change.

devin-ai-integration Bot and others added 9 commits September 23, 2026 01:17
… streaming and cost calc

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>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
…ansform_response

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>
…es output

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>
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
devin-ai-integration Bot force-pushed the litellm_databricks_service_tier_pricing branch from 03cfabf to b2e29f8 Compare September 23, 2026 01:25
@litellm-stampbot

Copy link
Copy Markdown
Contributor

@greptileai

@litellm-stampbot

Copy link
Copy Markdown
Contributor

bugbot run

@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.

Stale Bugbot comment from a previous run.

@devin-ai-integration

Copy link
Copy Markdown
Contributor Author

Re the precedence discussion: scoping served-over-requested pricing to Databricks is a small change in _resolve_service_tier, happy to make it if maintainers prefer

…vice tier tests

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

Copy link
Copy Markdown
Contributor

@greptileai

@litellm-stampbot

Copy link
Copy Markdown
Contributor

bugbot run

@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 76cd3e9. Configure here.

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>
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>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@github-actions

Copy link
Copy Markdown
Contributor

#41834 was fixed by #46111 on main, so this pull request is closed. Reopen it if something was missed.

@github-actions github-actions Bot closed this Oct 13, 2026
@kerry-berri

Copy link
Copy Markdown
Contributor

Superseded by #46111, which landed the remaining service_tier changes on top of #42870. Closing this one.

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.

[Bug]: Databricks service_tier is dropped in request translation, streaming, and pricing

4 participants