From f6227ebf8aa145740d17dce73874ab795e40ef87 Mon Sep 17 00:00:00 2001 From: Tobias Goerke Date: Thu, 24 Sep 2026 16:58:33 +0200 Subject: [PATCH] feat(staan): add Staan search provider support with configuration and API integration feat(staan): enhance StaanSearchConfig with mutable request headers and domain validation feat(staan): update Staan documentation and enhance StaanSearchConfig with environment variable validation and defaults feat(staan): update StaanSearchConfig to enforce boolean type for extra_snippets and add corresponding test feat(staan): update __all__ to use tuple syntax and add comprehensive tests for StaanSearchConfig feat(staan): remove STAAN_MARKET environment variable dependency and update related tests --- litellm/llms/staan/__init__.py | 0 litellm/llms/staan/search/__init__.py | 3 + litellm/llms/staan/search/transformation.py | 373 +++++++++++++ litellm/types/utils.py | 1 + litellm/utils.py | 2 + provider_endpoints_support.json | 7 + .../llms/staan/search/test_transformation.py | 528 ++++++++++++++++++ 7 files changed, 914 insertions(+) create mode 100644 litellm/llms/staan/__init__.py create mode 100644 litellm/llms/staan/search/__init__.py create mode 100644 litellm/llms/staan/search/transformation.py create mode 100644 tests/test_litellm/llms/staan/search/test_transformation.py diff --git a/litellm/llms/staan/__init__.py b/litellm/llms/staan/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/staan/search/__init__.py b/litellm/llms/staan/search/__init__.py new file mode 100644 index 00000000000..6ed2c5b0f71 --- /dev/null +++ b/litellm/llms/staan/search/__init__.py @@ -0,0 +1,3 @@ +from litellm.llms.staan.search.transformation import StaanSearchConfig + +__all__ = ("StaanSearchConfig",) diff --git a/litellm/llms/staan/search/transformation.py b/litellm/llms/staan/search/transformation.py new file mode 100644 index 00000000000..e8a2332cfd7 --- /dev/null +++ b/litellm/llms/staan/search/transformation.py @@ -0,0 +1,373 @@ +from collections.abc import Collection, Iterable, Mapping, Sequence +from types import MappingProxyType +from typing import Annotated, Final, Literal +from urllib.parse import urlsplit + +import httpx +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig, + SearchResponse, + SearchResult, +) +from litellm.secret_managers.main import get_secret_str + + +class StaanSearchRequest(TypedDict, total=False): + q: ReadOnly[str] + market: ReadOnly[Literal["fr-fr", "en-us", "de-de"]] + offset: ReadOnly[Literal[0, 10, 20, 30]] + include_domains: ReadOnly[Sequence[str]] + exclude_domains: ReadOnly[Sequence[str]] + extra_snippets: ReadOnly[bool] + max_snippets: ReadOnly[Annotated[int, Field(ge=1, le=10)]] + min_score: ReadOnly[Annotated[float, Field(ge=0, le=1)]] + full_content: ReadOnly[Literal["markdown", "html"]] + + +class _StaanExtraSnippet(BaseModel): + model_config = ConfigDict(extra="allow") + + chunk: str + score: float + + +class _StaanFullContent(BaseModel): + model_config = ConfigDict(extra="allow") + + text: str + format: Literal["markdown", "html"] + length: int + + +class _StaanWebResult(BaseModel): + model_config = ConfigDict(extra="allow") + + title: str = "" + url: str = "" + snippet: str = "" + published_date: str | None = None + extra_snippets: Sequence[_StaanExtraSnippet] | None = None + full_content: _StaanFullContent | None = None + + +class _StaanWebResponse(BaseModel): + model_config = ConfigDict(extra="allow") + + results: Sequence[_StaanWebResult] = () + + +class _StaanSearchResponse(BaseModel): + model_config = ConfigDict(extra="allow") + + web: _StaanWebResponse = Field(default_factory=_StaanWebResponse) + + +_STAAN_REQUEST_ADAPTER: Final = TypeAdapter(StaanSearchRequest) +_STAAN_RESPONSE_ADAPTER: Final = TypeAdapter(_StaanSearchResponse) +_STAAN_DOMAIN_LIST_ADAPTER: Final = TypeAdapter( + Annotated[list[Annotated[str, Field(min_length=1)]], Field(max_length=10)] +) +_STAAN_MARKETS: Final = frozenset(("fr-fr", "en-us", "de-de")) +_STAAN_OFFSETS: Final = frozenset((0, 10, 20, 30)) +_COUNTRY_TO_MARKET: Final[Mapping[str, str]] = MappingProxyType( + { + "fr": "fr-fr", + "fr-fr": "fr-fr", + "de": "de-de", + "de-de": "de-de", + "us": "en-us", + "en-us": "en-us", + } +) +_TRUE_STRINGS: Final = frozenset(("true", "1", "yes")) +_FALSE_STRINGS: Final = frozenset(("false", "0", "no")) +_EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) +_NORMALIZED_RESPONSE_FIELDS: Final = frozenset(("object", "results")) + + +def _normalize_domains(*domain_groups: Iterable[str]) -> frozenset[str]: + return frozenset(domain.lower().removeprefix("www.") for group in domain_groups for domain in group) + + +def _hostname_matches(hostname: str, domains: Collection[str]) -> bool: + return any(hostname == domain or hostname.endswith(f".{domain}") for domain in domains) + + +class StaanSearchConfig(BaseSearchConfig): + STAAN_API_BASE = "https://api.staan.ai/v2/search/web" + + @staticmethod + def ui_friendly_name() -> str: + return "Staan" + + def get_http_method(self) -> Literal["GET", "POST"]: + return "POST" + + def validate_environment( + self, + headers: Mapping[str, str], + api_key: str | None = None, + api_base: str | None = None, + **kwargs: object, # kwargs-ok: BaseSearchConfig hook accepts provider keyword arguments + ) -> dict: # mutable-ok: BaseSearchConfig requires mutable request headers + resolved_api_key: Final = self.resolve_server_api_key( + caller_api_key=api_key, + caller_api_base=api_base, + key_env_vars=("STAAN_API_KEY",), + base_env_var="STAAN_API_BASE", + default_api_base=self.STAAN_API_BASE, + ) + if not resolved_api_key: + raise ValueError("STAAN_API_KEY is not set. Set `STAAN_API_KEY` environment variable.") + request_headers: Final[dict[str, str]] = { # mutable-ok: HTTP handler hooks require mutable request headers + **headers, + "Accept": "application/json", + "Authorization": f"Bearer {resolved_api_key}", + "Content-Type": "application/json", + } + return request_headers + + def get_complete_url( + self, + api_base: str | None, + optional_params: dict, # mutable-ok: BaseSearchConfig hook requires this signature + data: dict | list[dict] | None = None, # mutable-ok: BaseSearchConfig hook requires this signature + **kwargs: object, # kwargs-ok: BaseSearchConfig hook accepts provider keyword arguments + ) -> str: + return api_base or get_secret_str("STAAN_API_BASE") or self.STAAN_API_BASE + + def transform_search_request( + self, + query: str | Sequence[str], + optional_params: dict, # mutable-ok: BaseSearchConfig hook requires this signature + api_key: str | None = None, + api_base: str | None = None, + headers: Mapping[str, str] | None = None, + **kwargs: object, # kwargs-ok: BaseSearchConfig hook accepts provider keyword arguments + ) -> StaanSearchRequest: + search_params: Final = MappingProxyType(optional_params) + search_query: Final = " ".join(query) if not isinstance(query, str) else query + if len(search_query) > 400: + raise ValueError("query must be at most 400 characters") + self._validate_max_results(search_params.get("max_results")) + + market: Final = self._resolve_market(search_params) + offset: Final = search_params.get("offset") + if offset is not None and offset not in _STAAN_OFFSETS: + raise ValueError("offset must be one of: 0, 10, 20, 30") + self._validate_count(search_params.get("count")) + + domain_params: Final = self._build_domain_params(search_params) + enrichment_params: Final = self._build_enrichment_params(search_params) + request_fields: Final = ( + ("q", search_query), + *((("market", market),) if market is not None else ()), + *((("offset", offset),) if offset is not None else ()), + *domain_params.items(), + *enrichment_params.items(), + ) + request_data: Final[StaanSearchRequest] = MappingProxyType(dict(request_fields)) + try: + validated_request: Final = _STAAN_REQUEST_ADAPTER.validate_python(request_data) + return validated_request + except ValidationError as exc: + raise ValueError(f"Invalid Staan search parameters: {exc}") from exc + + @staticmethod + def _resolve_market(optional_params: Mapping[str, object]) -> str | None: + explicit_market: Final = optional_params.get("market") + if explicit_market: + if explicit_market not in _STAAN_MARKETS: + raise ValueError("market must be one of: fr-fr, en-us, de-de") + return explicit_market + + country: Final = optional_params.get("country") + if country is not None: + country_market: Final = _COUNTRY_TO_MARKET.get(country.lower()) + if country_market is None and explicit_market is None: + raise ValueError("country must map to a supported Staan market: FR, DE, or US") + if country_market is not None: + return country_market + + return None + + @staticmethod + def _validate_count(count: int | None) -> None: + if count is not None and count != 10: + raise ValueError("Staan fixes count at 10; use max_results to limit results") + + @staticmethod + def _validate_max_results(max_results: object) -> int | None: + if max_results is None: + return None + if isinstance(max_results, bool) or not isinstance(max_results, int) or max_results < 1: + raise ValueError("max_results must be a positive integer") + return max_results + + @classmethod + def _build_domain_params( + cls, + optional_params: Mapping[str, object], + ) -> StaanSearchRequest: + raw_include_domains: Final = optional_params.get("include_domains") + raw_search_domain_filter: Final = optional_params.get("search_domain_filter") + raw_exclude_domains: Final = optional_params.get("exclude_domains") + include_domains: Final = cls._validate_domain_group(raw_include_domains, "include_domains") + search_domain_filter: Final = cls._validate_domain_group(raw_search_domain_filter, "search_domain_filter") + exclude_domains: Final = cls._validate_domain_group(raw_exclude_domains, "exclude_domains") + merged_include_domains: Final = tuple(dict.fromkeys((*include_domains, *search_domain_filter))) + if merged_include_domains and exclude_domains: + raise ValueError("include_domains and exclude_domains are mutually exclusive in Staan") + try: + domain_entries: Final = ( + ( + ( + "include_domains", + _STAAN_DOMAIN_LIST_ADAPTER.validate_python(merged_include_domains), + ), + ) + if merged_include_domains + else () + ) + excluded_entry: Final = ( + ( + ( + "exclude_domains", + _STAAN_DOMAIN_LIST_ADAPTER.validate_python(exclude_domains), + ), + ) + if exclude_domains + else () + ) + domain_params: Final[StaanSearchRequest] = MappingProxyType(dict(domain_entries + excluded_entry)) + return domain_params + except ValidationError as exc: + raise ValueError(f"Invalid Staan domain filters: {exc}") from exc + + @staticmethod + def _validate_domain_group(value: object, name: str) -> tuple[str, ...]: + if value is None: + return () + try: + return tuple(_STAAN_DOMAIN_LIST_ADAPTER.validate_python(value)) + except ValidationError as exc: + raise ValueError(f"Invalid Staan domain filter `{name}`: {exc}") from exc + + @classmethod + def _build_enrichment_params( + cls, + optional_params: Mapping[str, object], + ) -> StaanSearchRequest: + extra_snippets: Final = optional_params.get("extra_snippets") + should_fetch_snippets: Final = cls._parse_bool( + optional_params.get("max_snippets") is not None if extra_snippets is None else extra_snippets + ) + max_snippets: Final = optional_params.get("max_snippets") + min_score: Final = optional_params.get("min_score") + full_content: Final = optional_params.get("full_content") + + if should_fetch_snippets and max_snippets is not None and not 1 <= max_snippets <= 10: + raise ValueError("max_snippets must be between 1 and 10") + if should_fetch_snippets and min_score is not None and not 0 <= min_score <= 1: + raise ValueError("min_score must be between 0 and 1") + + enrichment_entries: Final = (("extra_snippets", True),) if should_fetch_snippets else () + max_snippet_entry: Final = ( + (("max_snippets", max_snippets),) if should_fetch_snippets and max_snippets is not None else () + ) + score_entry: Final = (("min_score", min_score),) if should_fetch_snippets and min_score is not None else () + content_entry: Final = (("full_content", full_content),) if full_content is not None else () + enrichment_params: Final[StaanSearchRequest] = MappingProxyType( + dict(enrichment_entries + max_snippet_entry + score_entry + content_entry) + ) + return enrichment_params + + @staticmethod + def _parse_bool(value: object) -> bool: + if isinstance(value, bool): + return value + if not isinstance(value, str): + raise TypeError("extra_snippets must be a boolean") + normalized: Final = value.strip().lower() + if normalized in _TRUE_STRINGS: + return True + if normalized in _FALSE_STRINGS: + return False + raise ValueError("extra_snippets must be a boolean") + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + optional_params: dict | None = None, # mutable-ok: required by the BaseSearchConfig hook + **kwargs: object, # kwargs-ok: required by the BaseSearchConfig hook + ) -> SearchResponse: + response: Final = _STAAN_RESPONSE_ADAPTER.validate_python(raw_response.json()) + search_options: Final = MappingProxyType(optional_params or _EMPTY_MAPPING) + included_domains: Final = _normalize_domains( + search_options.get("search_domain_filter", ()), + search_options.get("include_domains", ()), + ) + excluded_domains: Final = _normalize_domains(search_options.get("exclude_domains", ())) + results: Final = tuple( + self._transform_result(result) + for result in response.web.results + if self._result_matches_domains(result.url, included_domains, excluded_domains) + ) + max_results: Final = self._validate_max_results(search_options.get("max_results")) + response_extra: Final = tuple((response.model_extra or _EMPTY_MAPPING).items()) + web_extra: Final = tuple((response.web.model_extra or _EMPTY_MAPPING).items()) + response_extra_fields: Final[Mapping[str, object]] = MappingProxyType( + {key: value for key, value in response_extra + web_extra if key not in _NORMALIZED_RESPONSE_FIELDS} + ) + normalized_response: Final[SearchResponse] = SearchResponse( + results=tuple(results[:max_results] if max_results is not None else results), + object="search", + ) + return normalized_response.model_copy(update=response_extra_fields) + + @staticmethod + def _result_matches_domains( + url: str, + included_domains: Collection[str], + excluded_domains: Collection[str], + ) -> bool: + hostname: Final = (urlsplit(url).hostname or "").lower().removeprefix("www.") + if included_domains and (not hostname or not _hostname_matches(hostname, included_domains)): + return False + return not (excluded_domains and _hostname_matches(hostname, excluded_domains)) + + @staticmethod + def _transform_result(result: _StaanWebResult) -> SearchResult: + extra_snippets: Final = result.extra_snippets or () + snippet: Final = "\n\n".join( + text + for text in ( + result.snippet, + *(extra.chunk for extra in extra_snippets), + ) + if text + ) + snippet_fields: Final = ( + (("extra_snippets", tuple(extra.model_dump() for extra in extra_snippets)),) + if result.extra_snippets is not None + else () + ) + content_fields: Final = ( + (("full_content", result.full_content.model_dump()),) if result.full_content is not None else () + ) + additional_fields: Final[Mapping[str, object]] = MappingProxyType( + dict(tuple((result.model_extra or _EMPTY_MAPPING).items()) + snippet_fields + content_fields) + ) + standard_result: Final = SearchResult( + title=result.title, + url=result.url, + snippet=snippet, + date=result.published_date, + last_updated=None, + ) + return standard_result.model_copy(update=additional_fields) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index caf88e5d517..790b7386429 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4348,6 +4348,7 @@ class SearchProviders(str, Enum): AGENTCORE = "agentcore" NIMBLE = "nimble" BING_GROUNDING = "bing_grounding" + STAAN = "staan" # Create a set of all search provider values for quick lookup diff --git a/litellm/utils.py b/litellm/utils.py index 64097021dff..319fb13b01d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9778,6 +9778,7 @@ class ProviderConfigManager: from litellm.llms.searchapi.search.transformation import SearchAPIConfig from litellm.llms.searxng.search.transformation import SearXNGSearchConfig from litellm.llms.serper.search.transformation import SerperSearchConfig + from litellm.llms.staan.search.transformation import StaanSearchConfig from litellm.llms.tavily.search.transformation import TavilySearchConfig from litellm.llms.tinyfish.search.transformation import TinyfishSearchConfig from litellm.llms.you_com.search.transformation import YouComSearchConfig @@ -9803,6 +9804,7 @@ class ProviderConfigManager: SearchProviders.AGENTCORE: AgentCoreSearchConfig, SearchProviders.NIMBLE: NimbleSearchConfig, SearchProviders.BING_GROUNDING: BingGroundingSearchConfig, + SearchProviders.STAAN: StaanSearchConfig, } config_class: Final = PROVIDER_TO_CONFIG_MAP.get(provider, None) if config_class is None: diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index b8d1621cde3..eeb4ce4d084 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2495,6 +2495,13 @@ "search": true } }, + "staan": { + "display_name": "Staan (`staan`)", + "url": "https://docs.litellm.ai/docs/search/staan", + "endpoints": { + "search": true + } + }, "you_com": { "display_name": "You.com (`you_com`)", "url": "https://docs.litellm.ai/docs/search/you_com" diff --git a/tests/test_litellm/llms/staan/search/test_transformation.py b/tests/test_litellm/llms/staan/search/test_transformation.py new file mode 100644 index 00000000000..bf8076f3841 --- /dev/null +++ b/tests/test_litellm/llms/staan/search/test_transformation.py @@ -0,0 +1,528 @@ +"""Tests for the Staan search provider transformation.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +import litellm +from litellm.llms.staan.search.transformation import StaanSearchConfig +from litellm.router_utils.search_api_router import SearchAPIRouter +from litellm.types.utils import SearchProviders +from litellm.utils import ProviderConfigManager + + +def test_staan_is_registered_as_a_search_provider() -> None: + config = ProviderConfigManager.get_provider_search_config(SearchProviders.STAAN) + + assert isinstance(config, StaanSearchConfig) + + +def test_staan_has_ui_friendly_provider_name() -> None: + assert StaanSearchConfig.ui_friendly_name() == "Staan" + + +def test_staan_environment_validation_returns_mutable_request_headers() -> None: + headers = StaanSearchConfig().validate_environment( + headers={"X-Request-ID": "request-id"}, + api_key="test-key", + ) + + assert type(headers) is dict + assert headers == { + "X-Request-ID": "request-id", + "Accept": "application/json", + "Authorization": "Bearer test-key", + "Content-Type": "application/json", + } + + +def test_staan_api_defaults_resolve_from_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("STAAN_API_KEY", "environment-key") + monkeypatch.setenv("STAAN_API_BASE", "https://search.example/api") + config = StaanSearchConfig() + + assert config.validate_environment(headers={})["Authorization"] == "Bearer environment-key" + assert config.get_complete_url(None, {}) == "https://search.example/api" + assert config.transform_search_request("query", {}) == {"q": "query"} + + +def test_staan_market_is_not_read_from_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("STAAN_MARKET", "de-de") + + request = StaanSearchConfig().transform_search_request("query", {}) + + assert request == {"q": "query"} + + +def test_staan_explicit_api_base_takes_precedence_over_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("STAAN_API_BASE", "https://search.example/api") + + url = StaanSearchConfig().get_complete_url("https://caller.example/search", {}) + + assert url == "https://caller.example/search" + + +def test_staan_environment_validation_requires_api_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("STAAN_API_KEY", raising=False) + + with pytest.raises(ValueError, match="STAAN_API_KEY is not set"): + StaanSearchConfig().validate_environment(headers={}) + + +def test_staan_request_maps_enrichment_and_domain_parameters() -> None: + config = StaanSearchConfig() + + request = config.transform_search_request( + query="vector search", + optional_params={ + "market": "en-us", + "offset": 20, + "search_domain_filter": ["example.com"], + "extra_snippets": True, + "max_snippets": 5, + "min_score": 0.2, + "full_content": "markdown", + "max_results": 3, + }, + ) + + assert config.get_http_method() == "POST" + assert config.get_complete_url(None, {}, request) == config.STAAN_API_BASE + assert request == { + "q": "vector search", + "market": "en-us", + "offset": 20, + "include_domains": ["example.com"], + "extra_snippets": True, + "max_snippets": 5, + "min_score": 0.2, + "full_content": "markdown", + } + + +@pytest.mark.parametrize( + "optional_params", + [ + {"include_domains": [], "exclude_domains": ["blocked.example"]}, + {"search_domain_filter": [], "exclude_domains": ["blocked.example"]}, + { + "include_domains": [], + "search_domain_filter": [], + "exclude_domains": ["blocked.example"], + }, + ], +) +def test_staan_empty_include_filters_allow_exclude_domains(optional_params: dict[str, object]) -> None: + request = StaanSearchConfig().transform_search_request("query", optional_params) + + assert request == {"q": "query", "exclude_domains": ["blocked.example"]} + + +def test_staan_max_snippets_enables_enrichment_and_explicit_false_disables_it() -> None: + config = StaanSearchConfig() + + enabled = config.transform_search_request("query", {"max_snippets": 5}) + disabled = config.transform_search_request("query", {"extra_snippets": False, "max_snippets": 5}) + + assert enabled["extra_snippets"] is True + assert enabled["max_snippets"] == 5 + assert "extra_snippets" not in disabled + assert "max_snippets" not in disabled + + +def test_staan_maps_litellm_country_to_staan_market() -> None: + request = StaanSearchConfig().transform_search_request("query", {"country": "DE"}) + + assert request["market"] == "de-de" + + +def test_staan_empty_country_uses_explicit_market() -> None: + request = StaanSearchConfig().transform_search_request("query", {"country": "", "market": "en-us"}) + + assert request["market"] == "en-us" + + +def test_staan_country_and_explicit_market_use_explicit_market() -> None: + request = StaanSearchConfig().transform_search_request("query", {"country": "DE", "market": "en-us"}) + + assert request["market"] == "en-us" + + +def test_staan_joins_multi_part_query_and_keeps_count_fixed() -> None: + request = StaanSearchConfig().transform_search_request( + query=("vector", "search"), + optional_params={"count": 10}, + ) + + assert request == {"q": "vector search"} + + +def test_staan_rejects_queries_over_provider_limit() -> None: + with pytest.raises(ValueError, match="query must be at most 400 characters"): + StaanSearchConfig().transform_search_request("q" * 401, {}) + + +@pytest.mark.parametrize("value", ["true", "1", " YES "]) +def test_staan_accepts_true_string_values_for_extra_snippets(value: str) -> None: + request = StaanSearchConfig().transform_search_request("query", {"extra_snippets": value}) + + assert request["extra_snippets"] is True + + +@pytest.mark.parametrize("value", ["false", "0", " NO "]) +def test_staan_accepts_false_string_values_for_extra_snippets(value: str) -> None: + request = StaanSearchConfig().transform_search_request("query", {"extra_snippets": value, "max_snippets": 5}) + + assert "extra_snippets" not in request + assert "max_snippets" not in request + + +@pytest.mark.parametrize( + ("optional_params", "message"), + [ + ({"offset": 5}, "offset"), + ({"market": "es-es"}, "market"), + ({"country": "GB"}, "country"), + ({"count": 9}, "count"), + ({"max_results": 0}, "max_results"), + ({"max_results": -1}, "max_results"), + ({"max_results": 1.5}, "max_results"), + ({"max_results": True}, "max_results"), + ({"extra_snippets": True, "max_snippets": 11}, "max_snippets"), + ({"extra_snippets": True, "min_score": 1.1}, "min_score"), + ({"extra_snippets": "sometimes"}, "extra_snippets"), + ({"full_content": "text"}, "full_content"), + ( + {"include_domains": [f"{index}.example" for index in range(11)]}, + "include_domains", + ), + ( + {"include_domains": ["a.example"], "exclude_domains": ["b.example"]}, + "mutually exclusive", + ), + ( + {"search_domain_filter": ["a.example"], "exclude_domains": ["b.example"]}, + "mutually exclusive", + ), + ( + { + "include_domains": [f"include-{index}.example" for index in range(6)], + "search_domain_filter": [f"filter-{index}.example" for index in range(5)], + }, + "Invalid Staan domain filters", + ), + ({"include_domains": "example.com"}, "include_domains"), + ({"search_domain_filter": "example.com"}, "search_domain_filter"), + ({"exclude_domains": "example.com"}, "exclude_domains"), + ], +) +def test_staan_rejects_invalid_parameters(optional_params: dict[str, object], message: str) -> None: + with pytest.raises(ValueError, match=message): + StaanSearchConfig().transform_search_request("query", optional_params) + + +def test_staan_rejects_non_boolean_extra_snippets_type() -> None: + with pytest.raises(TypeError, match="extra_snippets must be a boolean"): + StaanSearchConfig().transform_search_request("query", {"extra_snippets": 1}) + + +def test_staan_response_preserves_all_enriched_result_fields() -> None: + response = httpx.Response( + 200, + json={ + "search_id": "search-id", + "query": { + "q": "vector database", + "market": "en-us", + "count": 10, + "offset": 0, + }, + "web": { + "results": [ + { + "title": "Search result", + "url": "https://example.com/page", + "snippet": "Provider preview", + "display_url": "example.com > page", + "hostname": "www.example.com", + "favicon_url": "https://example.com/favicon.ico", + "extra_snippets": [ + {"chunk": "First relevant passage", "score": 0.91}, + {"chunk": "Second relevant passage", "score": 0.82}, + {"chunk": "Third relevant passage", "score": 0.67}, + ], + "full_content": { + "text": "Complete page body", + "format": "markdown", + "length": 19, + }, + "published_date": "2026-09-24T00:00:00Z", + } + ], + "displayed_results": 1, + }, + }, + ) + + result_response = StaanSearchConfig().transform_search_response( + response, + logging_obj=None, + optional_params={}, + ) + + assert result_response.results[0].model_dump(exclude_none=True) == { + "title": "Search result", + "url": "https://example.com/page", + "snippet": ("Provider preview\n\nFirst relevant passage\n\nSecond relevant passage\n\nThird relevant passage"), + "date": "2026-09-24T00:00:00Z", + "display_url": "example.com > page", + "hostname": "www.example.com", + "favicon_url": "https://example.com/favicon.ico", + "extra_snippets": ( + {"chunk": "First relevant passage", "score": 0.91}, + {"chunk": "Second relevant passage", "score": 0.82}, + {"chunk": "Third relevant passage", "score": 0.67}, + ), + "full_content": { + "text": "Complete page body", + "format": "markdown", + "length": 19, + }, + } + assert result_response.search_id == "search-id" + assert result_response.displayed_results == 1 + assert result_response.query == { + "q": "vector database", + "market": "en-us", + "count": 10, + "offset": 0, + } + + +def test_staan_provider_metadata_cannot_replace_filtered_results() -> None: + response = httpx.Response( + 200, + json={ + "results": [{"title": "Provider override", "url": "https://blocked.example/"}], + "object": "provider-object", + "web": { + "results": [ + {"title": "Allowed", "url": "https://allowed.example/"}, + {"title": "Blocked", "url": "https://blocked.example/"}, + ] + }, + }, + ) + + result = StaanSearchConfig().transform_search_response( + response, + logging_obj=None, + optional_params={"search_domain_filter": ["allowed.example"], "max_results": 1}, + ) + + assert [item.title for item in result.results] == ["Allowed"] + assert result.object == "search" + + +@pytest.mark.parametrize("max_results", [0, -1]) +def test_staan_response_rejects_non_positive_max_results(max_results: int) -> None: + response = httpx.Response( + 200, + json={"web": {"results": [{"title": "Result", "url": "https://example.com/"}]}}, + ) + + with pytest.raises(ValueError, match="max_results must be a positive integer"): + StaanSearchConfig().transform_search_response( + response, + logging_obj=None, + optional_params={"max_results": max_results}, + ) + + +def test_staan_response_filters_results_by_domain_boundaries() -> None: + response = httpx.Response( + 200, + json={ + "web": { + "results": [ + {"title": "Root", "url": "https://example.com/"}, + {"title": "Subdomain", "url": "https://docs.example.com/"}, + {"title": "Lookalike", "url": "https://notexample.com/"}, + {"title": "No host", "url": "not a URL"}, + ] + } + }, + ) + + result = StaanSearchConfig().transform_search_response( + response, + logging_obj=None, + optional_params={"search_domain_filter": ["www.example.com"]}, + ) + + assert [item.title for item in result.results] == ["Root", "Subdomain"] + + +def test_staan_response_excludes_matching_domain_and_subdomains() -> None: + response = httpx.Response( + 200, + json={ + "web": { + "results": [ + {"title": "Excluded", "url": "https://www.example.com/"}, + {"title": "Excluded subdomain", "url": "https://docs.example.com/"}, + {"title": "Retained", "url": "https://other.com/"}, + ] + } + }, + ) + + result = StaanSearchConfig().transform_search_response( + response, + logging_obj=None, + optional_params={"exclude_domains": ["example.com"]}, + ) + + assert [item.title for item in result.results] == ["Retained"] + + +def test_staan_response_preserves_empty_enrichment_values() -> None: + response = httpx.Response( + 200, + json={ + "web": { + "results": [ + { + "title": "Search result", + "url": "https://example.com/page", + "snippet": "Provider preview", + "extra_snippets": [], + "full_content": {"text": "", "format": "markdown", "length": 0}, + } + ] + } + }, + ) + + result = ( + StaanSearchConfig() + .transform_search_response( + response, + logging_obj=None, + optional_params={}, + ) + .results[0] + ) + + assert result.snippet == "Provider preview" + assert result.extra_snippets == () + assert result.full_content == {"text": "", "format": "markdown", "length": 0} + + +def test_staan_max_results_does_not_truncate_extra_snippets() -> None: + response = httpx.Response( + 200, + json={ + "web": { + "results": [ + { + "title": "Search result", + "url": "https://example.com/page", + "snippet": "Base", + "extra_snippets": [ + {"chunk": "One", "score": 0.9}, + {"chunk": "Two", "score": 0.8}, + ], + }, + { + "title": "Second result", + "url": "https://example.com/second", + "snippet": "Second base", + }, + ] + } + }, + ) + + result = StaanSearchConfig().transform_search_response( + response, + logging_obj=None, + optional_params={"max_results": 1}, + ) + + assert len(result.results) == 1 + assert result.results[0].snippet == "Base\n\nOne\n\nTwo" + assert result.results[0].extra_snippets == ( + {"chunk": "One", "score": 0.9}, + {"chunk": "Two", "score": 0.8}, + ) + + +@pytest.mark.asyncio +async def test_search_tool_enrichment_defaults_reach_provider_call() -> None: + router = SimpleNamespace( + search_tools=[ + { + "search_tool_name": "search", + "litellm_params": { + "search_provider": "staan", + "api_key": "test-key", + "extra_snippets": True, + "max_snippets": 5, + }, + } + ] + ) + + response = httpx.Response( + 200, + json={ + "web": { + "results": [ + { + "title": "Search result", + "url": "https://example.com/page", + "snippet": "Preview", + "extra_snippets": [ + {"chunk": "Passage one", "score": 0.9}, + {"chunk": "Passage two", "score": 0.8}, + ], + } + ] + } + }, + ) + with ( + patch.dict("os.environ", {"STAAN_API_KEY": "test-key"}), + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + return_value=response, + ) as post, + ): + result = await SearchAPIRouter.async_search_with_fallbacks_helper( + router_instance=router, + model="search", + original_generic_function=litellm.asearch, + query="vector search", + ) + + assert post.call_args.kwargs["json"] == { + "q": "vector search", + "extra_snippets": True, + "max_snippets": 5, + } + assert result.results[0].snippet == "Preview\n\nPassage one\n\nPassage two" + + await SearchAPIRouter.async_search_with_fallbacks_helper( + router_instance=router, + model="search", + original_generic_function=litellm.asearch, + query="vector search", + max_snippets=3, + extra_snippets=False, + ) + + assert post.call_args.kwargs["json"] == {"q": "vector search"}