diff --git a/lumen/ai/agents.py b/lumen/ai/agents.py index 46d6cc30..a56fb6e9 100644 --- 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 index 989c9697..c5687486 100644 --- 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