mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fac0135d82
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
92 lines
3.4 KiB
Python
92 lines
3.4 KiB
Python
import datetime
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
|
|
import litellm
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
from litellm.search.cost_calculator import search_provider_cost_per_query
|
|
|
|
EXA_SEARCH_URL: Final = "https://api.exa.ai/search"
|
|
|
|
|
|
@pytest.fixture
|
|
def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
|
|
|
monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map())
|
|
|
|
|
|
def _proxy_style_logging_obj() -> Logging:
|
|
"""The proxy builds the logging object from the request body before the router adds the search tool's
|
|
litellm_params, so the per-tool price only reaches cost tracking if the search call itself records it."""
|
|
return Logging(
|
|
model="exa-search",
|
|
messages=[],
|
|
stream=False,
|
|
call_type="search",
|
|
start_time=datetime.datetime.now(),
|
|
litellm_call_id="search-tool-price",
|
|
function_id="f",
|
|
)
|
|
|
|
|
|
@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"}]})
|
|
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)
|
|
|
|
map_priced: 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 map_priced._hidden_params["response_cost"] == pytest.approx(map_price)
|
|
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")
|
|
def test_unpriced_fallback_search_tool_is_billed_at_its_own_price_not_the_failed_primarys(
|
|
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={"results": [{"title": "t", "url": "https://example.com"}]}),
|
|
)
|
|
)
|
|
map_price, _ = search_provider_cost_per_query(model="exa_ai/search", custom_llm_provider="exa_ai")
|
|
primary_price: Final = 0.05
|
|
assert map_price != primary_price
|
|
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(map_price)
|