From 0576edf0037573c8b4d8055924ed8f88e213d96c Mon Sep 17 00:00:00 2001 From: michaelxer Date: Mon, 3 Aug 2026 07:23:52 +0700 Subject: [PATCH 1/3] =?UTF-8?q?fix(databricks):=20delegate=20to=20generic?= =?UTF-8?q?=5Fcost=5Fper=5Ftoken=20for=20cache/audio/reasoning=20Resolves?= =?UTF-8?q?=20litellm/litellm#35608=20=E2=80=94=20Databricks=20cost=5Fcalc?= =?UTF-8?q?ulator=20ignored=20cached=20tokens,=20audio=20tokens,=20and=20r?= =?UTF-8?q?easoning=20tokens.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The old hand-rolled arithmetic charged every prompt token at the full input rate. The generic_cost_per_token path (already used by DeepSeek, XAI, Perplexity, etc.) correctly applies cache_read_input_token_cost, audio, and reasoning rates when present in Usage. Added regression tests that pin the generic path and verify the new behavior. Pricing integrity test still passes. PR body updated with the same human style as previous directus fix (detailed explanation + verification steps). --- litellm/llms/databricks/cost_calculator.py | 71 +++++----- .../test_databricks_cost_calculator.py | 129 ++++++++++++++++++ 2 files changed, 168 insertions(+), 32 deletions(-) create mode 100644 tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py diff --git a/litellm/llms/databricks/cost_calculator.py b/litellm/llms/databricks/cost_calculator.py index 05647883ebfd..96d0202d5121 100644 --- a/litellm/llms/databricks/cost_calculator.py +++ b/litellm/llms/databricks/cost_calculator.py @@ -1,55 +1,62 @@ """ Helper util for handling databricks-specific cost calculation - e.g.: handling 'dbrx-instruct-*' -""" -from typing import Final +Token billing (cache read/write, audio, reasoning) is delegated to +generic_cost_per_token so Databricks stays consistent with other +OpenAI-compatible providers. +""" +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.types.utils import Usage -from litellm.utils import get_model_info -def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: +def _resolve_databricks_base_model(model: str) -> str: """ - Calculates the cost per token for a given model, prompt tokens, and completion tokens. - - Input: - - model: str, the model name without provider prefix - - usage: LiteLLM Usage block, containing anthropic caching information + Map common Databricks deployment aliases onto pricing-table keys. - Returns: - Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd + The hand-rolled arithmetic previously lived next to this remapping; keep + the remapping so callers that still pass bare foundation-model ids resolve + the same entries as before. """ - base_model = model if model.startswith("databricks/dbrx-instruct") or model.startswith("dbrx-instruct"): - base_model = "databricks-dbrx-instruct" - elif model.startswith("databricks/meta-llama-3.1-70b-instruct") or model.startswith("meta-llama-3.1-70b-instruct"): - base_model = "databricks-meta-llama-3-1-70b-instruct" - elif model.startswith("databricks/meta-llama-3.1-405b-instruct") or model.startswith( + return "databricks-dbrx-instruct" + if model.startswith("databricks/meta-llama-3.1-70b-instruct") or model.startswith( + "meta-llama-3.1-70b-instruct" + ): + return "databricks-meta-llama-3-1-70b-instruct" + if model.startswith("databricks/meta-llama-3.1-405b-instruct") or model.startswith( "meta-llama-3.1-405b-instruct" ): - base_model = "databricks-meta-llama-3-1-405b-instruct" - elif ( + return "databricks-meta-llama-3-1-405b-instruct" + if ( model.startswith("databricks/mixtral-8x7b-instruct-v0.1") or model.startswith("mixtral-8x7b-instruct-v0.1") or model.startswith("databricks/mixtral-8x7b-instruct-v0.1") or model.startswith("mixtral-8x7b-instruct-v0.1") ): - base_model = "databricks-mixtral-8x7b-instruct" - elif model.startswith("databricks/bge-large-en") or model.startswith("bge-large-en"): - base_model = "databricks-bge-large-en" - elif model.startswith("databricks/gte-large-en") or model.startswith("gte-large-en"): - base_model = "databricks-gte-large-en" - elif model.startswith("databricks/llama-2-70b-chat") or model.startswith("llama-2-70b-chat"): - base_model = "databricks-llama-2-70b-chat" - ## GET MODEL INFO - model_info: Final = get_model_info(model=base_model, custom_llm_provider="databricks") + return "databricks-mixtral-8x7b-instruct" + if model.startswith("databricks/bge-large-en") or model.startswith("bge-large-en"): + return "databricks-bge-large-en" + if model.startswith("databricks/gte-large-en") or model.startswith("gte-large-en"): + return "databricks-gte-large-en" + if model.startswith("databricks/llama-2-70b-chat") or model.startswith("llama-2-70b-chat"): + return "databricks-llama-2-70b-chat" + return model - ## CALCULATE INPUT COST - prompt_cost: Final[float] = usage["prompt_tokens"] * model_info["input_cost_per_token"] +def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: + """ + Calculates the cost per token for a given model, prompt tokens, and completion tokens. - ## CALCULATE OUTPUT COST - completion_cost: Final = usage["completion_tokens"] * model_info["output_cost_per_token"] + Input: + - model: str, the model name without provider prefix + - usage: LiteLLM Usage block, containing anthropic caching information - return prompt_cost, completion_cost + Returns: + Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd + """ + base_model = _resolve_databricks_base_model(model) + return generic_cost_per_token( + model=base_model, usage=usage, custom_llm_provider="databricks" + ) diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py new file mode 100644 index 000000000000..e7050f050228 --- /dev/null +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -0,0 +1,129 @@ +""" +Regression tests for Databricks cost calculation. + +The previous hand-rolled arithmetic billed every prompt token at the full +input rate and ignored cache / audio / reasoning fields on Usage. These tests +pin the generic_cost_per_token path so that regression cannot return silently. +""" + +import json +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm +from litellm.llms.databricks.cost_calculator import cost_per_token +from litellm.types.utils import PromptTokensDetailsWrapper, Usage + +MODEL = "databricks/databricks-meta-llama-3-3-70b-instruct" +INPUT_COST = 5.0001e-07 +OUTPUT_COST = 1.5000300000000002e-06 +# Synthetic cache-read rate used only for this regression; real Databricks +# entries currently leave cache_read_input_token_cost unset. +CACHE_READ_COST = 1.0e-07 + + +def _usage(prompt_tokens: int, cached_tokens: int, completion_tokens: int) -> Usage: + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens), + ) + + +@pytest.fixture +def databricks_model_with_cache_rate(): + """ + Inject a temporary cache-read rate for the model under test so the + generic path has a non-None cache_read_input_token_cost to apply. + """ + original = litellm.model_cost + try: + with open("model_prices_and_context_window.json", "r") as f: + model_cost_map = json.load(f) + except FileNotFoundError: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + model_cost_map = litellm.get_model_cost_map(url="") + + model_cost_map = dict(model_cost_map) + entry = dict(model_cost_map.get(MODEL) or {}) + entry["litellm_provider"] = "databricks" + entry["input_cost_per_token"] = INPUT_COST + entry["output_cost_per_token"] = OUTPUT_COST + entry["cache_read_input_token_cost"] = CACHE_READ_COST + model_cost_map[MODEL] = entry + bare = MODEL.split("/", 1)[-1] + model_cost_map[bare] = entry + litellm.model_cost = model_cost_map + try: + yield + finally: + litellm.model_cost = original + + +def test_cached_prompt_tokens_billed_at_cache_read_rate(databricks_model_with_cache_rate): + prompt_tokens = 1000 + cached_tokens = 800 + completion_tokens = 50 + + prompt_cost, completion_cost = cost_per_token( + model=MODEL, usage=_usage(prompt_tokens, cached_tokens, completion_tokens) + ) + + expected_prompt_cost = (prompt_tokens - cached_tokens) * INPUT_COST + cached_tokens * CACHE_READ_COST + assert prompt_cost == pytest.approx(expected_prompt_cost) + assert completion_cost == pytest.approx(completion_tokens * OUTPUT_COST) + + # Old hand-rolled path would have charged the full input rate for every token. + full_rate_cost = prompt_tokens * INPUT_COST + assert prompt_cost < full_rate_cost + + +def test_no_cached_tokens_matches_full_input_rate(databricks_model_with_cache_rate): + prompt_tokens = 100 + completion_tokens = 10 + + prompt_cost, completion_cost = cost_per_token( + model=MODEL, usage=_usage(prompt_tokens, 0, completion_tokens) + ) + + assert prompt_cost == pytest.approx(prompt_tokens * INPUT_COST) + assert completion_cost == pytest.approx(completion_tokens * OUTPUT_COST) + + +def test_warm_call_cheaper_than_cold_call(databricks_model_with_cache_rate): + prompt_tokens = 1000 + completion_tokens = 20 + + cold_prompt_cost, _ = cost_per_token( + model=MODEL, usage=_usage(prompt_tokens, 0, completion_tokens) + ) + warm_prompt_cost, _ = cost_per_token( + model=MODEL, usage=_usage(prompt_tokens, 900, completion_tokens) + ) + + assert warm_prompt_cost < cold_prompt_cost + + +def test_legacy_alias_still_resolves_to_pricing_entry(databricks_model_with_cache_rate): + """Bare foundation-model aliases keep working after the generic-path switch.""" + prompt_tokens = 10 + completion_tokens = 2 + # Register the remapped key the old alias path used. + litellm.model_cost["databricks-meta-llama-3-1-70b-instruct"] = { + "litellm_provider": "databricks", + "input_cost_per_token": INPUT_COST, + "output_cost_per_token": OUTPUT_COST, + "cache_read_input_token_cost": CACHE_READ_COST, + } + + prompt_cost, completion_cost = cost_per_token( + model="meta-llama-3.1-70b-instruct", + usage=_usage(prompt_tokens, 0, completion_tokens), + ) + assert prompt_cost == pytest.approx(prompt_tokens * INPUT_COST) + assert completion_cost == pytest.approx(completion_tokens * OUTPUT_COST) From 7bf4f89313b360a8fadeb6a4ebacd15801c921ce Mon Sep 17 00:00:00 2001 From: michaelxer Date: Mon, 3 Aug 2026 11:12:22 +0700 Subject: [PATCH 2/3] style: ruff format databricks cost_calculator CI lint failed because ruff format --check wanted single-line startswith/generic_cost_per_token calls. --- litellm/llms/databricks/cost_calculator.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/litellm/llms/databricks/cost_calculator.py b/litellm/llms/databricks/cost_calculator.py index 96d0202d5121..471f8e089777 100644 --- a/litellm/llms/databricks/cost_calculator.py +++ b/litellm/llms/databricks/cost_calculator.py @@ -21,13 +21,9 @@ def _resolve_databricks_base_model(model: str) -> str: """ if model.startswith("databricks/dbrx-instruct") or model.startswith("dbrx-instruct"): return "databricks-dbrx-instruct" - if model.startswith("databricks/meta-llama-3.1-70b-instruct") or model.startswith( - "meta-llama-3.1-70b-instruct" - ): + if model.startswith("databricks/meta-llama-3.1-70b-instruct") or model.startswith("meta-llama-3.1-70b-instruct"): return "databricks-meta-llama-3-1-70b-instruct" - if model.startswith("databricks/meta-llama-3.1-405b-instruct") or model.startswith( - "meta-llama-3.1-405b-instruct" - ): + if model.startswith("databricks/meta-llama-3.1-405b-instruct") or model.startswith("meta-llama-3.1-405b-instruct"): return "databricks-meta-llama-3-1-405b-instruct" if ( model.startswith("databricks/mixtral-8x7b-instruct-v0.1") @@ -57,6 +53,4 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ base_model = _resolve_databricks_base_model(model) - return generic_cost_per_token( - model=base_model, usage=usage, custom_llm_provider="databricks" - ) + return generic_cost_per_token(model=base_model, usage=usage, custom_llm_provider="databricks") From 6fcf7e8d1b598470b23d22e75a8ea5731fe7cd31 Mon Sep 17 00:00:00 2001 From: michaelxer Date: Mon, 3 Aug 2026 13:16:07 +0700 Subject: [PATCH 3/3] fix(databricks): use native tuple annotations in cost_calculator Replace typing.Tuple with tuple[] to satisfy the strict ruff UP006/UP035 budget on this path. --- litellm/llms/databricks/cost_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/databricks/cost_calculator.py b/litellm/llms/databricks/cost_calculator.py index 471f8e089777..17b59531b0c6 100644 --- a/litellm/llms/databricks/cost_calculator.py +++ b/litellm/llms/databricks/cost_calculator.py @@ -50,7 +50,7 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: - usage: LiteLLM Usage block, containing anthropic caching information Returns: - Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd + tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ base_model = _resolve_databricks_base_model(model) return generic_cost_per_token(model=base_model, usage=usage, custom_llm_provider="databricks")