fix(cost_calculator): honor custom_pricing on rerank and embedding cost paths

litellm/__init__.py

@@ -245,6 +245,7 @@ use_chat_completions_url_for_anthropic_messages: bool = bool(
 # Or via `litellm_settings.strip_anthropic_total_tokens: true` in
 # config.yaml.
 strip_anthropic_total_tokens: bool = False
+custom_pricing: bool = bool(os.getenv("LITELLM_CUSTOM_PRICING", False))
 anthropic_sse_ping_interval_seconds: float = 15.0
 sse_keepalive_ping_interval_seconds: float | None = None
 route_all_chat_openai_to_responses: bool = (

litellm/cost_calculator.py

@@ -443,6 +443,7 @@ def cost_per_token(
     request_model: str | None = None,  # original request model for router detection
     custom_model_info: OCRPricing | None = None,
     bedrock_model_info: ModelInfo | None = None,
+    model_info: ModelInfo | None = None,
 ) -> tuple[float, float]:
     """
     Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@@ -649,6 +650,7 @@ def cost_per_token(
                 custom_llm_provider=custom_llm_provider,
                 service_tier=service_tier,
                 data_residency=data_residency,
+                model_info=model_info,
             )
 
         return prompt_cost, completion_cost
@@ -657,6 +659,7 @@ def cost_per_token(
             model=model,
             custom_llm_provider=custom_llm_provider,
             billed_units=rerank_billed_units,
+            model_info=model_info,
         )
     elif call_type == "avector_store_search" or call_type == "vector_store_search":
         return vector_store_search_cost(
@@ -761,6 +764,7 @@ def cost_per_token(
             usage=usage_block,
             service_tier=service_tier,
             data_residency=data_residency,
+            model_info=model_info,
         )
     elif custom_llm_provider == "databricks":
         return databricks_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
@@ -772,6 +776,7 @@ def cost_per_token(
             usage=usage_block,
             response_time_ms=response_time_ms,
             service_tier=service_tier,
+            model_info=model_info,
         )
     elif custom_llm_provider == "gemini":
         return gemini_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
@@ -800,14 +805,15 @@ def cost_per_token(
             service_tier=service_tier,
         )
     else:
-        model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
-        if _has_token_or_tiered_pricing(model_info):
+        resolved_fallback_model_info: Final = model_info or cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
+        if _has_token_or_tiered_pricing(resolved_fallback_model_info):
             return generic_cost_per_token(
                 model=model,
                 usage=usage_block,
                 custom_llm_provider=custom_llm_provider,
                 service_tier=service_tier,
                 data_residency=data_residency,
+                model_info=resolved_fallback_model_info,
             )
         verbose_logger.debug("No per-token, tiered, or per-second pricing for model=%s; cost is 0", model)
         return 0.0, 0.0
@@ -1545,6 +1551,8 @@ def completion_cost(
     or add it to model_prices_and_context_window.json.
     """
     try:
+        if custom_pricing is None:
+            custom_pricing = bool(getattr(litellm, "custom_pricing", False))
         call_type = infer_call_type(call_type, completion_response) or "completion"
         if pricing_provider is not None:
             custom_llm_provider = pricing_provider  # rebind-ok: model_info.pricing_provider overrides the routing provider for cost-map lookups
@@ -1864,19 +1872,21 @@ def completion_cost(
                         else 0.0
                     )
                 elif call_type in _RERANK_CALL_TYPES:
-                    if completion_response is not None and isinstance(completion_response, RerankResponse):
-                        meta_obj = completion_response.meta
+                    if completion_response is not None and isinstance(completion_response, (RerankResponse, dict)):
+                        if isinstance(completion_response, dict):
+                            meta_obj = completion_response.get("meta")
+                        else:
+                            meta_obj = completion_response.meta
                         if meta_obj is not None:
                             billed_units: RerankBilledUnits = meta_obj.get("billed_units") or {}
                         else:
                             billed_units = {}
 
+                        search_units = billed_units.get("search_units") or 1  # default to 1 query per rerank request
                         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
@@ -2048,6 +2058,17 @@ def completion_cost(
                     response=completion_response,
                     custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id),
                     bedrock_model_info=bedrock_model_info,
+                    model_info=(
+                        deployment_model_info
+                        if (
+                            (deployment_model_info := _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id))
+                            and (
+                                call_type in _RERANK_CALL_TYPES
+                                or deployment_model_info.get("input_cost_per_token") is not None
+                            )
+                        )
+                        else None
+                    ),
                 )
 
                 # Get additional costs from provider (e.g., routing fees, infrastructure costs)
