search across all sources (#676)

lumen/ai/agents.py

@@ -344,7 +344,6 @@ def _render_lumen(self, component: Component, message: pn.chat.ChatMessage = Non
         message_kwargs = dict(value=out, user=self.user)
         self.interface.stream(message=message, **message_kwargs, replace=True, max_width=self._max_width)
 
-
 class TableAgent(LumenBaseAgent):
     """
     Displays a single table / dataset. Does not discuss.
@@ -366,7 +365,22 @@ def _create_table_model(tables):
         return table_model
 
     async def answer(self, messages: list | str):
-        tables = tuple(memory["current_source"].get_tables())
+
+        if len(memory["available_sources"]) >= 1:
+            available_sources = memory["available_sources"]
+            tables_to_source = {}
+            tables_schema_str = "\nHere are the table schemas\n"
+            for source in available_sources:
+                for table in source.get_tables():
+                    tables_to_source[table] = source
+                    schema = get_schema(source, table, include_min_max=False, include_enum=False, limit=1)
+                    tables_schema_str += f"### {table}\n```yaml\n{yaml.safe_dump(schema)}```\n"
+        else:
+            source = memory["current_source"]
+            tables_to_source = {table: source for table in source.get_tables()}
+            tables_schema_str = ""
+
+        tables = tuple(tables_to_source)
         if len(tables) == 1:
             table = tables[0]
         else:
@@ -376,7 +390,7 @@ async def answer(self, messages: list | str):
                     tables = closest_tables
                 elif len(tables) > FUZZY_TABLE_LENGTH:
                     tables = await self._get_closest_tables(messages, tables)
-                system_prompt = await self._system_prompt_with_context(messages)
+                system_prompt = await self._system_prompt_with_context(messages) + tables_schema_str
                 if self.debug:
                     print(f"{self.name} is being instructed that it should {system_prompt}")
                 if len(tables) > 1:
@@ -391,6 +405,8 @@ async def answer(self, messages: list | str):
                 else:
                     table = tables[0]
                 step.stream(f"Selected table: {table}")
+
+        memory["current_source"] = tables_to_source[table]
         memory["current_table"] = table
         memory["current_pipeline"] = pipeline = Pipeline(
             source=memory["current_source"], table=table

lumen/ai/utils.py

@@ -98,12 +98,18 @@ def format_schema(schema):
 
 
 def get_schema(
-    source: Source | Pipeline, table: str | None = None, include_min_max: bool = True
+    source: Source | Pipeline,
+    table: str | None = None,
+    include_min_max: bool = True,
+    include_enum: bool = True,
+    **get_kwargs
 ):
     if isinstance(source, Pipeline):
         schema = source.get_schema()
     else:
-        schema = source.get_schema(table, limit=100)
+        if "limit" not in get_kwargs:
+            get_kwargs["limit"] = 100
+        schema = source.get_schema(table, **get_kwargs)
     schema = dict(schema)
 
     if include_min_max:
@@ -118,6 +124,16 @@ def get_schema(
                 spec.pop("inclusiveMinimum")
             if "inclusiveMaximum" in spec:
                 spec.pop("inclusiveMaximum")
+            if "min" in spec:
+                spec.pop("min")
+            if "max" in spec:
+                spec.pop("max")
+
+    if not include_enum:
+        for field, spec in schema.items():
+            if "enum" in spec:
+                spec.pop("enum")
+
     schema = format_schema(schema)
     return schema