517 lines
21 KiB
Python
517 lines
21 KiB
Python
import json
|
||
import logging
|
||
import re
|
||
from datetime import datetime, timezone
|
||
from typing import Any
|
||
from zoneinfo import ZoneInfo
|
||
|
||
from telegram import Update
|
||
from telegram.constants import ChatAction
|
||
from telegram.ext import ContextTypes
|
||
|
||
from .ai import AIClientError
|
||
from .config import CONVERSATION_HISTORY_LIMIT, MAX_AGENT_STEPS
|
||
from .models import AgentDecision, ReminderParseResult
|
||
from .prompts import ASSISTANT_SYSTEM_PROMPT
|
||
from .reminders import parse_reminder
|
||
from .storage import AssistantStorage
|
||
from .telegram_utils import (
|
||
edit_or_reply,
|
||
get_ai_client,
|
||
get_storage,
|
||
get_tz,
|
||
parse_positive_int,
|
||
require_user_id,
|
||
)
|
||
from .time_utils import format_local_dt, to_utc_iso, utc_now
|
||
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def build_assistant_context(storage: AssistantStorage, user_id: int, tz: ZoneInfo) -> str:
|
||
memories = storage.list_memories(user_id, limit=20)
|
||
notes = storage.list_notes(user_id, limit=10)
|
||
reminders = storage.list_reminders(user_id, limit=10)
|
||
tracked_items = storage.list_tracked_items(user_id, limit=20)
|
||
|
||
sections: list[str] = []
|
||
if memories:
|
||
sections.append(
|
||
"Долговременная память:\n"
|
||
+ "\n".join(f"- #{row['id']}: {row['text']}" for row in memories)
|
||
)
|
||
if notes:
|
||
sections.append(
|
||
"Последние заметки:\n"
|
||
+ "\n".join(f"- #{row['id']}: {row['text']}" for row in notes)
|
||
)
|
||
if reminders:
|
||
sections.append(
|
||
"Активные напоминания:\n"
|
||
+ "\n".join(
|
||
f"- #{row['id']} {format_local_dt(row['remind_at'], tz)}: {row['text']}"
|
||
for row in reminders
|
||
)
|
||
)
|
||
if tracked_items:
|
||
sections.append(
|
||
"Отслеживаемые статусы:\n"
|
||
+ "\n".join(
|
||
f"- #{row['id']} {row['title']}: {row['status']}"
|
||
for row in tracked_items
|
||
)
|
||
)
|
||
|
||
if not sections:
|
||
return "Долговременный контекст пока пуст."
|
||
return "\n\n".join(sections)
|
||
|
||
|
||
def json_dumps(data: Any) -> str:
|
||
return json.dumps(data, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def build_agent_tool_prompt(tz: ZoneInfo) -> str:
|
||
now_local = datetime.now(tz).strftime("%Y-%m-%d %H:%M")
|
||
return (
|
||
"Ты работаешь как агент с внутренними tools. Пользователь пишет свободным текстом, "
|
||
"а Telegram-команды ему не нужны.\n"
|
||
f"Текущее локальное время: {now_local}. Таймзона: {tz.key}.\n\n"
|
||
"Отвечай СТРОГО одним JSON-объектом без Markdown и без текста вокруг.\n"
|
||
"Если нужно выполнить действие, верни tool_calls. Если действие уже выполнено "
|
||
"или tool не нужен, верни final.\n\n"
|
||
"Форматы ответа:\n"
|
||
"{\"tool_calls\":[{\"name\":\"create_note\",\"arguments\":{\"text\":\"...\"}}],\"reset_context\":false}\n"
|
||
"{\"final\":\"Короткий ответ пользователю\",\"reset_context\":false}\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"
|
||
"Правила:\n"
|
||
"- Для просьб 'запомни', 'сохрани как факт', 'учти на будущее' используй remember.\n"
|
||
"- Для заметок используй create_note, для напоминаний create_reminder, "
|
||
"для контроля дел/заявок/ожиданий create_status или update_status.\n"
|
||
"- Если для действия не хватает данных, не вызывай tool, а задай уточняющий вопрос через final.\n"
|
||
"- Учитывай историю диалога для коротких ответов на уточняющие вопросы. "
|
||
"Если новый запрос явно начинает другую, не связанную с историей тему, не опирайся на старую тему "
|
||
"и верни reset_context=true. Для продолжения темы и сомнительных случаев верни false.\n"
|
||
"- Если пользователь просит найти, вспомнить или продолжить старое обсуждение, вызови "
|
||
"search_conversation. Передавай в query только ключевые слова темы, без общих слов. "
|
||
"Результаты поиска содержат соседние реплики; более свежие совпадения при прочих равных важнее.\n"
|
||
"- После результата tool верни final с кратким подтверждением или следующим tool_calls."
|
||
)
|
||
|
||
|
||
def extract_json_object(raw_text: str) -> dict[str, Any] | None:
|
||
candidates = [raw_text.strip()]
|
||
code_block = re.search(r"```(?:json)?\s*(.*?)```", raw_text, re.IGNORECASE | re.DOTALL)
|
||
if code_block:
|
||
candidates.insert(0, code_block.group(1).strip())
|
||
|
||
first_brace = raw_text.find("{")
|
||
last_brace = raw_text.rfind("}")
|
||
if 0 <= first_brace < last_brace:
|
||
candidates.append(raw_text[first_brace : last_brace + 1])
|
||
|
||
for candidate in candidates:
|
||
if not candidate:
|
||
continue
|
||
try:
|
||
parsed = json.loads(candidate)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
if isinstance(parsed, dict):
|
||
return parsed
|
||
|
||
return None
|
||
|
||
|
||
def parse_agent_decision(raw_text: str) -> AgentDecision:
|
||
parsed = extract_json_object(raw_text)
|
||
if parsed is None:
|
||
return AgentDecision(final=raw_text.strip(), tool_calls=[], reset_context=False)
|
||
|
||
raw_calls = parsed.get("tool_calls")
|
||
if raw_calls is None and parsed.get("tool"):
|
||
raw_calls = [
|
||
{
|
||
"name": parsed.get("tool"),
|
||
"arguments": parsed.get("arguments", {}),
|
||
}
|
||
]
|
||
|
||
tool_calls: list[dict[str, Any]] = []
|
||
if isinstance(raw_calls, list):
|
||
for raw_call in raw_calls:
|
||
if not isinstance(raw_call, dict):
|
||
continue
|
||
name = raw_call.get("name") or raw_call.get("tool")
|
||
arguments = raw_call.get("arguments", {})
|
||
if isinstance(name, str):
|
||
tool_calls.append(
|
||
{
|
||
"name": name.strip(),
|
||
"arguments": arguments if isinstance(arguments, dict) else {},
|
||
}
|
||
)
|
||
|
||
final = parsed.get("final") or parsed.get("answer")
|
||
reset_context = parsed.get("reset_context") is True
|
||
if isinstance(final, str) and final.strip():
|
||
return AgentDecision(
|
||
final=final.strip(),
|
||
tool_calls=tool_calls,
|
||
reset_context=reset_context,
|
||
)
|
||
return AgentDecision(final=None, tool_calls=tool_calls, reset_context=reset_context)
|
||
|
||
|
||
def save_conversation_exchange(
|
||
storage: AssistantStorage,
|
||
user_id: int,
|
||
chat_id: int,
|
||
prompt: str,
|
||
answer: str,
|
||
reset_context: bool,
|
||
) -> None:
|
||
if reset_context:
|
||
storage.start_new_conversation(user_id, chat_id)
|
||
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,
|
||
prompt: str,
|
||
) -> None:
|
||
if not update.effective_message:
|
||
return
|
||
|
||
user_id = require_user_id(update)
|
||
if user_id is None:
|
||
await update.effective_message.reply_text("Не могу определить пользователя.")
|
||
return
|
||
|
||
storage = get_storage(context)
|
||
ai_client = get_ai_client(context)
|
||
tz = get_tz(context)
|
||
model = ai_client.normalize_model(
|
||
storage.get_user_model(user_id, ai_client.provider)
|
||
)
|
||
chat_id = int(update.effective_message.chat_id)
|
||
assistant_context = build_assistant_context(storage, user_id, tz)
|
||
conversation_history = storage.list_conversation_messages(
|
||
user_id,
|
||
chat_id,
|
||
limit=CONVERSATION_HISTORY_LIMIT,
|
||
)
|
||
messages = [
|
||
{"role": "system", "content": ASSISTANT_SYSTEM_PROMPT},
|
||
{"role": "system", "content": build_agent_tool_prompt(tz)},
|
||
{"role": "system", "content": assistant_context},
|
||
{
|
||
"role": "system",
|
||
"content": (
|
||
"Далее идет недавняя история текущей темы от старых реплик к новым. "
|
||
"Чем старше реплика, тем меньше ее приоритет; при противоречии опирайся "
|
||
"на более новые сообщения. Архив других тем доступен через search_conversation."
|
||
),
|
||
},
|
||
*(
|
||
{"role": str(row["role"]), "content": str(row["content"])}
|
||
for row in conversation_history
|
||
),
|
||
{"role": "user", "content": prompt},
|
||
]
|
||
|
||
await context.bot.send_chat_action(
|
||
chat_id=chat_id,
|
||
action=ChatAction.TYPING,
|
||
)
|
||
placeholder = await update.effective_message.reply_text(
|
||
f"Думаю через {ai_client.display_name} / {model}..."
|
||
)
|
||
|
||
try:
|
||
last_tool_results: list[dict[str, Any]] = []
|
||
reset_context = False
|
||
for _step in range(MAX_AGENT_STEPS):
|
||
raw_answer = await ai_client.chat(model, messages, json_mode=True)
|
||
decision = parse_agent_decision(raw_answer)
|
||
reset_context = reset_context or decision.reset_context
|
||
|
||
if decision.tool_calls:
|
||
tool_results = []
|
||
for call in decision.tool_calls:
|
||
result = execute_agent_tool(
|
||
name=str(call.get("name", "")),
|
||
arguments=call.get("arguments", {}),
|
||
storage=storage,
|
||
user_id=user_id,
|
||
chat_id=chat_id,
|
||
tz=tz,
|
||
)
|
||
tool_results.append(
|
||
{
|
||
"name": call.get("name"),
|
||
"arguments": call.get("arguments", {}),
|
||
"result": result,
|
||
}
|
||
)
|
||
|
||
last_tool_results = tool_results
|
||
messages.append({"role": "assistant", "content": json_dumps({"tool_calls": decision.tool_calls})})
|
||
messages.append(
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"Результаты tools:\n"
|
||
f"{json_dumps({'tool_results': tool_results})}\n"
|
||
"Продолжи. Верни либо следующий tool_calls, либо final. "
|
||
"Ответ снова строго JSON."
|
||
),
|
||
}
|
||
)
|
||
continue
|
||
|
||
if decision.final:
|
||
await edit_or_reply(placeholder, decision.final)
|
||
save_conversation_exchange(
|
||
storage,
|
||
user_id,
|
||
chat_id,
|
||
prompt,
|
||
decision.final,
|
||
reset_context,
|
||
)
|
||
return
|
||
|
||
messages.append({"role": "assistant", "content": raw_answer})
|
||
|
||
final_answer = await ai_client.chat(
|
||
model,
|
||
[
|
||
*messages,
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"Лимит tool-шагов исчерпан. Больше не вызывай tools. "
|
||
f"Последние результаты tools: {json_dumps(last_tool_results)}. "
|
||
"Сформулируй короткий финальный ответ пользователю обычным текстом."
|
||
),
|
||
},
|
||
],
|
||
)
|
||
await edit_or_reply(placeholder, final_answer)
|
||
save_conversation_exchange(
|
||
storage,
|
||
user_id,
|
||
chat_id,
|
||
prompt,
|
||
final_answer,
|
||
reset_context,
|
||
)
|
||
except AIClientError as exc:
|
||
if ai_client.provider == "local":
|
||
hint = f"Проверь, что Ollama запущена и модель установлена: ollama pull {model}"
|
||
else:
|
||
hint = (
|
||
"Проверь YANDEX_CLOUD_FOLDER, YC_API_KEY "
|
||
"и доступ к выбранной модели."
|
||
)
|
||
await placeholder.edit_text(
|
||
f"Не получилось вызвать {ai_client.display_name}.\n{exc}\n\n{hint}"
|
||
)
|
||
return
|
||
except Exception:
|
||
logger.exception("Failed to run AI prompt")
|
||
await placeholder.edit_text("Произошла внутренняя ошибка при запросе к модели.")
|
||
|
||
|
||
async def run_ai_prompt(
|
||
update: Update,
|
||
context: ContextTypes.DEFAULT_TYPE,
|
||
prompt: str,
|
||
) -> None:
|
||
await run_agent_prompt(update, context, prompt)
|