从 Traces 中 redaction 敏感数据

Traces 捕获了用于调试和监控应用程序的强大洞察力，但是，它们可能包含您不想与他人共享的敏感数据，例如 个人身份信息 (PII)。MLflow 提供了一种完全可配置的方式，可以在将 Traces 保存到后端之前对其进行掩码处理，以隐藏敏感数据。

工作原理

MLflow 允许您配置一组后处理钩子，这些钩子将应用于 trace 中的每个 span。每个 span 处理器都是一个函数，它以 span 作为输入并就地更新它。

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

由于过滤器在将 span 发送到后端之前在 客户端 应用，因此敏感数据永远不会离开您的应用程序。

过滤函数

过滤函数必须接受单个参数，该参数是 Span 对象。它可以就地修改 span。它不得返回值。

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

示例 1：Redaction 电子邮件地址

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

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

生成的 trace 将在输入中 redaction 电子邮件地址

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

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

在下面的示例中，我们将 redaction trace 中的银行卡号，但根据 span 类型使用不同的过滤逻辑。

首先，我们定义一个简单的工具调用代理。

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 字段，我们可以将不同的过滤逻辑应用于不同的 span 类型。

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
    ]

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

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

生成的 trace 将在所有消息中 redaction 银行卡号

示例 3：使用 Microsoft Presidio Redaction PII

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

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

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 插入到 Traces 中以过滤敏感数据非常简单。

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

pip install presidio_analyzer presidio_anonymizer
python -m spacy download en_core_web_lg

然后，定义一个运行 Presidio 的分析器和匿名化器来处理 span 输入的过滤函数。

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})

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

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

生成的 trace 将 redaction PII

重置过滤器

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

mlflow.tracing.configure(span_processors=[])

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

mlflow.tracing.reset()