scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 359

Merge 5d7dcc2f08 into b7da9a2755

Browse source
This commit is contained in:
ASMAE EL HARTI 2026-10-29 14:55:55 +09:00 • committed by GitHub
commit 01b95edae1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 291 additions and 78 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/__init__.py
View file

@ -246,6 +246,7 @@ anthropic_messages_estimate_initial_input_tokens: bool = False
# Or via `litellm_settings.strip_anthropic_total_tokens: true` in
# config.yaml.
strip_anthropic_total_tokens: bool = False
custom_pricing: bool = False
anthropic_sse_ping_interval_seconds: float = 15.0
sse_keepalive_ping_interval_seconds: float | None = None
route_all_chat_openai_to_responses: bool = (

62
litellm/cost_calculator.py
View file

@ -444,6 +444,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.
@ -650,6 +651,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
@ -658,6 +660,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(
@ -762,6 +765,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)
@ -773,6 +777,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)
@ -801,14 +806,17 @@ 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
@ -1550,6 +1558,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 isinstance(completion_response, DecisionsResponse) and completion_response.usage is None:
raise ValueError("Cannot calculate Decisions cost: provider did not report usage")
@ -1864,19 +1874,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 +2060,27 @@ 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
and (
deployment_model_info.get("input_cost_per_query") is not None
or deployment_model_info.get("cost_per_query") is not None
)
)
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)
@ -2272,6 +2305,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
@ -2670,6 +2705,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
@ -2685,10 +2721,14 @@ 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 or (
model_info.get("input_cost_per_query") is None and model_info.get("cost_per_query") is None
):
try:
builtin_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
model_info = {**(model_info or {}), **builtin_info} if model_info else builtin_info
except Exception as e:
verbose_logger.debug("Failed to get model info in rerank_cost: %s", e)
return config.calculate_rerank_cost(
model=model,

24
litellm/litellm_core_utils/litellm_logging.py
View file

@ -2806,10 +2806,14 @@ 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.
@ -2817,6 +2821,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
@ -6374,7 +6381,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
@ -6382,9 +6389,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:

5
litellm/llms/azure/cost_calculation.py
View file

@ -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,
)

3
litellm/llms/openai/cost_calculation.py
View file

@ -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,
)

2
litellm/main.py
View file

@ -5772,6 +5772,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
@ -6791,6 +6792,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

18
litellm/rerank_api/main.py
View file

@ -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),
}

8
litellm/types/rerank.py
View file

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

16
litellm/types/utils.py
View file

@ -2192,6 +2192,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:
@ -2433,6 +2441,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

230
tests/unit/test_register_model_custom_pricing.py
View file

