diff --git a/python/openinference-instrumentation/pyproject.toml b/python/openinference-instrumentation/pyproject.toml index ffe2c4324..5424cfcaf 100644 --- a/python/openinference-instrumentation/pyproject.toml +++ b/python/openinference-instrumentation/pyproject.toml @@ -57,6 +57,7 @@ exclude = [ [tool.pytest.ini_options] asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" testpaths = [ "tests", ] diff --git a/python/openinference-instrumentation/src/openinference/instrumentation/config.py b/python/openinference-instrumentation/src/openinference/instrumentation/config.py index 1d441f6f8..dd52b8a30 100644 --- a/python/openinference-instrumentation/src/openinference/instrumentation/config.py +++ b/python/openinference-instrumentation/src/openinference/instrumentation/config.py @@ -70,6 +70,7 @@ def __aexit__(self, exc_type, exc_value, traceback) -> None: detach(self._token) +OPENINFERENCE_HIDE_LLM_INVOCATION_PARAMETERS = "OPENINFERENCE_HIDE_LLM_INVOCATION_PARAMETERS" OPENINFERENCE_HIDE_INPUTS = "OPENINFERENCE_HIDE_INPUTS" # Hides input value & messages OPENINFERENCE_HIDE_OUTPUTS = "OPENINFERENCE_HIDE_OUTPUTS" @@ -91,6 +92,7 @@ def __aexit__(self, exc_type, exc_value, traceback) -> None: REDACTED_VALUE = "__REDACTED__" # When a value is hidden, it will be replaced by this redacted value +DEFAULT_HIDE_LLM_INVOCATION_PARAMETERS = False DEFAULT_HIDE_INPUTS = False DEFAULT_HIDE_OUTPUTS = False @@ -118,6 +120,13 @@ class TraceConfig: observability. """ + hide_llm_invocation_parameters: Optional[bool] = field( + default=None, + metadata={ + "env_var": OPENINFERENCE_HIDE_LLM_INVOCATION_PARAMETERS, + "default_value": DEFAULT_HIDE_LLM_INVOCATION_PARAMETERS, + }, + ) hide_inputs: Optional[bool] = field( default=None, metadata={ @@ -208,7 +217,9 @@ def mask( key: str, value: Union[AttributeValue, Callable[[], AttributeValue]], ) -> Optional[AttributeValue]: - if self.hide_inputs and key == SpanAttributes.INPUT_VALUE: + if self.hide_llm_invocation_parameters and key == SpanAttributes.LLM_INVOCATION_PARAMETERS: + return + elif self.hide_inputs and key == SpanAttributes.INPUT_VALUE: value = REDACTED_VALUE elif self.hide_inputs and key == SpanAttributes.INPUT_MIME_TYPE: return diff --git a/python/openinference-instrumentation/tests/test_config.py b/python/openinference-instrumentation/tests/test_config.py index 9a7e279fd..abc5a8376 100644 --- a/python/openinference-instrumentation/tests/test_config.py +++ b/python/openinference-instrumentation/tests/test_config.py @@ -1,7 +1,7 @@ import os from contextlib import suppress from random import random -from typing import Dict, Optional +from typing import Any, Dict, Optional import pytest from opentelemetry.sdk import trace as trace_sdk @@ -19,6 +19,7 @@ DEFAULT_HIDE_INPUT_MESSAGES, DEFAULT_HIDE_INPUT_TEXT, DEFAULT_HIDE_INPUTS, + DEFAULT_HIDE_LLM_INVOCATION_PARAMETERS, DEFAULT_HIDE_OUTPUT_MESSAGES, DEFAULT_HIDE_OUTPUT_TEXT, DEFAULT_HIDE_OUTPUTS, @@ -33,10 +34,12 @@ REDACTED_VALUE, OITracer, ) +from openinference.semconv.trace import SpanAttributes def test_default_settings() -> None: config = TraceConfig() + assert config.hide_llm_invocation_parameters == DEFAULT_HIDE_LLM_INVOCATION_PARAMETERS assert config.hide_inputs == DEFAULT_HIDE_INPUTS assert config.hide_outputs == DEFAULT_HIDE_OUTPUTS assert config.hide_input_messages == DEFAULT_HIDE_INPUT_MESSAGES @@ -178,6 +181,45 @@ def test_settings_from_env_vars_and_code( assert config.base64_image_max_length == new_base64_image_max_length +@pytest.mark.parametrize( + "param,param_value,attr_key,attr_value,expected_value", + [ + ( + "hide_llm_invocation_parameters", + True, + SpanAttributes.LLM_INVOCATION_PARAMETERS, + "{api_key: '123'}", + None, + ), + ( + "hide_llm_invocation_parameters", + None, + SpanAttributes.LLM_INVOCATION_PARAMETERS, + "{api_key: '123'}", + "{api_key: '123'}", + ), + ], +) +def test_trace_config( + param: str, + param_value: Optional[bool], + attr_key: str, + attr_value: Any, + expected_value: Any, + tracer_provider: TracerProvider, + in_memory_span_exporter: InMemorySpanExporter, +) -> None: + config = TraceConfig(**({} if param_value is None else {param: param_value})) + tracer = OITracer(tracer_provider.get_tracer(__name__), config=config) + tracer.start_span("test", attributes={attr_key: attr_value}).end() + span = in_memory_span_exporter.get_finished_spans()[0] + assert span.attributes is not None + if expected_value is None: + assert span.attributes.get(attr_key) is None + else: + assert span.attributes.get(attr_key) == expected_value + + def parse_bool_from_env(env_var: str) -> Optional[bool]: env_value = os.getenv(env_var) if isinstance(env_value, str) and env_value.lower() == "true":