Lee-W commented on code in PR #72156:
URL: https://github.com/apache/airflow/pull/72156#discussion_r4045841090
##########
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:
Moved above the loop, and reworded to match — it says resolving rather than
resolved now. A test covers the line surviving a mid-chain
`AirflowNotFoundException`. One case it still cannot cover: a `ValueError` out
of `_get_fallback_conn_ids` itself, since the message needs the list that call
produces.
--
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]
