kaxil commented on code in PR #71437:
URL: https://github.com/apache/airflow/pull/71437#discussion_r4057057705
##########
providers/common/ai/docs/connections/pydantic_ai.rst:
##########
@@ -72,21 +97,31 @@ 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.
Fallback Connections
Other connection IDs to fail over to, in order, while this provider is
unavailable. Stored in ``extra["fallback_conn_ids"]``. Entries may name any
``pydanticai`` connection type, so one chain can span vendors. See
:doc:`/provider_fallback`.
+ Bedrock-specific fields include ``api_key``, ``base_url``, ``region_name``,
Review Comment:
This paragraph describes `extra` fields but is indented under the Fallback
Connections item that starts at :112, so it renders as part of the fallback
description. The Extra (JSON, optional) item it belongs under ends at :110.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -220,8 +257,92 @@ def _seed_connection(self, conn: Connection) -> None:
would fetch that same connection a second time the first time it runs,
doubling the
Execution API round trips a fallback chain costs.
"""
- self._conn = conn
- self._conn_extra_dejson = conn.extra_dejson
+ self._connections[conn.conn_id] = conn
+ self._connection_extra_dejson[conn.conn_id] = conn.extra_dejson
+
+ 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)
+ 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:')."
+ )
+ 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),
+ )
+ ignored_extra_fields = [field for field in
provider_config.ignored_extra_fields if extra.get(field)]
+ if ignored_extra_fields:
+ self.log.warning(
+ "Connection extra fields are ignored for provider %r on
connection %r: %s",
+ provider_name,
+ conn.conn_id,
+ ignored_extra_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)}: {e}"
+ ) 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)
Review Comment:
This one came out of my own suggestion on
https://github.com/apache/airflow/pull/71437#discussion_r4032377807, so the
regression is mine as much as yours. The prefixless `embed_model` half is
right; the `llm_provider is None` arm that came with it skips the credential
guard altogether.
Whether the guard runs now depends on whether the user typed a prefix that
the connection pages call optional. Measured against HEAD on pydantic-ai 2.31.1
with the real provider classes, on a generic `pydanticai` connection holding an
API key, `model: "gpt-4o"`, `embed_model: "cohere:embed-v4.0"`:
```
guard -> passes (llm_provider is None, early return)
kwargs -> {'api_key': 'sk-OPENAI-KEY'}
embedder-> CohereEmbeddingModel via CohereProvider,
base_url=https://api.cohere.com
```
The OpenAI key is sent to Cohere. Writing `openai:gpt-4o` on that same
connection raises correctly, and that is the case the guard exists for.
It rejects in the other direction too. A `pydanticai_bedrock` connection
with Model `us.anthropic.claude-opus-4-6-v1:0`, the bare native id
`pydantic_ai_bedrock.rst:65-74` documents, colon and all, plus Embedding Model
`bedrock:amazon.titan-embed-text-v2:0`, is refused with `configures different
LLM and embedding providers ('us.anthropic.claude-opus-4-6-v1' and 'bedrock')`,
because `parse_model_id` splits on the id's own colon. Both are Bedrock,
`get_conn()` resolves that same string to `BedrockConverseModel` behind
`BedrockProvider`, and the separate `embed_conn_id` the message prescribes is
not the fix.
Both directions have the same root: this line derives the LLM provider with
a bare `parse_model_id`, while every other path goes through
`_qualify_model_name` / `_has_recognized_provider_prefix`, which ignores a `:`
that is not a real provider name and falls back to `model_provider`. Reusing
that here, roughly `llm_provider = parse_model_id(llm_model_name)[0] if
_has_recognized_provider_prefix(llm_model_name) else self.model_provider`, and
returning early only when it is still `None`, covers both.
`test_prefixless_llm_model_skips_embedding_provider_validation` pins the
current arm, and no test sets a bare vendor model name together with an
`embed_model` on one connection.
While you are in here, the `embed_provider == "sentence-transformers"` arm
on the next line is unreachable: `_get_provider_factory_for_model` already
returns `None` for that prefix, and this validator only runs where the factory
is not `None`.
--
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]
