kaxil commented on code in PR #71437:
URL: https://github.com/apache/airflow/pull/71437#discussion_r3761967501
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -69,6 +73,7 @@ def __init__(
self,
llm_conn_id: str | None = None,
model_id: str | None = None,
+ embed_model_id: str | None = None,
Review Comment:
`LangChainHook` and `LlamaIndexHook` in this same provider both take an
`embed_conn_id` that falls back to `llm_conn_id`, because the embedding
provider often isn't the chat provider. Here `get_embedder()` always builds the
provider from `llm_conn_id`, so a connection with `model:
"anthropic:claude-opus-4-6"` plus `embed_model:
"openai:text-embedding-3-small"` hands the Anthropic key to `OpenAIProvider`,
and a `host` set for a self-hosted chat endpoint gets used as the OpenAI base
URL. The `TypeError` guard in `_create_provider_factory` won't catch it,
`OpenAIProvider(api_key=..., base_url=...)` constructs fine and just fails
later as a 401 or silently talks to the wrong endpoint. Worth adding
`embed_conn_id` here to match the other two hooks?
##########
providers/common/ai/docs/connections/pydantic_ai.rst:
##########
@@ -136,3 +141,8 @@ The hook reads the model from these sources in priority
order:
1. ``model_id`` parameter on the hook/operator
2. ``model`` in the connection's extra JSON (set by the "Model" conn-field in
the UI)
+
+The embedding model is resolved separately:
+
+1. ``embed_model_id`` parameter on the hook
+2. ``embed_model`` in the connection's extra JSON
Review Comment:
The `langchain` and `llamaindex` connection types already declare an
`embed_model` conn-field in `provider.yaml` (labelled "Embedding Model"), but
`pydanticai` only declares `model`. So on a new pydantic-ai connection the UI
gives you a dedicated Model input and nothing equivalent for the embedding
model, you have to type it into the Extra JSON by hand. Adding the conn-field
to the four `pydanticai*` connection types would line this up.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -160,30 +181,57 @@ 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.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
+ )
+
+ 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)
Review Comment:
`create_agent` applies `genai_instrumentation_settings()` so agent runs emit
GenAI spans, but the `Embedder` is built without it, so with `[common.ai]
otel_export_enabled` on you get spans for agents and nothing for embedding
calls. `Embedder.__init__` takes `instrument: InstrumentationSettings | bool |
None`, so `Embedder(embedding_model,
instrument=genai_instrumentation_settings())` covers it, and passing `None`
when tracing is off keeps the current zero-overhead behaviour.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -160,30 +181,57 @@ 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:
Review Comment:
`test_connection()` still only calls `get_conn()`, which raises `No model
specified` when the connection sets `embed_model` but no `model`. Now that an
embedding-only connection is a valid configuration, hitting Test in the UI on
one reports it as broken. Should `test_connection` fall back to resolving the
embedder when no LLM model is configured?
--
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]
