Add follow-up suggestion chips after successful queries

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 |

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

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
 

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

@@ -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")

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