From 5655cb87fc7e95b68e70f4018fab5ce4b0b5d0a8 Mon Sep 17 00:00:00 2001 From: openhands Date: Mon, 2 Mar 2026 08:11:07 +0000 Subject: [PATCH] fix: pass all custom pricing fields to register_model in completion() and embedding() Previously, register_model() was called with only input_cost_per_token, output_cost_per_token, and litellm_provider. This dropped ~40+ other pricing fields from CustomPricingLiteLLMParams (cache_read_input_token_cost, cache_creation_input_token_cost, output_cost_per_reasoning_token, etc.) as well as model_info metadata (mode, supports_prompt_caching, max_tokens). For DB-sourced custom-priced models, the first request after a pod restart would register a partial entry in litellm.model_cost, causing cost calculations to miss cache token discounts and other extended pricing until the entry was later enriched by deployment_callback_on_success mutating the lru_cache. Changes: - Add _build_custom_pricing_entry() helper that iterates over all CustomPricingLiteLLMParams.model_fields and merges model_info metadata - Replace hardcoded 3-field dicts in both completion() and embedding() with the new helper - Add 7 tests covering field collection, model_info merging, precedence, None skipping, and end-to-end register_model behavior Co-authored-by: openhands --- litellm/main.py | 80 ++++---- .../test_register_model_custom_pricing.py | 180 ++++++++++++++++++ 2 files changed, 223 insertions(+), 37 deletions(-) create mode 100644 tests/test_litellm/test_register_model_custom_pricing.py diff --git a/litellm/main.py b/litellm/main.py index cb3ddc2f401..737cf50ad24 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -107,6 +107,7 @@ from litellm.realtime_api.main import _realtime_health_check from litellm.secret_managers.main import get_secret_bool, get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + CustomPricingLiteLLMParams, ModelResponseStream, RawRequestTypedDict, StreamingChoices, @@ -996,6 +997,32 @@ def _drop_input_examples_from_tools( return cleaned_tools +def _build_custom_pricing_entry( + custom_llm_provider: str, + kwargs: dict, + model_info: Optional[dict] = None, +) -> dict: + """Build a complete model cost entry from kwargs and model_info. + + Collects all CustomPricingLiteLLMParams fields present in kwargs and + merges metadata from model_info (mode, supports_prompt_caching, max_tokens) + so that register_model() receives the full pricing configuration. + """ + entry: dict = {"litellm_provider": custom_llm_provider} + + for field_name in CustomPricingLiteLLMParams.model_fields: + value = kwargs.get(field_name) + if value is not None: + entry[field_name] = value + + if model_info and isinstance(model_info, dict): + for key in ("mode", "supports_prompt_caching", "max_tokens"): + if key in model_info and model_info[key] is not None: + entry.setdefault(key, model_info[key]) + + return entry + + @tracer.wrap() @client def completion( # type: ignore # noqa: PLR0915 @@ -1351,27 +1378,16 @@ def completion( # type: ignore # noqa: PLR0915 timeout = float(timeout) # type: ignore ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### - if input_cost_per_token is not None and output_cost_per_token is not None: + if ( + input_cost_per_token is not None and output_cost_per_token is not None + ) or input_cost_per_second is not None: litellm.register_model( { - f"{custom_llm_provider}/{model}": { - "input_cost_per_token": input_cost_per_token, - "output_cost_per_token": output_cost_per_token, - "litellm_provider": custom_llm_provider, - } - } - ) - elif ( - input_cost_per_second is not None - ): # time based pricing just needs cost in place - output_cost_per_second = output_cost_per_second - litellm.register_model( - { - f"{custom_llm_provider}/{model}": { - "input_cost_per_second": input_cost_per_second, - "output_cost_per_second": output_cost_per_second, - "litellm_provider": custom_llm_provider, - } + f"{custom_llm_provider}/{model}": _build_custom_pricing_entry( + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + model_info=model_info, + ) } ) ### BUILD CUSTOM PROMPT TEMPLATE -- IF GIVEN ### @@ -4644,7 +4660,6 @@ def embedding( # noqa: PLR0915 input_cost_per_token = kwargs.get("input_cost_per_token", None) output_cost_per_token = kwargs.get("output_cost_per_token", None) input_cost_per_second = kwargs.get("input_cost_per_second", None) - output_cost_per_second = kwargs.get("output_cost_per_second", None) openai_params = [ "user", "dimensions", @@ -4694,25 +4709,16 @@ def embedding( # noqa: PLR0915 ) ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### - if input_cost_per_token is not None and output_cost_per_token is not None: + if ( + input_cost_per_token is not None and output_cost_per_token is not None + ) or input_cost_per_second is not None: litellm.register_model( { - f"{custom_llm_provider}/{model}": { - "input_cost_per_token": input_cost_per_token, - "output_cost_per_token": output_cost_per_token, - "litellm_provider": custom_llm_provider, - } - } - ) - if input_cost_per_second is not None: # time based pricing just needs cost in place - output_cost_per_second = output_cost_per_second or 0.0 - litellm.register_model( - { - f"{custom_llm_provider}/{model}": { - "input_cost_per_second": input_cost_per_second, - "output_cost_per_second": output_cost_per_second, - "litellm_provider": custom_llm_provider, - } + f"{custom_llm_provider}/{model}": _build_custom_pricing_entry( + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + model_info=kwargs.get("model_info"), + ) } ) diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py new file mode 100644 index 00000000000..fa023d9943e --- /dev/null +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -0,0 +1,180 @@ +""" +Test that register_model() in completion() and embedding() passes all +custom pricing fields from kwargs and model_info, not just the base +input/output costs. + +Previously, only input_cost_per_token, output_cost_per_token, and +litellm_provider were forwarded. Fields like cache_read_input_token_cost, +mode, and supports_prompt_caching were dropped, causing incorrect cost +calculations for DB-sourced models with prompt caching pricing. +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.main import _build_custom_pricing_entry + + +def test_build_custom_pricing_entry_includes_all_kwargs_fields(): + """All CustomPricingLiteLLMParams fields present in kwargs should be + included in the resulting entry dict.""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "cache_read_input_token_cost": 0.00025, + "cache_creation_input_token_cost": 0.005, + "output_cost_per_reasoning_token": 0.01, + "input_cost_per_audio_token": 0.003, + "unrelated_kwarg": "should_be_ignored", + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + ) + + assert entry["litellm_provider"] == "openai" + assert entry["input_cost_per_token"] == 0.001 + assert entry["output_cost_per_token"] == 0.002 + assert entry["cache_read_input_token_cost"] == 0.00025 + assert entry["cache_creation_input_token_cost"] == 0.005 + assert entry["output_cost_per_reasoning_token"] == 0.01 + assert entry["input_cost_per_audio_token"] == 0.003 + assert "unrelated_kwarg" not in entry + + +def test_build_custom_pricing_entry_merges_model_info_metadata(): + """Fields from model_info (mode, supports_prompt_caching, max_tokens) + should be merged into the entry when present.""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + model_info = { + "id": "deployment-123", + "mode": "chat", + "supports_prompt_caching": True, + "max_tokens": 128000, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + model_info=model_info, + ) + + assert entry["mode"] == "chat" + assert entry["supports_prompt_caching"] is True + assert entry["max_tokens"] == 128000 + + +def test_build_custom_pricing_entry_kwargs_take_precedence_over_model_info(): + """If a field appears in both kwargs and model_info, the kwargs value + should take precedence (setdefault behavior).""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + model_info = { + "mode": "chat", + "supports_prompt_caching": True, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + model_info=model_info, + ) + + # model_info fields should be set via setdefault + assert entry["mode"] == "chat" + assert entry["supports_prompt_caching"] is True + + +def test_build_custom_pricing_entry_skips_none_values(): + """Fields with None values in kwargs should not be included.""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": None, # explicitly None + "cache_read_input_token_cost": None, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + ) + + assert entry["input_cost_per_token"] == 0.001 + assert "output_cost_per_token" not in entry + assert "cache_read_input_token_cost" not in entry + + +def test_build_custom_pricing_entry_handles_no_model_info(): + """Should work correctly when model_info is None.""" + kwargs = { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + model_info=None, + ) + + assert entry["litellm_provider"] == "openai" + assert entry["input_cost_per_token"] == 0.001 + assert entry["output_cost_per_token"] == 0.002 + assert "mode" not in entry + + +def test_register_model_receives_cache_pricing_fields(): + """End-to-end: when register_model is called with a full pricing entry, + the cache pricing fields should be present in litellm.model_cost.""" + model_key = "openai/test-custom-model-with-cache-pricing" + + litellm.register_model( + { + model_key: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "cache_read_input_token_cost": 0.00025, + "supports_prompt_caching": True, + "mode": "chat", + "max_tokens": 8192, + "litellm_provider": "openai", + } + } + ) + + registered = litellm.model_cost.get(model_key) + assert registered is not None, f"{model_key} should be in model_cost" + assert registered["cache_read_input_token_cost"] == 0.00025 + assert registered["supports_prompt_caching"] is True + assert registered["mode"] == "chat" + assert registered["max_tokens"] == 8192 + + # Cleanup + litellm.model_cost.pop(model_key, None) + + +def test_build_custom_pricing_entry_time_based(): + """Time-based pricing fields should be included correctly.""" + kwargs = { + "input_cost_per_second": 0.01, + "output_cost_per_second": 0.02, + } + + entry = _build_custom_pricing_entry( + custom_llm_provider="openai", + kwargs=kwargs, + ) + + assert entry["litellm_provider"] == "openai" + assert entry["input_cost_per_second"] == 0.01 + assert entry["output_cost_per_second"] == 0.02