diff --git a/lumen/ai/controls.py b/lumen/ai/controls.py index 6cf6f216..e8b56e62 100644 --- a/lumen/ai/controls.py +++ b/lumen/ai/controls.py @@ -17,7 +17,7 @@ class TableControls(Viewer): filename = param.String(default="", doc="Filename") - table = param.String(default="", doc="Table name") + table = param.String(default="", doc="What to name the uploaded file when querying it as a table") extension = param.String(default="", doc="File extension") sheet = param.Selector(default=None, objects=[], doc="Sheet") @@ -35,7 +35,7 @@ def __init__(self, file: io.BytesIO, **params): super().__init__(**params) self.file = file self._name_input = TextInput.from_param( - self.param.table, name="Table name" + self.param.table, name="Table alias", ) self._sheet_select = Select.from_param( self.param.sheet, name="Sheet", visible=False @@ -71,7 +71,7 @@ def __panel__(self): class SourceControls(Viewer): - add = param.Event(doc="Add tables") + add = param.Event(doc="Use table(s)") memory = param.ClassSelector(class_=_Memory, default=None, doc=""" Local memory which will be used to provide the agent context. @@ -121,7 +121,7 @@ def __init__(self, **params): self._add_button = Button.from_param( self.param.add, - name="Use tables", + name="Use table(s)", icon="table-plus", visible=False, button_type="success", diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py index 1da3f266..fe9c75a6 100644 --- a/lumen/ai/ui.py +++ b/lumen/ai/ui.py @@ -327,10 +327,11 @@ def _table_explorer(self): table_select = MultiChoice( sizing_mode='stretch_width', max_height=200, margin=(5, 0), max_items=5 ) - load_button = Button( - name='Load table(s)', icon='table-plus', button_type='primary', align='center', + explore_button = Button( + name='Explore table(s)', icon='chart-bar', button_type='primary', align='center', disabled=table_select.param.value.rx().rx.not_() ) + input_row = Row(table_select, explore_button) source_map = {} def update_source_map(_, __, sources, init=False): @@ -348,19 +349,31 @@ def update_source_map(_, __, sources, init=False): source_map.clear() source_map.update(new) table_select.param.update(options=list(source_map), value=selected) + input_row.visible = bool(source_map) memory.on_change('available_sources', update_source_map) update_source_map(None, None, memory['available_sources'], init=True) + def explore_table_if_single(event): + """ + If only one table is uploaded, help the user load it + without requiring them to click twice. This step + only triggers when the Upload in the Overview tab is used, + i.e. does not trigger with uploads through the SourceAgent + """ + if len(table_select.options) == 1: + explore_button.param.trigger("value") + controls = SourceControls(select_existing=False, name='Upload') + controls.param.watch(explore_table_if_single, "add") tabs = Tabs(controls, sizing_mode='stretch_both', design=Material) - @param.depends(table_select, load_button, watch=True) - def get_explorers(tables, load): + @param.depends(explore_button, watch=True) + def get_explorers(load): if not load: return - with load_button.param.update(loading=True): + with explore_button.param.update(loading=True): explorers = [] - for table in tables: + for table in table_select.value: source = source_map[table] if len(memory['available_sources']) > 1: _, table = table.rsplit(' : ', 1) @@ -369,16 +382,17 @@ def get_explorers(tables, load): ) walker = GraphicWalker( pipeline.param.data, sizing_mode='stretch_both', min_height=800, - kernel_computation=True, name=table, tab='data' + kernel_computation=True, name=f"View {table}", tab='data' ) explorers.append(walker) tabs.objects = explorers + [controls] + table_select.value = [] return Column( Markdown('### Start chatting or select an existing dataset or upload a .csv, .parquet, .xlsx file.', margin=(5, 0)), - Row(table_select, load_button), tabs, + input_row, sizing_mode='stretch_both', ) @@ -447,7 +461,19 @@ def render_plan(_, old, new): if any(step.expert == 'SQLAgent' for step in plan.steps): self._add_exploration(plan.title, local_memory) index += 1 + + def sync_available_sources_memory(_, __, sources): + """ + For cases when the user uploads a dataset through SourceAgent + this will update the available_sources in the global memory + so that the overview explorer can access it + """ + memory["available_sources"] += [ + source for source in sources if source not in memory["available_sources"] + ] + local_memory.on_change('plan', render_plan) + local_memory.on_change('available_sources', sync_available_sources_memory) def render_output(_, old, new): added = [out for out in new if out not in old]