From b20a14d6fe20399024d2df0df05f2aed96d1b1be Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Mon, 8 Apr 2024 10:51:25 +0200 Subject: [PATCH] fix docstrints format for tracing package (#7497) --- haystack/tracing/datadog.py | 14 ++++++++++++++ haystack/tracing/opentelemetry.py | 14 ++++++++++++++ haystack/tracing/tracer.py | 32 ++++++++++++++++++++++--------- haystack/tracing/utils.py | 3 ++- 4 files changed, 53 insertions(+), 10 deletions(-) diff --git a/haystack/tracing/datadog.py b/haystack/tracing/datadog.py index 6badbb107..a0d87fdae 100644 --- a/haystack/tracing/datadog.py +++ b/haystack/tracing/datadog.py @@ -14,13 +14,25 @@ class DatadogSpan(Span): 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 @@ class DatadogTracer(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.""" with self._tracer.trace(operation_name) as span: span = DatadogSpan(span) if tags: @@ -52,6 +65,7 @@ class DatadogTracer(Tracer): 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 78cafb699..4658b9e06 100644 --- a/haystack/tracing/opentelemetry.py +++ b/haystack/tracing/opentelemetry.py @@ -15,13 +15,25 @@ class OpenTelemetrySpan(Span): 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 @@ class OpenTelemetryTracer(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.""" with self._tracer.start_as_current_span(operation_name) as raw_span: span = OpenTelemetrySpan(raw_span) if tags: @@ -41,6 +54,7 @@ class OpenTelemetryTracer(Tracer): 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 f42443694..1197bedb0 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 @@ class Span(abc.ABC): 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 @@ class Span(abc.ABC): 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 @@ class Span(abc.ABC): 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 @@ class Span(abc.ABC): 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 @@ class Tracer(abc.ABC): @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 @@ class Tracer(abc.ABC): 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 @@ class ProxyTracer(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.""" 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 532a66ab4..2f461de8f 100644 --- a/haystack/tracing/utils.py +++ b/haystack/tracing/utils.py @@ -9,7 +9,8 @@ PRIMITIVE_TYPES = (bool, str, int, float) 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.