MLflow 与 Pydantic AI 集成指南：用 mlflow.pydantic_ai 自动追踪 Agent、工具与 LLM 调用

MLflow 与 Pydantic AI 集成指南：用 mlflow.pydantic_ai 自动追踪 Agent、工具与 LLM 调用

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow

导读

本文以 mlflow.pydantic_ai 这一 API 参考文档为骨架，系统讲解 MLflow 对 Pydantic AI（PydanticAI）框架的自动追踪（autolog tracing）集成：从mlflow.pydantic_ai.autolog()的函数签名与参数语义，到 Agent 调用、流式输出、工具执行、MCP 服务器与 LLM 请求的 Span 捕获原理，再到 1.x / 2.x 版本分派、token 用量解析与底层打桩（patching）实现。读完本文，你将掌握如何在 MLflow 中一键开启 Pydantic AI 可观测性，并理解其内部工作机制，从而在自己的 Agent 应用中正确配置与排查追踪数据。

一、模块定位：mlflow.pydantic_ai是什么

mlflow.pydantic_ai是 MLflow 仓库中专门面向 Pydantic AI（Python 生态中主打类型安全与生产级的 Agent / LLM 应用框架）的追踪集成模块，其 API 参考页见 mlflow.pydantic_ai.rst。该页面通过 Sphinx 的automodule指令把模块内所有公开成员（autolog等）的 docstring 渲染为 API 文档，其核心入口是autolog()函数——调用它即可开启 Pydantic AI 工作流的自动追踪，并把生成的嵌套 Trace 写入当前 MLflow Experiment。

从源码结构看，该模块由四个文件组成：

仓库的集成文档 pydantic_ai.mdx 将其定位概括为：通过一次调用启用自动追踪后，MLflow 会捕获 Pydantic AI 工作流执行的嵌套 Trace，并记录到当前活跃的 MLflow Experiment 中。

二、核心 API：autolog()函数签名与参数详解

模块唯一的顶层函数是autolog()，完整定义位于 mlflow/pydantic_ai/init.py：

@autologging_integration(FLAVOR_NAME) def autolog(log_traces: bool = True, disable: bool = False, silent: bool = False):

三个参数的语义如下：

该函数通过@autologging_integration装饰器注册为正式的 autologging 集成（flavor 名称为常量FLAVOR_NAME = "pydantic_ai"），并会通过_record_event上报AutologgingEvent遥测事件（包含 flavor、log_traces、disable三个字段），见init.py。

三、版本要求与分派逻辑

autolog()内部首先通过_get_pydantic_ai_version()检测安装的 Pydantic AI 版本（优先读取完整发行版pydantic-ai，其次回退到pydantic-ai-slim），然后按大版本走两条完全不同的打桩路径：

• Pydantic AI 2.x：要求版本>= 2.5.0（常量_PYDANTIC_AI_V2_MIN_VERSION = Version("2.5.0")），此时转调mlflow.pydantic_ai.autolog_v2.setup_autologging()；
• Pydantic AI 2.x 但版本 < 2.5.0：不会启用 autologging，而是打印警告日志，提示升级pydantic-ai；
• Pydantic AI 1.x：走 autolog.py 中的传统打桩路径。

该分派逻辑在测试 tests/pydantic_ai/test_pydanticai_autolog.py 中得到验证：2.5.0与2.15.0均断言setup_autologging被调用一次；而2.0.0、2.4.0则断言既不调用setup_autologging，也不探测旧版打桩，并输出版本升级警告。

此外，仓库的版本矩阵 ml-package-versions.yml 中专门配置了pydantic_ai一节，其 CI 验证命令为pytest tests/pydantic_ai，并注明了不同版本对依赖（如mcp、opentelemetry._events）的兼容性差异。

3.1 1.x 路径的补丁面（patch surface）

在 1.x 路径下，autolog()构建了一张类-方法映射表并逐个打桩，见init.py：

同时还会包装Agent.__init__，实现"自动开启 instrument"：当log_traces=True且用户未显式传入instrument参数时，自动补上instrument=True，见patched_agent_init（autolog.py）。这样用户不必手动设置instrument=True也能拿到 LLM 层级的 Span。

3.2 pydantic-ai >= 1.95 的 Instrumentation capability

