pankajastro commented on code in PR #74735:
URL: https://github.com/apache/airflow/pull/74735#discussion_r4282523256
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -1474,19 +1450,20 @@ def
test_code_mode_capability_turns_off_tool_approval(self, fake_harness, capabi
def test_harness_without_the_code_mode_extra_does_not_break_other_agents(
self, harness_without_code_mode_extra
):
- # Neither the durable check at construction nor the approval check at
run time may
- # touch the harness root, whose CodeMode export would raise
ImportError here.
- AgentOperator(task_id="t", prompt="p", llm_conn_id="llm",
durable=True, capabilities=[Thinking()])
+ # The approval check at run time must not touch the harness root,
whose CodeMode
+ # export would raise ImportError here.
op = AgentOperator(task_id="u", prompt="p", llm_conn_id="llm",
capabilities=[Thinking()])
assert op._supports_tool_approval() is True
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Per-tool approval
needs Airflow >= 3.3")
def test_capability_function_is_left_for_the_run_to_resolve(self,
fake_harness):
Review Comment:
This lost `durable=True` and is now skipped below Airflow 3.3, so it asserts
nothing there. Was that deliberate? Could the `op.capabilities == [build]`
check stay unskipped and only the `_supports_tool_approval()` assert be gated?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
`dev` now installs the harness, so should this be a plain import? A broken
install would skip these tests and CI would stay green.
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
In the durable case nothing sets `a_started`, so the test always waits the
full 1 s. Could it use a shorter timeout when `durable`?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
Replayed results are JSON-decoded, so `b"hi"` is `b'hi'` live but `'hi'` on
replay. Tuple, datetime and Decimal were identical. Is that worth naming here?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
Harness 0.35.0 requires `pydantic-ai-slim>=2.44.0`, but the base floor is
`>=2.36.0`. Should the base floor move, or is the gap intentional?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
--
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]
