mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 0d291d7435 into 7f1b0a4633
This commit is contained in:
Kes
GitHub
commit
e2aebc2bce
13 changed files with 1279 additions and 150 deletions
|
|
@ -42,6 +42,7 @@ from litellm._logging import suppress_startup_output as suppress_startup_output
|
|||
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",
|
||||
|
|
|
|||
|
|
@ -198,6 +198,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)
|
||||
|
|
|
|||
|
|
@ -811,6 +811,7 @@ def cost_per_token(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
number_of_queries=number_of_queries or 1,
|
||||
optional_params=(getattr(response, "_hidden_params", None) if response else None),
|
||||
response=response,
|
||||
)
|
||||
elif (
|
||||
per_second_cost := _per_second_pricing_cost(
|
||||
|
|
@ -2160,6 +2161,7 @@ def completion_cost(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
number_of_queries=number_of_queries,
|
||||
optional_params=optional_params,
|
||||
response=cast(object, completion_response), # cast-ok: callee validates legacy dict fields
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
49
litellm/llms/exa_ai/search/cost_calculator.py
Normal file
49
litellm/llms/exa_ai/search/cost_calculator.py
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
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
|
||||
EXA_AI_TYPES_WITHOUT_SUMMARIES: Final = frozenset({"lite"})
|
||||
_CONTENTS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _non_negative_int(value: object) -> int | None:
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
|
||||
return None
|
||||
return value
|
||||
|
||||
|
||||
def _requested_results(optional_params: Mapping[str, object]) -> int:
|
||||
for key in ("max_results", "numResults"):
|
||||
requested = _non_negative_int(optional_params.get(key))
|
||||
if requested is not None:
|
||||
return requested
|
||||
return EXA_AI_INCLUDED_RESULTS
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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, a per-result charge above the results that price includes,
|
||||
and a per-page charge for each requested result when summaries are requested on a type that offers them.
|
||||
"""
|
||||
results: Final = _requested_results(optional_params)
|
||||
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
|
||||
)
|
||||
|
|
@ -1,11 +1,11 @@
|
|||
"""
|
||||
Calls Exa AI's /search endpoint to search the web.
|
||||
|
||||
Exa AI API Reference: https://docs.exa.ai/reference/search
|
||||
Exa AI API Reference: https://exa.ai/docs/reference/search
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Final, TypedDict
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
|
@ -17,41 +17,16 @@ from litellm.llms.base_llm.search.transformation import (
|
|||
SearchResponse,
|
||||
SearchResult,
|
||||
)
|
||||
from litellm.llms.exa_ai.search.types import ExaAISearchRequest
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
__all__ = ["ExaAISearchConfig", "ExaAISearchRequest"]
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
_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"))
|
||||
|
||||
|
||||
class _ExaAISearchRequestRequired(TypedDict):
|
||||
"""Required fields for Exa AI Search API request."""
|
||||
|
||||
query: str # Required - search query
|
||||
|
||||
|
||||
class ExaAISearchRequest(_ExaAISearchRequestRequired, total=False):
|
||||
"""
|
||||
Exa AI Search API request format.
|
||||
Based on: https://docs.exa.ai/reference/search
|
||||
"""
|
||||
|
||||
type: str # Optional - search type ('keyword', 'neural', 'fast', 'auto'), default 'auto'
|
||||
category: str # Optional - data category ('company', 'research paper', 'news', 'pdf', 'github', 'tweet', 'personal site', 'linkedin profile', 'financial report')
|
||||
userLocation: str # Optional - two-letter ISO country code
|
||||
numResults: int # Optional - number of results (max 100), default 10
|
||||
includeDomains: list[str] # Optional - list of domains to include
|
||||
excludeDomains: list[str] # Optional - list of domains to exclude
|
||||
startCrawlDate: str # Optional - crawl date filter (ISO 8601 format)
|
||||
endCrawlDate: str # Optional - crawl date filter (ISO 8601 format)
|
||||
startPublishedDate: str # Optional - published date filter (ISO 8601 format)
|
||||
endPublishedDate: str # Optional - published date filter (ISO 8601 format)
|
||||
includeText: list[str] # Optional - strings that must be present in webpage text
|
||||
excludeText: list[str] # Optional - strings that must not be present in webpage text
|
||||
context: bool | dict # Optional - format results for LLMs
|
||||
moderation: bool # Optional - enable content moderation, default false
|
||||
contents: dict # Optional - content retrieval options
|
||||
_RENAMED_RESULT_FIELDS: Final = frozenset({"text", "publishedDate", "highlightScores"})
|
||||
|
||||
|
||||
class ExaAISearchConfig(BaseSearchConfig):
|
||||
|
|
@ -63,33 +38,31 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
headers: Mapping[str, object],
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
**kwargs,
|
||||
) -> dict:
|
||||
**kwargs: object,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Validate environment and return headers.
|
||||
"""
|
||||
api_key = self.resolve_server_api_key(
|
||||
resolved_key: Final = self.resolve_server_api_key(
|
||||
caller_api_key=api_key,
|
||||
caller_api_base=api_base,
|
||||
key_env_vars=("EXA_API_KEY",),
|
||||
base_env_var="EXA_API_BASE",
|
||||
default_api_base=self.EXA_AI_API_BASE,
|
||||
)
|
||||
if not api_key:
|
||||
if not resolved_key:
|
||||
raise ValueError("EXA_API_KEY is not set. Set `EXA_API_KEY` environment variable.")
|
||||
headers["x-api-key"] = api_key
|
||||
headers["Content-Type"] = "application/json"
|
||||
return headers
|
||||
return {**headers, "x-api-key": resolved_key, "Content-Type": "application/json"}
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
optional_params: dict,
|
||||
data: dict | list[dict] | None = None,
|
||||
**kwargs,
|
||||
optional_params: Mapping[str, object],
|
||||
data: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
|
||||
**kwargs: object,
|
||||
) -> str:
|
||||
"""
|
||||
Get complete URL for Search endpoint.
|
||||
|
|
@ -99,9 +72,9 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
def transform_search_request(
|
||||
self,
|
||||
query: str | list[str],
|
||||
optional_params: dict,
|
||||
**kwargs,
|
||||
) -> dict:
|
||||
optional_params: Mapping[str, object],
|
||||
**kwargs: object,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Transform Search request to Exa AI API format.
|
||||
|
||||
|
|
@ -121,44 +94,32 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
Returns:
|
||||
Dict with typed request data following ExaAISearchRequest spec
|
||||
"""
|
||||
if isinstance(query, list):
|
||||
# Exa AI only supports single string queries, join with spaces
|
||||
query = " ".join(query)
|
||||
|
||||
request_data: Final[ExaAISearchRequest] = {
|
||||
"query": query,
|
||||
mapped: Final = {
|
||||
target: optional_params[source]
|
||||
for source, target in (
|
||||
("max_results", "numResults"),
|
||||
("search_domain_filter", "includeDomains"),
|
||||
("country", "userLocation"),
|
||||
)
|
||||
if source in optional_params
|
||||
}
|
||||
provider_params: Final = {
|
||||
key: value
|
||||
for key, value in optional_params.items()
|
||||
if key not in self.get_supported_perplexity_optional_params() and key not in mapped
|
||||
}
|
||||
return {
|
||||
"query": " ".join(query) if isinstance(query, list) else query,
|
||||
**mapped,
|
||||
**provider_params,
|
||||
**({"contents": {"text": True}} if "contents" not in provider_params else {}),
|
||||
}
|
||||
|
||||
# Transform Perplexity unified spec parameters to Exa format
|
||||
if "max_results" in optional_params:
|
||||
request_data["numResults"] = optional_params["max_results"]
|
||||
|
||||
if "search_domain_filter" in optional_params:
|
||||
request_data["includeDomains"] = optional_params["search_domain_filter"]
|
||||
|
||||
if "country" in optional_params:
|
||||
request_data["userLocation"] = optional_params["country"]
|
||||
|
||||
# Convert to dict before dynamic key assignments
|
||||
result_data: Final = dict(request_data)
|
||||
|
||||
# pass through all other parameters as-is
|
||||
for param, value in optional_params.items():
|
||||
if param not in self.get_supported_perplexity_optional_params() and param not in result_data:
|
||||
result_data[param] = value
|
||||
|
||||
# By default, request text content if not explicitly specified
|
||||
# Exa AI doesn't return content/text unless explicitly requested
|
||||
if "contents" not in result_data:
|
||||
result_data["contents"] = {"text": True}
|
||||
|
||||
return result_data
|
||||
|
||||
def transform_search_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
**kwargs,
|
||||
**kwargs: object,
|
||||
) -> SearchResponse:
|
||||
"""
|
||||
Transform Exa AI API response to LiteLLM unified SearchResponse format.
|
||||
|
|
@ -173,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
|
||||
|
|
@ -184,29 +146,31 @@ class ExaAISearchConfig(BaseSearchConfig):
|
|||
"""
|
||||
response_json: Final = _JSON_OBJECT.validate_python(raw_response.json())
|
||||
|
||||
# Transform results to SearchResult objects
|
||||
results: Final = []
|
||||
for raw_result in _JSON_VALUES.validate_python(response_json.get("results", [])):
|
||||
result = _JSON_OBJECT.validate_python(raw_result)
|
||||
search_result = SearchResult.model_validate(
|
||||
{
|
||||
"title": result.get("title", ""),
|
||||
"url": result.get("url", ""),
|
||||
"snippet": result.get("text")
|
||||
or "\n\n".join(_STRS.validate_python(result.get("highlights") or []))
|
||||
or result.get("summary")
|
||||
or "",
|
||||
"date": result.get("publishedDate"), # ISO 8601 datetime string
|
||||
"last_updated": None, # Exa AI doesn't provide last_updated in response
|
||||
**{field: result[key] for key, field in _RESULT_EXTRAS if key in result},
|
||||
}
|
||||
)
|
||||
results.append(search_result)
|
||||
|
||||
return SearchResponse.model_validate(
|
||||
{
|
||||
"results": results,
|
||||
"results": [
|
||||
self._transform_result(value)
|
||||
for value in _JSON_VALUES.validate_python(response_json.get("results") or [])
|
||||
],
|
||||
"object": "search",
|
||||
"provider_metadata": {key: value for key, value in response_json.items() if key != "results"},
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _transform_result(value: object) -> SearchResult:
|
||||
result: Final = _JSON_OBJECT.validate_python(value)
|
||||
return SearchResult.model_validate(
|
||||
{
|
||||
**{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")
|
||||
or "\n\n".join(_STRS.validate_python(result.get("highlights") or []))
|
||||
or result.get("summary")
|
||||
or "",
|
||||
"date": result.get("publishedDate"),
|
||||
"last_updated": None,
|
||||
**{field: result[key] for key, field in _RESULT_EXTRAS if key in result},
|
||||
}
|
||||
)
|
||||
|
|
|
|||
99
litellm/llms/exa_ai/search/types.py
Normal file
99
litellm/llms/exa_ai/search/types.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
"""Exa /search request contract, https://exa.ai/docs/reference/search."""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
ExaSearchMode: TypeAlias = Literal["lite", "instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"]
|
||||
ExaPageSection: TypeAlias = Literal["header", "navigation", "banner", "body", "sidebar", "footer", "metadata"]
|
||||
|
||||
|
||||
class ExaTextOptions(TypedDict, total=False):
|
||||
maxCharacters: ReadOnly[int | None]
|
||||
includeHtmlTags: ReadOnly[bool | None]
|
||||
verbosity: ReadOnly[Literal["compact", "standard", "full"] | None]
|
||||
includeSections: ReadOnly[Sequence[ExaPageSection] | None]
|
||||
excludeSections: ReadOnly[Sequence[ExaPageSection] | None]
|
||||
|
||||
|
||||
class ExaHighlightsOptions(TypedDict, total=False):
|
||||
"""numSentences and highlightsPerUrl are deprecated compatibility fields."""
|
||||
|
||||
query: ReadOnly[str | None]
|
||||
verbosity: ReadOnly[Literal["low", "medium", "high"] | None]
|
||||
dynamic: ReadOnly[bool | None]
|
||||
maxCharacters: ReadOnly[int | None]
|
||||
numSentences: ReadOnly[int | None]
|
||||
highlightsPerUrl: ReadOnly[int | None]
|
||||
|
||||
|
||||
class ExaSummaryOptions(TypedDict, total=False):
|
||||
query: ReadOnly[str | None]
|
||||
schema: ReadOnly[Mapping[str, object] | None]
|
||||
|
||||
|
||||
class ExaExtrasOptions(TypedDict, total=False):
|
||||
links: ReadOnly[int | None]
|
||||
imageLinks: ReadOnly[int | None]
|
||||
richImageLinks: ReadOnly[int | None]
|
||||
richLinks: ReadOnly[int | None]
|
||||
codeBlocks: ReadOnly[int | None]
|
||||
|
||||
|
||||
class ExaContextOptions(TypedDict, total=False):
|
||||
maxCharacters: ReadOnly[int]
|
||||
|
||||
|
||||
class ExaContentsOptions(TypedDict, total=False):
|
||||
"""context, livecrawl, and crawledBeforeDate are deprecated compatibility fields."""
|
||||
|
||||
text: ReadOnly[bool | ExaTextOptions | None]
|
||||
highlights: ReadOnly[bool | ExaHighlightsOptions | None]
|
||||
summary: ReadOnly[ExaSummaryOptions | None]
|
||||
extras: ReadOnly[ExaExtrasOptions | None]
|
||||
livecrawlTimeout: ReadOnly[int | None]
|
||||
maxAgeHours: ReadOnly[int | None]
|
||||
snapshotAsOf: ReadOnly[str | None]
|
||||
subpages: ReadOnly[int | None]
|
||||
subpageTarget: ReadOnly[str | Sequence[str] | None]
|
||||
context: ReadOnly[bool | ExaContextOptions | None]
|
||||
livecrawl: ReadOnly[Literal["never", "always", "fallback", "preferred"] | None]
|
||||
crawledBeforeDate: ReadOnly[str | None]
|
||||
|
||||
|
||||
class ExaCountryLocation(TypedDict):
|
||||
country: ReadOnly[str]
|
||||
|
||||
|
||||
class ExaCoordinateLocation(TypedDict, total=False):
|
||||
country: ReadOnly[str]
|
||||
latitude: Required[ReadOnly[float]]
|
||||
longitude: Required[ReadOnly[float]]
|
||||
|
||||
|
||||
class ExaAISearchRequest(TypedDict, total=False):
|
||||
"""Crawl-date filters, includeText, excludeText, and context are deprecated."""
|
||||
|
||||
query: Required[ReadOnly[str]]
|
||||
type: ReadOnly[ExaSearchMode | None]
|
||||
category: ReadOnly[str | None]
|
||||
additionalQueries: ReadOnly[Sequence[str] | None]
|
||||
userLocation: ReadOnly[str | ExaCountryLocation | ExaCoordinateLocation | None]
|
||||
numResults: ReadOnly[int | None]
|
||||
includeDomains: ReadOnly[Sequence[str] | None]
|
||||
excludeDomains: ReadOnly[Sequence[str] | None]
|
||||
startPublishedDate: ReadOnly[str | None]
|
||||
endPublishedDate: ReadOnly[str | None]
|
||||
moderation: ReadOnly[bool | None]
|
||||
contents: ReadOnly[ExaContentsOptions | None]
|
||||
outputSchema: ReadOnly[Mapping[str, object] | None]
|
||||
objective: ReadOnly[str | None]
|
||||
systemPrompt: ReadOnly[str | None]
|
||||
userIdentifier: ReadOnly[str | None]
|
||||
compliance: ReadOnly[Literal["hipaa"] | None]
|
||||
startCrawlDate: ReadOnly[str | None]
|
||||
endCrawlDate: ReadOnly[str | None]
|
||||
includeText: ReadOnly[Sequence[str] | None]
|
||||
excludeText: ReadOnly[Sequence[str] | None]
|
||||
context: ReadOnly[bool | ExaContextOptions | None]
|
||||
|
|
@ -25044,11 +25044,15 @@
|
|||
"output_vector_size": 2560
|
||||
},
|
||||
"exa_ai/search": {
|
||||
"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, 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, with its first tier at the base price. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search",
|
||||
"tiered_pricing": [
|
||||
{
|
||||
"input_cost_per_query": 0.005,
|
||||
"input_cost_per_query": 0.007,
|
||||
"max_results_range": [
|
||||
0,
|
||||
25
|
||||
|
|
@ -25063,6 +25067,46 @@
|
|||
}
|
||||
]
|
||||
},
|
||||
"exa_ai/search-lite": {
|
||||
"input_cost_per_query": 0.004,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type lite, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-22)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"exa_ai/search-instant": {
|
||||
"input_cost_per_query": 0.004,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type instant, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"exa_ai/search-deep": {
|
||||
"input_cost_per_query": 0.012,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type deep, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"exa_ai/search-deep-lite": {
|
||||
"input_cost_per_query": 0.012,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type deep-lite, covering up to 10 results. Source: https://exa.ai/docs/reference/pricing (2026-10-27)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"exa_ai/search-deep-reasoning": {
|
||||
"input_cost_per_query": 0.015,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type deep-reasoning, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"firecrawl/search": {
|
||||
"litellm_provider": "firecrawl",
|
||||
"mode": "search",
|
||||
|
|
|
|||
|
|
@ -2,23 +2,38 @@
|
|||
Cost calculation for search providers.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import functools
|
||||
import hashlib
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from importlib.metadata import PackageNotFoundError, distribution
|
||||
from types import MappingProxyType
|
||||
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 (
|
||||
AliasedCostEntry,
|
||||
GetModelCostMap,
|
||||
get_model_cost_map_source_info,
|
||||
)
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
|
||||
tuple[Mapping[str, object], ...]
|
||||
)
|
||||
MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
SEARCH_MODE_PARAM_BY_PROVIDER: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{"tavily": "search_depth", "linkup": "depth"}
|
||||
{"tavily": "search_depth", "linkup": "depth", "exa_ai": "type"}
|
||||
)
|
||||
SEARCH_PRICE_FIELDS: Final = ("input_cost_per_query", "tiered_pricing")
|
||||
FLAT_PRICE_ONLY: Final = frozenset({"input_cost_per_query"})
|
||||
BUNDLED_COST_MAP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
|
||||
|
||||
|
||||
def _provider_usage(
|
||||
|
|
@ -33,10 +48,95 @@ def _provider_usage(
|
|||
return None
|
||||
|
||||
|
||||
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 _field(value: object, key: str) -> object:
|
||||
return MAPPING_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 _changed_price_fields(value: object, reference: object) -> frozenset[str]:
|
||||
"""Price fields `value` holds a price in that differ from `reference`, so a price an entry leaves out sets
|
||||
nothing."""
|
||||
return frozenset(
|
||||
name
|
||||
for name in SEARCH_PRICE_FIELDS
|
||||
if _field(value, name) is not None and _field(value, name) != _field(reference, name)
|
||||
)
|
||||
|
||||
|
||||
def _runtime_price_fields(entry: str) -> frozenset[str]:
|
||||
"""Price fields `register_model` set on `entry`, even to the catalog's own value, or that differ from the catalog
|
||||
as loaded."""
|
||||
live: Final = _live_entry(entry)
|
||||
registered: Final[frozenset[str]] = (
|
||||
live.registered_fields.intersection(SEARCH_PRICE_FIELDS) if isinstance(live, AliasedCostEntry) else frozenset()
|
||||
)
|
||||
return registered | _changed_price_fields(live, GetModelCostMap.loaded_model_cost_map().get(entry))
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _shipped_search_prices() -> Mapping[str, Mapping[str, object]]:
|
||||
shipped: Final = MAPPING_ADAPTER.validate_json(GetModelCostMap.read_local_model_cost_map_bytes())
|
||||
return MappingProxyType(
|
||||
{
|
||||
key: MappingProxyType({name: _field(entry, name) for name in SEARCH_PRICE_FIELDS})
|
||||
for key, entry in shipped.items()
|
||||
if _field(entry, "mode") == "search"
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _bundled_cost_map_edited() -> bool:
|
||||
"""Whether the bundled map's bytes differ from the sha256 its installer recorded, as when an image overwrites the
|
||||
file. An install that recorded no hash for it, like an editable one, counts as unedited."""
|
||||
try:
|
||||
installed: Final = distribution("litellm").files or ()
|
||||
except PackageNotFoundError:
|
||||
return False
|
||||
recorded: Final = next((path.hash for path in installed if str(path) == BUNDLED_COST_MAP_PATH), None)
|
||||
if recorded is None or recorded.mode != "sha256":
|
||||
return False
|
||||
digest: Final = hashlib.sha256(GetModelCostMap.read_local_model_cost_map_bytes()).digest()
|
||||
return base64.urlsafe_b64encode(digest).rstrip(b"=").decode() != recorded.value
|
||||
|
||||
|
||||
def _own_cost_map_price_fields(entry: str) -> frozenset[str]:
|
||||
"""Price fields the operator's own cost map sets. A map served from their own URL counts only where it differs
|
||||
from the map LiteLLM ships, so a mirror of LiteLLM's map prices nothing, and a bundled map their image edited
|
||||
counts in full, since nothing is left to compare it with."""
|
||||
source: Final = get_model_cost_map_source_info()
|
||||
loaded: Final = GetModelCostMap.loaded_model_cost_map().get(entry)
|
||||
if source["source"] != "remote":
|
||||
return _changed_price_fields(loaded, None) if _bundled_cost_map_edited() else frozenset()
|
||||
if source["url"] == DEFAULT_MODEL_COST_MAP_URL:
|
||||
return frozenset()
|
||||
return _changed_price_fields(loaded, _shipped_search_prices().get(entry))
|
||||
|
||||
|
||||
def _operator_price_fields(entry: str) -> frozenset[str]:
|
||||
return _runtime_price_fields(entry) | _own_cost_map_price_fields(entry)
|
||||
|
||||
|
||||
def _operator_priced(entry: str) -> bool:
|
||||
"""Whether an operator set a price on `entry`, with `register_model` or their own cost map."""
|
||||
return bool(_operator_price_fields(entry))
|
||||
|
||||
|
||||
def _operator_flat_price_replaces_tiers(entry: str) -> bool:
|
||||
"""Whether the operator's latest price on `entry` is a flat price without tiers, so the tiers already on the entry
|
||||
stop applying. `register_model` runs after the cost map loads, so its prices outrank the operator's own map."""
|
||||
return (_runtime_price_fields(entry) or _own_cost_map_price_fields(entry)) == FLAT_PRICE_ONLY
|
||||
|
||||
|
||||
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(
|
||||
|
|
@ -45,8 +145,8 @@ def _pricing_model(
|
|||
optional_params: Mapping[str, object] | None,
|
||||
) -> str:
|
||||
"""`<provider>/search-<mode>` when `model` is the generic `<provider>/search` entry (bare or provider-prefixed) and
|
||||
the request's mode parameter names a priced entry, else `model`. A price registered at runtime on the generic entry
|
||||
keeps applying until the mode entry is registered too."""
|
||||
the request's mode parameter names a priced entry, else `model`. An operator price on the generic entry keeps
|
||||
applying until the operator prices the mode entry too."""
|
||||
mode_param: Final = SEARCH_MODE_PARAM_BY_PROVIDER.get(custom_llm_provider or "")
|
||||
generic_model: Final = f"{custom_llm_provider}/search"
|
||||
if mode_param is None or model not in (generic_model, f"{custom_llm_provider}/{generic_model}"):
|
||||
|
|
@ -54,7 +154,7 @@ def _pricing_model(
|
|||
mode: Final = (optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS).get(mode_param)
|
||||
mode_model: Final = f"{custom_llm_provider}/search-{mode}"
|
||||
priced: Final = mode_model in litellm.model_cost # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed
|
||||
if not priced or (_registered_at_runtime(generic_model) and not _registered_at_runtime(mode_model)):
|
||||
if not priced or (_operator_priced(generic_model) and not _operator_priced(mode_model)):
|
||||
return model
|
||||
return mode_model
|
||||
|
||||
|
|
@ -64,19 +164,23 @@ def search_provider_cost_per_query(
|
|||
custom_llm_provider: str | None = None,
|
||||
number_of_queries: int = 1,
|
||||
optional_params: Mapping[str, object] | None = None,
|
||||
response: object = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculate cost for search-only providers.
|
||||
|
||||
Returns (input_cost, output_cost) where input_cost = queries * cost_per_query
|
||||
Supports tiered pricing based on max_results parameter, and per-mode price entries selected by the
|
||||
provider's mode parameter (Tavily `search_depth`, Linkup `depth`).
|
||||
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 an operator priced the entry. A flat price
|
||||
an operator set last, over tiers LiteLLM or their own cost map ships, replaces those tiers.
|
||||
|
||||
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 Exa's reported provider_metadata.costDollars.total, when available
|
||||
|
||||
Returns:
|
||||
Tuple of (input_cost, output_cost) where output_cost is always 0.0
|
||||
|
|
@ -97,10 +201,28 @@ 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 _operator_priced(model_info["key"]):
|
||||
reported_total: Final = _exa_reported_total(response)
|
||||
if reported_total is not None:
|
||||
return (reported_total, 0.0)
|
||||
|
||||
# Check for tiered pricing (e.g., Exa AI based on max_results)
|
||||
from litellm.llms.exa_ai.search.cost_calculator import exa_ai_list_price
|
||||
|
||||
return (
|
||||
exa_ai_list_price(
|
||||
base_price=float(model_info.get("input_cost_per_query") or 0.0),
|
||||
optional_params=optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS,
|
||||
),
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 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 _operator_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
|
||||
|
||||
|
|
|
|||
|
|
@ -25044,11 +25044,15 @@
|
|||
"output_vector_size": 2560
|
||||
},
|
||||
"exa_ai/search": {
|
||||
"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, 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, with its first tier at the base price. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search",
|
||||
"tiered_pricing": [
|
||||
{
|
||||
"input_cost_per_query": 0.005,
|
||||
"input_cost_per_query": 0.007,
|
||||
"max_results_range": [
|
||||
0,
|
||||
25
|
||||
|
|
@ -25063,6 +25067,46 @@
|
|||
}
|
||||
]
|
||||
},
|
||||
"exa_ai/search-lite": {
|
||||
"input_cost_per_query": 0.004,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type lite, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-22)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"exa_ai/search-instant": {
|
||||
"input_cost_per_query": 0.004,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type instant, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"exa_ai/search-deep": {
|
||||
"input_cost_per_query": 0.012,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type deep, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"exa_ai/search-deep-lite": {
|
||||
"input_cost_per_query": 0.012,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type deep-lite, covering up to 10 results. Source: https://exa.ai/docs/reference/pricing (2026-10-27)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"exa_ai/search-deep-reasoning": {
|
||||
"input_cost_per_query": 0.015,
|
||||
"litellm_provider": "exa_ai",
|
||||
"metadata": {
|
||||
"notes": "Base price per /search request with type deep-reasoning, covering up to 10 results. Source: https://exa.ai/pricing (2026-10-15)"
|
||||
},
|
||||
"mode": "search"
|
||||
},
|
||||
"firecrawl/search": {
|
||||
"litellm_provider": "firecrawl",
|
||||
"mode": "search",
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
117
tests/integration/sdk/test_exa_search_wire.py
Normal file
117
tests/integration/sdk/test_exa_search_wire.py
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
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)
|
||||
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
|
||||
@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 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: tuple[tuple[Mapping[str, object], SearchResponse], ...] = ()
|
||||
|
||||
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: 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()
|
||||
|
||||
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() == 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
|
||||
assert details["standard_logging_object"]["response_cost"] == expected_cost
|
||||
assert details["stream"] is False
|
||||
|
|
@ -122,35 +122,17 @@ def test_transform_search_response_maps_every_exa_result_in_provider_order():
|
|||
)
|
||||
|
||||
assert response.object == "search"
|
||||
assert response.results == [
|
||||
SearchResult.model_validate(
|
||||
{
|
||||
"title": "First",
|
||||
"url": "https://example.com/first",
|
||||
"snippet": "First text",
|
||||
"date": "2024-01-05T00:00:00.000Z",
|
||||
"highlights": ["ignored when text is present"],
|
||||
"highlight_scores": [0.5],
|
||||
"score": 0.42,
|
||||
}
|
||||
),
|
||||
SearchResult.model_validate(
|
||||
{
|
||||
"title": "Second",
|
||||
"url": "https://example.com/second",
|
||||
"snippet": "one\n\ntwo",
|
||||
"highlights": ["one", "two"],
|
||||
}
|
||||
expected: Final = [
|
||||
SearchResult(
|
||||
title="First", url="https://example.com/first", snippet="First text", date="2024-01-05T00:00:00.000Z"
|
||||
),
|
||||
SearchResult(title="Second", url="https://example.com/second", snippet="one\n\ntwo"),
|
||||
SearchResult(title="Third", url="https://example.com/third", snippet="Summary"),
|
||||
]
|
||||
assert response.model_extra == {
|
||||
"provider_metadata": {
|
||||
"requestId": "b5947044c4b78efa9552a7c89b306d95",
|
||||
"resolvedSearchType": "neural",
|
||||
"costDollars": {"total": 0.005},
|
||||
}
|
||||
}
|
||||
|
||||
assert [result.model_dump(include=set(SearchResult.model_fields)) for result in response.results] == [
|
||||
result.model_dump() for result in expected
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -162,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(
|
||||
|
|
@ -177,7 +167,8 @@ def test_transform_search_response_attaches_only_the_exa_result_fields_it_return
|
|||
}
|
||||
)
|
||||
|
||||
assert [result.model_extra for result in response.results] == [{}, expected_extras]
|
||||
assert response.results[0].model_extra == {}
|
||||
assert response.results[1].model_extra == expected_extras
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -198,7 +189,7 @@ def test_transform_search_response_keeps_the_exa_envelope_as_provider_metadata(
|
|||
):
|
||||
response: Final = _transform(payload)
|
||||
|
||||
assert response.model_extra == {"provider_metadata": expected_metadata}
|
||||
assert response.model_extra["provider_metadata"] == expected_metadata
|
||||
assert response.model_dump(mode="json")["provider_metadata"] == expected_metadata
|
||||
|
||||
|
||||
|
|
@ -283,3 +274,369 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str) ->
|
|||
ExaAISearchConfig().get_complete_url(api_base=api_base, optional_params={})
|
||||
== "https://gateway.example.com/exa/search"
|
||||
)
|
||||
|
||||
|
||||
OUTPUT_SCHEMA: Final = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
||||
OUTPUT: Final = {
|
||||
"content": {"answer": "Grounded answer"},
|
||||
"grounding": [
|
||||
{"field": "answer", "citations": [{"url": "https://example.com", "title": "Source"}], "confidence": "high"}
|
||||
],
|
||||
}
|
||||
FULL_RESULT: Final = {
|
||||
"title": "Source",
|
||||
"url": "https://example.com",
|
||||
"publishedDate": "2026-01-01",
|
||||
"author": None,
|
||||
"id": "document-id",
|
||||
"image": "https://example.com/image.png",
|
||||
"favicon": "https://example.com/icon.png",
|
||||
"text": "",
|
||||
"highlights": ["relevant passage", "another passage"],
|
||||
"highlightScores": [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]},
|
||||
}
|
||||
FULL_RESPONSE: Final = {
|
||||
"results": [FULL_RESULT],
|
||||
"output": OUTPUT,
|
||||
"requestId": "request-id",
|
||||
"costDollars": {"total": 0.012},
|
||||
"searchTime": 24,
|
||||
"resolvedSearchType": "auto",
|
||||
"context": "",
|
||||
"futureMetadata": [],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["instant", "fast", "auto", "deep-lite", "deep", "deep-reasoning"])
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.usefixtures("exa_api_key")
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_search_preserves_native_contract(respx_mock: respx.MockRouter, mode: str, asynchronous: bool):
|
||||
native: Final = {
|
||||
"type": mode,
|
||||
"additionalQueries": ["related question"],
|
||||
"outputSchema": OUTPUT_SCHEMA,
|
||||
"systemPrompt": "Use grounded sources",
|
||||
"objective": "Research",
|
||||
"userIdentifier": "opaque-id",
|
||||
"userLocation": {"latitude": 40.7, "longitude": -74.0, "country": "US"},
|
||||
"includeDomains": ["example.com"],
|
||||
"excludeDomains": ["other.example"],
|
||||
"startPublishedDate": "2026-01-01T00:00:00Z",
|
||||
"endPublishedDate": "2026-02-01T00:00:00Z",
|
||||
"moderation": False,
|
||||
"numResults": 3,
|
||||
"category": "publication",
|
||||
"compliance": "hipaa",
|
||||
"contents": {
|
||||
"highlights": {"query": "evidence", "dynamic": True, "verbosity": "medium"},
|
||||
"extras": {"links": 2, "imageLinks": 2, "richLinks": 2, "richImageLinks": 2, "codeBlocks": 2},
|
||||
"maxAgeHours": -1,
|
||||
"livecrawlTimeout": 5000,
|
||||
"subpages": 1,
|
||||
"subpageTarget": ["about"],
|
||||
},
|
||||
}
|
||||
route: Final = respx_mock.post(EXA_SEARCH_URL).respond(json=FULL_RESPONSE)
|
||||
headers: Final = {"Exa-Beta": "dynamic-highlights-2026-08-28"}
|
||||
result: Final = (
|
||||
await litellm.asearch(query="question", search_provider="exa_ai", extra_headers=headers, **native)
|
||||
if asynchronous
|
||||
else litellm.search(query="question", search_provider="exa_ai", extra_headers=headers, **native)
|
||||
)
|
||||
assert json.loads(route.calls.last.request.content) == {"query": "question", **native}
|
||||
assert route.calls.last.request.headers["exa-beta"] == headers["Exa-Beta"]
|
||||
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(
|
||||
"contents",
|
||||
[
|
||||
None,
|
||||
{},
|
||||
{"text": False},
|
||||
{"highlights": True},
|
||||
{
|
||||
"text": {
|
||||
"maxCharacters": 800,
|
||||
"includeHtmlTags": True,
|
||||
"verbosity": "full",
|
||||
"includeSections": ["body"],
|
||||
"excludeSections": ["footer"],
|
||||
}
|
||||
},
|
||||
{"summary": {"query": "summarize", "schema": OUTPUT_SCHEMA}},
|
||||
{"highlights": {"maxCharacters": 700, "query": "evidence"}, "snapshotAsOf": "2026-01-01"},
|
||||
{"text": True, "livecrawl": "never", "crawledBeforeDate": "2026-01-01", "context": {"maxCharacters": 200}},
|
||||
],
|
||||
)
|
||||
@pytest.mark.usefixtures("exa_api_key")
|
||||
def test_contents_are_sent_without_injecting_text(respx_mock: respx.MockRouter, contents: object):
|
||||
route: Final = respx_mock.post(EXA_SEARCH_URL).respond(json={"results": []})
|
||||
litellm.search(query="question", search_provider="exa_ai", contents=contents)
|
||||
assert json.loads(route.calls.last.request.content) == {"query": "question", "contents": contents}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"output", [None, {}, {"content": "", "grounding": []}, {"content": {}, "grounding": []}, OUTPUT]
|
||||
)
|
||||
def test_output_preserves_present_empty_and_null_values(output: object):
|
||||
response: Final = _transform({"results": [], "output": output})
|
||||
assert response.model_dump() == {"results": [], "object": "search", "provider_metadata": {"output": output}}
|
||||
|
||||
|
||||
def test_absent_output_stays_absent():
|
||||
assert _transform({"results": []}).model_dump() == {"results": [], "object": "search", "provider_metadata": {}}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_httpx_for_respx(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
|
||||
|
||||
def test_legacy_parameters_and_unified_precedence_are_preserved():
|
||||
legacy: Final = {
|
||||
"max_results": 4,
|
||||
"numResults": 9,
|
||||
"country": "US",
|
||||
"userLocation": "GB",
|
||||
"search_domain_filter": ["example.com"],
|
||||
"includeDomains": ["other.example"],
|
||||
"startCrawlDate": "2020-01-01",
|
||||
"endCrawlDate": "2021-01-01",
|
||||
"includeText": ["one"],
|
||||
"excludeText": ["two"],
|
||||
"context": True,
|
||||
"max_tokens_per_page": 100,
|
||||
}
|
||||
before: Final = json.dumps(legacy, sort_keys=True)
|
||||
request: Final = ExaAISearchConfig().transform_search_request(["first", "second"], legacy)
|
||||
assert request == {
|
||||
"query": "first second",
|
||||
"numResults": 4,
|
||||
"userLocation": "US",
|
||||
"includeDomains": ["example.com"],
|
||||
"startCrawlDate": "2020-01-01",
|
||||
"endCrawlDate": "2021-01-01",
|
||||
"includeText": ["one"],
|
||||
"excludeText": ["two"],
|
||||
"context": True,
|
||||
"contents": {"text": True},
|
||||
}
|
||||
assert json.dumps(legacy, sort_keys=True) == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("total", [0, 0.012, 2])
|
||||
@pytest.mark.parametrize("serialized", [False, True])
|
||||
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": 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,
|
||||
0.0,
|
||||
)
|
||||
assert (
|
||||
litellm.completion_cost(
|
||||
completion_response=payload,
|
||||
model="exa_ai/search",
|
||||
custom_llm_provider="exa_ai",
|
||||
call_type="search",
|
||||
optional_params={"query": ["one", "two"], "costDollars": {"total": 999}},
|
||||
)
|
||||
== total
|
||||
)
|
||||
assert litellm.cost_per_token(
|
||||
model="exa_ai/search", custom_llm_provider="exa_ai", call_type="search", response=payload
|
||||
) == (total, 0.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"reported",
|
||||
[
|
||||
None,
|
||||
{},
|
||||
{"total": None},
|
||||
{"total": -1},
|
||||
{"total": True},
|
||||
{"total": "0.01"},
|
||||
{"total": float("nan")},
|
||||
{"total": float("inf")},
|
||||
[],
|
||||
],
|
||||
)
|
||||
def test_exa_missing_or_invalid_total_uses_estimate(reported: object) -> None:
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
||||
estimate: Final = search_provider_cost_per_query("exa_ai/search", "exa_ai")
|
||||
assert (
|
||||
search_provider_cost_per_query(
|
||||
"exa_ai/search",
|
||||
"exa_ai",
|
||||
response={"costDollars": {"total": 999}, "provider_metadata": {"costDollars": reported}},
|
||||
optional_params={"costDollars": {"total": 999}},
|
||||
)
|
||||
== estimate
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
def test_exa_estimate_adds_each_result_above_the_included_ten() -> None:
|
||||
from litellm.llms.exa_ai.search.cost_calculator import EXA_AI_ADDITIONAL_RESULT_COST
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
||||
base, _ = search_provider_cost_per_query("exa_ai/search", "exa_ai")
|
||||
|
||||
assert base > 0.0
|
||||
assert search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params={"max_results": 5}) == (base, 0.0)
|
||||
assert search_provider_cost_per_query(
|
||||
"exa_ai/search", "exa_ai", optional_params={"max_results": 25}
|
||||
) == pytest.approx((base + 15 * EXA_AI_ADDITIONAL_RESULT_COST, 0.0))
|
||||
assert search_provider_cost_per_query(
|
||||
"exa_ai/search", "exa_ai", optional_params={"numResults": 12}
|
||||
) == 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_adds_no_summary_charge_on_lite() -> None:
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
||||
lite: Final = {"type": "lite", "max_results": 12}
|
||||
# https://exa.ai/docs/reference/pricing (2026-10-28) lists no AI page summary price for Lite
|
||||
assert search_provider_cost_per_query(
|
||||
"exa_ai/search", "exa_ai", optional_params={**lite, "contents": {"summary": True}}
|
||||
) == search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params=lite)
|
||||
|
||||
|
||||
@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
|
||||
|
||||
assert search_provider_cost_per_query(
|
||||
"exa_ai/search", "exa_ai", number_of_queries=3
|
||||
) == search_provider_cost_per_query("exa_ai/search", "exa_ai")
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
@pytest.mark.parametrize(
|
||||
("search_type", "pricing_model"),
|
||||
[
|
||||
("lite", "exa_ai/search-lite"),
|
||||
("instant", "exa_ai/search-instant"),
|
||||
("auto", "exa_ai/search"),
|
||||
("fast", "exa_ai/search"),
|
||||
("deep-lite", "exa_ai/search-deep-lite"),
|
||||
("deep", "exa_ai/search-deep"),
|
||||
("deep-reasoning", "exa_ai/search-deep-reasoning"),
|
||||
],
|
||||
)
|
||||
def test_exa_estimate_uses_the_price_of_its_search_type(search_type: str, pricing_model: str) -> None:
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
||||
model_info: Final = litellm.get_model_info(model=pricing_model, custom_llm_provider="exa_ai")
|
||||
|
||||
assert search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params={"type": search_type}) == (
|
||||
model_info["input_cost_per_query"],
|
||||
0.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("native_count", [1, 100])
|
||||
def test_exa_estimate_uses_the_result_count_sent_to_provider(native_count: int) -> None:
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
||||
params: Final = {"max_results": 25, "numResults": native_count}
|
||||
outgoing: Final = ExaAISearchConfig().transform_search_request("question", params)
|
||||
assert outgoing["numResults"] == 25
|
||||
assert search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params=params) == (
|
||||
search_provider_cost_per_query("exa_ai/search", "exa_ai", optional_params=outgoing)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("count", [True, -1, 1.5, float("nan"), float("inf"), "25"])
|
||||
def test_exa_estimate_ignores_invalid_result_counts(count: object) -> None:
|
||||
from litellm.search.cost_calculator import search_provider_cost_per_query
|
||||
|
||||
assert search_provider_cost_per_query(
|
||||
"exa_ai/search", "exa_ai", optional_params={"numResults": count}
|
||||
) == search_provider_cost_per_query("exa_ai/search", "exa_ai")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.usefixtures("exa_api_key")
|
||||
@pytest.mark.asyncio
|
||||
async def test_lite_search_preserves_mode_and_highlights(respx_mock: respx.MockRouter, asynchronous: bool) -> None:
|
||||
route: Final = respx_mock.post(EXA_SEARCH_URL).respond(json=FULL_RESPONSE)
|
||||
response: Final = (
|
||||
await litellm.asearch(query="question", search_provider="exa_ai", type="lite", contents={"highlights": True})
|
||||
if asynchronous
|
||||
else litellm.search(query="question", search_provider="exa_ai", type="lite", contents={"highlights": True})
|
||||
)
|
||||
assert json.loads(route.calls.last.request.content) == {
|
||||
"query": "question",
|
||||
"type": "lite",
|
||||
"contents": {"highlights": True},
|
||||
}
|
||||
assert response.results[0].snippet == "\n\n".join(FULL_RESULT["highlights"])
|
||||
|
|
|
|||
|
|
@ -1,16 +1,22 @@
|
|||
import base64
|
||||
import datetime
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from pathlib import Path
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Protocol
|
||||
|
||||
import httpx
|
||||
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
|
||||
from litellm.search.cost_calculator import _bundled_cost_map_edited, search_provider_cost_per_query
|
||||
|
||||
EXA_SEARCH_URL: Final = "https://api.exa.ai/search"
|
||||
TAVILY_SEARCH_URL: Final = "https://api.tavily.com/search"
|
||||
|
|
@ -279,3 +285,326 @@ async def test_async_linkup_deep_search_reports_the_deep_price(
|
|||
)
|
||||
assert json.loads(respx_mock.calls.last.request.content)["depth"] == "deep"
|
||||
assert response.hidden_params["response_cost"] == pytest.approx(0.011)
|
||||
|
||||
|
||||
EXA_REGISTERED_PRICE: Final = 0.123
|
||||
EXA_REPORTED_TOTAL: Final = 0.0137
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_loaded_catalog_prices")
|
||||
@pytest.mark.parametrize(
|
||||
("registered_entry", "search_type", "expected"),
|
||||
[
|
||||
pytest.param("exa_ai/search", None, EXA_REGISTERED_PRICE, id="generic-entry"),
|
||||
pytest.param("exa_ai/search", "deep", EXA_REGISTERED_PRICE, id="generic-entry-covers-unregistered-mode"),
|
||||
pytest.param("exa_ai/search-deep", "deep", EXA_REGISTERED_PRICE, id="mode-entry"),
|
||||
pytest.param("exa_ai/search-deep", "auto", EXA_REPORTED_TOTAL, id="other-mode-keeps-reported-total"),
|
||||
pytest.param(None, "deep", EXA_REPORTED_TOTAL, id="nothing-registered"),
|
||||
],
|
||||
)
|
||||
def test_a_price_registered_for_exa_wins_over_the_cost_exa_reports(
|
||||
respx_mock: respx.MockRouter, registered_entry: str | None, search_type: str | None, expected: float
|
||||
) -> None:
|
||||
respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}})
|
||||
if registered_entry is not None:
|
||||
litellm.register_model(
|
||||
{
|
||||
registered_entry: {
|
||||
"input_cost_per_query": EXA_REGISTERED_PRICE,
|
||||
"litellm_provider": "exa_ai",
|
||||
"mode": "search",
|
||||
}
|
||||
}
|
||||
)
|
||||
response: Final = litellm.search(
|
||||
query="litellm release notes",
|
||||
search_provider="exa_ai",
|
||||
api_key="test-exa-key",
|
||||
**({"type": search_type} if search_type is not None else {}),
|
||||
)
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(expected)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_loaded_catalog_prices")
|
||||
def test_registering_no_price_on_an_exa_entry_keeps_the_cost_exa_reports(respx_mock: respx.MockRouter) -> None:
|
||||
respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}})
|
||||
litellm.register_model(
|
||||
{"exa_ai/search": {"metadata": {"notes": "operator note"}, "litellm_provider": "exa_ai", "mode": "search"}}
|
||||
)
|
||||
|
||||
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(EXA_REPORTED_TOTAL)
|
||||
|
||||
|
||||
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
|
||||
ServedExaEntry = Callable[[Mapping[str, object]], Mapping[str, object]]
|
||||
NO_MODE_PRICES: Final[Mapping[str, float]] = MappingProxyType({})
|
||||
OWN_MAP_DEEP_PRICE: Final[Mapping[str, float]] = MappingProxyType({"exa_ai/search-deep": OPERATOR_TIER_PRICE})
|
||||
|
||||
|
||||
class _ServeCostMap(Protocol):
|
||||
def __call__(
|
||||
self, url: str, exa_entry: ServedExaEntry, mode_prices: Mapping[str, float] = NO_MODE_PRICES
|
||||
) -> None: ...
|
||||
|
||||
|
||||
@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[_ServeCostMap]:
|
||||
def serve(url: str, exa_entry: ServedExaEntry, mode_prices: Mapping[str, float] = NO_MODE_PRICES) -> None:
|
||||
monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP")
|
||||
shipped: Final = GetModelCostMap.load_local_model_cost_map()
|
||||
repriced_modes: Final = {
|
||||
entry: {**shipped[entry], "input_cost_per_query": price} for entry, price in mode_prices.items()
|
||||
}
|
||||
body: Final = json.dumps(
|
||||
{**shipped, "exa_ai/search": exa_entry(shipped["exa_ai/search"]), **repriced_modes}
|
||||
).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,
|
||||
lambda _: {"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,
|
||||
lambda shipped: {**shipped, "input_cost_per_query": OPERATOR_FLAT_PRICE},
|
||||
OPERATOR_FLAT_PRICE,
|
||||
id="own-map-flat-price-over-shipped-tiers",
|
||||
),
|
||||
pytest.param(
|
||||
OPERATOR_COST_MAP_URL,
|
||||
lambda _: {
|
||||
"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(OPERATOR_COST_MAP_URL, lambda shipped: shipped, EXA_REPORTED_TOTAL, id="own-map-mirrors-litellm"),
|
||||
pytest.param(
|
||||
OPERATOR_COST_MAP_URL,
|
||||
lambda shipped: {name: value for name, value in shipped.items() if name != "input_cost_per_query"},
|
||||
EXA_REPORTED_TOTAL,
|
||||
id="own-map-from-a-release-without-the-flat-price",
|
||||
),
|
||||
pytest.param(
|
||||
DEFAULT_MODEL_COST_MAP_URL,
|
||||
lambda _: {"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: _ServeCostMap,
|
||||
url: str,
|
||||
exa_entry: ServedExaEntry,
|
||||
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)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("search_type", "mode_prices", "expected"),
|
||||
[
|
||||
pytest.param("deep", NO_MODE_PRICES, OPERATOR_FLAT_PRICE, id="deep"),
|
||||
pytest.param("lite", NO_MODE_PRICES, OPERATOR_FLAT_PRICE, id="lite"),
|
||||
pytest.param("deep", OWN_MAP_DEEP_PRICE, OPERATOR_TIER_PRICE, id="own-map-prices-deep"),
|
||||
pytest.param("lite", OWN_MAP_DEEP_PRICE, OPERATOR_FLAT_PRICE, id="other-mode-repriced"),
|
||||
],
|
||||
)
|
||||
def test_an_exa_price_from_the_operators_own_cost_map_covers_every_search_type_it_leaves_as_shipped(
|
||||
respx_mock: respx.MockRouter,
|
||||
_remote_cost_map: _ServeCostMap,
|
||||
search_type: str,
|
||||
mode_prices: Mapping[str, float],
|
||||
expected: float,
|
||||
) -> None:
|
||||
_remote_cost_map(
|
||||
OPERATOR_COST_MAP_URL, lambda shipped: {**shipped, "input_cost_per_query": OPERATOR_FLAT_PRICE}, mode_prices
|
||||
)
|
||||
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", type=search_type
|
||||
)
|
||||
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(expected)
|
||||
|
||||
|
||||
def _sha256_record(body: bytes) -> str:
|
||||
return base64.urlsafe_b64encode(hashlib.sha256(body).digest()).rstrip(b"=").decode()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _installed_litellm(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str], None]]:
|
||||
def install(recorded_sha256: str) -> None:
|
||||
dist_info: Final = tmp_path / "litellm-0.0.0.dist-info"
|
||||
dist_info.mkdir()
|
||||
(dist_info / "METADATA").write_text("Metadata-Version: 2.1\nName: litellm\nVersion: 0.0.0\n")
|
||||
(dist_info / "RECORD").write_text(
|
||||
f"litellm/model_prices_and_context_window_backup.json,sha256={recorded_sha256},1\n"
|
||||
)
|
||||
(tmp_path / "litellm").mkdir()
|
||||
(tmp_path / "litellm" / "model_prices_and_context_window_backup.json").touch()
|
||||
monkeypatch.syspath_prepend(str(tmp_path))
|
||||
_bundled_cost_map_edited.cache_clear()
|
||||
|
||||
yield install
|
||||
_bundled_cost_map_edited.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_loaded_catalog_prices")
|
||||
@pytest.mark.parametrize("bundled_map_edited", [False, True], ids=["bundled-map-as-installed", "bundled-map-edited"])
|
||||
def test_a_bundled_cost_map_edited_after_install_prices_exa_from_the_map(
|
||||
respx_mock: respx.MockRouter, _installed_litellm: Callable[[str], None], bundled_map_edited: bool
|
||||
) -> None:
|
||||
bundled: Final = GetModelCostMap.read_local_model_cost_map_bytes()
|
||||
_installed_litellm(_sha256_record(bundled + b" " if bundled_map_edited else bundled))
|
||||
respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [], "costDollars": {"total": EXA_REPORTED_TOTAL}})
|
||||
tiers: Final = GetModelCostMap.loaded_model_cost_map()["exa_ai/search"]["tiered_pricing"]
|
||||
assert isinstance(tiers, list) and tiers
|
||||
map_price: Final = next(
|
||||
tier["input_cost_per_query"]
|
||||
for tier in tiers
|
||||
if tier["max_results_range"][0] <= 10 <= tier["max_results_range"][1]
|
||||
)
|
||||
|
||||
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(
|
||||
map_price if bundled_map_edited else EXA_REPORTED_TOTAL
|
||||
)
|
||||
|
||||
|
||||
REGISTERED_TIER_PRICE: Final = 0.0456
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("own_map_price", "registration", "expected"),
|
||||
[
|
||||
pytest.param(
|
||||
{"tiered_pricing": [{"input_cost_per_query": OPERATOR_TIER_PRICE, "max_results_range": [0, 100]}]},
|
||||
{"input_cost_per_query": EXA_REGISTERED_PRICE},
|
||||
EXA_REGISTERED_PRICE,
|
||||
id="flat-registered-over-own-map-tiers",
|
||||
),
|
||||
pytest.param(
|
||||
{"input_cost_per_query": OPERATOR_FLAT_PRICE},
|
||||
{"tiered_pricing": [{"input_cost_per_query": REGISTERED_TIER_PRICE, "max_results_range": [0, 100]}]},
|
||||
REGISTERED_TIER_PRICE,
|
||||
id="tiers-registered-over-own-map-flat-price",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_a_price_registered_over_the_operators_own_cost_map_bills_the_registered_price(
|
||||
respx_mock: respx.MockRouter,
|
||||
_remote_cost_map: _ServeCostMap,
|
||||
own_map_price: Mapping[str, object],
|
||||
registration: Mapping[str, object],
|
||||
expected: float,
|
||||
) -> None:
|
||||
_remote_cost_map(
|
||||
OPERATOR_COST_MAP_URL, lambda _: {**own_map_price, "litellm_provider": "exa_ai", "mode": "search"}
|
||||
)
|
||||
litellm.register_model({"exa_ai/search": dict(registration)})
|
||||
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)
|
||||
|
||||
|
||||
@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}
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_loaded_catalog_prices")
|
||||
def test_registering_the_catalogs_own_flat_price_over_its_tiers_still_bills_the_flat_price() -> None:
|
||||
catalog: Final = GetModelCostMap.loaded_model_cost_map()
|
||||
flat_and_tiered: Final = tuple(
|
||||
key
|
||||
for key, entry in catalog.items()
|
||||
if entry.get("mode") == "search" and entry.get("tiered_pricing") and entry.get("input_cost_per_query")
|
||||
)
|
||||
assert flat_and_tiered
|
||||
litellm.register_model(
|
||||
{key: {"input_cost_per_query": catalog[key]["input_cost_per_query"]} for key in flat_and_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 flat_and_tiered
|
||||
}
|
||||
|
||||
assert costs == {key: pytest.approx((catalog[key]["input_cost_per_query"], 0.0)) for key in flat_and_tiered}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue