341 lines
14 KiB
Python
341 lines
14 KiB
Python
import json
|
||
import logging
|
||
import re
|
||
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
|
||
from .prompts import ASSISTANT_SYSTEM_PROMPT
|
||
from .storage import AssistantStorage
|
||
from .telegram_utils import (
|
||
escape_markdown_text,
|
||
get_ai_client,
|
||
get_storage,
|
||
get_tz,
|
||
reply_markdown,
|
||
require_user_id,
|
||
typing_action,
|
||
)
|
||
from .time_utils import format_local_dt
|
||
|
||
|
||
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'
|
||
"Значение final оформляй обычным Markdown, не MarkdownV2. Умеренно используй "
|
||
"жирный и курсивный текст, списки, ссылки и блоки кода, когда они улучшают "
|
||
"читаемость. Для короткого простого ответа разметка не обязательна. "
|
||
"Не добавляй декоративные эмодзи чаще одного раза на сообщение.\n\n"
|
||
"Доступные tools:\n"
|
||
f"{render_agent_tool_catalog()}\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)
|
||
|
||
|
||
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 reply_markdown(
|
||
update.effective_message,
|
||
"⚠️ **Не могу определить пользователя.**",
|
||
)
|
||
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},
|
||
]
|
||
|
||
try:
|
||
last_tool_results: list[dict[str, Any]] = []
|
||
reset_context = False
|
||
final_answer: str | None = None
|
||
|
||
async with typing_action(context.bot, chat_id):
|
||
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:
|
||
final_answer = decision.final
|
||
break
|
||
|
||
messages.append({"role": "assistant", "content": raw_answer})
|
||
|
||
if final_answer is None:
|
||
final_answer = await ai_client.chat(
|
||
model,
|
||
[
|
||
*messages,
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"Лимит tool-шагов исчерпан. Больше не вызывай tools. "
|
||
f"Последние результаты tools: {json_dumps(last_tool_results)}. "
|
||
"Сформулируй короткий финальный ответ пользователю обычным Markdown."
|
||
),
|
||
},
|
||
],
|
||
)
|
||
|
||
await reply_markdown(update.effective_message, 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 reply_markdown(
|
||
update.effective_message,
|
||
f"⚠️ **Не получилось вызвать "
|
||
f"{escape_markdown_text(ai_client.display_name)}.**\n\n"
|
||
f"{escape_markdown_text(exc)}\n\n"
|
||
f"{escape_markdown_text(hint)}",
|
||
)
|
||
return
|
||
except Exception:
|
||
logger.exception("Failed to run AI prompt")
|
||
await reply_markdown(
|
||
update.effective_message,
|
||
"⚠️ **Произошла внутренняя ошибка** при запросе к модели.",
|
||
)
|