Properly working output handlers for Langchain & LlamaIndex agents (#54)

* add output handler to llama_index agent

* add llama_index_output_handler example

* add check llama_index module installed in run_step_decorator

* transfer run_step_decorator to the LlamaIndexMotleyAgent class

* implementation in a langchain output handler using decorators

* refactor add LangchainOutputHandlerMixin

* add output_handler for CrewAIMotleyAgentParent

* add crewai_output_handler example

* changing the way attributes are saved for decorated methods

* rename agent_plan_decorator

* prepare for merge

* Fix integration tests

* Minor improvements

* add use of output_handler via the agent_finish_blocker tool for the langchain library

* update a user's message change

* add processing of the final output_handler call

* add test langchain output_handler

* fix llama_index output_handler

* add test llama_index output_handler

* update test llama_index output_handler

* refactor langchain output handler runs methods decorated

* fix llama_index run_step_decorator

* elimination of comments

* remove _run_and_catch_output

* Raise NotImplementedError when using output handlers with CrewAI

---------

Co-authored-by: User <[email protected]>
Co-authored-by: whimo <[email protected]>

examples/__init__.py

@@ -0,0 +1 @@
+# Has to be here to make the examples package importable for integration tests

examples/crewai_output_handler.py

@@ -0,0 +1,75 @@
+from pathlib import Path
+import os
+import sys
+
+from dotenv import load_dotenv
+
+from motleycrew.agents.crewai import CrewAIMotleyAgent
+from motleycrew.common import configure_logging
+from motleycrew.tasks import SimpleTask
+from motleycrew.common.exceptions import InvalidOutput
+
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_core.tools import StructuredTool
+
+WORKING_DIR = Path(os.path.realpath(".."))
+
+try:
+    from motleycrew import MotleyCrew
+except ImportError:
+    # if we are running this from source
+    motleycrew_location = os.path.realpath(WORKING_DIR / "..")
+    sys.path.append(motleycrew_location)
+
+
+def main():
+    crew = MotleyCrew()
+
+    search_tool = DuckDuckGoSearchRun()
+
+    def check_output(output: str):
+        if "medicine" not in output.lower():
+            raise InvalidOutput("Add more information about AI applications in medicine.")
+        return {"checked_output": output.lower()}
+
+    output_handler = StructuredTool.from_function(
+        name="output_handler",
+        description="Output handler",
+        func=check_output,
+    )
+
+    researcher = CrewAIMotleyAgent(
+        role="Senior Research Analyst",
+        goal="Uncover cutting-edge developments in AI and data science, doing web search if necessary",
+        backstory="""You work at a leading tech think tank.
+    Your expertise lies in identifying emerging trends.
+    You have a knack for dissecting complex data and presenting actionable insights.""",
+        delegation=False,
+        output_handler=output_handler,
+        verbose=True,
+        tools=[search_tool],
+    )
+
+    # Create tasks for agent
+
+    analysis_report_task = SimpleTask(
+        crew=crew,
+        name="produce comprehensive analysis report on AI advancements",
+        description="""Conduct a comprehensive analysis of the latest advancements in AI in 2024.
+    Identify key trends, breakthrough technologies, and potential industry impacts.
+    Your final answer MUST be a full analysis report""",
+        agent=researcher,
+    )
+
+    # Get your crew to work!
+    result = crew.run()
+
+    # Get the outputs of the task
+    print(analysis_report_task.output)
+    return analysis_report_task.output
+
+
+if __name__ == "__main__":
+    configure_logging(verbose=True)
+    load_dotenv()
+    main()

examples/llama_index_output_handler.py

@@ -0,0 +1,65 @@
+from dotenv import load_dotenv
+from langchain_community.tools import DuckDuckGoSearchRun
+
+
+from motleycrew import MotleyCrew
+from motleycrew.agents.llama_index import ReActLlamaIndexMotleyAgent
+from motleycrew.common import configure_logging
+from motleycrew.tasks import SimpleTask
+from motleycrew.common.exceptions import InvalidOutput
+from motleycrew.common import AsyncBackend
+
+from langchain_core.tools import StructuredTool
+
+
+def main():
+    """Main function of running the example."""
+    search_tool = DuckDuckGoSearchRun()
+
+    def check_output(output: str):
+        if "medicine" not in output.lower():
+            raise InvalidOutput(
+                "Add more information about AI applications in medicine."
+            )
+
+        return {"checked_output": output}
+
+    output_handler = StructuredTool.from_function(
+        name="output_handler",
+        description="Output handler",
+        func=check_output,
+    )
+
+    # TODO: add LlamaIndex native tools
+    researcher = ReActLlamaIndexMotleyAgent(
+        description="Your goal is to uncover cutting-edge developments in AI and data science",
+        tools=[search_tool],
+        output_handler=output_handler,
+        verbose=True,
+        max_iterations=16,  # default is 10, we add more because the output handler may reject the output
+    )
+
+    crew = MotleyCrew(async_backend=AsyncBackend.NONE)
+
+    # Create tasks for your agents
+    task = SimpleTask(
+        crew=crew,
+        name="produce comprehensive analysis report on AI advancements",
+        description="""Conduct a comprehensive analysis of the latest advancements in AI in 2024.
+      Identify key trends, breakthrough technologies, and potential industry impacts.
+      Your final answer MUST be a full analysis report""",
+        agent=researcher,
+    )
+
+    # Get your crew to work!
+    crew.run()
+
+    print(task.output)
+    return task.output
+
+
+if __name__ == "__main__":
+    configure_logging(verbose=True)
+
+    load_dotenv()
+    main()

examples/old/single_llama_index.py

@@ -6,20 +6,25 @@
 from motleycrew.agents.llama_index import ReActLlamaIndexMotleyAgent
 from motleycrew.common import configure_logging
 from motleycrew.tasks import SimpleTask
+from motleycrew.common.exceptions import InvalidOutput
+from motleycrew.common import AsyncBackend
+
+from langchain_core.tools import StructuredTool
 
 
 def main():
     """Main function of running the example."""
     search_tool = DuckDuckGoSearchRun()
 
+
     # TODO: add LlamaIndex native tools
     researcher = ReActLlamaIndexMotleyAgent(
         description="Your goal is to uncover cutting-edge developments in AI and data science",
         tools=[search_tool],
         verbose=True,
     )
 
-    crew = MotleyCrew()
+    crew = MotleyCrew(async_backend=AsyncBackend.NONE)
 
     # Create tasks for your agents
     task = SimpleTask(

motleycrew/agents/crewai/crewai.py

@@ -4,12 +4,13 @@
 
 from langchain_core.runnables import RunnableConfig
 
-from motleycrew.agents.parent import MotleyAgentParent
 from motleycrew.agents.crewai import CrewAIAgentWithConfig
-from motleycrew.common import MotleySupportedTool
+from motleycrew.agents.parent import MotleyAgentParent
 from motleycrew.common import MotleyAgentFactory
-from motleycrew.tracking import add_default_callbacks_to_langchain_config
+from motleycrew.common import MotleySupportedTool
 from motleycrew.common.utils import ensure_module_is_installed
+from motleycrew.tools import MotleyTool
+from motleycrew.tracking import add_default_callbacks_to_langchain_config
 
 try:
     from crewai import Task as CrewAI__Task
@@ -24,6 +25,7 @@ def __init__(
         name: str | None = None,
         agent_factory: MotleyAgentFactory[CrewAIAgentWithConfig] | None = None,
         tools: Sequence[MotleySupportedTool] | None = None,
+        output_handler: MotleySupportedTool | None = None,
         verbose: bool = False,
     ):
         """Description
@@ -35,12 +37,20 @@ def __init__(
             tools (:obj:`Sequence[MotleySupportedTool]`, optional:
             verbose (bool):
         """
+
+        if output_handler:
+            raise NotImplementedError(
+                "Output handler is not supported for CrewAI agents "
+                "because of the specificity of CrewAi's prompts."
+            )
+
         ensure_module_is_installed("crewai")
         super().__init__(
             description=goal,
             name=name,
             agent_factory=agent_factory,
             tools=tools,
+            output_handler=output_handler,
             verbose=verbose,
         )
 
@@ -68,10 +78,16 @@ def invoke(
         crewai_task = CrewAI__Task(description=prompt)
 
         output = self.agent.execute_task(
-            task=crewai_task, context=input.get("context"), tools=langchain_tools, config=config
+            task=crewai_task,
+            context=input.get("context"),
+            tools=langchain_tools,
+            config=config,
         )
         return output
 
+    def materialize(self):
+        super().materialize()
+
     # TODO: what do these do?
     def set_cache_handler(self, cache_handler: Any) -> None:
         """Description

motleycrew/agents/crewai/crewai_agent.py

@@ -19,6 +19,7 @@ def __init__(
         delegation: bool = False,
         tools: Sequence[MotleySupportedTool] | None = None,
         llm: Optional[Any] = None,
+        output_handler: MotleySupportedTool | None = None,
         verbose: bool = False,
     ):
         """ Description
@@ -62,5 +63,6 @@ def agent_factory(tools: dict[str, MotleyTool]):
             name=role,
             agent_factory=agent_factory,
             tools=tools,
+            output_handler=output_handler,
             verbose=verbose,
         )