dabla commented on code in PR #74150:
URL: https://github.com/apache/airflow/pull/74150#discussion_r4368275181
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -206,7 +208,27 @@ def get_hook(cls, conn_id: str, hook_params: dict | None =
None):
f"one of {expected_types}. Change the connection type to
'pydanticai', put the API key "
"in the API Key field (stored as the password) and the model
(e.g. 'openai:gpt-5') in Model."
)
- return conn.get_hook(hook_params=hook_params)
+ # The check above guarantees a ``PydanticAIHook``. Seeding saves the
second fetch of
+ # this same connection its first ``get_conn`` would make.
+ hook = conn.get_hook(hook_params=hook_params)
+ hook._seed_connection(conn)
+ return hook
+
+ @classmethod
+ async def aget_hook(cls, conn_id: str, hook_params: dict | None = None):
+ """
+ Return the hook for ``conn_id``, built with ``hook_params``, for use
inside an async task.
+
+ The async counterpart of :meth:`get_hook`, with the same connection
type check and
+ the connection already cached on the hook. The lookup runs in a worker
thread:
+ ``get_connection`` sends to the supervisor synchronously, which raises
+ ``DeadlockImminentError`` on the event loop while another async call
is in flight,
+ but only waits its turn from any other thread.
+
+ :raises ValueError: if the connection's type is registered to a hook
that is not a
+ ``PydanticAIHook``.
+ """
+ return await asyncio.to_thread(cls.get_hook, conn_id, hook_params)
Review Comment:
Reproduced the mechanism in the source: `asend()` takes `_thread_lock`
through `run_in_executor(None, ...)` and reads with `asyncio.to_thread`, both
on the default executor, so hook lookups filling that pool left nothing to run
the read. The three methods now run on a single-worker executor of their own
(`_run_off_loop`, e72b7d64db), with the caller's context variables copied so
task logging keeps its context; one worker also serializes the hook's sends on
3.0/3.1. Tests: each hop completes while the default executor's only worker is
parked, and two gathered `aget_conn` calls never send at the same time.
`AirflowToolset.run_blocking`, `strands_storage`, `managed_agent` and
`sandbox` use `asyncio.to_thread` too and share the exposure; I can move
`run_blocking` to the same executor in a follow-up if you agree.
##########
providers/common/ai/docs/hooks/pydantic_ai.rst:
##########
@@ -97,3 +97,22 @@ The spec file's ``model`` is the last resort: ``model_id``
and the connection's
``model`` extra both take precedence over it, in the
:ref:`order the connection page lists <pydanticai-model-resolution>`. Passing
``instructions`` to ``create_agent`` when a ``spec_file`` is also given
appends additional instructions to the file value.
+
+Using the hook in an async task
+-------------------------------
+
+In an ``async def`` task, use the async counterparts ``aget_hook``,
``aget_conn``
+and ``acreate_agent``, and run the agent with ``await agent.run(...)``. They
take the
+same parameters as ``get_hook``, ``get_conn`` and ``create_agent`` and return
the same
+objects, fallback chain included. Connection lookups make synchronous calls to
the Task
+SDK, which raise ``DeadlockImminentError`` on the event loop while another
async call
+is in flight, so the async methods run them in a worker thread.
Review Comment:
Done: the section says Airflow 3.2.1 and later, and the sentence on the loop
now says the call blocks, raising `DeadlockImminentError` from 3.3.1
(e72b7d64db). The `aget_hook` and `aget_conn` docstrings use the same wording.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -2207,3 +2216,249 @@ def test_hook_resolves_from_environment_variable(self,
hook_class, serializer, m
assert resolved.conn_type == hook_class.conn_type
assert type(resolved.get_hook()) is hook_class
+
+
+def _off_the_event_loop(lookup):
+ """
+ Wrap a synchronous Task SDK lookup so it fails the test when it runs on
the event loop.
+
+ ``get_connection`` and the secret masking of ``extra_dejson`` send to the
supervisor
+ synchronously. On the event loop, with another async call in flight, that
raises
+ ``DeadlockImminentError``; from a worker thread it only waits its turn.
+ """
+
+ def guarded(*args, **kwargs):
+ try:
+ asyncio.get_running_loop()
+ except RuntimeError:
+ return lookup(*args, **kwargs)
+ raise AssertionError("synchronous Task SDK call on the event loop")
+
+ return guarded
+
+
[email protected]
+def loop_guarded_registry():
+ """The ``registry`` fixture, with every synchronous lookup refused on the
event loop."""
+ reg = _ConnRegistry()
+ with (
+ patch.object(
+ PydanticAIHook, "get_connection",
side_effect=_off_the_event_loop(reg.get_connection)
+ ) as mock_get_connection,
+ patch.object(
+ Connection,
+ "extra_dejson",
+ new=property(_off_the_event_loop(lambda conn:
json.loads(conn.extra) if conn.extra else {})),
+ ),
+ patch.object(Connection, "get_hook", side_effect=reg.get_hook,
autospec=True),
+ patch(
+
"airflow.providers.common.ai.hooks.pydantic_ai._get_registered_hooks",
+ autospec=True,
+ return_value=reg.hooks,
+ ),
+ ):
+ reg.get_connection_mock = mock_get_connection
+ yield reg
+
+
+class _ReadsAnotherConnectionHook(PydanticAIHook):
+ """A subclass whose provider kwargs come from another connection, as
Snowflake's do."""
+
+ def _get_provider_kwargs(self, api_key, base_url, extra):
+ return {"api_key":
self.get_connection("credentials").extra_dejson["token"]}
+
+
+class TestPydanticAIHookAsync:
+ @pytest.mark.asyncio
+ async def test_aget_conn_resolves_the_model_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+
+ model = await PydanticAIHook(llm_conn_id="primary").aget_conn()
+
+ assert model is infer_model_stub.models["openai:gpt-5.6-sol"]
+
loop_guarded_registry.get_connection_mock.assert_called_once_with("primary")
+
+ @pytest.mark.asyncio
+ async def test_aget_conn_resolves_the_fallback_chain_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add(
+ "primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second", "third"]}
+ )
+ loop_guarded_registry.add("second", extra={"model":
"anthropic:claude-opus-4-6"})
+ loop_guarded_registry.add("third", extra={"model": "groq:llama-4"})
+
+ model = await PydanticAIHook(llm_conn_id="primary").aget_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models == [
+ infer_model_stub.models["openai:gpt-5.6-sol"],
+ infer_model_stub.models["anthropic:claude-opus-4-6"],
+ infer_model_stub.models["groq:llama-4"],
+ ]
+ assert loop_guarded_registry.get_connection_mock.call_args_list == [
+ call("primary"),
+ call("second"),
+ call("third"),
+ ]
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("subclass_conn_id", ["primary", "second"],
ids=["primary", "fallback"])
+ async def
test_aget_conn_runs_a_subclass_provider_kwargs_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub, subclass_conn_id
+ ):
+ loop_guarded_registry.add("credentials", extra={"token":
"from-another-connection"})
+ loop_guarded_registry.add(
+ "primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second"]}
+ )
+ loop_guarded_registry.add(
+ "second", hook_class=_ReadsAnotherConnectionHook, extra={"model":
"anthropic:claude-opus-4-6"}
+ )
+ hook_class = _ReadsAnotherConnectionHook if subclass_conn_id ==
"primary" else PydanticAIHook
+
+ await hook_class(llm_conn_id="primary").aget_conn()
+
+
loop_guarded_registry.get_connection_mock.assert_any_call("credentials")
+
+ @pytest.mark.asyncio
+ async def test_aget_conn_is_cached(self, loop_guarded_registry,
infer_model_stub):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ assert await hook.aget_conn() is await hook.aget_conn()
+ assert hook.get_conn() is infer_model_stub.models["openai:gpt-5.6-sol"]
+ loop_guarded_registry.get_connection_mock.assert_called_once()
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def test_acreate_agent_forwards_every_argument(
+ self, mock_agent_cls, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+
+ agent = await PydanticAIHook(llm_conn_id="primary").acreate_agent(
+ int, instructions="Count.", retries=3
+ )
+
+ assert agent is mock_agent_cls.return_value
+ mock_agent_cls.assert_called_once_with(
+ infer_model_stub.models["openai:gpt-5.6-sol"],
+ output_type=int,
+ instructions="Count.",
+ retries=3,
+ capabilities=[MaskingCapability()],
+ )
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent")
Review Comment:
Added (e72b7d64db).
--
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]
