add typing simulation

2026-07-27 17:20:32 +03:00
parent 275eca8dc5
commit 4b7a50d0cc
3 changed files with 151 additions and 86 deletions
--- a/assistant_bot/agent.py
+++ b/assistant_bot/agent.py
@@ -6,7 +6,6 @@ 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
@@ -16,16 +15,14 @@ from .prompts import ASSISTANT_SYSTEM_PROMPT
 from .reminders import parse_reminder
 from .storage import AssistantStorage
 from .telegram_utils import (
-    edit_markdown,
-    edit_or_reply,
    escape_markdown_text,
    get_ai_client,
    get_storage,
    get_tz,
-    markdown_code,
    parse_positive_int,
    reply_markdown,
    require_user_id,
+    typing_action,
 )
 from .time_utils import format_local_dt, to_utc_iso, utc_now

@@ -435,92 +432,81 @@ async def run_agent_prompt(
        {"role": "user", "content": prompt},
    ]

-    await context.bot.send_chat_action(
-        chat_id=chat_id,
-        action=ChatAction.TYPING,
-    )
-    placeholder = await reply_markdown(
-        update.effective_message,
-        f"💭 *Думаю через {escape_markdown_text(ai_client.display_name)} / "
-        f"{markdown_code(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
+        final_answer: str | None = None

-            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(
+        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(
                        {
-                            "name": call.get("name"),
-                            "arguments": call.get("arguments", {}),
-                            "result": result,
+                            "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

-                last_tool_results = tool_results
-                messages.append(
-                    {
-                        "role": "assistant",
-                        "content": json_dumps({"tool_calls": decision.tool_calls}),
-                    }
+                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."
+                            ),
+                        },
+                    ],
                )
-                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)}. "
-                        "Сформулируй короткий финальный ответ пользователю обычным Markdown."
-                    ),
-                },
-            ],
-        )
-        await edit_or_reply(placeholder, final_answer)
+        await reply_markdown(update.effective_message, final_answer)
        save_conversation_exchange(
            storage,
            user_id,
@@ -536,8 +522,8 @@ async def run_agent_prompt(
            hint = (
                "Проверь YANDEX_CLOUD_FOLDER, YC_API_KEY и доступ к выбранной модели."
            )
-        await edit_markdown(
-            placeholder,
+        await reply_markdown(
+            update.effective_message,
            f"⚠️ **Не получилось вызвать "
            f"{escape_markdown_text(ai_client.display_name)}.**\n\n"
            f"{escape_markdown_text(exc)}\n\n"
@@ -546,8 +532,8 @@ async def run_agent_prompt(
        return
    except Exception:
        logger.exception("Failed to run AI prompt")
-        await edit_markdown(
-            placeholder,
+        await reply_markdown(
+            update.effective_message,
            "⚠️ **Произошла внутренняя ошибка** при запросе к модели.",
        )

--- a/assistant_bot/telegram_utils.py
+++ b/assistant_bot/telegram_utils.py
@@ -1,10 +1,13 @@
+import asyncio
 import logging
 import re
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager, suppress
 from uuid import uuid4
 from zoneinfo import ZoneInfo

 from telegram import Bot, InlineQueryResultArticle, InputTextMessageContent, Message, Update
-from telegram.constants import ParseMode
+from telegram.constants import ChatAction, ParseMode
 from telegram.error import BadRequest
 from telegram.ext import ContextTypes
 from telegramify_markdown import (
@@ -24,6 +27,7 @@ from .storage import AssistantStorage
 logger = logging.getLogger(__name__)

 FORMATTING_FALLBACK_NOTICE = "⚠️ Не удалось применить форматирование."
+TYPING_REFRESH_INTERVAL_SECONDS = 4.0
 _MARKDOWN_SPECIAL_CHARS = re.compile(r"([\\`*_[\]{}()#+\-.!|>~])")
 _RENDER_CONFIG = RenderConfig()
 _RENDER_CONFIG.markdown_symbol.heading_level_1 = ""
@@ -32,6 +36,51 @@ _RENDER_CONFIG.markdown_symbol.heading_level_3 = ""
 _RENDER_CONFIG.markdown_symbol.heading_level_4 = ""


+async def _refresh_typing(
+    bot: Bot,
+    chat_id: int,
+    interval: float,
+) -> None:
+    while True:
+        await asyncio.sleep(interval)
+        try:
+            await bot.send_chat_action(
+                chat_id=chat_id,
+                action=ChatAction.TYPING,
+            )
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.warning("Failed to refresh Telegram typing action", exc_info=True)
+
+
+@asynccontextmanager
+async def typing_action(
+    bot: Bot,
+    chat_id: int,
+    interval: float = TYPING_REFRESH_INTERVAL_SECONDS,
+) -> AsyncIterator[None]:
+    """Keep Telegram's typing indicator active while work is in progress."""
+    try:
+        await bot.send_chat_action(
+            chat_id=chat_id,
+            action=ChatAction.TYPING,
+        )
+    except Exception:
+        logger.warning("Failed to send Telegram typing action", exc_info=True)
+
+    refresh_task = asyncio.create_task(
+        _refresh_typing(bot, chat_id, interval),
+        name=f"telegram-typing-{chat_id}",
+    )
+    try:
+        yield
+    finally:
+        refresh_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await refresh_task
+
+
 def escape_markdown_text(value: object) -> str:
    """Escape dynamic text before embedding it into ordinary Markdown."""
    return _MARKDOWN_SPECIAL_CHARS.sub(r"\\\1", str(value))