feat: rebuild AI commands and migrate config to TOML
This commit is contained in:
Azimkin
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user