diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py index df04f9d98..248ebcce0 100644 --- a/langfuse/_client/client.py +++ b/langfuse/_client/client.py @@ -212,7 +212,7 @@ class Langfuse: release (Optional[str]): Release version/hash of your application. Used for grouping analytics by release. media_upload_thread_count (Optional[int]): Number of background threads for handling media uploads. Defaults to 1. Can also be set via LANGFUSE_MEDIA_UPLOAD_THREAD_COUNT environment variable. sample_rate (Optional[float]): Sampling rate for traces (0.0 to 1.0). Defaults to 1.0 (100% of traces are sampled). Can also be set via LANGFUSE_SAMPLE_RATE environment variable. - mask (Optional[MaskFunction]): Function to mask sensitive data synchronously when Langfuse SDK attributes are created. This applies only to data set through Langfuse SDK APIs such as `start_observation()` and `update()`. + mask (Optional[MaskFunction]): Function to mask sensitive data synchronously when Langfuse SDK attributes are created. This applies only to data set through Langfuse SDK APIs such as `start_observation()`, `update()`, `create_dataset_item()` and `create_score(comment=)`. mask_otel_spans (Optional[MaskOtelSpansFunction]): Synchronous export-stage hook for masking raw OpenTelemetry span attributes before this Langfuse client sends them to Langfuse. Use this for spans created by third-party OpenTelemetry instrumentations, or when you need to inspect final span attributes after export filtering and Langfuse media handling. It does not modify spans already exported through other OpenTelemetry exporters. The hook receives one OpenTelemetry export batch. A batch is not guaranteed to contain a complete trace, request, or Langfuse observation tree. The hook usually runs on the OpenTelemetry batch span processor worker thread; during `flush()` and shutdown it may run on the caller thread. Keep it synchronous, deterministic, and fast. @@ -1872,6 +1872,20 @@ def create_score( environment: Optional[str] = None, ) -> None: ... + def _apply_mask(self, data: Any) -> Any: + """Apply the configured mask to data sent outside a span, matching span masking.""" + if data is None or not self._mask: + return data + try: + return self._mask(data=data) + except Exception as e: + langfuse_logger.error( + "Masking error: Custom mask function threw exception when processing " + "data. Using fallback masking. Error: %s", + e, + ) + return "" + def create_score( self, *, @@ -1941,6 +1955,9 @@ def create_score( return score_id = score_id or self._create_observation_id() + comment = self._apply_mask(comment) + if comment is not None and not isinstance(comment, str): + comment = json.dumps(comment) # The API rejects timestamps without a UTC offset; naive values are UTC. if timestamp is not None and timestamp.tzinfo is None: @@ -3412,7 +3429,7 @@ def create_dataset( result = self.api.datasets.create( name=name, description=description, - metadata=metadata, + metadata=self._apply_mask(metadata), input_schema=input_schema, expected_output_schema=expected_output_schema, ) @@ -3511,9 +3528,9 @@ def create_dataset_item( result = self.api.dataset_items.create( dataset_name=dataset_name, - input=input, - expected_output=expected_output, - metadata=metadata, + input=self._apply_mask(input), + expected_output=self._apply_mask(expected_output), + metadata=self._apply_mask(metadata), source_trace_id=source_trace_id, source_observation_id=source_observation_id, status=status, diff --git a/tests/unit/test_mask_api_payloads.py b/tests/unit/test_mask_api_payloads.py new file mode 100644 index 000000000..6ab24ee79 --- /dev/null +++ b/tests/unit/test_mask_api_payloads.py @@ -0,0 +1,51 @@ +from unittest.mock import Mock + +from langfuse import Langfuse + + +def _client(): + client = Langfuse( + public_key="pk", + secret_key="sk", + host="https://mock-host.com", + tracing_enabled=False, + mask=lambda data: "masked", + ) + client.api = Mock() + client._resources.add_score_task = Mock() + return client + + +def test_create_dataset_item_masks_payload(): + client = _client() + client.create_dataset_item( + dataset_name="ds", input="secret", expected_output="secret", metadata="secret" + ) + kwargs = client.api.dataset_items.create.call_args.kwargs + assert (kwargs["input"], kwargs["expected_output"], kwargs["metadata"]) == ( + "masked", + "masked", + "masked", + ) + + +def test_create_dataset_masks_metadata(): + client = _client() + client.create_dataset(name="ds", metadata="secret") + assert client.api.datasets.create.call_args.kwargs["metadata"] == "masked" + + +def test_create_score_masks_comment(): + client = _client() + client._tracing_enabled = True # create_score is a no-op otherwise + client.create_score(name="s", value=1, trace_id="t" * 32, comment="secret") + assert client._resources.add_score_task.call_args.args[0]["comment"] == "masked" + + +def test_create_score_serializes_non_string_masked_comment(): + client = _client() + client._tracing_enabled = True + client._mask = lambda data: {"redacted": True} + client.create_score(name="s", value=1, trace_id="t" * 32, comment="secret") + score = client._resources.add_score_task.call_args.args[0] + assert score["comment"] == '{"redacted": true}'