Lee-W commented on code in PR #72156:
URL: https://github.com/apache/airflow/pull/72156#discussion_r4045796264
##########
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:
Pinned explicitly. the `kwarg` is there and its default is already
`(ModelAPIError,)`, so no version gate was needed.
--
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]
