Feature map¶

Most of the provider’s operators and decorators run Pydantic AI agents, with the model taken from an Airflow connection. A feature that calls a model through those connections also needs the extra for your model vendor (Supported model providers); the Extra column lists what it needs on top of that, and a bare name there is an extra of this provider. Installation lists the features that need a newer Airflow than the provider’s 2.11, and Stable and experimental features says which features are stable.

Run a model or an agent¶

Feature	What it does	Extra
@task.llm	Sends one prompt and returns the answer as text or a typed object.	None
@task.agent	Runs an agent that calls tools until it has an answer.	None
@task.llm_branch	Lets the model pick which downstream tasks run.	None
@task.llm_sql	Writes a SQL query from a question and checks it, without running it.	`sql`; for files, `apache-airflow-providers-common-sql[datafusion]`
@task.llm_file_analysis	Sends files from object storage, including images and PDFs, to the model with a prompt.	`parquet` or `avro` for those formats
@task.llm_schema_compare	Compares table schemas across databases and files and grades each mismatch.	`common.sql`; for files, `apache-airflow-providers-common-sql[datafusion]`
@task.llm_batch	Sends many prompts through the OpenAI or Anthropic batch API at about half the per-token price.	None
Document and RAG pipelines	Loads documents, embeds them and retrieves the passages closest to a question.	`llamaindex`, and install the vendor’s package for vendors other than OpenAI; `pdf` or `docx` for those formats

Choose a model¶

Feature	What it does	Extra
Models and providers	Takes the vendor, model and key from a connection, so changing vendor is a connection edit.	The vendor’s, such as `openai` or `anthropic`; the linked page lists each
Self-hosted models	Runs tasks against a model server you host, such as Ollama or vLLM.	`openai`
Decision models	Uses a model that cannot write text, such as TypeSafe’s hosted Jev or Ollama’s decision models, to answer a typed question with one value from a set you define, plus a confidence.	Jev: `typesafe`, and install `pydantic-ai-slim` 2.45+; others: install `pydantic-ai-slim` 2.53+
LlamaIndex models	Builds LlamaIndex language and embedding models from a connection.	`llamaindex`, and install the vendor’s package for vendors other than OpenAI

Give an agent tools¶

Toolsets compares what each toolset reaches and how you limit it; Securing agent tools covers what an agent holding them can reach, with a production checklist.

Feature	What it does	Extra
HookToolset	Turns the methods you list on any Airflow hook into tools.	None
SQLToolset	Queries a database, read-only by default, and only the tables in `allowed_tables` when you set it.	`sql`
ObjectStorageToolset	Lists and reads files under one object storage path.	None
MCPToolset	Gives the agent the tools of an MCP server.	`mcp`
DataFusionToolset	Runs SQL over Parquet, CSV and Avro files and Iceberg tables, in the worker.	`sql`, plus `apache-airflow-providers-common-sql[datafusion]`, and its `apache.iceberg` extra for Iceberg
AgentSkillsToolset	Loads `SKILL.md` instruction bundles when the model asks for one.	`skills`; on Airflow 2, install `pydantic-ai-skills`
SandboxToolset	Gives the agent a shell and a filesystem in an isolated sandbox: a hosted service, a server you run, or a microVM on your machine for local development.	`modal`, `islo`, `opensandbox`, `openshell` or `boat`, by backend; none for `sbx`
ManagedAgentToolset	Sends a prompt to an agent on Bedrock AgentCore or Vertex AI Agent Engine.	The Amazon or Google provider’s `common.ai` extra
Progressive tool disclosure	Hides a toolset’s tools from the model until it searches for them, with `.defer_loading()` on any toolset.	None

Control output and spend¶

Feature	What it does	Extra
Structured output	Returns a Pydantic object instead of text, on `@task.llm`, `@task.llm_file_analysis` and `@task.agent`.	None
Usage limits	Fails the task past a cap on requests, tokens, tool calls or estimated cost (for models pydantic-ai can price), on `@task.agent` and every `@task.llm*` task except `@task.llm_batch`.	None

Shape an agent run¶

These configure @task.agent and AgentOperator.

Feature	What it does	Extra
Message history	Continues a conversation from an earlier run.	None
Capabilities and guardrails	Turns on model reasoning (`Thinking`), web search (`WebSearch`) and other Pydantic AI capabilities, and input and output guards.	`shields` for guards
Code mode	Lets the model call several tools from one Python snippet.	`code-mode`
Agent memory	Pydantic AI’s `Memory` capability keeps notes the agent writes in one Dag run and reads in the next.	None; install `pydantic-ai-harness`
Prompt caching	Marks the repeated start of each request for caching on Anthropic, and on the Bedrock and OpenRouter models that support it, so later requests pay less for it; OpenAI and Gemini cache on their own. On by default.	None

Put a Human in the loop¶

Feature	What it does	Extra
Approve outputs	Pauses any `@task.llm*` task except `@task.llm_batch` until a person approves or rejects the output, optionally editing it.	None
Confidence gate	Sends the answer to a person, or fails the task, when the model’s confidence is under a bar you set. Only decision models report a confidence; with a text model every answer counts as uncertain. On `@task.llm` and `@task.llm_branch`.	As for decision models
Review agent sessions	Lets a reviewer send a `@task.agent` answer back with feedback, round after round.	None
Approve tool calls	Pauses `@task.agent` before a marked tool runs, until a person approves or rejects the call.	None

Recover from failures¶

Feature	What it does	Extra
Durable execution	With `durable=True` on `@task.agent`, or the `AirflowDurability` capability on your own agent, a retry replays the model and tool calls that finished instead of paying for them again.	None
Retry policies	Has a model read a failure and decide whether the task retries, and when. Works on any task.	None
Provider fallback	Moves to the next connection in `fallback_conn_ids` when a vendor’s API returns an error.	None

Bring your own agent¶

The Strands, ADK, Anthropic SDK, OpenAI Agents SDK and other-framework routes get every toolset except MCPToolset and AgentSkillsToolset.

Feature	What it does	Extra
Pydantic AI agent	Gives an agent you run in a `@task` a model from a connection and the toolsets.	None
Strands Agents	Gives a Strands agent the toolsets. Run through `invoke_durably`, it resumes on retry.	None; install `strands-agents`
Google ADK	Gives an ADK agent the toolsets.	None; install `google-adk`
Anthropic SDK tool runner	Gives a loop on the Anthropic SDK’s `tool_runner` the toolsets.	`anthropic`
OpenAI Agents SDK	Gives an OpenAI Agents SDK agent the toolsets.	None; install `openai-agents`
LangChain	Builds LangChain chat and embedding models from a connection, and turns a toolset into LangChain tools.	`langchain`, and install `langchain-<vendor>`
Any other framework	Gives an agent in any framework the toolsets, through a framework-neutral tool interface.	None

See what ran¶

Feature	What it does	Extra
Tracing	Exports model and tool calls as OpenTelemetry spans.	None
AI Model tab	Shows the model on the task instance page of `@task.llm` and `@task.agent` tasks, and for `@task.agent` the token counts and an estimated cost.	None