diff --git a/lumen/ai/coordinator.py b/lumen/ai/coordinator.py index a4c411558..3390e2832 100644 --- a/lumen/ai/coordinator.py +++ b/lumen/ai/coordinator.py @@ -203,6 +203,7 @@ async def hide_suggestions(_=None): if len(self.interface.objects) > num_objects: suggestion_buttons.visible = False + memory = self._memory async def use_suggestion(event): button = event.obj with button.param.update(loading=True), self.interface.active_widget.param.update(loading=True): @@ -219,10 +220,11 @@ async def use_suggestion(event): print("No analysis agent found.") return messages = [{'role': 'user', 'content': contents}] - await agent.respond( - messages, render_output=self.render_output, agents=self.agents - ) - await self._add_analysis_suggestions() + with agent.param.update(memory=memory): + await agent.respond( + messages, render_output=self.render_output, agents=self.agents + ) + await self._add_analysis_suggestions() else: self.interface.send(contents) diff --git a/lumen/ai/memory.py b/lumen/ai/memory.py index d9284ac48..4de3dde01 100644 --- a/lumen/ai/memory.py +++ b/lumen/ai/memory.py @@ -22,6 +22,10 @@ def __setitem__(self, key, new): super().__setitem__(key, new) self._trigger_update(key, old, new) + def cleanup(self): + self._callbacks.clear() + self._rx.clear() + def on_change(self, key, callback): self._callbacks[key].append(callback) diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py index 3466806c3..5a0866638 100644 --- a/lumen/ai/ui.py +++ b/lumen/ai/ui.py @@ -117,6 +117,13 @@ def __init__( ) self._resolve_data(data) self._main = Column(self._exports, self._coordinator, sizing_mode='stretch_both') + if state.curdoc and state.curdoc.session_context: + state.on_session_destroyed(self._destroy) + + def _destroy(self, session_context): + """ + Cleanup on session destroy + """ def _export_notebook(self): nb = export_notebook(self._coordinator.interface.objects, preamble=self.notebook_preamble) @@ -278,6 +285,13 @@ def __init__( ) self._output.param.watch(self._update_conversation, 'active') + def _destroy(self, session_context): + """ + Cleanup on session destroy + """ + for c in self._contexts: + c.cleanup() + def _update_conversation(self, event): if event.new: active = self._explorations.active @@ -292,12 +306,13 @@ def _update_conversation(self, event): self._coordinator.interface.objects = conversation def _cleanup_explorations(self, event): - if len(event.new) >= len(event.old): + if len(event.new) <= len(event.old): return for i, (old, new) in enumerate(zip(event.old, event.new)): if old is not new: self._contexts.pop(i) self._conversations.pop(i) + break def _set_context(self, event): if event.new == len(self._conversations): @@ -449,7 +464,6 @@ def render_output(_, old, new): await callback(contents, user, instance) finally: local_memory.remove_on_change('plan', render_plan) - local_memory.remove_on_change('outputs', render_output) if not outputs: prev_memory.update(local_memory) return wrapper