This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 d6bb744aea1 Validate the connection type in PydanticAIHook.get_hook
(#74626)
d6bb744aea1 is described below
commit d6bb744aea1ac9c1f5f43067bb62bbadaf148c6f
Author: Wei Lee <[email protected]>
AuthorDate: Wed Oct 14 17:29:32 2026 +0900
Validate the connection type in PydanticAIHook.get_hook (#74626)
* Validate the connection type in PydanticAIHook.get_hook
A connection of another type, such as openai, made the Pydantic AI
operators fail with "BaseHook.__init__() got an unexpected keyword
argument 'model_id'", because Connection.get_hook dispatched to a hook
from another provider.
PydanticAIHook.get_hook now raises a ValueError naming the connection,
its actual type and the registered pydanticai types when the type
resolves to a hook that is not a PydanticAIHook. Subclass types
(Azure, Bedrock, Vertex, Snowflake) pass, even if their module has not
been imported yet. A missing or unregistered type is still left to
Connection.get_hook.
* Reject connections without a hook class in PydanticAIHook.get_hook
Generic and email connections are registered with no hook class, so
`PydanticAIHook.get_hook` raised an `AttributeError` from `import_string`
instead of the new "needs one of [...]" message. They now get the same
message as any other non-Pydantic AI connection type.
Tighten the `get_hook` tests: add a `generic` rejection case, run
`_get_registered_hooks` unpatched once, accept a subclass registered
under a type without the `pydanticai` prefix, and assert on the
`Connection.get_hook` call instead of `get_connection`.
---
providers/common/ai/docs/troubleshooting.rst | 8 ++
.../providers/common/ai/hooks/pydantic_ai.py | 41 +++++++-
.../tests/unit/common/ai/hooks/test_pydantic_ai.py | 108 ++++++++++++++++++++-
3 files changed, 155 insertions(+), 2 deletions(-)
diff --git a/providers/common/ai/docs/troubleshooting.rst
b/providers/common/ai/docs/troubleshooting.rst
index 5e0413c8f47..f83a43cd9f3 100644
--- a/providers/common/ai/docs/troubleshooting.rst
+++ b/providers/common/ai/docs/troubleshooting.rst
@@ -66,6 +66,14 @@ An ``ImportError`` for ``pydantic_ai.models.<vendor>``, the
vendor SDK, or ``pyd
types. Chains are also not resolved recursively, so a fallback connection
may not
declare its own ``fallback_conn_ids``; list every vendor directly on the
primary.
+``Connection '...' has type '...', but a Pydantic AI hook needs one of [...]``
+ The connection named by ``llm_conn_id`` is not a ``pydanticai`` connection
type. This
+ is easy to hit by picking a similarly named type such as ``openai`` when
the OpenAI
+ provider is installed. Change the connection type to ``pydanticai``, put
the API key
+ in the **API Key** field (the password) and the model, for example
``openai:gpt-5``,
+ in **Model**. For Azure, Bedrock or Vertex use the matching
``pydanticai_azure``,
+ ``pydanticai_bedrock`` or ``pydanticai_vertex`` type instead; see
:doc:`model_providers`.
+
.. _troubleshooting-conn-type-rename:
A connection of type ``pydanticai-azure``, ``pydanticai-bedrock`` or
``pydanticai-vertex`` no longer resolves
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 6c3d4429db7..f44ae25ce54 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
@@ -27,6 +27,7 @@ from pydantic_ai.models.fallback import FallbackModel
from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.ai.observability import
genai_instrumentation_settings
+from airflow.providers.common.compat.module_loading import import_string
from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
@@ -39,6 +40,8 @@ _UNSET: Any = object()
FALLBACK_CONN_IDS_EXTRA_KEY = "fallback_conn_ids"
if TYPE_CHECKING:
+ from collections.abc import MutableMapping
+
from pydantic_ai.models import KnownModelName, Model
from airflow.providers.common.compat.sdk import Connection
@@ -86,6 +89,19 @@ def _looks_like_unrecognized_provider_prefix(prefix: str) ->
bool:
return bool(_PROVIDER_SLUG_RE.match(prefix))
+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,
+ # Airflow 2 included, go through ``ProvidersManager``. Both expose the
same ``hooks`` mapping.
+ try:
+ from airflow.sdk.providers_manager_runtime import
ProvidersManagerTaskRuntime
+ except ImportError:
+ from airflow.providers_manager import ProvidersManager
+
+ return ProvidersManager().hooks
+ return ProvidersManagerTaskRuntime().hooks
+
+
class PydanticAIHook(BaseHook):
"""
Hook for LLM access via pydantic-ai.
@@ -165,8 +181,31 @@ class PydanticAIHook(BaseHook):
Airflow 3's ``BaseHook.get_hook`` already takes ``hook_params``;
Airflow 2's does
not, so this mirrors the Airflow 3 body.
+
+ :raises ValueError: if the connection's type is registered to a hook
that is not a
+ ``PydanticAIHook``.
"""
- return cls.get_connection(conn_id).get_hook(hook_params=hook_params)
+ conn = cls.get_connection(conn_id)
+ hooks = _get_registered_hooks()
+ # Resolved from the registry rather than from ``cls.conn_type`` or the
subclasses
+ # imported so far: operators call this on the base class, and a
subclass from another
+ # provider (e.g. Snowflake's ``pydanticai_snowflake``) is only
imported on lookup.
+ # A missing or unregistered type is left to ``conn.get_hook``, whose
error explains a
+ # URI that parsed with no type, or the '-'/'_' spelling a renamed type
needs.
+ hook_info = hooks.get(conn.conn_type) if conn.conn_type else None
+ # ``generic`` and ``email`` are registered with
``hook_class_name=None``; that is not a
+ # Pydantic AI hook either.
+ if hook_info is not None and (
+ hook_info.hook_class_name is None
+ or not issubclass(import_string(hook_info.hook_class_name),
PydanticAIHook)
+ ):
+ expected_types = sorted(conn_type for conn_type in hooks if
conn_type.startswith("pydanticai"))
+ raise ValueError(
+ f"Connection {conn_id!r} has type {conn.conn_type!r}, but a
Pydantic AI hook needs "
+ 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)
@staticmethod
def get_ui_field_behaviour() -> dict[str, Any]:
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 f0c69bee7c7..4b703301b87 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
@@ -40,10 +40,13 @@ from airflow.providers.common.ai.hooks.pydantic_ai import (
PydanticAIBedrockHook,
PydanticAIHook,
PydanticAIVertexHook,
+ _get_registered_hooks,
_has_recognized_provider_prefix,
_looks_like_unrecognized_provider_prefix,
)
-from airflow.providers.common.compat.sdk import AirflowNotFoundException
+from airflow.providers.common.compat.module_loading import import_string,
qualname
+from airflow.providers.common.compat.sdk import AirflowNotFoundException,
BaseHook
+from airflow.providers_manager import HookInfo
# Matches the `google...` provider key pydantic-ai expects before the
`:model-name`
# separator, e.g. "google-cloud" out of "google-cloud:gemini-2.5-flash".
@@ -287,6 +290,23 @@ class _ConnRegistry:
def __init__(self) -> None:
self.conns: dict[str, Connection] = {}
self.hook_classes: dict[str, type[PydanticAIHook]] = {}
+ # Stands in for the providers manager's connection-type registry,
seeded with what
+ # this provider really registers; other providers' types are added per
test.
+ self.hooks: dict[str, HookInfo] = {}
+ for connection_type in get_provider_info()["connection-types"]:
+ self.register_hook(connection_type["connection-type"],
connection_type["hook-class-name"])
+
+ def register_hook(self, conn_type: str, hook_class_name: str | None) ->
None:
+ self.hooks[conn_type] = HookInfo(
+ # ``HookInfo`` types this as ``str``, but the providers manager
registers ``None`` for
+ # types like ``generic``.
+ hook_class_name=hook_class_name, # type: ignore[arg-type]
+ connection_id_attribute_name="llm_conn_id",
+ package_name="apache-airflow-providers-test",
+ hook_name=conn_type,
+ connection_type=conn_type,
+ connection_testable=False,
+ )
def add(
self,
@@ -333,10 +353,21 @@ def registry():
# argument -- without it, `conn.get_hook()` calls the mock with zero
arguments and
# `reg.get_hook` would have no way to know which connection dispatched
it.
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,
+ ),
):
yield reg
+class _OtherProviderHook(BaseHook):
+ """Another provider's hook, registered under a type that is easy to
confuse with ``pydanticai``."""
+
+ conn_type = "openai"
+
+
class TestPydanticAIHookGetHook:
def test_builds_the_connection_hook_with_hook_params(self, registry):
"""Airflow 2's ``BaseHook.get_hook`` takes no ``hook_params``; the
hook's own override does."""
@@ -351,6 +382,81 @@ class TestPydanticAIHookGetHook:
assert hook.model_id == "openai:gpt-5"
assert hook.fallback_conn_ids == []
+ def test_registers_the_pydanticai_hook_without_patching_the_registry(self):
+ assert _get_registered_hooks()["pydanticai"].hook_class_name ==
qualname(PydanticAIHook)
+
+ @pytest.mark.parametrize(
+ ("conn_type", "hook_class_name"),
+ [
+ pytest.param("openai", qualname(_OtherProviderHook),
id="other-provider-hook"),
+ pytest.param("generic", None, id="no-hook-class"),
+ ],
+ )
+ def test_rejects_a_connection_of_another_type_before_dispatching(
+ self, registry, conn_type, hook_class_name
+ ):
+ registry.add("llm", conn_type=conn_type)
+ registry.register_hook(conn_type, hook_class_name)
+ expected_types = ["pydanticai", "pydanticai_azure",
"pydanticai_bedrock", "pydanticai_vertex"]
+
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ f"Connection 'llm' has type {conn_type!r}, but a Pydantic AI
hook needs one of {expected_types}"
+ ),
+ ):
+ PydanticAIHook.get_hook("llm", hook_params={"model_id":
"openai:gpt-5"})
+
+ Connection.get_hook.assert_not_called()
+
+ def test_accepts_a_subclass_from_a_module_not_imported_yet(self, registry,
tmp_path, monkeypatch):
+ """Another provider's subclass (e.g. Snowflake's) is only imported
when its type is looked up."""
+ module_name = "pydanticai_subclass_not_imported_yet"
+ (tmp_path / f"{module_name}.py").write_text(
+ "from airflow.providers.common.ai.hooks.pydantic_ai import
PydanticAIHook\n\n\n"
+ "class PydanticAIElsewhereHook(PydanticAIHook):\n"
+ " conn_type = 'snowflake_llm'\n"
+ )
+ monkeypatch.syspath_prepend(tmp_path)
+ registry.add("llm", conn_type="snowflake_llm")
+ registry.register_hook("snowflake_llm",
f"{module_name}.PydanticAIElsewhereHook")
+ assert module_name not in sys.modules
+
+ try:
+ PydanticAIHook.get_hook("llm", hook_params={"model_id":
"openai:gpt-5"})
+ # The type has no ``pydanticai`` prefix, so only resolving its
hook class lets it through.
+ assert module_name in sys.modules
+ finally:
+ sys.modules.pop(module_name, None)
+
+ Connection.get_hook.assert_called_once_with(
+ registry.conns["llm"], hook_params={"model_id": "openai:gpt-5"}
+ )
+
+ def
test_leaves_an_unregistered_connection_type_to_connection_get_hook(self,
registry):
+ """``Connection.get_hook`` raises for it with a hint about the '-'/'_'
spelling."""
+ registry.add("llm", conn_type="pydanticai-azure")
+
+ PydanticAIHook.get_hook("llm")
+
+ Connection.get_hook.assert_called_once_with(registry.conns["llm"],
hook_params=None)
+
+ @pytest.mark.parametrize(
+ ("conn_type", "hook_class"),
+ [
+ (connection_type["connection-type"],
import_string(connection_type["hook-class-name"]))
+ for connection_type in get_provider_info()["connection-types"]
+ if connection_type["connection-type"].startswith("pydanticai")
+ ],
+ )
+ def test_accepts_every_pydanticai_connection_type(self, registry,
conn_type, hook_class):
+ registry.add("llm", conn_type=conn_type, hook_class=hook_class)
+
+ hook = PydanticAIHook.get_hook("llm")
+
+ assert type(hook) is hook_class
+ Connection.get_hook.assert_called_once_with(registry.conns["llm"],
hook_params=None)
+
class _InferModelStub:
"""Resolve every model string to its own recognisable model, and record
how it was built."""
