从追踪中遮蔽敏感数据

追踪（Traces）能为调试和监控应用程序提供强大的洞察力，但它们可能包含你不希望与他人共享的敏感数据，例如个人身份信息 (PII)。MLflow 提供了一种完全可配置的方法，在将追踪保存到后端之前对其敏感数据进行遮蔽。

工作原理

MLflow 允许你配置一系列后处理钩子（hooks），这些钩子会应用于追踪中的每个跨度（span）。每个跨度处理器都是一个以跨度作为输入并对其进行原地（in-place）更新的函数。

1. 定义一个自定义过滤函数，并调用 mlflow.tracing.configure 来注册它。
2. 每当创建一个新的跨度时，注册的过滤器就会按顺序应用于该跨度。
3. MLflow 将过滤后的跨度发送到后端。

由于过滤器是在客户端发送跨度到后端之前应用的，敏感数据永远不会离开你的应用程序。

过滤函数

过滤函数必须接受一个参数，即 Span 对象。它可以原地修改跨度。该函数不得返回值。

python

def filter_function(span: Span) -> None: ...

示例 1：遮蔽电子邮件地址

在此示例中，我们将使用简单的正则表达式来遮蔽用户输入中的电子邮件地址。

python

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


# 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 filtering 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")

生成的追踪中输入部分的电子邮件地址将被遮蔽。

示例 2：将过滤器应用于特定跨度

在 mlflow.tracing.configure 处注册的过滤函数会应用于所有跨度。如果你的追踪包含许多嵌套跨度，你可能只想将过滤器应用于某些特定跨度。此外，不同跨度类型的输入/输出格式通常不同，因此你可能需要应用不同的过滤逻辑。

在下面的示例中，我们将从追踪中遮蔽银行账号，但会根据跨度类型使用不同的过滤逻辑。

首先，让我们定义一个简单的工具调用代理（tool calling agent）。

python

import mlflow
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

# Enabling tracing for LangGraph
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 = ChatOpenAI(model="o4-mini")
tools = [get_bank_account_number]
graph = create_react_agent(llm, tools)

然后，定义一个过滤函数。通过检查 span_type 字段，我们可以对不同的跨度类型应用不同的过滤逻辑。

python

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

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


def filter_bank_account_number(span: Span) -> None:
    # Redact the output of the tool call span.
    if span.span_type == SpanType.TOOL:
        span.set_outputs("[REDACTED]")
        return

    # Redact the back account number from other spans.
    if isinstance(span.inputs, dict) and (messages := span.inputs.get("messages")):
        span.set_inputs({"messages": redact_messages(messages)})
    if isinstance(span.outputs, dict) and (messages := span.outputs.get("messages")):
        span.set_outputs({"messages": redact_messages(messages)})


def redact_messages(messages: list[dict]):
    if isinstance(messages, dict):
        messages = messages.get("messages")

    return [
        {**msg, "content": ACCOUNT_NUMBER_PATTERN.sub("[REDACTED]", msg["content"])}
        for msg in messages
    ]

现在，让我们注册过滤函数并运行应用程序。

python

# 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?"}]
})

生成的追踪中，所有消息里的银行账号都将被遮蔽。

示例 3：使用 Microsoft Presidio 遮蔽 PII

要超越简单的基于正则表达式的过滤，你可以使用更复杂的 PII 匿名化工具，例如 Microsoft Presidio。

在此示例中，我们运行一个模拟的客户支持代理，它接收用户请求，例如“我想取消我的信用卡 4095-2609-9393-4932”。请求中包含多种形式的敏感数据，如信用卡号、用户名、电子邮件地址，仅使用正则表达式覆盖所有这些内容并非易事。

python

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


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

使用 MLflow，接入 Presidio 来过滤追踪中的敏感数据非常简单。

首先，安装 Presidio 并下载分类器。

bash

pip install presidio_analyzer presidio_anonymizer
python -m spacy download en_core_web_lg

然后，定义一个过滤函数，在跨度输入上运行 Presidio 的分析器和匿名化器。

python

from presidio_anonymizer import AnonymizerEngine
from presidio_anonymizer.entities import RecognizerResult, OperatorConfig

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


# Define a filter function.
def filter_pii(span: Span) -> None:
    """Filter PII from the span input using Microsoft Presidio."""
    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})

最后，注册过滤函数并运行应用程序。

python

# 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."
)

生成的追踪中的 PII 将被遮蔽。

重置过滤器

要重置过滤器，请调用 mlflow.tracing.configure 并传入空的跨度处理器列表。

python

mlflow.tracing.configure(span_processors=[])

或者，你可以调用 mlflow.tracing.reset 来重置整个追踪配置。

python

mlflow.tracing.reset()