Lee-W commented on code in PR #72156:
URL: https://github.com/apache/airflow/pull/72156#discussion_r4045890749
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -240,6 +246,585 @@ def test_get_conn_caches_model(self, mock_infer_model):
mock_infer_model.assert_called_once()
+class _ConnRegistry:
+ """
+ In-memory stand-in for connection and hook lookup.
+
+ ``_resolve_fallback_models`` goes through ``BaseHook.get_hook``, which
needs both the
+ metadata DB and provider discovery; this resolves both from a dict instead.
+ """
+
+ def __init__(self) -> None:
+ self.conns: dict[str, Connection] = {}
+ self.hook_classes: dict[str, type[PydanticAIHook]] = {}
+
+ def add(
+ self,
+ conn_id: str,
+ *,
+ conn_type: str = "pydanticai",
+ hook_class: type[PydanticAIHook] = PydanticAIHook,
+ password: str | None = None,
+ extra: dict | None = None,
+ ) -> None:
+ self.conns[conn_id] = Connection(
+ conn_id=conn_id,
+ conn_type=conn_type,
+ password=password,
+ extra=json.dumps(extra) if extra else None,
+ )
+ self.hook_classes[conn_id] = hook_class
+
+ def get_connection(self, conn_id: str) -> Connection:
+ try:
+ return self.conns[conn_id]
+ except KeyError:
+ raise AirflowNotFoundException(f"The conn_id `{conn_id}` isn't
defined") from None
+
+ def get_hook(self, conn_id: str, hook_params: dict | None = None):
+ if conn_id not in self.conns:
+ raise AirflowNotFoundException(f"The conn_id `{conn_id}` isn't
defined")
+ hook_class = self.hook_classes[conn_id]
+ return hook_class(llm_conn_id=conn_id, **(hook_params or {}))
+
+
[email protected]
+def registry():
+ """Patch connection and hook lookup onto a registry the test populates."""
+ reg = _ConnRegistry()
+ with (
+ patch.object(PydanticAIHook, "get_connection",
side_effect=reg.get_connection),
+ patch.object(PydanticAIHook, "get_hook", side_effect=reg.get_hook),
+ ):
+ yield reg
+
+
+class _InferModelStub:
+ """Resolve every model string to its own recognisable model, and record
how it was built."""
+
+ def __init__(self, mock: MagicMock) -> None:
+ self.mock = mock
+ self.models: dict[str, MagicMock] = {}
+
+ def __call__(self, model_name: str, **kwargs) -> MagicMock:
+ return self.models.setdefault(model_name, MagicMock(spec=Model,
name=model_name))
+
+ def provider_kwargs_for(self, model_name: str, infer_provider_class:
MagicMock) -> dict:
+ """Return the kwargs the provider for *model_name* would be
constructed with."""
+ factory = next(
+ call.kwargs["provider_factory"] for call in
self.mock.call_args_list if call.args[0] == model_name
+ )
+ infer_provider_class.return_value.reset_mock()
+ factory(model_name.split(":")[0])
+ return infer_provider_class.return_value.call_args.kwargs
+
+
[email protected]
+def infer_model_stub():
+ """Patch ``infer_model`` so tests can tell the models of a chain apart."""
+ with patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True) as mock:
+ stub = _InferModelStub(mock)
+ mock.side_effect = stub
+ yield stub
+
+
+class TestPydanticAIHookModelProviderResolution:
+ """Bare model names get qualified with a connection's own platform
prefix."""
+
+ @pytest.mark.parametrize(
+ ("hook_class", "conn_type", "prefix"),
+ [
+ pytest.param(PydanticAIAzureHook, "pydanticai_azure", "azure",
id="azure"),
+ pytest.param(PydanticAIBedrockHook, "pydanticai_bedrock",
"bedrock", id="bedrock"),
+ pytest.param(PydanticAIVertexHook, "pydanticai_vertex",
"google-cloud", id="vertex"),
+ ],
+ )
+ def test_bare_model_id_gets_platform_prefix(
+ self, registry, infer_model_stub, hook_class, conn_type, prefix
+ ):
+ registry.add("primary", conn_type=conn_type, hook_class=hook_class,
extra={"model": "foo"})
+ hook = hook_class(llm_conn_id="primary")
+
+ assert hook.get_conn() is infer_model_stub.models[f"{prefix}:foo"]
+
+ def test_prefixed_model_id_used_verbatim(self, registry, infer_model_stub):
+ """A name that already contains ``:`` pins its own platform and is
never re-prefixed."""
+ registry.add(
+ "primary",
+ conn_type="pydanticai_azure",
+ hook_class=PydanticAIAzureHook,
+ extra={"model": "openai:gpt-4"},
+ )
+ hook = PydanticAIAzureHook(llm_conn_id="primary")
+
+ assert hook.get_conn() is infer_model_stub.models["openai:gpt-4"]
+
+ def test_generic_connection_bare_name_raises_actionable_error(self,
registry, infer_model_stub):
+ """The generic ``pydanticai`` connection type has no platform of its
own."""
+ registry.add("primary", extra={"model": "gpt-4"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ with pytest.raises(ValueError, match="primary") as exc_info:
+ hook.get_conn()
+
+ assert "gpt-4" in str(exc_info.value)
+
+ def test_vertex_bare_model_id_ignores_credential_shape(self, registry,
infer_model_stub):
+ """Vertex's default platform never depends on which credential fields
are set.
+
+ ``api_key`` in this hook's extra can mean either the Generative
Language API or
+ Vertex API-key auth, so it cannot decide the platform -- there is
deliberately no
+ inference here, only the class-level default.
+ """
+ registry.add(
+ "primary",
+ conn_type="pydanticai_vertex",
+ hook_class=PydanticAIVertexHook,
+ extra={"model": "gemini-2.0-flash", "api_key": "some-key"},
+ )
+ hook = PydanticAIVertexHook(llm_conn_id="primary")
+
+ assert hook.get_conn() is
infer_model_stub.models["google-cloud:gemini-2.0-flash"]
+
+ def
test_bedrock_bare_model_id_with_embedded_colon_gets_platform_prefix(self,
registry, infer_model_stub):
+ """A ``:`` alone doesn't pin a platform -- Bedrock's own ids contain
one.
+
+ Bedrock's version-suffixed ids (e.g.
``us.anthropic.claude-opus-4-6-v1:0``) contain
+ a ``:`` that is not a pydantic-ai provider name, so a bare copy of one
must still get
+ the ``bedrock:`` prefix, not be treated as already-qualified.
+
+ Mutation canary: reverting
``_qualify_model_name``/``_has_recognized_provider_prefix``
+ to the old ``":" in model_name`` check makes this resolve to the
unprefixed
+ ``"us.anthropic.claude-opus-4-6-v1:0"`` instead, failing the ``is``
identity assertion
+ (a different key in ``infer_model_stub.models``).
+ """
+ registry.add(
+ "primary",
+ conn_type="pydanticai_bedrock",
+ hook_class=PydanticAIBedrockHook,
+ extra={"model": "us.anthropic.claude-opus-4-6-v1:0"},
+ )
+ hook = PydanticAIBedrockHook(llm_conn_id="primary")
+
+ assert hook.get_conn() is
infer_model_stub.models["bedrock:us.anthropic.claude-opus-4-6-v1:0"]
+
+ def test_prefixed_model_id_with_embedded_colon_used_verbatim(self,
registry, infer_model_stub):
+ """A name already pinning a recognized platform is never re-prefixed,
even with an
+ embedded ``:`` of its own.
+
+ Mutation canary: dropping the ``infer_provider_class`` recognition
check (treating
+ every ``:`` split the same) has no effect on *this* test by itself
since the string
+ already starts with a recognized prefix -- what would catch a
regression here is a
+ mutation that re-adds prefixing unconditionally (e.g. always prepending
+ ``model_provider`` regardless of ``_has_recognized_provider_prefix``'s
result), which
+ would turn the resolved key into
+ ``"bedrock:bedrock:us.anthropic.claude-opus-4-6-v1:0"`` and fail the
identity assertion.
+ """
+ registry.add(
+ "primary",
+ conn_type="pydanticai_bedrock",
+ hook_class=PydanticAIBedrockHook,
+ extra={"model": "bedrock:us.anthropic.claude-opus-4-6-v1:0"},
+ )
+ hook = PydanticAIBedrockHook(llm_conn_id="primary")
+
+ assert hook.get_conn() is
infer_model_stub.models["bedrock:us.anthropic.claude-opus-4-6-v1:0"]
+
+ def
test_generic_connection_unrecognized_prefix_raises_actionable_error(self,
registry, infer_model_stub):
+ """A ``:`` whose left segment isn't a real provider must not slip past
as "prefixed".
+
+ On the generic connection type (no platform of its own) a name like
+ ``"us.anthropic.claude-opus-4-6-v1:0"`` must raise this hook's own
actionable error
+ naming the connection, not be forwarded to pydantic-ai's
``infer_model`` where it
+ would instead raise the less actionable ``UserError: Unknown model``.
+
+ Mutation canary: reverting to the old ``":" in model_name`` check
makes this string
+ look "already prefixed" (since it contains a ``:``) and skips the
``ValueError`` raise
+ entirely -- the ``pytest.raises(ValueError, match="primary")`` block
would then fail
+ because no exception is raised (the stubbed ``infer_model`` would
resolve it instead).
+ """
+ registry.add("primary", extra={"model":
"us.anthropic.claude-opus-4-6-v1:0"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ with pytest.raises(ValueError, match="primary") as exc_info:
+ hook.get_conn()
+
+ assert "us.anthropic.claude-opus-4-6-v1:0" in str(exc_info.value)
+
+
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class",
autospec=True)
+ def test_import_error_from_recognized_provider_counts_as_prefixed(self,
mock_infer_provider_class):
+ """A recognized provider name whose optional dependency isn't
installed still counts
+ as a platform prefix -- ``infer_provider_class`` raises
``ImportError`` (not
+ ``ValueError``) for a name it recognizes but can't import.
+
+ Mutation canary: changing the ``except ImportError`` branch to
``return False``
+ makes this assert ``True`` fail.
+ """
+ mock_infer_provider_class.side_effect = ImportError("Please install
the 'azure' extra")
+
+ assert _has_recognized_provider_prefix("azure:foo") is True
+
+
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class",
autospec=True)
+ def test_value_error_from_unknown_provider_counts_as_bare(self,
mock_infer_provider_class):
+ """An unrecognized name raises ``ValueError`` and is treated as a bare
model name --
+ the counterpart to the ``ImportError`` case above, proving the two
exceptions are
+ told apart rather than both mapping to the same answer.
+
+ Mutation canary: changing the ``except ValueError`` branch to ``return
True``
+ makes this assert ``False`` fail.
+ """
+ mock_infer_provider_class.side_effect = ValueError("Unknown provider:
bogus")
+
+ assert _has_recognized_provider_prefix("bogus:foo") is False
+
+
+class TestPydanticAIHookFallback:
+ def test_no_fallback_returns_the_bare_model(self, registry,
infer_model_stub):
+ """Without a chain the resolved model is not wrapped at all."""
+ registry.add("primary", extra={"model": "openai:gpt-5.6-sol"})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ assert hook.get_conn() is infer_model_stub.models["openai:gpt-5.6-sol"]
+
+ def test_param_builds_chain_in_order(self, registry, infer_model_stub):
+ registry.add("primary", extra={"model": "openai:gpt-5.6-sol"})
+ registry.add("second", extra={"model": "anthropic:claude-opus-4-6"})
+ registry.add("third", extra={"model": "groq:llama-4"})
+
+ hook = PydanticAIHook(llm_conn_id="primary",
fallback_conn_ids=["second", "third"])
+ model = hook.get_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"],
+ ]
+
+ def test_chain_from_connection_extra(self, registry, infer_model_stub):
+ """A deployment manager can configure failover without touching Dag
code."""
+ registry.add(
+ "primary",
+ extra={"model": "openai:gpt-5.6-sol", "fallback_conn_ids":
["second"]},
+ )
+ registry.add("second", extra={"model": "anthropic:claude-opus-4-6"})
+
+ model = PydanticAIHook(llm_conn_id="primary").get_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"],
+ ]
+
+ def test_param_overrides_extra(self, registry, infer_model_stub):
+ registry.add(
+ "primary",
+ extra={"model": "openai:gpt-5.6-sol", "fallback_conn_ids":
["ignored"]},
+ )
+ registry.add("ignored", extra={"model": "groq:llama-4"})
+ registry.add("second", extra={"model": "anthropic:claude-opus-4-6"})
+
+ model = PydanticAIHook(llm_conn_id="primary",
fallback_conn_ids=["second"]).get_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models[1] is
infer_model_stub.models["anthropic:claude-opus-4-6"]
+
+ def test_empty_list_param_disables_the_extra_chain(self, registry,
infer_model_stub):
+ """``[]`` is an explicit opt-out, distinct from ``None`` meaning "read
the extra"."""
+ registry.add(
+ "primary",
+ extra={"model": "openai:gpt-5.6-sol", "fallback_conn_ids":
["second"]},
+ )
+ registry.add("second", extra={"model": "anthropic:claude-opus-4-6"})
+
+ model = PydanticAIHook(llm_conn_id="primary",
fallback_conn_ids=[]).get_conn()
+
+ assert model is infer_model_stub.models["openai:gpt-5.6-sol"]
+
+ def test_chain_can_span_providers(self, registry, infer_model_stub):
+ """Each connection resolves through its own hook class, so credentials
differ per hop."""
+ registry.add("primary", password="sk-openai", extra={"model":
"openai:gpt-5.6-sol"})
+ registry.add(
+ "bedrock_dr",
+ conn_type="pydanticai_bedrock",
+ hook_class=PydanticAIBedrockHook,
+ extra={
+ "model": "bedrock:us.anthropic.claude-opus-4-5",
+ "region_name": "us-east-1",
+ "aws_access_key_id": "AKIA-test",
+ "aws_secret_access_key": "secret",
+ },
+ )
+
+ with patch(
+
"airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class",
autospec=True
+ ) as mock_infer_provider_class:
+ mock_infer_provider_class.return_value =
MagicMock(return_value=MagicMock())
+ hook = PydanticAIHook(llm_conn_id="primary",
fallback_conn_ids=["bedrock_dr"])
+ model = hook.get_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models == [
+ infer_model_stub.models["openai:gpt-5.6-sol"],
+
infer_model_stub.models["bedrock:us.anthropic.claude-opus-4-5"],
+ ]
+
+ # Each hop is built by its own hook's field mapping: the primary
from
+ # password/host, the Bedrock hop from its extra.
+ assert infer_model_stub.provider_kwargs_for("openai:gpt-5.6-sol",
mock_infer_provider_class) == {
+ "api_key": "sk-openai"
+ }
+ assert infer_model_stub.provider_kwargs_for(
+ "bedrock:us.anthropic.claude-opus-4-5",
mock_infer_provider_class
+ ) == {
+ "region_name": "us-east-1",
+ "aws_access_key_id": "AKIA-test",
+ "aws_secret_access_key": "secret",
+ }
+
+ def test_bare_model_id_forwarded_to_fallback_without_own_model(self,
registry, infer_model_stub):
+ """A bare ``model_id`` flows to a fallback with none, qualified with
*that* fallback's platform."""
+ registry.add("primary", conn_type="pydanticai_azure",
hook_class=PydanticAIAzureHook)
+ registry.add("bedrock_dr", conn_type="pydanticai_bedrock",
hook_class=PydanticAIBedrockHook)
+
+ hook = PydanticAIAzureHook(
+ llm_conn_id="primary", model_id="gpt-5-nano",
fallback_conn_ids=["bedrock_dr"]
+ )
+ model = hook.get_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models == [
+ infer_model_stub.models["azure:gpt-5-nano"],
+ infer_model_stub.models["bedrock:gpt-5-nano"],
+ ]
+
+ def
test_bare_connection_model_forwarded_to_fallback_without_own_model(self,
registry, infer_model_stub):
+ """A bare model from the primary's own ``extra`` forwards like a
``model_id`` argument.
+
+ This is the connection-driven shape the docs lead with: neither the
model nor the chain
+ is named in Dag code, so forwarding the constructor argument alone
never fires.
+
+ Mutation canary: forwarding ``self.model_id`` rather than the
primary's configured name
+ makes ``get_conn()`` raise "No model specified for connection
'bedrock_dr'" here, because
+ ``model_id`` is ``None`` in this shape -- failing before either
assertion is reached.
+ """
+ registry.add(
+ "primary",
+ conn_type="pydanticai_azure",
+ hook_class=PydanticAIAzureHook,
+ extra={"model": "gpt-5-nano", "fallback_conn_ids": ["bedrock_dr"]},
+ )
+ registry.add("bedrock_dr", conn_type="pydanticai_bedrock",
hook_class=PydanticAIBedrockHook)
+
+ model = PydanticAIAzureHook(llm_conn_id="primary").get_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models == [
+ infer_model_stub.models["azure:gpt-5-nano"],
+ infer_model_stub.models["bedrock:gpt-5-nano"],
+ ]
+
+ def test_fallback_own_model_overrides_forwarded(self, registry,
infer_model_stub):
+ """A fallback's own ``model`` extra wins over anything forwarded from
the primary."""
+ registry.add("primary", conn_type="pydanticai_azure",
hook_class=PydanticAIAzureHook)
+ registry.add(
+ "bedrock_dr",
+ conn_type="pydanticai_bedrock",
+ hook_class=PydanticAIBedrockHook,
+ extra={"model": "bedrock:us.anthropic.claude-opus-4-5"},
+ )
+
+ hook = PydanticAIAzureHook(
+ llm_conn_id="primary", model_id="gpt-5-nano",
fallback_conn_ids=["bedrock_dr"]
+ )
+ model = hook.get_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models[1] is
infer_model_stub.models["bedrock:us.anthropic.claude-opus-4-5"]
+
+ def test_prefixed_model_id_not_forwarded_to_fallback(self, registry,
infer_model_stub):
+ """A prefixed ``model_id`` pins the primary's own platform and is
unusable on a fallback."""
+ registry.add("primary", extra={"model": "openai:gpt-5.6-sol"})
+ registry.add("second") # no model of its own
+
+ hook = PydanticAIHook(llm_conn_id="primary", model_id="openai:gpt-5",
fallback_conn_ids=["second"])
+ with pytest.raises(ValueError, match="No model specified for
connection 'second'"):
+ hook.get_conn()
+
+ def
test_bedrock_style_bare_model_id_forwarded_to_fallback_without_own_model(
+ self, registry, infer_model_stub
+ ):
+ """A primary's bare model id with an embedded ``:`` of its own is
still forwardable.
+
+ The forwarding-priority check in ``_resolve_own_model`` has to use the
same
+ "recognized provider prefix" test as ``_qualify_model_name`` -- not
the old
+ ``":" in forwarded_model_id`` check -- or a Bedrock-style bare id
(which contains a
+ ``:`` from its own version suffix, not a platform prefix) would be
wrongly treated as
+ already-pinned and never forwarded. Using a different fallback
platform (Azure) makes
+ the resolved keys distinguishable so a wrong-prefix regression cannot
hide behind two
+ identical strings.
+
+ Mutation canary: reverting the forwarding check to ``":" not in
forwarded_model_id``
+ makes the fallback treat ``"us.anthropic.claude-opus-4-6-v1:0"`` as
already-pinned and
+ skip it, so ``hook.get_conn()`` itself raises "No model specified for
connection
+ 'azure_dr'" instead of returning -- failing this test before the
identity assertion is
+ even reached.
+ """
+ registry.add("primary", conn_type="pydanticai_bedrock",
hook_class=PydanticAIBedrockHook)
+ registry.add("azure_dr", conn_type="pydanticai_azure",
hook_class=PydanticAIAzureHook)
+
+ hook = PydanticAIBedrockHook(
+ llm_conn_id="primary",
+ model_id="us.anthropic.claude-opus-4-6-v1:0",
+ fallback_conn_ids=["azure_dr"],
+ )
+ model = hook.get_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models == [
+
infer_model_stub.models["bedrock:us.anthropic.claude-opus-4-6-v1:0"],
+ infer_model_stub.models["azure:us.anthropic.claude-opus-4-6-v1:0"],
+ ]
+
+ def test_non_pydanticai_fallback_raises(self, registry, infer_model_stub):
+ """``BaseHook.get_hook`` dispatches on conn_type alone and can return
anything."""
+ registry.add("primary", extra={"model": "openai:gpt-5.6-sol"})
+ registry.add("wrong_type", conn_type="langchain")
+ registry.hook_classes["wrong_type"] = MagicMock # type:
ignore[assignment]
+
+ hook = PydanticAIHook(llm_conn_id="primary",
fallback_conn_ids=["wrong_type"])
+ with pytest.raises(ValueError, match="not a PydanticAIHook"):
+ hook.get_conn()
+
+ def test_nested_chain_raises(self, registry, infer_model_stub):
+ registry.add("primary", extra={"model": "openai:gpt-5.6-sol"})
+ registry.add(
+ "second",
+ extra={"model": "anthropic:claude-opus-4-6", "fallback_conn_ids":
["third"]},
+ )
+ registry.add("third", extra={"model": "groq:llama-4"})
+
+ hook = PydanticAIHook(llm_conn_id="primary",
fallback_conn_ids=["second"])
+ with pytest.raises(ValueError, match="second.*not resolved
recursively"):
+ hook.get_conn()
+
+ @pytest.mark.parametrize(
+ ("fallback_conn_ids", "match"),
+ [
+ pytest.param(["second", "second"], "more than once",
id="repeated-fallback"),
+ pytest.param(["primary"], "as one of its own fallbacks",
id="primary-repeated"),
+ ],
+ )
+ def test_duplicate_conn_id_raises(self, registry, infer_model_stub,
fallback_conn_ids, match):
+ registry.add("primary", extra={"model": "openai:gpt-5.6-sol"})
+ registry.add("second", extra={"model": "anthropic:claude-opus-4-6"})
+
+ hook = PydanticAIHook(llm_conn_id="primary",
fallback_conn_ids=fallback_conn_ids)
+ with pytest.raises(ValueError, match=match):
+ hook.get_conn()
+
+ @pytest.mark.parametrize(
+ "fallback_conn_ids",
+ [
+ pytest.param("second,third", id="comma-separated-string"),
+ pytest.param(["second", ""], id="empty-entry"),
Review Comment:
Settled towards dropping them, and the case is rewritten rather than
removed: `["anthropic_prod", ""]` now asserts a working one-entry chain.
--
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]
