mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(search): add SerpApi search provider
This commit adds a production-quality SerpApi search provider to LiteLLM, conforming to the BaseSearchConfig abstraction. It supports: - Authentication via SERPAPI_API_KEY mapped to the api_key query parameter. - Dynamic engine selection (e.g. google, google_news, google_shopping), defaulting to google. - The LiteLLM unified search response structure mapped from organic_results. - Automatic pass-through of all extra supported parameters (e.g., num, gl, hl). - Comprehensive unit and integration test coverage.
This commit is contained in:
Gowtham
parent
a3236ba947
commit
5a7ef2cca2
8 changed files with 641 additions and 0 deletions
3
litellm/llms/serpapi/__init__.py
Normal file
3
litellm/llms/serpapi/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
"""
|
||||
SerpApi LLMs module.
|
||||
"""
|
||||
7
litellm/llms/serpapi/search/__init__.py
Normal file
7
litellm/llms/serpapi/search/__init__.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
"""
|
||||
SerpApi Search API module.
|
||||
"""
|
||||
|
||||
from litellm.llms.serpapi.search.transformation import SerpApiSearchConfig
|
||||
|
||||
__all__ = ["SerpApiSearchConfig"]
|
||||
174
litellm/llms/serpapi/search/transformation.py
Normal file
174
litellm/llms/serpapi/search/transformation.py
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
"""
|
||||
Calls SerpApi's search endpoint.
|
||||
|
||||
SerpApi API Reference: https://serpapi.com/search-api
|
||||
"""
|
||||
|
||||
from typing import Final, Literal, TypedDict
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
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 _SerpApiSearchRequestRequired(TypedDict):
|
||||
engine: ReadOnly[str]
|
||||
q: ReadOnly[str]
|
||||
api_key: ReadOnly[str]
|
||||
|
||||
|
||||
class SerpApiSearchRequest(_SerpApiSearchRequestRequired, total=False):
|
||||
"""
|
||||
SerpApi request format.
|
||||
Based on: https://serpapi.com/search-api
|
||||
"""
|
||||
|
||||
gl: ReadOnly[str]
|
||||
hl: ReadOnly[str]
|
||||
location: ReadOnly[str]
|
||||
device: ReadOnly[str]
|
||||
safe: ReadOnly[str]
|
||||
num: ReadOnly[int]
|
||||
start: ReadOnly[int]
|
||||
tbs: ReadOnly[str]
|
||||
tbm: ReadOnly[str]
|
||||
|
||||
|
||||
class SerpApiSearchConfig(BaseSearchConfig):
|
||||
SERPAPI_API_BASE: Final = "https://serpapi.com/search.json"
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "SerpApi"
|
||||
|
||||
def get_http_method(self) -> Literal["GET", "POST"]:
|
||||
return "GET"
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict[str, str], # mutable-ok: base class override
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
**kwargs: object, # kwargs-ok: base class override
|
||||
) -> dict[str, str]: # mutable-ok: base class override
|
||||
resolved_api_key: Final = self.resolve_server_api_key(
|
||||
caller_api_key=api_key,
|
||||
caller_api_base=api_base,
|
||||
key_env_vars=("SERPAPI_API_KEY",),
|
||||
base_env_var="SERPAPI_API_BASE",
|
||||
default_api_base=self.SERPAPI_API_BASE,
|
||||
)
|
||||
if not resolved_api_key:
|
||||
raise ValueError("SERPAPI_API_KEY is not set. Set `SERPAPI_API_KEY` environment variable.")
|
||||
return {
|
||||
**headers,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
optional_params: dict[str, object], # mutable-ok: base class override
|
||||
data: dict[str, object] | list[dict[str, object]] | None = None, # mutable-ok: base class override
|
||||
**kwargs: object, # kwargs-ok: base class override
|
||||
) -> str:
|
||||
base: Final = api_base or get_secret_str("SERPAPI_API_BASE") or self.SERPAPI_API_BASE
|
||||
if data and isinstance(data, dict) and "_serpapi_params" in data:
|
||||
params: Final = data["_serpapi_params"]
|
||||
if isinstance(params, dict):
|
||||
return f"{base}?{urlencode(params, doseq=True)}" # pyright: ignore[reportUnknownArgumentType] # params is a dict
|
||||
return base
|
||||
|
||||
def transform_search_request(
|
||||
self,
|
||||
query: str | list[str], # mutable-ok: base class override
|
||||
optional_params: dict[str, object], # mutable-ok: base class override
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
**kwargs: object, # kwargs-ok: base class override
|
||||
) -> dict[str, object]: # mutable-ok: base class override
|
||||
q_str: Final = " ".join(query) if isinstance(query, list) else query
|
||||
|
||||
resolved_api_key: Final = self.resolve_server_api_key(
|
||||
caller_api_key=api_key,
|
||||
caller_api_base=api_base,
|
||||
key_env_vars=("SERPAPI_API_KEY",),
|
||||
base_env_var="SERPAPI_API_BASE",
|
||||
default_api_base=self.SERPAPI_API_BASE,
|
||||
)
|
||||
if not resolved_api_key:
|
||||
raise ValueError("SERPAPI_API_KEY is not set. Set `SERPAPI_API_KEY` environment variable.")
|
||||
|
||||
engine: Final = str(optional_params.get("engine", "google"))
|
||||
|
||||
domains: Final = optional_params.get("search_domain_filter")
|
||||
q_final: Final = (
|
||||
f"({q_str}) ({" OR ".join(f"site:{d}" for d in domains)})" # pyright: ignore[reportUnknownVariableType] # items are strings
|
||||
if isinstance(domains, list) and len(domains) > 0 # pyright: ignore[reportUnknownArgumentType] # domains is a list
|
||||
else q_str
|
||||
)
|
||||
|
||||
base_params: Final[dict[str, object]] = { # mutable-ok: mapping type
|
||||
"engine": engine,
|
||||
"q": q_final,
|
||||
"api_key": resolved_api_key,
|
||||
}
|
||||
|
||||
optional_mappings: Final[dict[str, object]] = { # mutable-ok: mapping type
|
||||
**({"num": optional_params["max_results"]} if "max_results" in optional_params else {}),
|
||||
**({"gl": str(optional_params["country"]).lower()} if "country" in optional_params else {}),
|
||||
}
|
||||
|
||||
extra_params: Final[dict[str, object]] = { # mutable-ok: mapping type
|
||||
param: value
|
||||
for param, value in optional_params.items()
|
||||
if param not in self.get_supported_perplexity_optional_params()
|
||||
and param not in base_params
|
||||
and param not in optional_mappings
|
||||
}
|
||||
|
||||
request_data: Final[dict[str, object]] = {**base_params, **optional_mappings, **extra_params} # mutable-ok: mapping type
|
||||
|
||||
return {"_serpapi_params": request_data}
|
||||
|
||||
def transform_search_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
**kwargs: object, # kwargs-ok: base class override
|
||||
) -> SearchResponse:
|
||||
"""
|
||||
Transform SerpApi response to LiteLLM unified SearchResponse format.
|
||||
|
||||
SerpApi -> LiteLLM mappings:
|
||||
- organic_results[].title -> SearchResult.title
|
||||
- organic_results[].link -> SearchResult.url
|
||||
- organic_results[].snippet -> SearchResult.snippet
|
||||
- organic_results[].date -> SearchResult.date (optional)
|
||||
- organic_results[].position -> SearchResult (preserved via extra fields)
|
||||
"""
|
||||
response_json: Final = raw_response.json()
|
||||
|
||||
results: Final = [
|
||||
SearchResult(
|
||||
title=result.get("title", ""),
|
||||
url=result.get("link", ""),
|
||||
snippet=result.get("snippet", ""),
|
||||
date=result.get("date"),
|
||||
last_updated=None,
|
||||
)
|
||||
for result in response_json.get("organic_results", [])
|
||||
]
|
||||
|
||||
return SearchResponse(
|
||||
results=results,
|
||||
object="search",
|
||||
)
|
||||
|
|
@ -4317,6 +4317,7 @@ class SearchProviders(str, Enum):
|
|||
DUCKDUCKGO = "duckduckgo"
|
||||
SEARCHAPI = "searchapi"
|
||||
SERPER = "serper"
|
||||
SERPAPI = "serpapi"
|
||||
YOU_COM = "you_com"
|
||||
APISERPENT = "apiserpent"
|
||||
TINYFISH = "tinyfish"
|
||||
|
|
|
|||
|
|
@ -9950,6 +9950,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.serpapi.search.transformation import SerpApiSearchConfig
|
||||
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
|
||||
|
|
@ -9969,6 +9970,7 @@ class ProviderConfigManager:
|
|||
SearchProviders.DUCKDUCKGO: DuckDuckGoSearchConfig,
|
||||
SearchProviders.SEARCHAPI: SearchAPIConfig,
|
||||
SearchProviders.SERPER: SerperSearchConfig,
|
||||
SearchProviders.SERPAPI: SerpApiSearchConfig,
|
||||
SearchProviders.YOU_COM: YouComSearchConfig,
|
||||
SearchProviders.APISERPENT: APISerpentSearchConfig,
|
||||
SearchProviders.TINYFISH: TinyfishSearchConfig,
|
||||
|
|
|
|||
0
tests/unit/llms/serpapi/__init__.py
Normal file
0
tests/unit/llms/serpapi/__init__.py
Normal file
0
tests/unit/llms/serpapi/search/__init__.py
Normal file
0
tests/unit/llms/serpapi/search/__init__.py
Normal file
|
|
@ -0,0 +1,454 @@
|
|||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse
|
||||
from litellm.llms.serpapi.search.transformation import SerpApiSearchConfig
|
||||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
|
||||
|
||||
_MINIMAL_ORGANIC: Final = [
|
||||
{
|
||||
"position": 1,
|
||||
"title": "Coffee - Wikipedia",
|
||||
"link": "https://en.wikipedia.org/wiki/Coffee",
|
||||
"snippet": "Coffee is a beverage prepared from roasted coffee beans.",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
class TestSerpApiSearchConfig:
|
||||
def test_ui_friendly_name(self):
|
||||
assert SerpApiSearchConfig().ui_friendly_name() == "SerpApi"
|
||||
|
||||
def test_get_http_method(self):
|
||||
assert SerpApiSearchConfig().get_http_method() == "GET"
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_validate_environment_with_api_key(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
result = config.validate_environment({}, api_key="test_key")
|
||||
assert result["Content-Type"] == "application/json"
|
||||
|
||||
def test_validate_environment_without_api_key(self, monkeypatch):
|
||||
monkeypatch.delenv("SERPAPI_API_KEY", raising=False)
|
||||
config = SerpApiSearchConfig()
|
||||
with pytest.raises(ValueError, match="SERPAPI_API_KEY is not set"):
|
||||
config.validate_environment({})
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_basic(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="python programming",
|
||||
optional_params={},
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
assert "_serpapi_params" in result
|
||||
params = result["_serpapi_params"]
|
||||
assert params["q"] == "python programming"
|
||||
assert params["engine"] == "google"
|
||||
assert params["api_key"] == "test_key"
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_default_engine_is_google(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="test",
|
||||
optional_params={},
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
assert result["_serpapi_params"]["engine"] == "google"
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_explicit_engine_google_news(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="AI news",
|
||||
optional_params={"engine": "google_news"},
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
assert result["_serpapi_params"]["engine"] == "google_news"
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_explicit_engine_google_shopping(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="laptop",
|
||||
optional_params={"engine": "google_shopping"},
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
assert result["_serpapi_params"]["engine"] == "google_shopping"
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_with_max_results(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="test",
|
||||
optional_params={"max_results": 5},
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
assert result["_serpapi_params"]["num"] == 5
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_with_country(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="test",
|
||||
optional_params={"country": "US"},
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
assert result["_serpapi_params"]["gl"] == "us"
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_with_domain_filter(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="machine learning",
|
||||
optional_params={"search_domain_filter": ["arxiv.org", "nature.com"]},
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
q = result["_serpapi_params"]["q"]
|
||||
assert "site:arxiv.org" in q
|
||||
assert "site:nature.com" in q
|
||||
assert "machine learning" in q
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_with_list_query(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query=["python", "programming"],
|
||||
optional_params={},
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
assert result["_serpapi_params"]["q"] == "python programming"
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_passthrough_extra_params(self, mock_secret):
|
||||
mock_secret.return_value = "test_key"
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="test",
|
||||
optional_params={"hl": "fr", "device": "mobile"},
|
||||
api_key="test_key",
|
||||
)
|
||||
|
||||
params = result["_serpapi_params"]
|
||||
assert params["hl"] == "fr"
|
||||
assert params["device"] == "mobile"
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_get_complete_url_builds_query_string(self, mock_secret):
|
||||
mock_secret.return_value = None
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
data: Final = {
|
||||
"_serpapi_params": {
|
||||
"engine": "google",
|
||||
"q": "test query",
|
||||
"api_key": "test_key",
|
||||
}
|
||||
}
|
||||
|
||||
url = config.get_complete_url(api_base=None, optional_params={}, data=data)
|
||||
|
||||
assert "https://serpapi.com/search.json?" in url
|
||||
assert "engine=google" in url
|
||||
assert "q=test+query" in url
|
||||
assert "api_key=test_key" in url
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_get_complete_url_without_data(self, mock_secret):
|
||||
mock_secret.return_value = None
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
url = config.get_complete_url(api_base=None, optional_params={}, data=None)
|
||||
|
||||
assert url == "https://serpapi.com/search.json"
|
||||
|
||||
def test_transform_search_response_normalizes_organic_results(self):
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
mock_response = Mock(spec=httpx.Response)
|
||||
mock_response.json.return_value = {
|
||||
"organic_results": [
|
||||
{
|
||||
"position": 1,
|
||||
"title": "Test Result 1",
|
||||
"link": "https://example.com/1",
|
||||
"snippet": "Snippet for result 1",
|
||||
"date": "Jan 15, 2025",
|
||||
},
|
||||
{
|
||||
"position": 2,
|
||||
"title": "Test Result 2",
|
||||
"link": "https://example.com/2",
|
||||
"snippet": "Snippet for result 2",
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
result = config.transform_search_response(
|
||||
raw_response=mock_response,
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert isinstance(result, SearchResponse)
|
||||
assert result.object == "search"
|
||||
assert len(result.results) == 2
|
||||
|
||||
assert result.results[0].title == "Test Result 1"
|
||||
assert result.results[0].url == "https://example.com/1"
|
||||
assert result.results[0].snippet == "Snippet for result 1"
|
||||
assert result.results[0].date == "Jan 15, 2025"
|
||||
assert result.results[0].last_updated is None
|
||||
|
||||
assert result.results[1].title == "Test Result 2"
|
||||
assert result.results[1].url == "https://example.com/2"
|
||||
assert result.results[1].date is None
|
||||
|
||||
def test_transform_search_response_empty_results(self):
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
mock_response = Mock(spec=httpx.Response)
|
||||
mock_response.json.return_value = {"organic_results": []}
|
||||
|
||||
result = config.transform_search_response(
|
||||
raw_response=mock_response,
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert isinstance(result, SearchResponse)
|
||||
assert len(result.results) == 0
|
||||
|
||||
def test_transform_search_response_missing_organic_results_key(self):
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
mock_response = Mock(spec=httpx.Response)
|
||||
mock_response.json.return_value = {"search_metadata": {"status": "Success"}}
|
||||
|
||||
result = config.transform_search_response(
|
||||
raw_response=mock_response,
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert isinstance(result, SearchResponse)
|
||||
assert len(result.results) == 0
|
||||
|
||||
def test_transform_search_response_missing_optional_fields(self):
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
mock_response = Mock(spec=httpx.Response)
|
||||
mock_response.json.return_value = {
|
||||
"organic_results": [
|
||||
{
|
||||
"link": "https://example.com",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
result = config.transform_search_response(
|
||||
raw_response=mock_response,
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert result.results[0].title == ""
|
||||
assert result.results[0].snippet == ""
|
||||
assert result.results[0].url == "https://example.com"
|
||||
assert result.results[0].date is None
|
||||
|
||||
@patch("litellm.llms.serpapi.search.transformation.get_secret_str")
|
||||
def test_missing_api_key_in_transform_raises(self, mock_secret):
|
||||
mock_secret.return_value = None
|
||||
config = SerpApiSearchConfig()
|
||||
|
||||
with pytest.raises(ValueError, match="SERPAPI_API_KEY is not set"):
|
||||
config.transform_search_request(
|
||||
query="test",
|
||||
optional_params={},
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
|
||||
class TestSerpApiIntegration:
|
||||
@pytest.mark.asyncio
|
||||
async def test_asearch_endpoint_and_auth(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("SERPAPI_API_KEY", "test-api-key")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"organic_results": _MINIMAL_ORGANIC,
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get:
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
response = await litellm.asearch(
|
||||
query="coffee",
|
||||
search_provider="serpapi",
|
||||
max_results=1,
|
||||
)
|
||||
|
||||
assert mock_get.call_count == 1
|
||||
call_args = mock_get.call_args
|
||||
url: str = call_args.kwargs.get("url") or call_args.args[0]
|
||||
assert "serpapi.com/search.json" in url
|
||||
assert "api_key=test-api-key" in url
|
||||
assert "q=coffee" in url
|
||||
assert "engine=google" in url
|
||||
assert "num=1" in url
|
||||
|
||||
assert response.object == "search"
|
||||
assert len(response.results) == 1
|
||||
assert response.results[0].title == "Coffee - Wikipedia"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_asearch_explicit_engine_selection(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("SERPAPI_API_KEY", "test-api-key")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"organic_results": []}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get:
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
await litellm.asearch(
|
||||
query="AI news",
|
||||
search_provider="serpapi",
|
||||
engine="google_news",
|
||||
)
|
||||
|
||||
call_args = mock_get.call_args
|
||||
url: str = call_args.kwargs.get("url") or call_args.args[0]
|
||||
assert "engine=google_news" in url
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_asearch_google_shopping_engine(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("SERPAPI_API_KEY", "test-api-key")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"organic_results": []}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get:
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
await litellm.asearch(
|
||||
query="laptop",
|
||||
search_provider="serpapi",
|
||||
engine="google_shopping",
|
||||
)
|
||||
|
||||
call_args = mock_get.call_args
|
||||
url: str = call_args.kwargs.get("url") or call_args.args[0]
|
||||
assert "engine=google_shopping" in url
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_asearch_empty_results(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("SERPAPI_API_KEY", "test-api-key")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_get:
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
response = await litellm.asearch(
|
||||
query="xyznonexistent12345",
|
||||
search_provider="serpapi",
|
||||
)
|
||||
|
||||
assert response.object == "search"
|
||||
assert len(response.results) == 0
|
||||
|
||||
|
||||
class TestSerperRegressionWithSharedCode:
|
||||
@pytest.mark.asyncio
|
||||
async def test_serper_unaffected_by_serpapi_registration(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("SERPER_API_KEY", "test-serper-key")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"organic": [
|
||||
{
|
||||
"title": "Serper Result",
|
||||
"link": "https://example.com",
|
||||
"snippet": "Snippet",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_post:
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
response = await litellm.asearch(
|
||||
query="test",
|
||||
search_provider="serper",
|
||||
)
|
||||
|
||||
assert mock_post.call_count == 1
|
||||
call_args = mock_post.call_args
|
||||
url: str = call_args.kwargs.get("url") or call_args.args[0]
|
||||
assert "google.serper.dev/search" in url
|
||||
|
||||
headers = call_args.kwargs.get("headers", {})
|
||||
assert "X-API-KEY" in headers
|
||||
assert headers["X-API-KEY"] == "test-serper-key"
|
||||
|
||||
assert len(response.results) == 1
|
||||
assert response.results[0].title == "Serper Result"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _vcr_outcome_gate(request, vcr):
|
||||
install_live_call_probe(request, vcr)
|
||||
yield
|
||||
record_vcr_outcome(request, vcr)
|
||||
Loading…
Add table
Reference in a new issue