From 3c0926b6705861d281f1d71c11b19b5df8b45f8f Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 11 Sep 2025 18:17:53 +0100 Subject: [PATCH] .NET: Add support for IChatReducer to InMemoryChatMessageStore (#462) * Add support for IChatReducer to InMemoryChatMessageStore * Fix typos. --- .../InMemoryChatMessageStore.cs | 76 +++++++++- .../InMemoryChatMessageStoreTests.cs | 137 ++++++++++++++++++ 2 files changed, 207 insertions(+), 6 deletions(-) diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/InMemoryChatMessageStore.cs b/dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/InMemoryChatMessageStore.cs index 086a94b9c8..61e56c5f18 100644 --- a/dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/InMemoryChatMessageStore.cs +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/InMemoryChatMessageStore.cs @@ -2,6 +2,7 @@ using System.Collections; using System.Collections.Generic; +using System.Linq; using System.Text.Json; using System.Threading; using System.Threading.Tasks; @@ -12,9 +13,39 @@ namespace Microsoft.Extensions.AI.Agents; /// /// Represents an in-memory store for chat messages associated with a specific thread. /// -internal sealed class InMemoryChatMessageStore : IList, IChatMessageStore +public sealed class InMemoryChatMessageStore : IList, IChatMessageStore { - private readonly List _messages = new(); + private readonly IChatReducer? _chatReducer; + private readonly ChatReducerTriggerEvent _reducerTriggerEvent; + private List _messages = new(); + + /// + /// Initializes a new instance of the class. + /// + public InMemoryChatMessageStore() + { + } + + /// + /// Initializes a new instance of the class. + /// + /// An optional instance used to process or reduce chat messages. If null, no reduction logic will be applied. + /// The event that should trigger the reducer invocation. + public InMemoryChatMessageStore(IChatReducer chatReducer, ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent.BeforeMessagesRetrieval) + { + this._chatReducer = Throw.IfNull(chatReducer); + this._reducerTriggerEvent = reducerTriggerEvent; + } + + /// + /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied. + /// + public IChatReducer? ChatReducer => this._chatReducer; + + /// + /// Gets the event that triggers the reducer invocation in this store. + /// + public ChatReducerTriggerEvent ReducerTriggerEvent => this._reducerTriggerEvent; /// public int Count => this._messages.Count; @@ -30,17 +61,32 @@ internal sealed class InMemoryChatMessageStore : IList, IChatMessag } /// - public Task AddMessagesAsync(IReadOnlyCollection messages, CancellationToken cancellationToken) + public async Task AddMessagesAsync(IReadOnlyCollection messages, CancellationToken cancellationToken) { _ = Throw.IfNull(messages); + + if (messages.Count == 0) + { + return; + } + this._messages.AddRange(messages); - return Task.CompletedTask; + + if (this._reducerTriggerEvent == ChatReducerTriggerEvent.AfterMessageAdded && this._chatReducer is not null) + { + this._messages = (await this._chatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList(); + } } /// - public Task> GetMessagesAsync(CancellationToken cancellationToken) + public async Task> GetMessagesAsync(CancellationToken cancellationToken) { - return Task.FromResult>(this._messages); + if (this._reducerTriggerEvent == ChatReducerTriggerEvent.BeforeMessagesRetrieval && this._chatReducer is not null) + { + this._messages = (await this._chatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList(); + } + + return this._messages; } /// @@ -118,4 +164,22 @@ internal sealed class InMemoryChatMessageStore : IList, IChatMessag { public IList Messages { get; set; } = new List(); } + + /// + /// Defines the events that can trigger a reducer in the . + /// + public enum ChatReducerTriggerEvent + { + /// + /// Trigger the reducer when a new message is added. + /// will only complete when reducer processing is done. + /// + AfterMessageAdded, + + /// + /// Trigger the reducer before messages are retrieved from the store. + /// The reducer will process the messages before they are returned to the caller. + /// + BeforeMessagesRetrieval + } } diff --git a/dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs index d11174eb91..8b1503157f 100644 --- a/dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Extensions.AI.Agents.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs @@ -6,6 +6,7 @@ using System.Linq; using System.Text.Json; using System.Threading; using System.Threading.Tasks; +using Moq; namespace Microsoft.Extensions.AI.Agents.Abstractions.UnitTests; @@ -14,6 +15,36 @@ namespace Microsoft.Extensions.AI.Agents.Abstractions.UnitTests; /// public class InMemoryChatMessageStoreTests { + [Fact] + public void Constructor_Throws_ForNullReducer() + { + // Arrange & Act & Assert + Assert.Throws(() => new InMemoryChatMessageStore(null!)); + } + + [Fact] + public void Constructor_DefaultsToBeforeMessageRetrieval_ForNotProvidedTriggerEvent() + { + // Arrange & Act + var reducerMock = new Mock(); + var store = new InMemoryChatMessageStore(reducerMock.Object); + + // Assert + Assert.Equal(InMemoryChatMessageStore.ChatReducerTriggerEvent.BeforeMessagesRetrieval, store.ReducerTriggerEvent); + } + + [Fact] + public void Constructor_Arguments_SetOnPropertiesCorrectly() + { + // Arrange & Act + var reducerMock = new Mock(); + var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.AfterMessageAdded); + + // Assert + Assert.Same(reducerMock.Object, store.ChatReducer); + Assert.Equal(InMemoryChatMessageStore.ChatReducerTriggerEvent.AfterMessageAdded, store.ReducerTriggerEvent); + } + [Fact] public async Task AddMessagesAsyncAddsMessagesAndReturnsNullThreadIdAsync() { @@ -413,4 +444,110 @@ public class InMemoryChatMessageStoreTests Assert.Same(message1, messages[0]); Assert.Same(message2, messages[1]); } + + [Fact] + public async Task AddMessagesAsync_WithReducer_AfterMessageAdded_InvokesReducerAsync() + { + // Arrange + var originalMessages = new List + { + new(ChatRole.User, "Hello"), + new(ChatRole.Assistant, "Hi there!") + }; + var reducedMessages = new List + { + new(ChatRole.User, "Reduced") + }; + + var reducerMock = new Mock(); + reducerMock + .Setup(r => r.ReduceAsync(It.Is>(x => x.SequenceEqual(originalMessages)), It.IsAny())) + .ReturnsAsync(reducedMessages); + + var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.AfterMessageAdded); + + // Act + await store.AddMessagesAsync(originalMessages, CancellationToken.None); + + // Assert + Assert.Single(store); + Assert.Equal("Reduced", store[0].Text); + reducerMock.Verify(r => r.ReduceAsync(It.Is>(x => x.SequenceEqual(originalMessages)), It.IsAny()), Times.Once); + } + + [Fact] + public async Task GetMessagesAsync_WithReducer_BeforeMessagesRetrieval_InvokesReducerAsync() + { + // Arrange + var originalMessages = new List + { + new(ChatRole.User, "Hello"), + new(ChatRole.Assistant, "Hi there!") + }; + var reducedMessages = new List + { + new(ChatRole.User, "Reduced") + }; + + var reducerMock = new Mock(); + reducerMock + .Setup(r => r.ReduceAsync(It.Is>(x => x.SequenceEqual(originalMessages)), It.IsAny())) + .ReturnsAsync(reducedMessages); + + var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.BeforeMessagesRetrieval); + await store.AddMessagesAsync(originalMessages, CancellationToken.None); + + // Act + var result = (await store.GetMessagesAsync(CancellationToken.None)).ToList(); + + // Assert + Assert.Single(result); + Assert.Equal("Reduced", result[0].Text); + reducerMock.Verify(r => r.ReduceAsync(It.Is>(x => x.SequenceEqual(originalMessages)), It.IsAny()), Times.Once); + } + + [Fact] + public async Task AddMessagesAsync_WithReducer_ButWrongTrigger_DoesNotInvokeReducerAsync() + { + // Arrange + var originalMessages = new List + { + new(ChatRole.User, "Hello") + }; + + var reducerMock = new Mock(); + + var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.BeforeMessagesRetrieval); + + // Act + await store.AddMessagesAsync(originalMessages, CancellationToken.None); + + // Assert + Assert.Single(store); + Assert.Equal("Hello", store[0].Text); + reducerMock.Verify(r => r.ReduceAsync(It.IsAny>(), It.IsAny()), Times.Never); + } + + [Fact] + public async Task GetMessagesAsync_WithReducer_ButWrongTrigger_DoesNotInvokeReducerAsync() + { + // Arrange + var originalMessages = new List + { + new(ChatRole.User, "Hello") + }; + + var reducerMock = new Mock(); + + var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.AfterMessageAdded); + store.Add(originalMessages[0]); + + // Act + var result = (await store.GetMessagesAsync(CancellationToken.None)).ToList(); + + // Assert + Assert.Single(result); + Assert.Equal("Hello", result[0].Text); + reducerMock.Verify(r => r.ReduceAsync(It.IsAny>(), It.IsAny()), Times.Never); + } }