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 を使用しているプロジェクトのアップグレードにも対応しています。マルチモーダルプロンプト、toolsets、埋め込みオペレーター、構造化出力の XCom 動作など、複数の機能が 0.1.0 から 0.4.0 の間に変更されているためです。

重要：新しいプロバイダーには 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 バージョンを使います。利用可能な追加機能（0.4.0）：[openai]、[anthropic]、[google]、[bedrock]、[llamaindex]、[langchain]、[mcp] のほか、DocumentLoaderOperator 用のファイル形式追加機能（[pdf]、[docx]、[parquet]、[avro]）、SQL オペレーター用の [sql]／[common-sql] があります。[groq]／[mistral] の追加機能はありません。これらのプロバイダーを使用する場合は、対応する pydantic-ai-slim の追加機能を自分でインストールしてください。

プロジェクトで @task.embed を LlamaIndexEmbeddingOperator に移行する場合は、[llamaindex] を追加します（推奨。手順 3 を参照）。この場合、通常は 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（AWS 認証情報／リージョンを extra に設定）、pydanticai-vertex（GCP プロジェクト／ロケーションを extra に設定）です。LlamaIndex および LangChain のフックは、指定された接続 ID から API キー／ホスト／extra を読み取ります。そのため、1 つの pydanticai_default 接続を LLM の呼び出しと埋め込み処理の両方に利用できます。プロジェクト全体で API キーの登録は 1 か所だけで済みます。

---

手順 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	common-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 に渡す追加のキーワード引数
（利用不可）	serialize_output=True	BaseModel の出力を強制的に dict 形式にします

マルチモーダルプロンプト（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 のタスク SDK 1.2.x を含む）では、プロバイダーが自動的にdict（添字アクセス）に変換します。下流で属性アクセスと dict アクセスのどちらを使うか決める前に、実行時にどちらの形式で届くかを確認してください。または、どこでも dict 形式に統一するには 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 コンストラクターに直接渡されます。agent_params["tools"] より toolsets を優先してください。オペレーターは各ツールセットを LoggingToolset でラップするため、各ツール呼び出しの所要時間がタスクログに記録されます。

enable_hitl_review の動作： タスクは最初のドラフトを生成し、その後、人間が操作するまでブロックします。レビュー担当者はタスクインスタンスの HITL Review タブ／追加リンク（プロバイダーの自動登録プラグイン hitl_review のチャット UI）を使って、修正を依頼するか（エージェントはメッセージ履歴のフィードバックを使って再生成します）、承認します。制約：文字列プロンプトが必要で、durable=True とは併用できません。また、XCom に渡されるのは最終出力（再生成された場合は再生成後の出力）です。hitl_timeout が設定されていない場合、このタスクで Dag run が無期限に待機することをユーザーに伝えてください。ヘッドレステストでは、プラグインが /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（同等の機能なし — 置き換え方法は 3 つ）

新しいプロバイダーには embed デコレーターが含まれていません。プロジェクトの要件に応じて、置き換え方法を選んでください。

オプション A（推奨）：LlamaIndexEmbeddingOperator（0.4.0、[llamaindex] extra）。接続ベースで、1 つのタスクでドキュメントリスト全体を埋め込みます。また、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://, ...
)

オペレーターは {"chunks": [{"text", "metadata", "vector"}], ...} を返します。各ドキュメントの metadata に安定したキーを設定してください。チャンク分割後も引き継がれるため、ベクトルを元のレコードに対応付けられます。

オプション B：生のベクトルには LlamaIndexHook（オペレーターなし、インデックスの永続化なし）。ベクトルをデータベースに直接渡す場合に最短の方法です：

@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() するのではなく、1 つのタスク／バッチ呼び出しですべてのテキストを埋め込んでください。バッチ処理なら API の往復が 1 回で済み、タスクごとにモデルを読み込む必要もありません。

---

ステップ 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	@task.agent でのみ 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 接続からモデル設定を読み取るため、以前コード内のモデル設定に使用していた環境変数は通常不要になります。削除する前に、プロジェクト（および関連するスクリプトやサービス）を grep して、ほかで参照されていないことを確認してください。

OPENAI_API_KEY|OPENAI_BASE_URL|ANTHROPIC_API_KEY|GOOGLE_API_KEY

削除できる候補 （ほかのコードで参照されていない場合のみ）：

• OPENAI_API_KEY（現在は pydanticai 接続のパスワードフィールドに設定）
• OPENAI_BASE_URL（現在は接続のホストフィールドに設定）
• カスタムのモデル名変数（現在は connection の extra.model に設定）

移行した DAG 以外で使用されている場合（まだ移行していないほかの DAG、ヘルパースクリプト、.env を共有する Airflow 以外のサービスなど）は、そのまま残してください。

すべての接続の AIRFLOW_CONN_* 環境変数は残してください。

---

ステップ 7：検証する

移行後、コードベースを grep して古い参照が残っていないことを確認してください。

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 で文字列のみの戻り値に制限されていた場合、Vision 翻訳関数が再び [text, BinaryContent(...)] を返す
• output_type=BaseModel の結果を使う下流のコンシューマーが、実際に届く XCom の形式を使用している（古いコアでは dict、新しいコアではインスタンス。serialize_output=True を指定すると形式を固定できます）
• .env または connections.yaml に pydanticai 接続が設定されている
• requirements.txt で airflow-ai-sdk[...] ではなく apache-airflow-providers-common-ai[...] を使用している。不要になった場合は torch／sentence-transformers を削除する
• DAG をエンドツーエンドで実行する。enable_hitl_review=True または require_approval=True を指定したタスクは人間の入力を待つため、テスト計画にはそれらへの対応（UI タブまたは /hitl-review REST）を含める

---

クイックリファレンス：common-ai プロバイダーの新機能

これらの機能は移行後に利用できますが、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 を介してファイルや画像を分析します
NL-to-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 を通じて渡され、それ以外では辞書にフォールバックします
埋め込み	LlamaIndexEmbeddingOperator（および persist_dir）	0.4.0	接続ベースの埋め込みと、永続化されたベクトルインデックス
検索	LlamaIndexRetrievalOperator	0.4.0	永続化されたインデックスに対する Top-k 類似度検索