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 9e689545c94 Let LLMRetryPolicy run on a classifier model with
author-defined categories (#73450)
9e689545c94 is described below
commit 9e689545c94d200f853ceab7d3e4981da29fcc99
Author: Kaxil Naik <[email protected]>
AuthorDate: Mon Sep 21 19:09:04 2026 +0100
Let LLMRetryPolicy run on a classifier model with author-defined categories
(#73450)
The model now answers one question, which of the policy's categories the
failure belongs to, and reads each category's description from the output
schema. Whether the category is retried, after how long, and how sure the
model has to be come from the author's ``ErrorCategory`` table in the
worker. The ``should_retry``, ``suggested_delay_seconds`` and ``reasoning``
fields are gone: the first two were a lookup the prompt recited back to the
model, and a text field is what a classifier model refuses.
``min_confidence`` on the policy, and per category, discards an unsure
answer
and takes the existing fallback chain (``fallback_rules``, then the task's
own
retry behaviour). A bar with no confidence reported counts as unsure, so a
swap to a text model cannot silently switch the control off.
The ``Choices`` builder moves from the branch operator into the decision
utils
so both surfaces present described options the same way.
---
providers/common/ai/docs/changelog.rst | 27 +
providers/common/ai/docs/classifier_models.rst | 47 +-
providers/common/ai/docs/examples.rst | 4 +-
providers/common/ai/docs/retry_policies.rst | 314 ++++++--
.../ai/example_dags/example_llm_retry_policy.py | 94 ++-
.../providers/common/ai/operators/llm_branch.py | 55 +-
.../airflow/providers/common/ai/policies/retry.py | 328 ++++++--
.../airflow/providers/common/ai/utils/decision.py | 80 +-
.../unit/common/ai/operators/test_llm_branch.py | 6 +-
.../ai/tests/unit/common/ai/policies/test_retry.py | 859 ++++++++++++++++-----
.../ai/tests/unit/common/ai/utils/test_decision.py | 63 +-
11 files changed, 1473 insertions(+), 404 deletions(-)
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index 5fc707ae7c1..74b98db5408 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -25,6 +25,33 @@
Changelog
---------
+.. note::
+ ``LLMRetryPolicy`` now asks the model only which category a failure is;
whether that
+ category is retried, and after how long, comes from the policy's
``categories`` table
+ (``ErrorCategory(description, retry, delay, min_confidence)``), not from the
model.
+ ``ErrorClassification`` and its ``should_retry``,
``suggested_delay_seconds`` and
+ ``reasoning`` fields are removed, so a Dag file that imports the class fails
to parse,
+ and with it every Dag in that file. The new public names are
``ErrorCategory`` and
+ ``DEFAULT_CATEGORIES``.
+
+ A policy built with only ``llm_conn_id`` classifies into the same seven
categories with
+ the same retry/fail split and the same 60s/10s/30s delays. The delays are
now fixed by
+ the table rather than chosen by the model, so an error the model previously
answered
+ with its own delay now waits the category's. Custom ``instructions`` that
named a delay,
+ said "do NOT retry", or introduced category names outside the seven still
parse but no
+ longer steer anything: the model is constrained to ``categories``, so a name
of your own
+ is either mapped onto the nearest default or rejected by the schema and sent
to the
+ fallback path. The 0.9.0 guide's Snowflake example asked for ``rate_limit``
after 120s
+ and now gets the default 60s. Move each such rule into an ``ErrorCategory``
entry and
+ keep ``instructions`` for teaching the model your error strings; passing
custom
+ ``instructions`` without ``categories`` now raises a ``UserWarning`` at Dag
parse time
+ saying so.
+
+ The ``retry_reason`` written on a retry is now a generated line
+ (``category=... confidence=... threshold=... action=... delay=...``) rather
than the
+ model's prose, and a decision that came from ``fallback_rules`` has its
reason prefixed
+ with ``LLM classification not applied (<why>);``. See :doc:`retry_policies`.
+
.. note::
Configuring ``fallback_conn_ids`` on a connection (or the matching
operator/decorator
argument) changes what exception a task raises once every connection in the
chain fails:
diff --git a/providers/common/ai/docs/classifier_models.rst
b/providers/common/ai/docs/classifier_models.rst
index f53e36f6b2a..44fcdcf2ac3 100644
--- a/providers/common/ai/docs/classifier_models.rst
+++ b/providers/common/ai/docs/classifier_models.rst
@@ -102,8 +102,9 @@ Where it fits in this provider
- The downstream task ids are already presented to the model as a
constrained set of
choices, which is exactly the shape a classifier model answers. Setting
``model_id`` is the only change, as long as there are two or more
downstream tasks
- -- a one-option pick is refused. The other caveat is below: the
operator cannot see
- the confidence, so it will branch on an unsure answer without telling
you.
+ -- a one-option pick is refused. Describe each branch in ``branches``
and set a
+ ``decision_policy`` so an unsure pick goes to a person instead of
branching; see
+ :doc:`operators/llm_branch`.
* - :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` /
:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`
with a typed
``output_type``
@@ -112,12 +113,13 @@ Where it fits in this provider
becomes the question, and describe each option, which is what tells
them apart. An
option with no description is read from its name alone.
* - :doc:`LLMRetryPolicy <retry_policies>`
- - No
- - Three of ``ErrorClassification``'s four fields are shapes a classifier
model
- refuses: ``category`` and ``reasoning`` are bare ``str``, and
- ``suggested_delay_seconds`` is an unbounded ``int``. Only
``should_retry`` fits.
- Making the class answerable would need a ``Literal`` category, a rubric
delay, and
- no explanation at all.
+ - Yes
+ - The model names one of the policy's ``categories`` and nothing else;
retry or
+ fail, the delay and the confidence bar come from each category's entry
in the
+ worker. Set ``model_id`` and ``min_confidence``, and an unsure answer
goes to
+ ``fallback_rules`` and then the task's own retry behaviour, instead of
ending the
+ task on the model's say-so. This is the surface where
+ the model's speed and price matter most: it runs on every task failure.
* - Agents with toolsets
- Partly
- Which tool the text calls for is itself a pick, so a classifier model
can make it.
@@ -129,20 +131,21 @@ Where it fits in this provider
Reading the confidence
----------------------
-The confidence lives in ``provider_details`` on the model response, which no
operator
-exposes on XCom by default -- ``AgentOperator`` carries it inside the
``message_history``
-transcript when that is enabled. It is reported per output field, and a bare
output type is
-reported under ``"response"``. A bounded ``float`` field reports none at all,
because there
-the probability is the answer.
-
-This matters most for branching, because a branch is taken inside the operator
before any
-downstream task can look at it. Given two repair paths that both fit the
evidence, a
-classifier model splits its probability between them, and the branch is then
decided by
-whichever side of a near-even split came out ahead. The confidence reports
that faithfully;
-it is just not somewhere Airflow can act on yet, so a near-tie routes on a
margin the
-operator cannot show you. Treat ``LLMBranchOperator`` on a classifier model as
suited to
-branches whose options are genuinely distinct, and put anything closer than
that behind a
-task that reads the confidence itself:
+The confidence lives in ``provider_details`` on the model response. It is
reported per
+output field, and a bare output type is reported under ``"response"``. A
bounded ``float``
+field reports none at all, because there the probability is the answer.
+
+Two surfaces act on it for you.
:class:`~airflow.providers.common.ai.operators.llm_branch.LLMBranchOperator`
+and :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` take a
+``decision_policy`` whose ``min_confidence`` sends an unsure answer to a
person, or fails
+the task, before anything downstream runs on it, and record the confidence, the
+probabilities and the bar in the ``decision`` XCom (see
:doc:`operators/llm_branch`).
+:doc:`LLMRetryPolicy <retry_policies>` takes the same ``min_confidence`` and
hands an unsure
+answer to its deterministic fallback rules. In the branch operator and the
retry policy, a
+per-option bar lets the choice whose wrong pick costs most demand more
certainty than the rest.
+
+Outside those, read it yourself. ``AgentOperator`` carries it inside the
``message_history``
+transcript when that is enabled, and a hook-level call has it on the result:
.. code-block:: python
diff --git a/providers/common/ai/docs/examples.rst
b/providers/common/ai/docs/examples.rst
index 94a26471a0d..af69eb85e8d 100644
--- a/providers/common/ai/docs/examples.rst
+++ b/providers/common/ai/docs/examples.rst
@@ -152,7 +152,9 @@ Reliability
* - Guide
- What it shows
* - :doc:`retry_policies`
- - Classifying task failures with an LLM to decide retry, fail, or delay.
Source:
+ - Classifying task failures with an LLM into categories you define, then
deriving
+ retry, fail, or delay from the category; and the same on a classifier
model with a
+ confidence bar. Source:
`example_llm_retry_policy.py
<https://github.com/apache/airflow/blob/providers-common-ai/|version|/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py>`__.
* - :doc:`provider_fallback`
- Failing over to another vendor inside one task attempt, and drilling
the chain
diff --git a/providers/common/ai/docs/retry_policies.rst
b/providers/common/ai/docs/retry_policies.rst
index 8d1d5f14156..91d3fc868a7 100644
--- a/providers/common/ai/docs/retry_policies.rst
+++ b/providers/common/ai/docs/retry_policies.rst
@@ -21,9 +21,11 @@ LLM Retry Policies
.. 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
-(OpenAI, Anthropic, Bedrock, Vertex, Ollama, etc.).
+``LLMRetryPolicy`` asks a model which kind of failure a task hit, then retries
or
+fails the task according to a table you own. It works with any LLM provider
+supported by pydantic-ai (OpenAI, Anthropic, Bedrock, Vertex, Ollama, etc.),
and
+with a classifier model such as TypeSafe's Jev, which answers the question in
a few
+hundred milliseconds and reports how sure it is.
For the core retry policy concepts, see
:doc:`apache-airflow:core-concepts/tasks`.
If the task also needs to survive a worker crash without losing its progress,
@@ -57,7 +59,7 @@ Usage
llm_policy = LLMRetryPolicy(
llm_conn_id="pydanticai_default",
timeout=30.0, # max seconds to wait for LLM response
- fallback_rules=[ # used when LLM call fails
+ fallback_rules=[ # used when the LLM call fails or the answer is
under its bar
RetryRule(exception=ConnectionError, action=RetryAction.RETRY,
retry_delay=timedelta(seconds=10)),
RetryRule(exception=PermissionError, action=RetryAction.FAIL),
],
@@ -76,10 +78,16 @@ When a task fails, ``LLMRetryPolicy``:
is first masked through Airflow's secrets masker (see ``redactor`` below)
and truncated to ``max_exception_length`` characters before it is added
to the prompt.
-2. The LLM classifies the error into a category (``rate_limit``, ``auth``,
- ``network``, ``data``, ``resource``, ``transient``, ``permanent``)
-3. Based on the classification, returns RETRY (with a suggested delay) or FAIL
-4. The classification reason is logged in the task logs
+2. The model picks one of the policy's ``categories``. It sees each category's
+ name and description in the output schema and nothing else about it, and it
+ cannot answer with a name outside the set.
+3. The policy looks the category up in the same table and returns RETRY, with
+ that category's delay, or FAIL. If the policy has a confidence bar and the
+ model's answer is under it, the answer is discarded instead (see
+ `Confidence`_ below).
+4. The decision is logged in the task logs and, on a RETRY, written to the task
+ instance's ``retry_reason``, as one line such as
+ ``category=network confidence=0.91 threshold=0.60 action=retry delay=10s``.
This classification call is a separate LLM request, made by ``LLMRetryPolicy``
itself rather than by an operator -- it is not subject to an operator's
@@ -87,14 +95,176 @@ itself rather than by an operator -- it is not subject to
an operator's
of any cost cap configured on the failing task. It is bounded by ``timeout``
and ``max_exception_length``, but not by a cost limit.
-If the LLM call fails (provider down, timeout, bad credentials), the policy
-falls back to ``fallback_rules`` if configured, or to the task's standard
-retry behaviour.
+If the LLM call fails (provider down, timeout, bad credentials), or the model
+cannot produce one of the categories even after pydantic-ai re-prompts it, or
+its answer is under the confidence bar, the policy falls back to
+``fallback_rules`` if configured, or to the task's standard retry behaviour.
+The category's action and delay are not applied in that case, and the
decision's
+reason starts with ``LLM classification not applied (model_error)``,
+``(below_threshold)`` or ``(missing_confidence)`` so a ``retry_reason`` read
later
+is not mistaken for a classifier decision or a plain rule match.
This policy decides *between* attempts. Failing over to another vendor *within*
an attempt is a separate mechanism on the connection — see
:doc:`provider_fallback`, which also sets out how the two layers compose.
+Categories
+----------
+
+``categories`` maps a category name to an
+:class:`~airflow.providers.common.ai.policies.retry.ErrorCategory`: what
+failures belong there (the ``description`` the model reads), whether it is
+retried, after what ``delay``, and how sure the model has to be
+(``min_confidence``, covered below). Everything the model is told about a
+category, and everything the policy does with it, sits in that one entry, so
+the two cannot drift apart.
+
+The default,
+:data:`~airflow.providers.common.ai.policies.retry.DEFAULT_CATEGORIES`, is an
+example taxonomy, not the taxonomy:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 20 55 25
+
+ * - Category
+ - What belongs there
+ - Action
+ * - ``rate_limit``
+ - API throttling or a quota exceeded.
+ - Retry after 60s
+ * - ``network``
+ - Transient connectivity issue: connection reset, DNS, TLS handshake.
+ - Retry after 10s
+ * - ``transient``
+ - Temporary issue likely to resolve on its own.
+ - Retry after 30s
+ * - ``auth``
+ - Credentials invalid, expired, or missing permissions.
+ - Fail
+ * - ``data``
+ - Schema validation, type mismatch, or bad input data.
+ - Fail
+ * - ``resource``
+ - Resource not found or unavailable, such as a missing table or bucket.
+ - Fail
+ * - ``permanent``
+ - Problem that will not resolve without a code or configuration change.
+ - Fail
+
+Pass your own to change any of it. It **replaces** the default rather than
+merging into it, so include every category you want offered -- or spread the
+default and edit what you need:
+
+.. code-block:: python
+
+ from dataclasses import replace
+ from datetime import timedelta
+
+ from airflow.providers.common.ai.policies.retry import DEFAULT_CATEGORIES,
ErrorCategory, LLMRetryPolicy
+
+ LLMRetryPolicy(
+ llm_conn_id="pydanticai_default",
+ categories={
+ **DEFAULT_CATEGORIES,
+ # The table is created upstream, so a missing one is worth another
look.
+ "resource": replace(DEFAULT_CATEGORIES["resource"], retry=True,
delay=timedelta(minutes=5)),
+ },
+ )
+
+Or write a taxonomy for your stack. A Snowflake pipeline has kinds of failure
the
+seven defaults can only approximate, and naming them is what lets the policy
act
+on them differently:
+
+.. code-block:: python
+
+ snowflake_policy = LLMRetryPolicy(
+ llm_conn_id="pydanticai_default",
+ categories={
+ "queued": ErrorCategory(
+ "Statement queued or a concurrency limit reached; the
warehouse is busy.",
+ delay=timedelta(seconds=120),
+ ),
+ "warehouse_suspended": ErrorCategory(
+ "The warehouse is suspended and will auto-resume.",
delay=timedelta(seconds=30)
+ ),
+ "token_expired": ErrorCategory(
+ "A JWT or session token expired; the token rotates on its
own.", delay=timedelta(seconds=30)
+ ),
+ "schema_drift": ErrorCategory(
+ "A referenced column, table or view does not exist; a person
has to fix the schema.",
+ retry=False,
+ ),
+ },
+ )
+
+Write descriptions as the boundary between categories: what belongs here and
what
+does not. That is the whole of what the model reads about a category; the name
+alone gives it very little, and two categories with similar names and no
+descriptions are indistinguishable to it. A pick is relative: the model chooses
+the best fit among the categories offered, not whether any fits. If "none of
+these" or "not enough in the message to tell" is a real outcome for your
+pipeline, add a category for it (``retry=True`` with no delay keeps the task's
+own behaviour) rather than hoping the model refuses. When you change a
+description or the set of categories, treat the confidence values you measured
+before as stale: the distribution the model returns is over the options it was
+given.
+
+A ``delay`` is ``timedelta``, not seconds, and is used as-is: a task's own
+``max_retry_delay`` does not clamp it. ``delay=None`` (the default) means no
+override, so the task's own ``retry_delay`` / ``retry_exponential_backoff`` /
+``max_retry_delay`` apply instead (see
:doc:`apache-airflow:core-concepts/tasks`).
+``retry=False`` ends the task straight away even when attempts were left, so a
+wrong classification into such a category costs the task the retries it would
+otherwise have had; that is what the confidence bar below is for. A bare int
+delay, a negative delay, a delay on a ``retry=False`` category, an empty
+description, or fewer than two categories raises when the policy is
+constructed, at Dag parse time, rather than on the first task failure.
+
+Confidence
+----------
+
+A classifier model reports how sure it is of its answer. ``min_confidence`` is
+the bar that answer needs for the policy to act on it; under the bar the policy
+discards the answer and takes the same path it takes when the model call fails:
+``fallback_rules`` if one matches, otherwise the task's own retry behaviour. It
+does not substitute a delay of its own.
+
+Each category can carry its own bar. The stakes differ: a wrong ``transient``
+costs one more attempt, while a wrong ``permanent`` costs the task every retry
it
+had left, so the category that ends the task deserves the higher bar.
+``jev_default`` is the classifier-model connection from
:doc:`classifier_models`.
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
+ :language: python
+ :start-after: [START howto_retry_policy_classifier]
+ :end-before: [END howto_retry_policy_classifier]
+
+A category bar needs a policy bar to inherit from; setting one without the
other
+raises at construction. With no bar at all the policy acts on every answer, as
+it always has, and the confidence is still logged.
+
+**A missing confidence counts as an unsure one.** A text model reports no
+confidence, and so does a response whose metadata was dropped along the way.
+With no bar configured that changes nothing. With a bar configured, every such
+answer is discarded and the fallback path decides, so swapping the connection
+from a classifier model to a text model does not silently switch off a control
+you set on purpose. To run a text model, remove the bar.
+
+The confidence is a statistic on the shape of the probability distribution the
+model returned: concentrated on one category is high, spread out is low. It is
+not the probability that the answer is correct. Pick the bar from the
confidence
+values your own failures produce: run the policy with no bar first, read the
+logged ``confidence=`` values per category, and set the bar where the wrong
+answers start. For orientation, a calibration run of the two example taxonomies
+on ``jev-1.13.0`` over 31 realistic exception messages put every correct pick
at
+0.89 or above and three of the four wrong picks between 0.47 and 0.69; the
+fourth wrong pick was a ``permanent`` at 0.90, which is why a bar reduces wrong
+actions and does not eliminate them. Pin the model version
+(``typesafe:jev-1.13.0``, not ``jev-latest``): a bar tuned against one release
+is not guaranteed to mean the same thing after the next. See
+:doc:`classifier_models` for what these models answer well and badly.
+
When the connection also carries a fallback chain
--------------------------------------------------
@@ -128,66 +298,60 @@ the chain is in play.
What the model can and cannot do
--------------------------------
-The model answers two questions: retry or not, and how long to wait. It is
+The model answers one question: which kind of failure is this. It is
given no tools and there is no way to attach any, so it cannot run code, call
an API, read a connection, or reach your data. Beyond your ``instructions``,
it sees only the exception's class name, the exception message (after
-redaction and truncation), and how many attempts are left. The prompt says
-``attempt {try_number} of {max_tries}``, so the model knows the limit, not
-just where it is right now -- that is what makes an instruction like "retry
-once, then stop" (see the Snowflake example below) actually work. It returns
-four fields: ``category``, ``should_retry``, ``suggested_delay_seconds``, and
-``reasoning``. Only ``should_retry`` and ``suggested_delay_seconds`` affect
-the run.
-
-``category`` and ``reasoning`` are only recorded on a RETRY. They are written
-to the task instance's ``retry_reason`` (truncated to 500 characters, see
-below), then cleared once the next attempt starts running. On a FAIL they are
-not written anywhere -- they only show up in the task log.
-
-Two limits are worth knowing about:
-
-* RETRY cannot give a task more attempts than ``retries`` allows. FAIL,
though, ends the task
- straight away even when attempts were left, so a wrong classification costs
- the task the retries it would otherwise have had.
-* A positive ``suggested_delay_seconds`` is used as returned. There is no
- upper limit -- a task's own ``max_retry_delay`` does not clamp it. But 0 or
- a negative value is not used as a delay at all: it is treated the same as
- no delay, so the task's own ``retry_delay`` / ``retry_exponential_backoff``
- / ``max_retry_delay`` apply instead (see
- :doc:`apache-airflow:core-concepts/tasks`). A model told to retry
- immediately can still wait out the task's default delay. If particular
- delays matter to you, state them in ``instructions`` as the examples below
- do.
+redaction and truncation), how many attempts are left, and the category names
+and descriptions. The prompt says ``attempt {try_number} of {max_tries}``, so
+the model knows the limit and not just where it is right now. That only moves
+the category -- 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.
+
+It returns the category name and nothing else. It does not decide whether to
+retry, it does not choose the delay, and it does not explain itself -- the
+first two come from the category's entry in the worker process, and the
+explanation is the generated line in the task log, which says what mattered
+(the category, the confidence, the bar, the action) rather than what the model
+felt like saying. A model cannot return a category the policy does not
+recognize, and it cannot return a category paired with an action that
+contradicts it.
+
+The decision line is only recorded on a RETRY. It is written to the task
+instance's ``retry_reason`` (truncated to 500 characters), then cleared once
the
+next attempt starts running. On a FAIL it is not written anywhere -- it only
+shows up in the task log.
+
+RETRY cannot give a task more attempts than ``retries`` allows. FAIL ends the
+task straight away even when attempts were left.
Custom instructions
-------------------
-The default classifier handles generic categories. For domain-specific
-behaviour, override ``instructions`` to inject your own taxonomy. The LLM still
-returns an
:class:`~airflow.providers.common.ai.policies.retry.ErrorClassification`
-(``category``, ``should_retry``, ``suggested_delay_seconds``, ``reasoning``)
--- only the prompt changes.
+The default instructions say only that the model is classifying a failed
+pipeline task and should pick the best-fitting category; the categories and
+their meanings travel in the output schema, from ``categories``. Override
+``instructions`` to teach the model your stack's error strings when a
+description is not enough on its own. Pass ``categories`` alongside: custom
+``instructions`` with the default table raises a ``UserWarning`` at Dag parse
+time, because a prompt that names categories or delays of its own no longer
+defines either.
.. code-block:: python
SNOWFLAKE_INSTRUCTIONS = (
- "You are an error classifier for Snowflake-backed data pipelines. "
- "Classify the error into one of: rate_limit, auth, network, data, "
- "transient, permanent.\n\n"
- "Snowflake-specific guidance:\n"
- "- 'Statement queued' or 'concurrency limit' -> rate_limit, retry
after 120s\n"
- "- 'JWT token expired' -> transient (token rotates), retry after 30s\n"
- "- 'Authentication token has expired' AFTER multiple retries -> auth,
do NOT retry\n"
- "- 'Column does not exist' -> data, do NOT retry (schema drift needs
human fix)\n"
- "- 'Warehouse suspended' -> transient, retry after 30s
(auto-resume)\n\n"
- "Set suggested_delay_seconds based on the error type. "
- "Set 0 for errors that should not retry."
+ "You are classifying failures from Snowflake-backed data pipelines.\n"
+ "- 'Statement queued' or 'concurrency limit' -> queued\n"
+ "- '000606' or 'is suspended' -> warehouse_suspended\n"
+ "- 'JWT token' or 'session token' with 'expired' -> token_expired\n"
+ "- '002003' or 'does not exist' -> schema_drift\n"
)
snowflake_policy = LLMRetryPolicy(
llm_conn_id="pydanticai_default",
instructions=SNOWFLAKE_INSTRUCTIONS,
+ categories={...}, # the same names the instructions use
fallback_rules=[
RetryRule(
exception=ConnectionError,
@@ -203,13 +367,20 @@ returns an
:class:`~airflow.providers.common.ai.policies.retry.ErrorClassificati
When writing custom instructions:
-- The LLM must return the same ``ErrorClassification`` schema (``category``,
- ``should_retry``, ``suggested_delay_seconds``, ``reasoning``). Mention the
- fields explicitly so the model fills them.
-- Be concrete with examples (``"'Warehouse suspended' -> transient"``) rather
- than vague rules ("treat warehouse issues as recoverable").
-- ``retry_reason`` is truncated to 500 chars in the audit log -- keep
- ``reasoning`` outputs concise.
+- Use the category names from ``categories`` as-is. The model is constrained to
+ them, so a name you invent cannot come back. A model that insists on one
+ anyway is re-prompted once by pydantic-ai and then gives up, which lands the
+ task on ``fallback_rules`` or on its own retry behaviour, having billed two
+ calls. Offering a name the schema rejects is therefore worse than offering
+ none.
+- Be concrete with examples (``"'Warehouse suspended' ->
warehouse_suspended"``)
+ rather than vague rules ("treat warehouse issues as recoverable").
+- Do not spell out delays or "do NOT retry" instructions. The model does not
+ decide either one, and telling it to only spends tokens. Put them in the
+ category instead.
+- A classifier model sends ``instructions`` as the question it scores the
+ exception text against, not as rules it follows step by step, so a long
rubric
+ buys less there than a better description on each category does.
Parameters
----------
@@ -229,13 +400,28 @@ Parameters
- Override the model from the connection (e.g., ``"openai:gpt-4o-mini"``).
* - ``instructions``
- (built-in)
- - Custom system prompt for error classification.
+ - Custom system prompt for error classification. Teaches the model your
+ error strings; the categories themselves come from ``categories``.
* - ``fallback_rules``
- None
- - List of ``RetryRule`` objects used when the LLM call fails.
+ - List of ``RetryRule`` objects used when the LLM call fails or the answer
+ is under its confidence bar.
* - ``timeout``
- 30.0
- Max seconds to wait for the LLM response before falling back.
+ * - ``categories``
+ - ``DEFAULT_CATEGORIES``
+ - Mapping of category name to ``ErrorCategory(description, retry, delay,
+ min_confidence)``: what the model chooses between and what the policy
+ does with each answer. **Replaces** the default mapping rather than
+ merging into it. At least two categories; validated at construction.
+ * - ``min_confidence``
+ - None
+ - The confidence, from 0 to 1, the model's answer needs for the policy to
+ act on it. Under the bar, or with a bar set and no confidence reported,
+ the answer is discarded and ``fallback_rules`` then the task's own retry
+ behaviour apply. A category's own ``min_confidence`` overrides it for
+ that category.
* - ``redactor``
- None (uses ``redact_registered_secrets``)
- Callable ``(str) -> str`` applied to the exception's string
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 bdd2528dd28..810b9e1950f 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
@@ -17,23 +17,28 @@
"""
Example DAG demonstrating LLM-powered retry policies.
-Uses an LLM (via PydanticAIHook) to classify errors and decide whether
-to retry, fail immediately, or retry with a custom delay.
+The model names the kind of failure from the categories the policy offers it,
each with a
+description. Whether that category is retried, after how long, and how sure
the model has
+to be all come from the policy's category table in the worker.
Prerequisites:
- Connection ``pydanticai_default`` with ``conn_type='pydanticai'``,
``password=<API key>``, ``extra='{"model":
"anthropic:claude-haiku-4-5-20251001"}'``
- ``pip install apache-airflow-providers-common-ai[anthropic]``
+ - For the classifier-model Dag: connection ``jev_default`` with
+ ``extra='{"model": "typesafe:jev-1.13.0"}'`` and
+ ``pip install 'apache-airflow-providers-common-ai[typesafe]'``
"""
from __future__ import annotations
+from dataclasses import replace
from datetime import timedelta
from airflow.providers.common.compat.sdk import dag, task
try:
- from airflow.providers.common.ai.policies.retry import LLMRetryPolicy
+ from airflow.providers.common.ai.policies.retry import DEFAULT_CATEGORIES,
ErrorCategory, LLMRetryPolicy
from airflow.sdk.definitions.retry_policy import RetryAction, RetryRule
llm_policy = LLMRetryPolicy(
@@ -45,28 +50,105 @@ try:
],
)
+ # The default table fails ``resource``, on the grounds that a missing
table needs a
+ # human. Here the table is created upstream, so it is worth one more look.
+ patient_policy = LLMRetryPolicy(
+ llm_conn_id="pydanticai_default",
+ categories={
+ **DEFAULT_CATEGORIES,
+ "resource": replace(DEFAULT_CATEGORIES["resource"], retry=True,
delay=timedelta(minutes=5)),
+ },
+ )
+
@dag(catchup=False, tags=["example", "retry_policy", "llm"])
def example_llm_retry_policy():
@task(retries=3, retry_delay=timedelta(minutes=1),
retry_policy=llm_policy)
def task_auth_error():
- """LLM should classify as auth -> FAIL immediately."""
+ """Should classify as ``auth``, which the table fails -> FAIL
immediately."""
raise PermissionError("403 Forbidden: API key expired for service
account [email protected]")
@task(retries=3, retry_delay=timedelta(minutes=1),
retry_policy=llm_policy)
def task_rate_limit():
- """LLM should classify as rate_limit -> RETRY with ~60s delay."""
+ """Should classify as ``rate_limit``, which the table retries
after 60s."""
raise RuntimeError("429 Too Many Requests: Rate limit exceeded.
Retry after 60 seconds.")
@task(retries=3, retry_delay=timedelta(minutes=1),
retry_policy=llm_policy)
def task_data_error():
- """LLM should classify as data -> FAIL immediately."""
+ """Should classify as ``data``, which the table fails -> FAIL
immediately."""
raise ValueError("Column 'user_id' expected type INT but got
STRING in row 42.")
+ @task(retries=3, retry_delay=timedelta(minutes=1),
retry_policy=patient_policy)
+ def task_missing_table():
+ """Should classify as ``resource``, which this policy's table
retries after 5m."""
+ raise FileNotFoundError("Table 'analytics.daily_orders' does not
exist")
+
task_auth_error()
task_rate_limit()
task_data_error()
+ task_missing_table()
example_llm_retry_policy()
+
+ # [START howto_retry_policy_classifier]
+ # A classifier model answers the same question in a few hundred
milliseconds and reports
+ # how sure it is. The categories are this pipeline's own; their
descriptions are what the
+ # model reads. ``permanent`` ends the task and costs it every retry it had
left, so it
+ # 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.
+ snowflake_policy = LLMRetryPolicy(
+ llm_conn_id="jev_default",
+ min_confidence=0.8,
+ categories={
+ "queued": ErrorCategory(
+ "Statement queued or a concurrency limit reached; the
warehouse is busy.",
+ delay=timedelta(seconds=120),
+ ),
+ "warehouse_suspended": ErrorCategory(
+ "The warehouse is suspended and will auto-resume.",
delay=timedelta(seconds=30)
+ ),
+ "token_expired": ErrorCategory(
+ "A JWT or session token expired; the token rotates on its
own.", delay=timedelta(seconds=30)
+ ),
+ "schema_drift": ErrorCategory(
+ "A referenced column, table or view does not exist; a person
has to fix the schema.",
+ retry=False,
+ ),
+ "permanent": ErrorCategory(
+ "A code or configuration error that will fail identically on
every attempt.",
+ retry=False,
+ min_confidence=0.9,
+ ),
+ },
+ fallback_rules=[
+ RetryRule(exception=ConnectionError, action=RetryAction.RETRY,
retry_delay=timedelta(seconds=30)),
+ ],
+ )
+ # [END howto_retry_policy_classifier]
+
+ @dag(catchup=False, tags=["example", "retry_policy", "llm", "classifier"])
+ def example_llm_retry_policy_classifier():
+ @task(retries=3, retry_delay=timedelta(minutes=1),
retry_policy=snowflake_policy)
+ def task_warehouse_suspended():
+ """Should classify as ``warehouse_suspended`` -> RETRY after
30s."""
+ raise RuntimeError("000606 (57P03): Warehouse 'ANALYTICS_WH' is
suspended; auto-resume pending.")
+
+ @task(retries=3, retry_delay=timedelta(minutes=1),
retry_policy=snowflake_policy)
+ def task_schema_drift():
+ """Should classify as ``schema_drift`` -> FAIL immediately."""
+ raise RuntimeError("002003 (42S02): SQL compilation error: Object
'ORDERS_V2' does not exist.")
+
+ @task(retries=3, retry_delay=timedelta(minutes=1),
retry_policy=snowflake_policy)
+ def task_ambiguous():
+ """Reads as more than one category; under the bar the task's own
retry settings apply."""
+ raise RuntimeError("Query failed: an unexpected error occurred
while processing the request.")
+
+ task_warehouse_suspended()
+ task_schema_drift()
+ task_ambiguous()
+
+ example_llm_retry_policy_classifier()
except ImportError:
# RetryPolicy requires Airflow 3.3+; example DAG is skipped on older
versions.
pass
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
index 7c9cd655a47..7f6a271a082 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
@@ -20,22 +20,19 @@ from __future__ import annotations
import json
from collections.abc import Iterable, Mapping, Sequence
-from enum import Enum
from typing import TYPE_CHECKING, Any
-try:
- from pydantic_ai import Choice, Choices # type: ignore[attr-defined]
-except ImportError: # pydantic-ai < 2.46.0: build the same schema from an
Enum instead
- Choice = Choices = None # type: ignore[assignment,misc]
-
from airflow.providers.common.ai.operators.llm import LLMOperator
from airflow.providers.common.ai.policies.decision import BranchOption,
DecisionPolicy
from airflow.providers.common.ai.utils.decision import (
+ BARE_OUTPUT_FIELD,
ModelConfidence,
check_uncertain_action,
decision_record,
describe_confidence,
+ described_choices,
initial_decided_by,
+ picked_key,
policy_record,
review_reason,
threshold_for,
@@ -61,14 +58,10 @@ def _branch_choices(
Build the type the model picks a branch from: one option per downstream
task ID.
Sorted so every worker sends the model the same option order:
``downstream_task_ids``
- is a set, and set order follows string hashing, which differs between
processes.
-
- With a description on any branch the options render as ``anyOf`` of
``{const,
- description}`` instead of a bare ``enum`` list. That is the one JSON
Schema shape that
- carries a description per value, and it is what both a text model's tool
schema and
- pydantic-ai's TypeSafe adapter read an option's meaning from. On
pydantic-ai 2.46+ the
- type is its ``Choices``; before that, an ``Enum`` whose schema hook emits
the same shape.
- Either way the model has to answer with one of the task IDs.
+ is a set, and set order follows string hashing, which differs between
processes. The
+ option type itself comes from
:func:`~airflow.providers.common.ai.utils.decision.described_choices`,
+ so the model has to answer with one of the task IDs and reads each one's
description
+ from the schema.
"""
task_ids = sorted(downstream_task_ids)
unknown = sorted(set(configured) - set(task_ids))
@@ -77,33 +70,7 @@ def _branch_choices(
f"branches for {task_id!r} names {unknown}, which are not
downstream tasks. "
f"Downstream tasks: {task_ids}."
)
- if Choices is not None:
- if not descriptions:
- return Choices(task_ids, name="DownstreamTasks")
- return Choices({name: Choice(descriptions.get(name)) for name in
task_ids}, name="DownstreamTasks")
-
- enum_cls: type[Enum] = Enum("DownstreamTasks", {name: name for name in
task_ids}) # type: ignore[misc]
- if not descriptions:
- return enum_cls
-
- def json_schema(cls: type[Enum], core_schema: Any, handler: Any) ->
dict[str, Any]:
- options: list[dict[str, Any]] = []
- for member in cls:
- option: dict[str, Any] = {"const": member.value, "type": "string"}
- if text := descriptions.get(member.value):
- option["description"] = text
- options.append(option)
- return {"anyOf": options, "title": cls.__name__}
-
- # pydantic looks this hook up on the type when it builds the schema, so
attaching it to the
- # functional-API enum is the same as defining it in a class body.
- setattr(enum_cls, "__get_pydantic_json_schema__", classmethod(json_schema))
- return enum_cls
-
-
-def _picked(value: Any) -> str:
- """Return the task ID a picked option stands for: ``Choices`` answers with
the key, the Enum fallback with a member."""
- return value.value if isinstance(value, Enum) else str(value)
+ return described_choices("DownstreamTasks", {name: descriptions.get(name)
for name in task_ids})
class LLMBranchOperator(LLMOperator, BranchMixIn):
@@ -268,7 +235,7 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
# The output type validated the pick, so it is a task ID (or IDs) in
either encoding.
branches: str | list[str] = (
- [_picked(item) for item in output] if isinstance(output, list)
else _picked(output)
+ [picked_key(item) for item in output] if isinstance(output, list)
else picked_key(output)
)
if not branches:
@@ -281,7 +248,7 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
picked = [branches] if isinstance(branches, str) else branches
policy = self.decision_policy
threshold = threshold_for(policy.min_confidence, self._branch_bars,
picked)
- confidence = model_confidence.confidence.get("response")
+ confidence = model_confidence.confidence.get(BARE_OUTPUT_FIELD)
review = review_reason(
require_approval=self.require_approval, threshold=threshold,
confidence=confidence
)
@@ -317,7 +284,7 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
f"```\nPrompt: {self.prompt}\n\nChosen branch(es):
{chosen}\n```"
)
if review != "require_approval" or model_confidence.confidence:
- body += "\n\n" + describe_confidence(model_confidence,
"response", threshold)
+ body += "\n\n" + describe_confidence(model_confidence,
BARE_OUTPUT_FIELD, threshold)
modification_schema = (
{"type": "array", "items": {"type": "string", "enum":
choices}, "examples": choices}
if self.allow_multiple_branches
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 864e5b6dad4..c41a731b66e 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
@@ -17,22 +17,40 @@
"""
LLM-powered retry policy using pydantic-ai for error classification.
+The model answers one question, which category the failure belongs to, and
everything
+else is derived in the worker from the author's :class:`ErrorCategory` table:
whether the
+category is retried, after how long, and how sure the model has to be. That
shape is what
+a classifier model such as TypeSafe's Jev answers, and a text model answers it
too.
+
Requires Airflow 3.3+ (RetryPolicy was added in AIP-105).
"""
from __future__ import annotations
import logging
+import warnings
+from collections.abc import Mapping
+from dataclasses import dataclass
from datetime import timedelta
+from types import MappingProxyType
from typing import TYPE_CHECKING, cast
-from pydantic import BaseModel
-
+from airflow.providers.common.ai.policies.decision import check_bar
+from airflow.providers.common.ai.utils.decision import (
+ BARE_OUTPUT_FIELD,
+ ModelConfidence,
+ ReviewReason,
+ described_choices,
+ picked_key,
+ review_reason,
+ threshold_for,
+)
from airflow.providers.common.compat.sdk import redact
try:
from airflow.sdk.definitions.retry_policy import (
ExceptionRetryPolicy,
+ RetryAction,
RetryDecision,
RetryPolicy,
)
@@ -50,35 +68,100 @@ if TYPE_CHECKING:
log = logging.getLogger(__name__)
-__all__ = ["ErrorClassification", "LLMRetryPolicy",
"redact_registered_secrets"]
+__all__ = [
+ "DEFAULT_CATEGORIES",
+ "ErrorCategory",
+ "LLMRetryPolicy",
+ "redact_registered_secrets",
+]
-DEFAULT_INSTRUCTIONS = (
- "You are an error classifier for a data pipeline system. "
- "Given an error message from a failed task, classify it into one of these
categories:\n\n"
- "- rate_limit: API throttling or quota exceeded. Should retry after a
delay.\n"
- "- auth: Credentials invalid, expired, or missing permissions. Should NOT
retry.\n"
- "- network: Transient connectivity issue. Should retry quickly.\n"
- "- data: Schema validation, type mismatch, or bad input data. Should NOT
retry.\n"
- "- resource: Resource not found or unavailable (e.g., missing table,
bucket). Should NOT retry.\n"
- "- transient: Temporary issue likely to resolve on its own. Should
retry.\n"
- "- permanent: Problem that won't resolve without code or config changes.
Should NOT retry.\n\n"
- "Set suggested_delay_seconds based on the error type: "
- "60 for rate limits, 10 for network, 30 for transient. "
- "Set 0 for errors that should not retry."
-)
+
+@dataclass(frozen=True)
+class ErrorCategory:
+ """
+ One kind of failure the model may name, and what the policy does when it
does.
+
+ The value of :class:`LLMRetryPolicy`'s ``categories`` mapping, keyed by
the category
+ name the model answers with.
+
+ :param description: What failures belong here. Sent to the model in the
output schema
+ next to the category name, so the model reads the option together with
its meaning
+ rather than guessing from the name.
+ :param retry: Whether a failure in this category is retried (``True``,
default) or fails
+ the task at once.
+ :param delay: How long to wait before the retry. ``None`` (default) keeps
the task's own
+ ``retry_delay`` and backoff. Only meaningful with ``retry=True``.
+ :param min_confidence: A bar for this category that differs from the
policy's
+ ``min_confidence``, so a category whose wrong pick costs more
(``permanent`` ends
+ the task) can demand more certainty. Needs a policy bar to differ
from; ``None``
+ inherits it.
+ """
+
+ description: str
+ retry: bool = True
+ delay: timedelta | None = None
+ min_confidence: float | None = None
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.description, str) or not
self.description.strip():
+ raise ValueError("ErrorCategory.description must be a non-empty
string.")
+ if not isinstance(self.retry, bool):
+ raise TypeError(f"ErrorCategory.retry must be True or False, got
{self.retry!r}.")
+ if self.delay is not None:
+ if not isinstance(self.delay, timedelta):
+ raise TypeError(
+ f"ErrorCategory.delay must be a timedelta or None, got
{self.delay!r}. "
+ "Use timedelta(seconds=60) rather than 60."
+ )
+ if self.delay < timedelta(0):
+ raise ValueError(f"ErrorCategory.delay must not be negative,
got {self.delay!r}.")
+ if not self.retry:
+ raise ValueError(
+ "ErrorCategory.delay is set but retry=False; a failed task
has no retry to delay."
+ )
+ if self.min_confidence is not None:
+ object.__setattr__(
+ self, "min_confidence", check_bar(self.min_confidence,
"ErrorCategory.min_confidence")
+ )
-class ErrorClassification(BaseModel):
- """Structured LLM output for error classification."""
+DEFAULT_CATEGORIES: Mapping[str, ErrorCategory] = MappingProxyType(
+ {
+ "rate_limit": ErrorCategory("API throttling or a quota exceeded.",
delay=timedelta(seconds=60)),
+ "auth": ErrorCategory("Credentials invalid, expired, or missing
permissions.", retry=False),
+ "network": ErrorCategory(
+ "Transient connectivity issue: connection reset, DNS, TLS
handshake.", delay=timedelta(seconds=10)
+ ),
+ "data": ErrorCategory("Schema validation, type mismatch, or bad input
data.", retry=False),
+ "resource": ErrorCategory(
+ "Resource not found or unavailable, such as a missing table or
bucket.", retry=False
+ ),
+ "transient": ErrorCategory(
+ "Temporary issue likely to resolve on its own.",
delay=timedelta(seconds=30)
+ ),
+ "permanent": ErrorCategory(
+ "Problem that will not resolve without a code or configuration
change.", retry=False
+ ),
+ }
+)
+"""
+The example taxonomy :class:`LLMRetryPolicy` classifies into when
``categories`` is not set.
- category: str
- """One of: rate_limit, auth, network, data, resource, transient,
permanent."""
- should_retry: bool
- """Whether the operation should be retried."""
- suggested_delay_seconds: int = 0
- """How long to wait before retrying (0 if should_retry is False)."""
- reasoning: str
- """Brief explanation of the classification decision."""
+Retried: ``rate_limit`` after 60s, ``network`` after 10s, ``transient`` after
30s. Failed at
+once: ``auth``, ``data``, ``resource``, ``permanent``. Read-only; spread it
into your own
+mapping to change one entry.
+"""
+
+DEFAULT_INSTRUCTIONS = (
+ "You are an error classifier for a data pipeline system. Given an error
message from a failed "
+ "task, pick the single category that best describes it. Each category's
description says what "
+ "it covers."
+)
+"""
+The default system prompt. It no longer recites the categories: those travel
in the output
+schema from ``categories``, with their descriptions, so a prompt built as
+``DEFAULT_INSTRUCTIONS + hints`` should teach error strings and not name
categories or delays.
+"""
def redact_registered_secrets(message: str) -> str:
@@ -89,27 +172,47 @@ def redact_registered_secrets(message: str) -> str:
class LLMRetryPolicy(RetryPolicy):
"""
- Retry policy that uses an LLM to classify errors and decide retry
behaviour.
+ Retry policy that uses an LLM to name the kind of failure, then acts on
the author's table.
Uses :class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook`
to call any configured LLM provider (OpenAI, Anthropic, Bedrock, Vertex,
- Ollama, etc.) for error classification with structured output.
+ Ollama, a classifier model such as TypeSafe's Jev, etc.).
- When the LLM call itself fails, the policy falls back to ``fallback_rules``
- (if provided) or returns DEFAULT to use the task's standard retry logic.
+ The model's only job is to pick one of ``categories``; it reads each one's
description
+ from the output schema. Whether that category is retried, after how long,
and how sure
+ the model has to be all come from the :class:`ErrorCategory` in the worker
process.
+
+ When the LLM call itself fails, or the model is not sure enough to act on,
the policy
+ falls back to ``fallback_rules`` (if provided) or returns DEFAULT to use
the task's
+ standard retry logic.
:param llm_conn_id: Airflow connection ID for the LLM provider.
:param model_id: Model identifier override (e.g. ``"openai:gpt-4o-mini"``
- for cost efficiency). If not set, uses the model from the connection.
- :param instructions: Custom system prompt for classification.
- Defaults to a general-purpose error classifier.
+ for cost efficiency, or ``"typesafe:jev-1.13.0"`` for a classifier
model).
+ If not set, uses the model from the connection.
+ :param instructions: Custom system prompt for classification. Defaults to
a general
+ error classifier. Instructions can teach the model your stack's error
strings; the
+ categories themselves, and what each one means, are ``categories``.
:param fallback_rules: Optional list of
- :class:`~airflow.sdk.definitions.retry_policy.RetryRule` applied when
the
- LLM call fails. Provides a deterministic safety net.
+ :class:`~airflow.sdk.definitions.retry_policy.RetryRule` applied when
the LLM call
+ fails or the answer is under its confidence bar. Provides a
deterministic safety net.
:param timeout: Maximum seconds to wait for the LLM response before
falling back. Defaults to 30s. The LLM provider's own timeout
(e.g. 600s for Anthropic) is much longer; this keeps the retry
decision path fast even when the provider is degraded.
+ :param categories: The failure kinds the model chooses between, each with
its
+ description, action, delay and bar. Defaults to
:data:`DEFAULT_CATEGORIES`. Passing
+ this **replaces** the default mapping rather than merging into it; at
least two
+ categories are required. The model is constrained to these names, so
an answer
+ outside them is rejected before the policy acts on it.
+ :param min_confidence: The confidence, from 0 to 1, the model's answer
needs for the
+ policy to act on it. ``None`` (default) is no bar: the answer is acted
on as it
+ always was. Confidence comes from models that report one, such as a
classifier
+ model, in ``provider_details``. Under the bar, or when a bar is set
and the model
+ reported no confidence, the answer is discarded and ``fallback_rules``
then the
+ task's own retry behaviour apply, so swapping the connection to a text
model does
+ not silently switch off a control the author set. A category's own
+ ``min_confidence`` overrides this one for that category.
:param redactor: Callable applied to the exception's string representation
before it is added to the classification prompt. Defaults to
:func:`~airflow.providers.common.ai.policies.retry.redact_registered_secrets`,
@@ -157,6 +260,8 @@ class LLMRetryPolicy(RetryPolicy):
fallback_rules: list[RetryRule] | None = None,
timeout: float = 30.0,
*,
+ categories: Mapping[str, ErrorCategory] | None = None,
+ min_confidence: float | None = None,
redactor: Callable[[str], str] | None = None,
redact_exception: bool = True,
max_exception_length: int = 4096,
@@ -172,15 +277,65 @@ class LLMRetryPolicy(RetryPolicy):
)
self.llm_conn_id = llm_conn_id
self.model_id = model_id
+ if instructions and categories is None:
+ # The one silent upgrade case: a 0.9.x prompt that named its own
categories or delays.
+ warnings.warn(
+ "LLMRetryPolicy: instructions are custom but categories is not
set, so the model chooses "
+ "between DEFAULT_CATEGORIES and their delays. Instructions no
longer define categories, "
+ "retry or delay; if yours name a category outside the defaults
or a delay, move it into "
+ "categories={name: ErrorCategory(...)}.",
+ UserWarning,
+ stacklevel=2,
+ )
self.instructions = instructions or DEFAULT_INSTRUCTIONS
self.fallback_rules = fallback_rules
self.timeout = timeout
+ self.min_confidence = None if min_confidence is None else
check_bar(min_confidence, "min_confidence")
+ self.categories: Mapping[str, ErrorCategory] =
self._validate_categories(
+ DEFAULT_CATEGORIES if categories is None else categories
+ )
self.redactor: Callable[[str], str] | None = (
None if not redact_exception else redactor if redactor is not None
else redact_registered_secrets
)
self.redact_exception = redact_exception
self.max_exception_length = max_exception_length
+ def _validate_categories(self, categories: Mapping[str, ErrorCategory]) ->
dict[str, ErrorCategory]:
+ if not isinstance(categories, Mapping):
+ raise TypeError(
+ f"categories must be a mapping of name to ErrorCategory, got
{type(categories).__name__}."
+ )
+ if len(categories) < 2:
+ raise ValueError(
+ f"categories needs at least two entries for the model to
choose between, got {len(categories)}."
+ )
+ for name, category in categories.items():
+ if not isinstance(name, str) or not name.strip():
+ raise ValueError(f"categories keys must be non-empty strings,
got {name!r}.")
+ if not isinstance(category, ErrorCategory):
+ raise TypeError(
+ f"categories[{name!r}] must be an ErrorCategory, got
{type(category).__name__}."
+ )
+ with_bar = sorted(
+ name for name, category in categories.items() if
category.min_confidence is not None
+ )
+ if with_bar and self.min_confidence is None:
+ raise ValueError(
+ f"categories {with_bar} set min_confidence but the policy has
none to inherit. "
+ "Set min_confidence on LLMRetryPolicy as the bar for every
other category."
+ )
+ # A fresh dict, not the caller's mapping: DEFAULT_CATEGORIES is a
mappingproxy, which
+ # deepcopy rejects, and TaskGroup(default_args=...) and operator
copies deep-copy the policy.
+ return dict(categories)
+
+ @property
+ def _category_bars(self) -> dict[str, float]:
+ return {
+ name: category.min_confidence
+ for name, category in self.categories.items()
+ if category.min_confidence is not None
+ }
+
def evaluate(
self,
exception: BaseException,
@@ -188,29 +343,62 @@ class LLMRetryPolicy(RetryPolicy):
max_tries: int,
context: Context | None = None,
) -> RetryDecision:
+ outcome: RetryDecision | str
try:
- return self._classify(exception, try_number, max_tries)
+ outcome = self._classify(exception, try_number, max_tries)
except Exception:
log.exception("LLM retry classification failed, using fallback")
- if self.fallback_rules:
- return
ExceptionRetryPolicy(rules=self.fallback_rules).evaluate(
- exception, try_number, max_tries, context
+ outcome = "model_error"
+ if isinstance(outcome, RetryDecision):
+ return outcome
+ return self._fall_back(exception, try_number, max_tries, context,
why=outcome)
+
+ def _fall_back(
+ self,
+ exception: BaseException,
+ try_number: int,
+ max_tries: int,
+ context: Context | None,
+ *,
+ why: str,
+ ) -> RetryDecision:
+ """
+ Take the deterministic path: ``fallback_rules`` when one matches, else
the task's own retry behaviour.
+
+ The decision's reason says the model's answer was not applied and why,
so a ``retry_reason``
+ read later is not mistaken for a plain rule match or for a classifier
decision.
+ """
+ prefix = f"LLM classification not applied ({why})"
+ if self.fallback_rules:
+ ruled = ExceptionRetryPolicy(rules=self.fallback_rules).evaluate(
+ exception, try_number, max_tries, context
+ )
+ if ruled.action is not RetryAction.DEFAULT:
+ return RetryDecision(
+ action=ruled.action, retry_delay=ruled.retry_delay,
reason=f"{prefix}; {ruled.reason}"
)
- return RetryDecision.default()
+ return RetryDecision(action=RetryAction.DEFAULT, reason=f"{prefix};
task retry settings apply")
def _classify(
self,
exception: BaseException,
try_number: int,
max_tries: int,
- ) -> RetryDecision:
+ ) -> RetryDecision | ReviewReason:
+ """
+ Ask the model which category the failure is and act on the answer.
+
+ Returns the reason (``"below_threshold"`` or ``"missing_confidence"``)
instead of a
+ decision when the answer is under its confidence bar, so
:meth:`evaluate` takes the
+ same fallback path it takes when the model call fails and can say why.
+ """
from airflow.providers.common.ai.hooks.pydantic_ai import
PydanticAIHook
hook = PydanticAIHook(llm_conn_id=self.llm_conn_id,
model_id=self.model_id)
- agent = hook.create_agent(
- output_type=ErrorClassification,
- instructions=self.instructions,
+ output_type = described_choices(
+ "ErrorCategory", {name: category.description for name, category in
self.categories.items()}
)
+ agent = hook.create_agent(output_type=output_type,
instructions=self.instructions)
# Redact before truncating -- truncating first could cut a registered
secret in half.
message = self.redactor(str(exception)) if self.redactor is not None
else str(exception)
@@ -228,25 +416,37 @@ class LLMRetryPolicy(RetryPolicy):
prompt,
model_settings=ModelSettings(timeout=self.timeout),
)
- classification = result.output
-
- log.info(
- "LLM error classification: category=%s, should_retry=%s,
delay=%ds, reasoning=%s",
- classification.category,
- classification.should_retry,
- classification.suggested_delay_seconds,
- classification.reasoning,
+ # The output type validated the answer, so it is one of the configured
names.
+ name = picked_key(result.output)
+ category = self.categories[name]
+
+ # A bare output type is one field, ``response``; its confidence is
what the bar is compared against.
+ model_confidence = ModelConfidence.from_result(result)
+ confidence = model_confidence.confidence.get(BARE_OUTPUT_FIELD)
+ threshold = threshold_for(self.min_confidence, self._category_bars,
[name])
+ uncertain = review_reason(require_approval=False, threshold=threshold,
confidence=confidence)
+
+ summary = (
+ f"category={name}"
+ f" confidence={'n/a' if confidence is None else
f'{confidence:.2f}'}"
+ f" threshold={'n/a' if threshold is None else f'{threshold:.2f}'}"
)
+ if uncertain is not None:
+ log.info(
+ "LLM error classification not acted on (%s): %s, using
fallback. model=%s probabilities=%s",
+ uncertain,
+ summary,
+ model_confidence.model or "n/a",
+ model_confidence.probabilities.get(BARE_OUTPUT_FIELD) or "n/a",
+ )
+ return uncertain
- if not classification.should_retry:
- return RetryDecision.fail(reason=f"{classification.category}:
{classification.reasoning}")
+ if not category.retry:
+ reason = f"{summary} action=fail"
+ log.info("LLM error classification: %s model=%s", reason,
model_confidence.model or "n/a")
+ return RetryDecision.fail(reason=reason)
- delay = (
- timedelta(seconds=classification.suggested_delay_seconds)
- if classification.suggested_delay_seconds > 0
- else None
- )
- return RetryDecision.retry(
- delay=delay,
- reason=f"{classification.category}: {classification.reasoning}",
- )
+ delay_text = "task default" if category.delay is None else
f"{category.delay.total_seconds():g}s"
+ reason = f"{summary} action=retry delay={delay_text}"
+ log.info("LLM error classification: %s model=%s", reason,
model_confidence.model or "n/a")
+ return RetryDecision.retry(delay=category.delay, reason=reason)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/decision.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/decision.py
index 302e35228ae..aaf1bc504a0 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/decision.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/decision.py
@@ -30,7 +30,8 @@ probability that the answer is correct, and the gate does not
describe it as suc
The public configuration lives in
:mod:`airflow.providers.common.ai.policies.decision`
(:class:`DecisionPolicy`, :class:`BranchOption`); this module is the machinery
the
-operators call.
+operators and the retry policy call, including :func:`described_choices`,
which builds
+the option type a model picks from with each option's description in the
schema.
"""
from __future__ import annotations
@@ -38,12 +39,14 @@ from __future__ import annotations
import math
from collections.abc import Iterable, Mapping
from dataclasses import dataclass, field
+from enum import Enum
from typing import Any, Literal
from airflow.providers.common.ai.exceptions import LowConfidenceError
from airflow.providers.common.ai.policies.decision import DecisionPolicy,
UncertainAction
__all__ = [
+ "BARE_OUTPUT_FIELD",
"DECISION_XCOM_KEY",
"DecidedBy",
"DecisionPolicy",
@@ -53,8 +56,10 @@ __all__ = [
"check_uncertain_action",
"decision_record",
"describe_confidence",
+ "described_choices",
"finalize_record",
"initial_decided_by",
+ "picked_key",
"policy_fails",
"policy_record",
"review_reason",
@@ -65,11 +70,39 @@ __all__ = [
DECISION_XCOM_KEY = "decision"
+BARE_OUTPUT_FIELD = "response"
+"""The field name a classifier model reports a bare (non-object) output type's
confidence under."""
+
ReviewReason = Literal["require_approval", "below_threshold",
"missing_confidence"]
DecidedBy = Literal["model", "human", "timeout_default", "policy", "timeout"]
_UNCERTAIN: tuple[ReviewReason, ...] = ("below_threshold",
"missing_confidence")
+_NOT_LOADED = object()
+Choice: Any = _NOT_LOADED
+Choices: Any = _NOT_LOADED
+
+
+def _choice_types() -> tuple[Any, Any]:
+ """
+ Return pydantic-ai's ``(Choice, Choices)``, or ``(None, None)`` before
2.46.0.
+
+ Loaded on first use rather than at import: importing ``pydantic_ai`` costs
most of a second,
+ and this module is imported by the retry policy at Dag-parse time, where
nothing needs it yet.
+ Tests patch the module attributes directly, which this honours.
+ """
+ global Choice, Choices
+ if Choice is _NOT_LOADED or Choices is _NOT_LOADED:
+ try:
+ from pydantic_ai import ( # type: ignore[attr-defined]
+ Choice as loaded_choice,
+ Choices as loaded_choices,
+ )
+ except ImportError: # pydantic-ai < 2.46.0: build the same schema
from an Enum instead
+ loaded_choice = loaded_choices = None
+ Choice, Choices = loaded_choice, loaded_choices
+ return Choice, Choices
+
def validate_decision_policy(policy: Any) -> DecisionPolicy:
"""Accept a :class:`DecisionPolicy` or None (no gate); reject anything
else at construction."""
@@ -269,3 +302,48 @@ def describe_confidence(model_confidence: ModelConfidence,
key: str, threshold:
ranked = sorted(probabilities.items(), key=lambda item: item[1],
reverse=True)
line += "\n\nProbabilities: " + ", ".join(f"{option} {p:.2f}" for
option, p in ranked)
return line
+
+
+def described_choices(name: str, options: Mapping[str, str | None]) ->
type[Any]:
+ """
+ Build the type a model picks one of ``options`` from, each option carrying
its description.
+
+ With a description on any option the schema renders as ``anyOf`` of
``{const, description}``
+ instead of a bare ``enum`` list. That is the one JSON Schema shape that
carries a description
+ per value, and it is what both a text model's tool schema and
pydantic-ai's TypeSafe adapter
+ read an option's meaning from. On pydantic-ai 2.46+ the type is its
``Choices``; before that,
+ an ``Enum`` whose schema hook emits the same shape. Either way the model
has to answer with one
+ of the keys, in the order given, and :func:`picked_key` returns that key
whichever type answered.
+ """
+ keys = list(options)
+ descriptions = {key: text for key, text in options.items() if text}
+ choice, choices = _choice_types()
+ if choices is not None:
+ if not descriptions:
+ return choices(keys, name=name)
+ return choices({key: choice(descriptions.get(key)) for key in keys},
name=name)
+
+ # Generated member names: Enum reserves ``_sunder_`` names and ``mro``,
and a key need not be an
+ # identifier. The key is the member's value, which is what the schema, the
model and ``picked`` use.
+ enum_cls: type[Enum] = Enum(name, {f"option_{i}": key for i, key in
enumerate(keys)}) # type: ignore[misc]
+ if not descriptions:
+ return enum_cls
+
+ def json_schema(cls: type[Enum], core_schema: Any, handler: Any) ->
dict[str, Any]:
+ rendered: list[dict[str, Any]] = []
+ for member in cls:
+ option: dict[str, Any] = {"const": member.value, "type": "string"}
+ if text := descriptions.get(member.value):
+ option["description"] = text
+ rendered.append(option)
+ return {"anyOf": rendered, "title": cls.__name__}
+
+ # pydantic looks this hook up on the type when it builds the schema, so
attaching it to the
+ # functional-API enum is the same as defining it in a class body.
+ setattr(enum_cls, "__get_pydantic_json_schema__", classmethod(json_schema))
+ return enum_cls
+
+
+def picked_key(value: Any) -> str:
+ """Return the key a picked option stands for: ``Choices`` answers with the
key, the Enum fallback with a member."""
+ return value.value if isinstance(value, Enum) else str(value)
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
index 6f42a2ab6f9..835a4e49117 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
@@ -402,8 +402,8 @@ class TestLLMBranchOperator:
}
assert "template_fields" not in {f.name for f in
__import__("dataclasses").fields(BranchOption)}
- @patch("airflow.providers.common.ai.operators.llm_branch.Choices", None)
- @patch("airflow.providers.common.ai.operators.llm_branch.Choice", None)
+ @patch("airflow.providers.common.ai.utils.decision.Choices", None)
+ @patch("airflow.providers.common.ai.utils.decision.Choice", None)
@patch.object(LLMBranchOperator, "do_branch")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_enum_fallback_emits_the_same_schema_and_branches_on_the_member(
@@ -421,7 +421,7 @@ class TestLLMBranchOperator:
op.downstream_task_ids = {"handle_auth", "handle_billing"}
def capture(**kwargs):
- mock_agent.run_sync.return_value =
make_mock_run_result(kwargs["output_type"].handle_billing)
+ mock_agent.run_sync.return_value =
make_mock_run_result(kwargs["output_type"]("handle_billing"))
return mock_agent
mock_hook_cls.get_hook.return_value.create_agent.side_effect = capture
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 9b38eb243e1..1dc91c0e1ba 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
@@ -16,37 +16,75 @@
# under the License.
from __future__ import annotations
+import copy
+import logging
+import math
+import warnings
from datetime import timedelta
from unittest.mock import MagicMock, patch
import pytest
+from pydantic import TypeAdapter, ValidationError
+from pydantic_ai import Agent
+from pydantic_ai.agent import AgentRunResult
+from pydantic_ai.exceptions import UnexpectedModelBehavior
+from pydantic_ai.messages import ModelResponse
# LLMRetryPolicy depends on the RetryPolicy ABC introduced in Airflow 3.3
(AIP-105).
# Skip the entire test module on older Airflow versions tested in compat CI.
pytest.importorskip("airflow.sdk.definitions.retry_policy",
reason="RetryPolicy requires Airflow 3.3+")
from airflow.providers.common.ai.policies.retry import (
- ErrorClassification,
+ DEFAULT_CATEGORIES,
+ DEFAULT_INSTRUCTIONS,
+ ErrorCategory,
LLMRetryPolicy,
redact_registered_secrets,
)
+from airflow.providers.common.ai.utils.decision import picked_key
from airflow.sdk._shared.secrets_masker import reset_secrets_masker
from airflow.sdk.definitions.retry_policy import RetryAction, RetryRule
from airflow.sdk.log import mask_secret
+HOOK = "airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook"
+
+
+def _run_result(output, *, confidence=None, probabilities=None,
model_name="test-model"):
+ """A run result as pydantic-ai returns it: the pick in ``output``, the
confidence in ``provider_details``."""
+ details = {}
+ if confidence is not None:
+ details["confidence"] = {"response": confidence}
+ if probabilities is not None:
+ details["probabilities"] = {"response": probabilities}
+ result = MagicMock(spec=AgentRunResult)
+ result.output = output
+ result.response = MagicMock(spec=ModelResponse, model_name=model_name,
provider_details=details or None)
+ return result
+
+
+def _agent(output, **kwargs):
+ """A mock agent whose ``run_sync`` returns ``output``; ``kwargs`` go to
:func:`_run_result`."""
+ agent = MagicMock(spec=Agent)
+ agent.run_sync.return_value = _run_result(output, **kwargs)
+ return agent
+
+
+def _install(mock_hook_cls, agent):
+ """Wire a mock agent behind the patched hook class and return the mock
agent."""
+ mock_hook_cls.return_value.create_agent.return_value = agent
+ return agent
-def _make_mock_agent(category, should_retry, delay=0, reasoning="test"):
- """Create a mock agent that returns a canned ErrorClassification."""
- mock_result = MagicMock()
- mock_result.output = ErrorClassification(
- category=category,
- should_retry=should_retry,
- suggested_delay_seconds=delay,
- reasoning=reasoning,
- )
- mock_agent = MagicMock()
- mock_agent.run_sync.return_value = mock_result
- return mock_agent
+
+def _output_type(mock_hook_cls):
+ return
mock_hook_cls.return_value.create_agent.call_args.kwargs["output_type"]
+
+
+def _options(output_type):
+ """The ``(name, description)`` pairs the model is offered, as they appear
in the schema."""
+ schema = TypeAdapter(output_type).json_schema()
+ if "anyOf" in schema:
+ return [(o["const"], o.get("description")) for o in schema["anyOf"]]
+ return [(name, None) for name in schema["enum"]]
@pytest.mark.enable_redact
@@ -61,247 +99,664 @@ def
test_redact_registered_secrets_masks_only_registered_values():
)
-class TestLLMClassifyDecisions:
- """Test that _classify maps LLM classification to correct
RetryDecisions."""
-
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
- def test_auth_error_returns_fail(self, mock_hook_cls):
- mock_hook_cls.return_value.create_agent.return_value =
_make_mock_agent(
- "auth", should_retry=False, reasoning="API key expired"
+class TestErrorCategory:
+ @pytest.mark.parametrize("description", ["", " ", None, 3])
+ def test_description_must_be_a_non_empty_string(self, description):
+ with pytest.raises(ValueError, match="description must be a non-empty
string"):
+ ErrorCategory(description)
+
+ def test_int_seconds_delay_is_rejected(self):
+ """Int seconds was the shape of the removed
``suggested_delay_seconds`` field."""
+ with pytest.raises(TypeError, match="must be a timedelta or None, got
60"):
+ ErrorCategory("x", delay=60) # type: ignore[arg-type]
+
+ @pytest.mark.parametrize("retry", ["false", None, 0])
+ def test_retry_must_be_a_bool(self, retry):
+ """A truthy string here would retry a category the author meant to
fail."""
+ with pytest.raises(TypeError, match="retry must be True or False"):
+ ErrorCategory("x", retry=retry) # type: ignore[arg-type]
+
+ def test_negative_delay_is_rejected(self):
+ with pytest.raises(ValueError, match="must not be negative"):
+ ErrorCategory("x", delay=timedelta(seconds=-1))
+
+ def test_delay_without_retry_is_a_contradiction(self):
+ with pytest.raises(ValueError, match="retry=False; a failed task has
no retry to delay"):
+ ErrorCategory("x", retry=False, delay=timedelta(seconds=5))
+
+ def test_bar_is_normalized_and_validated(self):
+ assert ErrorCategory("x", min_confidence=1).min_confidence == 1.0
+ with pytest.raises(ValueError, match="ErrorCategory.min_confidence
must be between 0 and 1"):
+ ErrorCategory("x", min_confidence=1.5)
+ with pytest.raises(TypeError, match="ErrorCategory.min_confidence must
be a number"):
+ ErrorCategory("x", min_confidence="0.5") # type: ignore[arg-type]
+
+
+class TestDefaultCategories:
+ def test_cannot_be_mutated(self):
+ """The default table is shared by every policy instance, so it must be
read-only."""
+ with pytest.raises(TypeError):
+ DEFAULT_CATEGORIES["auth"] = ErrorCategory("x") # type:
ignore[index]
+
+ def test_documented_split(self):
+ """Spelled out here on purpose: the docs list this table, so a drift
shows up as a test failure."""
+ assert {name: (c.retry, c.delay, c.min_confidence) for name, c in
DEFAULT_CATEGORIES.items()} == {
+ "rate_limit": (True, timedelta(seconds=60), None),
+ "auth": (False, None, None),
+ "network": (True, timedelta(seconds=10), None),
+ "data": (False, None, None),
+ "resource": (False, None, None),
+ "transient": (True, timedelta(seconds=30), None),
+ "permanent": (False, None, None),
+ }
+
+ def test_default_instructions_do_not_recite_the_taxonomy(self):
+ """The categories travel in the schema with their descriptions, so the
prompt must not pin them."""
+ for name in DEFAULT_CATEGORIES:
+ assert f"- {name}:" not in DEFAULT_INSTRUCTIONS
+
+
+class TestLLMRetryPolicyConstruction:
+ @patch(HOOK, autospec=True)
+ def test_construction_makes_no_connection_or_network_call(self,
mock_hook_cls):
+ """The policy is instantiated at Dag parse time, so everything is
validated without a hook."""
+ LLMRetryPolicy(
+ llm_conn_id="test",
+ categories={"a": ErrorCategory("A", min_confidence=0.9), "b":
ErrorCategory("B", retry=False)},
+ min_confidence=0.6,
)
+
+ mock_hook_cls.assert_not_called()
+
+ def test_custom_instructions_without_categories_warn_at_parse_time(self):
+ """The 0.9.x pattern that put categories and delays in the prompt is
the one silent upgrade case."""
+ with pytest.warns(UserWarning, match="instructions are custom but
categories is not set"):
+ policy = LLMRetryPolicy(
+ llm_conn_id="test", instructions="'Statement queued' ->
rate_limit, retry after 120s"
+ )
+
+ assert policy.categories == dict(DEFAULT_CATEGORIES)
+
+ @pytest.mark.parametrize(
+ "kwargs",
+ [
+ pytest.param({}, id="defaults"),
+ pytest.param({"instructions": None}, id="explicit-none"),
+ pytest.param(
+ {"instructions": "hints", "categories": {"a":
ErrorCategory("A"), "b": ErrorCategory("B")}},
+ id="instructions-with-categories",
+ ),
+ pytest.param(
+ {"categories": {"a": ErrorCategory("A"), "b":
ErrorCategory("B")}}, id="categories-only"
+ ),
+ ],
+ )
+ def test_no_warning_when_the_taxonomy_is_explicit_or_default(self, kwargs):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error")
+ LLMRetryPolicy(llm_conn_id="test", **kwargs)
+
+ def test_categories_must_be_a_mapping(self):
+ with pytest.raises(TypeError, match="categories must be a mapping"):
+ LLMRetryPolicy(llm_conn_id="test",
categories=[ErrorCategory("x")]) # type: ignore[arg-type]
+
+ @pytest.mark.parametrize("categories", [{}, {"only": ErrorCategory("x")}])
+ def test_fewer_than_two_categories_is_rejected(self, categories):
+ with pytest.raises(ValueError, match="at least two entries"):
+ LLMRetryPolicy(llm_conn_id="test", categories=categories)
+
+ @pytest.mark.parametrize("name", ["", 3])
+ def test_category_names_must_be_non_empty_strings(self, name):
+ with pytest.raises(ValueError, match="keys must be non-empty strings"):
+ LLMRetryPolicy(llm_conn_id="test", categories={name:
ErrorCategory("x"), "b": ErrorCategory("y")})
+
+ def test_category_values_must_be_error_categories(self):
+ """A bare string is not a description here: retry and delay would have
nowhere to live."""
+ with pytest.raises(TypeError, match=r"categories\['a'\] must be an
ErrorCategory, got str"):
+ LLMRetryPolicy(llm_conn_id="test", categories={"a": "flaky", "b":
ErrorCategory("y")}) # type: ignore[dict-item]
+
+ def test_category_bar_needs_a_policy_bar_to_inherit(self):
+ with pytest.raises(
+ ValueError, match=r"categories \['a', 'b'\] set min_confidence but
the policy has none"
+ ):
+ LLMRetryPolicy(
+ llm_conn_id="test",
+ categories={
+ "a": ErrorCategory("A", min_confidence=0.9),
+ "b": ErrorCategory("B", min_confidence=0.8),
+ "c": ErrorCategory("C"),
+ },
+ )
+
+ def test_policy_bar_is_validated(self):
+ with pytest.raises(ValueError, match="min_confidence must be between 0
and 1"):
+ LLMRetryPolicy(llm_conn_id="test", min_confidence=2)
+
+ @pytest.mark.parametrize(
+ "categories",
+ [
+ pytest.param(None, id="default-table"),
+ pytest.param(
+ {"a": ErrorCategory("A", delay=timedelta(seconds=1)), "b":
ErrorCategory("B")}, id="own"
+ ),
+ ],
+ )
+ def test_policy_can_be_deep_copied(self, categories):
+ """``TaskGroup(default_args={"retry_policy": ...})`` deep-copies the
policy.
+
+ The default table is a mappingproxy, which cannot be pickled, so
holding it by
+ reference would make the documented configuration the one that fails
at Dag parse.
+ """
+ policy = LLMRetryPolicy(llm_conn_id="test", categories=categories)
+
+ assert copy.deepcopy(policy).categories == policy.categories
+
+ @patch(HOOK, autospec=True)
+ def test_caller_mapping_is_copied(self, mock_hook_cls):
+ """Mutating the caller's mapping after construction must not change
the policy."""
+ _install(mock_hook_cls, _agent("a"))
+ categories = {"a": ErrorCategory("A", delay=timedelta(seconds=7)),
"b": ErrorCategory("B")}
+ policy = LLMRetryPolicy(llm_conn_id="test", categories=categories)
+ categories["a"] = ErrorCategory("A", delay=timedelta(seconds=999))
+
+ decision = policy.evaluate(RuntimeError("x"), try_number=1,
max_tries=3)
+
+ assert decision.retry_delay == timedelta(seconds=7)
+
+ @pytest.mark.parametrize("max_exception_length", [0, -1])
+ def test_non_positive_max_exception_length_raises(self,
max_exception_length):
+ with pytest.raises(ValueError, match="max_exception_length must be a
positive integer"):
+ LLMRetryPolicy(llm_conn_id="test",
max_exception_length=max_exception_length)
+
+ def test_redact_exception_false_with_explicit_redactor_raises(self):
+ with pytest.raises(ValueError, match="redactor must not be set when
redact_exception=False"):
+ LLMRetryPolicy(llm_conn_id="test", redactor=lambda m: m,
redact_exception=False)
+
+
+class TestClassification:
+ """The model names the category; the table in the worker decides what
happens."""
+
+ @pytest.mark.parametrize(
+ ("category", "expected_action", "expected_delay"),
+ [
+ pytest.param("rate_limit", RetryAction.RETRY,
timedelta(seconds=60), id="rate_limit"),
+ pytest.param("network", RetryAction.RETRY, timedelta(seconds=10),
id="network"),
+ pytest.param("transient", RetryAction.RETRY,
timedelta(seconds=30), id="transient"),
+ pytest.param("auth", RetryAction.FAIL, None, id="auth"),
+ pytest.param("data", RetryAction.FAIL, None, id="data"),
+ pytest.param("resource", RetryAction.FAIL, None, id="resource"),
+ pytest.param("permanent", RetryAction.FAIL, None, id="permanent"),
+ ],
+ )
+ @patch(HOOK, autospec=True)
+ def test_default_table_covers_every_category(
+ self, mock_hook_cls, category, expected_action, expected_delay
+ ):
+ _install(mock_hook_cls, _agent(category))
policy = LLMRetryPolicy(llm_conn_id="test")
- decision = policy.evaluate(PermissionError("403"), try_number=1,
max_tries=3)
- assert decision.action == RetryAction.FAIL
- assert "auth" in decision.reason
- assert "API key expired" in decision.reason
+ decision = policy.evaluate(RuntimeError("boom"), try_number=1,
max_tries=3)
+
+ assert decision.action == expected_action
+ assert decision.retry_delay == expected_delay
+
+ @patch(HOOK, autospec=True)
+ def test_default_categories_reach_the_model_with_descriptions(self,
mock_hook_cls):
+ _install(mock_hook_cls, _agent("auth"))
+ LLMRetryPolicy(llm_conn_id="test").evaluate(RuntimeError("boom"),
try_number=1, max_tries=3)
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
- def test_rate_limit_returns_retry_with_delay(self, mock_hook_cls):
- mock_hook_cls.return_value.create_agent.return_value =
_make_mock_agent(
- "rate_limit", should_retry=True, delay=60, reasoning="429"
+ assert _options(_output_type(mock_hook_cls)) == [
+ (name, category.description) for name, category in
DEFAULT_CATEGORIES.items()
+ ]
+
+ @patch(HOOK, autospec=True)
+ def test_custom_categories_build_the_schema_in_the_authors_order(self,
mock_hook_cls):
+ """The author's names and descriptions are the whole taxonomy; the
prompt no longer carries one."""
+ _install(mock_hook_cls, _agent("warehouse_suspended"))
+ policy = LLMRetryPolicy(
+ llm_conn_id="test",
+ categories={
+ "warehouse_suspended": ErrorCategory(
+ "The warehouse auto-suspended.",
delay=timedelta(seconds=30)
+ ),
+ "schema_drift": ErrorCategory("A referenced column does not
exist.", retry=False),
+ },
)
+
+ decision = policy.evaluate(RuntimeError("Warehouse X is suspended"),
try_number=1, max_tries=3)
+
+ assert _options(_output_type(mock_hook_cls)) == [
+ ("warehouse_suspended", "The warehouse auto-suspended."),
+ ("schema_drift", "A referenced column does not exist."),
+ ]
+ assert decision.action == RetryAction.RETRY
+ assert decision.retry_delay == timedelta(seconds=30)
+
+ @patch(HOOK, autospec=True)
+ def
test_answer_outside_the_categories_is_rejected_by_the_output_type(self,
mock_hook_cls):
+ """A near-miss spelling or a sentence never reaches the table:
pydantic-ai re-prompts, then gives up."""
+ _install(mock_hook_cls, _agent("auth"))
+ LLMRetryPolicy(llm_conn_id="test").evaluate(RuntimeError("boom"),
try_number=1, max_tries=3)
+ adapter = TypeAdapter(_output_type(mock_hook_cls))
+
+ # ``Choices`` validates to the key; the Enum fallback to a member.
``picked_key`` reads both.
+ assert picked_key(adapter.validate_python("auth")) == "auth"
+ for bad in ("rate-limit", "RATE_LIMIT", "Rate limit exceeded, should
retry", ""):
+ with pytest.raises(ValidationError):
+ adapter.validate_python(bad)
+
+ @patch("airflow.providers.common.ai.utils.decision.Choices", None)
+ @patch("airflow.providers.common.ai.utils.decision.Choice", None)
+ @patch(HOOK, autospec=True)
+ def test_enum_fallback_emits_the_same_schema_and_acts_on_the_member(self,
mock_hook_cls):
+ """On pydantic-ai without ``Choices`` the Enum fallback carries the
same anyOf and the pick unwraps."""
+ agent = MagicMock(spec=Agent)
+
+ def capture(**kwargs):
+ agent.run_sync.return_value =
_run_result(kwargs["output_type"]("rate_limit"))
+ return agent
+
+ mock_hook_cls.return_value.create_agent.side_effect = capture
policy = LLMRetryPolicy(llm_conn_id="test")
+
decision = policy.evaluate(RuntimeError("429"), try_number=1,
max_tries=3)
+ assert _options(_output_type(mock_hook_cls))[0] == (
+ "rate_limit",
+ DEFAULT_CATEGORIES["rate_limit"].description,
+ )
assert decision.action == RetryAction.RETRY
assert decision.retry_delay == timedelta(seconds=60)
+ assert decision.reason.startswith("category=rate_limit ")
+
+ @pytest.mark.parametrize("name", ["_transient_", "mro", "rate limit",
"__weird__"])
+ @patch("airflow.providers.common.ai.utils.decision.Choices", None)
+ @patch("airflow.providers.common.ai.utils.decision.Choice", None)
+ @patch(HOOK, autospec=True)
+ def test_enum_fallback_survives_names_enum_reserves(self, mock_hook_cls,
name):
+ """A category name is a value, never an Enum member name:
``_sunder_``, ``mro`` and non-identifiers all work."""
+ agent = MagicMock(spec=Agent)
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
- def test_transient_retry_with_zero_delay_uses_default(self, mock_hook_cls):
- """suggested_delay_seconds=0 means use the task's default delay, not
override."""
- mock_hook_cls.return_value.create_agent.return_value =
_make_mock_agent(
- "transient", should_retry=True, delay=0
+ def capture(**kwargs):
+ agent.run_sync.return_value =
_run_result(kwargs["output_type"](name))
+ return agent
+
+ mock_hook_cls.return_value.create_agent.side_effect = capture
+ policy = LLMRetryPolicy(
+ llm_conn_id="test",
+ categories={
+ name: ErrorCategory("Odd name.", delay=timedelta(seconds=5)),
+ "other": ErrorCategory("O"),
+ },
)
+
+ decision = policy.evaluate(RuntimeError("x"), try_number=1,
max_tries=3)
+
+ assert _options(_output_type(mock_hook_cls)) == [(name, "Odd name."),
("other", "O")]
+ assert decision.action == RetryAction.RETRY
+ assert decision.retry_delay == timedelta(seconds=5)
+ assert decision.reason.startswith(f"category={name} ")
+
+ @pytest.mark.parametrize(
+ ("category", "confidence", "expected"),
+ [
+ pytest.param("auth", None, "category=auth confidence=n/a
threshold=n/a action=fail", id="fail"),
+ pytest.param(
+ "rate_limit",
+ 0.91,
+ "category=rate_limit confidence=0.91 threshold=n/a
action=retry delay=60s",
+ id="retry",
+ ),
+ ],
+ )
+ @patch(HOOK, autospec=True)
+ def test_reason_is_the_generated_summary(self, mock_hook_cls, category,
confidence, expected):
+ """No ``reasoning`` field: the reason a person reads is derived from
the decision itself."""
+ _install(mock_hook_cls, _agent(category, confidence=confidence))
policy = LLMRetryPolicy(llm_conn_id="test")
+
+ decision = policy.evaluate(RuntimeError("boom"), try_number=1,
max_tries=3)
+
+ assert decision.reason == expected
+
+ @patch(HOOK, autospec=True)
+ def test_none_delay_leaves_the_task_backoff_in_charge(self, mock_hook_cls):
+ _install(mock_hook_cls, _agent("flaky"))
+ policy = LLMRetryPolicy(
+ llm_conn_id="test",
+ categories={"flaky": ErrorCategory("F"), "broken":
ErrorCategory("B", retry=False)},
+ )
+
decision = policy.evaluate(RuntimeError("glitch"), try_number=1,
max_tries=3)
assert decision.action == RetryAction.RETRY
- assert decision.retry_delay is None # None = use task's default
+ assert decision.retry_delay is None
+ assert decision.reason.endswith("action=retry delay=task default")
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
- def test_negative_delay_treated_as_no_override(self, mock_hook_cls):
- """Negative delay from LLM should not produce a negative timedelta."""
- mock_hook_cls.return_value.create_agent.return_value =
_make_mock_agent(
- "transient", should_retry=True, delay=-5
+ @patch(HOOK, autospec=True)
+ def test_zero_delay_retries_without_waiting(self, mock_hook_cls):
+ """timedelta(0) is an override to retry at once, distinct from None."""
+ _install(mock_hook_cls, _agent("flaky"))
+ policy = LLMRetryPolicy(
+ llm_conn_id="test",
+ categories={
+ "flaky": ErrorCategory("F", delay=timedelta(0)),
+ "broken": ErrorCategory("B", retry=False),
+ },
)
+
+ decision = policy.evaluate(RuntimeError("glitch"), try_number=1,
max_tries=3)
+
+ assert decision.retry_delay == timedelta(0)
+ assert decision.reason.endswith("delay=0s")
+
+ @patch(HOOK, autospec=True)
+ def test_each_failure_is_classified_afresh(self, mock_hook_cls):
+ """Consecutive attempts with different exceptions each go to the
model; nothing is replayed."""
+ agent = _install(mock_hook_cls, MagicMock(spec=Agent))
+ agent.run_sync.side_effect = [_run_result("network"),
_run_result("auth")]
policy = LLMRetryPolicy(llm_conn_id="test")
- decision = policy.evaluate(RuntimeError("x"), try_number=1,
max_tries=3)
+
+ first = policy.evaluate(ConnectionError("reset"), try_number=1,
max_tries=3)
+ second = policy.evaluate(PermissionError("expired"), try_number=2,
max_tries=3)
+
+ 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 "PermissionError: expired" in prompts[1]
+ assert "attempt 2 of 3" in prompts[1]
+
+ @patch(HOOK, autospec=True)
+ def test_custom_instructions_forwarded_to_agent(self, mock_hook_cls):
+ _install(mock_hook_cls, _agent("auth"))
+ policy = LLMRetryPolicy(
+ llm_conn_id="test",
+ instructions="Snowflake errors only.",
+ categories={"auth": ErrorCategory("A", retry=False), "other":
ErrorCategory("O")},
+ )
+
+ policy.evaluate(RuntimeError("x"), try_number=1, max_tries=3)
+
+ assert (
+
mock_hook_cls.return_value.create_agent.call_args.kwargs["instructions"]
+ == "Snowflake errors only."
+ )
+
+ @patch(HOOK, autospec=True)
+ def test_model_id_and_connection_reach_the_hook(self, mock_hook_cls):
+ _install(mock_hook_cls, _agent("auth"))
+ policy = LLMRetryPolicy(llm_conn_id="jev_default",
model_id="typesafe:jev-1.13.0")
+
+ policy.evaluate(RuntimeError("x"), try_number=1, max_tries=3)
+
+ mock_hook_cls.assert_called_once_with(llm_conn_id="jev_default",
model_id="typesafe:jev-1.13.0")
+
+ @patch(HOOK, autospec=True)
+ def test_timeout_passed_via_model_settings(self, mock_hook_cls):
+ agent = _install(mock_hook_cls, _agent("auth"))
+ policy = LLMRetryPolicy(llm_conn_id="test", timeout=7.5)
+
+ policy.evaluate(RuntimeError("x"), try_number=1, max_tries=3)
+
+ assert agent.run_sync.call_args.kwargs["model_settings"] ==
{"timeout": 7.5}
+
+
+class TestConfidenceGate:
+ """Under the bar the answer is discarded and the deterministic fallback
path decides."""
+
+ RULES = [RetryRule(exception=RuntimeError, action=RetryAction.FAIL,
reason="rule")]
+
+ @pytest.mark.parametrize(
+ ("confidence", "expected_action", "expected_reason"),
+ [
+ pytest.param(
+ 0.60,
+ RetryAction.RETRY,
+ "category=network confidence=0.60 threshold=0.60 action=retry
delay=10s",
+ id="at-bar",
+ ),
+ pytest.param(
+ 0.59,
+ RetryAction.FAIL,
+ "LLM classification not applied (below_threshold); rule",
+ id="one-below",
+ ),
+ ],
+ )
+ @patch(HOOK, autospec=True)
+ def test_bar_boundary(self, mock_hook_cls, confidence, expected_action,
expected_reason):
+ _install(mock_hook_cls, _agent("network", confidence=confidence))
+ policy = LLMRetryPolicy(llm_conn_id="test", min_confidence=0.6,
fallback_rules=self.RULES)
+
+ decision = policy.evaluate(RuntimeError("reset"), try_number=1,
max_tries=3)
+
+ assert decision.action == expected_action
+ assert decision.reason == expected_reason
+
+ @patch(HOOK, autospec=True)
+ def test_under_the_bar_with_no_rules_keeps_the_task_behaviour(self,
mock_hook_cls, caplog):
+ """An unsure answer is an ordinary outcome: one INFO line with the
distribution, no ERROR traceback."""
+ _install(mock_hook_cls, _agent("auth", confidence=0.3,
probabilities={"auth": 0.3, "network": 0.28}))
+ policy = LLMRetryPolicy(llm_conn_id="test", min_confidence=0.6)
+
+ with caplog.at_level(logging.INFO,
logger="airflow.providers.common.ai.policies.retry"):
+ decision = policy.evaluate(RuntimeError("?"), try_number=1,
max_tries=3)
+
+ assert decision.action == RetryAction.DEFAULT
+ assert (
+ decision.reason == "LLM classification not applied
(below_threshold); task retry settings apply"
+ )
+ [record] = [r for r in caplog.records if r.name ==
"airflow.providers.common.ai.policies.retry"]
+ assert record.levelno == logging.INFO
+ assert (
+ "not acted on (below_threshold): category=auth confidence=0.30
threshold=0.60"
+ in record.getMessage()
+ )
+ assert "probabilities={'auth': 0.3, 'network': 0.28}" in
record.getMessage()
+
+ @patch(HOOK, autospec=True)
+ def test_under_the_bar_takes_fallback_rules_before_default(self,
mock_hook_cls):
+ """An unsure ``auth`` does not fail the task on the model's say-so;
the matching rule decides."""
+ _install(mock_hook_cls, _agent("auth", confidence=0.3))
+ policy = LLMRetryPolicy(
+ llm_conn_id="test",
+ min_confidence=0.6,
+ fallback_rules=[
+ RetryRule(
+ exception=ConnectionError, action=RetryAction.RETRY,
retry_delay=timedelta(seconds=5)
+ ),
+ RetryRule(exception=PermissionError, action=RetryAction.FAIL,
reason="rule"),
+ ],
+ )
+
+ assert policy.evaluate(ConnectionError("x"), try_number=1,
max_tries=3).retry_delay == timedelta(
+ seconds=5
+ )
+ assert (
+ policy.evaluate(PermissionError("x"), try_number=1,
max_tries=3).reason
+ == "LLM classification not applied (below_threshold); rule"
+ )
+ unmatched = policy.evaluate(ValueError("x"), try_number=1, max_tries=3)
+ assert unmatched.action == RetryAction.DEFAULT
+ assert (
+ unmatched.reason == "LLM classification not applied
(below_threshold); task retry settings apply"
+ )
+
+ @pytest.mark.parametrize(
+ ("category", "confidence", "expected_action"),
+ [
+ pytest.param("permanent", 0.80, RetryAction.DEFAULT,
id="permanent-under-its-own-bar"),
+ pytest.param("permanent", 0.85, RetryAction.FAIL,
id="permanent-at-its-own-bar"),
+ pytest.param("transient", 0.80, RetryAction.RETRY,
id="transient-inherits-the-policy-bar"),
+ pytest.param("network", 0.50, RetryAction.RETRY,
id="network-has-a-lower-bar-than-the-policy"),
+ pytest.param("transient", 0.50, RetryAction.DEFAULT,
id="transient-under-the-policy-bar"),
+ ],
+ )
+ @patch(HOOK, autospec=True)
+ def test_category_bar_overrides_the_policy_bar(
+ self, mock_hook_cls, category, confidence, expected_action
+ ):
+ _install(mock_hook_cls, _agent(category, confidence=confidence))
+ policy = LLMRetryPolicy(
+ llm_conn_id="test",
+ min_confidence=0.6,
+ categories={
+ **DEFAULT_CATEGORIES,
+ "permanent": ErrorCategory("Will fail identically.",
retry=False, min_confidence=0.85),
+ "network": ErrorCategory("Cheap to retry.",
delay=timedelta(seconds=10), min_confidence=0.4),
+ },
+ )
+
+ assert policy.evaluate(RuntimeError("x"), try_number=1,
max_tries=3).action == expected_action
+
+ @patch(HOOK, autospec=True)
+ def test_missing_confidence_with_a_bar_falls_back(self, mock_hook_cls):
+ """A text model reports no confidence; a model swap must not switch
off a bar the author set."""
+ _install(mock_hook_cls, _agent("network"))
+ policy = LLMRetryPolicy(llm_conn_id="test", min_confidence=0.6,
fallback_rules=self.RULES)
+
+ decision = policy.evaluate(RuntimeError("reset"), try_number=1,
max_tries=3)
+
+ assert decision.action == RetryAction.FAIL
+ assert decision.reason == "LLM classification not applied
(missing_confidence); rule"
+
+ @patch(HOOK, autospec=True)
+ def test_missing_confidence_without_a_bar_acts_as_before(self,
mock_hook_cls):
+ _install(mock_hook_cls, _agent("network"))
+ policy = LLMRetryPolicy(llm_conn_id="test", fallback_rules=self.RULES)
+
+ decision = policy.evaluate(RuntimeError("reset"), try_number=1,
max_tries=3)
assert decision.action == RetryAction.RETRY
- assert decision.retry_delay is None
+ assert decision.retry_delay == timedelta(seconds=10)
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
- def test_prompt_includes_exception_type_and_message(self, mock_hook_cls):
- mock_agent = _make_mock_agent("data", should_retry=False)
- mock_hook_cls.return_value.create_agent.return_value = mock_agent
+ @patch(HOOK, autospec=True)
+ def test_nan_confidence_counts_as_missing(self, mock_hook_cls):
+ _install(mock_hook_cls, _agent("network", confidence=math.nan))
+ policy = LLMRetryPolicy(llm_conn_id="test", min_confidence=0.6)
+
+ assert policy.evaluate(RuntimeError("reset"), try_number=1,
max_tries=3).action == RetryAction.DEFAULT
+ @patch(HOOK, autospec=True)
+ def test_confidence_without_a_bar_is_recorded_but_not_gated(self,
mock_hook_cls):
+ _install(mock_hook_cls, _agent("auth", confidence=0.05))
policy = LLMRetryPolicy(llm_conn_id="test")
+
+ decision = policy.evaluate(RuntimeError("x"), try_number=1,
max_tries=3)
+
+ assert decision.action == RetryAction.FAIL
+ assert decision.reason == "category=auth confidence=0.05 threshold=n/a
action=fail"
+
+
+class TestPrompt:
+ @patch(HOOK, autospec=True)
+ def test_prompt_includes_exception_type_and_message(self, mock_hook_cls):
+ agent = _install(mock_hook_cls, _agent("data"))
+ policy = LLMRetryPolicy(llm_conn_id="test")
+
policy.evaluate(ValueError("bad column type"), try_number=2,
max_tries=5)
- prompt = mock_agent.run_sync.call_args[0][0]
+ prompt = agent.run_sync.call_args.args[0]
assert "ValueError: bad column type" in prompt
assert "attempt 2 of 5" in prompt
@pytest.mark.enable_redact
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ @patch(HOOK, autospec=True)
def test_prompt_redacts_known_secrets(self, mock_hook_cls):
reset_secrets_masker()
secret_value = "super-secret-conn-password"
mask_secret(secret_value)
-
- mock_agent = _make_mock_agent("auth", should_retry=False)
- mock_hook_cls.return_value.create_agent.return_value = mock_agent
-
+ agent = _install(mock_hook_cls, _agent("auth"))
policy = LLMRetryPolicy(llm_conn_id="test")
+
policy.evaluate(
- ConnectionError(f"could not authenticate with password
{secret_value}"),
- try_number=1,
- max_tries=3,
+ ConnectionError(f"could not authenticate with password
{secret_value}"), try_number=1, max_tries=3
)
- prompt = mock_agent.run_sync.call_args[0][0]
- assert secret_value not in prompt
- assert prompt == (
+ assert agent.run_sync.call_args.args[0] == (
"Classify this error from a data pipeline task (attempt 1 of
3):\n\n"
"ConnectionError: could not authenticate with password ***"
)
@pytest.mark.enable_redact
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ @patch(HOOK, autospec=True)
def test_prompt_keeps_raw_message_when_redaction_disabled(self,
mock_hook_cls):
reset_secrets_masker()
secret_value = "super-secret-conn-password"
mask_secret(secret_value)
-
- mock_agent = _make_mock_agent("auth", should_retry=False)
- mock_hook_cls.return_value.create_agent.return_value = mock_agent
-
+ agent = _install(mock_hook_cls, _agent("auth"))
policy = LLMRetryPolicy(llm_conn_id="test", redact_exception=False)
- policy.evaluate(
- ConnectionError(f"could not authenticate with password
{secret_value}"),
- try_number=1,
- max_tries=3,
- )
- prompt = mock_agent.run_sync.call_args[0][0]
- assert secret_value in prompt
+ policy.evaluate(ConnectionError(f"password {secret_value}"),
try_number=1, max_tries=3)
+
+ assert secret_value in agent.run_sync.call_args.args[0]
@pytest.mark.enable_redact
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ @patch(HOOK, autospec=True)
def test_explicit_redactor_none_still_applies_default_masking(self,
mock_hook_cls):
- """redactor=None means "use the default masker" -- the same as
omitting it."""
reset_secrets_masker()
secret_value = "super-secret-conn-password"
mask_secret(secret_value)
-
- mock_agent = _make_mock_agent("auth", should_retry=False)
- mock_hook_cls.return_value.create_agent.return_value = mock_agent
-
+ agent = _install(mock_hook_cls, _agent("auth"))
policy = LLMRetryPolicy(llm_conn_id="test", redactor=None)
- policy.evaluate(
- ConnectionError(f"could not authenticate with password
{secret_value}"),
- try_number=1,
- max_tries=3,
- )
- prompt = mock_agent.run_sync.call_args[0][0]
- assert secret_value not in prompt
- assert "***" in prompt
+ policy.evaluate(ConnectionError(f"password {secret_value}"),
try_number=1, max_tries=3)
- def test_redact_exception_false_with_explicit_redactor_raises(self):
- with pytest.raises(ValueError, match="redactor must not be set when
redact_exception=False"):
- LLMRetryPolicy(llm_conn_id="test", redact_exception=False,
redactor=lambda message: message)
+ assert secret_value not in agent.run_sync.call_args.args[0]
@pytest.mark.enable_redact
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ @patch(HOOK, autospec=True)
def test_custom_redactor_replaces_masker_instead_of_stacking(self,
mock_hook_cls):
- """A custom redactor replaces the secrets masker entirely -- it is not
applied on top."""
reset_secrets_masker()
secret_value = "super-secret-conn-password"
mask_secret(secret_value)
-
- mock_agent = _make_mock_agent("auth", should_retry=False)
- mock_hook_cls.return_value.create_agent.return_value = mock_agent
-
- policy = LLMRetryPolicy(llm_conn_id="test", redactor=lambda s:
s.replace("authenticate", "REDACTED"))
- policy.evaluate(
- ConnectionError(f"could not authenticate with password
{secret_value}"),
- try_number=1,
- max_tries=3,
+ agent = _install(mock_hook_cls, _agent("auth"))
+ policy = LLMRetryPolicy(
+ llm_conn_id="test", redactor=lambda m:
m.replace("[email protected]", "<email>")
)
- prompt = mock_agent.run_sync.call_args[0][0]
- # The registered secret is untouched by the masker...
- assert secret_value in prompt
- # ...but the custom redactor's own transformation did apply.
- assert "REDACTED" in prompt
+ policy.evaluate(ConnectionError(f"[email protected] {secret_value}"),
try_number=1, max_tries=3)
+
+ prompt = agent.run_sync.call_args.args[0]
+ assert "<email>" in prompt
+ assert secret_value in prompt # the custom redactor replaced the
masker, so this is on the author
@pytest.mark.parametrize(
- ("max_exception_length", "message_length", "expect_truncated"),
+ ("length", "expected_tail"),
[
- pytest.param(4096, 4096, False, id="default-limit-exact-fit"),
- pytest.param(4096, 5000, True, id="default-limit-exceeded"),
- pytest.param(10, 20, True, id="custom-limit-exceeded"),
- pytest.param(10, 5, False, id="custom-limit-under"),
+ pytest.param(50, f"ValueError: {'x' * 10}... (truncated)",
id="over"),
+ pytest.param(10, f"ValueError: {'x' * 10}", id="exact-fit"),
],
)
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
- def test_message_truncated_when_over_max_exception_length(
- self, mock_hook_cls, max_exception_length, message_length,
expect_truncated
+ @patch(HOOK, autospec=True)
+ def test_message_truncated_only_when_over_max_exception_length(
+ self, mock_hook_cls, length, expected_tail
):
- mock_agent = _make_mock_agent("data", should_retry=False)
- mock_hook_cls.return_value.create_agent.return_value = mock_agent
+ agent = _install(mock_hook_cls, _agent("data"))
+ policy = LLMRetryPolicy(llm_conn_id="test", max_exception_length=10)
- policy = LLMRetryPolicy(
- llm_conn_id="test", redact_exception=False,
max_exception_length=max_exception_length
- )
- policy.evaluate(ValueError("x" * message_length), try_number=1,
max_tries=3)
+ policy.evaluate(ValueError("x" * length), try_number=1, max_tries=3)
- prompt = mock_agent.run_sync.call_args[0][0]
- assert ("... (truncated)" in prompt) is expect_truncated
- if expect_truncated:
- assert f"{'x' * max_exception_length}... (truncated)" in prompt
- else:
- assert "x" * message_length in prompt
+ assert agent.run_sync.call_args.args[0].endswith(expected_tail)
@pytest.mark.enable_redact
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ @patch(HOOK, autospec=True)
def test_truncation_happens_after_redaction(self, mock_hook_cls):
- """Redact-then-truncate must not equal truncate-then-redact for this
input.
-
- The secret sits right at the truncation boundary: truncating first
would slice
- it in half so the masker could no longer recognize and mask it.
- """
+ """Truncating first could cut a registered secret in half and leak the
head of it."""
reset_secrets_masker()
secret_value = "super-secret-conn-password"
mask_secret(secret_value)
- max_exception_length = 20
- # Padding places the secret so it straddles the truncation boundary.
- padding = "a" * (max_exception_length - 5)
- message = f"{padding}{secret_value}"
+ agent = _install(mock_hook_cls, _agent("auth"))
+ policy = LLMRetryPolicy(llm_conn_id="test", max_exception_length=12)
- mock_agent = _make_mock_agent("auth", should_retry=False)
- mock_hook_cls.return_value.create_agent.return_value = mock_agent
+ policy.evaluate(ConnectionError(f"pw {secret_value} tail"),
try_number=1, max_tries=3)
- policy = LLMRetryPolicy(llm_conn_id="test",
max_exception_length=max_exception_length)
- policy.evaluate(ConnectionError(message), try_number=1, max_tries=3)
+ assert agent.run_sync.call_args.args[0].endswith("ConnectionError: pw
*** tail")
- prompt = mock_agent.run_sync.call_args[0][0]
- assert secret_value not in prompt
- assert "***" in prompt
- @pytest.mark.parametrize("max_exception_length", [0, -1, -100])
- def test_non_positive_max_exception_length_raises(self,
max_exception_length):
- with pytest.raises(ValueError, match="max_exception_length must be a
positive integer"):
- LLMRetryPolicy(llm_conn_id="test",
max_exception_length=max_exception_length)
-
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
- def test_custom_instructions_forwarded_to_agent(self, mock_hook_cls):
- mock_hook_cls.return_value.create_agent.return_value =
_make_mock_agent("x", False)
-
- policy = LLMRetryPolicy(llm_conn_id="test", instructions="My custom
prompt")
- policy.evaluate(ValueError("x"), try_number=1, max_tries=3)
-
- mock_hook_cls.return_value.create_agent.assert_called_once_with(
- output_type=ErrorClassification,
- instructions="My custom prompt",
- )
-
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
- def test_timeout_passed_via_model_settings(self, mock_hook_cls):
- mock_agent = _make_mock_agent("auth", False)
- mock_hook_cls.return_value.create_agent.return_value = mock_agent
-
- policy = LLMRetryPolicy(llm_conn_id="test", timeout=15.0)
- policy.evaluate(ValueError("x"), try_number=1, max_tries=3)
-
- model_settings = mock_agent.run_sync.call_args.kwargs["model_settings"]
- assert model_settings["timeout"] == 15.0
-
-
-class TestLLMFallbackBehaviour:
- """Test fallback when the LLM call itself fails."""
+class TestFallbackBehaviour:
+ """When the LLM call itself fails the deterministic path decides,
unchanged."""
def test_falls_back_to_rules_when_connection_missing(self):
policy = LLMRetryPolicy(
@@ -313,53 +768,61 @@ class TestLLMFallbackBehaviour:
RetryRule(exception=PermissionError, action=RetryAction.FAIL,
reason="auth fallback"),
],
)
- d = policy.evaluate(ConnectionError("refused"), try_number=1,
max_tries=3)
- assert d.action == RetryAction.RETRY
- assert d.retry_delay == timedelta(seconds=10)
- d = policy.evaluate(PermissionError("denied"), try_number=1,
max_tries=3)
- assert d.action == RetryAction.FAIL
+ retry = policy.evaluate(ConnectionError("refused"), try_number=1,
max_tries=3)
+ fail = policy.evaluate(PermissionError("denied"), try_number=1,
max_tries=3)
+
+ assert (retry.action, retry.retry_delay) == (RetryAction.RETRY,
timedelta(seconds=10))
+ assert fail.action == RetryAction.FAIL
def test_falls_back_to_default_when_no_rules(self):
policy = LLMRetryPolicy(llm_conn_id="nonexistent")
- d = policy.evaluate(ValueError("bad"), try_number=1, max_tries=3)
- assert d.action == RetryAction.DEFAULT
+
+ assert policy.evaluate(ValueError("bad"), try_number=1,
max_tries=3).action == RetryAction.DEFAULT
def test_fallback_rules_no_match_returns_default(self):
- """When fallback rules exist but none match, DEFAULT is returned."""
policy = LLMRetryPolicy(
llm_conn_id="nonexistent",
- fallback_rules=[
- RetryRule(exception=PermissionError, action=RetryAction.FAIL),
- ],
+ fallback_rules=[RetryRule(exception=PermissionError,
action=RetryAction.FAIL)],
)
- # ValueError doesn't match the PermissionError rule
- d = policy.evaluate(ValueError("bad"), try_number=1, max_tries=3)
- assert d.action == RetryAction.DEFAULT
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
- def test_agent_run_sync_failure_triggers_fallback(self, mock_hook_cls):
- """Failure during run_sync (not hook creation) still triggers
fallback."""
- mock_agent = MagicMock()
- mock_agent.run_sync.side_effect = RuntimeError("network error
mid-call")
- mock_hook_cls.return_value.create_agent.return_value = mock_agent
+ assert policy.evaluate(ValueError("bad"), try_number=1,
max_tries=3).action == RetryAction.DEFAULT
+ @patch(HOOK, autospec=True)
+ def test_agent_run_sync_failure_triggers_fallback(self, mock_hook_cls):
+ """A model timeout or API failure surfaces here, as an exception from
run_sync."""
+ agent = _install(mock_hook_cls, MagicMock(spec=Agent))
+ agent.run_sync.side_effect = TimeoutError("model did not answer in
time")
policy = LLMRetryPolicy(
llm_conn_id="test",
fallback_rules=[RetryRule(exception=ValueError,
action=RetryAction.FAIL, reason="fallback")],
)
- d = policy.evaluate(ValueError("x"), try_number=1, max_tries=3)
- assert d.action == RetryAction.FAIL
- assert d.reason == "fallback"
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook",
autospec=True)
+ decision = policy.evaluate(ValueError("x"), try_number=1, max_tries=3)
+
+ assert decision.action == RetryAction.FAIL
+ assert decision.reason == "LLM classification not applied
(model_error); fallback"
+
+ @patch(HOOK, autospec=True)
def test_hook_creation_failure_triggers_fallback(self, mock_hook_cls):
- """Failure during hook.create_agent still triggers fallback."""
mock_hook_cls.return_value.create_agent.side_effect =
RuntimeError("unexpected")
-
policy = LLMRetryPolicy(
llm_conn_id="test",
fallback_rules=[RetryRule(exception=ValueError,
action=RetryAction.FAIL, reason="caught")],
)
- d = policy.evaluate(ValueError("x"), try_number=1, max_tries=3)
- assert d.action == RetryAction.FAIL
+
+ assert (
+ policy.evaluate(ValueError("x"), try_number=1, max_tries=3).reason
+ == "LLM classification not applied (model_error); caught"
+ )
+
+ @patch(HOOK, autospec=True)
+ def test_answer_the_type_rejected_triggers_fallback(self, mock_hook_cls):
+ """pydantic-ai raises once its re-prompts are exhausted; that is a
failed call like any other."""
+ agent = _install(mock_hook_cls, MagicMock(spec=Agent))
+ agent.run_sync.side_effect = UnexpectedModelBehavior(
+ "Exceeded maximum retries (1) for output validation"
+ )
+ policy = LLMRetryPolicy(llm_conn_id="test")
+
+ assert policy.evaluate(ValueError("x"), try_number=1,
max_tries=3).action == RetryAction.DEFAULT
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_decision.py
b/providers/common/ai/tests/unit/common/ai/utils/test_decision.py
index a9bca30eadf..8ac5e510b47 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_decision.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_decision.py
@@ -16,9 +16,12 @@
# under the License.
from __future__ import annotations
-from unittest.mock import MagicMock
+from dataclasses import dataclass
+from typing import Any
+from unittest.mock import MagicMock, patch
import pytest
+from pydantic import TypeAdapter
from pydantic_ai.messages import ModelResponse
from airflow.providers.common.ai.exceptions import LowConfidenceError as
ExportedLowConfidenceError
@@ -33,8 +36,10 @@ from airflow.providers.common.ai.utils.decision import (
check_uncertain_action,
decision_record,
describe_confidence,
+ described_choices,
finalize_record,
initial_decided_by,
+ picked_key,
policy_fails,
policy_record,
review_reason,
@@ -248,3 +253,59 @@ class TestRecords:
assert ExportedLowConfidenceError is LowConfidenceError
assert ExportedDecisionPolicy is DecisionPolicy
assert BranchOption("x").template_fields == ("description",)
+
+
+@dataclass(frozen=True)
+class _FakeChoice:
+ """pydantic-ai 2.46's ``Choice(description=None, *, value=UNSET)``, as far
as the builder uses it."""
+
+ description: str | None = None
+
+
+@dataclass(frozen=True)
+class _FakeChoices:
+ """pydantic-ai 2.46's ``Choices(choices, *, name=None,
description=None)``, recording the call shape."""
+
+ choices: Any
+ name: str | None = None
+
+
+class TestDescribedChoices:
+ """CI runs on a pydantic-ai without ``Choices``, so the 2.46+ call shape
is pinned against a stand-in."""
+
+ @patch("airflow.providers.common.ai.utils.decision.Choices", _FakeChoices)
+ @patch("airflow.providers.common.ai.utils.decision.Choice", _FakeChoice)
+ def test_choices_gets_the_keys_with_descriptions_in_order(self):
+ built = described_choices("Options", {"b": "Second.", "a": None})
+
+ assert built == _FakeChoices({"b": _FakeChoice("Second."), "a":
_FakeChoice(None)}, name="Options")
+
+ @patch("airflow.providers.common.ai.utils.decision.Choices", _FakeChoices)
+ @patch("airflow.providers.common.ai.utils.decision.Choice", _FakeChoice)
+ def test_choices_gets_a_plain_list_when_nothing_is_described(self):
+ assert described_choices("Options", {"b": None, "a": None}) ==
_FakeChoices(
+ ["b", "a"], name="Options"
+ )
+
+ @patch("airflow.providers.common.ai.utils.decision.Choices", None)
+ @patch("airflow.providers.common.ai.utils.decision.Choice", None)
+ def
test_enum_fallback_keeps_keys_as_values_and_emits_the_described_schema(self):
+ built = described_choices(
+ "Options", {"_sunder_": "Reserved by Enum.", "mro": None, "plain
key": None}
+ )
+
+ assert [member.value for member in built] == ["_sunder_", "mro",
"plain key"]
+ assert picked_key(built("mro")) == "mro"
+ assert TypeAdapter(built).json_schema()["anyOf"] == [
+ {"const": "_sunder_", "type": "string", "description": "Reserved
by Enum."},
+ {"const": "mro", "type": "string"},
+ {"const": "plain key", "type": "string"},
+ ]
+
+ @patch("airflow.providers.common.ai.utils.decision.Choices", None)
+ @patch("airflow.providers.common.ai.utils.decision.Choice", None)
+ def test_enum_fallback_without_descriptions_is_a_plain_enum(self):
+ built = described_choices("Options", {"x": None, "y": None})
+
+ assert TypeAdapter(built).json_schema()["enum"] == ["x", "y"]
+ assert picked_key("x") == "x"
