From 2c9c4eccd312dba2ebadd1faaf5e4268d1a1a07b Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 28 Oct 2026 22:49:56 -0700 Subject: [PATCH] refactor(exa): reuse pricing helpers and consolidate search tests --- litellm/llms/exa_ai/search/cost_calculator.py | 3 +- litellm/search/cost_calculator.py | 23 +-- .../sdk/test_exa_search_envelope_sdk.py | 116 +++++++++---- tests/integration/sdk/test_exa_search_wire.py | 105 ------------ .../llms/exa_ai/search/test_transformation.py | 54 +----- tests/unit/search/test_main.py | 155 +++++------------- 6 files changed, 140 insertions(+), 316 deletions(-) delete mode 100644 tests/integration/sdk/test_exa_search_wire.py diff --git a/litellm/llms/exa_ai/search/cost_calculator.py b/litellm/llms/exa_ai/search/cost_calculator.py index 8477a913754..44b72f9fb0d 100644 --- a/litellm/llms/exa_ai/search/cost_calculator.py +++ b/litellm/llms/exa_ai/search/cost_calculator.py @@ -100,7 +100,6 @@ EXA_AI_INCLUDED_RESULTS: Final = 10 EXA_AI_ADDITIONAL_RESULT_COST: Final = 0.001 EXA_AI_SUMMARY_COST_PER_PAGE: Final = 0.001 EXA_AI_TYPES_WITHOUT_SUMMARIES: Final = frozenset({"lite"}) -_CONTENTS: Final = TypeAdapter(Mapping[str, object]) def _non_negative_int(value: object) -> int | None: @@ -124,7 +123,7 @@ def _requests_summaries(optional_params: Mapping[str, object]) -> bool: if optional_params.get("type") in EXA_AI_TYPES_WITHOUT_SUMMARIES: return False contents: Final = optional_params.get("contents") - return isinstance(contents, Mapping) and _CONTENTS.validate_python(contents).get("summary") not in (None, False) + return isinstance(contents, Mapping) and _ENTRY.validate_python(contents).get("summary") not in (None, False) def exa_ai_list_price(base_price: float, optional_params: Mapping[str, object]) -> float: diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index df2c527e638..870aea0347a 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -39,25 +39,14 @@ def _provider_usage( return None -def _generic_search_entry(model: str, custom_llm_provider: str | None) -> str: - generic_model: Final = f"{custom_llm_provider}/search" - return generic_model if model == f"{custom_llm_provider}/{generic_model}" else model - - 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 - - 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) + return _live_entry(model) != GetModelCostMap.loaded_model_cost_map().get(model) def _pricing_model( @@ -122,14 +111,12 @@ def search_provider_cost_per_query( return (input_cost, 0.0) pricing_model: Final = _pricing_model(model, custom_llm_provider, optional_params) + exa_model: Final = ( + pricing_model.removeprefix("exa_ai/") if pricing_model == "exa_ai/exa_ai/search" else pricing_model + ) params: Final = optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS exa_price_fields: Final[frozenset[str]] = ( - 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" - else frozenset() + exa_ai_price_fields(exa_model, _live_entry(exa_model)) if custom_llm_provider == "exa_ai" else frozenset() ) reported_cost: Final = ( exa_ai_reported_cost(params, response) if custom_llm_provider == "exa_ai" and not exa_price_fields else None diff --git a/tests/integration/sdk/test_exa_search_envelope_sdk.py b/tests/integration/sdk/test_exa_search_envelope_sdk.py index 3ab2a1a4969..33bc0402399 100644 --- a/tests/integration/sdk/test_exa_search_envelope_sdk.py +++ b/tests/integration/sdk/test_exa_search_envelope_sdk.py @@ -1,5 +1,8 @@ +import asyncio import json -import uuid +import threading +from collections.abc import Mapping +from datetime import datetime from typing import Final import pytest @@ -7,15 +10,18 @@ from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.search.transformation import SearchResponse pytestmark: Final = pytest.mark.timeout(60) +_REPORTED_COST: Final = 0.005 _ENVELOPE: Final[dict[str, JsonValue]] = { "requestId": "exa-request-sdk", + "output": {"content": {"answer": "verified"}, "grounding": []}, "resolvedSearchType": "neural", "searchTime": 211.0, - "costDollars": {"total": 0.005, "search": {"neural": 0.005}}, + "costDollars": {"total": _REPORTED_COST, "search": {"neural": _REPORTED_COST}}, } _RESULT: Final[dict[str, JsonValue]] = { "id": "https://result.invalid/sdk", @@ -23,6 +29,7 @@ _RESULT: Final[dict[str, JsonValue]] = { "url": "https://result.invalid/sdk", "publishedDate": "2026-02-01T00:00:00.000Z", "score": 0.73, + "author": None, "highlights": ["First line.", "Second line."], "highlightScores": [0.61, 0.42], } @@ -39,6 +46,7 @@ _EXPECTED: Final[dict[str, JsonValue]] = { "highlights": ["First line.", "Second line."], "highlight_scores": [0.61, 0.42], "score": 0.73, + "author": None, } ], "provider_metadata": _ENVELOPE, @@ -50,34 +58,84 @@ def _respond(request: Request) -> Reply: return Reply(body=json.dumps({**_ENVELOPE, "results": [_RESULT]}).encode()) -def test_sync_exa_search_returns_highlights_scores_and_the_response_envelope() -> None: - with wire_server(_respond) as wire: - response: Final = litellm.search( - query=f"sdk sync {uuid.uuid4().hex}", - search_provider="exa_ai", - api_key="synthetic-exa-key", - api_base=f"{wire.url}/exa", - max_results=1, - contents={"highlights": {"numSentences": 2}}, - ) - assert len(wire.drain()) == 1 - - assert isinstance(response, SearchResponse), type(response) - assert response.model_dump(mode="json") == _EXPECTED +PARAMS: Final = { + "contents": {"highlights": {"dynamic": True, "verbosity": "low"}}, + "outputSchema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + "systemPrompt": "Use evidence", + "type": "auto", +} @pytest.mark.asyncio -async def test_async_exa_search_returns_highlights_scores_and_the_response_envelope() -> None: - with wire_server(_respond) as wire: - response: Final = await litellm.asearch( - query=f"sdk async {uuid.uuid4().hex}", - search_provider="exa_ai", - api_key="synthetic-exa-key", - api_base=f"{wire.url}/exa", - max_results=1, - contents={"highlights": {"numSentences": 2}}, - ) - assert len(wire.drain()) == 1 +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("schema", [False, True]) +@pytest.mark.parametrize("price_override", [None, 0.0, 0.123]) +@pytest.mark.parametrize("global_callback", [False, True]) +async def test_exa_search_request_response_and_logged_cost_cross_real_http( + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + schema: bool, + price_override: float | None, + global_callback: bool, +) -> None: + class Recorder(CustomLogger): + def __init__(self) -> None: + self.finished: Final = threading.Event() + self.records: tuple[tuple[Mapping[str, object], object], ...] = () - assert isinstance(response, SearchResponse), type(response) - assert response.model_dump(mode="json") == _EXPECTED + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, 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: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.log_success_event(kwargs, response_obj, start_time, end_time) + + recorder: Final = Recorder() + if global_callback: + monkeypatch.setattr(litellm, "callbacks", [recorder]) + provider_params: Final = { + **{key: value for key, value in PARAMS.items() if schema or key != "outputSchema"}, + "contents": PARAMS["contents"] if schema else {"highlights": {"numSentences": 2}}, + } + + def respond(request: Request) -> Reply: + assert request.headers["x-api-key"] == "test-key" + assert request.headers["exa-beta"] == "dynamic-highlights-2026-08-28" + assert json.loads(request.body) == {"query": "question", "numResults": 1, **provider_params} + return _respond(request) + + with wire_server(respond) as server: + params: Final = { + "query": "question", + "search_provider": "exa_ai", + "api_base": f"{server.url}/exa", + "max_results": 1, + "api_key": "test-key", + "extra_headers": {"Exa-Beta": "dynamic-highlights-2026-08-28"}, + **({"input_cost_per_query": price_override} if price_override is not None else {}), + **( + {} + if global_callback + else {"success_callback": [recorder.async_log_success_event if asynchronous else recorder]} + ), + **provider_params, + } + response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) + assert len(server.drain()) == 1 + assert isinstance(response, SearchResponse) + assert response.model_dump(mode="json") == _EXPECTED + assert await asyncio.to_thread(recorder.finished.wait, 5) + assert len(recorder.records) == 1 + details, logged = recorder.records[0] + assert isinstance(logged, SearchResponse) + assert logged.model_dump(mode="json") == _EXPECTED + assert isinstance(details["standard_logging_object"], Mapping) + expected_cost: Final = _REPORTED_COST if price_override is None else price_override + assert response.hidden_params["response_cost"] == expected_cost + assert details["response_cost"] == expected_cost + assert details["standard_logging_object"]["response_cost"] == expected_cost + assert details["stream"] is False diff --git a/tests/integration/sdk/test_exa_search_wire.py b/tests/integration/sdk/test_exa_search_wire.py deleted file mode 100644 index c27f0cfc548..00000000000 --- a/tests/integration/sdk/test_exa_search_wire.py +++ /dev/null @@ -1,105 +0,0 @@ -import json -import threading -from typing import Final - -import pytest - -import litellm -from tests.integration._support.wire import Reply, Request, wire_server - -PAYLOAD: Final = { - "results": [{"title": "Source", "url": "https://example.com", "highlights": ["evidence"], "author": None}], - "output": {"content": {"answer": "verified"}, "grounding": []}, - "requestId": "wire-request", - "searchTime": 14, - "costDollars": {"total": 0.01}, -} -PARAMS: Final = { - "contents": {"highlights": {"dynamic": True, "verbosity": "low"}}, - "outputSchema": {"type": "object", "properties": {"answer": {"type": "string"}}}, - "systemPrompt": "Use evidence", - "type": "auto", -} - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True]) -async def test_exa_search_request_and_full_response_cross_real_http(asynchronous: bool) -> None: - def respond(request: Request) -> Reply: - assert request.target == "/search" - assert request.headers["x-api-key"] == "test-key" - assert request.headers["exa-beta"] == "dynamic-highlights-2026-08-28" - assert json.loads(request.body) == {"query": "question", **PARAMS} - return Reply(body=json.dumps(PAYLOAD).encode()) - - with wire_server(respond) as server: - params: Final = { - "query": "question", - "search_provider": "exa_ai", - "api_base": server.url, - "api_key": "test-key", - "extra_headers": {"Exa-Beta": "dynamic-highlights-2026-08-28"}, - **PARAMS, - } - response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) - result: Final = json.loads(response.model_dump_json()) - assert result["provider_metadata"]["output"] == PAYLOAD["output"] - assert result["provider_metadata"]["requestId"] == PAYLOAD["requestId"] - assert result["provider_metadata"]["costDollars"] == PAYLOAD["costDollars"] - assert result["results"][0]["highlights"] == ["evidence"] - assert result["results"][0]["author"] is None - assert result["results"][0]["snippet"] == "evidence" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.parametrize("schema", [False, True]) -@pytest.mark.parametrize("price_override", [None, 0.0, 0.123]) -async def test_search_logs_one_completed_response_and_cost( - asynchronous: bool, schema: bool, price_override: float | None -) -> None: - import asyncio - from collections.abc import Mapping - - from litellm.integrations.custom_logger import CustomLogger - - class Recorder(CustomLogger): - def __init__(self) -> None: - self.finished = threading.Event() - self.records = () - - def log_success_event(self, kwargs, response_obj, start_time, end_time): - self.records = (*self.records, (kwargs, response_obj)) - self.finished.set() - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.log_success_event(kwargs, response_obj, start_time, end_time) - - recorder: Final = Recorder() - - def respond(request: Request) -> Reply: - assert "input_cost_per_query" not in json.loads(request.body) - return Reply(body=json.dumps(PAYLOAD).encode()) - - with wire_server(respond) as server: - params: Final = { - "query": "question", - "search_provider": "exa_ai", - "api_base": server.url, - "api_key": "test-key", - **({"input_cost_per_query": price_override} if price_override is not None else {}), - "success_callback": [recorder.async_log_success_event if asynchronous else recorder], - **{key: value for key, value in PARAMS.items() if schema or key != "outputSchema"}, - } - response: Final = await litellm.asearch(**params) if asynchronous else litellm.search(**params) - 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 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 - assert details["standard_logging_object"]["response_cost"] == expected_cost - assert details["stream"] is False diff --git a/tests/unit/llms/exa_ai/search/test_transformation.py b/tests/unit/llms/exa_ai/search/test_transformation.py index 389941dba4d..f5f88a6ee69 100644 --- a/tests/unit/llms/exa_ai/search/test_transformation.py +++ b/tests/unit/llms/exa_ai/search/test_transformation.py @@ -1,18 +1,13 @@ -import asyncio import json -import uuid -from collections.abc import Mapping from typing import Final from unittest.mock import Mock import httpx import pytest import respx -from pydantic import TypeAdapter, ValidationError +from pydantic import ValidationError import litellm -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_COST_DOLLARS_PARAM from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig @@ -20,7 +15,6 @@ 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.123456 -_LOGGED: Final = TypeAdapter(Mapping[str, object]) @pytest.fixture @@ -385,53 +379,21 @@ def test_a_caller_cannot_set_the_cost_exa_reports(respx_mock: respx.MockRouter) @pytest.mark.usefixtures("exa_api_key") -def test_costing_an_exa_reply_directly_returns_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None: +@pytest.mark.parametrize("serialized", (False, True)) +def test_costing_an_exa_reply_directly_returns_the_cost_exa_reports( + respx_mock: respx.MockRouter, serialized: bool +) -> None: respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_COST}}) response: Final = litellm.search(query="asyncio TaskGroup", search_provider="exa_ai", max_results=25) assert _catalog_price(25) != EXA_REPORTED_COST + costed_response: Final = response.model_dump() if serialized else response completion_cost: Final = litellm.completion_cost( - completion_response=response, model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search" + completion_response=costed_response, model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search" ) per_token: Final = litellm.cost_per_token( - model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search", response=response + model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search", response=costed_response ) assert completion_cost == pytest.approx(EXA_REPORTED_COST) assert per_token == pytest.approx((EXA_REPORTED_COST, 0.0)) - - -class _SpendCapture(CustomLogger): - def __init__(self, call_id: str) -> None: - super().__init__() - self.call_id: Final = call_id - self.costs: Final[list[object]] = [] # mutable-ok: the logging worker appends one entry per logged call - - async def async_log_success_event( - self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object - ) -> None: - if kwargs.get("litellm_call_id") == self.call_id: - self.costs.append(_LOGGED.validate_python(kwargs.get("standard_logging_object") or {}).get("response_cost")) - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("exa_api_key") -async def test_async_search_logs_the_cost_exa_reports_as_spend( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - litellm.in_memory_llm_clients_cache.flush_cache() - GLOBAL_LOGGING_WORKER.start() - capture: Final = _SpendCapture(call_id=f"exa-{uuid.uuid4()}") - monkeypatch.setattr(litellm, "callbacks", [capture]) - respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_COST}}) - assert _catalog_price(25) != EXA_REPORTED_COST - - response: Final = await litellm.asearch( - query="asyncio TaskGroup", search_provider="exa_ai", max_results=25, litellm_call_id=capture.call_id - ) - await asyncio.sleep(0) - await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) - - assert response.hidden_params["response_cost"] == EXA_REPORTED_COST - assert capture.costs == [EXA_REPORTED_COST] diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index 5ae3b25fcf8..e1f1330fa7e 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -1,7 +1,6 @@ import datetime import json -from collections.abc import Callable, Iterator, Mapping -from dataclasses import dataclass +from collections.abc import Iterator, Mapping from typing import Final import httpx @@ -64,8 +63,13 @@ def _proxy_style_logging_obj() -> Logging: @pytest.mark.usefixtures("_local_model_cost_map") -def test_search_bills_search_tool_input_cost_per_query(respx_mock: respx.MockRouter) -> None: - respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [{"title": "t", "url": "https://example.com"}]}) +@pytest.mark.parametrize("reported_cost", (None, EXA_REPORTED_COST)) +def test_search_bills_search_tool_input_cost_per_query( + respx_mock: respx.MockRouter, reported_cost: float | None +) -> None: + respx_mock.post(EXA_SEARCH_URL).respond( + json={"results": [], **({"costDollars": {"total": reported_cost}} if reported_cost is not None else {})} + ) map_price, _ = search_provider_cost_per_query(model="exa_ai/search", custom_llm_provider="exa_ai") configured_price: Final = 0.05 assert map_price not in (0.0, configured_price) @@ -84,19 +88,29 @@ def test_search_bills_search_tool_input_cost_per_query(respx_mock: respx.MockRou litellm_logging_obj=_proxy_style_logging_obj(), ) - assert map_priced._hidden_params["response_cost"] == pytest.approx(map_price) + assert map_priced._hidden_params["response_cost"] == pytest.approx( + map_price if reported_cost is None else reported_cost + ) assert tool_priced._hidden_params["response_cost"] == pytest.approx(configured_price) assert respx_mock.calls.last.request.content.find(b"input_cost_per_query") == -1 @pytest.mark.usefixtures("_local_model_cost_map") +@pytest.mark.parametrize("reported_cost", (None, EXA_REPORTED_COST)) def test_unpriced_fallback_search_tool_is_billed_at_its_own_price_not_the_failed_primarys( respx_mock: respx.MockRouter, + reported_cost: float | None, ) -> None: respx_mock.post(EXA_SEARCH_URL).mock( side_effect=( httpx.Response(500, json={"error": "primary down"}), - httpx.Response(200, json={"results": [{"title": "t", "url": "https://example.com"}]}), + httpx.Response( + 200, + json={ + "results": [], + **({"costDollars": {"total": reported_cost}} if reported_cost is not None else {}), + }, + ), ) ) map_price, _ = search_provider_cost_per_query(model="exa_ai/search", custom_llm_provider="exa_ai") @@ -119,103 +133,22 @@ def test_unpriced_fallback_search_tool_is_billed_at_its_own_price_not_the_failed litellm_logging_obj=shared_by_every_fallback_attempt, ) - assert fallback._hidden_params["response_cost"] == pytest.approx(map_price) - - -@pytest.mark.usefixtures("_local_model_cost_map") -def test_a_search_tool_price_overrides_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None: - respx_mock.post(EXA_SEARCH_URL).respond(json=EXA_REPLY_WITH_COST) - configured_price: Final = 0.05 - assert configured_price != EXA_REPORTED_COST - - reported: Final = litellm.search( - query="litellm release notes", - search_provider="exa_ai", - api_key="test-exa-key", - litellm_logging_obj=_proxy_style_logging_obj(), - ) - tool_priced: Final = litellm.search( - query="litellm release notes", - search_provider="exa_ai", - api_key="test-exa-key", - input_cost_per_query=configured_price, - litellm_logging_obj=_proxy_style_logging_obj(), - ) - - assert reported.hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_COST) - assert tool_priced.hidden_params["response_cost"] == pytest.approx(configured_price) - - -@pytest.mark.usefixtures("_local_model_cost_map") -def test_unpriced_fallback_search_tool_bills_the_cost_exa_reports_not_the_failed_primarys_price( - respx_mock: respx.MockRouter, -) -> None: - respx_mock.post(EXA_SEARCH_URL).mock( - side_effect=(httpx.Response(500, json={"error": "primary down"}), httpx.Response(200, json=EXA_REPLY_WITH_COST)) - ) - primary_price: Final = 0.05 - shared_by_every_fallback_attempt: Final = _proxy_style_logging_obj() - - with pytest.raises(litellm.InternalServerError): - litellm.search( - query="litellm release notes", - search_provider="exa_ai", - api_key="test-exa-key", - input_cost_per_query=primary_price, - litellm_logging_obj=shared_by_every_fallback_attempt, - ) - fallback: Final = litellm.search( - query="litellm release notes", - search_provider="exa_ai", - api_key="test-exa-key", - litellm_logging_obj=shared_by_every_fallback_attempt, - ) - - assert fallback.hidden_params["response_cost"] == pytest.approx(EXA_REPORTED_COST) - - -def _single_tier(price: float) -> list[dict[str, object]]: - return [{"input_cost_per_query": price, "max_results_range": [0, 100]}] - - -def _register_exa_price(monkeypatch: pytest.MonkeyPatch, price: float) -> None: - litellm.register_model( - {"exa_ai/search": {"input_cost_per_query": price, "litellm_provider": "exa_ai", "mode": "search"}} - ) - - -def _register_exa_tiers(monkeypatch: pytest.MonkeyPatch, price: float) -> None: - litellm.register_model( - {"exa_ai/search": {"tiered_pricing": _single_tier(price), "litellm_provider": "exa_ai", "mode": "search"}} - ) - - -def _serve_exa_price_in_cost_map(monkeypatch: pytest.MonkeyPatch, price: float) -> None: - monkeypatch.setitem( - litellm.model_cost, "exa_ai/search", {**litellm.model_cost["exa_ai/search"], "input_cost_per_query": price} - ) - - -def _serve_exa_tiers_in_cost_map(monkeypatch: pytest.MonkeyPatch, price: float) -> None: - monkeypatch.setitem( - litellm.model_cost, - "exa_ai/search", - {**litellm.model_cost["exa_ai/search"], "tiered_pricing": _single_tier(price)}, + assert fallback._hidden_params["response_cost"] == pytest.approx( + map_price if reported_cost is None else reported_cost ) @pytest.mark.usefixtures("_local_model_cost_map") -@pytest.mark.parametrize( - "set_exa_price", - [_register_exa_price, _register_exa_tiers, _serve_exa_price_in_cost_map, _serve_exa_tiers_in_cost_map], -) +@pytest.mark.parametrize("registered", (False, True)) +@pytest.mark.parametrize("price_field", ("input_cost_per_query", "tiered_pricing")) @pytest.mark.parametrize( "exa_reply", [EXA_REPLY_WITH_COST, {"results": []}], ids=["exa-reports-a-cost", "exa-reports-no-cost"] ) def test_an_exa_price_someone_set_overrides_the_cost_exa_reports_and_litellms_tiers( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, - set_exa_price: Callable[[pytest.MonkeyPatch, float], None], + registered: bool, + price_field: str, exa_reply: Mapping[str, object], ) -> None: respx_mock.post(EXA_SEARCH_URL).respond(json=exa_reply) @@ -224,7 +157,17 @@ def test_an_exa_price_someone_set_overrides_the_cost_exa_reports_and_litellms_ti ) set_price: Final = 0.003 assert set_price not in (tier_price, EXA_REPORTED_COST) - set_exa_price(monkeypatch, set_price) + price: Final = ( + set_price + if price_field == "input_cost_per_query" + else [{"input_cost_per_query": set_price, "max_results_range": [0, 100]}] + ) + if registered: + litellm.register_model({"exa_ai/search": {price_field: price, "litellm_provider": "exa_ai", "mode": "search"}}) + else: + monkeypatch.setitem( + litellm.model_cost, "exa_ai/search", {**litellm.model_cost["exa_ai/search"], price_field: price} + ) response: Final = litellm.search( query="litellm release notes", search_provider="exa_ai", api_key="test-exa-key", max_results=25 @@ -456,19 +399,6 @@ 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: @@ -513,14 +443,6 @@ def test_exa_native_result_count_bills_the_same_operator_tier(respx_mock: respx. 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 @@ -556,11 +478,12 @@ def test_exa_catalog_source_preserves_receipts_and_generic_operator_prices( **({"input_cost_per_query": price} if price is not None else {}), }, } + respx_mock.get(source).respond(json=served) 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())), + litellm.get_model_cost_map(url=source, max_attempts=1), ) 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)