kaxil commented on code in PR #71437:
URL: https://github.com/apache/airflow/pull/71437#discussion_r3968163728
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -198,6 +248,17 @@ def _get_conn_if_model_configured(self) -> Model | None:
return None
+ def _get_embedder_if_model_configured(self) -> Embedder | None:
+ """Return the embedder only when the hook or connection explicitly
configures one."""
+ if self.embed_model_id:
+ return self.get_embedder()
+
+ conn = self.get_connection(self.embed_conn_id)
Review Comment:
`_get_conn_if_model_configured` caches its fetch into `self._conn` so
`get_conn()` does not repeat it, but this one throws the connection away and
`get_embedder()` then fetches the same `embed_conn_id` again, so
`test_connection()` on a connection carrying both keys makes three
`get_connection()` calls where it used to make one. `secrets.use_cache` is off
by default, so from a task each of those is a real supervisor round trip plus a
secrets-backend read. Keying the cache by conn id and routing both paths
through it would remove the duplicate and stop a later caller picking up the
LLM connection for the embedder by accident, since `self._conn` currently
carries no conn id.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -160,30 +185,55 @@ def get_conn(self) -> Model:
provider_kwargs = self._get_provider_kwargs(api_key, base_url, extra)
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_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."
+ )
+
+ 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)
+ 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),
Review Comment:
`_get_provider_kwargs` is keyed to the LLM hook, so the embedding connection
gets the wrong mapping whenever the embed model's provider differs from the
chat one, which is the case `embed_conn_id` was added for:
`PydanticAIAzureHook(embed_conn_id=<a pydanticai conn>)` with `embed_model:
openai:text-embedding-3-small` yields `{api_key, azure_endpoint}`,
`OpenAIProvider` rejects `azure_endpoint`, the `except TypeError` above drops
both, and with `OPENAI_API_KEY` on the worker the embedder lands on
`https://api.openai.com/v1/` under the env key instead of the connection's key
and host (Bedrock and Vertex get there with no log line at all, since their
mapping reads only `extra` and returns `{}`). The base hook has a quieter
version that succeeds instead of erroring: one plain `pydanticai` connection
with a `host` plus `embed_model: cohere:embed-v4.0` builds a working
`CohereEmbeddingModel` on `api.cohere.com` under `CO_API_KEY`, so the task
embeds against the wrong account with a sing
le warning as the only trace. Could the mapping be selected from the embedding
model's provider, with the factory re-raising rather than retrying with the
credentials discarded, and the comment in
`PydanticAIVertexHook._get_provider_kwargs` (which already names this hazard)
updated to say the swallow now serves `get_embedder()` too?
##########
providers/common/ai/provider.yaml:
##########
@@ -322,6 +343,13 @@ connection-types:
type:
- string
- 'null'
+ embed_model:
Review Comment:
`test_conn_fields_model_description_prefix_is_valid_provider` reads only
`conn-fields["model"]["description"]`, so the four new `embed_model` prefixes
sit outside the tripwires that guard the same strings for `model`. Reusing
`_assert_prefix_is_known_provider` would not be enough on its own either, since
the provider registry is wider than what `infer_embedding_model` dispatches on:
`infer_provider_class("anthropic")` resolves while
`infer_embedding_model("anthropic:x")` raises `UserError: Unknown embeddings
model`. A parametrized test over each connection type's `embed_model`
description, validating through `infer_embedding_model`, would cover all four.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -160,30 +185,55 @@ def get_conn(self) -> Model:
provider_kwargs = self._get_provider_kwargs(api_key, base_url, extra)
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_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."
+ )
+
+ 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)
+ 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:
Nothing asserts this yet: the new tests only check `isinstance(result,
Embedder)`, which still passes with `instrument=` deleted, so 057ee9b is
uncovered. `TestPydanticAIHookCreateAgentInstrumentation` has three tests for
the same wiring on the agent side, and the embedder equivalent is a one-liner
patching `genai_instrumentation_settings` to a sentinel and asserting
`hook.get_embedder().instrument is sentinel`.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -89,8 +100,8 @@ def get_ui_field_behaviour() -> dict[str, Any]:
"hidden_fields": ["schema", "port", "login"],
"relabeling": {"password": "API Key"},
"placeholders": {
- "host": "https://api.openai.com/v1 (optional, for custom
endpoints / Ollama)",
- "extra": '{"model": "openai:gpt-5.6-sol"}',
+ "host": "https://api.openai.com/v1 (optional, for custom
endpoints / Ollama)",
Review Comment:
This picked up a second space before `(optional`, so it no longer matches
`provider.yaml:172` or its `get_provider_info.py:143` twin, both of which this
PR edits one line lower for the `extra` placeholder. #72087 lined those three
copies up on purpose, so this looks like a stray keystroke worth reverting.
--
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]
