diff --git a/providers/common/ai/docs/connections/mcp.rst b/providers/common/ai/docs/connections/mcp.rst index 07be621e9a186..a324bbdffe60f 100644 --- a/providers/common/ai/docs/connections/mcp.rst +++ b/providers/common/ai/docs/connections/mcp.rst @@ -150,9 +150,9 @@ into task logs) without ever being written to the connection: token_provider=mint_snowflake_jwt, ) -``token_provider`` is resolved in DAG code (it is a Python callable, not a stored +``token_provider`` is resolved in Dag code (it is a Python callable, not a stored connection field), so the signing key stays in your environment and is never baked -into the serialized DAG. +into the serialized Dag. Secrets in stdio subprocess environments ----------------------------------------- @@ -193,5 +193,5 @@ managed as its own Airflow connection: :start-after: [START howto_toolset_mcp_env_provider] :end-before: [END howto_toolset_mcp_env_provider] -Like ``token_provider``, ``env_provider`` is resolved in DAG code, so the secret is -fetched at task-execution time and never baked into the serialized DAG. +Like ``token_provider``, ``env_provider`` is resolved in Dag code, so the secret is +fetched at task-execution time and never baked into the serialized Dag. diff --git a/providers/common/ai/docs/connections/pydantic_ai.rst b/providers/common/ai/docs/connections/pydantic_ai.rst index a7aa0c85a5016..f3165379b0a07 100644 --- a/providers/common/ai/docs/connections/pydantic_ai.rst +++ b/providers/common/ai/docs/connections/pydantic_ai.rst @@ -38,7 +38,7 @@ Model dedicated input in the connection form (via ``conn-fields``) and stores its value in ``extra["model"]``. - Examples: ``openai:gpt-5.3``, ``anthropic:claude-sonnet-4-20250514``, + Examples: ``openai:gpt-5.3``, ``anthropic:claude-sonnet-5``, ``bedrock:us.anthropic.claude-opus-4-6-v1:0``, ``google:gemini-2.0-flash`` The model can also be overridden at the hook/operator level via the diff --git a/providers/common/ai/docs/hitl_review.rst b/providers/common/ai/docs/hitl_review.rst index 26735622e9623..699d51c5e5676 100644 --- a/providers/common/ai/docs/hitl_review.rst +++ b/providers/common/ai/docs/hitl_review.rst @@ -131,7 +131,7 @@ instance page. Use the **HITL Review** extra link on the task instance, or navigate to ``/dags/{dag_id}/runs/{run_id}/tasks/{task_id}/plugin/hitl-review``. -**Example DAG** +**Example Dag** .. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent.py :language: python @@ -150,8 +150,8 @@ The plugin exposes a FastAPI app at ``/hitl-review``. Base URL: **Common query parameters** (where applicable): -- ``dag_id`` — DAG ID. -- ``run_id`` — DAG run ID. +- ``dag_id`` — Dag ID. +- ``run_id`` — Dag run ID. - ``task_id`` — Task ID. - ``map_index`` — Map index for mapped tasks. Use ``-1`` for non-mapped tasks or index for dynamic mapping. diff --git a/providers/common/ai/docs/hooks/index.rst b/providers/common/ai/docs/hooks/index.rst index 2786cd1b0ced5..13455cdf89622 100644 --- a/providers/common/ai/docs/hooks/index.rst +++ b/providers/common/ai/docs/hooks/index.rst @@ -47,6 +47,10 @@ Choosing a hook by default). For non-OpenAI vendors, pass a pre-built ``BaseEmbedding`` / ``LLM`` instance straight to the operator and bypass the hook. + * - :class:`~airflow.providers.common.ai.hooks.mcp.MCPHook` + - Backs ``MCPToolset`` (see :doc:`../toolsets`) for agent tasks that call + tools on a remote MCP server. Configure the connection via + :doc:`../connections/mcp`. Hook guides ----------- diff --git a/providers/common/ai/docs/hooks/langchain.rst b/providers/common/ai/docs/hooks/langchain.rst index e7ee2e95d7d1e..5485f39ab2a26 100644 --- a/providers/common/ai/docs/hooks/langchain.rst +++ b/providers/common/ai/docs/hooks/langchain.rst @@ -57,7 +57,7 @@ Any model identifier accepted by works out of the box. Common identifiers: - ``openai:gpt-4o``, ``openai:gpt-4o-mini`` -- requires ``langchain-openai`` -- ``anthropic:claude-3-7-sonnet`` -- requires ``langchain-anthropic`` +- ``anthropic:claude-sonnet-5`` -- requires ``langchain-anthropic`` - ``groq:llama-3.3-70b-versatile`` -- requires ``langchain-groq`` - ``mistralai:mistral-large-latest`` -- requires ``langchain-mistralai`` - ``ollama:llama3`` -- requires ``langchain-ollama`` (point ``host`` at the Ollama URL) diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index 6854ca4e948ad..0b221cada7bcf 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -125,11 +125,11 @@ attribute access (``result.field``). The declared ``output_type`` (and any ``BaseModel`` reachable from ``Union``/``Optional``/``list`` shapes) is registered for XCom deserialization by -the worker when it loads the DAG, before any task runs. The Pydantic class must +the worker when it loads the Dag, before any task runs. The Pydantic class must be defined at **module scope** and bound to an attribute matching its -``__name__``. Same-DAG downstream tasks need no configuration. The UI's XCom +``__name__``. Same-Dag downstream tasks need no configuration. The UI's XCom viewer renders the value via the ``stringify`` path (no configuration needed; -see the ``LLMOperator`` guide for the exact representation). Cross-DAG +see the ``LLMOperator`` guide for the exact representation). Cross-Dag ``xcom_pull`` consumers still need the class ``qualname`` added to ``[core] allowed_deserialization_classes``. @@ -289,9 +289,9 @@ invalidate cached responses -- clear the cache to force a fully fresh run. After the run, a single INFO summary line reports how many steps were replayed vs executed fresh. Per-step detail is available at DEBUG level. -The cache is scoped to a single task instance (DAG id, run id, task id, and +The cache is scoped to a single task instance (Dag id, run id, task id, and map index), so each run replays only its own steps. On Airflow >= 3.3 the cache -lives in the task state store and is removed when the DAG run is cleaned up; on +lives in the task state store and is removed when the Dag run is cleaned up; on Airflow < 3.3 it is a JSON file named ``{dag_id}_{task_id}_{run_id}.json`` (with ``_{map_index}`` appended for mapped tasks) under the configured ``durable_cache_path``. @@ -299,8 +299,8 @@ Airflow < 3.3 it is a JSON file named ``{dag_id}_{task_id}_{run_id}.json`` (with .. note:: Runs that fail permanently (exhaust all retries) leave their cached steps - behind. These do not affect future DAG runs (each run is scoped separately). - On Airflow >= 3.3 they are reclaimed when the DAG run is removed; on Airflow + behind. These do not affect future Dag runs (each run is scoped separately). + On Airflow >= 3.3 they are reclaimed when the Dag run is removed; on Airflow < 3.3 the orphaned JSON files consume storage until cleaned up, so add a lifecycle policy to the storage backend or remove them periodically. @@ -360,7 +360,7 @@ Capabilities compose with toolsets -- pydantic-ai merges tools from both. ``agent_params`` is a templated field, which Airflow serializes by calling ``str()`` on values it doesn't natively understand. Capability instances - are not yet round-trip-safe through DAG serialization, so the examples + are not yet round-trip-safe through Dag serialization, so the examples below construct them inside the ``@dag`` function -- not at module level. First-class ``capabilities=`` support on ``AgentOperator`` (with proper serializer hooks) is tracked as a follow-up. @@ -469,6 +469,29 @@ Parameters When set, the post-run transcript is pushed to XCom under the key ``message_history`` for the next run to resume. Default ``None`` (single-turn). See `Multi-turn Sessions`_. +- ``serialize_output``: If ``True`` and ``output_type`` is a Pydantic + ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via + ``model_dump()`` before being pushed to XCom. Default ``False`` -- the + Pydantic instance flows through XCom unchanged. Set to ``True`` when a + downstream consumer needs the dict shape. + +**HITL Review parameters** (requires the ``hitl_review`` plugin -- see +:doc:`../hitl_review` for the full review workflow): + +- ``enable_hitl_review``: When ``True``, the operator enters an iterative + review loop after the first generation. A human reviewer can approve, + reject, or request changes via the plugin's REST API at ``/hitl-review`` + or through the **HITL Review** extra link on the task instance. Default + ``False``. +- ``max_hitl_iterations``: Maximum outputs shown to the reviewer (1 = initial + output). When the reviewer requests changes at iteration >= this limit, the + task fails with ``HITLMaxIterationsError`` without calling the LLM. E.g. 5 + allows changes at iterations 1-4. Default ``5``. +- ``hitl_timeout``: Maximum wall-clock time to wait for all review rounds + combined. ``None`` means no timeout (the operator blocks until a terminal + action). +- ``hitl_poll_interval``: Seconds between XCom polls while waiting for a + human response. Default ``10``. Logging diff --git a/providers/common/ai/docs/operators/index.rst b/providers/common/ai/docs/operators/index.rst index 7eaf414fe3a63..e8b043bf37476 100644 --- a/providers/common/ai/docs/operators/index.rst +++ b/providers/common/ai/docs/operators/index.rst @@ -62,7 +62,7 @@ via an ``output_type`` Pydantic model. **LLMFileAnalysisOperator / @task.llm_file_analysis** — stateless, single-turn file analysis. Use this when the prompt should reason over file contents or multimodal attachments already chosen -by the DAG author. The operator resolves files via ``ObjectStoragePath`` and keeps the interaction +by the Dag author. The operator resolves files via ``ObjectStoragePath`` and keeps the interaction read-only. **AgentOperator / @task.agent** — multi-turn tool-calling loop. The model decides which tools to diff --git a/providers/common/ai/docs/operators/llamaindex_retrieval.rst b/providers/common/ai/docs/operators/llamaindex_retrieval.rst index 6e0793604abab..3f064a2b825c2 100644 --- a/providers/common/ai/docs/operators/llamaindex_retrieval.rst +++ b/providers/common/ai/docs/operators/llamaindex_retrieval.rst @@ -40,7 +40,7 @@ Basic usage :start-after: [START howto_hook_llamaindex_retrieve] :end-before: [END howto_hook_llamaindex_retrieve] -``query`` is templated, so DAG-run params, XCom, and Variables all flow +``query`` is templated, so Dag-run params, XCom, and Variables all flow through cleanly. Cloud-persisted indexes diff --git a/providers/common/ai/docs/operators/llm.rst b/providers/common/ai/docs/operators/llm.rst index 1f5a375ca9003..188aa089da8b7 100644 --- a/providers/common/ai/docs/operators/llm.rst +++ b/providers/common/ai/docs/operators/llm.rst @@ -51,7 +51,7 @@ so downstream tasks can type-hint the class directly The declared ``output_type`` (and any ``BaseModel`` reachable from ``Union``/``Optional``/``list`` shapes) is registered for XCom deserialization by -the worker when it loads the DAG, before any task runs -- so no edit to +the worker when it loads the Dag, before any task runs -- so no edit to ``[core] allowed_deserialization_classes`` is needed. The Pydantic class must be defined at **module scope** and bound to an attribute matching its ``__name__``; classes nested inside a function or ``@dag``-decorated body, parameterized @@ -69,9 +69,9 @@ warning at worker startup and the value fails to deserialize at the consumer. :start-after: [START howto_operator_llm_structured] :end-before: [END howto_operator_llm_structured] -Registration covers downstream tasks in the **same DAG**: every worker walks the -loaded DAG's tasks at startup and registers each declared class, so it also works -for mapped producers (``.expand(...)``) and for workers that load DAGs from a +Registration covers downstream tasks in the **same Dag**: every worker walks the +loaded Dag's tasks at startup and registers each declared class, so it also works +for mapped producers (``.expand(...)``) and for workers that load Dags from a cache that bypasses operator construction. The Airflow UI's XCom viewer renders Pydantic instances via the @@ -80,8 +80,8 @@ The Airflow UI's XCom viewer renders Pydantic instances via the allow-list. It is not pretty (no field-by-field rendering today), but the value shows up; no configuration is required. -The remaining gap is **cross-DAG** ``xcom_pull`` -- a task in a different DAG -that pulls this XCom only parses its own DAG file, not the producer's, so the +The remaining gap is **cross-Dag** ``xcom_pull`` -- a task in a different Dag +that pulls this XCom only parses its own Dag file, not the producer's, so the class is not auto-registered. Add the class qualified name to ``[core] allowed_deserialization_classes`` (or a glob that matches it) to make that pattern work. diff --git a/providers/common/ai/docs/operators/llm_branch.rst b/providers/common/ai/docs/operators/llm_branch.rst index 3c9e64908238d..42db970427a6d 100644 --- a/providers/common/ai/docs/operators/llm_branch.rst +++ b/providers/common/ai/docs/operators/llm_branch.rst @@ -24,7 +24,7 @@ Use :class:`~airflow.providers.common.ai.operators.llm_branch.LLMBranchOperator` for LLM-driven branching — where the LLM decides which downstream task(s) to execute. -The operator discovers downstream tasks automatically from the DAG topology +The operator discovers downstream tasks automatically from the Dag topology and presents them to the LLM as a constrained enum via pydantic-ai structured output. No text parsing or manual validation is needed. @@ -80,7 +80,7 @@ How It Works At execution time, the operator: -1. Reads ``self.downstream_task_ids`` from the DAG topology. +1. Reads ``self.downstream_task_ids`` from the Dag topology. 2. Creates a dynamic ``Enum`` with one member per downstream task ID. 3. Passes that enum as ``output_type`` to ``pydantic-ai``, constraining the LLM to valid task IDs only. diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst b/providers/common/ai/docs/operators/llm_file_analysis.rst index 5e38851ef865f..8b2733ab05acb 100644 --- a/providers/common/ai/docs/operators/llm_file_analysis.rst +++ b/providers/common/ai/docs/operators/llm_file_analysis.rst @@ -80,12 +80,12 @@ back from the LLM instead of a plain string. The model instance is pushed to XCom unchanged so downstream tasks can type-hint the class directly. The declared ``output_type`` (and any ``BaseModel`` reachable from ``Union``/``Optional``/``list`` shapes) is registered for deserialization by the -worker when it loads the DAG. Define the class at **module scope** and bind it to +worker when it loads the Dag. Define the class at **module scope** and bind it to an attribute matching its ``__name__``: nested-in-function and dynamically-built classes cannot be re-imported, so they are skipped at worker startup and fail to -deserialize at the consumer. Same-DAG downstream tasks need no configuration; the +deserialize at the consumer. Same-Dag downstream tasks need no configuration; the UI XCom viewer renders the value -via the ``stringify`` path (no configuration needed). Cross-DAG ``xcom_pull`` +via the ``stringify`` path (no configuration needed). Cross-Dag ``xcom_pull`` consumers still need the class ``qualname`` added to ``[core] allowed_deserialization_classes`` (see the ``LLMOperator`` guide for details). @@ -140,6 +140,15 @@ Parameters ``BaseModel`` for structured output. - ``agent_params``: Additional keyword arguments passed to the pydantic-ai ``Agent`` constructor (e.g. ``retries``, ``model_settings``). +- ``serialize_output``: If ``True`` and ``output_type`` is a Pydantic + ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via + ``model_dump()`` before being pushed to XCom. Default ``False`` -- the + Pydantic instance flows through XCom unchanged. Set to ``True`` when a + downstream consumer needs the dict shape. + +This operator also inherits ``LLMOperator``'s HITL review parameters -- +``require_approval``, ``approval_timeout``, and ``allow_modifications`` -- see +:doc:`llm` for details. Supported Formats ----------------- diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index 036bc8e5908e7..8456b1262aa3f 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -18,7 +18,8 @@ LLM Retry Policies =================== -.. versionadded:: 3.3.0 +.. note:: + Requires Airflow >= 3.3.0. The ``LLMRetryPolicy`` uses an LLM to classify task errors and make intelligent retry decisions. It works with any LLM provider supported by pydantic-ai @@ -40,7 +41,7 @@ Setup - **Connection Id**: ``pydanticai_default`` - **Connection Type**: ``Pydantic AI`` - **Password**: Your API key - - **Extra**: ``{"model": "anthropic:claude-haiku-4-5-20251001"}`` + - **Extra**: ``{"model": "anthropic:claude-haiku-4-5"}`` Usage ----- diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index d97a8fece2bf9..faa407d6676db 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -364,7 +364,7 @@ For prototyping or when you want full PydanticAI control, you can pass ) This works because PydanticAI's ``MCPToolset`` implements ``AbstractToolset``. -The tradeoff: URLs and credentials are hardcoded in DAG code instead of being +The tradeoff: URLs and credentials are hardcoded in Dag code instead of being managed through Airflow connections and secret backends. @@ -392,8 +392,8 @@ extra to use it: Each source is a local directory or a connection-resolved :class:`~airflow.providers.common.ai.skills.GitSkills`. Sources are resolved when -the agent enters the toolset, on the worker -- never while the DAG processor -parses the file -- so a Git token is never baked into the serialized DAG, and +the agent enters the toolset, on the worker -- never while the Dag processor +parses the file -- so a Git token is never baked into the serialized Dag, and cloned repositories are removed when the run ends. A local directory of ``SKILL.md`` bundles: @@ -425,7 +425,7 @@ need strict isolation. Skill bundles can contain scripts that the agent may run on the worker via the ``run_skill_script`` tool. For a remote source, anyone who can modify the - repository can introduce code that executes on your worker, outside DAG + repository can introduce code that executes on your worker, outside Dag review and versioning. Point ``GitSkills`` at a trusted repository, pin ``branch`` to a trusted ref, and treat skill contents as code that runs in your environment. @@ -544,8 +544,8 @@ before it runs. If no registered tool can read the environment, the filesystem, or other connections, the model cannot reach them, regardless of what the prompt instructs it to do. -This is what "untrusted" means in this context. The DAG file itself is -author-written and trusted, exactly like any other DAG. What is untrusted is +This is what "untrusted" means in this context. The Dag file itself is +author-written and trusted, exactly like any other Dag. What is untrusted is the model's *output*: the tool-call requests and text it generates. That output is confined to your registered tools and bounded by the tool-call budget. An agent cannot create a new connection, read another connection's credentials, or @@ -570,13 +570,13 @@ No single layer is sufficient — they work together. - What it does - What it does NOT do * - **Airflow Connections** - - Credentials are stored in Airflow's secret backend, never in DAG code. + - Credentials are stored in Airflow's secret backend, never in Dag code. The LLM agent cannot see API keys or database passwords. - Does not prevent the agent from using the connection to access data the connection has access to. * - **HookToolset: explicit allow-list** - Only methods listed in ``allowed_methods`` are exposed as tools. - Auto-discovery is not supported. Methods are validated at DAG parse + Auto-discovery is not supported. Methods are validated at Dag parse time. - Does not restrict what arguments the agent passes to allowed methods. * - **SQLToolset: read-only by default** @@ -719,7 +719,7 @@ Production Checklist Before deploying an agent task to production: 1. **Connection credentials**: Use Airflow's secret backend. Never hardcode - API keys in DAG files. + API keys in Dag files. 2. **Database permissions**: Create a dedicated database user with minimum required grants. Don't reuse the admin connection. 3. **Tool allow-list**: Review ``allowed_methods`` / ``allowed_tables``. The diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index 0d63e6f4202dd..e5c26e60a55a9 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -158,7 +158,7 @@ connection-types: conn-fields: model: label: Model - description: "Model in provider:name format (e.g. anthropic:claude-sonnet-4-20250514, openai:gpt-5)" + description: "Model in provider:name format (e.g. anthropic:claude-sonnet-5, openai:gpt-5)" schema: type: - string @@ -388,7 +388,7 @@ connection-types: description: > Chat model in provider:name format dispatched via langchain.chat_models.init_chat_model - (e.g. openai:gpt-4o, anthropic:claude-3-7-sonnet). + (e.g. openai:gpt-4o, anthropic:claude-sonnet-5). schema: type: - string diff --git a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py index 4f31d62e84e7b..89d96737d94fa 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py +++ b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py @@ -128,7 +128,7 @@ def get_provider_info(): "conn-fields": { "model": { "label": "Model", - "description": "Model in provider:name format (e.g. anthropic:claude-sonnet-4-20250514, openai:gpt-5)", + "description": "Model in provider:name format (e.g. anthropic:claude-sonnet-5, openai:gpt-5)", "schema": {"type": ["string", "null"]}, } }, @@ -305,7 +305,7 @@ def get_provider_info(): "conn-fields": { "model": { "label": "Chat Model", - "description": "Chat model in provider:name format dispatched via langchain.chat_models.init_chat_model (e.g. openai:gpt-4o, anthropic:claude-3-7-sonnet).\n", + "description": "Chat model in provider:name format dispatched via langchain.chat_models.init_chat_model (e.g. openai:gpt-4o, anthropic:claude-sonnet-5).\n", "schema": {"type": ["string", "null"]}, }, "embed_model": {