diff --git a/lumen/ai/report.py b/lumen/ai/report.py index 802affb46..bf05f137a 100644 --- a/lumen/ai/report.py +++ b/lumen/ai/report.py @@ -23,8 +23,9 @@ from panel.viewable import Viewable, Viewer from panel_material_ui import ( Accordion, BreakpointSwitcher, Button, Card, ChatFeed, ChatMessage, - Container, Dialog, Divider, FileDownload, IconButton, Progress, Select, - SpeedDial, TextAreaInput, TextInput, Typography, + Container, Dialog, Divider, FileDownload, IconButton, MultiChoice, + Progress, Select, SpeedDial, TextAreaInput, TextInput, ToggleIcon, + Typography, ) from ..views.base import Panel, View @@ -136,7 +137,7 @@ def _add_step(self, title: str = "", **kwargs): return nullcontext(self._null_step) if self.interface is None else self.interface.add_step(title=title, **kwargs) def _populate_view(self): - self._view[:] = self._header + self._view[:] = self._header + self.views def _render_controls(self): return [ @@ -222,6 +223,8 @@ async def execute(self, context: TContext | None = None, **kwargs) -> tuple[list if not self._prepared: await self.prepare(context) views, out_context = await self._execute(context, **kwargs) + except asyncio.CancelledError: + raise except Exception: self.status = "error" raise @@ -229,6 +232,8 @@ async def execute(self, context: TContext | None = None, **kwargs) -> tuple[list if self.status != "error": self.status = "success" self.running = False + + self.views = views self.out_context = out_context return views, out_context @@ -237,9 +242,13 @@ async def prepare(self, context: TContext): def reset(self): """Resets the view, removing generated outputs.""" + self._prepared = False self.status = "idle" self._view[:] = [] - self.out_context.clear() + self.views = [] + self.out_context = {} + if self._header: + self._header = [] class TaskGroup(Task): @@ -282,10 +291,9 @@ def __init__(self, *tasks, **params): self._populate_view() def _init_views(self): - views = self.param._header.rx() + self.views = self.param._header.rx() for task in self: - views += task.param.views.rx() - self.views = views + self.views += task.param.views.rx() async def _sync_context(self, event: param.parameterized.Event): task = event.obj @@ -298,12 +306,20 @@ async def _sync_context(self, event: param.parameterized.Event): if index > self._current or self.running: # No need to invalidate if the task hasn't been executed return - changed = {k for k, v in event.new.items() if k in event.old and not is_equal(event.old.get(k), v)} + # Walk up to the root of the task tree and check if it is + # currently executing. The root's ``running`` flag stays True + # for the entire duration of an execution run, unlike the local + # flag which may have already been reset by the time a deferred + # async param-watcher callback is dispatched by the event loop. + root = self + while root.parent is not None: + root = root.parent + if root.running: + return + changed = {k for k, v in event.new.items() if not is_equal(event.old.get(k), v)} + if changed: self.invalidate(changed, start=index) - root = self - while root.parent is not None: - root = root.parent if not (root is self and index >= len(self)): await root.execute() @@ -327,12 +343,12 @@ def __repr__(self): return f"{self.__class__.__name__}({', '.join(params)}{''.join(tasks)})" def _populate_view(self): - self._view[:] = self._header + self._tasks + self._view[:] = self._header + list(self._tasks) async def _run_task( self, i: int, task: Self | Actor, context: TContext, **kwargs ) -> tuple[list[Any], TContext]: - if task.status == "success": + if task.status == "success" and task.views: return task.views, task.out_context await asyncio.sleep(0.01) messages = self._render_message_history(context) @@ -383,6 +399,14 @@ def _get_context( else: raise TypeError("Abstract Task does not implement _get_context.") + def _included_tasks(self) -> set[str] | None: + """Return the set of task titles to include, or None to include all. + + Subclasses can override this to filter which tasks are executed + and exported. + """ + return None + async def _execute(self, context: TContext, **kwargs): """ Executes the tasks. @@ -399,13 +423,26 @@ async def _execute(self, context: TContext, **kwargs): views = list(self._header) self.status = "running" + included_titles = self._included_tasks() for i, task in enumerate(self): + if included_titles is not None and task.title not in included_titles: + task.reset() + continue + + if task.status == 'success': + views += task.views + context = merge_contexts(LWW, [context, task.out_context]) + self._current = i + 1 + continue + new = [] try: new, new_context = await self._run_task(i, task, context, **kwargs) except MissingContextError: # Re-raise MissingContextError to allow retry logic at Plan level raise + except asyncio.CancelledError: + raise except Exception as e: tb.print_exception(e) self.status = "error" @@ -561,6 +598,12 @@ def invalidate( if invalidated: self.status = "idle" self._current = max(min(invalidated), 0) + # Keep parent views in sync with child invalidation state since + # execute() assigns a concrete list to self.views. + views = list(self._header) + for task in self: + views += task.views + self.views = views return [self[i] for i in invalidated], keys def merge(self, other: TaskGroup): @@ -585,6 +628,7 @@ def merge(self, other: TaskGroup): self.out_context = merge_contexts(LWW, [self.context, self.out_context]) self.history += [h for h in other.history if h not in self.history] self._init_views() + self._populate_view() return self async def prepare(self, context: TContext): @@ -596,7 +640,11 @@ def reset(self, start: int = 0): """ Resets the view, removing generated outputs. """ + self._prepared = False self.status = "idle" + self._view[:] = [] + self.views = [] + self.out_context = {} self._current = start with hold(): self._header = [] @@ -606,6 +654,14 @@ def reset(self, start: int = 0): task.reset() self._populate_view() + def reset_filters(self): + """ + Resets the filters for the task group and its subtasks. + """ + for task in self: + if isinstance(task, TaskGroup): + task.reset_filters() + def to_notebook(self): """ Returns the notebook representation of the tasks. @@ -614,8 +670,9 @@ def to_notebook(self): raise RuntimeError( "Report has not been executed, run report before exporting to_notebook." ) + included_views = self._collect_included_views() cells, extensions = [], ['tabulator'] - for out in self.views: + for out in included_views: ext = None if isinstance(out, Typography): level = int(out.variant[1:]) if out.variant and out.variant.startswith('h') else 0 @@ -625,14 +682,33 @@ def to_notebook(self): cell = make_md_cell(out.object) elif isinstance(out, LumenEditor): cell, ext = format_output(out) + elif isinstance(out, View): + cell, ext = format_output(out) elif isinstance(out, Viewable): cell, ext = format_output(Panel(object=out)) + else: + continue cells.append(cell) if ext and ext not in extensions: extensions.append(ext) cells = make_preamble("", extensions=extensions) + cells return write_notebook(cells) + def _collect_included_views(self) -> list[Any]: + included_titles = self._included_tasks() + included_views = list(self._header) + for task in self: + if included_titles is not None and task.title not in included_titles: + continue + if isinstance(task, TaskGroup): + included_views += task._collect_included_views() + else: + # Skip the task's own _header items (title Typography nodes) since + # they duplicate headings already emitted from self._header above + task_header_set = set(id(v) for v in task._header) + included_views += [v for v in task.views if id(v) not in task_header_set] + return included_views + def validate( self, context: TContext | None = None, @@ -720,13 +796,23 @@ def _init_view(self): "", variant="body2", margin=(10, 10), sx={"color": "text.secondary"} ) + self._export_selector = MultiChoice( + name="Include in process & export", + options=[], + value=[], + sizing_mode="stretch_width", + ) + self._selector_initialized = False + self._export_selector.param.watch(self._sync_placeholder_to_selector, 'value') self._view = Column(sizing_mode='stretch_width', styles={'min-height': 'unset'}, height_policy='fit') + self._populate_view() self._container = Column( Row( self._settings, styles={'position': 'absolute', 'top': '-57.5px', 'right': '20px'} ), self._dialog, + self._export_selector, self._view, sizing_mode='stretch_width', styles={'min-height': 'unset'}, @@ -736,6 +822,13 @@ def _init_view(self): def _open_settings(self, event): self._dialog.open = True + def _clear_progress_loaders(self): + # Progress widgets are transient execution indicators and should never + # remain in the section body once execution has finished. + for view in list(self._view): + if isinstance(view, Progress): + self._view.remove(view) + @param.depends('status', '_tasks', 'views', watch=True) def _update_placeholder(self): """Update placeholder text based on execution status.""" @@ -743,15 +836,64 @@ def _update_placeholder(self): return has_outputs = self.status != "idle" or bool(self.views) if not has_outputs: - n = len(self._tasks) - task_word = "task" if n == 1 else "tasks" - self._placeholder.object = f"{n} {task_word} ready" + n_total = len(self._tasks) + n_selected = len(self._export_selector.value) + task_word = "task" if n_total == 1 else "tasks" + if n_selected < n_total: + self._placeholder.object = f"{n_selected} selected / {n_total} {task_word} ready" + else: + self._placeholder.object = f"{n_total} {task_word} ready" self._placeholder.visible = True else: self._placeholder.visible = False + def _included_tasks(self) -> set[str] | None: + selected = self._export_selector.value + if len(selected) == len(self._tasks): + return None + return set(selected) + + def _sync_placeholder_to_selector(self, event=None): + # Changing selector membership invalidates section-level cached outputs. + self.status = "idle" + self._current = 0 + self._update_placeholder() + self._populate_view() + + def reset_filters(self): + """ + Resets the filters for the section. + """ + self._export_selector.value = self._export_selector.options + super().reset_filters() + def _populate_view(self): - self._view[:] = self._header + [self._placeholder] + list(self._tasks) + titles = [task.title for task in self._tasks] + prev_options = self._export_selector.options + self._export_selector.options = titles + if not self._selector_initialized: + # First populate: select all tasks by default + new_value = titles + self._selector_initialized = True + else: + # Keep existing selections that are still valid, and auto-select any + # newly added tasks (titles not previously in options) + kept = [t for t in self._export_selector.value if t in titles] + new_titles = [t for t in titles if t not in prev_options] + new_value = kept + [t for t in new_titles if t not in kept] + + # Guard against re-triggering the _sync_placeholder_to_selector + # watcher when the computed value hasn't actually changed. + if self._export_selector.value != new_value: + self._export_selector.value = new_value + + selected = self._export_selector.value + tasks = [task for task in self._tasks if task.title in selected] + # Note: self._header is intentionally excluded from the view here + # because the parent Report's Accordion already renders the section + # title as a panel header. _header is still kept for notebook export + # via _collect_included_views. + self._view[:] = [self._placeholder] + tasks async def _run_task(self, i: int, task: Task | Actor, context: TContext | None, **kwargs) -> list[Any]: if context is not None: @@ -762,11 +904,18 @@ async def _run_task(self, i: int, task: Task | Actor, context: TContext | None, context['reasoning'] = f"{self.title}\n\n{instructions}" return await super()._run_task(i, task, context, **kwargs) + async def _execute(self, context: TContext, **kwargs): + try: + return await super()._execute(context, **kwargs) + finally: + self._clear_progress_loaders() + @param.depends('running', watch=True) async def _running(self): await asyncio.sleep(0.05) if not self.running: return + self._clear_progress_loaders() # Start with indeterminate, switch to determinate after first task loader = Progress( variant='indeterminate', @@ -786,7 +935,8 @@ async def _running(self): while self.running: await asyncio.sleep(0.05) - self._view.remove(loader) + if loader in self._view: + self._view.remove(loader) class Report(TaskGroup): @@ -827,14 +977,22 @@ def _init_view(self): } } ) - self._run = IconButton( - icon="play_arrow", on_click=self._execute_event, margin=0, size="large", - description="Execute Report", loading=self.param.running, + self._abort = False + self._execute_task = None + self._run = ToggleIcon( + icon="play_circle_filled", active_icon="stop", + value=False, margin=0, size="large", + description="Execute / Stop Report", ) + self._run.param.watch(self._on_run_toggle, 'value') self._clear = IconButton( icon="clear", on_click=lambda _: self.reset(), margin=0, size="large", description="Clear outputs", visible=False ) + self._filter_reset = IconButton( + icon="restart_alt", on_click=lambda _: self.reset_filters(), margin=0, size="large", + description="Reset filters", visible=True + ) self._collapse = IconButton( styles={"margin-left": "auto"}, on_click=self._expand_all, icon="unfold_less", size="large", color="default", margin=(0, 0, 10, 0), @@ -861,6 +1019,7 @@ def _init_view(self): self._header_title, self._run, self._clear, + self._filter_reset, self._collapse, self._export, self._settings, @@ -871,6 +1030,7 @@ def _init_view(self): items=[ {"label": "Execute Report", "icon": "play_arrow"}, {"label": "Clear Report", "icon": "clear"}, + {"label": "Reset filters", "icon": "restart_alt"}, {"label": "Export Report to Notebook", "icon": "get_app"}, {"label": "Configure Report", "icon": "settings"} ], @@ -906,9 +1066,12 @@ def _init_view(self): async def _trigger_event(self, item: dict): icon = item["icon"] if icon == "play_arrow": - await self._execute_event() + self._run.value = not self._run.value + # watcher handles execution elif icon == "clear": self.reset() + elif icon == "restart_alt": + self.reset_filters() elif icon == "get_app": self._trigger_download() elif icon == "settings": @@ -923,14 +1086,14 @@ def _update_run_state(self): # Play button is always enabled - users can re-run after changing settings pass - @param.depends('status', 'views', watch=True) + @param.depends('status', 'views', 'running', watch=True) def _update_icon_visibility(self): """Show/hide icons based on whether report has outputs.""" has_outputs = self.status in ("success", "error") or bool(self.views) - self._clear.visible = has_outputs - self._collapse.visible = has_outputs - self._export.visible = has_outputs - self._settings.visible = has_outputs + self._clear.visible = has_outputs and not self.running + self._collapse.visible = has_outputs and not self.running + self._export.visible = has_outputs and not self.running + self._settings.visible = has_outputs and not self.running # Only animate play button when no outputs if has_outputs: self._run.sx = {} @@ -943,9 +1106,26 @@ def _update_icon_visibility(self): }, } + def _on_run_toggle(self, event): + if event.new: + pn.state.execute(self._execute_event) + # User clicked stop — cancel the running task + elif self._execute_task and not self._execute_task.done(): + self._execute_task.cancel() + async def _execute_event(self, event=None): - await asyncio.sleep(0.01) # yield the event loop to allow button loading state to update - await self.execute() + self._execute_task = asyncio.current_task() + try: + await asyncio.sleep(0.01) + await self.execute() + except asyncio.CancelledError: + self.running = False + for section in self: + if section.running: + section.running = False + finally: + self._run.value = False + self._execute_task = None async def _notebook_export(self): if len(self) and self.status != "success": @@ -1162,10 +1342,9 @@ def _add_outputs(self, views: list, context: TContext): self._view.extend(rendered) async def _execute(self, context: TContext, **kwargs) -> tuple[list[Any], TContext]: - views = [] + outputs = [] if self.title and not self._header: title = Typography(f"{'#'*self.level} {self.title}", margin=(10, 10, 0, 10)) - views.append(title) self.views = self._header = [title] messages = self._render_message_history(context) @@ -1193,7 +1372,7 @@ async def _execute(self, context: TContext, **kwargs) -> tuple[list[Any], TConte if unprovided: raise RuntimeError(f"{self.actor.__class__.__name__} failed to provide declared context: {', '.join(unprovided)}.") contexts = [self.context, out_context] if self.context else [out_context] - return outputs, merge_contexts(LWW, contexts) + return self.views, merge_contexts(LWW, contexts) def _actor_prompt(self, actor: Actor): prompt = Select( diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py index 02af91d22..eac21320d 100644 --- a/lumen/ai/ui.py +++ b/lumen/ai/ui.py @@ -2073,11 +2073,11 @@ def _cleanup_exploration(self, item): if exploration.plan is None: return for child in item.get('items', []): + child['view'].context.clear() if child['view'].plan is not None: child['view'].plan.cleanup() - child['view'].context.clear() - exploration.plan.cleanup() exploration.context.clear() + exploration.plan.cleanup() if item in self._explorations.items: # Top-level exploration: remove and switch to Home self._explorations.items = [ @@ -2365,7 +2365,6 @@ def _find_view_in_popped_out(self, exploration: Exploration, out: LumenEditor): def _add_views(self, exploration: Exploration, event: param.parameterized.Event | None = None, items: list | None = None): tabs = exploration.view[0] - content = [] new_items = items if event is None else event.new for view in new_items: if not isinstance(view, LumenEditor) or (event and view in event.old): @@ -2374,9 +2373,11 @@ def _add_views(self, exploration: Exploration, event: param.parameterized.Event tabs[0] = ("Data Source", view.render_explorer()) exploration.initialized = True title, vsplit = self._render_view(exploration, view) - content.append((title, vsplit)) - - tabs.extend(content) + if title in tabs._names: + tab_idx = tabs._names.index(title) + tabs[tab_idx] = (title, vsplit) + else: + tabs.append((title, vsplit)) tabs.active = max(tabs.active, len(tabs)-1) if self._exploration['view'] is exploration: self._update_main_view() @@ -2434,9 +2435,21 @@ def _reset_error(self, plan: Plan, exploration: Exploration, replan: bool = Fals if replan: tabs[:] = [("Data Source", Markdown("Waiting on data...", margin=(5, 20)))] tabs.active = 0 - elif len(tabs) > 1 and isinstance(tabs[-1], Markdown): - tabs.pop(-1) - tabs.active = len(tabs)-1 + elif len(tabs) > 1: + while len(tabs) > 1: + last_idx = len(tabs) - 1 + last_name = tabs._names[last_idx] if hasattr(tabs, "_names") else None + last_tab = tabs[last_idx] + is_error_markdown = ( + isinstance(last_tab, Markdown) + and isinstance(last_tab.object, str) + and last_tab.object.startswith("❌ **Unable to complete your request**") + ) + if last_name != "Error" and not is_error_markdown: + break + tabs.pop(last_idx) + if len(tabs): + tabs.active = len(tabs) - 1 async def _execute_plan(self, plan: Plan, rerun: bool = False, replan: bool = False): prev = self._explorations.value @@ -2472,7 +2485,8 @@ async def _execute_plan(self, plan: Plan, rerun: bool = False, replan: bool = Fa partial_plan = plan plan = parent.plan.merge(plan) partial_plan.cleanup() - if replan and new_exploration: + if replan: + plan.views = [] watcher = plan.param.watch(partial(self._add_views, exploration), "views") else: watcher = None diff --git a/lumen/tests/ai/test_report.py b/lumen/tests/ai/test_report.py index dddefe884..a72899867 100644 --- a/lumen/tests/ai/test_report.py +++ b/lumen/tests/ai/test_report.py @@ -289,6 +289,68 @@ async def test_taskgroup_invalidate(): assert isinstance(outs[0], Markdown) assert outs[0].object == "B" + +async def test_section_incremental_execution(): + order = [] + a = A(order=order, title="Action A") + b = B(order=order, title="Action B") + section = Section(a, b, title="Section 1") + report = Report(section, title="Selector Report") + + # Initial run with both A and B + await report.execute() + assert order == ["A", "B"] + assert a.status == "success" + assert b.status == "success" + assert len(a.views) == 2 # Title + Markdown + assert len(b.views) == 2 + + # Deselect B and run again + section._export_selector.value = ["Action A"] + await report.execute() + assert order == ["A", "B"] # A is cached, B is skipped (and reset) + assert b.status == "idle" + assert len(b.views) == 0 + assert len(section._view) == 2 # Placeholder + Action A (Action B is filtered out) + + # Select B again and run again + section._export_selector.value = ["Action A", "Action B"] + await report.execute() + assert order == ["A", "B", "B"] # A is cached, B is re-run + assert b.status == "success" + assert len(b.views) == 2 + + +async def test_section_reset_filters(): + section = Section(A(title="Action A"), B(title="Action B"), title="Section 1") + report = Report(section, title="Selector Report") + + section._export_selector.value = ["Action A"] + assert section._export_selector.value == ["Action A"] + + report.reset_filters() + assert section._export_selector.value == ["Action A", "Action B"] + + +async def test_section_selector_filters_ui(): + a = A(title="Action A") + b = B(title="Action B") + section = Section(a, b, title="Section 1") + + # Initial state: both selected and rendered. + assert section._export_selector.value == ["Action A", "Action B"] + assert len(section._view) == 3 + + # Remove B from selector + section._export_selector.value = ["Action A"] + assert section._export_selector.value == ["Action A"] + assert len(section._view) == 2 + + # Add B back + section._export_selector.value = ["Action A", "Action B"] + assert section._export_selector.value == ["Action A", "Action B"] + assert len(section._view) == 3 + async def test_taskgroup_invalidate_partial(): class BInputs(ContextModel): @@ -369,3 +431,151 @@ async def test_report_to_notebook(): assert cell3['cell_type'] == 'markdown' assert cell3['source'] == ["**Hello**"] + + +async def test_export_selector_filters_play_execution(): + order = [] + + section = Section( + A(order=order, title="Action A"), + B(order=order, title="Action B"), + title="Section 1", + ) + report = Report(section, title="Selector Report") + + section._export_selector.value = ["Action B"] + + await report.execute() + + assert order == ["B"] + assert len(section.views) == 3 + assert isinstance(section.views[0], Typography) + assert section.views[0].object == "## Section 1" + assert isinstance(section.views[1], Typography) + assert section.views[1].object == "### Action B" + assert isinstance(section.views[2], Markdown) + assert section.views[2].object == "B done" + + +async def test_sync_context_skipped_while_root_running(): + """_sync_context must not trigger re-execution while the root Report is running. + + Regression test: deferred async param-watcher callbacks for out_context + could fire after a Section finished but while the Report was still running, + causing spurious invalidation cascades and duplicate task execution. + """ + order = [] + + section = Section( + ASchema(title="Action A"), + BSchema(title="Action B"), + title="Section 1", + ) + report = Report(section, title="Test Report") + + await report.execute() + + assert section[0].out_context == {"a": "A"} + assert section[1].out_context == {"b": "AA"} + + # Simulate a deferred _sync_context callback arriving while the root + # is running. Before the fix, this would trigger invalidation and + # re-execution. + report.running = True + try: + section[0].out_context = {"a": "CHANGED"} + await asyncio.sleep(0.3) + + # B should NOT have been invalidated or re-run because the root + # was running — the _sync_context callback should have bailed out. + assert section[1].out_context == {"b": "AA"} + assert section[1].status == "success" + finally: + report.running = False + + +async def test_selector_subset_no_reexecution_loop(): + """Deselecting a task and running should not cause tasks to execute more than once. + + Regression test: when a task was deselected via _export_selector, the + reset of excluded tasks would trigger _sync_context callbacks that + re-ran the entire report in an infinite loop. + """ + order = [] + + section = Section( + A(order=order, title="Action A"), + B(order=order, title="Action B"), + title="Section 1", + ) + section_2 = Section( + A(order=order, title="Action C"), + title="Section 2", + ) + report = Report(section, section_2, title="Test Report") + + # Deselect Action B + section._export_selector.value = ["Action A"] + + await report.execute() + + # Allow any deferred callbacks to settle + await asyncio.sleep(0.5) + + # Each selected action should have run exactly once + assert order.count("A") == 2 # Action A in Section 1 + Action C (class A) in Section 2 + assert order.count("B") == 0 # Action B was deselected + + +async def test_populate_view_no_selector_retrigger(): + """Section._populate_view should not re-trigger _sync_placeholder_to_selector + when the selector value hasn't changed.""" + section = Section( + A(title="Action A"), + B(title="Action B"), + title="Section 1", + ) + + sync_calls = [] + original = section._sync_placeholder_to_selector + + def tracking_sync(*args, **kwargs): + sync_calls.append(1) + return original(*args, **kwargs) + + section._sync_placeholder_to_selector = tracking_sync + + # Calling _populate_view when value hasn't changed should NOT trigger the watcher + sync_calls.clear() + section._populate_view() + assert len(sync_calls) == 0 + + +async def test_export_selector_filters_notebook_export(): + order = [] + + section = Section( + A(order=order, title="Action A"), + B(order=order, title="Action B"), + title="Section 1", + ) + report = Report(section, title="Selector Report") + + await report.execute() + assert order == ["A", "B"] + + section._export_selector.value = ["Action A"] + + nb_string = report.to_notebook() + nb = json.loads(nb_string) + + markdown_cells = [ + "".join(cell["source"]) for cell in nb["cells"] if cell["cell_type"] == "markdown" + ] + + assert markdown_cells == [ + "# Selector Report", + "## Section 1", + "### Action A", + "A done", + ]