Fix workflow issue in tool calling

Author: leehuwuj

templates/components/agents/python/form_filling/app/engine/workflow.py

@@ -32,7 +32,7 @@ class ExtractMissingCellsEvent(Event):
 
 
 class FindAnswersEvent(Event):
-    tool_call: ToolSelection
+    missing_cells: list[MissingCell]
 
 
 class FillEvent(Event):
@@ -90,8 +90,7 @@ def __init__(
         self.query_engine_tool = query_engine_tool
         self.extractor_tool = extractor_tool
         self.filling_tool = filling_tool
-        self.tools = [self.query_engine_tool, self.extractor_tool, self.filling_tool]
-        self.llm: FunctionCallingLLM = llm or Settings.llm
+        self.llm = llm or Settings.llm
         if not isinstance(self.llm, FunctionCallingLLM):
             raise ValueError("FormFillingWorkflow only supports FunctionCallingLLM.")
         self.memory = ChatMemoryBuffer.from_defaults(
@@ -121,7 +120,7 @@ async def handle_llm_input(
         self,
         ctx: Context,
         ev: InputEvent,
-    ) -> ExtractMissingCellsEvent | FindAnswersEvent | FillEvent | StopEvent:
+    ) -> ExtractMissingCellsEvent | FillEvent | StopEvent:
         """
         Handle an LLM input and decide the next step.
         """
@@ -137,8 +136,6 @@ async def handle_llm_input(
             for tool_call in tool_calls:
                 if tool_call.tool_name == self.extractor_tool.metadata.get_name():
                     return ExtractMissingCellsEvent(tool_call=tool_call)
-                elif tool_call.tool_name == self.query_engine_tool.metadata.get_name():
-                    return FindAnswersEvent(tool_call=tool_call)
                 elif tool_call.tool_name == self.filling_tool.metadata.get_name():
                     return FillEvent(tool_call=tool_call)
         # If no tool call, return the generator
@@ -147,7 +144,7 @@ async def handle_llm_input(
     @step()
     async def extract_missing_cells(
         self, ctx: Context, ev: ExtractMissingCellsEvent
-    ) -> InputEvent:
+    ) -> InputEvent | FindAnswersEvent:
         """
         Extract missing cells in a CSV file and generate questions to fill them.
         """
@@ -168,7 +165,6 @@ async def extract_missing_cells(
             return InputEvent(input=self.memory.get())
 
         missing_cells = response.raw_output.get("missing_cells", [])
-        ctx.data["missing_cells"] = missing_cells
         message = ChatMessage(
             role=MessageRole.TOOL,
             content=str(missing_cells),
@@ -179,8 +175,8 @@ async def extract_missing_cells(
         )
         self.memory.put(message)
 
-        # send input event back with updated chat history
-        return InputEvent(input=self.memory.get())
+        # Forward missing cells information to find answers step
+        return FindAnswersEvent(missing_cells=missing_cells)
 
     @step()
     async def find_answers(self, ctx: Context, ev: FindAnswersEvent) -> InputEvent:
@@ -193,7 +189,7 @@ async def find_answers(self, ctx: Context, ev: FindAnswersEvent) -> InputEvent:
                 msg="Finding answers for missing cells",
             )
         )
-        missing_cells = ctx.data.get("missing_cells", None)
+        missing_cells = ev.missing_cells
         # If missing cells information is not found, fallback to other tools
         # It means that the extractor tool has not been called yet
         # Fallback to input
@@ -220,8 +216,7 @@ async def find_answers(self, ctx: Context, ev: FindAnswersEvent) -> InputEvent:
         # and stream the progress
         progress_id = str(uuid.uuid4())
         total_steps = len(missing_cells)
-        for i, missing_cell in enumerate(missing_cells):
-            cell = MissingCell(**missing_cell)
+        for i, cell in enumerate(missing_cells):
             if cell.question_to_answer is None:
                 continue
             ctx.write_event_to_stream(
@@ -248,15 +243,12 @@ async def find_answers(self, ctx: Context, ev: FindAnswersEvent) -> InputEvent:
                     value=str(answer),
                 )
             )
-        message = ChatMessage(
-            role=MessageRole.TOOL,
-            content=str(cell_values),
-            additional_kwargs={
-                "tool_call_id": ev.tool_call.tool_id,
-                "name": ev.tool_call.tool_name,
-            },
+        self.memory.put(
+            ChatMessage(
+                role=MessageRole.ASSISTANT,
+                content=str(cell_values),
+            )
         )
-        self.memory.put(message)
         return InputEvent(input=self.memory.get())
 
     @step()
@@ -295,7 +287,8 @@ async def _tool_call_generator(
         self, chat_history: list[ChatMessage]
     ) -> AsyncGenerator[ChatMessage | bool, None]:
         response_stream = await self.llm.astream_chat_with_tools(
-            self.tools, chat_history=chat_history
+            [self.extractor_tool, self.filling_tool],
+            chat_history=chat_history,
         )
 
         full_response = None
@@ -321,7 +314,6 @@ async def _tool_call_generator(
             self.memory.put(full_response.message)
             yield full_response
 
-    # TODO: Implement a _acall_tool method
     def _call_tool(
         self,
         ctx: Context,