From 35f5395c9daf8cec20c4904da34861ae8533c05f Mon Sep 17 00:00:00 2001 From: Charlie Freeman <243496084+cf3901646@users.noreply.github.com> Date: Wed, 21 Oct 2026 18:15:29 +0800 Subject: [PATCH 1/2] fix(cost_calculator): support custom_pricing on rerank and embedding cost paths - Honor global litellm.custom_pricing flag in use_custom_pricing_for_model - Accurately default search_units to 1 only when None, preventing 0-unit false billing - Streamline rerank_cost error handling and fallback to model_info pricing for unmapped providers - Add unit tests for rerank custom pricing, global flag detection, and zero-search-unit defense Closes #47798 Signed-off-by: Charlie Freeman <243496084+cf3901646@users.noreply.github.com> --- litellm/cost_calculator.py | 51 +++++++++++------- litellm/litellm_core_utils/litellm_logging.py | 2 + .../llms/base_llm/rerank/transformation.py | 7 +-- .../test_register_model_custom_pricing.py | 54 +++++++++++++++++++ 4 files changed, 89 insertions(+), 25 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 550bbfcde98..81ff8972033 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1871,12 +1871,12 @@ def completion_cost( else: billed_units = {} + raw_units = billed_units.get("search_units") + search_units = 1 if raw_units is None else raw_units # cohere charges per request by default. rerank_billed_units = RerankBilledUnits( - search_units=billed_units.get("search_units"), + search_units=search_units, total_tokens=billed_units.get("total_tokens"), ) - - search_units = billed_units.get("search_units") or 1 # cohere charges per request by default. completion_tokens = search_units elif call_type in _SEARCH_CALL_TYPES: from litellm.search import search_provider_cost_per_query @@ -2669,31 +2669,42 @@ def rerank_cost( ) -> tuple[float, float]: """ Returns - - float or None: cost of response OR none if error. + - tuple[float, float]: cost of response (prompt_cost, completion_cost) """ _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) try: - config: Final = ProviderConfigManager.get_provider_rerank_config( - model=model, - api_base=None, - present_version_params=[], - provider=LlmProviders(custom_llm_provider), - ) + model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + model_info = None + if custom_llm_provider and custom_llm_provider in LlmProvidersSet: try: - model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + config: Final = ProviderConfigManager.get_provider_rerank_config( + model=model, + api_base=None, + present_version_params=[], + provider=LlmProviders(custom_llm_provider), + ) + return config.calculate_rerank_cost( + model=model, + custom_llm_provider=custom_llm_provider, + billed_units=billed_units, + model_info=model_info, + ) except Exception: - model_info = None + pass - return config.calculate_rerank_cost( - model=model, - custom_llm_provider=custom_llm_provider, - billed_units=billed_units, - model_info=model_info, - ) - except Exception as e: - raise e + if ( + model_info is not None + and "input_cost_per_query" in model_info + and model_info["input_cost_per_query"] is not None + ): + raw_units = billed_units.get("search_units") if billed_units is not None else None + search_units: Final = 1 if raw_units is None else raw_units + return model_info["input_cost_per_query"] * search_units, 0.0 + + return 0.0, 0.0 def transcription_cost(model: str, custom_llm_provider: str | None, duration: float) -> tuple[float, float]: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d75d35e4531..7ece6328caf 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6313,6 +6313,8 @@ def use_custom_pricing_for_model(litellm_params: dict | None) -> bool: any custom pricing or deployment-scoped pricing field (such as ``off_peak_pricing``) is present in the metadata ``model_info`` """ + if getattr(litellm, "custom_pricing", False) is True: + return True if litellm_params is None: return False diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 5d2f92b5e82..a891b34b93d 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -122,14 +122,11 @@ class BaseRerankConfig(ABC): model_info is None or "input_cost_per_query" not in model_info or model_info["input_cost_per_query"] is None - or billed_units is None ): return 0.0, 0.0 - search_units: Final = billed_units.get("search_units") - - if search_units is None: - return 0.0, 0.0 + raw_units: Final = billed_units.get("search_units") if billed_units is not None else None + search_units: Final = 1 if raw_units is None else raw_units prompt_cost: Final = model_info["input_cost_per_query"] * search_units diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 35fbd28eae7..82a1bfd7b5c 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -1704,3 +1704,57 @@ def test_per_request_custom_pricing_scopes_a_colliding_deployment_id_to_its_prov finally: _restore_model_cost_entries(model_cost_entries) litellm.open_ai_chat_completion_models.intersection_update(openai_models_before) + + +def test_rerank_custom_pricing_registered_model(): + """Verify that a registered rerank model calculates cost with custom_pricing enabled (closes #47798).""" + model_name = "test-custom-rerank-model" + original_entries = _snapshot_model_cost_entries([model_name]) + original_custom_pricing = getattr(litellm, "custom_pricing", None) + try: + litellm.custom_pricing = True + litellm.register_model( + { + model_name: { + "input_cost_per_query": 0.001, + "litellm_provider": "openai", + "mode": "rerank", + } + } + ) + # 1. 默认无 billed_units 时，按单次请求默认 1 个 search unit 计费 + cost, comp_cost = litellm.cost_calculator.rerank_cost( + model=model_name, + custom_llm_provider="openai", + ) + assert cost == pytest.approx(0.001) + assert comp_cost == 0.0 + + # 2. 关键防御：真实 0 单元（search_units=0）绝不虚假计费，必须为 $0.0！ + zero_cost, _ = litellm.cost_calculator.rerank_cost( + model=model_name, + custom_llm_provider="openai", + billed_units={"search_units": 0}, + ) + assert zero_cost == 0.0 + finally: + _restore_model_cost_entries(original_entries) + if original_custom_pricing is None: + delattr(litellm, "custom_pricing") + else: + litellm.custom_pricing = original_custom_pricing + + +def test_custom_pricing_global_flag_detection(): + """Verify that use_custom_pricing_for_model honors the global litellm.custom_pricing flag (closes #47798).""" + from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model + original_custom_pricing = getattr(litellm, "custom_pricing", None) + try: + litellm.custom_pricing = True + assert use_custom_pricing_for_model(None) is True + assert use_custom_pricing_for_model({}) is True + finally: + if original_custom_pricing is None: + delattr(litellm, "custom_pricing") + else: + litellm.custom_pricing = original_custom_pricing From 0261269e9a0094c2c1d65c920bbf4170e20d03ab Mon Sep 17 00:00:00 2001 From: Charlie Freeman <243496084+cf3901646@users.noreply.github.com> Date: Thu, 22 Oct 2026 14:09:38 +0800 Subject: [PATCH 2/2] style(rerank): format transformation.py with ruff Signed-off-by: Charlie Freeman <243496084+cf3901646@users.noreply.github.com> --- litellm/llms/base_llm/rerank/transformation.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index a891b34b93d..415ee224435 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -118,11 +118,7 @@ class BaseRerankConfig(ABC): Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ - if ( - model_info is None - or "input_cost_per_query" not in model_info - or model_info["input_cost_per_query"] is None - ): + if model_info is None or "input_cost_per_query" not in model_info or model_info["input_cost_per_query"] is None: return 0.0, 0.0 raw_units: Final = billed_units.get("search_units") if billed_units is not None else None