kaxil commented on code in PR #72156:
URL: https://github.com/apache/airflow/pull/72156#discussion_r4042148747
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,33 +178,107 @@ def _get_provider_kwargs(
kwargs["base_url"] = base_url
return kwargs
+ def _get_conn_and_extra(self) -> tuple[Connection, dict[str, Any]]:
+ """Return this hook's connection and its deserialized extra, fetching
at most once."""
+ if self._conn is None:
+ self._conn = self.get_connection(self.llm_conn_id)
+ self._conn_extra_dejson = self._conn.extra_dejson
+ return self._conn, self._conn_extra_dejson
+
def get_conn(self) -> Model:
"""
Return a configured pydantic-ai ``Model``.
- Resolution order:
+ Resolution order for this hook's own connection:
1. **Explicit credentials** — when :meth:`_get_provider_kwargs` returns
a non-empty dict the provider class is instantiated with those
kwargs
and wrapped in a ``provider_factory``.
2. **Default resolution** — delegates to pydantic-ai ``infer_model``
which reads standard env vars (``OPENAI_API_KEY``, ``AWS_PROFILE``,
…).
+ A bare ``model_id`` (one with no recognized platform prefix) is
qualified with
+ this connection's own platform before either of the above -- see the
class
+ docstring's ``model_id`` entry for the resolution and
fallback-forwarding rules.
+
+ When ``fallback_conn_ids`` is configured (on the hook or in the
+ connection's extra) the resolved models are wrapped in a pydantic-ai
+ ``FallbackModel``, so a provider outage moves to the next connection
+ *within the same task attempt* instead of failing the task.
+
+ Two costs of that wrapping are worth knowing before configuring a long
+ chain. A ``timeout`` in ``ModelSettings`` is applied by pydantic-ai to
+ every model in the chain rather than to the chain as a whole, so the
+ worst-case wait is the timeout multiplied by the number of connections.
+ And there is no circuit breaker: every call retries the primary first,
+ so during an outage each task instance pays the primary's timeout
again.
+ Keep the primary's timeout short to bound both.
+
The resolved model is cached for the lifetime of this hook instance.
"""
if self._model is not None:
return self._model
- conn = self.get_connection(self.llm_conn_id) if self._conn is None
else self._conn
- extra: dict[str, Any] = (
- conn.extra_dejson if self._conn_extra_dejson is None else
self._conn_extra_dejson
- )
+ model = self._resolve_own_model()
+ fallback_models = self._resolve_fallback_models()
+ self._model = FallbackModel(model, *fallback_models) if
fallback_models else model
+ return self._model
+
+ def _qualify_model_name(self, model_name: str) -> str:
+ """
+ Prefix a bare model name with this connection's platform.
+
+ A name is treated as already pinning a platform only when the segment
before
+ its first ``:`` is itself a provider pydantic-ai recognizes (e.g.
+ ``"openai:gpt-4"``) -- see :func:`_has_recognized_provider_prefix`.
Everything
+ else is a bare name, even one that happens to contain a ``:`` of its
own (e.g.
+ Bedrock's version-suffixed ``"us.anthropic.claude-opus-4-6-v1:0"``),
and is
+ prefixed with :attr:`model_provider`; the generic ``pydanticai``
connection type
+ has no platform of its own (``model_provider`` is ``None``), so a bare
name there
+ raises instead of reaching pydantic-ai's own, less actionable
``Unknown model``
+ error.
+ """
+ if _has_recognized_provider_prefix(model_name):
Review Comment:
`"test"` has no colon, so it falls through to the `model_provider` branch: a
generic `pydanticai` connection with `extra={"model": "test"}` now raises
`ValueError` where `infer_model` returned pydantic-ai's `TestModel` at both
2.23.0 and current, and on the vendor subclasses it silently becomes
`azure:test` / `bedrock:test` / `google-cloud:test`. A typo'd prefix lands in
the same branch, where the raise calls `openi:gpt-5` a "bare model name" and
tells the user to set a `provider:model` string they already set, while
pydantic-ai names the real problem at both ends of the supported range --
`Unknown provider: openi` at the 2.23.0 floor, `Unknown model: openi:gpt-5. Did
you mean 'openai:gpt-5'?` at 2.42. An early `if model_name == "test": return
model_name` restores the sentinel on all four classes, and naming the
unrecognized segment would cover the typo.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,33 +178,107 @@ def _get_provider_kwargs(
kwargs["base_url"] = base_url
return kwargs
+ def _get_conn_and_extra(self) -> tuple[Connection, dict[str, Any]]:
+ """Return this hook's connection and its deserialized extra, fetching
at most once."""
+ if self._conn is None:
+ self._conn = self.get_connection(self.llm_conn_id)
+ self._conn_extra_dejson = self._conn.extra_dejson
+ return self._conn, self._conn_extra_dejson
+
def get_conn(self) -> Model:
"""
Return a configured pydantic-ai ``Model``.
- Resolution order:
+ Resolution order for this hook's own connection:
1. **Explicit credentials** — when :meth:`_get_provider_kwargs` returns
a non-empty dict the provider class is instantiated with those
kwargs
and wrapped in a ``provider_factory``.
2. **Default resolution** — delegates to pydantic-ai ``infer_model``
which reads standard env vars (``OPENAI_API_KEY``, ``AWS_PROFILE``,
…).
+ A bare ``model_id`` (one with no recognized platform prefix) is
qualified with
+ this connection's own platform before either of the above -- see the
class
+ docstring's ``model_id`` entry for the resolution and
fallback-forwarding rules.
+
+ When ``fallback_conn_ids`` is configured (on the hook or in the
+ connection's extra) the resolved models are wrapped in a pydantic-ai
+ ``FallbackModel``, so a provider outage moves to the next connection
+ *within the same task attempt* instead of failing the task.
+
+ Two costs of that wrapping are worth knowing before configuring a long
+ chain. A ``timeout`` in ``ModelSettings`` is applied by pydantic-ai to
+ every model in the chain rather than to the chain as a whole, so the
+ worst-case wait is the timeout multiplied by the number of connections.
+ And there is no circuit breaker: every call retries the primary first,
+ so during an outage each task instance pays the primary's timeout
again.
+ Keep the primary's timeout short to bound both.
+
The resolved model is cached for the lifetime of this hook instance.
"""
if self._model is not None:
return self._model
- conn = self.get_connection(self.llm_conn_id) if self._conn is None
else self._conn
- extra: dict[str, Any] = (
- conn.extra_dejson if self._conn_extra_dejson is None else
self._conn_extra_dejson
- )
+ model = self._resolve_own_model()
+ fallback_models = self._resolve_fallback_models()
+ self._model = FallbackModel(model, *fallback_models) if
fallback_models else model
Review Comment:
`fallback_on` is left to pydantic-ai's default, which is `(ModelAPIError,)`
at both 2.23.0 and current, so `UnexpectedModelBehavior`, `UsageLimitExceeded`
and `ContentFilterError` propagate without trying the next connection. That
scope is the thing the new docs pin in three places (the retry-layers table,
the changelog note, and retry_policies.rst), and pyproject says only
`pydantic-ai-slim>=2.23.0` with no upper bound, so an upstream default change
would move the documented behaviour silently. Passing
`fallback_on=(ModelAPIError,)` explicitly here keeps the contract in this repo,
and the kwarg exists at the 2.23.0 floor.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,33 +178,107 @@ def _get_provider_kwargs(
kwargs["base_url"] = base_url
return kwargs
+ def _get_conn_and_extra(self) -> tuple[Connection, dict[str, Any]]:
+ """Return this hook's connection and its deserialized extra, fetching
at most once."""
+ if self._conn is None:
+ self._conn = self.get_connection(self.llm_conn_id)
+ self._conn_extra_dejson = self._conn.extra_dejson
+ return self._conn, self._conn_extra_dejson
+
def get_conn(self) -> Model:
"""
Return a configured pydantic-ai ``Model``.
- Resolution order:
+ Resolution order for this hook's own connection:
1. **Explicit credentials** — when :meth:`_get_provider_kwargs` returns
a non-empty dict the provider class is instantiated with those
kwargs
and wrapped in a ``provider_factory``.
2. **Default resolution** — delegates to pydantic-ai ``infer_model``
which reads standard env vars (``OPENAI_API_KEY``, ``AWS_PROFILE``,
…).
+ A bare ``model_id`` (one with no recognized platform prefix) is
qualified with
+ this connection's own platform before either of the above -- see the
class
+ docstring's ``model_id`` entry for the resolution and
fallback-forwarding rules.
+
+ When ``fallback_conn_ids`` is configured (on the hook or in the
+ connection's extra) the resolved models are wrapped in a pydantic-ai
+ ``FallbackModel``, so a provider outage moves to the next connection
+ *within the same task attempt* instead of failing the task.
+
+ Two costs of that wrapping are worth knowing before configuring a long
+ chain. A ``timeout`` in ``ModelSettings`` is applied by pydantic-ai to
+ every model in the chain rather than to the chain as a whole, so the
+ worst-case wait is the timeout multiplied by the number of connections.
+ And there is no circuit breaker: every call retries the primary first,
+ so during an outage each task instance pays the primary's timeout
again.
+ Keep the primary's timeout short to bound both.
+
The resolved model is cached for the lifetime of this hook instance.
"""
if self._model is not None:
return self._model
- conn = self.get_connection(self.llm_conn_id) if self._conn is None
else self._conn
- extra: dict[str, Any] = (
- conn.extra_dejson if self._conn_extra_dejson is None else
self._conn_extra_dejson
- )
+ model = self._resolve_own_model()
+ fallback_models = self._resolve_fallback_models()
+ self._model = FallbackModel(model, *fallback_models) if
fallback_models else model
+ return self._model
+
+ def _qualify_model_name(self, model_name: str) -> str:
+ """
+ Prefix a bare model name with this connection's platform.
+
+ A name is treated as already pinning a platform only when the segment
before
+ its first ``:`` is itself a provider pydantic-ai recognizes (e.g.
+ ``"openai:gpt-4"``) -- see :func:`_has_recognized_provider_prefix`.
Everything
+ else is a bare name, even one that happens to contain a ``:`` of its
own (e.g.
+ Bedrock's version-suffixed ``"us.anthropic.claude-opus-4-6-v1:0"``),
and is
+ prefixed with :attr:`model_provider`; the generic ``pydanticai``
connection type
+ has no platform of its own (``model_provider`` is ``None``), so a bare
name there
+ raises instead of reaching pydantic-ai's own, less actionable
``Unknown model``
+ error.
+ """
+ if _has_recognized_provider_prefix(model_name):
+ return model_name
+ if self.model_provider is None:
+ raise ValueError(
+ f"Connection '{self.llm_conn_id}' has no default model
provider, so the bare "
Review Comment:
Beyond the typo'd-prefix case, this message can name a model the user never
put on the connection it blames: a Bedrock primary with a bare name plus a
generic `pydanticai` fallback with no `model` of its own produces "Connection
'openai_dr' has no default model provider, so the bare model name
'claude-opus-4-5' cannot be resolved", and nothing on `openai_dr` mentions that
name. It is also wrong for `policies/retry.py:209`, which builds the base
`PydanticAIHook` directly rather than through `get_hook`, so a
`pydanticai_bedrock` connection reached through `LLMRetryPolicy` is told its
connection type has no default provider when that type does. Saying the name
was forwarded from the primary would cover the first; the second looks like a
retry.py follow-up rather than a change here.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,33 +178,107 @@ def _get_provider_kwargs(
kwargs["base_url"] = base_url
return kwargs
+ def _get_conn_and_extra(self) -> tuple[Connection, dict[str, Any]]:
+ """Return this hook's connection and its deserialized extra, fetching
at most once."""
+ if self._conn is None:
+ self._conn = self.get_connection(self.llm_conn_id)
+ self._conn_extra_dejson = self._conn.extra_dejson
+ return self._conn, self._conn_extra_dejson
+
def get_conn(self) -> Model:
"""
Return a configured pydantic-ai ``Model``.
- Resolution order:
+ Resolution order for this hook's own connection:
1. **Explicit credentials** — when :meth:`_get_provider_kwargs` returns
a non-empty dict the provider class is instantiated with those
kwargs
and wrapped in a ``provider_factory``.
2. **Default resolution** — delegates to pydantic-ai ``infer_model``
which reads standard env vars (``OPENAI_API_KEY``, ``AWS_PROFILE``,
…).
+ A bare ``model_id`` (one with no recognized platform prefix) is
qualified with
+ this connection's own platform before either of the above -- see the
class
+ docstring's ``model_id`` entry for the resolution and
fallback-forwarding rules.
+
+ When ``fallback_conn_ids`` is configured (on the hook or in the
+ connection's extra) the resolved models are wrapped in a pydantic-ai
+ ``FallbackModel``, so a provider outage moves to the next connection
+ *within the same task attempt* instead of failing the task.
+
+ Two costs of that wrapping are worth knowing before configuring a long
+ chain. A ``timeout`` in ``ModelSettings`` is applied by pydantic-ai to
+ every model in the chain rather than to the chain as a whole, so the
+ worst-case wait is the timeout multiplied by the number of connections.
+ And there is no circuit breaker: every call retries the primary first,
+ so during an outage each task instance pays the primary's timeout
again.
+ Keep the primary's timeout short to bound both.
+
The resolved model is cached for the lifetime of this hook instance.
"""
if self._model is not None:
return self._model
- conn = self.get_connection(self.llm_conn_id) if self._conn is None
else self._conn
- extra: dict[str, Any] = (
- conn.extra_dejson if self._conn_extra_dejson is None else
self._conn_extra_dejson
- )
+ model = self._resolve_own_model()
+ fallback_models = self._resolve_fallback_models()
+ self._model = FallbackModel(model, *fallback_models) if
fallback_models else model
+ return self._model
+
+ def _qualify_model_name(self, model_name: str) -> str:
+ """
+ Prefix a bare model name with this connection's platform.
+
+ A name is treated as already pinning a platform only when the segment
before
+ its first ``:`` is itself a provider pydantic-ai recognizes (e.g.
+ ``"openai:gpt-4"``) -- see :func:`_has_recognized_provider_prefix`.
Everything
+ else is a bare name, even one that happens to contain a ``:`` of its
own (e.g.
+ Bedrock's version-suffixed ``"us.anthropic.claude-opus-4-6-v1:0"``),
and is
+ prefixed with :attr:`model_provider`; the generic ``pydanticai``
connection type
+ has no platform of its own (``model_provider`` is ``None``), so a bare
name there
+ raises instead of reaching pydantic-ai's own, less actionable
``Unknown model``
+ error.
+ """
+ if _has_recognized_provider_prefix(model_name):
+ return model_name
+ if self.model_provider is None:
+ raise ValueError(
+ f"Connection '{self.llm_conn_id}' has no default model
provider, so the bare "
+ f"model name '{model_name}' cannot be resolved. Use a vendor
connection type "
+ "(Azure/Bedrock/Vertex) or set an explicit 'provider:model'
string."
+ )
+ return f"{self.model_provider}:{model_name}"
Review Comment:
On a vendor connection this double-prefixes any name whose colon-prefix is
not a recognized provider, so `google-vertex:gemini-2.0-flash` on a
`pydanticai_vertex` connection becomes
`google-cloud:google-vertex:gemini-2.0-flash` and resolves cleanly into a
`GoogleModel` named `google-vertex:gemini-2.0-flash`, failing only once the
request reaches Vertex. That spelling is not hypothetical -- this provider
shipped `google-vertex:` and `google-gla:` in its own tests and docs until
#72011, so existing connections carry it, and `test_connection` goes from
reporting `Unknown provider: google-vertex` to "Model resolved successfully."
Keeping the Bedrock-native-id behaviour but logging a warning whenever a
prefixed name is treated as a bare vendor id would leave a stale prefix
discoverable.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -178,24 +303,101 @@ def _provider_factory(pname: str) -> Any:
)
return infer_provider(pname)
- self._model = infer_model(model_name,
provider_factory=_provider_factory)
- return self._model
+ return infer_model(model_name, provider_factory=_provider_factory)
- self._model = infer_model(model_name)
- return self._model
+ return infer_model(model_name)
- def _get_conn_if_model_configured(self) -> Model | None:
- """Return the hook model only when the hook or connection explicitly
configures one."""
- if self.model_id:
- return self.get_conn()
+ def _get_fallback_conn_ids(self) -> list[str]:
+ """Return the configured fallback connection IDs, hook argument
winning over the extra."""
+ if self.fallback_conn_ids is not None:
+ raw: Any = self.fallback_conn_ids
+ else:
+ _, extra = self._get_conn_and_extra()
+ raw = extra.get(FALLBACK_CONN_IDS_EXTRA_KEY)
+ if raw is None:
+ raw = []
+
+ if not isinstance(raw, (list, tuple)) or not all(isinstance(item, str)
and item for item in raw):
+ raise ValueError(
+ f"{FALLBACK_CONN_IDS_EXTRA_KEY} for connection
'{self.llm_conn_id}' must be a list "
+ f"of non-empty connection IDs, got {raw!r}."
+ )
+ return list(raw)
- conn = self.get_connection(self.llm_conn_id)
- self._conn = conn
- self._conn_extra_dejson = conn.extra_dejson
+ def _resolve_fallback_models(self) -> list[Model]:
+ """
+ Resolve one ``Model`` per fallback connection, in the configured order.
Review Comment:
This 20-line docstring restates the forwarding rules that the class
docstring's `model_id` entry and `_resolve_own_model`'s `:param
forwarded_model_id:` already state, and the Bedrock embedded-colon example
appears four times in this one file. Could the rule live in one place, with the
others pointing at it? The behaviour reads fine; it is the volume of duplicated
prose a reviewer has to diff against itself.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -178,24 +303,101 @@ def _provider_factory(pname: str) -> Any:
)
return infer_provider(pname)
- self._model = infer_model(model_name,
provider_factory=_provider_factory)
- return self._model
+ return infer_model(model_name, provider_factory=_provider_factory)
- self._model = infer_model(model_name)
- return self._model
+ return infer_model(model_name)
- def _get_conn_if_model_configured(self) -> Model | None:
- """Return the hook model only when the hook or connection explicitly
configures one."""
- if self.model_id:
- return self.get_conn()
+ def _get_fallback_conn_ids(self) -> list[str]:
+ """Return the configured fallback connection IDs, hook argument
winning over the extra."""
+ if self.fallback_conn_ids is not None:
+ raw: Any = self.fallback_conn_ids
+ else:
+ _, extra = self._get_conn_and_extra()
+ raw = extra.get(FALLBACK_CONN_IDS_EXTRA_KEY)
+ if raw is None:
+ raw = []
+
+ if not isinstance(raw, (list, tuple)) or not all(isinstance(item, str)
and item for item in raw):
Review Comment:
The new Fallback Connections conn-field renders as a textarea
(`FieldStringArray`, since the schema is `type: [array, null]` with
`items.type: string`) and splits on newline, and its blur handler only nulls
the all-empty case. So a trailing newline saves `["anthropic_prod", ""]` and
this check then rejects the whole chain at task runtime rather than at save
time. Dropping blank entries, or stripping them, would make the field behave
the way the textarea invites people to use it.
##########
providers/common/ai/docs/connections/pydantic_ai_bedrock.rst:
##########
@@ -107,6 +115,12 @@ more than one credential source is set at once, the bearer
token
- The environment-variable / instance-role credential chain
(``AWS_PROFILE``, IAM role, …) when none of the fields above are set.
+Fallback Connections
Review Comment:
This entry landed inside the Credentials section, after its bullet list,
rather than up in Configuring the Connection with the other conn-fields. The
generic and Azure pages put the same entry in Configuring the Connection, and a
fallback chain is not a credential. Same thing on the Vertex page at line 136.
##########
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),
Review Comment:
The "fetched at most once" promise on `_get_conn_and_extra` has no test, and
a fallback chain multiplies what a regression here would cost -- one Execution
API round trip per hop. The fixture mock already records its calls, so pinning
the count is a couple of lines.
##########
providers/common/ai/docs/provider_fallback.rst:
##########
@@ -0,0 +1,208 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Provider fallback
+=================
+
+A single ``llm_conn_id`` gives a task one provider. When that provider is
down, the task
+fails and retries into the same outage. ``fallback_conn_ids`` gives the
connection an
+ordered list of other connections to try, so a provider outage moves to the
next vendor
+inside the same task attempt.
+
+Configure it on the connection
+------------------------------
+
+Put the chain in the primary connection's extra:
+
+.. code-block:: json
+
+ {
+ "model": "openai:gpt-5",
+ "fallback_conn_ids": ["anthropic_prod", "bedrock_dr"]
+ }
+
+Every entry is an Airflow connection ID, resolved through the hook registered
for its own
+connection type. A chain can therefore mix vendors whose credentials live in
different
+connection fields — ``pydanticai`` for OpenAI, ``pydanticai_bedrock`` for a
Bedrock
+standby — without the Dag knowing anything about either.
+
+That is the point of configuring it here rather than in Dag code: the Dag
keeps naming one
+connection, and whoever administers the connections owns the failover
topology. Changing a
+standby provider is a connection edit, not a Dag deployment.
+
+A *bare* model name (e.g. ``"gpt-5"`` rather than ``"openai:gpt-5"``) is
forwarded down
+the chain as a logical model name: each connection that has no ``model`` of
its own
+resolves that name against its own platform, so one bare name can reach a
primary and
+every fallback without repeating it per connection. It does not matter where
the primary's
+name comes from -- the ``Model`` field on its connection and a ``model_id`` on
the operator
+or hook are forwarded alike. A fallback with its own ``model`` in
+extra always uses that instead -- this is how a fallback pins a spelling the
forwarded
+name would not produce, such as Bedrock's region-prefixed ``us.anthropic.``
model ids. A
+name that already pins a platform (its segment before the first ``:`` is
itself a
+recognized provider, e.g. ``"openai:gpt-5"``) is *not* forwarded; a fallback
with no
+``model`` of its own still raises "no model specified" rather than trying a
prefixed name
+meant for a different provider. See :doc:`connections/pydantic_ai_azure`,
+:doc:`connections/pydantic_ai_bedrock` and
:doc:`connections/pydantic_ai_vertex` for how
+each vendor connection resolves a bare name.
+
+Configure it on the operator
+-----------------------------
+
+``fallback_conn_ids`` is also a parameter on
+:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`,
+:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, their
subclasses,
+and the matching ``@task.llm`` / ``@task.agent`` decorators -- mirroring
``model_id``,
+which is settable at the same two layers:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_fallback.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_llm_fallback_operator_argument]
+ :end-before: [END howto_llm_fallback_operator_argument]
+
+The operator argument overrides the connection's extra field, and passing
``[]``
+explicitly disables a chain configured there -- ``None`` (the default) reads
whatever
+the connection says. Use this when a task should own its own failover order
instead of
+inheriting it from however the connection is configured.
+
+Configure it in code
+--------------------
+
+:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook` also
takes the list
+directly, which is what a task that constructs the hook itself (rather than
through an
+operator) should use:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_fallback.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_llm_fallback_hook_argument]
+ :end-before: [END howto_llm_fallback_hook_argument]
+
+The argument wins over the connection's extra, and passing ``[]`` explicitly
disables a
+chain configured there. Omitting it entirely (``None``) means "use whatever
the connection
+says", which is why the two are not interchangeable.
+
+Where this sits among the retry layers
+--------------------------------------
+
+Three mechanisms handle failure at different time scales, and they compose
rather than
+replace each other:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 25 40 35
+
+ * - Scope
+ - Mechanism
+ - Handles
+ * - Within one model call
+ - ``fallback_conn_ids``
+ - This vendor's API is returning errors; ask the next one (any
``ModelAPIError``,
+ transient or not)
+ * - Within one task attempt
+ - ``timeout`` in pydantic-ai's ``ModelSettings``
+ - This vendor is slow rather than down
+ * - Across task attempts
+ - :doc:`retry_policies` (including ``LLMRetryPolicy``)
+ - Whether this failure is worth retrying at all
+
+A chain does not remove the need for the outer layers. It covers the case
where another
+vendor can answer the same prompt now; a bad prompt, an exhausted quota on
every vendor, or
+a permanent data error still has to be decided by the retry policy.
+
+Adding a chain changes what the retry layer sees. When every connection in the
chain
+fails, the exception the task raises is
``pydantic_ai.exceptions.FallbackExceptionGroup``,
+not the last provider's own exception, so retry rules matched against a
provider-specific
+exception type stop matching. Before adding a chain to a connection that Dags
already use,
+read :doc:`retry_policies` -- the section "When the connection also carries a
fallback
+chain" spells out what to check.
+
+Costs to know before configuring a long chain
+---------------------------------------------
+
+**The timeout multiplies.** pydantic-ai applies a ``ModelSettings`` timeout to
each model
+in the chain, not to the chain as a whole. A 30-second timeout across three
connections is
+a 90-second worst case for one call.
+
+**There is no circuit breaker.** Every call tries the primary first. During an
outage each
+task instance pays the primary's timeout again before failing over, so 500
mapped tasks pay
+it 500 times. Keeping the primary's timeout short bounds both of these.
+
+**Chains are not resolved recursively.** If a connection listed as a fallback
declares its
+own ``fallback_conn_ids``, resolution fails with an error rather than
following it. List
+every provider directly on the primary; a flat chain is the one you can read
off a single
+connection.
+
+**Non-transient errors still walk the whole chain.** Failover triggers on
pydantic-ai's
+``ModelAPIError`` family, which includes ``ModelHTTPError`` -- raised for any
4xx as well as
+5xx. A malformed prompt, an expired key, or a misspelled model name is
therefore retried
Review Comment:
Of these three, only the malformed prompt is shared by every connection. An
expired key is per-connection, so the fallback presents its own and normally
answers, which is the whole point of the chain; and a misspelled model name is
only forwarded when it is bare, since a platform-pinned name is not forwarded
and each fallback uses its own model. The claim that a 4xx walks the chain
holds, but the "N billable calls for a request that was never going to succeed"
cost only follows for an error every connection shares.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -178,24 +303,101 @@ def _provider_factory(pname: str) -> Any:
)
return infer_provider(pname)
- self._model = infer_model(model_name,
provider_factory=_provider_factory)
- return self._model
+ return infer_model(model_name, provider_factory=_provider_factory)
- self._model = infer_model(model_name)
- return self._model
+ return infer_model(model_name)
- def _get_conn_if_model_configured(self) -> Model | None:
- """Return the hook model only when the hook or connection explicitly
configures one."""
- if self.model_id:
- return self.get_conn()
+ def _get_fallback_conn_ids(self) -> list[str]:
+ """Return the configured fallback connection IDs, hook argument
winning over the extra."""
+ if self.fallback_conn_ids is not None:
+ raw: Any = self.fallback_conn_ids
+ else:
+ _, extra = self._get_conn_and_extra()
+ raw = extra.get(FALLBACK_CONN_IDS_EXTRA_KEY)
+ if raw is None:
+ raw = []
+
+ if not isinstance(raw, (list, tuple)) or not all(isinstance(item, str)
and item for item in raw):
+ raise ValueError(
+ f"{FALLBACK_CONN_IDS_EXTRA_KEY} for connection
'{self.llm_conn_id}' must be a list "
+ f"of non-empty connection IDs, got {raw!r}."
+ )
+ return list(raw)
- conn = self.get_connection(self.llm_conn_id)
- self._conn = conn
- self._conn_extra_dejson = conn.extra_dejson
+ def _resolve_fallback_models(self) -> list[Model]:
+ """
+ Resolve one ``Model`` per fallback connection, in the configured order.
+
+ Each connection is resolved through the hook registered for its own
+ ``conn_type``, so a chain can mix providers whose credentials live in
+ different connection fields. The primary's configured model name --
its
+ ``model_id`` argument, or the ``model`` in its own ``extra`` -- is
forwarded
+ to each fallback as a logical model name: a fallback connection with
its own
+ ``model`` in ``extra`` uses that instead, but a fallback with none
falls
+ back to the forwarded name, qualified with *its own* platform prefix.
+ Only a *bare* forwarded name is usable this way -- a forwarded name
that
+ already pins a platform (e.g. ``"openai:gpt-5"``) names a model of the
+ primary's provider, not this fallback's, so it is not applied; that
+ fallback still raises "no model specified" unless its own ``extra``
sets
+ a ``model``. Whether a name already pins a platform is decided by
+ :func:`_has_recognized_provider_prefix`, not by whether it merely
contains a
+ ``:`` -- some vendors' native model ids contain one of their own (e.g.
Bedrock's
+ version-suffixed ``us.anthropic.claude-opus-4-6-v1:0``).
+ """
+ fallback_conn_ids = self._get_fallback_conn_ids()
+ if not fallback_conn_ids:
+ return []
+
+ forwarded_model_id = self._get_configured_model_name()
+
+ models: list[Model] = []
+ seen: set[str] = set()
+ for conn_id in fallback_conn_ids:
+ if conn_id == self.llm_conn_id:
+ raise ValueError(
+ f"Fallback chain for connection '{self.llm_conn_id}' lists
the primary "
+ "connection as one of its own fallbacks; every fallback
must differ from "
+ "the primary."
+ )
+ if conn_id in seen:
+ raise ValueError(
+ f"Fallback chain for connection '{self.llm_conn_id}' lists
'{conn_id}' more "
+ "than once; every entry must be distinct."
+ )
+ seen.add(conn_id)
+
+ # ``BaseHook.get_hook`` dispatches on the connection's
``conn_type`` and does not
+ # constrain the result to this class, so the type has to be
checked here.
+ hook = PydanticAIHook.get_hook(conn_id)
+ if not isinstance(hook, PydanticAIHook):
+ raise ValueError(
+ f"Fallback connection '{conn_id}' resolves to
{type(hook).__name__}, which is "
+ "not a PydanticAIHook. Only pydanticai connection types
can be used as "
+ f"fallbacks for '{self.llm_conn_id}'."
+ )
+ if hook._get_fallback_conn_ids():
+ raise ValueError(
+ f"Fallback connection '{conn_id}' declares its own "
+ f"{FALLBACK_CONN_IDS_EXTRA_KEY}. Chains are not resolved
recursively -- list "
+ f"every provider directly on '{self.llm_conn_id}' instead."
+ )
+
models.append(hook._resolve_own_model(forwarded_model_id=forwarded_model_id))
+
+ self.log.info("Resolved LLM fallback chain: %s", " ->
".join([self.llm_conn_id, *fallback_conn_ids]))
Review Comment:
This fires only after the whole chain resolves, so it is absent in exactly
the case where it would earn its keep. A typo'd fallback id raises
``AirflowNotFoundException("The conn_id `llm_fallback` isn't defined")``, and
on the connection-driven path the docs recommend, the DAG never mentions
`llm_fallback` at all, so the task log gives no hint where the name came from.
Moving the line above the loop would make every mid-resolution failure
attributable.
##########
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:
This parametrize case pins the all-or-nothing rejection as correct, and
`_get_fallback_conn_ids` applies the same `isinstance(item, str) and item`
check whether the list arrived as a hook kwarg or out of connection extra. So
the case a user actually hits -- the conn-field textarea leaving a trailing
blank, which saves `["anthropic_prod", ""]` -- is certified here as an opaque
`ValueError` rather than a working one-entry chain. Whichever way the
blank-entry question is settled, this line has to move with it, so it is worth
deciding now which behaviour is intended.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -65,6 +65,9 @@ class LLMSQLQueryOperator(LLMOperator):
:param llm_conn_id: Connection ID for the LLM provider.
:param model_id: Model identifier (e.g. ``"openai:gpt-4o"``).
Overrides the model stored in the connection's extra field.
+ :param fallback_conn_ids: Connection IDs to fail over to, in order, when
Review Comment:
This `:param:` stops after "Overrides the ``fallback_conn_ids`` set in the
connection's extra field", dropping the sentence `LLMOperator` and
`AgentOperator` both carry: `None` reads the connection's own extra, an
explicit `[]` disables a chain configured there. The behaviour is identical
here, since the kwarg reaches the same hook, so a reader of only this
operator's rendered docs cannot learn that `[]` is the off switch. Same
truncation on `llm_branch.py:48`, `llm_file_analysis.py:50` and
`llm_schema_compare.py:100`.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,33 +178,107 @@ def _get_provider_kwargs(
kwargs["base_url"] = base_url
return kwargs
+ def _get_conn_and_extra(self) -> tuple[Connection, dict[str, Any]]:
+ """Return this hook's connection and its deserialized extra, fetching
at most once."""
+ if self._conn is None:
+ self._conn = self.get_connection(self.llm_conn_id)
+ self._conn_extra_dejson = self._conn.extra_dejson
+ return self._conn, self._conn_extra_dejson
+
def get_conn(self) -> Model:
"""
Return a configured pydantic-ai ``Model``.
- Resolution order:
+ Resolution order for this hook's own connection:
1. **Explicit credentials** — when :meth:`_get_provider_kwargs` returns
a non-empty dict the provider class is instantiated with those
kwargs
and wrapped in a ``provider_factory``.
2. **Default resolution** — delegates to pydantic-ai ``infer_model``
which reads standard env vars (``OPENAI_API_KEY``, ``AWS_PROFILE``,
…).
+ A bare ``model_id`` (one with no recognized platform prefix) is
qualified with
+ this connection's own platform before either of the above -- see the
class
+ docstring's ``model_id`` entry for the resolution and
fallback-forwarding rules.
+
+ When ``fallback_conn_ids`` is configured (on the hook or in the
+ connection's extra) the resolved models are wrapped in a pydantic-ai
+ ``FallbackModel``, so a provider outage moves to the next connection
+ *within the same task attempt* instead of failing the task.
+
+ Two costs of that wrapping are worth knowing before configuring a long
+ chain. A ``timeout`` in ``ModelSettings`` is applied by pydantic-ai to
+ every model in the chain rather than to the chain as a whole, so the
+ worst-case wait is the timeout multiplied by the number of connections.
+ And there is no circuit breaker: every call retries the primary first,
+ so during an outage each task instance pays the primary's timeout
again.
+ Keep the primary's timeout short to bound both.
+
The resolved model is cached for the lifetime of this hook instance.
"""
if self._model is not None:
return self._model
- conn = self.get_connection(self.llm_conn_id) if self._conn is None
else self._conn
- extra: dict[str, Any] = (
- conn.extra_dejson if self._conn_extra_dejson is None else
self._conn_extra_dejson
- )
+ model = self._resolve_own_model()
+ fallback_models = self._resolve_fallback_models()
+ self._model = FallbackModel(model, *fallback_models) if
fallback_models else model
+ return self._model
+
+ def _qualify_model_name(self, model_name: str) -> str:
+ """
+ Prefix a bare model name with this connection's platform.
+
+ A name is treated as already pinning a platform only when the segment
before
+ its first ``:`` is itself a provider pydantic-ai recognizes (e.g.
+ ``"openai:gpt-4"``) -- see :func:`_has_recognized_provider_prefix`.
Everything
+ else is a bare name, even one that happens to contain a ``:`` of its
own (e.g.
+ Bedrock's version-suffixed ``"us.anthropic.claude-opus-4-6-v1:0"``),
and is
+ prefixed with :attr:`model_provider`; the generic ``pydanticai``
connection type
+ has no platform of its own (``model_provider`` is ``None``), so a bare
name there
+ raises instead of reaching pydantic-ai's own, less actionable
``Unknown model``
+ error.
+ """
+ if _has_recognized_provider_prefix(model_name):
+ return model_name
+ if self.model_provider is None:
+ raise ValueError(
+ f"Connection '{self.llm_conn_id}' has no default model
provider, so the bare "
+ f"model name '{model_name}' cannot be resolved. Use a vendor
connection type "
+ "(Azure/Bedrock/Vertex) or set an explicit 'provider:model'
string."
+ )
+ return f"{self.model_provider}:{model_name}"
+
+ def _get_configured_model_name(self) -> str | KnownModelName | None:
+ """Return the model name this connection configures, hook argument
winning over the extra."""
+ if self.model_id:
+ return self.model_id
+ _, extra = self._get_conn_and_extra()
+ return extra.get("model")
+
+ def _resolve_own_model(self, *, forwarded_model_id: str | None = None) ->
Model:
+ """
+ Resolve the ``Model`` for this hook's own connection, ignoring any
fallback chain.
+
+ :param forwarded_model_id: The primary connection's configured model
name,
+ forwarded down a fallback chain by
:meth:`_resolve_fallback_models`.
+ Used only when this
+ connection configures no ``model_id``/``model`` of its own, and
only when it
+ is a bare name: a name that already pins a platform (see
+ :func:`_has_recognized_provider_prefix`) names a model of the
*primary's*
+ provider, not this connection's, so it is not forwarded -- this
connection
+ still raises "no model specified" in that case.
+ """
+ conn, extra = self._get_conn_and_extra()
- model_name: str | KnownModelName = self.model_id or extra.get("model",
"")
+ model_name: str | KnownModelName | None =
self._get_configured_model_name()
+ if not model_name and forwarded_model_id and not
_has_recognized_provider_prefix(forwarded_model_id):
Review Comment:
A colon-bearing name is classified as bare precisely because it is a native
vendor id, and a native vendor id is never valid on another platform, yet it is
still forwarded across platforms. A Bedrock primary with
`us.anthropic.claude-opus-4-5-v1:0` hands an Azure fallback that has no `model`
of its own `azure:us.anthropic.claude-opus-4-5-v1:0`, which resolves without
complaint and then 404s during the outage the chain existed to cover, while
`test_connection` reports the chain healthy. Gating the forward on the two
hooks sharing a `model_provider` when the name contains a colon would refuse
only the combinations that provably cannot work, leaving Bedrock-to-Bedrock and
the colon-free cross-vendor case untouched.
##########
providers/common/ai/docs/provider_fallback.rst:
##########
@@ -0,0 +1,208 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Provider fallback
+=================
+
+A single ``llm_conn_id`` gives a task one provider. When that provider is
down, the task
+fails and retries into the same outage. ``fallback_conn_ids`` gives the
connection an
+ordered list of other connections to try, so a provider outage moves to the
next vendor
+inside the same task attempt.
+
+Configure it on the connection
+------------------------------
+
+Put the chain in the primary connection's extra:
+
+.. code-block:: json
+
+ {
+ "model": "openai:gpt-5",
+ "fallback_conn_ids": ["anthropic_prod", "bedrock_dr"]
+ }
+
+Every entry is an Airflow connection ID, resolved through the hook registered
for its own
+connection type. A chain can therefore mix vendors whose credentials live in
different
+connection fields — ``pydanticai`` for OpenAI, ``pydanticai_bedrock`` for a
Bedrock
+standby — without the Dag knowing anything about either.
+
+That is the point of configuring it here rather than in Dag code: the Dag
keeps naming one
+connection, and whoever administers the connections owns the failover
topology. Changing a
+standby provider is a connection edit, not a Dag deployment.
+
+A *bare* model name (e.g. ``"gpt-5"`` rather than ``"openai:gpt-5"``) is
forwarded down
+the chain as a logical model name: each connection that has no ``model`` of
its own
+resolves that name against its own platform, so one bare name can reach a
primary and
+every fallback without repeating it per connection. It does not matter where
the primary's
+name comes from -- the ``Model`` field on its connection and a ``model_id`` on
the operator
+or hook are forwarded alike. A fallback with its own ``model`` in
+extra always uses that instead -- this is how a fallback pins a spelling the
forwarded
+name would not produce, such as Bedrock's region-prefixed ``us.anthropic.``
model ids. A
+name that already pins a platform (its segment before the first ``:`` is
itself a
+recognized provider, e.g. ``"openai:gpt-5"``) is *not* forwarded; a fallback
with no
+``model`` of its own still raises "no model specified" rather than trying a
prefixed name
+meant for a different provider. See :doc:`connections/pydantic_ai_azure`,
+:doc:`connections/pydantic_ai_bedrock` and
:doc:`connections/pydantic_ai_vertex` for how
+each vendor connection resolves a bare name.
+
+Configure it on the operator
+-----------------------------
+
+``fallback_conn_ids`` is also a parameter on
+:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`,
+:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, their
subclasses,
+and the matching ``@task.llm`` / ``@task.agent`` decorators -- mirroring
``model_id``,
+which is settable at the same two layers:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_fallback.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_llm_fallback_operator_argument]
+ :end-before: [END howto_llm_fallback_operator_argument]
+
+The operator argument overrides the connection's extra field, and passing
``[]``
+explicitly disables a chain configured there -- ``None`` (the default) reads
whatever
+the connection says. Use this when a task should own its own failover order
instead of
+inheriting it from however the connection is configured.
+
+Configure it in code
+--------------------
+
+:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook` also
takes the list
+directly, which is what a task that constructs the hook itself (rather than
through an
+operator) should use:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_fallback.py
Review Comment:
The dag this pulls in passes `fallback_conn_ids=["llm_fallback"]` to a hook
whose connection `llm_primary_down` already carries that same chain in its
extra, so nothing in the snippet shows the argument winning over the extra that
the paragraph below then describes. The operator example above gets this right
by pointing at `llm_primary_down_no_chain`; the same switch here would make the
snippet carry its claim.
##########
providers/common/ai/docs/provider_fallback.rst:
##########
@@ -0,0 +1,208 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Provider fallback
+=================
+
+A single ``llm_conn_id`` gives a task one provider. When that provider is
down, the task
+fails and retries into the same outage. ``fallback_conn_ids`` gives the
connection an
+ordered list of other connections to try, so a provider outage moves to the
next vendor
+inside the same task attempt.
+
+Configure it on the connection
+------------------------------
+
+Put the chain in the primary connection's extra:
+
+.. code-block:: json
+
+ {
+ "model": "openai:gpt-5",
+ "fallback_conn_ids": ["anthropic_prod", "bedrock_dr"]
+ }
+
+Every entry is an Airflow connection ID, resolved through the hook registered
for its own
+connection type. A chain can therefore mix vendors whose credentials live in
different
+connection fields — ``pydanticai`` for OpenAI, ``pydanticai_bedrock`` for a
Bedrock
+standby — without the Dag knowing anything about either.
+
+That is the point of configuring it here rather than in Dag code: the Dag
keeps naming one
+connection, and whoever administers the connections owns the failover
topology. Changing a
+standby provider is a connection edit, not a Dag deployment.
+
+A *bare* model name (e.g. ``"gpt-5"`` rather than ``"openai:gpt-5"``) is
forwarded down
+the chain as a logical model name: each connection that has no ``model`` of
its own
+resolves that name against its own platform, so one bare name can reach a
primary and
+every fallback without repeating it per connection. It does not matter where
the primary's
+name comes from -- the ``Model`` field on its connection and a ``model_id`` on
the operator
+or hook are forwarded alike. A fallback with its own ``model`` in
+extra always uses that instead -- this is how a fallback pins a spelling the
forwarded
+name would not produce, such as Bedrock's region-prefixed ``us.anthropic.``
model ids. A
+name that already pins a platform (its segment before the first ``:`` is
itself a
+recognized provider, e.g. ``"openai:gpt-5"``) is *not* forwarded; a fallback
with no
+``model`` of its own still raises "no model specified" rather than trying a
prefixed name
+meant for a different provider. See :doc:`connections/pydantic_ai_azure`,
+:doc:`connections/pydantic_ai_bedrock` and
:doc:`connections/pydantic_ai_vertex` for how
+each vendor connection resolves a bare name.
+
+Configure it on the operator
+-----------------------------
+
+``fallback_conn_ids`` is also a parameter on
+:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`,
+:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, their
subclasses,
+and the matching ``@task.llm`` / ``@task.agent`` decorators -- mirroring
``model_id``,
+which is settable at the same two layers:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_fallback.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_llm_fallback_operator_argument]
+ :end-before: [END howto_llm_fallback_operator_argument]
+
+The operator argument overrides the connection's extra field, and passing
``[]``
+explicitly disables a chain configured there -- ``None`` (the default) reads
whatever
+the connection says. Use this when a task should own its own failover order
instead of
+inheriting it from however the connection is configured.
+
+Configure it in code
+--------------------
+
+:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook` also
takes the list
+directly, which is what a task that constructs the hook itself (rather than
through an
+operator) should use:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_fallback.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_llm_fallback_hook_argument]
+ :end-before: [END howto_llm_fallback_hook_argument]
+
+The argument wins over the connection's extra, and passing ``[]`` explicitly
disables a
+chain configured there. Omitting it entirely (``None``) means "use whatever
the connection
+says", which is why the two are not interchangeable.
+
+Where this sits among the retry layers
+--------------------------------------
+
+Three mechanisms handle failure at different time scales, and they compose
rather than
+replace each other:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 25 40 35
+
+ * - Scope
+ - Mechanism
+ - Handles
+ * - Within one model call
+ - ``fallback_conn_ids``
+ - This vendor's API is returning errors; ask the next one (any
``ModelAPIError``,
+ transient or not)
+ * - Within one task attempt
+ - ``timeout`` in pydantic-ai's ``ModelSettings``
+ - This vendor is slow rather than down
+ * - Across task attempts
+ - :doc:`retry_policies` (including ``LLMRetryPolicy``)
+ - Whether this failure is worth retrying at all
+
+A chain does not remove the need for the outer layers. It covers the case
where another
+vendor can answer the same prompt now; a bad prompt, an exhausted quota on
every vendor, or
+a permanent data error still has to be decided by the retry policy.
+
+Adding a chain changes what the retry layer sees. When every connection in the
chain
+fails, the exception the task raises is
``pydantic_ai.exceptions.FallbackExceptionGroup``,
+not the last provider's own exception, so retry rules matched against a
provider-specific
+exception type stop matching. Before adding a chain to a connection that Dags
already use,
+read :doc:`retry_policies` -- the section "When the connection also carries a
fallback
+chain" spells out what to check.
+
+Costs to know before configuring a long chain
+---------------------------------------------
+
+**The timeout multiplies.** pydantic-ai applies a ``ModelSettings`` timeout to
each model
+in the chain, not to the chain as a whole. A 30-second timeout across three
connections is
+a 90-second worst case for one call.
+
+**There is no circuit breaker.** Every call tries the primary first. During an
outage each
+task instance pays the primary's timeout again before failing over, so 500
mapped tasks pay
+it 500 times. Keeping the primary's timeout short bounds both of these.
+
+**Chains are not resolved recursively.** If a connection listed as a fallback
declares its
+own ``fallback_conn_ids``, resolution fails with an error rather than
following it. List
+every provider directly on the primary; a flat chain is the one you can read
off a single
+connection.
+
+**Non-transient errors still walk the whole chain.** Failover triggers on
pydantic-ai's
+``ModelAPIError`` family, which includes ``ModelHTTPError`` -- raised for any
4xx as well as
+5xx. A malformed prompt, an expired key, or a misspelled model name is
therefore retried
+against every connection in the chain before the task sees the failure: N
requests, N
+timeouts, and N billable calls for a request that was never going to succeed.
Keep chains
+short, and put deterministic rules for those errors in :doc:`retry_policies`.
+
+**Airflow's task-level** ``retries`` **multiplies on top of the chain.** A
task with
+``retries=5`` gets up to six attempts -- the initial attempt plus five retries
-- before
+Airflow marks it failed, and each attempt walks the whole chain again if every
connection
+is still down. Against the three-connection chain in the example above (the
primary plus
Review Comment:
Both exampleinclude blocks above use a one-primary-plus-one-fallback chain,
so a reader who checks "the example above" will conclude the arithmetic is
wrong. The only three-connection chain on the page is the JSON snippet near the
top, so naming that one explicitly would fix it. The 6 x 3 = 18 itself is right.
##########
providers/common/ai/docs/provider_fallback.rst:
##########
@@ -0,0 +1,208 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Provider fallback
+=================
+
+A single ``llm_conn_id`` gives a task one provider. When that provider is
down, the task
+fails and retries into the same outage. ``fallback_conn_ids`` gives the
connection an
+ordered list of other connections to try, so a provider outage moves to the
next vendor
+inside the same task attempt.
+
+Configure it on the connection
+------------------------------
+
+Put the chain in the primary connection's extra:
+
+.. code-block:: json
+
+ {
+ "model": "openai:gpt-5",
+ "fallback_conn_ids": ["anthropic_prod", "bedrock_dr"]
+ }
+
+Every entry is an Airflow connection ID, resolved through the hook registered
for its own
+connection type. A chain can therefore mix vendors whose credentials live in
different
+connection fields — ``pydanticai`` for OpenAI, ``pydanticai_bedrock`` for a
Bedrock
+standby — without the Dag knowing anything about either.
+
+That is the point of configuring it here rather than in Dag code: the Dag
keeps naming one
+connection, and whoever administers the connections owns the failover
topology. Changing a
+standby provider is a connection edit, not a Dag deployment.
+
+A *bare* model name (e.g. ``"gpt-5"`` rather than ``"openai:gpt-5"``) is
forwarded down
+the chain as a logical model name: each connection that has no ``model`` of
its own
+resolves that name against its own platform, so one bare name can reach a
primary and
+every fallback without repeating it per connection. It does not matter where
the primary's
+name comes from -- the ``Model`` field on its connection and a ``model_id`` on
the operator
+or hook are forwarded alike. A fallback with its own ``model`` in
+extra always uses that instead -- this is how a fallback pins a spelling the
forwarded
+name would not produce, such as Bedrock's region-prefixed ``us.anthropic.``
model ids. A
+name that already pins a platform (its segment before the first ``:`` is
itself a
+recognized provider, e.g. ``"openai:gpt-5"``) is *not* forwarded; a fallback
with no
+``model`` of its own still raises "no model specified" rather than trying a
prefixed name
+meant for a different provider. See :doc:`connections/pydantic_ai_azure`,
+:doc:`connections/pydantic_ai_bedrock` and
:doc:`connections/pydantic_ai_vertex` for how
+each vendor connection resolves a bare name.
+
+Configure it on the operator
+-----------------------------
+
+``fallback_conn_ids`` is also a parameter on
+:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`,
+:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, their
subclasses,
+and the matching ``@task.llm`` / ``@task.agent`` decorators -- mirroring
``model_id``,
+which is settable at the same two layers:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_fallback.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_llm_fallback_operator_argument]
+ :end-before: [END howto_llm_fallback_operator_argument]
+
+The operator argument overrides the connection's extra field, and passing
``[]``
+explicitly disables a chain configured there -- ``None`` (the default) reads
whatever
+the connection says. Use this when a task should own its own failover order
instead of
+inheriting it from however the connection is configured.
+
+Configure it in code
+--------------------
+
+:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook` also
takes the list
+directly, which is what a task that constructs the hook itself (rather than
through an
+operator) should use:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_fallback.py
+ :language: python
+ :dedent: 0
+ :start-after: [START howto_llm_fallback_hook_argument]
+ :end-before: [END howto_llm_fallback_hook_argument]
+
+The argument wins over the connection's extra, and passing ``[]`` explicitly
disables a
+chain configured there. Omitting it entirely (``None``) means "use whatever
the connection
+says", which is why the two are not interchangeable.
+
+Where this sits among the retry layers
+--------------------------------------
+
+Three mechanisms handle failure at different time scales, and they compose
rather than
+replace each other:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 25 40 35
+
+ * - Scope
+ - Mechanism
+ - Handles
+ * - Within one model call
+ - ``fallback_conn_ids``
+ - This vendor's API is returning errors; ask the next one (any
``ModelAPIError``,
+ transient or not)
+ * - Within one task attempt
+ - ``timeout`` in pydantic-ai's ``ModelSettings``
+ - This vendor is slow rather than down
+ * - Across task attempts
+ - :doc:`retry_policies` (including ``LLMRetryPolicy``)
+ - Whether this failure is worth retrying at all
+
+A chain does not remove the need for the outer layers. It covers the case
where another
+vendor can answer the same prompt now; a bad prompt, an exhausted quota on
every vendor, or
+a permanent data error still has to be decided by the retry policy.
+
+Adding a chain changes what the retry layer sees. When every connection in the
chain
+fails, the exception the task raises is
``pydantic_ai.exceptions.FallbackExceptionGroup``,
+not the last provider's own exception, so retry rules matched against a
provider-specific
+exception type stop matching. Before adding a chain to a connection that Dags
already use,
+read :doc:`retry_policies` -- the section "When the connection also carries a
fallback
+chain" spells out what to check.
+
+Costs to know before configuring a long chain
+---------------------------------------------
+
+**The timeout multiplies.** pydantic-ai applies a ``ModelSettings`` timeout to
each model
+in the chain, not to the chain as a whole. A 30-second timeout across three
connections is
+a 90-second worst case for one call.
+
+**There is no circuit breaker.** Every call tries the primary first. During an
outage each
+task instance pays the primary's timeout again before failing over, so 500
mapped tasks pay
+it 500 times. Keeping the primary's timeout short bounds both of these.
+
+**Chains are not resolved recursively.** If a connection listed as a fallback
declares its
+own ``fallback_conn_ids``, resolution fails with an error rather than
following it. List
+every provider directly on the primary; a flat chain is the one you can read
off a single
+connection.
+
+**Non-transient errors still walk the whole chain.** Failover triggers on
pydantic-ai's
+``ModelAPIError`` family, which includes ``ModelHTTPError`` -- raised for any
4xx as well as
+5xx. A malformed prompt, an expired key, or a misspelled model name is
therefore retried
+against every connection in the chain before the task sees the failure: N
requests, N
+timeouts, and N billable calls for a request that was never going to succeed.
Keep chains
+short, and put deterministic rules for those errors in :doc:`retry_policies`.
+
+**Airflow's task-level** ``retries`` **multiplies on top of the chain.** A
task with
+``retries=5`` gets up to six attempts -- the initial attempt plus five retries
-- before
+Airflow marks it failed, and each attempt walks the whole chain again if every
connection
+is still down. Against the three-connection chain in the example above (the
primary plus
+two fallbacks), that is up to 18 upstream calls, not 3, before the task is
finally marked
+failed.
+
+**A bad fallback connection fails the whole chain, including a healthy
primary.** The
+primary and every fallback are resolved eagerly, before any of them is called,
so a
+misspelled fallback ``conn_id`` or a fallback connection missing its ``model``
raises
+immediately -- the task never reaches the primary, even though the primary
itself would
+have answered fine. Run ``test_connection`` on the primary to catch this
before it costs a
+task; see *Verifying a chain* below.
+
+Verifying a chain
+-----------------
+
+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 not reported this way -- the hook catches that and retries
with the
+env-var-based provider constructor, logging a warning either way; that retry
still raises a
+``pydantic_ai.exceptions.UserError`` if the required env var is also missing,
so check the
+logs for that failure mode rather than relying on ``test_connection``. It also
does not call the provider, so a
Review Comment:
`test_connection` does surface that `UserError`: it wraps `get_conn()` in
`except Exception` and returns `(False, str(e))`, and `UserError` subclasses
`RuntimeError`. The case it genuinely cannot show is the opposite one, where
the `TypeError` is swallowed because the env-var retry succeeded, so the
connection tests green while the credentials you set were ignored. Worth
pointing the reader at the logs for that case instead.
--
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]