从 pydantic-ai 1.95 开始，模型调用不再经过InstrumentedModel，而是统一汇入Instrumentationcapability 的wrap_model_request。autolog()通过_has_instrumentation_capability()探测该模块是否可导入：

• 可导入：打桩Instrumentation.wrap_model_request，并从request_context上取具体模型实例（如OpenAIChatModel）来命名 Span 与记录模型属性，不再打InstrumentedModel，确保每次模型调用只产生一个 LLM Span；
• 不可导入：回退到InstrumentedModel.request / request_stream。

这一"两条路径二选一"的设计在源码注释中有明确说明（init.py），避免出现两个重叠的 LLM Span。

3.3 2.x 路径的补丁面

2.x 的setup_autologging()（autolog_v2.py）打桩面更细：

• Agent.__init__、Agent.run、Agent.run_sync、Agent.run_stream、Agent.run_stream_sync；
• Instrumentation.wrap_model_request（LLM Span）、on_tool_validate_error（PARSER类型 Span，Span 名为{tool_name}.validation）、wrap_tool_execute（TOOL类型 Span）；
• MCPToolset.list_tools与MCPToolset.direct_call_tool（TOOL类型 Span）。

其中工具校验与执行两个 hook 使用自定义的_safe_patch_async_hook包装：因为 Pydantic AI 会使用ModelRetry这类异常做控制流，通用的safe_patch会把共享的 autologging session 标记为失败从而抑制重试追踪，自定义包装则保证"原始调用已成功执行则不二次执行"，避免工具或传输操作被重复调用（autolog_v2.py）。

四、快速开始：一键开启自动追踪

按照官方集成文档 pydantic_ai.mdx 与仓库示例 examples/pydanticai/tracing.py，最快只需要两步：

import mlflow # 1. 开启自动追踪（等价于 mlflow.pydantic_ai.autolog(log_traces=True, disable=False)） mlflow.pydantic_ai.autolog() # 2.（可选）设置 Tracking URI 与 Experiment，便于集中管理 Trace mlflow.set_tracking_uri("http://localhost:5000") mlflow.set_experiment("PydanticAI")

仓库示例 examples/pydanticai/tracing.py 展示的正是这一标准姿势：先set_tracking_uri、再set_experiment("Pydantic AI Example")、最后mlflow.pydantic_ai.autolog(disable=False)。

4.1 一个带依赖注入与类型化输出的完整示例

下面是一个"银行客服"风格的多工具 Agent（源自 examples/pydanticai/tracing.py）：它使用deps_type注入依赖、output_type声明结构化输出，并注册了一个查询余额的工具函数。开启 autolog 后，每次run_sync都会在 MLflow 中形成完整的嵌套 Trace。

import mlflow import mlflow.pydantic_ai from dataclasses import dataclass from pydantic import BaseModel, Field from pydantic_ai import Agent, RunContext mlflow.set_tracking_uri("http://localhost:5000") mlflow.set_experiment("Pydantic AI Example") mlflow.pydantic_ai.autolog(disable=False) class DatabaseConn: """示例用的假数据库；真实场景可换成 PostgreSQL 等外部存储。""" @classmethod async def customer_name(cls, *, id: int) -> str | None: if id == 123: return "John" @classmethod async def customer_balance(cls, *, id: int, include_pending: bool) -> float: if id == 123 and include_pending: return 123.45 raise ValueError("Customer not found") @dataclass class SupportDependencies: customer_id: int db: DatabaseConn class SupportOutput(BaseModel): support_advice: str = Field(description="Advice returned to the customer") block_card: bool = Field(description="Whether to block their card or not") risk: int = Field(description="Risk level of query", ge=0, le=10) support_agent = Agent( "openai:gpt-4o", deps_type=SupportDependencies, output_type=SupportOutput, system_prompt=( "You are a support agent in our bank, give the " "customer support and judge the risk level of their query. " "Reply using the customer's name." ), instrument=True, ) @support_agent.tool async def customer_balance(ctx: RunContext[SupportDependencies], include_pending: bool) -> str: """Returns the customer's current account balance.""" balance = await ctx.deps.db.customer_balance( id=ctx.deps.customer_id, include_pending=include_pending ) return f"${balance:.2f}" if __name__ == "__main__": deps = SupportDependencies(customer_id=123, db=DatabaseConn()) result = support_agent.run_sync("What is my balance?", deps=deps) print(result.output)

即使代码里显式写了instrument=True，对 MLflow 自动追踪也不是必需的——patched_agent_init会在未指定时自动补上（详见 autolog.py）。

五、自动追踪会捕获什么

根据集成文档 pydantic_ai.mdx，开启mlflow.pydantic_ai.autolog()后，MLflow Trace 自动捕获以下信息：

• Agent 调用：prompt、kwargs 与输出响应；
• 流式操作：run_stream（异步）与run_stream_sync（同步）的完整执行；
• LLM 请求：模型名、prompt、参数与响应；
• 工具运行：工具名、参数与用量指标；
• MCP 服务器调用与列表：用于工具调用追踪；
• Span 元数据：延迟（latency）、错误与 run-ID 关联。

5.1 Span 类型与命名

从源码_get_span_type（autolog.py）与 2.x 各包装函数可以归纳出 Span 分类：

5.2 Span 属性与消息格式

Agent / 模型 Span 上会写入SpanAttributeKey.MESSAGE_FORMAT = "pydantic_ai"，并附带从实例上安全提取的公开属性（见_set_agent_attributes与_set_model_attributes）。对模型 Span 还会额外写入：

• SpanAttributeKey.MODEL：取model.model_name；
• SpanAttributeKey.MODEL_PROVIDER：优先取model.system（如"openai"）；若为None，则从"provider:model"格式的model_name前缀回退提取（如"anthropic:claude-3-5-haiku"→"anthropic"），见 autolog.py 与 2.x 中的同款逻辑。

工具列表会以[{"type": "function", "function": <tool 的 model_dumps 结果>}]的 OpenAI function-calling 风格序列化进 Span 属性（_parse_tools，autolog.py）。

5.3 Token 用量追踪

utils.py中的parse_usage()（utils.py）负责把 Pydantic AI 的 usage 对象转换成 MLflow 标准的 token 用量字典：

{ "input_tokens": <input_tokens 或兼容别名 request_tokens>, "output_tokens": <output_tokens 或兼容别名 response_tokens>, "total_tokens": <total_tokens，缺省时 = input + output>, # 存在时才写入： "cache_read_input_tokens": <cache_read_tokens>, # TokenUsageKey.CACHE_READ_INPUT_TOKENS "cache_creation_input_tokens": <cache_write_tokens>, # TokenUsageKey.CACHE_CREATION_INPUT_TOKENS }

它兼容三种形态：(result, usage)二元组、RunResult.usage属性、以及StreamedRunResult.usage()方法调用。每个 LLM Span 都会通过SpanAttributeKey.CHAT_USAGE写入该字典，从而支撑 UI 内置面板中的成本与时间趋势统计。

5.4 输出序列化的安全性

由于 Pydantic AI 的模型/运行结果中可能包含 httpx 客户端等不可序列化、且会干扰异步清理的对象，集成层采用"白名单式"序列化策略（extract_safe_attributes/is_safe_for_serialization，见 utils.py）：只保留str/int/float/bool、全安全的 dict/list/tuple、dataclass 实例与类型对象，跳过以下划线开头的私有属性与方法/函数；serialize_output()则优先把result.new_messages()序列化后以_new_messages_serialized键挂到输出上。

六、流式执行追踪

MLflow 同时支持 Pydantic AI 的异步与同步流式 API。集成文档明确指出：run_stream_sync需要Pydantic AI 1.10.0 或更高版本。

6.1 异步流式run_stream

import mlflow import asyncio from pydantic_ai import Agent mlflow.pydantic_ai.autolog() agent = Agent("openai:gpt-4o", instrument=True) async def main(): async with agent.run_stream("Tell me a joke") as response: async for chunk in response.stream_text(delta=True): print(chunk, end="", flush=True) print() asyncio.run(main())

run_stream是@asynccontextmanager异步上下文管理器，其包装器在with mlflow.start_span(...)内以async with消费原始流，并在流被完整消费后的finally块中序列化最终输出、解析 token 用量写入 Span（patched_async_stream_call，autolog.py）。

6.2 同步流式run_stream_sync

import mlflow from pydantic_ai import Agent mlflow.pydantic_ai.autolog() agent = Agent("openai:gpt-4o", instrument=True) result = agent.run_stream_sync("Tell me a joke") for chunk in result.stream_text(): print(chunk, end="", flush=True) print() print(f"Final output: {result.get_output()}")

同步流式的实现更有技巧性（autolog.py）：

• 由于run_stream_sync不是上下文管理器，返回的StreamedRunResult在被用户迭代期间 Span 必须保持开启，因此使用start_span_no_context（而非with start_span()）创建 Span，配合with_active_span(span)让子 Span（LLM 调用）正确挂到其下；
• 返回的_StreamedRunResultSyncWrapper拦截stream_text/stream_output/stream_responses/get_output，在迭代结束或get_output之后触发_finalize()：先手动结束尚未结束的直接子 Span（因为 pydantic-ai 的run_stream_sync内部使用会中途暂停的 async generator，导致异步上下文管理器永远无法正常退出），再写输出与用量、关闭根 Span（autolog.py）；
• 通过 contextvar_in_sync_stream_context标记同步流上下文，防止内部调用的run_stream再创建一层永远无法关闭的Agent.run_streamSpan。

两条流式路径最终都会产生：根Agent.run_stream/Agent.run_stream_syncSpan + 每次 LLM 调用的子 Span + 使用工具时的工具调用 Span。

七、MCP 服务器追踪

Pydantic AI 支持通过 MCP（Model Context Protocol）服务器扩展工具。MLflow 会自动捕获call_tool/list_tools等 MCP 交互并记录为独立 Span。集成文档给出的 MCP 示例（pydantic_ai.mdx）如下：

import mlflow import asyncio from pydantic_ai import Agent from pydantic_ai.mcp import MCPServerStdio mlflow.set_tracking_uri("http://localhost:5000") mlflow.set_experiment("MCP Server") mlflow.pydantic_ai.autolog() server = MCPServerStdio( "deno", args=[ "run", "-N", "-R=node_modules", "-W=node_modules", "--node-modules-dir=auto", "jsr:@pydantic/mcp-run-python", "stdio", ], ) agent = Agent("openai:gpt-4o", mcp_servers=[server], instrument=True) async def main(): async with agent.run_mcp_servers(): result = await agent.run("How many days between 2000-01-01 and 2025-03-18?") print(result.output) asyncio.run(main())

底层实现上，1.x 路径打桩pydantic_ai.mcp.MCPServer.call_tool / list_tools，并会把MCPServer实例上的安全属性（含序列化后的 tools 列表）写入 Span；2.x 路径则打桩MCPToolset.list_tools与MCPToolset.direct_call_tool。由于 MCP 属于可选依赖面（mcpextra），打桩过程做了降级处理：相关模块缺失时只打印 warning，不会导致整个autolog()失败（autolog_v2.py）。

八、禁用自动追踪

按集成文档 pydantic_ai.mdx，自动追踪可通过以下两种方式全局关闭：

mlflow.pydantic_ai.autolog(disable=True) # 仅关闭 Pydantic AI 集成 mlflow.autolog(disable=True) # 关闭所有 flavor 的 autolog

需要强调的是，2.x 路径下所有手工打桩的包装函数（_patch_streaming_method与_safe_patch_async_hook安装的）还会额外检查进程级全局开关autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED，确保在mlflow.autolog(disable=True)的全局抑制下不会继续泄漏 Span（如工具参数等输入数据），见 autolog_v2.py。

九、测试与验证体系

仓库为mlflow.pydantic_ai提供了完整的测试覆盖（tests/pydantic_ai/）：

这些测试由 ml-package-versions.yml 中的pydantic_ai矩阵驱动（run: pytest tests/pydantic_ai），并对 pydantic-ai 2.x 依赖的fastmcp替换等上游变化做了兼容性说明。

十、小结

mlflow.pydantic_ai是 MLflow 对 Pydantic AI 框架开箱即用的可观测性集成：只需一行mlflow.pydantic_ai.autolog()，即可获得覆盖 Agent 调用、同步/异步流式、LLM 请求、工具执行与 MCP 服务器的完整嵌套 Trace，并自动采集 token 用量与成本数据。其内部通过"版本分派 + 条件打桩"同时兼容 Pydantic AI 1.x 与 2.x（>= 2.5.0）两条演进路线，且在 1.95 版本引入 Instrumentation capability 后平滑切换到新的模型请求挂钩点。理解这些实现细节，将帮助你在升级 Pydantic AI 版本、排查 Span 缺失或数据序列化问题时快速定位根因。

</｜DSML｜parameter> </｜DSML｜invoke> </｜DSML｜tool_calls>

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow