Add follow-up suggestion chips after successful queries

docs/configuration/ui.md

@@ -78,6 +78,10 @@ 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. 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
 
 ### Enable chat logging

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,45 @@ async def _check_tool_relevance(self, tool: Tool, tool_output: str, actor: Actor
 
         return result.yes
 
+    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.
+        """
+        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 as e:
+            if state.notifications:
+                state.notifications.warning(
+                    f"Could not generate follow-up suggestion: {e}",
+                    duration=3000,
+                )
+            return None
+
     def __panel__(self):
         return self.interface
 

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.

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 %}

lumen/ai/ui.py

@@ -1596,6 +1596,58 @@ async def _add_analysis_suggestions(self, plan: Plan):
             num_objects=len(self.interface.objects),
         )
 
+    def _add_follow_up_icon(self, plan: Plan):
+        """Add follow-up suggestion icon to the last message footer.
+
+        The icon appears immediately. On click, it calls the coordinator
+        to generate a suggestion and populates the chat input.
+        """
+
+        async def _generate_follow_up():
+            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
+            finally:
+                follow_up_button.disabled = False
+                follow_up_button.description = "Suggest a follow-up question"
+
+        if not plan.out_context.get("pipeline"):
+            return
+        if not plan.out_context.get("data"):
+            return
+
+        follow_up_button = IconButton(
+            icon="lightbulb",
+            description="Suggest a follow-up question",
+            size="small",
+            icon_size="0.9em",
+            margin=(5, 0),
+            on_click=lambda _: state.execute(_generate_follow_up),
+            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 + [follow_up_button]
+            except AttributeError:
+                message.footer_objects = existing + [follow_up_button]
+
     def __panel__(self):
         return self._main
 
@@ -2506,6 +2558,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)
+                self._add_follow_up_icon(plan)
 
             if is_new:
                 plan.param.watch(partial(self._update_views, exploration), "views")

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_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])
+
+    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_follow_up(plan)
+    assert result == "Show value by category"
+
+
+async def test_suggest_follow_up_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_follow_up(plan)
+    assert result is None
+
+
+async def test_suggest_follow_up_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_follow_up(plan)
+    assert result is None
+
+
+async def test_suggest_follow_up_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_follow_up(plan)
+    assert result is None