.NET: [BREAKING] Add session StateBag for state storage and support multiple providers on the Agent

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

@@ -6,6 +6,7 @@
 
 using System.Runtime.CompilerServices;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 using Microsoft.Agents.AI;
 using Microsoft.Extensions.AI;
 using SampleApp;
@@ -28,6 +29,8 @@ internal sealed class UpperCaseParrotAgent : AIAgent
     {
         public override string? Name => "UpperCaseParrotAgent";
 
+        public readonly ChatHistoryProvider ChatHistoryProvider = new InMemoryChatHistoryProvider();
+
         protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default)
             => new(new CustomAgentSession());
 
@@ -38,11 +41,11 @@ protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession
                 throw new ArgumentException($"The provided session is not of type {nameof(CustomAgentSession)}.", nameof(session));
             }
 
-            return new(typedSession.Serialize(jsonSerializerOptions));
+            return new(JsonSerializer.SerializeToElement(typedSession, jsonSerializerOptions));
         }
 
         protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
-            => new(new CustomAgentSession(serializedState, jsonSerializerOptions));
+            => new(serializedState.Deserialize<CustomAgentSession>(jsonSerializerOptions)!);
 
         protected override async Task<AgentResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
         {
@@ -56,17 +59,14 @@ protected override async Task<AgentResponse> RunCoreAsync(IEnumerable<ChatMessag
 
             // Get existing messages from the store
             var invokingContext = new ChatHistoryProvider.InvokingContext(this, session, messages);
-            var storeMessages = await typedSession.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken);
+            var userAndChatHistoryMessages = await this.ChatHistoryProvider.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 session of the input and output messages.
-            var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, messages)
-            {
-                ResponseMessages = responseMessages
-            };
-            await typedSession.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken);
+            var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, userAndChatHistoryMessages, responseMessages);
+            await this.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken);
 
             return new AgentResponse
             {
@@ -88,17 +88,14 @@ protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingA
 
             // Get existing messages from the store
             var invokingContext = new ChatHistoryProvider.InvokingContext(this, session, messages);
-            var storeMessages = await typedSession.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken);
+            var userAndChatHistoryMessages = await this.ChatHistoryProvider.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 session of the input and output messages.
-            var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, messages)
-            {
-                ResponseMessages = responseMessages
-            };
-            await typedSession.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken);
+            var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, userAndChatHistoryMessages, responseMessages);
+            await this.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken);
 
             foreach (var message in responseMessages)
             {
@@ -140,15 +137,16 @@ private static IEnumerable<ChatMessage> CloneAndToUpperCase(IEnumerable<ChatMess
         /// <summary>
         /// A session type for our custom agent that only supports in memory storage of messages.
         /// </summary>
-        internal sealed class CustomAgentSession : InMemoryAgentSession
+        internal sealed class CustomAgentSession : AgentSession
         {
-            internal CustomAgentSession() { }
-
-            internal CustomAgentSession(JsonElement serializedSessionState, JsonSerializerOptions? jsonSerializerOptions = null)
-                : base(serializedSessionState, jsonSerializerOptions) { }
+            internal CustomAgentSession()
+            {
+            }
 
-            internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-                => base.Serialize(jsonSerializerOptions);
+            [JsonConstructor]
+            internal CustomAgentSession(AgentSessionStateBag stateBag) : base(stateBag)
+            {
+            }
         }
     }
 }

dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step01_ChatHistoryMemory/Program.cs

