.NET: Add per run / thread feature collection support and improved custom ChatMessageStore support

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

@@ -28,10 +28,10 @@ internal sealed class UpperCaseParrotAgent : AIAgent
     {
         public override string? Name => "UpperCaseParrotAgent";
 
-        public override AgentThread GetNewThread()
+        public override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
             => new CustomAgentThread();
 
-        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+        public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
             => new CustomAgentThread(serializedThread, jsonSerializerOptions);
 
         protected override async Task<AgentRunResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)

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

@@ -22,83 +22,237 @@
 // Replace this with a vector store implementation of your choice if you want to persist the chat history to disk.
 VectorStore vectorStore = new InMemoryVectorStore();
 
-// Create the agent
-AIAgent agent = new AzureOpenAIClient(
-    new Uri(endpoint),
-    new AzureCliCredential())
-    .GetChatClient(deploymentName)
-    .CreateAIAgent(new ChatClientAgentOptions
-    {
-        ChatOptions = new() { Instructions = "You are good at telling jokes." },
-        Name = "Joker",
-        ChatMessageStoreFactory = ctx =>
+// Execute various samples showing how to use a custom ChatMessageStore with an agent.
+await CustomChatMessageStore_UsingFactory_Async();
+await CustomChatMessageStore_UsingFactoryAndExistingExternalId_Async();
+await CustomChatMessageStore_PerThread_Async();
+await CustomChatMessageStore_PerRun_Async();
+
+// Here we can see how to create a custom ChatMessageStore using a factory method
+// provided to the agent via the ChatMessageStoreFactory option.
+// This allows us to use a custom chat message store, where the consumer of the agent
+// doesn't need to know anything about the storage mechanism used.
+async Task CustomChatMessageStore_UsingFactory_Async()
+{
+    Console.WriteLine("\n--- With Factory ---\n");
+
+    // Create the agent
+    AIAgent agent = new AzureOpenAIClient(
+        new Uri(endpoint),
+        new AzureCliCredential())
+        // Use a service that doesn't require storage of chat history in the service itself.
+        .GetChatClient(deploymentName)
+        .CreateAIAgent(new ChatClientAgentOptions
         {
-            // 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.
-            return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions);
-        }
-    });
+            ChatOptions = new() { Instructions = "You are good at telling jokes." },
+            Name = "Joker",
+            ChatMessageStoreFactory = ctx =>
+            {
+                // 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.
+                return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions, ctx.Features);
+            }
+        });
+
+    // Start a new thread for the agent conversation.
+    AgentThread thread = agent.GetNewThread();
+
+    // Run the agent with the thread that stores conversation history in the vector store.
+    Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
+
+    // Serialize the thread state, so it can be stored for later use.
+    // Since the chat history is stored in the vector store, the serialized thread
+    // only contains the guid that the messages are stored under in the vector store.
+    JsonElement serializedThread = thread.Serialize();
+
+    Console.WriteLine("\n--- Serialized thread ---\n");
+    Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
+
+    // The serialized thread can now be saved to a database, file, or any other storage mechanism
+    // and loaded again later.
+
+    // Deserialize the thread state after loading from storage.
+    AgentThread resumedThread = agent.DeserializeThread(serializedThread);
+
+    // Run the agent with the thread that stores conversation history in the vector store a second time.
+    Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
+}
+
+// Here we can see how to create a custom ChatMessageStore using a factory method
+// provided to the agent via the ChatMessageStoreFactory option.
+// It also shows how we can pass a custom storage id at runtime to the message store using
+// the VectorChatMessageStoreThreadDbKeyFeature.
+// Note that not all agents or chat message stores may support this feature.
+async Task CustomChatMessageStore_UsingFactoryAndExistingExternalId_Async()
+{
+    Console.WriteLine("\n--- With Factory and Existing External ID ---\n");
+
+    // Create the agent
+    AIAgent agent = new AzureOpenAIClient(
+        new Uri(endpoint),
+        new AzureCliCredential())
+        // Use a service that doesn't require storage of chat history in the service itself.
+        .GetChatClient(deploymentName)
+        .CreateAIAgent(new ChatClientAgentOptions
+        {
+            ChatOptions = new() { Instructions = "You are good at telling jokes." },
+            Name = "Joker",
+            ChatMessageStoreFactory = ctx =>
+            {
+                // 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.
+                return new VectorChatMessageStore(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions, ctx.Features);
+            }
+        });
+
+    // Start a new thread for the agent conversation.
+    AgentThread thread = agent.GetNewThread();
+
+    // Run the agent with the thread that stores conversation history in the vector store.
+    Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
+
+    // We can access the VectorChatMessageStore via the thread's GetService method if we need to read the key under which threads are stored.
+    var messageStoreFromFactory = thread.GetService<VectorChatMessageStore>()!;
+    Console.WriteLine($"\nThread is stored in vector store under key: {messageStoreFromFactory.ThreadDbKey}");
+
+    // It's possible to create a new thread that uses the same chat message store id by providing
+    // the VectorChatMessageStoreThreadDbKeyFeature in the feature collection when creating the new thread.
+    AgentThread resumedThread = agent.GetNewThread(
+        new AgentFeatureCollection().WithFeature(new VectorChatMessageStoreThreadDbKeyFeature(messageStoreFromFactory.ThreadDbKey!)));
 
-// Start a new thread for the agent conversation.
-AgentThread thread = agent.GetNewThread();
+    // Run the agent with the thread that stores conversation history in the vector store.
+    Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
+}
 
