Redact PII from traces before export

MLflow Tracing enables you to mask sensitive data in span inputs and outputs client-side, before MLflow exports a trace to the backend, so the raw PII never leaves your application. This is different from Redact PII from OpenTelemetry traces in Unity Catalog, which redacts OpenTelemetry (OTel) spans that Unity Catalog tables already store.

  • Use client-side redaction when you want to guarantee your application never transmits or persists sensitive values.
  • Use the OTel pipeline approach when you already store traces in Unity Catalog and you must redact them without changing your application code.

Note

MLflow redaction only affects what's recorded in the trace. The model or tool call itself still receives and returns the original, unredacted content. Use service policies to block PII from a Unity Catalog-registered AI service or to enforce a specific policy centrally across your organization.

How it works

Span processors implement redaction. These processors receive a span, mutate it in place, and return nothing. Register one or more span processors with mlflow.tracing.configure, and MLflow applies them to every span before exporting it.

from mlflow.entities.span import Span

def filter_function(span: Span) -> None:
    # Read span.inputs / span.outputs, redact, then write back.
    span.set_inputs(...)
    span.set_outputs(...)

mlflow.tracing.configure(span_processors=[filter_function])

Keep the following behavior in mind when you write span processors:

  • Filtering happens client-side, so your application never transmits or stores unredacted data.
  • Multiple span processors run in the order you register them. Each receives the span after the previous processor mutates it.
  • Span processors apply to every span in a trace, including nested spans created by framework integrations such as LangChain and LangGraph.
  • Use span.span_type to apply different redaction logic to different kinds of spans, such as LLM, TOOL, or AGENT.

Prerequisites

  • MLflow Tracing configured for your GenAI application. This solution works whether or not you store traces in Unity Catalog.

  • Install the MLflow Python library with the databricks extra along with databricks-langchain>=0.19.0 and langgraph>=1.1.0:

    pip install --upgrade "mlflow-skinny[databricks]" databricks-sdk "databricks-langchain>=0.19.0" "langgraph>=1.1.0"

  • To follow the Microsoft Presidio example, also install:

    pip install presidio_analyzer presidio_anonymizer
python -m spacy download en_core_web_lg

Redact PII with a regex

The following example matches email addresses in a span's inputs with a regular expression and replaces them with [REDACTED].

import re
import mlflow
from mlflow.entities.span import Span

# mlflow.set_experiment(experiment_id=experiment_id)


# Your application code (simplified).
@mlflow.trace
def predict(text: str):
    return "Answer"


# Regex pattern to match e-mail addresses.
EMAIL_PATTERN = r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}"


# Define a filter function that takes a span as input and mutates it in place.
def redact_email(span: Span) -> None:
    raw_input = span.inputs.get("text")
    redacted_input = re.sub(EMAIL_PATTERN, "[REDACTED]", raw_input)
    span.set_inputs({"text": redacted_input})


# Register the filter function.
mlflow.tracing.configure(span_processors=[redact_email])

# Run the application.
predict("My e-mail address is test@example.com")

Apply a filter to specific span types

Use span.span_type to redact only certain kinds of spans, such as tool calls. The following example redacts a bank account number from a LangGraph agent's tool-call spans. It checks every message and tool-call field in the payload, so it scrubs account numbers wherever they appear in the conversation.

import mlflow
from langchain_core.tools import tool
from databricks_langchain import ChatDatabricks
from langchain.agents import create_agent

# autolog() registers a LangChain callback so MLflow automatically captures spans
# for every LLM call, tool invocation, and agent step. Without this call,
# LangGraph operations would not appear in the trace.
mlflow.langchain.autolog()


@tool
def get_bank_account_number(user_name: str):
    """Return the bank account number for the given user name."""
    return "1234567890"


llm = ChatDatabricks(
    model="databricks-claude-opus-4-6",
    use_ai_gateway=True,
)
tools = [get_bank_account_number]
graph = create_agent(llm, tools)

Account numbers can appear in message content, tool call arguments, and model responses. The processor handles each shape with a dedicated helper function:

import re
from mlflow.entities.span import Span, SpanType

ACCOUNT_NUMBER_PATTERN = re.compile(r"\d{10}")


def _redact_str(s):
    """Replace every pattern match in s with [REDACTED]. Non-strings pass through unchanged."""
    return ACCOUNT_NUMBER_PATTERN.sub("[REDACTED]", s) if isinstance(s, str) else s


def _redact_content(content):
    """Redact pattern matches from a message content field.

    content is either a plain string or a multimodal list of content parts
    (e.g., [{"type": "text", "text": "..."}, {"type": "image_url", ...}]).
    Only text fields are redacted; non-text parts pass through unchanged."""
    if isinstance(content, str):
        return _redact_str(content)
    if isinstance(content, list):
        out = []
        for part in content:
            if isinstance(part, dict) and "text" in part:
                out.append({**part, "text": _redact_str(part.get("text"))})
            elif isinstance(part, str):
                out.append(_redact_str(part))
            else:
                out.append(part)
        return out
    return content


