FIX: Entra ID auth fallback for Azure endpoints in OpenAITextEmbedding#2235
Merged
romanlutz merged 3 commits intoJul 20, 2026
Merged
Conversation
OpenAITextEmbedding did not fall back to Entra ID authentication for Azure OpenAI endpoints when no API key was set, unlike every other OpenAI-family target: get_required_value raised a ValueError instead of minting a token. Mirror OpenAITarget's api_key resolution so a callable token provider is used directly, an env/param key is used when present, recognized Azure endpoints fall back to get_azure_openai_auth (Entra), and non-Azure endpoints without a key raise a clear ValueError. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c1fe01a4-bf26-4e1f-8a26-596f312001a8
hannahwestra25
approved these changes
Jul 20, 2026
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c1fe01a4-bf26-4e1f-8a26-596f312001a8
Update remaining tests to patch the shared OpenAI auth helper after auth resolution moved out of openai_target. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c1fe01a4-bf26-4e1f-8a26-596f312001a8
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
OpenAITextEmbeddingwas the only OpenAI-family class that could not authenticate to an Azure OpenAI endpoint via Entra ID. When no API key was set, its__init__calleddefault_values.get_required_value(...)for the key, which raisesValueErroron an empty/absent value, and it never minted an Entra token. SoOpenAITextEmbedding()against an Azure endpoint with no key just errored, even though every other target (viaOpenAITarget) falls back to Entra automatically.This change updates only the
api_keyresolution branch inOpenAITextEmbedding.__init__to mirrorOpenAITargetexactly:get_non_required_value(param then env var).is_azure_openai_endpoint), an Entra token provider is minted viaget_azure_openai_auth(endpoint).ValueErroris raised, with a message consistent withOpenAITarget(key required for non-Azure endpoints; Entra used automatically for recognized Azure endpoints).Endpoint/model resolution via
get_required_valueis unchanged, and the existingensure_async_token_providerwrapping plusAsyncOpenAI(api_key=..., base_url=endpoint)wiring are preserved. One small note for reviewers: the wrapped result is assigned to a separate local (async_api_key) rather than reusing the annotatedresolved_api_key, becauseensure_async_token_providerreturns... | Noneand reusing the variable trips a realtyinvalid-assignment error. This matches howOpenAITargetassigns the wrapped value to a distinct attribute.Tests and Documentation
Extended
tests/unit/embedding/test_azure_text_embedding.py, mirroring the auth-test style intests/unit/prompt_target/target/test_openai_target_auth.py(module helper clears the environment withpatch.dict(os.environ, {}, clear=True)soOPENAI_EMBEDDING_KEYcannot leak in):api_keyis passed to the client as-is.get_azure_openai_authwith the endpoint and wires the returned provider intoAsyncOpenAI.ValueError.The pre-existing
test_invalid_key_raiseswas updated: its old endpoint (mock.azure.com) is not a recognized Azure host, so an empty key now correctly raises ourValueError; it was pointed at a non-Azure endpoint to keep asserting the no-key error path.Validation (uv only):
uv run pytest tests/unit/embedding/test_azure_text_embedding.py -v(11 passed),uv run ruff check,uv run ruff format --check,uv run ty check pyrit/embedding/openai_text_embedding.py, anduv run pre-commit run --files <changed files>all pass.JupyText:
doc/code/memory/embeddings.pywas executed viajupytext --to ipynb --executeagainst a live Azure endpoint with noOPENAI_EMBEDDING_KEYset (exactly the Entra fallback path). It ran end-to-end and returned a real embedding, confirming the fix. The notebook source is unchanged, so no doc files are modified in this PR.