This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 190befa57c1 Ask the model again when a durable agent's retry would
replay its rejected answers (#74720)
190befa57c1 is described below
commit 190befa57c106a62a8101f240a0ebd3e2429f336
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 14 12:36:51 2026 +0100
Ask the model again when a durable agent's retry would replay its rejected
answers (#74720)
With durable=True, a run that pydantic-ai fails with UnexpectedModelBehavior
(a tool's ModelRetry or output validation past the retry limit) replayed
every recorded answer into the same error on each Airflow retry, without a
live model call. DurableRun.fail now takes model_at_fault: when no step
raised the error, the retry replays the steps before the model's last
request and asks the model again from there. A previous attempt's
divergence point also stays until an attempt records over it, so a re-asked
request that fails is not replayed from the older attempt.
---
providers/common/ai/docs/durable_execution.rst | 17 +++-
.../providers/common/ai/durable/capability.py | 13 ++-
.../airflow/providers/common/ai/durable/journal.py | 63 ++++++++++--
.../unit/common/ai/durable/test_capability.py | 106 ++++++++++++++++++++-
.../tests/unit/common/ai/durable/test_journal.py | 90 ++++++++++++++++-
5 files changed, 277 insertions(+), 12 deletions(-)
diff --git a/providers/common/ai/docs/durable_execution.rst
b/providers/common/ai/docs/durable_execution.rst
index 59d4aaa7408..5c670045998 100644
--- a/providers/common/ai/docs/durable_execution.rst
+++ b/providers/common/ai/docs/durable_execution.rst
@@ -151,6 +151,20 @@ other capabilities. Tool calls already running alongside
the failed one still re
step that raised and that the run recovered from, such as a tool error the
model was
told about, does not stop the steps after it from replaying.
+A run can also fail on the model's answers rather than in a step: pydantic-ai
raises
+``UnexpectedModelBehavior`` when the model keeps giving answers a tool rejects
with
+``ModelRetry``, or output that fails validation, past its retry limit. Those
answers are
+recorded, so replaying them would only reach the same error. The retry instead
replays
+the steps before the model's last request and asks the model again from there.
The
+earlier rejected answers still count against the retry limit, so each Airflow
retry gives
+the model one more answer, with the rejections it already got in its history.
Tool calls
+the last answer made run again if the model asks for them again, including
calls that
+succeeded alongside the one that was rejected. The task log names the step:
+
+.. code-block:: text
+
+ Durable: the run took a different path from the previous attempt; this
step and every step after it run again position=2 step=analyst__model.request
reason=the model's answers from this step on failed the previous attempt, so it
is asked again
+
Some steps have no request to fingerprint: another capability's
``@durable_operation``, or tool discovery. Such a step replays on its name,
and only when
the step before it also replayed.
@@ -236,7 +250,8 @@ Durable execution records what a tool returned. A replayed
tool's code does not
side effect it had (a file written, a message sent) is not repeated.
A tool that fails *before* its result is recorded runs again on retry, so a
tool that
sent an email and then raised sends it again. The same holds when the worker
dies
-between a tool's side effect and the write of its result. Make tools with side
effects
+between a tool's side effect and the write of its result, and for the tool
calls of a
+model answer that the retry asks for again after the model ran out of retries.
Make tools with side effects
idempotent: check whether the work already happened, or let a database
constraint
reject the duplicate.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
index 3a420b17adf..cabbb425336 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
@@ -37,7 +37,13 @@ from pydantic_ai.durable_exec import (
ToolsetCallToolId,
ToolsetValidateToolArgumentsId,
)
-from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry,
ToolFailed
+from pydantic_ai.exceptions import (
+ ApprovalRequired,
+ CallDeferred,
+ ModelRetry,
+ ToolFailed,
+ UnexpectedModelBehavior,
+)
from pydantic_ai.messages import ModelResponse, ToolReturn
from pydantic_ai.tools import AgentDepsT
from pydantic_ai.toolsets import DynamicToolset, FunctionToolset
@@ -227,7 +233,10 @@ class
AirflowDurability(BaseDurabilityCapability[AgentDepsT]):
with _active_run(_ActiveRun(durable_run, ledger)):
result = await super().wrap_run(ctx, handler=handler)
except BaseException as error:
- durable_run.fail(error)
+ # pydantic-ai raises UnexpectedModelBehavior when the model's
answers fail its checks
+ # too often (a tool's ModelRetry, output validation): the answers
are recorded, so a
+ # retry has to ask the model again rather than replay them into
the same error.
+ durable_run.fail(error, model_at_fault=isinstance(error,
UnexpectedModelBehavior))
raise
finally:
ledger.settle()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
index 4683634c1cd..ffaf955098e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
@@ -29,7 +29,8 @@ had the same name and the same fingerprint. The name says
what kind of step it w
fingerprint, computed by the framework adapter, says what it was asked (the
message
history, the tool arguments). The first mismatch means the run took a
different path
from the previous attempt, so that step and every step after it run live. So
does
-everything the previous attempt did after a step that raised.
+everything the previous attempt did after a step that raised, and, when the
previous
+attempt failed on the model's answers, everything from the model's last
request on.
A :class:`DurableJournal` belongs to one task attempt. Each agent run in it is
a
:class:`DurableRun`, numbered in the order the runs start, with positions of
its own,
@@ -206,8 +207,14 @@ class DurableRun:
self._peeked: dict[int, dict[str, Any] | None] = {}
# Exceptions steps raised this attempt, with the next position when
each was raised.
self._raised: list[tuple[BaseException, int]] = []
- # The furthest position this attempt recorded a step at.
+ # The furthest position this attempt recorded a step at, and every
position it recorded.
self._last_saved = -1
+ self._saved: set[int] = set()
+ # The position of the last model step this attempt claimed.
+ self._last_model_position: int | None = None
+ # Read with the rest of the run's entry: whether the previous attempt
failed because
+ # of the model's answers, which changes why the retry runs live from
``_diverge_from``.
+ self._model_at_fault = False
@property
def position(self) -> int:
@@ -222,6 +229,7 @@ class DurableRun:
diverge_from, high_water = meta.get("diverge_from"),
meta.get("high_water")
self._diverge_from = diverge_from if isinstance(diverge_from, int)
else None
self._high_water = high_water if isinstance(high_water, int) else 0
+ self._model_at_fault = meta.get("model_at_fault") is True
def peek(self, position: int) -> dict[str, Any] | None:
"""
@@ -256,9 +264,19 @@ class DurableRun:
position = self._position
self._position += 1
step = JournalStep(self, position, name, kind, fingerprint,
replayable=replayable)
+ if kind == "model":
+ self._last_model_position = position
self._load_meta()
if not self.diverged and self._diverge_from is not None and position
>= self._diverge_from:
- self._diverge(position, name, reason="the previous attempt failed
before reaching this step")
+ self._diverge(
+ position,
+ name,
+ reason=(
+ "the model's answers from this step on failed the previous
attempt, so it is asked again"
+ if self._model_at_fault
+ else "the previous attempt failed before reaching this
step"
+ ),
+ )
entry = self.peek(position) if replayable else None
self._peeked.pop(position, None)
previous_replayed, self._previous_replayed = self._previous_replayed,
False
@@ -303,26 +321,57 @@ class DurableRun:
def save(self, position: int, entry: dict[str, Any]) -> bool:
"""Store ``entry`` at ``position``; ``False`` when the storage skipped
it."""
self._last_saved = max(self._last_saved, position)
- return self.journal.storage.save_step(build_step_key(self.key,
position), entry)
+ if not self.journal.storage.save_step(build_step_key(self.key,
position), entry):
+ return False
+ self._saved.add(position)
+ return True
- def fail(self, error: BaseException) -> None:
+ def fail(self, error: BaseException, *, model_at_fault: bool = False) ->
None:
"""
Record that the run failed with ``error``, so a retry does not replay
its error path.
Steps the run started before the step that raised ``error`` (tool
calls running
alongside it, say) still replay on retry; everything from then on,
including what
error handling recorded, runs again.
+
+ :param model_at_fault: The run failed because of what the model
answered rather than
+ in a step, such as the framework giving up after the model's
answers failed its
+ checks too many times. Those answers were recorded and would
replay into the same
+ failure on every retry, so the retry asks the model again from its
last request.
+ Ignored when a step raised ``error``: that step decides where the
retry runs live.
"""
failed_at = self._position
+ raised_in_step = False
cause: BaseException | None = error
while cause is not None:
- failed_at = min([failed_at, *(position for raised, position in
self._raised if raised is cause)])
+ positions = [position for raised, position in self._raised if
raised is cause]
+ raised_in_step = raised_in_step or bool(positions)
+ failed_at = min([failed_at, *positions])
cause = cause.__cause__ or cause.__context__
self._load_meta()
# An attempt that was still replaying an earlier one when it failed,
and recorded
# nothing from the failure on, leaves that earlier attempt's steps as
they were.
diverge_from = failed_at if self.diverged or self._last_saved >=
failed_at else self._diverge_from
- meta = {"diverge_from": diverge_from, "high_water":
max(self._high_water, self._position)}
+ rerun_model = False
+ if model_at_fault and not raised_in_step and (last_model :=
self._last_model_position) is not None:
+ # Unlike a step that raised, the failure is known to come after
every recorded step,
+ # so this holds even for an attempt that was still replaying.
+ if diverge_from is None or last_model <= diverge_from:
+ diverge_from, rerun_model = last_model, True
+ previous = self._diverge_from
+ if (
+ previous is not None
+ and previous not in self._saved
+ and (diverge_from is None or previous < diverge_from)
+ ):
+ # The entry at the previous attempt's point is still that
attempt's, as this one
+ # recorded nothing over it (the step there raised, say), so a
retry must not replay it.
+ diverge_from, rerun_model = previous, self._model_at_fault
+ meta = {
+ "diverge_from": diverge_from,
+ "high_water": max(self._high_water, self._position),
+ "model_at_fault": rerun_model,
+ }
self.journal.storage.save_step(build_run_meta_key(self.key), meta)
def cleanup(self) -> None:
diff --git
a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
index 02fb114ee6b..c747afa147d 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
@@ -28,9 +28,10 @@ import dataclasses
from typing import TYPE_CHECKING, Any
import pytest
+from pydantic import BaseModel
from pydantic_ai import Agent, CancellationToken, RunContext
from pydantic_ai.capabilities import AbstractCapability, durable_operation
-from pydantic_ai.exceptions import ModelRetry
+from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior
from pydantic_ai.messages import (
ModelMessage,
ModelResponse,
@@ -494,6 +495,109 @@ class TestAirflowToolsets:
assert calls["model"] == 3
+class Answer(BaseModel):
+ rows: int
+
+
+class TestModelAtFault:
+ """A run pydantic-ai fails because of the model's answers is retried by
asking the model again."""
+
+ @pytest.mark.asyncio
+ async def
test_a_retry_asks_the_model_again_after_its_tool_retries_ran_out(self,
memory_storage):
+ calls = Calls()
+ sql = ["select nope"]
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) ->
ModelResponse:
+ calls.bump("model")
+ if returns := tool_returns(messages):
+ return ModelResponse(parts=[TextPart(f"answer:
{returns[-1].content}")])
+ return ModelResponse(parts=[ToolCallPart("query", {"sql":
sql[0]})])
+
+ def build() -> Agent[None, str]:
+ toolset = FunctionToolset(id="db", max_retries=1)
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ calls.bump("query")
+ if "nope" in sql:
+ raise ModelRetry("no column nope")
+ return "3 rows"
+
+ return Agent(
+ FunctionModel(model_fn),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(UnexpectedModelBehavior, match="exceeded max
retries"):
+ await attempt(memory_storage, build())
+ sql[0] = "select 1"
+ with journal_scope(journal := DurableJournal(memory_storage)):
+ result = await build().run("go")
+
+ assert result.output == "answer: 3 rows"
+ assert dict(journal.stats.replayed) == {"model": 1, "tool": 1}
+ # Attempt 1: two bad queries. Attempt 2 replays the first request and
its rejected
+ # query, asks the model again (good query), then for the answer.
+ assert (calls["model"], calls["query"]) == (4, 3)
+
+ @pytest.mark.asyncio
+ async def
test_a_retry_asks_the_model_again_after_its_output_retries_ran_out(self,
memory_storage):
+ calls = Calls()
+ rows: list[Any] = ["many"]
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) ->
ModelResponse:
+ calls.bump("model")
+ return
ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, {"rows":
rows[0]})])
+
+ def build() -> Agent[None, Answer]:
+ return Agent(
+ FunctionModel(model_fn),
+ name="analyst",
+ output_type=Answer,
+ retries=1,
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(UnexpectedModelBehavior, match="output retries"):
+ await attempt(memory_storage, build())
+ rows[0] = 3
+ result = await attempt(memory_storage, build())
+
+ assert result.output == Answer(rows=3)
+ # The first invalid answer replays; only the model's last request runs
again.
+ assert calls["model"] == 3
+
+ @pytest.mark.asyncio
+ async def test_a_tool_that_raises_keeps_the_models_request_replaying(self,
memory_storage):
+ """The tool failed, not the model: only the tool runs again."""
+ calls = Calls()
+
+ def build() -> Agent[None, str]:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ calls.bump("query")
+ if calls["query"] == 1:
+ raise RuntimeError("connection reset")
+ return "3 rows"
+
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query",
fail_final=[False])),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="connection reset"):
+ await attempt(memory_storage, build())
+ await attempt(memory_storage, build())
+
+ assert (calls["model"], calls["query"]) == (2, 2)
+
+
@pytest.mark.enable_redact
class TestMasking:
@pytest.mark.asyncio
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
b/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
index fe962087f08..41f70605409 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
@@ -25,7 +25,7 @@ from unittest import mock
import pytest
from airflow.providers.common.ai.durable import journal as journal_module
-from airflow.providers.common.ai.durable.base import build_step_key
+from airflow.providers.common.ai.durable.base import build_run_meta_key,
build_step_key
from airflow.providers.common.ai.durable.journal import (
DurableJournal,
current_journal,
@@ -202,6 +202,94 @@ class TestFailures:
assert await retry.run("a", kind="tool", fingerprint="a", body=lambda:
ok("live")) == "recorded"
+ @pytest.mark.asyncio
+ async def test_a_failure_the_model_caused_reruns_its_last_request(self,
memory_storage):
+ """The model's recorded answers would replay into the same failure, so
it is asked again."""
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("model", kind="model", fingerprint="m1", body=lambda:
ok("bad query"))
+ await first.run("query", kind="tool", fingerprint="q1", body=lambda:
ok("retry: no column"))
+ await first.run("model", kind="model", fingerprint="m2", body=lambda:
ok("bad query again"))
+ await first.run("query", kind="tool", fingerprint="q2", body=lambda:
ok("retry: no column"))
+ first.fail(RuntimeError("exceeded max retries"), model_at_fault=True)
+
+ retry = DurableJournal(memory_storage).start_run()
+ results = [
+ await retry.run("model", kind="model", fingerprint="m1",
body=lambda: ok("live")),
+ await retry.run("query", kind="tool", fingerprint="q1",
body=lambda: ok("live")),
+ await retry.run("model", kind="model", fingerprint="m2",
body=lambda: ok("good query")),
+ await retry.run("query", kind="tool", fingerprint="q2",
body=lambda: ok("3 rows")),
+ ]
+
+ assert results == ["bad query", "retry: no column", "good query", "3
rows"]
+ assert dict(retry.journal.stats.replayed) == {"model": 1, "tool": 1}
+
+ def test_a_failure_the_model_caused_is_noted_for_the_retrys_log(self,
memory_storage):
+ first = DurableJournal(memory_storage).start_run()
+ first.claim("model", kind="model", fingerprint="m1").record("bad
query")
+
+ first.fail(RuntimeError("exceeded max retries"), model_at_fault=True)
+
+ meta = memory_storage.load_step(build_run_meta_key(first.key))
+ assert (meta["diverge_from"], meta["model_at_fault"]) == (0, True)
+
+ @pytest.mark.asyncio
+ async def
test_a_failure_the_model_caused_while_replaying_reruns_its_last_request(self,
memory_storage):
+ """Otherwise an attempt that replayed everything would leave the next
one replaying it too."""
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("model", kind="model", fingerprint="m1", body=lambda:
ok("bad"))
+ await first.run("query", kind="tool", fingerprint="q1", body=lambda:
ok("retry"))
+ first.fail(RuntimeError("worker died"))
+ second = DurableJournal(memory_storage).start_run()
+ await second.run("model", kind="model", fingerprint="m1", body=lambda:
ok("live"))
+ await second.run("query", kind="tool", fingerprint="q1", body=lambda:
ok("live"))
+ second.fail(RuntimeError("exceeded max retries"), model_at_fault=True)
+
+ third = DurableJournal(memory_storage).start_run()
+ answer = await third.run("model", kind="model", fingerprint="m1",
body=lambda: ok("asked again"))
+
+ assert answer == "asked again"
+
+ @pytest.mark.asyncio
+ async def
test_a_model_request_asked_again_that_raises_is_asked_again_on_the_next_retry(
+ self, memory_storage
+ ):
+ """The rejected answer is still recorded at that position, so it must
not replay."""
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("model", kind="model", fingerprint="m1", body=lambda:
ok("bad query"))
+ await first.run("query", kind="tool", fingerprint="q1", body=lambda:
ok("retry: no column"))
+ first.fail(RuntimeError("exceeded max retries"), model_at_fault=True)
+ second = DurableJournal(memory_storage).start_run()
+ rate_limited = RuntimeError("429")
+ with pytest.raises(RuntimeError):
+ await second.run("model", kind="model", fingerprint="m1",
body=lambda: raise_(rate_limited))
+ second.fail(rate_limited)
+
+ third = DurableJournal(memory_storage).start_run()
+ answer = await third.run("model", kind="model", fingerprint="m1",
body=lambda: ok("good query"))
+
+ assert answer == "good query"
+
+ @pytest.mark.asyncio
+ async def
test_a_step_that_raised_the_error_decides_even_when_the_model_is_blamed(self,
memory_storage):
+ """An agent a tool runs can give up on its own model; the outer
model's answer was fine."""
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("model", kind="model", fingerprint="m1", body=lambda:
ok("call both tools"))
+ error = RuntimeError("nested agent exceeded max retries")
+ nested = first.claim("nested_agent", kind="tool", fingerprint="n")
+ sibling = first.claim("charge_card", kind="tool", fingerprint="c")
+ nested.fail(error)
+ sibling.record("charged")
+ first.fail(error, model_at_fault=True)
+
+ retry = DurableJournal(memory_storage).start_run()
+ model = await retry.run("model", kind="model", fingerprint="m1",
body=lambda: ok("asked again"))
+ await retry.run("nested_agent", kind="tool", fingerprint="n",
body=lambda: ok("ok"))
+ charge = await retry.run(
+ "charge_card", kind="tool", fingerprint="c", body=lambda:
ok("charged twice")
+ )
+
+ assert (model, charge) == ("call both tools", "charged")
+
@pytest.mark.asyncio
async def
test_a_step_without_a_fingerprint_does_not_replay_after_a_live_step(self,
memory_storage):
"""The live step may have produced something else, which the
unfingerprinted step would depend on."""
