feat: rebuild AI commands and migrate config to TOML
This commit is contained in:
113
app/agent_tools.py
Normal file
113
app/agent_tools.py
Normal file
@@ -0,0 +1,113 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Any
|
||||
|
||||
from micro_agents import Injected, SessionContext, Tool, ToolManager
|
||||
|
||||
|
||||
OWNER_ONLY_TOOL_NAMES = frozenset(
|
||||
{
|
||||
"list_files",
|
||||
"read_file",
|
||||
"attach_files",
|
||||
"write_file",
|
||||
"make_directory",
|
||||
"execute_python",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class AgentToolManager(ToolManager):
|
||||
async def is_allowed(self, tool: Tool, context: SessionContext) -> bool:
|
||||
if tool.name in OWNER_ONLY_TOOL_NAMES:
|
||||
return context.get("is_owner") is True
|
||||
return True
|
||||
|
||||
|
||||
class TelegramTools:
|
||||
def __init__(
|
||||
self,
|
||||
client: Any,
|
||||
chat_id: int | str,
|
||||
*,
|
||||
base_dir: Path,
|
||||
guest_base_dir: Path | None = None,
|
||||
reply_to_message_id: int | None = None,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self._chat_id = chat_id
|
||||
self._base_dir = base_dir.resolve()
|
||||
self._guest_base_dir = (guest_base_dir or base_dir).resolve()
|
||||
self._reply_to_message_id = reply_to_message_id
|
||||
|
||||
def tools(self) -> list[Tool]:
|
||||
return [
|
||||
Tool(self.send_message, name="send_message", strict=False),
|
||||
Tool(self.send_image, name="send_image", strict=False),
|
||||
Tool(self.send_file, name="send_file", strict=False),
|
||||
]
|
||||
|
||||
async def send_message(self, text: str) -> str:
|
||||
"""Send a text message to the current Telegram chat."""
|
||||
if not text.strip():
|
||||
raise ValueError("text must not be empty")
|
||||
message = await self._client.send_message(
|
||||
self._chat_id,
|
||||
text,
|
||||
reply_to_message_id=self._reply_to_message_id,
|
||||
)
|
||||
return self._result("message", message)
|
||||
|
||||
async def send_image(
|
||||
self,
|
||||
path: str,
|
||||
caption: str = "",
|
||||
context: Annotated[SessionContext | None, Injected] = None,
|
||||
) -> str:
|
||||
"""Send a local image from the agent workspace to the current Telegram chat."""
|
||||
target = self._resolve_file(path, context)
|
||||
message = await self._client.send_photo(
|
||||
self._chat_id,
|
||||
str(target),
|
||||
caption=caption or None,
|
||||
reply_to_message_id=self._reply_to_message_id,
|
||||
)
|
||||
return self._result("image", message)
|
||||
|
||||
async def send_file(
|
||||
self,
|
||||
path: str,
|
||||
caption: str = "",
|
||||
context: Annotated[SessionContext | None, Injected] = None,
|
||||
) -> str:
|
||||
"""Send a local file from the agent workspace to the current Telegram chat."""
|
||||
target = self._resolve_file(path, context)
|
||||
message = await self._client.send_document(
|
||||
self._chat_id,
|
||||
str(target),
|
||||
caption=caption or None,
|
||||
reply_to_message_id=self._reply_to_message_id,
|
||||
)
|
||||
return self._result("file", message)
|
||||
|
||||
def _resolve_file(self, value: str, context: SessionContext | None) -> Path:
|
||||
base_dir = (
|
||||
self._base_dir
|
||||
if context is not None and context.get("is_owner") is True
|
||||
else self._guest_base_dir
|
||||
)
|
||||
path = Path(value)
|
||||
if path.is_absolute():
|
||||
raise ValueError("Absolute paths are not allowed")
|
||||
target = (base_dir / path).resolve()
|
||||
if target != base_dir and base_dir not in target.parents:
|
||||
raise ValueError("Path escapes the agent workspace")
|
||||
if not target.is_file():
|
||||
raise ValueError(f"File does not exist: {value}")
|
||||
return target
|
||||
|
||||
@staticmethod
|
||||
def _result(kind: str, message: Any) -> str:
|
||||
message_id = getattr(message, "id", None)
|
||||
return f"Sent Telegram {kind}" + (f" (message_id={message_id})" if message_id else "")
|
||||
508
app/agentic.py
Normal file
508
app/agentic.py
Normal file
@@ -0,0 +1,508 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import mimetypes
|
||||
import re
|
||||
import tempfile
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from micro_agents import Aggregator, Media, OpenAIBasedLLM, UserMessage
|
||||
from micro_agents_tools import BraveSearchTools, FilesystemTools, PythonTools
|
||||
|
||||
from app.agent_tools import AgentToolManager, TelegramTools
|
||||
from app.config import AiConfig, AiProviderConfig
|
||||
|
||||
|
||||
class AgentError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class AgentSyntaxError(AgentError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AgentInvocation:
|
||||
prompt: str
|
||||
context_messages: int = 0
|
||||
use_deepseek: bool = False
|
||||
thinking: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PromptTemplates:
|
||||
master: str
|
||||
user: str
|
||||
|
||||
|
||||
@dataclass(slots=True, eq=False)
|
||||
class AgentSessionHandle:
|
||||
session: Any
|
||||
workspace: tempfile.TemporaryDirectory[str]
|
||||
response_keys: set[tuple[Any, Any]]
|
||||
|
||||
def close(self) -> None:
|
||||
self.workspace.cleanup()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AgentTurnResult:
|
||||
answer: str
|
||||
handle: AgentSessionHandle
|
||||
|
||||
|
||||
class AgentSessionStore:
|
||||
def __init__(self, max_sessions: int = 100) -> None:
|
||||
if max_sessions < 1:
|
||||
raise ValueError("max_sessions must be positive")
|
||||
self._max_sessions = max_sessions
|
||||
self._by_response: dict[tuple[Any, Any], AgentSessionHandle] = {}
|
||||
self._handles: OrderedDict[AgentSessionHandle, None] = OrderedDict()
|
||||
|
||||
def find(self, chat_id: Any, message_id: Any) -> AgentSessionHandle | None:
|
||||
handle = self._by_response.get((chat_id, message_id))
|
||||
if handle is not None and handle in self._handles:
|
||||
self._handles.move_to_end(handle)
|
||||
return handle
|
||||
|
||||
def register_response(self, handle: AgentSessionHandle, message: Any) -> None:
|
||||
key = _message_key(message)
|
||||
if key[0] is None or key[1] is None:
|
||||
return
|
||||
self._handles[handle] = None
|
||||
self._handles.move_to_end(handle)
|
||||
handle.response_keys.add(key)
|
||||
self._by_response[key] = handle
|
||||
self._evict()
|
||||
|
||||
def discard(self, handle: AgentSessionHandle) -> None:
|
||||
self._handles.pop(handle, None)
|
||||
for key in handle.response_keys:
|
||||
if self._by_response.get(key) is handle:
|
||||
self._by_response.pop(key, None)
|
||||
handle.response_keys.clear()
|
||||
handle.close()
|
||||
|
||||
def clear(self) -> None:
|
||||
for handle in list(self._handles):
|
||||
self.discard(handle)
|
||||
|
||||
def _evict(self) -> None:
|
||||
while len(self._handles) > self._max_sessions:
|
||||
oldest = next(iter(self._handles))
|
||||
self.discard(oldest)
|
||||
|
||||
|
||||
_PLACEHOLDER = re.compile(
|
||||
r"\{\{\s*([A-Za-z_]\w*(?:\.\w+)*)\s*\}\}|\{([A-Za-z_]\w*(?:\.\w+)*)\}"
|
||||
)
|
||||
_BLOCK_COMMENT = re.compile(r"/\*.*?\*/", re.DOTALL)
|
||||
_LINE_COMMENT = re.compile(r"(?m)(?<!\S)//.*$")
|
||||
|
||||
|
||||
def strip_prompt_comments(value: str) -> str:
|
||||
without_blocks = _BLOCK_COMMENT.sub(
|
||||
lambda match: "\n" * match.group(0).count("\n"),
|
||||
value,
|
||||
)
|
||||
return _LINE_COMMENT.sub("", without_blocks)
|
||||
|
||||
|
||||
def load_prompt_templates(prompt_dir: Path) -> PromptTemplates:
|
||||
try:
|
||||
master = strip_prompt_comments(
|
||||
(prompt_dir / "MASTER.txt").read_text(encoding="utf-8")
|
||||
)
|
||||
user = strip_prompt_comments(
|
||||
(prompt_dir / "USER.txt").read_text(encoding="utf-8")
|
||||
)
|
||||
except OSError as exc:
|
||||
raise AgentError("Could not load prompts/MASTER.txt and prompts/USER.txt") from exc
|
||||
if not master.strip() or not user.strip():
|
||||
raise AgentError("Prompt templates must not be empty")
|
||||
return PromptTemplates(master=master, user=user)
|
||||
|
||||
|
||||
def _template_value(context: Mapping[str, Any], path: str) -> Any:
|
||||
value: Any = context
|
||||
for part in path.split("."):
|
||||
if not isinstance(value, Mapping) or part not in value:
|
||||
raise AgentError(f"Unknown prompt placeholder: {path}")
|
||||
value = value[part]
|
||||
return value
|
||||
|
||||
|
||||
def render_prompt_template(template: str, context: Mapping[str, Any]) -> str:
|
||||
values = dict(context)
|
||||
if "source_text" in values:
|
||||
values["text"] = values["source_text"]
|
||||
|
||||
def replace(match: re.Match[str]) -> str:
|
||||
path = match.group(1) or match.group(2)
|
||||
value = _template_value(values, path)
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, bool):
|
||||
return str(value).lower()
|
||||
return str(value)
|
||||
|
||||
rendered = _PLACEHOLDER.sub(replace, template).strip()
|
||||
if not rendered:
|
||||
raise AgentError("Rendered prompt must not be empty")
|
||||
return rendered
|
||||
|
||||
|
||||
def parse_agent_invocation(value: str, config: AiConfig) -> AgentInvocation:
|
||||
tokens = value.split()
|
||||
if not tokens:
|
||||
raise AgentSyntaxError("AI prompt is required")
|
||||
|
||||
index = 0
|
||||
context_messages = 0
|
||||
if tokens[0].lstrip("+-").isdigit():
|
||||
try:
|
||||
context_messages = int(tokens[0])
|
||||
except ValueError as exc: # pragma: no cover - guarded by isdigit
|
||||
raise AgentSyntaxError("Invalid context message count") from exc
|
||||
if not 0 <= context_messages <= config.max_context_messages:
|
||||
raise AgentSyntaxError("Context message count is outside the configured range")
|
||||
index = 1
|
||||
|
||||
use_deepseek = False
|
||||
thinking = False
|
||||
while index < len(tokens) and tokens[index].startswith("."):
|
||||
flag = tokens[index].lower()
|
||||
if flag == ".ds":
|
||||
use_deepseek = True
|
||||
elif flag == ".th":
|
||||
thinking = True
|
||||
index += 1
|
||||
|
||||
prompt = " ".join(tokens[index:]).strip()
|
||||
if not prompt:
|
||||
raise AgentSyntaxError("AI prompt is required")
|
||||
|
||||
return AgentInvocation(
|
||||
prompt=prompt,
|
||||
context_messages=context_messages,
|
||||
use_deepseek=use_deepseek,
|
||||
thinking=thinking,
|
||||
)
|
||||
|
||||
|
||||
def _message_key(message: Any) -> tuple[Any, Any]:
|
||||
chat = getattr(message, "chat", None)
|
||||
return getattr(chat, "id", None), getattr(message, "id", None)
|
||||
|
||||
|
||||
def _user_context(user: Any | None) -> dict[str, Any]:
|
||||
if user is None:
|
||||
return {"id": None, "name": "unknown", "username": None}
|
||||
return {
|
||||
"id": getattr(user, "id", None),
|
||||
"name": (
|
||||
getattr(user, "full_name", None)
|
||||
or getattr(user, "first_name", None)
|
||||
or getattr(user, "username", None)
|
||||
or "unknown"
|
||||
),
|
||||
"username": getattr(user, "username", None),
|
||||
}
|
||||
|
||||
|
||||
def _chat_context(chat: Any | None) -> dict[str, Any]:
|
||||
chat_type = getattr(chat, "type", None)
|
||||
return {
|
||||
"id": getattr(chat, "id", None),
|
||||
"title": getattr(chat, "title", None),
|
||||
"type": getattr(chat_type, "value", chat_type),
|
||||
"username": getattr(chat, "username", None),
|
||||
}
|
||||
|
||||
|
||||
def _agent_context(command_message: Any, account: Any, invocation: AgentInvocation) -> dict[str, Any]:
|
||||
caller = getattr(command_message, "from_user", None)
|
||||
owner = getattr(account, "me", None)
|
||||
reply = getattr(command_message, "reply_to_message", None)
|
||||
account_config = getattr(account, "config", None)
|
||||
caller_id = getattr(caller, "id", None)
|
||||
owner_id = getattr(owner, "id", None)
|
||||
return {
|
||||
"caller_user": _user_context(caller),
|
||||
"owner": _user_context(owner),
|
||||
"is_owner": owner_id is not None and caller_id == owner_id,
|
||||
"text": invocation.prompt,
|
||||
"chat": _chat_context(getattr(command_message, "chat", None)),
|
||||
"message": _message_context(command_message, is_reply=False),
|
||||
"reply": {
|
||||
"id": getattr(reply, "id", None),
|
||||
"text": (
|
||||
getattr(reply, "text", None) or getattr(reply, "caption", None) or ""
|
||||
if reply is not None
|
||||
else ""
|
||||
),
|
||||
"user": _user_context(getattr(reply, "from_user", None)),
|
||||
},
|
||||
"account": {
|
||||
"name": getattr(account_config, "name", None),
|
||||
"session": getattr(account_config, "session", None),
|
||||
},
|
||||
"provider": "deepseek" if invocation.use_deepseek else "glm",
|
||||
"thinking": invocation.thinking,
|
||||
"context_messages": invocation.context_messages,
|
||||
}
|
||||
|
||||
|
||||
def _message_context(message: Any, *, is_reply: bool) -> dict[str, Any]:
|
||||
date = getattr(message, "date", None)
|
||||
return {
|
||||
"id": getattr(message, "id", None),
|
||||
"text": getattr(message, "text", None) or getattr(message, "caption", None) or "",
|
||||
"date": date.isoformat() if hasattr(date, "isoformat") else None,
|
||||
"user": _user_context(getattr(message, "from_user", None)),
|
||||
"is_reply": is_reply,
|
||||
}
|
||||
|
||||
|
||||
def _media_object(message: Any) -> Any | None:
|
||||
for field_name in (
|
||||
"photo",
|
||||
"document",
|
||||
"audio",
|
||||
"voice",
|
||||
"video",
|
||||
"video_note",
|
||||
"animation",
|
||||
"sticker",
|
||||
):
|
||||
value = getattr(message, field_name, None)
|
||||
if value is not None:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _media_mime_type(media_object: Any, path: Path) -> str:
|
||||
configured = getattr(media_object, "mime_type", None)
|
||||
if configured:
|
||||
return str(configured)
|
||||
guessed = mimetypes.guess_type(path.name)[0]
|
||||
return guessed or "application/octet-stream"
|
||||
|
||||
|
||||
def _media_kind(mime_type: str) -> str:
|
||||
if mime_type.startswith("image/"):
|
||||
return "image"
|
||||
if mime_type in {"audio/mpeg", "audio/wav", "audio/x-wav"}:
|
||||
return "audio"
|
||||
return "file"
|
||||
|
||||
|
||||
async def _download_message_media(
|
||||
client: Any,
|
||||
message: Any,
|
||||
directory: Path,
|
||||
max_bytes: int,
|
||||
) -> list[Media]:
|
||||
media_object = _media_object(message)
|
||||
if media_object is None:
|
||||
return []
|
||||
|
||||
declared_size = getattr(media_object, "file_size", None)
|
||||
if not isinstance(declared_size, int) or declared_size > max_bytes:
|
||||
return []
|
||||
|
||||
message_directory = directory / str(getattr(message, "id", "unknown"))
|
||||
message_directory.mkdir(parents=True, exist_ok=True)
|
||||
downloaded = await client.download_media(message, file_name=str(message_directory) + "/")
|
||||
if not isinstance(downloaded, str):
|
||||
raise AgentError("Telegram media download returned no file")
|
||||
|
||||
path = Path(downloaded)
|
||||
if path.stat().st_size > max_bytes:
|
||||
return []
|
||||
mime_type = _media_mime_type(media_object, path)
|
||||
return [
|
||||
Media(
|
||||
kind=_media_kind(mime_type),
|
||||
data=path.read_bytes(),
|
||||
filename=getattr(media_object, "file_name", None) or path.name,
|
||||
mime_type=mime_type,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
async def _get_history(client: Any, command_message: Any, limit: int) -> list[Any]:
|
||||
if limit == 0:
|
||||
return []
|
||||
chat_id = getattr(getattr(command_message, "chat", None), "id", None)
|
||||
if chat_id is None:
|
||||
raise AgentError("Command message has no chat")
|
||||
|
||||
history = client.get_chat_history(
|
||||
chat_id,
|
||||
limit=limit,
|
||||
offset_id=getattr(command_message, "id", 0),
|
||||
)
|
||||
if inspect.isawaitable(history):
|
||||
history = await history
|
||||
if hasattr(history, "__aiter__"):
|
||||
messages = [message async for message in history]
|
||||
else:
|
||||
messages = list(history)
|
||||
return list(reversed(messages))
|
||||
|
||||
|
||||
async def _build_context_messages(
|
||||
client: Any,
|
||||
command_message: Any,
|
||||
invocation: AgentInvocation,
|
||||
config: AiConfig,
|
||||
directory: Path,
|
||||
*,
|
||||
user_template: str,
|
||||
template_context: Mapping[str, Any],
|
||||
) -> list[UserMessage]:
|
||||
telegram_messages = await _get_history(client, command_message, invocation.context_messages)
|
||||
seen = {_message_key(message) for message in telegram_messages}
|
||||
reply = getattr(command_message, "reply_to_message", None)
|
||||
reply_key = _message_key(reply) if reply is not None else None
|
||||
if reply is not None and reply_key not in seen:
|
||||
telegram_messages.append(reply)
|
||||
|
||||
messages: list[UserMessage] = []
|
||||
for telegram_message in telegram_messages:
|
||||
media = await _download_message_media(
|
||||
client,
|
||||
telegram_message,
|
||||
directory,
|
||||
config.max_media_bytes,
|
||||
)
|
||||
values = dict(template_context)
|
||||
values["text"] = (
|
||||
getattr(telegram_message, "text", None)
|
||||
or getattr(telegram_message, "caption", None)
|
||||
or "[media without a caption]"
|
||||
)
|
||||
values["message"] = _message_context(
|
||||
telegram_message,
|
||||
is_reply=_message_key(telegram_message) == reply_key,
|
||||
)
|
||||
messages.append(
|
||||
UserMessage(render_prompt_template(user_template, values), media)
|
||||
)
|
||||
return messages
|
||||
|
||||
|
||||
def _provider(config: AiConfig, invocation: AgentInvocation) -> AiProviderConfig:
|
||||
return config.deepseek if invocation.use_deepseek else config.glm
|
||||
|
||||
|
||||
async def run_agent_turn(
|
||||
client: Any,
|
||||
command_message: Any,
|
||||
account: Any,
|
||||
session_store: AgentSessionStore,
|
||||
invocation: AgentInvocation,
|
||||
config: AiConfig,
|
||||
) -> AgentTurnResult:
|
||||
chat_id = getattr(getattr(command_message, "chat", None), "id", None)
|
||||
if chat_id is None:
|
||||
raise AgentError("Command message has no chat")
|
||||
|
||||
reply = getattr(command_message, "reply_to_message", None)
|
||||
existing = (
|
||||
session_store.find(chat_id, getattr(reply, "id", None))
|
||||
if reply is not None
|
||||
else None
|
||||
)
|
||||
if existing is not None:
|
||||
context = _agent_context(command_message, account, invocation)
|
||||
context["base_dir"] = config.base_dir
|
||||
existing.session.context.clear()
|
||||
existing.session.context.update(context)
|
||||
return AgentTurnResult(
|
||||
answer=await _run_session(existing.session, invocation.prompt),
|
||||
handle=existing,
|
||||
)
|
||||
|
||||
provider = _provider(config, invocation)
|
||||
if not provider.api_key:
|
||||
raise AgentError("The selected AI provider has no API key")
|
||||
|
||||
llm = OpenAIBasedLLM(
|
||||
api_url=provider.api_url,
|
||||
api_key=provider.api_key,
|
||||
model=provider.model,
|
||||
)
|
||||
templates = load_prompt_templates(config.prompt_dir)
|
||||
context = _agent_context(command_message, account, invocation)
|
||||
context["base_dir"] = config.base_dir
|
||||
|
||||
workspace = tempfile.TemporaryDirectory(prefix="multi-bot-ai-")
|
||||
try:
|
||||
temp_path = Path(workspace.name)
|
||||
messages = await _build_context_messages(
|
||||
client,
|
||||
command_message,
|
||||
invocation,
|
||||
config,
|
||||
temp_path,
|
||||
user_template=templates.user,
|
||||
template_context=context,
|
||||
)
|
||||
tools = [
|
||||
*BraveSearchTools(
|
||||
api_key=config.brave_api_key,
|
||||
timeout_seconds=config.brave_timeout_seconds,
|
||||
).tools(),
|
||||
*FilesystemTools(max_attachment_bytes=config.max_media_bytes).tools(),
|
||||
*PythonTools().tools(),
|
||||
*TelegramTools(
|
||||
client,
|
||||
chat_id,
|
||||
base_dir=config.base_dir,
|
||||
guest_base_dir=temp_path,
|
||||
reply_to_message_id=getattr(command_message, "id", None),
|
||||
).tools(),
|
||||
]
|
||||
agent = Aggregator(
|
||||
llm,
|
||||
AgentToolManager(tools),
|
||||
master_prompt=lambda values: render_prompt_template(
|
||||
load_prompt_templates(config.prompt_dir).master,
|
||||
values,
|
||||
),
|
||||
user_prompt=lambda values: render_prompt_template(
|
||||
load_prompt_templates(config.prompt_dir).user,
|
||||
values,
|
||||
),
|
||||
max_tool_rounds=config.max_tool_rounds,
|
||||
thinking=True if invocation.thinking else None,
|
||||
)
|
||||
session = agent.create_session(
|
||||
context=context,
|
||||
session_name=f"telegram:{chat_id}",
|
||||
messages=messages,
|
||||
)
|
||||
answer = await _run_session(session, invocation.prompt)
|
||||
return AgentTurnResult(
|
||||
answer=answer,
|
||||
handle=AgentSessionHandle(session, workspace, set()),
|
||||
)
|
||||
except Exception:
|
||||
workspace.cleanup()
|
||||
raise
|
||||
|
||||
|
||||
async def _run_session(session: Any, prompt: str) -> str:
|
||||
result = await session.run(prompt)
|
||||
if not result.succeeded or result.final_message is None:
|
||||
raise AgentError("Agent loop failed") from result.error
|
||||
answer = result.final_message.message().strip()
|
||||
if not answer:
|
||||
raise AgentError("Agent returned an empty answer")
|
||||
return answer
|
||||
@@ -1,9 +0,0 @@
|
||||
from app.ai.parser import AIInvocation, parse_ai_invocation
|
||||
from app.ai.service import AIResponse, AIService
|
||||
|
||||
__all__ = [
|
||||
"AIInvocation",
|
||||
"AIResponse",
|
||||
"AIService",
|
||||
"parse_ai_invocation",
|
||||
]
|
||||
@@ -1,310 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from html.parser import HTMLParser
|
||||
import ipaddress
|
||||
import re
|
||||
import socket
|
||||
from typing import Any
|
||||
from urllib.parse import urljoin, urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import AIPageExtractConfig
|
||||
|
||||
TEXT_BLOCK_TAGS = {
|
||||
"address",
|
||||
"article",
|
||||
"aside",
|
||||
"blockquote",
|
||||
"br",
|
||||
"dd",
|
||||
"div",
|
||||
"dl",
|
||||
"dt",
|
||||
"figcaption",
|
||||
"footer",
|
||||
"h1",
|
||||
"h2",
|
||||
"h3",
|
||||
"h4",
|
||||
"h5",
|
||||
"h6",
|
||||
"header",
|
||||
"li",
|
||||
"main",
|
||||
"nav",
|
||||
"ol",
|
||||
"p",
|
||||
"pre",
|
||||
"section",
|
||||
"table",
|
||||
"td",
|
||||
"th",
|
||||
"tr",
|
||||
"ul",
|
||||
}
|
||||
SKIP_TAGS = {"script", "style", "noscript", "template", "svg", "canvas"}
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PageExtractResult:
|
||||
url: str
|
||||
final_url: str
|
||||
title: str
|
||||
text: str
|
||||
method: str
|
||||
truncated: bool = False
|
||||
|
||||
|
||||
class ReadableTextParser(HTMLParser):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(convert_charrefs=True)
|
||||
self.title = ""
|
||||
self._in_title = False
|
||||
self._skip_depth = 0
|
||||
self._parts: list[str] = []
|
||||
|
||||
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
|
||||
tag = tag.lower()
|
||||
if tag in SKIP_TAGS:
|
||||
self._skip_depth += 1
|
||||
return
|
||||
if tag == "title":
|
||||
self._in_title = True
|
||||
return
|
||||
if tag in TEXT_BLOCK_TAGS:
|
||||
self._parts.append("\n")
|
||||
|
||||
def handle_endtag(self, tag: str) -> None:
|
||||
tag = tag.lower()
|
||||
if tag in SKIP_TAGS and self._skip_depth:
|
||||
self._skip_depth -= 1
|
||||
return
|
||||
if tag == "title":
|
||||
self._in_title = False
|
||||
return
|
||||
if tag in TEXT_BLOCK_TAGS:
|
||||
self._parts.append("\n")
|
||||
|
||||
def handle_data(self, data: str) -> None:
|
||||
if self._skip_depth:
|
||||
return
|
||||
text = data.strip()
|
||||
if not text:
|
||||
return
|
||||
if self._in_title:
|
||||
self.title = " ".join((self.title, text)).strip()
|
||||
return
|
||||
self._parts.append(text)
|
||||
self._parts.append(" ")
|
||||
|
||||
def get_text(self) -> str:
|
||||
text = "".join(self._parts)
|
||||
text = re.sub(r"[ \t\r\f\v]+", " ", text)
|
||||
text = re.sub(r"\n\s+", "\n", text)
|
||||
text = re.sub(r"\n{3,}", "\n\n", text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def extract_readable_text(html: str) -> tuple[str, str]:
|
||||
parser = ReadableTextParser()
|
||||
parser.feed(html)
|
||||
parser.close()
|
||||
return parser.title.strip(), parser.get_text()
|
||||
|
||||
|
||||
def normalize_url(raw_url: str) -> str:
|
||||
url = str(raw_url or "").strip()
|
||||
if not url:
|
||||
raise ValueError("extract_page: 'url' is required")
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme not in {"http", "https"}:
|
||||
raise ValueError("extract_page: only http and https URLs are allowed")
|
||||
if not parsed.hostname:
|
||||
raise ValueError("extract_page: URL must include a hostname")
|
||||
return url
|
||||
|
||||
|
||||
def validate_public_url(url: str, *, block_private_hosts: bool) -> None:
|
||||
if not block_private_hosts:
|
||||
return
|
||||
parsed = urlparse(url)
|
||||
host = parsed.hostname
|
||||
if not host:
|
||||
raise ValueError("extract_page: URL must include a hostname")
|
||||
|
||||
try:
|
||||
addresses = socket.getaddrinfo(host, parsed.port or 443, type=socket.SOCK_STREAM)
|
||||
except socket.gaierror as exc:
|
||||
raise ValueError(f"extract_page: could not resolve host {host!r}") from exc
|
||||
|
||||
for address in addresses:
|
||||
ip = ipaddress.ip_address(address[4][0])
|
||||
if (
|
||||
ip.is_private
|
||||
or ip.is_loopback
|
||||
or ip.is_link_local
|
||||
or ip.is_multicast
|
||||
or ip.is_reserved
|
||||
or ip.is_unspecified
|
||||
):
|
||||
raise ValueError("extract_page: private, local, and reserved hosts are blocked")
|
||||
|
||||
|
||||
def _truncate(text: str, max_chars: int) -> tuple[str, bool]:
|
||||
if len(text) <= max_chars:
|
||||
return text, False
|
||||
return text[:max_chars].rstrip() + "\n...(truncated)", True
|
||||
|
||||
|
||||
def _format_result(result: PageExtractResult) -> str:
|
||||
lines = [
|
||||
f"URL: {result.final_url or result.url}",
|
||||
f"Method: {result.method}",
|
||||
]
|
||||
if result.title:
|
||||
lines.append(f"Title: {result.title}")
|
||||
if result.truncated:
|
||||
lines.append("Note: content was truncated.")
|
||||
lines.append("")
|
||||
lines.append(result.text or "(no readable text extracted)")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
class PageExtractor:
|
||||
def __init__(self, config: AIPageExtractConfig) -> None:
|
||||
self.config = config
|
||||
|
||||
@property
|
||||
def available(self) -> bool:
|
||||
return self.config.enabled
|
||||
|
||||
async def extract(self, arguments: dict[str, Any]) -> PageExtractResult:
|
||||
if not self.config.enabled:
|
||||
raise ValueError("Page extraction is disabled")
|
||||
|
||||
url = normalize_url(str(arguments.get("url") or ""))
|
||||
validate_public_url(url, block_private_hosts=self.config.block_private_hosts)
|
||||
|
||||
render = str(arguments.get("render", "auto")).strip().lower()
|
||||
if render not in {"auto", "never", "always"}:
|
||||
raise ValueError("extract_page: 'render' must be auto, never, or always")
|
||||
|
||||
raw_max_chars = arguments.get("max_chars", self.config.max_chars)
|
||||
try:
|
||||
max_chars = int(raw_max_chars)
|
||||
except (TypeError, ValueError):
|
||||
max_chars = self.config.max_chars
|
||||
max_chars = min(max(max_chars, 1000), 50000)
|
||||
|
||||
if render == "always":
|
||||
return await self._extract_dynamic(url, max_chars=max_chars)
|
||||
|
||||
static_result = await self._extract_static(url, max_chars=max_chars)
|
||||
if render == "never" or not self.config.render_dynamic:
|
||||
return static_result
|
||||
if len(static_result.text) >= 700:
|
||||
return static_result
|
||||
|
||||
try:
|
||||
return await self._extract_dynamic(url, max_chars=max_chars)
|
||||
except ValueError:
|
||||
return static_result
|
||||
|
||||
async def extract_text(self, arguments: dict[str, Any]) -> str:
|
||||
return _format_result(await self.extract(arguments))
|
||||
|
||||
async def _extract_static(self, url: str, *, max_chars: int) -> PageExtractResult:
|
||||
headers = {
|
||||
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
|
||||
"User-Agent": (
|
||||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
||||
"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120 Safari/537.36"
|
||||
),
|
||||
}
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=self.config.timeout,
|
||||
headers=headers,
|
||||
) as client:
|
||||
current_url = url
|
||||
for _ in range(6):
|
||||
validate_public_url(
|
||||
current_url,
|
||||
block_private_hosts=self.config.block_private_hosts,
|
||||
)
|
||||
response = await client.get(current_url)
|
||||
if response.status_code not in {301, 302, 303, 307, 308}:
|
||||
break
|
||||
location = response.headers.get("location")
|
||||
if not location:
|
||||
break
|
||||
current_url = urljoin(str(response.url), location)
|
||||
else:
|
||||
raise ValueError("extract_page: too many redirects")
|
||||
except httpx.HTTPError as exc:
|
||||
raise ValueError(f"extract_page: request failed: {exc}") from exc
|
||||
|
||||
if response.status_code >= 400:
|
||||
raise ValueError(f"extract_page: HTTP {response.status_code}")
|
||||
|
||||
content_type = response.headers.get("content-type", "").lower()
|
||||
if "text/html" not in content_type and "application/xhtml" not in content_type:
|
||||
raise ValueError(f"extract_page: unsupported content type: {content_type or 'unknown'}")
|
||||
|
||||
title, text = extract_readable_text(response.text)
|
||||
text, truncated = _truncate(text, max_chars)
|
||||
return PageExtractResult(
|
||||
url=url,
|
||||
final_url=str(response.url),
|
||||
title=title,
|
||||
text=text,
|
||||
method="http",
|
||||
truncated=truncated,
|
||||
)
|
||||
|
||||
async def _extract_dynamic(self, url: str, *, max_chars: int) -> PageExtractResult:
|
||||
try:
|
||||
from playwright.async_api import async_playwright
|
||||
except ModuleNotFoundError as exc:
|
||||
raise ValueError(
|
||||
"extract_page: dynamic rendering requires playwright and browser install"
|
||||
) from exc
|
||||
|
||||
try:
|
||||
async with async_playwright() as playwright:
|
||||
browser = await playwright.chromium.launch(headless=True)
|
||||
try:
|
||||
page = await browser.new_page()
|
||||
if self.config.block_private_hosts:
|
||||
await page.route("**/*", self._route_public_requests_only)
|
||||
await page.goto(url, wait_until="networkidle", timeout=int(self.config.timeout * 1000))
|
||||
if self.config.dynamic_wait_ms:
|
||||
await page.wait_for_timeout(self.config.dynamic_wait_ms)
|
||||
title = (await page.title()).strip()
|
||||
text = await page.locator("body").inner_text(timeout=5000)
|
||||
final_url = page.url
|
||||
finally:
|
||||
await browser.close()
|
||||
except Exception as exc:
|
||||
raise ValueError(f"extract_page: dynamic render failed: {exc}") from exc
|
||||
|
||||
text = re.sub(r"\n{3,}", "\n\n", text).strip()
|
||||
text, truncated = _truncate(text, max_chars)
|
||||
return PageExtractResult(
|
||||
url=url,
|
||||
final_url=final_url,
|
||||
title=title,
|
||||
text=text,
|
||||
method="browser",
|
||||
truncated=truncated,
|
||||
)
|
||||
|
||||
async def _route_public_requests_only(self, route: Any) -> None:
|
||||
try:
|
||||
validate_public_url(route.request.url, block_private_hosts=True)
|
||||
except ValueError:
|
||||
await route.abort()
|
||||
return
|
||||
await route.continue_()
|
||||
@@ -1,44 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import shlex
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AIInvocation:
|
||||
context_messages: int
|
||||
model_flags: list[str]
|
||||
prompt: str
|
||||
|
||||
|
||||
def parse_ai_invocation(text: str, *, default_context_messages: int) -> AIInvocation:
|
||||
try:
|
||||
tokens = shlex.split(text)
|
||||
except ValueError:
|
||||
tokens = text.split()
|
||||
|
||||
context_messages = default_context_messages
|
||||
model_flags: list[str] = []
|
||||
|
||||
if tokens:
|
||||
first = tokens[0]
|
||||
if first.isdigit():
|
||||
context_messages = int(tokens.pop(0))
|
||||
|
||||
while tokens and tokens[0].startswith("--"):
|
||||
flag = tokens.pop(0)[2:]
|
||||
if "=" in flag:
|
||||
flag = flag.split("=", 1)[0]
|
||||
flag = flag.strip().lower()
|
||||
if flag:
|
||||
model_flags.append(flag)
|
||||
|
||||
prompt = " ".join(tokens).strip()
|
||||
if not prompt:
|
||||
raise ValueError("Usage: ..ai [N] [--model] prompt")
|
||||
|
||||
return AIInvocation(
|
||||
context_messages=context_messages,
|
||||
model_flags=model_flags,
|
||||
prompt=prompt,
|
||||
)
|
||||
@@ -1,137 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import signal
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ProcessRecord:
|
||||
name: str
|
||||
pid: int
|
||||
command: list[str]
|
||||
cwd: str
|
||||
log_file: str
|
||||
started_at: str
|
||||
|
||||
|
||||
class BackgroundProcessManager:
|
||||
def __init__(self, workspace_dir: Path) -> None:
|
||||
self.workspace_dir = workspace_dir.resolve()
|
||||
self.log_dir = self.workspace_dir / ".process-logs"
|
||||
self.state_file = self.workspace_dir / ".background-processes.json"
|
||||
self.processes: dict[str, ProcessRecord] = {}
|
||||
self._children: dict[str, asyncio.subprocess.Process] = {}
|
||||
self._load()
|
||||
|
||||
def _load(self) -> None:
|
||||
try:
|
||||
raw = json.loads(self.state_file.read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError):
|
||||
return
|
||||
for name, item in raw.items():
|
||||
try:
|
||||
item = dict(item)
|
||||
item.pop("name", None)
|
||||
self.processes[name] = ProcessRecord(name=name, **item)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
|
||||
def _save(self) -> None:
|
||||
self.workspace_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.state_file.write_text(
|
||||
json.dumps({name: asdict(record) for name, record in self.processes.items()}, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def _alive(self, pid: int) -> bool:
|
||||
if pid in {os.getpid(), 0}:
|
||||
return True
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except (OSError, ProcessLookupError):
|
||||
return False
|
||||
return True
|
||||
|
||||
async def start(self, name: str, command: Any, cwd: str | None = None) -> str:
|
||||
name = name.strip()
|
||||
if not name or not name.replace("-", "").replace("_", "").isalnum():
|
||||
raise ValueError("background_start: name must contain only letters, digits, '-' or '_'")
|
||||
if isinstance(command, str):
|
||||
command = shlex.split(command, posix=(os.name != "nt"))
|
||||
if not isinstance(command, list) or not command or not all(isinstance(x, str) for x in command):
|
||||
raise ValueError("background_start: command must be a non-empty argv list or shell-like string")
|
||||
if name in self.processes and self._alive(self.processes[name].pid):
|
||||
raise ValueError(f"background_start: process {name!r} is already running")
|
||||
|
||||
process_cwd = self.workspace_dir if not cwd else (self.workspace_dir / cwd).resolve()
|
||||
if process_cwd != self.workspace_dir and self.workspace_dir not in process_cwd.parents:
|
||||
raise ValueError("background_start: cwd must be inside workspace_dir")
|
||||
process_cwd.mkdir(parents=True, exist_ok=True)
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
log_file = self.log_dir / f"{name}.log"
|
||||
handle = log_file.open("ab")
|
||||
try:
|
||||
kwargs: dict[str, Any] = {
|
||||
"cwd": str(process_cwd),
|
||||
"stdout": handle,
|
||||
"stderr": asyncio.subprocess.STDOUT,
|
||||
}
|
||||
if os.name != "nt":
|
||||
kwargs["start_new_session"] = True
|
||||
process = await asyncio.create_subprocess_exec(*command, **kwargs)
|
||||
except Exception:
|
||||
handle.close()
|
||||
raise
|
||||
finally:
|
||||
handle.close()
|
||||
|
||||
record = ProcessRecord(
|
||||
name=name,
|
||||
pid=process.pid,
|
||||
command=command,
|
||||
cwd=str(process_cwd),
|
||||
log_file=str(log_file),
|
||||
started_at=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
self.processes[name] = record
|
||||
self._children[name] = process
|
||||
self._save()
|
||||
return f"Started {name} (pid={process.pid}); log={log_file}"
|
||||
|
||||
async def stop(self, name: str) -> str:
|
||||
record = self.processes.get(name)
|
||||
if record is None:
|
||||
raise ValueError(f"background_stop: unknown process {name!r}")
|
||||
process = self._children.get(name)
|
||||
if process is not None and process.returncode is None:
|
||||
process.terminate()
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=5)
|
||||
except asyncio.TimeoutError:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
elif self._alive(record.pid):
|
||||
try:
|
||||
os.kill(record.pid, signal.SIGTERM)
|
||||
except OSError as exc:
|
||||
raise ValueError(f"background_stop failed: {exc}") from exc
|
||||
self.processes.pop(name, None)
|
||||
self._children.pop(name, None)
|
||||
self._save()
|
||||
return f"Stopped {name} (pid={record.pid})"
|
||||
|
||||
def status(self) -> str:
|
||||
if not self.processes:
|
||||
return "No background processes"
|
||||
lines = []
|
||||
for name, record in sorted(self.processes.items()):
|
||||
state = "running" if self._alive(record.pid) else "exited"
|
||||
lines.append(f"{name}: {state}, pid={record.pid}, log={record.log_file}, command={record.command!r}")
|
||||
return "\n".join(lines)
|
||||
@@ -1,64 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
import re
|
||||
|
||||
|
||||
DEFAULT_SYSTEM_PROMPT = "You are a concise, accurate Telegram assistant."
|
||||
TEMPLATE_RE = re.compile(r"{{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*}}")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PromptBook:
|
||||
root: Path
|
||||
system_file: str = "system.md"
|
||||
parts_dir: str = "parts"
|
||||
|
||||
def _read_file(self, path: Path) -> str:
|
||||
return path.read_text(encoding="utf-8").strip()
|
||||
|
||||
def _template_context(self, now: datetime | None = None) -> dict[str, str]:
|
||||
current = now.astimezone() if now else datetime.now().astimezone()
|
||||
utc = current.astimezone(UTC)
|
||||
return {
|
||||
"current_date": current.date().isoformat(),
|
||||
"current_time": current.strftime("%H:%M:%S"),
|
||||
"current_datetime": current.isoformat(timespec="seconds"),
|
||||
"current_weekday": current.strftime("%A"),
|
||||
"timezone": current.tzname() or "",
|
||||
"utc_date": utc.date().isoformat(),
|
||||
"utc_time": utc.strftime("%H:%M:%S"),
|
||||
"utc_datetime": utc.isoformat(timespec="seconds"),
|
||||
"unix_timestamp": str(int(current.timestamp())),
|
||||
}
|
||||
|
||||
def _render_template(self, text: str, now: datetime | None = None) -> str:
|
||||
context = self._template_context(now)
|
||||
|
||||
def replace(match: re.Match[str]) -> str:
|
||||
name = match.group(1)
|
||||
return context.get(name, match.group(0))
|
||||
|
||||
return TEMPLATE_RE.sub(replace, text)
|
||||
|
||||
def load_system_prompt(self) -> str:
|
||||
sections: list[str] = []
|
||||
system_path = self.root / self.system_file
|
||||
if system_path.exists():
|
||||
text = self._read_file(system_path)
|
||||
if text:
|
||||
sections.append(text)
|
||||
|
||||
parts_root = self.root / self.parts_dir
|
||||
if parts_root.exists():
|
||||
for path in sorted(parts_root.rglob("*.md")):
|
||||
if path.is_file():
|
||||
text = self._read_file(path)
|
||||
if text:
|
||||
sections.append(text)
|
||||
|
||||
if not sections:
|
||||
return DEFAULT_SYSTEM_PROMPT
|
||||
return self._render_template("\n\n".join(sections)).strip()
|
||||
@@ -1,566 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from app.ai.parser import AIInvocation
|
||||
from app.ai.prompts import PromptBook
|
||||
from app.ai.page_extract import PageExtractor
|
||||
from app.ai.processes import BackgroundProcessManager
|
||||
from app.ai.transcript import AITranscriptLogger
|
||||
from app.ai.tools import (
|
||||
FILE_TOOL_DISPATCH,
|
||||
FILE_TOOL_SCHEMAS,
|
||||
PAGE_EXTRACT_SCHEMA,
|
||||
PYTHON_EXEC_SCHEMA,
|
||||
TEMPORARY_COMMAND_CREATE_SCHEMA,
|
||||
BACKGROUND_LIST_SCHEMA,
|
||||
BACKGROUND_START_SCHEMA,
|
||||
BACKGROUND_STOP_SCHEMA,
|
||||
WEB_SEARCH_SCHEMA,
|
||||
)
|
||||
from app.commands.temporary import TemporaryCommandStore
|
||||
from app.ai.web_search import BraveWebSearchClient, format_results
|
||||
from app.config import AIConfig, AIModelConfig, AIProviderConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.commands.core import CommandContext, CommandRegistry
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
ProgressCallback = Callable[[str], Awaitable[None] | None]
|
||||
|
||||
|
||||
async def _sanitize_response_encoding(response: httpx.Response) -> None:
|
||||
await response.aread()
|
||||
raw = response.content
|
||||
try:
|
||||
raw.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
response._content = raw.decode("utf-8", errors="replace").encode("utf-8")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AIResponse:
|
||||
text: str
|
||||
model: AIModelConfig
|
||||
provider: AIProviderConfig
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ResolvedAIModel:
|
||||
alias: str
|
||||
provider: AIProviderConfig
|
||||
model: AIModelConfig
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class HistoryMessage:
|
||||
message: Any
|
||||
role: str
|
||||
text: str
|
||||
|
||||
|
||||
class AIService:
|
||||
def __init__(self, config: AIConfig, command_registry: "CommandRegistry") -> None:
|
||||
self.config = config
|
||||
self.prompt_book = PromptBook(
|
||||
config.prompt.root,
|
||||
system_file=config.prompt.system_file,
|
||||
parts_dir=config.prompt.parts_dir,
|
||||
)
|
||||
self._clients: dict[str, AsyncOpenAI] = {}
|
||||
self.web_search = BraveWebSearchClient(config.web_search)
|
||||
self.page_extractor = PageExtractor(config.page_extract)
|
||||
self.processes = BackgroundProcessManager(config.workspace_dir)
|
||||
self.temporary_commands = TemporaryCommandStore(command_registry)
|
||||
self.transcript = AITranscriptLogger(config.workspace_dir)
|
||||
|
||||
async def initialize(self) -> None:
|
||||
return None
|
||||
|
||||
def close(self) -> None:
|
||||
self.transcript.close()
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return True
|
||||
|
||||
def available_model_aliases(self) -> list[str]:
|
||||
return sorted(self.config.models.keys())
|
||||
|
||||
def describe_models(self) -> str:
|
||||
lines = [f"default={self.config.default_model}"]
|
||||
for alias in self.available_model_aliases():
|
||||
model = self.config.models[alias]
|
||||
provider = self.config.providers[model.provider]
|
||||
lines.append(f"{alias} -> {provider.name}:{model.model}")
|
||||
return "; ".join(lines)
|
||||
|
||||
def select_model_alias(self, model_flags: list[str]) -> str:
|
||||
for flag in model_flags:
|
||||
if flag in self.config.models:
|
||||
return flag
|
||||
if model_flags:
|
||||
raise ValueError(
|
||||
"Unknown model flag: " + ", ".join(f"--{flag}" for flag in model_flags)
|
||||
)
|
||||
return self.config.default_model
|
||||
|
||||
def resolve_model(self, alias: str | None) -> ResolvedAIModel:
|
||||
selected_alias = alias or self.config.default_model
|
||||
model = self.config.models.get(selected_alias)
|
||||
if model is None:
|
||||
raise ValueError(f"Unknown AI model alias: {selected_alias}")
|
||||
provider = self.config.providers.get(model.provider)
|
||||
if provider is None:
|
||||
raise ValueError(
|
||||
f"AI model {selected_alias} references unknown provider: {model.provider}"
|
||||
)
|
||||
return ResolvedAIModel(alias=selected_alias, provider=provider, model=model)
|
||||
|
||||
def _get_client(self, provider: AIProviderConfig) -> AsyncOpenAI:
|
||||
client = self._clients.get(provider.name)
|
||||
if client is None:
|
||||
client = AsyncOpenAI(
|
||||
api_key=provider.api_key,
|
||||
base_url=provider.base_url,
|
||||
timeout=provider.timeout,
|
||||
http_client=httpx.AsyncClient(
|
||||
event_hooks={"response": [_sanitize_response_encoding]},
|
||||
),
|
||||
)
|
||||
self._clients[provider.name] = client
|
||||
return client
|
||||
|
||||
def load_system_prompt(self) -> str:
|
||||
return self.prompt_book.load_system_prompt()
|
||||
|
||||
async def _maybe_call_progress(self, progress: ProgressCallback | None, text: str) -> None:
|
||||
if progress is None:
|
||||
return
|
||||
result = progress(text)
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
|
||||
def _extract_message_text(self, message: Any) -> str:
|
||||
text = getattr(message, "text", None) or getattr(message, "caption", None)
|
||||
if not isinstance(text, str):
|
||||
return ""
|
||||
return text.strip()
|
||||
|
||||
def _is_own_message(self, message: Any, owner_id: int | None) -> bool:
|
||||
user = getattr(message, "from_user", None)
|
||||
user_id = getattr(user, "id", None)
|
||||
if owner_id is None:
|
||||
return bool(getattr(message, "outgoing", False))
|
||||
return user_id == owner_id or bool(getattr(message, "outgoing", False))
|
||||
|
||||
async def collect_history(self, ctx: "CommandContext", *, limit: int) -> list[HistoryMessage]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
collected: list[HistoryMessage] = []
|
||||
current_message_id = getattr(ctx.message, "id", 0)
|
||||
owner_id = ctx.owner_id
|
||||
|
||||
async for message in ctx.client.get_chat_history(
|
||||
ctx.message.chat.id,
|
||||
limit=limit,
|
||||
offset_id=current_message_id,
|
||||
):
|
||||
message_id = getattr(message, "id", 0)
|
||||
if message_id >= current_message_id:
|
||||
continue
|
||||
|
||||
text = self._extract_message_text(message)
|
||||
role = "assistant" if self._is_own_message(message, owner_id) else "user"
|
||||
collected.append(HistoryMessage(message=message, role=role, text=text))
|
||||
if len(collected) >= limit:
|
||||
break
|
||||
|
||||
collected.reverse()
|
||||
return collected
|
||||
|
||||
def _render_user_message(self, text: str, user: Any) -> str:
|
||||
template = self.config.prompt.user_template
|
||||
if not template or not text:
|
||||
return text
|
||||
first_name = str(getattr(user, "first_name", "") or "")
|
||||
last_name = str(getattr(user, "last_name", "") or "")
|
||||
username = str(getattr(user, "username", "") or "")
|
||||
user_id = str(getattr(user, "id", "") or "")
|
||||
full_name = f"{first_name} {last_name}".strip() or username or user_id
|
||||
try:
|
||||
return template.format_map({
|
||||
"text": text,
|
||||
"user_id": user_id,
|
||||
"first_name": first_name,
|
||||
"last_name": last_name,
|
||||
"username": username,
|
||||
"full_name": full_name,
|
||||
})
|
||||
except (KeyError, ValueError):
|
||||
return text
|
||||
|
||||
async def build_messages(
|
||||
self,
|
||||
ctx: "CommandContext",
|
||||
*,
|
||||
prompt: str,
|
||||
history: list[HistoryMessage],
|
||||
) -> list[dict[str, Any]]:
|
||||
messages: list[dict[str, Any]] = [{"role": "system", "content": self.load_system_prompt()}]
|
||||
for item in history:
|
||||
if not item.text:
|
||||
continue
|
||||
if item.role == "user":
|
||||
sender = getattr(item.message, "from_user", None)
|
||||
content = self._render_user_message(item.text, sender)
|
||||
else:
|
||||
content = item.text
|
||||
messages.append({"role": item.role, "content": content})
|
||||
messages.append({"role": "user", "content": self._render_user_message(prompt, ctx.actor)})
|
||||
return messages
|
||||
|
||||
async def _build_tool_list(self) -> list[dict]:
|
||||
tools = FILE_TOOL_SCHEMAS + [
|
||||
PYTHON_EXEC_SCHEMA,
|
||||
TEMPORARY_COMMAND_CREATE_SCHEMA,
|
||||
BACKGROUND_START_SCHEMA,
|
||||
BACKGROUND_STOP_SCHEMA,
|
||||
BACKGROUND_LIST_SCHEMA,
|
||||
]
|
||||
if self.web_search.available:
|
||||
tools.append(WEB_SEARCH_SCHEMA)
|
||||
if self.page_extractor.available:
|
||||
tools.append(PAGE_EXTRACT_SCHEMA)
|
||||
return tools
|
||||
|
||||
def _get_allowed_roots(self) -> tuple[Path, ...]:
|
||||
# [0] = base for resolution (project root); [1:] = permitted subdirs
|
||||
project_root = self.config.workspace_dir.parent
|
||||
return (project_root, self.config.workspace_dir, self.config.prompt.root)
|
||||
|
||||
def _build_exec_context(self, ctx: "CommandContext") -> dict[str, Any]:
|
||||
actor = ctx.actor
|
||||
replied = ctx.replied_message
|
||||
replied_user = ctx.replied_user
|
||||
return {
|
||||
"chat_id": getattr(getattr(ctx.message, "chat", None), "id", None),
|
||||
"message_id": getattr(ctx.message, "id", None),
|
||||
"message_text": getattr(ctx.message, "text", None) or "",
|
||||
"args": ctx.args,
|
||||
"from_user": {
|
||||
"id": getattr(actor, "id", None),
|
||||
"first_name": getattr(actor, "first_name", None),
|
||||
"last_name": getattr(actor, "last_name", None),
|
||||
"username": getattr(actor, "username", None),
|
||||
} if actor else None,
|
||||
"replied_message": {
|
||||
"id": getattr(replied, "id", None),
|
||||
"text": getattr(replied, "text", None) or getattr(replied, "caption", None) or "",
|
||||
"from_user": {
|
||||
"id": getattr(replied_user, "id", None),
|
||||
"first_name": getattr(replied_user, "first_name", None),
|
||||
"username": getattr(replied_user, "username", None),
|
||||
} if replied_user else None,
|
||||
} if replied else None,
|
||||
}
|
||||
|
||||
async def _execute_python_exec(
|
||||
self,
|
||||
arguments: dict,
|
||||
ctx: "CommandContext",
|
||||
) -> str:
|
||||
code = str(arguments.get("code", "")).strip()
|
||||
if not code:
|
||||
raise ValueError("python_exec: 'code' is required")
|
||||
|
||||
raw_timeout = arguments.get("timeout", 10)
|
||||
try:
|
||||
timeout = float(raw_timeout)
|
||||
except (TypeError, ValueError):
|
||||
timeout = 10.0
|
||||
timeout = min(max(timeout, 1.0), 30.0)
|
||||
|
||||
context = self._build_exec_context(ctx)
|
||||
context_json = json.dumps(context, ensure_ascii=False, default=str)
|
||||
context_line = f"import json as _json; ctx = _json.loads({context_json!r})\n"
|
||||
full_code = (context_line + code).encode("utf-8")
|
||||
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
sys.executable, "-",
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(
|
||||
proc.communicate(input=full_code), timeout=timeout
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
return f"Execution timed out after {timeout:.0f}s"
|
||||
except OSError as exc:
|
||||
raise ValueError(f"python_exec: failed to start interpreter: {exc}") from exc
|
||||
|
||||
out = stdout.decode("utf-8", errors="replace").rstrip()
|
||||
err = stderr.decode("utf-8", errors="replace").rstrip()
|
||||
|
||||
parts: list[str] = []
|
||||
if out:
|
||||
parts.append(out)
|
||||
if err:
|
||||
parts.append(f"[stderr]\n{err}")
|
||||
if not parts:
|
||||
parts.append(f"[exit {proc.returncode}]" if proc.returncode else "(no output)")
|
||||
|
||||
result = "\n".join(parts)
|
||||
if len(result) > 3000:
|
||||
result = result[:3000] + "\n…(truncated)"
|
||||
return result
|
||||
|
||||
async def _execute_send_file(
|
||||
self,
|
||||
arguments: dict,
|
||||
ctx: "CommandContext",
|
||||
allowed_roots: tuple[Path, ...],
|
||||
) -> str:
|
||||
from app.ai.tools import _resolve_sandboxed_path
|
||||
raw_path = str(arguments.get("path", "")).strip()
|
||||
caption = arguments.get("caption") or None
|
||||
if not raw_path:
|
||||
raise ValueError("send_file: 'path' is required")
|
||||
path = _resolve_sandboxed_path(raw_path, allowed_roots)
|
||||
if not path.exists():
|
||||
raise ValueError(f"send_file: file not found: {raw_path!r}")
|
||||
if not path.is_file():
|
||||
raise ValueError(f"send_file: not a file: {raw_path!r}")
|
||||
await ctx.client.send_document(
|
||||
ctx.message.chat.id,
|
||||
document=str(path),
|
||||
caption=caption,
|
||||
)
|
||||
return f"Sent {raw_path} to chat"
|
||||
|
||||
async def _execute_tool_call(
|
||||
self,
|
||||
tool_call: Any,
|
||||
ctx: "CommandContext",
|
||||
allowed_roots: tuple[Path, ...],
|
||||
progress: ProgressCallback | None,
|
||||
) -> str:
|
||||
name = getattr(getattr(tool_call, "function", None), "name", None) or ""
|
||||
raw_args = getattr(getattr(tool_call, "function", None), "arguments", None) or "{}"
|
||||
|
||||
try:
|
||||
arguments: dict = json.loads(raw_args)
|
||||
except json.JSONDecodeError as exc:
|
||||
return f"Tool {name!r}: invalid JSON arguments: {exc}"
|
||||
|
||||
try:
|
||||
if name in {
|
||||
"python_exec",
|
||||
"temporary_command_create",
|
||||
"background_start",
|
||||
"background_stop",
|
||||
"background_list",
|
||||
} and not ctx.is_owner:
|
||||
result = "Tool access denied: temporary commands, code execution, and process control are owner-only"
|
||||
elif name == "python_exec":
|
||||
result = await self._execute_python_exec(arguments, ctx)
|
||||
elif name == "temporary_command_create":
|
||||
result = self.temporary_commands.create(
|
||||
str(arguments.get("name", "")),
|
||||
str(arguments.get("response", "")),
|
||||
str(arguments.get("description", "")) or None,
|
||||
)
|
||||
elif name == "background_start":
|
||||
result = await self.processes.start(
|
||||
str(arguments.get("name", "")), arguments.get("command"), arguments.get("cwd")
|
||||
)
|
||||
elif name == "background_stop":
|
||||
result = await self.processes.stop(str(arguments.get("name", "")))
|
||||
elif name == "background_list":
|
||||
result = self.processes.status()
|
||||
elif name == "web_search":
|
||||
response = await self.web_search.search(arguments)
|
||||
result = format_results(response, include_heading=True)
|
||||
elif name == "extract_page":
|
||||
result = await self.page_extractor.extract_text(arguments)
|
||||
elif name == "send_file":
|
||||
result = await self._execute_send_file(arguments, ctx, allowed_roots)
|
||||
elif name in FILE_TOOL_DISPATCH:
|
||||
result = await FILE_TOOL_DISPATCH[name](arguments, allowed_roots)
|
||||
else:
|
||||
result = f"Unknown tool: {name!r}"
|
||||
except (ValueError, OSError) as exc:
|
||||
result = f"Tool {name!r} error: {exc}"
|
||||
|
||||
await self._maybe_call_progress(progress, f"🛠 Tool {name}: {result[:80]}")
|
||||
return result
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
ctx: "CommandContext",
|
||||
invocation: AIInvocation,
|
||||
*,
|
||||
progress: ProgressCallback | None = None,
|
||||
) -> AIResponse:
|
||||
model = self.resolve_model(invocation.model_flags[0] if invocation.model_flags else None)
|
||||
client = self._get_client(model.provider)
|
||||
|
||||
history = await self.collect_history(ctx, limit=invocation.context_messages)
|
||||
|
||||
await self._maybe_call_progress(progress, "🛠 Generating final answer")
|
||||
messages = await self.build_messages(
|
||||
ctx,
|
||||
prompt=invocation.prompt,
|
||||
history=history,
|
||||
)
|
||||
|
||||
tools = await self._build_tool_list()
|
||||
allowed_roots = self._get_allowed_roots()
|
||||
max_rounds = self.config.tool_max_rounds
|
||||
request_id = uuid4().hex
|
||||
repeated_calls: dict[str, int] = {}
|
||||
|
||||
for round_num in range(max_rounds + 1):
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": model.model.model,
|
||||
"messages": messages,
|
||||
}
|
||||
if tools:
|
||||
kwargs["tools"] = tools
|
||||
|
||||
self.transcript.record(
|
||||
"request",
|
||||
request_id=request_id,
|
||||
round=round_num,
|
||||
model_alias=model.alias,
|
||||
provider=model.provider.name,
|
||||
payload=kwargs,
|
||||
)
|
||||
try:
|
||||
response = await client.chat.completions.create(**kwargs)
|
||||
except Exception as exc:
|
||||
self.transcript.record(
|
||||
"request_error",
|
||||
request_id=request_id,
|
||||
round=round_num,
|
||||
error=str(exc),
|
||||
)
|
||||
log.exception("AI request failed for model %s", model.alias)
|
||||
raise RuntimeError(f"AI request failed: {exc}") from exc
|
||||
|
||||
self.transcript.record(
|
||||
"response",
|
||||
request_id=request_id,
|
||||
round=round_num,
|
||||
payload=response,
|
||||
)
|
||||
|
||||
choice = response.choices[0]
|
||||
finish_reason = getattr(choice, "finish_reason", None)
|
||||
|
||||
if finish_reason != "tool_calls":
|
||||
text = self._extract_response_text(response)
|
||||
if not text:
|
||||
raise RuntimeError("AI model returned an empty response")
|
||||
return AIResponse(text=text, model=model.model, provider=model.provider)
|
||||
|
||||
if round_num >= max_rounds:
|
||||
raise RuntimeError(
|
||||
f"Agentic loop exceeded the maximum of {max_rounds} tool rounds"
|
||||
)
|
||||
|
||||
tool_calls = getattr(choice.message, "tool_calls", None) or []
|
||||
if not tool_calls:
|
||||
raise RuntimeError("finish_reason was tool_calls but no tool_calls present")
|
||||
|
||||
messages.append({
|
||||
"role": "assistant",
|
||||
"content": getattr(choice.message, "content", None),
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
}
|
||||
for tc in tool_calls
|
||||
],
|
||||
})
|
||||
|
||||
for tc in tool_calls:
|
||||
tc_name = getattr(getattr(tc, "function", None), "name", "?")
|
||||
await self._maybe_call_progress(progress, f"🛠 Calling tool: {tc_name}")
|
||||
raw_arguments = getattr(getattr(tc, "function", None), "arguments", "{}") or "{}"
|
||||
try:
|
||||
signature_arguments = json.dumps(json.loads(raw_arguments), sort_keys=True)
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
signature_arguments = str(raw_arguments)
|
||||
signature = f"{tc_name}:{signature_arguments}"
|
||||
repeated_calls[signature] = repeated_calls.get(signature, 0) + 1
|
||||
if repeated_calls[signature] > 2:
|
||||
raise RuntimeError(
|
||||
f"Agentic loop repeated the identical {tc_name!r} tool call more than twice"
|
||||
)
|
||||
if repeated_calls[signature] > 1:
|
||||
result_text = (
|
||||
f"The identical {tc_name!r} call has already been performed. "
|
||||
"Do not call it again; provide the final answer now."
|
||||
)
|
||||
await self._maybe_call_progress(progress, f"🛠 Suppressed duplicate tool call: {tc_name}")
|
||||
else:
|
||||
result_text = await self._execute_tool_call(
|
||||
tc, ctx, allowed_roots, progress
|
||||
)
|
||||
messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.id,
|
||||
"content": result_text,
|
||||
})
|
||||
|
||||
raise RuntimeError("Agentic loop terminated without a response")
|
||||
|
||||
def _extract_response_text(self, response: Any) -> str:
|
||||
choices = getattr(response, "choices", None) or []
|
||||
if not choices:
|
||||
return ""
|
||||
choice = choices[0]
|
||||
message = getattr(choice, "message", None)
|
||||
if message is None:
|
||||
return ""
|
||||
content = getattr(message, "content", None)
|
||||
if isinstance(content, str):
|
||||
return self._strip_think(content)
|
||||
if isinstance(content, list):
|
||||
parts = []
|
||||
for part in content:
|
||||
text = getattr(part, "text", None)
|
||||
if isinstance(text, str):
|
||||
parts.append(text)
|
||||
return self._strip_think("\n".join(parts))
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def _strip_think(text: str) -> str:
|
||||
import re
|
||||
return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
|
||||
517
app/ai/tools.py
517
app/ai/tools.py
@@ -1,517 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
|
||||
def _resolve_sandboxed_path(raw_path: str, allowed_roots: tuple[Path, ...]) -> Path:
|
||||
"""
|
||||
allowed_roots[0] — base directory used for resolving raw_path.
|
||||
allowed_roots[1:] — permitted subdirectories; resolved path must fall inside one of them.
|
||||
When only one root is given it acts as both base and permitted directory.
|
||||
"""
|
||||
if raw_path.startswith("/"):
|
||||
raise ValueError(f"Absolute paths are not allowed: {raw_path!r}")
|
||||
base = allowed_roots[0]
|
||||
allowed_dirs = allowed_roots[1:] if len(allowed_roots) > 1 else allowed_roots
|
||||
candidate = (base / raw_path).resolve()
|
||||
for allowed in allowed_dirs:
|
||||
resolved = allowed.resolve()
|
||||
if candidate == resolved or str(candidate).startswith(str(resolved) + "/"):
|
||||
return candidate
|
||||
raise ValueError(
|
||||
f"Path {raw_path!r} is outside the allowed directories "
|
||||
f"({', '.join(a.name for a in allowed_dirs)})"
|
||||
)
|
||||
|
||||
|
||||
async def tool_write_file(arguments: dict, allowed_roots: tuple[Path, ...]) -> str:
|
||||
raw_path = str(arguments.get("path", "")).strip()
|
||||
content = str(arguments.get("content", ""))
|
||||
if not raw_path:
|
||||
raise ValueError("write_file: 'path' is required")
|
||||
try:
|
||||
path = _resolve_sandboxed_path(raw_path, allowed_roots)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(content, encoding="utf-8")
|
||||
return f"Written {len(content.encode('utf-8'))} bytes to {raw_path}"
|
||||
except ValueError:
|
||||
raise
|
||||
except OSError as exc:
|
||||
raise ValueError(f"write_file failed: {exc.strerror}") from exc
|
||||
|
||||
|
||||
async def tool_read_file(arguments: dict, allowed_roots: tuple[Path, ...]) -> str:
|
||||
raw_path = str(arguments.get("path", "")).strip()
|
||||
if not raw_path:
|
||||
raise ValueError("read_file: 'path' is required")
|
||||
try:
|
||||
path = _resolve_sandboxed_path(raw_path, allowed_roots)
|
||||
if not path.exists():
|
||||
raise ValueError(f"read_file: file not found: {raw_path!r}")
|
||||
if not path.is_file():
|
||||
raise ValueError(f"read_file: not a file: {raw_path!r}")
|
||||
size = path.stat().st_size
|
||||
if size > 100 * 1024:
|
||||
raise ValueError(f"read_file: file exceeds the 100 KB read limit ({size} bytes)")
|
||||
try:
|
||||
return path.read_text(encoding="utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return "File is binary (not readable as text)"
|
||||
except ValueError:
|
||||
raise
|
||||
except OSError as exc:
|
||||
raise ValueError(f"read_file failed: {exc.strerror}") from exc
|
||||
|
||||
|
||||
async def tool_create_file(arguments: dict, allowed_roots: tuple[Path, ...]) -> str:
|
||||
raw_path = str(arguments.get("path", "")).strip()
|
||||
content = str(arguments.get("content", ""))
|
||||
if not raw_path:
|
||||
raise ValueError("create_file: 'path' is required")
|
||||
try:
|
||||
path = _resolve_sandboxed_path(raw_path, allowed_roots)
|
||||
if path.exists():
|
||||
raise ValueError(f"create_file: file already exists: {raw_path!r}")
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(content, encoding="utf-8")
|
||||
return f"Created {raw_path}"
|
||||
except ValueError:
|
||||
raise
|
||||
except OSError as exc:
|
||||
raise ValueError(f"create_file failed: {exc.strerror}") from exc
|
||||
|
||||
|
||||
async def tool_modify_file(arguments: dict, allowed_roots: tuple[Path, ...]) -> str:
|
||||
raw_path = str(arguments.get("path", "")).strip()
|
||||
operation = str(arguments.get("operation", "")).strip().lower()
|
||||
if not raw_path:
|
||||
raise ValueError("modify_file: 'path' is required")
|
||||
if operation not in ("append", "replace"):
|
||||
raise ValueError("modify_file: 'operation' must be 'append' or 'replace'")
|
||||
try:
|
||||
path = _resolve_sandboxed_path(raw_path, allowed_roots)
|
||||
if operation == "append":
|
||||
text = str(arguments.get("text", ""))
|
||||
existing = path.read_text(encoding="utf-8") if path.exists() else ""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(existing + text, encoding="utf-8")
|
||||
return f"Appended {len(text.encode('utf-8'))} bytes to {raw_path}"
|
||||
else:
|
||||
old_text = arguments.get("old_text")
|
||||
new_text = arguments.get("new_text")
|
||||
if old_text is None or new_text is None:
|
||||
raise ValueError("modify_file replace: 'old_text' and 'new_text' are required")
|
||||
old_text = str(old_text)
|
||||
new_text = str(new_text)
|
||||
if not path.exists():
|
||||
raise ValueError(f"modify_file: file not found: {raw_path!r}")
|
||||
content = path.read_text(encoding="utf-8")
|
||||
if old_text not in content:
|
||||
raise ValueError(f"modify_file replace: 'old_text' not found in {raw_path!r}")
|
||||
path.write_text(content.replace(old_text, new_text, 1), encoding="utf-8")
|
||||
return f"Replaced text in {raw_path}"
|
||||
except ValueError:
|
||||
raise
|
||||
except OSError as exc:
|
||||
raise ValueError(f"modify_file failed: {exc.strerror}") from exc
|
||||
|
||||
|
||||
async def tool_delete_file(arguments: dict, allowed_roots: tuple[Path, ...]) -> str:
|
||||
raw_path = str(arguments.get("path", "")).strip()
|
||||
if not raw_path:
|
||||
raise ValueError("delete_file: 'path' is required")
|
||||
try:
|
||||
path = _resolve_sandboxed_path(raw_path, allowed_roots)
|
||||
if not path.exists():
|
||||
raise ValueError(f"delete_file: file not found: {raw_path!r}")
|
||||
path.unlink()
|
||||
return f"Deleted {raw_path}"
|
||||
except ValueError:
|
||||
raise
|
||||
except OSError as exc:
|
||||
raise ValueError(f"delete_file failed: {exc.strerror}") from exc
|
||||
|
||||
|
||||
async def tool_list_files(arguments: dict, allowed_roots: tuple[Path, ...]) -> str:
|
||||
raw_path = str(arguments.get("path", "")).strip()
|
||||
try:
|
||||
if raw_path:
|
||||
base = _resolve_sandboxed_path(raw_path, allowed_roots)
|
||||
else:
|
||||
# list all allowed dirs (skip the base root at index 0)
|
||||
dirs_to_list = allowed_roots[1:] if len(allowed_roots) > 1 else allowed_roots
|
||||
entries: list[str] = []
|
||||
for root in dirs_to_list:
|
||||
if not root.exists():
|
||||
continue
|
||||
resolved_root = root.resolve()
|
||||
entries.extend(_collect_files(root, resolved_root, max_depth=3))
|
||||
return "\n".join(sorted(entries)) or "(empty)"
|
||||
|
||||
if not base.exists():
|
||||
return "(empty)"
|
||||
if not base.is_dir():
|
||||
raise ValueError(f"list_files: not a directory: {raw_path!r}")
|
||||
resolved_base = base.resolve()
|
||||
entries = _collect_files(base, resolved_base, max_depth=3)
|
||||
return "\n".join(sorted(entries)) or "(empty)"
|
||||
except ValueError:
|
||||
raise
|
||||
except OSError as exc:
|
||||
raise ValueError(f"list_files failed: {exc.strerror}") from exc
|
||||
|
||||
|
||||
def _collect_files(directory: Path, resolved_root: Path, *, max_depth: int) -> list[str]:
|
||||
result: list[str] = []
|
||||
_walk(directory, resolved_root, directory, max_depth, 0, result)
|
||||
return result
|
||||
|
||||
|
||||
def _walk(
|
||||
current: Path,
|
||||
resolved_root: Path,
|
||||
base: Path,
|
||||
max_depth: int,
|
||||
depth: int,
|
||||
result: list[str],
|
||||
) -> None:
|
||||
if depth > max_depth:
|
||||
return
|
||||
try:
|
||||
for entry in current.iterdir():
|
||||
resolved_entry = entry.resolve()
|
||||
# skip symlinks that escape the root
|
||||
if not str(resolved_entry).startswith(str(resolved_root)):
|
||||
continue
|
||||
if entry.is_file():
|
||||
result.append(str(entry.relative_to(base)))
|
||||
elif entry.is_dir():
|
||||
_walk(entry, resolved_root, base, max_depth, depth + 1, result)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
PYTHON_EXEC_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "python_exec",
|
||||
"description": (
|
||||
"Execute Python code and return stdout/stderr output. "
|
||||
"A 'ctx' dict is pre-injected with message context: "
|
||||
"chat_id, message_id, message_text, args, from_user, replied_message. "
|
||||
"Use print() to produce output."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"code": {
|
||||
"type": "string",
|
||||
"description": "Python code to execute",
|
||||
},
|
||||
"timeout": {
|
||||
"type": "number",
|
||||
"description": "Execution timeout in seconds (default 10, max 30)",
|
||||
},
|
||||
},
|
||||
"required": ["code"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
TEMPORARY_COMMAND_CREATE_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "temporary_command_create",
|
||||
"description": (
|
||||
"Create a simple owner-only Telegram command that replies with fixed text. "
|
||||
"It exists only in memory and is removed when the bot restarts. Use this for "
|
||||
"test or short-lived commands instead of writing a persistent command module."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string", "description": "Command name without the leading dot"},
|
||||
"response": {"type": "string", "description": "Text returned when the command is used"},
|
||||
"description": {"type": "string", "description": "Optional help text for the command"},
|
||||
},
|
||||
"required": ["name", "response"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
BACKGROUND_START_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "background_start",
|
||||
"description": "Start an owner-approved long-running process without a shell. The process survives bot restarts; output is written to workspace/.process-logs/.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string", "description": "Stable process name"},
|
||||
"command": {"type": "array", "items": {"type": "string"}, "description": "Executable argv, e.g. [\"python\", \"server.py\"]"},
|
||||
"cwd": {"type": "string", "description": "Optional directory relative to workspace"},
|
||||
},
|
||||
"required": ["name", "command"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
BACKGROUND_STOP_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "background_stop",
|
||||
"description": "Stop a background process previously started by background_start.",
|
||||
"parameters": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]},
|
||||
},
|
||||
}
|
||||
|
||||
BACKGROUND_LIST_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "background_list",
|
||||
"description": "List tracked background processes and their logs.",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
|
||||
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"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
PAGE_EXTRACT_SCHEMA: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "extract_page",
|
||||
"description": (
|
||||
"Extract readable text from a web page URL. Use render='always' for dynamic "
|
||||
"JavaScript-heavy pages, render='never' for simple static pages, or render='auto' by default."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"url": {
|
||||
"type": "string",
|
||||
"description": "HTTP or HTTPS page URL to read.",
|
||||
},
|
||||
"render": {
|
||||
"type": "string",
|
||||
"enum": ["auto", "never", "always"],
|
||||
"description": "Whether to use browser rendering for dynamic pages.",
|
||||
},
|
||||
"max_chars": {
|
||||
"type": "integer",
|
||||
"description": "Maximum extracted characters to return, 1000-50000.",
|
||||
"minimum": 1000,
|
||||
"maximum": 50000,
|
||||
},
|
||||
},
|
||||
"required": ["url"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ToolExecutor = Callable[[dict, tuple[Path, ...]], Awaitable[str]]
|
||||
|
||||
FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = {
|
||||
"write_file": tool_write_file,
|
||||
"read_file": tool_read_file,
|
||||
"create_file": tool_create_file,
|
||||
"modify_file": tool_modify_file,
|
||||
"delete_file": tool_delete_file,
|
||||
"list_files": tool_list_files,
|
||||
}
|
||||
|
||||
FILE_TOOL_SCHEMAS: list[dict[str, Any]] = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "send_file",
|
||||
"description": (
|
||||
"Send a file from workspace/ or prompts/ to the current Telegram chat as a document. "
|
||||
"Use this when the user asks to receive, download, or share a file."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Relative file path"},
|
||||
"caption": {"type": "string", "description": "Optional caption for the file"},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "write_file",
|
||||
"description": (
|
||||
"Create or overwrite a file. Allowed path prefixes: workspace/ or prompts/. "
|
||||
"Example path: workspace/notes.txt or prompts/ai/parts/extra.md"
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Relative file path"},
|
||||
"content": {"type": "string", "description": "File content to write"},
|
||||
},
|
||||
"required": ["path", "content"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"description": "Read the content of a file (max 100 KB). Allowed path prefixes: workspace/ or prompts/.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Relative file path"},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "create_file",
|
||||
"description": (
|
||||
"Create a new file. Fails if the file already exists. "
|
||||
"Allowed path prefixes: workspace/ or prompts/."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Relative file path"},
|
||||
"content": {"type": "string", "description": "Initial file content"},
|
||||
},
|
||||
"required": ["path", "content"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "modify_file",
|
||||
"description": (
|
||||
"Modify an existing file. Two operations: "
|
||||
"'append' adds text at end (requires 'text'); "
|
||||
"'replace' substitutes first occurrence of old_text with new_text "
|
||||
"(requires 'old_text' and 'new_text', fails if old_text not found). "
|
||||
"Allowed path prefixes: workspace/ or prompts/."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Relative file path"},
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["append", "replace"],
|
||||
"description": "Operation to perform",
|
||||
},
|
||||
"text": {
|
||||
"type": "string",
|
||||
"description": "Text to append (used when operation=append)",
|
||||
},
|
||||
"old_text": {
|
||||
"type": "string",
|
||||
"description": "Text to find (used when operation=replace)",
|
||||
},
|
||||
"new_text": {
|
||||
"type": "string",
|
||||
"description": "Replacement text (used when operation=replace)",
|
||||
},
|
||||
},
|
||||
"required": ["path", "operation"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "delete_file",
|
||||
"description": "Delete a file. Allowed path prefixes: workspace/ or prompts/.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "Relative file path"},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "list_files",
|
||||
"description": (
|
||||
"List files recursively (up to 3 levels deep). "
|
||||
"Pass an empty path or omit to list all allowed directories. "
|
||||
"Allowed path prefixes: workspace/ or prompts/."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Directory path to list (relative). Leave empty to list all.",
|
||||
"default": "",
|
||||
},
|
||||
},
|
||||
"required": [],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
@@ -1,63 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import logging.handlers
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _json_value(value: Any) -> Any:
|
||||
"""Convert SDK objects to values suitable for a JSONL audit trail."""
|
||||
if value is None or isinstance(value, (str, int, float, bool)):
|
||||
return value
|
||||
if isinstance(value, Path):
|
||||
return str(value)
|
||||
if isinstance(value, dict):
|
||||
return {str(key): _json_value(item) for key, item in value.items()}
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [_json_value(item) for item in value]
|
||||
model_dump = getattr(value, "model_dump", None)
|
||||
if callable(model_dump):
|
||||
return _json_value(model_dump(exclude_none=True))
|
||||
return str(value)
|
||||
|
||||
|
||||
class AITranscriptLogger:
|
||||
"""Writes complete provider requests and responses to a dedicated rotating JSONL file."""
|
||||
|
||||
def __init__(self, workspace_dir: Path) -> None:
|
||||
self.path = workspace_dir / "ai-transcript.jsonl"
|
||||
self._logger = logging.getLogger(f"{__name__}.{id(self)}")
|
||||
self._logger.setLevel(logging.INFO)
|
||||
self._logger.propagate = False
|
||||
handler = logging.handlers.RotatingFileHandler(
|
||||
self.path,
|
||||
maxBytes=10 * 1024 * 1024,
|
||||
backupCount=3,
|
||||
encoding="utf-8",
|
||||
delay=True,
|
||||
)
|
||||
handler.setFormatter(logging.Formatter("%(message)s"))
|
||||
self._logger.addHandler(handler)
|
||||
|
||||
def record(self, event: str, **payload: Any) -> None:
|
||||
entry = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"event": event,
|
||||
**payload,
|
||||
}
|
||||
try:
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._logger.info(json.dumps(_json_value(entry), ensure_ascii=False, separators=(",", ":")))
|
||||
except OSError:
|
||||
log.exception("Failed to write AI transcript to %s", self.path)
|
||||
|
||||
def close(self) -> None:
|
||||
for handler in list(self._logger.handlers):
|
||||
handler.close()
|
||||
self._logger.removeHandler(handler)
|
||||
@@ -1,369 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
import logging
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.config import AIWebSearchConfig
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
SAFESEARCH_VALUES = {"off", "moderate", "strict"}
|
||||
COUNTRY_VALUES = {
|
||||
"AR", "AU", "AT", "BE", "BR", "CA", "CL", "DK", "FI", "FR", "DE", "GR", "HK",
|
||||
"IN", "ID", "IT", "JP", "KR", "MY", "MX", "NL", "NZ", "NO", "CN", "PL", "PT",
|
||||
"PH", "RU", "SA", "ZA", "ES", "SE", "CH", "TW", "TR", "GB", "US", "ALL",
|
||||
}
|
||||
UI_LANG_VALUES = {
|
||||
"es-AR", "en-AU", "de-AT", "nl-BE", "fr-BE", "pt-BR", "en-CA", "fr-CA",
|
||||
"es-CL", "da-DK", "fi-FI", "fr-FR", "de-DE", "el-GR", "zh-HK", "en-IN",
|
||||
"en-ID", "it-IT", "ja-JP", "ko-KR", "en-MY", "es-MX", "nl-NL", "en-NZ",
|
||||
"no-NO", "zh-CN", "pl-PL", "en-PH", "ru-RU", "en-ZA", "es-ES", "sv-SE",
|
||||
"fr-CH", "de-CH", "zh-TW", "tr-TR", "en-GB", "en-US", "es-US",
|
||||
}
|
||||
LANG_TO_UI_LANG = {
|
||||
"da": "da-DK",
|
||||
"de": "de-DE",
|
||||
"el": "el-GR",
|
||||
"en": "en-US",
|
||||
"es": "es-ES",
|
||||
"fi": "fi-FI",
|
||||
"fr": "fr-FR",
|
||||
"it": "it-IT",
|
||||
"ja": "ja-JP",
|
||||
"jp": "ja-JP",
|
||||
"ko": "ko-KR",
|
||||
"nl": "nl-NL",
|
||||
"no": "no-NO",
|
||||
"pl": "pl-PL",
|
||||
"pt": "pt-BR",
|
||||
"pt-br": "pt-BR",
|
||||
"ru": "ru-RU",
|
||||
"sv": "sv-SE",
|
||||
"tr": "tr-TR",
|
||||
"zh-hans": "zh-CN",
|
||||
"zh-hant": "zh-TW",
|
||||
}
|
||||
FRESHNESS_RE = re.compile(r"^(pd|pw|pm|py|\d{4}-\d{2}-\d{2}to\d{4}-\d{2}-\d{2})$")
|
||||
MAX_LOG_BODY_CHARS = 1000
|
||||
|
||||
|
||||
@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_country(value: Any, *, default: str) -> str:
|
||||
country = str(value or default).strip().upper() or default
|
||||
if country in COUNTRY_VALUES:
|
||||
return country
|
||||
log.info("Unsupported Brave search country %r; using ALL", country)
|
||||
return "ALL"
|
||||
|
||||
|
||||
def normalize_ui_lang(value: Any, *, search_lang: str, default: str) -> str:
|
||||
ui_lang = str(value or default).strip() or default
|
||||
if ui_lang in UI_LANG_VALUES:
|
||||
return ui_lang
|
||||
|
||||
normalized = ui_lang.replace("_", "-")
|
||||
language = normalized.split("-", 1)[0].lower()
|
||||
fallback = LANG_TO_UI_LANG.get(language) or LANG_TO_UI_LANG.get(search_lang) or "en-US"
|
||||
if fallback not in UI_LANG_VALUES:
|
||||
fallback = "en-US"
|
||||
log.info("Unsupported Brave search ui_lang %r; using %s", ui_lang, fallback)
|
||||
return fallback
|
||||
|
||||
|
||||
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:
|
||||
search_lang = str(arguments.get("search_lang") or config.search_lang).strip().lower()
|
||||
search_lang = search_lang or config.search_lang
|
||||
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=normalize_country(arguments.get("country"), default=config.country),
|
||||
search_lang=search_lang,
|
||||
ui_lang=normalize_ui_lang(
|
||||
arguments.get("ui_lang"),
|
||||
search_lang=search_lang,
|
||||
default=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,
|
||||
*,
|
||||
transport: httpx.AsyncBaseTransport | None = None,
|
||||
) -> None:
|
||||
self.config = config
|
||||
self.transport = transport
|
||||
|
||||
@property
|
||||
def available(self) -> bool:
|
||||
return self.config.available
|
||||
|
||||
def _build_headers(self) -> dict[str, str]:
|
||||
api_key = str(self.config.api_key or "").strip()
|
||||
if api_key.lower().startswith("bearer "):
|
||||
api_key = api_key[7:].strip()
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"Brave web search is not configured. Set BRAVE_API_KEY or BRAVE_SEARCH_API_KEY."
|
||||
)
|
||||
return {
|
||||
"Accept": "application/json",
|
||||
"Accept-Encoding": "gzip",
|
||||
"X-Subscription-Token": api_key,
|
||||
}
|
||||
|
||||
async def search(self, arguments: dict[str, Any]) -> WebSearchResponse:
|
||||
if not self.config.enabled:
|
||||
raise ValueError("Brave web search is disabled")
|
||||
headers = self._build_headers()
|
||||
|
||||
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
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=self.config.timeout,
|
||||
transport=self.transport,
|
||||
) 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:
|
||||
_log_error_response(response)
|
||||
raise ValueError(_format_error_response(response))
|
||||
|
||||
try:
|
||||
payload = response.json()
|
||||
except ValueError as exc:
|
||||
_log_error_response(response)
|
||||
raise ValueError("Brave web search returned invalid JSON") from exc
|
||||
if not isinstance(payload, dict):
|
||||
_log_response_summary(response, query.query, results_count=0)
|
||||
raise ValueError("Brave web search returned an unexpected response")
|
||||
|
||||
parsed = parse_response(query.query, payload)
|
||||
_log_response_summary(response, query.query, results_count=len(parsed.results))
|
||||
return parsed
|
||||
|
||||
|
||||
def _response_log_body(response: httpx.Response) -> str:
|
||||
body = response.text.replace("\n", "\\n").replace("\r", "\\r")
|
||||
if len(body) > MAX_LOG_BODY_CHARS:
|
||||
return body[:MAX_LOG_BODY_CHARS] + "...(truncated)"
|
||||
return body
|
||||
|
||||
|
||||
def _response_url(response: httpx.Response) -> str:
|
||||
try:
|
||||
return str(response.request.url)
|
||||
except RuntimeError:
|
||||
return str(response.url)
|
||||
|
||||
|
||||
def _log_response_summary(
|
||||
response: httpx.Response,
|
||||
query: str,
|
||||
*,
|
||||
results_count: int,
|
||||
) -> None:
|
||||
log.info(
|
||||
"Brave web search response status=%s query=%r results=%s url=%s",
|
||||
response.status_code,
|
||||
query,
|
||||
results_count,
|
||||
_response_url(response),
|
||||
)
|
||||
|
||||
|
||||
def _log_error_response(response: httpx.Response) -> None:
|
||||
log.warning(
|
||||
"Brave web search error status=%s url=%s body=%s",
|
||||
response.status_code,
|
||||
_response_url(response),
|
||||
_response_log_body(response),
|
||||
)
|
||||
|
||||
|
||||
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}"
|
||||
@@ -7,15 +7,12 @@ from app.commands.core import (
|
||||
CommandContext,
|
||||
CommandDefinition,
|
||||
CommandRegistry,
|
||||
combine_checks,
|
||||
owner_only_check,
|
||||
require_args_check,
|
||||
require_reply_check,
|
||||
shared_access_check,
|
||||
)
|
||||
from app.commands.dispatcher import CommandDispatcher
|
||||
from app.commands.persistent import load_persistent_commands
|
||||
from app.commands.temporary import TemporaryCommandStore
|
||||
|
||||
__all__ = [
|
||||
"CheckResult",
|
||||
@@ -28,11 +25,8 @@ __all__ = [
|
||||
"DEFAULT_REGISTRY",
|
||||
"RuntimeAclStore",
|
||||
"command_registry",
|
||||
"combine_checks",
|
||||
"owner_only_check",
|
||||
"require_args_check",
|
||||
"require_reply_check",
|
||||
"shared_access_check",
|
||||
"load_persistent_commands",
|
||||
"TemporaryCommandStore",
|
||||
]
|
||||
|
||||
@@ -1,59 +1,72 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
import signal
|
||||
from textwrap import wrap
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pyrogram import enums
|
||||
from pyrogram.enums import MessageEntityType
|
||||
from pyrogram.types import MessageEntity
|
||||
|
||||
from app.ai import AIInvocation, parse_ai_invocation
|
||||
from app.ai.web_search import format_results
|
||||
from app.agentic import parse_agent_invocation, run_agent_turn
|
||||
from app.commands.core import CommandRegistry, normalize_command_name
|
||||
from app.transcription import transcribe_file
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.commands.core import CommandContext
|
||||
|
||||
|
||||
_PH = "\x00{}\x00"
|
||||
_PH_RE = re.compile(r"\x00(\d+)\x00")
|
||||
|
||||
|
||||
def _md_to_tg(text: str) -> str:
|
||||
"""Convert markdown-ish text to Telegram HTML."""
|
||||
saved: list[str] = []
|
||||
|
||||
def _save(s: str) -> str:
|
||||
saved.append(s)
|
||||
return _PH.format(len(saved) - 1)
|
||||
|
||||
def _fence(m: re.Match) -> str:
|
||||
lang = (m.group(1) or "").strip()
|
||||
code = html.escape(m.group(2) or "")
|
||||
cls = f' class="language-{html.escape(lang)}"' if lang else ""
|
||||
return _save(f"<pre><code{cls}>{code}</code></pre>")
|
||||
|
||||
text = re.sub(r"```(\w*)\n?(.*?)```", _fence, text, flags=re.DOTALL)
|
||||
text = re.sub(r"`([^`\n]+)`", lambda m: _save(f"<code>{html.escape(m.group(1))}</code>"), text)
|
||||
|
||||
text = html.escape(text)
|
||||
|
||||
text = re.sub(r"^#{1,6}\s+(.+)$", r"<b>\1</b>", text, flags=re.MULTILINE)
|
||||
text = re.sub(r"\*\*(.+?)\*\*", r"<b>\1</b>", text, flags=re.DOTALL)
|
||||
text = re.sub(r"__(.+?)__", r"<b>\1</b>", text, flags=re.DOTALL)
|
||||
text = re.sub(r"(?<!\*)\*(?!\*)(.+?)(?<!\*)\*(?!\*)", r"<i>\1</i>", text, flags=re.DOTALL)
|
||||
text = re.sub(r"(?<!_)_(?!_)(.+?)(?<!_)_(?!_)", r"<i>\1</i>", text, flags=re.DOTALL)
|
||||
text = re.sub(r"~~(.+?)~~", r"<s>\1</s>", text, flags=re.DOTALL)
|
||||
text = re.sub(r"\[(.+?)\]\((.+?)\)", r'<a href="\2">\1</a>', text)
|
||||
|
||||
return _PH_RE.sub(lambda m: saved[int(m.group(1))], text)
|
||||
|
||||
|
||||
command_registry = CommandRegistry()
|
||||
DEFAULT_REGISTRY = command_registry
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _utf16_length(value: str) -> int:
|
||||
return len(value.encode("utf-16-le")) // 2
|
||||
|
||||
|
||||
def _split_transcript(value: str, *, limit: int = 4000) -> list[str]:
|
||||
chunks: list[str] = []
|
||||
remaining = value.strip()
|
||||
while remaining:
|
||||
units = 0
|
||||
cut = len(remaining)
|
||||
for index, character in enumerate(remaining):
|
||||
units += _utf16_length(character)
|
||||
if units > limit:
|
||||
cut = index
|
||||
break
|
||||
|
||||
if cut == len(remaining):
|
||||
chunks.append(remaining)
|
||||
break
|
||||
|
||||
whitespace = max(remaining.rfind("\n", 0, cut), remaining.rfind(" ", 0, cut))
|
||||
if whitespace > cut // 2:
|
||||
cut = whitespace
|
||||
|
||||
chunk = remaining[:cut].strip()
|
||||
if chunk:
|
||||
chunks.append(chunk)
|
||||
remaining = remaining[cut:].strip()
|
||||
return chunks
|
||||
|
||||
|
||||
async def _dislike(message: Any) -> None:
|
||||
try:
|
||||
await message.react("👎")
|
||||
except Exception:
|
||||
log.exception("Could not add a dislike reaction to message %s", getattr(message, "id", None))
|
||||
|
||||
|
||||
def _transcribable_media(message: Any) -> Any | None:
|
||||
for field in ("voice", "video_note", "video"):
|
||||
media = getattr(message, field, None)
|
||||
if media is not None:
|
||||
return media
|
||||
return None
|
||||
|
||||
|
||||
def _user_label(user: Any | None) -> str:
|
||||
@@ -186,60 +199,6 @@ def _format_id_report(ctx: "CommandContext") -> str:
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _chunk_text(text: str, limit: int = 3900) -> list[str]:
|
||||
if len(text) <= limit:
|
||||
return [text]
|
||||
chunks: list[str] = []
|
||||
current: list[str] = []
|
||||
current_len = 0
|
||||
for paragraph in text.split("\n"):
|
||||
paragraph = paragraph.rstrip()
|
||||
piece = paragraph if not current else "\n" + paragraph
|
||||
if current_len + len(piece) <= limit:
|
||||
current.append(piece.lstrip("\n") if not current else piece)
|
||||
current_len += len(piece)
|
||||
continue
|
||||
if current:
|
||||
chunks.append("".join(current))
|
||||
current = []
|
||||
current_len = 0
|
||||
if len(paragraph) <= limit:
|
||||
current = [paragraph]
|
||||
current_len = len(paragraph)
|
||||
continue
|
||||
for line in wrap(paragraph, width=limit):
|
||||
if len(line) <= limit:
|
||||
chunks.append(line)
|
||||
else:
|
||||
chunks.extend([line[i : i + limit] for i in range(0, len(line), limit)])
|
||||
if current:
|
||||
chunks.append("".join(current))
|
||||
return [chunk for chunk in chunks if chunk]
|
||||
|
||||
|
||||
async def _reply_long(ctx: "CommandContext", text: str) -> None:
|
||||
chunks = _chunk_text(text)
|
||||
if not chunks:
|
||||
await ctx.reply("")
|
||||
return
|
||||
html_chunks = [f"<blockquote expandable>{_md_to_tg(c)}</blockquote>" for c in chunks]
|
||||
first = await ctx.reply(html_chunks[0], parse_mode=enums.ParseMode.HTML)
|
||||
for chunk in html_chunks[1:]:
|
||||
await ctx.client.send_message(
|
||||
ctx.message.chat.id,
|
||||
chunk,
|
||||
parse_mode=enums.ParseMode.HTML,
|
||||
reply_to_message_id=getattr(first, "id", None),
|
||||
)
|
||||
|
||||
|
||||
def _format_ai_help(ctx: "CommandContext") -> str:
|
||||
ai_service = ctx.manager.ai_service
|
||||
if ai_service is None:
|
||||
return ""
|
||||
return "AI models: " + ai_service.describe_models()
|
||||
|
||||
|
||||
def _format_permissions(ctx: "CommandContext") -> str:
|
||||
store = ctx.acl_store
|
||||
target = ctx.replied_user
|
||||
@@ -372,9 +331,6 @@ async def status(ctx: "CommandContext") -> None:
|
||||
f"allow_rules={stats['allow_rules']}, "
|
||||
f"deny_rules={stats['deny_rules']}"
|
||||
)
|
||||
ai_service = ctx.manager.ai_service
|
||||
if ai_service is not None:
|
||||
lines.append("AI: " + ai_service.describe_models())
|
||||
await ctx.reply("\n".join(lines))
|
||||
|
||||
|
||||
@@ -406,12 +362,7 @@ async def id_command(ctx: "CommandContext") -> None:
|
||||
category="system",
|
||||
)
|
||||
async def help_command(ctx: "CommandContext") -> None:
|
||||
lines = [ctx.manager.command_registry.render_help(ctx)]
|
||||
ai_help = _format_ai_help(ctx)
|
||||
if ai_help:
|
||||
lines.append("")
|
||||
lines.append(ai_help)
|
||||
await ctx.reply("\n".join(lines))
|
||||
await ctx.reply(ctx.manager.command_registry.render_help(ctx))
|
||||
|
||||
|
||||
@command_registry.command(
|
||||
@@ -458,115 +409,66 @@ 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",
|
||||
usage=".ai / ..ai [N] [--model] prompt",
|
||||
category="ai",
|
||||
"Run an agent with optional Telegram history and media",
|
||||
usage=".ai / ..ai [messages] [.ds] [.th] <prompt>",
|
||||
category="shared",
|
||||
).shared()
|
||||
async def ai(ctx: "CommandContext") -> None:
|
||||
ai_service = ctx.manager.ai_service
|
||||
if ai_service is None or ctx.manager.config.ai is None:
|
||||
await ctx.reply("AI is not configured")
|
||||
return
|
||||
|
||||
try:
|
||||
invocation = parse_ai_invocation(
|
||||
ctx.args,
|
||||
default_context_messages=ctx.manager.config.ai.context_messages,
|
||||
invocation = parse_agent_invocation(ctx.args, ctx.manager.config.ai)
|
||||
turn = await run_agent_turn(
|
||||
ctx.client,
|
||||
ctx.message,
|
||||
ctx.account,
|
||||
ctx.account.ai_sessions,
|
||||
invocation,
|
||||
ctx.manager.config.ai,
|
||||
)
|
||||
except ValueError as exc:
|
||||
await ctx.reply(str(exc))
|
||||
for chunk in _split_transcript(turn.answer):
|
||||
sent = await ctx.reply(chunk)
|
||||
ctx.account.ai_sessions.register_response(turn.handle, sent)
|
||||
except Exception:
|
||||
log.exception("AI command failed for message %s", getattr(ctx.message, "id", None))
|
||||
await _dislike(ctx.message)
|
||||
|
||||
|
||||
@command_registry.command(
|
||||
"txt",
|
||||
"Transcribe a replied voice message, video note, or video",
|
||||
usage=".txt / ..txt (reply to media)",
|
||||
category="shared",
|
||||
).shared()
|
||||
async def txt(ctx: "CommandContext") -> None:
|
||||
reply = ctx.replied_message
|
||||
media = _transcribable_media(reply) if reply is not None else None
|
||||
if ctx.args.strip() or reply is None or media is None:
|
||||
await _dislike(ctx.message)
|
||||
return
|
||||
|
||||
try:
|
||||
selected_alias = ai_service.select_model_alias(invocation.model_flags)
|
||||
except ValueError as exc:
|
||||
await ctx.reply(str(exc))
|
||||
return
|
||||
with tempfile.TemporaryDirectory(prefix="multi-bot-txt-") as temp_dir:
|
||||
downloaded = await ctx.client.download_media(
|
||||
reply,
|
||||
file_name=f"{temp_dir}{os.sep}",
|
||||
)
|
||||
if not isinstance(downloaded, str):
|
||||
raise RuntimeError("Telegram media download returned no file")
|
||||
transcript = await transcribe_file(
|
||||
Path(downloaded),
|
||||
ctx.manager.config.elevenlabs,
|
||||
content_type=getattr(media, "mime_type", None),
|
||||
)
|
||||
|
||||
async def _react(emoji: str | None) -> None:
|
||||
try:
|
||||
await ctx.client.send_reaction(ctx.message.chat.id, ctx.message.id, emoji=emoji)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
await _react("🤔")
|
||||
|
||||
try:
|
||||
response = await ai_service.generate(
|
||||
ctx,
|
||||
AIInvocation(
|
||||
context_messages=invocation.context_messages,
|
||||
model_flags=[selected_alias],
|
||||
prompt=invocation.prompt,
|
||||
),
|
||||
)
|
||||
except (RuntimeError, ValueError) as exc:
|
||||
await _react("👎")
|
||||
await ctx.reply(str(exc))
|
||||
return
|
||||
|
||||
await _react("👍")
|
||||
await _reply_long(ctx, response.text)
|
||||
for chunk in _split_transcript(transcript):
|
||||
entity = MessageEntity(
|
||||
type=MessageEntityType.BLOCKQUOTE,
|
||||
offset=0,
|
||||
length=_utf16_length(chunk),
|
||||
expandable=True,
|
||||
)
|
||||
await ctx.reply(chunk, entities=[entity])
|
||||
except Exception:
|
||||
log.exception("Failed to transcribe media for message %s", getattr(ctx.message, "id", None))
|
||||
await _dislike(ctx.message)
|
||||
|
||||
@@ -106,18 +106,6 @@ def require_args_check() -> CommandCheck:
|
||||
return CommandCheck("require_args", _check, "Require non-empty arguments")
|
||||
|
||||
|
||||
def combine_checks(*checks: CommandCheck) -> CommandCheck:
|
||||
async def _check(ctx: "CommandContext") -> CheckResult:
|
||||
for check in checks:
|
||||
result = await check.run(ctx)
|
||||
if not result.allowed:
|
||||
return result
|
||||
return CheckResult.allow()
|
||||
|
||||
names = "+".join(check.name for check in checks) or "combined"
|
||||
return CommandCheck(names, _check, "Combined checks")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CommandContext:
|
||||
client: "Client"
|
||||
@@ -208,9 +196,6 @@ class CommandDefinition:
|
||||
return self.owner_prefix
|
||||
return self.owner_prefix if is_owner else self.guest_prefix
|
||||
|
||||
def accepts_prefix(self, prefix: str, *, is_owner: bool) -> bool:
|
||||
return prefix == self.expected_prefix(is_owner)
|
||||
|
||||
def help_line(self, ctx: CommandContext | None = None) -> str:
|
||||
usage = self.usage or self.display_name
|
||||
details: list[str] = []
|
||||
@@ -342,9 +327,6 @@ class CommandRegistry:
|
||||
canonical = self._aliases.get(normalized, normalized)
|
||||
return self._commands.get(canonical)
|
||||
|
||||
def has_name(self, token: str) -> bool:
|
||||
return self.resolve(token) is not None
|
||||
|
||||
def iter_commands(self) -> list[CommandDefinition]:
|
||||
return [self._commands[name] for name in self._order]
|
||||
|
||||
|
||||
@@ -1,32 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from app.commands.core import CommandRegistry
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def load_persistent_commands(commands_dir: Path, registry: CommandRegistry) -> list[str]:
|
||||
"""Load command modules from workspace/commands on every bot startup."""
|
||||
commands_dir.mkdir(parents=True, exist_ok=True)
|
||||
loaded: list[str] = []
|
||||
for path in sorted(commands_dir.glob("*.py")):
|
||||
if path.name.startswith("_"):
|
||||
continue
|
||||
module_name = f"bot_persistent_command_{path.stem}"
|
||||
try:
|
||||
spec = importlib.util.spec_from_file_location(module_name, path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError("cannot create import spec")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
loaded.append(path.name)
|
||||
except Exception:
|
||||
sys.modules.pop(module_name, None)
|
||||
log.exception("Failed to load persistent command module %s", path)
|
||||
return loaded
|
||||
@@ -1,47 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from app.commands.core import CommandRegistry
|
||||
|
||||
|
||||
_COMMAND_NAME_RE = re.compile(r"[a-z][a-z0-9_-]{0,31}")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class TemporaryCommandStore:
|
||||
"""Registers simple commands that exist only for this bot process."""
|
||||
|
||||
registry: CommandRegistry
|
||||
names: set[str] = field(default_factory=set)
|
||||
|
||||
def create(self, name: str, response: str, description: str | None = None) -> str:
|
||||
normalized_name = name.strip().lower()
|
||||
if not _COMMAND_NAME_RE.fullmatch(normalized_name):
|
||||
raise ValueError(
|
||||
"temporary_command_create: name must start with a letter and contain only "
|
||||
"letters, digits, '-' or '_' (up to 32 characters)"
|
||||
)
|
||||
if self.registry.has_name(normalized_name):
|
||||
raise ValueError(f"temporary_command_create: command .{normalized_name} already exists")
|
||||
|
||||
text = response.strip()
|
||||
if not text:
|
||||
raise ValueError("temporary_command_create: response must not be empty")
|
||||
if len(text) > 3900:
|
||||
raise ValueError("temporary_command_create: response must be at most 3900 characters")
|
||||
|
||||
command_description = (description or "Temporary AI-created command").strip()
|
||||
|
||||
async def handler(ctx: object) -> None:
|
||||
await ctx.reply(text) # type: ignore[attr-defined]
|
||||
|
||||
self.registry.command(
|
||||
normalized_name,
|
||||
command_description,
|
||||
usage=f".{normalized_name}",
|
||||
category="temporary",
|
||||
)(handler)
|
||||
self.names.add(normalized_name)
|
||||
return f"Created temporary owner-only command .{normalized_name}; it will be removed on bot restart"
|
||||
555
app/config.py
555
app/config.py
@@ -1,15 +1,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import tomllib
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
except ModuleNotFoundError: # pragma: no cover - optional dependency fallback
|
||||
@@ -36,71 +33,6 @@ class AccountConfig:
|
||||
hide_password: bool = False
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AIProviderConfig:
|
||||
name: str
|
||||
api_key: str
|
||||
base_url: str
|
||||
timeout: float | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AIModelConfig:
|
||||
alias: str
|
||||
provider: str
|
||||
model: str
|
||||
description: str | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AIPromptConfig:
|
||||
root: Path
|
||||
system_file: str = "system.md"
|
||||
parts_dir: str = "parts"
|
||||
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 AIPageExtractConfig:
|
||||
enabled: bool = True
|
||||
timeout: float = 15.0
|
||||
max_chars: int = 12000
|
||||
render_dynamic: bool = True
|
||||
dynamic_wait_ms: int = 1500
|
||||
block_private_hosts: bool = True
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AIConfig:
|
||||
default_model: str
|
||||
context_messages: int
|
||||
prompt: AIPromptConfig
|
||||
providers: dict[str, AIProviderConfig] = field(default_factory=dict)
|
||||
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)
|
||||
page_extract: AIPageExtractConfig = field(default_factory=AIPageExtractConfig)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class GitHubWebhookConfig:
|
||||
enabled: bool = False
|
||||
@@ -109,47 +41,58 @@ class GitHubWebhookConfig:
|
||||
secret: str | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ElevenLabsConfig:
|
||||
api_key: str | None = None
|
||||
model_id: str = "scribe_v2"
|
||||
language_code: str | None = None
|
||||
request_timeout_seconds: float = 300.0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AiProviderConfig:
|
||||
api_url: str
|
||||
api_key: str | None = None
|
||||
model: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AiConfig:
|
||||
glm: AiProviderConfig = field(
|
||||
default_factory=lambda: AiProviderConfig(
|
||||
api_url="https://api.z.ai/api/paas/v4",
|
||||
model="glm-5.3",
|
||||
)
|
||||
)
|
||||
deepseek: AiProviderConfig = field(
|
||||
default_factory=lambda: AiProviderConfig(
|
||||
api_url="https://api.deepseek.com",
|
||||
model="deepseek-v4-flash",
|
||||
)
|
||||
)
|
||||
prompt_dir: Path = Path("prompts")
|
||||
base_dir: Path = Path(".")
|
||||
max_context_messages: int = 100
|
||||
max_media_bytes: int = 5 * 1024 * 1024
|
||||
max_tool_rounds: int = 16
|
||||
brave_api_key: str | None = None
|
||||
brave_timeout_seconds: float = 10.0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AppConfig:
|
||||
api_id: int | None
|
||||
api_hash: str | None
|
||||
session_dir: Path
|
||||
accounts: list[AccountConfig]
|
||||
ai: AIConfig | None = None
|
||||
github_webhook: GitHubWebhookConfig = field(default_factory=GitHubWebhookConfig)
|
||||
elevenlabs: ElevenLabsConfig = field(default_factory=ElevenLabsConfig)
|
||||
ai: AiConfig = field(default_factory=AiConfig)
|
||||
|
||||
|
||||
def _read_json_file(path: Path) -> dict[str, Any]:
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def _write_json_file(path: Path, data: dict[str, Any]) -> None:
|
||||
serialized = json.dumps(data, ensure_ascii=False, indent=2)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
"w",
|
||||
encoding="utf-8",
|
||||
delete=False,
|
||||
dir=path.parent,
|
||||
prefix=path.name,
|
||||
suffix=".tmp",
|
||||
) as handle:
|
||||
handle.write(serialized)
|
||||
temp_path = Path(handle.name)
|
||||
temp_path.replace(path)
|
||||
|
||||
|
||||
def _merge_example_defaults(target: dict[str, Any], source: dict[str, Any], path: str = "") -> list[str]:
|
||||
"""Add scalar keys present in source but absent in target; recurse into existing dicts."""
|
||||
added: list[str] = []
|
||||
for key, value in source.items():
|
||||
full_path = f"{path}.{key}" if path else key
|
||||
if key not in target:
|
||||
if not isinstance(value, (dict, list)):
|
||||
target[key] = value
|
||||
added.append(full_path)
|
||||
elif isinstance(target[key], dict) and isinstance(value, dict):
|
||||
added.extend(_merge_example_defaults(target[key], value, full_path))
|
||||
return added
|
||||
def _read_toml_file(path: Path) -> dict[str, Any]:
|
||||
with path.open("rb") as handle:
|
||||
return tomllib.load(handle)
|
||||
|
||||
|
||||
def _parse_int(value: Any, *, field_name: str) -> int:
|
||||
@@ -159,13 +102,6 @@ def _parse_int(value: Any, *, field_name: str) -> int:
|
||||
raise ValueError(f"{field_name} must be an integer") from exc
|
||||
|
||||
|
||||
def _parse_float(value: Any, *, field_name: str) -> float:
|
||||
try:
|
||||
return float(value)
|
||||
except (TypeError, ValueError) as 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
|
||||
@@ -180,6 +116,13 @@ def _parse_bool(value: Any, *, field_name: str) -> bool:
|
||||
raise ValueError(f"{field_name} must be a boolean")
|
||||
|
||||
|
||||
def _parse_float(value: Any, *, field_name: str) -> float:
|
||||
try:
|
||||
return float(value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"{field_name} must be a number") from exc
|
||||
|
||||
|
||||
def _parse_mapping(value: Any, *, field_name: str) -> dict[str, Any]:
|
||||
if value is None:
|
||||
return {}
|
||||
@@ -188,125 +131,6 @@ def _parse_mapping(value: Any, *, field_name: str) -> dict[str, Any]:
|
||||
return value
|
||||
|
||||
|
||||
def _normalize_brave_api_key(value: Any) -> str | None:
|
||||
api_key = str(value or "").strip()
|
||||
if api_key.lower().startswith("bearer "):
|
||||
api_key = api_key[7:].strip()
|
||||
return api_key or None
|
||||
|
||||
|
||||
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 = _normalize_brave_api_key(raw.get("api_key"))
|
||||
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_search_api_key = os.environ.get("BRAVE_SEARCH_API_KEY")
|
||||
if env_search_api_key is not None:
|
||||
api_key = _normalize_brave_api_key(env_search_api_key)
|
||||
env_api_key = os.environ.get("BRAVE_API_KEY")
|
||||
if env_api_key is not None:
|
||||
api_key = _normalize_brave_api_key(env_api_key)
|
||||
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_page_extract_config(raw_ai: dict[str, Any]) -> AIPageExtractConfig:
|
||||
raw = _parse_mapping(raw_ai.get("page_extract"), field_name="ai.page_extract")
|
||||
|
||||
enabled = _parse_bool(raw.get("enabled", True), field_name="ai.page_extract.enabled")
|
||||
timeout = _parse_float(raw.get("timeout", 15.0), field_name="ai.page_extract.timeout")
|
||||
if timeout <= 0:
|
||||
raise ValueError("ai.page_extract.timeout must be > 0")
|
||||
|
||||
max_chars = _parse_int(raw.get("max_chars", 12000), field_name="ai.page_extract.max_chars")
|
||||
if max_chars < 1000:
|
||||
raise ValueError("ai.page_extract.max_chars must be >= 1000")
|
||||
|
||||
render_dynamic = _parse_bool(
|
||||
raw.get("render_dynamic", True),
|
||||
field_name="ai.page_extract.render_dynamic",
|
||||
)
|
||||
dynamic_wait_ms = _parse_int(
|
||||
raw.get("dynamic_wait_ms", 1500),
|
||||
field_name="ai.page_extract.dynamic_wait_ms",
|
||||
)
|
||||
dynamic_wait_ms = min(max(dynamic_wait_ms, 0), 10000)
|
||||
block_private_hosts = _parse_bool(
|
||||
raw.get("block_private_hosts", True),
|
||||
field_name="ai.page_extract.block_private_hosts",
|
||||
)
|
||||
|
||||
if os.environ.get("PAGE_EXTRACT_MAX_CHARS") is not None:
|
||||
max_chars = _parse_int(os.environ["PAGE_EXTRACT_MAX_CHARS"], field_name="PAGE_EXTRACT_MAX_CHARS")
|
||||
if max_chars < 1000:
|
||||
raise ValueError("PAGE_EXTRACT_MAX_CHARS must be >= 1000")
|
||||
if os.environ.get("PAGE_EXTRACT_RENDER_DYNAMIC") is not None:
|
||||
render_dynamic = _parse_bool(
|
||||
os.environ["PAGE_EXTRACT_RENDER_DYNAMIC"],
|
||||
field_name="PAGE_EXTRACT_RENDER_DYNAMIC",
|
||||
)
|
||||
|
||||
return AIPageExtractConfig(
|
||||
enabled=enabled,
|
||||
timeout=timeout,
|
||||
max_chars=max_chars,
|
||||
render_dynamic=render_dynamic,
|
||||
dynamic_wait_ms=dynamic_wait_ms,
|
||||
block_private_hosts=block_private_hosts,
|
||||
)
|
||||
|
||||
|
||||
def _parse_github_webhook_config(raw: Any) -> GitHubWebhookConfig:
|
||||
raw_config = _parse_mapping(raw, field_name="github_webhook")
|
||||
|
||||
@@ -353,6 +177,153 @@ def _parse_github_webhook_config(raw: Any) -> GitHubWebhookConfig:
|
||||
)
|
||||
|
||||
|
||||
def _parse_elevenlabs_config(raw: Any) -> ElevenLabsConfig:
|
||||
raw_config = _parse_mapping(raw, field_name="elevenlabs")
|
||||
|
||||
raw_api_key = raw_config.get("api_key")
|
||||
api_key = str(raw_api_key).strip() if raw_api_key is not None else None
|
||||
api_key = api_key or None
|
||||
model_id = str(raw_config.get("model_id", "scribe_v2") or "scribe_v2").strip()
|
||||
raw_language = raw_config.get("language_code")
|
||||
language_code = str(raw_language).strip() if raw_language is not None else None
|
||||
language_code = language_code or None
|
||||
request_timeout_seconds = _parse_float(
|
||||
raw_config.get("request_timeout_seconds", 300),
|
||||
field_name="elevenlabs.request_timeout_seconds",
|
||||
)
|
||||
|
||||
if os.environ.get("ELEVENLABS_API_KEY") is not None:
|
||||
api_key = os.environ["ELEVENLABS_API_KEY"].strip() or None
|
||||
if os.environ.get("ELEVENLABS_MODEL_ID") is not None:
|
||||
model_id = os.environ["ELEVENLABS_MODEL_ID"].strip()
|
||||
if os.environ.get("ELEVENLABS_LANGUAGE_CODE") is not None:
|
||||
language_code = os.environ["ELEVENLABS_LANGUAGE_CODE"].strip() or None
|
||||
if os.environ.get("ELEVENLABS_REQUEST_TIMEOUT_SECONDS") is not None:
|
||||
request_timeout_seconds = _parse_float(
|
||||
os.environ["ELEVENLABS_REQUEST_TIMEOUT_SECONDS"],
|
||||
field_name="ELEVENLABS_REQUEST_TIMEOUT_SECONDS",
|
||||
)
|
||||
|
||||
if not model_id:
|
||||
raise ValueError("elevenlabs.model_id is required")
|
||||
if request_timeout_seconds <= 0:
|
||||
raise ValueError("elevenlabs.request_timeout_seconds must be greater than zero")
|
||||
|
||||
return ElevenLabsConfig(
|
||||
api_key=api_key,
|
||||
model_id=model_id,
|
||||
language_code=language_code,
|
||||
request_timeout_seconds=request_timeout_seconds,
|
||||
)
|
||||
|
||||
|
||||
def _parse_ai_config(raw: Any) -> AiConfig:
|
||||
raw_config = _parse_mapping(raw, field_name="ai")
|
||||
|
||||
def parse_provider(
|
||||
name: str,
|
||||
*,
|
||||
default_url: str,
|
||||
default_model: str,
|
||||
env_prefix: str,
|
||||
) -> AiProviderConfig:
|
||||
provider = _parse_mapping(raw_config.get(name), field_name=f"ai.{name}")
|
||||
api_url = str(provider.get("api_url", default_url)).strip()
|
||||
raw_api_key = provider.get("api_key")
|
||||
api_key = str(raw_api_key).strip() if raw_api_key is not None else None
|
||||
api_key = api_key or None
|
||||
model = str(provider.get("model", default_model)).strip()
|
||||
|
||||
if os.environ.get(f"{env_prefix}_API_URL") is not None:
|
||||
api_url = os.environ[f"{env_prefix}_API_URL"].strip()
|
||||
if os.environ.get(f"{env_prefix}_API_KEY") is not None:
|
||||
api_key = os.environ[f"{env_prefix}_API_KEY"].strip() or None
|
||||
if os.environ.get(f"{env_prefix}_MODEL") is not None:
|
||||
model = os.environ[f"{env_prefix}_MODEL"].strip()
|
||||
|
||||
if not api_url:
|
||||
raise ValueError(f"ai.{name}.api_url is required")
|
||||
if not model:
|
||||
raise ValueError(f"ai.{name}.model is required")
|
||||
return AiProviderConfig(api_url=api_url, api_key=api_key, model=model)
|
||||
|
||||
glm = parse_provider(
|
||||
"glm",
|
||||
default_url="https://api.z.ai/api/paas/v4",
|
||||
default_model="glm-5.3",
|
||||
env_prefix="GLM",
|
||||
)
|
||||
deepseek = parse_provider(
|
||||
"deepseek",
|
||||
default_url="https://api.deepseek.com",
|
||||
default_model="deepseek-v4-flash",
|
||||
env_prefix="DEEPSEEK",
|
||||
)
|
||||
max_context_messages = _parse_int(
|
||||
raw_config.get("max_context_messages", 100),
|
||||
field_name="ai.max_context_messages",
|
||||
)
|
||||
max_media_bytes = _parse_int(
|
||||
raw_config.get("max_media_bytes", 5 * 1024 * 1024),
|
||||
field_name="ai.max_media_bytes",
|
||||
)
|
||||
max_tool_rounds = _parse_int(
|
||||
raw_config.get("max_tool_rounds", 16),
|
||||
field_name="ai.max_tool_rounds",
|
||||
)
|
||||
brave_timeout_seconds = _parse_float(
|
||||
raw_config.get("brave_timeout_seconds", 10),
|
||||
field_name="ai.brave_timeout_seconds",
|
||||
)
|
||||
|
||||
raw_brave_key = raw_config.get("brave_api_key")
|
||||
brave_api_key = str(raw_brave_key).strip() if raw_brave_key is not None else None
|
||||
brave_api_key = brave_api_key or None
|
||||
|
||||
if os.environ.get("BRAVE_SEARCH_API_KEY") is not None:
|
||||
brave_api_key = os.environ["BRAVE_SEARCH_API_KEY"].strip() or None
|
||||
|
||||
if os.environ.get("AI_MAX_CONTEXT_MESSAGES") is not None:
|
||||
max_context_messages = _parse_int(
|
||||
os.environ["AI_MAX_CONTEXT_MESSAGES"],
|
||||
field_name="AI_MAX_CONTEXT_MESSAGES",
|
||||
)
|
||||
if os.environ.get("AI_MAX_MEDIA_BYTES") is not None:
|
||||
max_media_bytes = _parse_int(
|
||||
os.environ["AI_MAX_MEDIA_BYTES"],
|
||||
field_name="AI_MAX_MEDIA_BYTES",
|
||||
)
|
||||
if os.environ.get("AI_MAX_TOOL_ROUNDS") is not None:
|
||||
max_tool_rounds = _parse_int(
|
||||
os.environ["AI_MAX_TOOL_ROUNDS"],
|
||||
field_name="AI_MAX_TOOL_ROUNDS",
|
||||
)
|
||||
if os.environ.get("BRAVE_SEARCH_TIMEOUT_SECONDS") is not None:
|
||||
brave_timeout_seconds = _parse_float(
|
||||
os.environ["BRAVE_SEARCH_TIMEOUT_SECONDS"],
|
||||
field_name="BRAVE_SEARCH_TIMEOUT_SECONDS",
|
||||
)
|
||||
|
||||
if max_context_messages < 0:
|
||||
raise ValueError("ai.max_context_messages must not be negative")
|
||||
if max_media_bytes < 1:
|
||||
raise ValueError("ai.max_media_bytes must be greater than zero")
|
||||
if max_tool_rounds < 1:
|
||||
raise ValueError("ai.max_tool_rounds must be greater than zero")
|
||||
if brave_timeout_seconds <= 0:
|
||||
raise ValueError("ai.brave_timeout_seconds must be greater than zero")
|
||||
|
||||
return AiConfig(
|
||||
glm=glm,
|
||||
deepseek=deepseek,
|
||||
max_context_messages=max_context_messages,
|
||||
max_media_bytes=max_media_bytes,
|
||||
max_tool_rounds=max_tool_rounds,
|
||||
brave_api_key=brave_api_key,
|
||||
brave_timeout_seconds=brave_timeout_seconds,
|
||||
)
|
||||
|
||||
|
||||
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 []
|
||||
@@ -418,109 +389,6 @@ def _parse_accounts(raw_accounts: Any, *, default_api_id: int | None, default_ap
|
||||
return accounts
|
||||
|
||||
|
||||
def _parse_ai_config(raw_ai: Any, *, root: Path) -> AIConfig | None:
|
||||
if raw_ai is None:
|
||||
return None
|
||||
if not isinstance(raw_ai, dict):
|
||||
raise ValueError("ai must be an object")
|
||||
|
||||
prompt_raw = _parse_mapping(raw_ai.get("prompt"), field_name="ai.prompt")
|
||||
prompt_root = Path(str(prompt_raw.get("root", "prompts/ai"))).expanduser()
|
||||
if not prompt_root.is_absolute():
|
||||
prompt_root = (root / prompt_root).resolve()
|
||||
raw_user_template = prompt_raw.get("user_template")
|
||||
user_template = str(raw_user_template).strip() if raw_user_template else None
|
||||
prompt = AIPromptConfig(
|
||||
root=prompt_root,
|
||||
system_file=str(prompt_raw.get("system_file", "system.md")),
|
||||
parts_dir=str(prompt_raw.get("parts_dir", "parts")),
|
||||
user_template=user_template or None,
|
||||
)
|
||||
|
||||
default_model = str(raw_ai.get("default_model", "")).strip()
|
||||
if not default_model:
|
||||
raise ValueError("ai.default_model is required")
|
||||
|
||||
context_messages = _parse_int(raw_ai.get("context_messages", 12), field_name="ai.context_messages")
|
||||
if context_messages < 0:
|
||||
raise ValueError("ai.context_messages must be >= 0")
|
||||
|
||||
providers_raw = _parse_mapping(raw_ai.get("providers"), field_name="ai.providers")
|
||||
models_raw = _parse_mapping(raw_ai.get("models"), field_name="ai.models")
|
||||
|
||||
providers: dict[str, AIProviderConfig] = {}
|
||||
for provider_name, provider_config in providers_raw.items():
|
||||
if not isinstance(provider_config, dict):
|
||||
raise ValueError(f"ai.providers.{provider_name} must be an object")
|
||||
api_key = str(provider_config.get("api_key", "")).strip()
|
||||
base_url = str(provider_config.get("base_url", "")).strip()
|
||||
if not api_key:
|
||||
raise ValueError(f"ai.providers.{provider_name}.api_key is required")
|
||||
if not base_url:
|
||||
raise ValueError(f"ai.providers.{provider_name}.base_url is required")
|
||||
timeout_value = provider_config.get("timeout")
|
||||
timeout = None
|
||||
if timeout_value is not None:
|
||||
timeout = _parse_float(timeout_value, field_name=f"ai.providers.{provider_name}.timeout")
|
||||
providers[provider_name] = AIProviderConfig(
|
||||
name=provider_name,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
models: dict[str, AIModelConfig] = {}
|
||||
for alias, model_config in models_raw.items():
|
||||
if not isinstance(model_config, dict):
|
||||
raise ValueError(f"ai.models.{alias} must be an object")
|
||||
provider = str(model_config.get("provider", "")).strip()
|
||||
model = str(model_config.get("model", "")).strip()
|
||||
if not provider:
|
||||
raise ValueError(f"ai.models.{alias}.provider is required")
|
||||
if not model:
|
||||
raise ValueError(f"ai.models.{alias}.model is required")
|
||||
description = model_config.get("description")
|
||||
models[alias] = AIModelConfig(
|
||||
alias=alias,
|
||||
provider=provider,
|
||||
model=model,
|
||||
description=str(description).strip() if description is not None else None,
|
||||
)
|
||||
|
||||
if default_model not in models:
|
||||
raise ValueError("ai.default_model must reference one of ai.models")
|
||||
|
||||
for alias, model_config in models.items():
|
||||
if model_config.provider not in providers:
|
||||
raise ValueError(
|
||||
f"ai.models.{alias}.provider references unknown provider: {model_config.provider}"
|
||||
)
|
||||
|
||||
workspace_dir_raw = raw_ai.get("workspace_dir", "workspace")
|
||||
workspace_dir = Path(str(workspace_dir_raw)).expanduser()
|
||||
if not workspace_dir.is_absolute():
|
||||
workspace_dir = (root / workspace_dir).resolve()
|
||||
|
||||
tool_max_rounds = _parse_int(raw_ai.get("tool_max_rounds", 10), field_name="ai.tool_max_rounds")
|
||||
if tool_max_rounds < 1:
|
||||
raise ValueError("ai.tool_max_rounds must be >= 1")
|
||||
|
||||
web_search = _parse_web_search_config(raw_ai)
|
||||
page_extract = _parse_page_extract_config(raw_ai)
|
||||
|
||||
return AIConfig(
|
||||
default_model=default_model,
|
||||
context_messages=context_messages,
|
||||
prompt=prompt,
|
||||
providers=providers,
|
||||
models=models,
|
||||
workspace_dir=workspace_dir,
|
||||
tool_max_rounds=tool_max_rounds,
|
||||
web_search=web_search,
|
||||
page_extract=page_extract,
|
||||
)
|
||||
|
||||
|
||||
def _accounts_from_env() -> list[dict[str, Any]] | None:
|
||||
raw_accounts = os.environ.get("ACCOUNTS_JSON")
|
||||
if raw_accounts:
|
||||
@@ -554,26 +422,15 @@ def load_app_config(root: Path) -> AppConfig:
|
||||
if not candidate.is_absolute():
|
||||
candidate = (root / candidate).resolve()
|
||||
else:
|
||||
candidate = root / "config.json"
|
||||
candidate = root / "config.toml"
|
||||
|
||||
raw: dict[str, Any] = {}
|
||||
if candidate.exists():
|
||||
raw = _read_json_file(candidate)
|
||||
example_path = root / "config.example.json"
|
||||
if example_path.exists():
|
||||
try:
|
||||
example_raw = _read_json_file(example_path)
|
||||
added = _merge_example_defaults(raw, example_raw)
|
||||
if added:
|
||||
_write_json_file(candidate, raw)
|
||||
log.info("Auto-updated %s with new fields: %s", candidate.name, ", ".join(added))
|
||||
except Exception as exc:
|
||||
log.warning("Could not auto-update config from example: %s", exc)
|
||||
raw = _read_toml_file(candidate)
|
||||
|
||||
env_api_id = os.environ.get("TG_API_ID")
|
||||
env_api_hash = os.environ.get("TG_API_HASH")
|
||||
env_session_dir = os.environ.get("SESSION_DIR")
|
||||
env_ai_config = os.environ.get("AI_CONFIG")
|
||||
|
||||
if env_api_id is not None:
|
||||
raw["api_id"] = env_api_id
|
||||
@@ -581,8 +438,6 @@ def load_app_config(root: Path) -> AppConfig:
|
||||
raw["api_hash"] = env_api_hash
|
||||
if env_session_dir is not None:
|
||||
raw["session_dir"] = env_session_dir
|
||||
if env_ai_config:
|
||||
raw["ai"] = json.loads(env_ai_config)
|
||||
|
||||
env_accounts = _accounts_from_env()
|
||||
if env_accounts is not None:
|
||||
@@ -590,7 +445,7 @@ def load_app_config(root: Path) -> AppConfig:
|
||||
|
||||
if "accounts" not in raw:
|
||||
raise ValueError(
|
||||
"Missing accounts configuration. Provide accounts in config.json or set ACCOUNTS_JSON / ACCOUNT_SESSIONS."
|
||||
"Missing accounts configuration. Provide accounts in config.toml or set ACCOUNTS_JSON / ACCOUNT_SESSIONS."
|
||||
)
|
||||
|
||||
api_id = _parse_int(raw["api_id"], field_name="api_id") if "api_id" in raw else None
|
||||
@@ -606,8 +461,11 @@ def load_app_config(root: Path) -> AppConfig:
|
||||
default_session_dir=session_dir,
|
||||
)
|
||||
|
||||
ai = _parse_ai_config(raw.get("ai"), root=root)
|
||||
github_webhook = _parse_github_webhook_config(raw.get("github_webhook"))
|
||||
elevenlabs = _parse_elevenlabs_config(raw.get("elevenlabs"))
|
||||
ai = _parse_ai_config(raw.get("ai"))
|
||||
ai.prompt_dir = root / "prompts"
|
||||
ai.base_dir = root
|
||||
|
||||
if not accounts:
|
||||
raise ValueError("At least one account must be configured")
|
||||
@@ -627,6 +485,7 @@ def load_app_config(root: Path) -> AppConfig:
|
||||
api_hash=api_hash,
|
||||
session_dir=session_dir,
|
||||
accounts=accounts,
|
||||
ai=ai,
|
||||
github_webhook=github_webhook,
|
||||
elevenlabs=elevenlabs,
|
||||
ai=ai,
|
||||
)
|
||||
|
||||
@@ -2,13 +2,12 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from pyrogram import Client, idle
|
||||
|
||||
from app.ai import AIService
|
||||
from app.commands import DEFAULT_REGISTRY, CommandRegistry, RuntimeAclStore, load_persistent_commands
|
||||
from app.agentic import AgentSessionStore
|
||||
from app.commands import DEFAULT_REGISTRY, CommandRegistry, RuntimeAclStore
|
||||
from app.config import AccountConfig, AppConfig
|
||||
from app.github_webhook import GitHubWebhookServer
|
||||
from app.handlers import register_common_handlers
|
||||
@@ -23,27 +22,19 @@ class AccountState:
|
||||
acl_store: RuntimeAclStore | None = None
|
||||
status: str = "pending"
|
||||
error: str | None = None
|
||||
started_at: datetime | None = None
|
||||
me: Any | None = None
|
||||
ai_sessions: AgentSessionStore = field(default_factory=AgentSessionStore)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AccountManager:
|
||||
config: AppConfig
|
||||
ai_service: AIService | None = field(init=False)
|
||||
command_registry: CommandRegistry = field(init=False)
|
||||
states: list[AccountState] = field(init=False)
|
||||
github_webhook_server: GitHubWebhookServer = field(init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.command_registry = DEFAULT_REGISTRY
|
||||
if self.config.ai is not None:
|
||||
loaded = load_persistent_commands(self.config.ai.workspace_dir / "commands", self.command_registry)
|
||||
if loaded:
|
||||
log.info("Loaded persistent commands: %s", ", ".join(loaded))
|
||||
self.ai_service = (
|
||||
AIService(self.config.ai, self.command_registry) if self.config.ai is not None else None
|
||||
)
|
||||
self.states = [AccountState(config=account) for account in self.config.accounts]
|
||||
self.github_webhook_server = GitHubWebhookServer(self, self.config.github_webhook)
|
||||
|
||||
@@ -92,7 +83,6 @@ class AccountManager:
|
||||
try:
|
||||
await client.start()
|
||||
state.me = await client.get_me()
|
||||
state.started_at = datetime.now(timezone.utc)
|
||||
state.status = "running"
|
||||
log.info("Started account %s (%s)", account.name, account.session)
|
||||
except Exception as exc:
|
||||
@@ -109,6 +99,7 @@ class AccountManager:
|
||||
|
||||
async def stop_all(self) -> None:
|
||||
for state in reversed(self.states):
|
||||
state.ai_sessions.clear()
|
||||
if state.client is None:
|
||||
continue
|
||||
if state.status not in {"running", "starting"}:
|
||||
@@ -133,8 +124,6 @@ class AccountManager:
|
||||
parts.append(
|
||||
f"acl={stats['allow_rules']} allow / {stats['deny_rules']} deny"
|
||||
)
|
||||
if self.ai_service is not None and self.config.ai is not None:
|
||||
parts.append(f"ai={self.config.ai.default_model}")
|
||||
if state.me is not None:
|
||||
display_name = (
|
||||
getattr(state.me, "full_name", None)
|
||||
@@ -157,12 +146,8 @@ class AccountManager:
|
||||
raise RuntimeError("No accounts were started successfully")
|
||||
|
||||
try:
|
||||
if self.ai_service is not None:
|
||||
await self.ai_service.initialize()
|
||||
await self.github_webhook_server.start()
|
||||
await idle()
|
||||
finally:
|
||||
await self.github_webhook_server.stop()
|
||||
await self.stop_all()
|
||||
if self.ai_service is not None:
|
||||
self.ai_service.close()
|
||||
|
||||
67
app/transcription.py
Normal file
67
app/transcription.py
Normal file
@@ -0,0 +1,67 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
|
||||
from app.config import ElevenLabsConfig
|
||||
|
||||
|
||||
ELEVENLABS_SPEECH_TO_TEXT_URL = "https://api.elevenlabs.io/v1/speech-to-text"
|
||||
|
||||
|
||||
class TranscriptionError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
async def transcribe_file(
|
||||
path: Path,
|
||||
config: ElevenLabsConfig,
|
||||
*,
|
||||
content_type: str | None = None,
|
||||
) -> str:
|
||||
if not config.api_key:
|
||||
raise TranscriptionError("ElevenLabs API key is not configured")
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("model_id", config.model_id)
|
||||
form.add_field("tag_audio_events", "false")
|
||||
if config.language_code:
|
||||
form.add_field("language_code", config.language_code)
|
||||
|
||||
timeout = aiohttp.ClientTimeout(total=config.request_timeout_seconds)
|
||||
headers = {"xi-api-key": config.api_key}
|
||||
|
||||
with path.open("rb") as media_file:
|
||||
form.add_field(
|
||||
"file",
|
||||
media_file,
|
||||
filename=path.name,
|
||||
content_type=content_type or "application/octet-stream",
|
||||
)
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.post(
|
||||
ELEVENLABS_SPEECH_TO_TEXT_URL,
|
||||
data=form,
|
||||
headers=headers,
|
||||
) as response:
|
||||
body = await response.text()
|
||||
if response.status >= 400:
|
||||
raise TranscriptionError(
|
||||
f"ElevenLabs returned HTTP {response.status}: {body[:500]}"
|
||||
)
|
||||
except (aiohttp.ClientError, TimeoutError) as exc:
|
||||
raise TranscriptionError("ElevenLabs request failed") from exc
|
||||
|
||||
try:
|
||||
payload: Any = json.loads(body)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise TranscriptionError("ElevenLabs returned invalid JSON") from exc
|
||||
|
||||
text = payload.get("text") if isinstance(payload, dict) else None
|
||||
if not isinstance(text, str) or not text.strip():
|
||||
raise TranscriptionError("ElevenLabs returned an empty transcript")
|
||||
return text.strip()
|
||||
Reference in New Issue
Block a user