refactor: centralize agent tool registry
This commit is contained in:
@@ -1,30 +1,29 @@
|
|||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime
|
||||||
from typing import Any
|
from typing import Any
|
||||||
from zoneinfo import ZoneInfo
|
from zoneinfo import ZoneInfo
|
||||||
|
|
||||||
from telegram import Update
|
from telegram import Update
|
||||||
from telegram.ext import ContextTypes
|
from telegram.ext import ContextTypes
|
||||||
|
|
||||||
|
from .agent_tools import execute_agent_tool, render_agent_tool_catalog
|
||||||
from .ai import AIClientError
|
from .ai import AIClientError
|
||||||
from .config import CONVERSATION_HISTORY_LIMIT, MAX_AGENT_STEPS
|
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 .prompts import ASSISTANT_SYSTEM_PROMPT
|
||||||
from .reminders import parse_reminder
|
|
||||||
from .storage import AssistantStorage
|
from .storage import AssistantStorage
|
||||||
from .telegram_utils import (
|
from .telegram_utils import (
|
||||||
escape_markdown_text,
|
escape_markdown_text,
|
||||||
get_ai_client,
|
get_ai_client,
|
||||||
get_storage,
|
get_storage,
|
||||||
get_tz,
|
get_tz,
|
||||||
parse_positive_int,
|
|
||||||
reply_markdown,
|
reply_markdown,
|
||||||
require_user_id,
|
require_user_id,
|
||||||
typing_action,
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -92,25 +91,7 @@ def build_agent_tool_prompt(tz: ZoneInfo) -> str:
|
|||||||
"читаемость. Для короткого простого ответа разметка не обязательна. "
|
"читаемость. Для короткого простого ответа разметка не обязательна. "
|
||||||
"Не добавляй декоративные эмодзи чаще одного раза на сообщение.\n\n"
|
"Не добавляй декоративные эмодзи чаще одного раза на сообщение.\n\n"
|
||||||
"Доступные tools:\n"
|
"Доступные tools:\n"
|
||||||
"- get_current_datetime {} - узнать текущее локальное и UTC-время.\n"
|
f"{render_agent_tool_catalog()}\n\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"
|
|
||||||
"Правила:\n"
|
"Правила:\n"
|
||||||
"- Для просьб 'запомни', 'сохрани как факт', 'учти на будущее' используй remember.\n"
|
"- Для просьб 'запомни', 'сохрани как факт', 'учти на будущее' используй remember.\n"
|
||||||
"- Для заметок используй create_note, для напоминаний create_reminder, "
|
"- Для заметок используй create_note, для напоминаний create_reminder, "
|
||||||
@@ -205,185 +186,6 @@ def save_conversation_exchange(
|
|||||||
storage.add_conversation_exchange(user_id, chat_id, prompt, answer)
|
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(
|
async def run_agent_prompt(
|
||||||
update: Update,
|
update: Update,
|
||||||
context: ContextTypes.DEFAULT_TYPE,
|
context: ContextTypes.DEFAULT_TYPE,
|
||||||
|
|||||||
417
assistant_bot/agent_tools.py
Normal file
417
assistant_bot/agent_tools.py
Normal file
@@ -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,
|
||||||
|
),
|
||||||
|
)
|
||||||
@@ -11,6 +11,7 @@ from assistant_bot.agent import (
|
|||||||
execute_agent_tool,
|
execute_agent_tool,
|
||||||
parse_agent_decision,
|
parse_agent_decision,
|
||||||
)
|
)
|
||||||
|
from assistant_bot.agent_tools import AGENT_TOOLS
|
||||||
from assistant_bot.reminders import parse_reminder, unit_to_timedelta
|
from assistant_bot.reminders import parse_reminder, unit_to_timedelta
|
||||||
from assistant_bot.storage import AssistantStorage
|
from assistant_bot.storage import AssistantStorage
|
||||||
|
|
||||||
@@ -78,6 +79,10 @@ class AgentToolContractTests(unittest.TestCase):
|
|||||||
def test_prompt_exposes_all_supported_tool_names(self) -> None:
|
def test_prompt_exposes_all_supported_tool_names(self) -> None:
|
||||||
prompt = build_agent_tool_prompt(ZoneInfo("UTC"))
|
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:
|
for tool_name in self.TOOL_NAMES:
|
||||||
with self.subTest(tool_name=tool_name):
|
with self.subTest(tool_name=tool_name):
|
||||||
self.assertIn(tool_name, prompt)
|
self.assertIn(tool_name, prompt)
|
||||||
|
|||||||
Reference in New Issue
Block a user