From 9b72287401d5c424a8951e1d6a15cca14fcd05cc Mon Sep 17 00:00:00 2001 From: Roger Yang <80478925+RogerHYang@users.noreply.github.com> Date: Wed, 14 Aug 2024 17:12:41 -0700 Subject: [PATCH] fix: mask attributes when starting span (#892) --- .../openinference/instrumentation/config.py | 36 ++++++++++++++----- .../tests/conftest.py | 17 +++++++++ .../tests/test_config.py | 15 ++++++++ python/tox.ini | 2 +- 4 files changed, 60 insertions(+), 10 deletions(-) create mode 100644 python/openinference-instrumentation/tests/conftest.py diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py index fc16f0efd..9dddc846b 100644 --- a/python/openinference-instrumentation/src/openinference/instrumentation/config.py +++ b/python/openinference-instrumentation/src/openinference/instrumentation/config.py @@ -1,6 +1,8 @@ import os from contextlib import contextmanager from dataclasses import dataclass, field, fields +from functools import partial +from inspect import signature from typing import ( Any, Callable, @@ -8,6 +10,7 @@ Iterator, Optional, Union, + cast, get_args, ) @@ -26,6 +29,7 @@ detach, set_value, ) +from opentelemetry.trace import Tracer from opentelemetry.util.types import AttributeValue from .logging import logger @@ -324,21 +328,35 @@ def set_attribute( span.set_attribute(key, value) +class _TracerSignatures: + # remove `self` via partial + start_as_current_span = signature(partial(Tracer.start_as_current_span, None)) + start_span = signature(partial(Tracer.start_span, None)) + + class OITracer(wrapt.ObjectProxy): # type: ignore[misc] def __init__(self, wrapped: trace_api.Tracer, config: TraceConfig) -> None: super().__init__(wrapped) self._self_config = config @contextmanager - def start_as_current_span(self, *args, **kwargs) -> Iterator[trace_api.Span]: - with self.__wrapped__.start_as_current_span(*args, **kwargs) as span: - yield _MaskedSpan(span, self._self_config) - - def start_span(self, *args, **kwargs) -> trace_api.Span: - return _MaskedSpan( - self.__wrapped__.start_span(*args, **kwargs), - config=self._self_config, - ) + def start_as_current_span(self, *args: Any, **kwargs: Any) -> Iterator[trace_api.Span]: + kwargs = _TracerSignatures.start_as_current_span.bind(*args, **kwargs).arguments + attributes = cast(Optional[Dict[str, AttributeValue]], kwargs.pop("attributes", None)) + with self.__wrapped__.start_as_current_span(**kwargs) as span: + span = _MaskedSpan(span, self._self_config) + if attributes: + span.set_attributes(attributes) + yield span + + def start_span(self, *args: Any, **kwargs: Any) -> trace_api.Span: + kwargs = _TracerSignatures.start_span.bind(*args, **kwargs).arguments + attributes = cast(Optional[Dict[str, AttributeValue]], kwargs.pop("attributes", None)) + span = self.__wrapped__.start_span(**kwargs) + span = _MaskedSpan(span, config=self._self_config) + if attributes: + span.set_attributes(attributes) + return span def is_base64_url(url: str) -> bool: diff --git a/python/openinference-instrumentation/tests/conftest.py b/python/openinference-instrumentation/tests/conftest.py new file mode 100644 index 000000000..5db5268ad --- /dev/null +++ b/python/openinference-instrumentation/tests/conftest.py @@ -0,0 +1,17 @@ +import pytest +from opentelemetry.sdk import trace as trace_sdk +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import TracerProvider + + +@pytest.fixture(scope="module") +def in_memory_span_exporter() -> InMemorySpanExporter: + return InMemorySpanExporter() + + +@pytest.fixture(scope="module") +def tracer_provider(in_memory_span_exporter: InMemorySpanExporter) -> TracerProvider: + tracer_provider = trace_sdk.TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(in_memory_span_exporter)) + return tracer_provider diff --git a/python/openinference-instrumentation/tests/test_config.py b/python/openinference-instrumentation/tests/test_config.py index 75eba2165..4f7aa6929 100644 --- a/python/openinference-instrumentation/tests/test_config.py +++ b/python/openinference-instrumentation/tests/test_config.py @@ -20,7 +20,11 @@ OPENINFERENCE_HIDE_OUTPUT_MESSAGES, OPENINFERENCE_HIDE_OUTPUT_TEXT, OPENINFERENCE_HIDE_OUTPUTS, + REDACTED_VALUE, + OITracer, ) +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import TracerProvider def test_default_settings() -> None: @@ -35,6 +39,17 @@ def test_default_settings() -> None: assert config.base64_image_max_length == DEFAULT_BASE64_IMAGE_MAX_LENGTH +def test_oi_tracer( + tracer_provider: TracerProvider, + in_memory_span_exporter: InMemorySpanExporter, +) -> None: + tracer = OITracer(tracer_provider.get_tracer(__name__), TraceConfig(hide_inputs=True)) + with tracer.start_as_current_span("a", attributes={"input.value": "c"}): + tracer.start_span("b", attributes={"input.value": "d"}).end() + for span in in_memory_span_exporter.get_finished_spans(): + assert (span.attributes or {}).get("input.value") == REDACTED_VALUE + + @pytest.mark.parametrize("hide_inputs", [False, True]) @pytest.mark.parametrize("hide_outputs", [False, True]) @pytest.mark.parametrize("hide_input_messages", [False, True]) diff --git a/python/tox.ini b/python/tox.ini index 42fdd35cb..306df4575 100644 --- a/python/tox.ini +++ b/python/tox.ini @@ -42,7 +42,7 @@ changedir = litellm: instrumentation/openinference-instrumentation-litellm/ instructor: instrumentation/openinference-instrumentation-instructor/ commands_pre = - instrumentation: uv pip install --reinstall {toxinidir}/openinference-instrumentation + instrumentation: uv pip install --reinstall {toxinidir}/openinference-instrumentation[test] semconv: uv pip install --reinstall {toxinidir}/openinference-semantic-conventions bedrock: uv pip install --reinstall {toxinidir}/instrumentation/openinference-instrumentation-bedrock[test] bedrock-latest: uv pip install -U boto3