Ensure Analyses are performed in context of the exploration (#775)

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)
 

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)
 

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