.NET: [BREAKING] Add GetService for AIContextProviders and ChatMessageStore (#874)

* Add GetService for AIContextProviders and ChatMessageStore

* Change styling and fix format issues

* Update sample code to allow for missing memory component.

dotnet/src/Microsoft.Agents.Workflows/WorkflowMessageStore.cs

@@ -10,7 +10,7 @@ using Microsoft.Extensions.AI.Agents;
 
 namespace Microsoft.Agents.Workflows;
 
-internal sealed class WorkflowMessageStore : IChatMessageStore
+internal sealed class WorkflowMessageStore : ChatMessageStore
 {
     private int _bookmark;
     private readonly List<ChatMessage> _chatMessages = [];
@@ -46,14 +46,14 @@ internal sealed class WorkflowMessageStore : IChatMessageStore
 
     internal void AddMessages(params ChatMessage[] messages) => this._chatMessages.AddRange(messages);
 
-    public Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
+    public override Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
     {
         this._chatMessages.AddRange(messages);
 
         return Task.CompletedTask;
     }
 
-    public Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken) => Task.FromResult<IEnumerable<ChatMessage>>(this._chatMessages.AsReadOnly());
+    public override Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken) => Task.FromResult<IEnumerable<ChatMessage>>(this._chatMessages.AsReadOnly());
 
     public IEnumerable<ChatMessage> GetFromBookmark()
     {
@@ -65,7 +65,7 @@ internal sealed class WorkflowMessageStore : IChatMessageStore
 
     public void UpdateBookmark() => this._bookmark = this._chatMessages.Count;
 
-    public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    public override ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
     {
         StoreState state = new()
         {

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/AIContextProvider.cs

@@ -5,6 +5,7 @@ using System.Collections.Generic;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Extensions.AI.Agents;
 
@@ -52,18 +53,35 @@ public abstract class AIContextProvider
         return default;
     }
 
-    /// <summary>
-    /// Deserializes the state contained in the provided <see cref="JsonElement"/> into the properties on this object.
-    /// </summary>
-    /// <param name="serializedState">A <see cref="JsonElement"/> representing the state of the object.</param>
-    /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A <see cref="ValueTask"/> that completes when the state has been deserialized.</returns>
-    public virtual ValueTask DeserializeAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    /// <summary>Asks the <see cref="AIContextProvider"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
+    /// <param name="serviceType">The type of object being requested.</param>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="serviceType"/> is <see langword="null"/>.</exception>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly-typed services that might be provided by the <see cref="AIContextProvider"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public virtual object? GetService(Type serviceType, object? serviceKey = null)
     {
-        return default;
+        _ = Throw.IfNull(serviceType);
+
+        return serviceKey is null && serviceType.IsInstanceOfType(this)
+            ? this
+            : null;
     }
 
+    /// <summary>Asks the <see cref="AIContextProvider"/> for an object of type <typeparamref name="TService"/>.</summary>
+    /// <typeparam name="TService">The type of the object to be retrieved.</typeparam>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly typed services that may be provided by the <see cref="AIContextProvider"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public TService? GetService<TService>(object? serviceKey = null)
+        => this.GetService(typeof(TService), serviceKey) is TService service ? service : default;
+
     /// <summary>
     /// Contains the event context provided to <see cref="InvokingAsync(InvokingContext, CancellationToken)"/>.
     /// </summary>

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/ChatMessageStore.cs

@@ -0,0 +1,85 @@
+// 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.Shared.Diagnostics;
+
+namespace Microsoft.Extensions.AI.Agents;
+
+/// <summary>
+/// Defines methods for storing and retrieving chat messages associated with a specific thread.
+/// </summary>
+/// <remarks>
+/// Implementations of this interface are responsible for managing the storage of chat messages,
+/// including handling large volumes of data by truncating or summarizing messages as necessary.
+/// </remarks>
+public abstract class ChatMessageStore
+{
+    /// <summary>
+    /// Gets all the messages from the store that should be used for the next agent invocation.
+    /// </summary>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>A collection of chat messages.</returns>
+    /// <remarks>
+    /// <para>
+    /// Messages are returned in ascending chronological order, with the oldest message first.
+    /// </para>
+    /// <para>
+    /// If the messages stored in the store become very large, it is up to the store to
+    /// truncate, summarize or otherwise limit the number of messages returned.
+    /// </para>
+    /// <para>
+    /// When using implementations of <see cref="ChatMessageStore"/>, a new one should be created for each thread
+    /// since they may contain state that is specific to a thread.
+    /// </para>
+    /// </remarks>
+    public abstract Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Adds messages to the store.
+    /// </summary>
+    /// <param name="messages">The messages to add.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>An async task.</returns>
+    public abstract Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default);
+
+    /// <summary>
+    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
+    /// </summary>
+    /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
+    /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
+    public abstract ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default);
+
+    /// <summary>Asks the <see cref="ChatMessageStore"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
+    /// <param name="serviceType">The type of object being requested.</param>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="serviceType"/> is <see langword="null"/>.</exception>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly-typed services that might be provided by the <see cref="ChatMessageStore"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public virtual object? GetService(Type serviceType, object? serviceKey = null)
+    {
+        _ = Throw.IfNull(serviceType);
+
+        return serviceKey is null && serviceType.IsInstanceOfType(this)
+            ? this
+            : null;
+    }
+
+    /// <summary>Asks the <see cref="ChatMessageStore"/> for an object of type <typeparamref name="TService"/>.</summary>
+    /// <typeparam name="TService">The type of the object to be retrieved.</typeparam>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly typed services that may be provided by the <see cref="ChatMessageStore"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public TService? GetService<TService>(object? serviceKey = null)
+        => this.GetService(typeof(TService), serviceKey) is TService service ? service : default;
+}

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/IChatMessageStore.cs

