Allow injecting embeddings and chat clients#27
Open
TheovanKraay wants to merge 1 commit into
Open
Conversation
Add optional embeddings_client and chat_client constructor arguments to CosmosMemoryClient and AsyncCosmosMemoryClient. When provided, the toolkit uses the injected client instead of building an Azure-backed one and does not close it (the caller owns its lifecycle, mirroring credential ownership). Enables OpenAI-compatible/self-hosted embedding and chat backends, reuse of a caller-configured client, and deterministic offline testing. Adds sync and async unit tests covering injection, ownership flags, and close() behavior.
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a supported dependency-injection seam for the model-layer clients used by the Cosmos memory clients, allowing callers to supply their own embeddings/chat backends (or fakes) while preserving existing default Azure-backed behavior and lifecycle semantics.
Changes:
- Add optional
embeddings_clientandchat_clientconstructor args toCosmosMemoryClientandAsyncCosmosMemoryClient, with ownership flags controlling whetherclose()tears them down. - Update
close()implementations to only close clients constructed by the toolkit. - Add sync/async unit tests and a CHANGELOG entry documenting the new injection capability.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
azure/cosmos/agent_memory/cosmos_memory_client.py |
Adds injectable embeddings/chat clients and ownership-aware closing for the sync client. |
azure/cosmos/agent_memory/aio/cosmos_memory_client.py |
Adds injectable embeddings/chat clients and ownership-aware closing for the async client. |
tests/unit/test_cosmos_memory_client.py |
Adds sync unit tests verifying injection behavior and lifecycle ownership. |
tests/unit/aio/test_cosmos_memory_client.py |
Adds async unit tests verifying injection behavior and lifecycle ownership. |
CHANGELOG.md |
Documents the new injection feature under Unreleased. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
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.
Summary
Adds optional
embeddings_clientandchat_clientconstructor arguments to bothCosmosMemoryClientandAsyncCosmosMemoryClient. When supplied, the toolkit usesthe injected client instead of constructing an Azure-backed one, and treats the
caller as the owner of its lifecycle (so
close()does not close it). When omitted,behavior is unchanged: the toolkit builds and owns the Azure OpenAI / AI Foundry
clients exactly as before.
Motivation
Today the embeddings and chat clients are hard-constructed inside
__init__, with nosupported way to supply your own. That couples the toolkit to Azure OpenAI / AI Foundry
and forces anyone who needs a different backend (or offline/deterministic behavior) to
reach into private internals (
_embeddings_client/_chat_client). This adds a clean,public seam instead.
Use cases this enables:
services (via an OpenAI-compatible client), gateways, or alternate providers.
custom retry/timeout/proxy/telemetry can hand the same instance to the toolkit
instead of it creating a second, differently-configured one.
Cosmos vector path (e.g. against the Cosmos DB emulator) without any Azure account,
and without monkeypatching private internals.
Changes
CosmosMemoryClient.__init__andAsyncCosmosMemoryClient.__init__acceptembeddings_clientandchat_client(both defaultNone)._owns_embeddings_client/_owns_chat_client, mirroring theexisting
_owns_*_credentialpattern.close()only closes clients the toolkit built.independent injection of one client, and that
close()does not close injected clients.Compatibility
Fully backward compatible. Both arguments default to
None; existing callers get thecurrent Azure-backed clients and lifecycle with no changes.
Testing
tests/unit/test_cosmos_memory_client.pyandtests/unit/aio/test_cosmos_memory_client.pypass (128 tests across both suites).
ruff checkandruff format --checkclean on the changed files.