
Apache Airflow Common AI Provider用 PydanticAIHook 构建 pydantic-ai LLM 任务的完整指南【免费下载链接】airflowApache Airflow - A platform to programmatically author, schedule, and monitor workflows项目地址: https://gitcode.com/GitHub_Trending/ai/airflow本篇技术指南围绕 Apache Airflow 的apache-airflow-providers-common-ai包中的PydanticAIHook展开讲解如何通过一次 Airflow Connection 配置即可对接任意 pydantic-ai 支持的 LLM 提供商如何在task中调用 LLM、如何用output_type获取结构化输出以及如何用 YAML/JSON AgentSpec 文件将提示词工程与 DAG 逻辑解耦。读完后你可以独立完成 Connection 配置、模型三级覆盖、结构化输出与 spec 文件加载四类实战场景并理解其底层解析链路。PydanticAIHook 的定位PydanticAIHook是 common-ai provider 提供的薄适配器它从 Airflow Connection 中读取凭证与配置然后构造出 pydantic-ai 原生的Model与Agent对象。同一个 Connection 类型pydanticai可以工作于 pydantic-ai 支持的所有提供商OpenAI、Anthropic、Groq、Mistral、DeepSeek、Ollama、vLLM 等因此你无需为每个 LLM 厂商单独建 hook。它在该 provider 中的角色定位可见 hooks 索引文档LLMOperator、AgentOperator、LLMBranchOperator等算子默认就是基于这个 hook 构建 pydantic-ai Agent 的参见 Common AI Hooks 索引。核心实现位于 hooks/pydantic_ai.pyprovider 元数据包名apache-airflow-providers-common-ai、当前版本线 0.9.x声明在 provider.yaml。Connection 字段映射hook 的__init__接受llm_conn_id默认回退到pydanticai_default与可选的model_id见 PydanticAIHook 构造函数。一个 Connection 的字段如何被消费源码中的类文档写得很清楚Connection 字段用途passwordLLM 提供商的 API KeyhostBase URL可选如https://api.openai.com/v1或本地 Ollama/vLLM 端点extraJSON{model: openai:gpt-5.6-sol}UI 中的 Model 专属输入框就写入这个位置完整的 Connection 配置说明含 OpenAI、Anthropic、Ollama、Bedrock、Gemini 的 JSON 示例见 Pydantic AI Connection 文档。以 OpenAI 为例等价配置为{ conn_type: pydanticai, password: sk-..., extra: {\model\: \openai:gpt-5.6-sol\} }hook 还通过get_ui_field_behaviour()定制了 Connection 表单隐藏schema/port/login字段、把password重命名为 API Key并为host与extra提供占位提示源码。基本用法在 task 中调用 LLM最典型的使用方式是在task函数内实例化 hook、创建 agent、同步运行。官方示例 DAG example_pydantic_ai_hook.py 中的基础用法如下dag(scheduleNone, tags[example]) def example_pydantic_ai_hook(): task def generate_summary(text: str) - str: hook PydanticAIHook(llm_conn_idpydanticai_default) agent hook.create_agent(output_typestr, instructionsSummarize concisely.) result agent.run_sync(text) return result.output generate_summary(Apache Airflow is a platform for programmatically authoring...)要点create_agent返回的是一个标准的 pydantic-aiAgent之后agent.run_sync(...)、result.output都是 pydantic-ai 的原生 API你可以直接享受该框架的全部能力重试、流式、运行事件等。不传llm_conn_id时使用默认连接pydanticai_default这一默认值由default_conn_name类属性决定且各云厂商子类使用各自的默认名如pydanticai_azure_default——这是源码里特意在运行时解析而非绑定默认参数值的原因见 构造函数注释。不传output_type时默认是str。create_agent 的两个入口分支从 create_agent 实现 看它内部有两条路径给了spec_file委托给 pydantic-ai 的Agent.from_file把 spec 文件YAML/JSON作为 Agent 配置来源没给spec_file则instructions为必填参数缺失会抛ValueError: instructions is required when spec_file is not provided.直接Agent(self.get_conn(), output_type..., instructions..., **agent_kwargs)。此外**agent_kwargs中未显式使用的关键字如retries会原样透传给Agent构造函数测试用例 test_create_agent_with_params 验证了retries3之类的参数会正确转发。模型覆盖的三级优先级文档明确给出的模型指定优先级从高到低hook 上的model_id参数ConnectionextraJSON 中的model键两者都未设置时无默认值——get_conn()抛出ValueError: No model specified. Set model_id on the hook or the Model field on the connection.# 使用 connection extra 中声明的模型 hook PydanticAIHook(llm_conn_idmy_llm) # 用指定模型覆盖 connection 中的配置 hook PydanticAIHook(llm_conn_idmy_llm, model_idanthropic:claude-opus-4-6)源码解析get_conn 的解析链路get_conn 方法 的完整逻辑是取模型名self.model_id or extra.get(model, )为空即抛错取凭证api_key来自conn.passwordbase_url来自conn.host显式凭证路径_get_provider_kwargs()返回非空 dict 时构造一个provider_factory闭包用infer_provider_class(pname)(**kwargs)实例化带凭证的 provider再交给infer_model(model_name, provider_factory...)。这里有一个防御性细节若 provider 构造函数拒绝这些 kwargsTypeErrorhook 会打 warning 并回退到环境变量认证源码默认路径无任何显式凭证例如 Bedrock 用AWS_PROFILE、Vertex 用GOOGLE_APPLICATION_CREDENTIALS、Ollama 本地服务时直接infer_model(model_name)由 pydantic-ai 自行读取标准环境变量缓存解析出的Model缓存在self._modelhook 生命周期内只解析一次测试 test_get_conn_caches_model 断言两次get_conn()返回同一对象且infer_model只调用一次。上述显式凭证注入、base_url 单独提供、无凭证走环境变量等分支都在 test_pydantic_ai.py 的TestPydanticAIHookGetConn中有对应测试覆盖包括 Ollama 场景只有 base_url 没有 api_key的断言。结构化输出output_typepydantic-ai 的结构化输出能力通过 hook 直接透传定义一个 Pydantic 模型描述期望的输出形状然后作为output_type传给create_agent。官方示例dag(scheduleNone, tags[example]) def example_pydantic_ai_structured_output(): task def generate_sql(prompt: str) - dict: class SQLResult(BaseModel): query: str explanation: str hook PydanticAIHook(llm_conn_idpydanticai_default) agent hook.create_agent( output_typeSQLResult, instructionsGenerate a SQL query and explain it., ) result agent.run_sync(prompt) return result.output.model_dump() generate_sql(Find the top 10 customers by revenue)create_agent用overload声明了两类签名带output_type时返回Agent[object, OutputT]不带时返回Agent[object, str]overload 声明因此静态类型检查可以识别result.output的具体类型。当任务返回给 DAG 的数据需要可序列化时如示例中的result.output.model_dump()把 Pydantic 模型转成 dict 再返回是稳妥做法。从 AgentSpec 文件加载 Agent 配置如果不想在 Python 里硬编码模型名、指令和采样参数可以把它们放进一个 YAML 或 JSON 的AgentSpec文件通过spec_file传入。这样做的好处提示词工程与 DAG 逻辑分离Agent 配置可以独立版本管理、独立评审。文档给出的 spec 文件示例# agent_spec.yaml model: openai:gpt-4o-mini instructions: You are a concise summarizer. Given any text, respond with a single paragraph that captures the key points. model_settings: temperature: 0.3 retries: 2仓库中真实存在的示例文件 example_agent_spec.yaml 与之基本一致还额外声明了end_strategy: early。对应的 DAG 用法来自官方示例 example_pydantic_ai_hook.pytask def summarize_from_spec(text: str) - str: spec_path Path(__file__).parent / example_agent_spec.yaml hook PydanticAIHook(llm_conn_idpydanticai_default) # 模型、instructions、temperature、retries 全部来自 YAML 文件 agent hook.create_agent(spec_filespec_path) result agent.run_sync(text) return result.output task def summarize_with_additional_instructions(text: str) - str: 调用时追加指令与 spec 文件中的指令合并 spec_path Path(__file__).parent / example_agent_spec.yaml hook PydanticAIHook(llm_conn_idpydanticai_default) agent hook.create_agent( spec_filespec_path, instructionsSummarize in exactly one sentence., ) result agent.run_sync(text) return result.output行为规则与文档描述一致且由源码与测试印证模型优先级spec 文件中的model只有在 hook 未显式配置模型时才生效。create_agent内部先调用_get_conn_if_model_configured只有model_id或 connection 的modelextra 存在时才返回 hook 模型并传给Agent.from_file(model...)否则from_file只带output_type由 spec 文件自带模型测试 test_create_agent_with_spec_file_uses_file_model_when_hook_model_not_configured 断言了infer_model根本不被调用。指令合并spec_file与instructions同时给出时调用时的instructions会被追加到文件里的指令之后create_agent 实现 将两者都传给Agent.from_file只给spec_file时则仅使用文件中的值。spec_file接受str或pathlib.Pathoutput_type以及其他agent_kwargs同样会转发给Agent.from_file测试见 test_create_agent_with_spec_file_custom_output_type。在 task 中直接使用 Toolset官方示例还展示了不经AgentOperator、直接在task中给 agent 挂工具集的写法示例代码工具集实现位于 toolsets 目录task def analyze_revenue() - str: from airflow.providers.common.ai.toolsets.sql import SQLToolset hook PydanticAIHook(llm_conn_idpydanticai_default) agent hook.create_agent( output_typestr, instructions( You are a sales analytics assistant. Use the SQL tools to explore the database schema and answer questions. ), toolsets[ SQLToolset( db_conn_idmy_database, allowed_tables[customers, orders], max_rows20, ), ], ) result agent.run_sync(Which customers have spent the most? Show the top 5.) return result.output这里的toolsets[...]正是create_agent中**agent_kwargs透传机制的体现——hook 本身不感知具体工具类型全部交给 pydantic-ai 的Agent处理。连接自检与可观测性test_connectionhook 实现了 Airflow 的 test_connection在 UI 中点击 Test 时它只做模型解析验证模型字符串合法、provider 类可用给定凭证实例化不会真正发起 LLM API 调用——因为真实调用既昂贵又会因配额、计费、限流等与连通性无关的原因误报失败。自动 GenAI 链路追踪create_agent的 docstring 提到当[common.ai] otel_export_enabled配置生效且 worker 进程中存在活的 OTLPTracerProvider时返回的 agent 会被自动注入InstrumentationSettings从而通过 Airflow 的 tracing 管道输出 GenAI span实现见 observability.py 的 genai_instrumentation_settings。两个值得注意的细节未开启时返回Noneagent 保持不插桩、零开销include_content默认关闭除非显式开启[common.ai] capture_content否则 prompt、completion 与工具 IO 不会被写入 span调用方如果自己在agent_kwargs里传了instrumenthook 会把它 pop 出来并在构建后赋值给agent.instrument即用户显式值优先于 provider 自动插桩create_agent 尾部逻辑。更多配置说明可参见 observability 文档。云厂商专用子类基类PydanticAIHook覆盖标准api_key 可选base_url的提供商对于认证方式非标准的云厂商同一文件提供了三个子类各自覆写_get_provider_kwargs()把 Connection 字段映射到对应 provider 的构造参数子类conn_type / 默认连接关键映射PydanticAIAzureHookpydanticai_azure/pydanticai_azure_defaultpassword→api_key、host→azure_endpoint、extra.api_version→api_versionPydanticAIBedrockHookpydanticai_bedrock/pydanticai_bedrock_default凭证全部走extraIAM 键、region_name、profile_name、bearer token、超时等UI 中隐藏host/password字段PydanticAIVertexHookpydanticai_vertex/pydanticai_vertex_default支持extra内嵌service_account_info惰性导入 google-auth 构建 Credentials、project/location/api_keyBedrock 子类的凭证解析顺序为extra 中的 bearer token 优先于 extra 中的 IAM 键两者都缺省时回退到 AWS 默认凭证链AWS_PROFILE、实例角色等。Vertex 子类还有一个向后兼容细节extra里的vertexai字段会被接受但忽略并打 warning——因为 pydantic-ai 现在完全依据模型前缀google-cloud:vsgoogle:来选择 Vertex AI 还是 Generative Language API源码注释。各子类的 UI 表单定制与字段示例分别见对应的 Azure、Bedrock、Vertex 连接文档。验证与延伸阅读本文所有行为性结论均可在仓库内对照单元测试默认连接名、模型解析、spec_file 路由、指令合并、kwargs 转发等test_pydantic_ai.py端到端示例 DAGexample_pydantic_ai_hook.pyConnection 配置参考connections/pydantic_ai.rst如果你要更进一步用算子LLMOperator、AgentOperator等而非裸 hook可参考该 provider 的 operators 文档目录common-ai provider 的其他 hookLangChain、LlamaIndex、MCP选型对比见 hooks 索引。适用前提提示本文基于当前仓库中apache-airflow-providers-common-ai0.9.x 代码spec_file、结构化输出等能力依赖 pydantic-ai 2.x源码中对instrument属性的处理方式即针对 pydantic-ai 2.x 的 API 变更。【免费下载链接】airflowApache Airflow - A platform to programmatically author, schedule, and monitor workflows项目地址: https://gitcode.com/GitHub_Trending/ai/airflow创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考