diff --git a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs index bac72d9b31..9311e56c88 100644 --- a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs @@ -86,27 +86,25 @@ namespace SampleApp /// /// Sample memory component that can remember a user's name and age. /// - internal sealed class UserInfoMemory : AIContextProvider + internal sealed class UserInfoMemory : AIContextProvider { private readonly IChatClient _chatClient; - private readonly Func _stateInitializer; public UserInfoMemory(IChatClient chatClient, Func? stateInitializer = null) + : base(stateInitializer ?? (_ => new UserInfo()), null, null, null, null) { this._chatClient = chatClient; - this._stateInitializer = stateInitializer ?? (_ => new UserInfo()); } public UserInfo GetUserInfo(AgentSession session) - => session.StateBag.GetValue(nameof(UserInfoMemory)) ?? new UserInfo(); + => this.GetOrInitializeState(session); public void SetUserInfo(AgentSession session, UserInfo userInfo) - => session.StateBag.SetValue(nameof(UserInfoMemory), userInfo); + => this.SaveState(session, userInfo); - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { - var userInfo = context.Session?.StateBag.GetValue(nameof(UserInfoMemory)) - ?? this._stateInitializer.Invoke(context.Session); + var userInfo = this.GetOrInitializeState(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 ((userInfo.UserName is null || userInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User)) @@ -123,20 +121,14 @@ protected override async ValueTask InvokedCoreAsync(InvokedContext context, Canc userInfo.UserAge ??= result.Result.UserAge; } - context.Session?.StateBag.SetValue(nameof(UserInfoMemory), userInfo); + this.SaveState(context.Session, userInfo); } - protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { - var inputContext = context.AIContext; - var userInfo = context.Session?.StateBag.GetValue(nameof(UserInfoMemory)) - ?? this._stateInitializer.Invoke(context.Session); + var userInfo = this.GetOrInitializeState(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 @@ -151,9 +143,7 @@ userInfo.UserAge is null ? return new ValueTask(new AIContext { - Instructions = instructions.ToString(), - Messages = inputContext.Messages, - Tools = inputContext.Tools + Instructions = instructions.ToString() }); } } diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs index cc4ca1a6ed..1c25c36d22 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs @@ -76,45 +76,23 @@ namespace SampleApp /// State (the session DB key) is stored in the so it roundtrips /// automatically with session serialization. /// - internal sealed class VectorChatHistoryProvider : ChatHistoryProvider + internal sealed class VectorChatHistoryProvider : ChatHistoryProvider { private readonly VectorStore _vectorStore; - private readonly Func _stateInitializer; - private readonly string _stateKey; - - /// - public override string StateKey => this._stateKey; public VectorChatHistoryProvider( VectorStore vectorStore, Func? stateInitializer = null, string? stateKey = null) + : base(stateInitializer: stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N"))), stateKey: stateKey, jsonSerializerOptions: null, provideOutputMessageFilter: null, storeInputMessageFilter: null) { this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore)); - this._stateInitializer = stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N"))); - this._stateKey = stateKey ?? base.StateKey; } public string GetSessionDbKey(AgentSession session) => this.GetOrInitializeState(session).SessionDbKey; - private State GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state); - } - - return state; - } - - protected override async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) { var state = this.GetOrInitializeState(context.Session); var collection = this._vectorStore.GetCollection("ChatHistory"); @@ -129,29 +107,17 @@ protected override async ValueTask> InvokingCoreAsync(I var messages = records.ConvertAll(x => JsonSerializer.Deserialize(x.SerializedMessage!)!); messages.Reverse(); - return messages - .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) - .Concat(context.RequestMessages); + return messages; } - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - // Don't store messages if the request failed. - if (context.InvokeException is not null) - { - return; - } - var state = this.GetOrInitializeState(context.Session); var collection = this._vectorStore.GetCollection("ChatHistory"); await collection.EnsureCollectionExistsAsync(cancellationToken); - // Add both request and response messages to the store, excluding messages that came from chat history. - // Optionally messages produced by the AIContextProvider can also be persisted (not shown). - var allNewMessages = context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory) - .Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); await collection.UpsertAsync(allNewMessages.Select(x => new ChatHistoryItem() { diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs index b201e9f8e1..023f6c8e5f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -10,7 +11,7 @@ namespace Microsoft.Agents.AI; /// -/// Provides an abstract base class for components that enhance AI context management during agent invocations. +/// Provides an abstract base class for components that enhance AI context during agent invocations. /// /// /// @@ -30,6 +31,25 @@ namespace Microsoft.Agents.AI; /// public abstract class AIContextProvider { + private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) + => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); + + private readonly Func, IEnumerable> _provideInputMessageFilter; + private readonly Func, IEnumerable> _storeInputMessageFilter; + + /// + /// Initializes a new instance of the class. + /// + /// An optional filter function to apply to input messages before providing context via . If not set, defaults to including only messages. + /// An optional filter function to apply to request messages before storing context via . If not set, defaults to including only messages. + protected AIContextProvider( + Func, IEnumerable>? provideInputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + { + this._provideInputMessageFilter = provideInputMessageFilter ?? DefaultExternalOnlyFilter; + this._storeInputMessageFilter = storeInputMessageFilter ?? DefaultExternalOnlyFilter; + } + /// /// Gets the key used to store the provider state in the . /// @@ -58,7 +78,7 @@ public abstract class AIContextProvider /// /// public ValueTask InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) - => this.InvokingCoreAsync(context, cancellationToken); + => this.InvokingCoreAsync(Throw.IfNull(context), cancellationToken); /// /// Called at the start of agent invocation to provide additional context. @@ -76,8 +96,96 @@ public ValueTask InvokingAsync(InvokingContext context, CancellationT /// Injecting contextual messages from conversation history /// /// + /// + /// The default implementation of this method filters the input messages using the configured provide-input message filter + /// (which defaults to including only messages), + /// then calls to get additional context, + /// stamps any messages from the returned context with source attribution, + /// and merges the returned context with the original (unfiltered) input context (concatenating instructions, messages, and tools). + /// For most scenarios, overriding is sufficient to provide additional context, + /// while still benefiting from the default filtering, merging and source stamping behavior. + /// However, for scenarios that require more control over context filtering, merging or source stamping, overriding this method + /// allows you to directly control the full returned for the invocation. + /// /// - protected abstract ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default); + protected virtual async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + var inputContext = context.AIContext; + + // Create a filtered context for ProvideAIContextAsync, filtering input messages + // to exclude non-external messages (e.g. chat history, other AI context provider messages). + var filteredContext = new InvokingContext( + context.Agent, + context.Session, + new AIContext + { + Instructions = inputContext.Instructions, + Messages = inputContext.Messages is not null ? this._provideInputMessageFilter(inputContext.Messages) : null, + Tools = inputContext.Tools + }); + + var provided = await this.ProvideAIContextAsync(filteredContext, cancellationToken).ConfigureAwait(false); + + var mergedInstructions = (inputContext.Instructions, provided.Instructions) switch + { + (null, null) => null, + (string a, null) => a, + (null, string b) => b, + (string a, string b) => a + "\n" + b + }; + + var providedMessages = provided.Messages is not null + ? provided.Messages.Select(m => m.WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!)) + : null; + + var mergedMessages = (inputContext.Messages, providedMessages) switch + { + (null, null) => null, + (var a, null) => a, + (null, var b) => b, + (var a, var b) => a.Concat(b) + }; + + var mergedTools = (inputContext.Tools, provided.Tools) switch + { + (null, null) => null, + (var a, null) => a, + (null, var b) => b, + (var a, var b) => a.Concat(b) + }; + + return new AIContext + { + Instructions = mergedInstructions, + Messages = mergedMessages, + Tools = mergedTools + }; + } + + /// + /// When overridden in a derived class, provides additional AI context to be merged with the input context for the current invocation. + /// + /// + /// + /// This method is called from . + /// Note that can be overridden to directly control context merging and source stamping, in which case + /// it is up to the implementer to call this method as needed to retrieve the additional context. + /// + /// + /// In contrast with , this method only returns additional context to be merged with the input, + /// while is responsible for returning the full merged for the invocation. + /// + /// + /// Contains the request context including the caller provided messages that will be used by the agent for this invocation. + /// The to monitor for cancellation requests. The default is . + /// + /// A task that represents the asynchronous operation. The task result contains an + /// with additional context to be merged with the input context. + /// + protected virtual ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + return new ValueTask(new AIContext()); + } /// /// Called at the end of the agent invocation to process the invocation results. @@ -106,7 +214,7 @@ public ValueTask InvokingAsync(InvokingContext context, CancellationT /// /// public ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) - => this.InvokedCoreAsync(context, cancellationToken); + => this.InvokedCoreAsync(Throw.IfNull(context), cancellationToken); /// /// Called at the end of the agent invocation to process the invocation results. @@ -128,9 +236,50 @@ public ValueTask InvokedAsync(InvokedContext context, CancellationToken cancella /// This method is called regardless of whether the invocation succeeded or failed. /// To check if the invocation was successful, inspect the property. /// + /// + /// The default implementation of this method skips execution for any invocation failures, + /// filters the request messages using the configured store-input message filter + /// (which defaults to including only messages), + /// and calls to process the invocation results. + /// For most scenarios, overriding is sufficient to process invocation results, + /// while still benefiting from the default error handling and filtering behavior. + /// However, for scenarios that require more control over error handling or message filtering, overriding this method + /// allows you to directly control the processing of invocation results. + /// /// protected virtual ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) - => default; + { + if (context.InvokeException is not null) + { + return default; + } + + var subContext = new InvokedContext(context.Agent, context.Session, this._storeInputMessageFilter(context.RequestMessages), context.ResponseMessages!); + return this.StoreAIContextAsync(subContext, cancellationToken); + } + + /// + /// When overridden in a derived class, processes invocation results at the end of the agent invocation. + /// + /// Contains the invocation context including request messages, response messages, and any exception that occurred. + /// The to monitor for cancellation requests. The default is . + /// A task that represents the asynchronous operation. + /// + /// + /// This method is called from . + /// Note that can be overridden to directly control error handling, in which case + /// it is up to the implementer to call this method as needed to process the invocation results. + /// + /// + /// In contrast with , this method only processes the invocation results, + /// while is also responsible for error handling. + /// + /// + /// The default implementation of only calls this method if the invocation succeeded. + /// + /// + protected virtual ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) => + default; /// Asks the for an object of the specified type . /// The type of object being requested. diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs new file mode 100644 index 0000000000..136724f44b --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Text.Json; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Provides an abstract base class for components that enhance AI context during agent invocations with support for maintaining provider state of type . +/// +/// The type of the state to be maintained by the context provider. Must be a reference type. +/// +/// This class extends by introducing a strongly-typed state management mechanism, allowing derived classes to maintain and persist custom state information across invocations. +/// The state is stored in the session's StateBag using a configurable key and JSON serialization options, enabling seamless integration with the agent session lifecycle. +/// +public abstract class AIContextProvider : AIContextProvider + where TState : class +{ + private readonly Func _stateInitializer; + private readonly string _stateKey; + private readonly JsonSerializerOptions _jsonSerializerOptions; + + /// + /// Initializes a new instance of the class. + /// + /// A function to initialize the state for the context provider. + /// The key used to store the state in the session's StateBag. + /// Options for JSON serialization and deserialization of the state. + /// An optional filter function to apply to input messages before providing context. If not set, defaults to including only messages. + /// An optional filter function to apply to request messages before storing context. If not set, defaults to including only messages. + protected AIContextProvider( + Func stateInitializer, + string? stateKey, + JsonSerializerOptions? jsonSerializerOptions, + Func, IEnumerable>? provideInputMessageFilter, + Func, IEnumerable>? storeInputMessageFilter) + : base(provideInputMessageFilter, storeInputMessageFilter) + { + this._stateInitializer = stateInitializer; + this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + this._stateKey = stateKey ?? this.GetType().Name; + } + + /// + public override string StateKey => this._stateKey; + + /// + /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. + /// + /// The agent session containing the StateBag. + /// The provider state. + protected virtual TState GetOrInitializeState(AgentSession? session) + { + if (session?.StateBag.TryGetValue(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null) + { + return state; + } + + state = this._stateInitializer(session); + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); + } + + return state; + } + + /// + /// Saves the specified state to the session's StateBag using the configured state key and JSON serializer options. + /// If the session is null, this method does nothing. + /// + /// + /// This method provides a convenient way for derived classes to persist state changes back to the session after processing. + /// It abstracts away the details of how state is stored in the session, allowing derived classes to focus on their specific logic. + /// + /// The agent session containing the StateBag. + /// The state to be saved. + protected virtual void SaveState(AgentSession? session, TState state) + { + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs index d16ca69528..ad3f3aacfb 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -39,6 +40,25 @@ namespace Microsoft.Agents.AI; /// public abstract class ChatHistoryProvider { + private static IEnumerable DefaultExcludeChatHistoryFilter(IEnumerable messages) + => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory); + + private readonly Func, IEnumerable>? _provideOutputMessageFilter; + private readonly Func, IEnumerable> _storeInputMessageFilter; + + /// + /// Initializes a new instance of the class. + /// + /// An optional filter function to apply to messages when retrieving them from the chat history. + /// An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type . + protected ChatHistoryProvider( + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + { + this._provideOutputMessageFilter = provideOutputMessageFilter; + this._storeInputMessageFilter = storeInputMessageFilter ?? DefaultExcludeChatHistoryFilter; + } + /// /// Gets the key used to store the provider state in the . /// @@ -50,20 +70,16 @@ public abstract class ChatHistoryProvider public virtual string StateKey => this.GetType().Name; /// - /// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation. + /// Called at the start of agent invocation to provide messages for the next agent invocation. /// /// Contains the request context including the caller provided messages that will be used by the agent for this invocation. /// The to monitor for cancellation requests. The default is . /// /// A task that represents the asynchronous operation. The task result contains a collection of - /// instances in ascending chronological order (oldest first). + /// instances that will be used for the agent invocation. /// /// /// - /// Messages are returned in chronological order to maintain proper conversation flow and context for the agent. - /// The oldest messages appear first in the collection, followed by more recent messages. - /// - /// /// If the total message history becomes very large, implementations should apply appropriate strategies to manage /// storage constraints, such as: /// @@ -75,23 +91,19 @@ public abstract class ChatHistoryProvider /// /// public ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) - => this.InvokingCoreAsync(context, cancellationToken); + => this.InvokingCoreAsync(Throw.IfNull(context), cancellationToken); /// - /// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation. + /// Called at the start of agent invocation to provide messages for the next agent invocation. /// /// Contains the request context including the caller provided messages that will be used by the agent for this invocation. /// The to monitor for cancellation requests. The default is . /// /// A task that represents the asynchronous operation. The task result contains a collection of - /// instances in ascending chronological order (oldest first). + /// instances that will be used for the agent invocation. /// /// /// - /// Messages are returned in chronological order to maintain proper conversation flow and context for the agent. - /// The oldest messages appear first in the collection, followed by more recent messages. - /// - /// /// If the total message history becomes very large, implementations should apply appropriate strategies to manage /// storage constraints, such as: /// @@ -102,11 +114,54 @@ public ValueTask> InvokingAsync(InvokingContext context /// /// /// - /// Each instance should be associated with a single to ensure proper message isolation - /// and context management. + /// The default implementation of this method, calls to get the chat history messages, applies the optional retrieval output filter, + /// and merges the returned messages with the caller provided messages (with chat history messages appearing first) before returning the full message list to be used for the invocation. + /// For most scenarios, overriding is sufficient to return the desired chat history messages, while still benefiting from the default merging and filtering behavior. + /// However, for scenarios that require more control over message filtering, merging or source stamping, overriding this method allows you to directly control the full set of messages returned for the invocation. /// /// - protected abstract ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default); + protected virtual async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + var output = await this.ProvideChatHistoryAsync(context, cancellationToken).ConfigureAwait(false); + + if (this._provideOutputMessageFilter is not null) + { + output = this._provideOutputMessageFilter(output); + } + + return output + .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) + .Concat(context.RequestMessages); + } + + /// + /// When overridden in a derived class, provides the chat history messages to be used for the current invocation. + /// + /// + /// + /// This method is called from . + /// Note that can be overridden to directly control message filtering, merging and source stamping, in which case + /// it is up to the implementer to call this method as needed to retrieve the unfiltered/unmerged chat history messages. + /// + /// + /// In contrast with , this method only returns additional messages to be added to the request, + /// while is responsible for returning the full set of messages to be used for the invocation (including caller provided messages). + /// + /// + /// Messages are returned in chronological order to maintain proper conversation flow and context for the agent. + /// The oldest messages appear first in the collection, followed by more recent messages. + /// + /// + /// Contains the request context including the caller provided messages that will be used by the agent for this invocation. + /// The to monitor for cancellation requests. The default is . + /// + /// A task that represents the asynchronous operation. The task result contains a collection of + /// instances in ascending chronological order (oldest first). + /// + protected virtual ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + return new ValueTask>([]); + } /// /// Called at the end of the agent invocation to add new messages to the chat history. @@ -134,7 +189,7 @@ public ValueTask> InvokingAsync(InvokingContext context /// /// public ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) => - this.InvokedCoreAsync(context, cancellationToken); + this.InvokedCoreAsync(Throw.IfNull(context), cancellationToken); /// /// Called at the end of the agent invocation to add new messages to the chat history. @@ -160,8 +215,59 @@ public ValueTask InvokedAsync(InvokedContext context, CancellationToken cancella /// This method is called regardless of whether the invocation succeeded or failed. /// To check if the invocation was successful, inspect the property. /// + /// + /// The default implementation of this method, skips execution for any invocation failures, filters messages using the optional storage input message filter + /// and calls to store new chat history messages. + /// For most scenarios, overriding is sufficient to store chat history messages, while still benefiting from the default error handling and filtering behavior. + /// However, for scenarios that require more control over error handling or message filtering, overriding this method allows you to directly control the messages that are stored for the invocation. + /// + /// + protected virtual ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + if (context.InvokeException is not null) + { + return default; + } + + var subContext = new InvokedContext(context.Agent, context.Session, this._storeInputMessageFilter(context.RequestMessages), context.ResponseMessages!); + return this.StoreChatHistoryAsync(subContext, cancellationToken); + } + + /// + /// When overridden in a derived class, adds new messages to the chat history at the end of the agent invocation. + /// + /// Contains the invocation context including request messages, response messages, and any exception that occurred. + /// The to monitor for cancellation requests. The default is . + /// A task that represents the asynchronous add operation. + /// + /// + /// Messages should be added in the order they were generated to maintain proper chronological sequence. + /// The is responsible for preserving message ordering and ensuring that subsequent calls to + /// return messages in the correct chronological order. + /// + /// + /// Implementations may perform additional processing during message addition, such as: + /// + /// Validating message content and metadata + /// Applying storage optimizations or compression + /// Triggering background maintenance operations + /// + /// + /// + /// This method is called from . + /// Note that can be overridden to directly control message filtering and error handling, in which case + /// it is up to the implementer to call this method as needed to store messages. + /// + /// + /// In contrast with , this method only stores messages, + /// while is also responsible for messages filtering and error handling. + /// + /// + /// The default implementation of only calls this method if the invocation succeeded. + /// /// - protected abstract ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default); + protected virtual ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) => + default; /// Asks the for an object of the specified type . /// The type of object being requested. diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs new file mode 100644 index 0000000000..ea1280eb1e --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Text.Json; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Provides an abstract base class for fetching chat messages from, and adding chat messages to, chat history for the purposes of agent execution with support for maintaining provider state of type . +/// +/// The type of the state to be maintained by the chat history provider. Must be a reference type. +/// +/// This class extends by introducing a strongly-typed state management mechanism, allowing derived classes to maintain and persist custom state information across invocations. +/// The state is stored in the session's StateBag using a configurable key and JSON serialization options, enabling seamless integration with the agent session lifecycle. +/// +public abstract class ChatHistoryProvider : ChatHistoryProvider + where TState : class +{ + private readonly Func _stateInitializer; + private readonly string _stateKey; + private readonly JsonSerializerOptions _jsonSerializerOptions; + + /// + /// Initializes a new instance of the class. + /// + /// A function to initialize the state for the chat history provider. + /// The key used to store the state in the session's StateBag. + /// Options for JSON serialization and deserialization of the state. + /// A filter function to apply to messages when retrieving them from the chat history. + /// A filter function to apply to messages before storing them in the chat history. + protected ChatHistoryProvider( + Func stateInitializer, + string? stateKey, + JsonSerializerOptions? jsonSerializerOptions, + Func, IEnumerable>? provideOutputMessageFilter, + Func, IEnumerable>? storeInputMessageFilter) + : base(provideOutputMessageFilter, storeInputMessageFilter) + { + this._stateInitializer = stateInitializer; + this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + this._stateKey = stateKey ?? this.GetType().Name; + } + + /// + public override string StateKey => this._stateKey; + + /// + /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. + /// + /// The agent session containing the StateBag. + /// The provider state, or null if no session is available. + protected virtual TState GetOrInitializeState(AgentSession? session) + { + if (session?.StateBag.TryGetValue(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null) + { + return state; + } + + state = this._stateInitializer(session); + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); + } + + return state; + } + + /// + /// Saves the specified state to the session's StateBag using the configured state key and JSON serializer options. + /// If the session is null, this method does nothing. + /// + /// + /// This method provides a convenient way for derived classes to persist state changes back to the session after processing. + /// It abstracts away the details of how state is stored in the session, allowing derived classes to focus on their specific logic. + /// + /// The agent session containing the StateBag. + /// The state to be saved. + protected virtual void SaveState(AgentSession? session, TState state) + { + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs index 9c535923f4..2517586c17 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs @@ -3,7 +3,6 @@ using System; using System.Collections.Generic; using System.Linq; -using System.Text.Json; using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; @@ -25,17 +24,8 @@ namespace Microsoft.Agents.AI; /// message reduction strategies or alternative storage implementations. /// /// -public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider +public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider { - private static IEnumerable DefaultExcludeChatHistoryFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory); - - private readonly string _stateKey; - private readonly Func _stateInitializer; - private readonly JsonSerializerOptions _jsonSerializerOptions; - private readonly Func, IEnumerable> _storageInputMessageFilter; - private readonly Func, IEnumerable>? _retrievalOutputMessageFilter; - /// /// Initializes a new instance of the class. /// @@ -44,19 +34,17 @@ private static IEnumerable DefaultExcludeChatHistoryFilter(IEnumera /// message reduction, and serialization settings. If , default settings will be used. /// public InMemoryChatHistoryProvider(InMemoryChatHistoryProviderOptions? options = null) + : base( + options?.StateInitializer ?? (_ => new State()), + options?.StateKey, + options?.JsonSerializerOptions, + options?.ProvideOutputMessageFilter, + options?.StorageInputMessageFilter) { - this._stateInitializer = options?.StateInitializer ?? (_ => new State()); this.ChatReducer = options?.ChatReducer; this.ReducerTriggerEvent = options?.ReducerTriggerEvent ?? InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval; - this._stateKey = options?.StateKey ?? base.StateKey; - this._jsonSerializerOptions = options?.JsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; - this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExcludeChatHistoryFilter; - this._retrievalOutputMessageFilter = options?.RetrievalOutputMessageFilter; } - /// - public override string StateKey => this._stateKey; - /// /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied. /// @@ -89,32 +77,9 @@ public void SetMessages(AgentSession? session, List messages) state.Messages = messages; } - /// - /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state, or null if no session is available. - private State GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); - } - - return state; - } - /// - protected override async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) { - _ = Throw.IfNull(context); - var state = this.GetOrInitializeState(context.Session); if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null) @@ -122,30 +87,16 @@ protected override async ValueTask> InvokingCoreAsync(I state.Messages = (await this.ChatReducer.ReduceAsync(state.Messages, cancellationToken).ConfigureAwait(false)).ToList(); } - IEnumerable output = state.Messages; - if (this._retrievalOutputMessageFilter is not null) - { - output = this._retrievalOutputMessageFilter(output); - } - return output - .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) - .Concat(context.RequestMessages); + return state.Messages; } /// - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - _ = Throw.IfNull(context); - - if (context.InvokeException is not null) - { - return; - } - var state = this.GetOrInitializeState(context.Session); // Add request and response messages to the provider - var allNewMessages = this._storageInputMessageFilter(context.RequestMessages).Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); state.Messages.AddRange(allNewMessages); if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs index 41ab46321f..ba24f55ded 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs @@ -71,7 +71,7 @@ public sealed class InMemoryChatHistoryProviderOptions /// /// When , no filtering is applied to the output messages. /// - public Func, IEnumerable>? RetrievalOutputMessageFilter { get; set; } + public Func, IEnumerable>? ProvideOutputMessageFilter { get; set; } /// /// Defines the events that can trigger a reducer in the . diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs index 265f3a3675..b12a465905 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs @@ -19,16 +19,11 @@ namespace Microsoft.Agents.AI; /// [RequiresUnreferencedCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with trimming.")] [RequiresDynamicCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with NativeAOT.")] -public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable +public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable { - private static IEnumerable DefaultExcludeChatHistoryFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory); - private readonly CosmosClient _cosmosClient; private readonly Container _container; private readonly bool _ownsClient; - private readonly string _stateKey; - private readonly Func _stateInitializer; private bool _disposed; /// @@ -46,9 +41,6 @@ private static JsonSerializerOptions CreateDefaultJsonOptions() return options; } - /// - public override string StateKey => this._stateKey; - /// /// Gets or sets the maximum number of messages to return in a single query batch. /// Default is 100 for optimal performance. @@ -84,25 +76,6 @@ private static JsonSerializerOptions CreateDefaultJsonOptions() /// public string ContainerId { get; init; } - /// - /// A filter function applied to request messages before they are stored - /// during . The default filter excludes messages with the - /// source type. - /// - public Func, IEnumerable> StorageInputMessageFilter { get; set { field = Throw.IfNull(value); } } = DefaultExcludeChatHistoryFilter; - - /// - /// Gets or sets an optional filter function applied to messages produced by this provider - /// during . - /// - /// - /// This filter is only applied to the messages that the provider itself produces (from its internal storage). - /// - /// - /// When , no filtering is applied to the output messages. - /// - public Func, IEnumerable>? RetrievalOutputMessageFilter { get; set; } - /// /// Initializes a new instance of the class. /// @@ -112,6 +85,8 @@ private static JsonSerializerOptions CreateDefaultJsonOptions() /// A delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId). /// Whether this instance owns the CosmosClient and should dispose it. /// An optional key to use for storing the state in the . + /// An optional filter function to apply to messages when retrieving them from the chat history. + /// An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type . /// Thrown when or is . /// Thrown when any string parameter is null or whitespace. public CosmosChatHistoryProvider( @@ -120,15 +95,16 @@ public CosmosChatHistoryProvider( string containerId, Func stateInitializer, bool ownsClient = false, - string? stateKey = null) + string? stateKey = null, + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : base(Throw.IfNull(stateInitializer), stateKey, null, provideOutputMessageFilter, storeInputMessageFilter) { this._cosmosClient = Throw.IfNull(cosmosClient); this.DatabaseId = Throw.IfNullOrWhitespace(databaseId); this.ContainerId = Throw.IfNullOrWhitespace(containerId); this._container = this._cosmosClient.GetContainer(databaseId, containerId); - this._stateInitializer = Throw.IfNull(stateInitializer); this._ownsClient = ownsClient; - this._stateKey = stateKey ?? base.StateKey; } /// @@ -139,6 +115,8 @@ public CosmosChatHistoryProvider( /// The identifier of the Cosmos DB container. /// A delegate that initializes the provider state on the first invocation. /// An optional key to use for storing the state in the . + /// An optional filter function to apply to messages when retrieving them from the chat history. + /// An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatHistoryProvider( @@ -146,8 +124,10 @@ public CosmosChatHistoryProvider( string databaseId, string containerId, Func stateInitializer, - string? stateKey = null) - : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey) + string? stateKey = null, + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputMessageFilter) { } @@ -160,6 +140,8 @@ public CosmosChatHistoryProvider( /// The identifier of the Cosmos DB container. /// A delegate that initializes the provider state on the first invocation. /// An optional key to use for storing the state in the . + /// An optional filter function to apply to messages when retrieving them from the chat history. + /// An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatHistoryProvider( @@ -168,30 +150,11 @@ public CosmosChatHistoryProvider( string databaseId, string containerId, Func stateInitializer, - string? stateKey = null) - : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey) - { - } - - /// - /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state, or null if no session is available. - private State GetOrInitializeState(AgentSession? session) + string? stateKey = null, + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputMessageFilter) { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, AgentAbstractionsJsonUtilities.DefaultOptions) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, AgentAbstractionsJsonUtilities.DefaultOptions); - } - - return state; } /// @@ -218,7 +181,7 @@ private static PartitionKey BuildPartitionKey(State state) } /// - protected override async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) { #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) @@ -227,8 +190,6 @@ protected override async ValueTask> InvokingCoreAsync(I } #pragma warning restore CA1513 - _ = Throw.IfNull(context); - var state = this.GetOrInitializeState(context.Session); var partitionKey = BuildPartitionKey(state); @@ -279,22 +240,12 @@ protected override async ValueTask> InvokingCoreAsync(I messages.Reverse(); } - return (this.RetrievalOutputMessageFilter is not null ? this.RetrievalOutputMessageFilter(messages) : messages) - .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) - .Concat(context.RequestMessages); + return messages; } /// - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - Throw.IfNull(context); - - if (context.InvokeException is not null) - { - // Do not store messages if there was an exception during invocation - return; - } - #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) { @@ -303,7 +254,7 @@ protected override async ValueTask InvokedCoreAsync(InvokedContext context, Canc #pragma warning restore CA1513 var state = this.GetOrInitializeState(context.Session); - var messageList = this.StorageInputMessageFilter(context.RequestMessages).Concat(context.ResponseMessages ?? []).ToList(); + var messageList = context.RequestMessages.Concat(context.ResponseMessages ?? []).ToList(); if (messageList.Count == 0) { return; diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs index aaf2333553..4ef096ba6e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs @@ -22,19 +22,12 @@ namespace Microsoft.Agents.AI.Mem0; /// for new invocations using a semantic search endpoint. Retrieved memories are injected as user messages /// to the model, prefixed by a configurable context prompt. /// -public sealed class Mem0Provider : AIContextProvider +public sealed class Mem0Provider : AIContextProvider { private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:"; - private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); - private readonly string _contextPrompt; private readonly bool _enableSensitiveTelemetryData; - private readonly string _stateKey; - private readonly Func _stateInitializer; - private readonly Func, IEnumerable> _searchInputMessageFilter; - private readonly Func, IEnumerable> _storageInputMessageFilter; private readonly Mem0Client _client; private readonly ILogger? _logger; @@ -58,6 +51,7 @@ private static IEnumerable DefaultExternalOnlyFilter(IEnumerable /// public Mem0Provider(HttpClient httpClient, Func stateInitializer, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + : base(ValidateStateInitializer(Throw.IfNull(stateInitializer)), options?.StateKey, Mem0JsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) { Throw.IfNull(httpClient); if (string.IsNullOrWhiteSpace(httpClient.BaseAddress?.AbsoluteUri)) @@ -65,63 +59,41 @@ public Mem0Provider(HttpClient httpClient, Func stateIniti throw new ArgumentException("The HttpClient BaseAddress must be set for Mem0 operations.", nameof(httpClient)); } - this._stateInitializer = Throw.IfNull(stateInitializer); this._logger = loggerFactory?.CreateLogger(); this._client = new Mem0Client(httpClient); this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt; this._enableSensitiveTelemetryData = options?.EnableSensitiveTelemetryData ?? false; - this._stateKey = options?.StateKey ?? base.StateKey; - this._searchInputMessageFilter = options?.SearchInputMessageFilter ?? DefaultExternalOnlyFilter; - this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExternalOnlyFilter; } - /// - public override string StateKey => this._stateKey; - - /// - /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state, or null if no session is available. - private State? GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, Mem0JsonUtilities.DefaultOptions) is true && state is not null) + private static Func ValidateStateInitializer(Func stateInitializer) => + session => { - return state; - } - - state = this._stateInitializer(session); - - if (state is null - || state.StorageScope is null - || (state.StorageScope.AgentId is null && state.StorageScope.ThreadId is null && state.StorageScope.UserId is null && state.StorageScope.ApplicationId is null) - || state.SearchScope is null - || (state.SearchScope.AgentId is null && state.SearchScope.ThreadId is null && state.SearchScope.UserId is null && state.SearchScope.ApplicationId is null)) - { - throw new InvalidOperationException("State initializer must return a non-null state with valid storage and search scopes, where at lest one scoping parameter is set for each."); - } + var state = stateInitializer(session); - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, Mem0JsonUtilities.DefaultOptions); - } + if (state is null + || state.StorageScope is null + || (state.StorageScope.AgentId is null && state.StorageScope.ThreadId is null && state.StorageScope.UserId is null && state.StorageScope.ApplicationId is null) + || state.SearchScope is null + || (state.SearchScope.AgentId is null && state.SearchScope.ThreadId is null && state.SearchScope.UserId is null && state.SearchScope.ApplicationId is null)) + { + throw new InvalidOperationException("State initializer must return a non-null state with valid storage and search scopes, where at least one scoping parameter is set for each."); + } - return state; - } + return state; + }; /// - protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { Throw.IfNull(context); - var inputContext = context.AIContext; var state = this.GetOrInitializeState(context.Session); - var searchScope = state?.SearchScope ?? new Mem0ProviderScope(); + var searchScope = state.SearchScope; string queryText = string.Join( Environment.NewLine, - this._searchInputMessageFilter(inputContext.Messages ?? []) + (context.AIContext.Messages ?? []) .Where(m => !string.IsNullOrWhiteSpace(m.Text)) .Select(m => m.Text)); @@ -138,9 +110,6 @@ protected override async ValueTask InvokingCoreAsync(InvokingContext var outputMessageText = memories.Count == 0 ? null : $"{this._contextPrompt}\n{string.Join(Environment.NewLine, memories)}"; - var outputMessage = memories.Count == 0 - ? null - : new ChatMessage(ChatRole.User, outputMessageText!).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!); if (this._logger?.IsEnabled(LogLevel.Information) is true) { @@ -167,11 +136,9 @@ protected override async ValueTask InvokingCoreAsync(InvokingContext return new AIContext { - Instructions = inputContext.Instructions, - Messages = - (inputContext.Messages ?? []) - .Concat(outputMessage is not null ? [outputMessage] : []), - Tools = inputContext.Tools + Messages = outputMessageText is not null + ? [new ChatMessage(ChatRole.User, outputMessageText)] + : null }; } catch (ArgumentException) @@ -190,27 +157,23 @@ protected override async ValueTask InvokingCoreAsync(InvokingContext searchScope.ThreadId, this.SanitizeLogData(searchScope.UserId)); } - return inputContext; + + return new AIContext(); } } /// - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { - if (context.InvokeException is not null) - { - return; // Do not update memory on failed invocations. - } - var state = this.GetOrInitializeState(context.Session); - var storageScope = state?.StorageScope ?? new Mem0ProviderScope(); + var storageScope = state.StorageScope; try { // Persist request and response messages after invocation. await this.PersistMessagesAsync( storageScope, - this._storageInputMessageFilter(context.RequestMessages) + context.RequestMessages .Concat(context.ResponseMessages ?? []), cancellationToken).ConfigureAwait(false); } @@ -238,12 +201,7 @@ public Task ClearStoredMemoriesAsync(AgentSession session, CancellationToken can { Throw.IfNull(session); var state = this.GetOrInitializeState(session); - var storageScope = state?.StorageScope; - - if (storageScope is null) - { - return Task.CompletedTask; // Nothing to clear if there is no state. - } + var storageScope = state.StorageScope; return this._client.ClearMemoryAsync( storageScope.ApplicationId, diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs index 6672b9e2a3..74e2fae3fa 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs @@ -9,10 +9,8 @@ namespace Microsoft.Agents.AI.Workflows; -internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider +internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider { - private readonly JsonSerializerOptions _jsonSerializerOptions; - /// /// Initializes a new instance of the class. /// @@ -22,8 +20,8 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider /// and source generated serializers are required, or Native AOT / Trimming is required. /// public WorkflowChatHistoryProvider(JsonSerializerOptions? jsonSerializerOptions = null) + : base(stateInitializer: _ => new StoreState(), stateKey: null, jsonSerializerOptions: jsonSerializerOptions, provideOutputMessageFilter: null, storeInputMessageFilter: null) { - this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; } internal sealed class StoreState @@ -32,43 +30,16 @@ internal sealed class StoreState public List Messages { get; set; } = []; } - private StoreState GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this.StateKey, out var state, this._jsonSerializerOptions) is true && state is not null) - { - return state; - } - - state = new(); - if (session is not null) - { - session.StateBag.SetValue(this.StateKey, state, this._jsonSerializerOptions); - } - - return state; - } - internal void AddMessages(AgentSession session, params IEnumerable messages) => this.GetOrInitializeState(session).Messages.AddRange(messages); - protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(this.GetOrInitializeState(context.Session) - .Messages - .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) - .Concat(context.RequestMessages)); + protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) + => new(this.GetOrInitializeState(context.Session).Messages); - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - if (context.InvokeException is not null) - { - return default; - } - - var allNewMessages = context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory) - .Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); this.GetOrInitializeState(context.Session).Messages.AddRange(allNewMessages); - return default; } diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs index 9d163f79cf..346365e8cc 100644 --- a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs @@ -24,8 +24,8 @@ namespace Microsoft.Agents.AI; /// abstractions to work with any compatible vector store implementation. /// /// -/// Messages are stored during the method and retrieved during the -/// method using semantic similarity search. +/// Messages are stored during the method and retrieved during the +/// method using semantic similarity search. /// /// /// Behavior is configurable through . When @@ -34,16 +34,13 @@ namespace Microsoft.Agents.AI; /// injecting them automatically on each invocation. /// /// -public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable +public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable { private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:"; private const int DefaultMaxResults = 3; private const string DefaultFunctionToolName = "Search"; private const string DefaultFunctionToolDescription = "Allows searching for related previous chat history to help answer the user question."; - private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); - #pragma warning disable CA2213 // VectorStore is not owned by this class - caller is responsible for disposal private readonly VectorStore _vectorStore; #pragma warning restore CA2213 @@ -55,10 +52,6 @@ private static IEnumerable DefaultExternalOnlyFilter(IEnumerable? _logger; - private readonly string _stateKey; - private readonly Func _stateInitializer; - private readonly Func, IEnumerable> _searchInputMessageFilter; - private readonly Func, IEnumerable> _storageInputMessageFilter; private bool _collectionInitialized; private readonly SemaphoreSlim _initializationLock = new(1, 1); @@ -81,21 +74,18 @@ public ChatHistoryMemoryProvider( Func stateInitializer, ChatHistoryMemoryProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + : base(Throw.IfNull(stateInitializer), options?.StateKey, AgentJsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) { this._vectorStore = Throw.IfNull(vectorStore); - this._stateInitializer = Throw.IfNull(stateInitializer); options ??= new ChatHistoryMemoryProviderOptions(); this._maxResults = options.MaxResults.HasValue ? Throw.IfLessThanOrEqual(options.MaxResults.Value, 0) : DefaultMaxResults; this._contextPrompt = options.ContextPrompt ?? DefaultContextPrompt; this._enableSensitiveTelemetryData = options.EnableSensitiveTelemetryData; this._searchTime = options.SearchTime; - this._stateKey = options.StateKey ?? base.StateKey; this._logger = loggerFactory?.CreateLogger(); this._toolName = options.FunctionToolName ?? DefaultFunctionToolName; this._toolDescription = options.FunctionToolDescription ?? DefaultFunctionToolDescription; - this._searchInputMessageFilter = options.SearchInputMessageFilter ?? DefaultExternalOnlyFilter; - this._storageInputMessageFilter = options.StorageInputMessageFilter ?? DefaultExternalOnlyFilter; // Create a definition so that we can use the dimensions provided at runtime. var definition = new VectorStoreCollectionDefinition @@ -120,37 +110,12 @@ public ChatHistoryMemoryProvider( } /// - public override string StateKey => this._stateKey; - - /// - /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state, or null if no session is available. - private State? GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, AgentJsonUtilities.DefaultOptions) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (state is not null && session is not null) - { - session.StateBag.SetValue(this._stateKey, state, AgentJsonUtilities.DefaultOptions); - } - - return state; - } - - /// - protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { _ = Throw.IfNull(context); - var inputContext = context.AIContext; var state = this.GetOrInitializeState(context.Session); - var searchScope = state?.SearchScope ?? new ChatHistoryMemoryProviderScope(); + var searchScope = state.SearchScope; if (this._searchTime == ChatHistoryMemoryProviderOptions.SearchBehavior.OnDemandFunctionCalling) { @@ -166,12 +131,10 @@ Task InlineSearchAsync(string userQuestion, CancellationToken ct) description: this._toolDescription) ]; - // Expose search tool for on-demand invocation by the model, accumulated with the input context + // Expose search tool for on-demand invocation by the model return new AIContext { - Instructions = inputContext.Instructions, - Messages = inputContext.Messages, - Tools = (inputContext.Tools ?? []).Concat(tools) + Tools = tools }; } @@ -179,13 +142,13 @@ Task InlineSearchAsync(string userQuestion, CancellationToken ct) { // Get the text from the current request messages var requestText = string.Join("\n", - this._searchInputMessageFilter(inputContext.Messages ?? []) + (context.AIContext.Messages ?? []) .Where(m => m != null && !string.IsNullOrWhiteSpace(m.Text)) .Select(m => m.Text)); if (string.IsNullOrWhiteSpace(requestText)) { - return inputContext; + return new AIContext(); } // Search for relevant chat history @@ -193,19 +156,12 @@ Task InlineSearchAsync(string userQuestion, CancellationToken ct) if (string.IsNullOrWhiteSpace(contextText)) { - return inputContext; + return new AIContext(); } return new AIContext { - Instructions = inputContext.Instructions, - Messages = - (inputContext.Messages ?? []) - .Concat( - [ - new ChatMessage(ChatRole.User, contextText).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!) - ]), - Tools = inputContext.Tools + Messages = [new ChatMessage(ChatRole.User, contextText)] }; } catch (Exception ex) @@ -221,30 +177,24 @@ Task InlineSearchAsync(string userQuestion, CancellationToken ct) this.SanitizeLogData(searchScope.UserId)); } - return inputContext; + return new AIContext(); } } /// - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { _ = Throw.IfNull(context); - // Only store if invocation was successful - if (context.InvokeException != null) - { - return; - } - var state = this.GetOrInitializeState(context.Session); - var storageScope = state?.StorageScope ?? new ChatHistoryMemoryProviderScope(); + var storageScope = state.StorageScope; try { // Ensure the collection is initialized var collection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false); - List> itemsToStore = this._storageInputMessageFilter(context.RequestMessages) + List> itemsToStore = context.RequestMessages .Concat(context.ResponseMessages ?? []) .Select(message => new Dictionary { diff --git a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs index f038fa3c38..099059d65e 100644 --- a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs @@ -32,16 +32,13 @@ namespace Microsoft.Agents.AI; /// multi-turn context to the retrieval layer without permanently altering the conversation history. /// /// -public sealed class TextSearchProvider : AIContextProvider +public sealed class TextSearchProvider : AIContextProvider { private const string DefaultPluginSearchFunctionName = "Search"; private const string DefaultPluginSearchFunctionDescription = "Allows searching for additional information to help answer the user question."; private const string DefaultContextPrompt = "## Additional Context\nConsider the following information from source documents when responding to the user:"; private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available."; - private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); - private readonly Func>> _searchAsync; private readonly ILogger? _logger; private readonly AITool[] _tools; @@ -50,10 +47,7 @@ private static IEnumerable DefaultExternalOnlyFilter(IEnumerable, string>? _contextFormatter; - private readonly Func, IEnumerable> _searchInputMessageFilter; - private readonly Func, IEnumerable> _storageInputMessageFilter; /// /// Initializes a new instance of the class. @@ -66,6 +60,7 @@ public TextSearchProvider( Func>> searchAsync, TextSearchProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + : base(_ => new TextSearchProviderState(), options?.StateKey, AgentJsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) { // Validate and assign parameters this._searchAsync = Throw.IfNull(searchAsync); @@ -75,10 +70,7 @@ public TextSearchProvider( this._searchTime = options?.SearchTime ?? TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke; this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt; this._citationsPrompt = options?.CitationsPrompt ?? DefaultCitationsPrompt; - this._stateKey = options?.StateKey ?? base.StateKey; this._contextFormatter = options?.ContextFormatter; - this._searchInputMessageFilter = options?.SearchInputMessageFilter ?? DefaultExternalOnlyFilter; - this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExternalOnlyFilter; // Create the on-demand search tool (only used if behavior is OnDemandFunctionCalling) this._tools = @@ -91,32 +83,25 @@ public TextSearchProvider( } /// - public override string StateKey => this._stateKey; - - /// - protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { - var inputContext = context.AIContext; - if (this._searchTime != TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke) { - // Expose the search tool for on-demand invocation, accumulated with the input context. + // Expose the search tool for on-demand invocation. return new AIContext { - Instructions = inputContext.Instructions, - Messages = inputContext.Messages, - Tools = (inputContext.Tools ?? []).Concat(this._tools) + Tools = this._tools }; } - // Retrieve recent messages from the session state bag. - var recentMessagesText = context.Session?.StateBag.GetValue(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText + // Retrieve recent messages from the session state. + var recentMessagesText = this.GetOrInitializeState(context.Session).RecentMessagesText ?? []; // Aggregate text from memory + current request messages. var sbInput = new StringBuilder(); var requestMessagesText = - this._searchInputMessageFilter(inputContext.Messages ?? []) + (context.AIContext.Messages ?? []) .Where(x => !string.IsNullOrWhiteSpace(x?.Text)).Select(x => x.Text); foreach (var messageText in recentMessagesText.Concat(requestMessagesText)) { @@ -142,7 +127,7 @@ protected override async ValueTask InvokingCoreAsync(InvokingContext if (materialized.Count == 0) { - return inputContext; + return new AIContext(); } // Format search results @@ -155,25 +140,18 @@ protected override async ValueTask InvokingCoreAsync(InvokingContext return new AIContext { - Instructions = inputContext.Instructions, - Messages = - (inputContext.Messages ?? []) - .Concat( - [ - new ChatMessage(ChatRole.User, formatted).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!) - ]), - Tools = inputContext.Tools + Messages = [new ChatMessage(ChatRole.User, formatted)] }; } catch (Exception ex) { this._logger?.LogError(ex, "TextSearchProvider: Failed to search for data due to error"); - return inputContext; + return new AIContext(); } } /// - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { int limit = this._recentMessageMemoryLimit; if (limit <= 0) @@ -186,16 +164,11 @@ protected override ValueTask InvokedCoreAsync(InvokedContext context, Cancellati return default; // No session to store state in. } - if (context.InvokeException is not null) - { - return default; // Do not update memory on failed invocations. - } - - // Retrieve existing recent messages from the session state bag. - var recentMessagesText = context.Session.StateBag.GetValue(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText + // Retrieve existing recent messages from the session state. + var recentMessagesText = this.GetOrInitializeState(context.Session).RecentMessagesText ?? []; - var newMessagesText = this._storageInputMessageFilter(context.RequestMessages) + var newMessagesText = context.RequestMessages .Concat(context.ResponseMessages ?? []) .Where(m => this._recentMessageRolesIncluded.Contains(m.Role) && @@ -208,11 +181,10 @@ protected override ValueTask InvokedCoreAsync(InvokedContext context, Cancellati ? allMessages.Skip(allMessages.Count - limit).ToList() : allMessages; - // Store updated state back to the session state bag. - context.Session.StateBag.SetValue( - this._stateKey, - new TextSearchProviderState { RecentMessagesText = updatedMessages }, - AgentJsonUtilities.DefaultOptions); + // Store updated state back to the session. + this.SaveState( + context.Session, + new TextSearchProviderState { RecentMessagesText = updatedMessages }); return default; } @@ -311,8 +283,14 @@ public sealed class TextSearchResult public object? RawRepresentation { get; set; } } - internal sealed class TextSearchProviderState + /// + /// Represents the per-session state of a stored in the . + /// + public sealed class TextSearchProviderState { + /// + /// Gets or sets the list of recent message texts retained for multi-turn search context. + /// public List? RecentMessagesText { get; set; } } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTStateTests.cs new file mode 100644 index 0000000000..b457bf1945 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTStateTests.cs @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Moq; + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Contains tests for the class. +/// +public class AIContextProviderTStateTests +{ + private static readonly AIAgent s_mockAgent = new Mock().Object; + + #region GetOrInitializeState Tests + + [Fact] + public void GetOrInitializeState_InitializesFromStateInitializerOnFirstCall() + { + // Arrange + var expectedState = new TestState { Value = "initialized" }; + var provider = new TestAIContextProvider(_ => expectedState); + var session = new TestAgentSession(); + + // Act + var state = provider.GetState(session); + + // Assert + Assert.Same(expectedState, state); + } + + [Fact] + public void GetOrInitializeState_ReturnsCachedStateFromStateBagOnSecondCall() + { + // Arrange + var callCount = 0; + var provider = new TestAIContextProvider(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }); + var session = new TestAgentSession(); + + // Act + var state1 = provider.GetState(session); + var state2 = provider.GetState(session); + + // Assert - initializer called only once; second call reads from StateBag + Assert.Equal(1, callCount); + Assert.Equal("init-1", state1.Value); + Assert.Equal("init-1", state2.Value); + } + + [Fact] + public void GetOrInitializeState_WorksWhenSessionIsNull() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState { Value = "no-session" }); + + // Act + var state = provider.GetState(null); + + // Assert + Assert.Equal("no-session", state.Value); + } + + [Fact] + public void GetOrInitializeState_ReInitializesWhenSessionIsNull() + { + // Arrange - without a session, state can't be cached in StateBag + var callCount = 0; + var provider = new TestAIContextProvider(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }); + + // Act + provider.GetState(null); + provider.GetState(null); + + // Assert - initializer called each time since there's no session to cache in + Assert.Equal(2, callCount); + } + + #endregion + + #region SaveState Tests + + [Fact] + public void SaveState_SavesToStateBag() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState()); + var session = new TestAgentSession(); + var state = new TestState { Value = "saved" }; + + // Act + provider.DoSaveState(session, state); + var retrieved = provider.GetState(session); + + // Assert + Assert.Equal("saved", retrieved.Value); + } + + [Fact] + public void SaveState_NoOpWhenSessionIsNull() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState { Value = "default" }); + + // Act - should not throw + provider.DoSaveState(null, new TestState { Value = "saved" }); + + // Assert - no exception; can't verify further without a session + } + + #endregion + + #region StateKey Tests + + [Fact] + public void StateKey_DefaultsToTypeName() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState()); + + // Act & Assert + Assert.Equal(nameof(TestAIContextProvider), provider.StateKey); + } + + [Fact] + public void StateKey_UsesCustomKeyWhenProvided() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState(), stateKey: "custom-key"); + + // Act & Assert + Assert.Equal("custom-key", provider.StateKey); + } + + #endregion + + #region Integration Tests + + [Fact] + public async Task InvokingCoreAsync_CanUseStateInProvideAIContextAsync() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState { Value = "state-value" }); + var session = new TestAgentSession(); + var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hi")] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, session, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - the provider uses state to produce context messages + var messages = result.Messages!.ToList(); + Assert.Equal(2, messages.Count); + Assert.Contains("state-value", messages[1].Text); + } + + #endregion + + public sealed class TestState + { + public string Value { get; set; } = string.Empty; + } + + private sealed class TestAIContextProvider : AIContextProvider + { + public TestAIContextProvider( + Func stateInitializer, + string? stateKey = null) + : base(stateInitializer, stateKey, null, null, null) + { + } + + public TestState GetState(AgentSession? session) => this.GetOrInitializeState(session); + + public void DoSaveState(AgentSession? session, TestState state) => this.SaveState(session, state); + + protected override ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + var state = this.GetOrInitializeState(context.Session); + return new(new AIContext + { + Messages = [new ChatMessage(ChatRole.System, $"Context from state: {state.Value}")] + }); + } + } + + private sealed class TestAgentSession : AgentSession; +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs index 14a0f81e08..811f9a3216 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Collections.ObjectModel; +using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -337,9 +338,314 @@ public void InvokedContext_FailureConstructor_ThrowsForNullException() #endregion + #region InvokingAsync / InvokedAsync Null Check Tests + + [Fact] + public async Task InvokingAsync_NullContext_ThrowsArgumentNullExceptionAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.InvokingAsync(null!).AsTask()); + } + + [Fact] + public async Task InvokedAsync_NullContext_ThrowsArgumentNullExceptionAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.InvokedAsync(null!).AsTask()); + } + + #endregion + + #region InvokingCoreAsync Tests + + [Fact] + public async Task InvokingCoreAsync_CallsProvideAIContextAndReturnsMergedContextAsync() + { + // Arrange + var providedMessages = new[] { new ChatMessage(ChatRole.System, "Context message") }; + var provider = new TestAIContextProvider(provideContext: new AIContext { Messages = providedMessages }); + var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "User input")] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - input messages + provided messages merged + var messages = result.Messages!.ToList(); + Assert.Equal(2, messages.Count); + Assert.Equal("User input", messages[0].Text); + Assert.Equal("Context message", messages[1].Text); + } + + [Fact] + public async Task InvokingCoreAsync_FiltersInputToExternalOnlyByDefaultAsync() + { + // Arrange + var provider = new TestAIContextProvider(captureFilteredContext: true); + var externalMsg = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMsg = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var contextProviderMsg = new ChatMessage(ChatRole.User, "ContextProvider") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "src"); + var inputContext = new AIContext { Messages = [externalMsg, chatHistoryMsg, contextProviderMsg] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + await provider.InvokingAsync(context); + + // Assert - ProvideAIContextAsync received only External messages + Assert.NotNull(provider.LastProvidedContext); + var filteredMessages = provider.LastProvidedContext!.AIContext.Messages!.ToList(); + Assert.Single(filteredMessages); + Assert.Equal("External", filteredMessages[0].Text); + } + + [Fact] + public async Task InvokingCoreAsync_StampsProvidedMessagesWithAIContextProviderSourceAsync() + { + // Arrange + var providedMessages = new[] { new ChatMessage(ChatRole.System, "Provided") }; + var provider = new TestAIContextProvider(provideContext: new AIContext { Messages = providedMessages }); + var inputContext = new AIContext { Messages = [] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert + var messages = result.Messages!.ToList(); + Assert.Single(messages); + Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, messages[0].GetAgentRequestMessageSourceType()); + } + + [Fact] + public async Task InvokingCoreAsync_MergesInstructionsAsync() + { + // Arrange + var provider = new TestAIContextProvider(provideContext: new AIContext { Instructions = "Provided instructions" }); + var inputContext = new AIContext { Instructions = "Input instructions" }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - instructions are joined with newline + Assert.Equal("Input instructions\nProvided instructions", result.Instructions); + } + + [Fact] + public async Task InvokingCoreAsync_MergesToolsAsync() + { + // Arrange + var inputTool = AIFunctionFactory.Create(() => "a", "inputTool"); + var providedTool = AIFunctionFactory.Create(() => "b", "providedTool"); + var provider = new TestAIContextProvider(provideContext: new AIContext { Tools = [providedTool] }); + var inputContext = new AIContext { Tools = [inputTool] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - both tools present + var tools = result.Tools!.ToList(); + Assert.Equal(2, tools.Count); + } + + [Fact] + public async Task InvokingCoreAsync_UsesCustomProvideInputFilterAsync() + { + // Arrange - filter that keeps all messages (not just External) + var provider = new TestAIContextProvider( + captureFilteredContext: true, + provideInputMessageFilter: msgs => msgs); + var externalMsg = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMsg = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var inputContext = new AIContext { Messages = [externalMsg, chatHistoryMsg] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + await provider.InvokingAsync(context); + + // Assert - ProvideAIContextAsync received ALL messages (custom filter keeps everything) + Assert.NotNull(provider.LastProvidedContext); + var filteredMessages = provider.LastProvidedContext!.AIContext.Messages!.ToList(); + Assert.Equal(2, filteredMessages.Count); + } + + [Fact] + public async Task InvokingCoreAsync_ReturnsEmptyContextByDefaultAsync() + { + // Arrange - provider that doesn't override ProvideAIContextAsync + var provider = new DefaultAIContextProvider(); + var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - only the input messages (no additional provided) + var messages = result.Messages!.ToList(); + Assert.Single(messages); + Assert.Equal("Hello", messages[0].Text); + } + + [Fact] + public async Task InvokingCoreAsync_MergesWithOriginalUnfilteredMessagesAsync() + { + // Arrange - default filter is External-only, but the MERGED result should include + // the original unfiltered input messages plus the provided messages + var providedMessages = new[] { new ChatMessage(ChatRole.System, "Provided") }; + var provider = new TestAIContextProvider(provideContext: new AIContext { Messages = providedMessages }); + var externalMsg = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMsg = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var inputContext = new AIContext { Messages = [externalMsg, chatHistoryMsg] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - original 2 input messages + 1 provided message + var messages = result.Messages!.ToList(); + Assert.Equal(3, messages.Count); + Assert.Equal("External", messages[0].Text); + Assert.Equal("History", messages[1].Text); + Assert.Equal("Provided", messages[2].Text); + } + + #endregion + + #region InvokedCoreAsync Tests + + [Fact] + public async Task InvokedCoreAsync_CallsStoreAIContextWithFilteredMessagesAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + var externalMessage = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMessage = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var responseMessages = new[] { new ChatMessage(ChatRole.Assistant, "Response") }; + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, new[] { externalMessage, chatHistoryMessage }, responseMessages); + + // Act + await provider.InvokedAsync(context); + + // Assert - default filter keeps only External messages + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("External", storedRequest[0].Text); + Assert.Same(responseMessages, provider.LastStoredContext.ResponseMessages); + } + + [Fact] + public async Task InvokedCoreAsync_SkipsStorageWhenInvokeExceptionIsNotNullAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "msg")], new InvalidOperationException("Failed")); + + // Act + await provider.InvokedAsync(context); + + // Assert - StoreAIContextAsync was NOT called + Assert.Null(provider.LastStoredContext); + } + + [Fact] + public async Task InvokedCoreAsync_UsesCustomStoreInputFilterAsync() + { + // Arrange - filter that only keeps System messages + var provider = new TestAIContextProvider( + storeInputMessageFilter: msgs => msgs.Where(m => m.Role == ChatRole.System)); + var messages = new[] + { + new ChatMessage(ChatRole.User, "User msg"), + new ChatMessage(ChatRole.System, "System msg") + }; + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, messages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(context); + + // Assert - only System messages were passed to store + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("System msg", storedRequest[0].Text); + } + + [Fact] + public async Task InvokedCoreAsync_DefaultFilterExcludesNonExternalMessagesAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + var external = new ChatMessage(ChatRole.User, "External"); + var fromHistory = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var fromContext = new ChatMessage(ChatRole.User, "Context") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "src"); + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [external, fromHistory, fromContext], []); + + // Act + await provider.InvokedAsync(context); + + // Assert - only External messages kept + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("External", storedRequest[0].Text); + } + + #endregion + private sealed class TestAIContextProvider : AIContextProvider { - protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(new AIContext()); + private readonly AIContext? _provideContext; + private readonly bool _captureFilteredContext; + + public InvokedContext? LastStoredContext { get; private set; } + + public InvokingContext? LastProvidedContext { get; private set; } + + public TestAIContextProvider( + AIContext? provideContext = null, + bool captureFilteredContext = false, + Func, IEnumerable>? provideInputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : base(provideInputMessageFilter, storeInputMessageFilter) + { + this._provideContext = provideContext; + this._captureFilteredContext = captureFilteredContext; + } + + protected override ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + if (this._captureFilteredContext) + { + this.LastProvidedContext = context; + } + + return new(this._provideContext ?? new AIContext()); + } + + protected override ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + this.LastStoredContext = context; + return default; + } } + + /// + /// A provider that uses only base class defaults (no overrides of ProvideAIContextAsync/StoreAIContextAsync). + /// + private sealed class DefaultAIContextProvider : AIContextProvider; } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs new file mode 100644 index 0000000000..e3c0507e8d --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs @@ -0,0 +1,195 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Moq; + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Contains tests for the class. +/// +public class ChatHistoryProviderTStateTests +{ + private static readonly AIAgent s_mockAgent = new Mock().Object; + + #region GetOrInitializeState Tests + + [Fact] + public void GetOrInitializeState_InitializesFromStateInitializerOnFirstCall() + { + // Arrange + var expectedState = new TestState { Value = "initialized" }; + var provider = new TestChatHistoryProvider(_ => expectedState); + var session = new TestAgentSession(); + + // Act + var state = provider.GetState(session); + + // Assert + Assert.Same(expectedState, state); + } + + [Fact] + public void GetOrInitializeState_ReturnsCachedStateFromStateBagOnSecondCall() + { + // Arrange + var callCount = 0; + var provider = new TestChatHistoryProvider(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }); + var session = new TestAgentSession(); + + // Act + var state1 = provider.GetState(session); + var state2 = provider.GetState(session); + + // Assert - initializer called only once; second call reads from StateBag + Assert.Equal(1, callCount); + Assert.Equal("init-1", state1.Value); + Assert.Equal("init-1", state2.Value); + } + + [Fact] + public void GetOrInitializeState_WorksWhenSessionIsNull() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState { Value = "no-session" }); + + // Act + var state = provider.GetState(null); + + // Assert + Assert.Equal("no-session", state.Value); + } + + [Fact] + public void GetOrInitializeState_ReInitializesWhenSessionIsNull() + { + // Arrange - without a session, state can't be cached in StateBag + var callCount = 0; + var provider = new TestChatHistoryProvider(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }); + + // Act + _ = provider.GetState(null); + provider.GetState(null); + + // Assert - initializer called each time since there's no session to cache in + Assert.Equal(2, callCount); + } + + #endregion + + #region SaveState Tests + + [Fact] + public void SaveState_SavesToStateBag() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState()); + var session = new TestAgentSession(); + var state = new TestState { Value = "saved" }; + + // Act + provider.DoSaveState(session, state); + var retrieved = provider.GetState(session); + + // Assert + Assert.Equal("saved", retrieved.Value); + } + + [Fact] + public void SaveState_NoOpWhenSessionIsNull() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState { Value = "default" }); + + // Act - should not throw + provider.DoSaveState(null, new TestState { Value = "saved" }); + + // Assert - no exception; can't verify further without a session + } + + #endregion + + #region StateKey Tests + + [Fact] + public void StateKey_DefaultsToTypeName() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState()); + + // Act & Assert + Assert.Equal(nameof(TestChatHistoryProvider), provider.StateKey); + } + + [Fact] + public void StateKey_UsesCustomKeyWhenProvided() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState(), stateKey: "custom-key"); + + // Act & Assert + Assert.Equal("custom-key", provider.StateKey); + } + + #endregion + + #region Integration Tests + + [Fact] + public async Task InvokingCoreAsync_CanUseStateInProvideChatHistoryAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState { Value = "state-value" }); + var session = new TestAgentSession(); + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Hi")]); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert - the provider uses state to produce history messages + Assert.Equal(2, result.Count); + Assert.Contains("state-value", result[0].Text); + } + + #endregion + + public sealed class TestState + { + public string Value { get; set; } = string.Empty; + } + + private sealed class TestChatHistoryProvider : ChatHistoryProvider + { + public TestChatHistoryProvider( + Func stateInitializer, + string? stateKey = null) + : base(stateInitializer, stateKey, null, null, null) + { + } + + public TestState GetState(AgentSession? session) => this.GetOrInitializeState(session); + + public void DoSaveState(AgentSession? session, TestState state) => this.SaveState(session, state); + + protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + var state = this.GetOrInitializeState(context.Session); + return new(new[] { new ChatMessage(ChatRole.System, $"History from state: {state.Value}") }); + } + } + + private sealed class TestAgentSession : AgentSession; +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs index 7fccca8d1b..5df661f009 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs @@ -274,12 +274,279 @@ public void InvokedContext_FailureConstructor_ThrowsForNullException() #endregion - private sealed class TestChatHistoryProvider : ChatHistoryProvider + #region InvokingAsync / InvokedAsync Null Check Tests + + [Fact] + public async Task InvokingAsync_NullContext_ThrowsArgumentNullExceptionAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.InvokingAsync(null!).AsTask()); + } + + [Fact] + public async Task InvokedAsync_NullContext_ThrowsArgumentNullExceptionAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.InvokedAsync(null!).AsTask()); + } + + #endregion + + #region InvokingCoreAsync Tests + + [Fact] + public async Task InvokingCoreAsync_CallsProvideChatHistoryAndReturnsMessagesAsync() + { + // Arrange + var historyMessages = new[] { new ChatMessage(ChatRole.User, "History message") }; + var provider = new TestChatHistoryProvider(provideMessages: historyMessages); + var requestMessages = new[] { new ChatMessage(ChatRole.User, "Request message") }; + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, requestMessages); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert + Assert.Equal(2, result.Count); + Assert.Equal("History message", result[0].Text); + Assert.Equal("Request message", result[1].Text); + } + + [Fact] + public async Task InvokingCoreAsync_HistoryAppearsBeforeRequestMessagesAsync() + { + // Arrange + var historyMessages = new[] + { + new ChatMessage(ChatRole.User, "Hist1"), + new ChatMessage(ChatRole.Assistant, "Hist2") + }; + var provider = new TestChatHistoryProvider(provideMessages: historyMessages); + var requestMessages = new[] { new ChatMessage(ChatRole.User, "Req1") }; + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, requestMessages); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert + Assert.Equal(3, result.Count); + Assert.Equal("Hist1", result[0].Text); + Assert.Equal("Hist2", result[1].Text); + Assert.Equal("Req1", result[2].Text); + } + + [Fact] + public async Task InvokingCoreAsync_StampsHistoryMessagesWithChatHistorySourceAsync() + { + // Arrange + var historyMessages = new[] { new ChatMessage(ChatRole.User, "History") }; + var provider = new TestChatHistoryProvider(provideMessages: historyMessages); + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert + Assert.Single(result); + Assert.Equal(AgentRequestMessageSourceType.ChatHistory, result[0].GetAgentRequestMessageSourceType()); + } + + [Fact] + public async Task InvokingCoreAsync_NoFilterAppliedWhenProvideOutputFilterIsNullAsync() { - protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(new ChatMessage[] { new(ChatRole.User, "Test Message") }.Concat(context.RequestMessages)); + // Arrange + var historyMessages = new[] + { + new ChatMessage(ChatRole.User, "User msg"), + new ChatMessage(ChatRole.System, "System msg"), + new ChatMessage(ChatRole.Assistant, "Assistant msg") + }; + var provider = new TestChatHistoryProvider(provideMessages: historyMessages); + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert - all 3 history messages returned (no filter) + Assert.Equal(3, result.Count); + } + + [Fact] + public async Task InvokingCoreAsync_AppliesProvideOutputFilterWhenProvidedAsync() + { + // Arrange + var historyMessages = new[] + { + new ChatMessage(ChatRole.User, "User msg"), + new ChatMessage(ChatRole.System, "System msg"), + new ChatMessage(ChatRole.Assistant, "Assistant msg") + }; + var provider = new TestChatHistoryProvider( + provideMessages: historyMessages, + provideOutputMessageFilter: msgs => msgs.Where(m => m.Role == ChatRole.User)); + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) - => default; + // Assert - only User messages remain after filter + Assert.Single(result); + Assert.Equal("User msg", result[0].Text); } + + [Fact] + public async Task InvokingCoreAsync_ReturnsEmptyHistoryByDefaultAsync() + { + // Arrange - provider that doesn't override ProvideChatHistoryAsync (uses base default) + var provider = new DefaultChatHistoryProvider(); + var requestMessages = new[] { new ChatMessage(ChatRole.User, "Hello") }; + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, requestMessages); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert - only the request message (no history) + Assert.Single(result); + Assert.Equal("Hello", result[0].Text); + } + + #endregion + + #region InvokedCoreAsync Tests + + [Fact] + public async Task InvokedCoreAsync_CallsStoreChatHistoryWithFilteredMessagesAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + var externalMessage = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMessage = new ChatMessage(ChatRole.User, "From history") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "source"); + var responseMessages = new[] { new ChatMessage(ChatRole.Assistant, "Response") }; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, new[] { externalMessage, chatHistoryMessage }, responseMessages); + + // Act + await provider.InvokedAsync(context); + + // Assert - default filter excludes ChatHistory-sourced messages + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("External", storedRequest[0].Text); + Assert.Same(responseMessages, provider.LastStoredContext.ResponseMessages); + } + + [Fact] + public async Task InvokedCoreAsync_SkipsStorageWhenInvokeExceptionIsNotNullAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "msg")], new InvalidOperationException("Failed")); + + // Act + await provider.InvokedAsync(context); + + // Assert - StoreChatHistoryAsync was NOT called + Assert.Null(provider.LastStoredContext); + } + + [Fact] + public async Task InvokedCoreAsync_UsesCustomStoreInputFilterAsync() + { + // Arrange - filter that only keeps System messages + var provider = new TestChatHistoryProvider( + storeInputMessageFilter: msgs => msgs.Where(m => m.Role == ChatRole.System)); + var messages = new[] + { + new ChatMessage(ChatRole.User, "User msg"), + new ChatMessage(ChatRole.System, "System msg") + }; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, messages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(context); + + // Assert - only System messages were passed to store + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("System msg", storedRequest[0].Text); + } + + [Fact] + public async Task InvokedCoreAsync_DefaultFilterExcludesChatHistorySourcedMessagesAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + var external = new ChatMessage(ChatRole.User, "External"); + var fromHistory = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var fromContext = new ChatMessage(ChatRole.User, "Context") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "src"); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [external, fromHistory, fromContext], []); + + // Act + await provider.InvokedAsync(context); + + // Assert - External and AIContextProvider messages kept, ChatHistory excluded + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Equal(2, storedRequest.Count); + Assert.Equal("External", storedRequest[0].Text); + Assert.Equal("Context", storedRequest[1].Text); + } + + [Fact] + public async Task InvokedCoreAsync_PassesResponseMessagesToStoreAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + var responseMessages = new[] { new ChatMessage(ChatRole.Assistant, "Resp1"), new ChatMessage(ChatRole.Assistant, "Resp2") }; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "msg")], responseMessages); + + // Act + await provider.InvokedAsync(context); + + // Assert + Assert.NotNull(provider.LastStoredContext); + Assert.Same(responseMessages, provider.LastStoredContext!.ResponseMessages); + } + + #endregion + + private sealed class TestChatHistoryProvider : ChatHistoryProvider + { + private readonly IEnumerable? _provideMessages; + + public InvokedContext? LastStoredContext { get; private set; } + + public TestChatHistoryProvider( + IEnumerable? provideMessages = null, + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : base(provideOutputMessageFilter, storeInputMessageFilter) + { + this._provideMessages = provideMessages; + } + + protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) + => new(this._provideMessages ?? []); + + protected override ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + this.LastStoredContext = context; + return default; + } + } + + /// + /// A provider that uses only base class defaults (no overrides of ProvideChatHistoryAsync/StoreChatHistoryAsync). + /// + private sealed class DefaultChatHistoryProvider : ChatHistoryProvider; } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs index b907529241..ebe1131ab7 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs @@ -446,7 +446,7 @@ public async Task InvokingAsync_OutputFilter_FiltersOutputMessagesAsync() var session = CreateMockSession(); var provider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions { - RetrievalOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User) + ProvideOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User) }); provider.SetMessages(session, [ diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs index 12bd467e71..4c142d3c1e 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs @@ -881,12 +881,12 @@ public async Task InvokedAsync_CustomStorageInputFilter_OverridesDefaultAsync() this.SkipIfEmulatorNotAvailable(); var session = CreateMockSession(); var conversationId = Guid.NewGuid().ToString(); - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, - _ => new CosmosChatHistoryProvider.State(conversationId)) - { - // Custom filter: only store External messages (also exclude AIContextProvider) - StorageInputMessageFilter = messages => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) - }; + using var provider = new CosmosChatHistoryProvider( + this._connectionString, + s_testDatabaseId, + TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId), + storeInputMessageFilter: messages => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External)); var requestMessages = new[] { @@ -919,12 +919,12 @@ public async Task InvokingAsync_RetrievalOutputFilter_FiltersRetrievedMessagesAs this.SkipIfEmulatorNotAvailable(); var session = CreateMockSession(); var conversationId = Guid.NewGuid().ToString(); - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, - _ => new CosmosChatHistoryProvider.State(conversationId)) - { - // Only return User messages when retrieving - RetrievalOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User) - }; + using var provider = new CosmosChatHistoryProvider( + this._connectionString, + s_testDatabaseId, + TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId), + provideOutputMessageFilter: messages => messages.Where(m => m.Role == ChatRole.User)); var requestMessages = new[] { @@ -943,7 +943,7 @@ public async Task InvokingAsync_RetrievalOutputFilter_FiltersRetrievedMessagesAs var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var messages = (await provider.InvokingAsync(invokingContext)).ToList(); - // Assert - Only User messages returned (System and Assistant filtered by RetrievalOutputMessageFilter) + // Assert - Only User messages returned (System and Assistant filtered by ProvideOutputMessageFilter) Assert.Single(messages); Assert.Equal("User message", messages[0].Text); Assert.Equal(ChatRole.User, messages[0].Role); diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs index f69fb3d636..b8e3c57af6 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs @@ -115,8 +115,8 @@ public void Clone_CreatesDeepCopyWithSameValues() const string Description = "Test description"; var tools = new List { AIFunctionFactory.Create(() => "test") }; - var mockChatHistoryProvider = new Mock().Object; - var mockAIContextProvider = new Mock().Object; + var mockChatHistoryProvider = new Mock(null, null).Object; + var mockAIContextProvider = new Mock(null, null).Object; var original = new ChatClientAgentOptions() { @@ -149,8 +149,8 @@ public void Clone_CreatesDeepCopyWithSameValues() public void Clone_WithoutProvidingChatOptions_ClonesCorrectly() { // Arrange - var mockChatHistoryProvider = new Mock().Object; - var mockAIContextProvider = new Mock().Object; + var mockChatHistoryProvider = new Mock(null, null).Object; + var mockAIContextProvider = new Mock(null, null).Object; var original = new ChatClientAgentOptions { diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs index d90fe5153a..12446c89c0 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs @@ -488,7 +488,7 @@ public async Task RunAsyncInvokesAIContextProviderAndUsesResultAsync() }) .ReturnsAsync(new ChatResponse(responseMessages)); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -559,7 +559,7 @@ public async Task RunAsyncInvokesAIContextProviderWhenGetResponseFailsAsync() It.IsAny())) .Throws(new InvalidOperationException("downstream failure")); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -617,7 +617,7 @@ public async Task RunAsyncInvokesAIContextProviderAndSucceedsWithEmptyAIContextA }) .ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")])); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -677,7 +677,7 @@ public async Task RunAsyncInvokesMultipleAIContextProvidersInOrderAsync() .ReturnsAsync(new ChatResponse(responseMessages)); // Provider 1: adds a system message and a tool - var mockProvider1 = new Mock(); + var mockProvider1 = new Mock(null, null); mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1"); mockProvider1 .Protected() @@ -696,7 +696,7 @@ public async Task RunAsyncInvokesMultipleAIContextProvidersInOrderAsync() // Provider 2: adds another system message and verifies it receives accumulated context from provider 1 AIContext? provider2ReceivedContext = null; - var mockProvider2 = new Mock(); + var mockProvider2 = new Mock(null, null); mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2"); mockProvider2 .Protected() @@ -784,7 +784,7 @@ public async Task RunAsyncInvokesMultipleAIContextProvidersOnFailureAsync() It.IsAny())) .ThrowsAsync(new InvalidOperationException("downstream failure")); - var mockProvider1 = new Mock(); + var mockProvider1 = new Mock(null, null); mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1"); mockProvider1 .Protected() @@ -801,7 +801,7 @@ public async Task RunAsyncInvokesMultipleAIContextProvidersOnFailureAsync() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(new ValueTask()); - var mockProvider2 = new Mock(); + var mockProvider2 = new Mock(null, null); mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2"); mockProvider2 .Protected() @@ -869,7 +869,7 @@ public async Task RunStreamingAsyncInvokesMultipleAIContextProvidersAsync() }) .Returns(ToAsyncEnumerableAsync(responseUpdates)); - var mockProvider1 = new Mock(); + var mockProvider1 = new Mock(null, null); mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1"); mockProvider1 .Protected() @@ -886,7 +886,7 @@ public async Task RunStreamingAsyncInvokesMultipleAIContextProvidersAsync() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(new ValueTask()); - var mockProvider2 = new Mock(); + var mockProvider2 = new Mock(null, null); mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2"); mockProvider2 .Protected() @@ -1828,7 +1828,7 @@ public async Task RunStreamingAsyncInvokesAIContextProviderAndUsesResultAsync() }) .Returns(ToAsyncEnumerableAsync(responseUpdates)); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -1907,7 +1907,7 @@ public async Task RunStreamingAsyncInvokesAIContextProviderWhenGetResponseFailsA It.IsAny())) .Throws(new InvalidOperationException("downstream failure")); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs index 84285ff9c4..64835f2b2f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs @@ -338,7 +338,7 @@ public async Task RunAsync_WhenContinuationTokenProvided_SkipsSessionMessagePopu List capturedMessages = []; // Create a mock chat history provider that would normally provide messages - var mockChatHistoryProvider = new Mock(); + var mockChatHistoryProvider = new Mock(null, null); mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() @@ -346,7 +346,7 @@ public async Task RunAsync_WhenContinuationTokenProvided_SkipsSessionMessagePopu .ReturnsAsync([new(ChatRole.User, "Message from chat history provider")]); // Create a mock AI context provider that would normally provide context - var mockContextProvider = new Mock(); + var mockContextProvider = new Mock(null, null); mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() @@ -407,7 +407,7 @@ public async Task RunStreamingAsync_WhenContinuationTokenProvided_SkipsSessionMe List capturedMessages = []; // Create a mock chat history provider that would normally provide messages - var mockChatHistoryProvider = new Mock(); + var mockChatHistoryProvider = new Mock(null, null); mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() @@ -415,7 +415,7 @@ public async Task RunStreamingAsync_WhenContinuationTokenProvided_SkipsSessionMe .ReturnsAsync([new(ChatRole.User, "Message from chat history provider")]); // Create a mock AI context provider that would normally provide context - var mockContextProvider = new Mock(); + var mockContextProvider = new Mock(null, null); mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() @@ -638,7 +638,7 @@ public async Task RunStreamingAsync_WhenResumingStreaming_UsesUpdatesFromInitial .Returns(ToAsyncEnumerableAsync(returnUpdates)); List capturedMessagesAddedToProvider = []; - var mockChatHistoryProvider = new Mock(); + var mockChatHistoryProvider = new Mock(null, null); mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() @@ -647,7 +647,7 @@ public async Task RunStreamingAsync_WhenResumingStreaming_UsesUpdatesFromInitial .Returns(new ValueTask()); AIContextProvider.InvokedContext? capturedInvokedContext = null; - var mockContextProvider = new Mock(); + var mockContextProvider = new Mock(null, null); mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() @@ -702,7 +702,7 @@ public async Task RunStreamingAsync_WhenResumingStreaming_UsesInputMessagesFromI .Returns(ToAsyncEnumerableAsync(Array.Empty())); List capturedMessagesAddedToProvider = []; - var mockChatHistoryProvider = new Mock(); + var mockChatHistoryProvider = new Mock(null, null); mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() @@ -711,7 +711,7 @@ public async Task RunStreamingAsync_WhenResumingStreaming_UsesInputMessagesFromI .Returns(new ValueTask()); AIContextProvider.InvokedContext? capturedInvokedContext = null; - var mockContextProvider = new Mock(); + var mockContextProvider = new Mock(null, null); mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs index 4731805af7..423c867abd 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs @@ -185,7 +185,7 @@ public async Task RunAsync_UsesChatHistoryProvider_WhenProvidedAndNoConversation It.IsAny(), It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")])); - Mock mockChatHistoryProvider = new(); + Mock mockChatHistoryProvider = new(null, null); mockChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -240,7 +240,7 @@ public async Task RunAsync_NotifiesChatHistoryProvider_OnFailureAsync() It.IsAny(), It.IsAny())).Throws(new InvalidOperationException("Test Error")); - Mock mockChatHistoryProvider = new(); + Mock mockChatHistoryProvider = new(null, null); mockChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -311,7 +311,7 @@ public async Task RunAsync_UsesOverrideChatHistoryProvider_WhenProvidedViaAdditi It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")])); // Arrange a chat history provider to override the factory provided one. - Mock mockOverrideChatHistoryProvider = new(); + Mock mockOverrideChatHistoryProvider = new(null, null); mockOverrideChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -324,7 +324,7 @@ public async Task RunAsync_UsesOverrideChatHistoryProvider_WhenProvidedViaAdditi // Arrange a chat history provider to provide to the agent at construction time. // This one shouldn't be used since it is being overridden. - Mock mockAgentOptionsChatHistoryProvider = new(); + Mock mockAgentOptionsChatHistoryProvider = new(null, null); mockAgentOptionsChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())