From 4b7a50d0cc3b2d1b8858bcee515d38e11c6946eb Mon Sep 17 00:00:00 2001 From: kandrusyak Date: Mon, 27 Jul 2026 17:20:32 +0300 Subject: [PATCH] add typing simulation --- assistant_bot/agent.py | 154 +++++++++++++++----------------- assistant_bot/telegram_utils.py | 51 ++++++++++- tests/test_telegram_utils.py | 32 ++++++- 3 files changed, 151 insertions(+), 86 deletions(-) diff --git a/assistant_bot/agent.py b/assistant_bot/agent.py index 222e733..156e3a7 100644 --- 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, "⚠️ **Произошла внутренняя ошибка** при запросе к модели.", ) diff --git a/assistant_bot/telegram_utils.py b/assistant_bot/telegram_utils.py index acb0b54..f387091 100644 --- 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)) diff --git a/tests/test_telegram_utils.py b/tests/test_telegram_utils.py index 5bba4c1..033a98a 100644 --- a/tests/test_telegram_utils.py +++ b/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()