kaxil commented on code in PR #72002:
URL: https://github.com/apache/airflow/pull/72002#discussion_r4004747745
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py:
##########
@@ -172,14 +181,42 @@ def get_embedding_model(self) -> BaseEmbedding:
except ImportError as e:
raise AirflowOptionalProviderFeatureException(e)
+ reserved_keys = sorted(self.embedding_kwargs.keys() & {"model",
"model_name"})
+ if reserved_keys:
+ raise ValueError(
+ f"embedding_kwargs must not contain reserved keys
{reserved_keys}; use embed_model instead"
+ )
+ additional_kwargs = self.embedding_kwargs.get("additional_kwargs")
+ if isinstance(additional_kwargs, dict):
+ reserved_request_keys = sorted(additional_kwargs.keys() &
{"input", "model", "model_name"})
+ if reserved_request_keys:
+ raise ValueError(
+ "embedding_kwargs['additional_kwargs'] must not contain
reserved keys "
+ f"{reserved_request_keys}; model identity and input are
managed by the hook"
+ )
+
conn = self.get_connection(self.embed_conn_id)
model_id = self._resolve_model(
conn.extra_dejson,
constructor_value=self.embed_model,
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())
Review Comment:
This warning is keyed on `self.embedding_kwargs.keys() &
connection_kwargs.keys()`, and `connection_kwargs` only ever holds
`api_key`/`api_base` -- so the one key that actually takes the connection's
credential away never intersects it. Measured at the declared
`llama-index-embeddings-openai==0.6.0` floor behind a mock transport, with the
connection password as `sk-CONNECTION`: `embedding_kwargs={"default_headers":
{"Authorization": "Bearer sk-OTHER"}}` puts `Authorization: Bearer sk-OTHER` on
the wire, and all three guards stay quiet. `default_headers` is a declared
field, so the unsupported-key warning doesn't fire; it isn't
`api_key`/`api_base`, so this warning doesn't fire; it isn't
`model`/`model_name`, so the raise doesn't fire. The control matters for the
fix: `{"X-Tenant": "team-a"}` through the same key leaves `Authorization`
untouched, which is the legitimate gateway-routing use, so the key can't simply
be reserved.
Mechanism: `_get_credential_kwargs` hands `default_headers` to the `OpenAI`
client, and the client composes its headers as `{..., **self.auth_headers,
**self._custom_headers}` -- caller headers merge last, over the bearer token.
`langchain.py:189` is the identical warning with the identical blind spot, and
that hook has no unsupported-key warning at all, so nothing catches it there
either. Both are new here: on `origin/main` only `_connection_kwargs` could
reach either constructor.
It isn't the only spelling. `embedding_kwargs={"additional_kwargs":
{"extra_body": {"model": "Qwen/Qwen3-Embedding-0.6B"}}}` succeeds, returns a
clean `list[float]`, puts the substituted model on the wire, and leaves
`embed_model` reading `text-embedding-3-small` in the rendered template fields
and the UI -- while the reserved spelling `additional_kwargs={"model": ...}`
raises loudly. `embeddings.create` accepts `extra_body` and merges it over the
named arguments (`make_request_options` maps it to `options["extra_json"]`,
then `_merge_mappings(json_data, options.extra_json)`, whose docstring reads
"In cases with duplicate keys the second mapping takes precedence"), and
`additional_kwargs={"extra_headers": {"Authorization": ...}}` does the same to
the credential. All of it behaves identically on `openai` 1.1.0, the declared
floor, and on current 3.x, so none of it is a recent SDK change.
Three channels, no shared name, and the benign uses sit in the same keys as
the harmful ones (`{"X-Tenant": ...}`, `{"truncate_prompt_tokens": 512}`, and
`extra_query`/`timeout`, which were measured additive-only), so no list of
names separates them. That's the case for validating contents instead: warning
when a headers mapping carries an `Authorization` key, naming
`self.embed_conn_id` which is already in scope on this line, would cover the
credential case. The alternative is to drop the denylist and say plainly in the
docs that `embedding_kwargs` is forwarded to the constructor and can override
anything the hook sets. Three docstrings and four rst pages currently promise
the opposite, and that's the part worth fixing either way. Two small things
riding along with whichever you pick: this warning doesn't say which connection
lost, and `self.embed_conn_id` is right there; and the two operator pages carry
the reserved-key list without the precedence rule the hook pages state, so t
hat row currently reads as if the reserved names were the whole policy.
##########
providers/common/ai/docs/operators/llamaindex_embedding.rst:
##########
@@ -86,22 +86,38 @@ Parameters
binding ``loader.output`` resolves to the native list before
execute.
* - ``embed_model``
- - String model name OR pre-built ``BaseEmbedding`` instance.
+ - String model name OR pre-built ``BaseEmbedding`` instance. Templated.
* - ``llm_conn_id``
- Airflow connection ID used when ``embed_model`` is a string. Falls
back to ``LlamaIndexHook.default_conn_name`` (``llamaindex_default``)
- when ``None``.
+ when ``None``. Templated.
* - ``embed_conn_id``
- Optional separate connection ID for the embedding provider. Falls
- back to ``llm_conn_id`` when ``None``.
+ back to ``llm_conn_id`` when ``None``. Templated.
+ * - ``embedding_kwargs``
+ - Additional keyword arguments passed to the embedding model constructor
+ when ``embed_model`` is a string or omitted, for example
+ ``{"dimensions": 128}``. Supports templating; set the Dag's
+ ``render_template_as_native_obj=True`` when templating typed values such
+ as ``dimensions`` so they remain integers instead of strings. When
Review Comment:
Since `embedding_kwargs` is a template field, this row can point at the
cheaper route the same table already documents three rows up: `documents` says
"binding ``loader.output`` resolves to the native list before execute", and
`embedding_kwargs` gets that for free, because `render_template` returns
`resolve(context)` for anything carrying a `.resolve` (`templater.py:264-265`).
An upstream `@task` returning `{"dimensions": 128}` stays a dict with no flag
set anywhere.
The Dag flag is a blunter instrument than the sentence suggests:
`render_template` sends every `str` template field through
`jinja_env.from_string` with no "contains `{{`" guard (`templater.py:254-258`),
and the environment is built once per Dag with `native=use_native`
(`dag.py:852-854`), so flipping it changes the rendered type of every string
template field of every task in that Dag, not just this one. The per-operator
`render_template_as_native_obj` would scope it, but it isn't available here: it
first appears in task-sdk 1.2.0 (Airflow 3.2.0) and this provider floors at
`apache-airflow>=3.0.0`, which is what makes the XComArg route the one that
works on every supported version. Same sentence at
`llamaindex_retrieval.rst:94`. This is me correcting advice I gave you in round
1, so treat it as a suggestion rather than a defect.
##########
providers/common/ai/tests/unit/common/ai/operators/test_llamaindex_embedding.py:
##########
@@ -132,11 +135,15 @@ def test_byo_embed_model_bypasses_hook(self, _li):
task_id="test",
documents=[{"text": "doc"}],
embed_model=byo,
+ embedding_kwargs={"dimensions": 128},
)
result = op.execute(context=MagicMock())
byo.get_text_embedding_batch.assert_called_once()
assert result["chunks"][0]["vector"] == [0.5]
+ assert (
+ "embedding_kwargs is ignored when embed_model is a pre-built
embedding model" in caplog.messages
Review Comment:
This is the only BYO construction in the file that passes
`embedding_kwargs`, and it asserts the warning is present but nothing asserts
when it should be absent -- the three other BYO constructions (`:171`, `:192`,
`:208`) say nothing about the log. I ran that mutation: delete the `if
self.embedding_kwargs:` at `llamaindex_embedding.py:211` and all 12 tests in
this file still pass, while every documented BYO task now logs a spurious
warning on every run. The sibling file already pins both directions on the
identical guard: `test_llamaindex_retrieval.py:150-157` parametrises `[(None,
False), ({"dimensions": 128}, True)]`, so this is a four-line copy.
--
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]
