kaxil commented on code in PR #72002:
URL: https://github.com/apache/airflow/pull/72002#discussion_r3997974027
##########
providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py:
##########
@@ -128,6 +129,43 @@ def test_dispatches_with_api_base(self, mock_get_conn,
mock_cls):
api_base="http://localhost:11434/v1";,
)
+ @patch("llama_index.embeddings.openai.OpenAIEmbedding")
+ @patch.object(LlamaIndexHook, "get_connection")
+ def test_dispatches_with_embedding_kwargs(self, mock_get_conn, mock_cls,
caplog):
+ mock_get_conn.return_value = _conn(password="sk-test")
+ mock_cls.model_fields = {"api_key": None, "dimensions": None,
"timeout": None}
+ hook = LlamaIndexHook(
+ embed_model="text-embedding-3-small",
+ embedding_kwargs={"api_key": "from-kwargs", "dimensions": 128,
"timeout": 30},
+ )
+
+ hook.get_embedding_model()
+
+ mock_cls.assert_called_once_with(
+ model="text-embedding-3-small",
+ api_key="sk-test",
+ dimensions=128,
+ timeout=30,
+ )
+ assert "Connection parameters override embedding_kwargs values:
['api_key']" in caplog.messages
+
+ @patch("llama_index.embeddings.openai.OpenAIEmbedding")
+ @patch.object(LlamaIndexHook, "get_connection")
+ def test_warns_about_unsupported_embedding_kwargs(self, mock_get_conn,
mock_cls, caplog):
Review Comment:
Both of these patch `OpenAIEmbedding` itself, so `supported_kwargs` is
computed from the mock rather than the class:
`inspect.signature(MagicMock().__init__).parameters` is `{'args', 'kw'}` and
`model_fields` is whatever the test assigned two lines up, so the assertion
checks the warning against its own stub. Deleting the `inspect.signature(...)`
half of the union leaves all six tests in this file green, and so does
replacing the whole union with `set()`, which would make the hook warn on every
correct call.
Keep the signature half, though: `mode`, `http_client` and
`async_http_client` are `__init__`-only and absent from `model_fields`, so
without it a user passing `http_client=` gets warned about a key that works.
One test against the real class, parametrized over a signature-only key, a
field-only key and a genuinely unsupported one, would pin it, and adding
`assert "ignores unsupported" not in caplog.text` to
`test_dispatches_with_embedding_kwargs` would catch the false-positive
direction.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py:
##########
@@ -179,7 +186,18 @@ 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}
+ supported_kwargs =
set(inspect.signature(OpenAIEmbedding.__init__).parameters) | set(
+ OpenAIEmbedding.model_fields
+ )
+ unsupported_keys = sorted(self.embedding_kwargs.keys() -
supported_kwargs)
Review Comment:
`model_name` passes this guard (it is in `model_fields`) and is not a
connection key, so neither new warning fires, but `OpenAIEmbedding.__init__`
pops it and uses it to set both `_query_engine` and `_text_engine`
(`llama_index/embeddings/openai/base.py:310-314` and `:332-333`), and those are
what go on the wire as `engine=` (`:398` plus five sibling call sites).
Captured the request at the declared 0.6.0 floor behind a mock transport:
`embed_model="text-embedding-3-small"` with `embedding_kwargs={"model_name":
"Qwen/Qwen3-Embedding-0.6B"}` sends `{"model": "Qwen/Qwen3-Embedding-0.6B",
"dimensions": 128, ...}` with an empty warning list, while `embed_model` still
reads `text-embedding-3-small` in the rendered template fields and in the UI.
That matters because `llamaindex_retrieval.rst` tells users `embed_model` is
the value that has to match between the embedding and retrieval tasks, and it
would now be naming a model that did not produce the vectors. It is also not
pathological input: your own test notes say to serve the vLLM model as
`text-embedding-3-small` "because `LlamaIndexHook` validates the model name",
so `model_name` is the next key a self-hosting user reaches for. Reserving
`{"model", "model_name"}` and raising would close it, or make it the documented
escape hatch.
Same expression, smaller point: the supported set also picks up `self` and
`kwargs` from the signature, so `embedding_kwargs={"kwargs": {...}}` (the
mis-nesting a parameter named `embedding_kwargs` invites) passes the guard
silently and is then dropped by pydantic's `extra="ignore"`, which is the exact
case the guard exists to announce. Filtering `VAR_KEYWORD`/`VAR_POSITIONAL` and
`self` out covers it.
##########
providers/common/ai/docs/operators/llamaindex_embedding.rst:
##########
@@ -86,22 +86,28 @@ 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.
Review Comment:
`{"dimensions": 128}` here creates an obligation that only the retrieval
page states ("must match those used to build the index"), which is the one page
where the reader can no longer act on it. Nothing records the value at persist
time, so a mismatch surfaces later as a shape error naming neither `dimensions`
nor `embedding_kwargs`. One clause on this side would carry it.
While you are here: none of the four new rows says where the valid keys come
from. `providers/openai` links the API reference for its own `embedding_kwargs`
(`operators/openai.py:44-47`), and a `.. seealso::` pointing at
`llama_index.embeddings.openai.OpenAIEmbedding` and
`langchain.embeddings.init_embeddings` would do the same job here. Worth the
line because the two parameters share a name but not a value space: theirs is a
per-request bag forwarded to `embeddings.create`, this one is a constructor bag.
--
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]
