Repository: deepset-ai/haystack

Issue:
Add internal tracing support to Agent
We would like to add tracing internally to the Agent's run method. Specifically we want to log the inputs and outputs each time the underlying chat generator and tool invoker is called. 


Discussion:

Patch:
diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py
index dce17d6a8f..a976ca2d7e 100644
--- a/haystack/components/agents/agent.py
+++ b/haystack/components/agents/agent.py
@@ -2,14 +2,15 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import asyncio
 import inspect
 from copy import deepcopy
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Iterator, List, Optional
 
-from haystack import component, default_from_dict, default_to_dict, logging
+from haystack import component, default_from_dict, default_to_dict, logging, tracing
 from haystack.components.generators.chat.types import ChatGenerator
 from haystack.components.tools import ToolInvoker
+from haystack.core.pipeline.async_pipeline import AsyncPipeline
+from haystack.core.pipeline.pipeline import Pipeline
 from haystack.core.serialization import component_to_dict
 from haystack.dataclasses import ChatMessage
 from haystack.dataclasses.state import State, _schema_from_dict, _schema_to_dict, _validate_schema
@@ -187,6 +188,26 @@ def from_dict(cls, data: Dict[str, Any]) -> "Agent":
 
         return default_from_dict(cls, data)
 
+    def _prepare_generator_inputs(self, streaming_callback: Optional[StreamingCallbackT] = None) -> Dict[str, Any]:
+        """Prepare inputs for the chat generator."""
+        generator_inputs = {"tools": self.tools}
+        selected_callback = streaming_callback or self.streaming_callback
+        if selected_callback is not None:
+            generator_inputs["streaming_callback"] = selected_callback
+        return generator_inputs
+
+    def _create_agent_span(self) -> Iterator[tracing.Span]:
+        """Create a span for the agent run."""
+        return tracing.tracer.trace(
+            "haystack.agent.run",
+            tags={
+                "haystack.agent.max_steps": self.max_agent_steps,
+                "haystack.agent.tools": self.tools,
+                "haystack.agent.exit_conditions": self.exit_conditions,
+                "haystack.agent.state_schema": self.state_schema,
+            },
+        )
+
     def run(
         self,
         messages: List[ChatMessage],
@@ -205,48 +226,66 @@ def run(
         if not self._is_warmed_up and hasattr(self.chat_generator, "warm_up"):
             raise RuntimeError("The component Agent wasn't warmed up. Run 'warm_up()' before calling 'run()'.")
 
-        state = State(schema=self.state_schema, data=kwargs)
-
         if self.system_prompt is not None:
             messages = [ChatMessage.from_system(self.system_prompt)] + messages
+
+        input_data = deepcopy({"messages": messages, "streaming_callback": streaming_callback, **kwargs})
+
+        state = State(schema=self.state_schema, data=kwargs)
         state.set("messages", messages)
 
-        generator_inputs: Dict[str, Any] = {"tools": self.tools}
+        generator_inputs = self._prepare_generator_inputs(streaming_callback=streaming_callback)
+
+        component_visits = dict.fromkeys(["chat_generator", "tool_invoker"], 0)
+        with self._create_agent_span() as span:
+            span.set_content_tag("haystack.agent.input", input_data)
+            counter = 0
+            while counter < self.max_agent_steps:
+                # 1. Call the ChatGenerator
+                llm_messages = Pipeline._run_component(
+                    component_name="chat_generator",
+                    component={"instance": self.chat_generator},
+                    inputs={"messages": messages, **generator_inputs},
+                    component_visits=component_visits,
+                    parent_span=span,
+                )["replies"]
+                state.set("messages", llm_messages)
+
+                # 2. Check if any of the LLM responses contain a tool call
+                if not any(msg.tool_call for msg in llm_messages):
+                    counter += 1
+                    break
 
-        selected_callback = streaming_callback or self.streaming_callback
-        if selected_callback is not None:
-            generator_inputs["streaming_callback"] = selected_callback
+                # 3. Call the ToolInvoker
+                # We only send the messages from the LLM to the tool invoker
+                tool_invoker_result = Pipeline._run_component(
+                    component_name="tool_invoker",
+                    component={"instance": self._tool_invoker},
+                    inputs={"messages": llm_messages, "state": state},
+                    component_visits=component_visits,
+                    parent_span=span,
+                )
+                tool_messages = tool_invoker_result["tool_messages"]
+                state = tool_invoker_result["state"]
+                state.set("messages", tool_messages)
 
-        # Repeat until the exit condition is met
-        counter = 0
-        while counter < self.max_agent_steps:
-            # 1. Call the ChatGenerator
-            llm_messages = self.chat_generator.run(messages=messages, **generator_inputs)["replies"]
-            state.set("messages", llm_messages)
-
-            # 2. Check if any of the LLM responses contain a tool call
-            if not any(msg.tool_call for msg in llm_messages):
-                return {**state.data}
-
-            # 3. Call the ToolInvoker
-            # We only send the messages from the LLM to the tool invoker
-            tool_invoker_result = self._tool_invoker.run(messages=llm_messages, state=state)
-            tool_messages = tool_invoker_result["tool_messages"]
-            state = tool_invoker_result["state"]
-            state.set("messages", tool_messages)
-
-            # 4. Check if any LLM message's tool call name matches an exit condition
-            if self.exit_conditions != ["text"] and self._check_exit_conditions(llm_messages, tool_messages):
-                return {**state.data}
-
-            # 5. Fetch the combined messages and send them back to the LLM
-            messages = state.get("messages")
-            counter += 1
-
-        logger.warning(
-            "Agent exceeded maximum agent steps of {max_agent_steps}, stopping.", max_agent_steps=self.max_agent_steps
-        )
-        return {**state.data}
+                # 4. Check if any LLM message's tool call name matches an exit condition
+                if self.exit_conditions != ["text"] and self._check_exit_conditions(llm_messages, tool_messages):
+                    counter += 1
+                    break
+
+                # 5. Fetch the combined messages and send them back to the LLM
+                messages = state.get("messages")
+                counter += 1
+
+            if counter >= self.max_agent_steps:
+                logger.warning(
+                    "Agent reached maximum agent steps of {max_agent_steps}, stopping.",
+                    max_agent_steps=self.max_agent_steps,
+                )
+            span.set_content_tag("haystack.agent.output", state.data)
+            span.set_tag("haystack.agent.steps_taken", counter)
+        return state.data
 
     async def run_async(
         self,
@@ -270,66 +309,70 @@ async def run_async(
         if not self._is_warmed_up and hasattr(self.chat_generator, "warm_up"):
             raise RuntimeError("The component Agent wasn't warmed up. Run 'warm_up()' before calling 'run_async()'.")
 
-        state = State(schema=self.state_schema, data=kwargs)
-
         if self.system_prompt is not None:
             messages = [ChatMessage.from_system(self.system_prompt)] + messages
-        state.set("messages", messages)
 
-        generator_inputs: Dict[str, Any] = {"tools": self.tools}
+        input_data = deepcopy({"messages": messages, "streaming_callback": streaming_callback, **kwargs})
 
-        selected_callback = streaming_callback or self.streaming_callback
-        if selected_callback is not None:
-            generator_inputs["streaming_callback"] = selected_callback
+        state = State(schema=self.state_schema, data=kwargs)
+        state.set("messages", messages)
 
-        # Repeat until the exit condition is met
-        counter = 0
-        while counter < self.max_agent_steps:
-            # 1. Call the ChatGenerator
-            # Check if the chat generator supports async execution
-            if getattr(self.chat_generator, "__haystack_supports_async__", False):
-                result = await self.chat_generator.run_async(messages=messages, **generator_inputs)  # type: ignore[attr-defined]
-                llm_messages = result["replies"]
-            else:
-                # Fall back to synchronous run if async is not available
-                loop = asyncio.get_running_loop()
-                result = await loop.run_in_executor(
-                    None, lambda: self.chat_generator.run(messages=messages, **generator_inputs)
+        generator_inputs = self._prepare_generator_inputs(streaming_callback=streaming_callback)
+
+        component_visits = dict.fromkeys(["chat_generator", "tool_invoker"], 0)
+        with self._create_agent_span() as span:
+            span.set_content_tag("haystack.agent.input", input_data)
+            counter = 0
+            while counter < self.max_agent_steps:
+                # 1. Call the ChatGenerator
+                result = await AsyncPipeline._run_component_async(
+                    component_name="chat_generator",
+                    component={"instance": self.chat_generator},
+                    component_inputs={"messages": messages, **generator_inputs},
+                    component_visits=component_visits,
+                    max_runs_per_component=self.max_agent_steps,
+                    parent_span=span,
                 )
                 llm_messages = result["replies"]
+                state.set("messages", llm_messages)
+
+                # 2. Check if any of the LLM responses contain a tool call
+                if not any(msg.tool_call for msg in llm_messages):
+                    counter += 1
+                    break
 
-            state.set("messages", llm_messages)
-
-            # 2. Check if any of the LLM responses contain a tool call
-            if not any(msg.tool_call for msg in llm_messages):
-                return {**state.data}
-
-            # 3. Call the ToolInvoker
-            # We only send the messages from the LLM to the tool invoker
-            # Check if the ToolInvoker supports async execution. Currently, it doesn't.
-            if getattr(self._tool_invoker, "__haystack_supports_async__", False):
-                tool_invoker_result = await self._tool_invoker.run_async(messages=llm_messages, state=state)  # type: ignore[attr-defined]
-            else:
-                loop = asyncio.get_running_loop()
-                tool_invoker_result = await loop.run_in_executor(
-                    None, lambda: self._tool_invoker.run(messages=llm_messages, state=state)
+                # 3. Call the ToolInvoker
+                # We only send the messages from the LLM to the tool invoker
+                # Check if the ToolInvoker supports async execution. Currently, it doesn't.
+                tool_invoker_result = await AsyncPipeline._run_component_async(
+                    component_name="tool_invoker",
+                    component={"instance": self._tool_invoker},
+                    component_inputs={"messages": llm_messages, "state": state},
+                    component_visits=component_visits,
+                    max_runs_per_component=self.max_agent_steps,
+                    parent_span=span,
                 )
-            tool_messages = tool_invoker_result["tool_messages"]
-            state = tool_invoker_result["state"]
-            state.set("messages", tool_messages)
+                tool_messages = tool_invoker_result["tool_messages"]
+                state = tool_invoker_result["state"]
+                state.set("messages", tool_messages)
 
-            # 4. Check if any LLM message's tool call name matches an exit condition
-            if self.exit_conditions != ["text"] and self._check_exit_conditions(llm_messages, tool_messages):
-                return {**state.data}
+                # 4. Check if any LLM message's tool call name matches an exit condition
+                if self.exit_conditions != ["text"] and self._check_exit_conditions(llm_messages, tool_messages):
+                    counter += 1
+                    break
 
-            # 5. Fetch the combined messages and send them back to the LLM
-            messages = state.get("messages")
-            counter += 1
+                # 5. Fetch the combined messages and send them back to the LLM
+                messages = state.get("messages")
+                counter += 1
 
-        logger.warning(
-            "Agent exceeded maximum agent steps of {max_agent_steps}, stopping.", max_agent_steps=self.max_agent_steps
-        )
-        return {**state.data}
+            if counter >= self.max_agent_steps:
+                logger.warning(
+                    "Agent reached maximum agent steps of {max_agent_steps}, stopping.",
+                    max_agent_steps=self.max_agent_steps,
+                )
+            span.set_content_tag("haystack.agent.output", state.data)
+            span.set_tag("haystack.agent.steps_taken", counter)
+        return state.data
 
     def _check_exit_conditions(self, llm_messages: List[ChatMessage], tool_messages: List[ChatMessage]) -> bool:
         """
diff --git a/haystack/core/pipeline/async_pipeline.py b/haystack/core/pipeline/async_pipeline.py
index f5e8b427dc..de2584cf4a 100644
--- a/haystack/core/pipeline/async_pipeline.py
+++ b/haystack/core/pipeline/async_pipeline.py
@@ -9,7 +9,13 @@
 from haystack import logging, tracing
 from haystack.core.component import Component
 from haystack.core.errors import PipelineMaxComponentRuns, PipelineRuntimeError
-from haystack.core.pipeline.base import ComponentPriority, PipelineBase
+from haystack.core.pipeline.base import (
+    _COMPONENT_INPUT,
+    _COMPONENT_OUTPUT,
+    _COMPONENT_VISITS,
+    ComponentPriority,
+    PipelineBase,
+)
 from haystack.telemetry import pipeline_running
 
 logger = logging.getLogger(__name__)
@@ -23,6 +29,57 @@ class AsyncPipeline(PipelineBase):
     This enables efficient processing of components by minimizing idle time and maximizing resource utilization.
     """
 
+    @staticmethod
+    async def _run_component_async(  # pylint: disable=too-many-positional-arguments
+        component_name: str,
+        component: Dict[str, Any],
+        component_inputs: Dict[str, Any],
+        component_visits: Dict[str, int],
+        max_runs_per_component: int = 100,
+        parent_span: Optional[tracing.Span] = None,
+    ) -> Dict[str, Any]:
+        """
+        Executes a single component asynchronously.
+
+        If the component supports async execution, it is awaited directly as it will run async;
+        otherwise the component is offloaded to executor.
+
+        The method also updates the `visits` count of the component, writes outputs to `inputs_state`,
+        and returns pruned outputs that get stored in `pipeline_outputs`.
+
+        :param component_name: The name of the component.
+        :param component_inputs: Inputs for the component.
+        :returns: Outputs from the component that can be yielded from run_async_generator.
+        """
+        if component_visits[component_name] > max_runs_per_component:
+            raise PipelineMaxComponentRuns(f"Max runs for '{component_name}' reached.")
+
+        instance: Component = component["instance"]
+        with PipelineBase._create_component_span(
+            component_name=component_name, instance=instance, inputs=component_inputs, parent_span=parent_span
+        ) as span:
+            span.set_content_tag(_COMPONENT_INPUT, deepcopy(component_inputs))
+            logger.info("Running component {component_name}", component_name=component_name)
+
+            if getattr(instance, "__haystack_supports_async__", False):
+                try:
+                    outputs = await instance.run_async(**component_inputs)  # type: ignore
+                except Exception as error:
+                    raise PipelineRuntimeError.from_exception(component_name, instance.__class__, error) from error
+            else:
+                loop = asyncio.get_running_loop()
+                outputs = await loop.run_in_executor(None, lambda: instance.run(**component_inputs))
+
+            component_visits[component_name] += 1
+
+            if not isinstance(outputs, dict):
+                raise PipelineRuntimeError.from_invalid_output(component_name, instance.__class__, outputs)
+
+            span.set_tag(_COMPONENT_VISITS, component_visits[component_name])
+            span.set_content_tag(_COMPONENT_OUTPUT, deepcopy(outputs))
+
+            return outputs
+
     async def run_async_generator(  # noqa: PLR0915,C901
         self, data: Dict[str, Any], include_outputs_from: Optional[Set[str]] = None, concurrency_limit: int = 4
     ) -> AsyncIterator[Dict[str, Any]]:
@@ -163,82 +220,6 @@ async def process_results():
             # We define some functions here so that they have access to local runtime state
             # (inputs, tasks, scheduled components) via closures.
             # -------------------------------------------------
-            async def _run_component_async(component_name: str, component_inputs: Dict[str, Any]) -> Dict[str, Any]:
-                """
-                Executes a single component asynchronously.
-
-                If the component supports async execution, it is awaited directly as it will run async;
-                otherwise the component is offloaded to executor.
-
-                The method also updates the `visits` count of the component, writes outputs to `inputs_state`,
-                and returns pruned outputs that get stored in `pipeline_outputs`.
-
-                :param component_name: The name of the component.
-                :param component_inputs: Inputs for the component.
-                :returns: Outputs from the component that can be yielded from run_async_generator.
-                """
-                if component_visits[component_name] > self._max_runs_per_component:
-                    raise PipelineMaxComponentRuns(f"Max runs for '{component_name}' reached.")
-
-                instance: Component = self.get_component(component_name)
-                with tracing.tracer.trace(
-                    "haystack.component.run",
-                    tags={
-                        "haystack.component.name": component_name,
-                        "haystack.component.type": instance.__class__.__name__,
-                        "haystack.component.input_types": {k: type(v).__name__ for k, v in component_inputs.items()},
-                        "haystack.component.input_spec": {
-                            key: {
-                                "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
-                                "senders": value.senders,
-                            }
-                            for key, value in instance.__haystack_input__._sockets_dict.items()  # type: ignore
-                        },
-                        "haystack.component.output_spec": {
-                            key: {
-                                "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
-                                "receivers": value.receivers,
-                            }
-                            for key, value in instance.__haystack_output__._sockets_dict.items()  # type: ignore
-                        },
-                    },
-                    parent_span=parent_span,
-                ) as span:
-                    span.set_content_tag("haystack.component.input", deepcopy(component_inputs))
-                    logger.info("Running component {component_name}", component_name=component_name)
-
-                    if getattr(instance, "__haystack_supports_async__", False):
-                        try:
-                            outputs = await instance.run_async(**component_inputs)  # type: ignore
-                        except Exception as error:
-                            raise PipelineRuntimeError.from_exception(
-                                component_name, instance.__class__, error
-                            ) from error
-                    else:
-                        loop = asyncio.get_running_loop()
-                        outputs = await loop.run_in_executor(None, lambda: instance.run(**component_inputs))
-
-                    component_visits[component_name] += 1
-
-                    if not isinstance(outputs, dict):
-                        raise PipelineRuntimeError.from_invalid_output(component_name, instance.__class__, outputs)
-
-                    span.set_tag("haystack.component.visits", component_visits[component_name])
-                    span.set_content_tag("haystack.component.output", deepcopy(outputs))
-
-                    # Distribute outputs to downstream inputs; also prune outputs based on `include_outputs_from`
-                    pruned = self._write_component_outputs(
-                        component_name=component_name,
-                        component_outputs=outputs,
-                        inputs=inputs_state,
-                        receivers=cached_receivers[component_name],
-                        include_outputs_from=include_outputs_from,
-                    )
-                    if pruned:
-                        pipeline_outputs[component_name] = pruned
-
-                    return pruned
-
             async def _run_highest_in_isolation(component_name: str) -> AsyncIterator[Dict[str, Any]]:
                 """
                 Runs a component with HIGHEST priority in isolation.
@@ -271,10 +252,29 @@ async def _run_highest_in_isolation(component_name: str) -> AsyncIterator[Dict[s
                 )
                 component_inputs = self._consume_component_inputs(component_name, comp_dict, inputs_state)
                 component_inputs = self._add_missing_input_defaults(component_inputs, comp_dict["input_sockets"])
-                result = await _run_component_async(component_name, component_inputs)
+                component_pipeline_outputs = await self._run_component_async(
+                    component_name=component_name,
+                    component=comp_dict,
+                    component_inputs=component_inputs,
+                    component_visits=component_visits,
+                    max_runs_per_component=self._max_runs_per_component,
+                    parent_span=parent_span,
+                )
+
+                # Distribute outputs to downstream inputs; also prune outputs based on `include_outputs_from`
+                pruned = self._write_component_outputs(
+                    component_name=component_name,
+                    component_outputs=component_pipeline_outputs,
+                    inputs=inputs_state,
+                    receivers=cached_receivers[component_name],
+                    include_outputs_from=include_outputs_from,
+                )
+                if pruned:
+                    pipeline_outputs[component_name] = pruned
+
                 scheduled_components.remove(component_name)
-                if result:
-                    yield {component_name: deepcopy(result)}
+                if pruned:
+                    yield {component_name: deepcopy(pruned)}
 
             async def _schedule_task(component_name: str) -> None:
                 """
@@ -298,10 +298,28 @@ async def _schedule_task(component_name: str) -> None:
 
                 async def _runner():
                     async with ready_sem:
-                        result = await _run_component_async(component_name, component_inputs)
+                        component_pipeline_outputs = await self._run_component_async(
+                            component_name=component_name,
+                            component=comp_dict,
+                            component_inputs=component_inputs,
+                            component_visits=component_visits,
+                            max_runs_per_component=self._max_runs_per_component,
+                            parent_span=parent_span,
+                        )
+
+                    # Distribute outputs to downstream inputs; also prune outputs based on `include_outputs_from`
+                    pruned = self._write_component_outputs(
+                        component_name=component_name,
+                        component_outputs=component_pipeline_outputs,
+                        inputs=inputs_state,
+                        receivers=cached_receivers[component_name],
+                        include_outputs_from=include_outputs_from,
+                    )
+                    if pruned:
+                        pipeline_outputs[component_name] = pruned
 
                     scheduled_components.remove(component_name)
-                    return result
+                    return pruned
 
                 task = asyncio.create_task(_runner())
                 running_tasks[task] = component_name
diff --git a/haystack/core/pipeline/base.py b/haystack/core/pipeline/base.py
index fcd51f6f20..38a2bc1481 100644
--- a/haystack/core/pipeline/base.py
+++ b/haystack/core/pipeline/base.py
@@ -12,7 +12,7 @@
 
 import networkx  # type:ignore
 
-from haystack import logging
+from haystack import logging, tracing
 from haystack.core.component import Component, InputSocket, OutputSocket, component
 from haystack.core.errors import (
     DeserializationError,
@@ -53,6 +53,12 @@
 logger = logging.getLogger(__name__)
 
 
+# Constants for tracing tags
+_COMPONENT_INPUT = "haystack.component.input"
+_COMPONENT_OUTPUT = "haystack.component.output"
+_COMPONENT_VISITS = "haystack.component.visits"
+
+
 class ComponentPriority(IntEnum):
     HIGHEST = 1
     READY = 2
@@ -768,6 +774,34 @@ def warm_up(self):
                 logger.info("Warming up component {node}...", node=node)
                 self.graph.nodes[node]["instance"].warm_up()
 
+    @staticmethod
+    def _create_component_span(
+        component_name: str, instance: Component, inputs: Dict[str, Any], parent_span: Optional[tracing.Span] = None
+    ):
+        return tracing.tracer.trace(
+            "haystack.component.run",
+            tags={
+                "haystack.component.name": component_name,
+                "haystack.component.type": instance.__class__.__name__,
+                "haystack.component.input_types": {k: type(v).__name__ for k, v in inputs.items()},
+                "haystack.component.input_spec": {
+                    key: {
+                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
+                        "senders": value.senders,
+                    }
+                    for key, value in instance.__haystack_input__._sockets_dict.items()  # type: ignore
+                },
+                "haystack.component.output_spec": {
+                    key: {
+                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
+                        "receivers": value.receivers,
+                    }
+                    for key, value in instance.__haystack_output__._sockets_dict.items()  # type: ignore
+                },
+            },
+            parent_span=parent_span,
+        )
+
     def _validate_input(self, data: Dict[str, Any]):
         """
         Validates pipeline input data.
diff --git a/haystack/core/pipeline/pipeline.py b/haystack/core/pipeline/pipeline.py
index a345dc599a..69bbd145ca 100644
--- a/haystack/core/pipeline/pipeline.py
+++ b/haystack/core/pipeline/pipeline.py
@@ -8,7 +8,13 @@
 from haystack import logging, tracing
 from haystack.core.component import Component
 from haystack.core.errors import PipelineRuntimeError
-from haystack.core.pipeline.base import ComponentPriority, PipelineBase
+from haystack.core.pipeline.base import (
+    _COMPONENT_INPUT,
+    _COMPONENT_OUTPUT,
+    _COMPONENT_VISITS,
+    ComponentPriority,
+    PipelineBase,
+)
 from haystack.telemetry import pipeline_running
 
 logger = logging.getLogger(__name__)
@@ -21,8 +27,9 @@ class Pipeline(PipelineBase):
     Orchestrates component execution according to the execution graph, one after the other.
     """
 
+    @staticmethod
     def _run_component(
-        self,
+        component_name: str,
         component: Dict[str, Any],
         inputs: Dict[str, Any],
         component_visits: Dict[str, int],
@@ -31,6 +38,7 @@ def _run_component(
         """
         Runs a Component with the given inputs.
 
+        :param component_name: Name of the Component.
         :param component: Component with component metadata.
         :param inputs: Inputs for the Component.
         :param component_visits: Current state of component visits.
@@ -40,44 +48,16 @@ def _run_component(
         :return: The output of the Component.
         """
         instance: Component = component["instance"]
-        component_name = self.get_component_name(instance)
-        component_inputs = self._consume_component_inputs(
-            component_name=component_name, component=component, inputs=inputs
-        )
-
-        # We need to add missing defaults using default values from input sockets because the run signature
-        # might not provide these defaults for components with inputs defined dynamically upon component initialization
-        component_inputs = self._add_missing_input_defaults(component_inputs, component["input_sockets"])
 
-        with tracing.tracer.trace(
-            "haystack.component.run",
-            tags={
-                "haystack.component.name": component_name,
-                "haystack.component.type": instance.__class__.__name__,
-                "haystack.component.input_types": {k: type(v).__name__ for k, v in component_inputs.items()},
-                "haystack.component.input_spec": {
-                    key: {
-                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
-                        "senders": value.senders,
-                    }
-                    for key, value in instance.__haystack_input__._sockets_dict.items()  # type: ignore
-                },
-                "haystack.component.output_spec": {
-                    key: {
-                        "type": (value.type.__name__ if isinstance(value.type, type) else str(value.type)),
-                        "receivers": value.receivers,
-                    }
-                    for key, value in instance.__haystack_output__._sockets_dict.items()  # type: ignore
-                },
-            },
-            parent_span=parent_span,
+        with PipelineBase._create_component_span(
+            component_name=component_name, instance=instance, inputs=inputs, parent_span=parent_span
         ) as span:
             # We deepcopy the inputs otherwise we might lose that information
             # when we delete them in case they're sent to other Components
-            span.set_content_tag("haystack.component.input", deepcopy(component_inputs))
+            span.set_content_tag(_COMPONENT_INPUT, deepcopy(inputs))
             logger.info("Running component {component_name}", component_name=component_name)
             try:
-                component_output = instance.run(**component_inputs)
+                component_output = instance.run(**inputs)
             except Exception as error:
                 raise PipelineRuntimeError.from_exception(component_name, instance.__class__, error) from error
             component_visits[component_name] += 1
@@ -85,8 +65,8 @@ def _run_component(
             if not isinstance(component_output, Mapping):
                 raise PipelineRuntimeError.from_invalid_output(component_name, instance.__class__, component_output)
 
-            span.set_tag("haystack.component.visits", component_visits[component_name])
-            span.set_content_tag("haystack.component.output", component_output)
+            span.set_tag(_COMPONENT_VISITS, component_visits[component_name])
+            span.set_content_tag(_COMPONENT_OUTPUT, component_output)
 
             return cast(Dict[Any, Any], component_output)
 
@@ -244,7 +224,21 @@ def run(  # noqa: PLR0915, PLR0912
                         component_name, component_visits[component_name]
                     )
 
-                component_outputs = self._run_component(component, inputs, component_visits, parent_span=span)
+                component_inputs = self._consume_component_inputs(
+                    component_name=component_name, component=component, inputs=inputs
+                )
+                # We need to add missing defaults using default values from input sockets because the run signature
+                # might not provide these defaults for components with inputs defined dynamically upon component
+                # initialization
+                component_inputs = self._add_missing_input_defaults(component_inputs, component["input_sockets"])
+
+                component_outputs = self._run_component(
+                    component_name=component_name,
+                    component=component,
+                    inputs=component_inputs,
+                    component_visits=component_visits,
+                    parent_span=span,
+                )
 
                 # Updates global input state with component outputs and returns outputs that should go to
                 # pipeline outputs.
diff --git a/releasenotes/notes/add-tracing-to-agent-da75b81fe68d33a2.yaml b/releasenotes/notes/add-tracing-to-agent-da75b81fe68d33a2.yaml
new file mode 100644
index 0000000000..02c5464648
--- /dev/null
+++ b/releasenotes/notes/add-tracing-to-agent-da75b81fe68d33a2.yaml
@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    We have added internal tracing support to Agent. It is now possible to track the internal loops within the agent by viewing the inputs and outputs each time the ChatGenerator and ToolInvoker is called.


Test patch:
diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py
index bc1e00a990..998fb41b45 100644
--- a/test/components/agents/test_agent.py
+++ b/test/components/agents/test_agent.py
@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 import os
 from datetime import datetime
 from typing import Iterator, Dict, Any, List, Optional, Union
@@ -12,11 +13,13 @@
 from openai import Stream
 from openai.types.chat import ChatCompletionChunk, chat_completion_chunk
 
+from haystack.tracing.logging_tracer import LoggingTracer
+from haystack import Pipeline, tracing
 from haystack.components.agents import Agent
+from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
 from haystack.components.builders.prompt_builder import PromptBuilder
 from haystack.components.generators.chat.openai import OpenAIChatGenerator
 from haystack.components.generators.chat.types import ChatGenerator
-from haystack import component
 from haystack.core.component.types import OutputSocket
 from haystack.dataclasses import ChatMessage, ToolCall
 from haystack.dataclasses.chat_message import ChatRole, TextContent
@@ -96,6 +99,9 @@ def openai_mock_chat_completion_chunk():
 class MockChatGeneratorWithoutTools(ChatGenerator):
     """A mock chat generator that implements ChatGenerator protocol but doesn't support tools."""
 
+    __haystack_input__ = MagicMock(_sockets_dict={})
+    __haystack_output__ = MagicMock(_sockets_dict={})
+
     def to_dict(self) -> Dict[str, Any]:
         return {"type": "MockChatGeneratorWithoutTools", "data": {}}
 
@@ -110,6 +116,27 @@ def run(self, messages: List[ChatMessage]) -> Dict[str, Any]:
 class MockChatGeneratorWithoutRunAsync(ChatGenerator):
     """A mock chat generator that implements ChatGenerator protocol but doesn't have run_async method."""
 
+    __haystack_input__ = MagicMock(_sockets_dict={})
+    __haystack_output__ = MagicMock(_sockets_dict={})
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}}
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "MockChatGeneratorWithoutRunAsync":
+        return cls()
+
+    def run(
+        self, messages: List[ChatMessage], tools: Optional[Union[List[Tool], Toolset]] = None, **kwargs
+    ) -> Dict[str, Any]:
+        return {"replies": [ChatMessage.from_assistant("Hello")]}
+
+
+class MockChatGeneratorWithRunAsync(ChatGenerator):
+    __haystack_supports_async__ = True
+    __haystack_input__ = MagicMock(_sockets_dict={})
+    __haystack_output__ = MagicMock(_sockets_dict={})
+
     def to_dict(self) -> Dict[str, Any]:
         return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}}
 
@@ -122,6 +149,11 @@ def run(
     ) -> Dict[str, Any]:
         return {"replies": [ChatMessage.from_assistant("Hello")]}
 
+    async def run_async(
+        self, messages: List[ChatMessage], tools: Optional[Union[List[Tool], Toolset]] = None, **kwargs
+    ) -> Dict[str, Any]:
+        return {"replies": [ChatMessage.from_assistant("Hello from run_async")]}
+
 
 class TestAgent:
     def test_output_types(self, weather_tool, component_tool, monkeypatch):
@@ -463,6 +495,27 @@ def test_multiple_llm_responses_with_tool_call(self, monkeypatch, weather_tool):
             == "{'weather': 'mostly sunny', 'temperature': 7, 'unit': 'celsius'}"
         )
 
+    def test_exceed_max_steps(self, monkeypatch, weather_tool, caplog):
+        monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
+        generator = OpenAIChatGenerator()
+
+        mock_messages = [
+            ChatMessage.from_assistant("First response"),
+            ChatMessage.from_assistant(
+                tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
+            ),
+        ]
+
+        agent = Agent(chat_generator=generator, tools=[weather_tool], max_agent_steps=0)
+        agent.warm_up()
+
+        # Patch agent.chat_generator.run to return mock_messages
+        agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages})
+
+        with caplog.at_level(logging.WARNING):
+            agent.run([ChatMessage.from_user("Hello")])
+            assert "Agent reached maximum agent steps" in caplog.text
+
     def test_exit_conditions_checked_across_all_llm_messages(self, monkeypatch, weather_tool):
         monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
         generator = OpenAIChatGenerator()
@@ -491,6 +544,20 @@ def test_exit_conditions_checked_across_all_llm_messages(self, monkeypatch, weat
             == "{'weather': 'mostly sunny', 'temperature': 7, 'unit': 'celsius'}"
         )
 
+    def test_run_with_system_prompt(self, weather_tool):
+        chat_generator = MockChatGeneratorWithoutRunAsync()
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.")
+        agent.warm_up()
+        response = agent.run([ChatMessage.from_user("What is the weather in Berlin?")])
+        assert response["messages"][0].text == "This is a system prompt."
+
+    def test_run_not_warmed_up(self, weather_tool):
+        chat_generator = MockChatGeneratorWithoutRunAsync()
+        chat_generator.warm_up = MagicMock()
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.")
+        with pytest.raises(RuntimeError, match="The component Agent wasn't warmed up."):
+            agent.run([ChatMessage.from_user("What is the weather in Berlin?")])
+
     @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
     @pytest.mark.integration
     def test_run(self, weather_tool):
@@ -523,6 +590,7 @@ def test_run(self, weather_tool):
     @pytest.mark.asyncio
     async def test_run_async_falls_back_to_run_when_chat_generator_has_no_run_async(self, weather_tool):
         chat_generator = MockChatGeneratorWithoutRunAsync()
+
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
         agent.warm_up()
 
@@ -544,27 +612,6 @@ async def test_run_async_falls_back_to_run_when_chat_generator_has_no_run_async(
 
     @pytest.mark.asyncio
     async def test_run_async_uses_chat_generator_run_async_when_available(self, weather_tool):
-        # Create a mock chat generator with run_async
-        # We need to use @component so that has_async_run is set
-        @component
-        class MockChatGeneratorWithRunAsync:
-            def to_dict(self) -> Dict[str, Any]:
-                return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}}
-
-            @classmethod
-            def from_dict(cls, data: Dict[str, Any]) -> "MockChatGeneratorWithoutRunAsync":
-                return cls()
-
-            def run(
-                self, messages: List[ChatMessage], tools: Optional[Union[List[Tool], Toolset]] = None, **kwargs
-            ) -> Dict[str, Any]:
-                return {"replies": [ChatMessage.from_assistant("Hello")]}
-
-            async def run_async(
-                self, messages: List[ChatMessage], tools: Optional[Union[List[Tool], Toolset]] = None, **kwargs
-            ) -> Dict[str, Any]:
-                return {"replies": [ChatMessage.from_assistant("Hello from run_async")]}
-
         chat_generator = MockChatGeneratorWithRunAsync()
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
         agent.warm_up()
@@ -586,3 +633,260 @@ async def run_async(
         assert len(result["messages"]) == 2
         assert [isinstance(reply, ChatMessage) for reply in result["messages"]]
         assert "Hello from run_async" in result["messages"][1].text
+
+
+class TestAgentTracing:
+    def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool):
+        chat_generator = MockChatGeneratorWithoutRunAsync()
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
+
+        tracing.tracer.is_content_tracing_enabled = True
+        tracing.enable_tracing(LoggingTracer())
+        caplog.set_level(logging.DEBUG)
+
+        _ = agent.run([ChatMessage.from_user("What's the weather in Paris?")])
+
+        # Ensure tracing span was emitted
+        assert any("Operation: haystack.component.run" in record.message for record in caplog.records)
+
+        # Check specific tags
+        tags_records = [r for r in caplog.records if hasattr(r, "tag_name")]
+
+        expected_tag_names = [
+            "haystack.component.name",
+            "haystack.component.type",
+            "haystack.component.input_types",
+            "haystack.component.input_spec",
+            "haystack.component.output_spec",
+            "haystack.component.input",
+            "haystack.component.visits",
+            "haystack.component.output",
+            "haystack.agent.max_steps",
+            "haystack.agent.tools",
+            "haystack.agent.exit_conditions",
+            "haystack.agent.state_schema",
+            "haystack.agent.input",
+            "haystack.agent.output",
+            "haystack.agent.steps_taken",
+        ]
+
+        expected_tag_values = [
+            "chat_generator",
+            "MockChatGeneratorWithoutRunAsync",
+            {"messages": "list", "tools": "list"},
+            {},
+            {},
+            {
+                "messages": [ChatMessage.from_user(text="What's the weather in Paris?")],
+                "tools": [
+                    Tool(
+                        name="weather_tool",
+                        description="Provides weather information for a given location.",
+                        parameters={
+                            "type": "object",
+                            "properties": {"location": {"type": "string"}},
+                            "required": ["location"],
+                        },
+                        function=weather_function,
+                        outputs_to_string=None,
+                        inputs_from_state=None,
+                        outputs_to_state=None,
+                    )
+                ],
+            },
+            1,
+            {"replies": [ChatMessage.from_assistant(text="Hello")]},
+            100,
+            [
+                Tool(
+                    name="weather_tool",
+                    description="Provides weather information for a given location.",
+                    parameters={
+                        "type": "object",
+                        "properties": {"location": {"type": "string"}},
+                        "required": ["location"],
+                    },
+                    function=weather_function,
+                    outputs_to_string=None,
+                    inputs_from_state=None,
+                    outputs_to_state=None,
+                )
+            ],
+            ["text"],
+            {"messages": {"type": List[ChatMessage], "handler": merge_lists}},
+            {"messages": [ChatMessage.from_user(text="What's the weather in Paris?")], "streaming_callback": None},
+            {
+                "messages": [
+                    ChatMessage.from_user(text="What's the weather in Paris?"),
+                    ChatMessage.from_assistant(text="Hello"),
+                ]
+            },
+            1,
+        ]
+        for idx, record in enumerate(tags_records):
+            assert record.tag_name == expected_tag_names[idx]
+            assert record.tag_value == expected_tag_values[idx]
+
+        # Clean up
+        tracing.tracer.is_content_tracing_enabled = False
+        tracing.disable_tracing()
+
+    @pytest.mark.asyncio
+    async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_tool):
+        chat_generator = MockChatGeneratorWithRunAsync()
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
+
+        tracing.tracer.is_content_tracing_enabled = True
+        tracing.enable_tracing(LoggingTracer())
+        caplog.set_level(logging.DEBUG)
+
+        _ = await agent.run_async([ChatMessage.from_user("What's the weather in Paris?")])
+
+        # Ensure tracing span was emitted
+        assert any("Operation: haystack.component.run" in record.message for record in caplog.records)
+
+        # Check specific tags
+        tags_records = [r for r in caplog.records if hasattr(r, "tag_name")]
+
+        expected_tag_names = [
+            "haystack.component.name",
+            "haystack.component.type",
+            "haystack.component.input_types",
+            "haystack.component.input_spec",
+            "haystack.component.output_spec",
+            "haystack.component.input",
+            "haystack.component.visits",
+            "haystack.component.output",
+            "haystack.agent.max_steps",
+            "haystack.agent.tools",
+            "haystack.agent.exit_conditions",
+            "haystack.agent.state_schema",
+            "haystack.agent.input",
+            "haystack.agent.output",
+            "haystack.agent.steps_taken",
+        ]
+
+        expected_tag_values = [
+            "chat_generator",
+            "MockChatGeneratorWithRunAsync",
+            {"messages": "list", "tools": "list"},
+            {},
+            {},
+            {
+                "messages": [ChatMessage.from_user(text="What's the weather in Paris?")],
+                "tools": [
+                    Tool(
+                        name="weather_tool",
+                        description="Provides weather information for a given location.",
+                        parameters={
+                            "type": "object",
+                            "properties": {"location": {"type": "string"}},
+                            "required": ["location"],
+                        },
+                        function=weather_function,
+                        outputs_to_string=None,
+                        inputs_from_state=None,
+                        outputs_to_state=None,
+                    )
+                ],
+            },
+            1,
+            {"replies": [ChatMessage.from_assistant(text="Hello from run_async")]},
+            100,
+            [
+                Tool(
+                    name="weather_tool",
+                    description="Provides weather information for a given location.",
+                    parameters={
+                        "type": "object",
+                        "properties": {"location": {"type": "string"}},
+                        "required": ["location"],
+                    },
+                    function=weather_function,
+                    outputs_to_string=None,
+                    inputs_from_state=None,
+                    outputs_to_state=None,
+                )
+            ],
+            ["text"],
+            {"messages": {"type": List[ChatMessage], "handler": merge_lists}},
+            {"messages": [ChatMessage.from_user(text="What's the weather in Paris?")], "streaming_callback": None},
+            {
+                "messages": [
+                    ChatMessage.from_user(text="What's the weather in Paris?"),
+                    ChatMessage.from_assistant(text="Hello from run_async"),
+                ]
+            },
+            1,
+        ]
+        for idx, record in enumerate(tags_records):
+            assert record.tag_name == expected_tag_names[idx]
+            assert record.tag_value == expected_tag_values[idx]
+
+        # Clean up
+        tracing.tracer.is_content_tracing_enabled = False
+        tracing.disable_tracing()
+
+    def test_agent_tracing_in_pipeline(self, caplog, monkeypatch, weather_tool):
+        chat_generator = MockChatGeneratorWithoutRunAsync()
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
+        agent.warm_up()
+
+        tracing.tracer.is_content_tracing_enabled = True
+        tracing.enable_tracing(LoggingTracer())
+        caplog.set_level(logging.DEBUG)
+
+        pipeline = Pipeline()
+        pipeline.add_component(
+            "prompt_builder", ChatPromptBuilder(template=[ChatMessage.from_user("Hello {{location}}")])
+        )
+        pipeline.add_component("agent", agent)
+        pipeline.connect("prompt_builder.prompt", "agent.messages")
+
+        pipeline.run(data={"prompt_builder": {"location": "Berlin"}})
+
+        assert any("Operation: haystack.pipeline.run" in record.message for record in caplog.records)
+        tags_records = [r for r in caplog.records if hasattr(r, "tag_name")]
+        expected_tag_names = [
+            "haystack.component.name",
+            "haystack.component.type",
+            "haystack.component.input_types",
+            "haystack.component.input_spec",
+            "haystack.component.output_spec",
+            "haystack.component.input",
+            "haystack.component.visits",
+            "haystack.component.output",
+            "haystack.component.name",
+            "haystack.component.type",
+            "haystack.component.input_types",
+            "haystack.component.input_spec",
+            "haystack.component.output_spec",
+            "haystack.component.input",
+            "haystack.component.visits",
+            "haystack.component.output",
+            "haystack.agent.max_steps",
+            "haystack.agent.tools",
+            "haystack.agent.exit_conditions",
+            "haystack.agent.state_schema",
+            "haystack.agent.input",
+            "haystack.agent.output",
+            "haystack.agent.steps_taken",
+            "haystack.component.name",
+            "haystack.component.type",
+            "haystack.component.input_types",
+            "haystack.component.input_spec",
+            "haystack.component.output_spec",
+            "haystack.component.input",
+            "haystack.component.visits",
+            "haystack.component.output",
+            "haystack.pipeline.input_data",
+            "haystack.pipeline.output_data",
+            "haystack.pipeline.metadata",
+            "haystack.pipeline.max_runs_per_component",
+        ]
+        for idx, record in enumerate(tags_records):
+            assert record.tag_name == expected_tag_names[idx]
+
+        # Clean up
+        tracing.tracer.is_content_tracing_enabled = False
+        tracing.disable_tracing()
diff --git a/test/core/pipeline/test_pipeline.py b/test/core/pipeline/test_pipeline.py
index ebd8e7c53d..f4c18c130a 100644
--- a/test/core/pipeline/test_pipeline.py
+++ b/test/core/pipeline/test_pipeline.py
@@ -40,25 +40,40 @@ def test_pipeline_thread_safety(self, waiting_component, spying_tracer):
         for span in component_spans:
             assert span.tags["haystack.component.visits"] == 1
 
+    def test_prepare_component_inputs(self):
+        joiner_1 = BranchJoiner(type_=str)
+        joiner_2 = BranchJoiner(type_=str)
+        pp = Pipeline()
+        component_name = "joiner_1"
+        pp.add_component(component_name, joiner_1)
+        pp.add_component("joiner_2", joiner_2)
+        pp.connect(component_name, "joiner_2")
+        inputs = {"joiner_1": {"value": [{"sender": None, "value": "test_value"}]}}
+        comp_dict = pp._get_component_with_graph_metadata_and_visits(component_name, 0)
+
+        _ = pp._consume_component_inputs(component_name=component_name, component=comp_dict, inputs=inputs)
+        # We remove input in greedy variadic sockets, even if they are from the user
+        assert inputs == {"joiner_1": {}}
+
     def test__run_component_success(self):
         """Test successful component execution"""
         joiner_1 = BranchJoiner(type_=str)
         joiner_2 = BranchJoiner(type_=str)
         pp = Pipeline()
-        pp.add_component("joiner_1", joiner_1)
+        component_name = "joiner_1"
+        pp.add_component(component_name, joiner_1)
         pp.add_component("joiner_2", joiner_2)
-        pp.connect("joiner_1", "joiner_2")
-        inputs = {"joiner_1": {"value": [{"sender": None, "value": "test_value"}]}}
+        pp.connect(component_name, "joiner_2")
+        inputs = {"value": ["test_value"]}
 
         outputs = pp._run_component(
-            component=pp._get_component_with_graph_metadata_and_visits("joiner_1", 0),
+            component_name=component_name,
+            component=pp._get_component_with_graph_metadata_and_visits(component_name, 0),
             inputs=inputs,
-            component_visits={"joiner_1": 0, "joiner_2": 0},
+            component_visits={component_name: 0, "joiner_2": 0},
         )
 
         assert outputs == {"value": "test_value"}
-        # We remove input in greedy variadic sockets, even if they are from the user
-        assert "value" not in inputs["joiner_1"]
 
     def test__run_component_fail(self):
         """Test error when component doesn't return a dictionary"""
@@ -72,11 +87,11 @@ def run(self, value: str):
         wrong = WrongOutput()
         pp = Pipeline()
         pp.add_component("wrong", wrong)
-
-        inputs = {"wrong": {"value": [{"sender": None, "value": "test_value"}]}}
+        inputs = {"value": "test_value"}
 
         with pytest.raises(PipelineRuntimeError) as exc_info:
             pp._run_component(
+                component_name="wrong",
                 component=pp._get_component_with_graph_metadata_and_visits("wrong", 0),
                 inputs=inputs,
                 component_visits={"wrong": 0},
@@ -101,6 +116,7 @@ def run(self):
 
         with pytest.raises(PipelineRuntimeError) as exc_info:
             pp._run_component(
+                component_name="erroring_component",
                 component=pp._get_component_with_graph_metadata_and_visits("erroring_component", 0),
                 inputs=inputs,
                 component_visits={"erroring_component": 0},