plugins/web_search.py
Rama master · réplica de lectura del repositorio original.
"""Web search plugin with multiple API backends and graceful fallbacks."""
from __future__ import annotations
import inspect
import logging
import os
import re
from collections.abc import Mapping
from typing import TYPE_CHECKING
from urllib.parse import quote_plus
from rex.config import settings
from rex.plugins import Plugin
if TYPE_CHECKING:
import requests
from bs4 import BeautifulSoup as BeautifulSoupType
else:
try:
import requests
except ImportError as exc:
requests = None
_REQUESTS_IMPORT_ERROR: Exception | None = exc
else:
_REQUESTS_IMPORT_ERROR = None
try:
from bs4 import BeautifulSoup as BeautifulSoupType
except ImportError:
BeautifulSoupType = None
logger = logging.getLogger(__name__)
_MAX_QUERY_LENGTH = 500
_EXPLICIT_SEARCH_RE = re.compile(
r"^\s*(?:"
r"search(?:\s+(?:the\s+web|online))?(?:\s+for)?|"
r"web\s+search(?:\s+for)?|"
r"google|"
r"look\s+(?:up|online\s+for|on\s+the\s+web\s+for)|"
r"lookup|"
r"find(?:\s+me)?"
r")\s+(.+?)\s*$",
re.IGNORECASE,
)
def _sanitize_query(query: str) -> str:
"""Strip control characters and truncate the query to a safe length."""
sanitized = "".join(ch if ch >= " " or ch in "\t\n" else " " for ch in query)
sanitized = " ".join(sanitized.split())
return sanitized[:_MAX_QUERY_LENGTH]
def _extract_explicit_search_query(text: str) -> str | None:
match = _EXPLICIT_SEARCH_RE.match(text)
if match is None:
return None
query = _sanitize_query(match.group(1))
return query or None
def _credential_value(
names: tuple[str, ...],
*,
config: object | None = None,
config_attr: str | None = None,
environ: Mapping[str, str] | None = None,
) -> str | None:
"""Resolve a search credential from runtime config, env, or the household vault."""
if config_attr and config is not None:
value = getattr(config, config_attr, None)
if isinstance(value, str) and value.strip():
return value.strip()
env = os.environ if environ is None else environ
for name in names:
value = env.get(name)
if value:
return value
try:
from rex.credentials import get_persisted_credential
for name in names:
value = get_persisted_credential(name)
if value:
return value
except Exception as exc:
logger.debug("Search credential lookup failed: %s", exc)
return None
def configured_search_providers(
config: object | None = None,
*,
environ: Mapping[str, str] | None = None,
) -> list[str]:
"""Return providers that are both enabled and actually usable right now."""
env = os.environ if environ is None else environ
source = config if config is not None else settings
configured = str(getattr(source, "search_providers", "") or "")
if config is None:
# Legacy direct-plugin compatibility only. Canonical assistant execution
# passes AppConfig explicitly, where non-secret provider authority lives.
configured = env.get("REX_SEARCH_PROVIDERS") or configured
providers = [item.strip().lower() for item in configured.split(",") if item.strip()]
usable: list[str] = []
for provider in providers:
if provider == "duckduckgo":
if requests is not None and BeautifulSoupType is not None:
usable.append(provider)
elif provider == "brave":
if _credential_value(
("BRAVE_API_KEY",), config=source, config_attr="brave_api_key", environ=env
):
usable.append(provider)
elif provider == "serpapi":
if _credential_value(("SERPAPI_KEY", "SERPAPI_API_KEY"), config=source, environ=env):
usable.append(provider)
elif provider == "google":
if _credential_value(("GOOGLE_API_KEY",), config=source, environ=env) and env.get(
"GOOGLE_CSE_ID"
):
usable.append(provider)
elif provider == "browserless":
if env.get("BROWSERLESS_API_KEY") and BeautifulSoupType is not None:
usable.append(provider)
return usable
# API Endpoints
SERPAPI_URL = "https://serpapi.com/search"
BRAVE_URL = "https://api.search.brave.com/res/v1/web/search"
GOOGLE_URL = "https://www.googleapis.com/customsearch/v1"
BROWSERLESS_URL = "https://chrome.browserless.io/content"
def _require_requests() -> None:
if requests is None:
raise RuntimeError(
"The web_search plugin requires the 'requests' package. "
"Install with: pip install requests"
) from globals().get("_REQUESTS_IMPORT_ERROR")
def _create_session() -> requests.Session:
_require_requests()
from requests.adapters import HTTPAdapter, Retry
session = requests.Session()
retries = Retry(total=3, backoff_factor=0.5, status_forcelist=[429, 502, 503, 504])
adapter = HTTPAdapter(max_retries=retries)
session.mount("http://", adapter)
session.mount("https://", adapter)
return session
class WebSearchPlugin:
name = "web_search"
def __init__(self) -> None:
self._session: requests.Session | None = None
def initialize(self) -> None:
pass
def shutdown(self) -> None:
if self._session is not None:
self._session.close()
def process(self, query: str) -> str | None:
explicit_query = _extract_explicit_search_query(query)
if explicit_query is None:
logger.debug("Skipping web search plugin; no explicit search intent")
return None
return self.search(explicit_query)
def search(self, query: str, *, config: object | None = None) -> str | None:
query = _sanitize_query(query)
if not query:
return None
source = config if config is not None else settings
for provider in configured_search_providers(source):
method = getattr(self, f"_search_{provider}", None)
if callable(method):
try:
parameters = inspect.signature(method).parameters
accepts_config = "config" in parameters or any(
parameter.kind == inspect.Parameter.VAR_KEYWORD
for parameter in parameters.values()
)
result = method(query, config=source) if accepts_config else method(query)
if isinstance(result, str):
return result
if result is not None:
return str(result)
except Exception as e:
logger.warning("Provider '%s' failed: %s", provider, e)
logger.warning("All search providers failed")
return None
def _format_result(self, title: str, url: str, snippet: str) -> str:
return f"{title} - {url}\n{snippet}"
def _search_serpapi(self, query: str, *, config: object | None = None) -> str | None:
api_key = _credential_value(("SERPAPI_KEY", "SERPAPI_API_KEY"), config=config)
if not api_key:
return None
params: dict[str, str] = {
"q": query,
"api_key": api_key,
"num": "3",
"engine": os.getenv("SERPAPI_ENGINE", "google"),
}
headers = {"X-Serpapi-Privacy": "true"}
try:
session = self._get_session()
resp = session.get(SERPAPI_URL, params=params, headers=headers, timeout=10)
resp.raise_for_status()
results = resp.json().get("organic_results", [])
if not results:
return None
top = results[0]
return self._format_result(
str(top["title"]), str(top["link"]), str(top.get("snippet", ""))
)
except Exception as e:
if isinstance(e, RuntimeError):
raise
logger.warning("SerpAPI search failed: %s", e)
return None
def _search_brave(self, query: str, *, config: object | None = None) -> str | None:
api_key = _credential_value(("BRAVE_API_KEY",), config=config, config_attr="brave_api_key")
if not api_key:
return None
headers = {"X-Subscription-Token": api_key}
params: dict[str, str | int] = {"q": query, "count": 3}
try:
session = self._get_session()
resp = session.get(BRAVE_URL, headers=headers, params=params, timeout=10)
resp.raise_for_status()
results = resp.json().get("web", {}).get("results", [])
if not results:
return None
top = results[0]
return self._format_result(
str(top["title"]), str(top["url"]), str(top.get("description", ""))
)
except Exception as e:
if isinstance(e, RuntimeError):
raise
logger.warning("Brave search failed: %s", e)
return None
def _search_duckduckgo(self, query: str, *, config: object | None = None) -> str | None:
if BeautifulSoupType is None:
logger.warning("BeautifulSoup is required for DuckDuckGo scraping")
return None
try:
url = f"https://html.duckduckgo.com/html/?q={quote_plus(query)}"
headers = {"User-Agent": "Mozilla/5.0"}
session = self._get_session()
resp = session.get(url, headers=headers, timeout=10)
resp.raise_for_status()
soup = BeautifulSoupType(resp.text, "html.parser")
result = soup.find("a", class_="result__a")
snippet = soup.find("a", class_="result__snippet")
if result:
href = result.get("href") if hasattr(result, "get") else None
if href is None:
return None
snippet_text = snippet.text if snippet is not None else ""
return self._format_result(str(result.text), str(href), str(snippet_text))
return None
except Exception as e:
if isinstance(e, RuntimeError):
raise
logger.warning("DuckDuckGo search failed: %s", e)
return None
def _search_google(self, query: str, *, config: object | None = None) -> str | None:
api_key = _credential_value(("GOOGLE_API_KEY",), config=config)
engine_id = os.getenv("GOOGLE_CSE_ID")
if not api_key or not engine_id:
return None
params: dict[str, str | int] = {"q": query, "key": api_key, "cx": engine_id, "num": 3}
try:
session = self._get_session()
resp = session.get(GOOGLE_URL, params=params, timeout=10)
resp.raise_for_status()
items = resp.json().get("items", [])
if not items:
return None
top = items[0]
return self._format_result(
str(top["title"]), str(top["link"]), str(top.get("snippet", ""))
)
except Exception as e:
if isinstance(e, RuntimeError):
raise
logger.warning("Google CSE search failed: %s", e)
return None
def _search_browserless(self, query: str, *, config: object | None = None) -> str | None:
token = os.getenv("BROWSERLESS_API_KEY")
if not token or BeautifulSoupType is None:
return None
payload = {
"url": f"https://duckduckgo.com/?q={quote_plus(query)}",
"gotoOptions": {"waitUntil": "networkidle2"},
}
headers = {"Cache-Control": "no-cache", "Content-Type": "application/json"}
try:
session = self._get_session()
resp = session.post(
BROWSERLESS_URL, headers=headers, params={"token": token}, json=payload, timeout=20
)
resp.raise_for_status()
soup = BeautifulSoupType(resp.text, "html.parser")
result = soup.find("a", class_="result__a")
snippet = soup.find("a", class_="result__snippet")
if result:
href = result.get("href") if hasattr(result, "get") else None
if href is None:
return None
snippet_text = snippet.text if snippet is not None else ""
return self._format_result(str(result.text), str(href), str(snippet_text))
return None
except Exception as e:
if isinstance(e, RuntimeError):
raise
logger.warning("Browserless search failed: %s", e)
return None
def _get_session(self) -> requests.Session:
if self._session is None:
self._session = _create_session()
return self._session
# Singleton for use across the app
_PLUGIN_SINGLETON: WebSearchPlugin | None = None
def _get_plugin() -> WebSearchPlugin:
global _PLUGIN_SINGLETON
if _PLUGIN_SINGLETON is None:
_PLUGIN_SINGLETON = WebSearchPlugin()
_PLUGIN_SINGLETON.initialize()
return _PLUGIN_SINGLETON
# --- Public API ---
def search_web(query: str, *, config: object | None = None) -> str | None:
"""Search the web using only providers that are configured and usable."""
plugin = _get_plugin()
if config is None:
return plugin.search(query)
return plugin.search(query, config=config)
def search_serpapi(query: str) -> str | None:
"""Force a SerpAPI search (if available)."""
return _get_plugin()._search_serpapi(query)
def search_duckduckgo(query: str) -> str | None:
"""Perform a DuckDuckGo HTML scrape."""
return _get_plugin()._search_duckduckgo(query)
def register() -> Plugin:
"""Plugin entry point for dynamic loading."""
return WebSearchPlugin()