This is an automated email from the ASF dual-hosted git repository.
dabla 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 49bdd34d922 Add async aget_conn and acreate_agent to PydanticAIHook
(#74150)
49bdd34d922 is described below
commit 49bdd34d9222182e137f7a30a03e738ff5d8abde
Author: David Blain <[email protected]>
AuthorDate: Wed Oct 28 16:19:33 2026 +0100
Add async aget_conn and acreate_agent to PydanticAIHook (#74150)
* Add async aget_conn and acreate_agent to PydanticAIHook
get_conn() resolves the hook's connection and every fallback connection
with get_connection() and reads extra_dejson, whose secret masking sends
to the supervisor synchronously; inside an async task with another async
SDK call in flight both raise DeadlockImminentError, although the agent
the hook builds is async-native. aget_conn() fetches the connection and
its fallbacks with get_async_connection() (extra from
Connection.aextra_dejson(), #71890, Airflow 3.3.2+, or the raw extra via
json.loads on older versions as in #72130), seeds the existing
connection cache and builds the model with the unchanged get_conn();
acreate_agent() does the same before create_agent(), spec_file included.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
* Use get_async_extra_dejson from common.compat
Replace the local aextra_dejson fallback with the common.compat helper
added in #74147, which masks the extra's secrets asynchronously on
Airflow 3.3.2+ and reads extra_dejson from a worker thread on older
versions.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
* Resolve the async PydanticAIHook paths in a worker thread
Prefetching the connections only covered the base hook's own lookups: a
subclass's _get_provider_kwargs (Snowflake's builds a SnowflakeHook)
still made synchronous Task SDK calls on the event loop, and on Airflow
3.1 to 3.3.1 get_async_connection itself reads extra_dejson
synchronously. aget_conn() and acreate_agent() now run the unchanged
get_conn() / create_agent() in asyncio.to_thread: from any thread other
than the loop's, a synchronous send() waits for the in-flight asend() to
release the comms lock instead of raising DeadlockImminentError. This
drops the prefetch, the prefetched fallback cache and the extra
parameter of _seed_connection.
aget_hook() is the async counterpart of get_hook(), with the same
connection type check, so an async operator can dispatch to the
pydanticai* subclasses. get_hook() now caches the connection it fetched
on the hook it returns, so the hook's first get_conn() does not fetch it
again.
The tests refuse the synchronous lookups only on the event loop, drive a
subclass that reads another connection as primary and as fallback, and
check that every argument of acreate_agent reaches the agent. The hook
docs gain a section on using it from an async task, and uv.lock is back
to main's.
* Publish the PydanticAIHook connection and its extra together
_get_conn_and_extra set _conn before reading extra_dejson, whose secret
masking is a supervisor round trip. A second acreate_agent gathered on
the same hook found _conn set and read the empty default extra, so it
raised "No model specified" or dropped the fallback chain; a failed
extra read left the hook half-seeded the same way. Both are now
published once both are in hand.
get_hook's type check already guarantees a PydanticAIHook, so the
isinstance guard before seeding is gone, and _seed_connection's
docstring describes its contract instead of one caller.
* Run the async PydanticAIHook lookups on a worker thread of their own
asyncio.to_thread shared the default executor with asend(), which holds
the comms lock while a pool thread reads its reply. Enough gathered hook
calls filled the pool: every thread blocked in send() on the lock and
nothing was left to run the read, so the task hung. The three async
methods now run on a single-worker executor of the hook module, with the
caller's context variables copied because task logging binds them. One
worker also serializes the hook's sends on Airflow 3.0 and 3.1, where
the channel to the supervisor has no lock of its own.
The docs section states Airflow 3.2.1 and later (3.2.0 reads the reply
outside the lock) and that the loop blocks before 3.3.1, where
DeadlockImminentError arrived.
* Test the PydanticAIHook context copy and the two-thread cache publish
The context copy in _run_off_loop had no test: a plain functools.partial
kept every async test green, while it is what carries the task instance
fields the task runner binds for logging into the worker thread. A test
now sets a context variable before each of the three async methods and
checks the worker sees it.
With one worker, the two gathered acreate_agent calls ran one after the
other, so that test no longer exercised the publish order of
_get_conn_and_extra: the old order kept it green. Two threads calling
get_conn() on one hook replace it, and the note in _get_conn_and_extra
names that case instead of gathered calls.
---------
Co-authored-by: Claude Opus 5.5 <[email protected]>
---
providers/common/ai/docs/hooks/pydantic_ai.rst | 20 ++
.../providers/common/ai/hooks/pydantic_ai.py | 135 ++++++++-
.../tests/unit/common/ai/hooks/test_pydantic_ai.py | 336 +++++++++++++++++++++
3 files changed, 481 insertions(+), 10 deletions(-)
diff --git a/providers/common/ai/docs/hooks/pydantic_ai.rst
b/providers/common/ai/docs/hooks/pydantic_ai.rst
index dc90cd91380..0ac7ef1336e 100644
--- a/providers/common/ai/docs/hooks/pydantic_ai.rst
+++ b/providers/common/ai/docs/hooks/pydantic_ai.rst
@@ -97,3 +97,23 @@ 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 on Airflow 3.2.1 and later, 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 block the event loop while
another async
+call is in flight (from 3.3.1 they raise ``DeadlockImminentError`` instead),
so the async
+methods run them on a worker thread of their own.
+
+.. code-block:: python
+
+ @task
+ async def summarize(text: str) -> str:
+ hook = await PydanticAIHook.aget_hook("my_llm")
+ agent = await hook.acreate_agent(instructions="Summarize the text in
one sentence.")
+ result = await agent.run(text)
+ return result.output
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
index 741dfe5a30d..8fee9c1da69 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
@@ -16,6 +16,10 @@
# under the License.
from __future__ import annotations
+import asyncio
+import concurrent.futures
+import contextvars
+import functools
import re
from pathlib import Path
from typing import TYPE_CHECKING, Any, TypeVar, overload
@@ -26,6 +30,7 @@ from pydantic_ai.exceptions import ModelAPIError
from pydantic_ai.models import infer_model
from pydantic_ai.models.fallback import FallbackModel
from pydantic_ai.providers import infer_provider_class
+from typing_extensions import ParamSpec
from airflow.providers.common.ai.observability import
genai_instrumentation_settings
from airflow.providers.common.ai.utils.toolset_base import MaskingCapability
@@ -37,6 +42,8 @@ from airflow.providers.common.compat.sdk import BaseHook
pydantic_ai.BANNER_ENABLED = False
OutputT = TypeVar("OutputT")
+P = ParamSpec("P")
+R = TypeVar("R")
# Sentinel distinguishing "caller did not pass ``instrument``" from an explicit
# ``instrument=None`` / ``instrument=False`` (which mean "do not instrument,
and
@@ -46,7 +53,7 @@ _UNSET: Any = object()
FALLBACK_CONN_IDS_EXTRA_KEY = "fallback_conn_ids"
if TYPE_CHECKING:
- from collections.abc import MutableMapping
+ from collections.abc import Callable, MutableMapping
from pydantic_ai.models import KnownModelName, Model
@@ -95,6 +102,20 @@ def _looks_like_unrecognized_provider_prefix(prefix: str)
-> bool:
return bool(_PROVIDER_SLUG_RE.match(prefix))
+# The async methods run their blocking lookups here, never on asyncio's
default executor:
+# ``asend()`` reads its reply on that executor while it holds the comms lock,
so blocking
+# lookups that fill it leave nothing to run the read and the task hangs. One
worker also
+# serializes this process's hook lookups on Airflow 3.0 and 3.1, where the
channel to the
+# supervisor has no lock of its own.
+_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1,
thread_name_prefix="pydantic-ai-hook")
+
+
+async def _run_off_loop(fn: Callable[P, R], /, *args: P.args, **kwargs:
P.kwargs) -> R:
+ """Run ``fn`` on the hook's own worker thread, with the caller's context
variables (task logging binds them)."""
+ call = functools.partial(contextvars.copy_context().run, fn, *args,
**kwargs)
+ return await asyncio.get_running_loop().run_in_executor(_executor, call)
+
+
def _get_registered_hooks() -> MutableMapping[str, Any]:
"""Return the connection-type-to-``HookInfo`` registry that
``Connection.get_hook`` dispatches through."""
# Airflow 3.2+ resolves hooks through the task runtime's manager; earlier
versions,
@@ -211,7 +232,27 @@ class PydanticAIHook(BaseHook):
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. ``get_connection`` sends to
the supervisor
+ synchronously, which blocks the event loop while another async call is
in flight (from
+ Airflow 3.3.1 it raises ``DeadlockImminentError`` instead), but only
waits its turn from
+ any other thread, so the lookup runs on the hook's own worker thread.
+
+ :raises ValueError: if the connection's type is registered to a hook
that is not a
+ ``PydanticAIHook``.
+ """
+ return await _run_off_loop(cls.get_hook, conn_id, hook_params)
@staticmethod
def get_ui_field_behaviour() -> dict[str, Any]:
@@ -261,19 +302,23 @@ class PydanticAIHook(BaseHook):
def _get_conn_and_extra(self) -> tuple[Connection, dict[str, Any]]:
"""Return this hook's connection and its deserialized extra, fetching
at most once."""
if self._conn is None:
- self._conn = self.get_connection(self.llm_conn_id)
- self._conn_extra_dejson = self._conn.extra_dejson
+ # Published together, once both are in hand: a reader never sees
the connection
+ # next to the empty default extra, whether ``extra_dejson`` is
still masking in
+ # another thread (``get_conn()`` called from two threads on one
hook) or raised.
+ conn = self.get_connection(self.llm_conn_id)
+ extra = conn.extra_dejson
+ self._conn, self._conn_extra_dejson = conn, extra
return self._conn, self._conn_extra_dejson
def _seed_connection(self, conn: Connection) -> None:
"""
- Prime this hook's connection cache with an already-fetched
``Connection``.
+ Prime this hook's connection cache with a ``Connection`` the caller
already fetched.
- Used by :meth:`_resolve_fallback_models`, which must call
``conn.get_hook()`` to
- dispatch a fallback's hook class from its ``conn_type`` and so already
holds the
- ``Connection`` that call built the hook from. Without this,
:meth:`_get_conn_and_extra`
- would fetch that same connection a second time the first time it runs,
doubling the
- Execution API round trips a fallback chain costs.
+ A caller that dispatches the hook class from ``conn.get_hook()`` holds
the
+ ``Connection`` it built the hook from: :meth:`get_hook` for the hook
it returns,
+ :meth:`_resolve_fallback_models` for each fallback. Without this,
+ :meth:`_get_conn_and_extra` would fetch that same connection a second
time the first
+ time it runs.
"""
self._conn = conn
self._conn_extra_dejson = conn.extra_dejson
@@ -323,6 +368,22 @@ class PydanticAIHook(BaseHook):
)
return self._model
+ async def aget_conn(self) -> Model:
+ """
+ Return the configured pydantic-ai ``Model``, for use inside an async
task.
+
+ The same model as :meth:`get_conn`, fallback chain included. Resolving
it reads
+ connections with ``get_connection`` and ``extra_dejson``, and a
subclass's
+ :meth:`_get_provider_kwargs` may read more (Snowflake's builds a
``SnowflakeHook``);
+ each of those sends to the supervisor synchronously, which blocks the
event loop while
+ another async call is in flight (from Airflow 3.3.1 it raises
``DeadlockImminentError``
+ instead), but only waits its turn from any other thread, so the
resolution runs on the
+ hook's own worker thread.
+ """
+ if self._model is not None:
+ return self._model
+ return await _run_off_loop(self.get_conn)
+
def _qualify_model_name(self, model_name: str, *, forwarded_from_conn_id:
str | None = None) -> str:
"""
Prefix a bare model name with this connection's platform.
@@ -673,6 +734,60 @@ class PydanticAIHook(BaseHook):
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, built on the
+ hook's own worker thread for the reason given in :meth:`aget_conn`.
Run the agent with
+ ``await agent.run(...)``.
+ """
+ if spec_file is not None:
+ build = functools.partial(
+ self.create_agent, output_type, spec_file=spec_file,
instructions=instructions, **agent_kwargs
+ )
+ elif instructions is None:
+ raise ValueError("instructions is required when spec_file is not
provided.")
+ else:
+ build = functools.partial(
+ self.create_agent, output_type, instructions=instructions,
**agent_kwargs
+ )
+ return await _run_off_loop(build)
+
def test_connection(self) -> tuple[bool, str]:
"""
Test connection by resolving the model.
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
index d86365ed00a..596cbd27a76 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
@@ -16,11 +16,16 @@
# under the License.
from __future__ import annotations
+import asyncio
import contextlib
import json
import logging
import re
import sys
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import ContextVar
from pathlib import Path
from unittest.mock import MagicMock, call, patch
@@ -385,6 +390,13 @@ class TestPydanticAIHookGetHook:
assert hook.model_id == "openai:gpt-5"
assert hook.fallback_conn_ids == []
+ def test_seeds_the_hook_with_the_connection_it_fetched(self, registry,
infer_model_stub):
+ registry.add("llm", extra={"model": "openai:gpt-5"})
+
+ PydanticAIHook.get_hook("llm").get_conn()
+
+ PydanticAIHook.get_connection.assert_called_once_with("llm")
+
def test_registers_the_pydanticai_hook_without_patching_the_registry(self):
assert _get_registered_hooks()["pydanticai"].hook_class_name ==
qualname(PydanticAIHook)
@@ -2282,3 +2294,327 @@ class TestConnTypeResolution:
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
+
+
+# Stands in for the task instance fields the task runner binds for logging.
+_caller_context: ContextVar[str | None] = ContextVar("caller_context",
default=None)
+
+# The three async entry points, each run against a hook for the "primary"
connection.
+_ASYNC_HOPS = [
+ pytest.param(lambda hook: type(hook).aget_hook("primary"), id="aget_hook"),
+ pytest.param(lambda hook: hook.aget_conn(), id="aget_conn"),
+ pytest.param(lambda hook: hook.acreate_agent(instructions="Be helpful."),
id="acreate_agent"),
+]
+
+
[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",
autospec=True)
+ async def test_acreate_agent_with_spec_file_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"})
+
+ await PydanticAIHook(llm_conn_id="primary").acreate_agent(
+ int, spec_file="/path/to/agent.yaml", instructions="Be brief.",
retries=3
+ )
+
+ mock_agent_cls.from_file.assert_called_once_with(
+ "/path/to/agent.yaml",
+ output_type=int,
+ model=infer_model_stub.models["openai:gpt-5.6-sol"],
+ instructions="Be brief.",
+ retries=3,
+ capabilities=[MaskingCapability()],
+ )
+
+ @pytest.mark.asyncio
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def test_acreate_agent_twice_fetches_each_connection_once(
+ self, mock_agent_cls, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add(
+ "primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second"]}
+ )
+ loop_guarded_registry.add("second", extra={"model":
"anthropic:claude-opus-4-6"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ await hook.acreate_agent(instructions="Be helpful.")
+ await hook.acreate_agent(instructions="Be helpful.")
+
+ assert loop_guarded_registry.get_connection_mock.call_args_list ==
[call("primary"), call("second")]
+
+ @pytest.mark.asyncio
+ async def
test_acreate_agent_without_instructions_or_spec_file_raises_before_fetching(
+ self, loop_guarded_registry
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+
+ with pytest.raises(ValueError, match="instructions is required when
spec_file is not provided"):
+ await PydanticAIHook(llm_conn_id="primary").acreate_agent()
+
+ loop_guarded_registry.get_connection_mock.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_aget_hook_builds_the_connection_hook_off_the_event_loop(
+ self, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("llm", extra={"model": "openai:gpt-5.6-sol"})
+
+ hook = await PydanticAIHook.aget_hook("llm", hook_params={"model_id":
"anthropic:claude-opus-4-6"})
+
+ assert hook.model_id == "anthropic:claude-opus-4-6"
+ assert await hook.aget_conn() is
infer_model_stub.models["anthropic:claude-opus-4-6"]
+
loop_guarded_registry.get_connection_mock.assert_called_once_with("llm")
+
+ @pytest.mark.asyncio
+ async def test_aget_hook_rejects_a_connection_of_another_type(self,
loop_guarded_registry):
+ loop_guarded_registry.add("llm", conn_type="openai")
+ loop_guarded_registry.register_hook("openai",
qualname(_OtherProviderHook))
+
+ with pytest.raises(ValueError, match="Connection 'llm' has type
'openai'"):
+ await PydanticAIHook.aget_hook("llm")
+
+ Connection.get_hook.assert_not_called()
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("hop", _ASYNC_HOPS)
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def test_async_hops_do_not_need_the_default_executor(
+ self, mock_agent_cls, hop, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+ loop = asyncio.get_running_loop()
+ default_executor = ThreadPoolExecutor(max_workers=1)
+ loop.set_default_executor(default_executor)
+ release = threading.Event()
+ # The only default worker is taken, as asend() takes it to read its
reply.
+ parked = loop.run_in_executor(None, release.wait)
+
+ try:
+ await asyncio.wait_for(hop(hook), timeout=5)
+ finally:
+ release.set()
+ await parked
+ default_executor.shutdown()
+
+
loop_guarded_registry.get_connection_mock.assert_called_once_with("primary")
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("hop", _ASYNC_HOPS)
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent",
autospec=True)
+ async def test_async_hops_carry_the_callers_context_variables(
+ self, mock_agent_cls, hop, loop_guarded_registry, infer_model_stub
+ ):
+ loop_guarded_registry.add("primary", extra={"model":
"openai:gpt-5.6-sol"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+ fetch = loop_guarded_registry.get_connection_mock.side_effect
+ seen_in_the_worker: list[str | None] = []
+
+ def fetch_and_record_the_context(conn_id):
+ seen_in_the_worker.append(_caller_context.get())
+ return fetch(conn_id)
+
+ loop_guarded_registry.get_connection_mock.side_effect =
fetch_and_record_the_context
+ _caller_context.set("task-run-1")
+
+ await hop(hook)
+
+ assert seen_in_the_worker == ["task-run-1"]
+
+ @pytest.mark.asyncio
+ async def test_gathered_hooks_never_send_at_the_same_time(self,
loop_guarded_registry, infer_model_stub):
+ loop_guarded_registry.add("first", extra={"model":
"openai:gpt-5.6-sol"})
+ loop_guarded_registry.add("second", extra={"model":
"anthropic:claude-opus-4-6"})
+ active = 0
+ overlapped = False
+
+ def slow_extra(conn):
+ nonlocal active, overlapped
+ active += 1
+ overlapped = overlapped or active > 1
+ time.sleep(0.05)
+ active -= 1
+ return json.loads(conn.extra) if conn.extra else {}
+
+ with patch.object(Connection, "extra_dejson",
new=property(_off_the_event_loop(slow_extra))):
+ await asyncio.gather(
+ PydanticAIHook(llm_conn_id="first").aget_conn(),
+ PydanticAIHook(llm_conn_id="second").aget_conn(),
+ )
+
+ assert not overlapped
+
+
+class TestPydanticAIHookConnectionCache:
+ def test_get_conn_from_two_threads_each_sees_a_complete_connection(self,
registry, infer_model_stub):
+ registry.add("primary", extra={"model": "openai:gpt-5.6-sol",
"fallback_conn_ids": ["second"]})
+ registry.add("second", extra={"model": "anthropic:claude-opus-4-6"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ def slow_extra(conn):
+ time.sleep(0.05) # long enough for the other thread to reach the
cache meanwhile
+ return json.loads(conn.extra) if conn.extra else {}
+
+ with (
+ patch.object(Connection, "extra_dejson", new=property(slow_extra)),
+ ThreadPoolExecutor(max_workers=2) as threads,
+ ):
+ first_model, second_model = (
+ f.result() for f in [threads.submit(hook.get_conn) for _ in
range(2)]
+ )
+
+ assert isinstance(first_model, FallbackModel)
+ assert isinstance(second_model, FallbackModel)
+ assert (
+ first_model.models
+ == second_model.models
+ == [
+ infer_model_stub.models["openai:gpt-5.6-sol"],
+ infer_model_stub.models["anthropic:claude-opus-4-6"],
+ ]
+ )
+
+ def test_get_conn_fetches_again_after_the_extra_failed(self, registry,
infer_model_stub):
+ registry.add("primary", extra={"model": "openai:gpt-5.6-sol"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+ failing_extra = property(MagicMock(side_effect=RuntimeError("masking
failed")))
+
+ with patch.object(Connection, "extra_dejson", new=failing_extra),
pytest.raises(RuntimeError):
+ hook.get_conn()
+
+ assert hook.get_conn() is infer_model_stub.models["openai:gpt-5.6-sol"]
+ assert PydanticAIHook.get_connection.call_args_list ==
[call("primary"), call("primary")]