def _redact_tool_calls(tool_calls):
    """Redact pattern matches from an OpenAI-format tool call list.

    In the OpenAI schema, each tool call's `function.arguments` value is a
    JSON-encoded string, not a parsed dict. This function redacts it as raw
    text to avoid parsing and re-serializing the JSON."""
    if not isinstance(tool_calls, list):
        return tool_calls
    out = []
    for tc in tool_calls:
        if not isinstance(tc, dict):
            out.append(tc)
            continue
        fn = tc.get("function")
        if isinstance(fn, dict) and isinstance(fn.get("arguments"), str):
            # arguments is a JSON string — redact the raw text.
            tc = {**tc, "function": {**fn, "arguments": _redact_str(fn["arguments"])}}
        out.append(tc)
    return out


def _redact_message(msg):
    """Redact pattern matches from a single chat message dict.

    A message is a dict like {"role": "user", "content": "..."}.
    Assistant messages may also include a `tool_calls` key, which is
    also redacted."""
    if not isinstance(msg, dict):
        return msg
    new = {**msg}
    if "content" in new:
        new["content"] = _redact_content(new.get("content"))
    if "tool_calls" in new:
        new["tool_calls"] = _redact_tool_calls(new.get("tool_calls"))
    return new


def redact_messages(messages):
    """Works for both input and output message lists."""
    if isinstance(messages, dict):  # {"messages": [...]}
        messages = messages.get("messages")
    if not isinstance(messages, list):
        return messages
    return [_redact_message(m) for m in messages]


def _redact_choices(choices):
    """OpenAI Chat Completions: choices[i].message / .delta / .text"""
    if not isinstance(choices, list):
        return choices
    out = []
    for ch in choices:
        if not isinstance(ch, dict):
            out.append(ch)
            continue
        new = dict(ch)
        for key in ("message", "delta"):
            if isinstance(new.get(key), dict):
                new[key] = _redact_message(new[key])
        if isinstance(new.get("text"), str):  # legacy completions
            new["text"] = _redact_str(new["text"])
        out.append(new)
    return out


def _redact_payload(payload):
    """Apply to either span.inputs or span.outputs uniformly."""
    if isinstance(payload, str):
        return _redact_str(payload)
    if isinstance(payload, list):  # bare list of messages
        return redact_messages(payload)
    if not isinstance(payload, dict):
        return payload
    new = dict(payload)
    if isinstance(new.get("messages"), list):
        new["messages"] = redact_messages(new["messages"])
    if "choices" in new:  # OpenAI chat / completions
        new["choices"] = _redact_choices(new["choices"])
    if isinstance(new.get("content"), (str, list)):  # Anthropic Messages
        new["content"] = _redact_content(new["content"])
    if isinstance(new.get("output_text"), str):  # OpenAI Responses API
        new["output_text"] = _redact_str(new["output_text"])
    if isinstance(new.get("prompt"), str):  # legacy completion input
        new["prompt"] = _redact_str(new["prompt"])
    return new


def filter_bank_account_number(span: Span) -> None:
    if span.span_type == SpanType.TOOL:
        span.set_outputs("[REDACTED]")
        return

    if span.inputs is not None:
        new_in = _redact_payload(span.inputs)
        if new_in != span.inputs:
            span.set_inputs(new_in)

    if span.outputs is not None:
        new_out = _redact_payload(span.outputs)
        if new_out != span.outputs:
            span.set_outputs(new_out)

Register the processor and invoke the agent:

# Register the filter function.
mlflow.tracing.configure(span_processors=[filter_bank_account_number])

# Run the application.
result = graph.invoke(
    {
        "messages": [
            {"role": "user", "content": "What is the bank account number for John Doe?"}
        ]
    }
)

Redact PII with Microsoft Presidio

To go beyond regex-based filtering, use a dedicated PII detection and anonymization library such as Microsoft Presidio. An AnalyzerEngine detects entities such as names, credit card numbers, and email addresses, and an AnonymizerEngine rewrites them.

import mlflow
from mlflow.entities.span import Span, SpanType


# Dummy application code for a customer support agent.
@mlflow.trace(span_type=SpanType.AGENT)
def customer_support_agent(request: str):
    return "Yes"

Initialize Presidio's analyzer and anonymizer, then define the span processor:

from presidio_analyzer import AnalyzerEngine
from presidio_anonymizer import AnonymizerEngine

# Initialize the analyzer and anonymizer.
analyzer = AnalyzerEngine()
anonymizer = AnonymizerEngine()


# Define a filter function.
def filter_pii(span: Span) -> None:
    text = span.inputs.get("request")

    results = analyzer.analyze(
        text=text,
        entities=["PERSON", "CREDIT_CARD", "EMAIL_ADDRESS", "LOCATION", "DATE_TIME"],
        language="en",
    )
    anonymized_text = anonymizer.anonymize(text=text, analyzer_results=results)

    span.set_inputs({"request": anonymized_text.text})

Register the processor and run the agent:

# Register the filter function.
mlflow.tracing.configure(span_processors=[filter_pii])

# Run the application.
customer_support_agent(
    "Please cancel my credit card effective September 19th. My name is John Doe and my credit "
    "card number is 4095-2609-9393-4932. My email is john.doe@example.com and I live in Amsterdam."
)

Reset the span processors

To stop redacting spans, clear all registered span processors.

mlflow.tracing.configure(span_processors=[])

Or use reset to clear the entire tracing configuration.

mlflow.tracing.reset()

Additional resources

  • Last updated on