This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch durable-toolsets-without-id in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 91ea6817d9d365f5f349f5d468cecf36f6086ede Author: Kaxil Naik <[email protected]> AuthorDate: Wed Nov 11 06:31:44 2026 +0000 Name the pydantic-ai-slim release the system-one decision model prefix needs Rename the durable tool step-mode helpers to say what they return (_is_tool_journaled, _resolve_tool_step_mode) and the test toolset builders to _build_dynamic_toolset and _build_ops_toolset. --- providers/common/ai/docs/decision_models.rst | 4 ++-- .../airflow/providers/common/ai/durable/capability.py | 16 +++++++++++----- .../ai/tests/unit/common/ai/durable/test_capability.py | 16 ++++++++++------ 3 files changed, 23 insertions(+), 13 deletions(-) diff --git a/providers/common/ai/docs/decision_models.rst b/providers/common/ai/docs/decision_models.rst index 88ad02f6623..95e8da5039e 100644 --- a/providers/common/ai/docs/decision_models.rst +++ b/providers/common/ai/docs/decision_models.rst @@ -54,8 +54,8 @@ model, so a model id is the whole integration. it. `pydantic-ai's System One page <https://pydantic.dev/docs/ai/models/system-one/>`__ lists more. -``system-one:`` needs a newer ``pydantic-ai-slim`` than the provider's minimum, so check the -release you have installed. +``system-one:`` needs ``pydantic-ai-slim`` 2.53.0 or later, newer than the provider's minimum, +so check the release you have installed. Setup ----- 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 a6a8a3e3b68..e3c242e0102 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 @@ -105,7 +105,7 @@ _UNLOADABLE = "the recorded result no longer loads, so it was written by another _StepMode: TypeAlias = Literal["journal", "live"] -def _journals_tool(tool_def: ToolDefinition) -> bool: +def _is_tool_journaled(tool_def: ToolDefinition) -> bool: """Whether calls of a tool are journaled; ``metadata={"airflow": False}`` runs them live on every attempt.""" setting = (tool_def.metadata or {}).get(TOOL_CONFIG_KEY) if setting is None: @@ -118,16 +118,22 @@ def _journals_tool(tool_def: ToolDefinition) -> bool: ) -def _tool_step_mode(operation_id: DurableOperationId, tool: object | None, tool_name: str) -> _StepMode: +def _resolve_tool_step_mode( + operation_id: DurableOperationId, tool: object | None, tool_name: str +) -> _StepMode: # Never pydantic-ai's ``False``, which would take the call out of the backend: a live call still # takes its position, so the steps after it keep theirs, and a failure is still the step's own. - if isinstance(tool, ToolsetTool) and not _journals_tool(tool.tool_def): + if isinstance(tool, ToolsetTool) and not _is_tool_journaled(tool.tool_def): return "live" return "journal" _STEP_MODES: RoleBasedOperationConfig[_StepMode] = RoleBasedOperationConfig( - model="journal", event="journal", capability="journal", tool="journal", resolve_tool=_tool_step_mode + model="journal", + event="journal", + capability="journal", + tool="journal", + resolve_tool=_resolve_tool_step_mode, ) @@ -564,7 +570,7 @@ class _JournaledToolset(WrapperToolset[Any]): # A toolset whose calls act on a system Airflow cannot observe, such as a managed # agent, runs them again on every attempt, and so does a tool opted out by its metadata. replayable=(not isinstance(leaf, AirflowToolset) or leaf.replayable) - and _journals_tool(tool.tool_def), + and _is_tool_journaled(tool.tool_def), ) if step.replayed: if isinstance(leaf, AirflowToolset): 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 9f8db919b92..5d4284ee0e9 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 @@ -477,7 +477,7 @@ def _function_toolset(calls: Calls, result: Any) -> FunctionToolset[None]: return toolset -def _dynamic_toolset(calls: Calls, result: Any) -> DynamicToolset[None]: +def _build_dynamic_toolset(calls: Calls, result: Any) -> DynamicToolset[None]: """A dynamic toolset that builds :func:`_function_toolset` per run step.""" async def build(ctx: RunContext[None]) -> FunctionToolset[None]: @@ -896,7 +896,9 @@ class TestAirflowToolsets: assert calls["model"] == 3 -def _ops_toolset(calls: Calls, *, clock_metadata: dict[str, Any] | None = None, clock_changes: bool = False): +def _build_ops_toolset( + calls: Calls, *, clock_metadata: dict[str, Any] | None = None, clock_changes: bool = False +) -> FunctionToolset[None]: toolset = FunctionToolset(id="ops") @toolset.tool_plain(metadata=clock_metadata) @@ -966,7 +968,9 @@ class TestToolOptOut: return Agent( FunctionModel(clock_then_lookup(calls, fail_final=fail_final)), name="ops", - toolsets=[_ops_toolset(calls, clock_metadata=clock_metadata, clock_changes=clock_changes)], + toolsets=[ + _build_ops_toolset(calls, clock_metadata=clock_metadata, clock_changes=clock_changes) + ], capabilities=[AirflowDurability(), *extra_capabilities], ) @@ -992,12 +996,12 @@ class TestToolOptOut: opt_out = [SetToolMetadata(tools=["clock"], airflow=False)] async def ops(ctx: RunContext[None]) -> FunctionToolset[None]: - return _ops_toolset(calls) + return _build_ops_toolset(calls) return Agent( FunctionModel(clock_then_lookup(calls, fail_final=fail_final)), name="ops", - toolsets=[DynamicToolset(ops, id="dynamic_ops") if dynamic else _ops_toolset(calls)], + toolsets=[DynamicToolset(ops, id="dynamic_ops") if dynamic else _build_ops_toolset(calls)], capabilities=[ AirflowDurability(), *(opt_out if attempt_number == opt_out_on_attempt else []), @@ -1019,7 +1023,7 @@ class TestToolOptOut: ) @pytest.mark.parametrize( "build_toolset", - [_function_toolset, _WarehouseToolset, _dynamic_toolset], + [_function_toolset, _WarehouseToolset, _build_dynamic_toolset], ids=["function", "airflow", "dynamic"], ) async def test_set_tool_metadata_decides_whether_a_tool_replays(
