Genkit Python 中间件实战：一个类实现模型调用前的 PII 邮箱脱敏

本篇以 Genkit Python 仓库中的 py/samples/middleware 官方示例为主体，演示如何只写一个类并通过 ai.generate(..., use=[...]) 挂载自定义中间件，在请求到达模型之前把工单文本中的邮箱地址替换为 [REDACTED_EMAIL]。读完后你可以掌握 Genkit Python 中间件体系的完整机制：BaseMiddleware 基类、wrap_model / wrap_generate / wrap_tool 三类钩子、基于 Pydantic 的配置模型、@ai.middleware 注册装饰器，以及中间件的执行顺序规则，并能够独立编写和运行自己的 PII 脱敏、审计、重试等横切逻辑。

示例定位：一条工单的 PII 脱敏

该示例的核心目标用一句话概括（见 示例 README）：

One class on ai.generate(..., use=[...]). This one redacts emails before the model sees the ticket.

即：把「脱敏」实现为一个继承 BaseMiddleware 的类，在每次模型 API 调用发出前，改写 params.request 里的消息内容。典型场景是客服工单：工单里带有客户的邮箱、电话等个人信息，直接发给模型会产生数据外泄风险；而在发送前把邮箱替换成占位符，可以让模型照常完成「起草支持回复」的任务，同时避免 PII 离开本地进程。

运行方式

示例 README 给出的完整运行步骤（在 py/samples/middleware 目录下执行）：

export GEMINI_API_KEY=your-api-key
uv sync
uv run src/main.py

其 pyproject.toml 声明了三个运行时依赖与 Python 版本要求，这是运行本示例的适用前提：

[project]
name = "middleware"
version = "0.2.0"
requires-python = ">=3.10"
dependencies = [
  "genkit",
  "genkit-google-genai",
  "pydantic>=2.0.0",
]

• genkit：核心框架，提供 Genkit 应用入口与 middleware 抽象；
• genkit-google-genai：Google AI 模型插件，提供 GoogleAI 与 GoogleAI.gemini_model(...) 模型引用；
• pydantic>=2.0.0：中间件配置模型（BaseModel）的底层依赖。

完整示例源码

下面是 示例主程序 的完整代码（含版权头之外的有效部分）：

"""One class on generate(..., use=[...]). This one redacts emails first."""

import re
from collections.abc import Awaitable, Callable

from genkit_google_genai import GoogleAI
from pydantic import BaseModel

from genkit import Genkit, ModelResponse, Part, TextPart
from genkit.middleware import BaseMiddleware, GenerateMiddlewareContext, ModelHookParams

ai = Genkit(plugins=[GoogleAI()], model=GoogleAI.gemini_model('gemini-flash-latest'))

_EMAIL = re.compile(r'[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}')


class PiiRedactConfig(BaseModel):
    pass


@ai.middleware(name='pii_redact')
class PiiRedact(BaseMiddleware[PiiRedactConfig]):
    # The provider sees whatever we put on params.request, including on
    # retries. This is the last place a ticket's email can be removed.
    async def wrap_model(
        self,
        params: ModelHookParams,
        ctx: GenerateMiddlewareContext,
        next_fn: Callable[[ModelHookParams, GenerateMiddlewareContext], Awaitable[ModelResponse]],
    ) -> ModelResponse:
        new_messages = []
        for message in params.request.messages:
            new_parts = []
            for part in message.content:
                root = part.root
                if isinstance(root, TextPart):
                    redacted = _EMAIL.sub('[REDACTED_EMAIL]', root.text)
                    new_parts.append(Part(root=root.model_copy(update={'text': redacted})))
                else:
                    new_parts.append(part)
            new_messages.append(message.model_copy(update={'content': new_parts}))
        params.request = params.request.model_copy(update={'messages': new_messages})
        return await next_fn(params, ctx)


async def main() -> None:
    ticket = 'Charged twice. Email me at [e-mail] when the refund lands.'
    response = await ai.generate(
        prompt=ticket,
        system='Draft a short support reply. Do not ask the customer to repeat contact details.',
        use=[PiiRedact()],
    )
    print(response.text)


if __name__ == '__main__':
    ai.run_main(main())

下面按「初始化 → 中间件类 → 调用链 → 入口」的顺序逐段拆解。

初始化 Genkit 应用与模型

ai = Genkit(plugins=[GoogleAI()], model=GoogleAI.gemini_model('gemini-flash-latest'))

• plugins=[GoogleAI()] 注册 Google AI 插件，使模型引用（如 googleai/gemini-flash-latest）能够被解析为真实的模型调用；
• model=... 设置应用级默认模型，后续 generate 调用可以不再显式传 model；
• 本示例使用 gemini-flash-latest，这是示例仓库中多个 Python 示例（如 genkit-middleware 插件 README）通用的默认模型引用。

Genkit 应用对象还承担了「注册表」的职责：@ai.middleware 装饰器最终会把中间件类注册进该 registry，从而让 Dev UI 与 Model Runner 能看到它（见下文「注册与 Dev UI」一节）。

