mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(databricks): delegate to generic_cost_per_token for cache/audio/reasoning
Resolves litellm/litellm#35608 — Databricks cost_calculator ignored cached tokens, audio tokens, and reasoning tokens. 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).
This commit is contained in:
michaelxer
parent
e2950a8995
commit
0576edf003
2 changed files with 172 additions and 36 deletions
|
|
@ -1,12 +1,48 @@
|
|||
"""
|
||||
Helper util for handling databricks-specific cost calculation
|
||||
- e.g.: handling 'dbrx-instruct-*'
|
||||
|
||||
Token billing (cache read/write, audio, reasoning) is delegated to
|
||||
generic_cost_per_token so Databricks stays consistent with other
|
||||
OpenAI-compatible providers.
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
|
||||
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 _resolve_databricks_base_model(model: str) -> str:
|
||||
"""
|
||||
Map common Databricks deployment aliases onto pricing-table keys.
|
||||
|
||||
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.
|
||||
"""
|
||||
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"
|
||||
):
|
||||
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"
|
||||
):
|
||||
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")
|
||||
):
|
||||
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
|
||||
|
||||
|
||||
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
|
||||
|
|
@ -20,36 +56,7 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
|
|||
Returns:
|
||||
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
|
||||
"""
|
||||
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(
|
||||
"meta-llama-3.1-405b-instruct"
|
||||
):
|
||||
base_model = "databricks-meta-llama-3-1-405b-instruct"
|
||||
elif (
|
||||
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")
|
||||
|
||||
## CALCULATE INPUT COST
|
||||
|
||||
prompt_cost: Final[float] = usage["prompt_tokens"] * model_info["input_cost_per_token"]
|
||||
|
||||
## CALCULATE OUTPUT COST
|
||||
completion_cost: Final = usage["completion_tokens"] * model_info["output_cost_per_token"]
|
||||
|
||||
return prompt_cost, completion_cost
|
||||
base_model = _resolve_databricks_base_model(model)
|
||||
return generic_cost_per_token(
|
||||
model=base_model, usage=usage, custom_llm_provider="databricks"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
Loading…
Add table
Reference in a new issue