migrating-ai-sdk-to-common-ai

将 airflow-ai-sdk 迁移到 apache-airflow-providers-common-ai

此技能可将 Airflow 项目从 airflow-ai-sdk 迁移到 apache-airflow-providers-common-ai（目标版本为 0.4.0+），这是基于 PydanticAI 构建的官方 Airflow AI 提供程序。此技能也涵盖已使用 common-ai 0.1.x 的项目升级，因为 0.1.0 到 0.4.0 之间有多项功能发生了变化（多模态提示词、toolsets、嵌入算子、结构化输出的 XCom 行为）。

重要：新提供程序要求使用 Airflow 3.0+，并且（对于 0.4.0）要求使用 pydantic-ai-slim >= 1.71.0。API 已发生变化：LLM 配置从代码（模型字符串／对象）移到了 Airflow 连接（pydanticai 类型）。新提供程序没有 @task.embed；嵌入功能改用 LlamaIndex 集成或普通的 @task（见步骤 3）。

开始之前

使用 Grep 工具和以下模式，盘点所有需要迁移的内容：

airflow_ai_sdk|airflow-ai-sdk|ai_sdk|@task\.llm|@task\.agent|@task\.llm_branch|@task\.embed

根据结果，记录：

1. 所有导入 airflow-ai-sdk / airflow_ai_sdk 的文件
2. 使用了哪些装饰器：@task.llm、@task.agent、@task.llm_branch、@task.embed
3. 模型配置模式（如 "gpt-5" 这样的字符串名称，或 OpenAIModel(...) 对象）
4. 所有用作 output_type 的 airflow_ai_sdk.BaseModel 子类

按照以下步骤，使用这份清单进行迁移。

---

步骤 1：更新 requirements.txt

移除：

airflow-ai-sdk[openai]
# or any variant: airflow-ai-sdk[openai]==0.1.7, airflow-ai-sdk[anthropic], etc.

添加：

apache-airflow-providers-common-ai[openai]>=0.4.0

除非用户指定了固定版本，否则请使用最新的 0.x 版本。可用的 extras（0.4.0）：[openai]、[anthropic]、[google]、[bedrock]、[llamaindex]、[langchain]、[mcp]，以及供 DocumentLoaderOperator 使用的文件格式 extras（[pdf]、[docx]、[parquet]、[avro]），还有供 SQL 算子使用的 [sql]／[common-sql]。没有 [groq]／[mistral] extras；对于这些提供程序，请自行安装相应的 pydantic-ai-slim extra。

如果项目要将 @task.embed 迁移到 LlamaIndexEmbeddingOperator（推荐，见步骤 3），请添加 [llamaindex]。在这种情况下，通常可以移除 sentence-transformers 和 torch，从而大幅缩小镜像。仅当项目通过普通的 @task 继续使用本地 sentence-transformers 嵌入时，才保留它们。

---

步骤 2：创建 PydanticAI 连接

新提供程序使用 Airflow 连接，而非代码中的模型字符串或对象。

连接类型： pydanticai 默认连接 ID： pydanticai_default

通过环境变量（.env）

AIRFLOW_CONN_PYDANTICAI_DEFAULT='{
    "conn_type": "pydanticai",
    "password": "<api-key>",
    "extra": {
        "model": "<provider>:<model-name>"
    }
}'

模型格式

模型字段使用 provider:model 格式：

提供程序	模型值示例
OpenAI	openai:gpt-5
Anthropic	anthropic:claude-sonnet-4-20250514
Google	google:gemini-2.5-pro
Groq	groq:llama-3.3-70b-versatile
Mistral	mistral:mistral-large-latest
Bedrock	bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0

自定义端点（Ollama、vLLM、Snowflake Cortex 等）

将 host 设置为基础 URL：

AIRFLOW_CONN_PYDANTICAI_CORTEX='{
    "conn_type": "pydanticai",
    "password": "<api-key>",
    "host": "https://my-endpoint.com/v1",
    "extra": {
        "model": "openai:<model-name>"
    }
}'

