a84467c385
OpenRouter is used when no provider flag is set. Claude is selected with .cl, and GLM stays available via .glm.
570 lines
18 KiB
Python
570 lines
18 KiB
Python
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
|
|
from app.llm import AnthropicMessagesLLM, ChatCompletionsLLM, OpenRouterChatLLM
|
|
from app.web_tools import WebPageTools
|
|
|
|
|
|
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
|
|
use_claude: bool = False
|
|
use_openrouter: 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(
|
|
f"Could not load {prompt_dir / 'MASTER.txt'} and {prompt_dir / '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
|
|
|
|
provider = "openrouter"
|
|
thinking = False
|
|
while index < len(tokens) and tokens[index].startswith("."):
|
|
flag = tokens[index].lower()
|
|
if flag == ".glm":
|
|
provider = "glm"
|
|
elif flag == ".ds":
|
|
provider = "deepseek"
|
|
elif flag in {".cl", ".claude"}:
|
|
provider = "claude"
|
|
elif flag in {".or", ".openrouter"}:
|
|
provider = "openrouter"
|
|
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=provider == "deepseek",
|
|
use_claude=provider == "claude",
|
|
use_openrouter=provider == "openrouter",
|
|
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": _provider_name(invocation),
|
|
"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_name(invocation: AgentInvocation) -> str:
|
|
if invocation.use_openrouter:
|
|
return "openrouter"
|
|
if invocation.use_claude:
|
|
return "claude"
|
|
if invocation.use_deepseek:
|
|
return "deepseek"
|
|
return "glm"
|
|
|
|
|
|
def _provider(config: AiConfig, invocation: AgentInvocation) -> AiProviderConfig:
|
|
if invocation.use_openrouter:
|
|
return config.openrouter
|
|
if invocation.use_claude:
|
|
return config.claude
|
|
if invocation.use_deepseek:
|
|
return config.deepseek
|
|
return config.glm
|
|
|
|
|
|
def _create_llm(provider: AiProviderConfig, invocation: AgentInvocation) -> Any:
|
|
if invocation.use_openrouter:
|
|
llm_class: type[Any] = OpenRouterChatLLM
|
|
elif invocation.use_claude:
|
|
llm_class = AnthropicMessagesLLM
|
|
elif invocation.use_deepseek:
|
|
llm_class = OpenAIBasedLLM
|
|
else:
|
|
llm_class = ChatCompletionsLLM
|
|
return llm_class(
|
|
api_url=provider.api_url,
|
|
api_key=provider.api_key,
|
|
model=provider.model,
|
|
)
|
|
|
|
|
|
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)
|
|
media = await _download_message_media(
|
|
client,
|
|
command_message,
|
|
Path(existing.workspace.name),
|
|
config.max_media_bytes,
|
|
)
|
|
return AgentTurnResult(
|
|
answer=await _run_session(existing.session, invocation.prompt, media=media),
|
|
handle=existing,
|
|
)
|
|
|
|
provider = _provider(config, invocation)
|
|
if not provider.api_key:
|
|
raise AgentError("The selected AI provider has no API key")
|
|
|
|
llm = _create_llm(provider, invocation)
|
|
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,
|
|
)
|
|
media = await _download_message_media(
|
|
client,
|
|
command_message,
|
|
temp_path,
|
|
config.max_media_bytes,
|
|
)
|
|
tools = [
|
|
*WebPageTools(
|
|
timeout_seconds=config.web_page_timeout_seconds,
|
|
max_chars=config.web_page_max_chars,
|
|
).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, media=media)
|
|
return AgentTurnResult(
|
|
answer=answer,
|
|
handle=AgentSessionHandle(session, workspace, set()),
|
|
)
|
|
except Exception:
|
|
workspace.cleanup()
|
|
raise
|
|
|
|
|
|
async def _run_session(session: Any, prompt: str, *, media: list[Media] | None = None) -> str:
|
|
if media:
|
|
result = await session.run(prompt, media=media)
|
|
else:
|
|
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
|