chore: import upstream snapshot with attribution

haystack/tracing/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# ruff: noqa: F401
+
+from haystack.tracing.tracer import (  # noqa: I001 (otherwise we end up with partial imports)
+    Span,
+    Tracer,
+    disable_tracing,
+    enable_tracing,
+    is_tracing_enabled,
+    tracer,
+)

haystack/tracing/logging_tracer.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import contextlib
+import dataclasses
+import threading
+from collections.abc import Iterator
+from typing import Any
+
+from haystack import logging
+from haystack.tracing import Span, Tracer
+from haystack.tracing.utils import coerce_tag_value
+
+logger = logging.getLogger(__name__)
+
+RESET_COLOR = "\033[0m"
+
+
+@dataclasses.dataclass
+class LoggingSpan(Span):
+    operation_name: str
+    tags: dict[str, Any] = dataclasses.field(default_factory=dict)
+
+    def set_tag(self, key: str, value: Any) -> None:
+        """
+        Set a single tag on the span.
+
+        :param key: the name of the tag.
+        :param value: the value of the tag.
+        """
+        self.tags[key] = value
+
+
+class LoggingTracer(Tracer):
+    """
+    A simple tracer that logs the operation name and tags of a span.
+    """
+
+    def __init__(self, tags_color_strings: dict[str, str] | None = None) -> None:
+        """
+        Initialize the LoggingTracer.
+
+        :param tags_color_strings:
+            A dictionary that maps tag names to color strings that should be used when logging the tags.
+            The color strings should be in the format of
+            [ANSI escape codes](https://en.wikipedia.org/wiki/ANSI_escape_code#Colors).
+            For example, to color the tag "haystack.component.input" in red, you would pass
+            `tags_color_strings={"haystack.component.input": "\x1b[1;31m"}`.
+        """
+
+        self.tags_color_strings = tags_color_strings or {}
+        # Spans can be created and closed from parallel worker threads (e.g. one span per tool call in the Agent).
+        # A span's operation name and its tags are emitted as separate log records, so we serialize the emission to
+        # keep each span's records contiguous instead of interleaved with those of concurrently-closing spans.
+        self._emit_lock = threading.Lock()
+
+    @contextlib.contextmanager
+    def trace(
+        self,
+        operation_name: str,
+        tags: dict[str, Any] | None = None,
+        parent_span: Span | None = None,  # noqa: ARG002
+    ) -> Iterator[Span]:
+        """
+        Trace the execution of a block of code.
+
+        :param operation_name: the name of the operation being traced.
+        :param tags: tags to apply to the newly created span.
+        :param parent_span: the parent span to use for the newly created span. Not used in this simple tracer.
+        :returns: the newly created span.
+        """
+
+        custom_span = LoggingSpan(operation_name, tags=tags or {})
+
+        try:
+            yield custom_span
+        except Exception as e:  # noqa: TRY203
+            raise e
+        # we make sure to log the operation name and tags of the span when the context manager exits
+        # both in case of success and error
+        finally:
+            operation_name = custom_span.operation_name
+            tags = custom_span.tags or {}
+            with self._emit_lock:
+                logger.debug("Operation: {operation_name}", operation_name=operation_name)
+                for tag_name, tag_value in tags.items():
+                    color_string = self.tags_color_strings.get(tag_name, "")
+                    coerced_value = coerce_tag_value(tag_value)
+                    logger.debug(
+                        color_string + "{tag_name}={tag_value}" + RESET_COLOR,
+                        tag_name=tag_name,
+                        tag_value=coerced_value,
+                    )
+
+    def current_span(self) -> Span | None:
+        """Return the current active span, if any."""
+        # we don't store spans in this simple tracer
+        return None

haystack/tracing/tracer.py

