add processing of the final output_handler call

motleycrew/agents/parent.py

@@ -29,6 +29,31 @@ def __init__(self, output: Any):
         self.output = output
 
 
+def _run_tool_direct_decorator(func: Callable):
+    """Decorator of the tool's _run method, for intercepting a DirectOutput exception"""
+
+    def wrapper(*args, **kwargs):
+        try:
+            result = func(*args, **kwargs)
+        except DirectOutput as direct_exc:
+            return direct_exc
+        return result
+
+    return wrapper
+
+
+def run_tool_direct_decorator(func: Callable):
+    """Decorator of the tool's run method, for intercepting a DirectOutput exception"""
+
+    def wrapper(*args, **kwargs):
+        result = func(*args, **kwargs)
+        if isinstance(result, DirectOutput):
+            raise result
+        return result
+
+    return wrapper
+
+
 class MotleyAgentParent(MotleyAgentAbstractParent, Runnable):
     def __init__(
         self,
@@ -150,6 +175,16 @@ def handle_agent_output(*args, **kwargs):
             args_schema=self.output_handler.args_schema,
         )
 
+        object.__setattr__(
+            prepared_output_handler,
+            "_run",
+            _run_tool_direct_decorator(prepared_output_handler._run),
+        )
+
+        object.__setattr__(
+            prepared_output_handler, "run", run_tool_direct_decorator(prepared_output_handler.run)
+        )
+
         return MotleyTool.from_langchain_tool(prepared_output_handler)
 
     @staticmethod