@@ -1,54 +0,0 @@
-// Copyright (c) Microsoft. All rights reserved.
-
-using System.Collections.Generic;
-using System.Text.Json;
-using System.Threading;
-using System.Threading.Tasks;
-
-namespace Microsoft.Extensions.AI.Agents;
-
-/// <summary>
-/// Defines methods for storing and retrieving chat messages associated with a specific thread.
-/// </summary>
-/// <remarks>
-/// Implementations of this interface are responsible for managing the storage of chat messages,
-/// including handling large volumes of data by truncating or summarizing messages as necessary.
-/// </remarks>
-public interface IChatMessageStore
-{
-    /// <summary>
-    /// Gets all the messages from the store that should be used for the next agent invocation.
-    /// </summary>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A collection of chat messages.</returns>
-    /// <remarks>
-    /// <para>
-    /// Messages are returned in ascending chronological order, with the oldest message first.
-    /// </para>
-    /// <para>
-    /// If the messages stored in the store become very large, it is up to the store to
-    /// truncate, summarize or otherwise limit the number of messages returned.
-    /// </para>
-    /// <para>
-    /// When using implementations of <see cref="IChatMessageStore"/>, a new one should be created for each thread
-    /// since they may contain state that is specific to a thread.
-    /// </para>
-    /// </remarks>
-    Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken = default);
-
-    /// <summary>
-    /// Adds messages to the store.
-    /// </summary>
-    /// <param name="messages">The messages to add.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>An async task.</returns>
-    Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default);
-
-    /// <summary>
-    /// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
-    /// </summary>
-    /// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
-    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
-    /// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
-    ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default);
-}

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/InMemoryAgentThread.cs

@@ -85,6 +85,10 @@ public abstract class InMemoryAgentThread : AgentThread
         return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(InMemoryAgentThreadState)));
     }
 
+    /// <inheritdoc/>
+    public override object? GetService(Type serviceType, object? serviceKey = null) =>
+        base.GetService(serviceType, serviceKey) ?? this.MessageStore?.GetService(serviceType, serviceKey);
+
     /// <inheritdoc />
     protected internal override Task MessagesReceivedAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken = default)
         => this.MessageStore.AddMessagesAsync(newMessages, cancellationToken);

