Skip to content

Commit a057cd5

Browse files
TheovanKraayTheo van Kraay
TheovanKraay
and
Theo van Kraay
authored
Python: Add agent-framework-azure-cosmos-memory context provider (microsoft#6719)
* Add agent-framework-azure-cosmos-memory context provider (draft) 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. * ci: exclude azure-cosmos-memory from uv workspace resolution 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. * ci: fix code-quality failures for azure-cosmos-memory - 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 * Exclude azure-cosmos-memory README from markdown-code-lint task * Address PR review comments on cosmos-memory context provider - 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. * Include azure-cosmos-memory in the uv workspace 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. * Address review feedback on cosmos-memory provider 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. * Add emulator-backed vector search integration test 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. * Fix CI typing and package checks for azure-cosmos-memory 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. * Apply pyupgrade: single-arg AsyncGenerator in test_integration * Make Cosmos memory extraction drain transparently on provider exit 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. * Add custom extraction-prompt seam and sample to cosmos-memory provider 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. * docs: document prompts_dir custom-extraction seam in cosmos-memory README 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. * Address review: rename _new_session, drop defensive toolkit import guard 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. * Pass cadence via cadence_thresholds instead of mutating os.environ * Mark package alpha and drop private naming in samples * Require Python 3.11 and inject user summary as untrusted context * CI: exclude azure-cosmos-memory from uv sync on Python 3.10 * Re-trigger CI (flaky external link check) * Require chat/embedding models instead of silent defaults * Fix pyright: narrow resolved chat/embedding models to str --------- Co-authored-by: Theo van Kraay <thvankra@microsoft.com>
1 parent c66bb39 commit a057cd5

17 files changed

Lines changed: 4084 additions & 1953 deletions

File tree

‎.github/workflows/python-lab-tests.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ jobs:
7171
with:
7272
python-version: ${{ matrix.python-version }}
7373
os: ${{ runner.os }}
74-
exclude-packages: ${{ matrix.python-version == '3.10' && 'agent-framework-github-copilot' || '' }}
74+
exclude-packages: ${{ matrix.python-version == '3.10' && 'agent-framework-github-copilot agent-framework-azure-cosmos-memory' || '' }}
7575
env:
7676
# Configure a constant location for the uv cache
7777
UV_CACHE_DIR: /tmp/.uv-cache

‎.github/workflows/python-tests.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ jobs:
3838
with:
3939
python-version: ${{ matrix.python-version }}
4040
os: ${{ runner.os }}
41-
exclude-packages: ${{ matrix.python-version == '3.10' && 'agent-framework-github-copilot' || '' }}
41+
exclude-packages: ${{ matrix.python-version == '3.10' && 'agent-framework-github-copilot agent-framework-azure-cosmos-memory' || '' }}
4242
env:
4343
# Configure a constant location for the uv cache
4444
UV_CACHE_DIR: /tmp/.uv-cache

‎python/packages/azure-cosmos-memory/AGENTS.md‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
# Azure Cosmos DB Memory Package (agent-framework-azure-cosmos-memory)
2+
3+
Long-term semantic memory for agents, backed by Azure Cosmos DB via the
4+
[Azure Cosmos DB Agent Memory Toolkit](https://github.com/AzureCosmosDB/AgentMemoryToolkit).
5+
6+
## Main Classes
7+
8+
- **`CosmosMemoryContextProvider`** - Context provider that integrates Cosmos DB-backed
9+
semantic memory (facts, procedural/episodic memories, and user/thread summaries) into agents.
10+
11+
## Usage
12+
13+
```python
14+
from azure.identity.aio import DefaultAzureCredential
15+
from agent_framework_azure_cosmos_memory import CosmosMemoryContextProvider
16+
17+
provider = CosmosMemoryContextProvider(
18+
cosmos_endpoint="https://<account>.documents.azure.com:443/",
19+
cosmos_database="ai_memory",
20+
foundry_endpoint="https://<project>.services.ai.azure.com",
21+
credential=DefaultAzureCredential(),
22+
)
23+
```
24+
25+
## Import Path
26+
27+
```python
28+
from agent_framework_azure_cosmos_memory import CosmosMemoryContextProvider
29+
```
30+
31+
## Notes
32+
33+
- Requires the `azure-cosmos-agent-memory` toolkit and an AI Foundry endpoint (used for both
34+
embeddings and fact extraction).
35+
- Set a stable `user_id` in `state["user_id"]` or `session.state["user_id"]` for long-term,
36+
cross-session memory. Without it, memory scopes to the ephemeral session id and the provider
37+
logs a one-time warning.
38+
- Background fact extraction runs out-of-band after each turn. Call `provider.flush()` before
39+
shutdown so in-flight extraction completes before the client closes.
40+
- See `README.md` for full configuration, authentication, and processor-tuning options.

‎python/packages/azure-cosmos-memory/LICENSE‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
MIT License
2+
3+
Copyright (c) Microsoft Corporation.
4+
5+
Permission is hereby granted, free of charge, to any person obtaining a copy
6+
of this software and associated documentation files (the "Software"), to deal
7+
in the Software without restriction, including without limitation the rights
8+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9+
copies of the Software, and to permit persons to whom the Software is
10+
furnished to do so, subject to the following conditions:
11+
12+
The above copyright notice and this permission notice shall be included in all
13+
copies or substantial portions of the Software.
14+
15+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21+
SOFTWARE

0 commit comments

Comments
 (0)