@@ -0,0 +1,177 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+import contextlib
+import os
+from collections.abc import Iterator
+from typing import Any
+
+HAYSTACK_CONTENT_TRACING_ENABLED_ENV_VAR = "HAYSTACK_CONTENT_TRACING_ENABLED"
+
+
+class Span(abc.ABC):
+    """Interface for an instrumented operation."""
+
+    @abc.abstractmethod
+    def set_tag(self, key: str, value: Any) -> None:
+        """
+        Set a single tag on the span.
+
+        Note that the value will be serialized to a string, so it's best to use simple types like strings, numbers, or
+        booleans.
+
+        :param key: the name of the tag.
+        :param value: the value of the tag.
+        """
+        pass
+
+    def set_tags(self, tags: dict[str, Any]) -> None:
+        """
+        Set multiple tags on the span.
+
+        :param tags: a mapping of tag names to tag values.
+        """
+        for key, value in tags.items():
+            self.set_tag(key, value)
+
+    def raw_span(self) -> Any:
+        """
+        Provides access to the underlying span object of the tracer.
+
+        Use this if you need full access to the underlying span object.
+
+        :return: The underlying span object.
+        """
+        return self
+
+    def set_content_tag(self, key: str, value: Any) -> None:
+        """
+        Set a single tag containing content information.
+
+        Content is sensitive information such as
+        - the content of a query
+        - the content of a document
+        - the content of an answer
+
+        By default, this behavior is disabled. To enable it
+        - set the environment variable `HAYSTACK_CONTENT_TRACING_ENABLED` to `true` or
+        - override the `set_content_tag` method in a custom tracer implementation.
+
+        :param key: the name of the tag.
+        :param value: the value of the tag.
+        """
+        if tracer.is_content_tracing_enabled:
+            self.set_tag(key, value)
+
+    def get_correlation_data_for_logs(self) -> dict[str, Any]:
+        """
+        Return a dictionary with correlation data for logs.
+
+        This is useful if you want to correlate logs with traces.
+        """
+        return {}
+
+
+class Tracer(abc.ABC):
+    """Interface for instrumenting code by creating and submitting spans."""
+
+    @abc.abstractmethod
+    @contextlib.contextmanager
+    def trace(
+        self, operation_name: str, tags: dict[str, Any] | None = None, parent_span: Span | None = None
+    ) -> Iterator[Span]:
+        """
+        Trace the execution of a block of code.
+
+        :param operation_name: the name of the operation being traced.
+        :param tags: tags to apply to the newly created span.
+        :param parent_span: the parent span to use for the newly created span.
+            If `None`, the newly created span will be a root span.
+        :return: the newly created span.
+        """
+        pass
+
+    @abc.abstractmethod
+    def current_span(self) -> Span | None:
+        """
+        Returns the currently active span. If no span is active, returns `None`.
+
+        :return: Currently active span or `None` if no span is active.
+        """
+        pass
+
+
+class ProxyTracer(Tracer):
+    """
+    Container for the actual tracer instance.
+
+    This eases
+    - replacing the actual tracer instance without having to change the global tracer instance
+    - implementing default behavior for the tracer
+    """
+
+    def __init__(self, provided_tracer: Tracer) -> None:
+        """Creates an instance of ProxyTracer."""
+        self.actual_tracer: Tracer = provided_tracer
+        self.is_content_tracing_enabled = os.getenv(HAYSTACK_CONTENT_TRACING_ENABLED_ENV_VAR, "false").lower() == "true"
+
+    @contextlib.contextmanager
+    def trace(
+        self, operation_name: str, tags: dict[str, Any] | None = None, parent_span: Span | None = None
+    ) -> Iterator[Span]:
+        """Activate and return a new span that inherits from the current active span."""
+        with self.actual_tracer.trace(operation_name, tags=tags, parent_span=parent_span) as span:
+            yield span
+
+    def current_span(self) -> Span | None:
+        """Return the current active span"""
+        return self.actual_tracer.current_span()
+
+
+class NullSpan(Span):
+    """A no-op implementation of the `Span` interface. This is used when tracing is disabled."""
+
+    def set_tag(self, key: str, value: Any) -> None:
+        """Set a single tag on the span."""
+        pass
+
+
+class NullTracer(Tracer):
+    """A no-op implementation of the `Tracer` interface. This is used when tracing is disabled."""
+
+    @contextlib.contextmanager
+    def trace(
+        self,
+        operation_name: str,  # noqa: ARG002
+        tags: dict[str, Any] | None = None,  # noqa: ARG002
+        parent_span: Span | None = None,  # noqa: ARG002
+    ) -> Iterator[Span]:
+        """Activate and return a new span that inherits from the current active span."""
+        yield NullSpan()
+
+    def current_span(self) -> Span | None:
+        """Return the current active span"""
+        return NullSpan()
+
+
+# We use the proxy pattern to allow for easy enabling and disabling of tracing without having to change the global
+# tracer instance. That's especially convenient if users import the object directly
+# (in that case we'd have to monkey-patch it in all of these modules).
+tracer: ProxyTracer = ProxyTracer(provided_tracer=NullTracer())
+
+
+def enable_tracing(provided_tracer: Tracer) -> None:
+    """Enable tracing by setting the global tracer instance."""
+    tracer.actual_tracer = provided_tracer
+
+
+def disable_tracing() -> None:
+    """Disable tracing by setting the global tracer instance to a no-op tracer."""
+    tracer.actual_tracer = NullTracer()
+
+
+def is_tracing_enabled() -> bool:
+    """Return whether tracing is enabled."""
+    return not isinstance(tracer.actual_tracer, NullTracer)

haystack/tracing/utils.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+from typing import Any
+
+from haystack import logging
+
+logger = logging.getLogger(__name__)
+
+PRIMITIVE_TYPES = (bool, str, int, float)
+
+
+def coerce_tag_value(value: Any) -> bool | str | int | float:
+    """
+    Coerces span tag values to compatible types for the tracing backend.
+
+    Most tracing libraries don't support sending complex types to the backend. Hence, we need to convert them to
+    compatible types.
+
+    :param value: an arbitrary value which should be coerced to a compatible type
+    :return: the value coerced to a compatible type
+    """
+    if isinstance(value, PRIMITIVE_TYPES):
+        return value
+
+    if value is None:
+        return ""
+
+    try:
+        # do that with-in try-except because who knows what kind of objects are being passed
+        serializable = _serializable_value(value=value, use_placeholders=True)
+        return json.dumps(serializable)
+    except Exception as error:
+        logger.debug("Failed to coerce tag value to string: {error}", error=error)
+
+        # Our last resort is to convert the value to a string
+        return str(value)
+
+
+def _serializable_value(value: Any, use_placeholders: bool = True) -> Any:
+    """
+    Serializes a value into a format suitable for tracing.
+
+    One-way only: this is never deserialized, so it's allowed to lose data (e.g. replace large
+    objects with a placeholder). Don't swap it for `base_serialization`'s schema serializer, which
+    is built to round-trip exactly.
+
+    :param value:
+        The value to serialize.
+    :param use_placeholders:
+        Whether to use string placeholders for large objects like ByteStream and ImageContent.
+    :returns:
+        The serialized value.
+    """
+    if isinstance(value, list):
+        return [_serializable_value(value=v, use_placeholders=use_placeholders) for v in value]
+
+    if isinstance(value, dict):
+        return {k: _serializable_value(value=v, use_placeholders=use_placeholders) for k, v in value.items()}
+
+    if use_placeholders and getattr(value, "_to_trace_dict", None):
+        return _serializable_value(value._to_trace_dict())
+
+    if getattr(value, "to_dict", None):
+        return _serializable_value(value.to_dict())
+
+    return value