deepset-ai--haystack/haystack/tracing/logging_tracer.py

# 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