@@ -2270,6 +2291,8 @@ def response_cost_calculator(
     - float or None: cost of response
     """
     try:
+        if custom_pricing is None:
+            custom_pricing = bool(getattr(litellm, "custom_pricing", False))
         response_cost: float = 0.0
         if cache_hit is not None and cache_hit is True:
             response_cost = 0.0
@@ -2666,6 +2689,7 @@ def rerank_cost(
     model: str,
     custom_llm_provider: str | None,
     billed_units: RerankBilledUnits | None = None,
+    model_info: ModelInfo | None = None,
 ) -> tuple[float, float]:
     """
     Returns
@@ -2681,10 +2705,11 @@ def rerank_cost(
             provider=LlmProviders(custom_llm_provider),
         )
 
-        try:
-            model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
-        except Exception:
-            model_info = None
+        if model_info is None:
+            try:
+                model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
+            except Exception:
+                model_info = None
 
         return config.calculate_rerank_cost(
             model=model,

litellm/litellm_core_utils/litellm_logging.py

@@ -2762,10 +2762,10 @@ class Logging(LiteLLMLoggingBaseClass):
 
         if self.model_call_details.get("cache_hit") is True:
             self.model_call_details["response_cost"] = 0.0
-        elif "response_cost" in hidden_params:
+        elif "response_cost" in hidden_params and not use_custom_pricing_for_model(self):
             self.model_call_details["response_cost"] = hidden_params["response_cost"]
             self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"])
-        elif (existing_cost := self.model_call_details.get("response_cost")) is not None and existing_cost != 0:
+        elif (existing_cost := self.model_call_details.get("response_cost")) is not None and existing_cost != 0 and not use_custom_pricing_for_model(self):
             # Preserve response_cost if already calculated (e.g., by pass-through
             # handlers like Gemini/Vertex which call completion_cost directly).
             # Do not preserve 0 from failure_handler on intermediate router retries.
@@ -2773,6 +2773,9 @@ class Logging(LiteLLMLoggingBaseClass):
         else:
             self.model_call_details["response_cost"] = self.response_cost_calculator(result=logging_result)
 
+        if isinstance(hidden_params, dict) and self.model_call_details.get("response_cost") is not None:
+            hidden_params["response_cost"] = self.model_call_details["response_cost"]
+
         if not build_logging_payload:
             return None
 
@@ -6305,7 +6308,7 @@ def _deployment_model_info(litellm_params: dict | None) -> Mapping[str, object]:
     return MappingProxyType({})
 
 
-def use_custom_pricing_for_model(litellm_params: dict | None) -> bool:
+def use_custom_pricing_for_model(litellm_params: Any) -> bool:
     """
     Check if the model uses custom pricing
 
@@ -6313,9 +6316,20 @@ 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
 
+    if hasattr(litellm_params, "litellm_params"):
+        litellm_params = getattr(litellm_params, "litellm_params", None)
+        if litellm_params is None:
+            return False
+
+    if not isinstance(litellm_params, Mapping):
+        return False
+
     # Check litellm_params using set intersection (only check keys that exist in both)
     matching_keys = _CUSTOM_PRICING_KEYS & litellm_params.keys()
     for key in matching_keys:

litellm/llms/azure/cost_calculation.py

@@ -4,7 +4,7 @@ Helper util for handling azure openai-specific cost calculation
 """
 
 from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
-from litellm.types.utils import Usage
+from litellm.types.utils import ModelInfo, Usage
 
 
 def cost_per_token(
@@ -12,6 +12,7 @@ def cost_per_token(
     usage: Usage,
     response_time_ms: float | None = 0.0,
     service_tier: str | None = None,
+    model_info: ModelInfo | None = None,
 ) -> tuple[float, float]:
     """
     Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@@ -19,6 +20,7 @@ def cost_per_token(
     Input:
         - model: str, the model name without provider prefix
         - usage: LiteLLM Usage block, containing caching and audio token information
+        - model_info: optional pre-resolved ModelInfo dictionary
 
     Returns:
         Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
@@ -30,4 +32,5 @@ def cost_per_token(
         usage=usage,
         custom_llm_provider="azure",
         service_tier=service_tier,
+        model_info=model_info,
     )

litellm/llms/openai/cost_calculation.py

@@ -24,6 +24,7 @@ def cost_per_token(
     usage: Usage,
     service_tier: str | None = None,
     data_residency: str | None = None,
+    model_info: ModelInfo | None = None,
 ) -> tuple[float, float]:
     """
     Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@@ -34,6 +35,7 @@ def cost_per_token(
         - data_residency: optional OpenAI data-residency region (e.g. "eu", "us"),
           inferred from api_base. Applies the model's regional-processing
           uplift multiplier when set.
+        - model_info: optional pre-resolved ModelInfo dictionary
 
     Returns:
         Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
@@ -44,6 +46,7 @@ def cost_per_token(
         custom_llm_provider="openai",
         service_tier=service_tier,
         data_residency=data_residency,
+        model_info=model_info,
     )
 
 

litellm/main.py

@@ -5724,6 +5724,7 @@ def completion(
         ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
         if (
             (input_cost_per_token is not None and output_cost_per_token is not None)
+            or kwargs.get("input_cost_per_query") is not None
             or input_cost_per_second is not None
             or output_cost_per_second is not None
             or cost_per_second is not None
@@ -6718,6 +6719,7 @@ 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 kwargs.get("input_cost_per_query") is not None
         or input_cost_per_second is not None
         or output_cost_per_second is not None
         or cost_per_second is not None

litellm/rerank_api/main.py

@@ -182,12 +182,30 @@ def rerank(
 
         model_response: Final = RerankResponse()
 
+        ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
+        if (
+            kwargs.get("input_cost_per_query") is not None
+            or (kwargs.get("input_cost_per_token") is not None and kwargs.get("output_cost_per_token") is not None)
+            or kwargs.get("input_cost_per_second") is not None
+            or kwargs.get("output_cost_per_second") is not None
+            or kwargs.get("cost_per_second") is not None
+        ):
+            from litellm.main import _register_custom_pricing_for_request
+
+            _register_custom_pricing_for_request(
+                model=model,
+                custom_llm_provider=_custom_llm_provider or "cohere",
+                kwargs=kwargs,
+                model_info=model_info,
+            )
+
         rerank_litellm_params: Final = {
             "litellm_call_id": litellm_call_id,
             "proxy_server_request": proxy_server_request,
             "model_info": model_info,
             "preset_cache_key": None,
             "stream_response": {},
+            "custom_pricing": kwargs.get("custom_pricing"),
             **optional_params.model_dump(exclude_unset=True),
         }
 

litellm/types/rerank.py

@@ -95,6 +95,14 @@ class RerankResponse(LiteLLMBaseModel):
     def hidden_params(self, hidden_params: dict[str, object]) -> None:  # mutable-ok: API requires mutation
         self._hidden_params = hidden_params
 
+    @property
+    def response_cost(self) -> float | None:
+        if isinstance(self._hidden_params, dict):
+            cost = self._hidden_params.get("response_cost")
+            if isinstance(cost, (int, float)):
+                return float(cost)
+        return None
+
     def __getitem__(self, key):
         return self.__dict__[key]
 

litellm/types/utils.py

@@ -2173,6 +2173,14 @@ class ModelResponseBase(OpenAIObject):
     def hidden_params(self, hidden_params: dict[str, object]) -> None:  # mutable-ok: API requires mutation
         self._hidden_params = hidden_params
 
+    @property
+    def response_cost(self) -> float | None:
+        if isinstance(self._hidden_params, dict):
+            cost = self._hidden_params.get("response_cost")
+            if isinstance(cost, (int, float)):
+                return float(cost)
+        return None
+
     _response_headers: dict | None = None
 
     def set_provider_response_headers(self, headers: httpx.Headers) -> None:
@@ -2414,6 +2422,14 @@ class EmbeddingResponse(OpenAIObject):
     def hidden_params(self, hidden_params: dict[str, object]) -> None:  # mutable-ok: API requires mutation
         self._hidden_params = hidden_params
 
+    @property
+    def response_cost(self) -> float | None:
+        if isinstance(self._hidden_params, dict):
+            cost = self._hidden_params.get("response_cost")
+            if isinstance(cost, (int, float)):
+                return float(cost)
+        return None
+
     _response_headers: dict | None = None
     _response_ms: float | None = None
 

tests/unit/test_register_model_custom_pricing.py

@@ -16,7 +16,7 @@ from collections.abc import Callable, Mapping
 from typing import Final, NoReturn
 
 import pytest
-
+import respx
 
 import litellm
 from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases, adopt_model_cost_map
@@ -1704,3 +1704,151 @@ 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_global_and_register_model(monkeypatch: pytest.MonkeyPatch):
+    import httpx
+    import respx
+
+    monkeypatch.setattr(litellm, "custom_pricing", True)
+    custom_model = "cohere/test-custom-rerank"
+    snapshot = _snapshot_model_cost_entries([custom_model])
+    try:
+        litellm.register_model(
+            {
+                custom_model: {
+                    "input_cost_per_query": 0.005,
+                    "litellm_provider": "cohere",
+                    "mode": "rerank",
+                }
+            },
+            persist_across_reloads=False,
+        )
+
+        with respx.mock(assert_all_called=False) as respx_mock:
+            respx_mock.post("https://api.cohere.com/v2/rerank").mock(
+                return_value=httpx.Response(
+                    200,
+                    json={
+                        "id": "rerank-test-1",
+                        "results": [
+                            {"index": 0, "relevance_score": 0.99, "document": {"text": "litellm"}},
+                        ],
+                        "meta": {"billed_units": {"search_units": 1}},
+                    },
+                )
+            )
+            response = litellm.rerank(
+                model=custom_model,
+                query="what is litellm?",
+                documents=["litellm"],
+                api_key="fake-cohere-key",
+            )
+            assert response.response_cost == pytest.approx(0.005)
+            calc_cost = litellm.completion_cost(
+                completion_response=response, model=custom_model, call_type="rerank"
+            )
+            assert calc_cost == pytest.approx(0.005)
+    finally:
+        _restore_model_cost_entries(snapshot)
+
+
+def test_embedding_custom_pricing_global_and_register_model(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(litellm, "custom_pricing", True)
+    custom_model = "test-custom-embedding-unregistered"
+    snapshot = _snapshot_model_cost_entries([custom_model])
+    try:
+        litellm.register_model(
+            {
+                custom_model: {
+                    "input_cost_per_token": 0.0001,
+                    "litellm_provider": "openai",
+                    "mode": "embedding",
+                }
+            },
+            persist_across_reloads=False,
+        )
+
+        response = litellm.embedding(
+            model=custom_model,
+            input=["hello custom embedding world"],
+            api_key="fake-key",
+            mock_response=[0.1, 0.2, 0.3],
+        )
+        assert response.response_cost is not None
+        assert response.response_cost > 0
+        calc_cost = litellm.completion_cost(
+            completion_response=response, model=custom_model, call_type="embedding"
+        )
+        assert calc_cost == pytest.approx(response.response_cost)
+    finally:
+        _restore_model_cost_entries(snapshot)
+
+
+def test_rerank_per_request_custom_pricing():
+    import httpx
+    import respx
+
+    custom_model = "cohere/test-request-rerank"
+    with respx.mock(assert_all_called=False) as respx_mock:
+        respx_mock.post("https://api.cohere.com/v2/rerank").mock(
+            return_value=httpx.Response(
+                200,
+                json={
+                    "id": "rerank-req-1",
+                    "results": [
+                        {"index": 0, "relevance_score": 0.99, "document": {"text": "litellm"}},
+                    ],
+                    "meta": {"billed_units": {"search_units": 1}},
+                },
+            )
+        )
+        response = litellm.rerank(
+            model=custom_model,
+            query="test query",
+            documents=["litellm"],
+            api_key="fake-key",
+            input_cost_per_query=0.007,
+        )
+        assert response.response_cost == pytest.approx(0.007)
+
+
+@pytest.mark.asyncio
+async def test_rerank_router_deployment_custom_pricing(
+    monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
+):
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    litellm.in_memory_llm_clients_cache.flush_cache()
+    try:
+        router = litellm.Router(
+            model_list=[
+                {
+                    "model_name": "cohere-rerank",
+                    "litellm_params": {
+                        "model": "cohere/rerank-english-v3.0",
+                        "api_key": "sk-fake",
+                        "input_cost_per_query": 0.004,
+                    },
+                    "model_info": {"id": "rerank-dep-1", "input_cost_per_query": 0.004},
+                }
+            ]
+        )
+        respx_mock.post("https://api.cohere.com/v2/rerank").respond(
+            200,
+            json={
+                "id": "rerank-1",
+                "results": [{"index": 0, "relevance_score": 0.9}],
+                "meta": {"billed_units": {"search_units": 1}},
+            },
+        )
+
+        public_response = await router.arerank(
+            model="cohere-rerank",
+            query="hello",
+            documents=["hello", "world"],
+            top_n=1,
+        )
+        assert public_response.response_cost == pytest.approx(0.004)
+    finally:
+        litellm.in_memory_llm_clients_cache.flush_cache()
+