Add Brave web search support

app/ai/service.py

@@ -14,7 +14,8 @@ from openai import AsyncOpenAI
 
 from app.ai.parser import AIInvocation
 from app.ai.prompts import PromptBook
-from app.ai.tools import FILE_TOOL_DISPATCH, FILE_TOOL_SCHEMAS, PYTHON_EXEC_SCHEMA
+from app.ai.tools import FILE_TOOL_DISPATCH, FILE_TOOL_SCHEMAS, PYTHON_EXEC_SCHEMA, WEB_SEARCH_SCHEMA
+from app.ai.web_search import BraveWebSearchClient, format_results
 from app.config import AIConfig, AIModelConfig, AIProviderConfig
 
 if TYPE_CHECKING:
@@ -62,6 +63,7 @@ class AIService:
         self.prompt_book = PromptBook(config.prompt.root)
         self._clients: dict[str, AsyncOpenAI] = {}
         self._system_prompt: str | None = None
+        self.web_search = BraveWebSearchClient(config.web_search)
 
     async def initialize(self) -> None:
         return None
@@ -210,7 +212,10 @@ class AIService:
         return messages
 
     async def _build_tool_list(self) -> list[dict]:
-        return FILE_TOOL_SCHEMAS + [PYTHON_EXEC_SCHEMA]
+        tools = FILE_TOOL_SCHEMAS + [PYTHON_EXEC_SCHEMA]
+        if self.web_search.available:
+            tools.append(WEB_SEARCH_SCHEMA)
+        return tools
 
     def _get_allowed_roots(self) -> tuple[Path, ...]:
         # [0] = base for resolution (project root); [1:] = permitted subdirs
@@ -339,6 +344,9 @@ class AIService:
         try:
             if name == "python_exec":
                 result = await self._execute_python_exec(arguments, ctx)
+            elif name == "web_search":
+                response = await self.web_search.search(arguments)
+                result = format_results(response, include_heading=True)
             elif name == "send_file":
                 result = await self._execute_send_file(arguments, ctx, allowed_roots)
             elif name in FILE_TOOL_DISPATCH:

app/ai/tools.py

@@ -219,6 +219,60 @@ PYTHON_EXEC_SCHEMA: dict[str, Any] = {
     },
 }
 
