Python: Add agent-framework-azure-cosmos-memory context provider#6719
Open
TheovanKraay wants to merge 4 commits into
Open
Python: Add agent-framework-azure-cosmos-memory context provider#6719TheovanKraay wants to merge 4 commits into
TheovanKraay wants to merge 4 commits into
Conversation
moonbox3
added
documentation
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new Python integration package, agent-framework-azure-cosmos-memory, introducing a CosmosMemoryContextProvider that persists and recalls long-term memories via Azure Cosmos DB using the azure-cosmos-agent-memory toolkit (including user-summary injection and retrieval-time context augmentation).
Changes:
- Introduces
CosmosMemoryContextProvider(async context manager +before_run/after_runhooks) plus package exports. - Adds documentation and samples demonstrating basic usage and an interactive Foundry-backed chat experience.
- Adds unit tests (mocked client) and live-Azure integration tests (pytest markers).
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py | Implements the Cosmos-backed memory context provider (retrieval + storage + user-summary injection + flush). |
| python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/init.py | Exports CosmosMemoryContextProvider and package version. |
| python/packages/azure-cosmos-memory/pyproject.toml | Defines the new package, dependencies, and pytest/tooling configuration (Python >=3.11). |
| python/packages/azure-cosmos-memory/README.md | End-user documentation, configuration guidance, and usage examples. |
| python/packages/azure-cosmos-memory/AGENTS.md | Package-level developer guidance and key behaviors (user_id/thread_id, flush). |
| python/packages/azure-cosmos-memory/LICENSE | Package license. |
| python/packages/azure-cosmos-memory/samples/basic_usage.py | Minimal “raw hooks” sample calling before_run()/after_run() directly. |
| python/packages/azure-cosmos-memory/samples/interactive_chat.py | Interactive CLI sample demonstrating an agent wired with Foundry + Cosmos memory provider. |
| python/packages/azure-cosmos-memory/tests/test_context_provider.py | Unit tests for provider behavior with a mocked memory client (incl. context manager + flush). |
| python/packages/azure-cosmos-memory/tests/test_integration.py | Live-Azure integration tests gated by env vars and markers. |
| python/packages/azure-cosmos-memory/tests/conftest.py | Pytest marker registration for the package’s tests. |
Comment on lines
+156
to
+169
| # Create Azure credential using the standard chain: EnvironmentCredential → | ||
| # ManagedIdentityCredential → AzureCliCredential → InteractiveBrowserCredential. | ||
| # This works seamlessly in production (via ManagedIdentity) and local dev (via az login). | ||
| if credential is None: | ||
| credential = DefaultAzureCredential() # type: ignore | ||
|
|
||
| memory_client = AsyncCosmosMemoryClient( | ||
| cosmos_endpoint=cosmos_endpoint, | ||
| cosmos_database=cosmos_database, | ||
| ai_foundry_endpoint=ai_foundry_endpoint, | ||
| embedding_deployment_name=embedding_deployment_name, | ||
| chat_deployment_name=chat_deployment_name, | ||
| use_default_credential=True, | ||
| ) |
Comment on lines
+318
to
+362
| """Store conversation turns and optionally trigger memory extraction. | ||
|
|
||
| Args: | ||
| agent: The agent that ran this invocation. | ||
| session: The current session. | ||
| context: The invocation context with response populated. | ||
| state: Provider-scoped mutable state. | ||
| """ | ||
| # Get user_id and thread_id from state or session (warns once if no stable user_id) | ||
| user_id = self._resolve_user_id(state, session) | ||
| thread_id = state.get("thread_id") or session.state.get("thread_id") or session.session_id or "default" | ||
|
|
||
| try: | ||
| # Store input messages | ||
| for msg in context.input_messages: | ||
| if hasattr(msg, "role") and hasattr(msg, "text") and msg.text: | ||
| role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role) | ||
| if role_value in {"user", "assistant", "system"}: | ||
| await self.memory_client.add_cosmos( | ||
| user_id=user_id, | ||
| thread_id=thread_id, | ||
| role=self._ROLE_MAP.get(role_value, role_value), | ||
| content=msg.text, | ||
| ) | ||
|
|
||
| # Store response messages | ||
| if context.response and context.response.messages: | ||
| for msg in context.response.messages: | ||
| if hasattr(msg, "role") and hasattr(msg, "text") and msg.text: | ||
| role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role) | ||
| if role_value in {"user", "assistant", "system"}: | ||
| await self.memory_client.add_cosmos( | ||
| user_id=user_id, | ||
| thread_id=thread_id, | ||
| role=self._ROLE_MAP.get(role_value, role_value), | ||
| content=msg.text, | ||
| ) | ||
|
|
||
| # Auto-extraction and processing: | ||
| # The AsyncCosmosMemoryClient uses an InProcessProcessor that runs in the background | ||
| # and automatically extracts facts, generates summaries, and reconciles memories based on | ||
| # configured thresholds (FACT_EXTRACTION_EVERY_N, DEDUP_EVERY_N, etc.). | ||
| # This happens asynchronously after add_cosmos() completes, so no explicit process_now() call is needed. | ||
| # To disable auto-extraction, set auto_extract=False and call memory_client.process_now() manually. | ||
|
|
Comment on lines
+331
to
+341
| # Store input messages | ||
| for msg in context.input_messages: | ||
| if hasattr(msg, "role") and hasattr(msg, "text") and msg.text: | ||
| role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role) | ||
| if role_value in {"user", "assistant", "system"}: | ||
| await self.memory_client.add_cosmos( | ||
| user_id=user_id, | ||
| thread_id=thread_id, | ||
| role=self._ROLE_MAP.get(role_value, role_value), | ||
| content=msg.text, | ||
| ) |
Comment on lines
+343
to
+355
| # Store response messages | ||
| if context.response and context.response.messages: | ||
| for msg in context.response.messages: | ||
| if hasattr(msg, "role") and hasattr(msg, "text") and msg.text: | ||
| role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role) | ||
| if role_value in {"user", "assistant", "system"}: | ||
| await self.memory_client.add_cosmos( | ||
| user_id=user_id, | ||
| thread_id=thread_id, | ||
| role=self._ROLE_MAP.get(role_value, role_value), | ||
| content=msg.text, | ||
| ) | ||
|
|
Comment on lines
+381
to
+389
| content = memory.get("content", "") | ||
| memory_type = memory.get("memory_type", "") | ||
| confidence = memory.get("confidence", 0.0) | ||
|
|
||
| # Format: [Type] Content (confidence: X.XX) | ||
| if memory_type and confidence: | ||
| formatted.append(f"[{memory_type}] {content} (confidence: {confidence:.2f})") | ||
| else: | ||
| formatted.append(content) |
Comment on lines
+8
to
+10
| def pytest_configure(config: pytest.Config) -> None: | ||
| """Register custom markers.""" | ||
| config.addinivalue_line("markers", "integration: mark test as integration test requiring live Azure accounts") |
Comment on lines
+170
to
+181
| 1. **Complete [Development Setup](#development-setup)** - Create venv and install package **with sample dependencies**: | ||
| ```bash | ||
| pip install -e ".[dev,samples]" | ||
| ``` | ||
| Or install separately: | ||
| ```bash | ||
| pip install -e ".[dev]" | ||
| pip install -e ".[samples]" | ||
| ``` with sample dependencies: | ||
| ```bash | ||
| pip install -e ".[dev,samples]" | ||
| ``` |
| description = "Azure Cosmos DB Agent Memory Toolkit integration for Microsoft Agent Framework - semantic memory with fact extraction and user profiles." | ||
| authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] | ||
| readme = "README.md" | ||
| requires-python = ">=3.11" |
Contributor
Author
There was a problem hiding this comment.
Excluded package.
Contributor
Python Test Coverage Report •
Python Unit Test Overview
|
Introduces CosmosMemoryContextProvider, a ContextProvider that wraps the azure-cosmos-agent-memory toolkit to give agents long-term, Cosmos DB-backed memory (fact/procedural recall + user summaries). Includes package scaffolding, unit tests (mocked client), live Azure integration tests (marked), samples, README, and AGENTS.md. Draft: uv.lock is intentionally left unchanged. This package depends on azure-cosmos-agent-memory (requires Python >=3.11), which is unsatisfiable against the workspace's current >=3.10 floor, so adding it to the shared lock requires a workspace decision (raise floor to 3.11 or exclude from workspace). Test coverage to be expanded.
The package depends on azure-cosmos-agent-memory which requires Python >=3.11 and a prompty pre-release (>=2.0.0a9). Both are unsatisfiable against the workspace's >=3.10 floor and pre-release policy, causing uv sync to fail in every Python CI job. Exclude the package from the shared workspace so it is resolved and tested as a standalone package.
- Strip trailing whitespace from package files (pre-commit trailing-whitespace hook) - Exclude the package README from markdown-code-lint: the package is excluded from the uv workspace, so its README snippets import a module that is not installed in the workspace env and Pyright cannot resolve it
TheovanKraay
force-pushed
the
feature/cosmos-memory-context-provider
branch
from
ab30779 to
4bc5c8e
Compare
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.
Motivation & Context
Right now there's no first-party way to give agents long-term, per-user memory backed by Azure Cosmos DB. This PR adds an optional package,
agent-framework-azure-cosmos-memory, that plugs Cosmos-backed memory into the standardContextProviderextension point. On each run it pulls the user's relevant past memories and running summary into the agent's context, then saves new memories after the run. So you get the "agent that remembers the user across sessions" scenario without coupling core to Cosmos.It builds on the
azure-cosmos-agent-memorytoolkit for storage, embeddings, and reconciliation, and installs as a standalone, opt-in package.Description & Review Guide
What are the major changes?
python/packages/azure-cosmos-memory/withCosmosMemoryContextProvider, implementing the framework'sContextProvidercontract.before_runretrieves per-user memories (semantic search) plus a user summary, injecting memories as context messages and the summary as instructions. Search and summary failures are handled separately so one doesn't block the other, and it warns once if there's nouser_id.after_runpersists new memories from the turn.basic_usage.py,interactive_chat.py), 32 unit tests, and an integration test scaffold.What is the impact of these changes?
python/pyproject.tomlvia[tool.uv.workspace] exclude. That's needed because its dependencyazure-cosmos-agent-memoryrequires Python>=3.11and apromptyprerelease (>=2.0.0a9), both of which are unsatisfiable against the monorepo's>=3.10floor andif-necessary-or-explicitprerelease policy. Excluding it keeps rootuv syncresolvable and leaves the committeduv.lockuntouched, while the package is built and tested standalone.3.11and enabling the prompty prerelease. Happy to go whichever way you prefer.What do you want reviewers to focus on?
python/pyproject.toml, and whether you want this standalone or as a workspace member.ContextProviderintegration semantics (how memories vs. summary are injected, and the independent failure handling).Related Issue
Fixes #
Contribution Checklist