From f8a1d4f41457a3be845a5e6d5c520e7e872bf299 Mon Sep 17 00:00:00 2001
From: Chris Pang <chrispangg@gmail.com>
Date: Tue, 12 Mar 2024 06:40:10 +1100
Subject: [PATCH] added langchain callback to agents (#333)

Co-authored-by: Chris Pang <chris_pang@racv.com.au>
---
 src/crewai/agent.py | 7 +++++++
 src/crewai/crew.py  | 6 ++++++
 2 files changed, 13 insertions(+)

diff --git a/src/crewai/agent.py b/src/crewai/agent.py
index a26085d48b..1f1de0ae7e 100644
--- a/src/crewai/agent.py
+++ b/src/crewai/agent.py
@@ -8,6 +8,8 @@
 from langchain.tools.render import render_text_description
 from langchain_core.agents import AgentAction
 from langchain_openai import ChatOpenAI
+from langchain_core.callbacks import BaseCallbackHandler
+
 from pydantic import (
     UUID4,
     BaseModel,
@@ -46,6 +48,7 @@ class Agent(BaseModel):
             allow_delegation: Whether the agent is allowed to delegate tasks to other agents.
             tools: Tools at agents disposal
             step_callback: Callback to be executed after each step of the agent execution.
+            callbacks: A list of callback functions from the langchain library that are triggered during the agent's execution process
     """
 
     __hash__ = object.__hash__  # type: ignore
@@ -110,6 +113,9 @@ class Agent(BaseModel):
     function_calling_llm: Optional[Any] = Field(
         description="Language model that will run the agent.", default=None
     )
+    callbacks: Optional[List[InstanceOf[BaseCallbackHandler]]] = Field(
+        default=None, description="Callback to be executed"
+    )
 
     def __init__(__pydantic_self__, **data):
         config = data.pop("config", {})
@@ -245,6 +251,7 @@ def create_agent_executor(self, tools=None) -> None:
             "step_callback": self.step_callback,
             "tools_handler": self.tools_handler,
             "function_calling_llm": self.function_calling_llm,
+            "callbacks": self.callbacks
         }
 
         if self._rpm_controller:
diff --git a/src/crewai/crew.py b/src/crewai/crew.py
index 0eeeb991e8..85e83f228e 100644
--- a/src/crewai/crew.py
+++ b/src/crewai/crew.py
@@ -2,6 +2,8 @@
 import uuid
 from typing import Any, Dict, List, Optional, Union
 
+from langchain_core.callbacks import BaseCallbackHandler
+
 from pydantic import (
     UUID4,
     BaseModel,
@@ -32,6 +34,7 @@ class Crew(BaseModel):
         tasks: List of tasks assigned to the crew.
         agents: List of agents part of this crew.
         manager_llm: The language model that will run manager agent.
+        manager_callbacks: The callback handlers to be executed by the manager agent when hierarchical process is used
         function_calling_llm: The language model that will run the tool calling for all the agents.
         process: The process flow that the crew will follow (e.g., sequential).
         verbose: Indicates the verbosity level for logging during execution.
@@ -64,6 +67,9 @@ class Crew(BaseModel):
     manager_llm: Optional[Any] = Field(
         description="Language model that will run the agent.", default=None
     )
+    manager_callbacks: Optional[List[InstanceOf[BaseCallbackHandler]]] = Field(
+        default=None, description="A list of callback handlers to be executed by the manager agent when hierarchical process is used"
+    )
     function_calling_llm: Optional[Any] = Field(
         description="Language model that will run the agent.", default=None
     )