-// Run the agent with the thread that stores conversation history in the vector store.
-Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
+// Here we can see how to create a custom ChatMessageStore and pass it to the thread
+// when creating a new thread.
+async Task CustomChatMessageStore_PerThread_Async()
+{
+    Console.WriteLine("\n--- Per Thread ---\n");
 
-// Serialize the thread state, so it can be stored for later use.
-// Since the chat history is stored in the vector store, the serialized thread
-// only contains the guid that the messages are stored under in the vector store.
-JsonElement serializedThread = thread.Serialize();
+    // We can also create an agent without a factory that provides a ChatMessageStore.
+    AIAgent agent = new AzureOpenAIClient(
+        new Uri(endpoint),
+        new AzureCliCredential())
+        // Use a service that doesn't require storage of chat history in the service itself.
+        .GetChatClient(deploymentName)
+        .CreateAIAgent(new ChatClientAgentOptions
+        {
+            ChatOptions = new() { Instructions = "You are good at telling jokes." },
+            Name = "Joker"
+        });
 
-Console.WriteLine("\n--- Serialized thread ---\n");
-Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
+    // Instead of using a factory on the agent to create the ChatMessageStore, we can
+    // create a VectorChatMessageStore ourselves and register it in a feature collection.
+    // We can then pass the feature collection when creating a new thread.
+    // We also have the opportunity here to pass any id that we want for storing the chat history in the vector store.
+    VectorChatMessageStore perThreadMessageStore = new(vectorStore, "chat-history-1");
+    AgentThread thread = agent.GetNewThread(new AgentFeatureCollection().WithFeature<ChatMessageStore>(perThreadMessageStore));
 
-// The serialized thread can now be saved to a database, file, or any other storage mechanism
-// and loaded again later.
+    Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", thread));
 
-// Deserialize the thread state after loading from storage.
-AgentThread resumedThread = agent.DeserializeThread(serializedThread);
+    // When serializing this thread, we'll see that it has the id from the message store stored in its state.
+    JsonElement serializedThread = thread.Serialize();
 
-// Run the agent with the thread that stores conversation history in the vector store a second time.
-Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
+    Console.WriteLine("\n--- Serialized thread ---\n");
+    Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
+}
 
