Oleksandr Kozachuk
94 lines
3.9 KiB
Python
94 lines
3.9 KiB
Python
"""Web search tool via Exa (SPEC-015, FDB-022).
|
|
|
|
A `web_search` function tool: the model looks things up on the open web
|
|
when a general "look it up" question is not covered by IGDB, the codex,
|
|
the news store, or a URL the user pasted. Results are external text, so
|
|
titles and snippets are sanitized (SAF-03) before they reach the prompt.
|
|
The Exa API key lives in host config (or the `EXA_API_KEY` env), never in
|
|
the repo.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
from typing import Any, Callable, Dict, List
|
|
|
|
import aiohttp
|
|
|
|
from .ai_responder import sanitize_external_text
|
|
|
|
EXA_SEARCH_URL = "https://api.exa.ai/search"
|
|
DEFAULT_RESULTS = 5
|
|
MAX_RESULTS = 10
|
|
DEFAULT_SNIPPET_CHARS = 400
|
|
FETCH_TIMEOUT_S = 15
|
|
|
|
WEB_SEARCH_TOOL = {
|
|
"name": "web_search",
|
|
"description": "Search the open web for current information when the user asks you to look something up and it is not "
|
|
"covered by game info (IGDB), the Codex, the news store, or a URL they pasted. Returns result titles, URLs, and a short "
|
|
"snippet; follow up with fetch_url on a result link for the full article.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "What to search the web for."},
|
|
"num_results": {"type": "integer", "description": "How many results to return (default 5, max 10)."},
|
|
},
|
|
"required": ["query"],
|
|
},
|
|
}
|
|
|
|
|
|
def _format_results(data: Any, snippet_chars: int) -> List[Dict[str, str]]:
|
|
"""Reduce an Exa response to sanitized {title, url, snippet, published} rows (WEB-02)."""
|
|
results = data.get("results", []) if isinstance(data, dict) else []
|
|
out: List[Dict[str, str]] = []
|
|
for item in results:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
out.append(
|
|
{
|
|
"title": sanitize_external_text(str(item.get("title") or ""), 200),
|
|
"url": str(item.get("url") or ""),
|
|
"snippet": sanitize_external_text(str(item.get("text") or item.get("snippet") or ""), snippet_chars),
|
|
"published": str(item.get("publishedDate") or ""),
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
class WebSearch:
|
|
def __init__(self, config_getter: Callable[[], Dict[str, Any]]) -> None:
|
|
self._config = config_getter
|
|
|
|
def _api_key(self) -> str:
|
|
return str(self._config().get("exa-api-key") or os.environ.get("EXA_API_KEY", ""))
|
|
|
|
def enabled(self) -> bool:
|
|
return bool(self._config().get("enable-web-search", False)) and bool(self._api_key())
|
|
|
|
async def _post(self, payload: Dict[str, Any], headers: Dict[str, str]) -> Any:
|
|
timeout = aiohttp.ClientTimeout(total=FETCH_TIMEOUT_S)
|
|
async with aiohttp.ClientSession(timeout=timeout) as session:
|
|
async with session.post(EXA_SEARCH_URL, json=payload, headers=headers) as response:
|
|
response.raise_for_status()
|
|
return await response.json()
|
|
|
|
async def search(self, query: str, num_results: int = DEFAULT_RESULTS) -> Dict[str, Any]:
|
|
"""Return sanitized web results, or an error dict — never raise (WEB-04)."""
|
|
key = self._api_key()
|
|
if not key:
|
|
return {"error": "web search unavailable: no api key"}
|
|
query = (query or "").strip()
|
|
if not query:
|
|
return {"query": "", "results": []}
|
|
num = max(1, min(int(num_results or DEFAULT_RESULTS), MAX_RESULTS)) # WEB-03
|
|
snippet_chars = int(self._config().get("web-snippet-chars", DEFAULT_SNIPPET_CHARS))
|
|
payload = {"query": query, "numResults": num, "type": "auto", "contents": {"text": {"maxCharacters": max(snippet_chars, 200)}}}
|
|
headers = {"x-api-key": key, "Content-Type": "application/json"}
|
|
try:
|
|
data = await self._post(payload, headers)
|
|
except Exception as err:
|
|
logging.warning(f"web search failed: {err!r}")
|
|
return {"error": "web search failed"}
|
|
return {"query": query, "results": _format_results(data, snippet_chars)}
|