@ -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
@ -236,10 +236,7 @@ def test_register_model_strips_none_litellm_provider():
registered = litellm.model_cost.get(model_key)
assert registered is not None, f"{model_key} should be in model_cost"
# The key may be absent entirely, but if present it must not be None.
assert (
"litellm_provider" not in registered
or registered["litellm_provider"] is not None
)
assert "litellm_provider" not in registered or registered["litellm_provider"] is not None
# Downstream consumers must accept this entry for any provider,
# mirroring what the cost calculator does.
assert _check_provider_match(registered, "openai") is True
@ -297,9 +294,7 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp
# The strip must have removed the None-valued provider that
# ``get_model_info`` returned. The key may be absent entirely, but
# it must never be present with value ``None``.
assert "litellm_provider" not in registered or (
registered["litellm_provider"] is not None
), (
assert "litellm_provider" not in registered or (registered["litellm_provider"] is not None), (
"register_model failed to strip litellm_provider=None returned "
f"by get_model_info, got {registered.get('litellm_provider')!r}"
)
@ -360,14 +355,8 @@ def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(monkeypa
)
registered = litellm.model_cost[registered_key]
assert (
registered.get("cache_creation_input_token_cost")
== builtin["cache_creation_input_token_cost"]
)
assert (
registered.get("cache_read_input_token_cost")
== builtin["cache_read_input_token_cost"]
)
assert registered.get("cache_creation_input_token_cost") == builtin["cache_creation_input_token_cost"]
assert registered.get("cache_read_input_token_cost") == builtin["cache_read_input_token_cost"]
assert registered["litellm_provider"] == "bedrock"
usage = Usage(
@ -431,8 +420,7 @@ def test_register_model_warns_when_no_builtin_match_for_cache_pricing(caplog):
)
assert any(
registered_key in record.message
and "cache_creation_input_token_cost" in record.message
registered_key in record.message and "cache_creation_input_token_cost" in record.message
for record in caplog.records
), "expected a warning naming the unmapped key and the cache cost fields"
finally:
@ -619,17 +607,14 @@ def test_register_model_router_add_deployment_custom_pricing_applies():
# the deployment's provider.
from litellm.utils import _check_provider_match
registered_keys = [
k for k in (deployment_model, model_key) if k in litellm.model_cost
]
registered_keys = [k for k in (deployment_model, model_key) if k in litellm.model_cost]
assert registered_keys, (
"Router.add_deployment did not register custom pricing for "
f"{model_key} / {deployment_model}"
f"Router.add_deployment did not register custom pricing for {model_key} / {deployment_model}"
)
for k in registered_keys:
assert (
_check_provider_match(litellm.model_cost[k], "openai") is True
), f"custom pricing for {k} was dropped by _check_provider_match"
assert _check_provider_match(litellm.model_cost[k], "openai") is True, (
f"custom pricing for {k} was dropped by _check_provider_match"
)
finally:
litellm.model_cost.pop(model_key, None)
litellm.model_cost.pop(deployment_model, None)
@ -646,12 +631,8 @@ def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing():
"""
shared_key = "openai/text-embedding-3-small"
deployment_id = "lit3991-wildcard-embed-zero"
snapshot = _snapshot_model_cost_entries(
[shared_key, "text-embedding-3-small", deployment_id]
)
builtin_input_cost = litellm.get_model_info(model=shared_key)[
"input_cost_per_token"
]
snapshot = _snapshot_model_cost_entries([shared_key, "text-embedding-3-small", deployment_id])
builtin_input_cost = litellm.get_model_info(model=shared_key)["input_cost_per_token"]
assert builtin_input_cost > 0
try:
@ -666,10 +647,9 @@ def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing():
mock_response=[0.1, 0.2],
)
assert (
litellm.get_model_info(model=shared_key)["input_cost_per_token"]
== builtin_input_cost
), "wildcard deployment's zero pricing leaked into the shared model_cost key"
assert litellm.get_model_info(model=shared_key)["input_cost_per_token"] == builtin_input_cost, (
"wildcard deployment's zero pricing leaked into the shared model_cost key"
)
assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.0
assert litellm.model_cost[deployment_id]["output_cost_per_token"] == 0.0
@ -679,9 +659,7 @@ def test_embedding_router_zero_pricing_does_not_clobber_builtin_pricing():
api_key="fake-key",
mock_response=[0.1, 0.2],
)
sibling_cost = litellm.completion_cost(
completion_response=sibling_response, call_type="embedding"
)
sibling_cost = litellm.completion_cost(completion_response=sibling_response, call_type="embedding")
assert sibling_cost == pytest.approx(10 * builtin_input_cost)
finally:
_restore_model_cost_entries(snapshot)
@ -695,12 +673,8 @@ def test_embedding_router_custom_pricing_costs_request_via_deployment_id():
shared_key = "openai/text-embedding-3-small"
deployment_id = "lit3991-wildcard-embed-custom"
override_input_cost = 5e-05
snapshot = _snapshot_model_cost_entries(
[shared_key, "text-embedding-3-small", deployment_id]
)
builtin_input_cost = litellm.get_model_info(model=shared_key)[
"input_cost_per_token"
]
snapshot = _snapshot_model_cost_entries([shared_key, "text-embedding-3-small", deployment_id])
builtin_input_cost = litellm.get_model_info(model=shared_key)["input_cost_per_token"]
assert builtin_input_cost != override_input_cost
try:
@ -724,10 +698,7 @@ def test_embedding_router_custom_pricing_costs_request_via_deployment_id():
router_model_id=deployment_id,
)
assert request_cost == pytest.approx(10 * override_input_cost)
assert (
litellm.get_model_info(model=shared_key)["input_cost_per_token"]
== builtin_input_cost
)
assert litellm.get_model_info(model=shared_key)["input_cost_per_token"] == builtin_input_cost
finally:
_restore_model_cost_entries(snapshot)
@ -736,12 +707,8 @@ def test_completion_router_zero_pricing_does_not_clobber_builtin_pricing():
"""Same isolation as the embedding path, exercised through completion()."""
shared_key = "openai/gpt-4o-mini"
deployment_id = "lit3991-wildcard-chat-zero"
snapshot = _snapshot_model_cost_entries(
[shared_key, "gpt-4o-mini", deployment_id]
)
builtin_input_cost = litellm.get_model_info(model=shared_key)[
"input_cost_per_token"
]
snapshot = _snapshot_model_cost_entries([shared_key, "gpt-4o-mini", deployment_id])
builtin_input_cost = litellm.get_model_info(model=shared_key)["input_cost_per_token"]
assert builtin_input_cost > 0
try:
@ -756,10 +723,9 @@ def test_completion_router_zero_pricing_does_not_clobber_builtin_pricing():
mock_response="hello back",
)
assert (
litellm.get_model_info(model=shared_key)["input_cost_per_token"]
== builtin_input_cost
), "wildcard deployment's zero pricing leaked into the shared model_cost key"
assert litellm.get_model_info(model=shared_key)["input_cost_per_token"] == builtin_input_cost, (
"wildcard deployment's zero pricing leaked into the shared model_cost key"
)
assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.0
finally:
_restore_model_cost_entries(snapshot)
@ -782,10 +748,7 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key():
mock_response=[0.1, 0.2],
)
assert (
litellm.model_cost[model_key]["input_cost_per_token"]
== override_input_cost
)
assert litellm.model_cost[model_key]["input_cost_per_token"] == override_input_cost
cost = litellm.completion_cost(
completion_response=response,
model=model_key,
@ -1704,3 +1667,144 @@ 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()