react v2 with ablity to train functiontool

SylphAI-Inc · Dec 26, 2024 · 4db576e · 4db576e
1 parent 117ae1b
commit 4db576e
Show file tree

Hide file tree

Showing 11 changed files with 1,215 additions and 124 deletions.
diff --git a/adalflow/adalflow/components/agent/react.py b/adalflow/adalflow/components/agent/react.py
@@ -143,6 +143,26 @@ def call(
         return output
 
 
+class FunctionOutputToStepOutput(GradComponent):
+    def __init__(self):
+        super().__init__()
+        self.name = "FunctionOutputToStepOutput"
+        self._component_desc = "Convert the FunctionOutput to StepOutput."
+
+    def call(self, output: FunctionOutput, step_output: StepOutput) -> StepOutput:
+        """Convert the FunctionOutput to StepOutput."""
+
+        temp_result = output.output
+        if isinstance(temp_result, Parameter):
+            step_output.observation = temp_result.data
+        else:
+            step_output.observation = temp_result
+        return step_output
+        # step_output = StepOutput(step=step)
+        # step_output.observation = output.output
+        # return step_output
+
+
 # TODO: make execute_action_fn to a GradComponent to enable the training of the tools too.
 def execute_action_fn(
     x: GeneratorOutput, step_output: StepOutput, step: int, execute_action: Any, id=None
@@ -298,6 +318,7 @@ def __init__(
         # added this component to the computation graph
         self.append_step_history = AppendStepHistory()
         self.execute_action = ExecuteAction()
+        self.function_output_to_step_output = FunctionOutputToStepOutput()
 
     def _init_tools(
         self,
@@ -353,13 +374,19 @@ def _execute_action(
             # replace the id
             fun.kwargs["id"] = id
 
-            result: FunctionOutput = self.tool_manager.execute_func(fun)
-            # TODO: optimize the action_step
+            result: Union[FunctionOutput, Parameter] = self.tool_manager(fun)
             action_step.function = fun
-            if isinstance(result.output, Parameter):
-                action_step.observation = result.output.data
+            if isinstance(result, Parameter):
+                result.add_successor_map_fn(
+                    successor=self.function_output_to_step_output,
+                    map_fn=lambda x: x.data,
+                )
+                action_step: StepOutput = self.function_output_to_step_output(
+                    output=result, step_output=action_step
+                )
             else:
                 action_step.observation = result.output
+
             return action_step
         except Exception as e:
             log.error(f"Error executing {action}: {e}")