diff --git a/docs/configuration/ui.md b/docs/configuration/ui.md index 548566375..713e45f4d 100644 --- a/docs/configuration/ui.md +++ b/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 diff --git a/lumen/ai/coordinator/base.py b/lumen/ai/coordinator/base.py index e79f04950..0048d24db 100644 --- a/lumen/ai/coordinator/base.py +++ b/lumen/ai/coordinator/base.py @@ -27,7 +27,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 ..tools.document_llm_tools import make_document_vector_llm_tools @@ -301,6 +301,11 @@ class Coordinator(Viewer, VectorLookupToolUser): "template": PROMPTS_DIR / "Coordinator" / "tool_relevance.jinja2", "response_model": ThinkingYesNo, }, + "follow_up": { + "template": PROMPTS_DIR / "Coordinator" / "follow_up_suggestions.jinja2", + "response_model": FollowUpSuggestion, + "llm_spec": "ui", + }, }, ) @@ -554,6 +559,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", + 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 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 0629e03a1..0bee6eefe 100644 --- a/lumen/ai/ui.py +++ b/lumen/ai/ui.py @@ -1664,6 +1664,50 @@ 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: + self._chat_input.value = 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] + existing = [ + obj for obj in (message.footer_objects or []) + if getattr(obj, 'name', '') != "FollowUp" + ] + message.footer_objects = existing + [follow_up_button] + def __panel__(self): return self._main @@ -2566,6 +2610,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") diff --git a/lumen/tests/ai/test_coordinator.py b/lumen/tests/ai/test_coordinator.py index 04510c08e..934ced906 100644 --- a/lumen/tests/ai/test_coordinator.py +++ b/lumen/tests/ai/test_coordinator.py @@ -24,7 +24,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 @@ -566,6 +566,84 @@ async def test_process_files_duplicate_table_name(): 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 + + # ------------------------------------------------------------------- # _serialize — image propagation # ------------------------------------------------------------------- diff --git a/lumen/tests/ai/test_ui.py b/lumen/tests/ai/test_ui.py index f07d818bf..41fdefde4 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 @@ -26,7 +27,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 @@ -1610,3 +1611,128 @@ 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_objects(msg): + """Get footer objects from message.""" + return msg.footer_objects or [] + + +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_follow_up_suggestion_model_max_length(): + """FollowUpSuggestion should reject queries over 100 characters.""" + with pytest.raises(Exception): + FollowUpSuggestion(query="x" * 101) + + +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")]) + + 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) + + last_msg = ui.interface.objects[-1] + follow_up_icons = [ + f for f in _get_footer_objects(last_msg) + if getattr(f, 'name', '') == "FollowUp" + ] + assert len(follow_up_icons) == 1 + assert follow_up_icons[0].icon == "lightbulb" + + +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] + follow_up_icons = [ + f for f in _get_footer_objects(last_msg) + if getattr(f, 'name', '') == "FollowUp" + ] + assert len(follow_up_icons) == 0 + + +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 + + chat_agent = ChatAgent(llm=ui.llm) + ui.llm.set_responses(["Here is my answer"]) + plan = Plan( + ActorTask(chat_agent), + history=[{"content": "What is lumen?", "role": "user"}], + title="Chat only", + context=ui.context, + ) + await ui._execute_plan(plan) + + last_msg = ui.interface.objects[-1] + follow_up_icons = [ + f for f in _get_footer_objects(last_msg) + if getattr(f, 'name', '') == "FollowUp" + ] + assert len(follow_up_icons) == 0 + + +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")]) + + 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) + + last_msg = ui.interface.objects[-1] + follow_up_icons = [ + f for f in _get_footer_objects(last_msg) + if getattr(f, 'name', '') == "FollowUp" + ] + assert len(follow_up_icons) == 1 + assert follow_up_icons[0].description == "Suggest a follow-up question" + assert follow_up_icons[0].icon == "lightbulb" diff --git a/pixi.toml b/pixi.toml index 18b60b578..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.8.1" +panel-material-ui = ">=0.9.1" hvplot = "*" intake = "<2" jinja2 = ">3.0" diff --git a/pyproject.toml b/pyproject.toml index 2e34a4009..512bf8568 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",