Python: Add agent-framework-azure-cosmos-memory context provider#6719
Open
TheovanKraay wants to merge 20 commits into
Open
Python: Add agent-framework-azure-cosmos-memory context provider#6719TheovanKraay wants to merge 20 commits into
TheovanKraay wants to merge 20 commits into
Conversation
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. |
Contributor
Python Test Coverage Report •
Python Unit Test Overview
|
||||||||||||||||||||||||||||||
ab30779 to
4bc5c8e
Compare
1150d87 to
f110b8a
Compare
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
- Wire credential into Cosmos and AI Foundry clients; let toolkit own DefaultAzureCredential when none supplied (remove dead import). - Honor auto_extract=False by zeroing extraction/summary cadence thresholds. - Skip whitespace-only conversation turns and store stripped content. - Show confidence 0.0 and coerce confidence to float in _format_memories. - Register both 'integration' and 'azure' pytest markers accurately. - Fix duplicated install block in README. - Update and extend unit tests for new credential wiring and fixes.
Follow the github_copilot pattern for a package with a Python 3.11-only dependency: lower requires-python to >=3.10 and gate azure-cosmos-agent-memory behind a python_version >= '3.11' marker. Add a direct, gated prompty pre-release dependency so the workspace's if-necessary-or-explicit prerelease policy permits the toolkit's transitive prompty requirement. Guard the test modules with pytest.importorskip so the 3.10 CI leg skips cleanly. Remove the workspace exclude and the markdown-code-lint exclude, and regenerate uv.lock.
Rename provider parameters to match Agent Framework conventions: foundry_endpoint (was ai_foundry_endpoint) and embedding_model/chat_model (were *_deployment_name). Move DEFAULT_* to module-level constants, type memory_types as a Literal, use DEFAULT_CONTEXT_PROMPT as the default value, and add ProcessorConfig/CosmosMemorySettings TypedDicts. Resolve connection settings via agent_framework load_settings with required-field validation, replacing the manual getenv/raise blocks. Scope user_id/thread_id to the provider state and drop the unpreventable first-turn warning. Rewrite the samples around Agent (not raw SessionContext), provider-scoped state, and session-id threading; use PEP 723 inline dependencies instead of a samples dependency group; use a plain input() loop; remove the dead custom processor stub. Update README/AGENTS for the renamed parameters and env vars. Add a samples ruff per-file-ignores entry now that the package is linted in CI.
Bump azure-cosmos-agent-memory to >=0.2.0b2 (adds the embeddings/chat client injection seam) and add tests/test_emulator.py: an integration (not azure) suite that exercises real Cosmos vector search with a quantizedFlat index against a local Cosmos DB emulator, using deterministic in-memory fakes for embeddings and chat so no Azure AI Foundry account or LLM is required. To run on a stock emulator the fixture strips the toolkit's full-text index (the provider only does pure vector search) and requests provisioned autoscale throughput instead of serverless. The suite skips cleanly when no emulator is reachable.
The package recently joined the uv workspace, so its source and tests are now covered by the Test Typing Checks and Package Checks gates for the first time. tests: rename stale constructor kwargs to the current provider API (foundry_endpoint/embedding_model/chat_model); use a typed _STUB_AGENT for the unused agent param so pyright/pyrefly/ty/zuban all accept it; make processor_config values ints; assert non-None memory_client in the emulator tests. source: relax reportUnknown*/reportOptional* for this package only (the toolkit ships no py.typed; mirrors the hosting-telegram precedent); decouple the conditional toolkit import from the annotation type; use settings.get(); fix memory_types list invariance; drop a redundant None guard; read role via getattr.
The provider now drains in-flight background memory extraction in __aexit__, so applications no longer need to call flush() in their own control flow; the client's close() would otherwise cancel pending extraction tasks. flush() is hardened against clients that expose no usable background-task registry. sample: interactive_chat reads input via asyncio.to_thread so the event loop stays free and background extraction runs during the session; removes the manual flush now that the provider drains on exit. tests: add explicit transparent-extraction integration tests (emulator: after_run schedules extraction and __aexit__ drains it; live Azure: a fact is extracted and recalled in a later session with no manual flush). Emulator tests reuse a single fixed database to avoid exhausting the emulator's partition budget across runs.
Adds a prompts_dir option to CosmosMemoryContextProvider that points the Agent Memory Toolkit pipeline at a caller-supplied directory of Prompty templates, so callers can override extract_memories.prompty to control what the extraction LLM produces. The toolkit exposes no public prompts-directory seam, so the provider contains the one internal touch (swapping the pipeline's template loader after the store connects); applies to both provider-built and supplied clients. sample: interactive_chat_custom_extraction.py - the interactive chat wired with a custom coding-assistant extraction rubric. It derives a complete prompts directory at runtime (copies the bundled templates and augments extract_memories.prompty) so it stays schema-compatible with the installed toolkit. tests: unit tests assert the provider redirects the pipeline loader only when prompts_dir is set; an emulator integration test proves end to end that a unique marker in a custom extract_memories.prompty reaches the extraction LLM call.
…ADME Replaces the stale, non-functional CustomMemoryProcessor snippet with the working prompts_dir approach, lists the new interactive_chat_custom_extraction.py sample, and corrects the interactive-sample feature list.
Sample (comment): rename _new_thread to _new_session in both interactive samples (a new session is the new thread). Provider (comment): replace the _memory_toolkit_available flag + __init__ ImportError guard with a plain guarded import that re-raises a clear ImportError, matching the github_copilot package's pattern for its 3.11-only SDK. Kept requires-python >=3.10 (bumping this one workspace member to 3.11 would force the entire uv workspace lock floor to 3.11). Tests now run importorskip before importing the package, mirroring github_copilot.
36e7769 to
cb70fcc
Compare
moonbox3
reviewed
Jul 13, 2026
added 3 commits
July 15, 2026 11:30
…ry-context-provider # Conflicts: # python/uv.lock
…ry-context-provider
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