kaxil commented on code in PR #74150:
URL: https://github.com/apache/airflow/pull/74150#discussion_r4289946839
##########
uv.lock:
##########
@@ -1,5 +1,5 @@
version = 1
-revision = 5
+revision = 3
Review Comment:
This looks like a re-lock with an older uv: the revision drops from 5 to 3
and the `[options]` exclude-newer block is gone. Nothing in the PR changes a
`pyproject.toml`, so I think `uv.lock` can come out of the diff (`git checkout
main -- uv.lock` after a rebase).
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -318,6 +326,31 @@ def get_conn(self) -> Model:
)
return self._model
+ async def aget_conn(self) -> Model:
+ """
+ Return the configured pydantic-ai ``Model``, for use inside an async
task or trigger.
+
+ The same model as :meth:`get_conn`, fallback chain included, but this
hook's connection
+ and its fallback connections are fetched with
``get_async_connection``, so resolving
+ them makes no synchronous call to the Task SDK. ``get_connection`` and
the secret
+ masking of ``extra_dejson`` send to the supervisor synchronously,
which raises
+ ``DeadlockImminentError`` on an event loop with another async call in
flight.
+ """
+ if self._model is not None:
+ return self._model
+ await self._aprefetch_connections()
+ return self.get_conn()
Review Comment:
After the prefetch this still runs the sync `get_conn()` on the loop, so a
subclass whose `_get_provider_kwargs` does its own I/O still makes sync SDK
calls. `PydanticAISnowflakeHook` is one: `_get_token_provider` builds a
`SnowflakeHook` and reads `_get_static_conn_params`, which calls
`get_connection(snowflake_conn_id)` and `extra_dejson` synchronously
(`snowflake.py:681-682`). That's hit for a Snowflake primary and for a
Snowflake connection in `fallback_conn_ids`, since fallbacks go through
`conn.get_hook()`. Same for `create_agent` in `acreate_agent`.
Would `return await asyncio.to_thread(self.get_conn)` be simpler? A `send()`
from a pool thread just blocks on `_thread_lock` until the in-flight `asend`
releases it, so that covers every subclass, including ones in other providers,
and `_prefetched_fallbacks` wouldn't be needed. Either way, a test that goes
through a subclass would catch this; all the new tests use the base hook.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -2207,3 +2208,136 @@ 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 _extra_dejson_off_the_event_loop(conn: Connection) -> dict:
+ try:
+ asyncio.get_running_loop()
+ except RuntimeError:
+ return json.loads(conn.extra) if conn.extra else {}
+ raise AssertionError("extra_dejson masks secrets synchronously on the
event loop")
+
+
[email protected]
+def async_registry():
+ """
+ Serve connections only through ``get_async_connection``: the synchronous
lookups fail on the loop.
+
+ ``get_connection`` and the secret masking of ``extra_dejson`` send to the
supervisor
+ synchronously, which raises ``DeadlockImminentError`` inside an async task
with another
+ async call in flight. ``get_connection`` must not be called at all, and
``extra_dejson``
+ only off the event loop, where ``get_async_extra_dejson`` reads it on
Airflow < 3.3.2.
+ """
+ reg = _ConnRegistry()
+ with (
+ patch(
+
"airflow.providers.common.ai.hooks.pydantic_ai.get_async_connection",
+ new=AsyncMock(side_effect=reg.get_connection),
+ ) as mock_get_async_connection,
+ patch.object(PydanticAIHook, "get_connection",
side_effect=AssertionError("sync get_connection")),
+ patch.object(Connection, "extra_dejson",
new=property(_extra_dejson_off_the_event_loop)),
+ patch.object(Connection, "get_hook", side_effect=reg.get_hook,
autospec=True),
+ ):
+ reg.get_async_connection = mock_get_async_connection
+ yield reg
+
+
+class TestPydanticAIHookAsync:
+ @pytest.mark.asyncio
+ async def test_aget_conn_resolves_the_model_asynchronously(self,
async_registry, infer_model_stub):
+ async_registry.add("primary", extra={"model": "openai:gpt-5.6-sol"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ assert await hook.aget_conn() is
infer_model_stub.models["openai:gpt-5.6-sol"]
+ async_registry.get_async_connection.assert_awaited_once_with("primary")
+
+ @pytest.mark.asyncio
+ async def test_aget_conn_resolves_the_fallback_chain_asynchronously(
+ self, async_registry, infer_model_stub
+ ):
+ async_registry.add(
+ "primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second", "third"]}
+ )
+ async_registry.add("second", extra={"model":
"anthropic:claude-opus-4-6"})
+ async_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 async_registry.get_async_connection.await_args_list == [
+ call("primary"),
+ call("second"),
+ call("third"),
+ ]
+
+ @pytest.mark.asyncio
+ async def test_aget_conn_is_cached(self, async_registry, infer_model_stub):
+ async_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"]
+ async_registry.get_async_connection.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def test_acreate_agent(self, mock_agent_cls, async_registry,
infer_model_stub):
Review Comment:
These pass with default arguments only. Dropping `output_type` and
`**agent_kwargs` from the `create_agent` call in `acreate_agent`, or
`instructions` in the `spec_file` branch, keeps all 7 tests green, and so does
replacing either fetch-once guard in `_aprefetch_connections` with `if True:`.
Could this one use a non-`str` `output_type` plus an agent kwarg, the
`spec_file` test pass `instructions`, and one test call `acreate_agent` twice
on a hook with fallbacks and assert each connection is awaited once?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -318,6 +326,31 @@ def get_conn(self) -> Model:
)
return self._model
+ async def aget_conn(self) -> Model:
Review Comment:
Not blocking, but how will an async caller get the hook in the first place?
The operators build it with `PydanticAIHook.get_hook(conn_id)` (`agent.py:775`,
`llm.py:318`), which calls `get_connection` synchronously, and constructing
`PydanticAIHook(llm_conn_id=...)` directly skips the subclass dispatch for the
other `pydanticai*` connection types. If this is the first step toward a
native-async `AgentOperator`, an `aget_hook` that fetches with
`get_async_connection`, does the same type check as `get_hook`, and seeds the
hook with `_seed_connection` probably belongs here too.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -669,6 +705,57 @@ def create_agent(
agent.instrument = settings
return agent
+ @overload
+ async def acreate_agent(
+ self, output_type: type[OutputT], *, instructions: str, **agent_kwargs
+ ) -> Agent[object, OutputT]: ...
+
+ @overload
+ async def acreate_agent(self, *, instructions: str, **agent_kwargs) ->
Agent[object, str]: ...
+
+ @overload
+ async def acreate_agent(
+ self,
+ output_type: type[OutputT],
+ *,
+ spec_file: str | Path,
+ instructions: str | None = ...,
+ **agent_kwargs,
+ ) -> Agent[object, OutputT]: ...
+
+ @overload
+ async def acreate_agent(
+ self,
+ *,
+ spec_file: str | Path,
+ instructions: str | None = ...,
+ **agent_kwargs,
+ ) -> Agent[object, str]: ...
+
+ async def acreate_agent(
Review Comment:
These are new public methods, but `docs/hooks/pydantic_ai.rst` only covers
`create_agent`. A short section on using the hook from an async task, with
`agent = await hook.acreate_agent(...)` and `await agent.run(...)`, would make
them findable.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -318,6 +326,31 @@ def get_conn(self) -> Model:
)
return self._model
+ async def aget_conn(self) -> Model:
+ """
+ Return the configured pydantic-ai ``Model``, for use inside an async
task or trigger.
+
+ The same model as :meth:`get_conn`, fallback chain included, but this
hook's connection
+ and its fallback connections are fetched with
``get_async_connection``, so resolving
+ them makes no synchronous call to the Task SDK. ``get_connection`` and
the secret
+ masking of ``extra_dejson`` send to the supervisor synchronously,
which raises
+ ``DeadlockImminentError`` on an event loop with another async call in
flight.
+ """
+ if self._model is not None:
+ return self._model
+ await self._aprefetch_connections()
+ return self.get_conn()
+
+ async def _aprefetch_connections(self) -> None:
+ """Fetch this hook's connection and its fallback connections
asynchronously, once."""
+ if self._conn is None:
+ conn = await get_async_connection(self.llm_conn_id)
Review Comment:
On Airflow 3.1 to 3.3.1 this fetch isn't fully async: after a secrets
backend returns the connection, `_async_get_connection` calls `conn.get_uri()`
synchronously to fill the secret cache, and `get_uri()` reads `extra_dejson`,
which masks with a sync `send()`. 3.3.2 switched that to `aget_uri()`. Every
pydanticai connection has an extra (the model), so on those versions the first
fetch can still raise `DeadlockImminentError`. Could the docstrings scope
"makes no synchronous call" to 3.3.2+? Running the build in a thread, as
suggested on `aget_conn`, would sidestep it too.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -669,6 +705,57 @@ def create_agent(
agent.instrument = settings
return agent
+ @overload
+ async def acreate_agent(
+ self, output_type: type[OutputT], *, instructions: str, **agent_kwargs
+ ) -> Agent[object, OutputT]: ...
+
+ @overload
+ async def acreate_agent(self, *, instructions: str, **agent_kwargs) ->
Agent[object, str]: ...
+
+ @overload
+ async def acreate_agent(
+ self,
+ output_type: type[OutputT],
+ *,
+ spec_file: str | Path,
+ instructions: str | None = ...,
+ **agent_kwargs,
+ ) -> Agent[object, OutputT]: ...
+
+ @overload
+ async def acreate_agent(
+ self,
+ *,
+ spec_file: str | Path,
+ instructions: str | None = ...,
+ **agent_kwargs,
+ ) -> Agent[object, str]: ...
+
+ async def acreate_agent(
+ self,
+ output_type: type[Any] = str,
+ *,
+ instructions: str | None = None,
+ spec_file: str | Path | None = None,
+ **agent_kwargs,
+ ) -> Agent[object, Any]:
+ """
+ Create a pydantic-ai Agent configured with this hook's model, for use
inside an async task.
+
+ The async counterpart of :meth:`create_agent`, with the same
parameters: the
+ connections are fetched as in :meth:`aget_conn`, so building the agent
makes no
+ synchronous call to the Task SDK. Run it with ``await agent.run(...)``.
+ """
+ if spec_file is None and instructions is None:
+ raise ValueError("instructions is required when spec_file is not
provided.")
+ await self._aprefetch_connections()
Review Comment:
`aget_conn` returns early when `_model` is already set, but this doesn't.
After a sync `get_conn()` or `test_connection()` on the same hook,
`_prefetched_fallbacks` is still empty, so every fallback is fetched again and
then thrown away, since `create_agent` returns the cached model. Moving the `if
self._model is not None: return` check to the top of `_aprefetch_connections`
would cover both entry points.
--
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]
