mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
fix(exa): unify receipt and operator pricing precedence
Keep receipts usable after serialization and without catalog entries. Preserve request modes for direct costing, apply generic custom-map prices across modes, and use the same result count for native and unified requests. The latest Exa price registration replaces the previous price kind. Adapted the public receipt lookup and canonical catalog URL from96a2b14005andffe8a669b0. Keep pricing detection local to Exa and compare catalog entries individually rather than fingerprinting the entire installed catalog. Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
shrey kharbanda
parent
e9727d2738
commit
ebee797b71
12 changed files with 298 additions and 56 deletions
|
|
@ -41,6 +41,7 @@ from litellm._logging import (
|
|||
)
|
||||
import re
|
||||
from litellm.constants import (
|
||||
DEFAULT_MODEL_COST_MAP_URL,
|
||||
DEFAULT_BATCH_SIZE,
|
||||
DEFAULT_FLUSH_INTERVAL_SECONDS,
|
||||
ROUTER_MAX_FALLBACKS,
|
||||
|
|
@ -410,7 +411,7 @@ aclient_session: Optional[Union[httpx.AsyncClient, DefaultAsyncHttpxClient]] = N
|
|||
model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks'
|
||||
model_cost_map_url: str = os.getenv(
|
||||
"LITELLM_MODEL_COST_MAP_URL",
|
||||
"https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json",
|
||||
DEFAULT_MODEL_COST_MAP_URL,
|
||||
)
|
||||
blog_posts_url: str = os.getenv(
|
||||
"LITELLM_BLOG_POSTS_URL",
|
||||
|
|
|
|||
|
|
@ -1738,6 +1738,9 @@ ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS: Final = {
|
|||
PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-"
|
||||
TYPESAFE_DANGER_ZONE_HEADER_PREFIX: Final = "typesafe-danger-zone-"
|
||||
INTERNAL_KWARG_PREFIX: Final = "_litellm_"
|
||||
DEFAULT_MODEL_COST_MAP_URL: Final = (
|
||||
"https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
|
||||
)
|
||||
CONTROL_OPTIONS_KEY: Final = f"{INTERNAL_KWARG_PREFIX}control"
|
||||
SESSION_ID_FROM_TRACE_ID_PARAM_KEY: Final = f"{INTERNAL_KWARG_PREFIX}session_id_from_trace_id"
|
||||
|
||||
|
|
|
|||
|
|
@ -709,6 +709,7 @@ def cost_per_token(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
number_of_queries=number_of_queries or 1,
|
||||
optional_params=(getattr(response, "_hidden_params", None) if response else None),
|
||||
response=response,
|
||||
)
|
||||
elif (
|
||||
per_second_cost := _per_second_pricing_cost(
|
||||
|
|
@ -1913,6 +1914,7 @@ def completion_cost(
|
|||
model=search_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
number_of_queries=number_of_queries,
|
||||
response=completion_response,
|
||||
optional_params={
|
||||
**(get_hidden_params(completion_response) or {}),
|
||||
**(optional_params or {}),
|
||||
|
|
|
|||
|
|
@ -1,14 +1,30 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.constants import INTERNAL_KWARG_PREFIX
|
||||
from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry
|
||||
from litellm.constants import DEFAULT_MODEL_COST_MAP_URL, INTERNAL_KWARG_PREFIX
|
||||
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
|
||||
|
||||
EXA_AI_COST_DOLLARS_PARAM: Final = f"{INTERNAL_KWARG_PREFIX}exa_ai_cost_dollars"
|
||||
_PRICE_FIELDS: Final = frozenset({"input_cost_per_query", "tiered_pricing"})
|
||||
_ENTRY: Final = TypeAdapter(Mapping[str, object])
|
||||
EXA_AI_CATALOG_PRICES: Final[Mapping[str, float]] = MappingProxyType(
|
||||
{
|
||||
"exa_ai/search": 0.007,
|
||||
"exa_ai/search-lite": 0.004,
|
||||
"exa_ai/search-instant": 0.004,
|
||||
"exa_ai/search-deep": 0.012,
|
||||
"exa_ai/search-deep-lite": 0.012,
|
||||
"exa_ai/search-deep-reasoning": 0.015,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _ExaCostDollars(BaseModel):
|
||||
|
|
@ -33,10 +49,15 @@ EXA_AI_LEGACY_TIERS: Final = (
|
|||
"""The tiers LiteLLM ships on `exa_ai/search`. Only older releases bill them, so they never change."""
|
||||
|
||||
|
||||
def exa_ai_reported_cost(params: Mapping[str, object]) -> float | None:
|
||||
def exa_ai_reported_cost(params: Mapping[str, object], response: object = None) -> float | None:
|
||||
"""The `costDollars.total` Exa reported for the request, when it is a finite non-negative number."""
|
||||
try:
|
||||
return _ExaCostDollars.model_validate(params.get(EXA_AI_COST_DOLLARS_PARAM)).total
|
||||
payload: Final = response.model_extra if isinstance(response, SearchResponse) else response
|
||||
metadata: Final = _ENTRY.validate_python(payload or {}).get("provider_metadata")
|
||||
receipt: Final = _ENTRY.validate_python(metadata or {}).get("costDollars")
|
||||
return _ExaCostDollars.model_validate(
|
||||
receipt if receipt is not None else params.get(EXA_AI_COST_DOLLARS_PARAM)
|
||||
).total
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
|
@ -48,17 +69,29 @@ def _legacy_tiers(tiers: object) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def exa_ai_price_set(entry: object) -> bool:
|
||||
"""Whether an Exa price entry holds a price someone set, which replaces the cost Exa reports. LiteLLM ships
|
||||
Exa with its legacy tiers only, so anything else there was set by `register_model()`, an admin's own cost map,
|
||||
or an edited bundled file."""
|
||||
if isinstance(entry, AliasedCostEntry) and entry.registered_fields & _PRICE_FIELDS:
|
||||
return True
|
||||
def exa_ai_price_fields(model: str, entry: object) -> frozenset[str]:
|
||||
if not isinstance(entry, Mapping):
|
||||
return False
|
||||
return frozenset()
|
||||
fields: Final = _ENTRY.validate_python(entry)
|
||||
tiers: Final = fields.get("tiered_pricing")
|
||||
return fields.get("input_cost_per_query") is not None or (tiers is not None and not _legacy_tiers(tiers))
|
||||
registered: Final = entry.registered_fields & _PRICE_FIELDS if isinstance(entry, AliasedCostEntry) else frozenset()
|
||||
if registered:
|
||||
return registered
|
||||
source: Final = get_model_cost_map_source_info()
|
||||
if source["source"] == "remote" and source["url"] == DEFAULT_MODEL_COST_MAP_URL:
|
||||
loaded: Final = GetModelCostMap.loaded_model_cost_map().get(model, {})
|
||||
return frozenset(
|
||||
name for name in _PRICE_FIELDS if fields.get(name) is not None and fields[name] != loaded.get(name)
|
||||
)
|
||||
return frozenset(
|
||||
name
|
||||
for name in _PRICE_FIELDS
|
||||
if fields.get(name) is not None
|
||||
and (
|
||||
fields[name] != EXA_AI_CATALOG_PRICES.get(model)
|
||||
if name == "input_cost_per_query"
|
||||
else not _legacy_tiers(fields[name])
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
EXA_AI_INCLUDED_RESULTS: Final = 10
|
||||
|
|
@ -74,7 +107,7 @@ def _non_negative_int(value: object) -> int | None:
|
|||
return value
|
||||
|
||||
|
||||
def _requested_results(optional_params: Mapping[str, object]) -> int:
|
||||
def exa_ai_requested_results(optional_params: Mapping[str, object]) -> int:
|
||||
for key in ("max_results", "numResults"):
|
||||
requested: Final = _non_negative_int(optional_params.get(key))
|
||||
if requested is not None:
|
||||
|
|
@ -97,7 +130,7 @@ def exa_ai_list_price(base_price: float, optional_params: Mapping[str, object])
|
|||
the search type's base price, a per-result charge above the results that price includes,
|
||||
and a per-page charge for each requested result when summaries are requested on a type that offers them.
|
||||
"""
|
||||
results: Final = _requested_results(optional_params)
|
||||
results: Final = exa_ai_requested_results(optional_params)
|
||||
additional_results: Final = max(results - EXA_AI_INCLUDED_RESULTS, 0)
|
||||
summarized_pages: Final = results if _requests_summaries(optional_params) else 0
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -122,7 +122,7 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
**kwargs,
|
||||
**kwargs: object,
|
||||
) -> SearchResponse:
|
||||
"""
|
||||
Transform Exa AI API response to LiteLLM unified SearchResponse format.
|
||||
|
|
@ -159,7 +159,18 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
"provider_metadata": {key: value for key, value in response_json.items() if key != "results"},
|
||||
}
|
||||
)
|
||||
set_hidden_params(response, {EXA_AI_COST_DOLLARS_PARAM: response_json.get("costDollars")})
|
||||
params: Final = _JSON_OBJECT.validate_python(kwargs.get("optional_params") or {})
|
||||
set_hidden_params(
|
||||
response,
|
||||
{
|
||||
**{
|
||||
key: value
|
||||
for key, value in params.items()
|
||||
if key in ("type", "max_results", "numResults", "contents")
|
||||
},
|
||||
EXA_AI_COST_DOLLARS_PARAM: response_json.get("costDollars"),
|
||||
},
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,12 @@ from pydantic import TypeAdapter, ValidationError
|
|||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.llms.exa_ai.search.cost_calculator import exa_ai_price_set, exa_ai_reported_cost
|
||||
from litellm.llms.exa_ai.search.cost_calculator import (
|
||||
exa_ai_list_price,
|
||||
exa_ai_price_fields,
|
||||
exa_ai_reported_cost,
|
||||
exa_ai_requested_results,
|
||||
)
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
|
||||
|
|
@ -49,6 +54,12 @@ def _registered_at_runtime(entry: str) -> bool:
|
|||
return live != GetModelCostMap.loaded_model_cost_map().get(entry) # pyright: ignore[reportUnknownVariableType] # `live` comes from the untyped registry
|
||||
|
||||
|
||||
def _operator_priced(model: str, provider: str | None) -> bool:
|
||||
if provider == "exa_ai":
|
||||
return bool(exa_ai_price_fields(model, _live_entry(model)))
|
||||
return _registered_at_runtime(model)
|
||||
|
||||
|
||||
def _pricing_model(
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
|
|
@ -61,10 +72,14 @@ def _pricing_model(
|
|||
generic_model: Final = f"{custom_llm_provider}/search"
|
||||
if mode_param is None or model not in (generic_model, f"{custom_llm_provider}/{generic_model}"):
|
||||
return model
|
||||
mode: Final = (optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS).get(mode_param)
|
||||
mode: Final = (optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS).get(
|
||||
mode_param, "auto" if custom_llm_provider == "exa_ai" else None
|
||||
)
|
||||
mode_model: Final = f"{custom_llm_provider}/search-{mode}"
|
||||
priced: Final = mode_model in litellm.model_cost # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed
|
||||
if not priced or (_registered_at_runtime(generic_model) and not _registered_at_runtime(mode_model)):
|
||||
if not priced or (
|
||||
_operator_priced(generic_model, custom_llm_provider) and not _operator_priced(mode_model, custom_llm_provider)
|
||||
):
|
||||
return model
|
||||
return mode_model
|
||||
|
||||
|
|
@ -74,15 +89,15 @@ def search_provider_cost_per_query(
|
|||
custom_llm_provider: str | None = None,
|
||||
number_of_queries: int = 1,
|
||||
optional_params: Mapping[str, object] | None = None,
|
||||
response: object = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculate cost for search-only providers.
|
||||
|
||||
Returns (input_cost, output_cost) where input_cost = queries * 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`). A flat price beats the tiers on
|
||||
the same entry, since LiteLLM never ships both, so the flat one is someone's override. An Exa request is billed
|
||||
the `costDollars.total` Exa reported for it unless its entry holds a price someone set.
|
||||
provider's mode parameter (Tavily `search_depth`, Linkup `depth`, Exa `type`). An Exa request uses an operator's
|
||||
price, otherwise its reported total, otherwise estimated list price.
|
||||
|
||||
Args:
|
||||
model: Model name (e.g., "exa_ai/search", "tavily/search")
|
||||
|
|
@ -106,21 +121,36 @@ def search_provider_cost_per_query(
|
|||
return (input_cost, 0.0)
|
||||
|
||||
pricing_model: Final = _pricing_model(model, custom_llm_provider, optional_params)
|
||||
reported_cost: Final = (
|
||||
exa_ai_reported_cost(optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS)
|
||||
params: Final = optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS
|
||||
exa_price_fields: Final = (
|
||||
exa_ai_price_fields(
|
||||
_generic_search_entry(pricing_model, custom_llm_provider),
|
||||
_live_entry(_generic_search_entry(pricing_model, custom_llm_provider)),
|
||||
)
|
||||
if custom_llm_provider == "exa_ai"
|
||||
and not exa_ai_price_set(_live_entry(_generic_search_entry(pricing_model, custom_llm_provider)))
|
||||
else None
|
||||
else frozenset()
|
||||
)
|
||||
reported_cost: Final = (
|
||||
exa_ai_reported_cost(params, response) if custom_llm_provider == "exa_ai" and not exa_price_fields else None
|
||||
)
|
||||
if reported_cost is not None:
|
||||
return (reported_cost, 0.0)
|
||||
|
||||
model_info: Final = get_model_info(model=pricing_model, custom_llm_provider=custom_llm_provider)
|
||||
if custom_llm_provider == "exa_ai" and not exa_price_fields:
|
||||
return (exa_ai_list_price(float(model_info.get("input_cost_per_query") or 0.0), params), 0.0)
|
||||
|
||||
# Check for tiered pricing (e.g., Exa AI based on max_results)
|
||||
tiered_pricing: Final = model_info.get("tiered_pricing")
|
||||
if tiered_pricing and isinstance(tiered_pricing, list) and model_info.get("input_cost_per_query") is None:
|
||||
max_results: Final = (optional_params or {}).get("max_results", 10) # default 10 results
|
||||
use_tiers: Final = (
|
||||
"tiered_pricing" in exa_price_fields and "input_cost_per_query" not in exa_price_fields
|
||||
if custom_llm_provider == "exa_ai"
|
||||
else model_info.get("input_cost_per_query") is None
|
||||
)
|
||||
if tiered_pricing and isinstance(tiered_pricing, list) and use_tiers:
|
||||
max_results: Final = (
|
||||
exa_ai_requested_results(params) if custom_llm_provider == "exa_ai" else params.get("max_results", 10)
|
||||
)
|
||||
cost_per_query = 0.0
|
||||
|
||||
for tier in tiered_pricing:
|
||||
|
|
|
|||
|
|
@ -3515,9 +3515,13 @@ def _replace_cost_map_entry(
|
|||
previous_fields: Final[frozenset[str]] = (
|
||||
current.registered_fields if isinstance(current, AliasedCostEntry) else frozenset()
|
||||
)
|
||||
registered_fields: Final = previous_fields | frozenset(
|
||||
name for name, value in registration.items() if value is not None
|
||||
supplied_fields: Final = frozenset(name for name, value in registration.items() if value is not None)
|
||||
replaced_fields: Final = (
|
||||
frozenset({"input_cost_per_query", "tiered_pricing"})
|
||||
if merged.get("litellm_provider") == "exa_ai" and supplied_fields & {"input_cost_per_query", "tiered_pricing"}
|
||||
else frozenset()
|
||||
)
|
||||
registered_fields: Final = (previous_fields - replaced_fields) | supplied_fields
|
||||
sharing_names: Final = tuple(name for name in names if cost_map.get(name) is current)
|
||||
cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, names, registered_fields)))
|
||||
|
||||
|
|
|
|||
|
|
@ -95,7 +95,9 @@ async def test_search_logs_one_completed_response_and_cost(
|
|||
assert await asyncio.to_thread(recorder.finished.wait, 5)
|
||||
assert len(recorder.records) == 1
|
||||
details, logged = recorder.records[0]
|
||||
assert logged.model_dump()["provider_metadata"]["output"] == response.model_dump()["provider_metadata"]["output"]
|
||||
assert (
|
||||
logged.model_dump()["provider_metadata"]["output"] == response.model_dump()["provider_metadata"]["output"]
|
||||
)
|
||||
assert isinstance(details["standard_logging_object"], Mapping)
|
||||
expected_cost: Final = PAYLOAD["costDollars"]["total"] if price_override is None else price_override
|
||||
assert details["response_cost"] == expected_cost
|
||||
|
|
|
|||
|
|
@ -1,11 +1,20 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_LEGACY_TIERS, ExaAITier
|
||||
from litellm.llms.exa_ai.search.cost_calculator import (
|
||||
EXA_AI_ADDITIONAL_RESULT_COST,
|
||||
EXA_AI_CATALOG_PRICES,
|
||||
EXA_AI_INCLUDED_RESULTS,
|
||||
EXA_AI_LEGACY_TIERS,
|
||||
EXA_AI_SUMMARY_COST_PER_PAGE,
|
||||
ExaAITier,
|
||||
exa_ai_list_price,
|
||||
)
|
||||
|
||||
REPO_ROOT: Final = Path(__file__).resolve().parents[5]
|
||||
SHIPPED_COST_MAPS: Final = (
|
||||
|
|
@ -23,18 +32,34 @@ def _search_entries(path: str) -> dict[str, dict[str, object]]:
|
|||
|
||||
|
||||
@pytest.mark.parametrize("path", SHIPPED_COST_MAPS)
|
||||
def test_litellm_ships_exa_search_with_only_the_legacy_tiers_the_cost_calculator_treats_as_its_own(path: str) -> None:
|
||||
exa_entry: Final = _search_entries(path)["exa_ai/search"]
|
||||
|
||||
assert TypeAdapter(tuple[ExaAITier, ...]).validate_python(exa_entry["tiered_pricing"]) == EXA_AI_LEGACY_TIERS
|
||||
assert exa_entry.get("input_cost_per_query") is None
|
||||
assert [key for key in _search_entries(path) if key.startswith("exa_ai/search-")] == []
|
||||
def test_shipped_exa_prices_agree_with_the_receipt_override_baseline(path: str) -> None:
|
||||
entries: Final = _search_entries(path)
|
||||
assert {
|
||||
key: entry["input_cost_per_query"] for key, entry in entries.items() if key.startswith("exa_ai/")
|
||||
} == EXA_AI_CATALOG_PRICES
|
||||
assert (
|
||||
TypeAdapter(tuple[ExaAITier, ...]).validate_python(entries["exa_ai/search"]["tiered_pricing"])
|
||||
== EXA_AI_LEGACY_TIERS
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", SHIPPED_COST_MAPS)
|
||||
def test_litellm_ships_no_search_entry_with_both_a_flat_price_and_tiers(path: str) -> None:
|
||||
assert [
|
||||
key
|
||||
for key, entry in _search_entries(path).items()
|
||||
if entry.get("input_cost_per_query") is not None and entry.get("tiered_pricing") is not None
|
||||
] == []
|
||||
@pytest.mark.parametrize(
|
||||
("params", "extra_results", "summaries"),
|
||||
(
|
||||
({}, 0, 0),
|
||||
({"max_results": None}, 0, 0),
|
||||
({"max_results": True, "numResults": EXA_AI_INCLUDED_RESULTS + 3}, 3, 0),
|
||||
({"numResults": EXA_AI_INCLUDED_RESULTS + 3}, 3, 0),
|
||||
({"max_results": EXA_AI_INCLUDED_RESULTS + 2, "numResults": 100}, 2, 0),
|
||||
({"max_results": 3, "contents": {"summary": {}}}, 0, 3),
|
||||
({"max_results": 3, "contents": {"summary": False}}, 0, 0),
|
||||
({"type": "lite", "max_results": 3, "contents": {"summary": {}}}, 0, 0),
|
||||
),
|
||||
)
|
||||
def test_fallback_adds_only_requested_results_and_summaries(
|
||||
params: Mapping[str, object], extra_results: int, summaries: int
|
||||
) -> None:
|
||||
base_price: Final = 0.123
|
||||
assert exa_ai_list_price(base_price, params) == pytest.approx(
|
||||
base_price + extra_results * EXA_AI_ADDITIONAL_RESULT_COST + summaries * EXA_AI_SUMMARY_COST_PER_PAGE
|
||||
)
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig
|
|||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
||||
EXA_SEARCH_URL: Final = "https://api.exa.ai/search"
|
||||
EXA_REPORTED_COST: Final = 0.022
|
||||
EXA_REPORTED_COST: Final = 0.123456
|
||||
_LOGGED: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import datetime
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -8,6 +9,7 @@ import pytest
|
|||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import DEFAULT_MODEL_COST_MAP_URL
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
|
@ -15,7 +17,7 @@ from litellm.search.cost_calculator import search_provider_cost_per_query
|
|||
EXA_SEARCH_URL: Final = "https://api.exa.ai/search"
|
||||
TAVILY_SEARCH_URL: Final = "https://api.tavily.com/search"
|
||||
LINKUP_SEARCH_URL: Final = "https://api.linkup.so/v1/search"
|
||||
EXA_REPORTED_COST: Final = 0.022
|
||||
EXA_REPORTED_COST: Final = 0.123456
|
||||
EXA_REPLY_WITH_COST: Final = {
|
||||
"results": [{"title": "t", "url": "https://example.com"}],
|
||||
"costDollars": {"total": EXA_REPORTED_COST},
|
||||
|
|
@ -234,9 +236,7 @@ def test_an_exa_price_someone_set_overrides_the_cost_exa_reports_and_litellms_ti
|
|||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
def test_registering_litellms_own_exa_tiers_still_overrides_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None:
|
||||
respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST)
|
||||
tier_price, _ = search_provider_cost_per_query(
|
||||
model="exa_ai/search", custom_llm_provider="exa_ai", optional_params={"max_results": 25}
|
||||
)
|
||||
tier_price: Final = litellm.model_cost["exa_ai/search"]["tiered_pricing"][0]["input_cost_per_query"]
|
||||
assert tier_price != EXA_REPORTED_COST
|
||||
litellm.register_model(
|
||||
{
|
||||
|
|
@ -298,13 +298,20 @@ def test_a_flat_price_registered_on_a_tiered_search_entry_beats_its_tiers() -> N
|
|||
|
||||
|
||||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
@pytest.mark.parametrize("mode", ("auto", "deep"))
|
||||
def test_the_cost_exa_reports_is_billed_when_the_cost_map_has_no_exa_entry(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, mode: str
|
||||
) -> None:
|
||||
respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST)
|
||||
monkeypatch.delitem(litellm.model_cost, "exa_ai/search")
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"model_cost",
|
||||
{key: value for key, value in litellm.model_cost.items() if not key.startswith("exa_ai/")},
|
||||
)
|
||||
|
||||
response: Final = litellm.search(query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key")
|
||||
response: Final = litellm.search(
|
||||
query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key", type=mode
|
||||
)
|
||||
|
||||
assert response.hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_COST)
|
||||
|
||||
|
|
@ -449,6 +456,116 @@ def test_a_runtime_price_on_the_mode_entry_wins_over_a_runtime_price_on_the_gene
|
|||
assert cost == pytest.approx((0.011, 0.0))
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
def test_exa_receipt_survives_public_response_serialization(respx_mock: respx.MockRouter) -> None:
|
||||
respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST)
|
||||
response: Final = litellm.search(query="receipt", search_provider="exa_ai", api_key="test-key")
|
||||
|
||||
assert litellm.completion_cost(
|
||||
completion_response=response.model_dump(),
|
||||
model="exa_ai/search",
|
||||
custom_llm_provider="exa_ai",
|
||||
call_type="search",
|
||||
) == pytest.approx(EXA_REPORTED_COST)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
@pytest.mark.parametrize("params", ({"type": "deep"}, {}), ids=("explicit-mode", "default-mode"))
|
||||
def test_exa_pricing_apis_keep_the_requested_mode(respx_mock: respx.MockRouter, params: Mapping[str, object]) -> None:
|
||||
respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST)
|
||||
mode: Final = params.get("type", "auto")
|
||||
price: Final = 0.321
|
||||
litellm.register_model(
|
||||
{f"exa_ai/search-{mode}": {"input_cost_per_query": price, "litellm_provider": "exa_ai", "mode": "search"}},
|
||||
persist_across_reloads=False,
|
||||
)
|
||||
response: Final = litellm.search(query="mode", search_provider="exa_ai", api_key="test-key", **params)
|
||||
|
||||
assert response.hidden_params["response_cost"] == pytest.approx(price)
|
||||
assert sum(
|
||||
litellm.cost_per_token(
|
||||
model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search", response=response
|
||||
)
|
||||
) == pytest.approx(price)
|
||||
assert litellm.completion_cost(
|
||||
completion_response=response, model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search"
|
||||
) == pytest.approx(price)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_local_model_cost_map")
|
||||
def test_exa_native_result_count_bills_the_same_operator_tier(respx_mock: respx.MockRouter) -> None:
|
||||
route: Final = respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST)
|
||||
price: Final = 0.456
|
||||
litellm.register_model(
|
||||
{
|
||||
"exa_ai/search": {
|
||||
"tiered_pricing": [
|
||||
{"input_cost_per_query": 0.123, "max_results_range": [0, 10]},
|
||||
{"input_cost_per_query": price, "max_results_range": [11, 100]},
|
||||
]
|
||||
}
|
||||
},
|
||||
persist_across_reloads=False,
|
||||
)
|
||||
for params in ({"max_results": 50}, {"numResults": 50}, {"max_results": None, "numResults": 50}):
|
||||
response: Final = litellm.search(query="count", search_provider="exa_ai", api_key="test-key", **params)
|
||||
assert json.loads(route.calls.last.request.content)["numResults"] == 50
|
||||
assert response.hidden_params["response_cost"] == pytest.approx(price)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CostMapServer:
|
||||
body: bytes
|
||||
|
||||
def get(self, url: str, *, timeout: float | None = None) -> httpx.Response:
|
||||
return httpx.Response(200, content=self.body, request=httpx.Request("GET", url))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _restore_catalog_source(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
|
||||
yield
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=DEFAULT_MODEL_COST_MAP_URL))
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_restore_catalog_source")
|
||||
@pytest.mark.parametrize(
|
||||
("source", "price", "expected_cost"),
|
||||
(
|
||||
(DEFAULT_MODEL_COST_MAP_URL, 0.321, EXA_REPORTED_COST),
|
||||
("https://prices.invalid/map.json", 0.321, 0.321),
|
||||
("https://prices.invalid/map.json", None, EXA_REPORTED_COST),
|
||||
),
|
||||
ids=("updated-official-catalog", "operator-map", "mirror"),
|
||||
)
|
||||
@pytest.mark.parametrize("mode", ("auto", "lite", "deep"))
|
||||
def test_exa_catalog_source_preserves_receipts_and_generic_operator_prices(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
source: str,
|
||||
price: float | None,
|
||||
expected_cost: float,
|
||||
mode: str,
|
||||
) -> None:
|
||||
respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST)
|
||||
catalog: Final = GetModelCostMap.load_local_model_cost_map()
|
||||
served: Final = {
|
||||
**catalog,
|
||||
"exa_ai/search": {
|
||||
**catalog["exa_ai/search"],
|
||||
**({"input_cost_per_query": price} if price is not None else {}),
|
||||
},
|
||||
}
|
||||
monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False)
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"model_cost",
|
||||
litellm.get_model_cost_map(url=source, max_attempts=1, client=_CostMapServer(json.dumps(served).encode())),
|
||||
)
|
||||
response: Final = litellm.search(query="catalog", search_provider="exa_ai", api_key="test-key", type=mode)
|
||||
assert response.hidden_params["response_cost"] == pytest.approx(expected_cost)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_search_mode_prices")
|
||||
def test_sync_tavily_advanced_search_reports_the_advanced_price(respx_mock: respx.MockRouter) -> None:
|
||||
respx_mock.post(TAVILY_SEARCH_URL).respond(json={"results": [{"title": "t", "url": "https://example.com"}]})
|
||||
|
|
|
|||
|
|
@ -2325,6 +2325,20 @@ def test_register_model_price_is_used_for_completion_cost() -> None:
|
|||
assert cost == pytest.approx(0.02)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("flat_first", (True, False))
|
||||
def test_latest_exa_price_registration_replaces_the_previous_price_kind(
|
||||
local_model_cost_map: None, flat_first: bool
|
||||
) -> None:
|
||||
flat: Final = {"input_cost_per_query": 0.321}
|
||||
tiered: Final = {"tiered_pricing": [{"input_cost_per_query": 0.456, "max_results_range": [0, 100]}]}
|
||||
for price in (flat, tiered) if flat_first else (tiered, flat):
|
||||
litellm.register_model({"exa_ai/search": price}, persist_across_reloads=False)
|
||||
|
||||
assert litellm.completion_cost(
|
||||
model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search"
|
||||
) == pytest.approx(0.456 if flat_first else 0.321)
|
||||
|
||||
|
||||
@respx.mock
|
||||
@pytest.mark.usefixtures("isolated_openai_model_sets", "local_model_cost_map")
|
||||
def test_completion_registers_test_owned_model_prices() -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue