add typing simulation

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):
+        final_answer: str | None = None
-            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:
+        async with typing_action(context.bot, chat_id):
-                tool_results = []
+            for _step in range(MAX_AGENT_STEPS):
-                for call in decision.tool_calls:
+                raw_answer = await ai_client.chat(model, messages, json_mode=True)
-                    result = execute_agent_tool(
+                decision = parse_agent_decision(raw_answer)
-                        name=str(call.get("name", "")),
+                reset_context = reset_context or decision.reset_context
-                        arguments=call.get("arguments", {}),
+
-                        storage=storage,
+                if decision.tool_calls:
-                        user_id=user_id,
+                    tool_results = []
-                        chat_id=chat_id,
+                    for call in decision.tool_calls:
-                        tz=tz,
+                        result = execute_agent_tool(
-                    )
+                            name=str(call.get("name", "")),
-                    tool_results.append(
+                            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"),
+                            "role": "assistant",
-                            "arguments": call.get("arguments", {}),
+                            "content": json_dumps(
-                            "result": result,
+                                {"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
+                if decision.final:
-                messages.append(
+                    final_answer = decision.final
-                    {
+                    break
-                        "role": "assistant",
+
-                        "content": json_dumps({"tool_calls": decision.tool_calls}),
+                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 reply_markdown(update.effective_message, final_answer)
-                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)
         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(
+        await reply_markdown(
-            placeholder,
+            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(
+        await reply_markdown(
-            placeholder,
+            update.effective_message,
             "⚠️ **Произошла внутренняя ошибка** при запросе к модели.",
         )
 

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))

tests/test_telegram_utils.py

@@ -1,8 +1,9 @@
+import asyncio
 import unittest
 from types import SimpleNamespace
 from unittest.mock import AsyncMock
 
-from telegram.constants import ParseMode
+from telegram.constants import ChatAction, ParseMode
 from telegram.error import BadRequest
 
 from assistant_bot.telegram_utils import (
@@ -11,6 +12,7 @@ from assistant_bot.telegram_utils import (
     escape_markdown_text,
     render_markdown_chunks,
     reply_markdown,
+    typing_action,
 )
 
 
@@ -112,5 +114,33 @@ class MarkdownDeliveryTests(unittest.IsolatedAsyncioTestCase):
         )
 
 
+class TypingActionTests(unittest.IsolatedAsyncioTestCase):
+    async def test_refreshes_typing_until_context_exits(self) -> None:
+        refreshed = asyncio.Event()
+        call_count = 0
+
+        async def send_chat_action(**_kwargs) -> None:
+            nonlocal call_count
+            call_count += 1
+            if call_count >= 2:
+                refreshed.set()
+
+        bot = SimpleNamespace(
+            send_chat_action=AsyncMock(side_effect=send_chat_action)
+        )
+
+        async with typing_action(bot, chat_id=42, interval=0.001):
+            await asyncio.wait_for(refreshed.wait(), timeout=0.5)
+
+        calls_after_exit = call_count
+        await asyncio.sleep(0.01)
+
+        self.assertGreaterEqual(calls_after_exit, 2)
+        self.assertEqual(call_count, calls_after_exit)
+        for call in bot.send_chat_action.await_args_list:
+            self.assertEqual(call.kwargs["chat_id"], 42)
+            self.assertEqual(call.kwargs["action"], ChatAction.TYPING)
+
+
 if __name__ == "__main__":
     unittest.main()