From 34bb50e386e5f18622afb9d048c353de40bc43af Mon Sep 17 00:00:00 2001
From: Pratyush Shukla <ps4534@nyu.edu>
Date: Thu, 30 Jan 2025 22:22:16 +0530
Subject: [PATCH] streamlined support for agentops

---
 src/crewai/agents/tools_handler.py | 17 +++++++++
 src/crewai/tools/tool_usage.py     | 59 ++++++++++++++++++++++++++----
 2 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/src/crewai/agents/tools_handler.py b/src/crewai/agents/tools_handler.py
index fd4bec7ee..4f7139147 100644
--- a/src/crewai/agents/tools_handler.py
+++ b/src/crewai/agents/tools_handler.py
@@ -21,9 +21,26 @@ class ToolsHandler:
         calling: Union[ToolCalling, InstructorToolCalling],
         output: str,
         should_cache: bool = True,
+        agentops: Optional[Any] = None,
     ) -> Any:
         """Run when tool ends running."""
         self.last_used_tool = calling  # type: ignore # BUG?: Incompatible types in assignment (expression has type "Union[ToolCalling, InstructorToolCalling]", variable has type "ToolCalling")
+
+        if agentops:
+            agentops.record(
+                agentops.ActionEvent(
+                    name=calling.tool_name,
+                    action_type="on_tool_use",
+                    params=calling.arguments,
+                    returns=output,
+                    logs={
+                        "tool": calling.tool_name,
+                        "tool_calling": calling,
+                        "should_cache": should_cache,
+                    },
+                )
+            )
+
         if self.cache and should_cache and calling.tool_name != CacheTools().name:
             self.cache.add(
                 tool=calling.tool_name,
diff --git a/src/crewai/tools/tool_usage.py b/src/crewai/tools/tool_usage.py
index 218410ef7..9e11e87fd 100644
--- a/src/crewai/tools/tool_usage.py
+++ b/src/crewai/tools/tool_usage.py
@@ -116,7 +116,10 @@ class ToolUsage:
                 self._printer.print(content=f"\n\n{error}\n", color="red")
             return error
 
-        if isinstance(tool, CrewStructuredTool) and tool.name == self._i18n.tools("add_image")["name"]:  # type: ignore
+        if (
+            isinstance(tool, CrewStructuredTool)
+            and tool.name == self._i18n.tools("add_image")["name"]
+        ):  # type: ignore
             try:
                 result = self._use(tool_string=tool_string, tool=tool, calling=calling)
                 return result
@@ -136,7 +139,6 @@ class ToolUsage:
         tool: Any,
         calling: Union[ToolCalling, InstructorToolCalling],
     ) -> str:  # TODO: Fix this return type
-        tool_event = agentops.ToolEvent(name=calling.tool_name) if agentops else None  # type: ignore
         if self._check_tool_repeated_usage(calling=calling):  # type: ignore # _check_tool_repeated_usage of "ToolUsage" does not return a value (it only ever returns None)
             try:
                 result = self._i18n.errors("task_repeated_usage").format(
@@ -181,7 +183,9 @@ class ToolUsage:
 
                 if calling.arguments:
                     try:
-                        acceptable_args = tool.args_schema.model_json_schema()["properties"].keys()  # type: ignore
+                        acceptable_args = tool.args_schema.model_json_schema()[
+                            "properties"
+                        ].keys()  # type: ignore
                         arguments = {
                             k: v
                             for k, v in calling.arguments.items()
@@ -214,7 +218,19 @@ class ToolUsage:
                 self.task.increment_tools_errors()
                 if agentops:
                     agentops.record(
-                        agentops.ErrorEvent(exception=e, trigger_event=tool_event)
+                        agentops.ErrorEvent(
+                            exception=e,
+                            logs={
+                                "tool_string": tool_string,
+                                "tool": tool,
+                                "tool_calling": calling,
+                                "error": error,
+                                "run_attempts": self._run_attempts,
+                                "llm": self.function_calling_llm,
+                                "task": self.task,
+                                "agent": self.agent,
+                            },
+                        )
                     )
                 return self.use(calling=calling, tool_string=tool_string)  # type: ignore # No return value expected
 
@@ -229,11 +245,12 @@ class ToolUsage:
                     )
 
                 self.tools_handler.on_tool_use(
-                    calling=calling, output=result, should_cache=should_cache
+                    calling=calling,
+                    output=result,
+                    should_cache=should_cache,
+                    agentops=agentops,
                 )
 
-        if agentops:
-            agentops.record(tool_event)
         self._telemetry.tool_usage(
             llm=self.function_calling_llm,
             tool_name=tool.name,
@@ -460,6 +477,19 @@ class ToolUsage:
 
     def on_tool_error(self, tool: Any, tool_calling: ToolCalling, e: Exception) -> None:
         event_data = self._prepare_event_data(tool, tool_calling)
+
+        if agentops:
+            agentops.record(
+                agentops.ErrorEvent(
+                    exception=e,
+                    logs={
+                        "tool": tool,
+                        "tool_calling": tool_calling,
+                        "event_data": event_data,
+                    },
+                )
+            )
+
         events.emit(
             source=self, event=ToolUsageError(**{**event_data, "error": str(e)})
         )
@@ -476,6 +506,21 @@ class ToolUsage:
                 "from_cache": from_cache,
             }
         )
+
+        if agentops:
+            agentops.record(
+                agentops.ActionEvent(
+                    name=tool_calling.tool_name,
+                    action_type="on_tool_use_finished",
+                    params=tool_calling.arguments,
+                    returns=event_data,
+                    logs={
+                        "tool": tool,
+                        "tool_calling": tool_calling,
+                    },
+                )
+            )
+
         events.emit(source=self, event=ToolUsageFinished(**event_data))
 
     def _prepare_event_data(self, tool: Any, tool_calling: ToolCalling) -> dict: