From a17cef5b03f500dee1e9e9c3fa0bcdb0c03f3eee Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Thu, 28 May 2026 23:00:44 +0530 Subject: [PATCH] Group Telegram tool streams by response --- .../helpers/draft_stream.py | 21 ++++- tests/test_telegram_intermediate_response.py | 82 +++++++++++++++++++ 2 files changed, 101 insertions(+), 2 deletions(-) diff --git a/plugins/_telegram_integration/helpers/draft_stream.py b/plugins/_telegram_integration/helpers/draft_stream.py index b246f8ab3..a64cb47a6 100644 --- a/plugins/_telegram_integration/helpers/draft_stream.py +++ b/plugins/_telegram_integration/helpers/draft_stream.py @@ -29,10 +29,10 @@ TOOL_EMOJIS: dict[str, str] = { "code": "⌨️", "code_execution_tool": "⌨️", "duckduckgo_search": "🔎", + "read_file": "📖", "file": "📄", "knowledge_tool": "📚", "memory": "🧠", - "read_file": "📖", "search": "🔎", "search_engine": "🔎", "search_files": "🔎", @@ -89,6 +89,12 @@ async def send_intermediate_response( response_text: str, keyboard: list[list[dict]] | None = None, ) -> bool: + if context.data.get(CTX_TG_RESPONSE_MESSAGE_ID): + sent = await finalize_response(context, response_text, keyboard) + if sent: + _reset_progress_group(context) + return sent + html = _format_response(response_text) if not html: return False @@ -104,7 +110,10 @@ async def send_intermediate_response( parse_mode="HTML", reply_markup=_keyboard_markup(keyboard), ) - return bool(sent_id) + sent = bool(sent_id) + if sent: + _reset_progress_group(context) + return sent except Exception as e: PrintStyle.debug(f"Telegram intermediate response failed: {e}") return False @@ -142,6 +151,11 @@ def clear(context: AgentContext) -> None: context.data.pop(key, None) +def _reset_progress_group(context: AgentContext) -> None: + context.data.pop(CTX_TG_PROGRESS_LINES, None) + context.data.pop(CTX_TG_PROGRESS_MESSAGE_ID, None) + + def _stream_enabled(context: AgentContext) -> bool: value = context.get_data(CTX_TG_STREAM_ENABLED) return True if value is None else bool(value) @@ -208,6 +222,7 @@ async def _update_response_message( keyboard: list[list[dict]] | None = None, force: bool = False, ) -> bool: + had_message = bool(context.data.get(CTX_TG_RESPONSE_MESSAGE_ID)) message_id = await _ensure_response_message(context, text) bot = _bot_instance(context) chat_id = context.data.get(CTX_TG_CHAT_ID) @@ -215,6 +230,8 @@ async def _update_response_message( return False markup = _keyboard_markup(keyboard) html = _format_response(text) + if not had_message and not markup and not force: + return True try: ok = await tc.raw_edit_text( bot.bot.token, diff --git a/tests/test_telegram_intermediate_response.py b/tests/test_telegram_intermediate_response.py index cb56946ed..c8c572557 100644 --- a/tests/test_telegram_intermediate_response.py +++ b/tests/test_telegram_intermediate_response.py @@ -4,6 +4,8 @@ from plugins._telegram_integration.helpers import draft_stream from plugins._telegram_integration.helpers.constants import ( CTX_TG_BOT, CTX_TG_CHAT_ID, + CTX_TG_PROGRESS_LINES, + CTX_TG_PROGRESS_MESSAGE_ID, CTX_TG_REPLY_TO, CTX_TG_RESPONSE_MESSAGE_ID, ) @@ -68,3 +70,83 @@ def test_intermediate_response_sends_separate_non_reply_message(monkeypatch): } ] assert CTX_TG_RESPONSE_MESSAGE_ID not in context.data + + +def test_intermediate_response_finalizes_active_stream_and_starts_next_tool_group(monkeypatch): + calls = [] + next_message_id = 100 + + async def fake_send(token, chat_id, text, reply_to_message_id=None, parse_mode="HTML", reply_markup=None): + nonlocal next_message_id + calls.append( + { + "method": "send", + "text": text, + "reply_to_message_id": reply_to_message_id, + "parse_mode": parse_mode, + "reply_markup": reply_markup, + "message_id": next_message_id, + } + ) + next_message_id += 1 + return next_message_id - 1 + + async def fake_edit(token, chat_id, message_id, text, parse_mode="HTML", reply_markup=None): + calls.append( + { + "method": "edit", + "message_id": message_id, + "text": text, + "parse_mode": parse_mode, + "reply_markup": reply_markup, + } + ) + return True + + context = FakeContext() + monkeypatch.setattr(draft_stream, "_bot_instance", lambda ctx: FakeBot()) + monkeypatch.setattr(draft_stream.tc, "raw_send_text", fake_send) + monkeypatch.setattr(draft_stream.tc, "raw_edit_text", fake_edit) + + asyncio.run(draft_stream.add_tool_start(context, "search_engine", {"query": "telegram bot api"})) + asyncio.run(draft_stream.update_response(context, "Found the docs.")) + sent = asyncio.run(draft_stream.send_intermediate_response(context, "Found the docs.")) + asyncio.run(draft_stream.add_tool_start(context, "read_file", {"path": "notes.md"})) + + assert sent is True + assert calls == [ + { + "method": "send", + "text": "🔎 search engine: telegram bot api", + "reply_to_message_id": None, + "parse_mode": None, + "reply_markup": None, + "message_id": 100, + }, + { + "method": "send", + "text": "Found the docs.", + "reply_to_message_id": 456, + "parse_mode": "HTML", + "reply_markup": None, + "message_id": 101, + }, + { + "method": "edit", + "message_id": 101, + "text": "Found the docs.", + "parse_mode": "HTML", + "reply_markup": None, + }, + { + "method": "send", + "text": "📖 read file: notes.md", + "reply_to_message_id": None, + "parse_mode": None, + "reply_markup": None, + "message_id": 102, + }, + ] + assert context.data[CTX_TG_PROGRESS_MESSAGE_ID] == 102 + assert context.data[CTX_TG_PROGRESS_LINES] == ["📖 read file: notes.md"] + assert CTX_TG_RESPONSE_MESSAGE_ID not in context.data