Coverage for haystack/tracing/tracer.py: 93%
54 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import abc
6import contextlib
7import os
8from collections.abc import Iterator
9from typing import Any
11HAYSTACK_CONTENT_TRACING_ENABLED_ENV_VAR = "HAYSTACK_CONTENT_TRACING_ENABLED"
14class Span(abc.ABC):
15 """Interface for an instrumented operation."""
17 @abc.abstractmethod
18 def set_tag(self, key: str, value: Any) -> None:
19 """
20 Set a single tag on the span.
22 Note that the value will be serialized to a string, so it's best to use simple types like strings, numbers, or
23 booleans.
25 :param key: the name of the tag.
26 :param value: the value of the tag.
27 """
28 pass
30 def set_tags(self, tags: dict[str, Any]) -> None:
31 """
32 Set multiple tags on the span.
34 :param tags: a mapping of tag names to tag values.
35 """
36 for key, value in tags.items():
37 self.set_tag(key, value)
39 def raw_span(self) -> Any:
40 """
41 Provides access to the underlying span object of the tracer.
43 Use this if you need full access to the underlying span object.
45 :return: The underlying span object.
46 """
47 return self
49 def set_content_tag(self, key: str, value: Any) -> None:
50 """
51 Set a single tag containing content information.
53 Content is sensitive information such as
54 - the content of a query
55 - the content of a document
56 - the content of an answer
58 By default, this behavior is disabled. To enable it
59 - set the environment variable `HAYSTACK_CONTENT_TRACING_ENABLED` to `true` or
60 - override the `set_content_tag` method in a custom tracer implementation.
62 :param key: the name of the tag.
63 :param value: the value of the tag.
64 """
65 if tracer.is_content_tracing_enabled:
66 self.set_tag(key, value)
68 def get_correlation_data_for_logs(self) -> dict[str, Any]:
69 """
70 Return a dictionary with correlation data for logs.
72 This is useful if you want to correlate logs with traces.
73 """
74 return {}
77class Tracer(abc.ABC):
78 """Interface for instrumenting code by creating and submitting spans."""
80 @abc.abstractmethod
81 @contextlib.contextmanager
82 def trace(
83 self, operation_name: str, tags: dict[str, Any] | None = None, parent_span: Span | None = None
84 ) -> Iterator[Span]:
85 """
86 Trace the execution of a block of code.
88 :param operation_name: the name of the operation being traced.
89 :param tags: tags to apply to the newly created span.
90 :param parent_span: the parent span to use for the newly created span.
91 If `None`, the newly created span will be a root span.
92 :return: the newly created span.
93 """
94 pass
96 @abc.abstractmethod
97 def current_span(self) -> Span | None:
98 """
99 Returns the currently active span. If no span is active, returns `None`.
101 :return: Currently active span or `None` if no span is active.
102 """
103 pass
106class ProxyTracer(Tracer):
107 """
108 Container for the actual tracer instance.
110 This eases
111 - replacing the actual tracer instance without having to change the global tracer instance
112 - implementing default behavior for the tracer
113 """
115 def __init__(self, provided_tracer: Tracer) -> None:
116 """Creates an instance of ProxyTracer."""
117 self.actual_tracer: Tracer = provided_tracer
118 self.is_content_tracing_enabled = os.getenv(HAYSTACK_CONTENT_TRACING_ENABLED_ENV_VAR, "false").lower() == "true"
120 @contextlib.contextmanager
121 def trace(
122 self, operation_name: str, tags: dict[str, Any] | None = None, parent_span: Span | None = None
123 ) -> Iterator[Span]:
124 """Activate and return a new span that inherits from the current active span."""
125 with self.actual_tracer.trace(operation_name, tags=tags, parent_span=parent_span) as span:
126 yield span
128 def current_span(self) -> Span | None:
129 """Return the current active span"""
130 return self.actual_tracer.current_span()
133class NullSpan(Span):
134 """A no-op implementation of the `Span` interface. This is used when tracing is disabled."""
136 def set_tag(self, key: str, value: Any) -> None:
137 """Set a single tag on the span."""
138 pass
141class NullTracer(Tracer):
142 """A no-op implementation of the `Tracer` interface. This is used when tracing is disabled."""
144 @contextlib.contextmanager
145 def trace(
146 self,
147 operation_name: str, # noqa: ARG002
148 tags: dict[str, Any] | None = None, # noqa: ARG002
149 parent_span: Span | None = None, # noqa: ARG002
150 ) -> Iterator[Span]:
151 """Activate and return a new span that inherits from the current active span."""
152 yield NullSpan()
154 def current_span(self) -> Span | None:
155 """Return the current active span"""
156 return NullSpan()
159# We use the proxy pattern to allow for easy enabling and disabling of tracing without having to change the global
160# tracer instance. That's especially convenient if users import the object directly
161# (in that case we'd have to monkey-patch it in all of these modules).
162tracer: ProxyTracer = ProxyTracer(provided_tracer=NullTracer())
165def enable_tracing(provided_tracer: Tracer) -> None:
166 """Enable tracing by setting the global tracer instance."""
167 tracer.actual_tracer = provided_tracer
170def disable_tracing() -> None:
171 """Disable tracing by setting the global tracer instance to a no-op tracer."""
172 tracer.actual_tracer = NullTracer()
175def is_tracing_enabled() -> bool:
176 """Return whether tracing is enabled."""
177 return not isinstance(tracer.actual_tracer, NullTracer)