From 09941dd1d1bd9950adbefb9b58cdfc2ea5d980aa Mon Sep 17 00:00:00 2001 From: Sampson Date: Tue, 20 Jan 2026 21:23:29 -0600 Subject: [PATCH] add search provider for brave search api (#19433) * add search provider for brave search api Introduces a minimal implementation of the Brave Search API as a search provider. Additionally, this PR introduces a test file to ensure the provider works properly, and numerous other smaller changes (e.g., changes to docs to mention the new option). * Update transformation.py --- docs/my-website/docs/search/brave.md | 55 ++++ docs/my-website/docs/search/index.md | 10 +- docs/my-website/sidebars.js | 1 + litellm/llms/brave/search/__init__.py | 7 + litellm/llms/brave/search/transformation.py | 307 ++++++++++++++++++ litellm/types/utils.py | 1 + litellm/utils.py | 2 + provider_endpoints_support.json | 17 + .../enforce_llms_folder_style.py | 1 + tests/search_tests/test_brave_search.py | 98 ++++++ 10 files changed, 497 insertions(+), 2 deletions(-) create mode 100644 docs/my-website/docs/search/brave.md create mode 100644 litellm/llms/brave/search/__init__.py create mode 100644 litellm/llms/brave/search/transformation.py create mode 100644 tests/search_tests/test_brave_search.py diff --git a/docs/my-website/docs/search/brave.md b/docs/my-website/docs/search/brave.md new file mode 100644 index 0000000000..d43efd47cd --- /dev/null +++ b/docs/my-website/docs/search/brave.md @@ -0,0 +1,55 @@ +# Brave Search + +Get started by creating a free API key via https://brave.com/search/api/. + +For documentation on other parameters supported by the Brave Search API, visit https://api-dashboard.search.brave.com/api-reference/web/search. + +## LiteLLM Python SDK + +```python showLineNumbers title="Brave Search" +import os +from litellm import search + +os.environ["BRAVE_API_KEY"] = "BSATzx..." + +response = search( + query="Brave browser features", + search_provider="brave", + max_results=5 +) +``` + +## 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: brave-search + litellm_params: + search_provider: brave + api_key: os.environ/BRAVE_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/brave-search \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ "query": "Brave browser features", "max_results": 5 }' +``` diff --git a/docs/my-website/docs/search/index.md b/docs/my-website/docs/search/index.md index 037a1b5938..551a495261 100644 --- a/docs/my-website/docs/search/index.md +++ b/docs/my-website/docs/search/index.md @@ -2,7 +2,7 @@ | Feature | Supported | |---------|-----------| -| Supported Providers | `perplexity`, `tavily`, `parallel_ai`, `exa_ai`, `google_pse`, `dataforseo`, `firecrawl`, `searxng`, `linkup` | +| Supported Providers | `perplexity`, `tavily`, `parallel_ai`, `exa_ai`, `brave`, `google_pse`, `dataforseo`, `firecrawl`, `searxng`, `linkup` | | Cost Tracking | ✅ | | Logging | ✅ | | Load Balancing | ❌ | @@ -162,6 +162,11 @@ search_tools: search_provider: exa_ai api_key: os.environ/EXA_API_KEY + - search_tool_name: my-search + litellm_params: + search_provider: brave + api_key: os.environ/BRAVE_API_KEY + router_settings: routing_strategy: simple-shuffle # or 'least-busy', 'latency-based-routing' ``` @@ -205,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"`, `"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"`, or `"linkup"` | | `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) | @@ -264,6 +269,7 @@ The response follows Perplexity's search format with the following structure: | Perplexity AI | `PERPLEXITYAI_API_KEY` | `perplexity` | | Tavily | `TAVILY_API_KEY` | `tavily` | | Exa AI | `EXA_API_KEY` | `exa_ai` | +| Brave Search | `BRAVE_API_KEY` | `brave` | | Parallel AI | `PARALLEL_AI_API_KEY` | `parallel_ai` | | Google PSE | `GOOGLE_PSE_API_KEY`, `GOOGLE_PSE_ENGINE_ID` | `google_pse` | | DataForSEO | `DATAFORSEO_LOGIN`, `DATAFORSEO_PASSWORD` | `dataforseo` | diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index dc4d7ce619..1a8a53774c 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -562,6 +562,7 @@ const sidebars = { "search/perplexity", "search/tavily", "search/exa_ai", + "search/brave", "search/parallel_ai", "search/google_pse", "search/dataforseo", diff --git a/litellm/llms/brave/search/__init__.py b/litellm/llms/brave/search/__init__.py new file mode 100644 index 0000000000..cc1168d7ef --- /dev/null +++ b/litellm/llms/brave/search/__init__.py @@ -0,0 +1,7 @@ +""" +Brave Search API module. +""" + +from litellm.llms.brave.search.transformation import BraveSearchConfig + +__all__ = ["BraveSearchConfig"] diff --git a/litellm/llms/brave/search/transformation.py b/litellm/llms/brave/search/transformation.py new file mode 100644 index 0000000000..a73029b040 --- /dev/null +++ b/litellm/llms/brave/search/transformation.py @@ -0,0 +1,307 @@ +""" +Brave Search /web/search endpoint. +Documentation: https://api-dashboard.search.brave.com/app/documentation/web-search/get-started +""" + +from __future__ import annotations +from datetime import datetime, timezone +from dateutil import parser +from typing import Dict, List, Literal, Optional, TypedDict, Union +import httpx +import re + +_ISO_YMD = re.compile(r"^\s*\d{4}[-/]\d{1,2}[-/]\d{1,2}\s*$") +_UNIX_TIMESTAMP = re.compile(r"^\s*-?\d+(\.\d+)?\s*$") +BRAVE_SECTIONS = ["web", "discussions", "faqs", "faq", "news", "videos"] + +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 + + +def to_yyyy_mm_dd( + s: Union[str, int, float, None], + *, + dayfirst: bool = False, + yearfirst: bool = False, +) -> Optional[str]: + """ + Convert a string/int/float to YYYY-MM-DD; return None if parsing fails. + """ + if not s: + return None + + s = str(s).strip() + + # Handle Unix timestamps (seconds or milliseconds). + if _UNIX_TIMESTAMP.match(s): + try: + ts_float = float(s) + # Treat large values as milliseconds. + if ts_float > 1e11 or ts_float < -1e11: + ts_float /= 1000.0 + return datetime.fromtimestamp(ts_float, tz=timezone.utc).date().isoformat() + except Exception: + return None + + # If it looks like YYYY-M-D (ISO-ish), force yearfirst to avoid surprises. + try: + if _ISO_YMD.match(s): + dt = parser.parse(s, yearfirst=True, dayfirst=False, fuzzy=True) + else: + dt = parser.parse(s, yearfirst=yearfirst, dayfirst=dayfirst, fuzzy=True) + return dt.date().isoformat() + except Exception: + return None + + +class _BraveSearchRequestRequired(TypedDict): + """Required fields for Brave Search API request.""" + + q: str # Required - search query + + +class BraveSearchRequest(_BraveSearchRequestRequired, total=False): + """ + Brave Search API request format. + Based on: https://api-dashboard.search.brave.com/app/documentation/web-search/get-started + """ + + count: int # Optional - number of web results to return (Brave max is 20) + offset: int # Optional - pagination offset + country: str # Optional - two-letter ISO country code + search_lang: str # Optional - language to bias results + ui_lang: str # Optional - language for UI strings + freshness: str # Optional - Brave freshness window (e.g., "pd", "pw", "pm") + safesearch: str # Optional - "off" | "moderate" | "strict" + spellcheck: str # Optional - "strict" | "moderate" | "off" + text_decorations: bool # Optional - enable/disable text decorations + result_filter: str # Optional - e.g., "web" + units: str # Optional - measurement units + goggles_id: str # Optional - Brave Goggles id + goggles: str # Optional - Brave Goggles DSL + extra_snippets: bool # Optional - request extra snippets + summary: bool # Optional - include summary block + enable_rich_callback: bool # Optional - structured result blocks + include_fetch_metadata: bool # Optional - include fetch metadata + operators: bool # Optional - enable advanced operators + + +class BraveSearchConfig(BaseSearchConfig): + BRAVE_API_BASE = "https://api.search.brave.com/res/v1/web/search" + + @staticmethod + def ui_friendly_name() -> str: + return "Brave Search" + + def get_http_method(self) -> Literal["GET", "POST"]: + """ + Brave Search API 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("BRAVE_API_KEY") + + if not api_key: + raise ValueError( + "BRAVE_API_KEY is not set. Set `BRAVE_API_KEY` environment variable." + ) + + headers["X-Subscription-Token"] = api_key + headers["Accept"] = "application/json" + headers["Accept-Encoding"] = "gzip" + 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. + + The Brave Search API uses GET requests and therefore needs the request + body (data) to construct query parameters in the URL. + """ + from urllib.parse import urlencode + + api_base = api_base or get_secret_str("BRAVE_API_BASE") or self.BRAVE_API_BASE + + # Build query parameters from the transformed request body + if data and isinstance(data, dict) and "_brave_params" in data: + params = data["_brave_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 Brave Search API format. + + Transforms Perplexity unified spec parameters: + - query → q (same) + - max_results → count + - search_domain_filter → q (append domain filters) + - country → country + - max_tokens_per_page → (not applicable, ignored) + + All other Brave Search API-specific parameters are passed through as-is. + + Args: + query: Search query (string or list of strings). Brave Search API supports single string queries. + optional_params: Optional parameters for the request + + Returns: + Dict with typed request data following Brave Search API spec + """ + if isinstance(query, list): + # Brave Search API only supports single string queries + query = " ".join(query) + + request_data: BraveSearchRequest = { + "q": query, + } + + # Only include "include_fetch_metadata" if it is not explicitly set to False + # This parameter results (more often than not) in a timestamp which we can use for last_updated + if ( + "include_fetch_metadata" in optional_params + and optional_params["include_fetch_metadata"] is False + ): + request_data["include_fetch_metadata"] = False + else: + request_data["include_fetch_metadata"] = True + + # Transform unified spec parameters to Brave Search API format + if "max_results" in optional_params: + # Brave Search API supports 1-20 results per /web/search request + num_results = min(optional_params["max_results"], 20) + request_data["count"] = num_results + + if "search_domain_filter" in optional_params: + # Convert to multiple "site:domain" clauses, joined by OR + domains = optional_params["search_domain_filter"] + if isinstance(domains, list) and len(domains) > 0: + request_data["q"] = self._append_domain_filters( + request_data["q"], domains + ) + + # Convert to dict before dynamic key assignments + result_data = 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 + + # Store params in special key for URL building (Brave Search API uses GET not POST) + # Return a wrapper dict that stores params for get_complete_url to use + return { + "_brave_params": result_data, + } + + @staticmethod + def _append_domain_filters(query: str, domains: List[str]) -> str: + """ + Add site: filters to emulate domain restriction in Brave. + """ + 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 Brave Search API response to LiteLLM unified SearchResponse format. + """ + response_json = raw_response.json() + + # Transform results to SearchResult objects + results: List[SearchResult] = [] + + query_params = raw_response.request.url.params if raw_response.request else {} + sections_to_process = self._sections_from_params(dict(query_params)) + max_results = max(1, min(int(query_params.get("count", 20)), 20)) + + for section in sections_to_process: + for result in response_json.get(section, {}).get("results", []): + # Because the `max_results`/`count` parameters do not affect + # the number of "discussion", "faq", "news", or "videos" + # results, we need to manually limit the number of results + # returned when an explicit limit has been provided. + if len(results) >= max_results: + break + + title = result.get("title", "") + url = result.get("url", "") + snippet = result.get("description", "") + date = to_yyyy_mm_dd(result.get("page_age") or result.get("age")) + last_updated = to_yyyy_mm_dd( + result.get("fetched_content_timestamp", "") + ) + + search_result = SearchResult( + title=title, + url=url, + snippet=snippet, + date=date, + last_updated=last_updated, + ) + + results.append(search_result) + + return SearchResponse( + results=results, + object="search", + ) + + @staticmethod + def _sections_from_params(query_params: dict) -> List[str]: + """ + Returns a list of sections the user has requested via the Brave Search + API's `result_filter` parameter. If no `result_filter` parameter is + provided, returns all sections. + """ + raw_filter = query_params.get("result_filter") + requested_filters: List[str] = [] + + if raw_filter and isinstance(raw_filter, str): + requested_filters = [part.strip() for part in raw_filter.split(",")] + + sections = [s.lower() for s in requested_filters if s.lower() in BRAVE_SECTIONS] + return sections or BRAVE_SECTIONS diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f5e217d8b4..53e0a8f185 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3117,6 +3117,7 @@ class SearchProviders(str, Enum): TAVILY = "tavily" PARALLEL_AI = "parallel_ai" EXA_AI = "exa_ai" + BRAVE = "brave" GOOGLE_PSE = "google_pse" DATAFORSEO = "dataforseo" FIRECRAWL = "firecrawl" diff --git a/litellm/utils.py b/litellm/utils.py index 3e88c6fe9e..3dea642fc6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8648,6 +8648,7 @@ class ProviderConfigManager: """ from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig + from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.firecrawl.search.transformation import FirecrawlSearchConfig from litellm.llms.google_pse.search.transformation import GooglePSESearchConfig from litellm.llms.linkup.search.transformation import LinkupSearchConfig @@ -8663,6 +8664,7 @@ class ProviderConfigManager: SearchProviders.TAVILY: TavilySearchConfig, SearchProviders.PARALLEL_AI: ParallelAISearchConfig, SearchProviders.EXA_AI: ExaAISearchConfig, + SearchProviders.BRAVE: BraveSearchConfig, SearchProviders.GOOGLE_PSE: GooglePSESearchConfig, SearchProviders.DATAFORSEO: DataForSEOSearchConfig, SearchProviders.FIRECRAWL: FirecrawlSearchConfig, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index b7892a67a1..441e27af98 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -759,6 +759,23 @@ "search": true } }, + "brave": { + "display_name": "Brave Search (`brave`)", + "url": "https://docs.litellm.ai/docs/search/brave", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true + } + }, "empower": { "display_name": "Empower (`empower`)", "url": "https://docs.litellm.ai/docs/providers/empower", diff --git a/tests/code_coverage_tests/enforce_llms_folder_style.py b/tests/code_coverage_tests/enforce_llms_folder_style.py index 715e0258f0..f684d884a6 100644 --- a/tests/code_coverage_tests/enforce_llms_folder_style.py +++ b/tests/code_coverage_tests/enforce_llms_folder_style.py @@ -12,6 +12,7 @@ SEARCH_PROVIDERS = [ "google_pse", "parallel_ai", "exa_ai", + "brave", "firecrawl", "searxng", "linkup", diff --git a/tests/search_tests/test_brave_search.py b/tests/search_tests/test_brave_search.py new file mode 100644 index 0000000000..382dab2b3c --- /dev/null +++ b/tests/search_tests/test_brave_search.py @@ -0,0 +1,98 @@ +""" +Tests for Brave Search API integration. +""" + +import os +import pytest +from urllib.parse import urlparse, parse_qs +from unittest.mock import AsyncMock, patch, MagicMock + +import litellm +from tests.search_tests.base_search_unit_tests import BaseSearchTest + + +class TestBraveSearch(BaseSearchTest): + """ + Tests for Brave Search functionality with mocked network responses. + """ + + def get_search_provider(self) -> str: + """Return the search provider name""" + return "brave" + + @pytest.mark.asyncio + async def test_basic_search(self): + """ + Test basic search functionality with a simple query. + """ + os.environ["BRAVE_API_KEY"] = "test-api-key" + + # Create a mock response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "web": { + "results": [ + { + "title": "Test Result 1", + "url": "https://example.com/1", + "description": "This is a test snippet for result 1", + } + ] + } + } + + # Mock the httpx AsyncClient get method + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", + new_callable=AsyncMock, + ) as mock_get: + mock_get.return_value = mock_response + + # Make the search call + response = await litellm.asearch( + query="Brave browser features", + search_provider="brave", + max_results=5, + result_filter="web", + ) + + # Verify the get method was called once + assert mock_get.call_count == 1 + + # Get the actual call arguments + call_args = mock_get.call_args + + # Verify URL (include_fetch_metadata=True is added by default) + parsed_url = urlparse(call_args.kwargs["url"]) + assert parsed_url.scheme == "https" + assert parsed_url.netloc == "api.search.brave.com" + assert parsed_url.path == "/res/v1/web/search" + + query_params = parse_qs(parsed_url.query) + assert query_params == { + "q": ["Brave browser features"], + "include_fetch_metadata": ["True"], + "count": ["5"], + "result_filter": ["web"], + } + + # Verify headers contains X-Subscription-Token + headers = call_args.kwargs.get("headers", {}) + assert "X-Subscription-Token" in headers + assert headers["X-Subscription-Token"] == "test-api-key" + + # Note: Brave uses GET requests, so parameters are in the URL, not in JSON body + # The URL already contains all the parameters we need to verify + + # Verify response structure + assert hasattr(response, "results") + assert hasattr(response, "object") + assert response.object == "search" + assert len(response.results) == 1 + + # Verify first result + first_result = response.results[0] + assert first_result.title == "Test Result 1" + assert first_result.url == "https://example.com/1" + assert first_result.snippet == "This is a test snippet for result 1"