From 20e00cf9f1313d2469b1f10e32b8c36f8d51fbf6 Mon Sep 17 00:00:00 2001
From: Philipp Rudiger <prudiger@anaconda.com>
Date: Mon, 16 Dec 2024 14:27:22 +0100
Subject: [PATCH] Add button to export all explorations from ExplorerUI (#870)

---
 lumen/ai/export.py | 15 ++++++++++++---
 lumen/ai/ui.py     | 43 +++++++++++++++++++++++++++----------------
 2 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/lumen/ai/export.py b/lumen/ai/export.py
index d677e1ca5..95263fffe 100644
--- a/lumen/ai/export.py
+++ b/lumen/ai/export.py
@@ -14,9 +14,12 @@
 from lumen.views import View
 
 
+def make_md_cell(text: str):
+    return nbformat.v4.new_markdown_cell(source=text)
+
 def make_preamble(preamble: str):
     now = dt.datetime.now()
-    header = nbformat.v4.new_markdown_cell(source=f'# Lumen.ai - Chat Logs {now}')
+    header = make_md_cell(f'# Lumen.ai - Chat Logs {now}')
     source = (preamble + dedent(
         """
         import lumen as lm
@@ -55,8 +58,8 @@ def format_output(msg: ChatMessage):
         ])
     return [nbformat.v4.new_code_cell(source='\n'.join(code))]
 
-def export_notebook(messages: list[ChatMessage], preamble: str = ""):
-    cells = make_preamble(preamble)
+def render_cells(messages: list[ChatMessage]):
+    cells = []
     for msg in messages:
         if msg.user == 'Help':
             continue
@@ -69,6 +72,12 @@ def export_notebook(messages: list[ChatMessage], preamble: str = ""):
                 if isinstance(obj, ChatStep):
                     continue
                 cells += format_output(obj)
+    return cells
 
+def write_notebook(cells):
     nb = nbformat.v4.new_notebook(cells=cells)
     return nbformat.v4.writes(nb)
+
+def export_notebook(messages: list[ChatMessage], preamble: str = ""):
+    cells = make_preamble(preamble) + render_cells(messages)
+    return write_notebook(cells)
diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py
index 33f918a0c..98cc8ab55 100644
--- a/lumen/ai/ui.py
+++ b/lumen/ai/ui.py
@@ -36,7 +36,9 @@
 from .components import SplitJS
 from .controls import SourceControls
 from .coordinator import Coordinator, Planner
-from .export import export_notebook
+from .export import (
+    export_notebook, make_md_cell, make_preamble, render_cells, write_notebook,
+)
 from .llm import Llm, OpenAI
 from .memory import _Memory, memory
 
@@ -129,7 +131,7 @@ def __init__(
             icon_size="1.5em",
             button_type="primary",
             callback=self._export_notebook,
-            filename=" ",
+            filename=f"{self.title.replace(' ', '_')}.ipynb",
             stylesheets=['.bk-btn a { padding: 0 6px; }'],
         )
         self._exports = Row(
@@ -143,8 +145,7 @@ def __init__(
                 for label, e in self.export_functions.items()
             ),
             styles={'position': 'relative', 'right': '20px', 'top': '-5px'},
-            sizing_mode='stretch_width',
-            visible=False
+            sizing_mode='stretch_width'
         )
         self._main = Column(self._exports, self._coordinator, sizing_mode='stretch_both')
         if state.curdoc and state.curdoc.session_context:
@@ -245,21 +246,12 @@ class ChatUI(UI):
     ```python
     import lumen.ai as lmai
 
-    lmai.ChatUI('~/data.csv').servable()
+    lmai.ChatUI(data='~/data.csv').servable()
     ```
     """
 
     title = param.String(default='Lumen ChatUI', doc="Title of the app.")
 
-    def __init__(
-        self,
-        data: DataT | list[DataT] | None = None,
-        **params
-    ):
-        super().__init__(data, **params)
-        self._notebook_export.filename = f"{self.title.replace(' ', '_')}.ipynb"
-        self._exports.visible = True
-
 
 class ExplorerUI(UI):
     """
@@ -275,7 +267,7 @@ class ExplorerUI(UI):
     ```python
     import lumen.ai as lmai
 
-    lmai.ExplorerUI('~/data.csv').servable()
+    lmai.ExplorerUI(data='~/data.csv').servable()
     ```
     """
 
@@ -294,6 +286,16 @@ def __init__(
         self._explorations = Tabs(sizing_mode='stretch_both', closable=True)
         self._explorations.param.watch(self._cleanup_explorations, ['objects'])
         self._explorations.param.watch(self._set_context, ['active'])
+        self._global_notebook_export = FileDownload(
+            icon="notebook",
+            icon_size="1.5em",
+            button_type="primary",
+            callback=self._global_export_notebook,
+            filename=f"{self.title.replace(' ', '_')}.ipynb",
+            stylesheets=['.bk-btn a { padding: 0 6px; }'],
+            styles={'position': 'absolute', 'right': '0px', 'top': '-5px', 'z-index': '100'},
+        )
+        self._exports.visible = False
         self._titles = []
         self._contexts = []
         self._root_conversation = self._coordinator.interface.objects
@@ -305,7 +307,7 @@ def __init__(
         )
         self._main = Column(
             SplitJS(
-                left=Column(self._output, styles={'overflow-x': 'auto'}, sizing_mode='stretch_both'),
+                left=Column(self._global_notebook_export, self._output, styles={'overflow-x': 'auto'}, sizing_mode='stretch_both'),
                 right=Column(self._exports, self._coordinator),
                 sizing_mode='stretch_both'
             )
@@ -322,6 +324,15 @@ def _destroy(self, session_context):
         for c in self._contexts:
             c.cleanup()
 
+    def _global_export_notebook(self):
+        cells = make_preamble(self.notebook_preamble)
+        for i, objects in enumerate(self._conversations):
+            title = self._titles[i]
+            header = make_md_cell(f'## {title}')
+            cells.append(header)
+            cells += render_cells(objects)
+        return StringIO(write_notebook(cells))
+
     async def _update_conversation(self, event=None, tab=None):
         active = self._explorations.active
         if tab is None: