From 058ec809ff49016f273eeb404776738eb01e453a Mon Sep 17 00:00:00 2001
From: h <h@kotikot.com>
Date: Wed, 21 Jan 2026 02:51:59 +0100
Subject: [PATCH] fix(bot): better buttons

---
 backend/src/bot/handlers/message/handler.py |  4 ++--
 backend/src/bot/modules/ai/agent.py         | 10 ++++++----
 backend/src/bot/modules/ai/prompts.py       |  5 ++++-
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/backend/src/bot/handlers/message/handler.py b/backend/src/bot/handlers/message/handler.py
index 9558093..902193c 100644
--- a/backend/src/bot/handlers/message/handler.py
+++ b/backend/src/bot/handlers/message/handler.py
@@ -202,7 +202,7 @@ async def process_message_from_web(  # noqa: C901, PLR0912, PLR0915
         follow_up_agent = create_follow_up_agent(
             api_key=api_key, model_name=follow_up_model, system_prompt=follow_up_prompt
         )
-        follow_ups = await get_follow_ups(follow_up_agent, full_history)
+        follow_ups = await get_follow_ups(follow_up_agent, full_history, chat_images)
 
         if state:
             await state.stop_typing()
@@ -336,7 +336,7 @@ async def process_message(
         follow_up_agent = create_follow_up_agent(
             api_key=api_key, model_name=follow_up_model, system_prompt=follow_up_prompt
         )
-        follow_ups = await get_follow_ups(follow_up_agent, full_history)
+        follow_ups = await get_follow_ups(follow_up_agent, full_history, chat_images)
 
         await state.stop_typing()
 
diff --git a/backend/src/bot/modules/ai/agent.py b/backend/src/bot/modules/ai/agent.py
index 0603bd8..92f9a2f 100644
--- a/backend/src/bot/modules/ai/agent.py
+++ b/backend/src/bot/modules/ai/agent.py
@@ -99,15 +99,17 @@ async def stream_response(  # noqa: PLR0913
 async def get_follow_ups(
     follow_up_agent: Agent[None, FollowUpOptions],
     history: list[dict[str, str]],
-    image: ImageData | None = None,
+    images: list[ImageData] | None = None,
 ) -> list[str]:
     message_history = build_message_history(history) if history else None
 
-    if image:
+    if images:
         prompt: list[str | BinaryContent] = [
-            "Suggest follow-up options based on this conversation and image.",
-            BinaryContent(data=image.data, media_type=image.media_type),
+            "Suggest follow-up options based on this conversation and images.",
         ]
+        prompt.extend(
+            BinaryContent(data=img.data, media_type=img.media_type) for img in images
+        )
     else:
         prompt = "Suggest follow-up questions based on this conversation."  # type: ignore[assignment]
 
diff --git a/backend/src/bot/modules/ai/prompts.py b/backend/src/bot/modules/ai/prompts.py
index c65b784..b1e9bca 100644
--- a/backend/src/bot/modules/ai/prompts.py
+++ b/backend/src/bot/modules/ai/prompts.py
@@ -15,7 +15,10 @@ When asked for DETAILS on a specific problem (or a problem number):
 EXAM_FOLLOW_UP = """You see a problem set image. List available problem numbers.
 Output only the numbers that exist in the image, like: 1, 2, 3, 4, 5
 If problems have letters (a, b, c), list them as: 1a, 1b, 2a, etc.
-Keep it minimal - just the identifiers."""
+Keep it minimal - just the identifiers.
+Then, if applicable, output some possible followups of conversation
+If there is a problem set, you MUST output ALL of problem numbers/letters,
+before another followups"""
 
 DEFAULT_FOLLOW_UP = (
     "Based on the conversation, suggest 3 short follow-up questions "