fix: Fix tracing of a pipelines outputs (#12090)

Author: sjrl

haystack/core/pipeline/pipeline.py

@@ -386,7 +386,6 @@ def run(  # noqa: PLR0915, PLR0912, C901
             "haystack.pipeline.run",
             tags={
                 "haystack.pipeline.input_data": data,
-                "haystack.pipeline.output_data": pipeline_outputs,
                 "haystack.pipeline.metadata": self.metadata,
                 "haystack.pipeline.max_runs_per_component": self._max_runs_per_component,
                 "haystack.pipeline.execution_mode": "sync",
@@ -517,6 +516,9 @@ def run(  # noqa: PLR0915, PLR0912, C901
                     break_point=break_point,
                 )
 
+            # Set here so the tag reflects the final outputs.
+            span.set_content_tag("haystack.pipeline.output_data", pipeline_outputs)
+
             return pipeline_outputs
 
     @staticmethod
@@ -907,7 +909,6 @@ async def process_results():
             "haystack.pipeline.run",
             tags={
                 "haystack.pipeline.input_data": data,
-                "haystack.pipeline.output_data": pipeline_outputs,
                 "haystack.pipeline.metadata": self.metadata,
                 "haystack.pipeline.max_runs_per_component": self._max_runs_per_component,
                 "haystack.pipeline.execution_mode": "async",
@@ -1051,6 +1052,9 @@ async def process_results():
                 ):
                     yield partial_outputs
 
+                # Set here so the tag reflects the final outputs.
+                parent_span.set_content_tag("haystack.pipeline.output_data", pipeline_outputs)
+
                 # Yield the final pipeline outputs.
                 yield pipeline_outputs
             finally:

releasenotes/notes/fix-pipeline-output-data-trace-tag-6f9e2c1a4b8d3e70.yaml

@@ -0,0 +1,8 @@
+---
+fixes:
+  - |
+    Fixed the ``haystack.pipeline.output_data`` tracing tag being empty. The tag was set at the start of
+    ``Pipeline.run``/``run_async`` from the still-empty outputs, and since tracing backends coerce a tag value
+    when it is set, the recorded output was always an empty dictionary. It is now set once the run completes so
+    it reflects the final pipeline outputs. The tag is also gated behind content tracing
+    (``HAYSTACK_CONTENT_TRACING_ENABLED``), consistent with the component-level input/output tags.

test/conftest.py

@@ -15,7 +15,7 @@
 from haystack.core.serialization import allow_deserialization_module
 from haystack.document_stores.in_memory import InMemoryDocumentStore
 from haystack.testing.test_utils import set_all_seeds
-from test.tracing.utils import SpyingTracer
+from test.tracing.utils import EagerSpyingTracer, SpyingTracer
 
 set_all_seeds(0)
 
@@ -108,6 +108,18 @@ def spying_tracer() -> Generator[SpyingTracer, None, None]:
     tracing.disable_tracing()
 
 
+@pytest.fixture()
+def eager_spying_tracer() -> Generator[EagerSpyingTracer, None, None]:
+    # Coerces tags when set, mirroring real backends. Content tracing is left to the test to toggle.
+    tracer = EagerSpyingTracer()
+    tracing.enable_tracing(tracer)
+
+    yield tracer
+
+    # Make sure to disable tracing after the test to avoid affecting other tests
+    tracing.disable_tracing()
+
+
 @pytest.fixture()
 def base64_image_string():
     return "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+ip1sAAAAASUVORK5CYII="

test/core/pipeline/test_tracing.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import asyncio
+import json
 from dataclasses import replace
 from unittest.mock import ANY
 
@@ -156,6 +157,30 @@ def test_with_enabled_content_tracing(
             ),
         ]
 
+    @pytest.mark.parametrize("run_async", [False, True])
+    @pytest.mark.parametrize("content_tracing", [True, False])
+    def test_pipeline_output_data_trace_tag(
+        self, pipeline, run_async, content_tracing, eager_spying_tracer, monkeypatch
+    ):
+        """
+        output_data holds the final outputs and is gated behind content tracing.
+
+        EagerSpyingTracer coerces tags when set (like real backends), so it catches output_data being set from the
+        still-empty outputs dict at span start.
+        """
+        monkeypatch.setattr(tracer, "is_content_tracing_enabled", content_tracing)
+        if run_async:
+            asyncio.run(pipeline.run_async(data={"word": "world"}))
+        else:
+            pipeline.run(data={"word": "world"})
+
+        tags = eager_spying_tracer.spans[0].tags
+        assert json.loads(tags["haystack.pipeline.input_data"]) == {"hello": {"word": "world"}}
+        if content_tracing:
+            assert json.loads(tags["haystack.pipeline.output_data"]) == {"hello2": {"output": "Hello, Hello, world!!"}}
+        else:
+            assert "haystack.pipeline.output_data" not in tags
+
     @pytest.mark.parametrize("run_async", [False, True])
     def test_span_input_not_affected_by_component_mutation(self, run_async, spying_tracer, monkeypatch):
         """

test/tracing/utils.py

@@ -9,6 +9,7 @@
 from typing import Any
 
 from haystack.tracing import Span, Tracer
+from haystack.tracing.utils import coerce_tag_value
 
 
 @dataclasses.dataclass
@@ -52,3 +53,41 @@ def trace(
         self.spans.append(new_span)
 
         yield new_span
+
+
+@dataclasses.dataclass
+class EagerSpyingSpan(Span):
+    """
+    A span that coerces tag values when `set_tag` is called.
+
+    This mirrors real tracing backends (e.g. OpenTelemetry, Datadog), which serialize a tag value eagerly
+    when it is set instead of holding a live reference to it. It is useful to catch tags that are set from a
+    mutable object before that object has been populated.
+    """
+
+    operation_name: str
+    parent_span: Span | None = None
+    tags: dict[str, Any] = dataclasses.field(default_factory=dict)
+
+    def set_tag(self, key: str, value: Any) -> None:
+        self.tags[key] = coerce_tag_value(value)
+
+
+class EagerSpyingTracer(Tracer):
+    def current_span(self) -> Span | None:
+        return self.spans[-1] if self.spans else None
+
+    def __init__(self) -> None:
+        self.spans: list[EagerSpyingSpan] = []
+
+    @contextlib.contextmanager
+    def trace(
+        self, operation_name: str, tags: dict[str, Any] | None = None, parent_span: Span | None = None
+    ) -> Iterator[Span]:
+        new_span = EagerSpyingSpan(operation_name, parent_span)
+        for key, value in (tags or {}).items():
+            new_span.set_tag(key, value)
+
+        self.spans.append(new_span)
+
+        yield new_span