Repository navigation
fix(databricks): carry service_tier through request, streaming and cost calc - #41837
devin-ai-integration[bot] wants to merge 15 commits into
Conversation
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
|
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
| usage_no_tier = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) | ||
| priority_reference = completion_cost( |
There was a problem hiding this comment.
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!
|
@greptileai please re-review, the streaming wrapper, Responses output, and test convention findings are addressed in the latest commits |
|
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 For Databricks the distinction matters materially because priority can fall back: The user's billing/UX expectation is:
That avoids needing a global rule like: which is exactly what creates the cross-provider compatibility change called out in review. Conceptually, the lifecycle becomes: 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 Regression shapeA compact matrix would protect both Databricks and compatibility:
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. |
… 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>
03cfabf to
b2e29f8
Compare
|
bugbot run |
|
Re the precedence discussion: scoping served-over-requested pricing to Databricks is a small change in |
…vice tier tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
|
bugbot run |
There was a problem hiding this comment.
✅ 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>
TLDR
Problem this solves:
service_tier: priorityis dropped from the outgoing requestservice_tierin the request bodyservice_tierin the stream wrapper and aggregationservice_tierdefaulttier did not override a requestedprioritytier in cost calcHow it solves it:
service_tierto Databricks supported params and carry it on responses and stream chunksservice_tierfrom the Responses request body, then kwargsservice_tierthrough the Databricks cost path to tiered pricing keysservice_tieron emitted and usage chunks;stream_chunk_builderkeeps itservice_tierfrom the chat completiontests/integration/pricing/test_service_tier_pricing.py, red when_resolve_service_tierreturns None)User Flow
Before: a developer sends a Databricks priority request through the proxy and is billed the standard price
"model": "databricks/<endpoint>"and"service_tier": "priority"service_tier, so it is served on the standard tierservice_tierfield andx-litellm-response-costmatches standard pricingAfter: the same request is served and billed as priority, and a fallback to default is billed as default
"service_tier": "priority"service_tier: priority"service_tier": "priority"(also on streamed chunks) andx-litellm-response-costis the priority price"service_tier": "default", the cost header and the logs page show the default price"service_tier": "priority"in the bodyRelevant 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
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@greptileaito 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=TrueShared setup: register
databricks/offline-tier-testwithinput_cost_per_token=2e-6,output_cost_per_token=4e-6,cache_read_input_token_cost=1e-6and the_prioritykeys at 2x, then callget_optional_params(custom_llm_provider="databricks", service_tier="priority", drop_params=True)andlitellm.completion_cost(...)Before (4e21178)
get_optional_paramsoutput containsservice_tier:Falsecompletion_costwith responseservice_tier="default":0.00038completion_costwith responseservice_tier="priority":0.00038(priority ignored)completion_costwith requestedpriorityand returneddefault:0.00038After (211a545)
get_optional_paramsoutput containsservice_tier:Truecompletion_costwith responseservice_tier="default":0.00038completion_costwith responseservice_tier="priority":0.00076completion_costwith requestedpriorityand returneddefault:0.00038Type
🐛 Bug Fix
Caveats (if any)
Medium
completion_costtier precedence changed for every provider: response tier, then usage tier, then requested tierservice_tier=oroptional_params) won over the tier on the responsetest_completion_cost_service_tier_prioritywas rewritten to the new orderLow
trafficTypemapping moved into the resolver as the last fallback, same effective orderFinal Attestation
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_costservice-tier precedence and spend headers/logs for every provider; incorrect tier resolution would mis-price requests, especially priority Databricks traffic.Overview
Fixes Databricks
service_tierbeing dropped so priority (and other tiered) traffic is forwarded, surfaced on responses/streams, and billed with the right tier keys.Databricks & Responses API:
service_tieris now a supported Databricks param, copied onto non-streaming and streaming chunks, and passed intodatabricks_cost_per_tokenfor_prioritypricing. The Responses API bridge readsservice_tierfrom the request body (then kwargs) and copies the served tier from the chat completion onto the Responses output.Streaming:
CustomStreamWrapperremembers the latest chunkservice_tierand stamps it on emitted chunks (including usage-only frames);stream_chunk_builderpicks 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_paramstier, with Geminitraffic_typeas a final fallback. Logging and token breakdowns use this resolved tier—so a returneddefaulttier is charged at default rates even when the client asked forpriority.Reviewed by Cursor Bugbot for commit 76cd3e9. Bugbot is set up for automated code reviews on this repo. Configure here.