无论实际提供程序是什么，任何兼容 OpenAI 的 API 都使用 openai: 前缀。

连接 ID 命名约定

环境变量名称决定连接 ID：

• AIRFLOW_CONN_PYDANTICAI_DEFAULT 会创建 pydanticai_default
• AIRFLOW_CONN_PYDANTICAI_CORTEX 会创建 pydanticai_cortex

模型解析优先级

1. 装饰器／算子上的 model_id 参数（优先级最高）
2. 连接 extra JSON 中的 model（备用）

其他连接类型（0.4.0）

除了 pydanticai，该提供程序还注册了特定于供应商的连接类型：pydanticai-azure（Azure OpenAI：host = 端点，extra 为 api_version）、pydanticai-bedrock（extra 中的 AWS 凭证／区域）和 pydanticai-vertex（extra 中的 GCP 项目／位置）。LlamaIndex 和 LangChain 钩子会从传入的连接 ID 对应的连接中读取 API 密钥、host 和 extra，因此一个 pydanticai_default 连接可同时用于 LLM 调用和嵌入：整个项目只需配置一项 API 密钥。

---

步骤 3：迁移装饰器

@task.llm

# BEFORE (airflow-ai-sdk)
import airflow_ai_sdk as ai_sdk

class MyOutput(ai_sdk.BaseModel):
    field: str

@task.llm(
    model="gpt-5",                    # or model=OpenAIModel(...)
    system_prompt="You are helpful.",
    output_type=MyOutput,
)
def my_task(text: str) -> str:
    return text

# AFTER (apache-airflow-providers-common-ai)
from pydantic import BaseModel

class MyOutput(BaseModel):
    field: str

@task.llm(
    llm_conn_id="pydanticai_default",  # Airflow connection ID
    system_prompt="You are helpful.",
    output_type=MyOutput,
)
def my_task(text: str) -> str:
    return text

参数映射：

airflow-ai-sdk	通用 AI 提供方	备注
model="gpt-5"	llm_conn_id="pydanticai_default"	在连接中指定模型
model=OpenAIModel(...)	llm_conn_id="pydanticai_default"	在连接中指定模型和端点
system_prompt="..."	system_prompt="..."	不变
output_type=MyModel	output_type=MyModel	不变
result_type=MyModel	output_type=MyModel	result_type 已弃用
（不可用）	model_id="openai:gpt-5"	覆盖连接中的模型
（不可用）	require_approval=True	内置 HITL 审核
（不可用）	agent_params={...}	传递给 pydantic-ai Agent 的额外 kwargs
（不可用）	serialize_output=True	强制将 BaseModel 输出转换为字典格式

多模态提示词（0.4.0+）： 翻译函数可以返回 Sequence[UserContent]，而不是字符串，例如用于视觉处理：

@task.llm(llm_conn_id="pydanticai_default", system_prompt="...", output_type=ReviewAnalysis)
def analyze(text: str, image_path: str | None = None):
    if image_path:
        with open(image_path, "rb") as f:
            return [text, BinaryContent(data=f.read(), media_type="image/jpeg")]
    return text

这与旧版 airflow-ai-sdk 的视觉处理模式一致，因此视觉处理代码无需更改即可迁移。注意：common-ai 0.1.x 仅接受字符串——如果项目为了迁移到 0.1.0 而禁用了视觉处理，请在升级到 0.4.0 时重新启用。非字符串提示词与 require_approval=True / enable_hitl_review=True 不兼容（两者都会将提示词渲染为文本）。

