mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
refactor(exa): reuse pricing helpers and consolidate search tests
This commit is contained in:
shrey kharbanda
parent
a27503f67f
commit
2c9c4eccd3
6 changed files with 140 additions and 316 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue