Merge pull request #824 from westey-m/chatmessagestore-serialize-updates

Update ChatMessageStore and Serialization related code samples

Author: westey-m

agent-framework/integrations/ag-ui/getting-started.md

@@ -178,7 +178,7 @@ AIAgent agent = chatClient.AsAIAgent(
     name: "agui-client",
     description: "AG-UI Client Agent");
 
-AgentThread thread = agent.GetNewThread();
+AgentThread thread = await agent.GetNewThreadAsync();
 List<ChatMessage> messages =
 [
     new(ChatRole.System, "You are a helpful assistant.")

agent-framework/migration-guide/from-semantic-kernel/index.md

@@ -103,7 +103,7 @@ The agent is responsible for creating the thread.
 
 ```csharp
 // New.
-AgentThread thread = agent.GetNewThread();
+AgentThread thread = await agent.GetNewThreadAsync();
 ```
 
 ## 4. Hosted Agent Thread Cleanup

agent-framework/tutorials/agents/function-tools-approvals.md

@@ -66,7 +66,7 @@ Since you now have a function that requires approval, the agent might respond wi
 You can check the response content for any `FunctionApprovalRequestContent` instances, which indicates that the agent requires user approval for a function.
 
 ```csharp
-AgentThread thread = agent.GetNewThread();
+AgentThread thread = await agent.GetNewThreadAsync();
 AgentResponse response = await agent.RunAsync("What is the weather like in Amsterdam?", thread);
 
 var functionApprovalRequests = response.Messages

agent-framework/tutorials/agents/memory.md

@@ -146,6 +146,8 @@ To use the custom `AIContextProvider`, you need to provide an `AIContextProvider
 
 When creating a `ChatClientAgent` it is possible to provide a `ChatClientAgentOptions` object that allows providing the `AIContextProviderFactory` in addition to all other agent options.
 
+The factory is an async function that receives a context object and a cancellation token.
+
 ```csharp
 using System;
 using Azure.AI.OpenAI;
@@ -161,19 +163,20 @@ ChatClient chatClient = new AzureOpenAIClient(
 AIAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions()
 {
     ChatOptions = new() { Instructions = "You are a friendly assistant. Always address the user by their name." },
-    AIContextProviderFactory = ctx => new UserInfoMemory(
-        chatClient.AsIChatClient(),
-        ctx.SerializedState,
-        ctx.JsonSerializerOptions)
+    AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(
+        new UserInfoMemory(
+            chatClient.AsIChatClient(),
+            ctx.SerializedState,
+            ctx.JsonSerializerOptions))
 });
 ```
 
-When creating a new thread, the `AIContextProvider` will be created by `GetNewThread`
+When creating a new thread, the `AIContextProvider` will be created by `GetNewThreadAsync`
 and attached to the thread. Once memories are extracted it is therefore possible to access the memory component via the thread's `GetService` method and inspect the memories.
 
 ```csharp
 // Create a new thread for the conversation.
-AgentThread thread = agent.GetNewThread();
+AgentThread thread = await agent.GetNewThreadAsync();
 
 Console.WriteLine(await agent.RunAsync("Hello, what is the square root of 9?", thread));
 Console.WriteLine(await agent.RunAsync("My name is Ruaidhrí", thread));

agent-framework/tutorials/agents/multi-turn-conversation.md

@@ -27,10 +27,10 @@ For prerequisites and creating the agent, see the [Create and run a simple agent
 Agents are stateless and do not maintain any state internally between calls.
 To have a multi-turn conversation with an agent, you need to create an object to hold the conversation state and pass this object to the agent when running it.
 
-To create the conversation state object, call the `GetNewThread` method on the agent instance.
+To create the conversation state object, call the `GetNewThreadAsync` method on the agent instance.
 
 ```csharp
-AgentThread thread = agent.GetNewThread();
+AgentThread thread = await agent.GetNewThreadAsync();
 ```
 
 You can then pass this thread object to the `RunAsync` and `RunStreamingAsync` methods on the agent instance, along with the user input.
@@ -52,8 +52,8 @@ These threads can then be used to maintain separate conversation states for each
 The conversations will be fully independent of each other, since the agent does not maintain any state internally.
 
 ```csharp
-AgentThread thread1 = agent.GetNewThread();
-AgentThread thread2 = agent.GetNewThread();
+AgentThread thread1 = await agent.GetNewThreadAsync();
+AgentThread thread2 = await agent.GetNewThreadAsync();
 Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread1));
 Console.WriteLine(await agent.RunAsync("Tell me a joke about a robot.", thread2));
 Console.WriteLine(await agent.RunAsync("Now add some emojis to the joke and tell it in the voice of a pirate's parrot.", thread1));

agent-framework/tutorials/agents/persisted-conversation.md

@@ -38,7 +38,7 @@ AIAgent agent = new AzureOpenAIClient(
      .GetChatClient("gpt-4o-mini")
      .AsAIAgent(instructions: "You are a helpful assistant.", name: "Assistant");
 
-AgentThread thread = agent.GetNewThread();
+AgentThread thread = await agent.GetNewThreadAsync();
 ```
 
 Run the agent, passing in the thread, so that the `AgentThread` includes this exchange.
@@ -75,7 +75,7 @@ string loadedJson = await File.ReadAllTextAsync(filePath);
 JsonElement reloaded = JsonSerializer.Deserialize<JsonElement>(loadedJson, JsonSerializerOptions.Web);
 
 // Deserialize the thread into an AgentThread tied to the same agent type
-AgentThread resumedThread = agent.DeserializeThread(reloaded, JsonSerializerOptions.Web);
+AgentThread resumedThread = await agent.DeserializeThreadAsync(reloaded, JsonSerializerOptions.Web);
 ```
 
 Use the resumed thread to continue the conversation.

agent-framework/tutorials/agents/third-party-chat-history-storage.md

@@ -47,32 +47,32 @@ To create a custom `ChatMessageStore`, you need to implement the abstract `ChatM
 
 The most important methods to implement are:
 
-- `AddMessagesAsync` - called to add new messages to the store.
-- `GetMessagesAsync` - called to retrieve the messages from the store.
+- `InvokingAsync` - called at the start of agent invocation to retrieve messages from the store that should be provided as context.
+- `InvokedAsync` - called at the end of agent invocation to add new messages to the store.
 
-`GetMessagesAsync` should return the messages in ascending chronological order. All messages returned by it will be used by the `ChatClientAgent` when making calls to the underlying <xref:Microsoft.Extensions.AI.IChatClient>.  It's therefore important that this method considers the limits of the underlying model, and only returns as many messages as can be handled by the model.
+`InvokingAsync` should return the messages in ascending chronological order (oldest first). All messages returned by it will be used by the `ChatClientAgent` when making calls to the underlying <xref:Microsoft.Extensions.AI.IChatClient>. It's therefore important that this method considers the limits of the underlying model, and only returns as many messages as can be handled by the model.
 
-Any chat history reduction logic, such as summarization or trimming, should be done before returning messages from `GetMessagesAsync`.
+Any chat history reduction logic, such as summarization or trimming, should be done before returning messages from `InvokingAsync`.
 
 ### Serialization
 
 `ChatMessageStore` instances are created and attached to an `AgentThread` when the thread is created, and when a thread is resumed from a serialized state.
 
 While the actual messages making up the chat history are stored externally, the `ChatMessageStore` instance might need to store keys or other state to identify the chat history in the external store.
 
-To allow persisting threads, you need to implement the `SerializeStateAsync` method of the `ChatMessageStore` class. You also need to provide a constructor that takes a <xref:System.Text.Json.JsonElement> parameter, which can be used to deserialize the state when resuming a thread.
+To allow persisting threads, you need to implement the `Serialize` method of the `ChatMessageStore` class. This method should return a `JsonElement` containing the state needed to restore the store later. When deserializing, the agent framework will pass this serialized state to the ChatMessageStoreFactory, allowing you to use it to recreate the store.
 
 ### Sample ChatMessageStore implementation
 
 The following sample implementation stores chat messages in a vector store.
 
-`AddMessagesAsync` upserts messages into the vector store, using a unique key for each message.
+`InvokedAsync` upserts messages into the vector store, using a unique key for each message. It stores both the request messages and response messages from the invocation context.
 
-`GetMessagesAsync` retrieves the messages for the current thread from the vector store, orders them by timestamp, and returns them in ascending order.
+`InvokingAsync` retrieves the messages for the current thread from the vector store, orders them by timestamp, and returns them in ascending chronological order (oldest first).
 
-When the first message is received, the store generates a unique key for the thread, which is then used to identify the chat history in the vector store for subsequent calls.
+When the first invocation occurs, the store generates a unique key for the thread, which is then used to identify the chat history in the vector store for subsequent calls.
 
-The unique key is stored in the `ThreadDbKey` property, which is serialized and deserialized using the `SerializeStateAsync` method and the constructor that takes a `JsonElement`.
+The unique key is stored in the `ThreadDbKey` property, which is serialized using the `Serialize` method and deserialized via the constructor that takes a `JsonElement`.
 This key will therefore be persisted as part of the `AgentThread` state, allowing the thread to be resumed later and continue using the same chat history.
 
 ```csharp
@@ -105,31 +105,22 @@ internal sealed class VectorChatMessageStore : ChatMessageStore
 
     public string? ThreadDbKey { get; private set; }
 
-    public override async Task AddMessagesAsync(
-        IEnumerable<ChatMessage> messages,
-        CancellationToken cancellationToken)
+    public override async ValueTask<IEnumerable<ChatMessage>> InvokingAsync(
+        InvokingContext context,
+        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()
+        if (this.ThreadDbKey is null)
         {
-            Key = this.ThreadDbKey + x.MessageId,
-            Timestamp = DateTimeOffset.UtcNow,
-            ThreadId = this.ThreadDbKey,
-            SerializedMessage = JsonSerializer.Serialize(x),
-            MessageText = x.Text
-        }), cancellationToken);
-    }
+            // No thread key yet, so no messages to retrieve
+            return [];
+        }
 
-    public override async Task<IEnumerable<ChatMessage>> GetMessagesAsync(
-        CancellationToken cancellationToken)
-    {
         var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
         await collection.EnsureCollectionExistsAsync(cancellationToken);
         var records = collection
             .GetAsync(
-                x => x.ThreadId == this.ThreadDbKey, 10,
+                x => x.ThreadId == this.ThreadDbKey, 
+                10,
                 new() { OrderBy = x => x.Descending(y => y.Timestamp) },
                 cancellationToken);
 
@@ -139,10 +130,41 @@ internal sealed class VectorChatMessageStore : ChatMessageStore
             messages.Add(JsonSerializer.Deserialize<ChatMessage>(record.SerializedMessage!)!);
         }
 
+        // Reverse to return in ascending chronological order (oldest first)
         messages.Reverse();
         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);
+        
+        // Store request messages, response messages, and optionally AIContextProvider messages
+        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 you can retrieve the messages using the same thread id.
         JsonSerializer.SerializeToElement(this.ThreadDbKey);
@@ -169,28 +191,46 @@ To use the custom `ChatMessageStore`, you need to provide a `ChatMessageStoreFac
 
 When creating a `ChatClientAgent` it is possible to provide a `ChatClientAgentOptions` object that allows providing the `ChatMessageStoreFactory` in addition to all other agent options.
 
+The factory is an async function that receives a context object and a cancellation token, and returns a `ValueTask<ChatMessageStore>`.
+
 ```csharp
 using Azure.AI.OpenAI;
 using Azure.Identity;
-using OpenAI;
+using Microsoft.Extensions.VectorData;
+using Microsoft.SemanticKernel.Connectors.InMemory;
+
+// Create a vector store to store the chat messages in.
+VectorStore vectorStore = new InMemoryVectorStore();
 
 AIAgent agent = new AzureOpenAIClient(
     new Uri("https://<myresource>.openai.azure.com"),
     new AzureCliCredential())
-    .GetChatClient("gpt-4o-mini")
-    .AsAIAgent(new ChatClientAgentOptions
-    {
-        Name = "Joker",
-        ChatOptions = new() { Instructions = "You are good at telling jokes." },
-        ChatMessageStoreFactory = ctx =>
-        {
-            // Create a new chat message store for this agent that stores the messages in a vector store.
-            return new VectorChatMessageStore(
-                new InMemoryVectorStore(),
+     .GetChatClient("gpt-4o-mini")
+     .AsAIAgent(new ChatClientAgentOptions
+     {
+         Name = "Joker",
+         ChatOptions = new() { Instructions = "You are good at telling jokes." },
+         ChatMessageStoreFactory = (ctx, ct) => new ValueTask<ChatMessageStore>(
+             // Create a new chat message store for this agent that stores the messages in a vector store.
+             // Each thread must get its own copy of the VectorChatMessageStore, since the store
+             // also contains the id that the thread is stored under.
+             new VectorChatMessageStore(
+                vectorStore,
                 ctx.SerializedState,
-                ctx.JsonSerializerOptions);
-        }
-    });
+                ctx.JsonSerializerOptions))
+     });
+
+// Start a new thread for the agent conversation.
+AgentThread thread = await agent.GetNewThreadAsync();
+
+// Run the agent with the thread
+var response = await agent.RunAsync("Tell me a joke about a pirate.", thread);
+
+// The thread state can be serialized for storage
+JsonElement serializedThread = thread.Serialize();
+
+// Later, deserialize the thread to resume the conversation
+AgentThread resumedThread = await agent.DeserializeThreadAsync(serializedThread);
 ```
 
 ::: zone-end

agent-framework/user-guide/agents/agent-background-responses.md

@@ -66,7 +66,7 @@ AgentRunOptions options = new()
     AllowBackgroundResponses = true
 };
 
-AgentThread thread = agent.GetNewThread();
+AgentThread thread = await agent.GetNewThreadAsync();
 
 // Get initial response - may return with or without a continuation token
 AgentResponse response = await agent.RunAsync("Write a very long novel about otters in space.", thread, options);
@@ -108,7 +108,7 @@ AgentRunOptions options = new()
     AllowBackgroundResponses = true
 };
 
-AgentThread thread = agent.GetNewThread();
+AgentThread thread = await agent.GetNewThreadAsync();
 
 AgentResponseUpdate? latestReceivedUpdate = null;