Source code for aidputils.agents.tools.experimental.web_search
from __future__ import annotations
from dataclasses import asdict
from typing import Any
from aidputils.agents.auth.client.generic_rest_client import GenericRestClient
from aidputils.agents.auth.strategy.bearer_strategy import BearerTokenAuthStrategy
from aidputils.agents.toolkit import tool_common
from aidputils.agents.toolkit.metrics_util import MetricUtility
from aidputils.agents.tools.base_tool import BaseTool
from aidputils.agents.tools.experimental.websearch.base import BaseWebSearchHandler, normalize_generic_results
from aidputils.agents.tools.experimental.websearch.brave import BraveWebSearchHandler
from aidputils.agents.tools.experimental.websearch.tavily import TavilyWebSearchHandler
_as_str = tool_common.as_str
_error = tool_common.error
_sanitize_params = tool_common.sanitize_params
[docs]
@BaseTool.register
class WebSearchTool(BaseTool):
"""
WEB_SEARCH tool (runtime in aidp-utils).
Dynamic tool that delegates to a configured HTTP endpoint (proxy) which performs the actual
web search using the specified provider.
This implementation uses `GenericRestClient` (requests-based) to ensure consistent signing/auth,
retries, and response handling across tools.
Required conf:
- search_engine: provider identifier (e.g., BRAVE, TAVILY, AZURE, CUSTOM_HTTP)
- endpoint: full URL for the proxy/provider endpoint to call
Optional conf:
- auth: {"authType": "BEARER_TOKEN", "token": "<token or resolver expression>"}
- customHeaders: { ... }
- defaultTopK, defaultLanguage
Runtime params:
- query (required) OR runtime_params.values.query OR runtime_params.paramValues.values.query
- topK/top_k optional
Request payload sent to endpoint:
{
"search_engine": "<SEARCH_ENGINE>",
"query": "<query>",
"top_k": <int>,
"params": <runtime_params>
}
Expected response:
- JSON object with "results": [ {title,url,snippet,source,publishedTime|...}, ... ]
(other shapes are tolerated and normalized)
"""
metricsUtil = MetricUtility("WEB_SEARCH_TOOL")
success_counter = metricsUtil.create_counter("WEB_SEARCH_TOOL_SUCCESS_COUNTER")
failure_counter = metricsUtil.create_counter("WEB_SEARCH_TOOL_FAILURE_COUNTER")
@classmethod
def _invoke_tool(cls, conf, runtime_params, **context_vars) -> dict:
if not isinstance(conf, dict):
return _error("Invalid conf: must be a dict", 400)
if not isinstance(runtime_params, dict):
return _error("Invalid runtime_params: must be a dict", 400)
# Optional injection for tests / advanced callers: allow providing a pre-built client.
# If present, it must provide `.post(url, json_body=..., timeout=...)` matching GenericRestClient.
injected_client = runtime_params.get("_generic_rest_client")
search_engine = _as_str(conf.get("search_engine")).strip().upper()
endpoint = _as_str(conf.get("endpoint")).strip()
auth = conf.get("auth") or {}
custom_headers = conf.get("customHeaders") or {}
# Resolve handler for engine (used for endpoint resolution and response parsing)
handler: BaseWebSearchHandler | None = None
if search_engine == "BRAVE":
handler = BraveWebSearchHandler()
elif search_engine == "TAVILY":
handler = TavilyWebSearchHandler()
# Derive endpoint based on handler (or fall back to endpoint in conf)
if not endpoint and handler is not None:
endpoint = handler.get_endpoint(conf)
# Query comes from runtime params "query" (preferred) or paramValues.values.query patterns
query = runtime_params.get("query")
if not query and isinstance(runtime_params.get("values"), dict):
query = runtime_params["values"].get("query")
if not query and isinstance(runtime_params.get("paramValues"), dict):
values = runtime_params["paramValues"].get("values")
if isinstance(values, dict):
query = values.get("query")
query = _as_str(query).strip()
if not query:
return _error("Missing required query (runtime_params.query)", 400)
top_k = (
runtime_params.get("topK")
or runtime_params.get("top_k")
or 5
)
try:
top_k = int(top_k)
except Exception:
top_k = 5
# Endpoint is required (either explicitly provided, or derived for supported engines).
if not endpoint:
return _error(
"WEB_SEARCH tool requires toolConfig.endpoint (or a supported search_engine with a derived endpoint).",
400,
{"search_engine": search_engine or None},
)
# Build headers (bearer auth + custom headers)
headers: dict[str, str] = {}
if isinstance(custom_headers, dict):
headers.update({str(k): str(v) for k, v in custom_headers.items()})
signer = None
if isinstance(auth, dict) and auth.get("auth_type") == "BEARER_TOKEN":
try:
signer = BearerTokenAuthStrategy(auth).get_signer()
except Exception as ex:
return _error(f"Invalid bearer token auth: {ex}", 400)
elif isinstance(auth, dict) and auth:
# Not supported in v1; keep explicit
return _error(
"Unsupported authType for WEB_SEARCH (only BEARER_TOKEN supported)",
400,
{"authType": auth.get("authType")},
)
# GenericRestClient requires a signer callable; if no auth, use a no-op signer
if signer is None:
signer = lambda req: req # noqa: E731
client = (
injected_client
if injected_client is not None
else GenericRestClient(endpoint="", signer=signer, default_headers=headers)
)
# Build payload shape based on engine expectations.
if search_engine == "TAVILY":
payload = {
"api_key": conf.get("api_key") or conf.get("tavily_api_key"),
"query": query,
"max_results": top_k,
"include_answer": bool(conf.get("include_answer", True)),
"include_raw_content": bool(conf.get("include_raw_content", False)),
}
else:
# Proxy mode payload (dynamic provider routing).
payload = {
"search_engine": search_engine,
"query": query,
"top_k": top_k,
"params": _sanitize_params(runtime_params),
}
timeout_secs = runtime_params.get("timeout_secs") or runtime_params.get("timeoutSecs") or 30
try:
timeout_secs = float(timeout_secs)
except Exception:
timeout_secs = 30.0
try:
resp = client.post(endpoint, json_body=payload, timeout=timeout_secs)
except Exception as ex:
# generic_rest_client already logs response content on failures
msg = str(ex)
# best-effort mapping
if "401" in msg:
return _error("Unauthorized", 401, {"endpoint": endpoint})
if "403" in msg:
return _error("Forbidden", 403, {"endpoint": endpoint})
return _error(f"WEB_SEARCH request failed: {ex}", 500, {"endpoint": endpoint})
# generic_rest_client returns SnakeAttrDict({'data': ...}) for JSON
data = getattr(resp, "data", None) if hasattr(resp, "data") else None
if data is None and isinstance(resp, dict):
data = resp.get("data")
if data is None:
return _error("WEB_SEARCH provider returned empty response", 502)
normalized_results = handler.parse(data) if handler is not None else normalize_generic_results(
data, source=search_engine or "UNKNOWN"
)
out: dict[str, Any] = {"query": query, "results": [asdict(r) for r in normalized_results]}
# Preserve Tavily synthesized answer if present
if search_engine == "TAVILY" and isinstance(data, dict) and data.get("answer"):
out["answer"] = data.get("answer")
return out
@classmethod
def _format_output_as_mcp(cls, response):
return response