Lee-W commented on code in PR #72156:
URL: https://github.com/apache/airflow/pull/72156#discussion_r4045820280
##########
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:
`_qualify_model_name` now takes a `forwarded_from_conn_id`, so a name that
arrived from the primary is attributed to it instead of to the connection that
never carried it.
Left `policies/retry.py` for a follow-up as suggested.
--
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]
