Add join chain of thought and fix scroll (#597)

holoviz · Jul 8, 2024 · f3e5609 · f3e5609
1 parent a2e74f5
commit f3e5609
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/lumen/ai/agents.py b/lumen/ai/agents.py
@@ -566,15 +566,16 @@ async def answer(self, messages: list | str):
             return None
 
         with self.interface.add_step(title="Checking if join is required") as step:
-            join_required = (await self.llm.invoke(
+            join = (await self.llm.invoke(
                 messages,
                 system="Determine whether a table join is required to answer the user's query.",
                 response_model=JoinRequired,
                 allow_partial=False,
-            )).join_required
-            step.success_title = 'Query requires join' if join_required else 'No join required'
+            ))
+            step.stream(join.chain_of_thought)
+            step.success_title = 'Query requires join' if join.join_required else 'No join required'
 
-        if join_required:
+        if join.join_required:
             available_tables = source.get_tables()
             with self.interface.add_step(title="Determining tables required for join") as step:
                 tables = (await self.llm.invoke(
@@ -583,7 +584,7 @@ async def answer(self, messages: list | str):
                     response_model=TableJoins,
                     allow_partial=False,
                 )).tables
-                step.stream(f'Join requires following tables: {tables}')
+                step.stream(f'\nJoin requires following tables: {tables}')
                 step.success_title = 'Found tables required for join'
         else:
             tables = [table]

diff --git a/lumen/ai/assistant.py b/lumen/ai/assistant.py
@@ -118,7 +118,7 @@ def download_notebook():
 
         if interface is None:
             interface = ChatInterface(
-                callback=self._chat_invoke, load_buffer=5,
+                callback=self._chat_invoke, load_buffer=10,
             )
         else:
             interface.callback = self._chat_invoke