From 160cd32725316f3ab785deea13dd5969bd92c939 Mon Sep 17 00:00:00 2001
From: Artemy Belousov <whimox@gmail.com>
Date: Thu, 19 Sep 2024 12:53:07 +0400
Subject: [PATCH] Fix CrewAI delegation (#79)

---
 motleycrew/agents/abstract_parent.py |  9 ++++--
 motleycrew/agents/crewai/crewai.py   |  5 ++--
 motleycrew/agents/parent.py          | 42 ++++++++++++++++++++++------
 motleycrew/tools/tool.py             | 16 ++---------
 4 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/motleycrew/agents/abstract_parent.py b/motleycrew/agents/abstract_parent.py
index 53ab36e..e509148 100644
--- a/motleycrew/agents/abstract_parent.py
+++ b/motleycrew/agents/abstract_parent.py
@@ -23,6 +23,11 @@ def invoke(
         pass
 
     @abstractmethod
-    def call_as_tool(self, *args, **kwargs) -> Any:
-        """Method that is called when the agent is used as a tool by another agent."""
+    def as_tool(self, **kwargs) -> Any:
+        """Convert the agent to a tool to be used by other agents via delegation.
+
+        Args:
+            kwargs: Additional arguments to pass to the tool.
+                See :class:`motleycrew.tools.tool.MotleyTool` for more details.
+        """
         pass
diff --git a/motleycrew/agents/crewai/crewai.py b/motleycrew/agents/crewai/crewai.py
index cd9e700..346857f 100644
--- a/motleycrew/agents/crewai/crewai.py
+++ b/motleycrew/agents/crewai/crewai.py
@@ -157,7 +157,7 @@ def from_agent(
         wrapped_agent._agent = agent
         return wrapped_agent
 
-    def as_tool(self) -> MotleyTool:
+    def as_tool(self, **kwargs) -> MotleyTool:
         if not self.description:
             raise ValueError("Agent must have a description to be called as a tool")
 
@@ -182,5 +182,6 @@ def call_agent(prompt: str, expected_output: str):
                 description=self.description,
                 func=call_agent,
                 args_schema=CrewAIAgentInputSchema,
-            )
+            ),
+            **kwargs,
         )
diff --git a/motleycrew/agents/parent.py b/motleycrew/agents/parent.py
index e687844..3af2496 100644
--- a/motleycrew/agents/parent.py
+++ b/motleycrew/agents/parent.py
@@ -6,6 +6,7 @@
 from langchain_core.messages import BaseMessage
 from langchain_core.prompts.chat import ChatPromptTemplate, HumanMessage, SystemMessage
 from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import Tool
 
 from motleycrew.agents.abstract_parent import MotleyAgentAbstractParent
 from motleycrew.common import MotleyAgentFactory, MotleySupportedTool, logger
@@ -111,7 +112,9 @@ def compose_prompt(
                 prompt_messages += self.prompt_prefix.invoke(input_dict).to_messages()
 
             elif isinstance(self.prompt_prefix, str):
-                prompt_messages.append(SystemMessage(content=self.prompt_prefix.format(**input_dict)))
+                prompt_messages.append(
+                    SystemMessage(content=self.prompt_prefix.format(**input_dict))
+                )
 
             else:
                 raise ValueError("Agent description must be a string or a ChatPromptTemplate")
@@ -211,15 +214,36 @@ def add_tools(self, tools: Sequence[MotleySupportedTool]):
                     )
                 self.tools[motley_tool.name] = motley_tool
 
-    def call_as_tool(self, *args, **kwargs):
-        """Method that is called when the agent is used as a tool by another agent."""
+    def as_tool(self, **kwargs) -> MotleyTool:
+        """Convert the agent to a tool to be used by other agents via delegation.
 
-        # TODO: this thing is hacky, we should have a better way to pass structured input
-        if args:
-            return self.invoke({"prompt": args[0]})
-        if len(kwargs) == 1:
-            return self.invoke({"prompt": list(kwargs.values())[0]})
-        return self.invoke(kwargs)
+        Args:
+            kwargs: Additional arguments to pass to the tool.
+                See :class:`motleycrew.tools.tool.MotleyTool` for more details.
+        """
+
+        if not getattr(self, "name", None) or not getattr(self, "description", None):
+            raise ValueError("Agent must have a name and description to be called as a tool")
+
+        def call_as_tool(self, *args, **kwargs):
+            # TODO: this thing is hacky, we should have a better way to pass structured input
+            if args:
+                return self.invoke({"prompt": args[0]})
+            if len(kwargs) == 1:
+                return self.invoke({"prompt": list(kwargs.values())[0]})
+            return self.invoke(kwargs)
+
+        # To be specialized if we expect structured input
+        return MotleyTool.from_langchain_tool(
+            Tool(
+                name=self.name.replace(
+                    " ", "_"
+                ).lower(),  # OpenAI doesn't accept spaces in function names
+                description=self.description,
+                func=call_as_tool,
+            ),
+            **kwargs,
+        )
 
     @abstractmethod
     def invoke(
diff --git a/motleycrew/tools/tool.py b/motleycrew/tools/tool.py
index a43b72b..5bd33be 100644
--- a/motleycrew/tools/tool.py
+++ b/motleycrew/tools/tool.py
@@ -247,20 +247,8 @@ def from_motley_agent(
             The tool representation of the agent.
         """
 
-        if not getattr(agent, "name", None) or not getattr(agent, "description", None):
-            raise ValueError("Agent must have a name and description to be called as a tool")
-
-        # To be specialized if we expect structured input
-        return MotleyTool.from_langchain_tool(
-            Tool(
-                name=agent.name.replace(
-                    " ", "_"
-                ).lower(),  # OpenAI doesn't accept spaces in function names
-                description=agent.description,
-                func=agent.call_as_tool,
-            ),
-            return_direct=return_direct,
-            exceptions_to_reflect=exceptions_to_reflect,
+        return agent.as_tool(
+            return_direct=return_direct, exceptions_to_reflect=exceptions_to_reflect
         )
 
     @staticmethod