From 88cae9eb2cc172d426c140a166c13181f0602e8b Mon Sep 17 00:00:00 2001 From: Genmin Date: Mon, 11 May 2026 16:56:13 -0700 Subject: [PATCH] fix(proxy): honor local litellm_proxy cost overrides --- litellm/cost_calculator.py | 54 ++++++++++++++--- tests/test_litellm/test_cost_calculator.py | 69 ++++++++++++++++++++++ 2 files changed, 115 insertions(+), 8 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9b4dd80265c..285aa2173ee 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1689,6 +1689,36 @@ def get_response_cost_from_hidden_params( return None +def _should_use_provider_response_cost( + provider_response_cost: Optional[float], + model: str, + custom_llm_provider: Optional[str], + custom_pricing: Optional[bool], + hidden_params: Union[dict, BaseModel], +) -> bool: + if provider_response_cost is None: + return False + + if custom_pricing is not True: + return True + + if isinstance(hidden_params, BaseModel): + hidden_params_dict = cast(BaseModel, hidden_params).model_dump() + else: + hidden_params_dict = hidden_params + + hidden_custom_llm_provider = hidden_params_dict.get("custom_llm_provider") + is_litellm_proxy_request = ( + custom_llm_provider == "litellm_proxy" + or hidden_custom_llm_provider == "litellm_proxy" + or model.startswith("litellm_proxy/") + ) + if is_litellm_proxy_request: + return False + + return True + + def response_cost_calculator( response_object: Union[ ModelResponse, @@ -1747,14 +1777,22 @@ def response_cost_calculator( if cache_hit is not None and cache_hit is True: response_cost = 0.0 else: - if isinstance(response_object, BaseModel): - if hasattr(response_object, "_hidden_params"): - response_object._hidden_params["optional_params"] = optional_params - provider_response_cost = get_response_cost_from_hidden_params( - response_object._hidden_params - ) - if provider_response_cost is not None: - return provider_response_cost + if hasattr(response_object, "_hidden_params"): + hidden_params = response_object._hidden_params + if isinstance(hidden_params, dict): + hidden_params["optional_params"] = optional_params + provider_response_cost = get_response_cost_from_hidden_params( + hidden_params + ) + if _should_use_provider_response_cost( + provider_response_cost=provider_response_cost, + model=model, + custom_llm_provider=custom_llm_provider, + custom_pricing=custom_pricing, + hidden_params=hidden_params, + ): + assert provider_response_cost is not None + return provider_response_cost response_cost = completion_cost( completion_response=response_object, diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index e53484dd287..2598dbe1a81 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -67,6 +67,75 @@ def test_cost_calculator_with_response_cost_in_additional_headers(): assert result == 1000 +def test_litellm_proxy_custom_pricing_overrides_response_cost_header(): + model_id = "proxy-zero-cost-deployment" + original_model_cost = litellm.model_cost.get(model_id) + litellm.register_model( + model_cost={ + model_id: { + "litellm_provider": "litellm_proxy", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + } + ) + + try: + response = ModelResponse( + model="hosted_vllm/glm-4.7-fp8", + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + response._hidden_params = { + "additional_headers": {"llm_provider-x-litellm-response-cost": "0.123"}, + "custom_llm_provider": "litellm_proxy", + } + + result = response_cost_calculator( + response_object=response, + model="glm-4.7", + custom_llm_provider="litellm_proxy", + call_type="completion", + optional_params={}, + cache_hit=None, + base_model=None, + custom_pricing=True, + router_model_id=model_id, + ) + + assert result == 0.0 + finally: + if original_model_cost is None: + litellm.model_cost.pop(model_id, None) + else: + litellm.model_cost[model_id] = original_model_cost + + +def test_litellm_proxy_uses_response_cost_header_without_custom_pricing(): + response = ModelResponse( + model="hosted_vllm/glm-4.7-fp8", + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + response._hidden_params = { + "additional_headers": {"llm_provider-x-litellm-response-cost": "0.123"}, + "custom_llm_provider": "litellm_proxy", + } + + result = response_cost_calculator( + response_object=response, + model="glm-4.7", + custom_llm_provider="litellm_proxy", + call_type="completion", + optional_params={}, + cache_hit=None, + base_model=None, + custom_pricing=False, + ) + + assert result == 0.123 + + def test_baseten_model_api_pricing_entries(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="")