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 c245439c5aa Add an example of a decision model approving an agent's
tool calls (#74697)
c245439c5aa is described below
commit c245439c5aa978870372863fd27823265918bfab
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 14 09:41:24 2026 +0100
Add an example of a decision model approving an agent's tool calls (#74697)
A HandleDeferredToolCalls handler asks a decision model whether each call
that needs approval may run. A confident allow runs it, a confident deny
refuses it, and anything else, including a judge that cannot be reached,
goes to AgentOperator's existing per-tool approval pause.
---
providers/common/ai/docs/decision_models.rst | 6 +
providers/common/ai/docs/tool_approval.rst | 70 ++++++
.../ai/example_dags/example_agent_tool_judge.py | 188 +++++++++++++++++
.../common/ai/operators/test_agent_tool_judge.py | 234 +++++++++++++++++++++
4 files changed, 498 insertions(+)
diff --git a/providers/common/ai/docs/decision_models.rst
b/providers/common/ai/docs/decision_models.rst
index 3f32da16aa2..337b306d1a6 100644
--- a/providers/common/ai/docs/decision_models.rst
+++ b/providers/common/ai/docs/decision_models.rst
@@ -194,6 +194,12 @@ Where it fits in this provider
taking arguments raises ``ToolCallProposed`` after the request, which
is a
``ModelAPIError`` rather than a refusal, so ``FallbackModel`` hands
those requests to
a text model behind it and only they cost a full call.
+ * - Approving an agent's tool calls
+ - Yes
+ - Whether a proposed call may run without a person is a three-way pick:
allow,
+ deny, or ask. A ``HandleDeferredToolCalls`` handler asks the decision
model for
+ each call that needs approval and sends an unsure answer to the task's
existing
+ approval pause; see :ref:`howto/tool_approval_judge`.
Reading the confidence
----------------------
diff --git a/providers/common/ai/docs/tool_approval.rst
b/providers/common/ai/docs/tool_approval.rst
index 744e9f36b1d..9baa0662483 100644
--- a/providers/common/ai/docs/tool_approval.rst
+++ b/providers/common/ai/docs/tool_approval.rst
@@ -130,6 +130,76 @@ The arguments shown to the reviewer pass through Airflow's
secrets masker first,
so a value under a key such as ``api_key`` or ``password``, or a value already
registered as a secret, is shown masked.
+.. _howto/tool_approval_judge:
+
+Let a decision model approve the routine calls
+----------------------------------------------
+
+Most refunds are routine, and a person who approves every one of them stops
reading
+them. A :doc:`decision model <decision_models>` can answer "may this call run
without
+a person?" for each one, and pydantic-ai's ``HandleDeferredToolCalls``
capability is
+the place to ask: its handler sees every call waiting for approval, with its
+validated arguments, before the task pauses. It approves the calls the model
is sure
+about, denies the ones the model is sure should not run, and leaves the rest.
The
+task pauses on whatever it leaves, exactly as above.
+
+A limit you can write down, such as "refunds over $50", needs no model: the
+``.approval_required()`` function above can check the arguments itself. The
model
+earns its call on the question a rule cannot answer, whether the call fits the
+request: the right order, an amount the customer's problem supports, an action
+somebody asked for.
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_agent_tool_judge.py
+ :language: python
+ :start-after: [START howto_agent_tool_judge]
+ :end-before: [END howto_agent_tool_judge]
+
+What happens to each call:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 30 35 35
+
+ * - The model answers
+ - Confidence at or above ``MIN_CONFIDENCE``
+ - Below it, or none reported
+ * - ``allow``
+ - The call runs.
+ - A person decides.
+ * - ``deny``
+ - The call does not run; the agent is told, as after a Reject.
+ - A person decides.
+ * - ``review``
+ - A person decides.
+ - A person decides.
+ * - The model call raises
+ - A person decides.
+ - A person decides.
+
+Once a person approves, the call runs without the model being asked again.
Once a
+call to a tool has been refused, by a person or by the model, later calls to
that
+tool in the same run go to a person: the model does not see why the earlier
call was
+refused, so it does not get to overrule that.
+
+**Keep hard limits in the tool.** A decision model reads a tool argument as
data, not
+as an attempt to manipulate it, so an argument written to argue for its own
approval
+can move the answer (see :doc:`decision_models`). ``within_order_total`` is the
+tool's ``args_validator``: it runs in the worker after the arguments validate
and
+before the call is marked for approval, so a refund above the order total goes
back
+to the agent without the model or a person being asked about it.
+
+**Set the bar from your own calls.** Run the handler over calls your agent
actually
+makes, and raise ``MIN_CONFIDENCE`` until the calls it lets through are ones
you
+would have approved. A model that hesitates sends more calls to a person, which
+costs time; one that is trusted too far costs whatever the tool does.
+
+The Required Actions page shows the call, not the model's answer; the handler
logs
+the answer and its confidence to the task log, where the reviewer can look. A
call the
+handler leaves for a person follows every rule on this page, including how
often one
+task instance can ask in a Dag run. Where tool approval is not available, such
a call
+fails the task. An exception the handler does not catch, such as a missing
+``decision_default`` connection, fails the task as well.
+
Requirements and limits
-----------------------
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_tool_judge.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_tool_judge.py
new file mode 100644
index 00000000000..bcaa10226ef
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_tool_judge.py
@@ -0,0 +1,188 @@
+# 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.
+"""
+Example Dag: a decision model decides the agent's refunds, and a person
decides the ones it is unsure of.
+
+Prerequisites: Airflow 3.3 or later, ``pydantic-ai-slim`` 2.46.0 or later, a
+``pydanticai_default`` connection for the agent's text model, and a
``decision_default``
+connection whose model is a decision model, such as TypeSafe Jev:
``password=<API key>``,
+``extra='{"model": "typesafe:jev-1.13.0"}'``, and the provider's ``typesafe``
extra.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import timedelta
+from enum import Enum
+
+from pydantic_ai import (
+ DeferredToolRequests,
+ DeferredToolResults,
+ ModelRetry,
+ RunContext,
+ Tool,
+ ToolDenied,
+ UseEnumMemberDocstrings,
+)
+from pydantic_ai.capabilities import HandleDeferredToolCalls
+from pydantic_ai.messages import ToolReturnPart, UserPromptPart
+from pydantic_ai.toolsets.function import FunctionToolset
+
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.operators.agent import AgentOperator
+from airflow.providers.common.ai.utils.decision import ModelConfidence
+from airflow.providers.common.compat.sdk import dag
+
+log = logging.getLogger(__name__)
+
+
+# [START howto_agent_tool_judge]
+ORDERS = {7: 42.00}
+
+
+def lookup_order(order_id: int) -> str:
+ """Look up an order."""
+ return f"order {order_id}: ${ORDERS[order_id]:.2f}, carrier marked it lost"
+
+
+def refund_order(order_id: int, amount: float, reason: str) -> str:
+ """Refund an order. Irreversible: the money leaves the account."""
+ return f"refunded ${amount:.2f} on order {order_id}"
+
+
+def within_order_total(ctx: RunContext, order_id: int, amount: float, reason:
str) -> None:
+ """A rule, not a judgment: a refund it refuses never reaches the model or
a person."""
+ if order_id not in ORDERS:
+ raise ModelRetry(f"There is no order {order_id}.")
+ if amount > ORDERS[order_id]:
+ raise ModelRetry(f"Order {order_id} cost ${ORDERS[order_id]:.2f};
refund at most that.")
+
+
+shop = FunctionToolset(
+ tools=[
+ lookup_order,
+ Tool(refund_order, requires_approval=True,
args_validator=within_order_total),
+ ]
+)
+TOOL_DESCRIPTIONS = {tool.__name__: tool.__doc__ for tool in (lookup_order,
refund_order)}
+
+
+class Verdict(UseEnumMemberDocstrings, str, Enum):
+ """Whether a proposed tool call may run without a person checking it
first."""
+
+ allow = "allow"
+ """The call does what the request asks, on the right record, at a size the
request supports."""
+
+ review = "review"
+ """The call may be right, but a person should check it: an unusual amount,
an irreversible step the request does not clearly ask for, or context that is
missing."""
+
+ deny = "deny"
+ """The call does not match the request: the wrong record, an amount far
beyond what was asked, or an action nobody asked for."""
+
+
+# Below this, the answer is not acted on and the call goes to a person.
+MIN_CONFIDENCE = 0.7
+
+
+def customer_request(ctx: RunContext) -> str:
+ """The request the agent is working on. A run resumed after a pause has no
prompt, only history."""
+ if isinstance(ctx.prompt, str):
+ return ctx.prompt
+ prompts = [
+ part.content
+ for message in ctx.messages
+ for part in message.parts
+ if isinstance(part, UserPromptPart) and isinstance(part.content, str)
+ ]
+ return prompts[-1] if prompts else ""
+
+
+def refused_earlier(ctx: RunContext, tool_name: str) -> bool:
+ """Whether a person or the model refused a call to this tool earlier in
the run."""
+ return any(
+ isinstance(part, ToolReturnPart) and part.tool_name == tool_name and
part.outcome == "denied"
+ for message in ctx.messages
+ for part in message.parts
+ )
+
+
+async def judge_tool_calls(ctx: RunContext, requests: DeferredToolRequests) ->
DeferredToolResults:
+ """Approve or deny the calls the decision model is sure about; leave the
rest for a person."""
+ judge = PydanticAIHook(llm_conn_id="decision_default").create_agent(
+ output_type=Verdict,
+ instructions=(
+ "A support agent proposed the tool call below while handling the
customer request. "
+ "Decide whether it may run without a person checking it first."
+ ),
+ )
+ results = DeferredToolResults()
+ for call in requests.approvals:
+ if refused_earlier(ctx, call.tool_name):
+ # The model does not see why the earlier call was refused, so it
must not overrule that.
+ log.info(
+ "%s (%s) was refused earlier in this run; asking a person",
call.tool_name, call.tool_call_id
+ )
+ continue
+ question = (
+ f"Customer request:\n{customer_request(ctx)}\n\n"
+ f"Proposed tool call: {call.tool_name}\n"
+ f"What the tool does: {TOOL_DESCRIPTIONS.get(call.tool_name,
'unknown')}\n"
+ f"Arguments: {json.dumps(call.args_as_dict())}"
+ )
+ try:
+ answer = await judge.run(question)
+ except Exception:
+ # An unreachable judge is not a verdict: leave the call for a
person.
+ log.warning("Judge failed on %s; asking a person", call.tool_name,
exc_info=True)
+ continue
+ confidence =
ModelConfidence.from_result(answer).confidence.get("response")
+ acted_on = confidence is not None and confidence >= MIN_CONFIDENCE
+ # No arguments here: they can carry secrets, which the review page
masks and a log line would not.
+ log.info(
+ "Judge on %s (%s): %s at confidence %s, bar %s; %s",
+ call.tool_name,
+ call.tool_call_id,
+ answer.output.value,
+ confidence,
+ MIN_CONFIDENCE,
+ "acting on it" if acted_on and answer.output is not Verdict.review
else "asking a person",
+ )
+ if not acted_on:
+ continue
+ if answer.output is Verdict.allow:
+ results.approvals[call.tool_call_id] = True
+ elif answer.output is Verdict.deny:
+ results.approvals[call.tool_call_id] = ToolDenied("An automated
check refused this call.")
+ return results
+
+
+@dag(tags=["example"])
+def example_agent_tool_judge():
+ AgentOperator(
+ task_id="handle_ticket",
+ llm_conn_id="pydanticai_default",
+ prompt="The customer says order 7 never arrived. Resolve it.",
+ toolsets=[shop],
+ capabilities=[HandleDeferredToolCalls(handler=judge_tool_calls)],
+ tool_approval_timeout=timedelta(hours=4),
+ )
+
+
+# [END howto_agent_tool_judge]
+
+example_agent_tool_judge()
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_agent_tool_judge.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent_tool_judge.py
new file mode 100644
index 00000000000..3f9596a3a0d
--- /dev/null
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_agent_tool_judge.py
@@ -0,0 +1,234 @@
+# 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.
+"""``example_agent_tool_judge``: a decision model judging the calls that need
approval in AgentOperator."""
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+# The example's Verdict uses UseEnumMemberDocstrings, new in pydantic-ai-slim
2.46.0.
+pytest.importorskip("pydantic_ai", minversion="2.46.0")
+
+from pydantic_ai import Agent
+from pydantic_ai.messages import (
+ ModelResponse,
+ RetryPromptPart,
+ TextPart,
+ ToolCallPart,
+ ToolReturnPart,
+ UserPromptPart,
+)
+from pydantic_ai.models.function import AgentInfo, FunctionModel
+
+from airflow.providers.common.ai.example_dags import example_agent_tool_judge
as example
+from airflow.providers.common.ai.operators.agent import AgentOperator
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
+from unit.common.ai.operators.test_agent_tool_approval import (
+ APPROVE,
+ UPSERT,
+ _context,
+ _FakeTaskStateStore,
+ _returns,
+)
+
+pytestmark = pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Per-tool
approval needs Airflow 3.3+.")
+
+if AIRFLOW_V_3_3_PLUS:
+ from airflow.sdk.exceptions import TaskAwaitingInput
+
+JUDGE_HOOK =
"airflow.providers.common.ai.example_dags.example_agent_tool_judge.PydanticAIHook"
+REQUEST = "The customer says order 7 never arrived. Resolve it."
+
+
+def _refunds_of(*amounts: float):
+ """An agent that proposes one refund of order 7 per step, in ``amounts``
order, then reports the results."""
+
+ def model(messages, info: AgentInfo) -> ModelResponse:
+ answered = [p for m in messages for p in m.parts if isinstance(p,
(ToolReturnPart, RetryPromptPart))]
+ if len(answered) < len(amounts):
+ step = len(answered)
+ args = {"order_id": 7, "amount": amounts[step], "reason": "lost
parcel"}
+ return ModelResponse(parts=[ToolCallPart("refund_order", args,
tool_call_id=f"c-{step}")])
+ return ModelResponse(parts=[TextPart(" | ".join(_returns(messages)))])
+
+ return model
+
+
+def _judge(answers: list[tuple[str, float | None]], questions: list[str]):
+ """A decision model giving ``answers`` in order, as (verdict, confidence),
reported the way Jev does."""
+ pending = list(answers)
+
+ def model(messages, info: AgentInfo) -> ModelResponse:
+ questions.extend(str(p.content) for m in messages for p in m.parts if
isinstance(p, UserPromptPart))
+ verdict, confidence = pending.pop(0)
+ return ModelResponse(
+ parts=[ToolCallPart(info.output_tools[0].name, {"response":
verdict})],
+ provider_details=None if confidence is None else {"confidence":
{"response": confidence}},
+ )
+
+ return model
+
+
+def _unreachable(messages, info: AgentInfo) -> ModelResponse:
+ raise ConnectionError("decision model unreachable")
+
+
+def _operator(agent_model) -> AgentOperator:
+ """The example Dag's task, with the agent's text model replaced."""
+ op = example.example_agent_tool_judge().get_task("handle_ticket")
+ assert isinstance(op, AgentOperator)
+ hook = MagicMock(spec=["create_agent"])
+ hook.create_agent.side_effect = lambda **kw:
Agent(FunctionModel(agent_model), **kw)
+ op.llm_hook = hook
+ return op
+
+
[email protected]
+def judge():
+ """Stand in for the ``decision_default`` connection; set ``judge.model``
to a FunctionModel callable."""
+ with patch(JUDGE_HOOK, autospec=True) as hook_cls:
+ state = MagicMock(spec=["model"])
+ hook_cls.return_value.create_agent.side_effect = lambda **kw:
Agent(FunctionModel(state.model), **kw)
+ yield state
+
+
[email protected]
+def refunds():
+ """The arguments of every refund whose body ran."""
+ calls: list[dict[str, Any]] = []
+ schema = example.shop.tools["refund_order"].function_schema
+
+ @functools.wraps(example.refund_order)
+ def recording_refund(**kwargs):
+ calls.append(kwargs)
+ return example.refund_order(**kwargs)
+
+ with patch.object(schema, "function", recording_refund):
+ yield calls
+
+
+def _execute(op: AgentOperator, store: _FakeTaskStateStore | None = None):
+ with patch(UPSERT, autospec=True) as upsert:
+ output = op.execute(_context(store or _FakeTaskStateStore()))
+ upsert.assert_not_called()
+ return output
+
+
+def _pause(op: AgentOperator, store: _FakeTaskStateStore) -> TaskAwaitingInput:
+ with patch(UPSERT, autospec=True) as upsert,
pytest.raises(TaskAwaitingInput) as paused:
+ op.execute(_context(store))
+ assert "**refund_order**" in upsert.call_args.kwargs["body"]
+ return paused.value
+
+
+def test_a_confident_allow_runs_the_refund_without_asking_a_person(judge,
refunds):
+ questions: list[str] = []
+ judge.model = _judge([("allow", 0.95)], questions)
+
+ output = _execute(_operator(_refunds_of(42.0)))
+
+ assert refunds == [{"order_id": 7, "amount": 42.0, "reason": "lost
parcel"}]
+ assert output == "refunded $42.00 on order 7"
+ assert REQUEST in questions[0]
+ assert '"amount": 42.0' in questions[0]
+ assert "Irreversible" in questions[0]
+
+
+def test_a_confident_deny_never_runs_the_refund(judge, refunds):
+ judge.model = _judge([("deny", 0.95)], [])
+
+ output = _execute(_operator(_refunds_of(42.0)))
+
+ assert refunds == []
+ assert output == "An automated check refused this call."
+
+
[email protected](
+ "answer",
+ [
+ pytest.param(("allow", 0.5), id="unsure-allow"),
+ pytest.param(("deny", 0.5), id="unsure-deny"),
+ pytest.param(("review", 0.99), id="review"),
+ pytest.param(("allow", None), id="no-confidence-reported"),
+ pytest.param(None, id="judge-unreachable"),
+ ],
+)
+def test_an_unsure_missing_or_failed_answer_asks_a_person(judge, refunds,
answer):
+ judge.model = _unreachable if answer is None else _judge([answer], [])
+
+ paused = _pause(_operator(_refunds_of(42.0)), _FakeTaskStateStore())
+
+ assert refunds == []
+ assert paused.kwargs["tool_call_ids"] == ["c-0"]
+
+
+def test_the_judge_is_not_asked_again_about_a_call_a_person_approved(judge,
refunds):
+ questions: list[str] = []
+ judge.model = _judge([("review", 0.99)], questions)
+ store = _FakeTaskStateStore()
+ paused = _pause(_operator(_refunds_of(42.0)), store)
+
+ output = _operator(_refunds_of(42.0)).resume_after_tool_approval(
+ _context(store), **paused.kwargs, event=APPROVE
+ )
+
+ assert len(questions) == 1
+ assert refunds == [{"order_id": 7, "amount": 42.0, "reason": "lost
parcel"}]
+ assert output == "refunded $42.00 on order 7"
+
+
+def test_a_resumed_run_still_shows_the_judge_the_customer_request(judge,
refunds):
+ questions: list[str] = []
+ judge.model = _judge([("review", 0.99), ("allow", 0.95)], questions)
+ store = _FakeTaskStateStore()
+ paused = _pause(_operator(_refunds_of(42.0, 5.0)), store)
+
+ _operator(_refunds_of(42.0, 5.0)).resume_after_tool_approval(
+ _context(store), **paused.kwargs, event=APPROVE
+ )
+
+ assert [r["amount"] for r in refunds] == [42.0, 5.0]
+ assert REQUEST in questions[1]
+ assert '"amount": 5.0' in questions[1]
+
+
+def
test_a_refund_above_the_order_total_reaches_neither_the_judge_nor_a_person(judge,
refunds):
+ questions: list[str] = []
+ judge.model = _judge([("allow", 0.95)], questions)
+
+ output = _execute(_operator(_refunds_of(4200.0, 42.0)))
+
+ assert refunds == [{"order_id": 7, "amount": 42.0, "reason": "lost
parcel"}]
+ assert len(questions) == 1
+ assert '"amount": 42.0' in questions[0]
+ assert output == "refunded $42.00 on order 7"
+
+
+def test_a_tool_refused_earlier_in_the_run_goes_to_a_person(judge, refunds):
+ questions: list[str] = []
+ judge.model = _judge([("deny", 0.95), ("allow", 0.99)], questions)
+
+ paused = _pause(_operator(_refunds_of(42.0, 30.0)), _FakeTaskStateStore())
+
+ assert refunds == []
+ assert len(questions) == 1
+ assert paused.kwargs["tool_call_ids"] == ["c-1"]
