diff --git a/langfuse/_client/client.py b/langfuse/_client/client.py index 42d861fd4..1fca04f7c 100644 --- a/langfuse/_client/client.py +++ b/langfuse/_client/client.py @@ -4,6 +4,7 @@ """ import asyncio +import json import logging import os import re @@ -1949,6 +1950,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, *, @@ -2018,6 +2033,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) try: new_body = ScoreBody( @@ -3550,7 +3568,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, ) @@ -3649,9 +3667,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..a0237c656 --- /dev/null +++ b/tests/unit/test_mask_api_payloads.py @@ -0,0 +1,53 @@ +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]["body"].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") + body = client._resources.add_score_task.call_args.args[0]["body"] + assert body.comment == '{"redacted": true}'