Fix model name

VRSEN · Dec 26, 2024 · fc8ca07 · fc8ca07
1 parent a7f02ee
commit fc8ca07
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/agency_swarm/threads/thread.py b/agency_swarm/threads/thread.py
@@ -22,7 +22,6 @@
 from agency_swarm.util.tracking import get_callback_handler
 from agency_swarm.util.tracking.langchain_types import (
     AgentAction,
-    AgentFinish,
     HumanMessage,
 )
 
@@ -297,6 +296,11 @@ def _create_run(
         response_format: dict | None = None,
     ):
         try:
+            model_name = (
+                self.agent.model
+                if isinstance(self.agent, Agent)
+                else recipient_agent.model
+            )
             if event_handler:
                 with self.client.beta.threads.runs.stream(
                     thread_id=self.id,
@@ -315,6 +319,7 @@ def _create_run(
                         "sender_agent_name": self.agent.name,
                         "recipient_agent_name": recipient_agent.name,
                         "run_status": "created",
+                        "ls_model_name": model_name,
                     },
                     response_format=response_format,
                 ) as stream:
@@ -336,6 +341,7 @@ def _create_run(
                         "sender_agent_name": self.agent.name,
                         "recipient_agent_name": recipient_agent.name,
                         "run_status": "created",
+                        "ls_model_name": model_name,
                     },
                 )
                 self._run = self.client.beta.threads.runs.poll(
@@ -500,6 +506,7 @@ def execute_tool(
                     "agent_name": self.agent.name,
                     "recipient_agent_name": recipient_agent.name,
                     "run_status": self._run.status,
+                    "ls_model_name": self._run.model,
                 },
             )
 
@@ -522,6 +529,7 @@ def execute_tool(
                     "agent_name": self.agent.name,
                     "recipient_agent_name": recipient_agent.name,
                     "run_status": self._run.status,
+                    "ls_model_name": self._run.model,
                 },
             )
 
@@ -914,6 +922,7 @@ def _on_chain_start(
                     "agent_name": self.agent.name,
                     "recipient_agent_name": recipient_agent.name,
                     "run_status": self._run.status,
+                    "ls_model_name": self._run.model,
                 },
             )
 
@@ -939,6 +948,7 @@ def _on_chat_model_start(
                     "agent_name": self.agent.name,
                     "recipient_agent_name": recipient_agent.name,
                     "run_status": self._run.status,
+                    "ls_model_name": self._run.model,
                 },
                 invocation_params={
                     "_type": "openai",

diff --git a/agency_swarm/util/tracking/__init__.py b/agency_swarm/util/tracking/__init__.py
@@ -43,7 +43,6 @@ def init_tracking(tracker_name: SUPPORTED_TRACKERS_TYPE, **kwargs):
         from agentops import LangchainCallbackHandler
 
         handler_class = LangchainCallbackHandler
-        kwargs["ignore_chat_model"] = True
 
     elif tracker_name == "langfuse":
         from langfuse.callback import CallbackHandler