kaxil commented on code in PR #71437:
URL: https://github.com/apache/airflow/pull/71437#discussion_r3979085074
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -535,3 +597,14 @@ def _get_provider_kwargs(
)
return kwargs
+
+
+_PROVIDER_KWARGS_MAPPER_BY_MODEL_PREFIX: dict[
+ str | None, Callable[[str | None, str | None, dict[str, Any]], dict[str,
Any]]
+] = {
+ "azure": PydanticAIAzureHook._get_provider_kwargs,
+ "azure-responses": PydanticAIAzureHook._get_provider_kwargs,
+ "bedrock": PydanticAIBedrockHook._get_provider_kwargs,
+ "google": PydanticAIVertexHook._get_provider_kwargs,
Review Comment:
`google` and `google-cloud` share the Vertex mapping, but
`infer_provider_class` resolves them to different classes:
`GoogleCloudProvider` accepts `project`/`location`/`credentials`, while
`GoogleProvider` accepts only
`api_key`/`base_url`/`client`/`http_client`/`retry_options`. So a
`pydanticai-vertex` connection carrying the `project` and `location` its own
conn-fields document, plus a `google:` model, raises `TypeError:
GoogleProvider.__init__() got an unexpected keyword argument 'project'` on
pydantic-ai 2.31.1 -- and with the `except TypeError` gone there is nothing
left to catch it. This is not embeddings-only: `get_conn()` with `model:
"google:gemini-2.0-flash"` returned a `GoogleModel` at d195f1e and raises the
same `TypeError` now, and `test_get_provider_kwargs_api_key_gla_mode` already
exercises the Vertex hook with a `google:` model. `google-cloud:` on the same
connection is fine, so `google` looks like it needs its own entry emitting only
`GoogleProvider`'s surface rat
her than sharing the Vertex one.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,6 +138,35 @@ 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 _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) -> dict[str, Any]:
+ provider_name, _ = parse_model_id(model_name)
+ provider_kwargs_mapper = _PROVIDER_KWARGS_MAPPER_BY_MODEL_PREFIX.get(
+ provider_name, PydanticAIHook._get_provider_kwargs
+ )
+ extra = conn.extra_dejson
+ self._warn_if_vertexai_field_ignored(extra)
Review Comment:
The mapper is chosen from the model prefix alone, with no reference to the
connection or hook it belongs to, and the Bedrock/Vertex mappers were written
for connection types that hide `password` and `host`, so they read only
`extra`. On a generic `pydanticai` connection those two fields are visible and
documented as the API key and the base URL, and they are now dropped for
`bedrock:`, `google:` and `google-cloud:` (`azure:` and `openai:` are
unaffected). Concretely: `model: "bedrock:us.anthropic.claude-opus-4-6-v1:0"`
with a bearer token in API Key and a custom endpoint in Host forwarded both to
`BedrockProvider` at d195f1e (it takes `api_key` and `base_url`) and now yields
`{}`, so the task silently runs under the worker's AWS identity against the
default endpoint, with no trace because the info log only fires for a non-empty
dict. Since the mappers hide those fields on their own conn types, a
`conn.password`/`conn.host` fallback when `extra` supplies neither can only
fire for g
eneric connections. Related: because dispatch never consults `self`, the
`_get_provider_kwargs` docstring's "subclasses override this method" no longer
describes what happens -- the three shipped subclasses are reached only because
this dict re-registers them. Worth saying explicitly that the fix there is the
wording (or declaring the prefixes on each hook class and populating the table
from them), not restoring `self`-dispatch -- that would undo the cross-provider
mapping `test_embedding_uses_model_provider_mapping_instead_of_hook_type` pins.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,6 +138,35 @@ 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 _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) -> dict[str, Any]:
+ provider_name, _ = parse_model_id(model_name)
+ provider_kwargs_mapper = _PROVIDER_KWARGS_MAPPER_BY_MODEL_PREFIX.get(
+ provider_name, PydanticAIHook._get_provider_kwargs
+ )
+ extra = conn.extra_dejson
+ self._warn_if_vertexai_field_ignored(extra)
+ return provider_kwargs_mapper(conn.password, conn.host, extra)
+
+ def _create_provider_factory(self, provider_kwargs: dict[str, Any]) ->
Callable[[str], Any]:
+ def _create_provider(provider_name: str) -> Any:
+ if provider_name.startswith("gateway/"):
Review Comment:
`gateway/` is a new user-facing model-id syntax and it is documented
nowhere: grepping the provider's docs, `provider.yaml` and `src` finds it only
at this branch and the import above. Meanwhile
`docs/self_hosted_models.rst:259-281` documents a different AI-gateway wiring
-- `host` as the gateway base URL, `password` as the bearer token,
`extra["model"]` as `openai:<route>` -- so there are now two gateway mechanisms
with incompatible `host` semantics and only the older one written down. The
branch is also unreachable from the Bedrock and Vertex connection types:
`parse_model_id` returns `"gateway/google"`, which matches no key in the mapper
table, so the base mapper runs and those types hide `password`/`host`, giving
`{}` and no `provider_factory` at all (verified: generic conn yields `{api_key,
base_url}`, vertex and bedrock conns yield `{}`). Normalising the prefix before
the table lookup would fix the second half.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -144,60 +184,85 @@ def get_conn(self) -> Model:
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
- )
+ conn = self._get_cached_connection(self.llm_conn_id)
+ extra = conn.extra_dejson
model_name: str | KnownModelName = self.model_id or extra.get("model",
"")
if not model_name:
raise ValueError(
"No model specified. Set model_id on the hook or the Model
field on the connection."
)
- api_key: str | None = conn.password or None
- base_url: str | None = conn.host or None
-
- provider_kwargs = self._get_provider_kwargs(api_key, base_url, extra)
+ provider_kwargs = self._get_provider_kwargs_for_model(conn, model_name)
if provider_kwargs:
- _kwargs = provider_kwargs # capture for closure
self.log.info(
"Using explicit credentials for provider with model '%s': %s",
model_name,
list(provider_kwargs),
)
-
- def _provider_factory(pname: str) -> Any:
- try:
- return infer_provider_class(pname)(**_kwargs)
- except TypeError:
- self.log.warning(
- "Provider '%s' rejected kwargs %s; falling back to
env-var auth",
- pname,
- list(_kwargs),
- )
- return infer_provider(pname)
-
- self._model = infer_model(model_name,
provider_factory=_provider_factory)
+ self._model = infer_model(
+ model_name,
+
provider_factory=self._create_provider_factory(provider_kwargs),
+ )
return self._model
self._model = infer_model(model_name)
return self._model
+ def get_embedder(self) -> Embedder:
+ """Return a pydantic-ai ``Embedder`` using this connection's
credentials."""
+ if self._embedder is not None:
+ return self._embedder
+
+ conn = self._get_cached_connection(self.embed_conn_id)
+ extra: dict[str, Any] = conn.extra_dejson
+
+ embed_model_name: str = self.embed_model_id or
extra.get("embed_model", "")
+ if not embed_model_name:
+ raise ValueError(
+ "No embedding model specified. Set embed_model_id on the hook
or the embed_model field "
+ "on the connection."
+ )
+
+ provider_kwargs = self._get_provider_kwargs_for_model(conn,
embed_model_name)
+ if provider_kwargs:
+ self.log.info(
+ "Using explicit credentials for provider with embedding model
'%s': %s",
+ embed_model_name,
+ list(provider_kwargs),
+ )
+ embedding_model = infer_embedding_model(
+ embed_model_name,
+
provider_factory=self._create_provider_factory(provider_kwargs),
+ )
+ else:
+ embedding_model = infer_embedding_model(embed_model_name)
+
+ self._embedder = Embedder(embedding_model,
instrument=genai_instrumentation_settings())
Review Comment:
`get_embedder()` takes no arguments, so `EmbeddingSettings` is unreachable
-- at 2.31.1 `Embedder.__init__` is `(model, *, settings, defer_model_check,
instrument)` and `EmbeddingSettings` carries
`dimensions`/`truncate`/`extra_headers`/`extra_body`. `dimensions` is the one
people reach for (text-embedding-3-small at 512 dims for a fixed-width vector
column), and setting it today means building `Embedder` by hand, which forfeits
exactly the credential resolution this PR adds. `create_agent` in this same
class already solves the shape with `**agent_kwargs` and the `_UNSET` sentinel
letting a caller-supplied `instrument` win over auto-instrumentation; `def
get_embedder(self, **embedder_kwargs)` reusing that pattern would match it,
with the cache either keyed on the kwargs or refusing them after the first call.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -127,6 +138,35 @@ 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 _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) -> dict[str, Any]:
+ provider_name, _ = parse_model_id(model_name)
+ provider_kwargs_mapper = _PROVIDER_KWARGS_MAPPER_BY_MODEL_PREFIX.get(
+ provider_name, PydanticAIHook._get_provider_kwargs
+ )
+ extra = conn.extra_dejson
+ self._warn_if_vertexai_field_ignored(extra)
+ return provider_kwargs_mapper(conn.password, conn.host, extra)
+
+ def _create_provider_factory(self, provider_kwargs: dict[str, Any]) ->
Callable[[str], Any]:
+ def _create_provider(provider_name: str) -> Any:
+ if provider_name.startswith("gateway/"):
+ return
gateway_provider(provider_name.removeprefix("gateway/"), **provider_kwargs)
+ return infer_provider_class(provider_name)(**provider_kwargs)
Review Comment:
`infer_embedding_model` calls `provider_factory(provider_name)`
unconditionally before its dispatch chain, including for the one branch that
never uses a provider: `SentenceTransformerEmbeddingModel(model_name)` takes no
`provider=`. Since `embed_conn_id` defaults to `llm_conn_id`, a hosted-chat
connection with an API key plus `embed_model:
"sentence-transformers:all-MiniLM-L6-v2"` now dies on `TypeError:
SentenceTransformersProvider() takes no arguments` -- it reached the model at
d195f1e, because `infer_provider("sentence-transformers")` succeeds and the old
swallow routed to it, so hosted chat plus local embeddings on one connection is
newly impossible. Separately, the base mapper's fixed `{api_key, base_url}`
shape is wider than several constructors (`CohereProvider(api_key,
cohere_client, http_client)` and `DeepSeekProvider` reject `base_url`,
`LiteLLMProvider` spells it `api_base`, `SnowflakeProvider` has no `api_key`),
and failing loudly there is the improvement I asked for
-- but neither `infer_model` nor `infer_embedding_model` wraps the factory,
so the user gets a bare `TypeError` from pydantic-ai internals naming no
connection field. Re-raising as something that names the conn id and the kwargs
supplied would keep the loud failure and restore the diagnostic the old warning
carried.
##########
providers/common/ai/provider.yaml:
##########
@@ -322,6 +343,13 @@ connection-types:
type:
- string
- 'null'
+ embed_model:
+ label: Embedding Model
+ description: "Google embedding model identifier (e.g.
google-cloud:text-embedding-005)"
Review Comment:
Still open from the last round: the three prefix tripwires
(`test_conn_fields_model_description_prefix_is_valid_provider` and its two
siblings) read only `conn-fields["model"]["description"]` for
`pydanticai-vertex`, so the four new `embed_model` prefixes have no drift
guard. All four resolve today -- I ran `openai:`, `azure:`, `bedrock:` and
`google-cloud:` through `infer_embedding_model` on 2.31.1 -- so this is a guard
rather than a live bug. It needs the same shape as
`_assert_prefix_is_known_provider` with a wider suppress: an unrecognised
prefix gives `ValueError: Unknown provider: <prefix>`, while a recognised one
with no key on the test box gives `ImportError` or a credential `UserError`, so
only the `ValueError` should fail the test.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -240,6 +251,279 @@ def test_get_conn_caches_model(self, mock_infer_model):
mock_infer_model.assert_called_once()
+class TestPydanticAIHookGetEmbedder:
+
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_embedding_model",
autospec=True)
+
@patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class",
autospec=True)
+ def test_embedding_uses_model_provider_mapping_instead_of_hook_type(
+ self, mock_infer_provider_class, mock_infer_embedding_model
+ ):
+ mock_embedding_model = MagicMock(spec=EmbeddingModel)
+ mock_infer_embedding_model.return_value = mock_embedding_model
+ mock_provider = mock_infer_provider_class.return_value.return_value
+ hook = PydanticAIAzureHook(embed_conn_id="test_conn",
embed_model_id="openai:text-embedding-3-small")
+ conn = Connection(
+ conn_id="test_conn",
+ conn_type="pydanticai",
+ password="sk-test-key",
+ host="https://api.openai.com/v1";,
+ )
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ result = hook.get_embedder()
+
+ assert isinstance(result, Embedder)
+ assert result.model is mock_embedding_model
+ call_args = mock_infer_embedding_model.call_args
+ assert call_args.args == ("openai:text-embedding-3-small",)
+ provider_factory = call_args.kwargs["provider_factory"]
+ assert provider_factory("openai") is mock_provider
+ mock_infer_provider_class.return_value.assert_called_once_with(
+ api_key="sk-test-key", base_url="https://api.openai.com/v1";
+ )
+
+ @pytest.mark.parametrize(
+ ("provider_name", "extra", "expected_provider_kwargs"),
+ [
+ ("openai", {}, {"api_key": "connection-key", "base_url":
"https://example.com"}),
+ (
+ "azure",
+ {"api_version": "2024-07-01-preview"},
+ {
+ "api_key": "connection-key",
+ "azure_endpoint": "https://example.com";,
+ "api_version": "2024-07-01-preview",
+ },
+ ),
+ (
+ "azure-responses",
+ {"api_version": "2024-07-01-preview"},
+ {
+ "api_key": "connection-key",
+ "azure_endpoint": "https://example.com";,
+ "api_version": "2024-07-01-preview",
+ },
+ ),
+ (
+ "bedrock",
+ {"region_name": "us-east-1"},
+ {"region_name": "us-east-1"},
+ ),
+ (
+ "google",
+ {
+ "api_key": "extra-key",
+ "base_url": "https://extra.example.com";,
+ "project": "project",
+ },
+ {
+ "api_key": "extra-key",
+ "base_url": "https://extra.example.com";,
+ "project": "project",
Review Comment:
Two rows here pin the shapes the code gets wrong. The `google` row asserts
`GoogleProvider(api_key=..., base_url=..., project="project")`, which the real
class rejects, and the `bedrock` row asserts only `{"region_name": ...}` while
the connection has `password` and `host` set -- so it encodes the discarded
credentials as intended behaviour. Both pass only because
`infer_provider_class` is patched, so `mock_infer_provider_class.return_value`
swallows any kwargs; `Non-DB-prov::3.10:common.ai` is green at HEAD.
`test_get_conn_vertexai_flag_is_not_forwarded` already shows the pattern that
would have caught it -- a `FakeGoogleCloudProvider` carrying the real
keyword-only signature -- and applying that to `GoogleProvider` for the
`google` row would close it. Worth noting the `google` and `google-cloud` rows
carry byte-identical expectations, so nothing here tests the one place those
two classes actually diverge.
--
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]
