从跟踪中编辑敏感数据

跟踪可以捕获强大的洞察力，用于调试和监控您的应用程序，但是，它们可能包含您不想与他人共享的敏感数据，例如个人身份信息（PII）。MLflow 提供了一种完全可配置的方式，可以在将敏感数据保存到后端之前，从跟踪中屏蔽这些数据。

工作原理

MLflow 允许您配置一个后处理钩子列表，这些钩子应用于跟踪中的每个 span。每个 span 处理器都是一个函数，它将一个 span 作为输入并返回一个 span。

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

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

过滤函数

过滤函数必须接受一个参数，即 Span 对象。它可以就地修改 span。它不能返回值。

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

示例 1：编辑电子邮件地址

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

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：将过滤器应用于特定 Span

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

在以下示例中，我们将从跟踪中编辑银行帐号，但根据 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?"}
        ]
    }
)

生成的跟踪将从所有消息中编辑银行帐号

示例 3：使用 Microsoft Presidio 编辑 PII

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

在此示例中，我们运行一个虚拟的自定义支持代理，它处理用户请求，例如“我想取消我的信用卡 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 插入以从跟踪中过滤敏感数据非常简单。

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

pip install presidio_analyzer presidio_anonymizer
python -m spacy download en_core_web_lg

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

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

生成的跟踪将编辑 PII

重置过滤器

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

mlflow.tracing.configure(span_processors=[])

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

mlflow.tracing.reset()