From 718b759e1c2e232825e45e57908034a1034486dd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 4 Jul 2026 17:49:32 -0700 Subject: [PATCH 1/2] fix(main): stop per-request custom pricing from clobbering shared model_cost pricing A request routed through a wildcard deployment with explicit zero pricing (e.g. openai/* with input_cost_per_token: 0) registered that pricing on the shared {provider}/{model} key in litellm.model_cost, so sibling deployments relying on built-in pricing logged $0 until process restart (LIT-3991). Request-time registration in completion()/embedding() now mirrors the router-startup isolation: router-originated requests register full pricing under the deployment's unique model id only, while the shared backend key receives the entry with custom pricing fields stripped. Direct SDK calls without a router deployment id keep the legacy shared-key registration. The stripping logic is shared via CustomPricingLiteLLMParams.strip_custom_pricing_fields and reused by Router._create_deployment and Router.add_deployment. --- litellm/main.py | 74 +++++-- litellm/router.py | 6 +- litellm/types/utils.py | 11 ++ .../test_register_model_custom_pricing.py | 180 ++++++++++++++++++ .../test_router_model_cost_isolation.py | 73 +++++++ 5 files changed, 324 insertions(+), 20 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index b3c176a13472..2f36f6cb3ff8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1081,6 +1081,54 @@ def _build_custom_pricing_entry( return entry +def _get_router_deployment_id(kwargs: dict) -> Optional[str]: + for metadata_key in ("litellm_metadata", "metadata"): + metadata = kwargs.get(metadata_key) or {} + if not isinstance(metadata, dict): + continue + deployment_model_info = metadata.get("model_info") or {} + if not isinstance(deployment_model_info, dict): + continue + deployment_id = deployment_model_info.get("id") + if deployment_id is not None: + return str(deployment_id) + return None + + +def _register_custom_pricing_for_request( + model: str, + custom_llm_provider: str, + kwargs: dict, + model_info: Optional[dict], +) -> None: + """Register per-request custom pricing in litellm.model_cost. + + Router-originated requests (identified by the deployment id the router puts + in metadata) get their full pricing registered under that unique id only; + the shared ``{provider}/{model}`` key receives the entry with pricing fields + stripped, mirroring Router._create_deployment. This keeps one deployment's + pricing overrides (e.g. a zero-cost wildcard) from clobbering built-in + pricing used by sibling deployments of the same backend model. Direct SDK + calls keep the legacy behavior of registering the shared key with pricing. + """ + entry = _build_custom_pricing_entry( + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + model_info=model_info, + ) + shared_key = f"{custom_llm_provider}/{model}" + deployment_id = _get_router_deployment_id(kwargs) + if deployment_id is None: + litellm.register_model({shared_key: entry}) + return + litellm.register_model( + { + deployment_id: entry, + shared_key: CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry), + } + ) + + def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: _azure_detection_model = ctx._azure_detection_model acompletion = ctx.acompletion @@ -5107,14 +5155,11 @@ def completion( # type: ignore 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}": _build_custom_pricing_entry( - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, - model_info=model_info, - ) - } + _register_custom_pricing_for_request( + model=model, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + model_info=model_info, ) ### BUILD CUSTOM PROMPT TEMPLATE -- IF GIVEN ### custom_prompt_dict = {} # type: ignore @@ -5957,14 +6002,11 @@ def embedding( ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### 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}": _build_custom_pricing_entry( - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, - model_info=kwargs.get("model_info"), - ) - } + _register_custom_pricing_for_request( + model=model, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + model_info=kwargs.get("model_info"), ) litellm_params_dict = get_litellm_params(**kwargs) diff --git a/litellm/router.py b/litellm/router.py index 9f333e471f30..99a9d3ca7217 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7354,8 +7354,7 @@ def _create_deployment( # deployment sharing the same backend model name. # Each deployment's full pricing is already stored under its # unique model_id above. - _custom_pricing_fields = CustomPricingLiteLLMParams.model_fields.keys() - _shared_model_info = {k: v for k, v in _model_info.items() if k not in _custom_pricing_fields} + _shared_model_info = CustomPricingLiteLLMParams.strip_custom_pricing_fields(_model_info) _existing_shared_mode = (cast(Optional[dict], litellm.model_cost.get(_model_name, {})) or {}).get("mode") _deployment_mode = _shared_model_info.get("mode") # Keep the built-in bridge mode stable for shared backend keys. @@ -8019,8 +8018,7 @@ def add_deployment(self, deployment: Deployment) -> Optional[Deployment]: # deployment sharing the same backend model name. # Each deployment's full pricing is already stored under its # unique model_id above (when present). - _custom_pricing_fields = CustomPricingLiteLLMParams.model_fields.keys() - _shared_model_info = {k: v for k, v in _model_info_dict.items() if k not in _custom_pricing_fields} + _shared_model_info = CustomPricingLiteLLMParams.strip_custom_pricing_fields(_model_info_dict) _backend_alias_cost = {_model_name: _shared_model_info} if "responses/" in _model_name: _stripped_model_name = _model_name.replace("responses/", "") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2e0eceaddd96..4b90317a50fc 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3048,6 +3048,17 @@ class CustomPricingLiteLLMParams(BaseModel): regional_processing_uplift_multiplier_eu: Optional[float] = None regional_processing_uplift_multiplier_us: Optional[float] = None + @classmethod + def strip_custom_pricing_fields(cls, model_info: Dict[str, Any]) -> Dict[str, Any]: + """Return a copy of ``model_info`` without per-deployment custom pricing fields. + + Used when registering a deployment's info under the shared + ``{provider}/{model}`` key in ``litellm.model_cost``, so one deployment's + pricing overrides don't pollute sibling deployments that share the same + backend model. Full pricing stays under the deployment's unique model id. + """ + return {k: v for k, v in model_info.items() if k not in cls.model_fields} + # Server-controlled fields that bound or drive an interceptor's agentic loop # (depth, cycle fingerprints, ceiling, code-interpreter sandbox state). Listed diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index e384d3e1161d..8c3f690982b8 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -9,15 +9,32 @@ calculations for DB-sourced models with prompt caching pricing. """ +import copy import os import sys +import pytest + 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 +from litellm.utils import _invalidate_model_cost_lowercase_map + + +def _snapshot_model_cost_entries(keys): + return {key: copy.deepcopy(litellm.model_cost.get(key)) for key in keys} + + +def _restore_model_cost_entries(original_entries): + for key, value in original_entries.items(): + if value is None: + litellm.model_cost.pop(key, None) + else: + litellm.model_cost[key] = value + _invalidate_model_cost_lowercase_map() def test_build_custom_pricing_entry_includes_all_kwargs_fields(): @@ -471,3 +488,166 @@ def test_register_model_router_add_deployment_custom_pricing_applies(): litellm.model_cost.pop(model_key, None) litellm.model_cost.pop(deployment_model, None) del router + + +def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing(): + """LIT-3991: a router-originated embedding request that carries explicit + zero custom pricing (e.g. resolved through an ``openai/*`` wildcard + deployment with ``input_cost_per_token: 0``) must not overwrite the shared + ``openai/text-embedding-3-small`` entry in ``litellm.model_cost``. Before + the fix, one call through the wildcard poisoned the shared key and every + sibling deployment relying on built-in pricing logged $0 until restart. + """ + shared_key = "openai/text-embedding-3-small" + deployment_id = "lit3991-wildcard-embed-zero" + snapshot = _snapshot_model_cost_entries( + [shared_key, "text-embedding-3-small", deployment_id] + ) + builtin_input_cost = litellm.get_model_info(model=shared_key)[ + "input_cost_per_token" + ] + assert builtin_input_cost > 0 + + try: + litellm.embedding( + model=shared_key, + input=["hello"], + api_key="fake-key", + input_cost_per_token=0.0, + output_cost_per_token=0.0, + model_info={"id": deployment_id}, + metadata={"model_info": {"id": deployment_id}}, + mock_response=[0.1, 0.2], + ) + + assert ( + litellm.get_model_info(model=shared_key)["input_cost_per_token"] + == builtin_input_cost + ), "wildcard deployment's zero pricing leaked into the shared model_cost key" + assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.0 + assert litellm.model_cost[deployment_id]["output_cost_per_token"] == 0.0 + + sibling_response = litellm.embedding( + model=shared_key, + input=["hello"], + api_key="fake-key", + mock_response=[0.1, 0.2], + ) + sibling_cost = litellm.completion_cost( + completion_response=sibling_response, call_type="embedding" + ) + assert sibling_cost == pytest.approx(10 * builtin_input_cost) + finally: + _restore_model_cost_entries(snapshot) + + +def test_embedding_router_custom_pricing_costs_request_via_deployment_id(): + """The request that carries custom pricing must still be costed with that + pricing (via its deployment id entry), while the shared backend key keeps + the built-in rate for siblings. + """ + shared_key = "openai/text-embedding-3-small" + deployment_id = "lit3991-wildcard-embed-custom" + override_input_cost = 5e-05 + snapshot = _snapshot_model_cost_entries( + [shared_key, "text-embedding-3-small", deployment_id] + ) + builtin_input_cost = litellm.get_model_info(model=shared_key)[ + "input_cost_per_token" + ] + assert builtin_input_cost != override_input_cost + + try: + response = litellm.embedding( + model=shared_key, + input=["hello"], + api_key="fake-key", + input_cost_per_token=override_input_cost, + output_cost_per_token=override_input_cost * 2, + model_info={"id": deployment_id}, + metadata={"model_info": {"id": deployment_id}}, + mock_response=[0.1, 0.2], + ) + + request_cost = litellm.completion_cost( + completion_response=response, + model=shared_key, + custom_llm_provider="openai", + call_type="embedding", + custom_pricing=True, + router_model_id=deployment_id, + ) + assert request_cost == pytest.approx(10 * override_input_cost) + assert ( + litellm.get_model_info(model=shared_key)["input_cost_per_token"] + == builtin_input_cost + ) + finally: + _restore_model_cost_entries(snapshot) + + +def test_completion_router_zero_pricing_does_not_clobber_builtin_pricing(): + """Same isolation as the embedding path, exercised through completion().""" + shared_key = "openai/gpt-4o-mini" + deployment_id = "lit3991-wildcard-chat-zero" + snapshot = _snapshot_model_cost_entries( + [shared_key, "gpt-4o-mini", deployment_id] + ) + builtin_input_cost = litellm.get_model_info(model=shared_key)[ + "input_cost_per_token" + ] + assert builtin_input_cost > 0 + + try: + litellm.completion( + model=shared_key, + messages=[{"role": "user", "content": "hello"}], + api_key="fake-key", + input_cost_per_token=0.0, + output_cost_per_token=0.0, + model_info={"id": deployment_id}, + metadata={"model_info": {"id": deployment_id}}, + mock_response="hello back", + ) + + assert ( + litellm.get_model_info(model=shared_key)["input_cost_per_token"] + == builtin_input_cost + ), "wildcard deployment's zero pricing leaked into the shared model_cost key" + assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.0 + finally: + _restore_model_cost_entries(snapshot) + + +def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key(): + """Direct SDK calls (no router deployment id in metadata) keep the legacy + behavior: custom pricing is registered under ``{provider}/{model}`` and the + request is costed with it. + """ + model_key = "openai/lit3991-direct-sdk-embed-model" + override_input_cost = 3e-05 + try: + response = litellm.embedding( + model=model_key, + input=["hello"], + api_key="fake-key", + input_cost_per_token=override_input_cost, + output_cost_per_token=override_input_cost * 2, + mock_response=[0.1, 0.2], + ) + + assert ( + litellm.model_cost[model_key]["input_cost_per_token"] + == override_input_cost + ) + cost = litellm.completion_cost( + completion_response=response, + model=model_key, + custom_llm_provider="openai", + call_type="embedding", + custom_pricing=True, + ) + assert cost == pytest.approx(10 * override_input_cost) + finally: + litellm.model_cost.pop(model_key, None) + _invalidate_model_cost_lowercase_map() diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index d4ac9659f003..6db7b04b3b7c 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -681,3 +681,76 @@ def test_custom_pricing_isolated_from_sibling_via_proxy_model_info_path(): assert resolved["gemini-2.5-flash"] != resolved["custom-priced-flash"] finally: _restore_model_cost_entries(model_keys) + + +def test_wildcard_zero_cost_request_does_not_poison_named_deployment_pricing(): + """LIT-3991 end to end: a proxy has a named text-embedding-3-small + deployment relying on built-in pricing plus an ``openai/*`` wildcard with + explicit zero pricing. One embedding call routed through the wildcard must + not clobber the shared ``openai/text-embedding-3-small`` pricing; requests + to the named deployment afterwards must still cost non-zero. + """ + shared_key = "openai/text-embedding-3-small" + model_keys = { + shared_key: copy.deepcopy(litellm.model_cost.get(shared_key)), + "text-embedding-3-small": copy.deepcopy( + litellm.model_cost.get("text-embedding-3-small") + ), + "openai/*": copy.deepcopy(litellm.model_cost.get("openai/*")), + "lit3991-named": litellm.model_cost.get("lit3991-named"), + "lit3991-wildcard": litellm.model_cost.get("lit3991-wildcard"), + } + builtin_input_cost = litellm.get_model_info(model=shared_key)[ + "input_cost_per_token" + ] + assert builtin_input_cost > 0 + + try: + router = Router( + model_list=[ + { + "model_name": "text-embedding-3-small", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "fake-key-named", + }, + "model_info": {"id": "lit3991-named"}, + }, + { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_key": "fake-key-wildcard", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": {"id": "lit3991-wildcard"}, + }, + ], + ) + + router.embedding( + model="openai/text-embedding-3-small", + input=["hello"], + mock_response=[0.1, 0.2], + ) + + assert ( + litellm.get_model_info(model=shared_key)["input_cost_per_token"] + == builtin_input_cost + ), ( + "one call through the zero-cost wildcard poisoned the shared " + f"{shared_key} pricing for the named deployment" + ) + + named_response = router.embedding( + model="text-embedding-3-small", + input=["hello"], + mock_response=[0.1, 0.2], + ) + named_cost = litellm.completion_cost( + completion_response=named_response, call_type="embedding" + ) + assert named_cost == pytest.approx(10 * builtin_input_cost) + finally: + _restore_model_cost_entries(model_keys) From 58df1167fc9ad9b55c384e85e0685332decbf960 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 4 Jul 2026 18:12:36 -0700 Subject: [PATCH 2/2] test: update legacy tests that asserted per-request pricing leaking into shared model_cost test_router_fallbacks_with_custom_model_costs asserted the shared claude-sonnet-4-5-20250929 entry ends up with the deployment's 30/60 pricing, which is exactly the cross-deployment leak this PR removes; it now asserts the shared key keeps the built-in pricing, matching the test's stated goal. test_cost_calc.py::test_run computed streaming cost via completion_cost(response), which only matched the non-stream cost while the shared gpt-3.5-turbo entry was poisoned with the per-request 2/token pricing; it now passes the request's custom pricing explicitly via custom_cost_per_token. --- tests/local_testing/test_cost_calc.py | 11 ++++++++++- tests/local_testing/test_router_fallbacks.py | 6 ++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/tests/local_testing/test_cost_calc.py b/tests/local_testing/test_cost_calc.py index ab4d44d22406..3623af598480 100644 --- a/tests/local_testing/test_cost_calc.py +++ b/tests/local_testing/test_cost_calc.py @@ -101,7 +101,16 @@ def test_run(model: str): pytest.skip( "LLM API returning inconsistent usage" ) # handles transient openai errors - streaming_cost_calc = completion_cost(response) * 100 + streaming_cost_calc = ( + completion_cost( + response, + custom_cost_per_token={ + "input_cost_per_token": kwargs["input_cost_per_token"], + "output_cost_per_token": kwargs["output_cost_per_token"], + }, + ) + * 100 + ) print(f"Stream output : {output}") print(f"Stream usage : {response.usage}") # type: ignore diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 6547a3eb6632..7c09c9780294 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -1330,6 +1330,8 @@ def test_router_fallbacks_with_custom_model_costs(): Goal: make sure custom model doesn't override default model costs. """ + default_model_info = litellm.get_model_info(model="claude-sonnet-4-5-20250929") + model_list = [ { "model_name": "claude-sonnet-4-5-20250929", @@ -1383,8 +1385,8 @@ def test_router_fallbacks_with_custom_model_costs(): print(f"key: {model_info['key']}") - assert model_info["input_cost_per_token"] == 30 - assert model_info["output_cost_per_token"] == 60 + assert model_info["input_cost_per_token"] == default_model_info["input_cost_per_token"] + assert model_info["output_cost_per_token"] == default_model_info["output_cost_per_token"] @pytest.mark.parametrize("sync_mode", [True, False])