From 1f6fa9cf9367c5938a8aefd259ce6c683e21cccc Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Wed, 3 Dec 2025 12:33:00 +0000 Subject: [PATCH 01/14] Refactor ChatMessageStore methods to be similar to AIContextProvider --- .../Program.cs | 20 ++- .../Program.cs | 44 ++++--- .../ChatMessageStore.cs | 116 ++++++++++++++++-- .../InMemoryChatMessageStore.cs | 29 +++-- .../WorkflowHostAgent.cs | 8 +- .../WorkflowMessageStore.cs | 19 ++- .../ChatClient/ChatClientAgent.cs | 32 +++-- .../AnthropicChatCompletionFixture.cs | 7 +- .../AIProjectClientFixture.cs | 7 +- .../ChatMessageStoreTests.cs | 8 +- .../InMemoryChatMessageStoreTests.cs | 64 +++++++--- .../ChatClient/ChatClientAgentTests.cs | 26 ++-- .../OpenAIChatCompletionFixture.cs | 7 +- .../OpenAIResponseFixture.cs | 7 +- 14 files changed, 302 insertions(+), 92 deletions(-) diff --git a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs index 8f1039251d..424232c9d2 100644 --- a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs +++ b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs @@ -44,11 +44,19 @@ public override async Task RunAsync(IEnumerable m throw new ArgumentException($"The provided thread is not of type {nameof(CustomAgentThread)}.", nameof(thread)); } + // Get existing messages from the store + var invokingContext = new ChatMessageStore.InvokingContext(messages); + var storeMessages = await typedThread.MessageStore.InvokingAsync(invokingContext, cancellationToken); + // Clone the input messages and turn them into response messages with upper case text. List responseMessages = CloneAndToUpperCase(messages, this.DisplayName).ToList(); // Notify the thread of the input and output messages. - await typedThread.MessageStore.AddMessagesAsync(messages.Concat(responseMessages), cancellationToken); + var invokedContext = new ChatMessageStore.InvokedContext(messages, storeMessages, null) + { + ResponseMessages = responseMessages + }; + await typedThread.MessageStore.InvokedAsync(invokedContext, cancellationToken); return new AgentRunResponse { @@ -68,11 +76,19 @@ public override async IAsyncEnumerable RunStreamingAsync throw new ArgumentException($"The provided thread is not of type {nameof(CustomAgentThread)}.", nameof(thread)); } + // Get existing messages from the store + var invokingContext = new ChatMessageStore.InvokingContext(messages); + var storeMessages = await typedThread.MessageStore.InvokingAsync(invokingContext, cancellationToken); + // Clone the input messages and turn them into response messages with upper case text. List responseMessages = CloneAndToUpperCase(messages, this.DisplayName).ToList(); // Notify the thread of the input and output messages. - await typedThread.MessageStore.AddMessagesAsync(messages.Concat(responseMessages), cancellationToken); + var invokedContext = new ChatMessageStore.InvokedContext(messages, storeMessages, null) + { + ResponseMessages = responseMessages + }; + await typedThread.MessageStore.InvokedAsync(invokedContext, cancellationToken); foreach (var message in responseMessages) { diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs index 8986734972..5722edf397 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs @@ -88,24 +88,7 @@ public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedSto public string? ThreadDbKey { get; private set; } - public override async Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default) - { - this.ThreadDbKey ??= Guid.NewGuid().ToString("N"); - - var collection = this._vectorStore.GetCollection("ChatHistory"); - await collection.EnsureCollectionExistsAsync(cancellationToken); - - await collection.UpsertAsync(messages.Select(x => new ChatHistoryItem() - { - Key = this.ThreadDbKey + x.MessageId, - Timestamp = DateTimeOffset.UtcNow, - ThreadId = this.ThreadDbKey, - SerializedMessage = JsonSerializer.Serialize(x), - MessageText = x.Text - }), cancellationToken); - } - - public override async Task> GetMessagesAsync(CancellationToken cancellationToken = default) + public override async ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) { var collection = this._vectorStore.GetCollection("ChatHistory"); await collection.EnsureCollectionExistsAsync(cancellationToken); @@ -123,6 +106,31 @@ public override async Task> GetMessagesAsync(Cancellati return messages; } + public override async ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + if (context.ResponseMessages is null) + { + return; + } + + this.ThreadDbKey ??= Guid.NewGuid().ToString("N"); + + var collection = this._vectorStore.GetCollection("ChatHistory"); + await collection.EnsureCollectionExistsAsync(cancellationToken); + + // Add both request and response messages to the store + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages); + + await collection.UpsertAsync(allNewMessages.Select(x => new ChatHistoryItem() + { + Key = this.ThreadDbKey + x.MessageId, + Timestamp = DateTimeOffset.UtcNow, + ThreadId = this.ThreadDbKey, + SerializedMessage = JsonSerializer.Serialize(x), + MessageText = x.Text + }), cancellationToken); + } + public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) => // We have to serialize the thread id, so that on deserialization we can retrieve the messages using the same thread id. JsonSerializer.SerializeToElement(this.ThreadDbKey); diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs index 9f89031464..ceb7782176 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs @@ -32,8 +32,9 @@ namespace Microsoft.Agents.AI; public abstract class ChatMessageStore { /// - /// Asynchronously retrieves all messages from the store that should be provided as context for the next agent invocation. + /// Called at the start of agent invocation to retrieve all messages from the store that should be provided as context for the next agent invocation. /// + /// Contains the request context including the caller provided messages that will be used by the agent for this invocation. /// The to monitor for cancellation requests. The default is . /// /// A task that represents the asynchronous operation. The task result contains a collection of @@ -59,20 +60,19 @@ public abstract class ChatMessageStore /// and context management. /// /// - public abstract Task> GetMessagesAsync(CancellationToken cancellationToken = default); + public abstract ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default); /// - /// Asynchronously adds new messages to the store. + /// Called at the end of the agent invocation to add new messages to the store. /// - /// The collection of chat messages to add to the store. + /// Contains the invocation context including request messages, response messages, and any exception that occurred. /// The to monitor for cancellation requests. The default is . /// A task that represents the asynchronous add operation. - /// is . /// /// /// Messages should be added in the order they were generated to maintain proper chronological sequence. /// The store is responsible for preserving message ordering and ensuring that subsequent calls to - /// return messages in the correct chronological order. + /// return messages in the correct chronological order. /// /// /// Implementations may perform additional processing during message addition, such as: @@ -83,8 +83,12 @@ public abstract class ChatMessageStore /// Updating indices or search capabilities /// /// + /// + /// This method is called regardless of whether the invocation succeeded or failed. + /// To check if the invocation was successful, inspect the property. + /// /// - public abstract Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default); + public abstract ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default); /// /// Serializes the current object's state to a using the specified serialization options. @@ -121,4 +125,102 @@ public abstract class ChatMessageStore /// public TService? GetService(object? serviceKey = null) => this.GetService(typeof(TService), serviceKey) is TService service ? service : default; + + /// + /// Contains the context information provided to . + /// + /// + /// This class provides context about the invocation before the messages are retrieved from the store, + /// including the new messages that will be used. Stores can use this information to determine what + /// messages should be retrieved for the invocation. + /// + public class InvokingContext + { + /// + /// Initializes a new instance of the class with the specified request messages. + /// + /// The new messages to be used by the agent for this invocation. + /// is . + public InvokingContext(IEnumerable requestMessages) + { + this.RequestMessages = requestMessages ?? throw new ArgumentNullException(nameof(requestMessages)); + } + + /// + /// Gets the caller provided messages that will be used by the agent for this invocation. + /// + /// + /// A collection of instances representing new messages that were provided by the caller. + /// + public IEnumerable RequestMessages { get; } + } + + /// + /// Contains the context information provided to . + /// + /// + /// This class provides context about a completed agent invocation, including both the + /// request messages that were used and the response messages that were generated. It also indicates + /// whether the invocation succeeded or failed. + /// + public class InvokedContext + { + /// + /// Initializes a new instance of the class with the specified request messages. + /// + /// The caller provided messages that were used by the agent for this invocation. + /// The messages retrieved from the for this invocation. + /// The messages provided by the for this invocation, if any. + /// is . + public InvokedContext(IEnumerable requestMessages, IEnumerable chatMessageStoreMessages, IEnumerable? aiContextProviderMessages) + { + this.RequestMessages = Throw.IfNull(requestMessages); + this.ChatMessageStoreMessages = chatMessageStoreMessages; + this.AIContextProviderMessages = aiContextProviderMessages; + } + + /// + /// Gets the caller provided messages that were used by the agent for this invocation. + /// + /// + /// A collection of instances representing new messages that were provided by the caller. + /// This does not include any supplied messages. + /// + public IEnumerable RequestMessages { get; } + + /// + /// Gets the messages retrieved from the for this invocation, if any. + /// + /// + /// A collection of instances that were retrieved from the , + /// and were used by the agent as part of the invocation. + /// + public IEnumerable ChatMessageStoreMessages { get; } + + /// + /// Gets the messages provided by the for this invocation, if any. + /// + /// + /// A collection of instances that were provided by the , + /// and were used by the agent as part of the invocation. + /// + public IEnumerable? AIContextProviderMessages { get; } + + /// + /// Gets the collection of response messages generated during this invocation if the invocation succeeded. + /// + /// + /// A collection of instances representing the response, + /// or if the invocation failed or did not produce response messages. + /// + public IEnumerable? ResponseMessages { get; set; } + + /// + /// Gets the that was thrown during the invocation, if the invocation failed. + /// + /// + /// The exception that caused the invocation to fail, or if the invocation succeeded. + /// + public Exception? InvokeException { get; set; } + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs index 79d303207c..01da79771a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs @@ -134,27 +134,36 @@ public ChatMessage this[int index] } /// - public override async Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default) + public override async ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) { - _ = Throw.IfNull(messages); + _ = Throw.IfNull(context); - this._messages.AddRange(messages); - - if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null) + if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null) { this._messages = (await this.ChatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList(); } + + return this._messages; } /// - public override async Task> GetMessagesAsync(CancellationToken cancellationToken = default) + public override async ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) { - if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null) + _ = Throw.IfNull(context); + + if (context.InvokeException is not null) { - this._messages = (await this.ChatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList(); + return; } - return this._messages; + // Add both request and response messages to the store + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); + this._messages.AddRange(allNewMessages); + + if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null) + { + this._messages = (await this.ChatReducer.ReduceAsync(this._messages, cancellationToken).ConfigureAwait(false)).ToList(); + } } /// @@ -221,7 +230,7 @@ public enum ChatReducerTriggerEvent { /// /// Trigger the reducer when a new message is added. - /// will only complete when reducer processing is done. + /// will only complete when reducer processing is done. /// AfterMessageAdded, diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs index 98dc5903bf..b1f973a461 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs @@ -66,7 +66,7 @@ private async ValueTask ValidateWorkflowAsync() public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null) => new WorkflowThread(this._workflow, serializedThread, this._executionEnvironment, this._checkpointManager, jsonSerializerOptions); - private async ValueTask UpdateThreadAsync(IEnumerable messages, AgentThread? thread = null, CancellationToken cancellationToken = default) + private ValueTask UpdateThreadAsync(IEnumerable messages, AgentThread? thread = null, CancellationToken cancellationToken = default) { thread ??= this.GetNewThread(); @@ -75,8 +75,10 @@ private async ValueTask UpdateThreadAsync(IEnumerable(workflowThread); } public override async diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs index 39c83bcadf..08d92157a6 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System.Collections.Generic; +using System.Linq; using System.Text.Json; using System.Threading; using System.Threading.Tasks; @@ -45,14 +46,22 @@ internal sealed class StoreState internal void AddMessages(params IEnumerable messages) => this._chatMessages.AddRange(messages); - public override Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default) + public override ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) + => new(this._chatMessages.AsReadOnly()); + + public override ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) { - this._chatMessages.AddRange(messages); + if (context.InvokeException is not null) + { + return default; + } - return Task.CompletedTask; - } + // Add both request and response messages to the store + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); + this._chatMessages.AddRange(allNewMessages); - public override Task> GetMessagesAsync(CancellationToken cancellationToken = default) => Task.FromResult>(this._chatMessages.AsReadOnly()); + return default; + } public IEnumerable GetFromBookmark() { diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index 30665fecf3..8a5c008dda 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -201,7 +201,7 @@ public override async IAsyncEnumerable RunStreamingAsync { var inputMessages = Throw.IfNull(messages) as IReadOnlyCollection ?? messages.ToList(); - (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List inputMessagesForChatClient, IList? aiContextProviderMessages) = + (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List inputMessagesForChatClient, IList? aiContextProviderMessages, IList? chatMessageStoreMessages) = await this.PrepareThreadAndMessagesAsync(thread, inputMessages, options, cancellationToken).ConfigureAwait(false); var chatClient = this.ChatClient; @@ -270,7 +270,7 @@ public override async IAsyncEnumerable RunStreamingAsync this.UpdateThreadWithTypeAndConversationId(safeThread, chatResponse.ConversationId); // To avoid inconsistent state we only notify the thread of the input messages if no error occurs after the initial request. - await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages.Concat(aiContextProviderMessages ?? []).Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false); + await NotifyMessageStoreOfNewMessagesAsync(safeThread, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false); // Notify the AIContextProvider of all new messages. await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, aiContextProviderMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false); @@ -377,7 +377,7 @@ private async Task RunCoreAsync ?? messages.ToList(); - (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List inputMessagesForChatClient, IList? aiContextProviderMessages) = + (ChatClientAgentThread safeThread, ChatOptions? chatOptions, List inputMessagesForChatClient, IList? aiContextProviderMessages, IList? chatMessageStoreMessages) = await this.PrepareThreadAndMessagesAsync(thread, inputMessages, options, cancellationToken).ConfigureAwait(false); var chatClient = this.ChatClient; @@ -413,7 +413,7 @@ private async Task RunCoreAsyncOptional parameters for agent invocation. /// The to monitor for cancellation requests. The default is . /// A tuple containing the thread, chat options, and thread messages. - private async Task<(ChatClientAgentThread AgentThread, ChatOptions? ChatOptions, List InputMessagesForChatClient, IList? AIContextProviderMessages)> PrepareThreadAndMessagesAsync( + private async Task<(ChatClientAgentThread AgentThread, ChatOptions? ChatOptions, List InputMessagesForChatClient, IList? AIContextProviderMessages, IList? ChatMessageStoreMessages)> PrepareThreadAndMessagesAsync( AgentThread? thread, IEnumerable inputMessages, AgentRunOptions? runOptions, @@ -623,6 +623,7 @@ await thread.AIContextProvider.InvokedAsync(new(inputMessages, aiContextProvider } List inputMessagesForChatClient = []; IList? aiContextProviderMessages = null; + IList? chatMessageStoreMessages = null; // Populate the thread messages only if we are not continuing an existing response as it's not allowed if (chatOptions?.ContinuationToken is null) @@ -630,7 +631,10 @@ await thread.AIContextProvider.InvokedAsync(new(inputMessages, aiContextProvider // Add any existing messages from the thread to the messages to be sent to the chat client. if (typedThread.MessageStore is not null) { - inputMessagesForChatClient.AddRange(await typedThread.MessageStore.GetMessagesAsync(cancellationToken).ConfigureAwait(false)); + var invokingContext = new ChatMessageStore.InvokingContext(inputMessages); + var storeMessages = await typedThread.MessageStore.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false); + inputMessagesForChatClient.AddRange(storeMessages); + chatMessageStoreMessages = storeMessages as IList ?? storeMessages.ToList(); } // If we have an AIContextProvider, we should get context from it, and update our @@ -690,7 +694,7 @@ await thread.AIContextProvider.InvokedAsync(new(inputMessages, aiContextProvider chatOptions.ConversationId = typedThread.ConversationId; } - return (typedThread, chatOptions, inputMessagesForChatClient, aiContextProviderMessages); + return (typedThread, chatOptions, inputMessagesForChatClient, aiContextProviderMessages, chatMessageStoreMessages); } private void UpdateThreadWithTypeAndConversationId(ChatClientAgentThread thread, string? responseConversationId) @@ -717,7 +721,13 @@ private void UpdateThreadWithTypeAndConversationId(ChatClientAgentThread thread, } } - private static Task NotifyMessageStoreOfNewMessagesAsync(ChatClientAgentThread thread, IEnumerable newMessages, CancellationToken cancellationToken) + private static Task NotifyMessageStoreOfNewMessagesAsync( + ChatClientAgentThread thread, + IEnumerable requestMessages, + IEnumerable? chatMessageStoreMessages, + IEnumerable? aiContextProviderMessages, + IEnumerable responseMessages, + CancellationToken cancellationToken) { var messageStore = thread.MessageStore; @@ -725,7 +735,11 @@ private static Task NotifyMessageStoreOfNewMessagesAsync(ChatClientAgentThread t // If we don't have one, it means that the chat history is service managed and the underlying service is responsible for storing messages. if (messageStore is not null) { - return messageStore.AddMessagesAsync(newMessages, cancellationToken); + var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages!, aiContextProviderMessages) + { + ResponseMessages = responseMessages + }; + return messageStore.InvokedAsync(invokedContext, cancellationToken).AsTask(); } return Task.CompletedTask; diff --git a/dotnet/tests/AnthropicChatCompletion.IntegrationTests/AnthropicChatCompletionFixture.cs b/dotnet/tests/AnthropicChatCompletion.IntegrationTests/AnthropicChatCompletionFixture.cs index 76ca18d3de..fb73ff021e 100644 --- a/dotnet/tests/AnthropicChatCompletion.IntegrationTests/AnthropicChatCompletionFixture.cs +++ b/dotnet/tests/AnthropicChatCompletion.IntegrationTests/AnthropicChatCompletionFixture.cs @@ -39,7 +39,12 @@ public async Task> GetChatHistoryAsync(AgentThread thread) { var typedThread = (ChatClientAgentThread)thread; - return typedThread.MessageStore is null ? [] : (await typedThread.MessageStore.GetMessagesAsync()).ToList(); + if (typedThread.MessageStore is null) + { + return []; + } + + return (await typedThread.MessageStore.InvokingAsync(new([]))).ToList(); } public Task CreateChatClientAgentAsync( diff --git a/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs b/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs index e982c8081f..cb3d4891e1 100644 --- a/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs +++ b/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs @@ -48,7 +48,12 @@ public async Task> GetChatHistoryAsync(AgentThread thread) return await this.GetChatHistoryFromResponsesChainAsync(chatClientThread.ConversationId); } - return chatClientThread.MessageStore is null ? [] : (await chatClientThread.MessageStore.GetMessagesAsync()).ToList(); + if (chatClientThread.MessageStore is null) + { + return []; + } + + return (await chatClientThread.MessageStore.InvokingAsync(new([]))).ToList(); } private async Task> GetChatHistoryFromResponsesChainAsync(string conversationId) diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreTests.cs index 4100b20f5a..883941458c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreTests.cs @@ -78,11 +78,11 @@ public void GetService_Generic_ReturnsNullForUnrelatedType() private sealed class TestChatMessageStore : ChatMessageStore { - public override Task> GetMessagesAsync(CancellationToken cancellationToken = default) - => Task.FromResult>([]); + public override ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) + => new(Array.Empty()); - public override Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default) - => Task.CompletedTask; + public override ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) + => default; public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) => default; diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs index 824fb62f6d..4e54d31e8b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs @@ -47,16 +47,23 @@ public void Constructor_Arguments_SetOnPropertiesCorrectly() } [Fact] - public async Task AddMessagesAsyncAddsMessagesAndReturnsNullThreadIdAsync() + public async Task InvokedAsyncAddsMessagesAsync() { var store = new InMemoryChatMessageStore(); - var messages = new List + var requestMessages = new List + { + new(ChatRole.User, "Hello") + }; + var responseMessages = new List { - new(ChatRole.User, "Hello"), new(ChatRole.Assistant, "Hi there!") }; - await store.AddMessagesAsync(messages, CancellationToken.None); + var context = new ChatMessageStore.InvokedContext(requestMessages, [], null) + { + ResponseMessages = responseMessages + }; + await store.InvokedAsync(context, CancellationToken.None); Assert.Equal(2, store.Count); Assert.Equal("Hello", store[0].Text); @@ -64,17 +71,21 @@ public async Task AddMessagesAsyncAddsMessagesAndReturnsNullThreadIdAsync() } [Fact] - public async Task AddMessagesAsyncWithEmptyDoesNotFailAsync() + public async Task InvokedAsyncWithEmptyDoesNotFailAsync() { var store = new InMemoryChatMessageStore(); - await store.AddMessagesAsync([], CancellationToken.None); + var context = new ChatMessageStore.InvokedContext([], [], null) + { + ResponseMessages = null + }; + await store.InvokedAsync(context, CancellationToken.None); Assert.Empty(store); } [Fact] - public async Task GetMessagesAsyncReturnsAllMessagesAsync() + public async Task InvokingAsyncReturnsAllMessagesAsync() { var store = new InMemoryChatMessageStore { @@ -82,7 +93,8 @@ public async Task GetMessagesAsyncReturnsAllMessagesAsync() new ChatMessage(ChatRole.Assistant, "Test2") }; - var result = (await store.GetMessagesAsync(CancellationToken.None)).ToList(); + var context = new ChatMessageStore.InvokingContext([]); + var result = (await store.InvokingAsync(context, CancellationToken.None)).ToList(); Assert.Equal(2, result.Count); Assert.Contains(result, m => m.Text == "Test1"); @@ -157,24 +169,28 @@ public async Task SerializeAndDeserializeWorksWithExperimentalContentTypesAsync( } [Fact] - public async Task AddMessagesAsyncWithEmptyMessagesDoesNotChangeStoreAsync() + public async Task InvokedAsyncWithEmptyMessagesDoesNotChangeStoreAsync() { var store = new InMemoryChatMessageStore(); var messages = new List(); - await store.AddMessagesAsync(messages, CancellationToken.None); + var context = new ChatMessageStore.InvokedContext(messages, [], null) + { + ResponseMessages = null + }; + await store.InvokedAsync(context, CancellationToken.None); Assert.Empty(store); } [Fact] - public async Task AddMessagesAsync_WithNullMessages_ThrowsArgumentNullExceptionAsync() + public async Task InvokedAsync_WithNullContext_ThrowsArgumentNullExceptionAsync() { // Arrange var store = new InMemoryChatMessageStore(); // Act & Assert - await Assert.ThrowsAsync(() => store.AddMessagesAsync(null!, CancellationToken.None)); + await Assert.ThrowsAsync(() => store.InvokedAsync(null!, CancellationToken.None).AsTask()); } [Fact] @@ -498,7 +514,11 @@ public async Task AddMessagesAsync_WithReducer_AfterMessageAdded_InvokesReducerA var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.AfterMessageAdded); // Act - await store.AddMessagesAsync(originalMessages, CancellationToken.None); + var context = new ChatMessageStore.InvokedContext(originalMessages, [], null) + { + ResponseMessages = Array.Empty() + }; + await store.InvokedAsync(context, CancellationToken.None); // Assert Assert.Single(store); @@ -526,10 +546,15 @@ public async Task GetMessagesAsync_WithReducer_BeforeMessagesRetrieval_InvokesRe .ReturnsAsync(reducedMessages); var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.BeforeMessagesRetrieval); - await store.AddMessagesAsync(originalMessages, CancellationToken.None); + // Add messages directly to the store for this test + foreach (var msg in originalMessages) + { + store.Add(msg); + } // Act - var result = (await store.GetMessagesAsync(CancellationToken.None)).ToList(); + var invokingContext = new ChatMessageStore.InvokingContext(Array.Empty()); + var result = (await store.InvokingAsync(invokingContext, CancellationToken.None)).ToList(); // Assert Assert.Single(result); @@ -551,7 +576,11 @@ public async Task AddMessagesAsync_WithReducer_ButWrongTrigger_DoesNotInvokeRedu var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.BeforeMessagesRetrieval); // Act - await store.AddMessagesAsync(originalMessages, CancellationToken.None); + var context = new ChatMessageStore.InvokedContext(originalMessages, [], null) + { + ResponseMessages = Array.Empty() + }; + await store.InvokedAsync(context, CancellationToken.None); // Assert Assert.Single(store); @@ -576,7 +605,8 @@ public async Task GetMessagesAsync_WithReducer_ButWrongTrigger_DoesNotInvokeRedu }; // Act - var result = (await store.GetMessagesAsync(CancellationToken.None)).ToList(); + var invokingContext = new ChatMessageStore.InvokingContext(Array.Empty()); + var result = (await store.InvokingAsync(invokingContext, CancellationToken.None)).ToList(); // Assert Assert.Single(result); diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs index 920f9f82ee..2ce0c133ae 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs @@ -518,7 +518,7 @@ public async Task RunAsyncUsesChatMessageStoreFactoryWhenProvidedAndNoConversati // Assert Assert.IsType(thread!.MessageStore, exactMatch: false); - mockChatMessageStore.Verify(s => s.AddMessagesAsync(It.Is>(x => x.Count() == 2), It.IsAny()), Times.Once); + mockChatMessageStore.Verify(s => s.InvokedAsync(It.Is(x => x.RequestMessages.Count() == 1 && x.ResponseMessages!.Count() == 1), It.IsAny()), Times.Once); mockFactory.Verify(f => f(It.IsAny()), Times.Once); } @@ -610,12 +610,12 @@ public async Task RunAsyncInvokesAIContextProviderAndUsesResultAsync() Assert.Contains(capturedTools, t => t.Name == "base function"); Assert.Contains(capturedTools, t => t.Name == "context provider function"); - // Verify that the thread was updated with the input, ai context and response messages + // Verify that the thread was updated with the input and response messages + // Note: AIContextProvider messages are not persisted to the message store, they're ephemeral context var messageStore = Assert.IsType(thread!.MessageStore); - Assert.Equal(3, messageStore.Count); + Assert.Equal(2, messageStore.Count); Assert.Equal("user message", messageStore[0].Text); - Assert.Equal("context provider message", messageStore[1].Text); - Assert.Equal("response", messageStore[2].Text); + Assert.Equal("response", messageStore[1].Text); mockProvider.Verify(p => p.InvokingAsync(It.IsAny(), It.IsAny()), Times.Once); mockProvider.Verify(p => p.InvokedAsync(It.Is(x => @@ -2068,12 +2068,12 @@ public async Task RunStreamingAsyncInvokesAIContextProviderAndUsesResultAsync() Assert.Contains(capturedTools, t => t.Name == "base function"); Assert.Contains(capturedTools, t => t.Name == "context provider function"); - // Verify that the thread was updated with the input, ai context and response messages + // Verify that the thread was updated with the input and response messages + // Note: AIContextProvider messages are not persisted to the message store, they're ephemeral context var messageStore = Assert.IsType(thread!.MessageStore); - Assert.Equal(3, messageStore.Count); + Assert.Equal(2, messageStore.Count); Assert.Equal("user message", messageStore[0].Text); - Assert.Equal("context provider message", messageStore[1].Text); - Assert.Equal("response", messageStore[2].Text); + Assert.Equal("response", messageStore[1].Text); mockProvider.Verify(p => p.InvokingAsync(It.IsAny(), It.IsAny()), Times.Once); mockProvider.Verify(p => p.InvokedAsync(It.Is(x => @@ -2454,7 +2454,7 @@ public async Task RunAsyncSkipsThreadMessagePopulationWithContinuationTokenAsync // Create a mock message store that would normally provide messages var mockMessageStore = new Mock(); mockMessageStore - .Setup(ms => ms.GetMessagesAsync(It.IsAny())) + .Setup(ms => ms.InvokingAsync(It.IsAny(), It.IsAny())) .ReturnsAsync([new(ChatRole.User, "Message from message store")]); // Create a mock AI context provider that would normally provide context @@ -2498,7 +2498,7 @@ public async Task RunAsyncSkipsThreadMessagePopulationWithContinuationTokenAsync // Verify that message store was never called due to continuation token mockMessageStore.Verify( - ms => ms.GetMessagesAsync(It.IsAny()), + ms => ms.InvokingAsync(It.IsAny(), It.IsAny()), Times.Never); // Verify that AI context provider was never called due to continuation token @@ -2516,7 +2516,7 @@ public async Task RunStreamingAsyncSkipsThreadMessagePopulationWithContinuationT // Create a mock message store that would normally provide messages var mockMessageStore = new Mock(); mockMessageStore - .Setup(ms => ms.GetMessagesAsync(It.IsAny())) + .Setup(ms => ms.InvokingAsync(It.IsAny(), It.IsAny())) .ReturnsAsync([new(ChatRole.User, "Message from message store")]); // Create a mock AI context provider that would normally provide context @@ -2560,7 +2560,7 @@ public async Task RunStreamingAsyncSkipsThreadMessagePopulationWithContinuationT // Verify that message store was never called due to continuation token mockMessageStore.Verify( - ms => ms.GetMessagesAsync(It.IsAny()), + ms => ms.InvokingAsync(It.IsAny(), It.IsAny()), Times.Never); // Verify that AI context provider was never called due to continuation token diff --git a/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionFixture.cs b/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionFixture.cs index f98540d8cc..0bc8895c2d 100644 --- a/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionFixture.cs +++ b/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionFixture.cs @@ -32,7 +32,12 @@ public async Task> GetChatHistoryAsync(AgentThread thread) { var typedThread = (ChatClientAgentThread)thread; - return typedThread.MessageStore is null ? [] : (await typedThread.MessageStore.GetMessagesAsync()).ToList(); + if (typedThread.MessageStore is null) + { + return []; + } + + return (await typedThread.MessageStore.InvokingAsync(new([]))).ToList(); } public Task CreateChatClientAgentAsync( diff --git a/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseFixture.cs b/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseFixture.cs index fbb087a153..f3c428ca9b 100644 --- a/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseFixture.cs +++ b/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseFixture.cs @@ -50,7 +50,12 @@ public async Task> GetChatHistoryAsync(AgentThread thread) return [.. previousMessages, responseMessage]; } - return typedThread.MessageStore is null ? [] : (await typedThread.MessageStore.GetMessagesAsync()).ToList(); + if (typedThread.MessageStore is null) + { + return []; + } + + return (await typedThread.MessageStore.InvokingAsync(new([]))).ToList(); } private static ChatMessage ConvertToChatMessage(ResponseItem item) From fdee4912340bade790d9262375c4a5f127fded29 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 4 Dec 2025 09:15:31 +0000 Subject: [PATCH 02/14] Fix file encoding --- .../CosmosChatMessageStoreTests.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 1d39a7cb9e..0166c8e703 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; From 137a514de9d94e61a6177300a0258db94e04b1ca Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 4 Dec 2025 12:05:20 +0000 Subject: [PATCH 03/14] Ensure that AIContextProvider messages area also persisted. --- .../Program.cs | 6 ++++-- .../InMemoryChatMessageStore.cs | 4 ++-- .../CosmosChatMessageStore.cs | 2 +- .../ChatClient/ChatClientAgentTests.cs | 18 +++++++++--------- 4 files changed, 16 insertions(+), 14 deletions(-) diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs index 64e3538f60..e12172d397 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs @@ -108,7 +108,8 @@ public override async ValueTask> InvokingAsync(Invoking public override async ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) { - if (context.ResponseMessages is null) + // Don't store messages if the request failed. + if (context.InvokeException is not null) { return; } @@ -119,7 +120,8 @@ public override async ValueTask InvokedAsync(InvokedContext context, Cancellatio await collection.EnsureCollectionExistsAsync(cancellationToken); // Add both request and response messages to the store - var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages); + // Optionally messages produced by the AIContextProvider can also be persisted (not shown). + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); await collection.UpsertAsync(allNewMessages.Select(x => new ChatHistoryItem() { diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs index 01da79771a..cca2adc5e4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs @@ -156,8 +156,8 @@ public override async ValueTask InvokedAsync(InvokedContext context, Cancellatio return; } - // Add both request and response messages to the store - var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); + // Add ai context provider, request and response messages to the store + var allNewMessages = (context.AIContextProviderMessages ?? []).Concat(context.RequestMessages).Concat(context.ResponseMessages ?? []); this._messages.AddRange(allNewMessages); if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 1d092bdc62..d9e0a1e17d 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -364,7 +364,7 @@ public override async ValueTask InvokedAsync(InvokedContext context, Cancellatio } #pragma warning restore CA1513 - var messageList = (context.AIContextProviderMessages ?? []).Concat(context.RequestMessages).ToList(); + var messageList = (context.AIContextProviderMessages ?? []).Concat(context.RequestMessages).Concat(context.ResponseMessages ?? []).ToList(); if (messageList.Count == 0) { return; diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs index 8508d92c9d..c25bf0732a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs @@ -610,12 +610,12 @@ public async Task RunAsyncInvokesAIContextProviderAndUsesResultAsync() Assert.Contains(capturedTools, t => t.Name == "base function"); Assert.Contains(capturedTools, t => t.Name == "context provider function"); - // Verify that the thread was updated with the input and response messages - // Note: AIContextProvider messages are not persisted to the message store, they're ephemeral context + // Verify that the thread was updated with the ai context provider, input and response messages var messageStore = Assert.IsType(thread!.MessageStore); - Assert.Equal(2, messageStore.Count); - Assert.Equal("user message", messageStore[0].Text); - Assert.Equal("response", messageStore[1].Text); + Assert.Equal(3, messageStore.Count); + Assert.Equal("context provider message", messageStore[0].Text); + Assert.Equal("user message", messageStore[1].Text); + Assert.Equal("response", messageStore[2].Text); mockProvider.Verify(p => p.InvokingAsync(It.IsAny(), It.IsAny()), Times.Once); mockProvider.Verify(p => p.InvokedAsync(It.Is(x => @@ -2068,11 +2068,11 @@ public async Task RunStreamingAsyncInvokesAIContextProviderAndUsesResultAsync() Assert.Contains(capturedTools, t => t.Name == "context provider function"); // Verify that the thread was updated with the input and response messages - // Note: AIContextProvider messages are not persisted to the message store, they're ephemeral context var messageStore = Assert.IsType(thread!.MessageStore); - Assert.Equal(2, messageStore.Count); - Assert.Equal("user message", messageStore[0].Text); - Assert.Equal("response", messageStore[1].Text); + Assert.Equal(3, messageStore.Count); + Assert.Equal("context provider message", messageStore[0].Text); + Assert.Equal("user message", messageStore[1].Text); + Assert.Equal("response", messageStore[2].Text); mockProvider.Verify(p => p.InvokingAsync(It.IsAny(), It.IsAny()), Times.Once); mockProvider.Verify(p => p.InvokedAsync(It.Is(x => From ba63612649cd2979007efc378bba0e1ef0d3a04b Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 4 Dec 2025 12:45:08 +0000 Subject: [PATCH 04/14] Update formatting and seal context classes --- .../Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs | 4 ++-- .../Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs | 9 ++++++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs index ceb7782176..ece9ff8584 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs @@ -134,7 +134,7 @@ public abstract class ChatMessageStore /// including the new messages that will be used. Stores can use this information to determine what /// messages should be retrieved for the invocation. /// - public class InvokingContext + public sealed class InvokingContext { /// /// Initializes a new instance of the class with the specified request messages. @@ -163,7 +163,7 @@ public InvokingContext(IEnumerable requestMessages) /// request messages that were used and the response messages that were generated. It also indicates /// whether the invocation succeeded or failed. /// - public class InvokedContext + public sealed class InvokedContext { /// /// Initializes a new instance of the class with the specified request messages. diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index e2287f3787..b5faddbbb0 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -603,7 +603,14 @@ await thread.AIContextProvider.InvokedAsync(new(inputMessages, aiContextProvider /// Optional parameters for agent invocation. /// The to monitor for cancellation requests. The default is . /// A tuple containing the thread, chat options, and thread messages. - private async Task<(ChatClientAgentThread AgentThread, ChatOptions? ChatOptions, List InputMessagesForChatClient, IList? AIContextProviderMessages, IList? ChatMessageStoreMessages)> PrepareThreadAndMessagesAsync( + private async Task< + ( + ChatClientAgentThread AgentThread, + ChatOptions? ChatOptions, + List InputMessagesForChatClient, + IList? AIContextProviderMessages, + IList? ChatMessageStoreMessages + )> PrepareThreadAndMessagesAsync( AgentThread? thread, IEnumerable inputMessages, AgentRunOptions? runOptions, From cf6956dcf99e07bd223b10cfd1e0f0551caf21f8 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 4 Dec 2025 12:48:00 +0000 Subject: [PATCH 05/14] Improve formatting --- dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index b5faddbbb0..f734e11de0 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -603,8 +603,8 @@ await thread.AIContextProvider.InvokedAsync(new(inputMessages, aiContextProvider /// Optional parameters for agent invocation. /// The to monitor for cancellation requests. The default is . /// A tuple containing the thread, chat options, and thread messages. - private async Task< - ( + private async Task + <( ChatClientAgentThread AgentThread, ChatOptions? ChatOptions, List InputMessagesForChatClient, From c894d619dc64ad9ed0700dd92b02b1e08f3333bf Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 5 Dec 2025 13:26:56 +0000 Subject: [PATCH 06/14] Remove optional messages from constructor and add unit test --- .../Program.cs | 4 +- .../ChatMessageStore.cs | 8 ++-- .../ChatClient/ChatClientAgent.cs | 3 +- .../InMemoryChatMessageStoreTests.cs | 42 +++++++++---------- .../CosmosChatMessageStoreTests.cs | 28 ++++++------- 5 files changed, 42 insertions(+), 43 deletions(-) diff --git a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs index 424232c9d2..0c65553807 100644 --- a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs +++ b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs @@ -52,7 +52,7 @@ public override async Task RunAsync(IEnumerable m List responseMessages = CloneAndToUpperCase(messages, this.DisplayName).ToList(); // Notify the thread of the input and output messages. - var invokedContext = new ChatMessageStore.InvokedContext(messages, storeMessages, null) + var invokedContext = new ChatMessageStore.InvokedContext(messages, storeMessages) { ResponseMessages = responseMessages }; @@ -84,7 +84,7 @@ public override async IAsyncEnumerable RunStreamingAsync List responseMessages = CloneAndToUpperCase(messages, this.DisplayName).ToList(); // Notify the thread of the input and output messages. - var invokedContext = new ChatMessageStore.InvokedContext(messages, storeMessages, null) + var invokedContext = new ChatMessageStore.InvokedContext(messages, storeMessages) { ResponseMessages = responseMessages }; diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs index ece9ff8584..d28cd191b7 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStore.cs @@ -170,13 +170,11 @@ public sealed class InvokedContext /// /// The caller provided messages that were used by the agent for this invocation. /// The messages retrieved from the for this invocation. - /// The messages provided by the for this invocation, if any. /// is . - public InvokedContext(IEnumerable requestMessages, IEnumerable chatMessageStoreMessages, IEnumerable? aiContextProviderMessages) + public InvokedContext(IEnumerable requestMessages, IEnumerable chatMessageStoreMessages) { this.RequestMessages = Throw.IfNull(requestMessages); this.ChatMessageStoreMessages = chatMessageStoreMessages; - this.AIContextProviderMessages = aiContextProviderMessages; } /// @@ -198,13 +196,13 @@ public InvokedContext(IEnumerable requestMessages, IEnumerable ChatMessageStoreMessages { get; } /// - /// Gets the messages provided by the for this invocation, if any. + /// Gets or sets the messages provided by the for this invocation, if any. /// /// /// A collection of instances that were provided by the , /// and were used by the agent as part of the invocation. /// - public IEnumerable? AIContextProviderMessages { get; } + public IEnumerable? AIContextProviderMessages { get; set; } /// /// Gets the collection of response messages generated during this invocation if the invocation succeeded. diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index f734e11de0..fc97e4c084 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -750,8 +750,9 @@ private static Task NotifyMessageStoreOfNewMessagesAsync( // If we don't have one, it means that the chat history is service managed and the underlying service is responsible for storing messages. if (messageStore is not null) { - var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages!, aiContextProviderMessages) + var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages!) { + AIContextProviderMessages = aiContextProviderMessages, ResponseMessages = responseMessages }; return messageStore.InvokedAsync(invokedContext, cancellationToken).AsTask(); diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs index 4e54d31e8b..3b7f1742dd 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs @@ -49,7 +49,6 @@ public void Constructor_Arguments_SetOnPropertiesCorrectly() [Fact] public async Task InvokedAsyncAddsMessagesAsync() { - var store = new InMemoryChatMessageStore(); var requestMessages = new List { new(ChatRole.User, "Hello") @@ -58,16 +57,29 @@ public async Task InvokedAsyncAddsMessagesAsync() { new(ChatRole.Assistant, "Hi there!") }; + var messageStoreMessages = new List() + { + new(ChatRole.System, "original instructions") + }; + var aiContextProviderMessages = new List() + { + new(ChatRole.System, "additional context") + }; - var context = new ChatMessageStore.InvokedContext(requestMessages, [], null) + var store = new InMemoryChatMessageStore(); + store.Add(messageStoreMessages[0]); + var context = new ChatMessageStore.InvokedContext(requestMessages, messageStoreMessages) { + AIContextProviderMessages = aiContextProviderMessages, ResponseMessages = responseMessages }; await store.InvokedAsync(context, CancellationToken.None); - Assert.Equal(2, store.Count); - Assert.Equal("Hello", store[0].Text); - Assert.Equal("Hi there!", store[1].Text); + Assert.Equal(4, store.Count); + Assert.Equal("original instructions", store[0].Text); + Assert.Equal("additional context", store[1].Text); + Assert.Equal("Hello", store[2].Text); + Assert.Equal("Hi there!", store[3].Text); } [Fact] @@ -75,10 +87,7 @@ public async Task InvokedAsyncWithEmptyDoesNotFailAsync() { var store = new InMemoryChatMessageStore(); - var context = new ChatMessageStore.InvokedContext([], [], null) - { - ResponseMessages = null - }; + var context = new ChatMessageStore.InvokedContext([], []); await store.InvokedAsync(context, CancellationToken.None); Assert.Empty(store); @@ -174,10 +183,7 @@ public async Task InvokedAsyncWithEmptyMessagesDoesNotChangeStoreAsync() var store = new InMemoryChatMessageStore(); var messages = new List(); - var context = new ChatMessageStore.InvokedContext(messages, [], null) - { - ResponseMessages = null - }; + var context = new ChatMessageStore.InvokedContext(messages, []); await store.InvokedAsync(context, CancellationToken.None); Assert.Empty(store); @@ -514,10 +520,7 @@ public async Task AddMessagesAsync_WithReducer_AfterMessageAdded_InvokesReducerA var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.AfterMessageAdded); // Act - var context = new ChatMessageStore.InvokedContext(originalMessages, [], null) - { - ResponseMessages = Array.Empty() - }; + var context = new ChatMessageStore.InvokedContext(originalMessages, []); await store.InvokedAsync(context, CancellationToken.None); // Assert @@ -576,10 +579,7 @@ public async Task AddMessagesAsync_WithReducer_ButWrongTrigger_DoesNotInvokeRedu var store = new InMemoryChatMessageStore(reducerMock.Object, InMemoryChatMessageStore.ChatReducerTriggerEvent.BeforeMessagesRetrieval); // Act - var context = new ChatMessageStore.InvokedContext(originalMessages, [], null) - { - ResponseMessages = Array.Empty() - }; + var context = new ChatMessageStore.InvokedContext(originalMessages, []); await store.InvokedAsync(context, CancellationToken.None); // Assert diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 0166c8e703..b40edeb7a3 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -213,7 +213,7 @@ public async Task InvokedAsync_WithSingleMessage_ShouldAddMessageAsync() using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); var message = new ChatMessage(ChatRole.User, "Hello, world!"); - var context = new ChatMessageStore.InvokedContext([message], [], null) + var context = new ChatMessageStore.InvokedContext([message], []) { ResponseMessages = [] }; @@ -284,7 +284,7 @@ public async Task InvokedAsync_WithMultipleMessages_ShouldAddAllMessagesAsync() new ChatMessage(ChatRole.User, "Third message") }; - var context = new ChatMessageStore.InvokedContext(messages, [], null) + var context = new ChatMessageStore.InvokedContext(messages, []) { ResponseMessages = [] }; @@ -334,8 +334,8 @@ public async Task InvokingAsync_WithConversationIsolation_ShouldOnlyReturnMessag using var store1 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversation1); using var store2 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversation2); - var context1 = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Message for conversation 1")], [], null); - var context2 = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Message for conversation 2")], [], null); + var context1 = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Message for conversation 1")], []); + var context2 = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Message for conversation 2")], []); await store1.InvokedAsync(context1); await store2.InvokedAsync(context2); @@ -379,7 +379,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() }; // Act 1: Add messages - var invokedContext = new ChatMessageStore.InvokedContext(messages, [], null); + var invokedContext = new ChatMessageStore.InvokedContext(messages, []); await originalStore.InvokedAsync(invokedContext); // Act 2: Verify messages were added @@ -533,7 +533,7 @@ public async Task InvokedAsync_WithHierarchicalPartitioning_ShouldAddMessageWith using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); var message = new ChatMessage(ChatRole.User, "Hello from hierarchical partitioning!"); - var context = new ChatMessageStore.InvokedContext([message], [], null); + var context = new ChatMessageStore.InvokedContext([message], []); // Act await store.InvokedAsync(context); @@ -590,7 +590,7 @@ public async Task InvokedAsync_WithHierarchicalMultipleMessages_ShouldAddAllMess new ChatMessage(ChatRole.User, "Third hierarchical message") }; - var context = new ChatMessageStore.InvokedContext(messages, [], null); + var context = new ChatMessageStore.InvokedContext(messages, []); // Act await store.InvokedAsync(context); @@ -625,8 +625,8 @@ public async Task InvokingAsync_WithHierarchicalPartitionIsolation_ShouldIsolate using var store2 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); // Add messages to both stores - var context1 = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Message from user 1")], [], null); - var context2 = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Message from user 2")], [], null); + var context1 = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Message from user 1")], []); + var context2 = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Message from user 2")], []); await store1.InvokedAsync(context1); await store2.InvokedAsync(context2); @@ -663,7 +663,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser using var originalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); - var context = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Test serialization message")], [], null); + var context = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Test serialization message")], []); await originalStore.InvokedAsync(context); // Act - Serialize the store state @@ -705,8 +705,8 @@ public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); // Add messages to both - var simpleContext = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Simple partitioning message")], [], null); - var hierarchicalContext = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Hierarchical partitioning message")], [], null); + var simpleContext = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Simple partitioning message")], []); + var hierarchicalContext = new ChatMessageStore.InvokedContext([new ChatMessage(ChatRole.User, "Hierarchical partitioning message")], []); await simpleStore.InvokedAsync(simpleContext); await hierarchicalStore.InvokedAsync(hierarchicalContext); @@ -748,7 +748,7 @@ public async Task MaxMessagesToRetrieve_ShouldLimitAndReturnMostRecentAsync() await Task.Delay(10); // Small delay to ensure different timestamps } - var context = new ChatMessageStore.InvokedContext(messages, [], null); + var context = new ChatMessageStore.InvokedContext(messages, []); await store.InvokedAsync(context); // Wait for eventual consistency @@ -786,7 +786,7 @@ public async Task MaxMessagesToRetrieve_Null_ShouldReturnAllMessagesAsync() messages.Add(new ChatMessage(ChatRole.User, $"Message {i}")); } - var context = new ChatMessageStore.InvokedContext(messages, [], null); + var context = new ChatMessageStore.InvokedContext(messages, []); await store.InvokedAsync(context); // Wait for eventual consistency From 3bd5fb963aaad9e34e8e19d2f6a9c94a76c4a36f Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 5 Dec 2025 14:41:14 +0000 Subject: [PATCH 07/14] Add ChatMessageStore filtering via a decorator --- .../Program.cs | 7 +- .../ChatMessageStoreExtensions.cs | 50 ++++ .../ChatMessageStoreMessageFilter.cs | 67 +++++ .../InMemoryChatMessageStore.cs | 2 +- .../ChatMessageStoreMessageFilterTests.cs | 230 ++++++++++++++++++ .../InMemoryChatMessageStoreTests.cs | 4 +- 6 files changed, 356 insertions(+), 4 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreMessageFilter.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreMessageFilterTests.cs diff --git a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs index 5470b1ed48..cb0b6bdea3 100644 --- a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs @@ -62,7 +62,12 @@ .CreateAIAgent(new ChatClientAgentOptions { ChatOptions = new() { Instructions = "You are a helpful support specialist for Contoso Outdoors. Answer questions using the provided context and cite the source document when available." }, - AIContextProviderFactory = ctx => new TextSearchProvider(SearchAdapter, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions) + AIContextProviderFactory = ctx => new TextSearchProvider(SearchAdapter, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions), + // Since we are using ChatCompletion which stores chat history locally, we can also add a message removal policy + // that removes messages produced by the TextSearchProvider before they are added to the chat history, so that + // we don't bloat chat history with all the search result messages. + ChatMessageStoreFactory = ctx => new InMemoryChatMessageStore(ctx.SerializedState, ctx.JsonSerializerOptions) + .WithAIContextProviderMessageRemoval(), }); AgentThread thread = agent.GetNewThread(); diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs new file mode 100644 index 0000000000..589325748e --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs @@ -0,0 +1,50 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Contains extension methods for the class. +/// +public static class ChatMessageStoreExtensions +{ + /// + /// Adds message filering to an existing store, so that messages passed to the store and messages produced by the store + /// can be filtered, updated or replaced. + /// + /// The store to add the message filter to. + /// An optional filter function to apply to messages before they are invoked. If null, no filter is applied at this + /// stage. + /// An optional filter function to apply to the invocation context after messages have been invoked. If null, no + /// filter is applied at this stage. + /// The with filtering applied. + public static ChatMessageStore WithMessageFilters( + this ChatMessageStore store, + Func, IEnumerable>? invokingMessagesFilter = null, + Func? invokedMessagesFilter = null) + { + return new ChatMessageStoreMessageFilter( + innerChatMessageStore: store, + invokingMessagesFilter: invokingMessagesFilter, + invokedMessagesFilter: invokedMessagesFilter); + } + + /// + /// Decorates the provided chat message store so that it does not store messages produced by any . + /// + /// The store to add the message filter to. + /// A new instance that filters out messages so they do not get stored. + public static ChatMessageStore WithAIContextProviderMessageRemoval(this ChatMessageStore store) + { + return new ChatMessageStoreMessageFilter( + innerChatMessageStore: store, + invokedMessagesFilter: (ctx) => + { + ctx.AIContextProviderMessages = null; + return ctx; + }); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreMessageFilter.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreMessageFilter.cs new file mode 100644 index 0000000000..d20e69f46b --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreMessageFilter.cs @@ -0,0 +1,67 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// A decorator that allows filtering the messages +/// passed into and out of an inner . +/// +public sealed class ChatMessageStoreMessageFilter : ChatMessageStore +{ + private readonly ChatMessageStore _innerChatMessageStore; + private readonly Func, IEnumerable>? _invokingMessagesFilter; + private readonly Func? _invokedMessagesFilter; + + /// + /// Initializes a new instance of the class. + /// + /// Use this constructor to customize how messages are filtered before and after invocation by + /// providing appropriate filter functions. If no filters are provided, the message store operates without + /// additional filtering. + /// The underlying chat message store to be wrapped. Cannot be null. + /// An optional filter function to apply to messages before they are invoked. If null, no filter is applied at this + /// stage. + /// An optional filter function to apply to the invocation context after messages have been invoked. If null, no + /// filter is applied at this stage. + /// Thrown if innerChatMessageStore is null. + public ChatMessageStoreMessageFilter( + ChatMessageStore innerChatMessageStore, + Func, IEnumerable>? invokingMessagesFilter = null, + Func? invokedMessagesFilter = null) + { + this._innerChatMessageStore = innerChatMessageStore ?? throw new ArgumentNullException(nameof(innerChatMessageStore)); + this._invokingMessagesFilter = invokingMessagesFilter; + this._invokedMessagesFilter = invokedMessagesFilter; + } + + /// + public override async ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + var messages = await this._innerChatMessageStore.InvokingAsync(context, cancellationToken).ConfigureAwait(false); + return this._invokingMessagesFilter != null ? this._invokingMessagesFilter(messages) : messages; + } + + /// + public override ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + if (this._invokedMessagesFilter != null) + { + context = this._invokedMessagesFilter(context); + } + + return this._innerChatMessageStore.InvokedAsync(context, cancellationToken); + } + + /// + public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + { + return this._innerChatMessageStore.Serialize(jsonSerializerOptions); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs index cca2adc5e4..1581f86ebc 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs @@ -157,7 +157,7 @@ public override async ValueTask InvokedAsync(InvokedContext context, Cancellatio } // Add ai context provider, request and response messages to the store - var allNewMessages = (context.AIContextProviderMessages ?? []).Concat(context.RequestMessages).Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages.Concat(context.AIContextProviderMessages ?? []).Concat(context.ResponseMessages ?? []); this._messages.AddRange(allNewMessages); if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null) diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreMessageFilterTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreMessageFilterTests.cs new file mode 100644 index 0000000000..3a0b06e40c --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreMessageFilterTests.cs @@ -0,0 +1,230 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Moq; + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Contains tests for the class. +/// +public sealed class ChatMessageStoreMessageFilterTests +{ + [Fact] + public void Constructor_WithNullInnerStore_ThrowsArgumentNullException() + { + // Arrange, Act & Assert + Assert.Throws(() => new ChatMessageStoreMessageFilter(null!)); + } + + [Fact] + public void Constructor_WithOnlyInnerStore_CreatesInstance() + { + // Arrange + var innerStoreMock = new Mock(); + + // Act + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object); + + // Assert + Assert.NotNull(filter); + } + + [Fact] + public void Constructor_WithAllParameters_CreatesInstance() + { + // Arrange + var innerStoreMock = new Mock(); + + IEnumerable InvokingFilter(IEnumerable msgs) => msgs; + ChatMessageStore.InvokedContext InvokedFilter(ChatMessageStore.InvokedContext ctx) => ctx; + + // Act + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object, InvokingFilter, InvokedFilter); + + // Assert + Assert.NotNull(filter); + } + + [Fact] + public async Task InvokingAsync_WithNoFilters_ReturnsInnerStoreMessagesAsync() + { + // Arrange + var innerStoreMock = new Mock(); + var expectedMessages = new List + { + new(ChatRole.User, "Hello"), + new(ChatRole.Assistant, "Hi there!") + }; + var context = new ChatMessageStore.InvokingContext([new ChatMessage(ChatRole.User, "Test")]); + + innerStoreMock + .Setup(s => s.InvokingAsync(context, It.IsAny())) + .ReturnsAsync(expectedMessages); + + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object); + + // Act + var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList(); + + // Assert + Assert.Equal(2, result.Count); + Assert.Equal("Hello", result[0].Text); + Assert.Equal("Hi there!", result[1].Text); + innerStoreMock.Verify(s => s.InvokingAsync(context, It.IsAny()), Times.Once); + } + + [Fact] + public async Task InvokingAsync_WithInvokingFilter_AppliesFilterAsync() + { + // Arrange + var innerStoreMock = new Mock(); + var innerMessages = new List + { + new(ChatRole.User, "Hello"), + new(ChatRole.Assistant, "Hi there!"), + new(ChatRole.User, "How are you?") + }; + var context = new ChatMessageStore.InvokingContext([new ChatMessage(ChatRole.User, "Test")]); + + innerStoreMock + .Setup(s => s.InvokingAsync(context, It.IsAny())) + .ReturnsAsync(innerMessages); + + // Filter to only user messages + IEnumerable InvokingFilter(IEnumerable msgs) => msgs.Where(m => m.Role == ChatRole.User); + + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object, InvokingFilter); + + // Act + var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList(); + + // Assert + Assert.Equal(2, result.Count); + Assert.All(result, msg => Assert.Equal(ChatRole.User, msg.Role)); + innerStoreMock.Verify(s => s.InvokingAsync(context, It.IsAny()), Times.Once); + } + + [Fact] + public async Task InvokingAsync_WithInvokingFilter_CanModifyMessagesAsync() + { + // Arrange + var innerStoreMock = new Mock(); + var innerMessages = new List + { + new(ChatRole.User, "Hello"), + new(ChatRole.Assistant, "Hi there!") + }; + var context = new ChatMessageStore.InvokingContext([new ChatMessage(ChatRole.User, "Test")]); + + innerStoreMock + .Setup(s => s.InvokingAsync(context, It.IsAny())) + .ReturnsAsync(innerMessages); + + // Filter that transforms messages + IEnumerable InvokingFilter(IEnumerable msgs) => + msgs.Select(m => new ChatMessage(m.Role, $"[FILTERED] {m.Text}")); + + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object, InvokingFilter); + + // Act + var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList(); + + // Assert + Assert.Equal(2, result.Count); + Assert.Equal("[FILTERED] Hello", result[0].Text); + Assert.Equal("[FILTERED] Hi there!", result[1].Text); + } + + [Fact] + public async Task InvokedAsync_WithNoFilters_CallsInnerStoreAsync() + { + // Arrange + var innerStoreMock = new Mock(); + var requestMessages = new List { new(ChatRole.User, "Hello") }; + var chatMessageStoreMessages = new List { new(ChatRole.System, "System") }; + var context = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages); + + innerStoreMock + .Setup(s => s.InvokedAsync(It.IsAny(), It.IsAny())) + .Returns(default(ValueTask)); + + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object); + + // Act + await filter.InvokedAsync(context, CancellationToken.None); + + // Assert + innerStoreMock.Verify(s => s.InvokedAsync(context, It.IsAny()), Times.Once); + } + + [Fact] + public async Task InvokedAsync_WithInvokedFilter_AppliesFilterAsync() + { + // Arrange + var innerStoreMock = new Mock(); + var requestMessages = new List { new(ChatRole.User, "Hello") }; + var chatMessageStoreMessages = new List { new(ChatRole.System, "System") }; + var responseMessages = new List { new(ChatRole.Assistant, "Response") }; + var context = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages) + { + ResponseMessages = responseMessages + }; + + ChatMessageStore.InvokedContext? capturedContext = null; + innerStoreMock + .Setup(s => s.InvokedAsync(It.IsAny(), It.IsAny())) + .Callback((ctx, ct) => capturedContext = ctx) + .Returns(default(ValueTask)); + + // Filter that modifies the context + ChatMessageStore.InvokedContext InvokedFilter(ChatMessageStore.InvokedContext ctx) + { + var modifiedRequestMessages = ctx.RequestMessages.Select(m => new ChatMessage(m.Role, $"[FILTERED] {m.Text}")).ToList(); + return new ChatMessageStore.InvokedContext(modifiedRequestMessages, ctx.ChatMessageStoreMessages) + { + ResponseMessages = ctx.ResponseMessages, + AIContextProviderMessages = ctx.AIContextProviderMessages, + InvokeException = ctx.InvokeException + }; + } + + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object, invokedMessagesFilter: InvokedFilter); + + // Act + await filter.InvokedAsync(context, CancellationToken.None); + + // Assert + Assert.NotNull(capturedContext); + Assert.Single(capturedContext.RequestMessages); + Assert.Equal("[FILTERED] Hello", capturedContext.RequestMessages.First().Text); + innerStoreMock.Verify(s => s.InvokedAsync(It.IsAny(), It.IsAny()), Times.Once); + } + + [Fact] + public void Serialize_DelegatesToInnerStore() + { + // Arrange + var innerStoreMock = new Mock(); + var expectedJson = JsonSerializer.SerializeToElement("data", TestJsonSerializerContext.Default.String); + + innerStoreMock + .Setup(s => s.Serialize(It.IsAny())) + .Returns(expectedJson); + + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object); + + // Act + var result = filter.Serialize(); + + // Assert + Assert.Equal(expectedJson.GetRawText(), result.GetRawText()); + innerStoreMock.Verify(s => s.Serialize(null), Times.Once); + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs index 3b7f1742dd..43bfacca79 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs @@ -77,8 +77,8 @@ public async Task InvokedAsyncAddsMessagesAsync() Assert.Equal(4, store.Count); Assert.Equal("original instructions", store[0].Text); - Assert.Equal("additional context", store[1].Text); - Assert.Equal("Hello", store[2].Text); + Assert.Equal("Hello", store[1].Text); + Assert.Equal("additional context", store[2].Text); Assert.Equal("Hi there!", store[3].Text); } From 446001b0ad943905afc717487a0dc5b02b0ca497 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 5 Dec 2025 15:02:12 +0000 Subject: [PATCH 08/14] Update sample and cosmos message store to store AIContextProvider messages in right order. Fix unit tests. --- .../Program.cs | 2 +- .../CosmosChatMessageStore.cs | 2 +- .../CosmosChatMessageStoreTests.cs | 19 +++++++++++++++---- .../ChatClient/ChatClientAgentTests.cs | 8 ++++---- 4 files changed, 21 insertions(+), 10 deletions(-) diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs index e12172d397..10df206d62 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs @@ -121,7 +121,7 @@ public override async ValueTask InvokedAsync(InvokedContext context, Cancellatio // Add both request and response messages to the store // Optionally messages produced by the AIContextProvider can also be persisted (not shown). - var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages.Concat(context.AIContextProviderMessages ?? []).Concat(context.ResponseMessages ?? []); await collection.UpsertAsync(allNewMessages.Select(x => new ChatHistoryItem() { diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index d9e0a1e17d..32bdcbcec8 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -364,7 +364,7 @@ public override async ValueTask InvokedAsync(InvokedContext context, Cancellatio } #pragma warning restore CA1513 - var messageList = (context.AIContextProviderMessages ?? []).Concat(context.RequestMessages).Concat(context.ResponseMessages ?? []).ToList(); + var messageList = context.RequestMessages.Concat(context.AIContextProviderMessages ?? []).Concat(context.ResponseMessages ?? []).ToList(); if (messageList.Count == 0) { return; diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index b40edeb7a3..9bc7e3793b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -277,16 +277,25 @@ public async Task InvokedAsync_WithMultipleMessages_ShouldAddAllMessagesAsync() this.SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); - var messages = new[] + var requestMessages = new[] { new ChatMessage(ChatRole.User, "First message"), new ChatMessage(ChatRole.Assistant, "Second message"), new ChatMessage(ChatRole.User, "Third message") }; + var aiContextProviderMessages = new[] + { + new ChatMessage(ChatRole.System, "System context message") + }; + var responseMessages = new[] + { + new ChatMessage(ChatRole.Assistant, "Response message") + }; - var context = new ChatMessageStore.InvokedContext(messages, []) + var context = new ChatMessageStore.InvokedContext(requestMessages, []) { - ResponseMessages = [] + AIContextProviderMessages = aiContextProviderMessages, + ResponseMessages = responseMessages }; // Act @@ -296,10 +305,12 @@ public async Task InvokedAsync_WithMultipleMessages_ShouldAddAllMessagesAsync() var invokingContext = new ChatMessageStore.InvokingContext([]); var retrievedMessages = await store.InvokingAsync(invokingContext); var messageList = retrievedMessages.ToList(); - Assert.Equal(3, messageList.Count); + Assert.Equal(5, messageList.Count); Assert.Equal("First message", messageList[0].Text); Assert.Equal("Second message", messageList[1].Text); Assert.Equal("Third message", messageList[2].Text); + Assert.Equal("System context message", messageList[3].Text); + Assert.Equal("Response message", messageList[4].Text); } #endregion diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs index c25bf0732a..0aef0f43b6 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs @@ -613,8 +613,8 @@ public async Task RunAsyncInvokesAIContextProviderAndUsesResultAsync() // Verify that the thread was updated with the ai context provider, input and response messages var messageStore = Assert.IsType(thread!.MessageStore); Assert.Equal(3, messageStore.Count); - Assert.Equal("context provider message", messageStore[0].Text); - Assert.Equal("user message", messageStore[1].Text); + Assert.Equal("user message", messageStore[0].Text); + Assert.Equal("context provider message", messageStore[1].Text); Assert.Equal("response", messageStore[2].Text); mockProvider.Verify(p => p.InvokingAsync(It.IsAny(), It.IsAny()), Times.Once); @@ -2070,8 +2070,8 @@ public async Task RunStreamingAsyncInvokesAIContextProviderAndUsesResultAsync() // Verify that the thread was updated with the input and response messages var messageStore = Assert.IsType(thread!.MessageStore); Assert.Equal(3, messageStore.Count); - Assert.Equal("context provider message", messageStore[0].Text); - Assert.Equal("user message", messageStore[1].Text); + Assert.Equal("user message", messageStore[0].Text); + Assert.Equal("context provider message", messageStore[1].Text); Assert.Equal("response", messageStore[2].Text); mockProvider.Verify(p => p.InvokingAsync(It.IsAny(), It.IsAny()), Times.Once); From 9f6d0f545ed1c2af00cf0b72264ad78d11311d8f Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 5 Dec 2025 15:09:55 +0000 Subject: [PATCH 09/14] Update Workflowmessage store to use aicontext provider messages. --- .../src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs index 08d92157a6..87cef04e76 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowMessageStore.cs @@ -56,8 +56,7 @@ public override ValueTask InvokedAsync(InvokedContext context, CancellationToken return default; } - // Add both request and response messages to the store - var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages.Concat(context.AIContextProviderMessages ?? []).Concat(context.ResponseMessages ?? []); this._chatMessages.AddRange(allNewMessages); return default; From cc4f51cec672abe8102e033d0fdfe744659e77fc Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 5 Dec 2025 19:30:32 +0000 Subject: [PATCH 10/14] Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../ChatMessageStoreExtensions.cs | 2 +- .../InMemoryChatMessageStore.cs | 2 +- .../ChatClient/ChatClientAgentTests.cs | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs index 589325748e..2c0dbf7220 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs @@ -12,7 +12,7 @@ namespace Microsoft.Agents.AI; public static class ChatMessageStoreExtensions { /// - /// Adds message filering to an existing store, so that messages passed to the store and messages produced by the store + /// Adds message filtering to an existing store, so that messages passed to the store and messages produced by the store /// can be filtered, updated or replaced. /// /// The store to add the message filter to. diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs index 1581f86ebc..f7f4522f8f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatMessageStore.cs @@ -156,7 +156,7 @@ public override async ValueTask InvokedAsync(InvokedContext context, Cancellatio return; } - // Add ai context provider, request and response messages to the store + // Add request, AI context provider, and response messages to the store var allNewMessages = context.RequestMessages.Concat(context.AIContextProviderMessages ?? []).Concat(context.ResponseMessages ?? []); this._messages.AddRange(allNewMessages); diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs index 0aef0f43b6..464481fb77 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs @@ -2067,7 +2067,7 @@ public async Task RunStreamingAsyncInvokesAIContextProviderAndUsesResultAsync() Assert.Contains(capturedTools, t => t.Name == "base function"); Assert.Contains(capturedTools, t => t.Name == "context provider function"); - // Verify that the thread was updated with the input and response messages + // Verify that the thread was updated with the input, ai context provider, and response messages var messageStore = Assert.IsType(thread!.MessageStore); Assert.Equal(3, messageStore.Count); Assert.Equal("user message", messageStore[0].Text); From 5bf38b1f61756ab781fd9ec56424a9efd72177c4 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Wed, 31 Dec 2025 10:59:26 +0000 Subject: [PATCH 11/14] Apply suggestions from code review Co-authored-by: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com> --- .../ChatMessageStoreExtensions.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs index 2c0dbf7220..4c0e128102 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs @@ -16,9 +16,9 @@ public static class ChatMessageStoreExtensions /// can be filtered, updated or replaced. /// /// The store to add the message filter to. - /// An optional filter function to apply to messages before they are invoked. If null, no filter is applied at this + /// An optional filter function to apply to messages produced by the store. If null, no filter is applied at this /// stage. - /// An optional filter function to apply to the invocation context after messages have been invoked. If null, no + /// An optional filter function to apply to the invocation context messages before they passed to the store. If null, no /// filter is applied at this stage. /// The with filtering applied. public static ChatMessageStore WithMessageFilters( From 5862cfe3543c202bf37fa9418faa5e8e314343bd Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Wed, 31 Dec 2025 11:02:52 +0000 Subject: [PATCH 12/14] Improve xml docs messaging --- .../ChatMessageStoreExtensions.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs index 4c0e128102..a205fc1d9e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreExtensions.cs @@ -18,7 +18,7 @@ public static class ChatMessageStoreExtensions /// The store to add the message filter to. /// An optional filter function to apply to messages produced by the store. If null, no filter is applied at this /// stage. - /// An optional filter function to apply to the invocation context messages before they passed to the store. If null, no + /// An optional filter function to apply to the invoked context messages before they are passed to the store. If null, no /// filter is applied at this stage. /// The with filtering applied. public static ChatMessageStore WithMessageFilters( From 32ac83cef6ffc2b58f87b374b26dc997d1446893 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Wed, 31 Dec 2025 11:20:59 +0000 Subject: [PATCH 13/14] Address code review comments. --- .../ChatMessageStoreMessageFilter.cs | 9 ++++- .../ChatMessageStoreMessageFilterTests.cs | 37 +++---------------- 2 files changed, 14 insertions(+), 32 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreMessageFilter.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreMessageFilter.cs index d20e69f46b..e58f233067 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreMessageFilter.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageStoreMessageFilter.cs @@ -6,6 +6,7 @@ using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; @@ -36,7 +37,13 @@ public ChatMessageStoreMessageFilter( Func, IEnumerable>? invokingMessagesFilter = null, Func? invokedMessagesFilter = null) { - this._innerChatMessageStore = innerChatMessageStore ?? throw new ArgumentNullException(nameof(innerChatMessageStore)); + this._innerChatMessageStore = Throw.IfNull(innerChatMessageStore); + + if (invokingMessagesFilter == null && invokedMessagesFilter == null) + { + throw new ArgumentException("At least one filter function, invokingMessagesFilter or invokedMessagesFilter, must be provided."); + } + this._invokingMessagesFilter = invokingMessagesFilter; this._invokedMessagesFilter = invokedMessagesFilter; } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreMessageFilterTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreMessageFilterTests.cs index 3a0b06e40c..ab10c377ae 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreMessageFilterTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageStoreMessageFilterTests.cs @@ -24,16 +24,13 @@ public void Constructor_WithNullInnerStore_ThrowsArgumentNullException() } [Fact] - public void Constructor_WithOnlyInnerStore_CreatesInstance() + public void Constructor_WithOnlyInnerStore_Throws() { // Arrange var innerStoreMock = new Mock(); - // Act - var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object); - - // Assert - Assert.NotNull(filter); + // Act & Assert + Assert.Throws(() => new ChatMessageStoreMessageFilter(innerStoreMock.Object)); } [Fact] @@ -53,7 +50,7 @@ public void Constructor_WithAllParameters_CreatesInstance() } [Fact] - public async Task InvokingAsync_WithNoFilters_ReturnsInnerStoreMessagesAsync() + public async Task InvokingAsync_WithNoOpFilters_ReturnsInnerStoreMessagesAsync() { // Arrange var innerStoreMock = new Mock(); @@ -68,7 +65,7 @@ public async Task InvokingAsync_WithNoFilters_ReturnsInnerStoreMessagesAsync() .Setup(s => s.InvokingAsync(context, It.IsAny())) .ReturnsAsync(expectedMessages); - var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object); + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object, x => x, x => x); // Act var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList(); @@ -142,28 +139,6 @@ IEnumerable InvokingFilter(IEnumerable msgs) => Assert.Equal("[FILTERED] Hi there!", result[1].Text); } - [Fact] - public async Task InvokedAsync_WithNoFilters_CallsInnerStoreAsync() - { - // Arrange - var innerStoreMock = new Mock(); - var requestMessages = new List { new(ChatRole.User, "Hello") }; - var chatMessageStoreMessages = new List { new(ChatRole.System, "System") }; - var context = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages); - - innerStoreMock - .Setup(s => s.InvokedAsync(It.IsAny(), It.IsAny())) - .Returns(default(ValueTask)); - - var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object); - - // Act - await filter.InvokedAsync(context, CancellationToken.None); - - // Assert - innerStoreMock.Verify(s => s.InvokedAsync(context, It.IsAny()), Times.Once); - } - [Fact] public async Task InvokedAsync_WithInvokedFilter_AppliesFilterAsync() { @@ -218,7 +193,7 @@ public void Serialize_DelegatesToInnerStore() .Setup(s => s.Serialize(It.IsAny())) .Returns(expectedJson); - var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object); + var filter = new ChatMessageStoreMessageFilter(innerStoreMock.Object, x => x, x => x); // Act var result = filter.Serialize(); From ba985a0111ce6c002ea04316914f8b90a0559a97 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Wed, 31 Dec 2025 13:13:21 +0000 Subject: [PATCH 14/14] Also notify message store on failure --- .../ChatClient/ChatClientAgent.cs | 30 ++++++++++ .../ChatClient/ChatClientAgentTests.cs | 59 ++++++++++++++++++- 2 files changed, 88 insertions(+), 1 deletion(-) diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index c0bb157478..9c5858b8e2 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -225,6 +225,7 @@ protected override async IAsyncEnumerable RunCoreStreami } catch (Exception ex) { + await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); throw; } @@ -239,6 +240,7 @@ protected override async IAsyncEnumerable RunCoreStreami } catch (Exception ex) { + await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); throw; } @@ -260,6 +262,7 @@ protected override async IAsyncEnumerable RunCoreStreami } catch (Exception ex) { + await NotifyMessageStoreOfFailureAsync(safeThread, ex, inputMessages, chatMessageStoreMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, aiContextProviderMessages, cancellationToken).ConfigureAwait(false); throw; } @@ -398,6 +401,7 @@ private async Task RunCoreAsync requestMessages, + IEnumerable? chatMessageStoreMessages, + IEnumerable? aiContextProviderMessages, + CancellationToken cancellationToken) + { + var messageStore = thread.MessageStore; + + // Only notify the message store if we have one. + // If we don't have one, it means that the chat history is service managed and the underlying service is responsible for storing messages. + if (messageStore is not null) + { + var invokedContext = new ChatMessageStore.InvokedContext(requestMessages, chatMessageStoreMessages!) + { + AIContextProviderMessages = aiContextProviderMessages, + InvokeException = ex + }; + + return messageStore.InvokedAsync(invokedContext, cancellationToken).AsTask(); + } + + return Task.CompletedTask; + } + private static Task NotifyMessageStoreOfNewMessagesAsync( ChatClientAgentThread thread, IEnumerable requestMessages, diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs index 464481fb77..5850bc56ba 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs @@ -502,6 +502,12 @@ public async Task RunAsyncUsesChatMessageStoreFactoryWhenProvidedAndNoConversati It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")])); Mock mockChatMessageStore = new(); + mockChatMessageStore.Setup(s => s.InvokingAsync( + It.IsAny(), + It.IsAny())).ReturnsAsync([new ChatMessage(ChatRole.User, "Existing Chat History")]); + mockChatMessageStore.Setup(s => s.InvokedAsync( + It.IsAny(), + It.IsAny())).Returns(new ValueTask()); Mock> mockFactory = new(); mockFactory.Setup(f => f(It.IsAny())).Returns(mockChatMessageStore.Object); @@ -518,7 +524,58 @@ public async Task RunAsyncUsesChatMessageStoreFactoryWhenProvidedAndNoConversati // Assert Assert.IsType(thread!.MessageStore, exactMatch: false); - mockChatMessageStore.Verify(s => s.InvokedAsync(It.Is(x => x.RequestMessages.Count() == 1 && x.ResponseMessages!.Count() == 1), It.IsAny()), Times.Once); + mockService.Verify( + x => x.GetResponseAsync( + It.Is>(msgs => msgs.Count() == 2 && msgs.Any(m => m.Text == "Existing Chat History") && msgs.Any(m => m.Text == "test")), + It.IsAny(), + It.IsAny()), + Times.Once); + mockChatMessageStore.Verify(s => s.InvokingAsync( + It.Is(x => x.RequestMessages.Count() == 1), + It.IsAny()), + Times.Once); + mockChatMessageStore.Verify(s => s.InvokedAsync( + It.Is(x => x.RequestMessages.Count() == 1 && x.ChatMessageStoreMessages.Count() == 1 && x.ResponseMessages!.Count() == 1), + It.IsAny()), + Times.Once); + mockFactory.Verify(f => f(It.IsAny()), Times.Once); + } + + /// + /// Verify that RunAsync notifies the ChatMessageStore on failure. + /// + [Fact] + public async Task RunAsyncNotifiesChatMessageStoreOnFailureAsync() + { + // Arrange + Mock mockService = new(); + mockService.Setup( + s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())).Throws(new InvalidOperationException("Test Error")); + + Mock mockChatMessageStore = new(); + + Mock> mockFactory = new(); + mockFactory.Setup(f => f(It.IsAny())).Returns(mockChatMessageStore.Object); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + ChatOptions = new() { Instructions = "test instructions" }, + ChatMessageStoreFactory = mockFactory.Object + }); + + // Act + ChatClientAgentThread? thread = agent.GetNewThread() as ChatClientAgentThread; + await Assert.ThrowsAsync(() => agent.RunAsync([new(ChatRole.User, "test")], thread)); + + // Assert + Assert.IsType(thread!.MessageStore, exactMatch: false); + mockChatMessageStore.Verify(s => s.InvokedAsync( + It.Is(x => x.RequestMessages.Count() == 1 && x.ResponseMessages == null && x.InvokeException!.Message == "Test Error"), + It.IsAny()), + Times.Once); mockFactory.Verify(f => f(It.IsAny()), Times.Once); }