Python: Add Azure Cosmos history provider package#4271
Python: Add Azure Cosmos history provider package#4271eavanvalkenburg merged 6 commits intomicrosoft:mainfrom
Conversation
There was a problem hiding this comment.
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-cosmospackage withCosmosHistoryProvider(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. |
python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py
Outdated
Show resolved
Hide resolved
python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py
Outdated
Show resolved
Hide resolved
python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py
Outdated
Show resolved
Hide resolved
python/packages/azure-cosmos/samples/cosmos_history_provider.py
Outdated
Show resolved
Hide resolved
python/packages/azure-cosmos/samples/cosmos_history_provider.py
Outdated
Show resolved
Hide resolved
python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py
Outdated
Show resolved
Hide resolved
moonbox3
left a comment
There was a problem hiding this comment.
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 Contentimports from test methods intest_aisearch_context_provider.py. The removal of these local imports is suspicious: there is no corresponding addition of a module-level import forContentanywhere in the diff, which would causeNameErrorat 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
Nonesession_id silently maps to a shared "default" partition, risking unintended cross-session data leakage; (2) partial batch failures insave_messages/clearare unhandled, leaving data in an inconsistent state; and (3) the__aexit__can swallow the original exception ifclose()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.pyremove the localfrom agent_framework import Contentimport but no module-level import forContentis added in this diff, which will causeNameError: name 'Content' is not definedwhen running these tests.
Suggestions
- Consider adding
from agent_framework import Contentas a module-level import at the top oftest_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 Nonein_session_partition_keyrather than silently falling back to a shared"default"partition, which can leak messages across unrelated callers. - Add error handling around
execute_item_batchinsave_messagesandclearso 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
python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py
Outdated
Show resolved
Hide resolved
python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py
Show resolved
Hide resolved
python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py
Outdated
Show resolved
Hide resolved
7845133 to
54e44d3
Compare
- 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>
c8e5c8f to
f852551
Compare
Summary
agent-framework-azure-cosmosCosmosHistoryProviderfor conversation history in Azure Cosmos DBexecute_item_batchand addlist_sessionsValidation
uv run --directory packages/azure-cosmos poe testuv run --directory packages/azure-cosmos poe lintuv run --directory packages/azure-cosmos poe pyrightuv run --directory packages/azure-cosmos poe mypyFixes #1390