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 4eda1ee76ea Make Common AI example Dags reach the behaviour they show
(#74901)
4eda1ee76ea is described below
commit 4eda1ee76eab10f09f36a57005b1dd18eb18103c
Author: Kaxil Naik <[email protected]>
AuthorDate: Sat Oct 17 17:51:50 2026 +0100
Make Common AI example Dags reach the behaviour they show (#74901)
---
providers/common/ai/docs/frameworks/anthropic.rst | 4 +-
providers/common/ai/docs/retry_policies.rst | 4 +-
.../common/ai/example_dags/example_adk_agent.py | 8 ++--
.../ai/example_dags/example_agent_tool_approval.py | 2 +-
.../example_dags/example_anthropic_tool_runner.py | 7 +++-
.../ai/example_dags/example_llm_retry_policy.py | 19 +++++++--
.../ai/example_dags/example_strands_agent.py | 6 +--
.../airflow/providers/common/ai/policies/retry.py | 5 ++-
.../ai/tests/unit/common/ai/policies/test_retry.py | 45 ++++++++++++++++++----
9 files changed, 76 insertions(+), 24 deletions(-)
diff --git a/providers/common/ai/docs/frameworks/anthropic.rst
b/providers/common/ai/docs/frameworks/anthropic.rst
index 214ee136a8d..11b7bd35186 100644
--- a/providers/common/ai/docs/frameworks/anthropic.rst
+++ b/providers/common/ai/docs/frameworks/anthropic.rst
@@ -68,7 +68,9 @@ the runner still hands an exception from one of your own
tools to the model. A r
reaches ``max_iterations`` while the model is still calling tools has no
answer, so
``tools.run`` raises
:class:`~airflow.providers.common.ai.exceptions.ToolRunnerMaxIterationsError`
and the
-task fails rather than returning an empty result.
+task fails rather than returning an empty result. A refusal has no answer
either, but
+``tools.run`` returns it like any final message, so the task above checks for
+``stop_reason == "refusal"`` and raises, which fails the task.
To reach Claude on Amazon Bedrock, Claude Platform on AWS, Google Vertex AI or
Microsoft
Foundry, build the
diff --git a/providers/common/ai/docs/retry_policies.rst
b/providers/common/ai/docs/retry_policies.rst
index 183657e740b..e8eeab44a46 100644
--- a/providers/common/ai/docs/retry_policies.rst
+++ b/providers/common/ai/docs/retry_policies.rst
@@ -326,7 +326,7 @@ Escalating to an LLM
A classifier is cheap and fast, and a text model can reason about a failure it
has never seen a category for. ``fallback_policy`` puts one behind the other.
``snowflake_policy`` is the classifier policy from the previous section; the
-chain reuses its category table:
+chain reuses its category table with a stricter bar on ``permanent``:
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
:language: python
@@ -398,7 +398,7 @@ Beyond your ``instructions``, it sees only the exception's
class name and
message (or, with ``include_traceback=True``, the formatted traceback; either
after redaction and truncation), how many attempts are left,
and, under ``ClassifierRetryPolicy``, the category names and descriptions. The
prompt says
-``attempt {try_number} of {max_tries}``, so the model knows the limit and not
+``attempt 2 of 4`` on the second try of a task with ``retries=3``, so the
model knows the limit and not
just where it is right now; an instruction like "after two attempts treat an
expired token as ``auth`` rather than ``transient``" works because the model
can see which attempt this is.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_adk_agent.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_adk_agent.py
index fc4bcf888d9..aef76608f44 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_adk_agent.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_adk_agent.py
@@ -25,10 +25,12 @@ results, and with Airflow's secret masker applied to every
tool result.
Before running:
-1. Install ADK: ``pip install "google-adk>=2.9.1"``, without Airflow's
constraints file.
+1. Install ADK and this provider's ``anthropic`` and ``sql`` extras, without
Airflow's
+ constraints file:
+ ``pip install "google-adk>=2.9.1"
"apache-airflow-providers-common-ai[anthropic,sql]"``.
ADK caps OpenTelemetry at 1.42.1, below the version that file pins; Airflow
itself
- accepts it. The Anthropic model below also needs ``anthropic``, which the
- ``anthropic`` extra of this provider installs.
+ accepts it. The Anthropic model below needs ``anthropic`` and
``SQLToolset`` needs
+ ``sqlglot``.
2. Create a connection (``LLM_CONN_ID``, default ``anthropic_default``) whose
password is an Anthropic API key. Set its host only to route through a
gateway.
3. Create a database connection (``DB_CONN_ID``, default ``sql_default``) whose
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_tool_approval.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_tool_approval.py
index 4edcef5a9ca..2ca8f8a6e75 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_tool_approval.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_tool_approval.py
@@ -29,7 +29,7 @@ from airflow.providers.common.compat.sdk import dag
# [START howto_agent_tool_approval]
def lookup_order(order_id: int) -> str:
"""Look up an order."""
- return f"order {order_id}: $42, delivered"
+ return f"order {order_id}: $42, carrier marked it lost"
def refund_order(order_id: int) -> str:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_anthropic_tool_runner.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_anthropic_tool_runner.py
index cebaeab96ef..499697cc456 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_anthropic_tool_runner.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_anthropic_tool_runner.py
@@ -25,8 +25,8 @@ masker applied to every result.
Before running:
-1. Install this provider's ``anthropic`` extra:
- ``pip install "apache-airflow-providers-common-ai[anthropic]"``.
+1. Install this provider's ``anthropic`` and ``sql`` extras:
+ ``pip install "apache-airflow-providers-common-ai[anthropic,sql]"``.
2. Create a connection (``LLM_CONN_ID``, default ``anthropic_default``) whose
password is
an Anthropic API key. Set its host only to route through a gateway.
3. Create a database connection (``DB_CONN_ID``, default ``sql_default``)
whose hook is a
@@ -71,6 +71,9 @@ def example_anthropic_tool_runner():
max_iterations=10,
)
message = tools.run(runner)
+ # A refusal ends the run with no answer, so fail the task rather than
return "".
+ if message.stop_reason == "refusal":
+ raise RuntimeError(f"The model refused to answer: {question!r}")
return "".join(block.text for block in message.content if block.type
== "text")
ask_the_warehouse()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
index a2dd6a52444..8a9f17ce845 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
@@ -102,7 +102,8 @@ try:
# demands more certainty than the rest. Under a bar the answer is
discarded and the
# fallback rules, then the task's own retry settings, decide instead. The
bars come from
# a calibration run on jev-1.13.0: correct picks landed at 0.89 and above,
wrong ones
- # at 0.47 to 0.69, with one wrong ``permanent`` at 0.90 that no sensible
bar catches.
+ # at 0.47 to 0.69, with one wrong ``permanent`` at 0.90 that no bar
catches here without
+ # also discarding correct picks.
snowflake_policy = ClassifierRetryPolicy(
llm_conn_id="decision_default",
min_confidence=0.8,
@@ -136,11 +137,17 @@ try:
# [START howto_retry_policy_escalation]
# The three layers chained. The decision model answers the clear cases
cheaply. When it is
# unsure or unreachable, a text model reasons about the failure and
decides retry and delay
- # itself. When that model is unreachable too, the rules decide.
+ # itself. When that model is unreachable too, the rules decide. With a
text model to ask, a
+ # wrong ``permanent`` costs more than an escalation, so its bar goes up to
0.95. That catches
+ # the wrong ``permanent`` at 0.90 from the calibration run, and a correct
pick under 0.95
+ # costs one text-model call rather than the task.
escalating_policy = ClassifierRetryPolicy(
llm_conn_id="decision_default",
min_confidence=0.8,
- categories=snowflake_policy.categories,
+ categories={
+ **snowflake_policy.categories,
+ "permanent": replace(snowflake_policy.categories["permanent"],
min_confidence=0.95),
+ },
fallback_policy=LLMRetryPolicy(llm_conn_id="pydanticai_default",
timeout=30.0),
fallback_rules=[
RetryRule(exception=ConnectionError, action=RetryAction.RETRY,
retry_delay=timedelta(seconds=30)),
@@ -162,7 +169,11 @@ try:
@task(retries=3, retry_delay=timedelta(minutes=1),
retry_policy=escalating_policy)
def task_ambiguous():
- """Reads as more than one category; under the bar the text model
decides, then the rules."""
+ """
+ Names no symptom, so the classifier is unsure and the text model
decides.
+
+ jev-1.13.0 picks ``permanent`` for it at about 0.8 to 0.9, under
that category's 0.95 bar here.
+ """
raise RuntimeError("Query failed: an unexpected error occurred
while processing the request.")
task_warehouse_suspended()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_strands_agent.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_strands_agent.py
index faebd874dba..c7a5cb1f07b 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_strands_agent.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_strands_agent.py
@@ -26,9 +26,9 @@ Airflow's secret masker applied to every tool result.
Before running:
-1. Install Strands: ``pip install "strands-agents>=1.56.0"``. The Anthropic
model
- below also needs ``anthropic``, which the ``anthropic`` extra of this
provider
- installs.
+1. Install Strands and this provider's ``anthropic`` and ``sql`` extras:
+ ``pip install "strands-agents>=1.56.0"
"apache-airflow-providers-common-ai[anthropic,sql]"``.
+ The Anthropic model below needs ``anthropic`` and ``SQLToolset`` needs
``sqlglot``.
2. Create a connection (``LLM_CONN_ID``, default ``anthropic_default``) whose
password is an Anthropic API key. Set its host only to route through a
gateway.
3. Create a database connection (``DB_CONN_ID``, default ``sql_default``) whose
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
index 06768fe522c..637679b1177 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
@@ -335,8 +335,11 @@ class _ModelRetryPolicy(RetryPolicy):
text = f"{text[: self.max_exception_length]}... (truncated)"
if not self.include_traceback:
text = f"{type(exception).__name__}: {text}"
+ # ``max_tries`` is the last try number that may still retry (core adds
to it on a clear),
+ # so the final attempt is ``max_tries + 1``, as core logs "Starting
attempt N of M".
return (
- f"Classify this error from a data pipeline task (attempt
{try_number} of {max_tries}):\n\n{text}"
+ "Classify this error from a data pipeline task "
+ f"(attempt {try_number} of {max_tries + 1}):\n\n{text}"
)
def _run(
diff --git a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
index e4c25f3e76e..254f4105091 100644
--- a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
+++ b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
@@ -473,9 +473,9 @@ class TestClassification:
assert (first.action, second.action) == (RetryAction.RETRY,
RetryAction.FAIL)
prompts = [call.args[0] for call in agent.run_sync.call_args_list]
assert "ConnectionError: reset" in prompts[0]
- assert "attempt 1 of 3" in prompts[0]
+ assert "attempt 1 of 4" in prompts[0]
assert "PermissionError: expired" in prompts[1]
- assert "attempt 2 of 3" in prompts[1]
+ assert "attempt 2 of 4" in prompts[1]
@patch(HOOK, autospec=True)
def test_custom_instructions_forwarded_to_agent(self, mock_hook_cls):
@@ -661,7 +661,7 @@ class TestConfidenceGate:
assert decision.reason == "category=auth confidence=0.05 threshold=n/a
action=fail"
-PROMPT_HEADER = "Classify this error from a data pipeline task (attempt 2 of
4):\n\n"
+PROMPT_HEADER = "Classify this error from a data pipeline task (attempt 2 of
5):\n\n"
POLICY_CLASSES = [
pytest.param(ClassifierRetryPolicy, id="classifier"),
@@ -692,13 +692,15 @@ def _truncated_read(*, explicit: bool = True) ->
Exception:
)
-def _prompt_for(mock_hook_cls, policy_cls, exception, **kwargs) -> str:
+def _prompt_for(
+ mock_hook_cls, policy_cls, exception, *, try_number: int = 2, max_tries:
int = 4, **kwargs
+) -> str:
"""Evaluate ``exception`` under a ``policy_cls`` built with ``kwargs`` and
return the prompt the model got."""
if policy_cls is ClassifierRetryPolicy:
agent = _install(mock_hook_cls, _agent("data"))
else:
agent = _install(mock_hook_cls, _open_agent("data",
should_retry=False))
- policy_cls(llm_conn_id="test", **kwargs).evaluate(exception, try_number=2,
max_tries=4)
+ policy_cls(llm_conn_id="test", **kwargs).evaluate(exception,
try_number=try_number, max_tries=max_tries)
return agent.run_sync.call_args.args[0]
@@ -712,7 +714,36 @@ class TestPrompt:
prompt = agent.run_sync.call_args.args[0]
assert "ValueError: bad column type" in prompt
- assert "attempt 2 of 5" in prompt
+ assert "attempt 2 of 6" in prompt
+
+ @pytest.mark.parametrize(
+ ("try_number", "max_tries", "expected"),
+ [
+ pytest.param(1, 1, "attempt 1 of 2",
id="first-try-with-one-retry-left"),
+ pytest.param(2, 1, "attempt 2 of 2", id="last-try"),
+ pytest.param(3, 3, "attempt 3 of 4", id="after-clear"),
+ ],
+ )
+ @pytest.mark.parametrize("policy_cls", POLICY_CLASSES)
+ @patch(HOOK, autospec=True)
+ def test_prompt_counts_the_first_try_in_the_attempt_total(
+ self, mock_hook_cls, policy_cls, try_number, max_tries, expected
+ ):
+ """
+ ``max_tries`` is the last try number that may still retry, so the
final attempt is one more.
+
+ It equals ``retries`` on a fresh run; a clear adds the tries already
used, so a ``retries=1``
+ task cleared after two tries runs try 3 with ``max_tries=3``.
+ """
+ prompt = _prompt_for(
+ mock_hook_cls,
+ policy_cls,
+ ValueError("bad column type"),
+ try_number=try_number,
+ max_tries=max_tries,
+ )
+
+ assert f"({expected})" in prompt
@pytest.mark.enable_redact
@patch(HOOK, autospec=True)
@@ -728,7 +759,7 @@ class TestPrompt:
)
assert agent.run_sync.call_args.args[0] == (
- "Classify this error from a data pipeline task (attempt 1 of
3):\n\n"
+ "Classify this error from a data pipeline task (attempt 1 of
4):\n\n"
"ConnectionError: could not authenticate with password ***"
)
