kaxil commented on code in PR #72002:
URL: https://github.com/apache/airflow/pull/72002#discussion_r4005066797
##########
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:
Stepping back from individual keys, because I have been feeding you one at a
time and that is my fault. Here is the whole surface in one place so you can
settle it in a single pass.
**The fix I suggested above does not hold.** I proposed warning when a
headers mapping carries an `Authorization` key. HTTP header names are
case-insensitive, so that check is defeated by spelling. Measured at the
declared floor (`llama-index-embeddings-openai==0.6.0`), connection password
`sk-CONN`, request body and headers captured behind a mock transport:
| `embedding_kwargs` | credential | model | input | warnings |
|---|---|---|---|---|
| `{"default_headers": {"Authorization": ...}}` | replaced | ok | ok | none |
| `{"default_headers": {"authorization": ...}}` | replaced | ok | ok | none |
| `{"default_headers": {"AUTHORIZATION": ...}}` | replaced | ok | ok | none |
| `{"additional_kwargs": {"extra_headers": {"Authorization": ...}}}` |
replaced | ok | ok | none |
| `{"additional_kwargs": {"extra_body": {"model": ...}}}` | ok | replaced |
ok | none |
| `{"additional_kwargs": {"extra_body": {"input": ...}}}` | ok | ok |
replaced | none |
| `{"dimensions": 128}` (control) | ok | ok | ok | none |
Two of those rows defeat guards already in this PR. The case variants defeat
any exact-name key check. And `extra_body.input` defeats the `input`
reservation at `llamaindex.py:191`: the key is reserved at the top level of
`additional_kwargs` and passes straight through one level deeper.
The LangChain hook has the same exposure with fewer guards. At
`langchain==1.0.0`, its declared floor: `default_headers` replaces the
credential in either spelling, `model_kwargs={"extra_body": {"model": ...}}`
swaps the model past the top-level `{"model", "model_name", "provider"}` list
at `langchain.py:174`, and against a connection with no host, `{"base_url":
"http://elsewhere/v1"}` sends the connection's key to that host. That hook has
no unsupported-key warning at all, while `docs/hooks/langchain.rst` makes the
same promise as the LlamaIndex page.
The shape of the problem is that each guard is a denylist of exact names at
a fixed nesting depth, while the value is an arbitrary mapping merged over
hook-owned values at any depth. Reserving a name leaves the next spelling open,
which is what my last few comments have been doing to you. The benign uses also
sit in the same keys as the harmful ones: `{"X-Tenant": "team-a"}` through
`default_headers` is ordinary gateway routing, and `{"truncate_prompt_tokens":
512}` through `extra_body` is an ordinary self-hosted option, so no list of
names separates them.
My suggestion is to stop denylisting and say what the parameter is:
`embedding_kwargs` is forwarded to the constructor and takes effect over
anything the hook sets, credentials included. That is one paragraph and it is
accurate. The part worth changing either way is the promise in the other
direction, since three docstrings and four rst pages currently say the
connection wins, and the warning at `llamaindex.py:206` and `langchain.py:189`
only fires for `api_key`/`api_base`/`base_url`, the keys where it already is
true.
If you would rather keep a guard, the one that matches the problem is a
single case-insensitive check over any headers mapping at any depth, described
as advisory rather than as a boundary. I would not extend the reserved-name
lists further.
Sorry for arriving at this piecemeal.
--
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]
