migrating-ai-sdk-to-common-ai
This skill automates migration of Airflow projects from the deprecated airflow-ai-sdk to apache-airflow-providers-common-ai (0.1.0+), the official Airflow AI provider built on PydanticAI. Use it when replacing airflow-ai-sdk, migrating LLM decorators like @task.llm and @task.agent, converting model strings to connection-based configuration, or updating imports to the new provider package.
git clone --depth 1 https://github.com/astronomer/agents /tmp/migrating-ai-sdk-to-common-ai && cp -r /tmp/migrating-ai-sdk-to-common-ai/skills/migrating-ai-sdk-to-common-ai ~/.claude/skills/migrating-ai-sdk-to-common-aiSKILL.md
# Migrate airflow-ai-sdk to apache-airflow-providers-common-ai
This skill migrates Airflow projects from `airflow-ai-sdk` to `apache-airflow-providers-common-ai` (0.1.0+), the official Airflow AI provider built on PydanticAI.
> **CRITICAL**: The new provider requires **Airflow 3.0+** and **pydantic-ai-slim >= 1.34.0**. The API surface has changed: LLM configuration moves from code (model strings/objects) to Airflow connections (`pydanticai` type). There is no `@task.embed` in the new provider.
## Before starting
Use the Grep tool with the pattern below to inventory everything that needs to migrate:
```
airflow_ai_sdk|airflow-ai-sdk|ai_sdk|@task\.llm|@task\.agent|@task\.llm_branch|@task\.embed
```
From the results, capture:
1. All files importing `airflow-ai-sdk` / `airflow_ai_sdk`
2. Which decorators are in use: `@task.llm`, `@task.agent`, `@task.llm_branch`, `@task.embed`
3. The model configuration pattern (string names like `"gpt-5"`, or `OpenAIModel(...)` objects)
4. Any `airflow_ai_sdk.BaseModel` subclasses used as `output_type`
Use this inventory to drive the steps below.
---
## Step 1: Update requirements.txt
**Remove:**
```
airflow-ai-sdk[openai]
# or any variant: airflow-ai-sdk[openai]==0.1.7, airflow-ai-sdk[anthropic], etc.
```
**Add:**
```
apache-airflow-providers-common-ai[openai]>=0.1.0
```
Use the latest available 0.x version unless the user has pinned a specific one. Available extras match the LLM provider: `[openai]`, `[anthropic]`, `[google]`, `[bedrock]`, `[groq]`, `[mistral]`, `[mcp]`.
Keep `sentence-transformers` and `torch` if the project uses embeddings (they now run via plain `@task` instead of `@task.embed`).
---
## Step 2: Create PydanticAI connection
The new provider uses an Airflow connection instead of model strings or objects in code.
**Connection type:** `pydanticai`
**Default connection ID:** `pydanticai_default`
### Via environment variable (.env)
```bash
AIRFLOW_CONN_PYDANTICAI_DEFAULT='{
"conn_type": "pydanticai",
"password": "<api-key>",
"extra": {
"model": "<provider>:<model-name>"
}
}'
```
### Model format
The model field uses `provider:model` format:
| Provider | Example model value |
|----------|-------------------|
| 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` |
### Custom endpoints (Ollama, vLLM, Snowflake Cortex, etc.)
Set `host` to the base URL:
```bash
AIRFLOW_CONN_PYDANTICAI_CORTEX='{
"conn_type": "pydanticai",
"password": "<api-key>",
"host": "https://my-endpoint.com/v1",
"extra": {
"model": "openai:<model-name>"
}
}'
```
Use the `openai:` prefix for any OpenAI-compatible API, regardless of the actual provider.
### Connection ID convention
The env var name determines the connection ID:
- `AIRFLOW_CONN_PYDANTICAI_DEFAULT` creates `pydanticai_default`
- `AIRFLOW_CONN_PYDANTICAI_CORTEX` creates `pydanticai_cortex`
### Model resolution priority
1. `model_id` parameter on the decorator/operator (highest)
2. `model` in connection's extra JSON (fallback)
---
## Step 3: Migrate decorators
### @task.llm
```python
# 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
```
**Parameter mapping:**
| airflow-ai-sdk | common-ai provider | Notes |
|----------------|-------------------|-------|
| `model="gpt-5"` | `llm_conn_id="pydanticai_default"` | Model specified in connection |
| `model=OpenAIModel(...)` | `llm_conn_id="pydanticai_default"` | Model + endpoint in connection |
| `system_prompt="..."` | `system_prompt="..."` | Unchanged |
| `output_type=MyModel` | `output_type=MyModel` | Unchanged |
| `result_type=MyModel` | `output_type=MyModel` | `result_type` was already deprecated |
| (not available) | `model_id="openai:gpt-5"` | Override connection's model |
| (not available) | `require_approval=True` | Built-in HITL review |
| (not available) | `agent_params={...}` | Extra kwargs for pydantic-ai Agent |
### @task.llm_branch
```python
# 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
```
Only change: `model=` becomes `llm_conn_id=`.
### @task.agent
This has the biggest API change. The Agent is no longer pre-built in user code.
```python
# 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
@task.agent(
llm_conn_id="pydanticai_default",
system_prompt="You are a research assistant.",
agent_params={"tools": [search_tool, lookup_tool]},
)
def research(question: str) -> str:
return question
```
**Parameter mapping:**
| airflow-ai-sdk | common-ai provider | Notes |
|----------------|-------------------|-------|
| `agent=Agent(model, ...)` | `llm_conn_id="...Add a new method to both Airflow adapters
Add a new MCP tool to server.py
Verify code works with both Airflow 2.x and 3.x
Airflow adapter pattern for v2/v3 API compatibility. Use when working with adapters, version detection, or adding new API methods that need to work across Airflow 2.x and 3.x.
Use when the user needs human-in-the-loop workflows in Airflow (approval/reject, form input, or human-driven branching). Covers ApprovalOperator, HITLOperator, HITLBranchOperator, HITLEntryOperator, HITLTrigger. Requires Airflow 3.1+. Does not cover AI/LLM calls (see airflow-ai).
Build Airflow 3.1+ plugins that embed FastAPI apps, custom UI pages, React components, middleware, macros, and operator links directly into the Airflow UI. Use this skill whenever the user wants to create an Airflow plugin, add a custom UI page or nav entry to Airflow, build FastAPI-backed endpoints inside Airflow, serve static assets from a plugin, embed a React app in the Airflow UI, add middleware to the Airflow API server, create custom operator extra links, or call the Airflow REST API from inside a plugin. Also trigger when the user mentions AirflowPlugin, fastapi_apps, external_views, react_apps, plugin registration, or embedding a web app in Airflow 3.1+. If someone is building anything custom inside Airflow 3.1+ that involves Python and a browser-facing interface, this skill almost certainly applies.
Queries, manages, and troubleshoots Apache Airflow using the af CLI. Covers listing DAGs, triggering runs, reading task logs, diagnosing failures, debugging DAG import errors, checking connections, variables, pools, and monitoring health. Also routes to sub-skills for writing DAGs, debugging, deploying, and migrating Airflow 2 to 3. Use when user mentions "Airflow", "DAG", "DAG run", "task log", "import error", "parse error", "broken DAG", or asks to "trigger a pipeline", "debug import errors", "check Airflow health", "list connections", "retry a run", or any Airflow operation. Do NOT use for warehouse/SQL analytics on Airflow metadata tables — use analyzing-data instead.
Queries data warehouse and answers business questions about data. Handles questions requiring database/warehouse queries including "who uses X", "how many Y", "show me Z", "find customers", "what is the count", data lookups, metrics, trends, or SQL analysis.