-// We can access the VectorChatMessageStore via the thread's GetService method if we need to read the key under which threads are stored.
-var messageStore = resumedThread.GetService<VectorChatMessageStore>()!;
-Console.WriteLine($"\nThread is stored in vector store under key: {messageStore.ThreadDbKey}");
+// Here we can see how to create a custom ChatMessageStore for a single run using the Features option
+// passed when we run the agent.
+// Note that if the agent doesn't support a chat message store, it would be ignored.
+async Task CustomChatMessageStore_PerRun_Async()
+{
+    Console.WriteLine("\n--- Per Run ---\n");
+
+    // We can also create an agent without a factory that provides a ChatMessageStore.
+    AIAgent agent = new AzureOpenAIClient(
+        new Uri(endpoint),
+        new AzureCliCredential())
+        // Use a service that doesn't require storage of chat history in the service itself.
+        .GetChatClient(deploymentName)
+        .CreateAIAgent(new ChatClientAgentOptions
+        {
+            ChatOptions = new() { Instructions = "You are good at telling jokes." },
+            Name = "Joker"
+        });
+
+    // Start a new thread for the agent conversation.
+    AgentThread thread = agent.GetNewThread();
+
+    // Instead of using a factory on the agent to create the ChatMessageStore, we can
+    // create a VectorChatMessageStore ourselves and register it in a feature collection.
+    // We can then pass the feature collection to the agent when running it by using the Features option.
+    // The message store would only be used for the run that it's passed to.
+    // If the agent doesn't support a message store, it would be ignored.
+    // We also have the opportunity here to pass any id that we want for storing the chat history in the vector store.
+    VectorChatMessageStore perRunMessageStore = new(vectorStore, "chat-history-1");
+    Console.WriteLine(await agent.RunAsync(
+        "Tell me a joke about a pirate.",
+        thread,
+        options: new AgentRunOptions()
+        {
+            Features = new AgentFeatureCollection().WithFeature<ChatMessageStore>(perRunMessageStore)
+        }));
+
+    // When serializing this thread, we'll see that it has no messagestore state, since the messagestore was not attached to the thread,
+    // but just provided for the single run. Note that, depending on the circumstances, the thread may still contain other state, e.g. Memories,
+    // if an AIContextProvider is attached which adds memory to an agent.
+    JsonElement serializedThread = thread.Serialize();
+
+    Console.WriteLine("\n--- Serialized thread ---\n");
+    Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerOptions { WriteIndented = true }));
+}
 
 namespace SampleApp
 {
+    /// <summary>
+    /// A feature that allows providing the thread database key for the <see cref="VectorChatMessageStore"/>.
+    /// </summary>
+    internal sealed class VectorChatMessageStoreThreadDbKeyFeature(string threadDbKey)
+    {
+        public string ThreadDbKey { get; } = threadDbKey;
+    }
+
     /// <summary>
     /// A sample implementation of <see cref="ChatMessageStore"/> that stores chat messages in a vector store.
     /// </summary>
     internal sealed class VectorChatMessageStore : ChatMessageStore
     {
         private readonly VectorStore _vectorStore;
 
-        public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
+        public VectorChatMessageStore(VectorStore vectorStore, string threadDbKey)
+        {
+            this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
+            this.ThreadDbKey = threadDbKey ?? throw new ArgumentNullException(nameof(threadDbKey));
+        }
+
+        public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? features = null)
         {
             this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
 
-            if (serializedStoreState.ValueKind is JsonValueKind.String)
-            {
-                // Here we can deserialize the thread id so that we can access the same messages as before the suspension.
-                this.ThreadDbKey = serializedStoreState.Deserialize<string>();
-            }
+            // Here we can deserialize the thread id so that we can access the same messages as before the suspension, or if
+            // a user provided a ConversationIdAgentFeature in the features collection, we can use that
+            // or finally we can generate one ourselves.
+            this.ThreadDbKey = serializedStoreState.ValueKind is JsonValueKind.String
+                ? serializedStoreState.Deserialize<string>()
+                : features?.TryGet<VectorChatMessageStoreThreadDbKeyFeature>(out var threadDbKeyFeature) is true
+                    ? threadDbKeyFeature.ThreadDbKey
+                    : Guid.NewGuid().ToString("N");
         }
 
-        public string? ThreadDbKey { get; private set; }
+        public string? ThreadDbKey { get; }
 
         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,
+                Key = this.ThreadDbKey + (string.IsNullOrWhiteSpace(x.MessageId) ? Guid.NewGuid().ToString("N") : x.MessageId),
                 Timestamp = DateTimeOffset.UtcNow,
                 ThreadId = this.ThreadDbKey,
                 SerializedMessage = JsonSerializer.Serialize(x),

dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs

@@ -52,8 +52,13 @@ public A2AAgent(A2AClient a2aClient, string? id = null, string? name = null, str
     }
 
     /// <inheritdoc/>
-    public sealed override AgentThread GetNewThread()
-        => new A2AAgentThread();
+    public sealed override AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null)
+        => new A2AAgentThread()
+        {
+            ContextId = featureCollection?.TryGet<ConversationIdAgentFeature>(out var conversationIdFeature) is true
+                ? conversationIdFeature.ConversationId
+                : null
+        };
 
     /// <summary>
     /// Get a new <see cref="AgentThread"/> instance using an existing context id, to continue that conversation.
@@ -64,7 +69,7 @@ public AgentThread GetNewThread(string contextId)
         => new A2AAgentThread() { ContextId = contextId };
 
     /// <inheritdoc/>
-    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
+    public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null)
         => new A2AAgentThread(serializedThread, jsonSerializerOptions);
 
     /// <inheritdoc/>

dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgent.cs

@@ -105,6 +105,7 @@ public abstract class AIAgent
     /// <summary>
     /// Creates a new conversation thread that is compatible with this agent.
     /// </summary>
+    /// <param name="featureCollection">An optional feature collection to override or provide additional context or capabilities to the thread where the thread supports these features.</param>
     /// <returns>A new <see cref="AgentThread"/> instance ready for use with this agent.</returns>
     /// <remarks>
     /// <para>
@@ -118,13 +119,14 @@ public abstract class AIAgent
     /// may be deferred until first use to optimize performance.
     /// </para>
     /// </remarks>
-    public abstract AgentThread GetNewThread();
+    public abstract AgentThread GetNewThread(IAgentFeatureCollection? featureCollection = null);
 
     /// <summary>
     /// Deserializes an agent thread from its JSON serialized representation.
     /// </summary>
     /// <param name="serializedThread">A <see cref="JsonElement"/> containing the serialized thread state.</param>
     /// <param name="jsonSerializerOptions">Optional settings to customize the deserialization process.</param>
+    /// <param name="featureCollection">An optional feature collection to override or provide additional context or capabilities to the thread where the thread supports these features.</param>
     /// <returns>A restored <see cref="AgentThread"/> instance with the state from <paramref name="serializedThread"/>.</returns>
     /// <exception cref="ArgumentException">The <paramref name="serializedThread"/> is not in the expected format.</exception>
     /// <exception cref="JsonException">The serialized data is invalid or cannot be deserialized.</exception>
@@ -133,7 +135,7 @@ public abstract class AIAgent
     /// allowing conversations to resume across application restarts or be migrated between
     /// different agent instances.
     /// </remarks>
-    public abstract AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null);
+    public abstract AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, IAgentFeatureCollection? featureCollection = null);
 
     /// <summary>
     /// Run the agent with no message assuming that all required instructions are already provided to the agent or on the thread.