Skip to content

Python: Add Azure Cosmos history provider package#4271

Merged
eavanvalkenburg merged 6 commits intomicrosoft:mainfrom
eavanvalkenburg:cosmos_history_provider
Mar 3, 2026
Merged

Python: Add Azure Cosmos history provider package#4271
eavanvalkenburg merged 6 commits intomicrosoft:mainfrom
eavanvalkenburg:cosmos_history_provider

Conversation

@eavanvalkenburg
Copy link
Member

@eavanvalkenburg eavanvalkenburg commented Feb 25, 2026

Summary

  • add new Python package agent-framework-azure-cosmos
  • implement CosmosHistoryProvider for conversation history in Azure Cosmos DB
  • add batching for save/clear operations with execute_item_batch and add list_sessions
  • add package-local sample and documentation updates

Validation

  • uv run --directory packages/azure-cosmos poe test
  • uv run --directory packages/azure-cosmos poe lint
  • uv run --directory packages/azure-cosmos poe pyright
  • uv run --directory packages/azure-cosmos poe mypy

Fixes #1390

Copilot AI review requested due to automatic review settings February 25, 2026 19:21
@markwallace-microsoft markwallace-microsoft added documentation Improvements or additions to documentation python labels Feb 25, 2026
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds a new Python workspace package (agent-framework-azure-cosmos) that implements an Azure Cosmos DB–backed BaseHistoryProvider for persisting conversation history, along with unit tests and a runnable sample. This supports the Agent Framework’s pluggable “context/history provider” story similarly to existing integrations (e.g., Redis).

Changes:

  • Introduce agent-framework-azure-cosmos package with CosmosHistoryProvider (Cosmos DB transactional batch writes + session partitioning).
  • Add unit tests and package-local sample/README for the Cosmos history provider.
  • Wire the new package into the Python workspace (pyproject + uv.lock) and apply minor formatting cleanups in existing tests/modules.

Reviewed changes

Copilot reviewed 15 out of 16 changed files in this pull request and generated 7 comments.

Show a summary per file
File Description
python/uv.lock Adds the new workspace member and locks azure-cosmos dependency.
python/pyproject.toml Registers agent-framework-azure-cosmos as a uv workspace source.
python/packages/core/tests/core/test_skills.py Minor whitespace-only change.
python/packages/core/tests/core/test_function_invocation_logic.py Formatting of assertion messages (single-line f-strings).
python/packages/core/agent_framework/_skills.py Minor formatting/line-wrapping adjustments.
python/packages/azure-cosmos/pyproject.toml New package metadata, dependencies, and tooling config.
python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py Implements CosmosHistoryProvider (get/save/clear/list + batching + container creation).
python/packages/azure-cosmos/agent_framework_azure_cosmos/init.py Exports CosmosHistoryProvider + version discovery.
python/packages/azure-cosmos/tests/test_cosmos_history_provider.py New unit test suite for the provider (mocked Cosmos client/container).
python/packages/azure-cosmos/samples/cosmos_history_provider.py Runnable sample demonstrating agent usage with Cosmos-backed history.
python/packages/azure-cosmos/samples/init.py Samples package marker + docstring.
python/packages/azure-cosmos/samples/README.md Sample runner documentation.
python/packages/azure-cosmos/README.md Package-level “getting started” documentation.
python/packages/azure-cosmos/LICENSE Package license file.
python/packages/azure-cosmos/AGENTS.md Package agent documentation (class + import path).
python/packages/azure-ai-search/tests/test_aisearch_context_provider.py Removes redundant local imports + minor formatting.

Copy link
Contributor

@moonbox3 moonbox3 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated Code Review

Reviewers: 3 | Confidence: 80%

✗ Correctness

The diff adds a new Azure Cosmos DB history provider package with solid implementation and tests, applies formatting-only changes to several existing files, and removes six local from agent_framework import Content imports from test methods in test_aisearch_context_provider.py. The removal of these local imports is suspicious: there is no corresponding addition of a module-level import for Content anywhere in the diff, which would cause NameError at runtime in all six affected test methods. The rest of the changes (new Cosmos package, formatting, pytest marker additions) are correct.

✓ Security Reliability

New Azure Cosmos DB history provider is generally well-structured: parameterized queries prevent injection, credential handling uses SecretString properly, and resource ownership is tracked for cleanup. The main concerns are: (1) a None session_id silently maps to a shared "default" partition, risking unintended cross-session data leakage; (2) partial batch failures in save_messages/clear are unhandled, leaving data in an inconsistent state; and (3) the __aexit__ can swallow the original exception if close() also raises. The remaining changes are formatting-only or test import hoisting and carry no risk.

