feat: trace through tool call

Signed-off-by: Casper Guldbech Nielsen <scni@novonordisk.com>

dapr_agents/tool/base.py

@@ -8,6 +8,15 @@ from dapr_agents.tool.utils.tool import ToolHelper
 from dapr_agents.tool.utils.function_calling import to_function_call_definition
 from dapr_agents.types import ToolError
 
+from pydantic import PrivateAttr
+from dapr_agents.agent.telemetry import (
+    span_decorator,
+    async_span_decorator,
+)
+
+from opentelemetry import trace
+from opentelemetry.trace import Tracer, Status, StatusCode
+
 logger = logging.getLogger(__name__)
 
 
@@ -37,6 +46,7 @@ class AgentTool(BaseModel):
     )
 
     _is_async: bool = PrivateAttr(default=False)
+    _tracer: Optional[Tracer] = PrivateAttr(default=None)
 
     @model_validator(mode="before")
     @classmethod
@@ -76,6 +86,17 @@ class AgentTool(BaseModel):
             self._initialize_from_func(self.func)
         else:
             self._initialize_from_run()
+
+        try:
+            provider = provider = trace.get_tracer_provider()
+
+            self._tracer = provider.get_tracer(f"{self.name}_tracer")
+        except Exception as e:
+            logger.warning(
+                f"OpenTelemetry initialization failed: {e}. Continuing without telemetry."
+            )
+            self._tracer = None
+
         return super().model_post_init(__context)
 
     def _initialize_from_func(self, func: Callable) -> None:
@@ -88,6 +109,7 @@ class AgentTool(BaseModel):
         if self.args_model is None:
             self.args_model = ToolHelper.infer_func_schema(self._run)
 
+    @span_decorator("validate_and_prep_args")
     def _validate_and_prepare_args(
         self, func: Callable, *args, **kwargs
     ) -> Dict[str, Any]:
@@ -138,15 +160,19 @@ class AgentTool(BaseModel):
         except Exception as e:
             self._log_and_raise_error(e)
 
+    @async_span_decorator("arun_tool")
     async def arun(self, *args, **kwargs) -> Any:
         """
         Execute the tool asynchronously (whether it's sync or async under the hood).
         """
+        span = trace.get_current_span()
         try:
             func = self.func or self._run
             kwargs = self._validate_and_prepare_args(func, *args, **kwargs)
             return await func(**kwargs) if self._is_async else func(**kwargs)
         except Exception as e:
+            span.set_status(Status(StatusCode.ERROR))
+            span.record_exception(e)
             self._log_and_raise_error(e)
 
     def _run(self, *args, **kwargs) -> Any:

dapr_agents/tool/executor.py

@@ -7,6 +7,14 @@ from rich.console import Console
 from dapr_agents.tool import AgentTool
 from dapr_agents.types import AgentToolExecutorError, ToolError
 
+from pydantic import PrivateAttr
+from dapr_agents.agent.telemetry import (
+    async_span_decorator,
+)
+
+from opentelemetry import trace
+from opentelemetry.trace import Tracer, Status, StatusCode
+
 logger = logging.getLogger(__name__)
 
 
@@ -22,12 +30,24 @@ class AgentToolExecutor(BaseModel):
         default_factory=list, description="List of tools to register and manage."
     )
     _tools_map: Dict[str, AgentTool] = PrivateAttr(default_factory=dict)
+    _tracer: Optional[Tracer] = PrivateAttr(default=None)
 
     def model_post_init(self, __context: Any) -> None:
         """Initializes the internal tools map after model creation."""
         for tool in self.tools:
             self.register_tool(tool)
         logger.info(f"Tool Executor initialized with {len(self._tools_map)} tool(s).")
+
+        try:
+            provider = provider = trace.get_tracer_provider()
+
+            self._tracer = provider.get_tracer(f"agent_tool_exec_tracer")
+        except Exception as e:
+            logger.warning(
+                f"OpenTelemetry initialization failed: {e}. Continuing without telemetry."
+            )
+            self._tracer = None
+
         super().model_post_init(__context)
 
     def register_tool(self, tool: AgentTool) -> None:
@@ -88,6 +108,7 @@ class AgentToolExecutor(BaseModel):
             for tool in self._tools_map.values()
         )
 
+    @async_span_decorator("run_tool")
     async def run_tool(self, tool_name: str, *args, **kwargs) -> Any:
         """
         Executes a tool by name, automatically handling both sync and async tools.
@@ -103,6 +124,7 @@ class AgentToolExecutor(BaseModel):
         Raises:
             AgentToolExecutorError: If the tool is not found or execution fails.
         """
+        span = trace.get_current_span()
         tool = self.get_tool(tool_name)
         if not tool:
             logger.error(f"Tool not found: {tool_name}")
@@ -114,9 +136,13 @@ class AgentToolExecutor(BaseModel):
             return tool(*args, **kwargs)
         except ToolError as e:
             logger.error(f"Tool execution error in '{tool_name}': {e}")
+            span.set_status(Status(StatusCode.ERROR))
+            span.record_exception(e)
             raise AgentToolExecutorError(str(e)) from e
         except Exception as e:
             logger.error(f"Unexpected error in '{tool_name}': {e}")
+            span.set_status(Status(StatusCode.ERROR))
+            span.record_exception(e)
             raise AgentToolExecutorError(
                 f"Unexpected error in tool '{tool_name}': {e}"
             ) from e

dapr_agents/workflow/agents/assistant/agent.py

@@ -377,15 +377,22 @@ class AssistantAgent(AgentWorkflowBase):
         Raises:
             AgentError: If the tool call is malformed or execution fails.
         """
+        span = trace.get_current_span()
+        span.set_attribute("workflow.id", instance_id)
+
         function_details = tool_call.get("function", {})
         function_name = function_details.get("name")
+        span.set_attribute("tool.call.name", function_name)
+        span.set_attribute("tool.call.details", str(function_details))
 
         if not function_name:
+            span.set_attribute("error.type", type(e).__name__)
             raise AgentError("Missing function name in tool execution request.")
 
         try:
             function_args = function_details.get("arguments", "")
             function_args_as_dict = json.loads(function_args) if function_args else {}
+            span.set_attributes("tool.call.args", str(function_args_as_dict))
 
             # Execute tool function
             result = await self.tool_executor.run_tool(
@@ -407,7 +414,6 @@ class AssistantAgent(AgentWorkflowBase):
         except (ToolError, AgentToolExecutorError) as e:
             logger.info(e)
 
-            span = trace.get_current_span()
             span.set_status(Status(StatusCode.ERROR))
             span.record_exception(e)
 
@@ -431,6 +437,7 @@ class AssistantAgent(AgentWorkflowBase):
 
         except Exception as e:
             logger.error(f"Error executing tool '{function_name}': {e}", exc_info=True)
+            span.set_attribute("error.type", type(e).__name__)
             raise AgentError(f"Error executing tool '{function_name}': {e}") from e
 
     @task