Skip to content

Adding embedding policy changes for EGS#46870

Merged
aayush3011 merged 9 commits into
Azure:mainfrom
aayush3011:users/akataria/EGSPolicy
May 20, 2026
Merged

Adding embedding policy changes for EGS#46870
aayush3011 merged 9 commits into
Azure:mainfrom
aayush3011:users/akataria/EGSPolicy

Conversation

@aayush3011
Copy link
Copy Markdown
Member

@aayush3011 aayush3011 commented May 13, 2026

Description

Adds documentation and test coverage for the new optional embeddingSource object on vectorEmbeddingPolicy, which allows the service to compute embeddings on the server side from one or more source paths.

New embeddingSource fields:

  • sourcePaths – list of document paths to embed
  • deploymentName – embedding model deployment name
  • modelName – embedding model name
  • endpoint – embedding service endpoint
  • authType – ApiKey or Entra

No SDK code changes were required: vector_embedding_policy is passed through as a dict, so the new fields flow to the service without modification.

All SDK Contribution checklist:

  • The pull request does not introduce [breaking changes]
  • CHANGELOG is updated for new features, bug fixes or other significant changes.
  • I have read the contribution guidelines.

General Guidelines and Best Practices

  • Title of the pull request is clear and informative.
  • There are a small number of commits, each of which have an informative message. This means that previously merged commits do not appear in the history of the PR. For more information on cleaning up the commits in your PR, see this page.

Testing Guidelines

  • Pull request includes test coverage for the included changes.

@aayush3011 aayush3011 marked this pull request as ready for review May 13, 2026 21:02
@aayush3011 aayush3011 requested a review from a team as a code owner May 13, 2026 21:02
Copilot AI review requested due to automatic review settings May 13, 2026 21:02
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Documents and adds test coverage for the new optional embeddingSource object on entries in vector_embedding_policy.vectorEmbeddings for Azure Cosmos DB. No SDK code logic changes were needed since vector_embedding_policy is passed through as a dict.

Changes:

  • Adds README section and docstring updates in sync/async database.py describing the new embeddingSource fields (sourcePaths, deploymentName, modelName, endpoint, authType).
  • Adds new sync and async tests verifying round-trip of an embeddingSource policy for both ApiKey and Entra auth types, plus a negative test for an invalid authType.
  • Adds a CHANGELOG entry under the 4.16.0b3 (Unreleased) section.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
sdk/cosmos/azure-cosmos/CHANGELOG.md Documents the new optional embeddingSource field.
sdk/cosmos/azure-cosmos/README.md Adds a sample policy snippet showing embeddingSource usage.
sdk/cosmos/azure-cosmos/azure/cosmos/database.py Extends docstrings of all create_container/create_container_if_not_exists overloads.
sdk/cosmos/azure-cosmos/azure/cosmos/aio/_database.py Same docstring extensions for the async client.
sdk/cosmos/azure-cosmos/tests/test_vector_policy.py Adds positive (ApiKey/Entra) and negative (invalid authType) sync tests.
sdk/cosmos/azure-cosmos/tests/test_vector_policy_async.py Adds equivalent async tests.

Comment thread sdk/cosmos/azure-cosmos/CHANGELOG.md Outdated
@aayush3011 aayush3011 requested a review from Copilot May 13, 2026 21:07
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Comment thread sdk/cosmos/azure-cosmos/tests/test_vector_policy.py Outdated
Comment thread sdk/cosmos/azure-cosmos/tests/test_vector_policy.py Outdated
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

@aayush3011
Copy link
Copy Markdown
Member Author

/azp run python - cosmos - tests

@azure-pipelines
Copy link
Copy Markdown

Azure Pipelines successfully started running 1 pipeline(s).

Copy link
Copy Markdown
Member

@FabianMeiswinkel FabianMeiswinkel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Copy link
Copy Markdown
Member

@kushagraThapar kushagraThapar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @aayush3011 for documenting the EGS surface so cleanly — the dict-passthrough approach makes this a remarkably small change for a meaningful new capability, and I love that the framing in the PR body is honest about "no SDK code changes were required."