✓ Test Coverage

The new CosmosHistoryProvider has solid test coverage for core CRUD operations (get, save, clear, list_sessions), initialization variants, lifecycle management (close, async context manager), and before/after run hooks. However, several non-trivial code paths lack tests: the batch-splitting logic when operations exceed _BATCH_OPERATION_LIMIT (100), the _session_partition_key fallback to 'default' when session_id is None, the _resolve_credential ValueError path when neither credential nor env key is provided, the RuntimeError in _get_container when database_client is None, and the filtering of non-dict message payloads in get_messages. The existing tests in test_aisearch_context_provider.py had only import-cleanup changes (moving Content import to module level) and formatting, which are fine.

Blocking Issues

  • Six test methods in test_aisearch_context_provider.py remove the local from agent_framework import Content import but no module-level import for Content is added in this diff, which will cause NameError: name 'Content' is not defined when running these tests.

Suggestions

  • Consider adding from agent_framework import Content as a module-level import at the top of test_aisearch_context_provider.py (alongside the existing imports) to replace the removed local imports.
  • Consider raising an error or generating a unique key when session_id is None in _session_partition_key rather than silently falling back to a shared "default" partition, which can leak messages across unrelated callers.
  • Add error handling around execute_item_batch in save_messages and clear so a partial batch failure doesn't silently leave data in an inconsistent state (e.g., log the error, raise, or retry the remaining batch).
  • Guard close() inside __aexit__ with a try/except so that a failure during cleanup doesn't mask the original exception that triggered the exit.
  • Add a test for batch splitting in save_messages and clear when the number of items exceeds _BATCH_OPERATION_LIMIT (100). This is a non-trivial loop that could silently lose data if broken.
  • Add a test for _session_partition_key returning 'default' when session_id is None — this is a meaningful behavioral contract that callers may depend on.
  • Add a test for _resolve_credential raising ValueError when both credential and settings key are None (the third branch in that method).
  • Add a test for _get_container raising RuntimeError when _database_client is None, to ensure the guard clause works correctly.
  • Add a test for get_messages skipping non-dict message payloads (e.g., a malformed item with a string or None 'message' field), since the code explicitly checks isinstance(message_payload, dict).
  • The after_run test (test_after_run_stores_input_and_response) only asserts on len(batch_operations) == 2 but doesn't verify the content of the stored messages (roles, text). Consider adding assertions on the actual message payloads to match the thoroughness of test_saves_messages.

Automated review by moonbox3's agents

@eavanvalkenburg eavanvalkenburg force-pushed the cosmos_history_provider branch from 7845133 to 54e44d3 Compare March 2, 2026 09:06
@markwallace-microsoft
Copy link
Member

markwallace-microsoft commented Mar 2, 2026

Python Test Coverage

Python Test Coverage Report •
FileStmtsMissCoverMissing
TOTAL22253275787% 
report-only-changed-files is enabled. No files were changed during this commit :)

Python Unit Test Overview

Tests Skipped Failures Errors Time
4523 25 💤 0 ❌ 0 🔥 1m 20s ⏱️

@eavanvalkenburg eavanvalkenburg added this pull request to the merge queue Mar 3, 2026
@github-merge-queue github-merge-queue bot removed this pull request from the merge queue due to failed status checks Mar 3, 2026
@eavanvalkenburg eavanvalkenburg added this pull request to the merge queue Mar 3, 2026
@github-merge-queue github-merge-queue bot removed this pull request from the merge queue due to failed status checks Mar 3, 2026
eavanvalkenburg and others added 6 commits March 3, 2026 11:55
- address provider/test/sample review feedback and cleanup typing
- add cosmos integration test coverage and skip gating
- add dedicated cosmos emulator jobs to python merge/integration workflows
- switch cosmos workflow execution to package poe integration-tests task

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- replace default partition fallback for empty session_id
- log warning and generate GUID when session_id is empty
- update unit tests to validate GUID fallback behavior

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@eavanvalkenburg eavanvalkenburg force-pushed the cosmos_history_provider branch from c8e5c8f to f852551 Compare March 3, 2026 10:55
@eavanvalkenburg eavanvalkenburg added this pull request to the merge queue Mar 3, 2026
Merged via the queue into microsoft:main with commit c37f74f Mar 3, 2026
31 of 35 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

.NET: Python: Add support for Azure Cosmos DB as a chat message store and a provider

6 participants