kaxil commented on code in PR #74735: URL: https://github.com/apache/airflow/pull/74735#discussion_r4283248644
########## providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py: ########## @@ -0,0 +1,230 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +AirflowDurability with pydantic-ai-harness's ``CodeMode``, through a real agent loop. + +``run_code`` is not a step of its own: a retry runs each snippet again, and the tool +calls the snippet makes replay from the journal one by one. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +from collections.abc import AsyncIterator +from typing import Any + +import pytest +from pydantic_ai import Agent, RunContext +from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart +from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel +from pydantic_ai.toolsets import FunctionToolset + +from airflow.providers.common.ai.durable import AirflowDurability +from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope + +from unit.common.ai.durable.test_capability import Calls, attempt, responses_so_far, tool_returns + +code_mode = pytest.importorskip("pydantic_ai_harness.code_mode") Review Comment: Yes, switched to a plain import in 9583d595a94. Every CI job has the harness. The CI image syncs the provider `dev` group, which now includes `pydantic-ai-harness==0.35.0`. The compat jobs only uninstall editable packages and the installed providers before installing the target Airflow. The lowest-deps job installs `--all-extras`. ########## providers/common/ai/pyproject.toml: ########## @@ -127,7 +127,9 @@ dependencies = [ # Provides the CodeMode capability for AgentOperator(capabilities=[CodeMode()]). # pydantic-ai-harness is pre-1.0; an opt-in extra so its churn never breaks the # base provider install. -"code-mode" = ["pydantic-ai-harness[codemode]>=0.24.0"] +# 0.35.0 is the first release that caps pydantic-monty below 1.0; earlier releases also +# accept monty 1.x, which they cannot run code with. +"code-mode" = ["pydantic-ai-harness[codemode]>=0.35.0"] Review Comment: Intentional. The base floor stays at what the provider itself needs, and only the `code-mode` extra needs 2.44. The harness's own requirement raises it when the extra is installed. The lowest-deps job shows it working: it resolves harness 0.35.0 with pydantic-ai-slim 2.44.0, and the suite passes. ########## providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py: ########## @@ -0,0 +1,230 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +AirflowDurability with pydantic-ai-harness's ``CodeMode``, through a real agent loop. + +``run_code`` is not a step of its own: a retry runs each snippet again, and the tool +calls the snippet makes replay from the journal one by one. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +from collections.abc import AsyncIterator +from typing import Any + +import pytest +from pydantic_ai import Agent, RunContext +from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart +from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel +from pydantic_ai.toolsets import FunctionToolset + +from airflow.providers.common.ai.durable import AirflowDurability +from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope + +from unit.common.ai.durable.test_capability import Calls, attempt, responses_so_far, tool_returns + +code_mode = pytest.importorskip("pydantic_ai_harness.code_mode") +CodeMode = code_mode.CodeMode + +GATHER = "import asyncio\na, b = await asyncio.gather(fetch(name='a'), fetch(name='b'))\n[a, b]" + + +def fetch_toolset(calls: Calls, started: dict[str, asyncio.Event] | None = None) -> FunctionToolset[Any]: + toolset = FunctionToolset(id="db") + + @toolset.tool_plain + async def fetch(name: str) -> str: + """Fetch a record.""" + calls.bump(f"fetch:{name}") + if started is not None and name in started: + started[name].set() + return f"{name}-data" + + return toolset + + +def run_code_then_answer(calls: Calls, snippets: list[str], *, fail_final: list[bool]): + """A model that sends each snippet to ``run_code`` in turn, then answers with the last result.""" + + def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + calls.bump("model") + if (sent := responses_so_far(messages)) < len(snippets): + return ModelResponse(parts=[ToolCallPart("run_code", {"code": snippets[sent]})]) + if fail_final[0]: + raise RuntimeError("worker died") + return ModelResponse(parts=[TextPart(f"answer: {tool_returns(messages)[-1].content}")]) + + return model_fn + + +def streamed(model_fn): + """The streaming form of ``model_fn``, for runs a capability streams (``eager``, ``speculate``).""" + + async def stream_fn(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str | DeltaToolCalls]: + for part in model_fn(messages, info).parts: + if isinstance(part, ToolCallPart): + yield {0: DeltaToolCall(name=part.tool_name, json_args=part.args_as_json_str())} + elif isinstance(part, TextPart): + yield part.content + + return stream_fn + + +class _HoldBack(AbstractCapability[Any]): + """Holds one tool call back until another has started, so ``asyncio.gather`` starts them in another order.""" + + def __init__(self, held_back: str, until_started: asyncio.Event) -> None: + self.held_back = held_back + self.until_started = until_started + + async def before_tool_execute(self, ctx: RunContext[Any], *, call, tool_def, args): + if args.get("name") == self.held_back: + await asyncio.wait_for(self.until_started.wait(), timeout=5) + return args + + +class TestReplay: + @pytest.mark.asyncio + async def test_a_retry_replays_the_steps_after_a_run_code_call(self, memory_storage): + """ + The snippet runs again and its nested calls replay, so the ``run_code`` return the + next model request sees differs only in the timestamps of those nested calls. + """ + calls = Calls() + fail_final = [True] + + def build() -> Agent[None, str]: + return Agent( + FunctionModel( + run_code_then_answer(calls, [GATHER, "await fetch(name='c')"], fail_final=fail_final) + ), + name="analyst", + toolsets=[fetch_toolset(calls)], + capabilities=[CodeMode(), AirflowDurability()], + ) + + with pytest.raises(RuntimeError, match="worker died"): + await attempt(memory_storage, build()) + fail_final[0] = False + with journal_scope(journal := DurableJournal(memory_storage)): + result = await build().run("go") + + assert result.output == "answer: c-data" + assert dict(journal.stats.replayed) == {"model": 2, "tool": 3} + assert dict(journal.stats.recorded) == {"model": 1} + assert (calls["fetch:a"], calls["fetch:b"], calls["fetch:c"]) == (1, 1, 1) + + @pytest.mark.asyncio + async def test_nested_calls_started_in_another_order_run_live(self, memory_storage): + """ + A nested call's fingerprint holds its arguments and its call id, which code mode + numbers in the order the snippet makes the calls. A call that reaches the journal + in another position never replays the result recorded for its neighbour. + """ + calls = Calls() + fail_final = [True] + + def build(*capabilities: AbstractCapability[Any], started=None) -> Agent[None, str]: + return Agent( + FunctionModel(run_code_then_answer(calls, [GATHER], fail_final=fail_final)), + name="analyst", + toolsets=[fetch_toolset(calls, started)], + capabilities=[CodeMode(), *capabilities, AirflowDurability()], + ) + + with pytest.raises(RuntimeError, match="worker died"): + await attempt(memory_storage, build()) + fail_final[0] = False + b_started = asyncio.Event() + with journal_scope(journal := DurableJournal(memory_storage)): + result = await build(_HoldBack("a", until_started=b_started), started={"b": b_started}).run("go") + + assert result.output == "answer: ['a-data', 'b-data']" + assert dict(journal.stats.replayed) == {"model": 1} + assert (calls["fetch:a"], calls["fetch:b"]) == (2, 2) + + +class TestEarlyExecution: + """CodeMode's ``eager`` and ``speculate`` run tool calls ahead of the snippet; not while durable.""" + + @pytest.mark.parametrize("durable", [False, True]) + @pytest.mark.asyncio + async def test_speculation_is_off_while_durable(self, memory_storage, durable): + calls = Calls() + speculating = CodeMode(speculate=["fetch"]) + model_fn = run_code_then_answer(calls, ["await fetch(name='a')"], fail_final=[False]) + agent = Agent( + FunctionModel(model_fn, stream_function=streamed(model_fn)), + name="analyst", + toolsets=[fetch_toolset(calls)], + capabilities=[speculating, AirflowDurability()], + ) + + if durable: + await attempt(memory_storage, agent) + else: + await agent.run("go") + + assert (speculating.speculation_stats.launched > 0) is not durable + assert calls["fetch:a"] == 1 + + @pytest.mark.parametrize("durable", [False, True]) + @pytest.mark.asyncio + async def test_eager_execution_is_off_while_durable(self, memory_storage, durable): + calls = Calls() + seen_mid_stream: list[int] = [] + a_started = asyncio.Event() + code = "a = await fetch(name='a')\nb = await fetch(name='b')\n[a, b]" + args = json.dumps({"code": code}) + # A statement counts as complete once the next one has streamed in full. + first_statement = args.index("\\n", args.index("\\n") + 2) + 2 + + async def stream_fn( + messages: list[ModelMessage], info: AgentInfo + ) -> AsyncIterator[str | DeltaToolCalls]: + if responses_so_far(messages): + yield "done" + return + yield {0: DeltaToolCall(name="run_code", json_args=args[:first_statement], tool_call_id="rc")} + # Give an eager pump time to run the first statement before the call completes. + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(a_started.wait(), timeout=1) Review Comment: Done in 9583d595a94. It waits 0.5 s when durable and up to 5 s otherwise, and the non-durable case returns as soon as the eager call starts. The durable case now takes about 0.5 s. ########## providers/common/ai/docs/durable_execution.rst: ########## @@ -193,6 +193,33 @@ In a plain ``@task``, which has no such moment, the capability deletes a run's j as soon as that run succeeds: after a failure, only the run that failed replays, and the agents that ran before it run again. +.. _durable-code-mode: + +Code mode +--------- + +With pydantic-ai-harness's ``CodeMode`` (:doc:`code_mode`), the model sends Python to one +``run_code`` tool, and the snippet calls the agent's tools as functions. ``run_code`` is +not a step of its own: the retry runs each snippet again, and the tool calls the snippet +makes replay from the journal one by one, so a snippet that called three tools is three +recorded steps. When the snippet returns what it returned before, the model request after +it replays too. + +A call the snippet makes replays only when the call at that position in the journal has +the same tool, arguments and call id. Code mode numbers a snippet's calls in the order the +snippet makes them, so when calls under ``asyncio.gather`` reach the journal in another +order on the retry, the first call out of place and every step after it run live. A call +never gets the result recorded for another one. + +The snippet's own code is not recorded. Whatever it does outside a tool call happens again Review Comment: Not inside a snippet, at least on harness 0.35.0. `CodeMode` turns each nested result into JSON-compatible data before the snippet sees it, on the live run too. I tried a nested tool returning `b"hi"`, `(1, 2)`, a datetime and a `Decimal`. The snippet saw `str 'hi'`, `list [1, 2]`, `str '2026-01-01T00:00:00Z'` and `str '1.5'`, the same live and on replay. So I left the docs as they are. Where did you see `b'hi'` on the live run? If it was a tool called outside code mode, that's the general replay path rather than something this PR changes. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