A few things I verified up front that I think are worth mentioning in the review summary:

  • Wire-shape parity with the .NET counterpart (Azure/azure-cosmos-dotnet-v3#5848) is exactsourcePaths, deploymentName, modelName, endpoint, authType ∈ {ApiKey, Entra} all match the typed EmbeddingSource / EmbeddingAuthType shape introduced there. Your negative test (test_fail_create_vector_embedding_source_invalid_auth_type) confirms the service literal is the bare-word Entra (not EntraID / MicrosoftEntra / AAD), which is the same enum member name the .NET PR uses. So whatever EGS work lands later in this PR is contract-faithful.
  • Cross-SDK landscape: Python is the second SDK to expose any EGS surface (after .NET). I verified Java / JS / Go / Rust currently have no embeddingSource field on their typed VectorEmbedding models and no open PRs to add one. Java in particular would silently drop the field on round-trip today (Jackson, no @JsonAnySetter on CosmosVectorEmbedding) — meaning a container created via Python or .NET with EGS would not round-trip correctly through Java. That's not blocking on this PR — those are work items I'll file separately against the peer SDKs — but worth being aware of for cross-customer scenarios.
  • Design-doc follow-up: there is no chapter in cosmosdb-design-docs/ on vector embedding policy / EGS today (the existing 08-indexing-engine.md covers only Inverted/Range/Spatial/Composite indexes). I'll open a separate PR against tvaron3/cosmosdb-design-docs proposing a ## Vector Indexing & Vector Embedding Policy section so the next SDK doing this work has a canonical reference to cite. Also separate from this PR.

Approving on the basis that nothing here blocks merge — all my notes are docs / test improvements that you could fold into this PR or a follow-up, your call. There are five themes I'd love your thoughts on; I'll drop them inline. A handful of true nits I'm bundling here so the inline thread doesn't get noisy:

  • NIT-1Wording consistency: the README says the embedding is "generated by the service" (which I find materially more informative because it makes the server-side-generation contract explicit), but the docstrings say "produced". Could we align the docstrings to use the README's phrasing across all 12 sites?
  • NIT-2Tests use strict == round-trip: assert properties["vectorEmbeddingPolicy"] == vector_embedding_policy at test_vector_policy.py:633 will be fragile if the service ever stamps back default fields on embeddingSource (e.g. a future authVersion). It's an inherited pre-existing pattern, not something this PR introduced — just flagging so we have it in mind once EGS GAs and we un-skip the tests.
  • NIT-3PR title oversells scope: the title is "Adding embedding policy changes for EGS" but the body itself says "no SDK code changes were required". Could we retitle to something like "[Cosmos] Document optional embeddingSource on vector_embedding_policy (docs + tests only)" to make the git log / release-notes search clearer?
  • NIT-4Agentic surface (optional): two low-cost ergonomic improvements for LLM-driven consumers — could you consider embedding the README's code-fence example directly into the docstring (LLMs extract JSON-shape from triple-back-ticked examples more reliably than from prose), and/or adding an EMBEDDING_AUTH_TYPES: Final = ("ApiKey", "Entra") constants tuple to _constants.py so agents have a stable symbol to reference? Both are optional and consistent with Python's dict-passthrough idiom.
  • NIT-5Unregistered pytest marker: the @pytest.mark.cosmosSearchQuery class marker the EGS tests inherit isn't registered in sdk/cosmos/azure-cosmos/pytest.ini (siblings like semanticReranker are). It emits a PytestUnknownMarkWarning on every collection even though the tests themselves are @unittest.skip'd at HEAD and don't actually run. Since pytest.ini isn't part of this PR I can't drop an inline, but if you're touching the file anyway, could you please add a registration line parallel to semanticReranker: cosmosSearchQuery: marks tests requiring a Cosmos account with vector search / EGS preview enabled.

Thanks again — this is a tight PR and I appreciate the cleanup-hygiene try/finally refactor in the test file too. 🙏

Comment thread sdk/cosmos/azure-cosmos/CHANGELOG.md Outdated
Comment thread sdk/cosmos/azure-cosmos/README.md
Comment thread sdk/cosmos/azure-cosmos/azure/cosmos/database.py
Comment thread sdk/cosmos/azure-cosmos/README.md Outdated
Comment thread sdk/cosmos/azure-cosmos/README.md
Comment thread sdk/cosmos/azure-cosmos/README.md
Comment thread sdk/cosmos/azure-cosmos/tests/test_vector_policy.py
Comment thread sdk/cosmos/azure-cosmos/tests/test_vector_policy.py Outdated
Comment thread sdk/cosmos/azure-cosmos/tests/test_vector_policy.py Outdated
Comment thread sdk/cosmos/azure-cosmos/tests/test_vector_policy.py Outdated
Copy link
Copy Markdown
Member

@simorenoh simorenoh left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM pending Kushagra's comments on general supportability

Copy link
Copy Markdown
Member

@kushagraThapar kushagraThapar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM ✅ — thanks for the thorough turn on the resolution commits. All 13 prior review threads are addressed (CHANGELOG/README preview-marker wording, all 12 docstring **provisional** sites, the EGS heading, endpoint/Entra clarifications, the three service-side invariants, the offline contract-test sentinel, the skipUnless(os.getenv("COSMOS_ENABLE_EGS_TESTS")) switch, the expanded negative-fields test, and the shared VectorPolicyTestData entry).

Two trivial follow-ups worth fixing in a quick commit before merge (non-blocking — happy for either of these to ride in a follow-up PR if easier):

  1. Duplicate import copy in sdk/cosmos/azure-cosmos/tests/test_vector_policy.py (lines 4–5). Pylint will flag it:

    import copy
    import copy
    import os

    Just delete the second occurrence.

  2. TestVectorEmbeddingPolicyOffline is defined after if __name__ == '__main__': unittest.main() at the bottom of test_vector_policy.py. Pytest still discovers it (which is why CI is green), but python test_vector_policy.py won't execute the new offline regression sentinel via the unittest.main() path, and the placement reads oddly. Please move the class TestVectorEmbeddingPolicyOffline(...) block above the if __name__ == '__main__': guard.

Optional (truly non-blocking, file under "if you're already touching the file"): the new test_fail_create_vector_embedding_source_missing_required_fields covers sourcePaths/endpoint/deploymentName but not modelName or authType missing. Two more entries in the loop would round out the required-field contract pinning. Fine to defer.

@aayush3011 aayush3011 merged commit 2dc9ce7 into Azure:main May 20, 2026
35 checks passed
ninghu pushed a commit to ninghu/azure-sdk-for-python that referenced this pull request May 22, 2026
* Adding embedding policy changes for EGS

* Updating changelog

* Updating changelog

* Resolving comments

* Resolving comments

* Fixing builds

---------

Co-authored-by: Aayush Kataria <aayushkataria@Aayushs-MacBook-Pro-2.local>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

5 participants