From 582b0e17c9b826f41e4cbd5415dc5fb29c0f9fa4 Mon Sep 17 00:00:00 2001
From: ghostiee-11 <168410465+ghostiee-11@users.noreply.github.com>
Date: Tue, 24 Mar 2026 12:44:02 +0530
Subject: [PATCH 1/7] Add follow-up suggestion icon button after successful
 queries

Add a lightbulb icon to message footer that populates the chat input
with an AI-generated follow-up suggestion on click. Uses the cheapest
'ui' LLM spec for efficient generation.

- Add FollowUpSuggestion model in models.py
- Add suggest_followup() to Coordinator with llm_spec='ui'
- Add follow_up_suggestions.jinja2 extending Actor/main.jinja2
- Add lightbulb icon to footer_actions (fallback to footer_objects)
- Add follow_up_suggestions settings toggle
- 14 tests (10 UI + 4 coordinator) + docs

Depends on panel-extensions/panel-material-ui#605 for footer_actions.
---
 docs/configuration/ui.md                      |  14 +
 docs/getting_started/navigating_the_ui.md     |   2 +
 lumen/ai/coordinator/base.py                  |  41 ++-
 lumen/ai/models.py                            |   7 +
 .../Coordinator/follow_up_suggestions.jinja2  |  31 ++
 lumen/ai/ui.py                                |  74 +++++
 lumen/tests/ai/test_coordinator.py            |  80 ++++-
 lumen/tests/ai/test_ui.py                     | 292 +++++++++++++++++-
 8 files changed, 538 insertions(+), 3 deletions(-)
 create mode 100644 lumen/ai/prompts/Coordinator/follow_up_suggestions.jinja2

diff --git a/docs/configuration/ui.md b/docs/configuration/ui.md
index 548566375..5f3653101 100644
--- a/docs/configuration/ui.md
+++ b/docs/configuration/ui.md
@@ -78,6 +78,19 @@ ui = lmai.ExplorerUI(
 
 1. Tuples of (Material icon name, button text)
 
+### Follow-up suggestions
+
+After each successful query, a lightbulb icon appears in the message footer. Clicking it populates the chat input with an AI-generated follow-up suggestion that references actual column names from the data.
+
+``` py title="Disable follow-up suggestions"
+ui = lmai.ExplorerUI(
+    data='penguins.csv',
+    follow_up_suggestions=False
+)
+```
+
+Users can also toggle this at runtime via **Settings > Follow-Up Suggestions**.
+
 ## Advanced parameters
 
 ### Enable chat logging
@@ -203,6 +216,7 @@ Quick reference:
 | `notebook_preamble` | str | Export header |
 | `provider_choices` | dict | LLM providers shown in Settings |
 | `source_controls` | list | Source control components for data |
+| `follow_up_suggestions` | bool | AI follow-up suggestion icon after queries (default: True) |
 | `suggestions` | list | Quick action buttons |
 | `title` | str | App title |
 | `tools` | list | Custom tools |
diff --git a/docs/getting_started/navigating_the_ui.md b/docs/getting_started/navigating_the_ui.md
index 004bf5818..408b85599 100644
--- a/docs/getting_started/navigating_the_ui.md
+++ b/docs/getting_started/navigating_the_ui.md
@@ -110,6 +110,8 @@ Use **Settings** in the left sidebar to configure analysis behavior:
 
 **Validation Step** — AI validates results for correctness (enabled by default)
 
+**Follow-Up Suggestions** — After each successful query, a lightbulb icon appears in the message footer. Click it to populate the chat input with an AI-generated follow-up question. Toggle off in Settings if not needed.
+
 **Code Execution** — If enabled by administrator, controls how visualizations are generated. See [Code Execution for Visualizations](using_lumen_ai.md#code-execution-for-visualizations) for security implications.
 
 **LLM Configuration** — Choose and configure your language model
diff --git a/lumen/ai/coordinator/base.py b/lumen/ai/coordinator/base.py
index bef93e379..50eb332ec 100644
--- a/lumen/ai/coordinator/base.py
+++ b/lumen/ai/coordinator/base.py
@@ -25,7 +25,7 @@
 from ..config import PROMPTS_DIR, MissingContextError
 from ..context import TContext
 from ..llm import LlamaCpp, Llm, Message
-from ..models import ThinkingYesNo
+from ..models import FollowUpSuggestion, ThinkingYesNo
 from ..report import ActorTask, Section, TaskGroup
 from ..tools import MetadataLookup, Tool, VectorLookupToolUser
 from ..utils import (
@@ -284,6 +284,11 @@ class Coordinator(Viewer, VectorLookupToolUser):
                 "template": PROMPTS_DIR / "Coordinator" / "tool_relevance.jinja2",
                 "response_model": ThinkingYesNo,
             },
+            "follow_up_suggestions": {
+                "template": PROMPTS_DIR / "Coordinator" / "follow_up_suggestions.jinja2",
+                "response_model": FollowUpSuggestion,
+                "llm_spec": "ui",
+            },
         },
     )
 
