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 6e990ea1927 Make common.ai hooks report misconfiguration instead of
failing silently (#74451)
6e990ea1927 is described below
commit 6e990ea1927c71aedb34e842d27a4780049d297f
Author: Wei Lee <[email protected]>
AuthorDate: Mon Oct 26 14:17:57 2026 +0800
Make common.ai hooks report misconfiguration instead of failing silently
(#74451)
---
providers/common/ai/docs/changelog.rst | 32 ++++++
providers/common/ai/docs/connections/mcp.rst | 3 +-
.../ai/docs/connections/pydantic_ai_azure.rst | 13 ++-
.../ai/docs/connections/pydantic_ai_vertex.rst | 10 +-
providers/common/ai/docs/provider_fallback.rst | 16 +--
providers/common/ai/docs/self_hosted_models.rst | 10 +-
providers/common/ai/docs/toolsets/index.rst | 9 +-
providers/common/ai/provider.yaml | 2 +-
.../ai/example_dags/example_llm_retry_policy.py | 7 +-
.../providers/common/ai/get_provider_info.py | 2 +-
.../airflow/providers/common/ai/hooks/langchain.py | 29 ++++-
.../providers/common/ai/hooks/llamaindex.py | 46 ++++++--
.../src/airflow/providers/common/ai/hooks/mcp.py | 55 +++++++--
.../providers/common/ai/hooks/pydantic_ai.py | 46 ++++----
.../airflow/providers/common/ai/policies/retry.py | 5 +-
.../tests/unit/common/ai/hooks/test_langchain.py | 74 ++++++++++++
.../tests/unit/common/ai/hooks/test_llamaindex.py | 84 ++++++++++++++
.../ai/tests/unit/common/ai/hooks/test_mcp.py | 128 +++++++++++++++++++++
.../tests/unit/common/ai/hooks/test_pydantic_ai.py | 71 ++++++++++--
.../ai/tests/unit/common/ai/policies/test_retry.py | 65 ++++++++---
.../tests/unit/common/ai/test_provider_metadata.py | 7 +-
.../unit/snowflake/hooks/test_cortex_model.py | 18 ++-
22 files changed, 618 insertions(+), 114 deletions(-)
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index 4093a86db8f..ee9c5ac9552 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -48,6 +48,38 @@ Changelog
``"PT1M30S"`` for a ``timedelta``. Tuples and sets in that model now arrive
as lists. See
:doc:`structured_output`.
+.. warning::
+ The following changes are breaking for code written against 1.0.0:
+
+ * ``PydanticAIHook``: when a connection supplies credentials (password or
host) and the provider class
+ rejects the resulting keyword arguments (``TypeError``), the hook now
raises ``ValueError``. In 1.0.0
+ it logged a warning and fell back to environment-variable authentication,
which could authenticate
+ as a different identity than the connection. Fix the connection fields for
that provider, or remove
+ the credentials from the connection if environment-variable authentication
is what you want.
+ * ``airflow.providers.common.ai.policies.retry``: importing it without the
Airflow ``RetryPolicy``
+ support (Airflow older than 3.3) now raises
``AirflowOptionalProviderFeatureException`` instead of
+ ``ImportError``. It is not an ``ImportError`` subclass (it derives from
``AirflowException``), so a
+ Dag that guards the import with ``except ImportError:`` as in the 1.0.0
example Dag now fails to
+ parse on those versions. Catch both, importing the exception from
+ ``airflow.providers.common.compat.sdk`` (works on Airflow 2.11 and 3.x)::
+
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+ try:
+ from airflow.providers.common.ai.policies.retry import LLMRetryPolicy
+ except (ImportError, AirflowOptionalProviderFeatureException):
+ LLMRetryPolicy = None
+
+.. note::
+ ``LlamaIndexHook`` and ``LangChainHook`` (both experimental, see
:doc:`stability`):
+ passing ``embed_model`` and ``llm_model`` positionally is deprecated. In
1.0.0 they were the third and
+ fourth positional parameters (after ``llm_conn_id`` and ``embed_conn_id``),
and the two hooks listed
+ them in opposite order (``embed_model`` then ``llm_model`` for
``LlamaIndexHook``, ``llm_model`` then
+ ``embed_model`` for ``LangChainHook``). Positional calls still work and emit
an
+ ``AirflowProviderDeprecationWarning``, but support will be removed in a
future release. Pass them by
+ keyword, e.g. ``LlamaIndexHook(llm_model="gpt-5")``, which also avoids the
meanings swapping between
+ the two hooks.
+
1.1.0
.....
diff --git a/providers/common/ai/docs/connections/mcp.rst
b/providers/common/ai/docs/connections/mcp.rst
index dc51df1cada..04cb73dcd34 100644
--- a/providers/common/ai/docs/connections/mcp.rst
+++ b/providers/common/ai/docs/connections/mcp.rst
@@ -70,7 +70,8 @@ Environment (Extra field)
Timeout (Extra field)
Connection init timeout in seconds for the ``stdio`` transport. Ignored
- for ``http``/``sse``. Default: ``10``.
+ for ``http``/``sse``. Must be a non-negative finite number; ``0`` or
``null``
+ disables the timeout. Default: ``10``.
Examples: ``30``
diff --git a/providers/common/ai/docs/connections/pydantic_ai_azure.rst
b/providers/common/ai/docs/connections/pydantic_ai_azure.rst
index 0332412d3b4..d26c416bca8 100644
--- a/providers/common/ai/docs/connections/pydantic_ai_azure.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai_azure.rst
@@ -52,10 +52,15 @@ Model
this connection type's own platform, so nothing else needs naming it
explicitly. Writing the ``azure:`` prefix yourself has the same effect and
is
still accepted. A name prefixed with a *different*, recognized platform
(e.g.
- ``openai:gpt-5``) is used verbatim instead, pinning that platform and
- bypassing Azure OpenAI entirely -- a name is only treated as already
prefixed
- when the segment before its first ``:`` is itself a real pydantic-ai
- provider, not merely present.
+ ``openai:gpt-5``) is used verbatim instead -- a name is only treated as
+ already prefixed when the segment before its first ``:`` is itself a real
+ pydantic-ai provider, not merely present. That other platform's provider
+ class is built from the fields this connection forwards (``API Key``,
+ ``Azure Endpoint`` as ``azure_endpoint``, ``API Version``), so it works
only
+ if it accepts all of them. ``OpenAIProvider`` does not accept
+ ``azure_endpoint`` or ``api_version``, so ``openai:`` with an endpoint or
API
+ version set fails with a ``ValueError`` naming the fields passed; it never
+ silently falls back to environment-variable credentials.
API Key (Password field)
The Azure OpenAI API key.
diff --git a/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
b/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
index 03eb1290da4..3cd9e20852a 100644
--- a/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
@@ -57,8 +57,12 @@ Model
mean Vertex AI Express Mode credentials (see "Credentials" below), so its
presence alone cannot tell the two platforms apart. To reach the
Generative Language API instead, prefix the model explicitly with
- ``google:`` -- that spelling routes to a different provider regardless of
- which fields this connection sets.
+ ``google:``. The ``google:`` provider accepts only the ``API Key`` and
+ ``Custom Endpoint URL`` fields (without an API key it reads
``GOOGLE_API_KEY``), so on this
+ connection type that spelling must not be combined with Vertex-only fields
+ (``GCP Project``, ``Location / Region``, ``Service Account Info``): the
+ connection fails with a ``ValueError`` instead of silently falling back to
+ environment-variable credentials.
GCP Project
Google Cloud project ID. Falls back to the ``GOOGLE_CLOUD_PROJECT``
@@ -87,7 +91,7 @@ API Key
``project`` / ``location`` / ``service_account_info`` (those select the
credentials/ADC path instead, which takes precedence and nulls the API
key). For the Generative Language API
- (non-Vertex, API-key-only), use the ``google:`` prefix on the generic
+ (non-Vertex; API key or ``GOOGLE_API_KEY``), use the ``google:`` prefix on
the generic
:doc:`pydantic_ai` connection instead.
Service Account Info
diff --git a/providers/common/ai/docs/provider_fallback.rst
b/providers/common/ai/docs/provider_fallback.rst
index a8a74574b5c..7b2821c1054 100644
--- a/providers/common/ai/docs/provider_fallback.rst
+++ b/providers/common/ai/docs/provider_fallback.rst
@@ -198,16 +198,12 @@ Two checks, neither of which requires waiting for a real
outage:
*Test the connection.* ``test_connection`` on the primary resolves every
connection in the
chain, so a fallback with a missing ``model`` or an unknown connection ID is
reported by name
-there rather than discovered mid-incident. Credential fields a provider class
rejects with a
-``TypeError`` are caught by the hook, which retries with the env-var-based
provider
-constructor and logs a warning either way; if the required env var is also
missing, that
-retry raises ``pydantic_ai.exceptions.UserError``, which ``test_connection``
does surface
-since it wraps the whole resolution in a broad exception handler. What it
cannot show is the
-opposite case: the env var *is* set on the worker, the retry quietly succeeds,
and
-``test_connection`` reports success even though the credentials you configured
on the
-connection were silently ignored -- check the logs for that warning rather
than relying on
-``test_connection`` alone. It also does not call the provider, so a
well-formed but revoked
-key still passes -- that is what the drill below is for.
+there rather than discovered mid-incident. Fields a provider class cannot be
built from (a
+``TypeError``) fail the resolution with a ``ValueError`` naming the
connection, the provider and
+the fields passed, with the provider's ``TypeError`` text -- the hook never
falls back to
+environment-variable auth, which would authenticate as a different identity --
and
+``test_connection`` reports it as a failure. It does not call the provider, so
a well-formed
+but revoked key still passes -- that is what the drill below is for.
*Drill it.* Point the primary at an endpoint nothing listens on and run the
Dag. The task
should still succeed, and the run summary in its log names the model that
answered:
diff --git a/providers/common/ai/docs/self_hosted_models.rst
b/providers/common/ai/docs/self_hosted_models.rst
index f454d0a2a2b..ce449cd65fa 100644
--- a/providers/common/ai/docs/self_hosted_models.rst
+++ b/providers/common/ai/docs/self_hosted_models.rst
@@ -317,11 +317,11 @@ its OpenAI provider, and there's no separate ``ollama``
extra.
``LiteLLMProvider`` is an exception: its constructor takes ``api_base``
instead of ``base_url``, so passing ``host`` raises a ``TypeError`` that
- the hook catches, logs a warning, and falls back to environment-variable
- auth, ignoring ``host`` entirely. Point a LiteLLM proxy's OpenAI-compatible
- endpoint via ``openai:<model>`` instead -- this is the same proxy role
- used in the gateway section above, just addressed through the working
- prefix rather than the rejected ``litellm:`` one.
+ the hook re-raises as a ``ValueError``: the connection is rejected rather
+ than falling back to environment-variable auth. Point a LiteLLM proxy's
+ OpenAI-compatible endpoint via ``openai:<model>`` instead -- this is the
+ same proxy role used in the gateway section above, just addressed through
+ the working prefix rather than the rejected ``litellm:`` one.
.. _self_hosted_models_other_servers:
diff --git a/providers/common/ai/docs/toolsets/index.rst
b/providers/common/ai/docs/toolsets/index.rst
index 8be3079b1d6..1e9351cbb4b 100644
--- a/providers/common/ai/docs/toolsets/index.rst
+++ b/providers/common/ai/docs/toolsets/index.rst
@@ -113,10 +113,11 @@ Every toolset imports from the
``airflow.providers.common.ai.toolsets`` package
SQLToolset,
)
-``SQLToolset`` and ``DataFusionToolset`` depend on optional extras. If the
extra is not
-installed, the first access to the name raises
``AirflowOptionalProviderFeatureException``.
-``MCPToolset`` imports without its extra; the missing dependency raises
``ImportError``
-when the toolset first connects to the MCP server.
+``SQLToolset``, ``DataFusionToolset`` and ``MCPToolset`` depend on optional
extras. If the
+extra is not installed, they raise
``AirflowOptionalProviderFeatureException``. For
+``SQLToolset`` and ``DataFusionToolset`` this happens on the first access to
the name.
+``MCPToolset`` imports without its extra; it raises only when the toolset
first connects
+to the MCP server.
Every toolset here implements pydantic-ai's
`AbstractToolset <https://ai.pydantic.dev/toolsets/>`__ interface, so it works
in any
diff --git a/providers/common/ai/provider.yaml
b/providers/common/ai/provider.yaml
index a985704dbea..814a850a20b 100644
--- a/providers/common/ai/provider.yaml
+++ b/providers/common/ai/provider.yaml
@@ -426,7 +426,7 @@ connection-types:
- 'null'
vertexai:
label: Force Vertex AI Mode
- description: "Ignored (kept for compatibility); mode is now selected
via the Model field's prefix."
+ description: "Deprecated and ignored; the Model field's prefix selects
the mode."
schema:
type:
- boolean
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
index 8a9f17ce845..3da0d71b1fd 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_retry_policy.py
@@ -36,7 +36,7 @@ from __future__ import annotations
from dataclasses import replace
from datetime import timedelta
-from airflow.providers.common.compat.sdk import dag, task
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, dag, task
try:
from airflow.providers.common.ai.policies.retry import (
@@ -181,6 +181,7 @@ try:
task_ambiguous()
example_llm_retry_policy_classifier()
-except ImportError:
- # RetryPolicy requires Airflow 3.3+; example DAG is skipped on older
versions.
+except (ImportError, AirflowOptionalProviderFeatureException):
+ # RetryPolicy requires Airflow 3.3+; policies/retry.py raises
AirflowOptionalProviderFeatureException
+ # (not an ImportError) on older versions, so the example Dag is skipped
there.
pass
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
index 19775160600..97364eb0c50 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
@@ -341,7 +341,7 @@ def get_provider_info():
},
"vertexai": {
"label": "Force Vertex AI Mode",
- "description": "Ignored (kept for compatibility); mode
is now selected via the Model field's prefix.",
+ "description": "Deprecated and ignored; the Model
field's prefix selects the mode.",
"schema": {"type": ["boolean", "null"]},
},
"api_key": {
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py
index 1d6122359a1..8c7b37c777f 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py
@@ -18,8 +18,10 @@
from __future__ import annotations
+import warnings
from typing import TYPE_CHECKING, Any
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.compat.sdk import BaseHook
if TYPE_CHECKING:
@@ -68,12 +70,13 @@ class LangChainHook(BaseHook):
single hook instance.
:param llm_model: Chat model identifier in ``provider:name`` format
(e.g. ``"openai:gpt-5"``, ``"anthropic:claude-sonnet-5"``).
- Overrides ``extra["model"]`` on the connection.
+ Overrides ``extra["model"]`` on the connection. Passing it positionally
+ (1.0.0 order: ``llm_model, embed_model``) is deprecated.
:param embed_model: Embedding model identifier in ``provider:name`` format
(e.g. ``"openai:text-embedding-3-small"``). Overrides
``extra["embed_model"]`` on the connection. When ``embedding_kwargs``
supplies ``provider`` explicitly, use a model name without the provider
- prefix.
+ prefix. Passing it positionally (1.0.0 order: ``llm_model,
embed_model``) is deprecated.
:param embedding_kwargs: Additional keyword arguments to pass to the
embedding
model constructor without filtering. Values can override hook-provided
settings, including the endpoint and credentials. In particular,
@@ -96,12 +99,32 @@ class LangChainHook(BaseHook):
self,
llm_conn_id: str | None = None,
embed_conn_id: str | None = None,
+ *args: Any,
llm_model: str | None = None,
embed_model: str | None = None,
- *,
embedding_kwargs: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
+ # ``*args`` only exists to keep the 1.0.0 positional order working:
``llm_model``, ``embed_model``.
+ if len(args) > 2:
+ raise TypeError(
+ f"{type(self).__name__}.__init__() takes from 1 to 5
positional arguments "
+ f"but {3 + len(args)} were given"
+ )
+ legacy_values = dict(zip(("llm_model", "embed_model"), args))
+ keyword_values = {"llm_model": llm_model, "embed_model": embed_model}
+ for name in legacy_values:
+ # A keyword of None counts as not provided, so ``name=None`` plus
a positional keeps the positional.
+ if keyword_values[name] is not None:
+ raise TypeError(f"{type(self).__name__}.__init__() got
multiple values for argument '{name}'")
+ warnings.warn(
+ f"Parameter `{name}` passed positionally is deprecated. Use
`{name}=...` instead; "
+ f"positional support will be removed in a future release.",
+ category=AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ llm_model = legacy_values.get("llm_model", llm_model)
+ embed_model = legacy_values.get("embed_model", embed_model)
super().__init__(**kwargs)
# Resolve at runtime so a future subclass with its own
# `default_conn_name` is honoured. A bare `llm_conn_id: str =
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py
index 6a3b01f9edd..ede87d85922 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py
@@ -19,8 +19,10 @@
from __future__ import annotations
import inspect
+import warnings
from typing import TYPE_CHECKING, Any
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.compat.sdk import (
AirflowOptionalProviderFeatureException,
BaseHook,
@@ -75,7 +77,8 @@ class LlamaIndexHook(BaseHook):
still fails on the model name, not on connectivity.
``get_embedding_model()`` raises immediately at construction;
``get_llm()`` defers the error until the first call that reads
- ``.metadata`` (``.chat()`` / ``.complete()``).
+ ``.metadata`` (``.chat()`` / ``.complete()``); ``test_connection()``
+ reads it, so it reports such a model name as a failure.
Connection fields:
@@ -95,10 +98,10 @@ class LlamaIndexHook(BaseHook):
embedding provider. Falls back to ``llm_conn_id`` when not set.
:param embed_model: Embedding model name (e.g.
``"text-embedding-3-small"``). Overrides ``extra["embed_model"]``
- on the connection.
+ on the connection. Passing it positionally (1.0.0 order:
``embed_model, llm_model``) is deprecated.
:param llm_model: LLM model name (e.g. ``"gpt-5"``). Overrides
``extra["llm_model"]`` on the connection. Required when calling
- :meth:`get_llm`.
+ :meth:`get_llm`. Passing it positionally (1.0.0 order: ``embed_model,
llm_model``) is deprecated.
:param embedding_kwargs: Additional keyword arguments to pass to the
embedding
model constructor without filtering. Connection ``api_key`` and
``api_base``
values take precedence at the top level, but nested options supported
by the
@@ -115,12 +118,32 @@ class LlamaIndexHook(BaseHook):
self,
llm_conn_id: str | None = None,
embed_conn_id: str | None = None,
- embed_model: str | None = None,
+ *args: Any,
llm_model: str | None = None,
- *,
+ embed_model: str | None = None,
embedding_kwargs: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
+ # ``*args`` only exists to keep the 1.0.0 positional order working:
``embed_model``, ``llm_model``.
+ if len(args) > 2:
+ raise TypeError(
+ f"{type(self).__name__}.__init__() takes from 1 to 5
positional arguments "
+ f"but {3 + len(args)} were given"
+ )
+ legacy_values = dict(zip(("embed_model", "llm_model"), args))
+ keyword_values = {"embed_model": embed_model, "llm_model": llm_model}
+ for name in legacy_values:
+ # A keyword of None counts as not provided, so ``name=None`` plus
a positional keeps the positional.
+ if keyword_values[name] is not None:
+ raise TypeError(f"{type(self).__name__}.__init__() got
multiple values for argument '{name}'")
+ warnings.warn(
+ f"Parameter `{name}` passed positionally is deprecated. Use
`{name}=...` instead; "
+ f"positional support will be removed in a future release.",
+ category=AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ embed_model = legacy_values.get("embed_model", embed_model)
+ llm_model = legacy_values.get("llm_model", llm_model)
super().__init__(**kwargs)
# Resolve at runtime so a future per-vendor subclass with its own
# ``default_conn_name`` is honoured.
@@ -233,15 +256,16 @@ class LlamaIndexHook(BaseHook):
def test_connection(self) -> tuple[bool, str]:
"""
- Test connection by resolving the LLM.
+ Test connection by resolving the LLM and reading its metadata.
- Validates that the model identifier is valid and the provider can be
- instantiated with the supplied credentials. Does NOT make an LLM API
- call -- that would be expensive and fail for reasons unrelated to
- connectivity (quotas, billing, rate limits).
+ Reading ``.metadata`` makes LlamaIndex validate the model name against
its
+ OpenAI-only allowlist, so an Ollama or vLLM model name fails here
instead of at the
+ first ``.chat()`` call. Does NOT make an LLM API call -- that would be
expensive and
+ fail for reasons unrelated to connectivity (quotas, billing, rate
limits), so a
+ revoked API key or an unreachable **host** still passes.
"""
try:
- self.get_llm()
+ self.get_llm().metadata
return True, "Model resolved successfully."
except Exception as e:
return False, str(e)
diff --git a/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
index 46c81d3fcc0..810b0fccbc4 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py
@@ -16,9 +16,14 @@
# under the License.
from __future__ import annotations
+import json
from typing import TYPE_CHECKING, Any
-from airflow.providers.common.compat.sdk import BaseHook, mask_secret
+from airflow.providers.common.compat.sdk import (
+ AirflowOptionalProviderFeatureException,
+ BaseHook,
+ mask_secret,
+)
if TYPE_CHECKING:
from collections.abc import Callable
@@ -203,10 +208,10 @@ class MCPHook(BaseHook):
try:
from fastmcp.client.transports import SSETransport,
StdioTransport, StreamableHttpTransport
from pydantic_ai.mcp import MCPToolset
- except ImportError:
- raise ImportError(
+ except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(
'MCP support requires the `mcp` package. Install it with: pip
install "pydantic-ai-slim[mcp]"'
- )
+ ) from e
conn = self.get_connection(self.mcp_conn_id)
extra = conn.extra_dejson
@@ -226,13 +231,47 @@ class MCPHook(BaseHook):
raise ValueError(
f"Connection {self.mcp_conn_id!r} requires 'command' in
extra for stdio transport."
)
- args = extra.get("args", [])
- if isinstance(args, str):
- args = [args]
+ args = extra.get("args")
+ if args is None:
+ args = []
+ elif isinstance(args, str):
+ if args.strip().startswith("["):
+ try:
+ args = json.loads(args)
+ except ValueError as e:
+ raise ValueError(
+ f"'args' in extra for connection
{self.mcp_conn_id!r} looks like a JSON array "
+ "but is not valid JSON."
+ ) from e
+ else:
+ args = [args]
+ if not isinstance(args, list):
+ raise ValueError(
+ f"'args' in extra for connection {self.mcp_conn_id!r} must
be a string, a JSON "
+ f"array of strings or a list of strings, got
{type(args).__name__}."
+ )
+ for index, arg in enumerate(args):
+ if not isinstance(arg, str):
+ raise ValueError(
+ f"'args' in extra for connection {self.mcp_conn_id!r}
must contain only strings, "
+ f"got {type(arg).__name__} at index {index}."
+ )
timeout = extra.get("timeout", 10)
+ # ``None`` and ``0`` are both valid "no timeout" values, but
fastmcp replaces
+ # ``init_timeout=None`` with its global ``client_init_timeout``
setting, so ``None``
+ # is mapped to ``0`` to keep ``FASTMCP_CLIENT_INIT_TIMEOUT`` from
re-enabling a timeout.
+ if timeout is not None and (
+ isinstance(timeout, bool)
+ or not isinstance(timeout, (int, float))
+ or not 0 <= timeout < float("inf")
+ ):
+ raise ValueError(
+ f"'timeout' in extra for connection {self.mcp_conn_id!r}
must be a non-negative "
+ f"finite number of seconds, got {timeout!r}."
+ )
toolset = MCPToolset(
StdioTransport(command=command, args=args,
env=self._stdio_env(extra)),
- init_timeout=timeout,
+ init_timeout=0 if timeout is None else timeout,
)
else:
raise ValueError(
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 146560881e4..741dfe5a30d 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
@@ -25,7 +25,7 @@ from pydantic_ai import Agent
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, infer_provider_class
+from pydantic_ai.providers import infer_provider_class
from airflow.providers.common.ai.observability import
genai_instrumentation_settings
from airflow.providers.common.ai.utils.toolset_base import MaskingCapability
@@ -451,13 +451,12 @@ class PydanticAIHook(BaseHook):
def _provider_factory(pname: str) -> Any:
try:
return infer_provider_class(pname)(**_kwargs)
- except TypeError:
- self.log.warning(
- "Provider '%s' rejected kwargs %s; falling back to
env-var auth",
- pname,
- list(_kwargs),
- )
- return infer_provider(pname)
+ except TypeError as e:
+ raise ValueError(
+ f"Could not build provider '{pname}' from the fields
on connection "
+ f"'{self.llm_conn_id}' (fields passed: {',
'.join(_kwargs)}): {e}. "
+ "Not falling back to environment-variable auth, which
would authenticate as a different identity."
+ ) from e
return infer_model(model_name, provider_factory=_provider_factory)
@@ -678,11 +677,10 @@ class PydanticAIHook(BaseHook):
"""
Test connection by resolving the model.
- A success here can come from this connection's own credentials, or --
when a
- provider class rejects them with a ``TypeError`` -- from a silent
retry against
- the standard environment variables, which ignores those credentials
entirely.
- See :doc:`/provider_fallback`'s *Verifying a chain* section for how to
tell the
- two apart. Does NOT make an LLM API call — that would be expensive and
fail for
+ A provider class that cannot be built from the connection's fields is
reported as a failure
+ naming the provider, the fields passed and the provider's
``TypeError`` text.
+
+ Does NOT make an LLM API call — that would be expensive and fail for
reasons unrelated to connectivity (quotas, billing, rate limits).
Every connection in ``fallback_conn_ids`` is resolved too, so a
@@ -691,9 +689,9 @@ class PydanticAIHook(BaseHook):
"""
try:
self.get_conn()
- return True, "Model resolved successfully."
except Exception as e:
return False, str(e)
+ return True, "Model resolved successfully."
class PydanticAIAzureHook(PydanticAIHook):
@@ -848,7 +846,7 @@ class PydanticAIVertexHook(PydanticAIHook):
1. ``service_account_info`` (JSON object) in ``extra``
— loaded into a ``google.auth.credentials.Credentials``
- object and passed as ``credentials`` to ``GoogleProvider``.
+ object and passed as ``credentials`` to ``GoogleCloudProvider``.
2. ``api_key`` in ``extra`` — for Generative Language API (non-Vertex) or
Vertex API-key auth.
3. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
@@ -880,8 +878,11 @@ class PydanticAIVertexHook(PydanticAIHook):
above), so its presence alone cannot tell the two platforms apart; guessing
would risk silently authenticating against the wrong endpoint. To use the
Generative Language API, set an explicit ``google:``-prefixed model id (on
- the hook or the connection's ``model`` extra) -- that spelling already
- works today.
+ the hook or the connection's ``model`` extra). ``GoogleProvider`` accepts
only
+ ``api_key`` and ``base_url`` (without a key it reads ``GOOGLE_API_KEY``),
so on
+ this connection type that spelling must not set Vertex-only fields
+ (``project``, ``location``, ``service_account_info``): they are rejected
+ (fail closed).
:param llm_conn_id: Airflow connection ID.
:param model_id: Model identifier, e.g.
``"google-cloud:gemini-2.5-flash"``. A
@@ -919,7 +920,7 @@ class PydanticAIVertexHook(PydanticAIHook):
sa_info = extra.get("service_account_info")
kwargs: dict[str, Any] = {}
- # Direct GoogleProvider scalar kwargs.
+ # Scalar kwargs: GoogleCloudProvider takes all four; GoogleProvider
only api_key/base_url.
for _key in ("api_key", "project", "location", "base_url"):
if extra.get(_key):
kwargs[_key] = extra[_key]
@@ -927,11 +928,10 @@ class PydanticAIVertexHook(PydanticAIHook):
# "vertexai" predates pydantic-ai splitting GoogleProvider (Generative
Language API)
# from GoogleCloudProvider (Vertex AI, which hardcodes vertexai=True
internally and
# accepts no such constructor kwarg) in pydantic/pydantic-ai#5336.
Forwarding it would
- # raise TypeError, which the base hook's `except TypeError` in
get_conn() would then
- # catch: it logs a warning and falls back to env-var auth rather than
raising --
- # authenticating as the wrong identity (with *all* other kwargs
discarded). Accept
- # the field for backward compatibility but never forward it: which API
is used is
- # now controlled by the model prefix.
+ # raise TypeError, which the base hook's `except TypeError` in
get_conn() would turn
+ # into a ValueError rejecting the whole connection. Accept the field
for backward
+ # compatibility but never forward it: which API is used is now
controlled by the
+ # model prefix.
if extra.get("vertexai") is not None:
self.log.warning(
"The 'vertexai' connection field is ignored; Vertex AI vs.
Generative Language "
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
index 637679b1177..4f14db9dd0d 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
@@ -57,6 +57,7 @@ from airflow.providers.common.ai.utils.decision import (
threshold_for,
)
from airflow.providers.common.ai.utils.masking import mask_secrets
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
try:
from airflow.sdk.definitions.retry_policy import (
@@ -66,7 +67,7 @@ try:
RetryPolicy,
)
except ImportError:
- raise ImportError(
+ raise AirflowOptionalProviderFeatureException(
"LLMRetryPolicy requires Airflow 3.3+ which includes RetryPolicy
support. "
"Please upgrade apache-airflow-core."
) from None
@@ -317,7 +318,7 @@ class _ModelRetryPolicy(RetryPolicy):
def _hook(self) -> PydanticAIHook:
from airflow.providers.common.ai.hooks.pydantic_ai import
PydanticAIHook
- return PydanticAIHook(llm_conn_id=self.llm_conn_id,
model_id=self.model_id)
+ return PydanticAIHook.get_hook(self.llm_conn_id,
hook_params={"model_id": self.model_id})
def _prompt(self, exception: BaseException, try_number: int, max_tries:
int) -> str:
if self.include_traceback:
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py
index 58fb66f1bf2..17feedcdca1 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_langchain.py
@@ -17,10 +17,12 @@
from __future__ import annotations
import sys
+import warnings
from unittest.mock import MagicMock, patch
import pytest
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.ai.hooks.langchain import LangChainHook
@@ -72,6 +74,78 @@ class TestLangChainHookInit:
assert hook.llm_conn_id == "my_conn"
assert hook.embed_conn_id == "my_conn"
+ @pytest.mark.parametrize(
+ ("positional", "expected"),
+ [
+ pytest.param(("llm-pos",), {"llm_model": "llm-pos", "embed_model":
None}, id="one"),
+ pytest.param(
+ ("llm-pos", "embed-pos"),
+ {"llm_model": "llm-pos", "embed_model": "embed-pos"},
+ id="two",
+ ),
+ ],
+ )
+ def test_positional_model_emits_deprecation_warning(self, positional,
expected):
+ with pytest.warns(AirflowProviderDeprecationWarning) as record:
+ hook = LangChainHook("c1", "c2", *positional)
+ assert len(record) == len(positional)
+ assert "Parameter `llm_model` passed positionally is deprecated" in
str(record[0].message)
+ assert hook.llm_conn_id == "c1"
+ assert hook.embed_conn_id == "c2"
+ assert hook.llm_model == expected["llm_model"]
+ assert hook.embed_model == expected["embed_model"]
+
+ def test_positional_warning_stacklevel_points_to_caller(self):
+ with pytest.warns(AirflowProviderDeprecationWarning) as record:
+ LangChainHook("c1", "c2", "llm-pos")
+ assert record[0].filename == __file__
+
+ def test_one_warning_per_positional_argument(self):
+ with pytest.warns(AirflowProviderDeprecationWarning) as record:
+ LangChainHook("c1", "c2", "llm-pos", "embed-pos")
+ messages = [str(w.message) for w in record]
+ assert len(messages) == 2
+ assert "`llm_model`" in messages[0]
+ assert "`embed_model`" in messages[1]
+
+ def test_keyword_model_arguments_do_not_warn(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", AirflowProviderDeprecationWarning)
+ hook = LangChainHook("c1", "c2", llm_model="llm-pos",
embed_model="embed-pos")
+ assert hook.llm_model == "llm-pos"
+ assert hook.embed_model == "embed-pos"
+
+ def test_connection_positional_arguments_do_not_warn(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", AirflowProviderDeprecationWarning)
+ hook = LangChainHook("c1", "c2")
+ assert hook.embed_model is None
+ assert hook.llm_model is None
+
+ def test_too_many_positional_arguments_raise_type_error(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", AirflowProviderDeprecationWarning)
+ with pytest.raises(TypeError, match="positional argument"):
+ LangChainHook("c1", "c2", "llm-pos", "embed-pos", "extra")
+
+ def test_positional_and_keyword_same_parameter_raise_type_error(self):
+ with pytest.raises(TypeError, match="multiple values for argument
'llm_model'"):
+ LangChainHook("c1", "c2", "llm-pos", llm_model="other")
+
+ def
test_second_positional_and_keyword_same_parameter_raise_type_error(self):
+ with (
+ pytest.warns(AirflowProviderDeprecationWarning),
+ pytest.raises(TypeError, match="multiple values for argument
'embed_model'"),
+ ):
+ LangChainHook("c1", "c2", "L", "E", embed_model="X")
+
+ def test_positional_and_keyword_different_parameters_combine(self):
+ with pytest.warns(AirflowProviderDeprecationWarning) as record:
+ hook = LangChainHook("c1", "c2", "llm-pos",
embed_model="embed-pos")
+ assert len(record) == 1
+ assert hook.llm_model == "llm-pos"
+ assert hook.embed_model == "embed-pos"
+
def test_explicit_separate_conns_and_models(self):
hook = LangChainHook(
llm_conn_id="chat_conn",
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py
index ac0fcfd71cb..3274519629c 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py
@@ -16,10 +16,12 @@
# under the License.
from __future__ import annotations
+import warnings
from unittest.mock import MagicMock, patch
import pytest
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.ai.hooks.llamaindex import LlamaIndexHook
@@ -56,6 +58,78 @@ class TestLlamaIndexHookInit:
assert hook.embed_model == "text-embedding-3-large"
assert hook.llm_model == "gpt-4o"
+ @pytest.mark.parametrize(
+ ("positional", "expected"),
+ [
+ pytest.param(("embed-pos",), {"embed_model": "embed-pos",
"llm_model": None}, id="one"),
+ pytest.param(
+ ("embed-pos", "llm-pos"),
+ {"embed_model": "embed-pos", "llm_model": "llm-pos"},
+ id="two",
+ ),
+ ],
+ )
+ def test_positional_model_emits_deprecation_warning(self, positional,
expected):
+ with pytest.warns(AirflowProviderDeprecationWarning) as record:
+ hook = LlamaIndexHook("c1", "c2", *positional)
+ assert len(record) == len(positional)
+ assert "Parameter `embed_model` passed positionally is deprecated" in
str(record[0].message)
+ assert hook.llm_conn_id == "c1"
+ assert hook.embed_conn_id == "c2"
+ assert hook.embed_model == expected["embed_model"]
+ assert hook.llm_model == expected["llm_model"]
+
+ def test_positional_warning_stacklevel_points_to_caller(self):
+ with pytest.warns(AirflowProviderDeprecationWarning) as record:
+ LlamaIndexHook("c1", "c2", "embed-pos")
+ assert record[0].filename == __file__
+
+ def test_one_warning_per_positional_argument(self):
+ with pytest.warns(AirflowProviderDeprecationWarning) as record:
+ LlamaIndexHook("c1", "c2", "embed-pos", "llm-pos")
+ messages = [str(w.message) for w in record]
+ assert len(messages) == 2
+ assert "`embed_model`" in messages[0]
+ assert "`llm_model`" in messages[1]
+
+ def test_keyword_model_arguments_do_not_warn(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", AirflowProviderDeprecationWarning)
+ hook = LlamaIndexHook("c1", "c2", embed_model="embed-pos",
llm_model="llm-pos")
+ assert hook.embed_model == "embed-pos"
+ assert hook.llm_model == "llm-pos"
+
+ def test_connection_positional_arguments_do_not_warn(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", AirflowProviderDeprecationWarning)
+ hook = LlamaIndexHook("c1", "c2")
+ assert hook.embed_model is None
+ assert hook.llm_model is None
+
+ def test_too_many_positional_arguments_raise_type_error(self):
+ with warnings.catch_warnings():
+ warnings.simplefilter("error", AirflowProviderDeprecationWarning)
+ with pytest.raises(TypeError, match="positional argument"):
+ LlamaIndexHook("c1", "c2", "embed-pos", "llm-pos", "extra")
+
+ def test_positional_and_keyword_same_parameter_raise_type_error(self):
+ with pytest.raises(TypeError, match="multiple values for argument
'embed_model'"):
+ LlamaIndexHook("c1", "c2", "embed-pos", embed_model="other")
+
+ def
test_second_positional_and_keyword_same_parameter_raise_type_error(self):
+ with (
+ pytest.warns(AirflowProviderDeprecationWarning),
+ pytest.raises(TypeError, match="multiple values for argument
'llm_model'"),
+ ):
+ LlamaIndexHook("c1", "c2", "E", "L", llm_model="X")
+
+ def test_positional_and_keyword_different_parameters_combine(self):
+ with pytest.warns(AirflowProviderDeprecationWarning) as record:
+ hook = LlamaIndexHook("c1", "c2", "embed-pos", llm_model="llm-pos")
+ assert len(record) == 1
+ assert hook.embed_model == "embed-pos"
+ assert hook.llm_model == "llm-pos"
+
def test_conn_type_is_llamaindex(self):
assert LlamaIndexHook.conn_type == "llamaindex"
assert LlamaIndexHook.default_conn_name == "llamaindex_default"
@@ -266,6 +340,16 @@ class TestConnectionTest:
assert success is False
assert "Invalid API key" in message
+ @patch.object(LlamaIndexHook, "get_connection")
+ def test_non_openai_model_name_fails(self, mock_get_conn):
+ """The real ``OpenAI`` class only rejects a model name once
``.metadata`` is read."""
+ mock_get_conn.return_value = _conn(password="sk-test",
extra={"llm_model": "llama3.2"})
+
+ success, message = LlamaIndexHook().test_connection()
+
+ assert success is False
+ assert "Unknown model 'llama3.2'" in message
+
@patch.object(LlamaIndexHook, "get_connection")
def test_failed_connection_no_model(self, mock_get_conn):
mock_get_conn.return_value = _conn()
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
index d226d585dd7..d0efd7f64dd 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_mcp.py
@@ -17,12 +17,14 @@
from __future__ import annotations
import json
+import sys
from unittest.mock import MagicMock, patch
import pytest
from airflow.models.connection import Connection
from airflow.providers.common.ai.hooks.mcp import MCPHook
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
# The hook imports these lazily inside get_conn(), so we patch them at their
# source modules rather than on the hook module. MCPToolset is the unified
@@ -178,6 +180,31 @@ class TestMCPHookGetConn:
mock_transport_cls.assert_called_once_with(command="uvx",
args=["mcp-run-python"], env=None)
+ @pytest.mark.parametrize(
+ ("raw_args", "expected_args"),
+ [
+ (None, []),
+ ('["mcp-run-python"]', ["mcp-run-python"]),
+ (' ["-m", "server"]', ["-m", "server"]),
+ ],
+ ids=["null", "json_array_string", "json_array_string_leading_space"],
+ )
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_STDIO_TRANSPORT, autospec=True)
+ def test_args_from_connection_form_shapes(
+ self, mock_transport_cls, mock_toolset_cls, raw_args, expected_args
+ ):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command": "uvx", "args":
raw_args}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.get_conn()
+
+ mock_transport_cls.assert_called_once_with(command="uvx",
args=expected_args, env=None)
+
def test_http_without_host_raises(self):
hook = MCPHook(mcp_conn_id="test_conn")
conn = Connection(conn_id="test_conn", conn_type="mcp")
@@ -207,6 +234,107 @@ class TestMCPHookGetConn:
with pytest.raises(ValueError, match="requires 'command'"):
hook.get_conn()
+ def test_missing_mcp_dependency_raises_optional_feature_exception(self):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ with patch.dict(sys.modules, {"fastmcp.client.transports": None}):
+ with pytest.raises(AirflowOptionalProviderFeatureException,
match="pydantic-ai-slim"):
+ hook.get_conn()
+
+ @pytest.mark.parametrize(
+ "bad_args",
+ [{"a": 1}, ["-m", 1], 5, '["-m", 1]', '["-m"', '["-m"] trailing'],
+ ids=["dict", "non_string_item", "int", "json_non_string_item",
"malformed_json", "trailing_text"],
+ )
+ def test_stdio_invalid_args_raises(self, bad_args):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command": "uvx", "args":
bad_args}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="'args' in extra"):
+ hook.get_conn()
+
+ @pytest.mark.parametrize(
+ "bad_args",
+ [["--header", "Authorization: Bearer s3cr3t", 1], '["--header",
"Bearer s3cr3t", 1]'],
+ ids=["list", "json_string"],
+ )
+ def
test_stdio_non_string_arg_error_reports_type_and_index_not_values(self,
bad_args):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command": "uvx", "args":
bad_args}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match=r"int at index 2") as
exc_info:
+ hook.get_conn()
+ assert "s3cr3t" not in str(exc_info.value)
+
+ def test_stdio_non_list_args_error_does_not_echo_value(self):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command": "uvx", "args":
{"token": "s3cr3t"}}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="got dict") as exc_info:
+ hook.get_conn()
+ assert "s3cr3t" not in str(exc_info.value)
+
+ def test_stdio_invalid_json_args_error_does_not_echo_value(self):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command": "uvx", "args":
'["Bearer s3cr3t"'}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="looks like a JSON array") as
exc_info:
+ hook.get_conn()
+ assert "s3cr3t" not in str(exc_info.value)
+
+ @pytest.mark.parametrize(
+ "bad_timeout",
+ ["30", -5, -0.5, float("nan"), float("inf"), True, [10]],
+ ids=["str", "neg", "neg_float", "nan", "inf", "bool", "list"],
+ )
+ def test_stdio_invalid_timeout_raises(self, bad_timeout):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command": "uvx", "args":
[], "timeout": bad_timeout}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="'timeout' in extra"):
+ hook.get_conn()
+
+ @pytest.mark.parametrize(
+ ("timeout", "expected_init_timeout"),
+ [(2.5, 2.5), (0, 0), (0.0, 0.0), (None, 0)],
+ ids=["float", "zero", "zero_float", "null_maps_to_zero"],
+ )
+ @patch(_MCP_TOOLSET, autospec=True)
+ @patch(_STDIO_TRANSPORT, autospec=True)
+ def test_stdio_float_zero_and_null_timeout_accepted(
+ self, mock_transport_cls, mock_toolset_cls, timeout,
expected_init_timeout
+ ):
+ hook = MCPHook(mcp_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="mcp",
+ extra=json.dumps({"transport": "stdio", "command": "uvx", "args":
[], "timeout": timeout}),
+ )
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.get_conn()
+ mock_toolset_cls.assert_called_once_with(
+ mock_transport_cls.return_value, init_timeout=expected_init_timeout
+ )
+
def test_unknown_transport_raises(self):
hook = MCPHook(mcp_conn_id="test_conn")
conn = Connection(
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 2340189bb91..d86365ed00a 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
@@ -1561,6 +1561,20 @@ class TestPydanticAIHookCreateAgentInstrumentation:
mock_settings.assert_not_called()
+class _RejectsKwargsProvider:
+ """Stand-in for a provider class whose constructor takes none of the
connection's fields."""
+
+ def __init__(self):
+ pass
+
+
+def _call_provider_factory(model_name, provider_factory=None):
+ """Stand-in for ``infer_model`` that, like the real one, builds the
provider eagerly."""
+ if provider_factory is not None:
+ provider_factory(model_name.split(":", 1)[0])
+ return MagicMock(spec=Model)
+
+
class TestPydanticAIHookTestConnection:
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
def test_successful_connection(self, mock_infer_model):
@@ -1610,6 +1624,52 @@ class TestPydanticAIHookTestConnection:
assert success is False
assert "No model specified" in message
+
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class",
autospec=True)
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
+ def test_get_conn_raises_when_provider_class_cannot_be_built(
+ self, mock_infer_model, mock_infer_provider_class, registry
+ ):
+ mock_infer_provider_class.return_value = _RejectsKwargsProvider
+ mock_infer_model.side_effect = _call_provider_factory
+ registry.add("primary", password="sk-key", extra={"model":
"openai:gpt-5.6-sol"})
+
+ with pytest.raises(ValueError, match="Could not build provider") as
exc_info:
+ PydanticAIHook(llm_conn_id="primary").get_conn()
+
+ assert isinstance(exc_info.value.__cause__, TypeError)
+
+ @pytest.mark.parametrize(
+ ("fallback_conn_ids", "expected_conn_id"),
+ [
+ pytest.param(None, "primary", id="primary"),
+ pytest.param(["second"], "second", id="fallback_connection"),
+ ],
+ )
+
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class",
autospec=True)
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
+ def test_reports_provider_class_build_failure(
+ self, mock_infer_model, mock_infer_provider_class, registry,
fallback_conn_ids, expected_conn_id
+ ):
+ mock_infer_provider_class.return_value = _RejectsKwargsProvider
+ mock_infer_model.side_effect = _call_provider_factory
+ registry.add(
+ "primary",
+ password="sk-key" if fallback_conn_ids is None else None,
+ extra={"model": "openai:gpt-5.6-sol"},
+ )
+ registry.add("second", password="sk-key", extra={"model":
"openai:gpt-5.6-sol"})
+
+ success, message = PydanticAIHook(
+ llm_conn_id="primary", fallback_conn_ids=fallback_conn_ids
+ ).test_connection()
+
+ assert success is False
+ assert (
+ f"Could not build provider 'openai' from the fields on connection
'{expected_conn_id}'" in message
+ )
+ assert "(fields passed: api_key)" in message
+ assert "unexpected keyword argument 'api_key'" in message
+
# ---------------------------------------------------------------------------
# Subclass hook tests
@@ -2020,13 +2080,12 @@ class TestPydanticAIVertexHook:
mock_provider_cls.assert_called_with(project="my-project",
location="europe-west4")
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
- @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider",
autospec=True)
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class",
autospec=True)
- def test_get_conn_vertexai_flag_does_not_trigger_typeerror_fallback(
- self, mock_infer_provider_class, mock_infer_provider, mock_infer_model
+ def test_get_conn_vertexai_flag_does_not_trigger_typeerror_rejection(
+ self, mock_infer_provider_class, mock_infer_model
):
- """Setting ``vertexai`` must not push ``get_conn`` onto the ``except
TypeError``
- fallback path, which would silently discard
project/location/credentials.
+ """Setting ``vertexai`` must not make the provider class reject the
connection's
+ kwargs with ``TypeError``, which ``get_conn`` turns into a
``ValueError``.
The stand-in below has the exact keyword-only signature of the real
``GoogleCloudProvider.__init__`` (verified against the installed
pydantic-ai) so
@@ -2080,8 +2139,6 @@ class TestPydanticAIVertexHook:
assert isinstance(provider, FakeGoogleCloudProvider)
assert provider.kwargs["project"] == "my-project"
assert provider.kwargs["location"] == "us-central1"
- # The TypeError fallback must never have been reached.
- mock_infer_provider.assert_not_called()
def test_documented_model_prefix_is_a_valid_pydantic_ai_provider(self):
"""Regression test: the model-prefix documented in the connection form
and
diff --git a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
index a317988efe8..33e638eeb4d 100644
--- a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
+++ b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
@@ -17,10 +17,12 @@
from __future__ import annotations
import copy
+import importlib.util
import json
import logging
import math
import re
+import sys
import traceback
import warnings
from datetime import timedelta
@@ -37,6 +39,9 @@ from pydantic_ai.messages import ModelResponse
# Skip the entire test module on older Airflow versions tested in compat CI.
pytest.importorskip("airflow.sdk.definitions.retry_policy",
reason="RetryPolicy requires Airflow 3.3+")
+from airflow.models.connection import Connection
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIAzureHook,
PydanticAIHook
+from airflow.providers.common.ai.policies import retry as retry_module
from airflow.providers.common.ai.policies.retry import (
CLASSIFIER_INSTRUCTIONS,
DEFAULT_CATEGORIES,
@@ -48,6 +53,7 @@ from airflow.providers.common.ai.policies.retry import (
redact_registered_secrets,
)
from airflow.providers.common.ai.utils.decision import picked_key
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
from airflow.sdk._shared.secrets_masker import reset_secrets_masker
from airflow.sdk.definitions.retry_policy import RetryAction, RetryDecision,
RetryPolicy, RetryRule
from airflow.sdk.log import mask_secret
@@ -79,14 +85,21 @@ def _agent(output, **kwargs):
return agent
+def _stub_hook(mock_hook_cls):
+ """Make ``get_hook`` return a hook instance that keeps the real method
signatures."""
+ hook = create_autospec(PydanticAIHook, instance=True)
+ mock_hook_cls.get_hook.return_value = hook
+ return hook
+
+
def _install(mock_hook_cls, agent):
"""Wire a mock agent behind the patched hook class and return the mock
agent."""
- mock_hook_cls.return_value.create_agent.return_value = agent
+ _stub_hook(mock_hook_cls).create_agent.return_value = agent
return agent
def _output_type(mock_hook_cls):
- return
mock_hook_cls.return_value.create_agent.call_args.kwargs["output_type"]
+ return
mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs["output_type"]
def _options(output_type):
@@ -109,6 +122,15 @@ def
test_redact_registered_secrets_masks_only_registered_values():
)
+def
test_import_without_retry_policy_support_raises_optional_feature_exception():
+ """A fresh copy of the module is executed so the real one in
``sys.modules`` stays intact."""
+ spec = importlib.util.spec_from_file_location("_retry_without_support",
retry_module.__file__)
+ module = importlib.util.module_from_spec(spec)
+ with patch.dict(sys.modules, {"airflow.sdk.definitions.retry_policy":
None}):
+ with pytest.raises(AirflowOptionalProviderFeatureException,
match="Airflow 3.3"):
+ spec.loader.exec_module(module)
+
+
class TestErrorCategory:
@pytest.mark.parametrize("description", ["", " ", None, 3])
def test_description_must_be_a_non_empty_string(self, description):
@@ -177,7 +199,20 @@ class TestClassifierRetryPolicyConstruction:
min_confidence=0.6,
)
- mock_hook_cls.assert_not_called()
+ mock_hook_cls.get_hook.assert_not_called()
+
+ def test_hook_resolves_to_the_connection_types_subclass(self):
+ """An Azure connection must get ``PydanticAIAzureHook``, not the base
hook's ``base_url`` kwargs."""
+ policy = ClassifierRetryPolicy(
+ llm_conn_id="azure_conn", model_id="azure:gpt-5",
categories=DEFAULT_CATEGORIES
+ )
+ conn = Connection(conn_id="azure_conn", conn_type="pydanticai_azure")
+
+ with patch.object(PydanticAIHook, "get_connection", return_value=conn):
+ hook = policy._hook()
+
+ assert isinstance(hook, PydanticAIAzureHook)
+ assert hook.model_id == "azure:gpt-5"
def test_categories_must_be_a_mapping(self):
with pytest.raises(TypeError, match="categories must be a mapping"):
@@ -366,7 +401,7 @@ class TestClassification:
agent.run_sync.return_value =
_run_result(kwargs["output_type"]("rate_limit"))
return agent
- mock_hook_cls.return_value.create_agent.side_effect = capture
+ _stub_hook(mock_hook_cls).create_agent.side_effect = capture
policy = ClassifierRetryPolicy(llm_conn_id="test",
categories=DEFAULT_CATEGORIES)
decision = policy.evaluate(RuntimeError("429"), try_number=1,
max_tries=3)
@@ -391,7 +426,7 @@ class TestClassification:
agent.run_sync.return_value =
_run_result(kwargs["output_type"](name))
return agent
- mock_hook_cls.return_value.create_agent.side_effect = capture
+ _stub_hook(mock_hook_cls).create_agent.side_effect = capture
policy = ClassifierRetryPolicy(
llm_conn_id="test",
categories={
@@ -489,7 +524,7 @@ class TestClassification:
policy.evaluate(RuntimeError("x"), try_number=1, max_tries=3)
assert (
-
mock_hook_cls.return_value.create_agent.call_args.kwargs["instructions"]
+
mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs["instructions"]
== "Snowflake errors only."
)
@@ -502,7 +537,9 @@ class TestClassification:
policy.evaluate(RuntimeError("x"), try_number=1, max_tries=3)
- mock_hook_cls.assert_called_once_with(llm_conn_id="jev_default",
model_id="typesafe:jev-1.13.0")
+ mock_hook_cls.get_hook.assert_called_once_with(
+ "jev_default", hook_params={"model_id": "typesafe:jev-1.13.0"}
+ )
@patch(HOOK, autospec=True)
def test_timeout_passed_via_model_settings(self, mock_hook_cls):
@@ -971,7 +1008,7 @@ class TestFallbackBehaviour:
@patch(HOOK, autospec=True)
def test_hook_creation_failure_triggers_fallback(self, mock_hook_cls):
- mock_hook_cls.return_value.create_agent.side_effect =
RuntimeError("unexpected")
+ _stub_hook(mock_hook_cls).create_agent.side_effect =
RuntimeError("unexpected")
policy = ClassifierRetryPolicy(
llm_conn_id="test",
fallback_rules=[RetryRule(exception=ValueError,
action=RetryAction.FAIL, reason="caught")],
@@ -1058,7 +1095,7 @@ class TestOnUncertain:
def create_agent(**kwargs):
return llm if kwargs["output_type"] is ErrorClassification else
classifier
- mock_hook_cls.return_value.create_agent.side_effect = create_agent
+ _stub_hook(mock_hook_cls).create_agent.side_effect = create_agent
return ClassifierRetryPolicy(
llm_conn_id="jev",
min_confidence=bar,
@@ -1097,7 +1134,7 @@ class TestOnUncertain:
assert decision.action == RetryAction.RETRY
assert decision.reason == "category=network confidence=0.95
threshold=0.80 action=retry delay=10s"
- assert mock_hook_cls.return_value.create_agent.call_count == 1
+ assert mock_hook_cls.get_hook.return_value.create_agent.call_count == 1
@patch(HOOK, autospec=True)
def test_unsure_classifier_escalates_to_the_llm(self, mock_hook_cls):
@@ -1256,7 +1293,7 @@ class TestOnUncertain:
@patch(HOOK, autospec=True)
def
test_nested_classifiers_both_down_still_reach_the_outer_fail_rule(self,
mock_hook_cls):
"""A ClassifierRetryPolicy as fallback_policy falls back with a
DEFAULT of its own; the outer rules must still run."""
-
mock_hook_cls.return_value.create_agent.return_value.run_sync.side_effect =
TimeoutError("down")
+
_stub_hook(mock_hook_cls).create_agent.return_value.run_sync.side_effect =
TimeoutError("down")
inner = ClassifierRetryPolicy(llm_conn_id="jev_b", min_confidence=0.5)
policy = ClassifierRetryPolicy(
llm_conn_id="jev_a",
@@ -1271,7 +1308,7 @@ class TestOnUncertain:
assert decision.action == RetryAction.FAIL
assert decision.reason == "classifier answer not applied
(model_error); never retry 403"
- assert mock_hook_cls.return_value.create_agent.call_count == 2
+ assert mock_hook_cls.get_hook.return_value.create_agent.call_count == 2
def _open_agent(category, should_retry, delay=0, reasoning="test"):
@@ -1316,7 +1353,7 @@ class TestLLMRetryPolicy:
_install(mock_hook_cls, _open_agent("auth", should_retry=False))
LLMRetryPolicy(llm_conn_id="test").evaluate(PermissionError("403"),
try_number=1, max_tries=3)
- kwargs = mock_hook_cls.return_value.create_agent.call_args.kwargs
+ kwargs =
mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs
assert kwargs["output_type"] is ErrorClassification
assert kwargs["instructions"] == DEFAULT_INSTRUCTIONS
@@ -1368,7 +1405,7 @@ class TestLLMRetryPolicy:
assert decision.action == RetryAction.RETRY
assert decision.retry_delay == timedelta(seconds=30)
assert decision.reason == "warehouse_suspended: auto-resume"
- assert
mock_hook_cls.return_value.create_agent.call_args.kwargs["instructions"].startswith(
+ assert
mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs["instructions"].startswith(
"'Warehouse suspended'"
)
diff --git a/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
b/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
index 16a3aa58232..400829ba0ff 100644
--- a/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
+++ b/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
@@ -132,15 +132,16 @@ DEPRECATED_UPSTREAM_MODULES = {"github"}
# only `credentials`, `credential_source`, `openai_client`, and `http_client`
-- no parameter
# exists to receive what the hook sends; confirmed empirically (as of
pydantic-ai-slim 2.44.0):
# `OpenAICodexProvider(api_key=..., base_url=...)` raises
-# ``TypeError: unexpected keyword argument 'api_key'``.
`PydanticAIHook.get_conn()`'s
-# `_provider_factory` catches that `TypeError` and retries via
`infer_provider(pname)`, i.e.
+# ``TypeError: unexpected keyword argument 'api_key'``.
`PydanticAIHook.get_conn()`
+# re-raises that `TypeError` as a `ValueError` when the connection has a
password or host. Only
+# a connection with neither falls through to `infer_model(model_name)`, i.e.
# `OpenAICodexProvider()` with no arguments, which falls through to its
# `_read_codex_cli_credentials()` helper: a read-only load of the Codex CLI's
on-disk
# `~/.codex/auth.json`, populated by the OAuth authorization-code + PKCE login
flow the Codex CLI
# itself performs (see `pydantic_ai.providers.openai_codex`'s module docstring
and
# `OpenAICodexOAuthFlow`) -- not by anything `hooks/pydantic_ai.py` can
express. So credentials
# configured on a `pydanticai` connection (password/host) can never drive this
provider; only
-# the ambient Codex CLI login can, via that fallback. The services page lists
only vendors a
+# the ambient Codex CLI login can, with no password/host set. The services
page lists only vendors a
# connection's credentials can drive, so these modules stay off it.
# `test_hook_incompatible_auth_modules_stay_incompatible` below guards the
reversal: if
# pydantic-ai gives one of these providers an `api_key`/`base_url` constructor
parameter, this
diff --git
a/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_model.py
b/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_model.py
index 8371c9505f8..c3fcbaf937f 100644
--- a/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_model.py
+++ b/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_model.py
@@ -65,8 +65,8 @@ class TestGetProviderKwargs:
def test_ignores_connection_password_and_omits_api_key(self,
mock_hook_cls, mock_provider_cls):
"""Regression: a leftover Password on the pydanticai_snowflake
connection must not be
forwarded as api_key -- the base hook's generic handling raises
TypeError against
- SnowflakeProvider (which takes `token`, not `api_key`) and silently
falls back to
- env-var auth; this hook's own kwargs must never trigger that path."""
+ SnowflakeProvider (which takes `token`, not `api_key`), which fails
the connection
+ with a ValueError; this hook's own kwargs must never trigger that
path."""
mock_hook_cls.return_value._get_static_conn_params = {"account":
"airflow"}
hook = PydanticAISnowflakeHook(snowflake_conn_id="snowflake_conn")
@@ -248,17 +248,14 @@ class TestSnowflakeCortexAuth:
class TestGetConn:
@mock.patch(f"{MODULE_PATH}.SnowflakeRestTokenProvider", autospec=True)
@mock.patch(f"{MODULE_PATH}.SnowflakeHook", autospec=True)
-
@mock.patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider",
autospec=True)
def test_get_conn_uses_explicit_credentials_despite_connection_password(
- self, mock_infer_provider, mock_hook_cls, mock_provider_cls
+ self, mock_hook_cls, mock_provider_cls
):
- """End-to-end: get_conn() must not raise and must not fall back to
env-var auth.
+ """End-to-end: get_conn() must not raise when the connection also
carries a password.
- ``infer_provider`` is only reached by the base hook's ``except
TypeError`` fallback
- (``pydantic_ai.py``'s ``_provider_factory``), so asserting it was
never called pins that
- the explicit-credentials path was taken -- ``model is not None`` alone
would also pass
- if this hook's kwargs were rejected and pydantic-ai silently fell back
to
- ``SNOWFLAKE_ACCOUNT``/``SNOWFLAKE_TOKEN`` env vars instead.
+ The base hook raises ``ValueError`` (it never falls back to env-var
auth) when the
+ provider class rejects the kwargs it derives from the connection, so a
model that
+ resolves means the explicit-credentials path was taken.
"""
mock_hook_cls.return_value._get_static_conn_params = {"account":
"airflow"}
@@ -276,7 +273,6 @@ class TestGetConn:
assert model is not None
mock_hook_cls.assert_called_once_with(snowflake_conn_id="snowflake_conn")
- mock_infer_provider.assert_not_called()
@mock.patch(f"{MODULE_PATH}.SnowflakeRestTokenProvider", autospec=True)
@mock.patch(f"{MODULE_PATH}.SnowflakeHook", autospec=True)
