Merge 0261269e9a into 1cc7da6274

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]:

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
 

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
 

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