Add Brave web search support
This commit is contained in:
@@ -2,4 +2,10 @@ TG_API_ID=123456
|
|||||||
TG_API_HASH=your_api_hash_here
|
TG_API_HASH=your_api_hash_here
|
||||||
SESSION_DIR=sessions
|
SESSION_DIR=sessions
|
||||||
ACCOUNTS_JSON=[{"name":"personal","session":"personal"},{"name":"work","session":"work"}]
|
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"}}}
|
# 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"}}}
|
||||||
|
|||||||
@@ -14,7 +14,8 @@ from openai import AsyncOpenAI
|
|||||||
|
|
||||||
from app.ai.parser import AIInvocation
|
from app.ai.parser import AIInvocation
|
||||||
from app.ai.prompts import PromptBook
|
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
|
from app.config import AIConfig, AIModelConfig, AIProviderConfig
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -62,6 +63,7 @@ class AIService:
|
|||||||
self.prompt_book = PromptBook(config.prompt.root)
|
self.prompt_book = PromptBook(config.prompt.root)
|
||||||
self._clients: dict[str, AsyncOpenAI] = {}
|
self._clients: dict[str, AsyncOpenAI] = {}
|
||||||
self._system_prompt: str | None = None
|
self._system_prompt: str | None = None
|
||||||
|
self.web_search = BraveWebSearchClient(config.web_search)
|
||||||
|
|
||||||
async def initialize(self) -> None:
|
async def initialize(self) -> None:
|
||||||
return None
|
return None
|
||||||
@@ -210,7 +212,10 @@ class AIService:
|
|||||||
return messages
|
return messages
|
||||||
|
|
||||||
async def _build_tool_list(self) -> list[dict]:
|
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, ...]:
|
def _get_allowed_roots(self) -> tuple[Path, ...]:
|
||||||
# [0] = base for resolution (project root); [1:] = permitted subdirs
|
# [0] = base for resolution (project root); [1:] = permitted subdirs
|
||||||
@@ -339,6 +344,9 @@ class AIService:
|
|||||||
try:
|
try:
|
||||||
if name == "python_exec":
|
if name == "python_exec":
|
||||||
result = await self._execute_python_exec(arguments, ctx)
|
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":
|
elif name == "send_file":
|
||||||
result = await self._execute_send_file(arguments, ctx, allowed_roots)
|
result = await self._execute_send_file(arguments, ctx, allowed_roots)
|
||||||
elif name in FILE_TOOL_DISPATCH:
|
elif name in FILE_TOOL_DISPATCH:
|
||||||
|
|||||||
@@ -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]]
|
ToolExecutor = Callable[[dict, tuple[Path, ...]], Awaitable[str]]
|
||||||
|
|
||||||
FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = {
|
FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = {
|
||||||
|
|||||||
244
app/ai/web_search.py
Normal file
244
app/ai/web_search.py
Normal 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}"
|
||||||
@@ -2,12 +2,14 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import html
|
import html
|
||||||
import re
|
import re
|
||||||
|
import shlex
|
||||||
from textwrap import wrap
|
from textwrap import wrap
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from pyrogram import enums
|
from pyrogram import enums
|
||||||
|
|
||||||
from app.ai import AIInvocation, parse_ai_invocation
|
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
|
from app.commands.core import CommandRegistry, normalize_command_name
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -322,6 +324,67 @@ async def echo(ctx: "CommandContext") -> None:
|
|||||||
await ctx.reply(text)
|
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(
|
@command_registry.command(
|
||||||
"ai",
|
"ai",
|
||||||
"Ask a configurable AI model",
|
"Ask a configurable AI model",
|
||||||
|
|||||||
@@ -60,6 +60,24 @@ class AIPromptConfig:
|
|||||||
user_template: str | None = None
|
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)
|
@dataclass(slots=True)
|
||||||
class AIConfig:
|
class AIConfig:
|
||||||
default_model: str
|
default_model: str
|
||||||
@@ -69,6 +87,7 @@ class AIConfig:
|
|||||||
models: dict[str, AIModelConfig] = field(default_factory=dict)
|
models: dict[str, AIModelConfig] = field(default_factory=dict)
|
||||||
workspace_dir: Path = field(default_factory=lambda: Path("workspace"))
|
workspace_dir: Path = field(default_factory=lambda: Path("workspace"))
|
||||||
tool_max_rounds: int = 10
|
tool_max_rounds: int = 10
|
||||||
|
web_search: AIWebSearchConfig = field(default_factory=AIWebSearchConfig)
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@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
|
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]:
|
def _parse_mapping(value: Any, *, field_name: str) -> dict[str, Any]:
|
||||||
if value is None:
|
if value is None:
|
||||||
return {}
|
return {}
|
||||||
@@ -135,6 +168,69 @@ def _parse_mapping(value: Any, *, field_name: str) -> dict[str, Any]:
|
|||||||
return value
|
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]:
|
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:
|
if raw_accounts is None:
|
||||||
return []
|
return []
|
||||||
@@ -287,6 +383,8 @@ def _parse_ai_config(raw_ai: Any, *, root: Path) -> AIConfig | None:
|
|||||||
if tool_max_rounds < 1:
|
if tool_max_rounds < 1:
|
||||||
raise ValueError("ai.tool_max_rounds must be >= 1")
|
raise ValueError("ai.tool_max_rounds must be >= 1")
|
||||||
|
|
||||||
|
web_search = _parse_web_search_config(raw_ai)
|
||||||
|
|
||||||
return AIConfig(
|
return AIConfig(
|
||||||
default_model=default_model,
|
default_model=default_model,
|
||||||
context_messages=context_messages,
|
context_messages=context_messages,
|
||||||
@@ -295,6 +393,7 @@ def _parse_ai_config(raw_ai: Any, *, root: Path) -> AIConfig | None:
|
|||||||
models=models,
|
models=models,
|
||||||
workspace_dir=workspace_dir,
|
workspace_dir=workspace_dir,
|
||||||
tool_max_rounds=tool_max_rounds,
|
tool_max_rounds=tool_max_rounds,
|
||||||
|
web_search=web_search,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -17,6 +17,18 @@
|
|||||||
"context_messages": 12,
|
"context_messages": 12,
|
||||||
"workspace_dir": "workspace",
|
"workspace_dir": "workspace",
|
||||||
"tool_max_rounds": 10,
|
"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": {
|
"prompt": {
|
||||||
"root": "prompts/ai",
|
"root": "prompts/ai",
|
||||||
"system_file": "system.md",
|
"system_file": "system.md",
|
||||||
|
|||||||
@@ -2,4 +2,5 @@ kurigram
|
|||||||
tgcrypto
|
tgcrypto
|
||||||
python-dotenv
|
python-dotenv
|
||||||
openai
|
openai
|
||||||
|
httpx
|
||||||
elevenlabs
|
elevenlabs
|
||||||
|
|||||||
67
tests/test_config_web_search.py
Normal file
67
tests/test_config_web_search.py
Normal 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
84
tests/test_web_search.py
Normal 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()
|
||||||
Reference in New Issue
Block a user