@@ -37,16 +37,21 @@
     {
         ChatOptions = new() { Instructions = "You are good at telling jokes." },
         Name = "Joker",
-        AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new ChatHistoryMemoryProvider(
+        AIContextProviders = [new ChatHistoryMemoryProvider(
             vectorStore,
             collectionName: "chathistory",
             vectorDimensions: 3072,
-            // Configure the scope values under which chat messages will be stored.
-            // In this case, we are using a fixed user ID and a unique session ID for each new session.
-            storageScope: new() { UserId = "UID1", SessionId = Guid.NewGuid().ToString() },
-            // Configure the scope which would be used to search for relevant prior messages.
-            // In this case, we are searching for any messages for the user across all sessions.
-            searchScope: new() { UserId = "UID1" }))
+            // Callback to configure the initial state of the ChatHistoryMemoryProvider.
+            // The ChatHistoryMemoryProvider stores its state in the AgentSession and this callback
+            // will be called whenever the ChatHistoryMemoryProvider cannot find existing state in the session,
+            // typically the first time it is used with a new session.
+            session => new ChatHistoryMemoryProvider.State(
+                // Configure the scope values under which chat messages will be stored.
+                // In this case, we are using a fixed user ID and a unique session ID for each new session.
+                storageScope: new() { UserId = "UID1", SessionId = Guid.NewGuid().ToString() },
+                // Configure the scope which would be used to search for relevant prior messages.
+                // In this case, we are searching for any messages for the user across all sessions.
+                searchScope: new() { UserId = "UID1" }))]
     });
 
 // Start a new session for the agent conversation.

dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step02_MemoryUsingMem0/Program.cs

@@ -34,20 +34,21 @@
     .AsAIAgent(new ChatClientAgentOptions()
     {
         ChatOptions = new() { Instructions = "You are a friendly travel assistant. Use known memories about the user when responding, and do not invent details." },
-        AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(ctx.SerializedState.ValueKind is not JsonValueKind.Null and not JsonValueKind.Undefined
-            // If each session should have its own Mem0 scope, you can create a new id per session here:
-            // ? new Mem0Provider(mem0HttpClient, new Mem0ProviderScope() { ThreadId = Guid.NewGuid().ToString() })
-            // In this case we are storing memories scoped by application and user instead so that memories are retained across threads.
-            ? new Mem0Provider(mem0HttpClient, new Mem0ProviderScope() { ApplicationId = "getting-started-agents", UserId = "sample-user" })
-            // For cases where we are restoring from serialized state:
-            : new Mem0Provider(mem0HttpClient, ctx.SerializedState, ctx.JsonSerializerOptions))
+        // The stateInitializer can be used to customize the Mem0 scope per session and it will be called each time a session
+        // is encountered by the Mem0Provider that does not already have Mem0Provider state stored on the session.
+        // If each session should have its own Mem0 scope, you can create a new id per session via the stateInitializer, e.g.:
+        // new Mem0Provider(mem0HttpClient, stateInitializer: _ => new(new Mem0ProviderScope() { ThreadId = Guid.NewGuid().ToString() }))
+        // In our case we are storing memories scoped by application and user instead so that memories are retained across threads.
+        AIContextProviders = [new Mem0Provider(mem0HttpClient, stateInitializer: _ => new(new Mem0ProviderScope() { ApplicationId = "getting-started-agents", UserId = "sample-user" }))]
     });
 
 AgentSession session = await agent.CreateSessionAsync();
 
 // Clear any existing memories for this scope to demonstrate fresh behavior.
-Mem0Provider mem0Provider = session.GetService<Mem0Provider>()!;
-await mem0Provider.ClearStoredMemoriesAsync();
+// Note that the ClearStoredMemoriesAsync method will clear memories
+// using the scope stored in the session, or provided via the stateInitializer.
+Mem0Provider mem0Provider = agent.GetService<Mem0Provider>()!;
+await mem0Provider.ClearStoredMemoriesAsync(session);
 
 Console.WriteLine(await agent.RunAsync("Hi there! My name is Taylor and I'm planning a hiking trip to Patagonia in November.", session));
 Console.WriteLine(await agent.RunAsync("I'm travelling with my sister and we love finding scenic viewpoints.", session));

dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs

