Merge pull request #22752 from BerriAI/litellm_search_api_add
[Feat] Add Google Search API Integration
This commit is contained in:
Sameer Kankute
GitHub
commit
b5183e9f3b
@ -2,7 +2,7 @@
|
||||
|
||||
| Feature | Supported |
|
||||
|---------|-----------|
|
||||
| Supported Providers | `perplexity`, `tavily`, `parallel_ai`, `exa_ai`, `brave`, `google_pse`, `dataforseo`, `firecrawl`, `searxng`, `linkup` |
|
||||
| Supported Providers | `perplexity`, `tavily`, `parallel_ai`, `exa_ai`, `brave`, `google_pse`, `dataforseo`, `firecrawl`, `searxng`, `linkup`, `duckduckgo`, `searchapi` |
|
||||
| Cost Tracking | ✅ |
|
||||
| Logging | ✅ |
|
||||
| Load Balancing | ❌ |
|
||||
@ -210,7 +210,7 @@ See the [official Perplexity Search documentation](https://docs.perplexity.ai/ap
|
||||
| Parameter | Type | Required | Description |
|
||||
|-----------|------|----------|-------------|
|
||||
| `query` | string or array | Yes | Search query. Can be a single string or array of strings |
|
||||
| `search_provider` | string | Yes (SDK) | The search provider to use: `"perplexity"`, `"tavily"`, `"parallel_ai"`, `"exa_ai"`, `"brave"`, `"google_pse"`, `"dataforseo"`, `"firecrawl"`, `"searxng"`, or `"linkup"` |
|
||||
| `search_provider` | string | Yes (SDK) | The search provider to use: `"perplexity"`, `"tavily"`, `"parallel_ai"`, `"exa_ai"`, `"brave"`, `"google_pse"`, `"dataforseo"`, `"firecrawl"`, `"searxng"`, `"linkup"`, `"duckduckgo"`, or `"searchapi"` |
|
||||
| `search_tool_name` | string | Yes (Proxy) | Name of the search tool configured in `config.yaml` |
|
||||
| `max_results` | integer | No | Maximum number of results to return (1-20). Default: 10 |
|
||||
| `search_domain_filter` | array | No | List of domains to filter results (max 20 domains) |
|
||||
@ -276,7 +276,8 @@ The response follows Perplexity's search format with the following structure:
|
||||
| Firecrawl | `FIRECRAWL_API_KEY` | `firecrawl` |
|
||||
| SearXNG | `SEARXNG_API_BASE` (required) | `searxng` |
|
||||
| Linkup | `LINKUP_API_KEY` | `linkup` |
|
||||
| DuckDuckGo | `DUCKDUCKGO_API_BASE` | `duckduckgo` |
|
||||
| DuckDuckGo | `DUCKDUCKGO_API_BASE` | `duckduckgo` |
|
||||
| SearchAPI.io | `SEARCHAPI_API_KEY` | `searchapi` |
|
||||
|
||||
See the individual provider documentation for detailed setup instructions and provider-specific parameters.
|
||||
|
||||
|
||||
197
docs/my-website/docs/search/searchapi.md
Normal file
197
docs/my-website/docs/search/searchapi.md
Normal file
@ -0,0 +1,197 @@
|
||||
# SearchAPI.io (Google Search)
|
||||
|
||||
Get started by creating a free API key via https://www.searchapi.io/.
|
||||
|
||||
SearchAPI.io provides access to Google Search results with a simple API. It supports all Google Search parameters including location, language, time filters, and more.
|
||||
|
||||
For complete documentation on all supported parameters, visit https://www.searchapi.io/docs/google.
|
||||
|
||||
## LiteLLM Python SDK
|
||||
|
||||
```python showLineNumbers title="SearchAPI.io Search"
|
||||
import os
|
||||
from litellm import search
|
||||
|
||||
os.environ["SEARCHAPI_API_KEY"] = "your-api-key"
|
||||
|
||||
response = search(
|
||||
query="latest AI developments",
|
||||
search_provider="searchapi",
|
||||
max_results=10
|
||||
)
|
||||
|
||||
# Access search results
|
||||
for result in response.results:
|
||||
print(f"{result.title}: {result.url}")
|
||||
print(f"Snippet: {result.snippet}\n")
|
||||
```
|
||||
|
||||
### Advanced Usage with SearchAPI.io Parameters
|
||||
|
||||
SearchAPI.io supports many Google Search-specific parameters:
|
||||
|
||||
```python showLineNumbers title="Advanced SearchAPI.io Parameters"
|
||||
import os
|
||||
from litellm import search
|
||||
|
||||
os.environ["SEARCHAPI_API_KEY"] = "your-api-key"
|
||||
|
||||
response = search(
|
||||
query="machine learning research",
|
||||
search_provider="searchapi",
|
||||
max_results=10,
|
||||
# Unified parameters
|
||||
country="US",
|
||||
search_domain_filter=["arxiv.org", "nature.com"],
|
||||
# SearchAPI.io specific parameters
|
||||
gl="us", # Country code
|
||||
hl="en", # Interface language
|
||||
time_period="last_month", # Time filter
|
||||
safe="active", # SafeSearch
|
||||
device="desktop", # Device type
|
||||
location="New York" # Geographic location
|
||||
)
|
||||
```
|
||||
|
||||
## LiteLLM AI Gateway
|
||||
|
||||
### 1. Setup config.yaml
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
model_list:
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
search_tools:
|
||||
- search_tool_name: google-search
|
||||
litellm_params:
|
||||
search_provider: searchapi
|
||||
api_key: os.environ/SEARCHAPI_API_KEY
|
||||
```
|
||||
|
||||
### 2. Start the proxy
|
||||
|
||||
```bash
|
||||
litellm --config /path/to/config.yaml
|
||||
|
||||
# RUNNING on http://0.0.0.0:4000
|
||||
```
|
||||
|
||||
### 3. Test the search endpoint
|
||||
|
||||
```bash showLineNumbers title="Test Request"
|
||||
curl http://0.0.0.0:4000/v1/search/google-search \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"query": "latest AI developments",
|
||||
"max_results": 10,
|
||||
"country": "US"
|
||||
}'
|
||||
```
|
||||
|
||||
## SearchAPI.io Specific Parameters
|
||||
|
||||
SearchAPI.io supports many Google Search parameters. Here are some commonly used ones:
|
||||
|
||||
| Parameter | Type | Description |
|
||||
|-----------|------|-------------|
|
||||
| `gl` | string | Country code (e.g., 'us', 'uk', 'de') |
|
||||
| `hl` | string | Interface language (e.g., 'en', 'es', 'fr') |
|
||||
| `location` | string | Geographic location (e.g., 'New York', 'London') |
|
||||
| `device` | string | Device type: 'desktop', 'mobile', 'tablet' |
|
||||
| `time_period` | string | Time filter: 'last_hour', 'last_day', 'last_week', 'last_month', 'last_year' |
|
||||
| `time_period_min` | string | Start date (MM/DD/YYYY) |
|
||||
| `time_period_max` | string | End date (MM/DD/YYYY) |
|
||||
| `safe` | string | SafeSearch: 'active' or 'off' |
|
||||
| `lr` | string | Language restriction (e.g., 'lang_en', 'lang_es') |
|
||||
| `cr` | string | Country restriction |
|
||||
| `page` | integer | Page number for pagination |
|
||||
|
||||
### Example with Time Filters
|
||||
|
||||
```python showLineNumbers title="Search with Time Filter"
|
||||
response = search(
|
||||
query="AI breakthroughs",
|
||||
search_provider="searchapi",
|
||||
max_results=10,
|
||||
time_period="last_month"
|
||||
)
|
||||
```
|
||||
|
||||
### Example with Custom Date Range
|
||||
|
||||
```python showLineNumbers title="Search with Custom Date Range"
|
||||
response = search(
|
||||
query="AI research papers",
|
||||
search_provider="searchapi",
|
||||
max_results=10,
|
||||
time_period_min="01/01/2024",
|
||||
time_period_max="03/01/2024"
|
||||
)
|
||||
```
|
||||
|
||||
### Example with Location
|
||||
|
||||
```python showLineNumbers title="Search with Location"
|
||||
response = search(
|
||||
query="AI conferences",
|
||||
search_provider="searchapi",
|
||||
max_results=10,
|
||||
location="San Francisco",
|
||||
gl="us"
|
||||
)
|
||||
```
|
||||
|
||||
## Response Format
|
||||
|
||||
SearchAPI.io returns results in the standard LiteLLM search format:
|
||||
|
||||
```json
|
||||
{
|
||||
"object": "search",
|
||||
"results": [
|
||||
{
|
||||
"title": "Latest AI Developments",
|
||||
"url": "https://example.com/ai-news",
|
||||
"snippet": "Recent breakthroughs in artificial intelligence...",
|
||||
"date": "2024-01-15"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
## Rate Limits
|
||||
|
||||
SearchAPI.io has different rate limits based on your plan:
|
||||
- Free tier: 100 requests/month
|
||||
- Paid plans: Higher limits available
|
||||
|
||||
Check your current usage at https://www.searchapi.io/dashboard.
|
||||
|
||||
## Error Handling
|
||||
|
||||
```python showLineNumbers title="Error Handling"
|
||||
from litellm import search
|
||||
import os
|
||||
|
||||
os.environ["SEARCHAPI_API_KEY"] = "your-api-key"
|
||||
|
||||
try:
|
||||
response = search(
|
||||
query="test query",
|
||||
search_provider="searchapi",
|
||||
max_results=10
|
||||
)
|
||||
print(f"Found {len(response.results)} results")
|
||||
except Exception as e:
|
||||
print(f"Search failed: {str(e)}")
|
||||
```
|
||||
|
||||
## Additional Resources
|
||||
|
||||
- SearchAPI.io Documentation: https://www.searchapi.io/docs
|
||||
- API Dashboard: https://www.searchapi.io/dashboard
|
||||
- Pricing: https://www.searchapi.io/pricing
|
||||
1
litellm/llms/searchapi/__init__.py
Normal file
1
litellm/llms/searchapi/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""SearchAPI.io integration for LiteLLM."""
|
||||
4
litellm/llms/searchapi/search/__init__.py
Normal file
4
litellm/llms/searchapi/search/__init__.py
Normal file
@ -0,0 +1,4 @@
|
||||
"""SearchAPI.io search integration for LiteLLM."""
|
||||
from litellm.llms.searchapi.search.transformation import SearchAPIConfig
|
||||
|
||||
__all__ = ["SearchAPIConfig"]
|
||||
232
litellm/llms/searchapi/search/transformation.py
Normal file
232
litellm/llms/searchapi/search/transformation.py
Normal file
@ -0,0 +1,232 @@
|
||||
"""
|
||||
Calls SearchAPI.io's Google Search API endpoint.
|
||||
|
||||
SearchAPI.io API Reference: https://www.searchapi.io/docs/google
|
||||
"""
|
||||
from typing import Dict, List, Literal, Optional, TypedDict, Union
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import httpx
|
||||
|
||||
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 _SearchAPIRequestRequired(TypedDict):
|
||||
"""Required fields for SearchAPI.io request."""
|
||||
engine: str # Required - search engine (e.g., 'google')
|
||||
q: str # Required - search query
|
||||
|
||||
|
||||
class SearchAPIRequest(_SearchAPIRequestRequired, total=False):
|
||||
"""
|
||||
SearchAPI.io request format for Google Search.
|
||||
Based on: https://www.searchapi.io/docs/google
|
||||
"""
|
||||
kgmid: str # Optional - Knowledge Graph identifier
|
||||
device: str # Optional - device type ('desktop', 'mobile', 'tablet')
|
||||
location: str # Optional - geographic location
|
||||
uule: str # Optional - Google-encoded location
|
||||
google_domain: str # Optional - Google domain (deprecated)
|
||||
gl: str # Optional - country code (e.g., 'us', 'uk')
|
||||
hl: str # Optional - interface language (e.g., 'en', 'es')
|
||||
lr: str # Optional - language restriction (e.g., 'lang_en')
|
||||
cr: str # Optional - country restriction
|
||||
nfpr: int # Optional - exclude auto-corrected results (0 or 1)
|
||||
filter: int # Optional - duplicate/host crowding filter (0 or 1)
|
||||
safe: str # Optional - SafeSearch ('active', 'off')
|
||||
time_period: str # Optional - time period ('last_hour', 'last_day', 'last_week', 'last_month', 'last_year')
|
||||
time_period_min: str # Optional - start date (MM/DD/YYYY)
|
||||
time_period_max: str # Optional - end date (MM/DD/YYYY)
|
||||
num: int # Optional - number of results (phased out by Google, constant 10)
|
||||
page: int # Optional - page number for pagination
|
||||
optimization_strategy: str # Optional - 'performance' or 'ads'
|
||||
|
||||
|
||||
class SearchAPIConfig(BaseSearchConfig):
|
||||
SEARCHAPI_API_BASE = "https://www.searchapi.io/api/v1/search"
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "SearchAPI.io (Google Search)"
|
||||
|
||||
def get_http_method(self) -> Literal["GET", "POST"]:
|
||||
"""
|
||||
SearchAPI.io uses GET requests for search.
|
||||
"""
|
||||
return "GET"
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Dict,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> Dict:
|
||||
"""
|
||||
Validate environment and return headers.
|
||||
"""
|
||||
api_key = api_key or get_secret_str("SEARCHAPI_API_KEY")
|
||||
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable."
|
||||
)
|
||||
|
||||
headers["Content-Type"] = "application/json"
|
||||
|
||||
return headers
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
optional_params: dict,
|
||||
data: Optional[Union[Dict, List[Dict]]] = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""
|
||||
Get complete URL for Search endpoint with query parameters.
|
||||
|
||||
SearchAPI.io uses GET requests and includes api_key in query params.
|
||||
"""
|
||||
api_base = api_base or get_secret_str("SEARCHAPI_API_BASE") or self.SEARCHAPI_API_BASE
|
||||
|
||||
# Build query parameters from the transformed request body
|
||||
if data and isinstance(data, dict) and "_searchapi_params" in data:
|
||||
params = data["_searchapi_params"]
|
||||
query_string = urlencode(params, doseq=True)
|
||||
return f"{api_base}?{query_string}"
|
||||
|
||||
return api_base
|
||||
|
||||
def transform_search_request(
|
||||
self,
|
||||
query: Union[str, List[str]],
|
||||
optional_params: dict,
|
||||
api_key: Optional[str] = None,
|
||||
search_engine_id: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> Dict:
|
||||
"""
|
||||
Transform Search request to SearchAPI.io format.
|
||||
|
||||
Transforms unified spec parameters:
|
||||
- query → q
|
||||
- max_results → num (limited to 10 by Google)
|
||||
- search_domain_filter → q (append site: filters)
|
||||
- country → gl
|
||||
|
||||
Args:
|
||||
query: Search query (string or list of strings)
|
||||
optional_params: Optional parameters for the request
|
||||
api_key: API key for authentication
|
||||
|
||||
Returns:
|
||||
Dict with typed request data following SearchAPI.io spec
|
||||
"""
|
||||
if isinstance(query, list):
|
||||
query = " ".join(query)
|
||||
|
||||
# Get API key from parameter or environment
|
||||
api_key = api_key or get_secret_str("SEARCHAPI_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"SEARCHAPI_API_KEY is not set. Set `SEARCHAPI_API_KEY` environment variable."
|
||||
)
|
||||
|
||||
request_data: SearchAPIRequest = {
|
||||
"engine": "google",
|
||||
"q": query,
|
||||
}
|
||||
|
||||
# Add API key to request
|
||||
result_data = dict(request_data)
|
||||
result_data["api_key"] = api_key
|
||||
|
||||
# Transform unified spec parameters to SearchAPI.io format
|
||||
if "max_results" in optional_params:
|
||||
# Google now returns constant 10 results, but we can still set num
|
||||
num_results = min(optional_params["max_results"], 10)
|
||||
result_data["num"] = num_results
|
||||
|
||||
if "search_domain_filter" in optional_params:
|
||||
# Convert to multiple "site:domain" clauses
|
||||
domains = optional_params["search_domain_filter"]
|
||||
if isinstance(domains, list) and len(domains) > 0:
|
||||
result_data["q"] = self._append_domain_filters(
|
||||
result_data["q"], domains
|
||||
)
|
||||
|
||||
if "country" in optional_params:
|
||||
# Map to gl parameter
|
||||
result_data["gl"] = optional_params["country"].lower()
|
||||
|
||||
# Pass through all other SearchAPI.io-specific parameters
|
||||
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
|
||||
|
||||
# Store params in special key for URL building (GET request)
|
||||
return {
|
||||
"_searchapi_params": result_data,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _append_domain_filters(query: str, domains: List[str]) -> str:
|
||||
"""
|
||||
Add site: filters to restrict search to specific domains.
|
||||
"""
|
||||
domain_clauses = [f"site:{domain}" for domain in domains]
|
||||
domain_query = " OR ".join(domain_clauses)
|
||||
|
||||
return f"({query}) AND ({domain_query})"
|
||||
|
||||
def transform_search_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: Optional[LiteLLMLoggingObj],
|
||||
**kwargs,
|
||||
) -> SearchResponse:
|
||||
"""
|
||||
Transform SearchAPI.io response to LiteLLM unified SearchResponse format.
|
||||
|
||||
SearchAPI.io → LiteLLM mappings:
|
||||
- organic_results[].title → SearchResult.title
|
||||
- organic_results[].link → SearchResult.url
|
||||
- organic_results[].snippet → SearchResult.snippet
|
||||
- organic_results[].date → SearchResult.date
|
||||
"""
|
||||
response_json = raw_response.json()
|
||||
|
||||
# Transform results to SearchResult objects
|
||||
results: List[SearchResult] = []
|
||||
|
||||
# Process organic results
|
||||
for result in response_json.get("organic_results", []):
|
||||
title = result.get("title", "")
|
||||
url = result.get("link", "")
|
||||
snippet = result.get("snippet", "")
|
||||
date = result.get("date") # SearchAPI.io provides date in some results
|
||||
|
||||
search_result = SearchResult(
|
||||
title=title,
|
||||
url=url,
|
||||
snippet=snippet,
|
||||
date=date,
|
||||
last_updated=None, # SearchAPI.io doesn't provide last_updated
|
||||
)
|
||||
|
||||
results.append(search_result)
|
||||
|
||||
return SearchResponse(
|
||||
results=results,
|
||||
object="search",
|
||||
)
|
||||
@ -3238,6 +3238,7 @@ class SearchProviders(str, Enum):
|
||||
SEARXNG = "searxng"
|
||||
LINKUP = "linkup"
|
||||
DUCKDUCKGO = "duckduckgo"
|
||||
SEARCHAPI = "searchapi"
|
||||
|
||||
|
||||
# Create a set of all search provider values for quick lookup
|
||||
|
||||
@ -8861,6 +8861,7 @@ class ProviderConfigManager:
|
||||
ParallelAISearchConfig,
|
||||
)
|
||||
from litellm.llms.perplexity.search.transformation import PerplexitySearchConfig
|
||||
from litellm.llms.searchapi.search.transformation import SearchAPIConfig
|
||||
from litellm.llms.searxng.search.transformation import SearXNGSearchConfig
|
||||
from litellm.llms.tavily.search.transformation import TavilySearchConfig
|
||||
|
||||
@ -8876,6 +8877,7 @@ class ProviderConfigManager:
|
||||
SearchProviders.SEARXNG: SearXNGSearchConfig,
|
||||
SearchProviders.LINKUP: LinkupSearchConfig,
|
||||
SearchProviders.DUCKDUCKGO: DuckDuckGoSearchConfig,
|
||||
SearchProviders.SEARCHAPI: SearchAPIConfig,
|
||||
}
|
||||
config_class = PROVIDER_TO_CONFIG_MAP.get(provider, None)
|
||||
if config_class is None:
|
||||
|
||||
@ -17,6 +17,7 @@ SEARCH_PROVIDERS = [
|
||||
"searxng",
|
||||
"linkup",
|
||||
"duckduckgo",
|
||||
"searchapi",
|
||||
]
|
||||
|
||||
ALLOWED_FILES_IN_LLMS_FOLDER = [
|
||||
|
||||
270
tests/search_tests/test_searchapi_search.py
Normal file
270
tests/search_tests/test_searchapi_search.py
Normal file
@ -0,0 +1,270 @@
|
||||
"""
|
||||
Tests for SearchAPI.io (Google Search) integration.
|
||||
|
||||
Tests the SearchAPI.io search provider implementation including:
|
||||
- Request transformation
|
||||
- Response transformation
|
||||
- Parameter mapping
|
||||
- Error handling
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
)
|
||||
|
||||
from litellm.llms.searchapi.search.transformation import SearchAPIConfig
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
|
||||
|
||||
|
||||
class TestSearchAPIConfig:
|
||||
"""Test SearchAPI.io configuration and transformations."""
|
||||
|
||||
def test_ui_friendly_name(self):
|
||||
"""Test that UI friendly name is returned correctly."""
|
||||
config = SearchAPIConfig()
|
||||
assert config.ui_friendly_name() == "SearchAPI.io (Google Search)"
|
||||
|
||||
def test_get_http_method(self):
|
||||
"""Test that HTTP method is GET."""
|
||||
config = SearchAPIConfig()
|
||||
assert config.get_http_method() == "GET"
|
||||
|
||||
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
|
||||
def test_validate_environment_with_api_key(self, mock_get_secret):
|
||||
"""Test environment validation with API key."""
|
||||
mock_get_secret.return_value = "test_api_key"
|
||||
config = SearchAPIConfig()
|
||||
headers = {}
|
||||
|
||||
result = config.validate_environment(headers, api_key="test_api_key")
|
||||
|
||||
assert result["Content-Type"] == "application/json"
|
||||
|
||||
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
|
||||
def test_validate_environment_without_api_key(self, mock_get_secret):
|
||||
"""Test environment validation without API key raises error."""
|
||||
mock_get_secret.return_value = None
|
||||
config = SearchAPIConfig()
|
||||
headers = {}
|
||||
|
||||
with pytest.raises(ValueError, match="SEARCHAPI_API_KEY is not set"):
|
||||
config.validate_environment(headers)
|
||||
|
||||
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_basic(self, mock_get_secret):
|
||||
"""Test basic search request transformation."""
|
||||
mock_get_secret.return_value = "test_api_key"
|
||||
config = SearchAPIConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="test query",
|
||||
optional_params={},
|
||||
api_key="test_api_key"
|
||||
)
|
||||
|
||||
assert "_searchapi_params" in result
|
||||
params = result["_searchapi_params"]
|
||||
assert params["engine"] == "google"
|
||||
assert params["q"] == "test query"
|
||||
assert params["api_key"] == "test_api_key"
|
||||
|
||||
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_with_max_results(self, mock_get_secret):
|
||||
"""Test search request transformation with max_results parameter."""
|
||||
mock_get_secret.return_value = "test_api_key"
|
||||
config = SearchAPIConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="test query",
|
||||
optional_params={"max_results": 5},
|
||||
api_key="test_api_key"
|
||||
)
|
||||
|
||||
params = result["_searchapi_params"]
|
||||
assert params["num"] == 5
|
||||
|
||||
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_with_country(self, mock_get_secret):
|
||||
"""Test search request transformation with country parameter."""
|
||||
mock_get_secret.return_value = "test_api_key"
|
||||
config = SearchAPIConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="test query",
|
||||
optional_params={"country": "US"},
|
||||
api_key="test_api_key"
|
||||
)
|
||||
|
||||
params = result["_searchapi_params"]
|
||||
assert params["gl"] == "us"
|
||||
|
||||
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_with_domain_filter(self, mock_get_secret):
|
||||
"""Test search request transformation with domain filter."""
|
||||
mock_get_secret.return_value = "test_api_key"
|
||||
config = SearchAPIConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query="test query",
|
||||
optional_params={"search_domain_filter": ["example.com", "test.com"]},
|
||||
api_key="test_api_key"
|
||||
)
|
||||
|
||||
params = result["_searchapi_params"]
|
||||
assert "site:example.com" in params["q"]
|
||||
assert "site:test.com" in params["q"]
|
||||
|
||||
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
|
||||
def test_transform_search_request_with_list_query(self, mock_get_secret):
|
||||
"""Test search request transformation with list query."""
|
||||
mock_get_secret.return_value = "test_api_key"
|
||||
config = SearchAPIConfig()
|
||||
|
||||
result = config.transform_search_request(
|
||||
query=["test", "query"],
|
||||
optional_params={},
|
||||
api_key="test_api_key"
|
||||
)
|
||||
|
||||
params = result["_searchapi_params"]
|
||||
assert params["q"] == "test query"
|
||||
|
||||
@patch("litellm.llms.searchapi.search.transformation.get_secret_str")
|
||||
def test_get_complete_url(self, mock_get_secret):
|
||||
"""Test URL construction with query parameters."""
|
||||
mock_get_secret.return_value = None
|
||||
config = SearchAPIConfig()
|
||||
|
||||
data = {
|
||||
"_searchapi_params": {
|
||||
"engine": "google",
|
||||
"q": "test query",
|
||||
"api_key": "test_key"
|
||||
}
|
||||
}
|
||||
|
||||
url = config.get_complete_url(
|
||||
api_base=None,
|
||||
optional_params={},
|
||||
data=data
|
||||
)
|
||||
|
||||
assert "https://www.searchapi.io/api/v1/search?" in url
|
||||
assert "engine=google" in url
|
||||
assert "q=test+query" in url
|
||||
assert "api_key=test_key" in url
|
||||
|
||||
def test_transform_search_response(self):
|
||||
"""Test search response transformation."""
|
||||
config = SearchAPIConfig()
|
||||
|
||||
# Mock response
|
||||
mock_response = Mock(spec=httpx.Response)
|
||||
mock_response.json.return_value = {
|
||||
"organic_results": [
|
||||
{
|
||||
"title": "Test Result 1",
|
||||
"link": "https://example.com/1",
|
||||
"snippet": "This is a test snippet 1",
|
||||
"date": "2024-01-01"
|
||||
},
|
||||
{
|
||||
"title": "Test Result 2",
|
||||
"link": "https://example.com/2",
|
||||
"snippet": "This is a test snippet 2"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
result = config.transform_search_response(
|
||||
raw_response=mock_response,
|
||||
logging_obj=None
|
||||
)
|
||||
|
||||
assert isinstance(result, SearchResponse)
|
||||
assert result.object == "search"
|
||||
assert len(result.results) == 2
|
||||
|
||||
# Check first result
|
||||
assert result.results[0].title == "Test Result 1"
|
||||
assert result.results[0].url == "https://example.com/1"
|
||||
assert result.results[0].snippet == "This is a test snippet 1"
|
||||
assert result.results[0].date == "2024-01-01"
|
||||
assert result.results[0].last_updated is None
|
||||
|
||||
# Check second result
|
||||
assert result.results[1].title == "Test Result 2"
|
||||
assert result.results[1].url == "https://example.com/2"
|
||||
assert result.results[1].snippet == "This is a test snippet 2"
|
||||
assert result.results[1].date is None
|
||||
|
||||
def test_transform_search_response_empty(self):
|
||||
"""Test search response transformation with no results."""
|
||||
config = SearchAPIConfig()
|
||||
|
||||
mock_response = Mock(spec=httpx.Response)
|
||||
mock_response.json.return_value = {
|
||||
"organic_results": []
|
||||
}
|
||||
|
||||
result = config.transform_search_response(
|
||||
raw_response=mock_response,
|
||||
logging_obj=None
|
||||
)
|
||||
|
||||
assert isinstance(result, SearchResponse)
|
||||
assert len(result.results) == 0
|
||||
|
||||
def test_append_domain_filters(self):
|
||||
"""Test domain filter appending logic."""
|
||||
config = SearchAPIConfig()
|
||||
|
||||
query = "test query"
|
||||
domains = ["example.com", "test.com"]
|
||||
|
||||
result = config._append_domain_filters(query, domains)
|
||||
|
||||
assert "(test query)" in result
|
||||
assert "site:example.com" in result
|
||||
assert "site:test.com" in result
|
||||
assert "OR" in result
|
||||
assert "AND" in result
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("SEARCHAPI_API_KEY") is None,
|
||||
reason="SEARCHAPI_API_KEY not set in environment"
|
||||
)
|
||||
class TestSearchAPIIntegration:
|
||||
"""Integration tests for SearchAPI.io (requires API key)."""
|
||||
|
||||
def test_real_search_request(self):
|
||||
"""
|
||||
Test a real search request to SearchAPI.io.
|
||||
This test is skipped if SEARCHAPI_API_KEY is not set.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
response = litellm.search(
|
||||
query="Python programming",
|
||||
search_provider="searchapi",
|
||||
max_results=5
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert hasattr(response, "results")
|
||||
assert len(response.results) > 0
|
||||
assert all(hasattr(r, "title") for r in response.results)
|
||||
assert all(hasattr(r, "url") for r in response.results)
|
||||
assert all(hasattr(r, "snippet") for r in response.results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Loading…
Reference in New Issue
Block a user