通过 XCom 输出结构化数据（0.4.0 行为变更）： 当使用 output_type=<BaseModel subclass> 时，在任务 SDK 支持 SUPPORTS_OPERATOR_DESERIALIZATION_WALKER 属性的 Airflow 核心版本中，模型实例会通过 XCom 传递（下游使用属性访问）；在较旧的核心版本中（包括 Astro Runtime 3.2 task SDK 1.2.x），提供方会自动将其转换为字典（使用下标访问）。在选择下游使用属性访问还是字典访问之前，请检查运行时收到的数据格式；也可以设置 serialize_output=True，以便始终强制使用字典格式。output_type 类必须定义在模块级别（嵌套类无法从 XCom 反序列化）。

@task.llm_branch

# BEFORE
@task.llm_branch(
    model="gpt-5",
    system_prompt="Choose a team...",
    allow_multiple_branches=False,
)
def route(text: str) -> str:
    return text

# AFTER
@task.llm_branch(
    llm_conn_id="pydanticai_default",
    system_prompt="Choose a team...",
    allow_multiple_branches=False,    # same parameter, unchanged
)
def route(text: str) -> str:
    return text

唯一的变更：将 model= 改为 llm_conn_id=。

@task.agent

这是 API 变更最大的部分。用户代码中不再预先构建 Agent。

# BEFORE (airflow-ai-sdk) - Agent built at module level
from pydantic_ai import Agent

my_agent = Agent(
    "gpt-5",
    system_prompt="You are a research assistant.",
    tools=[search_tool, lookup_tool],
)

@task.agent(agent=my_agent)
def research(question: str) -> str:
    return question

# AFTER (common-ai provider) - No Agent object, config via parameters
from pydantic_ai.toolsets import FunctionToolset

@task.agent(
    llm_conn_id="pydanticai_default",
    system_prompt="You are a research assistant.",
    toolsets=[FunctionToolset(tools=[search_tool, lookup_tool])],
)
def research(question: str) -> str:
    return question

参数映射：

airflow-ai-sdk	common-ai provider	说明
agent=Agent(model, ...)	llm_conn_id="..."	使用连接中的模型
Agent 的 system_prompt	system_prompt="..."	现在是装饰器参数
Agent 的 tools=[...]	toolsets=[FunctionToolset(tools=[...])]	推荐使用：会自动记录工具调用日志
Agent 的 tools=[...]	agent_params={"tools": [...]}	也可使用，但不会记录工具调用日志
Agent 的 output_type	output_type=MyModel	现在是装饰器参数
（不可用）	durable=True	步骤级缓存（需要 [common.ai] durable_cache_path）
（不可用）	enable_hitl_review=True	迭代式人工审核循环（见下文）

关键要点： 之前在 Agent() 构造函数中配置的所有内容，现在都应放入顶层装饰器参数或 agent_params。agent_params 字典会直接传给 pydantic-ai 的 Agent 构造函数。优先使用 toolsets，而不是 agent_params["tools"]：Operator 会将每个工具集封装在 LoggingToolset 中，因此每次工具调用都会连同耗时一起显示在任务日志中。

enable_hitl_review 的行为： 任务会先生成初稿，然后阻塞，直到有人采取操作。审核者可通过任务实例上的 HITL Review 选项卡/额外链接（来自 provider 自动注册的 hitl_review 插件的聊天界面）请求修改（Agent 会根据消息历史中的反馈重新生成内容）或批准。限制：需要字符串提示词；与 durable=True 不兼容；最终输出（可能经过重新生成）会传入 XCom。请提醒用户，除非设置 hitl_timeout，否则 DAG 运行会无限期地停在此任务上。对于无头测试，该插件会在 /hitl-review 下提供 REST 端点：GET /sessions/find、POST /sessions/feedback、POST /sessions/approve、POST /sessions/reject（查询参数为 dag_id、task_id、run_id、map_index）。

@task.embed（无等效功能——三种替代方案）

新 provider 不包含 embed 装饰器。请根据项目需求选择替代方案：

选项 A（推荐）：LlamaIndexEmbeddingOperator（0.4.0，[llamaindex] extra）。基于连接，一个任务即可嵌入整个文档列表；设置 persist_dir 后，生成的向量索引会持久化，以供检索使用（可与 LlamaIndexRetrievalOperator 配合）：

