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.
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_typeto apply different redaction logic to different kinds of spans, such asLLM,TOOL, orAGENT.
Prerequisites
-
MLflow Tracing configured for your GenAI application. This solution works whether or not you store traces in Unity Catalog.
- See MLflow Tracing - GenAI observability to setup tracing.
- If you store traces in Unity Catalog, see Setup: Create an experiment with a Unity Catalog trace location to create an experiment with a Unity Catalog trace location before running the examples in this page.
-
Install the MLflow Python library with the
databricksextra along withdatabricks-langchain>=0.19.0andlanggraph>=1.1.0:Bashpip install --upgrade "mlflow-skinny[databricks]" databricks-sdk "databricks-langchain>=0.19.0" "langgraph>=1.1.0" -
To follow the Microsoft Presidio example, also install:
Bashpip 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()