PiiRedact：一个只覆写 wrap_model 的中间件

配置模型与泛型基类

class PiiRedactConfig(BaseModel):
    pass

@ai.middleware(name='pii_redact')
class PiiRedact(BaseMiddleware[PiiRedactConfig]):

BaseMiddleware<a href="https://link.gitcode.com/i/1486a137154db5fdf7fc3aaaec17b304" target="_blank">TConfig] 是一个泛型类，类型参数是它的 Pydantic 配置模型（定义于 [核心中间件模块）。本示例的 PiiRedactConfig 是空的——脱敏逻辑没有可调参数，但保留配置模型可以让类结构符合标准写法：BaseMiddleware.__init_subclass__ 会从 BaseMiddleware<a href="https://link.gitcode.com/i/95738f2482a7faaf9569e1cbea27ff8a" target="_blank">PiiRedactConfig] 的 __orig_bases__ 中提取配置类型并缓存到 cls.Config（源码见 [py/packages/genkit/src/genkit/_core/_middleware.py#L214-L243），这使得：

• PiiRedact() 与 PiiRedact(config=PiiRedactConfig(...)) 两种构造方式都合法；
• 注册框架能自动派生配置的 JSON Schema，供 Dev UI 生成配置表单（若配置类有字段，表单会自动出现对应输入项；空配置则退化为空对象）。

wrap_model 钩子：在模型看到请求前改写它

PiiRedact 只覆写了 wrap_model 一个钩子。基类的默认实现直接透传：

# py/packages/genkit/src/genkit/_core/_middleware.py
async def wrap_model(
    self,
    params: ModelHookParams,
    ctx: GenerateMiddlewareContext,
    next_fn: Callable[[ModelHookParams, GenerateMiddlewareContext], Awaitable[ModelResponse]],
) -> ModelResponse:
    """Wrap each model API call."""
    return await next_fn(params, ctx)

三个参数的含义：

参数	类型	含义
params	ModelHookParams	携带本次「原始模型 API 调用」的请求体 request: ModelRequest（含 messages 等），是中间件唯一能改写发送内容的入口
ctx	GenerateMiddlewareContext	本次 generate() 调用范围内的共享运行时服务：作用域化的 ai 视图、abort_signal、custom_context、流式回调 on_chunk、遥测标签等（见 GenerateMiddlewareContext 定义）
next_fn	可调用对象	调用链的「下一环」。在 next_fn 之前执行的是「请求前处理」，之后执行的是「响应后处理」

PiiRedact.wrap_model 的处理流程：

1. 遍历消息：params.request.messages 逐条取出，message.content 逐 Part 处理；
2. 区分 Part 类型：part.root 解包后若是 TextPart，用正则 _EMAIL = re.compile(r'[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}') 做全量替换 re.sub('[REDACTED_EMAIL]', ...)；非文本 Part（图片等）原样保留；
3. 不可变更新：全程使用 Pydantic 的 model_copy(update={...}) 生成新对象而非原地修改——Part(root=root.model_copy(update={'text': redacted})) 重建 Part，message.model_copy(update={'content': new_parts}) 重建消息，最后 params.request = params.request.model_copy(update={'messages': new_messages}) 整体替换请求对象；
4. 放行：return await next_fn(params, ctx) 把改写后的请求交给下一环（最终是模型提供方）。

源码中的注释点出了为什么选择 wrap_model 这一层：

The provider sees whatever we put on params.request, including on retries. This is the last place a ticket's email can be removed.

即：模型提供方最终看到的就是 params.request 里的内容，包括任何中间件触发重试时的重发请求。wrap_model 包裹的是每一次原始模型 API 调用，因此它是「最后一道能删掉邮箱的位置」——比在 generate 入口预处理 prompt 更彻底，因为它覆盖了工具循环多轮迭代中产生的每一次模型调用。

挂载中间件：use=[PiiRedact()]

async def main() -> None:
    ticket = 'Charged twice. Email me at [e-mail] when the refund lands.'
    response = await ai.generate(
        prompt=ticket,
        system='Draft a short support reply. Do not ask the customer to repeat contact details.',
        use=[PiiRedact()],
    )
    print(response.text)

• prompt 是含真实邮箱的原始工单；system 指示模型起草简短支持回复且不要向客户索要联系方式；
• use=[PiiRedact()] 传入的是中间件的实例（可带配置，如 Retry(max_retries=3)），也可以不传配置直接实例化；
• 系统提示词特意要求「不要让客户重复提供联系方式」，与脱敏逻辑形成配套：模型看到 [REDACTED_EMAIL] 占位符后，不会在回复里编造或索要邮箱。

中间件机制深挖：钩子、注册与执行顺序

示例只用了 wrap_model，但理解整个钩子体系有助于判断「我的逻辑该放在哪一层」。genkit/middleware 包的模块文档 明确了 BaseMiddleware 提供四类可覆写钩子：

钩子	作用范围	典型用途
wrap_generate(params, ctx, next_fn)	工具循环的每一次迭代（一次模型调用 + 可选的工具解析），params 为 GenerateHookParams（含 options、iteration、message_index）	记录整轮迭代日志、按迭代计数做限流
wrap_model(params, ctx, next_fn)	每一次原始模型 API 调用，params 为 ModelHookParams	PII 脱敏（本示例）、审计、重试、降级
wrap_tool(params, ctx, next_fn)	每一次工具执行，params 为 ToolHookParams；可抛出 Interrupt(metadata) 中断工具调用	工具审批、工具结果改写
tools(ctx) -> list[Action]	非钩子方法	为本次 generate 动态暴露额外工具（如 Filesystem、Artifacts 插件中间件即靠它注册工具）

执行顺序规则同样来自该模块文档：use=[...] 中靠前的中间件先被调用，形成洋葱模型：

use = [A(), B()]

执行序列为：

A_before()
    B_before()
        model_call()
    B_after()
A_after()

这对组合多个中间件时的顺序敏感逻辑（例如「脱敏」必须在「日志」内层，日志才能记录脱敏后的请求；或「重试」必须包在「降级」外层）有直接指导意义。

注册与 Dev UI：@ai.middleware 装饰器

示例中的 @ai.middleware(name='pii_redact') 并非必需（直接 use=[PiiRedact()] 即可运行），但它有两个实际作用：

1. 注册进应用 registry：装饰器内部调用 define_middleware，将类封装为 GenerateMiddleware 并通过 self.registry.register_value('middleware', name, desc) 注册（实现见 py/packages/genkit/src/genkit/_ai/_aio.py#L344-L356）。一旦注册，中间件会出现在本地 Dev UI 中，可以在 Model Runner 里可视化地组合、调试不同中间件对响应的影响（middleware 模块文档 原话："Once registered, the middleware is visible in the Dev UI"）；
2. 名称合法性校验：注册的 name 必须是单个路径安全的 token——不能为空、不能含首尾空白，不能含 /（该形态保留给模型等 action）、空白、:、反斜杠或控制字符（校验逻辑见 _validate_middleware_key_segment），例如 myorg_logging_mw 是合法命名。

description 参数可选，不传时会取类的 docstring（GenerateMiddleware.__init__ 中 inspect.cleandoc(cls.__doc__)），因此给中间件类写 docstring 会直接反映到 Dev UI 的描述里。

配置使用方式与注意事项

BaseMiddleware 的 docstring（源码注释）还展示了两种等价用法，对含字段的配置尤其有用：

# 关键字直传（推荐，IDE 自动补全）
use=[Retry(max_retries=5)]

# 显式传 config 对象
use=[Retry(config=RetryConfig(max_retries=5))]

两种传法不可混用，混用会抛 TypeError('pass either config= or keyword config fields, not both')（见 BaseMiddleware.__init__，py/packages/genkit/src/genkit/_core/_middleware.py#L245-L253）。同时基类 docstring 明确提醒：config 不应在钩子内部被修改——中间件实例按配置构造，钩子只读配置。

与官方 genkit-middleware 插件的关系

Genkit Python 还提供了一个开箱即用的中间件插件包 genkit-middleware（插件 README），内置六种实现：Retry（指数退避重试）、Fallback（模型降级）、ToolApproval（工具审批中断）、Skills（技能库）、Filesystem（沙箱文件工具）、Artifacts（会话工件工具）。其用法与本示例完全同构——实例直接放进 use=[...]：

from genkit_middleware import Retry, Fallback

response = await ai.generate(
    model=GoogleAI.gemini_model('gemini-flash-latest'),
    prompt='Hello!',
    use=[
        Retry(max_retries=5),
        Fallback(models=['googleai/gemini-2.5-pro']),
    ],
)

从源码结构看，本示例中的 PII 脱敏与 Retry 同属 wrap_model 层逻辑，因此可以叠加使用，例如 use=[Retry(max_retries=3), PiiRedact()]：注意顺序语义——列表靠前的 Retry 是外层洋葱，重试发生时每次重发都会再次经过 PiiRedact.wrap_model，这与示例注释强调的「包括 retries」的防护目标一致。若你的场景已有现成需求（重试、降级、审批），优先用插件；只有像脱敏这类定制规则才需要按本文方式自研一个类。

小结

本示例用不到 40 行有效代码展示了 Genkit Python 自定义中间件的完整闭环：

1. 在 py/samples/middleware/src/main.py 中定义空的 PiiRedactConfig(BaseModel) 并让 PiiRedact(BaseMiddleware[PiiRedactConfig]) 覆写 wrap_model；
2. 钩子内用正则替换 TextPart 文本、用 model_copy 不可变重建请求后调用 next_fn(params, ctx) 放行；
3. @ai.middleware(name='pii_redact') 注册后进入 Dev UI，use=[PiiRedact()] 挂载到具体 generate 调用；
4. 按 export GEMINI_API_KEY=...、uv sync、uv run src/main.py 三步即可本地运行验证。

掌握这一模式后，你可以把审计日志、请求改写、重试、工具拦截等横切逻辑都收敛到 use=[...] 的洋葱链中，而不必侵入每次 generate 的业务代码。