Merge pull request #2 from bonk1t/dev/integrate-langchain-callbacks

Integrate langchain callbacks
VRSEN · Dec 19, 2024 · baced53 · baced53
2 parents d5ca39f + e3f5e31
commit baced53
Show file tree

Hide file tree

Showing 27 changed files with 996 additions and 1,047 deletions.
diff --git a/agency_swarm/__init__.py b/agency_swarm/__init__.py
@@ -2,28 +2,26 @@
 from .agents import Agent
 from .tools import BaseTool
 from .util import (
+    get_callback_handler,
     get_openai_client,
-    get_tracker,
+    init_tracking,
     llm_validator,
+    set_callback_handler,
     set_openai_client,
     set_openai_key,
-    set_tracker,
-)
-from .util.streaming import (
-    AgencyEventHandler,
-    AgencyEventHandlerWithTracking,
 )
+from .util.streaming import AgencyEventHandler
 
 __all__ = [
     "Agency",
     "Agent",
     "BaseTool",
     "AgencyEventHandler",
-    "AgencyEventHandlerWithTracking",
     "get_openai_client",
     "set_openai_client",
     "set_openai_key",
     "llm_validator",
-    "set_tracker",
-    "get_tracker",
+    "init_tracking",
+    "get_callback_handler",
+    "set_callback_handler",
 ]
diff --git a/agency_swarm/agency/agency.py b/agency_swarm/agency/agency.py
@@ -30,7 +30,6 @@
 from agency_swarm.user import User
 from agency_swarm.util.errors import RefusalError
 from agency_swarm.util.files import get_file_purpose, get_tools
-from agency_swarm.util.oai import get_tracker
 from agency_swarm.util.shared_state import SharedState
 from agency_swarm.util.streaming import (
     AgencyEventHandler,
@@ -68,6 +67,7 @@ def __init__(
         max_prompt_tokens: int = None,
         max_completion_tokens: int = None,
         truncation_strategy: dict = None,
+        callback_handler: Optional[Any] = None,
     ):
         """
         Initializes the Agency object, setting up agents, threads, and core functionalities.
@@ -86,6 +86,7 @@ def __init__(
             max_prompt_tokens (int, optional): The maximum number of tokens allowed in the prompt for each agent. Agent-specific values will override this. Defaults to None.
             max_completion_tokens (int, optional): The maximum number of tokens allowed in the completion for each agent. Agent-specific values will override this. Defaults to None.
             truncation_strategy (dict, optional): The truncation strategy to use for the completion for each agent. Agent-specific values will override this. Defaults to None.
+            callback_handler (Any, optional): The callback handler to use for tracking events. Defaults to None.
 
         This constructor initializes various components of the Agency, including CEO, agents, threads, and user interactions. It parses the agency chart to set up the organizational structure and initializes the messaging tools, agents, and threads necessary for the operation of the agency. Additionally, it prepares a main thread for user interactions.
         """
@@ -107,6 +108,7 @@ def __init__(
         self.max_prompt_tokens = max_prompt_tokens
         self.max_completion_tokens = max_completion_tokens
         self.truncation_strategy = truncation_strategy
+        self.callback_handler = callback_handler
 
         # set thread type based send_message_tool_class async mode
         if (
@@ -154,7 +156,6 @@ def __init__(
         self._create_special_tools()
         self._init_agents()
 
-    @get_tracker().get_observe_decorator()
     def get_completion(
         self,
         message: str,
@@ -210,7 +211,6 @@ def get_completion(
 
         return res
 
-    @get_tracker().get_observe_decorator()
     def get_completion_stream(
         self,
         message: str,

diff --git a/agency_swarm/agency/genesis/GenesisCEO/tools/FinalizeAgency.py b/agency_swarm/agency/genesis/GenesisCEO/tools/FinalizeAgency.py
@@ -34,7 +34,7 @@ def run(self):
             f.close()
 
         res = client.chat.completions.create(
-            model="gpt-3.5-turbo",
+            model="gpt-4o-mini",
             messages=examples
             + [
                 {"role": "user", "content": agency_py},

diff --git a/agency_swarm/agents/BrowsingAgent/tools/WebPageSummarizer.py b/agency_swarm/agents/BrowsingAgent/tools/WebPageSummarizer.py
@@ -22,7 +22,7 @@ def run(self):
         content = " ".join(content.split()[:10000])
 
         completion = client.chat.completions.create(
-            model="gpt-3.5-turbo",
+            model="gpt-4o-mini",
             messages=[
                 {
                     "role": "system",

diff --git a/agency_swarm/agents/agent.py b/agency_swarm/agents/agent.py
@@ -18,7 +18,7 @@
 )
 from agency_swarm.tools.oai.FileSearch import FileSearchConfig
 from agency_swarm.util.constants import DEFAULT_MODEL
-from agency_swarm.util.oai import get_openai_client, get_tracker
+from agency_swarm.util.oai import get_openai_client
 from agency_swarm.util.openapi import validate_openapi_spec
 from agency_swarm.util.shared_state import SharedState
 
@@ -295,7 +295,6 @@ def init_oai(self):
 
         return self
 
-    @get_tracker().get_observe_decorator()
     def _create_assistant(self):
         """Creates a new OpenAI assistant with the agent's current configuration."""
         return self.client.beta.assistants.create(