diff --git a/litellm/__init__.py b/litellm/__init__.py index 1784deeb797..6adc8b45f1d 100644 --- a/litellm/__init__.py +++ b/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 = ( diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 550bbfcde98..d32c44d9bad 100644 --- a/litellm/cost_calculator.py +++ b/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, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d75d35e4531..039d2c37dd4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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: diff --git a/litellm/llms/azure/cost_calculation.py b/litellm/llms/azure/cost_calculation.py index 057e9dbb9d9..ea7cabb0d41 100644 --- a/litellm/llms/azure/cost_calculation.py +++ b/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, ) diff --git a/litellm/llms/openai/cost_calculation.py b/litellm/llms/openai/cost_calculation.py index 36d7e4d95de..189e85f947f 100644 --- a/litellm/llms/openai/cost_calculation.py +++ b/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, ) diff --git a/litellm/main.py b/litellm/main.py index 27700cf9def..4c76f3511bd 100644 --- a/litellm/main.py +++ b/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 diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index f7e2616324f..baea20edf3b 100644 --- a/litellm/rerank_api/main.py +++ b/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), } diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index 6d1ab85012b..439a16209e5 100644 --- a/litellm/types/rerank.py +++ b/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] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5939a350970..4cd1ba64d4f 100644 --- a/litellm/types/utils.py +++ b/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 diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 35fbd28eae7..a84c2d52e3f 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/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() +