diff --git a/assistant_bot/agent.py b/assistant_bot/agent.py index d76e74d..c9feefd 100644 --- a/assistant_bot/agent.py +++ b/assistant_bot/agent.py @@ -1,30 +1,29 @@ import json import logging import re -from datetime import datetime, timezone +from datetime import datetime from typing import Any from zoneinfo import ZoneInfo from telegram import Update from telegram.ext import ContextTypes +from .agent_tools import execute_agent_tool, render_agent_tool_catalog from .ai import AIClientError from .config import CONVERSATION_HISTORY_LIMIT, MAX_AGENT_STEPS -from .models import AgentDecision, ReminderParseResult +from .models import AgentDecision from .prompts import ASSISTANT_SYSTEM_PROMPT -from .reminders import parse_reminder from .storage import AssistantStorage from .telegram_utils import ( escape_markdown_text, get_ai_client, get_storage, get_tz, - parse_positive_int, reply_markdown, require_user_id, typing_action, ) -from .time_utils import format_local_dt, to_utc_iso, utc_now +from .time_utils import format_local_dt logger = logging.getLogger(__name__) @@ -92,25 +91,7 @@ def build_agent_tool_prompt(tz: ZoneInfo) -> str: "читаемость. Для короткого простого ответа разметка не обязательна. " "Не добавляй декоративные эмодзи чаще одного раза на сообщение.\n\n" "Доступные tools:\n" - "- get_current_datetime {} - узнать текущее локальное и UTC-время.\n" - '- remember {"text": string} - сохранить важный долгосрочный факт о пользователе.\n' - '- list_memory {"limit": number} - показать сохраненную память.\n' - '- delete_memory {"id": number} - удалить запись памяти.\n' - '- create_note {"text": string} - сохранить заметку.\n' - '- list_notes {"limit": number} - показать заметки.\n' - '- delete_note {"id": number} - удалить заметку.\n' - '- create_reminder {"when": string, "text": string} - поставить напоминание. ' - "when можно указывать как '30m', 'через 2 часа', '18:30', '2026-07-16 18:30'. " - "Если пользователь говорит 'завтра/послезавтра/через неделю', сам рассчитай " - "дату от текущего локального времени и передай 'YYYY-MM-DD HH:MM'.\n" - '- list_reminders {"limit": number} - показать активные напоминания.\n' - '- cancel_reminder {"id": number} - отменить напоминание.\n' - '- create_status {"title": string, "status": string} - начать отслеживать статус.\n' - '- list_statuses {"limit": number} - показать отслеживаемые статусы.\n' - '- update_status {"id": number, "status": string} - обновить статус.\n' - '- delete_status {"id": number} - удалить отслеживаемый объект.\n\n' - '- search_conversation {"query": string, "limit": number} - найти старое обсуждение ' - "во всей сохраненной переписке по содержательным ключевым словам.\n\n" + f"{render_agent_tool_catalog()}\n\n" "Правила:\n" "- Для просьб 'запомни', 'сохрани как факт', 'учти на будущее' используй remember.\n" "- Для заметок используй create_note, для напоминаний create_reminder, " @@ -205,185 +186,6 @@ def save_conversation_exchange( storage.add_conversation_exchange(user_id, chat_id, prompt, answer) -def coerce_positive_int(value: Any, default: int, maximum: int = 50) -> int: - try: - parsed = int(value) - except (TypeError, ValueError): - return default - if parsed <= 0: - return default - return min(parsed, maximum) - - -def coerce_required_text(arguments: dict[str, Any], key: str) -> str: - value = arguments.get(key, "") - return str(value).strip() - - -def parse_agent_reminder( - when: str, text: str, tz: ZoneInfo -) -> ReminderParseResult | None: - parsed = parse_reminder(f"{when} {text}".strip(), tz) - if parsed: - return parsed - - normalized_when = when.strip().replace("T", " ") - try: - parsed_dt = datetime.fromisoformat(normalized_when) - except ValueError: - return None - - if parsed_dt.tzinfo is None: - parsed_dt = parsed_dt.replace(tzinfo=tz) - return ReminderParseResult(parsed_dt.astimezone(timezone.utc), text) - - -def execute_agent_tool( - name: str, - arguments: dict[str, Any], - storage: AssistantStorage, - user_id: int, - chat_id: int, - tz: ZoneInfo, -) -> dict[str, Any]: - tool_name = name.strip().lower() - - if tool_name == "get_current_datetime": - now = utc_now() - return { - "ok": True, - "local": now.astimezone(tz).strftime("%Y-%m-%d %H:%M:%S"), - "utc": to_utc_iso(now), - "timezone": tz.key, - } - - if tool_name == "remember": - text = coerce_required_text(arguments, "text") - if not text: - return {"ok": False, "error": "text is required"} - memory_id = storage.add_memory(user_id, text) - return {"ok": True, "id": memory_id, "text": text} - - if tool_name == "list_memory": - rows = storage.list_memories( - user_id, coerce_positive_int(arguments.get("limit"), 20) - ) - return {"ok": True, "items": rows} - - if tool_name == "delete_memory": - memory_id = parse_positive_int(str(arguments.get("id", ""))) - if memory_id is None: - return {"ok": False, "error": "valid id is required"} - return {"ok": storage.delete_memory(user_id, memory_id), "id": memory_id} - - if tool_name == "create_note": - text = coerce_required_text(arguments, "text") - if not text: - return {"ok": False, "error": "text is required"} - note_id = storage.add_note(user_id, text) - return {"ok": True, "id": note_id, "text": text} - - if tool_name == "list_notes": - rows = storage.list_notes( - user_id, coerce_positive_int(arguments.get("limit"), 20) - ) - return {"ok": True, "items": rows} - - if tool_name == "delete_note": - note_id = parse_positive_int(str(arguments.get("id", ""))) - if note_id is None: - return {"ok": False, "error": "valid id is required"} - return {"ok": storage.delete_note(user_id, note_id), "id": note_id} - - if tool_name == "create_reminder": - when = coerce_required_text(arguments, "when") - text = coerce_required_text(arguments, "text") - if not when or not text: - return {"ok": False, "error": "when and text are required"} - parsed = parse_agent_reminder(when, text, tz) - if not parsed: - return {"ok": False, "error": "could not parse reminder time"} - if parsed.remind_at_utc <= utc_now(): - return {"ok": False, "error": "reminder time must be in the future"} - - reminder_id = storage.add_reminder( - user_id, chat_id, parsed.text, parsed.remind_at_utc - ) - return { - "ok": True, - "id": reminder_id, - "text": parsed.text, - "remind_at": format_local_dt(parsed.remind_at_utc, tz), - } - - if tool_name == "list_reminders": - rows = storage.list_reminders( - user_id, coerce_positive_int(arguments.get("limit"), 20) - ) - for row in rows: - row["remind_at_local"] = format_local_dt(row["remind_at"], tz) - return {"ok": True, "items": rows} - - if tool_name == "cancel_reminder": - reminder_id = parse_positive_int(str(arguments.get("id", ""))) - if reminder_id is None: - return {"ok": False, "error": "valid id is required"} - return {"ok": storage.cancel_reminder(user_id, reminder_id), "id": reminder_id} - - if tool_name == "create_status": - title = coerce_required_text(arguments, "title") - status = coerce_required_text(arguments, "status") or "open" - if not title: - return {"ok": False, "error": "title is required"} - item_id = storage.add_tracked_item(user_id, title, status) - return {"ok": True, "id": item_id, "title": title, "status": status} - - if tool_name == "list_statuses": - rows = storage.list_tracked_items( - user_id, coerce_positive_int(arguments.get("limit"), 30) - ) - return {"ok": True, "items": rows} - - if tool_name == "update_status": - item_id = parse_positive_int(str(arguments.get("id", ""))) - status = coerce_required_text(arguments, "status") - if item_id is None or not status: - return {"ok": False, "error": "valid id and status are required"} - return { - "ok": storage.set_tracked_status(user_id, item_id, status), - "id": item_id, - "status": status, - } - - if tool_name == "delete_status": - item_id = parse_positive_int(str(arguments.get("id", ""))) - if item_id is None: - return {"ok": False, "error": "valid id is required"} - return {"ok": storage.delete_tracked_item(user_id, item_id), "id": item_id} - - if tool_name == "search_conversation": - query = coerce_required_text(arguments, "query") - if not query: - return {"ok": False, "error": "query is required"} - limit = coerce_positive_int(arguments.get("limit"), 5, maximum=10) - hits = storage.search_conversation_messages(user_id, chat_id, query, limit) - discussions = [ - { - "matched_message_id": hit["id"], - "score": hit["score"], - "messages": storage.conversation_message_window( - user_id, - chat_id, - int(hit["id"]), - ), - } - for hit in hits - ] - return {"ok": True, "query": query, "discussions": discussions} - - return {"ok": False, "error": f"unknown tool: {name}"} - - async def run_agent_prompt( update: Update, context: ContextTypes.DEFAULT_TYPE, diff --git a/assistant_bot/agent_tools.py b/assistant_bot/agent_tools.py new file mode 100644 index 0000000..49370c6 --- /dev/null +++ b/assistant_bot/agent_tools.py @@ -0,0 +1,417 @@ +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any, Callable +from zoneinfo import ZoneInfo + +from .models import ReminderParseResult +from .reminders import parse_reminder +from .storage import AssistantStorage +from .telegram_utils import parse_positive_int +from .time_utils import format_local_dt, to_utc_iso, utc_now + + +ToolResult = dict[str, Any] +ToolExecutor = Callable[[dict[str, Any], "ToolContext"], ToolResult] + + +@dataclass(frozen=True) +class ToolContext: + storage: AssistantStorage + user_id: int + chat_id: int + tz: ZoneInfo + + +@dataclass(frozen=True) +class AgentTool: + name: str + usage: str + description: str + execute: ToolExecutor + paragraph_before: bool = False + + def prompt_line(self) -> str: + suffix = f" {self.usage}" if self.usage else " {}" + prefix = "\n" if self.paragraph_before else "" + return f"{prefix}- {self.name}{suffix} - {self.description}" + + +def coerce_positive_int(value: Any, default: int, maximum: int = 50) -> int: + try: + parsed = int(value) + except (TypeError, ValueError): + return default + if parsed <= 0: + return default + return min(parsed, maximum) + + +def coerce_required_text(arguments: dict[str, Any], key: str) -> str: + value = arguments.get(key, "") + return str(value).strip() + + +def parse_agent_reminder( + when: str, + text: str, + tz: ZoneInfo, +) -> ReminderParseResult | None: + parsed = parse_reminder(f"{when} {text}".strip(), tz) + if parsed: + return parsed + + normalized_when = when.strip().replace("T", " ") + try: + parsed_dt = datetime.fromisoformat(normalized_when) + except ValueError: + return None + + if parsed_dt.tzinfo is None: + parsed_dt = parsed_dt.replace(tzinfo=tz) + return ReminderParseResult(parsed_dt.astimezone(timezone.utc), text) + + +def get_current_datetime( + _arguments: dict[str, Any], + context: ToolContext, +) -> ToolResult: + now = utc_now() + return { + "ok": True, + "local": now.astimezone(context.tz).strftime("%Y-%m-%d %H:%M:%S"), + "utc": to_utc_iso(now), + "timezone": context.tz.key, + } + + +def remember(arguments: dict[str, Any], context: ToolContext) -> ToolResult: + text = coerce_required_text(arguments, "text") + if not text: + return {"ok": False, "error": "text is required"} + memory_id = context.storage.add_memory(context.user_id, text) + return {"ok": True, "id": memory_id, "text": text} + + +def list_memory(arguments: dict[str, Any], context: ToolContext) -> ToolResult: + rows = context.storage.list_memories( + context.user_id, + coerce_positive_int(arguments.get("limit"), 20), + ) + return {"ok": True, "items": rows} + + +def delete_memory(arguments: dict[str, Any], context: ToolContext) -> ToolResult: + memory_id = parse_positive_int(str(arguments.get("id", ""))) + if memory_id is None: + return {"ok": False, "error": "valid id is required"} + return { + "ok": context.storage.delete_memory(context.user_id, memory_id), + "id": memory_id, + } + + +def create_note(arguments: dict[str, Any], context: ToolContext) -> ToolResult: + text = coerce_required_text(arguments, "text") + if not text: + return {"ok": False, "error": "text is required"} + note_id = context.storage.add_note(context.user_id, text) + return {"ok": True, "id": note_id, "text": text} + + +def list_notes(arguments: dict[str, Any], context: ToolContext) -> ToolResult: + rows = context.storage.list_notes( + context.user_id, + coerce_positive_int(arguments.get("limit"), 20), + ) + return {"ok": True, "items": rows} + + +def delete_note(arguments: dict[str, Any], context: ToolContext) -> ToolResult: + note_id = parse_positive_int(str(arguments.get("id", ""))) + if note_id is None: + return {"ok": False, "error": "valid id is required"} + return { + "ok": context.storage.delete_note(context.user_id, note_id), + "id": note_id, + } + + +def create_reminder( + arguments: dict[str, Any], + context: ToolContext, +) -> ToolResult: + when = coerce_required_text(arguments, "when") + text = coerce_required_text(arguments, "text") + if not when or not text: + return {"ok": False, "error": "when and text are required"} + parsed = parse_agent_reminder(when, text, context.tz) + if not parsed: + return {"ok": False, "error": "could not parse reminder time"} + if parsed.remind_at_utc <= utc_now(): + return {"ok": False, "error": "reminder time must be in the future"} + + reminder_id = context.storage.add_reminder( + context.user_id, + context.chat_id, + parsed.text, + parsed.remind_at_utc, + ) + return { + "ok": True, + "id": reminder_id, + "text": parsed.text, + "remind_at": format_local_dt(parsed.remind_at_utc, context.tz), + } + + +def list_reminders( + arguments: dict[str, Any], + context: ToolContext, +) -> ToolResult: + rows = context.storage.list_reminders( + context.user_id, + coerce_positive_int(arguments.get("limit"), 20), + ) + for row in rows: + row["remind_at_local"] = format_local_dt( + row["remind_at"], + context.tz, + ) + return {"ok": True, "items": rows} + + +def cancel_reminder( + arguments: dict[str, Any], + context: ToolContext, +) -> ToolResult: + reminder_id = parse_positive_int(str(arguments.get("id", ""))) + if reminder_id is None: + return {"ok": False, "error": "valid id is required"} + return { + "ok": context.storage.cancel_reminder( + context.user_id, + reminder_id, + ), + "id": reminder_id, + } + + +def create_status(arguments: dict[str, Any], context: ToolContext) -> ToolResult: + title = coerce_required_text(arguments, "title") + status = coerce_required_text(arguments, "status") or "open" + if not title: + return {"ok": False, "error": "title is required"} + item_id = context.storage.add_tracked_item( + context.user_id, + title, + status, + ) + return { + "ok": True, + "id": item_id, + "title": title, + "status": status, + } + + +def list_statuses( + arguments: dict[str, Any], + context: ToolContext, +) -> ToolResult: + rows = context.storage.list_tracked_items( + context.user_id, + coerce_positive_int(arguments.get("limit"), 30), + ) + return {"ok": True, "items": rows} + + +def update_status(arguments: dict[str, Any], context: ToolContext) -> ToolResult: + item_id = parse_positive_int(str(arguments.get("id", ""))) + status = coerce_required_text(arguments, "status") + if item_id is None or not status: + return {"ok": False, "error": "valid id and status are required"} + return { + "ok": context.storage.set_tracked_status( + context.user_id, + item_id, + status, + ), + "id": item_id, + "status": status, + } + + +def delete_status(arguments: dict[str, Any], context: ToolContext) -> ToolResult: + item_id = parse_positive_int(str(arguments.get("id", ""))) + if item_id is None: + return {"ok": False, "error": "valid id is required"} + return { + "ok": context.storage.delete_tracked_item( + context.user_id, + item_id, + ), + "id": item_id, + } + + +def search_conversation( + arguments: dict[str, Any], + context: ToolContext, +) -> ToolResult: + query = coerce_required_text(arguments, "query") + if not query: + return {"ok": False, "error": "query is required"} + limit = coerce_positive_int(arguments.get("limit"), 5, maximum=10) + hits = context.storage.search_conversation_messages( + context.user_id, + context.chat_id, + query, + limit, + ) + discussions = [ + { + "matched_message_id": hit["id"], + "score": hit["score"], + "messages": context.storage.conversation_message_window( + context.user_id, + context.chat_id, + int(hit["id"]), + ), + } + for hit in hits + ] + return { + "ok": True, + "query": query, + "discussions": discussions, + } + + +AGENT_TOOLS = ( + AgentTool( + "get_current_datetime", + "", + "узнать текущее локальное и UTC-время.", + get_current_datetime, + ), + AgentTool( + "remember", + '{"text": string}', + "сохранить важный долгосрочный факт о пользователе.", + remember, + ), + AgentTool( + "list_memory", + '{"limit": number}', + "показать сохраненную память.", + list_memory, + ), + AgentTool( + "delete_memory", + '{"id": number}', + "удалить запись памяти.", + delete_memory, + ), + AgentTool( + "create_note", + '{"text": string}', + "сохранить заметку.", + create_note, + ), + AgentTool( + "list_notes", + '{"limit": number}', + "показать заметки.", + list_notes, + ), + AgentTool( + "delete_note", + '{"id": number}', + "удалить заметку.", + delete_note, + ), + AgentTool( + "create_reminder", + '{"when": string, "text": string}', + ( + "поставить напоминание. when можно указывать как '30m', " + "'через 2 часа', '18:30', '2026-07-16 18:30'. Если пользователь " + "говорит 'завтра/послезавтра/через неделю', сам рассчитай дату от " + "текущего локального времени и передай 'YYYY-MM-DD HH:MM'." + ), + create_reminder, + ), + AgentTool( + "list_reminders", + '{"limit": number}', + "показать активные напоминания.", + list_reminders, + ), + AgentTool( + "cancel_reminder", + '{"id": number}', + "отменить напоминание.", + cancel_reminder, + ), + AgentTool( + "create_status", + '{"title": string, "status": string}', + "начать отслеживать статус.", + create_status, + ), + AgentTool( + "list_statuses", + '{"limit": number}', + "показать отслеживаемые статусы.", + list_statuses, + ), + AgentTool( + "update_status", + '{"id": number, "status": string}', + "обновить статус.", + update_status, + ), + AgentTool( + "delete_status", + '{"id": number}', + "удалить отслеживаемый объект.", + delete_status, + ), + AgentTool( + "search_conversation", + '{"query": string, "limit": number}', + ( + "найти старое обсуждение во всей сохраненной переписке по " + "содержательным ключевым словам." + ), + search_conversation, + paragraph_before=True, + ), +) +AGENT_TOOL_BY_NAME = {tool.name: tool for tool in AGENT_TOOLS} + + +def render_agent_tool_catalog() -> str: + return "\n".join(tool.prompt_line() for tool in AGENT_TOOLS) + + +def execute_agent_tool( + name: str, + arguments: dict[str, Any], + storage: AssistantStorage, + user_id: int, + chat_id: int, + tz: ZoneInfo, +) -> ToolResult: + tool = AGENT_TOOL_BY_NAME.get(name.strip().lower()) + if tool is None: + return {"ok": False, "error": f"unknown tool: {name}"} + return tool.execute( + arguments, + ToolContext( + storage=storage, + user_id=user_id, + chat_id=chat_id, + tz=tz, + ), + ) diff --git a/tests/test_core.py b/tests/test_core.py index 87dd823..4740726 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -11,6 +11,7 @@ from assistant_bot.agent import ( execute_agent_tool, parse_agent_decision, ) +from assistant_bot.agent_tools import AGENT_TOOLS from assistant_bot.reminders import parse_reminder, unit_to_timedelta from assistant_bot.storage import AssistantStorage @@ -78,6 +79,10 @@ class AgentToolContractTests(unittest.TestCase): def test_prompt_exposes_all_supported_tool_names(self) -> None: prompt = build_agent_tool_prompt(ZoneInfo("UTC")) + self.assertEqual( + tuple(tool.name for tool in AGENT_TOOLS), + self.TOOL_NAMES, + ) for tool_name in self.TOOL_NAMES: with self.subTest(tool_name=tool_name): self.assertIn(tool_name, prompt)