kaxil commented on code in PR #72002:
URL: https://github.com/apache/airflow/pull/72002#discussion_r3979564655
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py:
##########
@@ -179,7 +185,12 @@ def get_embedding_model(self) -> BaseEmbedding:
extra_key="embed_model",
kind="embedding",
)
- return OpenAIEmbedding(model=model_id, **self._connection_kwargs(conn))
+ connection_kwargs = self._connection_kwargs(conn)
+ overridden_keys = sorted(self.embedding_kwargs.keys() &
connection_kwargs.keys())
+ if overridden_keys:
+ self.log.warning("Connection parameters override embedding_kwargs
values: %s", overridden_keys)
+ kwargs = {**self.embedding_kwargs, **connection_kwargs}
+ return OpenAIEmbedding(model=model_id, **kwargs)
Review Comment:
Checked this at the declared floor (`llama-index-embeddings-openai==0.6.0`):
`OpenAIEmbedding.model_config` sets no `extra`, so pydantic's default
`extra="ignore"` applies and anything the class doesn't declare is dropped
without a word. Constructing it with `base_url="http://vllm.internal/v1"`
leaves `api_base` at `https://api.openai.com/v1`, and `dimension=128`
(singular) leaves `dimensions=None`, both with an empty warning list. Until
this PR only `api_key` and `api_base` could reach that constructor and both
were hook-controlled, so this is the change that makes arbitrary keys reachable.
The spelling is what makes it reachable rather than theoretical: this
provider's own LangChain hook emits `base_url` (`hooks/langchain.py:139`) while
this one emits `api_base` (`:162`), and both hook doc rows say only "the base
URL" without naming either. A user carrying `embedding_kwargs={"base_url":
...}` across from the LangChain hook, against a connection with no host, gets
it dropped while `api_key` still applies, so the request goes to api.openai.com
and the task succeeds. LangChain is loud about the same mistake
(`init_embeddings` emits `UserWarning: WARNING! api_base is not default
parameter` and moves it into `model_kwargs`), so this hook is the one that
needs the signal.
A second warning next to the collision one you just added would cover it,
keyed on `set(kwargs) -
(set(inspect.signature(OpenAIEmbedding.__init__).parameters) |
set(OpenAIEmbedding.model_fields))`; naming `api_base` in the two hook doc rows
would be the cheap minimum. Worth noting no hook test can catch this today,
since they all patch `OpenAIEmbedding` itself and a `MagicMock` accepts every
key.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llamaindex_embedding.py:
##########
@@ -190,6 +195,7 @@ def _resolve_embed_model(self) -> BaseEmbedding:
llm_conn_id=self.llm_conn_id,
embed_conn_id=self.embed_conn_id,
embed_model=self.embed_model,
+ embedding_kwargs=self.embedding_kwargs,
Review Comment:
`embedding_kwargs` reaches the hook only on this branch. When `embed_model`
is a pre-built `BaseEmbedding`, the duck-typed branch below returns it
untouched and `self.embedding_kwargs` is never read again, with no log line,
and because `embed_model` is a template field precisely so it can be bound to
an XComArg the author can't always tell at authoring time which branch runs.
The docstring does scope it to "when `embed_model` is a string or omitted",
so this is a diagnostics gap rather than a doc error, and it's the same
silent-drop shape you just added a warning for in the hook. One
`self.log.warning` before the duck-typed return would close it, same spot in
`llamaindex_retrieval.py`, and setting `embedding_kwargs` in the existing
`test_byo_embed_model_bypasses_hook` would pin it.
--
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]