dabla commented on code in PR #74201:
URL: https://github.com/apache/airflow/pull/74201#discussion_r4552269372
##########
providers/common/ai/src/airflow/providers/common/ai/decorators/agent.py:
##########
@@ -106,6 +117,19 @@ def execute(self, context: Context) -> Any:
self.prompt = self.render_template(self.prompt, context)
return AgentOperator.execute(self, context)
+ async def aexecute(self, context: Context) -> Any:
+ context_merge(context, self.op_kwargs)
+ kwargs = determine_kwargs(self.python_callable, self.op_args, context)
+
+ self.prompt = await self.python_callable(*self.op_args, **kwargs)
+
+ validate_prompt(self.prompt, decorator_name="@task.agent")
+
+ # Rendering a template can read a Variable or a Connection with a
blocking call to
+ # the supervisor, which must not happen on the event loop.
+ await asyncio.to_thread(self.render_template_fields, context)
Review Comment:
You're right, and the merge of main made it worse: `execute` got #76244 and
`aexecute` kept the old line. Fixed in 1fdc7937f4: `aexecute` now merges
`call_arguments(...)` into the context and renders only the returned prompt,
with `asyncio.to_thread(self.render_template, self.prompt, context)`. Two tests
pin it, the twin of `test_execute_renders_only_the_returned_prompt` and one
that calls the async function with a positional and a keyword argument.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1085,16 +1134,69 @@ def execute(self, context: Context) -> Any:
# Without a history to inherit one from, pydantic-ai starts a
new conversation
# on every attempt; keep one for the durable run, for the same
reason.
run_kwargs["conversation_id"] = run_kwargs["run_id"]
+ return run_kwargs
+
+ def _reject_unsupported_on_async_path(self) -> None:
+ """Refuse the features that block the event loop: both do synchronous
I/O during the run."""
+ if self.durable or self.enable_hitl_review:
+ raise ValueError(
+ f"{type(self).__name__}: durable=True and
enable_hitl_review=True are not supported "
+ "with an async callable. Use a synchronous callable for them."
+ )
+
+ async def aexecute(self, context: Context) -> Any:
+ """
+ Run the agent on the worker's event loop.
+
+ The async version of :meth:`execute`: the agent run is awaited on the
event loop, and the
+ bookkeeping of a successful run (usage budget, run metadata, message
history) makes no
+ blocking call to the supervisor on it. What has no async counterpart
(the connection
+ lookup, the task state store before Airflow 3.4, the usage report of a
failed run, the
+ pause for a tool approval) runs the synchronous code in a worker
thread.
+ """
+ self._reject_unsupported_on_async_path()
+ usage_limits = coerce_usage_limits(self.usage_limits)
+
+ if self._supports_tool_approval() and (store :=
context.get("task_state_store")) is not None:
+ await self._adelete_approval_transcript(store)
+
+ ti = context["task_instance"]
+ self._durable_journal = None
+ self._usage_budget = self._build_usage_budget(context, usage_limits,
ti=ti)
+ self._run_usage = await self._usage_budget.aload() if
self._usage_budget else RunUsage()
+ run_usage: RunUsage = self._run_usage
+ self._run_usage_base = copy_run_usage(run_usage)
+
+ agent = await self._abuild_agent()
+ run_kwargs = self._run_kwargs(agent, ti, usage_limits)
- # A killed run raises RunCancelled (see run_agent_sync), which
propagates to fail the
- # task. The journal cleanup in _complete_run is skipped on the raise,
keeping it for the retry.
try:
- result, attempt_usage = self._run_and_report_on_failure(context,
agent, run_usage, run_kwargs)
- finally:
- # Also on a raise: the failed attempt is the one Airflow retries.
- if self._durable_journal is not None:
- self._durable_journal.log_summary(self.log)
- return self._complete_run(context, result, attempt_usage=attempt_usage)
+ result, attempt_usage = await self._arun_tracked(
+ agent, self.prompt, run_usage=run_usage, **run_kwargs
+ )
+ except BaseException:
+ await asyncio.to_thread(self._report_failed_run, context,
run_usage, run_id=run_kwargs["run_id"])
+ raise
+ return await self._acomplete_run(context, result,
attempt_usage=attempt_usage)
+
+ async def _acomplete_run(self, context: Context, result: Any, *,
attempt_usage: RunUsage) -> Any:
+ """Async version of :meth:`_complete_run`, without the features
:meth:`aexecute` refuses."""
+ log_run_summary(self.log, result, usage=attempt_usage)
+ if isinstance(result.output, DeferredToolRequests):
Review Comment:
Yes, same guard now, in 4889639dd8: `_acomplete_run` catches `BaseException`
around the pause, runs `_report_failed_run` in a worker thread unless the error
is `TaskAwaitingInput`, and re-raises. The new test ends a run on
`UnsupportedToolDeferralError` and checks the `run_id` and `usage` XComs.
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -4161,3 +4170,295 @@ def
test_a_toolset_that_masks_its_own_output_is_not_wrapped_again(self):
op._build_agent()
assert hook.create_agent.call_args.kwargs["toolsets"] == [sql]
+
+
+def _refused_on_the_event_loop(call):
+ """Wrap a blocking call so it fails the test when it runs on the event
loop."""
+
+ def guarded(*args, **kwargs):
+ try:
+ asyncio.get_running_loop()
+ except RuntimeError:
+ return call(*args, **kwargs)
+ raise AssertionError("blocking call to the supervisor on the event
loop")
+
+ return guarded
+
+
+def _make_async_task_state_store_accessor(*, async_methods: bool = True):
+ """
+ A task state store backed by ``accessor.data``, whose blocking methods
fail on the event loop.
+
+ Every write and delete, async or blocking, is recorded in
``accessor.ops``. With
+ ``async_methods=False`` the store has no async methods, as before Airflow
3.4, and the operator
+ has to call the blocking ones from a worker thread.
+ """
+ from airflow.sdk.execution_time.context import TaskStateStoreAccessor
+
+ accessor = MagicMock(spec=TaskStateStoreAccessor)
+ data: dict[str, Any] = {}
+ ops: list[tuple[str, str]] = []
+ accessor.data = data
+ accessor.ops = ops
+
+ def get(key, default=None):
+ return data.get(key, default)
+
+ def set_(key, value, retention=None):
+ ops.append(("set", key))
+ data[key] = value
+
+ def delete(key):
+ ops.append(("delete", key))
+ data.pop(key, None)
+
+ accessor.get.side_effect = _refused_on_the_event_loop(get)
+ accessor.set.side_effect = _refused_on_the_event_loop(set_)
+ accessor.delete.side_effect = _refused_on_the_event_loop(delete)
+ if async_methods:
+ accessor.aget = AsyncMock(side_effect=get)
+ accessor.aset = AsyncMock(side_effect=set_)
+ accessor.adelete = AsyncMock(side_effect=delete)
+ else:
+ del accessor.aget, accessor.aset, accessor.adelete
+ return accessor
+
+
+def _make_async_context(task_state_store=None):
+ """A context for ``aexecute``: XComs go through ``axcom_push`` on Airflow
3.3+, ``xcom_push`` on 3.2."""
+ ti = _make_ti()
+ ti.axcom_push = AsyncMock()
+ context = {"task_instance": ti}
+ if task_state_store is not None:
+ context["task_state_store"] = task_state_store
+ return context
+
+
+def _make_async_mock_agent(mock_hook_cls, **run):
+ """Wire a mock agent whose ``run`` is awaitable into the hook's async
path, and return it."""
+ mock_agent = MagicMock(spec=["run", "instrument"])
+ mock_agent.run = AsyncMock(**run)
+ mock_hook_cls.aget_hook = AsyncMock()
+ mock_hook_cls.aget_hook.return_value.acreate_agent =
AsyncMock(return_value=mock_agent)
+ return mock_agent
+
+
[email protected](not AIRFLOW_V_3_2_PLUS, reason="async tasks need Airflow
>= 3.2")
+class TestAgentOperatorAsync:
+ def test_operator_is_synchronous_unless_a_subclass_says_otherwise(self):
+ assert AgentOperator(task_id="test", prompt="run",
llm_conn_id="my_llm").is_async is False
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def test_aexecute_awaits_the_run_through_the_async_hook_path(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ mock_agent = _make_async_mock_agent(mock_hook_cls,
return_value=make_mock_run_result("done"))
+
+ op = AgentOperator(task_id="test", prompt="run", llm_conn_id="my_llm")
+ result = await op.aexecute(_make_async_context())
+
+ assert result == "done"
+ mock_hook_cls.aget_hook.assert_awaited_once_with(
+ "my_llm", hook_params={"model_id": None, "fallback_conn_ids": None}
+ )
+ mock_hook_cls.get_hook.assert_not_called()
+ mock_agent.run.assert_awaited_once_with(
+ "run", usage_limits=None, run_id="ti-1", cancellation_token=ANY,
usage=ANY
+ )
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task state store needs
Airflow >= 3.3")
+ @pytest.mark.parametrize("async_store", [True, False], ids=["async-store",
"store-before-3.4"])
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def
test_successful_run_makes_no_blocking_call_to_the_supervisor(self,
mock_hook_cls, async_store):
+ """A real run with a usage budget and a message history: no blocking
call reaches the event
+ loop, whether the task state store has async methods (Airflow 3.4+) or
not."""
+ mock_hook_cls.aget_hook = AsyncMock()
+ mock_hook_cls.aget_hook.return_value.acreate_agent = AsyncMock(
+ side_effect=lambda **kw:
Agent(FunctionModel(_build_priced_response), **kw)
+ )
+ store =
_make_async_task_state_store_accessor(async_methods=async_store)
+ context = _make_async_context(task_state_store=store)
+ context["task_instance"].xcom_push.side_effect =
AssertionError("blocking xcom_push")
+
+ op = AgentOperator(
+ task_id="test",
+ prompt="run",
+ llm_conn_id="my_llm",
+ usage_limits={"cost_limit": str(PRICED_COST * 2)},
+ message_history=[],
+ )
+ with (
+
patch("airflow.providers.common.ai.operators.agent.AIRFLOW_V_3_4_PLUS",
async_store),
+
patch("airflow.providers.common.ai.utils.usage_budget.AIRFLOW_V_3_4_PLUS",
async_store),
+ ):
+ result = await op.aexecute(context)
+
+ assert result == "the answer"
+ pushed = [c.kwargs["key"] for c in
context["task_instance"].axcom_push.await_args_list]
+ assert pushed == ["run_id", "usage", MODEL_NAME_XCOM_KEY,
"message_history"]
+ assert ("delete", _TOOL_APPROVAL_TRANSCRIPT_KEY) in store.ops
+ assert ("set", USAGE_BUDGET_KEY) in store.ops
+ assert store.ops[-1] == ("delete", USAGE_BUDGET_KEY)
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task state store needs
Airflow >= 3.3")
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def test_failed_run_saves_the_budget_and_reports_its_usage(self,
mock_hook_cls):
+ _make_async_mock_agent(mock_hook_cls, side_effect=RuntimeError("boom"))
+ store = _make_async_task_state_store_accessor()
+ context = _make_async_context(task_state_store=store)
+
+ op = AgentOperator(
+ task_id="test", prompt="run", llm_conn_id="my_llm",
usage_limits={"request_limit": 5}
+ )
+ with pytest.raises(RuntimeError, match="boom"):
+ await op.aexecute(context)
+
+ assert ("set", USAGE_BUDGET_KEY) in store.ops
+ # The failure report reuses the synchronous code, from a worker thread.
+ pushed = {c.kwargs["key"] for c in
context["task_instance"].xcom_push.call_args_list}
+ assert pushed == {"run_id", "usage"}
+
+ @pytest.mark.parametrize("feature", ["durable", "enable_hitl_review"])
+ @pytest.mark.asyncio
+ async def test_aexecute_rejects_features_that_block_the_event_loop(self,
feature):
+ op = AgentOperator(task_id="test", prompt="run", llm_conn_id="my_llm",
**{feature: True})
+
+ with pytest.raises(ValueError, match="not supported with an async
callable"):
+ await op.aexecute(_make_async_context())
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def
test_run_waiting_on_a_tool_approval_pauses_through_the_synchronous_path(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ class Paused(Exception):
+ pass
+
+ _make_async_mock_agent(mock_hook_cls,
return_value=make_mock_run_result(DeferredToolRequests()))
+ context = _make_async_context()
+
+ op = AgentOperator(task_id="test", prompt="run", llm_conn_id="my_llm")
+ with patch.object(
+ AgentOperator,
+ "_pause_for_tool_approval",
+ side_effect=_refused_on_the_event_loop(Mock(side_effect=Paused)),
+ ) as pause:
+ with pytest.raises(Paused):
+ await op.aexecute(context)
+
+ pause.assert_called_once()
+ context["task_instance"].axcom_push.assert_not_awaited()
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task state store needs
Airflow >= 3.3")
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def test_a_retry_starts_from_the_usage_of_earlier_tries(self,
mock_hook_cls, make_mock_run_result):
+ mock_agent = _make_async_mock_agent(mock_hook_cls,
return_value=make_mock_run_result("done"))
+ store = _make_async_task_state_store_accessor()
+ context = _make_async_context(task_state_store=store)
+ store.data[USAGE_BUDGET_KEY] = TaskStateStoreUsageBudget(
+ store, max_tries=context["task_instance"].max_tries
+ )._record(RunUsage(requests=3))
+
+ op = AgentOperator(
+ task_id="test", prompt="run", llm_conn_id="my_llm",
usage_limits={"request_limit": 5}
+ )
+ await op.aexecute(context)
+
+ assert mock_agent.run.await_args.kwargs["usage"].requests == 3
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task state store needs
Airflow >= 3.3")
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def test_the_approval_transcript_of_an_earlier_try_is_deleted(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ _make_async_mock_agent(mock_hook_cls,
return_value=make_mock_run_result("done"))
+ store = _make_async_task_state_store_accessor()
+ store.data[_TOOL_APPROVAL_TRANSCRIPT_KEY] = "[]"
+
+ op = AgentOperator(task_id="test", prompt="run", llm_conn_id="my_llm")
+ await op.aexecute(_make_async_context(task_state_store=store))
+
+ assert ("delete", _TOOL_APPROVAL_TRANSCRIPT_KEY) in store.ops
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def test_a_model_output_is_returned_as_a_dict(self, mock_hook_cls,
make_mock_run_result):
+ class Answer(BaseModel):
+ text: str
+
+ _make_async_mock_agent(mock_hook_cls,
return_value=make_mock_run_result(Answer(text="done")))
+
+ op = AgentOperator(
+ task_id="test", prompt="run", llm_conn_id="my_llm",
output_type=Answer, serialize_output=True
+ )
+
+ assert await op.aexecute(_make_async_context()) == {"text": "done"}
+
+ @requires_typed_xcom
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def test_a_model_the_worker_does_not_register_is_returned_as_json(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ """The async path converts the output the way ``execute`` does: a
model inside a list output is
+ dumped, since the worker registers no class for an output marker."""
+ _make_async_mock_agent(mock_hook_cls,
return_value=make_mock_run_result([Summary(text="ok")]))
+
+ op = AgentOperator(
+ task_id="test", prompt="run", llm_conn_id="my_llm",
output_type=ToolOutput(list[Summary])
+ )
+
+ assert await op.aexecute(_make_async_context()) == [{"text": "ok",
"score": 0.0}]
+
+ @pytest.mark.parametrize(
+ ("kwargs", "emits"),
+ [
+ pytest.param({"message_history": []}, True, id="seeded"),
+ pytest.param({"emit_message_history": True}, True, id="opt-in"),
+ pytest.param({"message_history": [], "emit_message_history":
False}, False, id="opt-out"),
+ ],
+ )
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ async def test_the_transcript_push_follows_emit_message_history(
+ self, mock_hook_cls, kwargs, emits, make_mock_run_result
+ ):
+ mock_agent = _make_async_mock_agent(mock_hook_cls,
return_value=make_mock_run_result("ok"))
+ mock_agent.run.return_value.all_messages.return_value =
_sample_history()
+ context = _make_async_context()
+
+ op = AgentOperator(task_id="test", prompt="run", llm_conn_id="my_llm",
**kwargs)
+ await op.aexecute(context)
+
+ pushes = {
+ c.kwargs["key"]: c.kwargs["value"] for c in
context["task_instance"].axcom_push.await_args_list
Review Comment:
Fixed in 9caace3c3a, both ways: the tests that read the pushes now collect
them from the mock the running version uses, and `_axcom_push` has its own test
with the 3.3 flag patched on and off. The off case also fails if `xcom_push`
runs on the event loop, so the fallback branch is covered on every version.
##########
providers/common/ai/docs/operators/agent.rst:
##########
@@ -83,6 +83,30 @@ the prompt string; all other parameters are passed to the
operator.
:start-after: [START howto_decorator_agent]
:end-before: [END howto_decorator_agent]
+.. _howto/operator:agent-async:
+
+Async callables
+^^^^^^^^^^^^^^^
+
+On Airflow 3.2+ the decorated function can be an ``async def``. It is awaited
for the prompt, and
+the agent run is then awaited on the task's event loop. The function can await
async hooks. What
+has no async counterpart, such as the connection lookup, runs in a worker
thread, so nothing blocks
+the event loop.
+
+.. code-block:: python
+
+ @task.agent(llm_conn_id="pydanticai_default", system_prompt="You are a
support analyst.")
+ async def summarize_ticket(ticket_id: str) -> str:
+ ticket = await fetch_ticket(ticket_id)
+ return f"Summarize this ticket: {ticket}"
+
+A single run is not faster this way, as its duration is the model's. Tools
that call blocking hooks,
+such as ``SQLToolset`` and ``HookToolset``, run in a worker thread with either
kind of function.
+
+``durable=True`` and ``enable_hitl_review=True`` do blocking I/O during the
run and are not
+supported with an ``async def`` function: the Dag fails to parse with a
``ValueError``. Use a regular
Review Comment:
Good catch. I went for the docs in 18c8c6148e: the page now says a mapped
task is only created when it runs, so it fails there with the same
`ValueError`. A parse-time check on the mapped path would have to live in the
Task SDK's `_TaskDecorator.expand()`, since the provider's operator is not
built at that point, and I'd rather not add a core change to this PR. Happy to
look at it as a follow-up if you think it is worth it.
--
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]
