dabla commented on code in PR #74150:
URL: https://github.com/apache/airflow/pull/74150#discussion_r4293018469
##########
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:
Switched to `asyncio.to_thread` for `get_conn` and `create_agent` in
158be99963, so subclass overrides are covered and the prefetch is gone. A new
test goes through a subclass that reads another connection.
##########
uv.lock:
##########
@@ -1,5 +1,5 @@
version = 1
-revision = 5
+revision = 3
Review Comment:
Rebased and restored `uv.lock` from main.
##########
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:
No longer applies: the async fetch is gone, so this holds on every version.
##########
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:
Gone with the prefetch. `create_agent` finds the cached model; a test calls
`acreate_agent` twice and checks each connection is fetched once.
##########
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:
Done: non-`str` `output_type`, an agent kwarg, `instructions` with
`spec_file`, and `acreate_agent` called twice. Each test fails if its argument
or thread hop is dropped.
--
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]
