This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f10a3de1136 Add a guide to developing and testing Common AI tasks
locally (#73602)
f10a3de1136 is described below
commit f10a3de113680203c72f732cbd7807cebd18443c
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Sep 23 12:35:22 2026 +0100
Add a guide to developing and testing Common AI tasks locally (#73602)
Covers iterating on an agent in a notebook with the same hook and
toolsets the task uses, testing an agent task offline with a scripted
FunctionModel and dag.test(), and why setting a connection's model to
"test" does not work for an agent with a SQLToolset. The agent operator
page gains a section on defining an agent once and reusing it across
tasks, and the quick start links to the new page.
---
providers/common/ai/docs/index.rst | 1 +
providers/common/ai/docs/local_development.rst | 154 +++++++++++++++++++++++++
providers/common/ai/docs/operators/agent.rst | 81 +++++++++++++
providers/common/ai/docs/quickstart.rst | 2 +
4 files changed, 238 insertions(+)
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index 247ce5c73ac..033593b19fa 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -94,6 +94,7 @@ waits for the result, use that vendor's provider.
Installation <installation>
Quick start <quickstart>
Core concepts <concepts>
+ Develop and test locally <local_development>
.. toctree::
:titlesonly:
diff --git a/providers/common/ai/docs/local_development.rst
b/providers/common/ai/docs/local_development.rst
new file mode 100644
index 00000000000..64884ebfc23
--- /dev/null
+++ b/providers/common/ai/docs/local_development.rst
@@ -0,0 +1,154 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/local_development:
+
+Develop and test AI tasks locally
+=================================
+
+Getting an agent's prompt and tools right takes many reruns. Do them in a
+notebook, with no Airflow running, then test the finished Dag without an API
key.
+
+Iterate in a notebook
+---------------------
+
+``AgentOperator`` builds its agent from a
+:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook` and the
+toolsets you pass it. Build the same agent yourself and call it directly. Hooks
+read connections from ``AIRFLOW_CONN_<ID>`` environment variables, so no
+scheduler, metadata database, or ``airflow db migrate`` is involved:
+
+.. code-block:: bash
+
+ export AIRFLOW_CONN_MY_LLM='{"conn_type": "pydanticai", "password":
"sk-...", "extra": {"model": "openai:gpt-5"}}'
+ export AIRFLOW_CONN_ORDERS_DB='{"conn_type": "sqlite", "host":
"/tmp/orders.db"}'
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ agent = PydanticAIHook.get_hook("my_llm").create_agent(
+ instructions="You answer questions about the orders database.",
+ toolsets=[SQLToolset(db_conn_id="orders_db",
allowed_tables=["orders"])],
+ )
+ result = agent.run_sync("How many orders are there?")
+ print(result.output)
+
+Change the instructions, tools, or prompt and rerun the cell. To talk to the
agent
+instead, pydantic-ai's ``agent.to_cli_sync()`` opens a chat in the terminal (it
+needs the ``pydantic-ai-slim[cli]`` extra).
+
+When the answers look right, copy the arguments into the task. ``instructions``
+becomes ``system_prompt``, the connection ID goes in ``llm_conn_id``,
``toolsets``
+and ``output_type`` keep their names, and any other agent argument, such as
+``retries``, goes in ``agent_params``.
+
+.. code-block:: python
+
+ # dags/orders_report.py
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+ from airflow.sdk import dag, task
+
+
+ @dag(schedule=None)
+ def orders_report():
+ @task.agent(
+ llm_conn_id="my_llm",
+ system_prompt="You answer questions about the orders database.",
+ toolsets=[SQLToolset(db_conn_id="orders_db",
allowed_tables=["orders"])],
+ )
+ def orders_question() -> str:
+ return "How many orders are there?"
+
+ orders_question()
+
+
+ orders_report()
+
+To keep prompt data on your machine while iterating, point ``my_llm`` at a
model
+you serve yourself; see :ref:`howto/self_hosted_models`.
+
+Test an agent task without calling a model
+------------------------------------------
+
+Script the model's replies with pydantic-ai's
+`FunctionModel <https://pydantic.dev/docs/ai/guides/testing/>`__ and patch
+``PydanticAIHook.get_conn``, which is where ``AgentOperator`` and
``@task.agent``
+get their model. ``dag.test()`` then runs the real task, including template
rendering,
+toolset calls against your test database and XCom, without network access:
+
+.. code-block:: python
+
+ # tests/test_orders_report.py
+ from unittest import mock
+
+ import pydantic_ai.models
+ from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart,
ToolReturnPart
+ from pydantic_ai.models.function import FunctionModel
+
+ from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+
+ from orders_report import orders_report # dags/orders_report.py from above
+
+ # Fail the test instead of calling a real model by accident.
+ pydantic_ai.models.ALLOW_MODEL_REQUESTS = False
+
+
+ def scripted_model(messages, info):
+ returns = [p for m in messages for p in m.parts if isinstance(p,
ToolReturnPart)]
+ if not returns:
+ return ModelResponse(parts=[ToolCallPart("query", {"sql": "SELECT
COUNT(*) AS n FROM orders"})])
+ return ModelResponse(parts=[TextPart(f"Result:
{returns[-1].content}")])
+
+
+ def test_report_counts_orders():
+ with mock.patch.object(
+ PydanticAIHook, "get_conn", autospec=True,
return_value=FunctionModel(scripted_model)
+ ):
+ dag_run = orders_report().test()
+
+ assert dag_run.state == "success"
+
+The first reply calls the ``query`` tool, which runs against the test database
+behind ``orders_db``; the second turns the tool result into the answer. The
task
+still looks up ``my_llm``, so define it in the test environment too; it needs
no
+key:
+
+.. code-block:: bash
+
+ export AIRFLOW_CONN_MY_LLM='{"conn_type": "pydanticai", "extra": {"model":
"openai:gpt-5"}}'
+
+``dag.test()`` needs a metadata database, so run ``airflow db migrate`` once
in the
+test environment. It also looks the Dag up in your Dags folder, so import the
Dag
+from its file instead of defining it in the test, and put the Dags folder on
the
+import path (pytest's ``pythonpath`` setting).
+
+Setting a ``pydanticai`` connection's model to ``test`` swaps in pydantic-ai's
+``TestModel`` instead, with no patching. It calls *every* tool once with
generated
+arguments, against your real connections. That suits ``@task.llm`` and agents
whose
+tools accept any input, but a ``SQLToolset`` rejects the generated SQL until
its
+retries run out, and the task fails.
+
+Check the answers themselves
+----------------------------
+
+The test above scripts the model's replies, so it can't tell you whether a
changed
+prompt still gives good answers. For that, run a set of cases with expected
+outcomes against the notebook agent from the first section, using a real model
and
+a library such as `pydantic-evals
<https://pydantic.dev/docs/ai/evals/evals/>`__,
+before you ship the change.
diff --git a/providers/common/ai/docs/operators/agent.rst
b/providers/common/ai/docs/operators/agent.rst
index 1dfe9daa663..f443150b5ac 100644
--- a/providers/common/ai/docs/operators/agent.rst
+++ b/providers/common/ai/docs/operators/agent.rst
@@ -160,6 +160,87 @@ Open the **Rendered Template** tab on the task instance to
see the
substituted ``system_prompt`` after Jinja fills in ``classify``'s XCom
values.
+.. _howto/operator:agent-reuse:
+
+Reuse one agent across tasks
+----------------------------
+
+When several tasks, or several Dags, run the same agent, define it once and
+import it. ``AgentOperator`` and ``@task.agent`` take the whole agent
definition
+as keyword arguments, so a dict in a module next to your Dags is enough. A task
+that needs something different overrides single keys.
+
+.. code-block:: python
+
+ # dags/shared_agents/__init__.py
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ ORDERS_ANALYST = {
+ "llm_conn_id": "pydanticai_default",
+ "system_prompt": "You are the orders analyst. Answer only from the
orders database.",
+ "toolsets": [SQLToolset(db_conn_id="orders_db",
allowed_tables=["orders"])],
+ "agent_params": {"name": "orders_analyst"},
+ }
+
+.. code-block:: python
+
+ # dags/orders.py
+ from shared_agents import ORDERS_ANALYST
+
+ from airflow.sdk import dag, task
+
+
+ @dag(schedule=None)
+ def orders():
+ @task.agent(**ORDERS_ANALYST)
+ def weekly_summary() -> str:
+ return "Summarize this week's orders."
+
+ @task.agent(**{**ORDERS_ANALYST, "system_prompt": "Answer in one
sentence."})
+ def one_liner() -> str:
+ return "How many orders are there?"
+
+ weekly_summary()
+ one_liner()
+
+
+ orders()
+
+When span export is on (see :doc:`../observability`), the ``name`` in
+``agent_params`` becomes the ``gen_ai.agent.name`` attribute on each agent
run's
+span, so traces from every task that uses the definition group under one agent.
+
+To keep the definition out of Python, for example to share it with a program
that
+does not run on Airflow, write a pydantic-ai
+`agent spec <https://pydantic.dev/docs/ai/core-concepts/agent-spec/>`__ file
and pass its path through
+``agent_params``. A model set on the connection wins over a ``model`` in the
+file, and ``system_prompt`` is added to the file's ``instructions``:
+
+.. code-block:: yaml
+
+ # dags/shared_agents/orders_analyst.yaml
+ name: orders_analyst
+ instructions: >
+ You are the orders analyst. Answer only from the orders database.
+ retries: 2
+
+.. code-block:: python
+
+ from pathlib import Path
+
+ AgentOperator(
+ task_id="orders_question",
+ llm_conn_id="pydanticai_default",
+ prompt="How many orders are there?",
+ agent_params={"spec_file": Path(__file__).parent / "shared_agents" /
"orders_analyst.yaml"},
+ )
+
+Build the path from ``__file__``: a relative path resolves against the worker's
+working directory, not the Dag file.
+
+With ``durable=True``, tools from capabilities declared in the spec file are
not
+replayed on retry; they run again. Pass tools you need replayed in
``toolsets=``.
+
Agent features
--------------
diff --git a/providers/common/ai/docs/quickstart.rst
b/providers/common/ai/docs/quickstart.rst
index 0ee55c3cd8e..a21393ed190 100644
--- a/providers/common/ai/docs/quickstart.rst
+++ b/providers/common/ai/docs/quickstart.rst
@@ -121,3 +121,5 @@ Where to go next
files, generating SQL, batch processing.
- :doc:`operators/agent` gives the model tools built from Airflow hooks, SQL
databases or
MCP servers, so it can act instead of only answering.
+- :doc:`local_development` shows how to iterate on an agent in a notebook and
test the
+ Dag without an API key.
