Add Brave web search support

This commit is contained in:
2026-05-18 22:38:54 +02:00
parent 1ef2893345
commit 7147c657ab
10 changed files with 640 additions and 2 deletions

View File

@@ -2,4 +2,10 @@ TG_API_ID=123456
TG_API_HASH=your_api_hash_here
SESSION_DIR=sessions
ACCOUNTS_JSON=[{"name":"personal","session":"personal"},{"name":"work","session":"work"}]
BRAVE_API_KEY=your_brave_search_api_key_here
BRAVE_SEARCH_COUNT=5
BRAVE_SEARCH_COUNTRY=US
BRAVE_SEARCH_LANG=en
BRAVE_SEARCH_UI_LANG=en-US
BRAVE_SEARCH_SAFESEARCH=moderate
# AI_CONFIG={"default_model":"mm","context_messages":12,"prompt":{"root":"prompts/ai","system_file":"system.md","parts_dir":"parts"},"providers":{"minimax":{"api_key":"your_minimax_api_key_here","base_url":"https://api.minimax.io/v1"},"openrouter":{"api_key":"your_openrouter_api_key_here","base_url":"https://openrouter.ai/api/v1"}},"models":{"mm":{"provider":"minimax","model":"MiniMax-Text-01"},"or":{"provider":"openrouter","model":"openai/gpt-4o-mini"}}}

View File

@@ -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:

View File

@@ -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] = {

244
app/ai/web_search.py Normal file
View File

@@ -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}"

View File

@@ -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",

View File

@@ -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,
)

View File

@@ -17,6 +17,18 @@
"context_messages": 12,
"workspace_dir": "workspace",
"tool_max_rounds": 10,
"web_search": {
"enabled": true,
"api_key": "",
"endpoint": "https://api.search.brave.com/res/v1/web/search",
"count": 5,
"country": "US",
"search_lang": "en",
"ui_lang": "en-US",
"safesearch": "moderate",
"timeout": 10.0,
"extra_snippets": false
},
"prompt": {
"root": "prompts/ai",
"system_file": "system.md",

View File

@@ -2,4 +2,5 @@ kurigram
tgcrypto
python-dotenv
openai
httpx
elevenlabs

View File

@@ -0,0 +1,67 @@
from __future__ import annotations
from pathlib import Path
import unittest
from unittest.mock import patch
from app.config import _parse_ai_config
def _raw_ai_config() -> dict:
return {
"default_model": "test",
"context_messages": 1,
"prompt": {"root": "prompts/ai"},
"providers": {
"test": {
"api_key": "model-key",
"base_url": "https://example.test/v1",
}
},
"models": {
"test": {
"provider": "test",
"model": "test-model",
}
},
"web_search": {
"api_key": "config-key",
"count": 99,
"safesearch": "strict",
},
}
class ConfigWebSearchTests(unittest.TestCase):
def test_web_search_config_clamps_count(self) -> None:
with patch.dict("os.environ", {}, clear=True):
config = _parse_ai_config(_raw_ai_config(), root=Path.cwd())
assert config is not None
self.assertEqual(config.web_search.api_key, "config-key")
self.assertEqual(config.web_search.count, 20)
self.assertEqual(config.web_search.safesearch, "strict")
def test_web_search_env_overrides(self) -> None:
env = {
"BRAVE_API_KEY": "env-key",
"BRAVE_SEARCH_COUNT": "3",
"BRAVE_SEARCH_COUNTRY": "pl",
"BRAVE_SEARCH_LANG": "pl",
"BRAVE_SEARCH_UI_LANG": "pl-PL",
"BRAVE_SEARCH_SAFESEARCH": "off",
}
with patch.dict("os.environ", env, clear=True):
config = _parse_ai_config(_raw_ai_config(), root=Path.cwd())
assert config is not None
self.assertEqual(config.web_search.api_key, "env-key")
self.assertEqual(config.web_search.count, 3)
self.assertEqual(config.web_search.country, "PL")
self.assertEqual(config.web_search.search_lang, "pl")
self.assertEqual(config.web_search.ui_lang, "pl-PL")
self.assertEqual(config.web_search.safesearch, "off")
if __name__ == "__main__":
unittest.main()

84
tests/test_web_search.py Normal file
View File

@@ -0,0 +1,84 @@
from __future__ import annotations
import unittest
import httpx
from app.ai.web_search import (
_format_error_response,
build_query,
format_results,
parse_response,
)
from app.config import AIWebSearchConfig
class WebSearchTests(unittest.TestCase):
def test_build_query_clamps_and_defaults(self) -> None:
config = AIWebSearchConfig(api_key="key", count=5)
query = build_query(
config,
{
"query": " python web framework ",
"count": 99,
"offset": -1,
"freshness": "pw",
},
)
self.assertEqual(query.query, "python web framework")
self.assertEqual(query.count, 20)
self.assertEqual(query.offset, 0)
self.assertEqual(query.freshness, "pw")
self.assertEqual(query.safesearch, "moderate")
def test_build_query_rejects_invalid_freshness(self) -> None:
with self.assertRaisesRegex(ValueError, "freshness"):
build_query(AIWebSearchConfig(), {"query": "news", "freshness": "yesterday"})
def test_parse_and_format_results(self) -> None:
response = parse_response(
"openai",
{
"query": {"altered": "open ai", "more_results_available": True},
"web": {
"results": [
{
"title": "OpenAI",
"url": "https://openai.com",
"description": "AI research and products.",
"age": "May 18, 2026",
"extra_snippets": ["Models and APIs.", ""],
},
{"title": "", "url": ""},
]
},
},
)
self.assertEqual(len(response.results), 1)
self.assertTrue(response.more_results_available)
formatted = format_results(response)
self.assertIn("Search results for: openai", formatted)
self.assertIn("https://openai.com", formatted)
self.assertIn("Models and APIs.", formatted)
def test_format_empty_results(self) -> None:
formatted = format_results(parse_response("nothing", {"web": {"results": []}}))
self.assertIn("No results found.", formatted)
def test_format_error_response_uses_brave_detail(self) -> None:
request = httpx.Request("GET", "https://api.search.brave.com/res/v1/web/search")
response = httpx.Response(
429,
request=request,
json={"error": {"detail": "Rate limit exceeded"}},
)
self.assertIn("Rate limit exceeded", _format_error_response(response))
if __name__ == "__main__":
unittest.main()