dotnet/src/Microsoft.Extensions.AI.Agents.Abstractions/InMemoryChatMessageStore.cs

@@ -13,7 +13,7 @@ namespace Microsoft.Extensions.AI.Agents;
 /// <summary>
 /// Represents an in-memory store for chat messages associated with a specific thread.
 /// </summary>
-public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageStore
+public sealed class InMemoryChatMessageStore : ChatMessageStore, IList<ChatMessage>
 {
     private List<ChatMessage> _messages;
 
@@ -96,7 +96,7 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
     }
 
     /// <inheritdoc />
-    public async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
+    public override async Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken)
     {
         _ = Throw.IfNull(messages);
 
@@ -109,7 +109,7 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
     }
 
     /// <inheritdoc />
-    public async Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken)
+    public override async Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken)
     {
         if (this.ReducerTriggerEvent is ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null)
         {
@@ -120,7 +120,7 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
     }
 
     /// <inheritdoc />
-    public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
+    public override ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
     {
         StoreState state = new()
         {

dotnet/src/Microsoft.Extensions.AI.Agents/ChatCompletion/ChatClientAgent.cs

@@ -262,7 +262,7 @@ public sealed class ChatClientAgent : AIAgent
     /// <inheritdoc/>
     public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
     {
-        Func<JsonElement, JsonSerializerOptions?, IChatMessageStore>? chatMessageStoreFactory = this._agentOptions?.ChatMessageStoreFactory is null ?
+        Func<JsonElement, JsonSerializerOptions?, ChatMessageStore>? chatMessageStoreFactory = this._agentOptions?.ChatMessageStoreFactory is null ?
             null :
             (jse, jso) => this._agentOptions.ChatMessageStoreFactory.Invoke(new() { SerializedState = jse, JsonSerializerOptions = jso });
 

dotnet/src/Microsoft.Extensions.AI.Agents/ChatCompletion/ChatClientAgentOptions.cs

@@ -76,10 +76,10 @@ public class ChatClientAgentOptions
     public ChatOptions? ChatOptions { get; set; }
 
     /// <summary>
-    /// Gets or sets a factory function to create an instance of <see cref="IChatMessageStore"/>
+    /// Gets or sets a factory function to create an instance of <see cref="ChatMessageStore"/>
     /// which will be used to store chat messages for this agent.
     /// </summary>
-    public Func<ChatMessageStoreFactoryContext, IChatMessageStore>? ChatMessageStoreFactory { get; set; }
+    public Func<ChatMessageStoreFactoryContext, ChatMessageStore>? ChatMessageStoreFactory { get; set; }
 
     /// <summary>
     /// Gets or sets a factory function to create an instance of <see cref="AIContextProvider"/>
@@ -135,14 +135,14 @@ public class ChatClientAgentOptions
     }
 
     /// <summary>
-    /// Context object passed to the <see cref="ChatMessageStoreFactory"/> to create a new instance of <see cref="IChatMessageStore"/>.
+    /// Context object passed to the <see cref="ChatMessageStoreFactory"/> to create a new instance of <see cref="ChatMessageStore"/>.
     /// </summary>
     public class ChatMessageStoreFactoryContext
     {
         /// <summary>
         /// Gets or sets the serialized state of the chat message store, if any.
         /// </summary>
-        /// <value><see langword="default"/> if there is no state, e.g. when the <see cref="IChatMessageStore"/> is first created.</value>
+        /// <value><see langword="default"/> if there is no state, e.g. when the <see cref="ChatMessageStore"/> is first created.</value>
         public JsonElement SerializedState { get; set; }
 
         /// <summary>

dotnet/src/Microsoft.Extensions.AI.Agents/ChatCompletion/ChatClientAgentThread.cs

@@ -16,7 +16,7 @@ namespace Microsoft.Extensions.AI.Agents;
 public class ChatClientAgentThread : AgentThread
 {
     private string? _conversationId;
-    private IChatMessageStore? _messageStore;
+    private ChatMessageStore? _messageStore;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="ChatClientAgentThread"/> class.
@@ -30,12 +30,12 @@ public class ChatClientAgentThread : AgentThread
     /// </summary>
     /// <param name="serializedThreadState">A <see cref="JsonElement"/> representing the serialized state of the thread.</param>
     /// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
-    /// <param name="chatMessageStoreFactory">An optional factory function to create a custom <see cref="IChatMessageStore"/>.</param>
+    /// <param name="chatMessageStoreFactory">An optional factory function to create a custom <see cref="ChatMessageStore"/>.</param>
     /// <param name="aiContextProviderFactory">An optional factory function to create a custom <see cref="AIContextProvider"/>.</param>
     internal ChatClientAgentThread(
         JsonElement serializedThreadState,
         JsonSerializerOptions? jsonSerializerOptions = null,
-        Func<JsonElement, JsonSerializerOptions?, IChatMessageStore>? chatMessageStoreFactory = null,
+        Func<JsonElement, JsonSerializerOptions?, ChatMessageStore>? chatMessageStoreFactory = null,
         Func<JsonElement, JsonSerializerOptions?, AIContextProvider>? aiContextProviderFactory = null)
     {
         if (serializedThreadState.ValueKind != JsonValueKind.Object)
@@ -77,7 +77,7 @@ public class ChatClientAgentThread : AgentThread
     /// <para>
     /// This property may be null in the following cases:
     /// <list type="bullet">
-    /// <item>The thread stores messages via the <see cref="IChatMessageStore"/> and not in the agent service.</item>
+    /// <item>The thread stores messages via the <see cref="ChatMessageStore"/> and not in the agent service.</item>
     /// <item>This thread object is new and a server managed thread has not yet been created in the agent service.</item>
     /// </list>
     /// </para>
@@ -110,7 +110,7 @@ public class ChatClientAgentThread : AgentThread
     }
 
     /// <summary>
-    /// Gets or sets the <see cref="IChatMessageStore"/> used by this thread, for cases where messages should be stored in a custom location.
+    /// Gets or sets the <see cref="ChatMessageStore"/> used by this thread, for cases where messages should be stored in a custom location.
     /// </summary>
     /// <remarks>
     /// <para>
@@ -121,12 +121,12 @@ public class ChatClientAgentThread : AgentThread
     /// <para>
     /// This property may be null in the following cases:
     /// <list type="bullet">
-    /// <item>The thread stores messages in the agent service and just has an id to the remove thread, instead of in an <see cref="IChatMessageStore"/>.</item>
-    /// <item>This thread object is new it is not yet clear whether it will be backed by a server managed thread or an <see cref="IChatMessageStore"/>.</item>
+    /// <item>The thread stores messages in the agent service and just has an id to the remove thread, instead of in an <see cref="ChatMessageStore"/>.</item>
+    /// <item>This thread object is new it is not yet clear whether it will be backed by a server managed thread or an <see cref="ChatMessageStore"/>.</item>
     /// </list>
     /// </para>
     /// </remarks>
-    public IChatMessageStore? MessageStore
+    public ChatMessageStore? MessageStore
     {
         get => this._messageStore;
         internal set
@@ -179,12 +179,12 @@ public class ChatClientAgentThread : AgentThread
     }
 
     /// <inheritdoc/>
-    public override object? GetService(Type serviceType, object? serviceKey = null)
-    {
-        return serviceType == typeof(AgentThreadMetadata) ?
-            new AgentThreadMetadata(this.ConversationId) :
-            base.GetService(serviceType, serviceKey);
-    }
+    public override object? GetService(Type serviceType, object? serviceKey = null) =>
+        serviceType == typeof(AgentThreadMetadata)
+            ? new AgentThreadMetadata(this.ConversationId)
+            : base.GetService(serviceType, serviceKey)
+            ?? this.AIContextProvider?.GetService(serviceType, serviceKey)
+            ?? this.MessageStore?.GetService(serviceType, serviceKey);
 
     /// <inheritdoc />
     protected override async Task MessagesReceivedAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken = default)