-
Notifications
You must be signed in to change notification settings - Fork 945
.NET: Add Cosmos DB implementations for ChatMessageStore and CheckpointStore. #1838
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
.NET: Add Cosmos DB implementations for ChatMessageStore and CheckpointStore. #1838
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull Request Overview
This PR adds Azure Cosmos DB integration to the Agent Framework, providing persistent storage for chat messages and workflow checkpoints. The implementation includes both simple and hierarchical partition key support, along with comprehensive test coverage.
Key Changes
- New Cosmos DB implementations for chat message storage (
CosmosChatMessageStore) and checkpoint storage (CosmosCheckpointStore) - Extension methods for easy configuration of Cosmos DB stores
- Support for connection string, managed identity, and CosmosClient-based authentication
- Comprehensive unit tests for all Cosmos DB components
Reviewed Changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 33 comments.
Show a summary per file
| File | Description |
|---|---|
| dotnet/Directory.Packages.props | Added package references for Microsoft.Azure.Cosmos (3.53.1) and Newtonsoft.Json (13.0.3) |
| dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs | Implements chat message storage with Cosmos DB, supporting both simple and hierarchical partitioning |
| dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosAgentThread.cs | Abstract base class for agent threads using Cosmos DB for persistence |
| dotnet/src/Microsoft.Agents.AI/CosmosDBChatExtensions.cs | Extension methods for configuring Cosmos DB message stores in ChatClientAgentOptions |
| dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs | Implements checkpoint storage with Cosmos DB for workflow persistence |
| dotnet/src/Microsoft.Agents.AI.Workflows/CosmosDBWorkflowExtensions.cs | Extension methods for creating Cosmos DB checkpoint stores |
| dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs | Comprehensive tests for CosmosChatMessageStore including hierarchical partitioning |
| dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs | Tests for CosmosAgentThread functionality and message operations |
| dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs | Tests for CosmosCheckpointStore covering checkpoint operations and error handling |
| Multiple .csproj files | Added package references to Azure.Identity, Microsoft.Azure.Cosmos, and Newtonsoft.Json |
Motivation and Context
This change adds Cosmos DB implementations for Agent Framework's storage abstractions to provide persistent storage options for chat messages and workflow checkpoints.
Partially addresses issue (.NET only): #1390
Description
Added Cosmos DB implementations for three core storage abstractions:
CosmosChatMessageStore- ImplementsChatMessageStorefor chat message persistenceCosmosCheckpointStore- ImplementsJsonCheckpointStorefor workflow state persistenceAdded convenience extensions:
CosmosDBChatExtensions- Extension methods for configuring chat agents with Cosmos DBCosmosDBWorkflowExtensions- Factory methods for creating checkpoint storesFeatures implemented:
/tenantId/userId/sessionId)Added corresponding unit tests.
Contribution Checklist