kaxil commented on code in PR #71437:
URL: https://github.com/apache/airflow/pull/71437#discussion_r3980258809
##########
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:
You're right, drop the fallback as I put it. On that connection `password`
and `host` are the OpenAI ones, and handing them to `GoogleProvider` is the
same wrong-identity bug I flagged last round, just moved.
The narrower version I had in mind doesn't work either. Fall back only when
the connection names a single provider, and
`PydanticAIHook(model_id="bedrock:...")` against a connection storing `model:
openai:...` still passes: one prefix in play, so the OpenAI key goes to
`BedrockProvider` as a bearer token. `model_id` is a documented override, so
that isn't a corner case.
Your change did close a real leak, for the five mapped prefixes. Before it
`get_embedder()` used the hook-bound mapper, so a generic connection handed
`{api_key, base_url}` straight to `GoogleProvider`, which accepts both, and the
OpenAI key went out as an `x-goog-api-key` header against `api.openai.com`. No
`TypeError`, so the old swallow never saw it. That only ever existed inside
this PR's own history though: `main` has no `get_embedder`.
What still blocks is a regression against `main`, and it's in `get_conn()`
rather than embeddings. Generic `pydanticai` connection, `model:
bedrock:us.anthropic.claude-opus-4-6-v1:0`, bearer token in API Key, custom
runtime endpoint in Host:
```
main: {"api_key": "bearer-tok", "base_url": "https://custom-bedrock"}
HEAD: {}
```
`BedrockProvider` takes both of those. At HEAD the prefix table picks the
extra-only mapper, so that same connection now runs under the worker's AWS
identity against the default endpoint. There's only one provider in play, so
nothing is ambiguous, and nothing errors.
So: keep extra-only precedence, but say something when a mapper throws away
a non-empty `conn.password` or `conn.host`, and name the `extra` keys that
replace them. Those keys (`region_name`, `api_key`, `base_url`, `project`,
`location`) do work on a generic connection. They're just documented only on
the `pydanticai-bedrock` and `pydanticai-vertex` pages, which both open by
saying those credential shapes don't fit this connection type.
`pydantic_ai.rst` tells you to leave API Key empty for Bedrock and Vertex,
which reads as "not needed" rather than "ignored", and says nothing at all
about Host.
One more that sits outside the table: `cohere:` is one of the embedding
providers pydantic-ai ships, and on a connection with API Key set and Host
empty, `model: openai:...` plus `embed_model: cohere:embed-english-v3.0` hands
the OpenAI key to `CohereProvider`. Set Host as well and it's an uncaught
`TypeError`, since `CohereProvider` has no `base_url`. And yes, `embed_conn_id`
is the answer for genuinely mixed providers.
--
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]
