kaxil commented on code in PR #71437:
URL: https://github.com/apache/airflow/pull/71437#discussion_r4032377795
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -540,3 +697,21 @@ def _get_provider_kwargs(
)
return kwargs
+
+
+_PROVIDER_CONNECTION_CONFIGS: dict[str | None, _ProviderConnectionConfig] = {
+ "azure":
_ProviderConnectionConfig(PydanticAIAzureHook._get_provider_kwargs),
+ "azure-responses":
_ProviderConnectionConfig(PydanticAIAzureHook._get_provider_kwargs),
+ "bedrock": _ProviderConnectionConfig(
+ PydanticAIBedrockHook._get_provider_kwargs,
+ ("api_key", "base_url", "region_name"),
+ ),
+ "google": _ProviderConnectionConfig(
+ PydanticAIVertexHook._get_google_provider_kwargs,
+ ("api_key", "base_url"),
Review Comment:
`replacement_fields` drives a warning that inspects only `conn.password` and
`conn.host`, and those two are `hidden_fields` on `pydanticai_vertex`, so on
that connection type it can never fire. Meanwhile `_get_google_provider_kwargs`
reads just `api_key` and `base_url` out of `extra` (correctly, `GoogleProvider`
takes no `project`, `location` or `credentials` on 2.31.1). A Vertex connection
that fills in the GCP Project and Location conn-fields and sets `model:
"google:gemini-2.0-flash"` therefore gets `{}`, no `provider_factory`, and ADC,
with two documented form fields discarded and an empty task log.
`service_account_info` goes the same way.
That combination is first-class here, not a misconfiguration:
`test_get_conn_google_uses_google_provider_signature` exercises exactly it. And
`main` was louder on this path, since the old self-dispatch produced `{project,
location}`, hit `GoogleProvider`'s TypeError and logged `Provider 'google'
rejected kwargs ['project', 'location']; falling back to env-var auth`.
A second tuple on `_ProviderConnectionConfig` for the extra keys a mapper is
known to drop, warned about next to the existing one, closes it inside the
structure already here.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -540,3 +697,21 @@ def _get_provider_kwargs(
)
return kwargs
+
+
+_PROVIDER_CONNECTION_CONFIGS: dict[str | None, _ProviderConnectionConfig] = {
+ "azure":
_ProviderConnectionConfig(PydanticAIAzureHook._get_provider_kwargs),
+ "azure-responses":
_ProviderConnectionConfig(PydanticAIAzureHook._get_provider_kwargs),
+ "bedrock": _ProviderConnectionConfig(
+ PydanticAIBedrockHook._get_provider_kwargs,
+ ("api_key", "base_url", "region_name"),
Review Comment:
The Bedrock tuple lists `api_key`, `base_url` and `region_name`, but the
mapper also reads `aws_access_key_id`, `aws_secret_access_key`,
`aws_session_token` and `profile_name`. The warning presents this list as the
provider-specific values that replace the ignored ones, so a Bedrock user whose
Password held an access key is steered to `api_key`, which is the bearer token
and takes precedence over IAM keys once set. Listing all seven keeps the
message from suggesting a different auth mode.
##########
providers/common/ai/docs/hooks/pydantic_ai.rst:
##########
@@ -56,6 +57,46 @@ The model can be specified at three levels (highest priority
first):
# Override with a specific model
hook = PydanticAIHook(llm_conn_id="my_llm",
model_id="anthropic:claude-opus-4-6")
+Embedding Models
+----------------
+
+Set ``embed_model_id`` on the hook or ``embed_model`` in the connection's
extra JSON,
+then call ``get_embedder()``. Use ``embed_conn_id`` when the embedding
provider uses
+different credentials or an endpoint from the LLM provider; it defaults to
+``llm_conn_id``. Different LLM and embedding providers require separate
connections
+so credentials cannot be reused for the wrong provider. Equivalent OpenAI and
Azure
+chat/response prefixes can share their provider's embedding connection. Local
+``sentence-transformers:`` embeddings can also share the LLM connection
because they
+do not use provider credentials. The resolved ``Embedder`` is cached on the
hook instance.
+
+.. code-block:: python
+
+ hook = PydanticAIHook(
+ llm_conn_id="my_llm",
+ embed_conn_id="my_embeddings",
+ embed_model_id="openai:text-embedding-3-small",
+ )
+ embedder = hook.get_embedder()
+ result = embedder.embed_query_sync("Apache Airflow orchestrates
workflows.")
+ embedding = result.embeddings[0]
+
+Keyword arguments accepted by pydantic-ai's `Embedder constructor
+<https://ai.pydantic.dev/api/embeddings/#pydantic_ai.embeddings.Embedder.__init__>`__
+can be passed directly to ``get_embedder()``. These currently include
``settings``,
+``defer_model_check``, and ``instrument``. Caller-supplied ``instrument`` takes
+precedence over Airflow's automatic instrumentation. Repeated calls with the
same
+arguments return the cached instance; passing different arguments creates and
caches
+a new instance.
+
+.. code-block:: python
+
+ from pydantic_ai.embeddings import EmbeddingSettings
+
+ embedder = hook.get_embedder(
+ settings=EmbeddingSettings(dimensions=512),
+ defer_model_check=False,
Review Comment:
`defer_model_check=False` is a no-op through `get_embedder()`.
`Embedder.__init__` does `self._model = model if defer_model_check else
infer_embedding_model(model)`, and `infer_embedding_model` returns an
`EmbeddingModel` instance unchanged, which is always what the hook passes. The
flag only does anything when `Embedder` is given a model string. Dropping it
from the example avoids teaching a setting that cannot fire here.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -239,6 +277,411 @@ def test_get_conn_caches_model(self, mock_infer_model):
assert first is second
mock_infer_model.assert_called_once()
+ @pytest.mark.parametrize(
+ ("model_name", "provider_name", "replacement_fields"),
+ [
+ (
+ "bedrock:us.anthropic.claude-opus-4-6-v1:0",
+ "bedrock",
+ ["api_key", "base_url", "region_name"],
+ ),
+ ("google:gemini-2.0-flash", "google", ["api_key", "base_url"]),
+ (
+ "google-cloud:gemini-2.0-flash",
+ "google-cloud",
+ ["api_key", "base_url", "project", "location"],
+ ),
+ ],
+ )
+ @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model",
autospec=True)
+ def test_generic_connection_warns_when_password_and_host_are_ignored(
+ self, mock_infer_model, model_name, provider_name, replacement_fields
+ ):
+ mock_infer_model.return_value = MagicMock(spec=Model)
+ hook = PydanticAIHook(llm_conn_id="test_conn")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="pydanticai",
+ password="provider-key",
+ host="https://provider.example.com";,
+ extra=json.dumps({"model": model_name}),
+ )
+
+ with (
+ patch.object(hook, "get_connection", return_value=conn),
+ patch.object(hook.log, "warning", autospec=True) as mock_warning,
+ ):
+ hook.get_conn()
+
+ mock_warning.assert_called_once_with(
Review Comment:
This is one assertion away from also pinning extra-only precedence, which is
the thing round 3 settled. Nothing in the file currently fails if the Google
mappers grow the fill-in-when-absent `conn.password` / `conn.host` fallback:
the two direct mapper tests pass `None, None` for those parameters, the
`google` and `google-cloud` rows of
`test_embedding_credentials_are_mapped_by_model_provider` already carry
`api_key` and `base_url` in `extra` so a fill-if-absent branch never fires, and
this test sets the connection fields but checks only the warning.
All three rows here limit `extra` to `model`, so each resolves to `{}` and
no `provider_factory` today. Adding
`mock_infer_model.assert_called_once_with(model_name)` below covers bedrock,
google and google-cloud in one line and turns that mutation red.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,6 +159,96 @@ def _get_provider_kwargs(
kwargs["base_url"] = base_url
return kwargs
+ def _get_cached_connection(self, conn_id: str) -> Connection:
+ if conn_id not in self._connections:
+ self._connections[conn_id] = self.get_connection(conn_id)
+ return self._connections[conn_id]
+
+ def _get_cached_connection_extra_dejson(self, conn_id: str) -> dict[str,
Any]:
+ if conn_id not in self._connection_extra_dejson:
+ conn = self._get_cached_connection(conn_id)
+ self._connection_extra_dejson[conn_id] = conn.extra_dejson
+ return self._connection_extra_dejson[conn_id]
+
+ def _warn_if_vertexai_field_ignored(self, extra: dict[str, Any]) -> None:
+ if extra.get("vertexai") is not None:
+ self.log.warning(
+ "The 'vertexai' connection field is ignored; Vertex AI vs.
Generative Language "
+ "API mode is now selected via the model prefix
('google-cloud:' vs. 'google:')."
+ )
+
+ def _get_provider_kwargs_for_model(
+ self, conn: Connection, model_name: str, extra: dict[str, Any]
+ ) -> dict[str, Any]:
+ provider_name, _ = parse_model_id(model_name)
+ provider_config = _PROVIDER_CONNECTION_CONFIGS.get(provider_name)
+ self._warn_if_vertexai_field_ignored(extra)
+ if provider_config is None:
+ return PydanticAIHook._get_provider_kwargs(conn.password,
conn.host, extra)
+ if provider_config.replacement_fields:
+ ignored_fields = [
+ field for field, value in (("password", conn.password),
("host", conn.host)) if value
+ ]
+ if ignored_fields:
+ self.log.warning(
+ "Connection fields are ignored for provider %r on
connection %r; "
+ "ignored fields: %s; configure these provider-specific
values in extra: %s",
+ provider_name,
+ conn.conn_id,
+ ignored_fields,
+ list(provider_config.replacement_fields),
+ )
+ return provider_config.get_kwargs(conn.password, conn.host, extra)
+
+ def _get_provider_factory_for_model(
+ self, conn: Connection, model_name: str, extra: dict[str, Any]
+ ) -> Callable[[str], Any] | None:
+ provider_name, _ = parse_model_id(model_name)
+ if provider_name == "sentence-transformers":
+ return None
+
+ provider_kwargs = self._get_provider_kwargs_for_model(conn,
model_name, extra)
+ if not provider_kwargs:
+ return None
+
+ self.log.info(
+ "Using explicit connection credentials for model '%s': %s",
+ model_name,
+ list(provider_kwargs),
+ )
+
+ def create_provider(provider: str) -> Any:
+ try:
+ return infer_provider_class(provider)(**provider_kwargs)
+ except TypeError as e:
+ raise TypeError(
+ f"Provider {provider!r} rejected connection
{conn.conn_id!r} fields "
+ f"mapped to kwargs {sorted(provider_kwargs)}"
Review Comment:
This names the kwargs that were supplied but not the one the provider
actually rejected, and the cause does not reach the user. `test_connection` at
:461 returns `str(e)`, not a formatted traceback, so the Test button in the UI
shows only this line. Reproduced against 2.31.1 with a generic connection
carrying an API key and a Host, `embed_model: cohere:embed-v4.0`:
```
str(e): Provider 'cohere' rejected connection 'my_conn' fields mapped
to kwargs ['api_key', 'base_url']
str(e.__cause__): CohereProvider.__init__() got an unexpected keyword
argument 'base_url'
```
Both kwargs came from the connection and only `base_url` is wrong, so the
message as rendered points at the key. `self_hosted_models.rst:320`, added by
this PR, tells the reader the TypeError "identifies the rejected connection
fields", which describes the cause line rather than this one. Appending `: {e}`
to the message restores it.
The same one-liner covers the other half: the `try` wraps the whole
constructor call, so a TypeError raised inside a provider's body (an SDK
version skew, say) currently gets relabelled as a connection-field problem with
no hint otherwise.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,6 +159,96 @@ def _get_provider_kwargs(
kwargs["base_url"] = base_url
return kwargs
+ def _get_cached_connection(self, conn_id: str) -> Connection:
+ if conn_id not in self._connections:
+ self._connections[conn_id] = self.get_connection(conn_id)
+ return self._connections[conn_id]
+
+ def _get_cached_connection_extra_dejson(self, conn_id: str) -> dict[str,
Any]:
+ if conn_id not in self._connection_extra_dejson:
+ conn = self._get_cached_connection(conn_id)
+ self._connection_extra_dejson[conn_id] = conn.extra_dejson
+ return self._connection_extra_dejson[conn_id]
+
+ def _warn_if_vertexai_field_ignored(self, extra: dict[str, Any]) -> None:
+ if extra.get("vertexai") is not None:
+ self.log.warning(
+ "The 'vertexai' connection field is ignored; Vertex AI vs.
Generative Language "
+ "API mode is now selected via the model prefix
('google-cloud:' vs. 'google:')."
+ )
+
+ def _get_provider_kwargs_for_model(
+ self, conn: Connection, model_name: str, extra: dict[str, Any]
+ ) -> dict[str, Any]:
+ provider_name, _ = parse_model_id(model_name)
+ provider_config = _PROVIDER_CONNECTION_CONFIGS.get(provider_name)
+ self._warn_if_vertexai_field_ignored(extra)
+ if provider_config is None:
+ return PydanticAIHook._get_provider_kwargs(conn.password,
conn.host, extra)
+ if provider_config.replacement_fields:
+ ignored_fields = [
+ field for field, value in (("password", conn.password),
("host", conn.host)) if value
+ ]
+ if ignored_fields:
+ self.log.warning(
+ "Connection fields are ignored for provider %r on
connection %r; "
+ "ignored fields: %s; configure these provider-specific
values in extra: %s",
+ provider_name,
+ conn.conn_id,
+ ignored_fields,
+ list(provider_config.replacement_fields),
+ )
+ return provider_config.get_kwargs(conn.password, conn.host, extra)
+
+ def _get_provider_factory_for_model(
+ self, conn: Connection, model_name: str, extra: dict[str, Any]
+ ) -> Callable[[str], Any] | None:
+ provider_name, _ = parse_model_id(model_name)
+ if provider_name == "sentence-transformers":
+ return None
+
+ provider_kwargs = self._get_provider_kwargs_for_model(conn,
model_name, extra)
+ if not provider_kwargs:
+ return None
+
+ self.log.info(
+ "Using explicit connection credentials for model '%s': %s",
+ model_name,
+ list(provider_kwargs),
+ )
+
+ def create_provider(provider: str) -> Any:
+ try:
+ return infer_provider_class(provider)(**provider_kwargs)
+ except TypeError as e:
+ raise TypeError(
+ f"Provider {provider!r} rejected connection
{conn.conn_id!r} fields "
+ f"mapped to kwargs {sorted(provider_kwargs)}"
+ ) from e
+
+ return create_provider
+
+ def _validate_embedding_connection_provider(self, embed_model_name: str,
extra: dict[str, Any]) -> None:
+ if self.embed_conn_id != self.llm_conn_id:
+ return
+
+ llm_model_name = self.model_id or extra.get("model", "")
+ if not llm_model_name:
+ return
+
+ llm_provider, _ = parse_model_id(llm_model_name)
+ embed_provider, _ = parse_model_id(embed_model_name)
Review Comment:
`parse_model_id` returns `(None, name)` for a value with no prefix, so
`{"model": "openai:gpt-4o", "embed_model": "text-embedding-3-small"}` falls
into the mismatch branch and reports `configures different LLM and embedding
providers ('openai' and None). Set embed_conn_id to a separate connection for
the embedding provider.` A separate connection does not help; drop the `model`
key and pydantic-ai gives the accurate `ValueError: You must provide a provider
prefix when specifying an embedding model name` instead.
The prefixless form is a plausible thing to type into this field: the
`llamaindex` conn type in this same `provider.yaml` documents `embed_model` as
`text-embedding-3-small`. `if embed_provider is None: return` hands it back to
pydantic-ai.
##########
providers/common/ai/docs/connections/pydantic_ai.rst:
##########
@@ -67,14 +88,24 @@ Host (optional)
Extra (JSON, optional)
A JSON object with additional configuration. Programmatic users can set the
- model directly in extra:
+ LLM and embedding models directly in extra:
.. code-block:: json
- {"model": "openai:gpt-5.6-sol"}
+ {
+ "model": "openai:gpt-5.6-sol",
+ "embed_model": "openai:text-embedding-3-small"
+ }
- When using the UI, the "Model" field above writes to this same location
- automatically.
+ When using the UI, the "Model" and "Embedding Model" fields above write to
+ this same location automatically.
+
+ Bedrock-specific fields include ``api_key``, ``base_url``, ``region_name``,
+ AWS credentials and profile fields, and read/connect timeouts. ``google:``
+ accepts ``api_key`` and ``base_url``; ``google-cloud:`` additionally
accepts
+ ``project``, ``location``, and ``service_account_info``. See the dedicated
+ :doc:`pydantic_ai_bedrock` and :doc:`pydantic_ai_vertex` connection pages
for
Review Comment:
This sends the reader to two pages that now say the opposite of what the
hook does. `pydantic_ai_bedrock.rst:26-29` and `pydantic_ai_vertex.rst:25-29`
both open with those credential shapes being "none of which fit the plain
``api_key`` + ``base_url`` shape that the generic :doc:`pydantic_ai` connection
assumes", which held while the mapper was picked by hook class. With the prefix
table, a generic `pydanticai` connection carrying `model: "bedrock:..."` plus
`region_name` and IAM keys in `extra` builds `BedrockProvider` with all of
them. What is left of the difference is that the dedicated types give those
fields their own form inputs and hide `password`/`host`, not that the generic
one cannot reach the provider. Worth a sentence on each opener, since this new
cross-reference is what will send people there.
--
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]
