mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
fix(exa): keep Exa's envelope in provider_metadata, price summaries, and honor operator prices
Exa's request envelope (requestId, costDollars, output and the rest) now lives only under provider_metadata like the other search providers, and native result fields that LiteLLM renames are no longer duplicated on each result. Fallback estimates add Exa's per-page summary charge. A price the operator registered at runtime or served from their own cost map wins over Exa's reported total, and a flat registered price replaces the catalog's tiers for every tiered search entry. exa_ai/search keeps its tiered prices for older releases that read the remote cost map.
This commit is contained in:
mateo-berri
parent
43c59c2cd7
commit
96a2b14005
11 changed files with 322 additions and 58 deletions
|
|
@ -42,6 +42,7 @@ from litellm._logging import (
|
|||
import re
|
||||
from litellm.constants import (
|
||||
DEFAULT_BATCH_SIZE,
|
||||
DEFAULT_MODEL_COST_MAP_URL,
|
||||
DEFAULT_FLUSH_INTERVAL_SECONDS,
|
||||
ROUTER_MAX_FALLBACKS,
|
||||
DEFAULT_MAX_RETRIES,
|
||||
|
|
@ -408,10 +409,7 @@ add_function_to_prompt: bool = (
|
|||
client_session: Optional[Union[httpx.Client, DefaultHttpxClient]] = None
|
||||
aclient_session: Optional[Union[httpx.AsyncClient, DefaultAsyncHttpxClient]] = None
|
||||
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",
|
||||
)
|
||||
model_cost_map_url: str = os.getenv("LITELLM_MODEL_COST_MAP_URL", DEFAULT_MODEL_COST_MAP_URL)
|
||||
blog_posts_url: str = os.getenv(
|
||||
"LITELLM_BLOG_POSTS_URL",
|
||||
"https://docs.litellm.ai/blog/rss.xml",
|
||||
|
|
|
|||
|
|
@ -189,6 +189,9 @@ MODEL_COST_MAP_MIN_MODEL_COUNT: Final = get_env_int(
|
|||
MODEL_COST_MAP_MAX_SHRINK_RATIO: Final = get_env_float(
|
||||
"MODEL_COST_MAP_MAX_SHRINK_RATIO", 0.5
|
||||
) # Maximum allowed shrinkage ratio vs local backup (0.5 = reject if fetched map is <50% of backup)
|
||||
DEFAULT_MODEL_COST_MAP_URL: Final = (
|
||||
"https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
|
||||
)
|
||||
DEFAULT_IMAGE_WIDTH: Final = get_env_int("DEFAULT_IMAGE_WIDTH", 300)
|
||||
DEFAULT_IMAGE_HEIGHT: Final = get_env_int("DEFAULT_IMAGE_HEIGHT", 300)
|
||||
# Maximum size for image URL downloads in MB (default 50MB, set to 0 to disable limit)
|
||||
|
|
|
|||
|
|
@ -1,8 +1,12 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
EXA_AI_INCLUDED_RESULTS: Final = 10
|
||||
EXA_AI_ADDITIONAL_RESULT_COST: Final = 0.001
|
||||
EXA_AI_SUMMARY_COST_PER_PAGE: Final = 0.001
|
||||
_CONTENTS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _non_negative_int(value: object) -> int | None:
|
||||
|
|
@ -19,12 +23,24 @@ def _requested_results(optional_params: Mapping[str, object]) -> int:
|
|||
return EXA_AI_INCLUDED_RESULTS
|
||||
|
||||
|
||||
def _requests_summaries(optional_params: Mapping[str, object]) -> bool:
|
||||
contents: Final = optional_params.get("contents")
|
||||
return isinstance(contents, Mapping) and _CONTENTS.validate_python(contents).get("summary") not in (None, False)
|
||||
|
||||
|
||||
def exa_ai_list_price(base_price: float, optional_params: Mapping[str, object]) -> float:
|
||||
"""
|
||||
List price of one Exa /search request, for responses that carry no `costDollars.total`.
|
||||
|
||||
Multi-query searches are joined into a single request, so this prices one request:
|
||||
the search type's base price plus a per-result charge above the results that price includes.
|
||||
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.
|
||||
"""
|
||||
additional_results: Final = max(_requested_results(optional_params) - EXA_AI_INCLUDED_RESULTS, 0)
|
||||
return base_price + additional_results * EXA_AI_ADDITIONAL_RESULT_COST
|
||||
results: Final = _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 (
|
||||
base_price
|
||||
+ additional_results * EXA_AI_ADDITIONAL_RESULT_COST
|
||||
+ summarized_pages * EXA_AI_SUMMARY_COST_PER_PAGE
|
||||
)
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i
|
|||
_JSON_VALUES: Final = TypeAdapter(Iterable[object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_STRS: Final = TypeAdapter(Iterable[str], config=ConfigDict(hide_input_in_errors=True))
|
||||
_RESULT_EXTRAS: Final = (("highlights", "highlights"), ("highlightScores", "highlight_scores"), ("score", "score"))
|
||||
_RENAMED_RESULT_FIELDS: Final = frozenset({"text", "publishedDate", "highlightScores"})
|
||||
|
||||
|
||||
class ExaAISearchConfig(BaseSearchConfig):
|
||||
|
|
@ -133,7 +134,8 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
- No last_updated field in Exa AI response (set to None)
|
||||
- results[].highlights / highlightScores / score → SearchResult.highlights / highlight_scores / score,
|
||||
each set only when Exa returned it
|
||||
- every top-level key except results (requestId, costDollars, ...) → SearchResponse.provider_metadata
|
||||
- every other results[] field Exa returned (id, author, image, summary, ...) → the SearchResult as is
|
||||
- every top-level key except results (requestId, costDollars, output, ...) → SearchResponse.provider_metadata
|
||||
|
||||
Args:
|
||||
raw_response: Raw httpx response from Exa AI API
|
||||
|
|
@ -146,7 +148,6 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
|
||||
return SearchResponse.model_validate(
|
||||
{
|
||||
**response_json,
|
||||
"results": [
|
||||
self._transform_result(value)
|
||||
for value in _JSON_VALUES.validate_python(response_json.get("results") or [])
|
||||
|
|
@ -161,7 +162,7 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
result: Final = _JSON_OBJECT.validate_python(value)
|
||||
return SearchResult.model_validate(
|
||||
{
|
||||
**result,
|
||||
**{key: value for key, value in result.items() if key not in _RENAMED_RESULT_FIELDS},
|
||||
"title": result.get("title", ""),
|
||||
"url": result.get("url", ""),
|
||||
"snippet": result.get("text")
|
||||
|
|
|
|||
|
|
@ -23627,9 +23627,25 @@
|
|||
"input_cost_per_query": 0.007,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present, unless this price is registered at runtime. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
"notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search"
|
||||
"mode": "search",
|
||||
"tiered_pricing": [
|
||||
{
|
||||
"input_cost_per_query": 0.005,
|
||||
"max_results_range": [
|
||||
0,
|
||||
25
|
||||
]
|
||||
},
|
||||
{
|
||||
"input_cost_per_query": 0.025,
|
||||
"max_results_range": [
|
||||
26,
|
||||
100
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"exa_ai/search-lite": {
|
||||
"input_cost_per_query": 0.004,
|
||||
|
|
|
|||
|
|
@ -10,7 +10,8 @@ from typing import Final
|
|||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
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.llms.base_llm.search.transformation import SearchResponse
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
|
|
@ -36,10 +37,40 @@ def _provider_usage(
|
|||
return None
|
||||
|
||||
|
||||
def _field(value: object, key: str) -> object:
|
||||
return PROVIDER_COST_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 _registered_at_runtime(entry: str) -> bool:
|
||||
"""Whether `entry`'s live price differs from the catalog as loaded, i.e. `register_model` changed it."""
|
||||
live: Final[object] = litellm.model_cost.get(entry) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # registry is dynamically typed
|
||||
return live != GetModelCostMap.loaded_model_cost_map().get(entry) # pyright: ignore[reportUnknownVariableType] # `live` comes from the untyped registry
|
||||
return _live_entry(entry) != GetModelCostMap.loaded_model_cost_map().get(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)
|
||||
|
||||
|
||||
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 _exa_reported_total(response: object) -> float | None:
|
||||
payload: Final = response.model_extra if isinstance(response, SearchResponse) else response
|
||||
total: Final = _field(_field(_field(payload, "provider_metadata"), "costDollars"), "total")
|
||||
if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0:
|
||||
return float(total)
|
||||
return None
|
||||
|
||||
|
||||
def _pricing_model(
|
||||
|
|
@ -75,14 +106,15 @@ def search_provider_cost_per_query(
|
|||
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`). Exa bills one request at the
|
||||
`costDollars.total` it reports, else at its list price, unless the price entry was registered at runtime.
|
||||
`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.
|
||||
|
||||
Args:
|
||||
model: Model name (e.g., "exa_ai/search", "tavily/search")
|
||||
custom_llm_provider: Provider name (e.g., "exa_ai", "tavily")
|
||||
number_of_queries: Number of search queries performed (default: 1)
|
||||
optional_params: Optional parameters including max_results for tiered pricing
|
||||
response: Provider response with reported Exa costDollars.total, when available
|
||||
response: Provider response with Exa's reported provider_metadata.costDollars.total, when available
|
||||
|
||||
Returns:
|
||||
Tuple of (input_cost, output_cost) where output_cost is always 0.0
|
||||
|
|
@ -103,16 +135,10 @@ def search_provider_cost_per_query(
|
|||
model=_pricing_model(model, custom_llm_provider, optional_params),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
if custom_llm_provider == "exa_ai" and not _registered_at_runtime(model_info["key"]):
|
||||
payload: Final = response.model_extra if isinstance(response, SearchResponse) else response
|
||||
reported: Final = (
|
||||
PROVIDER_COST_ADAPTER.validate_python(payload).get("costDollars") if isinstance(payload, Mapping) else None
|
||||
)
|
||||
total: Final = (
|
||||
PROVIDER_COST_ADAPTER.validate_python(reported).get("total") if isinstance(reported, Mapping) else None
|
||||
)
|
||||
if isinstance(total, (int, float)) and not isinstance(total, bool) and math.isfinite(total) and total >= 0:
|
||||
return (float(total), 0.0)
|
||||
if custom_llm_provider == "exa_ai" and not _operator_priced(model_info["key"]):
|
||||
reported_total: Final = _exa_reported_total(response)
|
||||
if reported_total is not None:
|
||||
return (reported_total, 0.0)
|
||||
|
||||
from litellm.llms.exa_ai.search.cost_calculator import exa_ai_list_price
|
||||
|
||||
|
|
@ -126,7 +152,11 @@ def search_provider_cost_per_query(
|
|||
|
||||
# Check for tiered pricing based on max_results
|
||||
tiered_pricing: Final = model_info.get("tiered_pricing")
|
||||
if tiered_pricing and isinstance(tiered_pricing, list):
|
||||
if (
|
||||
tiered_pricing
|
||||
and isinstance(tiered_pricing, list)
|
||||
and not _registered_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
|
||||
|
||||
|
|
|
|||
|
|
@ -23627,9 +23627,25 @@
|
|||
"input_cost_per_query": 0.007,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001. Exa's reported costDollars.total takes precedence when present, unless this price is registered at runtime. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
"notes": "Base price per /search request, covering up to 10 results. Each result above 10 adds $0.001, and each summarized page adds $0.001. Exa's reported costDollars.total takes precedence when present, unless an operator priced this entry. tiered_pricing keeps what older releases bill for this entry. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search"
|
||||
"mode": "search",
|
||||
"tiered_pricing": [
|
||||
{
|
||||
"input_cost_per_query": 0.005,
|
||||
"max_results_range": [
|
||||
0,
|
||||
25
|
||||
]
|
||||
},
|
||||
{
|
||||
"input_cost_per_query": 0.025,
|
||||
"max_results_range": [
|
||||
26,
|
||||
100
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"exa_ai/search-lite": {
|
||||
"input_cost_per_query": 0.004,
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ _EXPECTED: Final[dict[str, JsonValue]] = {
|
|||
"object": "search",
|
||||
"results": [
|
||||
{
|
||||
"id": "https://result.invalid/sdk",
|
||||
"title": "SDK result",
|
||||
"url": "https://result.invalid/sdk",
|
||||
"snippet": "First line.\n\nSecond line.",
|
||||
|
|
|
|||
|
|
@ -42,13 +42,21 @@ async def test_exa_search_request_and_full_response_cross_real_http(asynchronous
|
|||
**PARAMS,
|
||||
}
|
||||
response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params)
|
||||
result: Final = json.loads(response.model_dump_json())
|
||||
assert result["output"] == PAYLOAD["output"]
|
||||
assert result["requestId"] == PAYLOAD["requestId"]
|
||||
assert result["costDollars"] == PAYLOAD["costDollars"]
|
||||
assert result["results"][0]["highlights"] == ["evidence"]
|
||||
assert result["results"][0]["author"] is None
|
||||
assert result["results"][0]["snippet"] == "evidence"
|
||||
assert json.loads(response.model_dump_json()) == {
|
||||
"object": "search",
|
||||
"results": [
|
||||
{
|
||||
"title": "Source",
|
||||
"url": "https://example.com",
|
||||
"snippet": "evidence",
|
||||
"date": None,
|
||||
"last_updated": None,
|
||||
"highlights": ["evidence"],
|
||||
"author": None,
|
||||
}
|
||||
],
|
||||
"provider_metadata": {key: value for key, value in PAYLOAD.items() if key != "results"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -60,19 +68,25 @@ async def test_search_logs_one_completed_response_and_cost(
|
|||
) -> None:
|
||||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse
|
||||
|
||||
class Recorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
self.finished = threading.Event()
|
||||
self.records = ()
|
||||
self.records: tuple[tuple[Mapping[str, object], SearchResponse], ...] = ()
|
||||
|
||||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
def log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: SearchResponse, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.records = (*self.records, (kwargs, response_obj))
|
||||
self.finished.set()
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: SearchResponse, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.log_success_event(kwargs, response_obj, start_time, end_time)
|
||||
|
||||
recorder: Final = Recorder()
|
||||
|
|
@ -95,7 +109,7 @@ 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()["output"] == response.model_dump()["output"]
|
||||
assert logged.model_dump() == response.model_dump()
|
||||
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
|
||||
|
|
|
|||
|
|
@ -144,10 +144,18 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order():
|
|||
({"highlightScores": [0.1]}, {"highlight_scores": [0.1]}),
|
||||
({"score": None}, {"score": None}),
|
||||
({"highlights": [], "highlightScores": [], "score": 0}, {"highlights": [], "highlight_scores": [], "score": 0}),
|
||||
({"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"}, {}),
|
||||
(
|
||||
{"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"},
|
||||
{"id": "https://example.com", "author": "Ada", "image": None, "summary": "a summary"},
|
||||
),
|
||||
({"publishedDate": "2026-01-01"}, {}),
|
||||
(
|
||||
{"text": "full text", "highlightScores": [0.2], "publishedDate": None, "favicon": "f"},
|
||||
{"highlight_scores": [0.2], "favicon": "f"},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_transform_search_response_attaches_only_the_exa_result_fields_it_returned(
|
||||
def test_transform_search_response_carries_each_exa_result_field_once(
|
||||
exa_fields: dict[str, object], expected_extras: dict[str, object]
|
||||
):
|
||||
response: Final = _transform(
|
||||
|
|
@ -159,8 +167,8 @@ def test_transform_search_response_attaches_only_the_exa_result_fields_it_return
|
|||
}
|
||||
)
|
||||
|
||||
assert {key: response.results[1].model_extra[key] for key in expected_extras} == expected_extras
|
||||
assert {key: response.results[1].model_extra[key] for key in exa_fields} == exa_fields
|
||||
assert response.results[0].model_extra == {}
|
||||
assert response.results[1].model_extra == expected_extras
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -346,15 +354,34 @@ async def test_public_search_preserves_native_contract(respx_mock: respx.MockRou
|
|||
)
|
||||
assert json.loads(route.calls.last.request.content) == {"query": "question", **native}
|
||||
assert route.calls.last.request.headers["exa-beta"] == headers["Exa-Beta"]
|
||||
dumped: Final = result.model_dump()
|
||||
for key, value in FULL_RESPONSE.items():
|
||||
if key != "results":
|
||||
assert dumped[key] == value
|
||||
for key, value in FULL_RESULT.items():
|
||||
assert dumped["results"][0][key] == value
|
||||
assert dumped["results"][0]["snippet"] == "\n\n".join(FULL_RESULT["highlights"])
|
||||
assert dumped["results"][0]["date"] == FULL_RESULT["publishedDate"]
|
||||
assert json.loads(result.model_dump_json()) == dumped
|
||||
assert result.model_dump() == {
|
||||
"object": "search",
|
||||
"results": [
|
||||
{
|
||||
"title": "Source",
|
||||
"url": "https://example.com",
|
||||
"snippet": "relevant passage\n\nanother passage",
|
||||
"date": "2026-01-01",
|
||||
"last_updated": None,
|
||||
"author": None,
|
||||
"id": "document-id",
|
||||
"image": "https://example.com/image.png",
|
||||
"favicon": "https://example.com/icon.png",
|
||||
"highlights": ["relevant passage", "another passage"],
|
||||
"highlight_scores": [0.8, 0.7],
|
||||
"summary": "summary",
|
||||
"subpages": [{"title": "Subpage", "url": "https://example.com/sub"}],
|
||||
"extras": {
|
||||
"links": ["https://example.com/link"],
|
||||
"richLinks": [{"url": "https://example.com/link", "anchor": "Link"}],
|
||||
},
|
||||
"entities": [{"id": "entity", "type": "company", "properties": {"name": "Example"}}],
|
||||
"futureField": {"nested": [False, 0, None]},
|
||||
}
|
||||
],
|
||||
"provider_metadata": {key: value for key, value in FULL_RESPONSE.items() if key != "results"},
|
||||
}
|
||||
assert json.loads(result.model_dump_json()) == result.model_dump()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -390,7 +417,7 @@ def test_contents_are_sent_without_injecting_text(respx_mock: respx.MockRouter,
|
|||
)
|
||||
def test_output_preserves_present_empty_and_null_values(output: object):
|
||||
response: Final = _transform({"results": [], "output": output})
|
||||
assert response.model_dump()["output"] == output
|
||||
assert response.model_dump() == {"results": [], "object": "search", "provider_metadata": {"output": output}}
|
||||
|
||||
|
||||
def test_absent_output_stays_absent():
|
||||
|
|
@ -439,7 +466,9 @@ def test_legacy_parameters_and_unified_precedence_are_preserved():
|
|||
def test_exa_reported_total_prices_the_whole_request(total: float, serialized: bool) -> None:
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
||||
response: Final = SearchResponse.model_validate({"results": [], "costDollars": {"total": total}})
|
||||
response: Final = SearchResponse.model_validate(
|
||||
{"results": [], "costDollars": {"total": 999}, "provider_metadata": {"costDollars": {"total": total}}}
|
||||
)
|
||||
payload: Final = response.model_dump() if serialized else response
|
||||
assert search_provider_cost_per_query("exa_ai/search", "exa_ai", number_of_queries=3, response=payload) == (
|
||||
total,
|
||||
|
|
@ -482,7 +511,7 @@ def test_exa_missing_or_invalid_total_uses_estimate(reported: object) -> None:
|
|||
search_provider_cost_per_query(
|
||||
"exa_ai/search",
|
||||
"exa_ai",
|
||||
response={"costDollars": reported},
|
||||
response={"costDollars": {"total": 999}, "provider_metadata": {"costDollars": reported}},
|
||||
optional_params={"costDollars": {"total": 999}},
|
||||
)
|
||||
== estimate
|
||||
|
|
@ -506,6 +535,30 @@ def test_exa_estimate_adds_each_result_above_the_included_ten() -> None:
|
|||
) == pytest.approx((base + 2 * EXA_AI_ADDITIONAL_RESULT_COST, 0.0))
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
@pytest.mark.parametrize(
|
||||
("contents", "summarized_pages"),
|
||||
[
|
||||
(None, 0),
|
||||
({"text": True}, 0),
|
||||
({"summary": False}, 0),
|
||||
({"summary": True}, 12),
|
||||
({"summary": {}}, 12),
|
||||
({"summary": {"query": "key facts"}}, 12),
|
||||
],
|
||||
)
|
||||
def test_exa_estimate_adds_each_summarized_page(contents: object, summarized_pages: int) -> None:
|
||||
from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_SUMMARY_COST_PER_PAGE
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
||||
unsummarized, _ = search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params={"max_results": 12})
|
||||
params: Final = {"max_results": 12, **({"contents": contents} if contents is not None else {})}
|
||||
|
||||
assert search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params=params) == pytest.approx(
|
||||
(unsummarized + summarized_pages * EXA_AI_SUMMARY_COST_PER_PAGE, 0.0)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
def test_exa_estimate_prices_one_request_however_many_queries_it_joins() -> None:
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import datetime
|
||||
import json
|
||||
from collections.abc import 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
|
||||
|
|
@ -317,3 +319,117 @@ def test_a_price_registered_for_exa_wins_over_the_cost_exa_reports(
|
|||
**({"type": search_type} if search_type is not None else {}),
|
||||
)
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(expected)
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ServedCostMap:
|
||||
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 _remote_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, Mapping[str, object]], None]]:
|
||||
def serve(url: str, exa_entry: Mapping[str, object]) -> None:
|
||||
monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP")
|
||||
body: Final = json.dumps({**GetModelCostMap.load_local_model_cost_map(), "exa_ai/search": exa_entry}).encode()
|
||||
monkeypatch.setattr(
|
||||
litellm, "model_cost", litellm.get_model_cost_map(url=url, max_attempts=1, client=_ServedCostMap(body))
|
||||
)
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
yield serve
|
||||
litellm.get_model_cost_map(url="")
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("url", "exa_entry", "expected"),
|
||||
[
|
||||
pytest.param(
|
||||
OPERATOR_COST_MAP_URL,
|
||||
{"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,
|
||||
{
|
||||
"tiered_pricing": [{"input_cost_per_query": OPERATOR_TIER_PRICE, "max_results_range": [0, 100]}],
|
||||
"litellm_provider": "exa_ai",
|
||||
"mode": "search",
|
||||
},
|
||||
OPERATOR_TIER_PRICE,
|
||||
id="own-map-tiers",
|
||||
),
|
||||
pytest.param(
|
||||
DEFAULT_MODEL_COST_MAP_URL,
|
||||
{"input_cost_per_query": OPERATOR_FLAT_PRICE, "litellm_provider": "exa_ai", "mode": "search"},
|
||||
EXA_REPORTED_TOTAL,
|
||||
id="published-map",
|
||||
),
|
||||
],
|
||||
)
|
||||
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],
|
||||
url: str,
|
||||
exa_entry: Mapping[str, object],
|
||||
expected: float,
|
||||
) -> None:
|
||||
_remote_cost_map(url, exa_entry)
|
||||
respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}})
|
||||
|
||||
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(expected)
|
||||
|
||||
|
||||
REGISTERED_TIER_PRICE: Final = 0.0456
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_loaded_catalog_prices")
|
||||
@pytest.mark.parametrize(
|
||||
("registration", "expected"),
|
||||
[
|
||||
pytest.param({"input_cost_per_query": EXA_REGISTERED_PRICE}, EXA_REGISTERED_PRICE, id="flat"),
|
||||
pytest.param(
|
||||
{"tiered_pricing": [{"input_cost_per_query": REGISTERED_TIER_PRICE, "max_results_range": [0, 100]}]},
|
||||
REGISTERED_TIER_PRICE,
|
||||
id="tiers",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"input_cost_per_query": EXA_REGISTERED_PRICE,
|
||||
"tiered_pricing": [{"input_cost_per_query": REGISTERED_TIER_PRICE, "max_results_range": [0, 100]}],
|
||||
},
|
||||
REGISTERED_TIER_PRICE,
|
||||
id="flat-and-tiers",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_a_price_registered_over_catalog_tiers_bills_every_tiered_search_entry(
|
||||
registration: Mapping[str, object], expected: float
|
||||
) -> None:
|
||||
catalog: Final = GetModelCostMap.loaded_model_cost_map()
|
||||
tiered: Final = tuple(
|
||||
key for key, entry in catalog.items() if entry.get("mode") == "search" and entry.get("tiered_pricing")
|
||||
)
|
||||
assert tiered
|
||||
litellm.register_model({key: dict(registration) for key in 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 tiered
|
||||
}
|
||||
|
||||
assert costs == {key: pytest.approx((expected, 0.0)) for key in tiered}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue