
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 AIPydanticAI框架的自动追踪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 AIPython 生态中主打类型安全与生产级的 Agent / LLM 应用框架的追踪集成模块其 API 参考页见 mlflow.pydantic_ai.rst。该页面通过 Sphinx 的automodule指令把模块内所有公开成员autolog等的 docstring 渲染为 API 文档其核心入口是autolog()函数——调用它即可开启 Pydantic AI 工作流的自动追踪并把生成的嵌套 Trace 写入当前 MLflow Experiment。从源码结构看该模块由四个文件组成文件职责mlflow/pydantic_ai/init.py公开入口autolog()、版本检测与 1.x/2.x 分派逻辑mlflow/pydantic_ai/autolog.pyPydantic AI 1.x 的各个打桩包装函数Agent / Tool / MCP / InstrumentedModelmlflow/pydantic_ai/autolog_v2.pyPydantic AI 2.x 2.5.0的setup_autologging与全新包装函数mlflow/pydantic_ai/utils.py跨版本复用的序列化、安全属性提取、token 用量解析工具仓库的集成文档 pydantic_ai.mdx 将其定位概括为通过一次调用启用自动追踪后MLflow 会捕获 Pydantic AI 工作流执行的嵌套 Trace并记录到当前活跃的 MLflow Experiment 中。二、核心 APIautolog()函数签名与参数详解模块唯一的顶层函数是autolog()完整定义位于 mlflow/pydantic_ai/init.pyautologging_integration(FLAVOR_NAME) def autolog(log_traces: bool True, disable: bool False, silent: bool False):三个参数的语义如下参数类型默认值作用log_tracesboolTrue是否捕获 Agent 调用与模型调用的 Span。设为False时所有包装函数会直接透传原始调用不产生任何追踪数据disableboolFalse是否禁用自动打桩。为True时不启用补丁silentboolFalse是否抑制 MLflow 的 warning/info 日志输出该函数通过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-aiPydantic 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目标类打桩方法说明pydantic_ai.Agentrun、run_sync、run_stream 1.10.0 起追加run_stream_syncAgent 调用的根 Spanpydantic_ai.ToolManagerexecute_tool_call 1.63.0或handle_call旧版工具执行 Span1.63.0 起内部图直接调用execute_tool_call因此必须改打这个新入口pydantic_ai.mcp.MCPServercall_tool、list_toolsMCP 服务器工具调用与列表pydantic_ai.Toolrun仅当该方法存在工具定义调用pydantic_ai.models.instrumented.InstrumentedModelrequest、request_streamLLM 请求 Span仅在无 Instrumentation capability 时打桩同时还会包装Agent.__init__实现自动开启 instrument当log_tracesTrue且用户未显式传入instrument参数时自动补上instrumentTrue见patched_agent_initautolog.py。这样用户不必手动设置instrumentTrue也能拿到 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_syncInstrumentation.wrap_model_requestLLM Span、on_tool_validate_errorPARSER类型 SpanSpan 名为{tool_name}.validation、wrap_tool_executeTOOL类型 SpanMCPToolset.list_tools与MCPToolset.direct_call_toolTOOL类型 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_tracesTrue, disableFalse) 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(disableFalse)。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(disableFalse) 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(descriptionAdvice returned to the customer) block_card: bool Field(descriptionWhether to block their card or not) risk: int Field(descriptionRisk level of query, ge0, le10) support_agent Agent( openai:gpt-4o, deps_typeSupportDependencies, output_typeSupportOutput, 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 customers name. ), instrumentTrue, ) support_agent.tool async def customer_balance(ctx: RunContext[SupportDependencies], include_pending: bool) - str: Returns the customers current account balance. balance await ctx.deps.db.customer_balance( idctx.deps.customer_id, include_pendinginclude_pending ) return f${balance:.2f} if __name__ __main__: deps SupportDependencies(customer_id123, dbDatabaseConn()) result support_agent.run_sync(What is my balance?, depsdeps) print(result.output)即使代码里显式写了instrumentTrue对 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_typeautolog.py与 2.x 各包装函数可以归纳出 Span 分类对象Span 类型SpanType典型 Span 名Agentrun / run_sync / run_stream / run_stream_syncAGENTAgent.run、Agent.run_sync、Agent.run_stream、Agent.run_stream_sync模型实例InstrumentedModel或具体 Provider 模型LLMOpenAIChatModel.request取自type(model).__name__Tool/ToolManager/MCPServer/MCPToolsetTOOL工具名、MCPToolset.list_tools、MCPToolset.direct_call_tool工具参数校验失败2.xPARSER{tool_name}.validation5.2 Span 属性与消息格式Agent / 模型 Span 上会写入SpanAttributeKey.MESSAGE_FORMAT pydantic_ai并附带从实例上安全提取的公开属性见_set_agent_attributes与_set_model_attributes。对模型 Span 还会额外写入SpanAttributeKey.MODEL取model.model_nameSpanAttributeKey.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_toolsautolog.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_streamimport mlflow import asyncio from pydantic_ai import Agent mlflow.pydantic_ai.autolog() agent Agent(openai:gpt-4o, instrumentTrue) async def main(): async with agent.run_stream(Tell me a joke) as response: async for chunk in response.stream_text(deltaTrue): print(chunk, end, flushTrue) print() asyncio.run(main())run_stream是asynccontextmanager异步上下文管理器其包装器在with mlflow.start_span(...)内以async with消费原始流并在流被完整消费后的finally块中序列化最终输出、解析 token 用量写入 Spanpatched_async_stream_callautolog.py。6.2 同步流式run_stream_syncimport mlflow from pydantic_ai import Agent mlflow.pydantic_ai.autolog() agent Agent(openai:gpt-4o, instrumentTrue) result agent.run_stream_sync(Tell me a joke) for chunk in result.stream_text(): print(chunk, end, flushTrue) print() print(fFinal output: {result.get_output()})同步流式的实现更有技巧性autolog.py由于run_stream_sync不是上下文管理器返回的StreamedRunResult在被用户迭代期间 Span 必须保持开启因此使用start_span_no_context而非with start_span()创建 Span配合with_active_span(span)让子 SpanLLM 调用正确挂到其下返回的_StreamedRunResultSyncWrapper拦截stream_text/stream_output/stream_responses/get_output在迭代结束或get_output之后触发_finalize()先手动结束尚未结束的直接子 Span因为 pydantic-ai 的run_stream_sync内部使用会中途暂停的 async generator导致异步上下文管理器永远无法正常退出再写输出与用量、关闭根 Spanautolog.py通过 contextvar_in_sync_stream_context标记同步流上下文防止内部调用的run_stream再创建一层永远无法关闭的Agent.run_streamSpan。两条流式路径最终都会产生根Agent.run_stream/Agent.run_stream_syncSpan 每次 LLM 调用的子 Span 使用工具时的工具调用 Span。七、MCP 服务器追踪Pydantic AI 支持通过 MCPModel 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, -Rnode_modules, -Wnode_modules, --node-modules-dirauto, jsr:pydantic/mcp-run-python, stdio, ], ) agent Agent(openai:gpt-4o, mcp_servers[server], instrumentTrue) 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 列表写入 Span2.x 路径则打桩MCPToolset.list_tools与MCPToolset.direct_call_tool。由于 MCP 属于可选依赖面mcpextra打桩过程做了降级处理相关模块缺失时只打印 warning不会导致整个autolog()失败autolog_v2.py。八、禁用自动追踪按集成文档 pydantic_ai.mdx自动追踪可通过以下两种方式全局关闭mlflow.pydantic_ai.autolog(disableTrue) # 仅关闭 Pydantic AI 集成 mlflow.autolog(disableTrue) # 关闭所有 flavor 的 autolog需要强调的是2.x 路径下所有手工打桩的包装函数_patch_streaming_method与_safe_patch_async_hook安装的还会额外检查进程级全局开关autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED确保在mlflow.autolog(disableTrue)的全局抑制下不会继续泄漏 Span如工具参数等输入数据见 autolog_v2.py。九、测试与验证体系仓库为mlflow.pydantic_ai提供了完整的测试覆盖tests/pydantic_ai/测试文件覆盖内容test_pydanticai_autolog.pyautolog()的版本分派、2.x 最小版本门槛2.5.0、pydantic-ai-slim回退检测test_pydanticai_tracing.py1.x 路径下 Agent / LLM / 工具调用的一般追踪行为test_pydanticai_v2_tracing.py2.x 路径Instrumentation capability、工具校验/执行 Spantest_pydanticai_fluent_tracing.py流式调用下 Span 的完整性含patched_capability_model_request对流式全生命周期捕获的验证test_pydanticai_mcp_tracing.pyMCP 服务器工具调用与列表追踪test_utils.pyserialize_output、parse_usage等工具函数这些测试由 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 缺失或数据序列化问题时快速定位根因。/DSMLparameter /DSMLinvoke /DSMLtool_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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考