+WEB_SEARCH_SCHEMA: dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "web_search",
+        "description": (
+            "Search the web with Brave Search for current or source-backed information. "
+            "Returns result titles, URLs, snippets, dates when available, and optional extra snippets."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "Search query, max 400 characters / 50 words.",
+                },
+                "count": {
+                    "type": "integer",
+                    "description": "Number of web results to return, 1-20.",
+                    "minimum": 1,
+                    "maximum": 20,
+                },
+                "country": {
+                    "type": "string",
+                    "description": "Two-letter result country code, e.g. US, PL, DE.",
+                },
+                "search_lang": {
+                    "type": "string",
+                    "description": "Search language code, e.g. en, pl, de.",
+                },
+                "ui_lang": {
+                    "type": "string",
+                    "description": "UI language for response metadata, e.g. en-US.",
+                },
+                "safesearch": {
+                    "type": "string",
+                    "enum": ["off", "moderate", "strict"],
+                    "description": "Adult content filtering level.",
+                },
+                "freshness": {
+                    "type": "string",
+                    "description": (
+                        "Optional recency filter: pd, pw, pm, py, or YYYY-MM-DDtoYYYY-MM-DD."
+                    ),
+                },
+                "extra_snippets": {
+                    "type": "boolean",
+                    "description": "Include additional excerpts per result when available.",
+                },
+            },
+            "required": ["query"],
+        },
+    },
+}
+
 ToolExecutor = Callable[[dict, tuple[Path, ...]], Awaitable[str]]
 
 FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = {

app/ai/web_search.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import re
+from typing import Any
+
+import httpx
+
+from app.config import AIWebSearchConfig
+
+SAFESEARCH_VALUES = {"off", "moderate", "strict"}
+FRESHNESS_RE = re.compile(r"^(pd|pw|pm|py|\d{4}-\d{2}-\d{2}to\d{4}-\d{2}-\d{2})$")
+
+
+@dataclass(slots=True)
+class WebSearchQuery:
+    query: str
+    count: int
+    offset: int
+    country: str
+    search_lang: str
+    ui_lang: str
+    safesearch: str
+    freshness: str | None = None
+    extra_snippets: bool = False
+
+
+@dataclass(slots=True)
+class WebSearchResult:
+    title: str
+    url: str
+    description: str = ""
+    age: str = ""
+    page_age: str = ""
+    extra_snippets: list[str] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class WebSearchResponse:
+    query: str
+    results: list[WebSearchResult]
+    altered_query: str | None = None
+    more_results_available: bool = False
+
+
+def clamp_int(value: Any, *, default: int, minimum: int, maximum: int) -> int:
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError):
+        parsed = default
+    return min(max(parsed, minimum), maximum)
+
+
+def sanitize_query(query: str) -> str:
+    normalized = " ".join(str(query).split())
+    if not normalized:
+        raise ValueError("web_search: 'query' is required")
+    words = normalized.split()
+    if len(words) > 50:
+        normalized = " ".join(words[:50])
+    if len(normalized) > 400:
+        normalized = normalized[:400].rstrip()
+    if not normalized:
+        raise ValueError("web_search: 'query' is required")
+    return normalized
+
+
+def normalize_safesearch(value: Any, *, default: str) -> str:
+    safesearch = str(value or default).strip().lower()
+    if safesearch not in SAFESEARCH_VALUES:
+        raise ValueError("web_search: 'safesearch' must be one of: off, moderate, strict")
+    return safesearch
+
+
+def normalize_freshness(value: Any) -> str | None:
+    if value is None:
+        return None
+    freshness = str(value).strip()
+    if not freshness:
+        return None
+    if not FRESHNESS_RE.match(freshness):
+        raise ValueError(
+            "web_search: 'freshness' must be pd, pw, pm, py, or YYYY-MM-DDtoYYYY-MM-DD"
+        )
+    return freshness
+
+
+def build_query(config: AIWebSearchConfig, arguments: dict[str, Any]) -> WebSearchQuery:
+    return WebSearchQuery(
+        query=sanitize_query(str(arguments.get("query") or arguments.get("q") or "")),
+        count=clamp_int(arguments.get("count"), default=config.count, minimum=1, maximum=20),
+        offset=clamp_int(arguments.get("offset"), default=0, minimum=0, maximum=9),
+        country=str(arguments.get("country") or config.country).strip().upper() or config.country,
+        search_lang=str(arguments.get("search_lang") or config.search_lang).strip().lower()
+        or config.search_lang,
+        ui_lang=str(arguments.get("ui_lang") or config.ui_lang).strip() or config.ui_lang,
+        safesearch=normalize_safesearch(arguments.get("safesearch"), default=config.safesearch),
+        freshness=normalize_freshness(arguments.get("freshness")),
+        extra_snippets=bool(arguments.get("extra_snippets", config.extra_snippets)),
+    )
+
+
+def parse_response(query: str, payload: dict[str, Any]) -> WebSearchResponse:
+    raw_query = payload.get("query") if isinstance(payload.get("query"), dict) else {}
+    raw_web = payload.get("web") if isinstance(payload.get("web"), dict) else {}
+    raw_results = raw_web.get("results") if isinstance(raw_web.get("results"), list) else []
+
+    results: list[WebSearchResult] = []
+    for item in raw_results:
+        if not isinstance(item, dict):
+            continue
+        title = str(item.get("title") or "").strip()
+        url = str(item.get("url") or "").strip()
+        if not title and not url:
+            continue
+        snippets = item.get("extra_snippets")
+        extra_snippets = [
+            str(snippet).strip()
+            for snippet in snippets
+            if isinstance(snippet, str) and snippet.strip()
+        ] if isinstance(snippets, list) else []
+        results.append(
+            WebSearchResult(
+                title=title or url,
+                url=url,
+                description=str(item.get("description") or "").strip(),
+                age=str(item.get("age") or "").strip(),
+                page_age=str(item.get("page_age") or "").strip(),
+                extra_snippets=extra_snippets,
+            )
+        )
+
+    altered = raw_query.get("altered") if isinstance(raw_query, dict) else None
+    return WebSearchResponse(
+        query=query,
+        results=results,
+        altered_query=str(altered).strip() if altered else None,
+        more_results_available=bool(raw_query.get("more_results_available", False))
+        if isinstance(raw_query, dict)
+        else False,
+    )
+
+
+def format_results(response: WebSearchResponse, *, include_heading: bool = True) -> str:
+    lines: list[str] = []
+    if include_heading:
+        heading = f"Search results for: {response.query}"
+        if response.altered_query:
+            heading += f" (spellchecked: {response.altered_query})"
+        lines.append(heading)
+
+    if not response.results:
+        lines.append("No results found.")
+        return "\n".join(lines)
+
+    for index, result in enumerate(response.results, start=1):
+        lines.append(f"{index}. {result.title}")
+        if result.url:
+            lines.append(result.url)
+        meta = result.age or result.page_age
+        if meta:
+            lines.append(f"Date: {meta}")
+        if result.description:
+            lines.append(result.description)
+        for snippet in result.extra_snippets[:5]:
+            lines.append(f"- {snippet}")
+        lines.append("")
+
+    if not lines[-1]:
+        lines.pop()
+    return "\n".join(lines)
+
+
+class BraveWebSearchClient:
+    def __init__(self, config: AIWebSearchConfig) -> None:
+        self.config = config
+
+    @property
+    def available(self) -> bool:
+        return self.config.available
+
+    async def search(self, arguments: dict[str, Any]) -> WebSearchResponse:
+        if not self.config.enabled:
+            raise ValueError("Brave web search is disabled")
+        if not self.config.api_key:
+            raise ValueError("Brave web search is not configured. Set BRAVE_API_KEY.")
+
+        query = build_query(self.config, arguments)
+        params: dict[str, Any] = {
+            "q": query.query,
+            "count": query.count,
+            "offset": query.offset,
+            "country": query.country,
+            "search_lang": query.search_lang,
+            "ui_lang": query.ui_lang,
+            "safesearch": query.safesearch,
+            "text_decorations": False,
+            "spellcheck": True,
+        }
+        if query.freshness:
+            params["freshness"] = query.freshness
+        if query.extra_snippets:
+            params["extra_snippets"] = True
+
+        headers = {
+            "Accept": "application/json",
+            "Accept-Encoding": "gzip",
+            "X-Subscription-Token": self.config.api_key,
+        }
+
+        try:
+            async with httpx.AsyncClient(timeout=self.config.timeout) as client:
+                response = await client.get(self.config.endpoint, params=params, headers=headers)
+        except httpx.HTTPError as exc:
+            raise ValueError(f"Brave web search request failed: {exc}") from exc
+
+        if response.status_code >= 400:
+            raise ValueError(_format_error_response(response))
+
+        try:
+            payload = response.json()
+        except ValueError as exc:
+            raise ValueError("Brave web search returned invalid JSON") from exc
+        if not isinstance(payload, dict):
+            raise ValueError("Brave web search returned an unexpected response")
+
+        return parse_response(query.query, payload)
+
+
+def _format_error_response(response: httpx.Response) -> str:
+    try:
+        payload = response.json()
+    except ValueError:
+        return f"Brave web search failed with HTTP {response.status_code}: {response.text[:200]}"
+    if isinstance(payload, dict):
+        error = payload.get("error")
+        if isinstance(error, dict):
+            detail = error.get("detail") or error.get("message") or error.get("code")
+            if detail:
+                return f"Brave web search failed with HTTP {response.status_code}: {detail}"
+        message = payload.get("message")
+        if message:
+            return f"Brave web search failed with HTTP {response.status_code}: {message}"
+    return f"Brave web search failed with HTTP {response.status_code}"

app/commands/builtins.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 import html
 import re
+import shlex
 from textwrap import wrap
 from typing import TYPE_CHECKING, Any
 
 from pyrogram import enums
 
 from app.ai import AIInvocation, parse_ai_invocation
+from app.ai.web_search import format_results
 from app.commands.core import CommandRegistry, normalize_command_name
 
 if TYPE_CHECKING:
@@ -322,6 +324,67 @@ async def echo(ctx: "CommandContext") -> None:
     await ctx.reply(text)
 
 
+def _parse_search_args(text: str) -> dict[str, Any]:
+    try:
+        tokens = shlex.split(text)
+    except ValueError:
+        tokens = text.split()
+
+    arguments: dict[str, Any] = {}
+    query_parts: list[str] = []
+    index = 0
+    while index < len(tokens):
+        token = tokens[index]
+        if token == "--count":
+            index += 1
+            if index >= len(tokens):
+                raise ValueError("Usage: .search / ..search [--count N] [--freshness pd|pw|pm|py|YYYY-MM-DDtoYYYY-MM-DD] <query>")
+            arguments["count"] = tokens[index]
+        elif token.startswith("--count="):
+            arguments["count"] = token.split("=", 1)[1]
+        elif token == "--freshness":
+            index += 1
+            if index >= len(tokens):
+                raise ValueError("Usage: .search / ..search [--count N] [--freshness pd|pw|pm|py|YYYY-MM-DDtoYYYY-MM-DD] <query>")
+            arguments["freshness"] = tokens[index]
+        elif token.startswith("--freshness="):
+            arguments["freshness"] = token.split("=", 1)[1]
+        else:
+            query_parts.append(token)
+        index += 1
+
+    query = " ".join(query_parts).strip()
+    if not query:
+        raise ValueError("Usage: .search / ..search [--count N] [--freshness pd|pw|pm|py|YYYY-MM-DDtoYYYY-MM-DD] <query>")
+    arguments["query"] = query
+    return arguments
+
+
+@command_registry.command(
+    "search",
+    "Search the web with Brave Search",
+    usage=".search / ..search [--count N] [--freshness pd|pw|pm|py|YYYY-MM-DDtoYYYY-MM-DD] <query>",
+    category="ai",
+).alias("web").shared()
+async def search(ctx: "CommandContext") -> None:
+    ai_service = ctx.manager.ai_service
+    if ai_service is None:
+        await ctx.reply("Brave web search is not configured. Configure ai.web_search and BRAVE_API_KEY.")
+        return
+    if not ai_service.web_search.available:
+        await ctx.reply("Brave web search is not configured. Set BRAVE_API_KEY.")
+        return
+
+    try:
+        arguments = _parse_search_args(ctx.args)
+        response = await ai_service.web_search.search(arguments)
+    except ValueError as exc:
+        await ctx.reply(str(exc))
+        return
+
+    await _reply_long(ctx, format_results(response, include_heading=True))
+
+
 @command_registry.command(
     "ai",
     "Ask a configurable AI model",

app/config.py

@@ -60,6 +60,24 @@ class AIPromptConfig:
     user_template: str | None = None
 
 
+@dataclass(slots=True)
+class AIWebSearchConfig:
+    enabled: bool = True
+    api_key: str | None = None
+    endpoint: str = "https://api.search.brave.com/res/v1/web/search"
+    count: int = 5
+    country: str = "US"
+    search_lang: str = "en"
+    ui_lang: str = "en-US"
+    safesearch: str = "moderate"
+    timeout: float = 10.0
+    extra_snippets: bool = False
+
+    @property
+    def available(self) -> bool:
+        return self.enabled and bool(self.api_key)
+
+
 @dataclass(slots=True)
 class AIConfig:
     default_model: str
@@ -69,6 +87,7 @@ class AIConfig:
     models: dict[str, AIModelConfig] = field(default_factory=dict)
     workspace_dir: Path = field(default_factory=lambda: Path("workspace"))
     tool_max_rounds: int = 10
+    web_search: AIWebSearchConfig = field(default_factory=AIWebSearchConfig)
 
 
 @dataclass(slots=True)
@@ -127,6 +146,20 @@ def _parse_float(value: Any, *, field_name: str) -> float:
         raise ValueError(f"{field_name} must be a number") from exc
 
 
+def _parse_bool(value: Any, *, field_name: str) -> bool:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        normalized = value.strip().lower()
+        if normalized in {"1", "true", "yes", "on"}:
+            return True
+        if normalized in {"0", "false", "no", "off"}:
+            return False
+    if isinstance(value, int):
+        return bool(value)
+    raise ValueError(f"{field_name} must be a boolean")
+
+
 def _parse_mapping(value: Any, *, field_name: str) -> dict[str, Any]:
     if value is None:
         return {}
@@ -135,6 +168,69 @@ def _parse_mapping(value: Any, *, field_name: str) -> dict[str, Any]:
     return value
 
 
+def _parse_web_search_config(raw_ai: dict[str, Any]) -> AIWebSearchConfig:
+    raw = _parse_mapping(raw_ai.get("web_search"), field_name="ai.web_search")
+
+    enabled = _parse_bool(raw.get("enabled", True), field_name="ai.web_search.enabled")
+    api_key = str(raw.get("api_key") or "").strip() or None
+    endpoint = str(
+        raw.get("endpoint") or "https://api.search.brave.com/res/v1/web/search"
+    ).strip()
+    if not endpoint:
+        raise ValueError("ai.web_search.endpoint is required")
+
+    count = _parse_int(raw.get("count", 5), field_name="ai.web_search.count")
+    count = min(max(count, 1), 20)
+
+    country = str(raw.get("country", "US") or "US").strip().upper()
+    search_lang = str(raw.get("search_lang", "en") or "en").strip().lower()
+    ui_lang = str(raw.get("ui_lang", "en-US") or "en-US").strip()
+    safesearch = str(raw.get("safesearch", "moderate") or "moderate").strip().lower()
+    if safesearch not in {"off", "moderate", "strict"}:
+        raise ValueError("ai.web_search.safesearch must be one of: off, moderate, strict")
+
+    timeout = _parse_float(raw.get("timeout", 10.0), field_name="ai.web_search.timeout")
+    if timeout <= 0:
+        raise ValueError("ai.web_search.timeout must be > 0")
+
+    extra_snippets = _parse_bool(
+        raw.get("extra_snippets", False),
+        field_name="ai.web_search.extra_snippets",
+    )
+
+    env_api_key = os.environ.get("BRAVE_API_KEY")
+    if env_api_key is not None:
+        api_key = env_api_key.strip() or None
+    if os.environ.get("BRAVE_SEARCH_COUNT") is not None:
+        count = min(
+            max(_parse_int(os.environ["BRAVE_SEARCH_COUNT"], field_name="BRAVE_SEARCH_COUNT"), 1),
+            20,
+        )
+    if os.environ.get("BRAVE_SEARCH_COUNTRY") is not None:
+        country = os.environ["BRAVE_SEARCH_COUNTRY"].strip().upper() or country
+    if os.environ.get("BRAVE_SEARCH_LANG") is not None:
+        search_lang = os.environ["BRAVE_SEARCH_LANG"].strip().lower() or search_lang
+    if os.environ.get("BRAVE_SEARCH_UI_LANG") is not None:
+        ui_lang = os.environ["BRAVE_SEARCH_UI_LANG"].strip() or ui_lang
+    if os.environ.get("BRAVE_SEARCH_SAFESEARCH") is not None:
+        safesearch = os.environ["BRAVE_SEARCH_SAFESEARCH"].strip().lower()
+        if safesearch not in {"off", "moderate", "strict"}:
+            raise ValueError("BRAVE_SEARCH_SAFESEARCH must be one of: off, moderate, strict")
+
+    return AIWebSearchConfig(
+        enabled=enabled,
+        api_key=api_key,
+        endpoint=endpoint,
+        count=count,
+        country=country,
+        search_lang=search_lang,
+        ui_lang=ui_lang,
+        safesearch=safesearch,
+        timeout=timeout,
+        extra_snippets=extra_snippets,
+    )
+
+
 def _parse_accounts(raw_accounts: Any, *, default_api_id: int | None, default_api_hash: str | None, default_session_dir: Path) -> list[AccountConfig]:
     if raw_accounts is None:
         return []
@@ -287,6 +383,8 @@ def _parse_ai_config(raw_ai: Any, *, root: Path) -> AIConfig | None:
     if tool_max_rounds < 1:
         raise ValueError("ai.tool_max_rounds must be >= 1")
 
+    web_search = _parse_web_search_config(raw_ai)
+
     return AIConfig(
         default_model=default_model,
         context_messages=context_messages,
@@ -295,6 +393,7 @@ def _parse_ai_config(raw_ai: Any, *, root: Path) -> AIConfig | None:
         models=models,
         workspace_dir=workspace_dir,
         tool_max_rounds=tool_max_rounds,
+        web_search=web_search,
     )