@@ -509,6 +514,40 @@ async def _check_tool_relevance(self, tool: Tool, tool_output: str, actor: Actor
 
         return result.yes
 
+    async def suggest_followup(self, plan: Plan) -> str | None:
+        """Generate a single follow-up suggestion from a completed plan.
+
+        Returns a suggestion query string, or None on failure.
+        """
+        pipeline = plan.out_context.get("pipeline")
+        if pipeline is None:
+            return None
+
+        data_summary = plan.out_context.get("data", "")
+        if not data_summary:
+            return None
+
+        sql = plan.out_context.get("sql", "")
+        previous_queries = [
+            m["content"] for m in plan.history
+            if m.get("role") == "user"
+        ]
+        output_type = "chart" if plan.out_context.get("view") else "table"
+
+        try:
+            result = await self._invoke_prompt(
+                "follow_up_suggestions",
+                messages=[{"role": "user", "content": "Generate a follow-up suggestion."}],
+                context=plan.out_context,
+                data_summary=data_summary,
+                sql=sql,
+                output_type=output_type,
+                previous_queries=previous_queries,
+            )
+            return result.query
+        except Exception:
+            return None
+
     def __panel__(self):
         return self.interface
 
diff --git a/lumen/ai/models.py b/lumen/ai/models.py
index 271dd1abb..451dc632c 100644
--- a/lumen/ai/models.py
+++ b/lumen/ai/models.py
@@ -29,6 +29,13 @@ def model_post_init(self, __context):
             raise MissingContextError(self.insufficient_context_reason)
 
 
+class FollowUpSuggestion(BaseModel):
+    query: str = Field(
+        description="A single dataset-specific follow-up question referencing actual column names",
+        max_length=100
+    )
+
+
 class ErrorDescription(BaseModel):
     """
     Represents a user-facing error explanation.
diff --git a/lumen/ai/prompts/Coordinator/follow_up_suggestions.jinja2 b/lumen/ai/prompts/Coordinator/follow_up_suggestions.jinja2
new file mode 100644
index 000000000..79d7749f3
--- /dev/null
+++ b/lumen/ai/prompts/Coordinator/follow_up_suggestions.jinja2
@@ -0,0 +1,31 @@
+{% extends 'Actor/main.jinja2' %}
+
+{% block instructions %}
+Suggest ONE concise follow-up question based on the user's query and the resulting data.
+
+Rules:
+- MUST reference specific column names from the data summary
+- Keep it short and actionable (under 100 characters)
+- Do NOT repeat any previous queries
+- Do NOT suggest analyses the data cannot support
+{% endblock %}
+
+{% block examples %}
+Examples:
+- "Show revenue trend over date by product"
+- "Compare revenue across region"
+- "Filter to top 5 products by revenue"
+{% endblock %}
+
+{% block context %}
+Data Summary:
+{{ data_summary }}
+{% if sql %}
+SQL Executed: {{ sql }}
+{% endif %}
+Output Type: {{ output_type }}
+{% if previous_queries %}
+Previous Queries (do NOT repeat):
+{% for q in previous_queries %}- {{ q }}
+{% endfor %}{% endif %}
+{% endblock %}
diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py
index 02af91d22..3097995bf 100644
--- a/lumen/ai/ui.py
+++ b/lumen/ai/ui.py
@@ -410,6 +410,12 @@ class UI(Viewer):
     export_functions = param.Dict(default={}, doc="""
        Dictionary mapping from name of exporter to export function.""")
 
+    follow_up_suggestions = param.Boolean(default=True, doc="""
+        Whether to generate AI-powered follow-up suggestions after each
+        successful query. When enabled, a lightbulb icon button appears
+        in the message footer that populates the input with a follow-up
+        suggestion on click. Requires an additional LLM call per query.""")
+
     interface = param.ClassSelector(class_=ChatFeed, doc="""
         The interface for the Coordinator to interact with.""")
 
@@ -1042,6 +1048,7 @@ def _configure_coordinator(self):
         if self.analyses:
             agents.append(AnalysisAgent(analyses=self.analyses))
 
+        self._background_tasks = set()
         self._coordinator = self.coordinator(
             agents=agents,
             context=self.context,
@@ -1596,6 +1603,67 @@ async def _add_analysis_suggestions(self, plan: Plan):
             num_objects=len(self.interface.objects),
         )
 
+    def _add_followup_icon_to_footer(self, suggestion: str, num_objects: int):
+        """Add a lightbulb icon button to the message footer actions.
+
+        On click, populates the ChatAreaInput with the suggestion
+        without auto-submitting, so the user can review or edit before sending.
+        """
+
+        def on_click(event):
+            with edit_readonly(self._chat_input):
+                self._chat_input.value_input = suggestion
+            followup_button.visible = False
+
+        def hide_followup(_=None):
+            if len(self.interface.objects) > num_objects:
+                followup_button.visible = False
+
+        followup_button = IconButton(
+            icon="lightbulb",
+            description=f"Suggest: {suggestion}",
+            size="small",
+            icon_size="0.9em",
+            margin=(5, 0),
+            on_click=on_click,
+            name="FollowUp",
+            disabled=self.interface.param.loading,
+            color="default",
+            sx={"color": "#FFD700", "padding": "0 0.1em"},
+        )
+
+        if len(self.interface):
+            message = self.interface.objects[-1]
+            try:
+                existing = message.footer_actions or []
+            except AttributeError:
+                existing = message.footer_objects or []
+            existing = [
+                obj for obj in existing
+                if getattr(obj, 'name', '') != "FollowUp"
+            ]
+            try:
+                message.footer_actions = existing + [followup_button]
+            except AttributeError:
+                message.footer_objects = existing + [followup_button]
+            self.interface.param.watch(hide_followup, "objects")
+
+    async def _generate_followup_suggestions(self, plan: Plan):
+        """Render follow-up suggestion from the coordinator."""
+        if not self.follow_up_suggestions:
+            return
+
+        num_objects = len(self.interface.objects)
+        suggestion = await self._coordinator.suggest_followup(plan)
+        if not suggestion:
+            return
+
+        # Guard against stale message if user sent a new query while we were generating
+        if len(self.interface.objects) != num_objects:
+            return
+
+        self._add_followup_icon_to_footer(suggestion, num_objects)
+
     def __panel__(self):
         return self._main
 
@@ -1799,6 +1867,9 @@ def _render_sidebar(self) -> list[Viewable]:
         validation = Switch(label='Validation Step', description='Check if the response fully answered your question', value=True)
         self._coordinator.validation_enabled = validation
         switches.append(validation)
+        followup = Switch(label='Follow-Up Suggestions', description='Show AI-generated follow-up questions after each query', value=self.follow_up_suggestions)
+        followup.link(self, value='follow_up_suggestions', bidirectional=True)
+        switches.append(followup)
 
         # Add code execution selector if not hidden
         code_agents = False
@@ -2506,6 +2577,9 @@ async def _postprocess_exploration(
             await self._sync_sources(SimpleNamespace(new=plan.out_context), global_context=plan.out_context)
             if "pipeline" in plan.out_context:
                 await self._add_analysis_suggestions(plan)
+                task = asyncio.create_task(self._generate_followup_suggestions(plan))
+                self._background_tasks.add(task)
+                task.add_done_callback(self._background_tasks.discard)
 
             if is_new:
                 plan.param.watch(partial(self._update_views, exploration), "views")
diff --git a/lumen/tests/ai/test_coordinator.py b/lumen/tests/ai/test_coordinator.py
index 3ef3f5f5a..b4d29930c 100644
--- a/lumen/tests/ai/test_coordinator.py
+++ b/lumen/tests/ai/test_coordinator.py
@@ -19,7 +19,7 @@
 from lumen.ai.coordinator import Coordinator, Plan, Planner
 from lumen.ai.coordinator.planner import Reasoning, make_plan_model
 from lumen.ai.editors import SQLEditor
-from lumen.ai.models import ReplaceLine, RetrySpec
+from lumen.ai.models import FollowUpSuggestion, ReplaceLine, RetrySpec
 from lumen.ai.report import ActorTask
 from lumen.ai.schemas import get_metaset
 from lumen.ai.tools import FunctionTool, define_tool
@@ -587,3 +587,81 @@ async def test_process_files_duplicate_table_name():
     df = sources[0].get("data")
     # Should have the latest data
     assert df.iloc[0, 0] == 99
+
+
+# --- Follow-up suggestion coordinator tests ---
+
+async def test_suggest_followup_returns_suggestion(llm):
+    """Coordinator.suggest_followup should return a query string."""
+    suggestion_response = FollowUpSuggestion(query="Show value by category")
+    llm.set_responses([suggestion_response])
+
+    coordinator = Coordinator(llm=llm)
+    plan = Plan(
+        ActorTask(ChatAgent(llm=llm)),
+        history=[{"content": "Sum value column", "role": "user"}],
+        title="test",
+        context={},
+    )
+    plan.out_context = {
+        "pipeline": object(),
+        "data": "columns: category, value",
+        "sql": "SELECT * FROM t",
+    }
+
+    result = await coordinator.suggest_followup(plan)
+    assert result == "Show value by category"
+
+
+async def test_suggest_followup_no_pipeline(llm):
+    """Should return None when plan has no pipeline."""
+    coordinator = Coordinator(llm=llm)
+    plan = Plan(
+        ActorTask(ChatAgent(llm=llm)),
+        history=[],
+        title="test",
+        context={},
+    )
+    plan.out_context = {}
+
+    result = await coordinator.suggest_followup(plan)
+    assert result is None
+
+
+async def test_suggest_followup_no_data(llm):
+    """Should return None when plan has no data summary."""
+    coordinator = Coordinator(llm=llm)
+    plan = Plan(
+        ActorTask(ChatAgent(llm=llm)),
+        history=[],
+        title="test",
+        context={},
+    )
+    plan.out_context = {"pipeline": object(), "data": ""}
+
+    result = await coordinator.suggest_followup(plan)
+    assert result is None
+
+
+async def test_suggest_followup_llm_failure(llm):
+    """Should return None when LLM call fails."""
+    def raise_error():
+        raise RuntimeError("LLM unavailable")
+
+    llm.set_responses([raise_error])
+
+    coordinator = Coordinator(llm=llm)
+    plan = Plan(
+        ActorTask(ChatAgent(llm=llm)),
+        history=[{"content": "test", "role": "user"}],
+        title="test",
+        context={},
+    )
+    plan.out_context = {
+        "pipeline": object(),
+        "data": "columns: a, b",
+        "sql": "SELECT * FROM t",
+    }
+
+    result = await coordinator.suggest_followup(plan)
+    assert result is None
diff --git a/lumen/tests/ai/test_ui.py b/lumen/tests/ai/test_ui.py
index f07d818bf..790eced19 100644
--- a/lumen/tests/ai/test_ui.py
+++ b/lumen/tests/ai/test_ui.py
@@ -26,7 +26,7 @@
 from lumen.ai.config import PROVIDED_SOURCE_NAME
 from lumen.ai.coordinator import Plan
 from lumen.ai.editors import SQLEditor
-from lumen.ai.models import ErrorDescription
+from lumen.ai.models import ErrorDescription, FollowUpSuggestion
 from lumen.ai.report import ActorTask
 from lumen.ai.schemas import get_metaset
 from lumen.ai.ui import UI, Exploration, ExplorerUI
@@ -1464,6 +1464,7 @@ async def test_edit_on_second_exploration_preserves_first(explorer_ui):
     second one, only the second exploration (sql_2 + spec_2) should be
     removed. The first exploration (sql_1 + spec_1) must remain intact."""
     ui = explorer_ui
+    ui.follow_up_suggestions = False
     test_source = ui.context["source"]
     SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
 
@@ -1610,3 +1611,292 @@ async def noop_callback(contents, user, instance):
     # Parent and Home still intact
     assert len(ui._explorations.items) == 2
     assert parent_exploration.plan is not None
+
+
+# --- Follow-up suggestions tests ---
+
+def _get_footer_actions(msg):
+    """Get footer actions from message, falling back to footer_objects for older pmui."""
+    try:
+        return msg.footer_actions or []
+    except AttributeError:
+        return msg.footer_objects or []
+
+
+def test_followup_suggestion_model():
+    """FollowUpSuggestion model should validate correctly."""
+    result = FollowUpSuggestion(query="Show value by category")
+    assert result.query == "Show value by category"
+
+
+def test_followup_suggestion_model_max_length():
+    """FollowUpSuggestion should reject queries over 100 characters."""
+    with pytest.raises(Exception):
+        FollowUpSuggestion(query="x" * 101)
+
+
+async def test_followup_suggestion_after_query(explorer_ui):
+    """After a successful SQL query, a follow-up suggestion icon should appear in footer."""
+    ui = explorer_ui
+    test_source = ui.context["source"]
+    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
+
+    suggestion_response = FollowUpSuggestion(query="Show value by category")
+
+    ui.llm.set_responses([
+        SQLQueryWithTables(
+            query="SELECT SUM(value) as value_sum FROM test_table",
+            table_slug="test_sql_agg",
+            tables=["test_table"]
+        ),
+        "Summary of results",
+        suggestion_response,
+    ])
+
+    sql_agent = SQLAgent(llm=ui.llm)
+    chat_agent = ChatAgent(llm=ui.llm)
+    plan = Plan(
+        ActorTask(sql_agent),
+        ActorTask(chat_agent),
+        history=[{"content": "Sum value column", "role": "user"}],
+        title="Sum of value",
+        context=ui.context
+    )
+    await ui._execute_plan(plan)
+
+    # Wait for async follow-up suggestion task to complete
+    await asyncio.sleep(0.5)
+
+    last_msg = ui.interface.objects[-1]
+    followup_icons = [
+        f for f in _get_footer_actions(last_msg)
+        if getattr(f, 'name', '') == "FollowUp"
+    ]
+    assert len(followup_icons) == 1
+    assert followup_icons[0].icon == "lightbulb"
+
+
+async def test_followup_suggestion_disabled(explorer_ui):
+    """No suggestion should appear when follow_up_suggestions=False."""
+    ui = explorer_ui
+    ui.follow_up_suggestions = False
+    test_source = ui.context["source"]
+    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
+
+    ui.llm.set_responses([
+        SQLQueryWithTables(
+            query="SELECT SUM(value) as value_sum FROM test_table",
+            table_slug="test_sql_agg",
+            tables=["test_table"]
+        ),
+        "Summary of results",
+    ])
+
+    sql_agent = SQLAgent(llm=ui.llm)
+    chat_agent = ChatAgent(llm=ui.llm)
+    plan = Plan(
+        ActorTask(sql_agent),
+        ActorTask(chat_agent),
+        history=[{"content": "Sum value column", "role": "user"}],
+        title="Sum of value",
+        context=ui.context
+    )
+    await ui._execute_plan(plan)
+    await asyncio.sleep(0.3)
+
+    last_msg = ui.interface.objects[-1]
+    followup_icons = [
+        f for f in _get_footer_actions(last_msg)
+        if getattr(f, 'name', '') == "FollowUp"
+    ]
+    assert len(followup_icons) == 0
+
+
+async def test_followup_suggestion_not_on_error(explorer_ui_with_error):
+    """Follow-up suggestion should NOT appear when plan errors."""
+    ui = explorer_ui_with_error
+    await asyncio.sleep(0.3)
+
+    last_msg = ui.interface.objects[-1]
+    followup_icons = [
+        f for f in _get_footer_actions(last_msg)
+        if getattr(f, 'name', '') == "FollowUp"
+    ]
+    assert len(followup_icons) == 0
+
+
+async def test_followup_suggestion_llm_failure_silent(explorer_ui):
+    """If the LLM call for suggestions fails, it should fail silently."""
+    ui = explorer_ui
+    test_source = ui.context["source"]
+    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
+
+    def raise_error():
+        raise RuntimeError("LLM unavailable")
+
+    ui.llm.set_responses([
+        SQLQueryWithTables(
+            query="SELECT SUM(value) as value_sum FROM test_table",
+            table_slug="test_sql_agg",
+            tables=["test_table"]
+        ),
+        "Summary of results",
+        raise_error,
+    ])
+
+    sql_agent = SQLAgent(llm=ui.llm)
+    chat_agent = ChatAgent(llm=ui.llm)
+    plan = Plan(
+        ActorTask(sql_agent),
+        ActorTask(chat_agent),
+        history=[{"content": "Sum value column", "role": "user"}],
+        title="Sum of value",
+        context=ui.context
+    )
+    await ui._execute_plan(plan)
+    await asyncio.sleep(0.3)
+
+    # Should not crash, no followup suggestion, but plan succeeded
+    assert len(ui._explorations.items) == 2
+    last_msg = ui.interface.objects[-1]
+    followup_icons = [
+        f for f in _get_footer_actions(last_msg)
+        if getattr(f, 'name', '') == "FollowUp"
+    ]
+    assert len(followup_icons) == 0
+
+
+async def test_followup_suggestion_click_populates_input(explorer_ui):
+    """Clicking a follow-up suggestion button should set chat input and hide the button."""
+    ui = explorer_ui
+    test_source = ui.context["source"]
+    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
+
+    suggestion_response = FollowUpSuggestion(query="Show value by category")
+
+    ui.llm.set_responses([
+        SQLQueryWithTables(
+            query="SELECT SUM(value) as value_sum FROM test_table",
+            table_slug="test_sql_agg",
+            tables=["test_table"]
+        ),
+        "Summary of results",
+        suggestion_response,
+    ])
+
+    sql_agent = SQLAgent(llm=ui.llm)
+    chat_agent = ChatAgent(llm=ui.llm)
+    plan = Plan(
+        ActorTask(sql_agent),
+        ActorTask(chat_agent),
+        history=[{"content": "Sum value column", "role": "user"}],
+        title="Sum of value",
+        context=ui.context
+    )
+    await ui._execute_plan(plan)
+    await asyncio.sleep(0.5)
+
+    # Get the follow-up icon button
+    last_msg = ui.interface.objects[-1]
+    followup_icons = [
+        f for f in _get_footer_actions(last_msg)
+        if getattr(f, 'name', '') == "FollowUp"
+    ]
+    assert len(followup_icons) == 1
+
+    # Simulate clicking the icon button -- should populate input and hide button
+    followup_icons[0].clicks += 1
+    await asyncio.sleep(0.3)
+
+    assert ui._chat_input.value_input == "Show value by category"
+    assert not followup_icons[0].visible
+
+
+async def test_followup_suggestion_stale_message_guard(explorer_ui):
+    """If a new message arrives while suggestion is generating,
+    it should NOT be appended to the (now-stale) message."""
+    ui = explorer_ui
+    test_source = ui.context["source"]
+    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
+
+    suggestion_response = FollowUpSuggestion(query="Show value by category")
+
+    ui.llm.set_responses([
+        SQLQueryWithTables(
+            query="SELECT SUM(value) as value_sum FROM test_table",
+            table_slug="test_sql_agg",
+            tables=["test_table"]
+        ),
+        "Summary of results",
+        suggestion_response,
+    ])
+
+    sql_agent = SQLAgent(llm=ui.llm)
+    chat_agent = ChatAgent(llm=ui.llm)
+    plan = Plan(
+        ActorTask(sql_agent),
+        ActorTask(chat_agent),
+        history=[{"content": "Sum value column", "role": "user"}],
+        title="Sum of value",
+        context=ui.context
+    )
+
+    # Monkey-patch llm.invoke to add a message AFTER num_objects is captured
+    # but BEFORE the suggestion is appended
+    original_invoke = ui.llm.invoke
+    async def invoke_with_extra_message(*args, **kwargs):
+        result = await original_invoke(*args, **kwargs)
+        # Only intercept the suggestions call (FollowUpSuggestion response_model)
+        if kwargs.get("response_model") is FollowUpSuggestion:
+            ui.interface.send("New question arrived", user="User", respond=False)
+        return result
+    ui.llm.invoke = invoke_with_extra_message
+
+    await ui._execute_plan(plan)
+    await asyncio.sleep(0.5)
+
+    # The extra message means num_objects changed, so no FollowUp should appear
+    footer_with_followup = []
+    for msg in ui.interface.objects:
+        footer_with_followup.extend(
+            f for f in _get_footer_actions(msg)
+            if getattr(f, 'name', '') == "FollowUp"
+        )
+    assert len(footer_with_followup) == 0
+    ui.llm.invoke = original_invoke
+
+
+async def test_followup_icon_appears_in_footer(explorer_ui):
+    """Follow-up icon should appear in footer after successful query."""
+    ui = explorer_ui
+    test_source = ui.context["source"]
+    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
+
+    ui.llm.set_responses([
+        SQLQueryWithTables(
+            query="SELECT SUM(value) as value_sum FROM test_table",
+            table_slug="test_sql_agg",
+            tables=["test_table"]
+        ),
+        "Summary of results",
+        FollowUpSuggestion(query="test query"),
+    ])
+
+    sql_agent = SQLAgent(llm=ui.llm)
+    chat_agent = ChatAgent(llm=ui.llm)
+    plan = Plan(
+        ActorTask(sql_agent),
+        ActorTask(chat_agent),
+        history=[{"content": "Sum value column", "role": "user"}],
+        title="Sum of value",
+        context=ui.context
+    )
+    await ui._execute_plan(plan)
+    await asyncio.sleep(0.5)
+
+    last_msg = ui.interface.objects[-1]
+    followup_icons = [
+        f for f in _get_footer_actions(last_msg)
+        if getattr(f, 'name', '') == "FollowUp"
+    ]
+    assert len(followup_icons) == 1

From 869fcb0f0ad80bffea27d2a30e7c12614c73c8af Mon Sep 17 00:00:00 2001
From: ghostiee-11 <168410465+ghostiee-11@users.noreply.github.com>
Date: Fri, 27 Mar 2026 01:31:43 +0530
Subject: [PATCH 2/7] Address review: on-demand follow-up, remove toggle, fix
 naming

- Follow-up suggestion is no longer auto-generated after every query.
  The lightbulb icon always appears and triggers the LLM call on click.
- Remove follow_up_suggestions param and Settings toggle. Toggles are
  reserved for auto-generating features (Validation, SQL Planning).
- Fix naming: suggest_followup -> suggest_follow_up, followup_button ->
  follow_up_button for consistency. is_followup on Plan is unchanged
  (separate exploration feature).
- Remove _background_tasks (no longer needed).
- Update tests and docs accordingly.
---
 docs/configuration/ui.md                  |  12 +-
 docs/getting_started/navigating_the_ui.md |   2 -
 lumen/ai/coordinator/base.py              |   2 +-
 lumen/ai/ui.py                            |  76 ++++----
 lumen/tests/ai/test_coordinator.py        |  18 +-
 lumen/tests/ai/test_ui.py                 | 213 +++-------------------
 6 files changed, 68 insertions(+), 255 deletions(-)

diff --git a/docs/configuration/ui.md b/docs/configuration/ui.md
index 5f3653101..713e45f4d 100644
--- a/docs/configuration/ui.md
+++ b/docs/configuration/ui.md
@@ -80,16 +80,7 @@ ui = lmai.ExplorerUI(
 
 ### Follow-up suggestions
 
-After each successful query, a lightbulb icon appears in the message footer. Clicking it populates the chat input with an AI-generated follow-up suggestion that references actual column names from the data.
-
-``` py title="Disable follow-up suggestions"
-ui = lmai.ExplorerUI(
-    data='penguins.csv',
-    follow_up_suggestions=False
-)
-```
-
-Users can also toggle this at runtime via **Settings > Follow-Up Suggestions**.
+After each successful query, a lightbulb icon appears in the message footer. Click it to generate an AI-powered follow-up question that references actual column names from your data. The suggestion populates the chat input so you can review or edit before sending.
 
 ## Advanced parameters
 
@@ -216,7 +207,6 @@ Quick reference:
 | `notebook_preamble` | str | Export header |
 | `provider_choices` | dict | LLM providers shown in Settings |
 | `source_controls` | list | Source control components for data |
-| `follow_up_suggestions` | bool | AI follow-up suggestion icon after queries (default: True) |
 | `suggestions` | list | Quick action buttons |
 | `title` | str | App title |
 | `tools` | list | Custom tools |
diff --git a/docs/getting_started/navigating_the_ui.md b/docs/getting_started/navigating_the_ui.md
index 408b85599..004bf5818 100644
--- a/docs/getting_started/navigating_the_ui.md
+++ b/docs/getting_started/navigating_the_ui.md
@@ -110,8 +110,6 @@ Use **Settings** in the left sidebar to configure analysis behavior:
 
 **Validation Step** — AI validates results for correctness (enabled by default)
 
-**Follow-Up Suggestions** — After each successful query, a lightbulb icon appears in the message footer. Click it to populate the chat input with an AI-generated follow-up question. Toggle off in Settings if not needed.
-
 **Code Execution** — If enabled by administrator, controls how visualizations are generated. See [Code Execution for Visualizations](using_lumen_ai.md#code-execution-for-visualizations) for security implications.
 
 **LLM Configuration** — Choose and configure your language model
diff --git a/lumen/ai/coordinator/base.py b/lumen/ai/coordinator/base.py
index 50eb332ec..1ceebb631 100644
--- a/lumen/ai/coordinator/base.py
+++ b/lumen/ai/coordinator/base.py
@@ -514,7 +514,7 @@ async def _check_tool_relevance(self, tool: Tool, tool_output: str, actor: Actor
 
         return result.yes
 
-    async def suggest_followup(self, plan: Plan) -> str | None:
+    async def suggest_follow_up(self, plan: Plan) -> str | None:
         """Generate a single follow-up suggestion from a completed plan.
 
         Returns a suggestion query string, or None on failure.
diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py
index 3097995bf..e58c6a979 100644
--- a/lumen/ai/ui.py
+++ b/lumen/ai/ui.py
@@ -410,12 +410,6 @@ class UI(Viewer):
     export_functions = param.Dict(default={}, doc="""
        Dictionary mapping from name of exporter to export function.""")
 
-    follow_up_suggestions = param.Boolean(default=True, doc="""
-        Whether to generate AI-powered follow-up suggestions after each
-        successful query. When enabled, a lightbulb icon button appears
-        in the message footer that populates the input with a follow-up
-        suggestion on click. Requires an additional LLM call per query.""")
-
     interface = param.ClassSelector(class_=ChatFeed, doc="""
         The interface for the Coordinator to interact with.""")
 
@@ -1048,7 +1042,6 @@ def _configure_coordinator(self):
         if self.analyses:
             agents.append(AnalysisAgent(analyses=self.analyses))
 
-        self._background_tasks = set()
         self._coordinator = self.coordinator(
             agents=agents,
             context=self.context,
@@ -1603,29 +1596,43 @@ async def _add_analysis_suggestions(self, plan: Plan):
             num_objects=len(self.interface.objects),
         )
 
-    def _add_followup_icon_to_footer(self, suggestion: str, num_objects: int):
-        """Add a lightbulb icon button to the message footer actions.
+    def _add_follow_up_icon(self, plan: Plan):
+        """Add follow-up suggestion icon to the last message footer.
 
-        On click, populates the ChatAreaInput with the suggestion
-        without auto-submitting, so the user can review or edit before sending.
+        The icon appears immediately. On click, it calls the coordinator
+        to generate a suggestion and populates the chat input.
         """
+        if not plan.out_context.get("pipeline"):
+            return
+        if not plan.out_context.get("data"):
+            return
 
-        def on_click(event):
-            with edit_readonly(self._chat_input):
-                self._chat_input.value_input = suggestion
-            followup_button.visible = False
+        num_objects = len(self.interface.objects)
+
+        async def _generate_suggestion(_=None):
+            follow_up_button.disabled = True
+            follow_up_button.description = "Generating suggestion..."
+            try:
+                suggestion = await self._coordinator.suggest_follow_up(plan)
+                if suggestion:
+                    with edit_readonly(self._chat_input):
+                        self._chat_input.value_input = suggestion
+                follow_up_button.visible = False
+            except Exception:
+                follow_up_button.disabled = False
+                follow_up_button.description = "Click to retry"
 
-        def hide_followup(_=None):
+        def hide_follow_up(_=None):
             if len(self.interface.objects) > num_objects:
-                followup_button.visible = False
+                follow_up_button.visible = False
 
-        followup_button = IconButton(
+        follow_up_button = IconButton(
             icon="lightbulb",
-            description=f"Suggest: {suggestion}",
+            description="Suggest a follow-up question",
             size="small",
             icon_size="0.9em",
             margin=(5, 0),
-            on_click=on_click,
+            on_click=lambda _: state.execute(_generate_suggestion),
             name="FollowUp",
             disabled=self.interface.param.loading,
             color="default",
@@ -1643,26 +1650,10 @@ def hide_followup(_=None):
                 if getattr(obj, 'name', '') != "FollowUp"
             ]
             try:
-                message.footer_actions = existing + [followup_button]
+                message.footer_actions = existing + [follow_up_button]
             except AttributeError:
-                message.footer_objects = existing + [followup_button]
-            self.interface.param.watch(hide_followup, "objects")
-
-    async def _generate_followup_suggestions(self, plan: Plan):
-        """Render follow-up suggestion from the coordinator."""
-        if not self.follow_up_suggestions:
-            return
-
-        num_objects = len(self.interface.objects)
-        suggestion = await self._coordinator.suggest_followup(plan)
-        if not suggestion:
-            return
-
-        # Guard against stale message if user sent a new query while we were generating
-        if len(self.interface.objects) != num_objects:
-            return
-
-        self._add_followup_icon_to_footer(suggestion, num_objects)
+                message.footer_objects = existing + [follow_up_button]
+            self.interface.param.watch(hide_follow_up, "objects")
 
     def __panel__(self):
         return self._main
@@ -1867,9 +1858,6 @@ def _render_sidebar(self) -> list[Viewable]:
         validation = Switch(label='Validation Step', description='Check if the response fully answered your question', value=True)
         self._coordinator.validation_enabled = validation
         switches.append(validation)
-        followup = Switch(label='Follow-Up Suggestions', description='Show AI-generated follow-up questions after each query', value=self.follow_up_suggestions)
-        followup.link(self, value='follow_up_suggestions', bidirectional=True)
-        switches.append(followup)
 
         # Add code execution selector if not hidden
         code_agents = False
@@ -2577,9 +2565,7 @@ async def _postprocess_exploration(
             await self._sync_sources(SimpleNamespace(new=plan.out_context), global_context=plan.out_context)
             if "pipeline" in plan.out_context:
                 await self._add_analysis_suggestions(plan)
-                task = asyncio.create_task(self._generate_followup_suggestions(plan))
-                self._background_tasks.add(task)
-                task.add_done_callback(self._background_tasks.discard)
+                self._add_follow_up_icon(plan)
 
             if is_new:
                 plan.param.watch(partial(self._update_views, exploration), "views")
diff --git a/lumen/tests/ai/test_coordinator.py b/lumen/tests/ai/test_coordinator.py
index b4d29930c..f7145e3f1 100644
--- a/lumen/tests/ai/test_coordinator.py
+++ b/lumen/tests/ai/test_coordinator.py
@@ -591,8 +591,8 @@ async def test_process_files_duplicate_table_name():
 
 # --- Follow-up suggestion coordinator tests ---
 
-async def test_suggest_followup_returns_suggestion(llm):
-    """Coordinator.suggest_followup should return a query string."""
+async def test_suggest_follow_up_returns_suggestion(llm):
+    """Coordinator.suggest_follow_up should return a query string."""
     suggestion_response = FollowUpSuggestion(query="Show value by category")
     llm.set_responses([suggestion_response])
 
@@ -609,11 +609,11 @@ async def test_suggest_followup_returns_suggestion(llm):
         "sql": "SELECT * FROM t",
     }
 
-    result = await coordinator.suggest_followup(plan)
+    result = await coordinator.suggest_follow_up(plan)
     assert result == "Show value by category"
 
 
-async def test_suggest_followup_no_pipeline(llm):
+async def test_suggest_follow_up_no_pipeline(llm):
     """Should return None when plan has no pipeline."""
     coordinator = Coordinator(llm=llm)
     plan = Plan(
@@ -624,11 +624,11 @@ async def test_suggest_followup_no_pipeline(llm):
     )
     plan.out_context = {}
 
-    result = await coordinator.suggest_followup(plan)
+    result = await coordinator.suggest_follow_up(plan)
     assert result is None
 
 
-async def test_suggest_followup_no_data(llm):
+async def test_suggest_follow_up_no_data(llm):
     """Should return None when plan has no data summary."""
     coordinator = Coordinator(llm=llm)
     plan = Plan(
@@ -639,11 +639,11 @@ async def test_suggest_followup_no_data(llm):
     )
     plan.out_context = {"pipeline": object(), "data": ""}
 
-    result = await coordinator.suggest_followup(plan)
+    result = await coordinator.suggest_follow_up(plan)
     assert result is None
 
 
-async def test_suggest_followup_llm_failure(llm):
+async def test_suggest_follow_up_llm_failure(llm):
     """Should return None when LLM call fails."""
     def raise_error():
         raise RuntimeError("LLM unavailable")
@@ -663,5 +663,5 @@ def raise_error():
         "sql": "SELECT * FROM t",
     }
 
-    result = await coordinator.suggest_followup(plan)
+    result = await coordinator.suggest_follow_up(plan)
     assert result is None
diff --git a/lumen/tests/ai/test_ui.py b/lumen/tests/ai/test_ui.py
index 790eced19..ad6e9ff52 100644
--- a/lumen/tests/ai/test_ui.py
+++ b/lumen/tests/ai/test_ui.py
@@ -15,6 +15,7 @@
 except ModuleNotFoundError:
     pytest.skip("lumen.ai could not be imported, skipping tests.", allow_module_level=True)
 
+from panel.io import state
 from panel.layout import Column, Row
 from panel.tests.util import async_wait_until
 from panel.util import edit_readonly
@@ -1464,7 +1465,6 @@ async def test_edit_on_second_exploration_preserves_first(explorer_ui):
     second one, only the second exploration (sql_2 + spec_2) should be
     removed. The first exploration (sql_1 + spec_1) must remain intact."""
     ui = explorer_ui
-    ui.follow_up_suggestions = False
     test_source = ui.context["source"]
     SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
 
@@ -1623,66 +1623,24 @@ def _get_footer_actions(msg):
         return msg.footer_objects or []
 
 
-def test_followup_suggestion_model():
+def test_follow_up_suggestion_model():
     """FollowUpSuggestion model should validate correctly."""
     result = FollowUpSuggestion(query="Show value by category")
     assert result.query == "Show value by category"
 
 
-def test_followup_suggestion_model_max_length():
+def test_follow_up_suggestion_model_max_length():
     """FollowUpSuggestion should reject queries over 100 characters."""
     with pytest.raises(Exception):
         FollowUpSuggestion(query="x" * 101)
 
 
-async def test_followup_suggestion_after_query(explorer_ui):
-    """After a successful SQL query, a follow-up suggestion icon should appear in footer."""
+async def test_follow_up_icon_after_query(explorer_ui):
+    """After a successful SQL query, a follow-up icon should appear in footer (no LLM call)."""
     ui = explorer_ui
     test_source = ui.context["source"]
     SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
 
-    suggestion_response = FollowUpSuggestion(query="Show value by category")
-
-    ui.llm.set_responses([
-        SQLQueryWithTables(
-            query="SELECT SUM(value) as value_sum FROM test_table",
-            table_slug="test_sql_agg",
-            tables=["test_table"]
-        ),
-        "Summary of results",
-        suggestion_response,
-    ])
-
-    sql_agent = SQLAgent(llm=ui.llm)
-    chat_agent = ChatAgent(llm=ui.llm)
-    plan = Plan(
-        ActorTask(sql_agent),
-        ActorTask(chat_agent),
-        history=[{"content": "Sum value column", "role": "user"}],
-        title="Sum of value",
-        context=ui.context
-    )
-    await ui._execute_plan(plan)
-
-    # Wait for async follow-up suggestion task to complete
-    await asyncio.sleep(0.5)
-
-    last_msg = ui.interface.objects[-1]
-    followup_icons = [
-        f for f in _get_footer_actions(last_msg)
-        if getattr(f, 'name', '') == "FollowUp"
-    ]
-    assert len(followup_icons) == 1
-    assert followup_icons[0].icon == "lightbulb"
-
-
-async def test_followup_suggestion_disabled(explorer_ui):
-    """No suggestion should appear when follow_up_suggestions=False."""
-    ui = explorer_ui
-    ui.follow_up_suggestions = False
-    test_source = ui.context["source"]
-    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
-
     ui.llm.set_responses([
         SQLQueryWithTables(
             query="SELECT SUM(value) as value_sum FROM test_table",
@@ -1702,172 +1660,53 @@ async def test_followup_suggestion_disabled(explorer_ui):
         context=ui.context
     )
     await ui._execute_plan(plan)
-    await asyncio.sleep(0.3)
 
     last_msg = ui.interface.objects[-1]
-    followup_icons = [
+    follow_up_icons = [
         f for f in _get_footer_actions(last_msg)
         if getattr(f, 'name', '') == "FollowUp"
     ]
-    assert len(followup_icons) == 0
+    assert len(follow_up_icons) == 1
+    assert follow_up_icons[0].icon == "lightbulb"
 
 
-async def test_followup_suggestion_not_on_error(explorer_ui_with_error):
-    """Follow-up suggestion should NOT appear when plan errors."""
+async def test_follow_up_icon_not_on_error(explorer_ui_with_error):
+    """Follow-up icon should NOT appear when plan errors."""
     ui = explorer_ui_with_error
     await asyncio.sleep(0.3)
 
     last_msg = ui.interface.objects[-1]
-    followup_icons = [
-        f for f in _get_footer_actions(last_msg)
-        if getattr(f, 'name', '') == "FollowUp"
-    ]
-    assert len(followup_icons) == 0
-
-
-async def test_followup_suggestion_llm_failure_silent(explorer_ui):
-    """If the LLM call for suggestions fails, it should fail silently."""
-    ui = explorer_ui
-    test_source = ui.context["source"]
-    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
-
-    def raise_error():
-        raise RuntimeError("LLM unavailable")
-
-    ui.llm.set_responses([
-        SQLQueryWithTables(
-            query="SELECT SUM(value) as value_sum FROM test_table",
-            table_slug="test_sql_agg",
-            tables=["test_table"]
-        ),
-        "Summary of results",
-        raise_error,
-    ])
-
-    sql_agent = SQLAgent(llm=ui.llm)
-    chat_agent = ChatAgent(llm=ui.llm)
-    plan = Plan(
-        ActorTask(sql_agent),
-        ActorTask(chat_agent),
-        history=[{"content": "Sum value column", "role": "user"}],
-        title="Sum of value",
-        context=ui.context
-    )
-    await ui._execute_plan(plan)
-    await asyncio.sleep(0.3)
-
-    # Should not crash, no followup suggestion, but plan succeeded
-    assert len(ui._explorations.items) == 2
-    last_msg = ui.interface.objects[-1]
-    followup_icons = [
+    follow_up_icons = [
         f for f in _get_footer_actions(last_msg)
         if getattr(f, 'name', '') == "FollowUp"
     ]
-    assert len(followup_icons) == 0
+    assert len(follow_up_icons) == 0
 
 
-async def test_followup_suggestion_click_populates_input(explorer_ui):
-    """Clicking a follow-up suggestion button should set chat input and hide the button."""
+async def test_follow_up_icon_not_shown_without_pipeline(explorer_ui):
+    """Follow-up icon should NOT appear when plan has no pipeline in context."""
     ui = explorer_ui
-    test_source = ui.context["source"]
-    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
 
-    suggestion_response = FollowUpSuggestion(query="Show value by category")
-
-    ui.llm.set_responses([
-        SQLQueryWithTables(
-            query="SELECT SUM(value) as value_sum FROM test_table",
-            table_slug="test_sql_agg",
-            tables=["test_table"]
-        ),
-        "Summary of results",
-        suggestion_response,
-    ])
-
-    sql_agent = SQLAgent(llm=ui.llm)
     chat_agent = ChatAgent(llm=ui.llm)
+    ui.llm.set_responses(["Here is my answer"])
     plan = Plan(
-        ActorTask(sql_agent),
         ActorTask(chat_agent),
-        history=[{"content": "Sum value column", "role": "user"}],
-        title="Sum of value",
-        context=ui.context
+        history=[{"content": "What is lumen?", "role": "user"}],
+        title="Chat only",
+        context=ui.context,
     )
     await ui._execute_plan(plan)
-    await asyncio.sleep(0.5)
 
-    # Get the follow-up icon button
     last_msg = ui.interface.objects[-1]
-    followup_icons = [
+    follow_up_icons = [
         f for f in _get_footer_actions(last_msg)
         if getattr(f, 'name', '') == "FollowUp"
     ]
-    assert len(followup_icons) == 1
-
-    # Simulate clicking the icon button -- should populate input and hide button
-    followup_icons[0].clicks += 1
-    await asyncio.sleep(0.3)
-
-    assert ui._chat_input.value_input == "Show value by category"
-    assert not followup_icons[0].visible
-
-
-async def test_followup_suggestion_stale_message_guard(explorer_ui):
-    """If a new message arrives while suggestion is generating,
-    it should NOT be appended to the (now-stale) message."""
-    ui = explorer_ui
-    test_source = ui.context["source"]
-    SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
-
-    suggestion_response = FollowUpSuggestion(query="Show value by category")
-
-    ui.llm.set_responses([
-        SQLQueryWithTables(
-            query="SELECT SUM(value) as value_sum FROM test_table",
-            table_slug="test_sql_agg",
-            tables=["test_table"]
-        ),
-        "Summary of results",
-        suggestion_response,
-    ])
-
-    sql_agent = SQLAgent(llm=ui.llm)
-    chat_agent = ChatAgent(llm=ui.llm)
-    plan = Plan(
-        ActorTask(sql_agent),
-        ActorTask(chat_agent),
-        history=[{"content": "Sum value column", "role": "user"}],
-        title="Sum of value",
-        context=ui.context
-    )
-
-    # Monkey-patch llm.invoke to add a message AFTER num_objects is captured
-    # but BEFORE the suggestion is appended
-    original_invoke = ui.llm.invoke
-    async def invoke_with_extra_message(*args, **kwargs):
-        result = await original_invoke(*args, **kwargs)
-        # Only intercept the suggestions call (FollowUpSuggestion response_model)
-        if kwargs.get("response_model") is FollowUpSuggestion:
-            ui.interface.send("New question arrived", user="User", respond=False)
-        return result
-    ui.llm.invoke = invoke_with_extra_message
-
-    await ui._execute_plan(plan)
-    await asyncio.sleep(0.5)
-
-    # The extra message means num_objects changed, so no FollowUp should appear
-    footer_with_followup = []
-    for msg in ui.interface.objects:
-        footer_with_followup.extend(
-            f for f in _get_footer_actions(msg)
-            if getattr(f, 'name', '') == "FollowUp"
-        )
-    assert len(footer_with_followup) == 0
-    ui.llm.invoke = original_invoke
+    assert len(follow_up_icons) == 0
 
 
-async def test_followup_icon_appears_in_footer(explorer_ui):
-    """Follow-up icon should appear in footer after successful query."""
+async def test_follow_up_icon_description(explorer_ui):
+    """Follow-up icon should have the correct description."""
     ui = explorer_ui
     test_source = ui.context["source"]
     SQLQueryWithTables = make_sql_model([(test_source.name, "test_table")])
@@ -1879,7 +1718,6 @@ async def test_followup_icon_appears_in_footer(explorer_ui):
             tables=["test_table"]
         ),
         "Summary of results",
-        FollowUpSuggestion(query="test query"),
     ])
 
     sql_agent = SQLAgent(llm=ui.llm)
@@ -1892,11 +1730,12 @@ async def test_followup_icon_appears_in_footer(explorer_ui):
         context=ui.context
     )
     await ui._execute_plan(plan)
-    await asyncio.sleep(0.5)
 
     last_msg = ui.interface.objects[-1]
-    followup_icons = [
+    follow_up_icons = [
         f for f in _get_footer_actions(last_msg)
         if getattr(f, 'name', '') == "FollowUp"
     ]
-    assert len(followup_icons) == 1
+    assert len(follow_up_icons) == 1
+    assert follow_up_icons[0].description == "Suggest a follow-up question"
+    assert follow_up_icons[0].icon == "lightbulb"

From e84ae85a835dcc00aae04b82c8e9c8d528f20f9b Mon Sep 17 00:00:00 2001
From: ghostiee-11 <168410465+ghostiee-11@users.noreply.github.com>
Date: Fri, 27 Mar 2026 04:52:59 +0530
Subject: [PATCH 3/7] Address review: notify on failure, fix async pattern,
 keep icon visible

- Replace silent except with state.notifications.warning in coordinator
- Rename _generate_suggestion to _generate_follow_up, move under signature
- Keep state.execute() for async on_click (matches replan/rerun pattern)
- Remove hide logic so button stays visible for repeated use
- Remove bare except Exception: pass, use try/finally instead
---
 lumen/ai/coordinator/base.py |  7 ++++++-
 lumen/ai/ui.py               | 23 ++++++++---------------
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/lumen/ai/coordinator/base.py b/lumen/ai/coordinator/base.py
index 1ceebb631..ce5a0a759 100644
--- a/lumen/ai/coordinator/base.py
+++ b/lumen/ai/coordinator/base.py
@@ -545,7 +545,12 @@ async def suggest_follow_up(self, plan: Plan) -> str | None:
                 previous_queries=previous_queries,
             )
             return result.query
-        except Exception:
+        except Exception as e:
+            if state.notifications:
+                state.notifications.warning(
+                    f"Could not generate follow-up suggestion: {e}",
+                    duration=3000,
+                )
             return None
 
     def __panel__(self):
diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py
index e58c6a979..43fe19324 100644
--- a/lumen/ai/ui.py
+++ b/lumen/ai/ui.py
@@ -1602,14 +1602,8 @@ def _add_follow_up_icon(self, plan: Plan):
         The icon appears immediately. On click, it calls the coordinator
         to generate a suggestion and populates the chat input.
         """
-        if not plan.out_context.get("pipeline"):
-            return
-        if not plan.out_context.get("data"):
-            return
-
-        num_objects = len(self.interface.objects)
 
-        async def _generate_suggestion(_=None):
+        async def _generate_follow_up():
             follow_up_button.disabled = True
             follow_up_button.description = "Generating suggestion..."
             try:
@@ -1617,14 +1611,14 @@ async def _generate_suggestion(_=None):
                 if suggestion:
                     with edit_readonly(self._chat_input):
                         self._chat_input.value_input = suggestion
-                follow_up_button.visible = False
-            except Exception:
+            finally:
                 follow_up_button.disabled = False
-                follow_up_button.description = "Click to retry"
+                follow_up_button.description = "Suggest a follow-up question"
 
-        def hide_follow_up(_=None):
-            if len(self.interface.objects) > num_objects:
-                follow_up_button.visible = False
+        if not plan.out_context.get("pipeline"):
+            return
+        if not plan.out_context.get("data"):
+            return
 
         follow_up_button = IconButton(
             icon="lightbulb",
@@ -1632,7 +1626,7 @@ def hide_follow_up(_=None):
             size="small",
             icon_size="0.9em",
             margin=(5, 0),
-            on_click=lambda _: state.execute(_generate_suggestion),
+            on_click=lambda _: state.execute(_generate_follow_up),
             name="FollowUp",
             disabled=self.interface.param.loading,
             color="default",
@@ -1653,7 +1647,6 @@ def hide_follow_up(_=None):
                 message.footer_actions = existing + [follow_up_button]
             except AttributeError:
                 message.footer_objects = existing + [follow_up_button]
-            self.interface.param.watch(hide_follow_up, "objects")
 
     def __panel__(self):
         return self._main

From d91b685b58b9a1760f19a54c88423773add66a15 Mon Sep 17 00:00:00 2001
From: ghostiee-11 <168410465+ghostiee-11@users.noreply.github.com>
Date: Sat, 28 Mar 2026 00:17:05 +0530
Subject: [PATCH 4/7] Address review: rename follow_up key, simplify
 footer_actions, bump pin

- Rename prompt key follow_up_suggestions -> follow_up
- Use self._chat_input.value instead of edit_readonly + value_input
- Remove try/except fallback for footer_actions, assume it exists
- Bump panel-material-ui pin to >=0.9.0 (has footer_actions)
---
 lumen/ai/coordinator/base.py |  4 ++--
 lumen/ai/ui.py               | 14 +++-----------
 pixi.toml                    |  2 +-
 pyproject.toml               |  2 +-
 4 files changed, 7 insertions(+), 15 deletions(-)

diff --git a/lumen/ai/coordinator/base.py b/lumen/ai/coordinator/base.py
index ce5a0a759..37ef9a0e5 100644
--- a/lumen/ai/coordinator/base.py
+++ b/lumen/ai/coordinator/base.py
@@ -284,7 +284,7 @@ class Coordinator(Viewer, VectorLookupToolUser):
                 "template": PROMPTS_DIR / "Coordinator" / "tool_relevance.jinja2",
                 "response_model": ThinkingYesNo,
             },
-            "follow_up_suggestions": {
+            "follow_up": {
                 "template": PROMPTS_DIR / "Coordinator" / "follow_up_suggestions.jinja2",
                 "response_model": FollowUpSuggestion,
                 "llm_spec": "ui",
@@ -536,7 +536,7 @@ async def suggest_follow_up(self, plan: Plan) -> str | None:
 
         try:
             result = await self._invoke_prompt(
-                "follow_up_suggestions",
+                "follow_up",
                 messages=[{"role": "user", "content": "Generate a follow-up suggestion."}],
                 context=plan.out_context,
                 data_summary=data_summary,
diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py
index 43fe19324..01b051357 100644
--- a/lumen/ai/ui.py
+++ b/lumen/ai/ui.py
@@ -1609,8 +1609,7 @@ async def _generate_follow_up():
             try:
                 suggestion = await self._coordinator.suggest_follow_up(plan)
                 if suggestion:
-                    with edit_readonly(self._chat_input):
-                        self._chat_input.value_input = suggestion
+                    self._chat_input.value = suggestion
             finally:
                 follow_up_button.disabled = False
                 follow_up_button.description = "Suggest a follow-up question"
@@ -1635,18 +1634,11 @@ async def _generate_follow_up():
 
         if len(self.interface):
             message = self.interface.objects[-1]
-            try:
-                existing = message.footer_actions or []
-            except AttributeError:
-                existing = message.footer_objects or []
             existing = [
-                obj for obj in existing
+                obj for obj in (message.footer_actions or [])
                 if getattr(obj, 'name', '') != "FollowUp"
             ]
-            try:
-                message.footer_actions = existing + [follow_up_button]
-            except AttributeError:
-                message.footer_objects = existing + [follow_up_button]
+            message.footer_actions = existing + [follow_up_button]
 
     def __panel__(self):
         return self._main
diff --git a/pixi.toml b/pixi.toml
index 8ee0ae7e8..ab5490b1b 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -20,7 +20,7 @@ lint = ["py312", "lint"]
 [dependencies]
 bokeh = "*"
 holoviews = ">=1.17.0"
-panel-material-ui = ">=0.8.1"
+panel-material-ui = ">=0.9.0"
 hvplot = "*"
 intake = "<2"
 jinja2 = ">3.0"
diff --git a/pyproject.toml b/pyproject.toml
index ead041c71..f1960e42e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,7 +42,7 @@ dependencies = [
     "pandas",
     "panel >=1.7.5",
     "panel-graphic-walker[kernel] >=0.6.4",
-    "panel-material-ui >=0.8.1",
+    "panel-material-ui >=0.9.0",
     "panel-splitjs >=0.3.2",
     "param >=2.2.1",
     "pyarrow",

From 17e88cfd9cb4633b4c7c0f1867e8da1d65eb2686 Mon Sep 17 00:00:00 2001
From: ghostiee-11 <168410465+ghostiee-11@users.noreply.github.com>
Date: Sat, 28 Mar 2026 02:57:35 +0530
Subject: [PATCH 5/7] Fix CI: fallback to footer_objects until footer_actions
 is released

footer_actions is not yet in any released panel-material-ui version.
Use hasattr to detect it and fall back to footer_objects for now.
Revert pin bump to >=0.8.1.
---
 lumen/ai/ui.py | 5 +++--
 pixi.toml      | 2 +-
 pyproject.toml | 2 +-
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py
index 01b051357..3bfa3a926 100644
--- a/lumen/ai/ui.py
+++ b/lumen/ai/ui.py
@@ -1634,11 +1634,12 @@ async def _generate_follow_up():
 
         if len(self.interface):
             message = self.interface.objects[-1]
+            attr = "footer_actions" if hasattr(message, "footer_actions") else "footer_objects"
             existing = [
-                obj for obj in (message.footer_actions or [])
+                obj for obj in (getattr(message, attr) or [])
                 if getattr(obj, 'name', '') != "FollowUp"
             ]
-            message.footer_actions = existing + [follow_up_button]
+            setattr(message, attr, existing + [follow_up_button])
 
     def __panel__(self):
         return self._main
diff --git a/pixi.toml b/pixi.toml
index ab5490b1b..8ee0ae7e8 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -20,7 +20,7 @@ lint = ["py312", "lint"]
 [dependencies]
 bokeh = "*"
 holoviews = ">=1.17.0"
-panel-material-ui = ">=0.9.0"
+panel-material-ui = ">=0.8.1"
 hvplot = "*"
 intake = "<2"
 jinja2 = ">3.0"
diff --git a/pyproject.toml b/pyproject.toml
index f1960e42e..ead041c71 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,7 +42,7 @@ dependencies = [
     "pandas",
     "panel >=1.7.5",
     "panel-graphic-walker[kernel] >=0.6.4",
-    "panel-material-ui >=0.9.0",
+    "panel-material-ui >=0.8.1",
     "panel-splitjs >=0.3.2",
     "param >=2.2.1",
     "pyarrow",

From 40df87c41cbc144e47462f1da343aad09ca9a5df Mon Sep 17 00:00:00 2001
From: ghostiee-11 <168410465+ghostiee-11@users.noreply.github.com>
Date: Thu, 2 Apr 2026 11:35:39 +0530
Subject: [PATCH 6/7] Address review: use footer_actions directly, restore
 panel-material-ui pin

Removed hasattr fallback for footer_actions/footer_objects since
panel-material-ui >=0.9.0 guarantees footer_actions. Restored the
pin that was accidentally lowered to >=0.8.1.
---
 lumen/ai/ui.py | 5 ++---
 pixi.toml      | 2 +-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py
index 3bfa3a926..01b051357 100644
--- a/lumen/ai/ui.py
+++ b/lumen/ai/ui.py
@@ -1634,12 +1634,11 @@ async def _generate_follow_up():
 
         if len(self.interface):
             message = self.interface.objects[-1]
-            attr = "footer_actions" if hasattr(message, "footer_actions") else "footer_objects"
             existing = [
-                obj for obj in (getattr(message, attr) or [])
+                obj for obj in (message.footer_actions or [])
                 if getattr(obj, 'name', '') != "FollowUp"
             ]
-            setattr(message, attr, existing + [follow_up_button])
+            message.footer_actions = existing + [follow_up_button]
 
     def __panel__(self):
         return self._main
diff --git a/pixi.toml b/pixi.toml
index 8ee0ae7e8..ab5490b1b 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -20,7 +20,7 @@ lint = ["py312", "lint"]
 [dependencies]
 bokeh = "*"
 holoviews = ">=1.17.0"
-panel-material-ui = ">=0.8.1"
+panel-material-ui = ">=0.9.0"
 hvplot = "*"
 intake = "<2"
 jinja2 = ">3.0"

From bd3e033c2f74c15fd7fdeb4799daed64d0399e98 Mon Sep 17 00:00:00 2001
From: ghostiee-11 <168410465+ghostiee-11@users.noreply.github.com>
Date: Sun, 5 Apr 2026 09:48:59 +0530
Subject: [PATCH 7/7] Add on-demand follow-up suggestion icon after successful
 queries

---
 lumen/ai/ui.py            |  4 ++--
 lumen/tests/ai/test_ui.py | 17 +++++++----------
 pixi.toml                 |  2 +-
 pyproject.toml            |  2 +-
 4 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py
index 01b051357..b95826d16 100644
--- a/lumen/ai/ui.py
+++ b/lumen/ai/ui.py
@@ -1635,10 +1635,10 @@ async def _generate_follow_up():
         if len(self.interface):
             message = self.interface.objects[-1]
             existing = [
-                obj for obj in (message.footer_actions or [])
+                obj for obj in (message.footer_objects or [])
                 if getattr(obj, 'name', '') != "FollowUp"
             ]
-            message.footer_actions = existing + [follow_up_button]
+            message.footer_objects = existing + [follow_up_button]
 
     def __panel__(self):
         return self._main
diff --git a/lumen/tests/ai/test_ui.py b/lumen/tests/ai/test_ui.py
index ad6e9ff52..41fdefde4 100644
--- a/lumen/tests/ai/test_ui.py
+++ b/lumen/tests/ai/test_ui.py
@@ -1615,12 +1615,9 @@ async def noop_callback(contents, user, instance):
 
 # --- Follow-up suggestions tests ---
 
-def _get_footer_actions(msg):
-    """Get footer actions from message, falling back to footer_objects for older pmui."""
-    try:
-        return msg.footer_actions or []
-    except AttributeError:
-        return msg.footer_objects or []
+def _get_footer_objects(msg):
+    """Get footer objects from message."""
+    return msg.footer_objects or []
 
 
 def test_follow_up_suggestion_model():
@@ -1663,7 +1660,7 @@ async def test_follow_up_icon_after_query(explorer_ui):
 
     last_msg = ui.interface.objects[-1]
     follow_up_icons = [
-        f for f in _get_footer_actions(last_msg)
+        f for f in _get_footer_objects(last_msg)
         if getattr(f, 'name', '') == "FollowUp"
     ]
     assert len(follow_up_icons) == 1
@@ -1677,7 +1674,7 @@ async def test_follow_up_icon_not_on_error(explorer_ui_with_error):
 
     last_msg = ui.interface.objects[-1]
     follow_up_icons = [
-        f for f in _get_footer_actions(last_msg)
+        f for f in _get_footer_objects(last_msg)
         if getattr(f, 'name', '') == "FollowUp"
     ]
     assert len(follow_up_icons) == 0
@@ -1699,7 +1696,7 @@ async def test_follow_up_icon_not_shown_without_pipeline(explorer_ui):
 
     last_msg = ui.interface.objects[-1]
     follow_up_icons = [
-        f for f in _get_footer_actions(last_msg)
+        f for f in _get_footer_objects(last_msg)
         if getattr(f, 'name', '') == "FollowUp"
     ]
     assert len(follow_up_icons) == 0
@@ -1733,7 +1730,7 @@ async def test_follow_up_icon_description(explorer_ui):
 
     last_msg = ui.interface.objects[-1]
     follow_up_icons = [
-        f for f in _get_footer_actions(last_msg)
+        f for f in _get_footer_objects(last_msg)
         if getattr(f, 'name', '') == "FollowUp"
     ]
     assert len(follow_up_icons) == 1
diff --git a/pixi.toml b/pixi.toml
index e95fb1752..489694b4d 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -20,7 +20,7 @@ lint = ["py312", "lint"]
 [dependencies]
 bokeh = "*"
 holoviews = ">=1.17.0"
-panel-material-ui = ">=0.9.0"
+panel-material-ui = ">=0.9.1"
 hvplot = "*"
 intake = "<2"
 jinja2 = ">3.0"
diff --git a/pyproject.toml b/pyproject.toml
index d6c3d6fc9..c81fbdf20 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,7 +42,7 @@ dependencies = [
     "pandas",
     "panel >=1.7.5",
     "panel-graphic-walker[kernel] >=0.6.4",
-    "panel-material-ui >=0.8.1",
+    "panel-material-ui >=0.9.1",
     "panel-splitjs >=0.3.2",
     "param >=2.2.1",
     "pyarrow",