Python: Add agent-framework-azure-cosmos-memory context provider#6719
Open
TheovanKraay wants to merge 1 commit into
Open
Python: Add agent-framework-azure-cosmos-memory context provider#6719TheovanKraay wants to merge 1 commit into
TheovanKraay wants to merge 1 commit into
Conversation
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.
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" |
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
The Agent Framework has no first-party way to give agents long-term, durable memory backed by Azure Cosmos DB. This change adds an integration package that lets agents persist and recall facts, procedural knowledge, and user summaries across sessions, scoped per user. It contributes to the long-term/persistent memory scenario, complementing the existing
mem0provider with a Cosmos DB backed option built on theazure-cosmos-agent-memorytoolkit.Description & Review Guide
What are the major changes?
agent-framework-azure-cosmos-memoryprovidingCosmosMemoryContextProvider, aContextProviderwrapping theazure-cosmos-agent-memorytoolkit (fact + procedural recall and user-summary injection).pyproject.toml,README.md,AGENTS.md,LICENSE), samples, unit tests (mocked client), and marked live-Azure integration tests.uv.lockis intentionally not modified (see impact below).What is the impact of these changes?
azure-cosmos-agent-memory, which requires Python >=3.11, while the workspace is >=3.10. Becausepackages/*are auto-included as workspace members, a shareduv.lockcan't be regenerated until maintainers choose either (recommended) excluding this package via[tool.uv.workspace] excludeto keep the monorepo at >=3.10, or raising the workspace floor to >=3.11. CI lock-check will be red until this is resolved.diskANNindex.What do you want reviewers to focus on?
CosmosMemoryContextProviderpublic surface/design and its alignment with the existingmem0provider, and a steer on the workspace/uv.lockdecision above.Related Issue
Fixes #
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix), a workflow keeps the label and title prefix in sync automatically.