diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 527981a1dc6..d84327c78c2 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2091,12 +2091,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 @@ -2946,31 +2946,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 97c019a2180..8b8dcde6f04 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6835,6 +6835,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 7f246bf33da..3dada159270 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -118,18 +118,11 @@ 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 - or billed_units 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 - 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