392 lines
14 KiB
Python
392 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from app.services.ai_tools.schemas import SearchEvidence, WebSearchConfig, WebSearchProviderConfig
|
|
|
|
|
|
WEB_SEARCH_PROVIDER_PRESETS: dict[str, dict[str, Any]] = {
|
|
"tavily": {
|
|
"provider": "tavily",
|
|
"label": "Tavily",
|
|
"api_key_env": "TAVILY_API_KEY",
|
|
"base_url": "https://api.tavily.com",
|
|
"endpoint_path": "/search",
|
|
"max_results": 5,
|
|
"timeout_seconds": 20,
|
|
"search_depth": "basic",
|
|
"include_answer": False,
|
|
"include_raw_content": False,
|
|
},
|
|
"brave": {
|
|
"provider": "brave",
|
|
"label": "Brave Search API",
|
|
"api_key_env": "BRAVE_SEARCH_API_KEY",
|
|
"base_url": "https://api.search.brave.com",
|
|
"endpoint_path": "/res/v1/web/search",
|
|
"max_results": 5,
|
|
"timeout_seconds": 20,
|
|
},
|
|
"serpapi": {
|
|
"provider": "serpapi",
|
|
"label": "SerpAPI",
|
|
"api_key_env": "SERPAPI_API_KEY",
|
|
"base_url": "https://serpapi.com",
|
|
"endpoint_path": "/search.json",
|
|
"engine": "google",
|
|
"max_results": 5,
|
|
"timeout_seconds": 20,
|
|
},
|
|
"exa": {
|
|
"provider": "exa",
|
|
"label": "Exa",
|
|
"api_key_env": "EXA_API_KEY",
|
|
"base_url": "https://api.exa.ai",
|
|
"endpoint_path": "/search",
|
|
"max_results": 5,
|
|
"timeout_seconds": 20,
|
|
"include_text": False,
|
|
},
|
|
"firecrawl": {
|
|
"provider": "firecrawl",
|
|
"label": "Firecrawl Search / Scrape",
|
|
"api_key_env": "FIRECRAWL_API_KEY",
|
|
"base_url": "https://api.firecrawl.dev",
|
|
"search_path": "/v2/search",
|
|
"scrape_path": "/v2/scrape",
|
|
"max_results": 5,
|
|
"timeout_seconds": 30,
|
|
"scrape_formats": ["markdown"],
|
|
},
|
|
"searxng": {
|
|
"provider": "searxng",
|
|
"label": "SearXNG",
|
|
"api_key_env": "SEARXNG_API_KEY",
|
|
"base_url": "http://localhost:8080",
|
|
"endpoint_path": "/",
|
|
"max_results": 5,
|
|
"timeout_seconds": 20,
|
|
"categories": "general",
|
|
"engines": [],
|
|
},
|
|
}
|
|
|
|
|
|
class WebSearchError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class WebSearchConfigurationError(WebSearchError):
|
|
pass
|
|
|
|
|
|
def normalize_web_search_provider(provider: str | None) -> str:
|
|
return (provider or "tavily").strip().lower() or "tavily"
|
|
|
|
|
|
def get_web_search_provider_preset(provider: str) -> dict[str, Any]:
|
|
provider_id = normalize_web_search_provider(provider)
|
|
preset = WEB_SEARCH_PROVIDER_PRESETS.get(provider_id)
|
|
if not preset:
|
|
raise ValueError(f"Unsupported web search provider: {provider}")
|
|
return deepcopy(preset)
|
|
|
|
|
|
def list_web_search_provider_presets() -> list[dict[str, Any]]:
|
|
return [get_web_search_provider_preset(provider) for provider in WEB_SEARCH_PROVIDER_PRESETS]
|
|
|
|
|
|
def provider_defaults(provider: str) -> WebSearchProviderConfig:
|
|
preset = get_web_search_provider_preset(provider)
|
|
return WebSearchProviderConfig(**{
|
|
key: value
|
|
for key, value in preset.items()
|
|
if key in WebSearchProviderConfig.model_fields
|
|
})
|
|
|
|
|
|
class WebSearchClient:
|
|
def __init__(self, config: WebSearchConfig) -> None:
|
|
self.config = config
|
|
|
|
async def search(
|
|
self,
|
|
query: str,
|
|
*,
|
|
max_results: int | None = None,
|
|
domains: list[str] | None = None,
|
|
freshness_days: int | None = None,
|
|
) -> list[SearchEvidence]:
|
|
if not self.config.enabled:
|
|
raise WebSearchConfigurationError("WebSearch is disabled.")
|
|
provider_config = self.config.active_provider_config
|
|
provider = normalize_web_search_provider(provider_config.provider)
|
|
if provider != "searxng" and not provider_config.api_key:
|
|
raise WebSearchConfigurationError(f"{provider} API key is not configured.")
|
|
query = " ".join(str(query or "").split())
|
|
if not query:
|
|
raise WebSearchConfigurationError("search query is required.")
|
|
limit = max_results or provider_config.max_results
|
|
if provider == "tavily":
|
|
return await self._search_tavily(provider_config, query, limit, domains, freshness_days)
|
|
if provider == "brave":
|
|
return await self._search_brave(provider_config, query, limit, domains)
|
|
if provider == "serpapi":
|
|
return await self._search_serpapi(provider_config, query, limit)
|
|
if provider == "exa":
|
|
return await self._search_exa(provider_config, query, limit, domains)
|
|
if provider == "firecrawl":
|
|
return await self._search_firecrawl(provider_config, query, limit)
|
|
if provider == "searxng":
|
|
return await self._search_searxng(provider_config, query, limit, domains)
|
|
raise WebSearchConfigurationError(f"Unsupported web search provider: {provider}")
|
|
|
|
async def test_connection(self) -> list[SearchEvidence]:
|
|
return await self.search("Planet WebSearch connectivity test", max_results=1)
|
|
|
|
async def _request_json(
|
|
self,
|
|
method: str,
|
|
url: str,
|
|
*,
|
|
provider_config: WebSearchProviderConfig,
|
|
headers: dict[str, str] | None = None,
|
|
params: dict[str, Any] | None = None,
|
|
json: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=provider_config.timeout_seconds) as client:
|
|
response = await client.request(
|
|
method,
|
|
url,
|
|
headers=headers,
|
|
params=params,
|
|
json=json,
|
|
)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
except httpx.HTTPStatusError as exc:
|
|
detail = exc.response.text or exc.response.reason_phrase
|
|
raise WebSearchError(f"{provider_config.provider} request failed: {detail}") from exc
|
|
except httpx.HTTPError as exc:
|
|
raise WebSearchError(f"{provider_config.provider} request failed: {exc}") from exc
|
|
except ValueError as exc:
|
|
raise WebSearchError(f"{provider_config.provider} returned invalid JSON") from exc
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
async def _search_tavily(
|
|
self,
|
|
config: WebSearchProviderConfig,
|
|
query: str,
|
|
max_results: int,
|
|
domains: list[str] | None,
|
|
freshness_days: int | None,
|
|
) -> list[SearchEvidence]:
|
|
body: dict[str, Any] = {
|
|
"api_key": config.api_key,
|
|
"query": query,
|
|
"max_results": max_results,
|
|
"search_depth": config.search_depth or "basic",
|
|
"include_answer": config.include_answer,
|
|
"include_raw_content": config.include_raw_content,
|
|
}
|
|
if domains:
|
|
body["include_domains"] = domains
|
|
if freshness_days:
|
|
body["days"] = freshness_days
|
|
data = await self._request_json(
|
|
"POST",
|
|
_join_url(config.base_url, config.endpoint_path or "/search"),
|
|
provider_config=config,
|
|
json=body,
|
|
)
|
|
return [
|
|
SearchEvidence(
|
|
title=str(item.get("title") or ""),
|
|
url=str(item.get("url") or ""),
|
|
snippet=str(item.get("content") or ""),
|
|
content=str(item.get("raw_content") or ""),
|
|
score=_float_or_none(item.get("score")),
|
|
source_provider="tavily",
|
|
metadata={"query": data.get("query") or query},
|
|
)
|
|
for item in data.get("results") or []
|
|
if isinstance(item, dict) and item.get("url")
|
|
]
|
|
|
|
async def _search_brave(
|
|
self,
|
|
config: WebSearchProviderConfig,
|
|
query: str,
|
|
max_results: int,
|
|
domains: list[str] | None,
|
|
) -> list[SearchEvidence]:
|
|
search_query = query
|
|
if domains:
|
|
search_query = f"{query} " + " ".join(f"site:{domain}" for domain in domains)
|
|
data = await self._request_json(
|
|
"GET",
|
|
_join_url(config.base_url, config.endpoint_path or "/res/v1/web/search"),
|
|
provider_config=config,
|
|
headers={"X-Subscription-Token": config.api_key},
|
|
params={"q": search_query, "count": max_results},
|
|
)
|
|
results = (data.get("web") or {}).get("results") or []
|
|
return [
|
|
SearchEvidence(
|
|
title=str(item.get("title") or ""),
|
|
url=str(item.get("url") or ""),
|
|
snippet=str(item.get("description") or ""),
|
|
source_provider="brave",
|
|
metadata={"age": item.get("age")},
|
|
)
|
|
for item in results
|
|
if isinstance(item, dict) and item.get("url")
|
|
]
|
|
|
|
async def _search_serpapi(
|
|
self,
|
|
config: WebSearchProviderConfig,
|
|
query: str,
|
|
max_results: int,
|
|
) -> list[SearchEvidence]:
|
|
data = await self._request_json(
|
|
"GET",
|
|
_join_url(config.base_url, config.endpoint_path or "/search.json"),
|
|
provider_config=config,
|
|
params={
|
|
"api_key": config.api_key,
|
|
"engine": config.engine or "google",
|
|
"q": query,
|
|
"num": max_results,
|
|
},
|
|
)
|
|
return [
|
|
SearchEvidence(
|
|
title=str(item.get("title") or ""),
|
|
url=str(item.get("link") or ""),
|
|
snippet=str(item.get("snippet") or ""),
|
|
source_provider="serpapi",
|
|
metadata={"position": item.get("position")},
|
|
)
|
|
for item in data.get("organic_results") or []
|
|
if isinstance(item, dict) and item.get("link")
|
|
]
|
|
|
|
async def _search_exa(
|
|
self,
|
|
config: WebSearchProviderConfig,
|
|
query: str,
|
|
max_results: int,
|
|
domains: list[str] | None,
|
|
) -> list[SearchEvidence]:
|
|
body: dict[str, Any] = {
|
|
"query": query,
|
|
"numResults": max_results,
|
|
}
|
|
if domains:
|
|
body["includeDomains"] = domains
|
|
if config.include_text:
|
|
body["contents"] = {"text": True}
|
|
data = await self._request_json(
|
|
"POST",
|
|
_join_url(config.base_url, config.endpoint_path or "/search"),
|
|
provider_config=config,
|
|
headers={"Authorization": f"Bearer {config.api_key}"},
|
|
json=body,
|
|
)
|
|
return [
|
|
SearchEvidence(
|
|
title=str(item.get("title") or ""),
|
|
url=str(item.get("url") or ""),
|
|
snippet=str(item.get("summary") or ""),
|
|
content=str(item.get("text") or ""),
|
|
score=_float_or_none(item.get("score")),
|
|
source_provider="exa",
|
|
metadata={"id": item.get("id")},
|
|
)
|
|
for item in data.get("results") or []
|
|
if isinstance(item, dict) and item.get("url")
|
|
]
|
|
|
|
async def _search_firecrawl(
|
|
self,
|
|
config: WebSearchProviderConfig,
|
|
query: str,
|
|
max_results: int,
|
|
) -> list[SearchEvidence]:
|
|
data = await self._request_json(
|
|
"POST",
|
|
_join_url(config.base_url, config.search_path or "/v2/search"),
|
|
provider_config=config,
|
|
headers={"Authorization": f"Bearer {config.api_key}"},
|
|
json={"query": query, "limit": max_results},
|
|
)
|
|
raw_results = data.get("data") or data.get("results") or []
|
|
return [
|
|
SearchEvidence(
|
|
title=str(item.get("title") or ""),
|
|
url=str(item.get("url") or item.get("sourceURL") or ""),
|
|
snippet=str(item.get("description") or item.get("markdown") or ""),
|
|
source_provider="firecrawl",
|
|
metadata={"status": item.get("status")},
|
|
)
|
|
for item in raw_results
|
|
if isinstance(item, dict) and (item.get("url") or item.get("sourceURL"))
|
|
]
|
|
|
|
async def _search_searxng(
|
|
self,
|
|
config: WebSearchProviderConfig,
|
|
query: str,
|
|
max_results: int,
|
|
domains: list[str] | None,
|
|
) -> list[SearchEvidence]:
|
|
search_query = query
|
|
if domains:
|
|
search_query = f"{query} " + " ".join(f"site:{domain}" for domain in domains)
|
|
params: dict[str, Any] = {
|
|
"q": search_query,
|
|
"format": "json",
|
|
"categories": config.categories or "general",
|
|
}
|
|
if config.engines:
|
|
params["engines"] = ",".join(config.engines)
|
|
headers = {"Authorization": f"Bearer {config.api_key}"} if config.api_key else None
|
|
data = await self._request_json(
|
|
"GET",
|
|
_join_url(config.base_url, config.endpoint_path or "/"),
|
|
provider_config=config,
|
|
headers=headers,
|
|
params=params,
|
|
)
|
|
results = data.get("results") or []
|
|
evidence = [
|
|
SearchEvidence(
|
|
title=str(item.get("title") or ""),
|
|
url=str(item.get("url") or ""),
|
|
snippet=str(item.get("content") or ""),
|
|
score=_float_or_none(item.get("score")),
|
|
source_provider="searxng",
|
|
metadata={"engine": item.get("engine")},
|
|
)
|
|
for item in results
|
|
if isinstance(item, dict) and item.get("url")
|
|
]
|
|
return evidence[:max_results]
|
|
|
|
|
|
def _join_url(base_url: str, path: str) -> str:
|
|
return f"{(base_url or '').rstrip('/')}/{(path or '').lstrip('/')}"
|
|
|
|
|
|
def _float_or_none(value: Any) -> float | None:
|
|
try:
|
|
return float(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|