scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

fix(exa): count only prices the operator set, so mirrors and catalog-equal registrations keep Exa's cost

Browse source 
A cost map served from the operator's own URL now prices Exa only where its
price fields differ from the map LiteLLM ships, so a mirror of LiteLLM's map
keeps billing Exa's reported total. A price register_model sets counts even at
the catalog's own value, a registration without a price field sets nothing,
and a flat price the operator set without tiers replaces the shipped tiers.
This commit is contained in:
mateo-berri 2026-10-28 13:54:14 -07:00
parent 96a2b14005
commit ffe8a669b0
2 changed files with 123 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

82
litellm/search/cost_calculator.py
View file

@ -2,6 +2,7 @@
Cost calculation for search providers.
"""
import functools
import math
from collections.abc import Mapping
from types import MappingProxyType
@ -11,18 +12,24 @@ from pydantic import TypeAdapter, ValidationError
import litellm
from litellm.constants import DEFAULT_MODEL_COST_MAP_URL
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap, get_model_cost_map_source_info
from litellm.litellm_core_utils.get_model_cost_map import (
AliasedCostEntry,
GetModelCostMap,
get_model_cost_map_source_info,
)
from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.utils import get_model_info
PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
tuple[Mapping[str, object], ...]
)
PROVIDER_COST_ADAPTER: Final = TypeAdapter(Mapping[str, object])
MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object])
EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
SEARCH_MODE_PARAM_BY_PROVIDER: Final[Mapping[str, str]] = MappingProxyType(
{"tavily": "search_depth", "linkup": "depth", "exa_ai": "type"}
)
SEARCH_PRICE_FIELDS: Final = ("input_cost_per_query", "tiered_pricing")
FLAT_PRICE_ONLY: Final = frozenset({"input_cost_per_query"})
def _provider_usage(
@ -38,31 +45,72 @@ def _provider_usage(
def _field(value: object, key: str) -> object:
return PROVIDER_COST_ADAPTER.validate_python(value).get(key) if isinstance(value, Mapping) else None
return MAPPING_ADAPTER.validate_python(value).get(key) if isinstance(value, Mapping) else None
def _live_entry(entry: str) -> object:
return litellm.model_cost.get(entry) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # registry is dynamically typed
def _changed_price_fields(value: object, reference: object) -> frozenset[str]:
"""Price fields `value` holds a price in that differ from `reference`, so a price an entry leaves out sets
nothing."""
return frozenset(
name
for name in SEARCH_PRICE_FIELDS
if _field(value, name) is not None and _field(value, name) != _field(reference, name)
)
def _runtime_price_fields(entry: str) -> frozenset[str]:
"""Price fields `register_model` set on `entry`, even to the catalog's own value, or that differ from the catalog
as loaded."""
live: Final = _live_entry(entry)
registered: Final[frozenset[str]] = (
live.registered_fields.intersection(SEARCH_PRICE_FIELDS) if isinstance(live, AliasedCostEntry) else frozenset()
)
return registered | _changed_price_fields(live, GetModelCostMap.loaded_model_cost_map().get(entry))
@functools.cache
def _shipped_search_prices() -> Mapping[str, Mapping[str, object]]:
shipped: Final = MAPPING_ADAPTER.validate_json(GetModelCostMap.read_local_model_cost_map_bytes())
return MappingProxyType(
{
key: MappingProxyType({name: _field(entry, name) for name in SEARCH_PRICE_FIELDS})
for key, entry in shipped.items()
if _field(entry, "mode") == "search"
}
)
def _own_cost_map_price_fields(entry: str) -> frozenset[str]:
"""Price fields a cost map served from the operator's own URL sets differently from the map LiteLLM ships, so a
mirror of LiteLLM's map prices nothing."""
source: Final = get_model_cost_map_source_info()
if source["source"] != "remote" or source["url"] == DEFAULT_MODEL_COST_MAP_URL:
return frozenset()
return _changed_price_fields(
GetModelCostMap.loaded_model_cost_map().get(entry), _shipped_search_prices().get(entry)
)
def _operator_price_fields(entry: str) -> frozenset[str]:
return _runtime_price_fields(entry) | _own_cost_map_price_fields(entry)
def _registered_at_runtime(entry: str) -> bool:
"""Whether `entry`'s live price differs from the catalog as loaded, i.e. `register_model` changed it."""
return _live_entry(entry) != GetModelCostMap.loaded_model_cost_map().get(entry)
return bool(_runtime_price_fields(entry))
def _operator_priced(entry: str) -> bool:
"""Whether an operator set `entry`'s price, with `register_model` or a cost map served from their own URL."""
source: Final = get_model_cost_map_source_info()
own_cost_map: Final = source["source"] == "remote" and source["url"] != DEFAULT_MODEL_COST_MAP_URL
return own_cost_map or _registered_at_runtime(entry)
"""Whether an operator set a price on `entry`, with `register_model` or a cost map served from their own URL."""
return bool(_operator_price_fields(entry))
def _registered_flat_price_replaces_tiers(entry: str) -> bool:
"""Whether `register_model` changed `entry`'s flat price but left the catalog's tiers it merged into in place."""
live: Final = _live_entry(entry)
catalog: Final = GetModelCostMap.loaded_model_cost_map().get(entry, EMPTY_OPTIONAL_PARAMS)
same_tiers: Final = _field(live, "tiered_pricing") == catalog.get("tiered_pricing")
return same_tiers and _field(live, "input_cost_per_query") != catalog.get("input_cost_per_query")
def _operator_flat_price_replaces_tiers(entry: str) -> bool:
"""Whether the operator set `entry`'s flat price but not its tiers, so the tiers LiteLLM ships stop applying."""
return _operator_price_fields(entry) == FLAT_PRICE_ONLY
def _exa_reported_total(response: object) -> float | None:
@ -107,7 +155,7 @@ def search_provider_cost_per_query(
Supports tiered pricing based on max_results parameter, and per-mode price entries selected by the
provider's mode parameter (Tavily `search_depth`, Linkup `depth`, Exa `type`). Exa bills one request at the
`costDollars.total` it reports, else at its list price, unless an operator priced the entry. A flat price
registered over an entry's catalog tiers replaces them.
an operator set over the tiers LiteLLM ships replaces them.
Args:
model: Model name (e.g., "exa_ai/search", "tavily/search")
@ -155,7 +203,7 @@ def search_provider_cost_per_query(
if (
tiered_pricing
and isinstance(tiered_pricing, list)
and not _registered_flat_price_replaces_tiers(model_info["key"])
and not _operator_flat_price_replaces_tiers(model_info["key"])
):
max_results: Final = (optional_params or {}).get("max_results", 10) # default 10 results
cost_per_query = 0.0

66
tests/unit/search/test_main.py
View file

@ -321,9 +321,22 @@ def test_a_price_registered_for_exa_wins_over_the_cost_exa_reports(
assert response._hidden_params["response_cost"] == pytest.approx(expected)
@pytest.mark.usefixtures("_loaded_catalog_prices")
def test_registering_no_price_on_an_exa_entry_keeps_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None:
respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}})
litellm.register_model(
{"exa_ai/search": {"metadata": {"notes": "operator note"}, "litellm_provider": "exa_ai", "mode": "search"}}
)
response: Final = litellm.search(query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key")
assert response._hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_TOTAL)
OPERATOR_COST_MAP_URL: Final = "https://costs.invalid/model_prices_and_context_window.json"
OPERATOR_FLAT_PRICE: Final = 0.003
OPERATOR_TIER_PRICE: Final = 0.0045
ServedExaEntry = Callable[[Mapping[str, object]], Mapping[str, object]]
@dataclass(frozen=True, slots=True)
@ -335,10 +348,11 @@ class _ServedCostMap:
@pytest.fixture
def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, Mapping[str, object]], None]]:
def serve(url: str, exa_entry: Mapping[str, object]) -> None:
def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, ServedExaEntry], None]]:
def serve(url: str, exa_entry: ServedExaEntry) -> None:
monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP")
body: Final = json.dumps({**GetModelCostMap.load_local_model_cost_map(), "exa_ai/search": exa_entry}).encode()
shipped: Final = GetModelCostMap.load_local_model_cost_map()
body: Final = json.dumps({**shipped, "exa_ai/search": exa_entry(shipped["exa_ai/search"])}).encode()
monkeypatch.setattr(
litellm, "model_cost", litellm.get_model_cost_map(url=url, max_attempts=1, client=_ServedCostMap(body))
)
@ -355,13 +369,19 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str,
[
pytest.param(
OPERATOR_COST_MAP_URL,
{"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"},
lambda _: {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"},
OPERATOR_FLAT_PRICE,
id="own-map-flat-price",
),
pytest.param(
OPERATOR_COST_MAP_URL,
{
lambda shipped: {**shipped, "input_cost_per_query": OPERATOR_FLAT_PRICE},
OPERATOR_FLAT_PRICE,
id="own-map-flat-price-over-shipped-tiers",
),
pytest.param(
OPERATOR_COST_MAP_URL,
lambda _: {
"tiered_pricing": [{"input_cost_per_query": OPERATOR_TIER_PRICE, "max_results_range": [0, 100]}],
"litellm_provider": "exa_ai",
"mode": "search",
@ -369,9 +389,16 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str,
OPERATOR_TIER_PRICE,
id="own-map-tiers",
),
pytest.param(OPERATOR_COST_MAP_URL, lambda shipped: shipped, EXA_REPORTED_TOTAL, id="own-map-mirrors-litellm"),
pytest.param(
OPERATOR_COST_MAP_URL,
lambda shipped: {name: value for name, value in shipped.items() if name != "input_cost_per_query"},
EXA_REPORTED_TOTAL,
id="own-map-from-a-release-without-the-flat-price",
),
pytest.param(
DEFAULT_MODEL_COST_MAP_URL,
{"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"},
lambda _: {"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"},
EXA_REPORTED_TOTAL,
id="published-map",
),
@ -379,9 +406,9 @@ def _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str,
)
def test_an_exa_price_from_the_operators_own_cost_map_wins_over_the_cost_exa_reports(
respx_mock: respx.MockRouter,
_remote_cost_map: Callable[[str, Mapping[str, object]], None],
_remote_cost_map: Callable[[str, ServedExaEntry], None],
url: str,
exa_entry: Mapping[str, object],
exa_entry: ServedExaEntry,
expected: float,
) -> None:
_remote_cost_map(url, exa_entry)
@ -433,3 +460,26 @@ def test_a_price_registered_over_catalog_tiers_bills_every_tiered_search_entry(
}
assert costs == {key: pytest.approx((expected, 0.0)) for key in tiered}
@pytest.mark.usefixtures("_loaded_catalog_prices")
def test_registering_the_catalogs_own_flat_price_over_its_tiers_still_bills_the_flat_price() -> None:
catalog: Final = GetModelCostMap.loaded_model_cost_map()
flat_and_tiered: Final = tuple(
key
for key, entry in catalog.items()
if entry.get("mode") == "search" and entry.get("tiered_pricing") and entry.get("input_cost_per_query")
)
assert flat_and_tiered
litellm.register_model(
{key: {"input_cost_per_query": catalog[key]["input_cost_per_query"]} for key in flat_and_tiered}
)
costs: Final = {
key: search_provider_cost_per_query(
model=key, custom_llm_provider=str(catalog[key]["litellm_provider"]), optional_params={"max_results": 50}
)
for key in flat_and_tiered
}
assert costs == {key: pytest.approx((catalog[key]["input_cost_per_query"], 0.0)) for key in flat_and_tiered}