Skip to content

FIX: Entra ID auth fallback for Azure endpoints in OpenAITextEmbedding#2235

Merged
romanlutz merged 3 commits into
microsoft:mainfrom
romanlutz:romanlutz-fix-embedding-entra-fallback
Jul 20, 2026
Merged

FIX: Entra ID auth fallback for Azure endpoints in OpenAITextEmbedding#2235
romanlutz merged 3 commits into
microsoft:mainfrom
romanlutz:romanlutz-fix-embedding-entra-fallback

Conversation

@romanlutz

Copy link
Copy Markdown
Contributor

Description

OpenAITextEmbedding was 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__ called default_values.get_required_value(...) for the key, which raises ValueError on an empty/absent value, and it never minted an Entra token. So OpenAITextEmbedding() against an Azure endpoint with no key just errored, even though every other target (via OpenAITarget) falls back to Entra automatically.

This change updates only the api_key resolution branch in OpenAITextEmbedding.__init__ to mirror OpenAITarget exactly:

  • A callable token provider is used directly.
  • Otherwise the key is resolved via get_non_required_value (param then env var).
  • A non-empty value is used as the key string.
  • Else, for a recognized Azure OpenAI / AI Foundry endpoint (is_azure_openai_endpoint), an Entra token provider is minted via get_azure_openai_auth(endpoint).
  • Else a ValueError is raised, with a message consistent with OpenAITarget (key required for non-Azure endpoints; Entra used automatically for recognized Azure endpoints).

Endpoint/model resolution via get_required_value is unchanged, and the existing ensure_async_token_provider wrapping plus AsyncOpenAI(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 annotated resolved_api_key, because ensure_async_token_provider returns ... | None and reusing the variable trips a real ty invalid-assignment error. This matches how OpenAITarget assigns the wrapped value to a distinct attribute.

Tests and Documentation

Extended tests/unit/embedding/test_azure_text_embedding.py, mirroring the auth-test style in tests/unit/prompt_target/target/test_openai_target_auth.py (module helper clears the environment with patch.dict(os.environ, {}, clear=True) so OPENAI_EMBEDDING_KEY cannot leak in):

  • Explicit string api_key is passed to the client as-is.
  • A callable token provider is used as-is (not overwritten by env resolution).
  • No key + recognized Azure endpoint invokes get_azure_openai_auth with the endpoint and wires the returned provider into AsyncOpenAI.
  • No key + non-Azure endpoint raises ValueError.

The pre-existing test_invalid_key_raises was updated: its old endpoint (mock.azure.com) is not a recognized Azure host, so an empty key now correctly raises our ValueError; 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, and uv run pre-commit run --files <changed files> all pass.

JupyText: doc/code/memory/embeddings.py was executed via jupytext --to ipynb --execute against a live Azure endpoint with no OPENAI_EMBEDDING_KEY set (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.

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
Comment thread pyrit/embedding/openai_text_embedding.py Outdated
Copilot AI added 2 commits July 20, 2026 09:50
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
@romanlutz
romanlutz added this pull request to the merge queue Jul 20, 2026
Merged via the queue into microsoft:main with commit be4c90e Jul 20, 2026
53 checks passed
@romanlutz
romanlutz deleted the romanlutz-fix-embedding-entra-fallback branch July 20, 2026 18:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants