.NET: [BREAKING] Refactor ChatMessageStore methods to be similar to AIContextProvider and add filtering support (#2604)

* Refactor ChatMessageStore methods to be similar to AIContextProvider

* Fix file encoding

* Ensure that AIContextProvider messages area also persisted.

* Update formatting and seal context classes

* Improve formatting

* Remove optional messages from constructor and add unit test

* Add ChatMessageStore filtering via a decorator

* Update sample and cosmos message store to store AIContextProvider messages in right order. Fix unit tests.

* Update Workflowmessage store to use aicontext provider messages.

* Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com>

* Improve xml docs messaging

* Address code review comments.

* Also notify message store on failure

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com>

dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs

@@ -44,11 +44,19 @@ namespace SampleApp
                 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<ChatMessage> responseMessages = CloneAndToUpperCase(messages, this.Name).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)
+            {
+                ResponseMessages = responseMessages
+            };
+            await typedThread.MessageStore.InvokedAsync(invokedContext, cancellationToken);
 
             return new AgentRunResponse
             {
@@ -68,11 +76,19 @@ namespace SampleApp
                 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<ChatMessage> responseMessages = CloneAndToUpperCase(messages, this.Name).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)
+            {
+                ResponseMessages = responseMessages
+            };
+            await typedThread.MessageStore.InvokedAsync(invokedContext, cancellationToken);
 
             foreach (var message in responseMessages)
             {

dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs

@@ -62,7 +62,12 @@ AIAgent agent = azureOpenAIClient
     .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();

dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyThreadStorage/Program.cs

@@ -89,24 +89,7 @@ namespace SampleApp
 
         public string? ThreadDbKey { get; private set; }
 
-        public override async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
-        {
-            this.ThreadDbKey ??= Guid.NewGuid().ToString("N");
-
-            var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("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<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default)
+        public override async ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
         {
             var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
             await collection.EnsureCollectionExistsAsync(cancellationToken);
@@ -124,6 +107,33 @@ namespace SampleApp
             return messages;
         }
 
+        public override async ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default)
+        {
+            // Don't store messages if the request failed.
+            if (context.InvokeException is not null)
+            {
+                return;
+            }
+
+            this.ThreadDbKey ??= Guid.NewGuid().ToString("N");
+
+            var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
+            await collection.EnsureCollectionExistsAsync(cancellationToken);
+
+            // 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.AIContextProviderMessages ?? []).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);