from airflow.providers.common.ai.operators.llamaindex_embedding import LlamaIndexEmbeddingOperator

_embeddings = LlamaIndexEmbeddingOperator(
    task_id="create_embeddings",
    documents=[{"text": "...", "metadata": {"id": 1}}, ...],  # templated, accepts XComArg
    llm_conn_id="pydanticai_default",   # reuses the same connection (API key only)
    embed_model="text-embedding-3-small",
    persist_dir=f"{AIRFLOW_HOME}/include/my_index",  # optional; local path or s3://, gs://, ...
)

该 Operator 会返回 {"chunks": [{"text", "metadata", "vector"}], ...}。请在每个文档的 metadata 中放入一个稳定的键——该键会在分块过程中保留，因此可以将向量映射回源记录。

选项 B：用于原始向量的 LlamaIndexHook（没有 Operator，也不会持久化索引）。当向量直接写入数据库时，这是最简便的方式：

@task
def create_embeddings(rows):
    from airflow.providers.common.ai.hooks.llamaindex import LlamaIndexHook
    embed_model = LlamaIndexHook(
        llm_conn_id="pydanticai_default",
        embed_model="text-embedding-3-small",
    ).get_embedding_model()
    vectors = embed_model.get_text_embedding_batch([r["text"] for r in rows])
    return list(zip([r["id"] for r in rows], vectors))

选项 C：使用 sentence-transformers 的普通 @task（保留旧的本地/离线行为，不产生 API 费用；需要在 requirements 中保留 sentence-transformers 和 torch）：

@task
def embed_texts(texts: list[str]) -> list[list[float]]:
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer("all-MiniLM-L6-v2")
    return model.encode(texts, normalize_embeddings=True).tolist()

关于维度的说明：从 all-MiniLM-L6-v2（384）切换到 text-embedding-3-small（1536）会改变向量大小——必须重新生成现有存储的嵌入，并且固定大小的向量列（例如 pgvector vector(384)）需要修改架构。请在一个任务/批次调用中嵌入所有文本，而不是针对每段文本使用 .expand()：批处理只需进行一次 API 往返，还可避免逐个任务加载模型。

---

第 4 步：更新导入语句

旧导入语句	新导入语句
import airflow_ai_sdk as ai_sdk	完全移除
from airflow_ai_sdk import BaseModel	from pydantic import BaseModel
from airflow_ai_sdk.models.base import BaseModel	from pydantic import BaseModel
class Foo(ai_sdk.BaseModel):	class Foo(BaseModel):
from pydantic_ai import Agent	如果 Agent 仅用于 @task.agent，则移除
from pydantic_ai.models.openai import OpenAIModel	移除（模型配置现在位于连接中）
（新增）	用于 @task.agent 工具集的 from pydantic_ai.toolsets import FunctionToolset

@task.llm、@task.agent、@task.llm_branch 装饰器会由提供程序自动注册。除 from airflow.sdk import task 外，无需显式导入。

对于非装饰器用途（例如，用于多模态的 BinaryContent），pydantic_ai 导入仍然有效，因为新的提供程序依赖 pydantic-ai-slim（提供程序 0.4.0 需要 >= 1.71.0）。

---

第 5 步：更新 connections.yaml（如果用于本地测试）

pydanticai_default:
  conn_type: pydanticai
  password: <api-key>
  extra:
    model: "openai:gpt-5"

对于自定义端点：

pydanticai_cortex:
  conn_type: pydanticai
  password: <api-key>
  host: https://my-endpoint.com/v1
  extra:
    model: "openai:llama3.1-8b"

---

第 6 步：清理环境变量

新的提供程序从 pydanticai 连接读取模型配置，因此，之前在代码中用于配置模型的环境变量通常是多余的。删除其中任何变量前，请搜索项目（以及任何同级脚本或服务），确认没有其他内容仍在引用它们：

