init

assistant_bot/agent.py

@@ -0,0 +1,516 @@
+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)