@@ -36,7 +36,7 @@
 AIAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions()
 {
     ChatOptions = new() { Instructions = "You are a friendly assistant. Always address the user by their name." },
-    AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new UserInfoMemory(chatClient.AsIChatClient(), ctx.SerializedState, ctx.JsonSerializerOptions))
+    AIContextProviders = [new UserInfoMemory(chatClient.AsIChatClient())]
 });
 
 // Create a new session for the conversation.
@@ -58,23 +58,23 @@
 var deserializedSession = await agent.DeserializeSessionAsync(sesionElement);
 Console.WriteLine(await agent.RunAsync("What is my name and age?", deserializedSession));
 
-Console.WriteLine("\n>> Read memories from memory component\n");
+Console.WriteLine("\n>> Read memories using memory component\n");
 
-// It's possible to access the memory component via the session's GetService method.
-var userInfo = deserializedSession.GetService<UserInfoMemory>()?.UserInfo;
+// It's possible to access the memory component via the agent's GetService method.
+var userInfo = agent.GetService<UserInfoMemory>()?.GetUserInfo(deserializedSession);
 
 // Output the user info that was captured by the memory component.
 Console.WriteLine($"MEMORY - User Name: {userInfo?.UserName}");
 Console.WriteLine($"MEMORY - User Age: {userInfo?.UserAge}");
 
 Console.WriteLine("\n>> Use new session with previously created memories\n");
 
-// It is also possible to set the memories in a memory component on an individual session.
+// It is also possible to set the memories using a memory component on an individual session.
 // This is useful if we want to start a new session, but have it share the same memories as a previous session.
 var newSession = await agent.CreateSessionAsync();
-if (userInfo is not null && newSession.GetService<UserInfoMemory>() is UserInfoMemory newSessionMemory)
+if (userInfo is not null && agent.GetService<UserInfoMemory>() is UserInfoMemory newSessionMemory)
 {
-    newSessionMemory.UserInfo = userInfo;
+    newSessionMemory.SetUserInfo(newSession, userInfo);
 }
 
 // Invoke the agent and output the text result.
@@ -89,28 +89,27 @@ namespace SampleApp
     internal sealed class UserInfoMemory : AIContextProvider
     {
         private readonly IChatClient _chatClient;
+        private readonly Func<AgentSession?, UserInfo> _stateInitializer;
 
-        public UserInfoMemory(IChatClient chatClient, UserInfo? userInfo = null)
+        public UserInfoMemory(IChatClient chatClient, Func<AgentSession?, UserInfo>? stateInitializer = null)
         {
             this._chatClient = chatClient;
-            this.UserInfo = userInfo ?? new UserInfo();
+            this._stateInitializer = stateInitializer ?? (_ => new UserInfo());
         }
 
-        public UserInfoMemory(IChatClient chatClient, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
-        {
-            this._chatClient = chatClient;
-
-            this.UserInfo = serializedState.ValueKind == JsonValueKind.Object ?
-                serializedState.Deserialize<UserInfo>(jsonSerializerOptions)! :
-                new UserInfo();
-        }
+        public UserInfo GetUserInfo(AgentSession session)
+            => session.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory)) ?? new UserInfo();
 
-        public UserInfo UserInfo { get; set; }
+        public void SetUserInfo(AgentSession session, UserInfo userInfo)
+            => session.StateBag.SetValue(nameof(UserInfoMemory), userInfo);
 
         protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
         {
+            var userInfo = context.Session?.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory))
+                ?? this._stateInitializer.Invoke(context.Session);
+
             // Try and extract the user name and age from the message if we don't have it already and it's a user message.
-            if ((this.UserInfo.UserName is null || this.UserInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User))
+            if ((userInfo.UserName is null || userInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User))
             {
                 var result = await this._chatClient.GetResponseAsync<UserInfo>(
                     context.RequestMessages,
@@ -120,36 +119,43 @@ protected override async ValueTask InvokedCoreAsync(InvokedContext context, Canc
                     },
                     cancellationToken: cancellationToken);
 
-                this.UserInfo.UserName ??= result.Result.UserName;
-                this.UserInfo.UserAge ??= result.Result.UserAge;
+                userInfo.UserName ??= result.Result.UserName;
+                userInfo.UserAge ??= result.Result.UserAge;
             }