OPENAI_API_KEY|OPENAI_BASE_URL|ANTHROPIC_API_KEY|GOOGLE_API_KEY

仅当没有其他代码引用时，才可考虑移除以下变量：

• OPENAI_API_KEY（现在位于 pydanticai 连接的密码字段中）
• OPENAI_BASE_URL（现在位于连接的主机字段中）
• 自定义模型名称变量（现在位于连接的 extra.model 中）

如果迁移后的 DAG 以外的内容仍在使用这些变量（其他尚未迁移的 DAG、辅助脚本、共享 .env 的非 Airflow 服务），请保留它们。

保留所有连接的 AIRFLOW_CONN_* 环境变量。

---

第 7 步：验证

迁移后，搜索代码库以确认没有遗留引用：

airflow_ai_sdk|airflow-ai-sdk|ai_sdk\.BaseModel|from pydantic_ai import Agent|from pydantic_ai.models

验证：

• 没有来自 airflow_ai_sdk 的导入语句
• 没有为 @task.agent 创建 Agent() 对象（除非在装饰器之外使用）
• LLM 装饰器上没有 model= 参数（应改为 llm_conn_id=）
• 所有 @task.embed 均已替换（使用 LlamaIndex 运算符/钩子或普通 @task）；如果模型或维度发生变化，请重新生成已存储的嵌入
• 如果视觉翻译函数在 common-ai 0.1.x 下被限制为仅返回字符串，请恢复为返回 [text, BinaryContent(...)]
• output_type=BaseModel 结果的下游使用方应采用实际收到的 XCom 形状（较旧核心版本中为字典，较新版本中为实例；serialize_output=True 可固定其形状）
• 已在 .env 或 connections.yaml 中配置 pydanticai 连接
• requirements.txt 使用 apache-airflow-providers-common-ai[...] 而非 airflow-ai-sdk[...]；如果不再使用，则移除 torch/sentence-transformers
• 端到端运行 DAG：启用了 enable_hitl_review=True 或 require_approval=True 的任务会等待人工输入，因此测试计划必须包含对这些任务的处理（通过 UI 标签页或 /hitl-review REST）

---

快速参考：common-ai provider 中的新功能

迁移后即可使用这些功能，但 airflow-ai-sdk 中没有对应功能：

功能	参数 / API	起始版本	说明
HITL 审批	@task.llm 上的 require_approval=True	0.1.0	返回结果前暂停以供人工审核
HITL 审核循环	@task.agent 上的 enable_hitl_review=True	0.1.0	通过重新生成进行迭代审核（通过 hitl_review 插件提供聊天 UI）
持久执行	@task.agent 上的 durable=True	0.1.0	按步骤缓存以增强韧性
工具日志记录	@task.agent 上的 enable_tool_logging=True	0.1.0	INFO 级别的工具调用日志（默认：开启；需要 toolsets）
模型覆盖	model_id="openai:gpt-5"	0.1.0	按任务覆盖连接中的模型
文件分析	@task.llm_file_analysis	0.1.0	通过 ObjectStoragePath 分析文件 / 图像
自然语言转 SQL	@task.llm_sql	0.1.0	根据自然语言生成 SQL
多模态提示词	翻译函数返回 Sequence[UserContent]	0.4.0	在 @task.llm / @task.agent / @task.llm_branch 中使用视觉及其他二进制内容
通过 XCom 传递 Pydantic 实例	output_type=BaseModel（可通过 serialize_output 选择退出）	0.4.0	在支持的核心中，实例通过 XCom 传递；否则回退为 dict
嵌入	LlamaIndexEmbeddingOperator（+ persist_dir）	0.4.0	基于连接的嵌入 + 持久化向量索引
检索	LlamaIndexRetrievalOperator	0.4.0	在持久化索引上执行 top-k 相似度搜索