kaxil commented on code in PR #72002:
URL: https://github.com/apache/airflow/pull/72002#discussion_r3973850415
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/llamaindex.py:
##########
@@ -179,7 +185,8 @@ def get_embedding_model(self) -> BaseEmbedding:
extra_key="embed_model",
kind="embedding",
)
- return OpenAIEmbedding(model=model_id, **self._connection_kwargs(conn))
+ kwargs = {**self.embedding_kwargs, **self._connection_kwargs(conn)}
Review Comment:
Keeping the connection ahead of `embedding_kwargs` looks like the right
call, since the only keys that can collide are credentials and `embed_conn_id`
is already the escape hatch. What's missing is any signal: a user who passes
`api_base` here against a connection that has a host gets it dropped silently.
`ClickHouseHook` handles the same shape by naming the managed keys and logging
the intersection
(`providers/clickhousedb/src/airflow/providers/clickhousedb/hooks/clickhouse.py:42`
and 228-237). One warning on the overlap here and at `hooks/langchain.py:179`
would cover it.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py:
##########
@@ -128,6 +129,24 @@ 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):
Review Comment:
Both new tests use a connection whose only populated field is `password`, so
no key in `embedding_kwargs` ever collides with `_connection_kwargs`. The merge
order is the one behaviour this PR actually decides, and all four rst files
document it, so a case passing `embedding_kwargs={"api_key": "from-kwargs"}`
against a connection with a password would pin it. Same gap in
`test_langchain.py`.
##########
providers/common/ai/docs/operators/llamaindex_embedding.rst:
##########
@@ -94,6 +94,10 @@ Parameters
* - ``embed_conn_id``
- Optional separate connection ID for the embedding provider. Falls
back to ``llm_conn_id`` when ``None``.
+ * - ``embedding_kwargs``
+ - Additional keyword arguments passed to the embedding model constructor
+ when ``embed_model`` is a string or omitted, for example
+ ``{"dimensions": 128}``. Templated.
Review Comment:
The "Templated." promise doesn't hold for `dimensions`, which is the option
this row uses as its own example. A literal `{"dimensions": 128}` works,
because `render_template` returns non-str values unchanged. A Jinja value
renders to a string, and LlamaIndex's `OpenAIEmbedding` copies `dimensions`
into `additional_kwargs` before pydantic validation
(`llama_index/embeddings/openai/base.py:297-299` at the declared 0.6.0 floor),
then spreads `**self.additional_kwargs` into the request (`:399` plus five
sibling call sites) rather than reading the validated field. Capturing the
request body at that floor behind a mock transport, `{"dimensions": "{{
params.dim }}"}` sends `"dimensions": "128"` while `self.dimensions` holds a
correctly coerced `128` that nothing reads, and the embeddings API requires an
integer there. Either coerce the known int keys in the hook, or drop
"Templated." here and at `llamaindex_retrieval.rst:93` and note that a
templated `dimensions` needs `render_template_a
s_native_obj=True`. LangChain is unaffected, since `langchain_openai` reads
its validated field for request params. Separately, this table now marks only
`documents` and `embedding_kwargs` as templated while `embed_model`,
`llm_conn_id`, `embed_conn_id`, `persist_dir` and `persist_conn_id` are equally
templated and unmarked, which reads as if those aren't.
--
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]
