feat(telegram): stop generation from the client, questions survive a restart, images reach the model inline, albums become one turn

tests/test_conversations.py

@@ -88,6 +88,7 @@ class ScriptedClient:
     def __init__(self, options: Any) -> None:
         self.options = options
         self.prompts: list[str] = []
+        self.contents: list[list[dict[str, Any]]] = []
         self.session_id = options.resume or str(uuid.uuid4())
         self.interrupted = False
         self.connected = False
@@ -96,7 +97,14 @@ class ScriptedClient:
     async def connect(self) -> None:
         self.connected = True
 
-    async def query(self, prompt: str) -> None:
+    async def query(self, prompt) -> None:
+        if not isinstance(prompt, str):
+            messages = [m async for m in prompt]
+            content = messages[-1]["message"]["content"]
+            self.contents.append(content)
+            prompt = "\n".join(
+                b.get("text", "") for b in content if b.get("type") == "text"
+            )
         self.prompts.append(prompt)
 
     async def receive_response(self):

tests/test_telegram.py

@@ -1,4 +1,6 @@
 import asyncio
+import base64
+import contextlib
 import tempfile
 from pathlib import Path
 from types import SimpleNamespace
@@ -919,3 +921,100 @@ async def test_gone_topic_unbinds_and_the_next_message_starts_afresh(
         frontend="telegram", external_id=f"{USER}/7"
     )
     assert fresh is not None and fresh.id != branch.id and fresh.status == "open"
+
+
+async def test_photo_reaches_the_model_as_an_image_block(stack: Stack) -> None:
+    stack.bot.media(_photo("что на фото"))
+    await stack.until(lambda: stack.sent_with("что на фото"), what="reply")
+    client = next(c for c in ScriptedClient.instances if c.contents)
+    content = client.contents[-1]
+    assert content[0]["type"] == "text" and "что на фото" in content[0]["text"]
+    image = content[1]
+    assert image["type"] == "image"
+    assert image["source"]["media_type"] == "image/jpeg"
+    assert base64.b64decode(image["source"]["data"]) == b"data"
+
+
+async def test_album_becomes_one_turn(stack: Stack) -> None:
+    stack.tg.album_delay = 0.2
+    stack.bot.media({**_photo("альбом"), "media_group_id": "g1"})
+    second = _photo("")
+    second["photo"][0].update(file_id="f2", file_unique_id="u2")
+    del second["caption"]
+    stack.bot.media({**second, "media_group_id": "g1"})
+    await stack.until(lambda: stack.sent_with("альбом"), what="reply")
+    await asyncio.sleep(0.3)
+    prompts = [
+        p for c in ScriptedClient.instances for p in c.prompts if "attachment" in p
+    ]
+    assert len(prompts) == 1
+    assert prompts[0].count("[attachment: photo") == 2
+    client = next(c for c in ScriptedClient.instances if c.contents)
+    assert [b["type"] for b in client.contents[-1]] == ["text", "image", "image"]
+
+
+async def test_stop_button_interrupts_the_turn(stack: Stack) -> None:
+    stack.bot.message("hi")
+    await stack.until(lambda: stack.sent_with("hi"), what="reply")
+    ScriptedClient.hold = asyncio.Event()
+    stack.bot.message("думай долго")
+    await stack.until(lambda: stack.tg._drafts, what="draft")  # noqa: SLF001
+    draft = next(iter(stack.tg._drafts.values()))  # noqa: SLF001
+    await stack.until(
+        lambda: any(d.get("draft_id") == draft.draft_id for d in stack.bot.drafts),
+        what="draft pushed",
+    )
+    assert stack.bot.drafts[-1]["can_stop"] is True
+    stack.bot.push(
+        {
+            "stopped_message_generation": {
+                "chat": {"id": USER, "type": "private"},
+                "draft_id": draft.draft_id,
+                "message_thread_id": stack.bot.drafts[-1].get("message_thread_id"),
+            }
+        }
+    )
+    await stack.until(lambda: stack.sent_with("думай долго"), what="reply after stop")
+    master = await stack.world.conversations.find_bound(
+        frontend="telegram", external_id=GENERAL
+    )
+    items = await stack.world.conversations.queue.recent(master.id, limit=1)
+    assert items[0].status == "interrupted"
+    ScriptedClient.hold = None
+
+
+async def test_question_survives_a_gateway_restart(stack: Stack) -> None:
+    stack.bot.message("hi")
+    await stack.until(lambda: stack.sent_with("hi"), what="reply")
+    master = await stack.world.conversations.find_bound(
+        frontend="telegram", external_id=GENERAL
+    )
+    payload = {
+        "questions": [
+            {
+                "header": "Цвет",
+                "question": "Какой цвет?",
+                "options": [{"label": "Синий"}, {"label": "Красный"}],
+                "multiSelect": False,
+            }
+        ]
+    }
+    asking = asyncio.create_task(
+        stack.world.conversations.ask(master.external_id, payload)
+    )
+    question = await stack.until(
+        lambda: stack.sent_with("Какой цвет?"), what="question"
+    )
+    pending = stack.world.conversations.pending_question(master.external_id)
+    assert pending is not None
+    row = await stack.world.conversations.get(master.external_id)
+    assert row.flags["ask"]["id"] == pending[0]
+    asking.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await asking
+    stack.tg._asks.clear()  # noqa: SLF001
+    stack.world.conversations._questions.clear()  # noqa: SLF001
+    stack.bot.callback(f"q:{pending[0]}:0:1", question["message_id"])
+    await stack.until(lambda: stack.sent_with("ok:Красный"), what="answer as message")
+    row = await stack.world.conversations.get(master.external_id)
+    assert row.flags.get("ask") is None