From 5de5258b099d5576ed911cddee6ba28216ff0fdb Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Mon, 21 Sep 2026 20:05:10 +0100 Subject: [PATCH 1/3] Add ClassifierRetryPolicy and keep LLMRetryPolicy as the text-model policy it was #73450 folded a classifier-model mode into LLMRetryPolicy, which changed what a 0.9.0 Dag does on a text model (the model no longer chose the retry or the delay, custom instructions that named their own categories were mapped onto the defaults, ErrorClassification disappeared) and left arguments on the class that mean nothing without a classifier. A back-compat check of the released 0.9.0 provider against main found those and a few smaller drifts. Retry policies are now one class per layer of a ladder from hardcoded to reasoning. Fallback rules (ExceptionRetryPolicy, or ``fallback_rules`` on either policy) are the floor. ``ClassifierRetryPolicy`` is the middle: the model names one of ``categories`` and the ``ErrorCategory`` table decides retry, delay and how sure the model has to be, tuned by descriptions and a confidence bar; it is the policy for a classifier model such as TypeSafe's Jev. ``LLMRetryPolicy`` is the top and is the 0.9.0 policy unchanged, with no new arguments. ``on_uncertain`` chains them: a classifier policy names a policy, typically an LLMRetryPolicy on a text model, to consult when it is unsure or unreachable, and whatever that decides nothing about falls to the rules. Smaller drifts fixed on the classifier path: a ``fallback_rules`` entry with ``action=DEFAULT`` keeps its delay and reason; an answer missing from the table and a subclass ``_classify`` returning a non-decision log what happened and fall back. ``LLMOperator._push_decision`` no longer raises when the context's ``task_instance`` is a plain dict. The changelog describes the branching changes (``decision`` XCom on every run, sorted option order, ``branches`` template field) and the ``decision`` keyword ``execute_complete`` grew, which breaks a subclass override with the old signature at resume. The branch guide gets the "add an option for none of these" guidance the retry guide has. --- providers/common/ai/docs/changelog.rst | 52 +- .../common/ai/docs/classifier_models.rst | 16 +- .../common/ai/docs/operators/llm_branch.rst | 9 + providers/common/ai/docs/retry_policies.rst | 310 +++++++---- .../example_dags/example_llm_retry_policy.py | 43 +- .../providers/common/ai/operators/llm.py | 7 +- .../providers/common/ai/policies/retry.py | 479 ++++++++++++----- .../unit/common/ai/operators/test_llm.py | 25 + .../unit/common/ai/policies/test_retry.py | 488 ++++++++++++++---- 9 files changed, 1083 insertions(+), 346 deletions(-) diff --git a/providers/common/ai/docs/changelog.rst b/providers/common/ai/docs/changelog.rst index 74b98db5408a3..fd36a8e2394ac 100644 --- a/providers/common/ai/docs/changelog.rst +++ b/providers/common/ai/docs/changelog.rst @@ -26,31 +26,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 ();``. See :doc:`retry_policies`. + ``LLMBranchOperator`` and ``LLMOperator`` now push a ``decision`` XCom on every run, next to + ``return_value``, with the model's pick, the action taken, the confidence and probabilities + when the model reports them, and the ``decision_policy`` that applied. ``LLMBranchOperator`` + also offers the downstream task IDs to the model in sorted order (it was set order, which + differed between workers), gained ``branches`` as a template field, and builds its option + type from pydantic-ai's ``Choices`` on 2.46+ or an equivalent enum whose member names are + generated; the option values are still the task IDs, so ``do_branch`` receives the same + strings as before. + +.. note:: + ``execute_complete`` on ``LLMOperator``, ``LLMBranchOperator``, ``LLMSQLQueryOperator`` and + ``LLMSchemaCompareOperator`` gained a keyword argument, ``decision``, and every review pause + now passes it on resume. A subclass that overrides ``execute_complete`` with the old + three-argument signature raises ``TypeError`` when the reviewed task resumes; add + ``decision=None`` to the override. A review that was already pending when you upgraded + resumes without it and is unaffected. + +.. note:: + New ``ClassifierRetryPolicy``: the model names one of the author's ``categories`` (a + table of ``ErrorCategory(description, retry, delay, min_confidence)``) and the table + decides whether to retry, after how long, and how sure the model has to be. It is the + policy for a classifier model such as TypeSafe's Jev, and ``on_uncertain`` lets it hand + an unsure or failed classification to another policy, typically an ``LLMRetryPolicy`` on + a text model, before ``fallback_rules``. ``LLMRetryPolicy`` itself is unchanged: the model + returns ``ErrorClassification`` and chooses the retry and the delay from your + ``instructions``, and it grows no new arguments. Other new public names: + ``ErrorCategory``, ``DEFAULT_CATEGORIES``, ``CLASSIFIER_INSTRUCTIONS``. .. note:: Configuring ``fallback_conn_ids`` on a connection (or the matching operator/decorator diff --git a/providers/common/ai/docs/classifier_models.rst b/providers/common/ai/docs/classifier_models.rst index 44fcdcf2ac354..80cac0e650a69 100644 --- a/providers/common/ai/docs/classifier_models.rst +++ b/providers/common/ai/docs/classifier_models.rst @@ -112,14 +112,16 @@ Where it fits in this provider - A ``Literal``, ``Enum``, ``bool`` or bounded number works. Describe the field, which 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 ` + * - :doc:`ClassifierRetryPolicy ` - 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. + worker. Set ``min_confidence`` and an unsure answer goes to ``on_uncertain`` + (typically an ``LLMRetryPolicy`` on a text model), then ``fallback_rules``, then + the task's own retry behaviour, instead of ending the task on the model's say-so. + ``LLMRetryPolicy`` itself asks for free text, which a classifier model refuses. + 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. @@ -140,8 +142,8 @@ 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 ` takes the same ``min_confidence`` and hands an unsure -answer to its deterministic fallback rules. In the branch operator and the retry policy, a +:doc:`ClassifierRetryPolicy ` takes the same ``min_confidence`` and hands an +unsure answer to ``on_uncertain``, then 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`` diff --git a/providers/common/ai/docs/operators/llm_branch.rst b/providers/common/ai/docs/operators/llm_branch.rst index f302b570f4c55..5d0ca54f9bd7d 100644 --- a/providers/common/ai/docs/operators/llm_branch.rst +++ b/providers/common/ai/docs/operators/llm_branch.rst @@ -87,6 +87,15 @@ and a text model's structured output carries no confidence to read. With a classifier model such as TypeSafe's, the descriptions become the criteria of its choice question, which is the text it weighs each option by. +A pick is relative: the model chooses the best fit among the downstream tasks +offered, not whether any of them fits. If "none of these" or "not enough to +tell" is a real outcome, give it a downstream task of its own (an +``EmptyOperator`` that ends the run, or a task that opens a ticket) and +describe it, rather than expecting the model to refuse. When you change a +description or the set of branches, treat confidence values measured before +as stale: the distribution the model returns is over the options it was +given. + Multiple Branches ----------------- diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index 91d3fc868a797..1d44d16a12a04 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -21,11 +21,45 @@ LLM Retry Policies .. note:: Requires Airflow >= 3.3.0. -``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. +Two policies ask a model about a task failure and turn the answer into a retry +decision. They are two layers of a ladder that runs from fully hardcoded to +fully model-driven, with the SDK's ``ExceptionRetryPolicy`` as the bottom rung: + +.. list-table:: + :header-rows: 1 + :widths: 22 40 38 + + * - Layer + - What decides + - How you tune it + * - **Fallback rules** + (``ExceptionRetryPolicy``, or ``fallback_rules`` on either policy below) + - ``RetryRule`` matches on the exception type, then the task's own + ``retries`` and ``retry_delay``. No model. + - Write the rules. Always the floor: whatever no layer above decides lands + here. + * - **Classifier** + (``ClassifierRetryPolicy``) + - The model names one of your ``categories``; the table says whether that + category is retried, after how long, and how sure the model has to be. + A classifier model such as TypeSafe's Jev answers in a few hundred + milliseconds and reports its confidence; a text model can sit here too. + - Category descriptions and the confidence bar. No reasoning, no prose. + * - **LLM** + (``LLMRetryPolicy``) + - A text model classifies the failure and decides whether to retry and how + long to wait, guided by ``instructions``, and explains itself. + - The instructions: your taxonomy, your retry rules, your delays, in + prose. + +Each class has only the arguments its layer needs. ``LLMRetryPolicy`` is the +policy this guide has always described and is unchanged. The layers chain: +``on_uncertain`` on a ``ClassifierRetryPolicy`` names the policy to consult when +the classifier is unsure or unreachable, typically an ``LLMRetryPolicy`` on a +text model, so the cheap typed model handles the clear cases and the reasoning +model the rest, and the rules catch what neither decides. See +`Escalating to an LLM`_ below. Both work with any LLM provider supported by +pydantic-ai (OpenAI, Anthropic, Bedrock, Vertex, Ollama, etc.). 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, @@ -72,56 +106,68 @@ Usage How it works ------------ -When a task fails, ``LLMRetryPolicy``: +When a task fails, either policy: 1. Sends the exception message to the configured LLM. By default, the message 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 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). +2. With ``LLMRetryPolicy``, the model returns an + :class:`~airflow.providers.common.ai.policies.retry.ErrorClassification`: a + category, whether to retry, a suggested delay, and its reasoning. With + ``ClassifierRetryPolicy``, it picks one of the names in ``categories``; it + sees each category's name and description in the output schema and cannot + answer with a name outside the set. +3. The policy returns RETRY or FAIL. For ``LLMRetryPolicy`` that is the model's + ``should_retry`` and ``suggested_delay_seconds``. For ``ClassifierRetryPolicy`` + it is the picked category's ``retry`` and ``delay``, unless the policy has a + confidence bar and the answer is under it, in which case the answer is + discarded (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``. + instance's ``retry_reason``: ``: `` from + ``LLMRetryPolicy``, or one line such as + ``category=network confidence=0.91 threshold=0.60 action=retry delay=10s`` + from ``ClassifierRetryPolicy``. -This classification call is a separate LLM request, made by ``LLMRetryPolicy`` +This classification call is a separate model request, made by the policy itself rather than by an operator -- it is not subject to an operator's ``usage_limits``, and it runs on every task failure regardless 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), 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. +If the model call fails (provider down, timeout, bad credentials), the policy +falls back to ``fallback_rules`` if configured, or to the task's standard +retry behaviour. ``ClassifierRetryPolicy`` does the same when the model cannot +produce one of the categories even after pydantic-ai re-prompts it, or when +its answer is under the confidence bar, after consulting ``on_uncertain`` if +set; its fallback decision's reason then starts with +``classifier answer 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. ``LLMRetryPolicy``'s fallback +decision is what ``fallback_rules`` returned, as it always was. 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 ----------- +ClassifierRetryPolicy +--------------------- +``ClassifierRetryPolicy`` takes the retry decision away from the model. Its ``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 two cannot drift apart. This is also the policy a classifier model needs: +such a model refuses the free-text fields of ``ErrorClassification``, so an +``LLMRetryPolicy`` pointed at one fails every classification and falls back, +with a log line saying to use ``ClassifierRetryPolicy``. -The default, -:data:`~airflow.providers.common.ai.policies.retry.DEFAULT_CATEGORIES`, is an -example taxonomy, not the taxonomy: +:data:`~airflow.providers.common.ai.policies.retry.DEFAULT_CATEGORIES` is the +default: the same seven categories the LLM policy's default instructions +describe, with the same retry/fail split and delays: .. list-table:: :header-rows: 1 @@ -161,9 +207,13 @@ default and edit what you need: from dataclasses import replace from datetime import timedelta - from airflow.providers.common.ai.policies.retry import DEFAULT_CATEGORIES, ErrorCategory, LLMRetryPolicy + from airflow.providers.common.ai.policies.retry import ( + ClassifierRetryPolicy, + DEFAULT_CATEGORIES, + ErrorCategory, + ) - LLMRetryPolicy( + ClassifierRetryPolicy( llm_conn_id="pydanticai_default", categories={ **DEFAULT_CATEGORIES, @@ -178,7 +228,7 @@ on them differently: .. code-block:: python - snowflake_policy = LLMRetryPolicy( + snowflake_policy = ClassifierRetryPolicy( llm_conn_id="pydanticai_default", categories={ "queued": ErrorCategory( @@ -241,8 +291,8 @@ had left, so the category that ends the task deserves the higher bar. :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. +raises at construction. With no bar the policy acts on every answer and still +logs the confidence. **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. @@ -265,6 +315,37 @@ actions and does not eliminate them. Pin the model version is not guaranteed to mean the same thing after the next. See :doc:`classifier_models` for what these models answer well and badly. +Escalating to an LLM +-------------------- + +A classifier is cheap and fast, and a text model can reason about a failure it +has never seen a category for. ``on_uncertain`` puts one behind the other: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py + :language: python + :start-after: [START howto_retry_policy_escalation] + :end-before: [END howto_retry_policy_escalation] + +The order of events on a failure: + +1. The classifier names a category. At or above the bar, its category's action + and delay apply and the text model is never called. +2. Under the bar, with no confidence reported, or if the classifier call fails, + the ``on_uncertain`` policy runs. A text-model ``LLMRetryPolicy`` there + classifies the failure with its own instructions and chooses retry and delay + itself. Its decision is used, with the reason prefixed by why the classifier's + answer was not: ``escalated (below_threshold); rate_limit: 429 with a + Retry-After header``. +3. If that policy decides nothing either (its model was unreachable and none of + its own ``fallback_rules`` matched), the outer ``fallback_rules`` apply, then + the task's own retry behaviour. + +``on_uncertain`` accepts any ``RetryPolicy``, so an ``ExceptionRetryPolicy`` +works there too, and needs ``min_confidence``: without a bar the classifier is +never unsure. Both model calls run on the worker at failure time, so a task +that escalates pays for two before its retry is scheduled; ``timeout`` on each +policy bounds that. + When the connection also carries a fallback chain -------------------------------------------------- @@ -298,49 +379,80 @@ the chain is in play. What the model can and cannot do -------------------------------- -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), 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. +Under either policy the model 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), how many attempts are left, +and, under ``ClassifierRetryPolicy``, the category names and descriptions. The prompt says +``attempt {try_number} of {max_tries}``, so the model knows the limit and not +just where it is right now; an instruction like "after two attempts treat an +expired token as ``auth`` rather than ``transient``" works because the model +can see which attempt this is. + +Under ``LLMRetryPolicy`` it answers four fields: ``category``, ``should_retry``, +``suggested_delay_seconds`` and ``reasoning``, and the first two after +``category`` decide the run. A positive delay is used as returned, with no +upper limit; zero or negative means no override, so the task's own +``retry_delay`` and backoff apply. ``category`` and ``reasoning`` become the +``retry_reason``. + +Under ``ClassifierRetryPolicy`` it answers 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, and the explanation is +the generated line in the task log, which says what mattered (the category, +the confidence, the bar, the action). 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 ``retry_reason`` is only recorded on a RETRY. It is written to the task +instance (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. +task straight away even when attempts were left, so a wrong classification into +a failing category costs the task the retries it would otherwise have had; +``ClassifierRetryPolicy``'s confidence bar exists for that. Custom instructions ------------------- -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. +For ``LLMRetryPolicy`` the instructions are the taxonomy, and +:data:`~airflow.providers.common.ai.policies.retry.DEFAULT_INSTRUCTIONS` shows +the shape: name the categories, say which to retry, and give the delays. Override +them to teach the model your stack, and name your own categories if the seven +defaults do not fit; the model returns whatever name you taught it, with its own +retry decision and delay: .. code-block:: python SNOWFLAKE_INSTRUCTIONS = ( + "You are an error classifier for Snowflake-backed data pipelines. " + "Classify the error into one of: queued, warehouse_suspended, token_expired, " + "schema_drift, permanent.\n\n" + "- 'Statement queued' or 'concurrency limit' -> queued, retry after 120s\n" + "- '000606' or 'is suspended' -> warehouse_suspended, retry after 30s\n" + "- 'JWT token' or 'session token' with 'expired' -> token_expired, retry after 30s\n" + "- '002003' or 'does not exist' -> schema_drift, do NOT retry\n" + "- anything else that will fail identically -> permanent, do NOT retry\n" + "Set suggested_delay_seconds as above and 0 for errors that should not retry." + ) + + snowflake_policy = LLMRetryPolicy(llm_conn_id="pydanticai_default", instructions=SNOWFLAKE_INSTRUCTIONS) + +For ``ClassifierRetryPolicy`` the same taxonomy lives in ``categories``, and +the default instructions +(:data:`~airflow.providers.common.ai.policies.retry.CLASSIFIER_INSTRUCTIONS`) say +only that the model is classifying a failed pipeline task and should pick the +best-fitting category. Override ``instructions`` there to teach the model your +error strings when a description is not enough on its own, and leave the +category names, actions and delays to the table; a prompt that names them only +spends tokens, because the model is constrained to the table's names and does +not decide the action: + +.. code-block:: python + + SNOWFLAKE_HINTS = ( "You are classifying failures from Snowflake-backed data pipelines.\n" "- 'Statement queued' or 'concurrency limit' -> queued\n" "- '000606' or 'is suspended' -> warehouse_suspended\n" @@ -348,10 +460,10 @@ defines either. "- '002003' or 'does not exist' -> schema_drift\n" ) - snowflake_policy = LLMRetryPolicy( + snowflake_policy = ClassifierRetryPolicy( llm_conn_id="pydanticai_default", - instructions=SNOWFLAKE_INSTRUCTIONS, - categories={...}, # the same names the instructions use + instructions=SNOWFLAKE_HINTS, + categories={...}, # the same names the hints use fallback_rules=[ RetryRule( exception=ConnectionError, @@ -367,17 +479,15 @@ defines either. When writing custom instructions: -- 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. +- For ``LLMRetryPolicy``, mention the four ``ErrorClassification`` fields so the + model fills them, and keep ``reasoning`` concise: ``retry_reason`` is truncated + to 500 characters. +- For ``ClassifierRetryPolicy``, use the table's names as-is. A name you invent cannot + come back: a model that insists on one 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. - 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. @@ -385,6 +495,10 @@ When writing custom instructions: Parameters ---------- +Both policies share every parameter below except ``categories``, +``min_confidence`` and ``on_uncertain``, which exist only on +``ClassifierRetryPolicy``. + .. list-table:: :header-rows: 1 :widths: 20 15 65 @@ -400,28 +514,38 @@ Parameters - Override the model from the connection (e.g., ``"openai:gpt-4o-mini"``). * - ``instructions`` - (built-in) - - Custom system prompt for error classification. Teaches the model your - error strings; the categories themselves come from ``categories``. + - Custom system prompt for error classification. On ``LLMRetryPolicy`` it + is the whole taxonomy (names, which to retry, delays); on + ``ClassifierRetryPolicy`` it can only teach the model your error strings. * - ``fallback_rules`` - None - - List of ``RetryRule`` objects used when the LLM call fails or the answer - is under its confidence bar. + - List of ``RetryRule`` objects used when the model call fails or, on + ``ClassifierRetryPolicy``, when the answer is under its confidence bar and + ``on_uncertain`` decided nothing. * - ``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. + - ``ClassifierRetryPolicy`` only. 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. + - ``ClassifierRetryPolicy`` only. 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 + ``on_uncertain`` if set, else ``fallback_rules`` then the task's own + retry behaviour, apply. A category's own ``min_confidence`` overrides it + for that category. + * - ``on_uncertain`` + - None + - ``ClassifierRetryPolicy`` only. A ``RetryPolicy`` to consult when the + classifier is under its bar, reports no confidence, or cannot be + reached; typically an ``LLMRetryPolicy`` on a text model. Its RETRY or FAIL is used; if it decides nothing, the outer + ``fallback_rules`` apply. Needs ``min_confidence``. * - ``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 810b9e1950fde..1aaaebacf1446 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,9 +17,10 @@ """ Example DAG demonstrating LLM-powered retry policies. -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. +``llm_policy`` is the plain form: a text model classifies the failure and decides whether to +retry and how long to wait, guided by its instructions. ``patient_policy`` and the +classifier-model policies are ``ClassifierRetryPolicy``: the model only names the kind of failure +and the policy's table decides the action, the delay and how sure the model has to be. Prerequisites: - Connection ``pydanticai_default`` with ``conn_type='pydanticai'``, @@ -38,7 +39,12 @@ from airflow.providers.common.compat.sdk import dag, task try: - from airflow.providers.common.ai.policies.retry import DEFAULT_CATEGORIES, ErrorCategory, LLMRetryPolicy + from airflow.providers.common.ai.policies.retry import ( + DEFAULT_CATEGORIES, + ClassifierRetryPolicy, + ErrorCategory, + LLMRetryPolicy, + ) from airflow.sdk.definitions.retry_policy import RetryAction, RetryRule llm_policy = LLMRetryPolicy( @@ -52,7 +58,7 @@ # 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( + patient_policy = ClassifierRetryPolicy( llm_conn_id="pydanticai_default", categories={ **DEFAULT_CATEGORIES, @@ -64,17 +70,17 @@ def example_llm_retry_policy(): @task(retries=3, retry_delay=timedelta(minutes=1), retry_policy=llm_policy) def task_auth_error(): - """Should classify as ``auth``, which the table fails -> FAIL immediately.""" + """The LLM should classify this as auth and decide not to retry -> FAIL immediately.""" raise PermissionError("403 Forbidden: API key expired for service account analytics@proj.iam") @task(retries=3, retry_delay=timedelta(minutes=1), retry_policy=llm_policy) def task_rate_limit(): - """Should classify as ``rate_limit``, which the table retries after 60s.""" + """The LLM should classify this as rate_limit and suggest about a 60s delay -> RETRY.""" 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(): - """Should classify as ``data``, which the table fails -> FAIL immediately.""" + """The LLM should classify this as data and decide not to retry -> 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) @@ -97,7 +103,7 @@ def task_missing_table(): # 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( + snowflake_policy = ClassifierRetryPolicy( llm_conn_id="jev_default", min_confidence=0.8, categories={ @@ -127,6 +133,21 @@ def task_missing_table(): ) # [END howto_retry_policy_classifier] + # [START howto_retry_policy_escalation] + # The three layers chained. The classifier answers the clear cases in a few hundred + # milliseconds. When it is unsure or unreachable, a text model reasons about the failure + # and decides retry and delay itself. When that model is unreachable too, the rules decide. + escalating_policy = ClassifierRetryPolicy( + llm_conn_id="jev_default", + min_confidence=0.8, + categories=snowflake_policy.categories, + on_uncertain=LLMRetryPolicy(llm_conn_id="pydanticai_default", timeout=30.0), + fallback_rules=[ + RetryRule(exception=ConnectionError, action=RetryAction.RETRY, retry_delay=timedelta(seconds=30)), + ], + ) + # [END howto_retry_policy_escalation] + @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) @@ -139,9 +160,9 @@ 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) + @task(retries=3, retry_delay=timedelta(minutes=1), retry_policy=escalating_policy) def task_ambiguous(): - """Reads as more than one category; under the bar the task's own retry settings apply.""" + """Reads as more than one category; under the bar the text model decides, then the rules.""" raise RuntimeError("Query failed: an unexpected error occurred while processing the request.") task_warehouse_suspended() diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py index 90606fe50a75b..54832926b2cbe 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py @@ -422,9 +422,14 @@ def _push_decision(self, context: Context, record: dict[str, Any]) -> None: try: ti = context["task_instance"] except (KeyError, TypeError): + ti = None + push = getattr(ti, "xcom_push", None) + if not callable(push): + # A hand-built context (a dict, or no task instance at all) has nowhere to push to; the + # record is inspection output, so the run goes on without it. self.log.warning("No task instance in the context; the decision record was not pushed to XCom.") return - ti.xcom_push(key=DECISION_XCOM_KEY, value=record) + push(key=DECISION_XCOM_KEY, value=record) def _finalize_decision( self, context: Context, event: dict[str, Any], decision: dict[str, Any] | None, *, action: Any 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 c41a731b66e8d..71f4677bb9dc1 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 @@ -15,12 +15,22 @@ # specific language governing permissions and limitations # under the License. """ -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. +Model-backed retry policies, one per layer of a ladder from hardcoded to reasoning. + +* **Fallback rules.** :class:`~airflow.sdk.definitions.retry_policy.ExceptionRetryPolicy` and + ``fallback_rules``: ``RetryRule`` matches on the exception type, then the task's own + ``retries`` and ``retry_delay``. No model. Always the floor. +* **Classifier.** :class:`ClassifierRetryPolicy`: the model names one of the author's + ``categories`` and the :class:`ErrorCategory` table decides whether that category is + retried, after how long, and how sure the model has to be. Tuned through descriptions and + a confidence bar, not through reasoning. A classifier model such as TypeSafe's Jev runs + here; a text model can too. +* **LLM.** :class:`LLMRetryPolicy`: a text model classifies the failure, decides whether to + retry and how long to wait from ``instructions``, and explains itself. + +They chain: ``ClassifierRetryPolicy(..., on_uncertain=LLMRetryPolicy(...))`` consults the +LLM when the classifier is unsure or unreachable, and whatever no layer decides falls to the +rules. Requires Airflow 3.3+ (RetryPolicy was added in AIP-105). """ @@ -28,12 +38,13 @@ 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 typing import TYPE_CHECKING, Any, ClassVar, Literal, cast + +from pydantic import BaseModel from airflow.providers.common.ai.policies.decision import check_bar from airflow.providers.common.ai.utils.decision import ( @@ -69,20 +80,53 @@ log = logging.getLogger(__name__) __all__ = [ + "CLASSIFIER_INSTRUCTIONS", "DEFAULT_CATEGORIES", + "DEFAULT_INSTRUCTIONS", + "ClassifierRetryPolicy", "ErrorCategory", + "ErrorClassification", "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." +) +"""The default system prompt of :class:`LLMRetryPolicy`: the taxonomy, retry rules and delays live here.""" + + +class ErrorClassification(BaseModel): + """Structured output of :class:`LLMRetryPolicy`: the model's category, decision, delay and reasoning.""" + + category: str + """One of the categories the instructions describe, by default: 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.""" + @dataclass(frozen=True) class ErrorCategory: """ - One kind of failure the model may name, and what the policy does when it does. + One kind of failure a :class:`ClassifierRetryPolicy` may name, and what it does when it does. - The value of :class:`LLMRetryPolicy`'s ``categories`` mapping, keyed by the category - name the model answers with. + The value of the policy'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 @@ -145,74 +189,33 @@ def __post_init__(self) -> None: } ) """ -The example taxonomy :class:`LLMRetryPolicy` classifies into when ``categories`` is not set. +The default ``categories`` of :class:`ClassifierRetryPolicy`: the same seven the +:data:`DEFAULT_INSTRUCTIONS` describe, with the same retry/fail split and delays. 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 = ( +CLASSIFIER_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. +The default system prompt of :class:`ClassifierRetryPolicy`. It does not recite the +categories: those travel in the output schema with their descriptions, so a prompt built as +``CLASSIFIER_INSTRUCTIONS + hints`` should teach error strings and not name categories or delays. """ def redact_registered_secrets(message: str) -> str: - """Mask values registered via ``mask_secret()``; the default ``redactor`` for :class:`LLMRetryPolicy`.""" + """Mask values registered via ``mask_secret()``; the default ``redactor`` for the policies here.""" # redact() is typed for arbitrary containers; a str in always yields a str out. return cast("str", redact(message)) -class LLMRetryPolicy(RetryPolicy): - """ - 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, a classifier model such as TypeSafe's Jev, etc.). - - 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, 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 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. +_REDACTION_PARAMS_DOC = """ :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`, @@ -237,8 +240,8 @@ class LLMRetryPolicy(RetryPolicy): external LLM provider (OpenAI, Anthropic, Bedrock, Vertex, Ollama, etc.) as part of the classification prompt, so it may leak whatever the failing task put in the exception message — connection strings, - credential fragments, PII, or other secrets. By default - ``_classify()`` runs the message through + credential fragments, PII, or other secrets. By default the message + is run through :func:`~airflow.providers.common.ai.policies.retry.redact_registered_secrets` via ``redactor``, which masks values already registered via ``mask_secret()`` (for example, connection passwords Airflow @@ -250,7 +253,13 @@ class LLMRetryPolicy(RetryPolicy): redaction altogether. You are still responsible for confirming that your task's exception messages are safe to send to a third-party LLM provider. - """ +""" + + +class _ModelRetryPolicy(RetryPolicy): + """What the two model-backed policies share: the connection, the prompt, redaction, and the rules floor.""" + + _default_instructions: ClassVar[str] def __init__( self, @@ -260,8 +269,6 @@ def __init__( 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, @@ -277,29 +284,231 @@ def __init__( ) 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.instructions = instructions or self._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 _hook(self) -> Any: + from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook + + return PydanticAIHook(llm_conn_id=self.llm_conn_id, model_id=self.model_id) + + def _prompt(self, exception: BaseException, try_number: int, max_tries: int) -> str: + # 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) + if len(message) > self.max_exception_length: + message = f"{message[: self.max_exception_length]}... (truncated)" + return ( + f"Classify this error from a data pipeline task " + f"(attempt {try_number} of {max_tries}):\n\n" + f"{type(exception).__name__}: {message}" + ) + + def _run(self, agent: Any, exception: BaseException, try_number: int, max_tries: int) -> Any: + from pydantic_ai.settings import ModelSettings + + return agent.run_sync( + self._prompt(exception, try_number, max_tries), + model_settings=ModelSettings(timeout=self.timeout), + ) + + def _rules_decision( + self, exception: BaseException, try_number: int, max_tries: int, context: Context | None + ) -> RetryDecision: + """Return the rules floor's decision: a ``fallback_rules`` match, else the task's own retry behaviour.""" + if self.fallback_rules: + return ExceptionRetryPolicy(rules=self.fallback_rules).evaluate( + exception, try_number, max_tries, context + ) + return RetryDecision.default() + + +class LLMRetryPolicy(_ModelRetryPolicy): + """ + Retry policy that uses an LLM to classify errors and decide retry behaviour. + + 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. The model + returns an :class:`ErrorClassification`: which category the error is, whether + to retry, how long to wait, and why, all steered by ``instructions``. This is + the reasoning layer; for a cheap typed decision from a classifier model such as + TypeSafe's Jev, use :class:`ClassifierRetryPolicy`, which can name this policy as + its ``on_uncertain``. + + 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. + + :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, :data:`DEFAULT_INSTRUCTIONS`. + The instructions are the whole taxonomy: the category names, which to + retry, and the delays. + :param fallback_rules: Optional list of + :class:`~airflow.sdk.definitions.retry_policy.RetryRule` applied when the + LLM call fails. 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. + """ + + __doc__ = __doc__ + _REDACTION_PARAMS_DOC + + _default_instructions = DEFAULT_INSTRUCTIONS + + def evaluate( + self, + exception: BaseException, + try_number: int, + max_tries: int, + context: Context | None = None, + ) -> RetryDecision: + try: + return self._classify(exception, try_number, max_tries) + except Exception: + log.exception("LLM retry classification failed, using fallback") + return self._rules_decision(exception, try_number, max_tries, context) + + def _classify( + self, + exception: BaseException, + try_number: int, + max_tries: int, + ) -> RetryDecision: + agent = self._hook().create_agent(output_type=ErrorClassification, instructions=self.instructions) + try: + result = self._run(agent, exception, try_number, max_tries) + except Exception as exc: + if "not supported by this model" in str(exc): + # A classifier model refuses ErrorClassification's free-text fields client-side. + log.error( + "This model cannot answer ErrorClassification; it needs a typed question. " + "Use ClassifierRetryPolicy for a classifier model such as TypeSafe's Jev." + ) + raise + 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, + ) + + if not classification.should_retry: + return RetryDecision.fail(reason=f"{classification.category}: {classification.reasoning}") + + 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}", + ) + + +class ClassifierRetryPolicy(_ModelRetryPolicy): + """ + Retry policy where the model names the kind of failure and the author's table decides. + + 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. + That is the shape a classifier model such as TypeSafe's Jev answers, in a few hundred + milliseconds and with a confidence; a text model answers it too. + + When the model call fails, or the answer is under its confidence bar, the policy + consults ``on_uncertain`` if set, then ``fallback_rules``, then returns DEFAULT to use + the task's standard retry logic. + + :param llm_conn_id: Airflow connection ID for the model. + :param model_id: Model identifier override (e.g. ``"typesafe:jev-1.13.0"``). + If not set, uses the model from the connection. + :param instructions: Custom system prompt. Defaults to :data:`CLASSIFIER_INSTRUCTIONS`. + 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 model + call fails or the answer is under its confidence bar and ``on_uncertain`` decided + nothing. Provides a deterministic safety net. + :param timeout: Maximum seconds to wait for the model response before + falling back. Defaults to 30s. + :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 whatever + the confidence. 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 ``on_uncertain``, then + ``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 on_uncertain: A policy to consult when the answer is under its bar, reports no + confidence, or the model call fails; typically an :class:`LLMRetryPolicy` on a text + model, so the classifier handles the clear cases and a reasoning model the rest. + Its RETRY or FAIL is used, with the reason prefixed by why the classifier's answer + was not. If it decides nothing either (its model was unreachable and none of its own + ``fallback_rules`` matched), this policy's ``fallback_rules`` and then the task's own + retry behaviour apply. Needs ``min_confidence``. + """ + + __doc__ = __doc__ + _REDACTION_PARAMS_DOC + + _default_instructions = CLASSIFIER_INSTRUCTIONS + + def __init__( + self, + llm_conn_id: str, + model_id: str | None = None, + instructions: str | None = None, + fallback_rules: list[RetryRule] | None = None, + timeout: float = 30.0, + *, + categories: Mapping[str, ErrorCategory] | None = None, + min_confidence: float | None = None, + on_uncertain: RetryPolicy | None = None, + redactor: Callable[[str], str] | None = None, + redact_exception: bool = True, + max_exception_length: int = 4096, + ) -> None: + super().__init__( + llm_conn_id, + model_id, + instructions, + fallback_rules, + timeout, + redactor=redactor, + redact_exception=redact_exception, + max_exception_length=max_exception_length, + ) + self.min_confidence = None if min_confidence is None else check_bar(min_confidence, "min_confidence") + self.categories: dict[str, ErrorCategory] = self._validate_categories( + DEFAULT_CATEGORIES if categories is None else categories + ) + if on_uncertain is not None: + if not isinstance(on_uncertain, RetryPolicy): + raise TypeError(f"on_uncertain must be a RetryPolicy, got {type(on_uncertain).__name__}.") + if self.min_confidence is None: + raise ValueError( + "on_uncertain needs min_confidence: it is consulted when the classifier's answer is " + "under the bar or the classifier could not answer, and without a bar there is no such case." + ) + self.on_uncertain = on_uncertain + def _validate_categories(self, categories: Mapping[str, ErrorCategory]) -> dict[str, ErrorCategory]: if not isinstance(categories, Mapping): raise TypeError( @@ -322,7 +531,7 @@ def _validate_categories(self, categories: Mapping[str, ErrorCategory]) -> dict[ 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." + "Set min_confidence on ClassifierRetryPolicy 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. @@ -347,12 +556,55 @@ def evaluate( try: outcome = self._classify(exception, try_number, max_tries) except Exception: - log.exception("LLM retry classification failed, using fallback") + log.exception("Classifier retry classification failed, using fallback") outcome = "model_error" if isinstance(outcome, RetryDecision): return outcome + if not isinstance(outcome, str): + # A subclass's _classify returned something else; say so rather than acting on its repr. + log.error( + "Classifier retry classification returned %s instead of a RetryDecision, using fallback", + type(outcome).__name__, + ) + outcome = "model_error" + if self.on_uncertain is not None: + escalated = self._escalate(exception, try_number, max_tries, context, why=outcome) + if escalated is not None: + return escalated return self._fall_back(exception, try_number, max_tries, context, why=outcome) + def _escalate( + self, + exception: BaseException, + try_number: int, + max_tries: int, + context: Context | None, + *, + why: str, + ) -> RetryDecision | None: + """ + Consult ``on_uncertain`` and return its decision, or None when it decided nothing. + + A DEFAULT decision with no reason is what a policy returns when its own model call failed + and no rule of its own matched, so that case falls through to this policy's + ``fallback_rules``. Any other decision is returned with the reason prefixed by why the + classifier's answer was not used, so a ``retry_reason`` shows the whole chain. + """ + policy = cast("RetryPolicy", self.on_uncertain) + log.info("Classifier answer not applied (%s), consulting %s", why, type(policy).__name__) + try: + decision = policy.evaluate(exception, try_number, max_tries, context) + except Exception: + log.exception("on_uncertain policy failed, using fallback rules") + return None + if decision.action is RetryAction.DEFAULT and decision.reason is None: + return None + return RetryDecision( + action=decision.action, + retry_delay=decision.retry_delay, + reason=f"escalated ({why}); {decision.reason}", + ) + def _fall_back( self, exception: BaseException, @@ -363,20 +615,18 @@ def _fall_back( why: str, ) -> RetryDecision: """ - Take the deterministic path: ``fallback_rules`` when one matches, else the task's own retry behaviour. + Take the rules floor, with a reason that says the classifier's answer was not applied and why. - 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. + A ``retry_reason`` read later is then not mistaken for a plain rule match or a classifier + decision. A matched rule always carries a reason; an unmatched evaluation is DEFAULT with + none. A matched rule keeps its action, delay and reason whatever the action. """ - prefix = f"LLM classification not applied ({why})" - if self.fallback_rules: - ruled = ExceptionRetryPolicy(rules=self.fallback_rules).evaluate( - exception, try_number, max_tries, context + prefix = f"classifier answer not applied ({why})" + ruled = self._rules_decision(exception, try_number, max_tries, context) + if ruled.action is not RetryAction.DEFAULT or ruled.reason is not None: + return RetryDecision( + action=ruled.action, retry_delay=ruled.retry_delay, reason=f"{prefix}; {ruled.reason}" ) - if ruled.action is not RetryAction.DEFAULT: - return RetryDecision( - action=ruled.action, retry_delay=ruled.retry_delay, reason=f"{prefix}; {ruled.reason}" - ) return RetryDecision(action=RetryAction.DEFAULT, reason=f"{prefix}; task retry settings apply") def _classify( @@ -384,41 +634,26 @@ def _classify( exception: BaseException, try_number: int, max_tries: int, - ) -> RetryDecision | ReviewReason: + ) -> RetryDecision | ReviewReason | Literal["model_error"]: """ 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. + Returns the reason (``"below_threshold"``, ``"missing_confidence"``, or ``"model_error"`` + for an answer the table does not know) instead of a decision when the answer is not + acted on, so :meth:`evaluate` takes the escalation and fallback path 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) 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) - if len(message) > self.max_exception_length: - message = f"{message[: self.max_exception_length]}... (truncated)" - prompt = ( - f"Classify this error from a data pipeline task " - f"(attempt {try_number} of {max_tries}):\n\n" - f"{type(exception).__name__}: {message}" - ) - - from pydantic_ai.settings import ModelSettings - - result = agent.run_sync( - prompt, - model_settings=ModelSettings(timeout=self.timeout), - ) + agent = self._hook().create_agent(output_type=output_type, instructions=self.instructions) + result = self._run(agent, exception, try_number, max_tries) # The output type validated the answer, so it is one of the configured names. name = picked_key(result.output) - category = self.categories[name] + category = self.categories.get(name) + if category is None: + # The output type validates the answer, so this needs the schema and the table to disagree. + log.error("Classifier answered %r, which is not a configured category", name) + return "model_error" # A bare output type is one field, ``response``; its confidence is what the bar is compared against. model_confidence = ModelConfidence.from_result(result) @@ -433,7 +668,7 @@ def _classify( ) if uncertain is not None: log.info( - "LLM error classification not acted on (%s): %s, using fallback. model=%s probabilities=%s", + "Classifier answer not acted on (%s): %s. model=%s probabilities=%s", uncertain, summary, model_confidence.model or "n/a", @@ -443,10 +678,10 @@ def _classify( if not category.retry: reason = f"{summary} action=fail" - log.info("LLM error classification: %s model=%s", reason, model_confidence.model or "n/a") + log.info("Classifier decision: %s model=%s", reason, model_confidence.model or "n/a") return RetryDecision.fail(reason=reason) 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") + log.info("Classifier decision: %s model=%s", reason, model_confidence.model or "n/a") return RetryDecision.retry(delay=category.delay, reason=reason) diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py index 9b84dea1e3b42..a6e0875183861 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import logging from datetime import timedelta from decimal import Decimal from unittest.mock import MagicMock, patch @@ -378,6 +379,30 @@ def _result(self, make_mock_run_result, output, details): result.response = ModelResponse(parts=[], model_name="jev-1.13.0", provider_details=details) return result + @pytest.mark.parametrize( + "context", + [ + pytest.param({}, id="no-task-instance"), + pytest.param({"task_instance": {"id": "not-a-ti"}}, id="task-instance-is-a-dict"), + pytest.param({"task_instance": None}, id="task-instance-is-none"), + ], + ) + @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + def test_hand_built_context_skips_the_decision_push_with_a_warning( + self, mock_hook_cls, make_mock_run_result, context, caplog + ): + """A dict-shaped or missing task instance (old tests, custom runners) must not fail the run.""" + mock_agent = MagicMock(spec=["run_sync"]) + mock_agent.run_sync.return_value = self._result(make_mock_run_result, Summary(text="t"), None) + mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c", output_type=Summary) + + with caplog.at_level(logging.WARNING): + output = op.execute(context) + + assert isinstance(output, (Summary, dict)) + assert "the decision record was not pushed to XCom" in caplog.text + @pytest.mark.skipif( not AIRFLOW_V_3_1_PLUS, reason="a reviewing decision_policy needs the HITL flow, Airflow >= 3.1" ) 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 1dc91c0e1bac9..012a0dadf2a1e 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 @@ -35,15 +35,18 @@ pytest.importorskip("airflow.sdk.definitions.retry_policy", reason="RetryPolicy requires Airflow 3.3+") from airflow.providers.common.ai.policies.retry import ( + CLASSIFIER_INSTRUCTIONS, DEFAULT_CATEGORIES, DEFAULT_INSTRUCTIONS, + ClassifierRetryPolicy, ErrorCategory, + ErrorClassification, 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.definitions.retry_policy import RetryAction, RetryDecision, RetryRule from airflow.sdk.log import mask_secret HOOK = "airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook" @@ -150,17 +153,18 @@ def test_documented_split(self): "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.""" + def test_category_instructions_do_not_recite_the_taxonomy(self): + """With categories the taxonomy travels in the schema, so the categories prompt must not pin it.""" for name in DEFAULT_CATEGORIES: - assert f"- {name}:" not in DEFAULT_INSTRUCTIONS + assert f"- {name}:" not in CLASSIFIER_INSTRUCTIONS + assert f"- {name}:" 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( + ClassifierRetryPolicy( llm_conn_id="test", categories={"a": ErrorCategory("A", min_confidence=0.9), "b": ErrorCategory("B", retry=False)}, min_confidence=0.6, @@ -168,58 +172,32 @@ def test_construction_makes_no_connection_or_network_call(self, mock_hook_cls): 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] + ClassifierRetryPolicy(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) + ClassifierRetryPolicy(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")}) + ClassifierRetryPolicy( + 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] + ClassifierRetryPolicy(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( + ClassifierRetryPolicy( llm_conn_id="test", categories={ "a": ErrorCategory("A", min_confidence=0.9), @@ -230,7 +208,7 @@ def test_category_bar_needs_a_policy_bar_to_inherit(self): 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) + ClassifierRetryPolicy(llm_conn_id="test", min_confidence=2) @pytest.mark.parametrize( "categories", @@ -247,7 +225,7 @@ def test_policy_can_be_deep_copied(self, categories): 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) + policy = ClassifierRetryPolicy(llm_conn_id="test", categories=categories) assert copy.deepcopy(policy).categories == policy.categories @@ -256,7 +234,7 @@ 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) + policy = ClassifierRetryPolicy(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) @@ -266,11 +244,18 @@ def test_caller_mapping_is_copied(self, mock_hook_cls): @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) + ClassifierRetryPolicy( + llm_conn_id="test", max_exception_length=max_exception_length, categories=DEFAULT_CATEGORIES + ) 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) + ClassifierRetryPolicy( + llm_conn_id="test", + redactor=lambda m: m, + redact_exception=False, + categories=DEFAULT_CATEGORIES, + ) class TestClassification: @@ -293,7 +278,7 @@ 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") + policy = ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES) decision = policy.evaluate(RuntimeError("boom"), try_number=1, max_tries=3) @@ -303,7 +288,9 @@ def test_default_table_covers_every_category( @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) + ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES).evaluate( + RuntimeError("boom"), try_number=1, max_tries=3 + ) assert _options(_output_type(mock_hook_cls)) == [ (name, category.description) for name, category in DEFAULT_CATEGORIES.items() @@ -313,7 +300,7 @@ def test_default_categories_reach_the_model_with_descriptions(self, mock_hook_cl 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( + policy = ClassifierRetryPolicy( llm_conn_id="test", categories={ "warehouse_suspended": ErrorCategory( @@ -336,7 +323,9 @@ def test_custom_categories_build_the_schema_in_the_authors_order(self, mock_hook 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) + ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES).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. @@ -357,7 +346,7 @@ def capture(**kwargs): return agent mock_hook_cls.return_value.create_agent.side_effect = capture - policy = LLMRetryPolicy(llm_conn_id="test") + policy = ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES) decision = policy.evaluate(RuntimeError("429"), try_number=1, max_tries=3) @@ -382,7 +371,7 @@ def capture(**kwargs): return agent mock_hook_cls.return_value.create_agent.side_effect = capture - policy = LLMRetryPolicy( + policy = ClassifierRetryPolicy( llm_conn_id="test", categories={ name: ErrorCategory("Odd name.", delay=timedelta(seconds=5)), @@ -413,7 +402,7 @@ def capture(**kwargs): 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") + policy = ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES) decision = policy.evaluate(RuntimeError("boom"), try_number=1, max_tries=3) @@ -422,7 +411,7 @@ def test_reason_is_the_generated_summary(self, mock_hook_cls, category, confiden @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( + policy = ClassifierRetryPolicy( llm_conn_id="test", categories={"flaky": ErrorCategory("F"), "broken": ErrorCategory("B", retry=False)}, ) @@ -437,7 +426,7 @@ def test_none_delay_leaves_the_task_backoff_in_charge(self, mock_hook_cls): 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( + policy = ClassifierRetryPolicy( llm_conn_id="test", categories={ "flaky": ErrorCategory("F", delay=timedelta(0)), @@ -455,7 +444,7 @@ 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") + policy = ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES) first = policy.evaluate(ConnectionError("reset"), try_number=1, max_tries=3) second = policy.evaluate(PermissionError("expired"), try_number=2, max_tries=3) @@ -470,7 +459,7 @@ def test_each_failure_is_classified_afresh(self, mock_hook_cls): @patch(HOOK, autospec=True) def test_custom_instructions_forwarded_to_agent(self, mock_hook_cls): _install(mock_hook_cls, _agent("auth")) - policy = LLMRetryPolicy( + policy = ClassifierRetryPolicy( llm_conn_id="test", instructions="Snowflake errors only.", categories={"auth": ErrorCategory("A", retry=False), "other": ErrorCategory("O")}, @@ -486,7 +475,9 @@ def test_custom_instructions_forwarded_to_agent(self, mock_hook_cls): @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 = ClassifierRetryPolicy( + llm_conn_id="jev_default", model_id="typesafe:jev-1.13.0", categories=DEFAULT_CATEGORIES + ) policy.evaluate(RuntimeError("x"), try_number=1, max_tries=3) @@ -495,7 +486,7 @@ def test_model_id_and_connection_reach_the_hook(self, mock_hook_cls): @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 = ClassifierRetryPolicy(llm_conn_id="test", timeout=7.5, categories=DEFAULT_CATEGORIES) policy.evaluate(RuntimeError("x"), try_number=1, max_tries=3) @@ -519,7 +510,7 @@ class TestConfidenceGate: pytest.param( 0.59, RetryAction.FAIL, - "LLM classification not applied (below_threshold); rule", + "classifier answer not applied (below_threshold); rule", id="one-below", ), ], @@ -527,7 +518,7 @@ class TestConfidenceGate: @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) + policy = ClassifierRetryPolicy(llm_conn_id="test", min_confidence=0.6, fallback_rules=self.RULES) decision = policy.evaluate(RuntimeError("reset"), try_number=1, max_tries=3) @@ -538,15 +529,13 @@ def test_bar_boundary(self, mock_hook_cls, confidence, expected_action, expected 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) + policy = ClassifierRetryPolicy(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" - ) + assert decision.reason == "classifier answer 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 ( @@ -559,7 +548,7 @@ def test_under_the_bar_with_no_rules_keeps_the_task_behaviour(self, mock_hook_cl 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( + policy = ClassifierRetryPolicy( llm_conn_id="test", min_confidence=0.6, fallback_rules=[ @@ -575,12 +564,12 @@ def test_under_the_bar_takes_fallback_rules_before_default(self, mock_hook_cls): ) assert ( policy.evaluate(PermissionError("x"), try_number=1, max_tries=3).reason - == "LLM classification not applied (below_threshold); rule" + == "classifier answer 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" + unmatched.reason == "classifier answer not applied (below_threshold); task retry settings apply" ) @pytest.mark.parametrize( @@ -598,7 +587,7 @@ 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( + policy = ClassifierRetryPolicy( llm_conn_id="test", min_confidence=0.6, categories={ @@ -614,17 +603,19 @@ def test_category_bar_overrides_the_policy_bar( 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) + policy = ClassifierRetryPolicy(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" + assert decision.reason == "classifier answer 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) + policy = ClassifierRetryPolicy( + llm_conn_id="test", fallback_rules=self.RULES, categories=DEFAULT_CATEGORIES + ) decision = policy.evaluate(RuntimeError("reset"), try_number=1, max_tries=3) @@ -634,14 +625,14 @@ def test_missing_confidence_without_a_bar_acts_as_before(self, mock_hook_cls): @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) + policy = ClassifierRetryPolicy(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") + policy = ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES) decision = policy.evaluate(RuntimeError("x"), try_number=1, max_tries=3) @@ -653,7 +644,7 @@ 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 = ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES) policy.evaluate(ValueError("bad column type"), try_number=2, max_tries=5) @@ -668,7 +659,7 @@ def test_prompt_redacts_known_secrets(self, mock_hook_cls): secret_value = "super-secret-conn-password" mask_secret(secret_value) agent = _install(mock_hook_cls, _agent("auth")) - policy = LLMRetryPolicy(llm_conn_id="test") + policy = ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES) policy.evaluate( ConnectionError(f"could not authenticate with password {secret_value}"), try_number=1, max_tries=3 @@ -686,7 +677,9 @@ def test_prompt_keeps_raw_message_when_redaction_disabled(self, mock_hook_cls): secret_value = "super-secret-conn-password" mask_secret(secret_value) agent = _install(mock_hook_cls, _agent("auth")) - policy = LLMRetryPolicy(llm_conn_id="test", redact_exception=False) + policy = ClassifierRetryPolicy( + llm_conn_id="test", redact_exception=False, categories=DEFAULT_CATEGORIES + ) policy.evaluate(ConnectionError(f"password {secret_value}"), try_number=1, max_tries=3) @@ -699,7 +692,7 @@ def test_explicit_redactor_none_still_applies_default_masking(self, mock_hook_cl secret_value = "super-secret-conn-password" mask_secret(secret_value) agent = _install(mock_hook_cls, _agent("auth")) - policy = LLMRetryPolicy(llm_conn_id="test", redactor=None) + policy = ClassifierRetryPolicy(llm_conn_id="test", redactor=None, categories=DEFAULT_CATEGORIES) policy.evaluate(ConnectionError(f"password {secret_value}"), try_number=1, max_tries=3) @@ -712,8 +705,10 @@ def test_custom_redactor_replaces_masker_instead_of_stacking(self, mock_hook_cls secret_value = "super-secret-conn-password" mask_secret(secret_value) agent = _install(mock_hook_cls, _agent("auth")) - policy = LLMRetryPolicy( - llm_conn_id="test", redactor=lambda m: m.replace("user@example.com", "") + policy = ClassifierRetryPolicy( + llm_conn_id="test", + redactor=lambda m: m.replace("user@example.com", ""), + categories=DEFAULT_CATEGORIES, ) policy.evaluate(ConnectionError(f"user@example.com {secret_value}"), try_number=1, max_tries=3) @@ -734,7 +729,9 @@ def test_message_truncated_only_when_over_max_exception_length( self, mock_hook_cls, length, expected_tail ): agent = _install(mock_hook_cls, _agent("data")) - policy = LLMRetryPolicy(llm_conn_id="test", max_exception_length=10) + policy = ClassifierRetryPolicy( + llm_conn_id="test", max_exception_length=10, categories=DEFAULT_CATEGORIES + ) policy.evaluate(ValueError("x" * length), try_number=1, max_tries=3) @@ -748,7 +745,9 @@ def test_truncation_happens_after_redaction(self, mock_hook_cls): secret_value = "super-secret-conn-password" mask_secret(secret_value) agent = _install(mock_hook_cls, _agent("auth")) - policy = LLMRetryPolicy(llm_conn_id="test", max_exception_length=12) + policy = ClassifierRetryPolicy( + llm_conn_id="test", max_exception_length=12, categories=DEFAULT_CATEGORIES + ) policy.evaluate(ConnectionError(f"pw {secret_value} tail"), try_number=1, max_tries=3) @@ -759,7 +758,7 @@ 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( + policy = ClassifierRetryPolicy( llm_conn_id="nonexistent", fallback_rules=[ RetryRule( @@ -767,6 +766,7 @@ def test_falls_back_to_rules_when_connection_missing(self): ), RetryRule(exception=PermissionError, action=RetryAction.FAIL, reason="auth fallback"), ], + categories=DEFAULT_CATEGORIES, ) retry = policy.evaluate(ConnectionError("refused"), try_number=1, max_tries=3) @@ -776,14 +776,15 @@ def test_falls_back_to_rules_when_connection_missing(self): assert fail.action == RetryAction.FAIL def test_falls_back_to_default_when_no_rules(self): - policy = LLMRetryPolicy(llm_conn_id="nonexistent") + policy = ClassifierRetryPolicy(llm_conn_id="nonexistent", categories=DEFAULT_CATEGORIES) assert policy.evaluate(ValueError("bad"), try_number=1, max_tries=3).action == RetryAction.DEFAULT def test_fallback_rules_no_match_returns_default(self): - policy = LLMRetryPolicy( + policy = ClassifierRetryPolicy( llm_conn_id="nonexistent", fallback_rules=[RetryRule(exception=PermissionError, action=RetryAction.FAIL)], + categories=DEFAULT_CATEGORIES, ) assert policy.evaluate(ValueError("bad"), try_number=1, max_tries=3).action == RetryAction.DEFAULT @@ -793,27 +794,29 @@ 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( + policy = ClassifierRetryPolicy( llm_conn_id="test", fallback_rules=[RetryRule(exception=ValueError, action=RetryAction.FAIL, reason="fallback")], + categories=DEFAULT_CATEGORIES, ) 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" + assert decision.reason == "classifier answer not applied (model_error); fallback" @patch(HOOK, autospec=True) def test_hook_creation_failure_triggers_fallback(self, mock_hook_cls): mock_hook_cls.return_value.create_agent.side_effect = RuntimeError("unexpected") - policy = LLMRetryPolicy( + policy = ClassifierRetryPolicy( llm_conn_id="test", fallback_rules=[RetryRule(exception=ValueError, action=RetryAction.FAIL, reason="caught")], + categories=DEFAULT_CATEGORIES, ) assert ( policy.evaluate(ValueError("x"), try_number=1, max_tries=3).reason - == "LLM classification not applied (model_error); caught" + == "classifier answer not applied (model_error); caught" ) @patch(HOOK, autospec=True) @@ -823,6 +826,317 @@ def test_answer_the_type_rejected_triggers_fallback(self, mock_hook_cls): agent.run_sync.side_effect = UnexpectedModelBehavior( "Exceeded maximum retries (1) for output validation" ) - policy = LLMRetryPolicy(llm_conn_id="test") + policy = ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES) assert policy.evaluate(ValueError("x"), try_number=1, max_tries=3).action == RetryAction.DEFAULT + + def test_matched_default_rule_keeps_its_delay_and_reason(self): + """A rule the author wrote with action=DEFAULT is a match like any other; 0.9.0 returned it verbatim.""" + policy = ClassifierRetryPolicy( + llm_conn_id="nonexistent", + fallback_rules=[ + RetryRule( + exception=ValueError, + action=RetryAction.DEFAULT, + retry_delay=timedelta(seconds=45), + reason="my rule reason", + ) + ], + categories=DEFAULT_CATEGORIES, + ) + + decision = policy.evaluate(ValueError("x"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.DEFAULT + assert decision.retry_delay == timedelta(seconds=45) + assert decision.reason == "classifier answer not applied (model_error); my rule reason" + + @patch(HOOK, autospec=True) + def test_answer_missing_from_the_table_falls_back_with_an_error_log(self, mock_hook_cls, caplog): + """Reachable only if the schema and the table disagree; it must not surface as a KeyError traceback.""" + _install(mock_hook_cls, _agent("not_a_category")) + policy = ClassifierRetryPolicy(llm_conn_id="test", categories=DEFAULT_CATEGORIES) + + with caplog.at_level(logging.ERROR, logger="airflow.providers.common.ai.policies.retry"): + decision = policy.evaluate(ValueError("x"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.DEFAULT + assert decision.reason == "classifier answer not applied (model_error); task retry settings apply" + assert "answered 'not_a_category', which is not a configured category" in caplog.text + assert "KeyError" not in caplog.text + + def test_subclass_classify_returning_the_wrong_type_is_reported_not_acted_on(self, caplog): + class Odd(ClassifierRetryPolicy): + def _classify(self, exception, try_number, max_tries): + return {"should_retry": True} # the 0.9.0 shape, not a RetryDecision + + policy = Odd(llm_conn_id="test") + + with caplog.at_level(logging.ERROR, logger="airflow.providers.common.ai.policies.retry"): + decision = policy.evaluate(ValueError("x"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.DEFAULT + assert decision.reason == "classifier answer not applied (model_error); task retry settings apply" + assert "returned dict instead of a RetryDecision" in caplog.text + + +class TestOnUncertain: + """Classifier first; an LLM policy when it is unsure or unreachable; the rules when neither decides.""" + + RULES = [ + RetryRule( + exception=RuntimeError, action=RetryAction.RETRY, retry_delay=timedelta(seconds=5), reason="rule" + ) + ] + + def _chain(self, mock_hook_cls, classifier_answer, confidence, llm_answer): + """Wire two agents behind one patched hook: the classifier by output type, the LLM by ErrorClassification.""" + classifier = ( + _agent(classifier_answer, confidence=confidence) if classifier_answer else MagicMock(spec=Agent) + ) + if classifier_answer is None: + classifier.run_sync.side_effect = TimeoutError("classifier unreachable") + llm = _agent(llm_answer) if isinstance(llm_answer, ErrorClassification) else MagicMock(spec=Agent) + if llm_answer is None: + llm.run_sync.side_effect = TimeoutError("llm unreachable") + + def create_agent(**kwargs): + return llm if kwargs["output_type"] is ErrorClassification else classifier + + mock_hook_cls.return_value.create_agent.side_effect = create_agent + return ClassifierRetryPolicy( + llm_conn_id="jev", + min_confidence=0.8, + on_uncertain=LLMRetryPolicy(llm_conn_id="text"), + fallback_rules=self.RULES, + ) + + def test_on_uncertain_needs_a_bar(self): + with pytest.raises(ValueError, match="on_uncertain needs min_confidence"): + ClassifierRetryPolicy( + llm_conn_id="jev", categories=DEFAULT_CATEGORIES, on_uncertain=LLMRetryPolicy(llm_conn_id="t") + ) + + def test_on_uncertain_must_be_a_retry_policy(self): + with pytest.raises(TypeError, match="on_uncertain must be a RetryPolicy"): + ClassifierRetryPolicy(llm_conn_id="jev", min_confidence=0.5, on_uncertain="llm") # type: ignore[arg-type] + + @patch(HOOK, autospec=True) + def test_confident_classifier_answer_never_consults_the_llm(self, mock_hook_cls): + policy = self._chain( + mock_hook_cls, + "network", + 0.95, + ErrorClassification(category="x", should_retry=False, reasoning="r"), + ) + + decision = policy.evaluate(RuntimeError("reset"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.reason == "category=network confidence=0.95 threshold=0.80 action=retry delay=10s" + assert mock_hook_cls.return_value.create_agent.call_count == 1 + + @patch(HOOK, autospec=True) + def test_unsure_classifier_escalates_to_the_llm(self, mock_hook_cls): + llm_answer = ErrorClassification( + category="rate_limit", + should_retry=True, + suggested_delay_seconds=120, + reasoning="429 with Retry-After", + ) + policy = self._chain(mock_hook_cls, "auth", 0.45, llm_answer) + + decision = policy.evaluate(RuntimeError("?"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.retry_delay == timedelta(seconds=120) + assert decision.reason == "escalated (below_threshold); rate_limit: 429 with Retry-After" + + @patch(HOOK, autospec=True) + def test_unreachable_classifier_escalates_to_the_llm(self, mock_hook_cls): + llm_answer = ErrorClassification(category="auth", should_retry=False, reasoning="expired key") + policy = self._chain(mock_hook_cls, None, None, llm_answer) + + decision = policy.evaluate(RuntimeError("?"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.FAIL + assert decision.reason == "escalated (model_error); auth: expired key" + + @patch(HOOK, autospec=True) + def test_unreachable_llm_falls_to_the_rules(self, mock_hook_cls): + policy = self._chain(mock_hook_cls, "auth", 0.45, None) + + decision = policy.evaluate(RuntimeError("?"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.retry_delay == timedelta(seconds=5) + assert decision.reason == "classifier answer not applied (below_threshold); rule" + + @patch(HOOK, autospec=True) + def test_nothing_decides_means_the_task_default(self, mock_hook_cls): + policy = self._chain(mock_hook_cls, "auth", 0.45, None) + policy.fallback_rules = None + + decision = policy.evaluate(ValueError("?"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.DEFAULT + assert decision.reason == "classifier answer not applied (below_threshold); task retry settings apply" + + def test_inner_policy_with_its_own_rules_is_used_as_is(self): + """The LLM policy's own fallback_rules count as its decision, so the outer rules are not reached.""" + inner = MagicMock(spec=LLMRetryPolicy) + inner.evaluate.return_value = RetryDecision.fail(reason="inner rule") + policy = ClassifierRetryPolicy( + llm_conn_id="nonexistent", min_confidence=0.8, on_uncertain=inner, fallback_rules=self.RULES + ) + + decision = policy.evaluate(RuntimeError("?"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.FAIL + assert decision.reason == "escalated (model_error); inner rule" + inner.evaluate.assert_called_once() + + +def _open_agent(category, should_retry, delay=0, reasoning="test"): + """A mock agent answering the 0.9.0 ``ErrorClassification`` shape, for a policy without categories.""" + return _agent( + ErrorClassification( + category=category, should_retry=should_retry, suggested_delay_seconds=delay, reasoning=reasoning + ) + ) + + +class TestLLMRetryPolicy: + """The LLM layer is the 0.9.0 policy: the text model decides retry and delay, and nothing classifier-shaped is on it.""" + + def test_defaults(self): + policy = LLMRetryPolicy(llm_conn_id="test") + + assert policy.instructions == DEFAULT_INSTRUCTIONS + assert not hasattr(policy, "categories") + assert not hasattr(policy, "min_confidence") + + @pytest.mark.parametrize("kwarg", ["categories", "min_confidence", "on_uncertain"]) + def test_has_no_classifier_arguments(self, kwarg): + """Those belong to ClassifierRetryPolicy; a text-model user never sees them.""" + with pytest.raises(TypeError, match="unexpected keyword argument"): + LLMRetryPolicy(llm_conn_id="test", **{kwarg: None}) + + def test_classifier_defaults(self): + policy = ClassifierRetryPolicy(llm_conn_id="test", min_confidence=0.6) + + assert policy.categories == dict(DEFAULT_CATEGORIES) + assert policy.instructions == CLASSIFIER_INSTRUCTIONS + + @patch(HOOK, autospec=True) + def test_model_output_type_is_error_classification(self, mock_hook_cls): + _install(mock_hook_cls, _open_agent("auth", should_retry=False)) + LLMRetryPolicy(llm_conn_id="test").evaluate(PermissionError("403"), try_number=1, max_tries=3) + + kwargs = mock_hook_cls.return_value.create_agent.call_args.kwargs + assert kwargs["output_type"] is ErrorClassification + assert kwargs["instructions"] == DEFAULT_INSTRUCTIONS + + @patch(HOOK, autospec=True) + def test_auth_error_returns_fail_with_the_models_reasoning(self, mock_hook_cls): + _install(mock_hook_cls, _open_agent("auth", should_retry=False, reasoning="API key expired")) + policy = LLMRetryPolicy(llm_conn_id="test") + + decision = policy.evaluate(PermissionError("403"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.FAIL + assert decision.reason == "auth: API key expired" + + @patch(HOOK, autospec=True) + def test_model_chosen_delay_is_used(self, mock_hook_cls): + _install(mock_hook_cls, _open_agent("rate_limit", should_retry=True, delay=120, reasoning="429")) + policy = LLMRetryPolicy(llm_conn_id="test") + + decision = policy.evaluate(RuntimeError("429"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.retry_delay == timedelta(seconds=120) + assert decision.reason == "rate_limit: 429" + + @pytest.mark.parametrize("delay", [0, -5]) + @patch(HOOK, autospec=True) + def test_zero_or_negative_delay_leaves_the_task_backoff_in_charge(self, mock_hook_cls, delay): + _install(mock_hook_cls, _open_agent("transient", should_retry=True, delay=delay)) + policy = LLMRetryPolicy(llm_conn_id="test") + + decision = policy.evaluate(RuntimeError("glitch"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.retry_delay is None + + @patch(HOOK, autospec=True) + def test_custom_taxonomy_in_instructions_keeps_working(self, mock_hook_cls): + """The 0.9.0 guide taught this: a category the prompt invents, with the model choosing the action.""" + _install( + mock_hook_cls, + _open_agent("warehouse_suspended", should_retry=True, delay=30, reasoning="auto-resume"), + ) + policy = LLMRetryPolicy( + llm_conn_id="test", instructions="'Warehouse suspended' -> warehouse_suspended, retry after 30s" + ) + + decision = policy.evaluate(RuntimeError("Warehouse X is suspended"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.retry_delay == timedelta(seconds=30) + assert decision.reason == "warehouse_suspended: auto-resume" + assert mock_hook_cls.return_value.create_agent.call_args.kwargs["instructions"].startswith( + "'Warehouse suspended'" + ) + + def test_fallback_decisions_are_the_rules_verbatim(self): + """No prefix on this path: the decision is what ExceptionRetryPolicy returns, as in 0.9.0.""" + policy = LLMRetryPolicy( + llm_conn_id="nonexistent", + fallback_rules=[ + RetryRule( + exception=ConnectionError, action=RetryAction.RETRY, retry_delay=timedelta(seconds=10) + ), + RetryRule(exception=PermissionError, action=RetryAction.FAIL, reason="auth fallback"), + ], + ) + + retry = policy.evaluate(ConnectionError("refused"), try_number=1, max_tries=3) + fail = policy.evaluate(PermissionError("denied"), try_number=1, max_tries=3) + unmatched = policy.evaluate(ValueError("x"), try_number=1, max_tries=3) + + assert (retry.action, retry.retry_delay, retry.reason) == ( + RetryAction.RETRY, + timedelta(seconds=10), + "Matched rule for ConnectionError", + ) + assert (fail.action, fail.reason) == (RetryAction.FAIL, "auth fallback") + assert (unmatched.action, unmatched.reason) == (RetryAction.DEFAULT, None) + + def test_no_rules_falls_back_to_default_with_no_reason(self): + decision = LLMRetryPolicy(llm_conn_id="nonexistent").evaluate( + ValueError("x"), try_number=1, max_tries=3 + ) + + assert decision == RetryDecision.default() + + @patch(HOOK, autospec=True) + def test_classifier_refusal_logs_the_categories_hint(self, mock_hook_cls, caplog): + """A classifier model refuses ErrorClassification's text fields; the log says what to do.""" + agent = _install(mock_hook_cls, MagicMock(spec=Agent)) + agent.run_sync.side_effect = RuntimeError("Output field 'reasoning' is not supported by this model") + policy = LLMRetryPolicy(llm_conn_id="test", model_id="typesafe:jev-1.13.0") + + with caplog.at_level(logging.ERROR, logger="airflow.providers.common.ai.policies.retry"): + decision = policy.evaluate(ValueError("x"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.DEFAULT + assert "Use ClassifierRetryPolicy" in caplog.text + + def test_no_warning_for_any_instructions(self): + with warnings.catch_warnings(): + warnings.simplefilter("error") + LLMRetryPolicy( + llm_conn_id="test", instructions="'Statement queued' -> rate_limit, retry after 120s" + ) + LLMRetryPolicy(llm_conn_id="test", instructions=DEFAULT_INSTRUCTIONS + " hints") From 70bd5528abbd410aa24a7dfdb8d144a6c64f18a5 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Mon, 21 Sep 2026 22:11:11 +0100 Subject: [PATCH 2/3] Treat any DEFAULT from on_uncertain as no decision so the outer fallback rules always run Only a RETRY or FAIL from the on_uncertain policy ends the chain. A DEFAULT, whatever reason it carries, hands back to the ClassifierRetryPolicy's own fallback_rules, so an outer PermissionError -> FAIL rule still holds when a nested classifier or an LLM policy has nothing to add. Previously a DEFAULT with a reason counted as a decision, which a ClassifierRetryPolicy used as on_uncertain always produced, so the outer rules never ran. A decision without a reason no longer renders as "escalated (...); None". on_uncertain no longer requires min_confidence: a classifier outage escalates without a bar. The "not supported by this model" hint is worded as a hint, since text models whose profile lacks structured output raise the same words. ErrorClassification keeps its 0.9.0 docstring, which pydantic sends to the model as the schema description, and the class docstring concatenation survives python -OO. The changelog note on execute_complete names the two operators that pass decision on resume; the additive ClassifierRetryPolicy note is dropped. --- providers/common/ai/docs/changelog.rst | 20 +- providers/common/ai/docs/retry_policies.rst | 39 ++-- .../providers/common/ai/policies/retry.py | 102 +++++----- .../unit/common/ai/operators/test_llm.py | 2 +- .../unit/common/ai/policies/test_retry.py | 186 +++++++++++++----- 5 files changed, 226 insertions(+), 123 deletions(-) diff --git a/providers/common/ai/docs/changelog.rst b/providers/common/ai/docs/changelog.rst index fd36a8e2394ac..4fb07be4a6048 100644 --- a/providers/common/ai/docs/changelog.rst +++ b/providers/common/ai/docs/changelog.rst @@ -37,23 +37,13 @@ Changelog .. note:: ``execute_complete`` on ``LLMOperator``, ``LLMBranchOperator``, ``LLMSQLQueryOperator`` and - ``LLMSchemaCompareOperator`` gained a keyword argument, ``decision``, and every review pause - now passes it on resume. A subclass that overrides ``execute_complete`` with the old - three-argument signature raises ``TypeError`` when the reviewed task resumes; add - ``decision=None`` to the override. A review that was already pending when you upgraded + ``LLMSchemaCompareOperator`` gained a keyword argument, ``decision``. ``LLMOperator`` and + ``LLMBranchOperator`` pass it on resume, so a subclass of either that overrides + ``execute_complete`` with the old three-argument signature raises ``TypeError`` when the + reviewed task resumes; add ``decision=None`` to the override. The other two accept the + keyword but do not pass it yet. A review that was already pending when you upgraded resumes without it and is unaffected. -.. note:: - New ``ClassifierRetryPolicy``: the model names one of the author's ``categories`` (a - table of ``ErrorCategory(description, retry, delay, min_confidence)``) and the table - decides whether to retry, after how long, and how sure the model has to be. It is the - policy for a classifier model such as TypeSafe's Jev, and ``on_uncertain`` lets it hand - an unsure or failed classification to another policy, typically an ``LLMRetryPolicy`` on - a text model, before ``fallback_rules``. ``LLMRetryPolicy`` itself is unchanged: the model - returns ``ErrorClassification`` and chooses the retry and the delay from your - ``instructions``, and it grows no new arguments. Other new public names: - ``ErrorCategory``, ``DEFAULT_CATEGORIES``, ``CLASSIFIER_INSTRUCTIONS``. - .. 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/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index 1d44d16a12a04..e61d287453405 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -43,7 +43,8 @@ fully model-driven, with the SDK's ``ExceptionRetryPolicy`` as the bottom rung: - The model names one of your ``categories``; the table says whether that category is retried, after how long, and how sure the model has to be. A classifier model such as TypeSafe's Jev answers in a few hundred - milliseconds and reports its confidence; a text model can sit here too. + milliseconds and reports its confidence; a text model can sit here too, + without a bar. - Category descriptions and the confidence bar. No reasoning, no prose. * - **LLM** (``LLMRetryPolicy``) @@ -93,7 +94,7 @@ Usage llm_policy = LLMRetryPolicy( llm_conn_id="pydanticai_default", timeout=30.0, # max seconds to wait for LLM response - fallback_rules=[ # used when the LLM call fails or the answer is under its bar + fallback_rules=[ # used when the LLM call fails RetryRule(exception=ConnectionError, action=RetryAction.RETRY, retry_delay=timedelta(seconds=10)), RetryRule(exception=PermissionError, action=RetryAction.FAIL), ], @@ -319,7 +320,9 @@ Escalating to an LLM -------------------- A classifier is cheap and fast, and a text model can reason about a failure it -has never seen a category for. ``on_uncertain`` puts one behind the other: +has never seen a category for. ``on_uncertain`` puts one behind the other. +``snowflake_policy`` is the classifier policy from the previous section; the +chain reuses its category table: .. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py :language: python @@ -336,20 +339,25 @@ The order of events on a failure: itself. Its decision is used, with the reason prefixed by why the classifier's answer was not: ``escalated (below_threshold); rate_limit: 429 with a Retry-After header``. -3. If that policy decides nothing either (its model was unreachable and none of - its own ``fallback_rules`` matched), the outer ``fallback_rules`` apply, then - the task's own retry behaviour. - -``on_uncertain`` accepts any ``RetryPolicy``, so an ``ExceptionRetryPolicy`` -works there too, and needs ``min_confidence``: without a bar the classifier is -never unsure. Both model calls run on the worker at failure time, so a task -that escalates pays for two before its retry is scheduled; ``timeout`` on each -policy bounds that. +3. If that policy returns DEFAULT, whatever reason it attached, it decided + nothing: the outer ``fallback_rules`` apply, then the task's own retry + behaviour. Only a RETRY or FAIL from ``on_uncertain`` ends the chain, so an + outer rule such as ``PermissionError -> FAIL`` still holds when both models + are unreachable. + +``on_uncertain`` accepts any ``RetryPolicy``. An ``ExceptionRetryPolicy`` works +there too; its ``default`` is what it returns when none of its rules match, so +``default=RetryAction.FAIL`` fails every unsure classification and the outer +rules never run. Without ``min_confidence`` the classifier's answer is always +acted on, and ``on_uncertain`` is consulted only when the classifier call itself +fails. Both model calls run on the worker at failure time, so a task that +escalates pays for two before its retry is scheduled; ``timeout`` on each policy +bounds that. When the connection also carries a fallback chain -------------------------------------------------- -``LLMRetryPolicy`` builds its classifier hook from ``llm_conn_id`` without passing +Either policy builds its hook from ``llm_conn_id`` without passing ``fallback_conn_ids``, so if that connection's extra configures a chain (see :doc:`provider_fallback`), the policy inherits it silently -- editing the connection changes retry behaviour with no change to the Dag. Two things follow: @@ -544,8 +552,9 @@ Both policies share every parameter below except ``categories``, - None - ``ClassifierRetryPolicy`` only. A ``RetryPolicy`` to consult when the classifier is under its bar, reports no confidence, or cannot be - reached; typically an ``LLMRetryPolicy`` on a text model. Its RETRY or FAIL is used; if it decides nothing, the outer - ``fallback_rules`` apply. Needs ``min_confidence``. + reached; typically an ``LLMRetryPolicy`` on a text model. Its RETRY or FAIL + is used; a DEFAULT, whatever its reason, means the outer ``fallback_rules`` + apply. * - ``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/policies/retry.py b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py index 71f4677bb9dc1..2f74b518e0f5b 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 @@ -74,6 +74,10 @@ if TYPE_CHECKING: from collections.abc import Callable + from pydantic_ai import Agent + from pydantic_ai.agent import AgentRunResult + + from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook from airflow.sdk.definitions.context import Context from airflow.sdk.definitions.retry_policy import RetryRule @@ -108,7 +112,7 @@ class ErrorClassification(BaseModel): - """Structured output of :class:`LLMRetryPolicy`: the model's category, decision, delay and reasoning.""" + """Structured LLM output for error classification.""" category: str """One of the categories the instructions describe, by default: rate_limit, auth, network, data, resource, transient, permanent.""" @@ -293,7 +297,7 @@ def __init__( self.redact_exception = redact_exception self.max_exception_length = max_exception_length - def _hook(self) -> Any: + def _hook(self) -> PydanticAIHook: from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook return PydanticAIHook(llm_conn_id=self.llm_conn_id, model_id=self.model_id) @@ -309,7 +313,9 @@ def _prompt(self, exception: BaseException, try_number: int, max_tries: int) -> f"{type(exception).__name__}: {message}" ) - def _run(self, agent: Any, exception: BaseException, try_number: int, max_tries: int) -> Any: + def _run( + self, agent: Agent[Any, Any], exception: BaseException, try_number: int, max_tries: int + ) -> AgentRunResult[Any]: from pydantic_ai.settings import ModelSettings return agent.run_sync( @@ -360,7 +366,7 @@ class LLMRetryPolicy(_ModelRetryPolicy): decision path fast even when the provider is degraded. """ - __doc__ = __doc__ + _REDACTION_PARAMS_DOC + __doc__ = (__doc__ or "") + _REDACTION_PARAMS_DOC # ``or ""`` keeps the import working under python -OO _default_instructions = DEFAULT_INSTRUCTIONS @@ -388,10 +394,11 @@ def _classify( result = self._run(agent, exception, try_number, max_tries) except Exception as exc: if "not supported by this model" in str(exc): - # A classifier model refuses ErrorClassification's free-text fields client-side. + # A classifier model refuses ErrorClassification's free-text fields client-side. A text + # model whose profile lacks structured output raises the same words, so this is a hint. log.error( - "This model cannot answer ErrorClassification; it needs a typed question. " - "Use ClassifierRetryPolicy for a classifier model such as TypeSafe's Jev." + "This model cannot answer ErrorClassification. If it is a classifier model such as " + "TypeSafe's Jev, use ClassifierRetryPolicy, which asks it a typed question." ) raise classification = result.output @@ -461,12 +468,13 @@ class ClassifierRetryPolicy(_ModelRetryPolicy): confidence, or the model call fails; typically an :class:`LLMRetryPolicy` on a text model, so the classifier handles the clear cases and a reasoning model the rest. Its RETRY or FAIL is used, with the reason prefixed by why the classifier's answer - was not. If it decides nothing either (its model was unreachable and none of its own - ``fallback_rules`` matched), this policy's ``fallback_rules`` and then the task's own - retry behaviour apply. Needs ``min_confidence``. + was not. A DEFAULT from it counts as no decision, whatever reason it carries, and this + policy's ``fallback_rules`` and then the task's own retry behaviour apply. Without + ``min_confidence`` the classifier's answer is always acted on, so this policy is + consulted only when the classifier call itself fails. """ - __doc__ = __doc__ + _REDACTION_PARAMS_DOC + __doc__ = (__doc__ or "") + _REDACTION_PARAMS_DOC # ``or ""`` keeps the import working under python -OO _default_instructions = CLASSIFIER_INSTRUCTIONS @@ -499,14 +507,8 @@ def __init__( self.categories: dict[str, ErrorCategory] = self._validate_categories( DEFAULT_CATEGORIES if categories is None else categories ) - if on_uncertain is not None: - if not isinstance(on_uncertain, RetryPolicy): - raise TypeError(f"on_uncertain must be a RetryPolicy, got {type(on_uncertain).__name__}.") - if self.min_confidence is None: - raise ValueError( - "on_uncertain needs min_confidence: it is consulted when the classifier's answer is " - "under the bar or the classifier could not answer, and without a bar there is no such case." - ) + if on_uncertain is not None and not isinstance(on_uncertain, RetryPolicy): + raise TypeError(f"on_uncertain must be a RetryPolicy, got {type(on_uncertain).__name__}.") self.on_uncertain = on_uncertain def _validate_categories(self, categories: Mapping[str, ErrorCategory]) -> dict[str, ErrorCategory]: @@ -560,21 +562,17 @@ def evaluate( outcome = "model_error" if isinstance(outcome, RetryDecision): return outcome - if not isinstance(outcome, str): - # A subclass's _classify returned something else; say so rather than acting on its repr. - log.error( - "Classifier retry classification returned %s instead of a RetryDecision, using fallback", - type(outcome).__name__, - ) - outcome = "model_error" if self.on_uncertain is not None: - escalated = self._escalate(exception, try_number, max_tries, context, why=outcome) + escalated = self._escalate( + self.on_uncertain, exception, try_number, max_tries, context, why=outcome + ) if escalated is not None: return escalated return self._fall_back(exception, try_number, max_tries, context, why=outcome) def _escalate( self, + policy: RetryPolicy, exception: BaseException, try_number: int, max_tries: int, @@ -583,26 +581,32 @@ def _escalate( why: str, ) -> RetryDecision | None: """ - Consult ``on_uncertain`` and return its decision, or None when it decided nothing. + Consult ``on_uncertain`` and return its RETRY or FAIL, or None when it decided nothing. - A DEFAULT decision with no reason is what a policy returns when its own model call failed - and no rule of its own matched, so that case falls through to this policy's - ``fallback_rules``. Any other decision is returned with the reason prefixed by why the - classifier's answer was not used, so a ``retry_reason`` shows the whole chain. + Only RETRY and FAIL are decisions. DEFAULT means the policy had nothing to add to the + task's own settings, whatever reason it attached (its own fallback message, or a matched + rule with ``action=DEFAULT``), so this policy's ``fallback_rules`` still get their say. + A decision comes back with the reason prefixed by why the classifier's answer was not + used, so a ``retry_reason`` shows the whole chain. """ - policy = cast("RetryPolicy", self.on_uncertain) log.info("Classifier answer not applied (%s), consulting %s", why, type(policy).__name__) try: decision = policy.evaluate(exception, try_number, max_tries, context) except Exception: log.exception("on_uncertain policy failed, using fallback rules") return None - if decision.action is RetryAction.DEFAULT and decision.reason is None: + if decision.action is RetryAction.DEFAULT: + log.info( + "%s decided nothing (%s), using fallback rules", + type(policy).__name__, + decision.reason or "no reason given", + ) return None + prefix = f"escalated ({why})" return RetryDecision( action=decision.action, retry_delay=decision.retry_delay, - reason=f"escalated ({why}); {decision.reason}", + reason=prefix if decision.reason is None else f"{prefix}; {decision.reason}", ) def _fall_back( @@ -619,7 +623,8 @@ def _fall_back( A ``retry_reason`` read later is then not mistaken for a plain rule match or a classifier decision. A matched rule always carries a reason; an unmatched evaluation is DEFAULT with - none. A matched rule keeps its action, delay and reason whatever the action. + none. A matched rule keeps its action and reason whatever the action; on DEFAULT the worker + applies the task's own settings and the reason reaches only the log. """ prefix = f"classifier answer not applied ({why})" ruled = self._rules_decision(exception, try_number, max_tries, context) @@ -647,16 +652,17 @@ def _classify( ) agent = self._hook().create_agent(output_type=output_type, instructions=self.instructions) result = self._run(agent, exception, try_number, max_tries) - # The output type validated the answer, so it is one of the configured names. name = picked_key(result.output) category = self.categories.get(name) if category is None: - # The output type validates the answer, so this needs the schema and the table to disagree. + # The output type constrains the answer to the configured names, so reaching this needs + # the schema and the table to disagree. log.error("Classifier answered %r, which is not a configured category", name) return "model_error" # A bare output type is one field, ``response``; its confidence is what the bar is compared against. model_confidence = ModelConfidence.from_result(result) + model_name = model_confidence.model or "n/a" 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) @@ -668,20 +674,20 @@ def _classify( ) if uncertain is not None: log.info( - "Classifier answer not acted on (%s): %s. model=%s probabilities=%s", + "Classifier answer not applied (%s): %s. model=%s probabilities=%s", uncertain, summary, - model_confidence.model or "n/a", + model_name, model_confidence.probabilities.get(BARE_OUTPUT_FIELD) or "n/a", ) return uncertain - if not category.retry: + if category.retry: + 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}" + decision = RetryDecision.retry(delay=category.delay, reason=reason) + else: reason = f"{summary} action=fail" - log.info("Classifier decision: %s model=%s", reason, model_confidence.model or "n/a") - return RetryDecision.fail(reason=reason) - - 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("Classifier decision: %s model=%s", reason, model_confidence.model or "n/a") - return RetryDecision.retry(delay=category.delay, reason=reason) + decision = RetryDecision.fail(reason=reason) + log.info("Classifier decision: %s model=%s", reason, model_name) + return decision diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py index a6e0875183861..827de2b0730fb 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py @@ -400,7 +400,7 @@ def test_hand_built_context_skips_the_decision_push_with_a_warning( with caplog.at_level(logging.WARNING): output = op.execute(context) - assert isinstance(output, (Summary, dict)) + assert Summary.model_validate(output).text == "t" assert "the decision record was not pushed to XCom" in caplog.text @pytest.mark.skipif( 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 012a0dadf2a1e..ad98fe52bef08 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 @@ -46,7 +46,7 @@ ) 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, RetryDecision, RetryRule +from airflow.sdk.definitions.retry_policy import RetryAction, RetryDecision, RetryPolicy, RetryRule from airflow.sdk.log import mask_secret HOOK = "airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook" @@ -160,7 +160,7 @@ def test_category_instructions_do_not_recite_the_taxonomy(self): assert f"- {name}:" in DEFAULT_INSTRUCTIONS -class TestLLMRetryPolicyConstruction: +class TestClassifierRetryPolicyConstruction: @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.""" @@ -229,6 +229,20 @@ def test_policy_can_be_deep_copied(self, categories): assert copy.deepcopy(policy).categories == policy.categories + def test_chained_policy_can_be_deep_copied(self): + """The example Dag ships a classifier with an LLMRetryPolicy behind it.""" + policy = ClassifierRetryPolicy( + llm_conn_id="jev", + min_confidence=0.8, + on_uncertain=LLMRetryPolicy(llm_conn_id="text", timeout=7.0), + ) + + copied = copy.deepcopy(policy) + + assert isinstance(copied.on_uncertain, LLMRetryPolicy) + assert copied.on_uncertain is not policy.on_uncertain + assert copied.on_uncertain.timeout == 7.0 + @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.""" @@ -539,7 +553,7 @@ def test_under_the_bar_with_no_rules_keeps_the_task_behaviour(self, mock_hook_cl [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" + "not applied (below_threshold): category=auth confidence=0.30 threshold=0.60" in record.getMessage() ) assert "probabilities={'auth': 0.3, 'network': 0.28}" in record.getMessage() @@ -865,20 +879,6 @@ def test_answer_missing_from_the_table_falls_back_with_an_error_log(self, mock_h assert "answered 'not_a_category', which is not a configured category" in caplog.text assert "KeyError" not in caplog.text - def test_subclass_classify_returning_the_wrong_type_is_reported_not_acted_on(self, caplog): - class Odd(ClassifierRetryPolicy): - def _classify(self, exception, try_number, max_tries): - return {"should_retry": True} # the 0.9.0 shape, not a RetryDecision - - policy = Odd(llm_conn_id="test") - - with caplog.at_level(logging.ERROR, logger="airflow.providers.common.ai.policies.retry"): - decision = policy.evaluate(ValueError("x"), try_number=1, max_tries=3) - - assert decision.action == RetryAction.DEFAULT - assert decision.reason == "classifier answer not applied (model_error); task retry settings apply" - assert "returned dict instead of a RetryDecision" in caplog.text - class TestOnUncertain: """Classifier first; an LLM policy when it is unsure or unreachable; the rules when neither decides.""" @@ -889,16 +889,21 @@ class TestOnUncertain: ) ] - def _chain(self, mock_hook_cls, classifier_answer, confidence, llm_answer): - """Wire two agents behind one patched hook: the classifier by output type, the LLM by ErrorClassification.""" - classifier = ( - _agent(classifier_answer, confidence=confidence) if classifier_answer else MagicMock(spec=Agent) - ) + def _chain(self, mock_hook_cls, classifier_answer, confidence, llm_answer, *, bar=0.8, rules=RULES): + """Wire two agents behind one patched hook: the classifier by output type, the LLM by ErrorClassification. + + ``None`` for either answer makes that model unreachable. + """ if classifier_answer is None: + classifier = MagicMock(spec=Agent) classifier.run_sync.side_effect = TimeoutError("classifier unreachable") - llm = _agent(llm_answer) if isinstance(llm_answer, ErrorClassification) else MagicMock(spec=Agent) + else: + classifier = _agent(classifier_answer, confidence=confidence) if llm_answer is None: + llm = MagicMock(spec=Agent) llm.run_sync.side_effect = TimeoutError("llm unreachable") + else: + llm = _agent(llm_answer) def create_agent(**kwargs): return llm if kwargs["output_type"] is ErrorClassification else classifier @@ -906,16 +911,24 @@ def create_agent(**kwargs): mock_hook_cls.return_value.create_agent.side_effect = create_agent return ClassifierRetryPolicy( llm_conn_id="jev", - min_confidence=0.8, + min_confidence=bar, on_uncertain=LLMRetryPolicy(llm_conn_id="text"), - fallback_rules=self.RULES, + fallback_rules=rules, ) - def test_on_uncertain_needs_a_bar(self): - with pytest.raises(ValueError, match="on_uncertain needs min_confidence"): - ClassifierRetryPolicy( - llm_conn_id="jev", categories=DEFAULT_CATEGORIES, on_uncertain=LLMRetryPolicy(llm_conn_id="t") - ) + @patch(HOOK, autospec=True) + def test_without_a_bar_the_llm_is_consulted_only_when_the_classifier_is_unreachable(self, mock_hook_cls): + llm_answer = ErrorClassification(category="auth", should_retry=False, reasoning="expired key") + + answered = self._chain(mock_hook_cls, "network", None, llm_answer, bar=None) + decision = answered.evaluate(RuntimeError("reset"), try_number=1, max_tries=3) + assert decision.action == RetryAction.RETRY + assert decision.reason == "category=network confidence=n/a threshold=n/a action=retry delay=10s" + + unreachable = self._chain(mock_hook_cls, None, None, llm_answer, bar=None) + decision = unreachable.evaluate(RuntimeError("reset"), try_number=1, max_tries=3) + assert decision.action == RetryAction.FAIL + assert decision.reason == "escalated (model_error); auth: expired key" def test_on_uncertain_must_be_a_retry_policy(self): with pytest.raises(TypeError, match="on_uncertain must be a RetryPolicy"): @@ -974,8 +987,7 @@ def test_unreachable_llm_falls_to_the_rules(self, mock_hook_cls): @patch(HOOK, autospec=True) def test_nothing_decides_means_the_task_default(self, mock_hook_cls): - policy = self._chain(mock_hook_cls, "auth", 0.45, None) - policy.fallback_rules = None + policy = self._chain(mock_hook_cls, "auth", 0.45, None, rules=None) decision = policy.evaluate(ValueError("?"), try_number=1, max_tries=3) @@ -996,6 +1008,86 @@ def test_inner_policy_with_its_own_rules_is_used_as_is(self): assert decision.reason == "escalated (model_error); inner rule" inner.evaluate.assert_called_once() + def test_inner_decision_without_a_reason_is_prefixed_without_a_none_suffix(self): + """An ExceptionRetryPolicy with default=RETRY returns RETRY and no reason when nothing matches.""" + inner = MagicMock(spec=LLMRetryPolicy) + inner.evaluate.return_value = RetryDecision.retry(delay=timedelta(seconds=9)) + policy = ClassifierRetryPolicy(llm_conn_id="nonexistent", min_confidence=0.8, on_uncertain=inner) + + decision = policy.evaluate(RuntimeError("?"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.retry_delay == timedelta(seconds=9) + assert decision.reason == "escalated (model_error)" + + @pytest.mark.parametrize( + "inner_decision", + [ + pytest.param(RetryDecision.default(), id="bare-default"), + pytest.param( + RetryDecision(action=RetryAction.DEFAULT, reason="matched a DEFAULT rule"), + id="default-with-reason", + ), + pytest.param( + RetryDecision(action=RetryAction.DEFAULT, retry_delay=timedelta(seconds=99), reason="x"), + id="default-with-delay", + ), + ], + ) + def test_any_default_from_on_uncertain_lets_the_outer_rules_run(self, inner_decision, caplog): + """Only RETRY or FAIL ends the chain; the reason text on a DEFAULT does not make it a decision.""" + inner = MagicMock(spec=LLMRetryPolicy) + inner.evaluate.return_value = inner_decision + policy = ClassifierRetryPolicy( + llm_conn_id="nonexistent", min_confidence=0.8, on_uncertain=inner, fallback_rules=self.RULES + ) + + 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.RETRY + assert decision.retry_delay == timedelta(seconds=5) + assert decision.reason == "classifier answer not applied (model_error); rule" + assert "decided nothing" in caplog.text + + def test_on_uncertain_raising_falls_to_the_outer_rules(self, caplog): + """A third-party policy that blows up must not take the classifier's rules floor with it.""" + + class Boom(RetryPolicy): + def evaluate(self, exception, try_number, max_tries, context=None): + raise RuntimeError("policy bug") + + policy = ClassifierRetryPolicy( + llm_conn_id="nonexistent", min_confidence=0.8, on_uncertain=Boom(), fallback_rules=self.RULES + ) + + with caplog.at_level(logging.ERROR, logger="airflow.providers.common.ai.policies.retry"): + decision = policy.evaluate(RuntimeError("?"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.reason == "classifier answer not applied (model_error); rule" + assert "on_uncertain policy failed" in caplog.text + + @patch(HOOK, autospec=True) + def test_nested_classifiers_both_down_still_reach_the_outer_fail_rule(self, mock_hook_cls): + """A ClassifierRetryPolicy as on_uncertain falls back with a DEFAULT of its own; the outer rules must still run.""" + mock_hook_cls.return_value.create_agent.return_value.run_sync.side_effect = TimeoutError("down") + inner = ClassifierRetryPolicy(llm_conn_id="jev_b", min_confidence=0.5) + policy = ClassifierRetryPolicy( + llm_conn_id="jev_a", + min_confidence=0.8, + on_uncertain=inner, + fallback_rules=[ + RetryRule(exception=PermissionError, action=RetryAction.FAIL, reason="never retry 403") + ], + ) + + decision = policy.evaluate(PermissionError("403"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.FAIL + assert decision.reason == "classifier answer not applied (model_error); never retry 403" + assert mock_hook_cls.return_value.create_agent.call_count == 2 + def _open_agent(category, should_retry, delay=0, reasoning="test"): """A mock agent answering the 0.9.0 ``ErrorClassification`` shape, for a policy without categories.""" @@ -1009,12 +1101,24 @@ def _open_agent(category, should_retry, delay=0, reasoning="test"): class TestLLMRetryPolicy: """The LLM layer is the 0.9.0 policy: the text model decides retry and delay, and nothing classifier-shaped is on it.""" - def test_defaults(self): - policy = LLMRetryPolicy(llm_conn_id="test") + def test_positional_arguments_keep_the_0_9_0_order(self): + rules = [RetryRule(exception=ValueError, action=RetryAction.FAIL)] + + policy = LLMRetryPolicy("conn", "openai:gpt-4o", "be brief", rules, 12.0) - assert policy.instructions == DEFAULT_INSTRUCTIONS - assert not hasattr(policy, "categories") - assert not hasattr(policy, "min_confidence") + assert ( + policy.llm_conn_id, + policy.model_id, + policy.instructions, + policy.fallback_rules, + policy.timeout, + ) == ( + "conn", + "openai:gpt-4o", + "be brief", + rules, + 12.0, + ) @pytest.mark.parametrize("kwarg", ["categories", "min_confidence", "on_uncertain"]) def test_has_no_classifier_arguments(self, kwarg): @@ -1022,12 +1126,6 @@ def test_has_no_classifier_arguments(self, kwarg): with pytest.raises(TypeError, match="unexpected keyword argument"): LLMRetryPolicy(llm_conn_id="test", **{kwarg: None}) - def test_classifier_defaults(self): - policy = ClassifierRetryPolicy(llm_conn_id="test", min_confidence=0.6) - - assert policy.categories == dict(DEFAULT_CATEGORIES) - assert policy.instructions == CLASSIFIER_INSTRUCTIONS - @patch(HOOK, autospec=True) def test_model_output_type_is_error_classification(self, mock_hook_cls): _install(mock_hook_cls, _open_agent("auth", should_retry=False)) @@ -1131,7 +1229,7 @@ def test_classifier_refusal_logs_the_categories_hint(self, mock_hook_cls, caplog decision = policy.evaluate(ValueError("x"), try_number=1, max_tries=3) assert decision.action == RetryAction.DEFAULT - assert "Use ClassifierRetryPolicy" in caplog.text + assert "use ClassifierRetryPolicy" in caplog.text def test_no_warning_for_any_instructions(self): with warnings.catch_warnings(): From 6260f11d73c2135ac4db5c1cfa58d0d695857c78 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Mon, 21 Sep 2026 22:25:49 +0100 Subject: [PATCH 3/3] Rename ClassifierRetryPolicy's on_uncertain to fallback_policy The policy behind a classifier is consulted on a low-confidence answer, a missing confidence and a model outage alike, and it can be plain rules rather than a stronger model, so "fallback" describes it and "uncertain" does not. The operators' DecisionPolicy keeps on_uncertain, where the value is an action ("review" or "fail"), not a policy; using one name for two different kinds of value in the same package invited confusion. The name has not shipped in a release. --- .../common/ai/docs/classifier_models.rst | 4 +-- providers/common/ai/docs/retry_policies.rst | 22 ++++++------ .../example_dags/example_llm_retry_policy.py | 2 +- .../providers/common/ai/policies/retry.py | 28 +++++++-------- .../unit/common/ai/policies/test_retry.py | 36 +++++++++---------- 5 files changed, 46 insertions(+), 46 deletions(-) diff --git a/providers/common/ai/docs/classifier_models.rst b/providers/common/ai/docs/classifier_models.rst index 80cac0e650a69..4d82f08d55995 100644 --- a/providers/common/ai/docs/classifier_models.rst +++ b/providers/common/ai/docs/classifier_models.rst @@ -116,7 +116,7 @@ Where it fits in this provider - 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 ``min_confidence`` and an unsure answer goes to ``on_uncertain`` + worker. Set ``min_confidence`` and an unsure answer goes to ``fallback_policy`` (typically an ``LLMRetryPolicy`` on a text model), then ``fallback_rules``, then the task's own retry behaviour, instead of ending the task on the model's say-so. ``LLMRetryPolicy`` itself asks for free text, which a classifier model refuses. @@ -143,7 +143,7 @@ and :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` take a 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:`ClassifierRetryPolicy ` takes the same ``min_confidence`` and hands an -unsure answer to ``on_uncertain``, then its deterministic fallback rules. In the branch operator and the retry policy, a +unsure answer to ``fallback_policy``, then 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`` diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index e61d287453405..f513479bc858e 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -55,7 +55,7 @@ fully model-driven, with the SDK's ``ExceptionRetryPolicy`` as the bottom rung: Each class has only the arguments its layer needs. ``LLMRetryPolicy`` is the policy this guide has always described and is unchanged. The layers chain: -``on_uncertain`` on a ``ClassifierRetryPolicy`` names the policy to consult when +``fallback_policy`` on a ``ClassifierRetryPolicy`` names the policy to consult when the classifier is unsure or unreachable, typically an ``LLMRetryPolicy`` on a text model, so the cheap typed model handles the clear cases and the reasoning model the rest, and the rules catch what neither decides. See @@ -140,7 +140,7 @@ If the model call fails (provider down, timeout, bad credentials), the policy falls back to ``fallback_rules`` if configured, or to the task's standard retry behaviour. ``ClassifierRetryPolicy`` does the same when the model cannot produce one of the categories even after pydantic-ai re-prompts it, or when -its answer is under the confidence bar, after consulting ``on_uncertain`` if +its answer is under the confidence bar, after consulting ``fallback_policy`` if set; its fallback decision's reason then starts with ``classifier answer not applied (model_error)``, ``(below_threshold)`` or ``(missing_confidence)`` so a ``retry_reason`` read later is not mistaken for a @@ -320,7 +320,7 @@ Escalating to an LLM -------------------- A classifier is cheap and fast, and a text model can reason about a failure it -has never seen a category for. ``on_uncertain`` puts one behind the other. +has never seen a category for. ``fallback_policy`` puts one behind the other. ``snowflake_policy`` is the classifier policy from the previous section; the chain reuses its category table: @@ -334,22 +334,22 @@ The order of events on a failure: 1. The classifier names a category. At or above the bar, its category's action and delay apply and the text model is never called. 2. Under the bar, with no confidence reported, or if the classifier call fails, - the ``on_uncertain`` policy runs. A text-model ``LLMRetryPolicy`` there + the ``fallback_policy`` policy runs. A text-model ``LLMRetryPolicy`` there classifies the failure with its own instructions and chooses retry and delay itself. Its decision is used, with the reason prefixed by why the classifier's answer was not: ``escalated (below_threshold); rate_limit: 429 with a Retry-After header``. 3. If that policy returns DEFAULT, whatever reason it attached, it decided nothing: the outer ``fallback_rules`` apply, then the task's own retry - behaviour. Only a RETRY or FAIL from ``on_uncertain`` ends the chain, so an + behaviour. Only a RETRY or FAIL from ``fallback_policy`` ends the chain, so an outer rule such as ``PermissionError -> FAIL`` still holds when both models are unreachable. -``on_uncertain`` accepts any ``RetryPolicy``. An ``ExceptionRetryPolicy`` works +``fallback_policy`` accepts any ``RetryPolicy``. An ``ExceptionRetryPolicy`` works there too; its ``default`` is what it returns when none of its rules match, so ``default=RetryAction.FAIL`` fails every unsure classification and the outer rules never run. Without ``min_confidence`` the classifier's answer is always -acted on, and ``on_uncertain`` is consulted only when the classifier call itself +acted on, and ``fallback_policy`` is consulted only when the classifier call itself fails. Both model calls run on the worker at failure time, so a task that escalates pays for two before its retry is scheduled; ``timeout`` on each policy bounds that. @@ -504,7 +504,7 @@ Parameters ---------- Both policies share every parameter below except ``categories``, -``min_confidence`` and ``on_uncertain``, which exist only on +``min_confidence`` and ``fallback_policy``, which exist only on ``ClassifierRetryPolicy``. .. list-table:: @@ -529,7 +529,7 @@ Both policies share every parameter below except ``categories``, - None - List of ``RetryRule`` objects used when the model call fails or, on ``ClassifierRetryPolicy``, when the answer is under its confidence bar and - ``on_uncertain`` decided nothing. + ``fallback_policy`` decided nothing. * - ``timeout`` - 30.0 - Max seconds to wait for the LLM response before falling back. @@ -545,10 +545,10 @@ Both policies share every parameter below except ``categories``, - ``ClassifierRetryPolicy`` only. 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 - ``on_uncertain`` if set, else ``fallback_rules`` then the task's own + ``fallback_policy`` if set, else ``fallback_rules`` then the task's own retry behaviour, apply. A category's own ``min_confidence`` overrides it for that category. - * - ``on_uncertain`` + * - ``fallback_policy`` - None - ``ClassifierRetryPolicy`` only. A ``RetryPolicy`` to consult when the classifier is under its bar, reports no confidence, or cannot be 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 1aaaebacf1446..869fee2e5499f 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 @@ -141,7 +141,7 @@ def task_missing_table(): llm_conn_id="jev_default", min_confidence=0.8, categories=snowflake_policy.categories, - on_uncertain=LLMRetryPolicy(llm_conn_id="pydanticai_default", timeout=30.0), + fallback_policy=LLMRetryPolicy(llm_conn_id="pydanticai_default", timeout=30.0), fallback_rules=[ RetryRule(exception=ConnectionError, action=RetryAction.RETRY, retry_delay=timedelta(seconds=30)), ], 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 2f74b518e0f5b..fffaa5e33c422 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 @@ -28,7 +28,7 @@ * **LLM.** :class:`LLMRetryPolicy`: a text model classifies the failure, decides whether to retry and how long to wait from ``instructions``, and explains itself. -They chain: ``ClassifierRetryPolicy(..., on_uncertain=LLMRetryPolicy(...))`` consults the +They chain: ``ClassifierRetryPolicy(..., fallback_policy=LLMRetryPolicy(...))`` consults the LLM when the classifier is unsure or unreachable, and whatever no layer decides falls to the rules. @@ -345,7 +345,7 @@ class LLMRetryPolicy(_ModelRetryPolicy): to retry, how long to wait, and why, all steered by ``instructions``. This is the reasoning layer; for a cheap typed decision from a classifier model such as TypeSafe's Jev, use :class:`ClassifierRetryPolicy`, which can name this policy as - its ``on_uncertain``. + its ``fallback_policy``. 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. @@ -436,7 +436,7 @@ class ClassifierRetryPolicy(_ModelRetryPolicy): milliseconds and with a confidence; a text model answers it too. When the model call fails, or the answer is under its confidence bar, the policy - consults ``on_uncertain`` if set, then ``fallback_rules``, then returns DEFAULT to use + consults ``fallback_policy`` if set, then ``fallback_rules``, then returns DEFAULT to use the task's standard retry logic. :param llm_conn_id: Airflow connection ID for the model. @@ -447,7 +447,7 @@ class ClassifierRetryPolicy(_ModelRetryPolicy): themselves, and what each one means, are ``categories``. :param fallback_rules: Optional list of :class:`~airflow.sdk.definitions.retry_policy.RetryRule` applied when the model - call fails or the answer is under its confidence bar and ``on_uncertain`` decided + call fails or the answer is under its confidence bar and ``fallback_policy`` decided nothing. Provides a deterministic safety net. :param timeout: Maximum seconds to wait for the model response before falling back. Defaults to 30s. @@ -460,11 +460,11 @@ class ClassifierRetryPolicy(_ModelRetryPolicy): policy to act on it. ``None`` (default) is no bar: the answer is acted on whatever the confidence. 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 ``on_uncertain``, then + reported no confidence, the answer is discarded and ``fallback_policy``, then ``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 on_uncertain: A policy to consult when the answer is under its bar, reports no + :param fallback_policy: A policy to consult when the answer is under its bar, reports no confidence, or the model call fails; typically an :class:`LLMRetryPolicy` on a text model, so the classifier handles the clear cases and a reasoning model the rest. Its RETRY or FAIL is used, with the reason prefixed by why the classifier's answer @@ -488,7 +488,7 @@ def __init__( *, categories: Mapping[str, ErrorCategory] | None = None, min_confidence: float | None = None, - on_uncertain: RetryPolicy | None = None, + fallback_policy: RetryPolicy | None = None, redactor: Callable[[str], str] | None = None, redact_exception: bool = True, max_exception_length: int = 4096, @@ -507,9 +507,9 @@ def __init__( self.categories: dict[str, ErrorCategory] = self._validate_categories( DEFAULT_CATEGORIES if categories is None else categories ) - if on_uncertain is not None and not isinstance(on_uncertain, RetryPolicy): - raise TypeError(f"on_uncertain must be a RetryPolicy, got {type(on_uncertain).__name__}.") - self.on_uncertain = on_uncertain + if fallback_policy is not None and not isinstance(fallback_policy, RetryPolicy): + raise TypeError(f"fallback_policy must be a RetryPolicy, got {type(fallback_policy).__name__}.") + self.fallback_policy = fallback_policy def _validate_categories(self, categories: Mapping[str, ErrorCategory]) -> dict[str, ErrorCategory]: if not isinstance(categories, Mapping): @@ -562,9 +562,9 @@ def evaluate( outcome = "model_error" if isinstance(outcome, RetryDecision): return outcome - if self.on_uncertain is not None: + if self.fallback_policy is not None: escalated = self._escalate( - self.on_uncertain, exception, try_number, max_tries, context, why=outcome + self.fallback_policy, exception, try_number, max_tries, context, why=outcome ) if escalated is not None: return escalated @@ -581,7 +581,7 @@ def _escalate( why: str, ) -> RetryDecision | None: """ - Consult ``on_uncertain`` and return its RETRY or FAIL, or None when it decided nothing. + Consult ``fallback_policy`` and return its RETRY or FAIL, or None when it decided nothing. Only RETRY and FAIL are decisions. DEFAULT means the policy had nothing to add to the task's own settings, whatever reason it attached (its own fallback message, or a matched @@ -593,7 +593,7 @@ def _escalate( try: decision = policy.evaluate(exception, try_number, max_tries, context) except Exception: - log.exception("on_uncertain policy failed, using fallback rules") + log.exception("fallback_policy failed, using fallback rules") return None if decision.action is RetryAction.DEFAULT: log.info( 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 ad98fe52bef08..5168b6b1c4f11 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 @@ -234,14 +234,14 @@ def test_chained_policy_can_be_deep_copied(self): policy = ClassifierRetryPolicy( llm_conn_id="jev", min_confidence=0.8, - on_uncertain=LLMRetryPolicy(llm_conn_id="text", timeout=7.0), + fallback_policy=LLMRetryPolicy(llm_conn_id="text", timeout=7.0), ) copied = copy.deepcopy(policy) - assert isinstance(copied.on_uncertain, LLMRetryPolicy) - assert copied.on_uncertain is not policy.on_uncertain - assert copied.on_uncertain.timeout == 7.0 + assert isinstance(copied.fallback_policy, LLMRetryPolicy) + assert copied.fallback_policy is not policy.fallback_policy + assert copied.fallback_policy.timeout == 7.0 @patch(HOOK, autospec=True) def test_caller_mapping_is_copied(self, mock_hook_cls): @@ -912,7 +912,7 @@ def create_agent(**kwargs): return ClassifierRetryPolicy( llm_conn_id="jev", min_confidence=bar, - on_uncertain=LLMRetryPolicy(llm_conn_id="text"), + fallback_policy=LLMRetryPolicy(llm_conn_id="text"), fallback_rules=rules, ) @@ -930,9 +930,9 @@ def test_without_a_bar_the_llm_is_consulted_only_when_the_classifier_is_unreacha assert decision.action == RetryAction.FAIL assert decision.reason == "escalated (model_error); auth: expired key" - def test_on_uncertain_must_be_a_retry_policy(self): - with pytest.raises(TypeError, match="on_uncertain must be a RetryPolicy"): - ClassifierRetryPolicy(llm_conn_id="jev", min_confidence=0.5, on_uncertain="llm") # type: ignore[arg-type] + def test_fallback_policy_must_be_a_retry_policy(self): + with pytest.raises(TypeError, match="fallback_policy must be a RetryPolicy"): + ClassifierRetryPolicy(llm_conn_id="jev", min_confidence=0.5, fallback_policy="llm") # type: ignore[arg-type] @patch(HOOK, autospec=True) def test_confident_classifier_answer_never_consults_the_llm(self, mock_hook_cls): @@ -999,7 +999,7 @@ def test_inner_policy_with_its_own_rules_is_used_as_is(self): inner = MagicMock(spec=LLMRetryPolicy) inner.evaluate.return_value = RetryDecision.fail(reason="inner rule") policy = ClassifierRetryPolicy( - llm_conn_id="nonexistent", min_confidence=0.8, on_uncertain=inner, fallback_rules=self.RULES + llm_conn_id="nonexistent", min_confidence=0.8, fallback_policy=inner, fallback_rules=self.RULES ) decision = policy.evaluate(RuntimeError("?"), try_number=1, max_tries=3) @@ -1012,7 +1012,7 @@ def test_inner_decision_without_a_reason_is_prefixed_without_a_none_suffix(self) """An ExceptionRetryPolicy with default=RETRY returns RETRY and no reason when nothing matches.""" inner = MagicMock(spec=LLMRetryPolicy) inner.evaluate.return_value = RetryDecision.retry(delay=timedelta(seconds=9)) - policy = ClassifierRetryPolicy(llm_conn_id="nonexistent", min_confidence=0.8, on_uncertain=inner) + policy = ClassifierRetryPolicy(llm_conn_id="nonexistent", min_confidence=0.8, fallback_policy=inner) decision = policy.evaluate(RuntimeError("?"), try_number=1, max_tries=3) @@ -1034,12 +1034,12 @@ def test_inner_decision_without_a_reason_is_prefixed_without_a_none_suffix(self) ), ], ) - def test_any_default_from_on_uncertain_lets_the_outer_rules_run(self, inner_decision, caplog): + def test_any_default_from_fallback_policy_lets_the_outer_rules_run(self, inner_decision, caplog): """Only RETRY or FAIL ends the chain; the reason text on a DEFAULT does not make it a decision.""" inner = MagicMock(spec=LLMRetryPolicy) inner.evaluate.return_value = inner_decision policy = ClassifierRetryPolicy( - llm_conn_id="nonexistent", min_confidence=0.8, on_uncertain=inner, fallback_rules=self.RULES + llm_conn_id="nonexistent", min_confidence=0.8, fallback_policy=inner, fallback_rules=self.RULES ) with caplog.at_level(logging.INFO, logger="airflow.providers.common.ai.policies.retry"): @@ -1050,7 +1050,7 @@ def test_any_default_from_on_uncertain_lets_the_outer_rules_run(self, inner_deci assert decision.reason == "classifier answer not applied (model_error); rule" assert "decided nothing" in caplog.text - def test_on_uncertain_raising_falls_to_the_outer_rules(self, caplog): + def test_fallback_policy_raising_falls_to_the_outer_rules(self, caplog): """A third-party policy that blows up must not take the classifier's rules floor with it.""" class Boom(RetryPolicy): @@ -1058,7 +1058,7 @@ def evaluate(self, exception, try_number, max_tries, context=None): raise RuntimeError("policy bug") policy = ClassifierRetryPolicy( - llm_conn_id="nonexistent", min_confidence=0.8, on_uncertain=Boom(), fallback_rules=self.RULES + llm_conn_id="nonexistent", min_confidence=0.8, fallback_policy=Boom(), fallback_rules=self.RULES ) with caplog.at_level(logging.ERROR, logger="airflow.providers.common.ai.policies.retry"): @@ -1066,17 +1066,17 @@ def evaluate(self, exception, try_number, max_tries, context=None): assert decision.action == RetryAction.RETRY assert decision.reason == "classifier answer not applied (model_error); rule" - assert "on_uncertain policy failed" in caplog.text + assert "fallback_policy failed" in caplog.text @patch(HOOK, autospec=True) def test_nested_classifiers_both_down_still_reach_the_outer_fail_rule(self, mock_hook_cls): - """A ClassifierRetryPolicy as on_uncertain falls back with a DEFAULT of its own; the outer rules must still run.""" + """A ClassifierRetryPolicy as fallback_policy falls back with a DEFAULT of its own; the outer rules must still run.""" mock_hook_cls.return_value.create_agent.return_value.run_sync.side_effect = TimeoutError("down") inner = ClassifierRetryPolicy(llm_conn_id="jev_b", min_confidence=0.5) policy = ClassifierRetryPolicy( llm_conn_id="jev_a", min_confidence=0.8, - on_uncertain=inner, + fallback_policy=inner, fallback_rules=[ RetryRule(exception=PermissionError, action=RetryAction.FAIL, reason="never retry 403") ], @@ -1120,7 +1120,7 @@ def test_positional_arguments_keep_the_0_9_0_order(self): 12.0, ) - @pytest.mark.parametrize("kwarg", ["categories", "min_confidence", "on_uncertain"]) + @pytest.mark.parametrize("kwarg", ["categories", "min_confidence", "fallback_policy"]) def test_has_no_classifier_arguments(self, kwarg): """Those belong to ClassifierRetryPolicy; a text-model user never sees them.""" with pytest.raises(TypeError, match="unexpected keyword argument"):