diff --git a/haystack/tracing/datadog.py b/haystack/tracing/datadog.py index 6badbb1070..a0d87fdaea 100644 --- a/haystack/tracing/datadog.py +++ b/haystack/tracing/datadog.py @@ -14,13 +14,25 @@ def __init__(self, span: "ddtrace.Span") -> None: self._span = span 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. + """ coerced_value = tracing_utils.coerce_tag_value(value) self._span.set_tag(key, coerced_value) def raw_span(self) -> Any: + """ + Provides access to the underlying span object of the tracer. + + :return: The underlying span object. + """ return self._span def get_correlation_data_for_logs(self) -> Dict[str, Any]: + """Return a dictionary with correlation data for logs.""" raw_span = self.raw_span() if not raw_span: return {} @@ -44,6 +56,7 @@ def __init__(self, tracer: "ddtrace.Tracer") -> None: @contextlib.contextmanager def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> Iterator[Span]: + """Activate and return a new span that inherits from the current active span.""" with self._tracer.trace(operation_name) as span: span = DatadogSpan(span) if tags: @@ -52,6 +65,7 @@ def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> I yield span def current_span(self) -> Optional[Span]: + """Return the current active span""" current_span = self._tracer.current_span() if current_span is None: return None diff --git a/haystack/tracing/opentelemetry.py b/haystack/tracing/opentelemetry.py index 78cafb6995..4658b9e06a 100644 --- a/haystack/tracing/opentelemetry.py +++ b/haystack/tracing/opentelemetry.py @@ -15,13 +15,25 @@ def __init__(self, span: "opentelemetry.trace.Span") -> None: self._span = span 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. + """ coerced_value = tracing_utils.coerce_tag_value(value) self._span.set_attribute(key, coerced_value) def raw_span(self) -> Any: + """ + Provides access to the underlying span object of the tracer. + + :return: The underlying span object. + """ return self._span def get_correlation_data_for_logs(self) -> Dict[str, Any]: + """Return a dictionary with correlation data for logs.""" span_context = self._span.get_span_context() return {"trace_id": span_context.trace_id, "span_id": span_context.span_id} @@ -33,6 +45,7 @@ def __init__(self, tracer: "opentelemetry.trace.Tracer") -> None: @contextlib.contextmanager def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> Iterator[Span]: + """Activate and return a new span that inherits from the current active span.""" with self._tracer.start_as_current_span(operation_name) as raw_span: span = OpenTelemetrySpan(raw_span) if tags: @@ -41,6 +54,7 @@ def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> I yield span def current_span(self) -> Optional[Span]: + """Return the current active span""" current_span = opentelemetry.trace.get_current_span() if isinstance(current_span, opentelemetry.trace.NonRecordingSpan): return None diff --git a/haystack/tracing/tracer.py b/haystack/tracing/tracer.py index f42443694b..1197bedb01 100644 --- a/haystack/tracing/tracer.py +++ b/haystack/tracing/tracer.py @@ -16,7 +16,8 @@ class Span(abc.ABC): @abc.abstractmethod def set_tag(self, key: str, value: Any) -> None: - """Set a single tag on the span. + """ + 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. @@ -27,7 +28,8 @@ def set_tag(self, key: str, value: Any) -> None: pass def set_tags(self, tags: Dict[str, Any]) -> None: - """Set multiple tags on the span. + """ + Set multiple tags on the span. :param tags: a mapping of tag names to tag values. """ @@ -35,7 +37,8 @@ def set_tags(self, tags: Dict[str, Any]) -> None: self.set_tag(key, value) def raw_span(self) -> Any: - """Provides access to the underlying span object of the tracer. + """ + Provides access to the underlying span object of the tracer. Use this if you need full access to the underlying span object. @@ -44,7 +47,8 @@ def raw_span(self) -> Any: return self def set_content_tag(self, key: str, value: Any) -> None: - """Set a single tag containing content information. + """ + Set a single tag containing content information. Content is sensitive information such as - the content of a query @@ -62,7 +66,8 @@ def set_content_tag(self, key: str, value: Any) -> None: self.set_tag(key, value) def get_correlation_data_for_logs(self) -> Dict[str, Any]: - """Return a dictionary with correlation data for logs. + """ + Return a dictionary with correlation data for logs. This is useful if you want to correlate logs with traces.""" return {} @@ -74,7 +79,8 @@ class Tracer(abc.ABC): @abc.abstractmethod @contextlib.contextmanager def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> Iterator[Span]: - """Trace the execution of a block of code. + """ + 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. @@ -84,7 +90,8 @@ def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> I @abc.abstractmethod def current_span(self) -> Optional[Span]: - """Returns the currently active span. If no span is active, returns `None`. + """ + Returns the currently active span. If no span is active, returns `None`. :return: Currently active span or `None` if no span is active. """ @@ -92,7 +99,8 @@ def current_span(self) -> Optional[Span]: class ProxyTracer(Tracer): - """Container for the actual tracer instance. + """ + Container for the actual tracer instance. This eases - replacing the actual tracer instance without having to change the global tracer instance @@ -105,10 +113,12 @@ def __init__(self, provided_tracer: Tracer) -> None: @contextlib.contextmanager def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = 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) as span: yield span def current_span(self) -> Optional[Span]: + """Return the current active span""" return self.actual_tracer.current_span() @@ -116,6 +126,7 @@ 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 @@ -124,9 +135,11 @@ class NullTracer(Tracer): @contextlib.contextmanager def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> Iterator[Span]: + """Activate and return a new span that inherits from the current active span.""" yield NullSpan() def current_span(self) -> Optional[Span]: + """Return the current active span""" return NullSpan() @@ -152,7 +165,8 @@ def is_tracing_enabled() -> bool: def auto_enable_tracing() -> None: - """Auto-enable the right tracing backend. + """ + Auto-enable the right tracing backend. This behavior can be disabled by setting the environment variable `HAYSTACK_AUTO_TRACE_ENABLED` to `false`. Note that it will only work correctly if tracing was configured _before_ Haystack is imported. diff --git a/haystack/tracing/utils.py b/haystack/tracing/utils.py index 532a66ab4f..2f461de8f7 100644 --- a/haystack/tracing/utils.py +++ b/haystack/tracing/utils.py @@ -9,7 +9,8 @@ def coerce_tag_value(value: Any) -> Union[bool, str, int, float]: - """Coerces span tag values to compatible types for the tracing backend. + """ + 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.