+
+            context.Session?.StateBag.SetValue(nameof(UserInfoMemory), userInfo);
         }
 
         protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
         {
+            var inputContext = context.AIContext;
+            var userInfo = context.Session?.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory))
+                ?? this._stateInitializer.Invoke(context.Session);
+
             StringBuilder instructions = new();
+            if (!string.IsNullOrEmpty(inputContext.Instructions))
+            {
+                instructions.AppendLine(inputContext.Instructions);
+            }
 
             // If we don't already know the user's name and age, add instructions to ask for them, otherwise just provide what we have to the context.
             instructions
                 .AppendLine(
-                    this.UserInfo.UserName is null ?
+                    userInfo.UserName is null ?
                         "Ask the user for their name and politely decline to answer any questions until they provide it." :
-                        $"The user's name is {this.UserInfo.UserName}.")
+                        $"The user's name is {userInfo.UserName}.")
                 .AppendLine(
-                    this.UserInfo.UserAge is null ?
+                    userInfo.UserAge is null ?
                         "Ask the user for their age and politely decline to answer any questions until they provide it." :
-                        $"The user's age is {this.UserInfo.UserAge}.");
+                        $"The user's age is {userInfo.UserAge}.");
 
             return new ValueTask<AIContext>(new AIContext
             {
-                Instructions = instructions.ToString()
+                Instructions = instructions.ToString(),
+                Messages = inputContext.Messages,
+                Tools = inputContext.Tools
             });
         }
-
-        public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
-        {
-            return JsonSerializer.SerializeToElement(this.UserInfo, jsonSerializerOptions);
-        }
     }
 
     internal sealed class UserInfo

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

@@ -65,12 +65,16 @@
     .AsAIAgent(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, ct) => new ValueTask<AIContextProvider>(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
+        AIContextProviders = [new TextSearchProvider(SearchAdapter, textSearchOptions)],
+        // Since we are using ChatCompletion which stores chat history locally, we can also add a message filter
         // 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.
-        ChatHistoryProviderFactory = (ctx, ct) => new ValueTask<ChatHistoryProvider>(new InMemoryChatHistoryProvider(ctx.SerializedState, ctx.JsonSerializerOptions)
-            .WithAIContextProviderMessageRemoval()),
+        // By default the chat history provider will store all messages, except for those that came from chat history in the first place.
+        // We also want to maintain that exclusion here.
+        ChatHistoryProvider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions
+        {
+            StorageInputMessageFilter = messages => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.AIContextProvider && m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory)
+        }),
     });
 
 AgentSession session = await agent.CreateSessionAsync();

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

@@ -74,7 +74,7 @@
     .AsAIAgent(new ChatClientAgentOptions
     {
         ChatOptions = new() { Instructions = "You are a helpful support specialist for the Microsoft Agent Framework. Answer questions using the provided context and cite the source document when available. Keep responses brief." },
-        AIContextProviderFactory = (ctx, ct) => new ValueTask<AIContextProvider>(new TextSearchProvider(SearchAdapter, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions))
+        AIContextProviders = [new TextSearchProvider(SearchAdapter, textSearchOptions)]
     });
 
 AgentSession session = await agent.CreateSessionAsync();

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

@@ -32,7 +32,7 @@
     .AsAIAgent(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, ct) => new ValueTask<AIContextProvider>(new TextSearchProvider(MockSearchAsync, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions))
+        AIContextProviders = [new TextSearchProvider(MockSearchAsync, textSearchOptions)]
     });
 
 AgentSession session = await agent.CreateSessionAsync();