.NET: [BREAKING] Update providers in such a way that they can participate in a pipeline (#3846)

* Make providers pipeline capable

* Fix unit tests

* Move source stamping to providers from base class

* Also update samples.

* Address PR comments

* Rename AsAgentRequestMessageSourcedMessage to WithAgentRequestMessageSource
This commit is contained in:
westey
2026-02-11 16:24:37 +00:00
committed by GitHub
Unverified
parent 65f7aff145
commit 1d158c24be
25 changed files with 332 additions and 549 deletions
@@ -125,10 +125,15 @@ namespace SampleApp
protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
var inputContext = context.AIContext;
var userInfo = context.Session?.StateBag.GetValue<UserInfo>(nameof(UserInfoMemory))
?? this._stateInitializer.Invoke(context.Session);
StringBuilder instructions = new();
if (!string.IsNullOrEmpty(inputContext.Instructions))
{
instructions.AppendLine(inputContext.Instructions);
}
// If we don't already know the user's name and age, add instructions to ask for them, otherwise just provide what we have to the context.
instructions
@@ -143,7 +148,9 @@ namespace SampleApp
return new ValueTask<AIContext>(new AIContext
{
Instructions = instructions.ToString()
Instructions = instructions.ToString(),
Messages = inputContext.Messages,
Tools = inputContext.Tools
});
}
}
@@ -126,7 +126,9 @@ namespace SampleApp
var messages = records.ConvertAll(x => JsonSerializer.Deserialize<ChatMessage>(x.SerializedMessage!)!);
messages.Reverse();
return messages;
return messages
.Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!))
.Concat(context.RequestMessages);
}
protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
@@ -87,6 +87,7 @@ namespace SampleApp
protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
var inputContext = context.AIContext;
var todoItems = GetTodoItems(context.Session);
StringBuilder outputMessageBuilder = new();
@@ -106,12 +107,19 @@ namespace SampleApp
return new ValueTask<AIContext>(new AIContext
{
Tools =
[
Instructions = inputContext.Instructions,
Tools = (inputContext.Tools ?? []).Concat(new AITool[]
{
AIFunctionFactory.Create((string item) => AddTodoItem(context.Session, item), "AddTodoItem", "Adds an item to the todo list."),
AIFunctionFactory.Create((int index) => RemoveTodoItem(context.Session, index), "RemoveTodoItem", "Removes an item from the todo list. Index is zero based.")
],
Messages = [new MEAI.ChatMessage(ChatRole.User, outputMessageBuilder.ToString())]
}).ToList(),
Messages =
(inputContext.Messages ?? [])
.Concat(
[
new MEAI.ChatMessage(ChatRole.User, outputMessageBuilder.ToString()).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!)
])
.ToList()
});
}
@@ -142,6 +150,7 @@ namespace SampleApp
{
protected override async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
var inputContext = context.AIContext;
var events = await loadNextThreeCalendarEvents();
StringBuilder outputMessageBuilder = new();
@@ -153,10 +162,15 @@ namespace SampleApp
return new()
{
Instructions = inputContext.Instructions,
Messages =
[
new MEAI.ChatMessage(ChatRole.User, outputMessageBuilder.ToString()),
]
(inputContext.Messages ?? [])
.Concat(
[
new MEAI.ChatMessage(ChatRole.User, outputMessageBuilder.ToString()).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!)
])
.ToList(),
Tools = inputContext.Tools
};
}
}
@@ -177,16 +191,13 @@ namespace SampleApp
protected override async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
// Invoke all the sub providers.
var tasks = this._providers.Select(provider => provider.InvokingAsync(context, cancellationToken).AsTask());
var results = await Task.WhenAll(tasks);
// Combine the results from each sub provider.
return new AIContext
var currentAIContext = context.AIContext;
foreach (var provider in this._providers)
{
Tools = results.SelectMany(r => r.Tools ?? []).ToList(),
Messages = results.SelectMany(r => r.Messages ?? []).ToList(),
Instructions = string.Join("\n", results.Select(r => r.Instructions).Where(s => !string.IsNullOrEmpty(s)))
};
currentAIContext = await provider.InvokingAsync(new InvokingContext(context.Agent, context.Session, currentAIContext), cancellationToken);
}
return currentAIContext;
}
}
}
@@ -2,7 +2,6 @@
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
@@ -31,25 +30,6 @@ namespace Microsoft.Agents.AI;
/// </remarks>
public abstract class AIContextProvider
{
private readonly string _sourceId;
/// <summary>
/// Initializes a new instance of the <see cref="AIContextProvider"/> class.
/// </summary>
protected AIContextProvider()
{
this._sourceId = this.GetType().FullName!;
}
/// <summary>
/// Initializes a new instance of the <see cref="AIContextProvider"/> class with the specified source id.
/// </summary>
/// <param name="sourceId">The source id to stamp on <see cref="ChatMessage.AdditionalProperties"/> for each messages produced by the <see cref="AIContextProvider"/>.</param>
protected AIContextProvider(string sourceId)
{
this._sourceId = sourceId;
}
/// <summary>
/// Gets the key used to store the provider state in the <see cref="AgentSession.StateBag"/>.
/// </summary>
@@ -77,20 +57,8 @@ public abstract class AIContextProvider
/// </list>
/// </para>
/// </remarks>
public async ValueTask<AIContext> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
var aiContext = await this.InvokingCoreAsync(context, cancellationToken).ConfigureAwait(false);
if (aiContext.Messages is null)
{
return aiContext;
}
aiContext.Messages = aiContext.Messages
.Select(message => message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.AIContextProvider, this._sourceId))
.ToList();
return aiContext;
}
public ValueTask<AIContext> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> this.InvokingCoreAsync(context, cancellationToken);
/// <summary>
/// Called at the start of agent invocation to provide additional context.
@@ -205,20 +173,20 @@ public abstract class AIContextProvider
public sealed class InvokingContext
{
/// <summary>
/// Initializes a new instance of the <see cref="InvokingContext"/> class with the specified request messages.
/// Initializes a new instance of the <see cref="InvokingContext"/> class.
/// </summary>
/// <param name="agent">The agent being invoked.</param>
/// <param name="session">The session associated with the agent invocation.</param>
/// <param name="requestMessages">The messages to be used by the agent for this invocation.</param>
/// <exception cref="ArgumentNullException"><paramref name="requestMessages"/> is <see langword="null"/>.</exception>
/// <param name="aiContext">The AI context to be used by the agent for this invocation.</param>
/// <exception cref="ArgumentNullException"><paramref name="agent"/> or <paramref name="aiContext"/> is <see langword="null"/>.</exception>
public InvokingContext(
AIAgent agent,
AgentSession? session,
IEnumerable<ChatMessage> requestMessages)
AIContext aiContext)
{
this.Agent = Throw.IfNull(agent);
this.Session = session;
this.RequestMessages = Throw.IfNull(requestMessages);
this.AIContext = Throw.IfNull(aiContext);
}
/// <summary>
@@ -232,12 +200,23 @@ public abstract class AIContextProvider
public AgentSession? Session { get; }
/// <summary>
/// Gets the caller provided messages that will be used by the agent for this invocation.
/// Gets the <see cref="AIContext"/> being built for the current invocation. Context providers can modify
/// and return or return a new <see cref="AIContext"/> instance to provide additional context for the invocation.
/// </summary>
/// <value>
/// A collection of <see cref="ChatMessage"/> instances representing new messages that were provided by the caller.
/// </value>
public IEnumerable<ChatMessage> RequestMessages { get; set { field = Throw.IfNull(value); } }
/// <remarks>
/// <para>
/// If multiple <see cref="AIContextProvider"/> instances are used in the same invocation, each <see cref="AIContextProvider"/>
/// will receive the context returned by the previous <see cref="AIContextProvider"/> allowing them to build on top of each other's context.
/// </para>
/// <para>
/// The first <see cref="AIContextProvider"/> in the invocation pipeline will receive an <see cref="AIContext"/> instance
/// that already contains the caller provided messages that will be used by the agent for this invocation.
/// </para>
/// <para>
/// It may also contain messages from chat history, if a <see cref="ChatHistoryProvider"/> is being used.
/// </para>
/// </remarks>
public AIContext AIContext { get; }
}
/// <summary>
@@ -255,7 +234,7 @@ public abstract class AIContextProvider
/// </summary>
/// <param name="agent">The agent being invoked.</param>
/// <param name="session">The session associated with the agent invocation.</param>
/// <param name="requestMessages">The caller provided messages that were used by the agent for this invocation.</param>
/// <param name="requestMessages">The messages that were used by the agent for this invocation.</param>
/// <exception cref="ArgumentNullException"><paramref name="requestMessages"/> is <see langword="null"/>.</exception>
public InvokedContext(
AIAgent agent,
@@ -278,11 +257,10 @@ public abstract class AIContextProvider
public AgentSession? Session { get; }
/// <summary>
/// Gets the caller provided messages that were used by the agent for this invocation.
/// Gets the messages that were used by the agent for this invocation.
/// </summary>
/// <value>
/// A collection of <see cref="ChatMessage"/> instances representing new messages that were provided by the caller.
/// This does not include any <see cref="AIContextProvider"/> supplied messages.
/// A collection of <see cref="ChatMessage"/> instances representing all messages that were used by the agent for this invocation.
/// </value>
public IEnumerable<ChatMessage> RequestMessages { get; set { field = Throw.IfNull(value); } }
@@ -2,7 +2,6 @@
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
@@ -40,25 +39,6 @@ namespace Microsoft.Agents.AI;
/// </remarks>
public abstract class ChatHistoryProvider
{
private readonly string _sourceId;
/// <summary>
/// Initializes a new instance of the <see cref="ChatHistoryProvider"/> class.
/// </summary>
protected ChatHistoryProvider()
{
this._sourceId = this.GetType().FullName!;
}
/// <summary>
/// Initializes a new instance of the <see cref="ChatHistoryProvider"/> class with the specified source id.
/// </summary>
/// <param name="sourceId">The source id to stamp on <see cref="ChatMessage.AdditionalProperties"/> for each messages produced by the <see cref="ChatHistoryProvider"/>.</param>
protected ChatHistoryProvider(string sourceId)
{
this._sourceId = sourceId;
}
/// <summary>
/// Gets the key used to store the provider state in the <see cref="AgentSession.StateBag"/>.
/// </summary>
@@ -94,12 +74,8 @@ public abstract class ChatHistoryProvider
/// </list>
/// </para>
/// </remarks>
public async ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
var messages = await this.InvokingCoreAsync(context, cancellationToken).ConfigureAwait(false);
return messages.Select(message => message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.ChatHistory, this._sourceId));
}
public ValueTask<IEnumerable<ChatMessage>> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> this.InvokingCoreAsync(context, cancellationToken);
/// <summary>
/// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation.
@@ -231,7 +207,7 @@ public abstract class ChatHistoryProvider
/// </summary>
/// <param name="agent">The agent being invoked.</param>
/// <param name="session">The session associated with the agent invocation.</param>
/// <param name="requestMessages">The new messages to be used by the agent for this invocation.</param>
/// <param name="requestMessages">The messages to be used by the agent for this invocation.</param>
/// <exception cref="ArgumentNullException"><paramref name="requestMessages"/> is <see langword="null"/>.</exception>
public InvokingContext(
AIAgent agent,
@@ -254,11 +230,22 @@ public abstract class ChatHistoryProvider
public AgentSession? Session { get; }
/// <summary>
/// Gets the caller provided messages that will be used by the agent for this invocation.
/// Gets the messages that will be used by the agent for this invocation. <see cref="ChatHistoryProvider"/> instances can modify
/// and return or return a new message list to add additional messages for the invocation.
/// </summary>
/// <value>
/// A collection of <see cref="ChatMessage"/> instances representing new messages that were provided by the caller.
/// A collection of <see cref="ChatMessage"/> instances representing the messages that will be used by the agent for this invocation.
/// </value>
/// <remarks>
/// <para>
/// If multiple <see cref="ChatHistoryProvider"/> instances are used in the same invocation, each <see cref="ChatHistoryProvider"/>
/// will receive the messages returned by the previous <see cref="ChatHistoryProvider"/> allowing them to build on top of each other's context.
/// </para>
/// <para>
/// The first <see cref="ChatHistoryProvider"/> in the invocation pipeline will receive the
/// caller provided messages.
/// </para>
/// </remarks>
public IEnumerable<ChatMessage> RequestMessages { get; set { field = Throw.IfNull(value); } }
}
@@ -54,7 +54,7 @@ public static class ChatMessageExtensions
/// If the message is already tagged with the provided source type and source id, it is returned as is.
/// Otherwise, a cloned message is returned with the appropriate tagging in the AdditionalProperties.
/// </remarks>
public static ChatMessage AsAgentRequestMessageSourcedMessage(this ChatMessage message, AgentRequestMessageSourceType sourceType, string? sourceId = null)
public static ChatMessage WithAgentRequestMessageSource(this ChatMessage message, AgentRequestMessageSourceType sourceType, string? sourceId = null)
{
if (message.AdditionalProperties != null
// Check if the message was already tagged with the required source type and source id
@@ -115,7 +115,9 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
state.Messages = (await this.ChatReducer.ReduceAsync(state.Messages, cancellationToken).ConfigureAwait(false)).ToList();
}
return state.Messages;
return state.Messages
.Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!))
.Concat(context.RequestMessages);
}
/// <inheritdoc />
@@ -257,7 +257,9 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
messages.Reverse();
}
return messages;
return messages
.Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!))
.Concat(context.RequestMessages);
}
/// <inheritdoc />
@@ -108,12 +108,13 @@ public sealed class Mem0Provider : AIContextProvider
{
Throw.IfNull(context);
var inputContext = context.AIContext;
var state = this.GetOrInitializeState(context.Session);
var searchScope = state?.SearchScope ?? new Mem0ProviderScope();
string queryText = string.Join(
Environment.NewLine,
context.RequestMessages
(inputContext.Messages ?? [])
.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External)
.Where(m => !string.IsNullOrWhiteSpace(m.Text))
.Select(m => m.Text));
@@ -131,6 +132,9 @@ public sealed class Mem0Provider : AIContextProvider
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)
{
@@ -157,7 +161,12 @@ public sealed class Mem0Provider : AIContextProvider
return new AIContext
{
Messages = [new ChatMessage(ChatRole.User, outputMessageText)]
Instructions = inputContext.Instructions,
Messages =
(inputContext.Messages ?? [])
.Concat(outputMessage is not null ? [outputMessage] : [])
.ToList(),
Tools = inputContext.Tools
};
}
catch (ArgumentException)
@@ -176,7 +185,7 @@ public sealed class Mem0Provider : AIContextProvider
searchScope.ThreadId,
this.SanitizeLogData(searchScope.UserId));
}
return new AIContext();
return inputContext;
}
}
@@ -52,7 +52,10 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
=> this.GetOrInitializeState(session).Messages.AddRange(messages);
protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new(this.GetOrInitializeState(context.Session).Messages.AsReadOnly());
=> new(this.GetOrInitializeState(context.Session)
.Messages
.Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!))
.Concat(context.RequestMessages));
protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
{
@@ -219,7 +219,6 @@ public sealed partial class ChatClientAgent : AIAgent
(ChatClientAgentSession safeSession,
ChatOptions? chatOptions,
List<ChatMessage> inputMessagesForProviders,
List<ChatMessage> inputMessagesForChatClient,
ChatClientAgentContinuationToken? continuationToken) =
await this.PrepareSessionAndMessagesAsync(session, inputMessages, options, cancellationToken).ConfigureAwait(false);
@@ -243,8 +242,8 @@ public sealed partial class ChatClientAgent : AIAgent
}
catch (Exception ex)
{
await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), cancellationToken).ConfigureAwait(false);
await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), cancellationToken).ConfigureAwait(false);
throw;
}
@@ -258,8 +257,8 @@ public sealed partial class ChatClientAgent : AIAgent
}
catch (Exception ex)
{
await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), cancellationToken).ConfigureAwait(false);
await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), cancellationToken).ConfigureAwait(false);
throw;
}
@@ -285,8 +284,8 @@ public sealed partial class ChatClientAgent : AIAgent
}
catch (Exception ex)
{
await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), cancellationToken).ConfigureAwait(false);
await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), cancellationToken).ConfigureAwait(false);
throw;
}
}
@@ -298,10 +297,10 @@ public sealed partial class ChatClientAgent : AIAgent
this.UpdateSessionConversationId(safeSession, chatResponse.ConversationId, cancellationToken);
// To avoid inconsistent state we only notify the session of the input messages if no error occurs after the initial request.
await this.NotifyChatHistoryProviderOfNewMessagesAsync(safeSession, GetInputMessages(inputMessagesForProviders, continuationToken), chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyChatHistoryProviderOfNewMessagesAsync(safeSession, GetInputMessages(inputMessagesForChatClient, continuationToken), chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false);
// Notify the AIContextProvider of all new messages.
await this.NotifyAIContextProviderOfSuccessAsync(safeSession, GetInputMessages(inputMessagesForProviders, continuationToken), chatResponse.Messages, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfSuccessAsync(safeSession, GetInputMessages(inputMessagesForChatClient, continuationToken), chatResponse.Messages, cancellationToken).ConfigureAwait(false);
}
/// <inheritdoc/>
@@ -382,7 +381,6 @@ public sealed partial class ChatClientAgent : AIAgent
(ChatClientAgentSession safeSession,
ChatOptions? chatOptions,
List<ChatMessage> inputMessagesForProviders,
List<ChatMessage> inputMessagesForChatClient,
ChatClientAgentContinuationToken? _) =
await this.PrepareSessionAndMessagesAsync(session, inputMessages, options, cancellationToken).ConfigureAwait(false);
@@ -403,8 +401,8 @@ public sealed partial class ChatClientAgent : AIAgent
}
catch (Exception ex)
{
await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, inputMessagesForProviders, chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, inputMessagesForProviders, cancellationToken).ConfigureAwait(false);
await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, inputMessagesForChatClient, chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, inputMessagesForChatClient, cancellationToken).ConfigureAwait(false);
throw;
}
@@ -421,10 +419,10 @@ public sealed partial class ChatClientAgent : AIAgent
}
// Only notify the session of new messages if the chatResponse was successful to avoid inconsistent message state in the session.
await this.NotifyChatHistoryProviderOfNewMessagesAsync(safeSession, inputMessagesForProviders, chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false);
await this.NotifyChatHistoryProviderOfNewMessagesAsync(safeSession, inputMessagesForChatClient, chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false);
// Notify the AIContextProvider of all new messages.
await this.NotifyAIContextProviderOfSuccessAsync(safeSession, inputMessagesForProviders, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
await this.NotifyAIContextProviderOfSuccessAsync(safeSession, inputMessagesForChatClient, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
var agentResponse = agentResponseFactoryFunc(chatResponse);
@@ -629,7 +627,6 @@ public sealed partial class ChatClientAgent : AIAgent
<(
ChatClientAgentSession AgentSession,
ChatOptions? ChatOptions,
List<ChatMessage> inputMessagesForProviders,
List<ChatMessage> InputMessagesForChatClient,
ChatClientAgentContinuationToken? ContinuationToken
)> PrepareSessionAndMessagesAsync(
@@ -659,7 +656,6 @@ public sealed partial class ChatClientAgent : AIAgent
throw new InvalidOperationException("Input messages are not allowed when continuing a background response using a continuation token.");
}
List<ChatMessage> inputMessagesForProviders = [];
List<ChatMessage> inputMessagesForChatClient = [];
// Populate the session messages only if we are not continuing an existing response as it's not allowed
@@ -668,43 +664,45 @@ public sealed partial class ChatClientAgent : AIAgent
ChatHistoryProvider? chatHistoryProvider = this.ResolveChatHistoryProvider(chatOptions, typedSession);
// Add any existing messages from the session to the messages to be sent to the chat client.
// The ChatHistoryProvider returns the merged result (history + input messages).
if (chatHistoryProvider is not null)
{
var invokingContext = new ChatHistoryProvider.InvokingContext(this, typedSession, inputMessages);
var providerMessages = await chatHistoryProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
inputMessagesForChatClient.AddRange(providerMessages);
}
// Add the input messages before getting context from AIContextProvider.
inputMessagesForProviders.AddRange(inputMessages);
inputMessagesForChatClient.AddRange(inputMessages);
else
{
inputMessagesForChatClient.AddRange(inputMessages);
}
// If we have an AIContextProvider, we should get context from it, and update our
// messages and options with the additional context.
// The AIContextProvider returns the accumulated AIContext (original + new contributions).
if (this._agentOptions?.AIContextProvider is { } aiContextProvider)
{
var invokingContext = new AIContextProvider.InvokingContext(this, typedSession, inputMessages);
var aiContext = await aiContextProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
if (aiContext.Messages is { Count: > 0 })
var aiContext = new AIContext
{
inputMessagesForProviders.AddRange(aiContext.Messages);
inputMessagesForChatClient.AddRange(aiContext.Messages);
}
Instructions = chatOptions?.Instructions,
Messages = inputMessagesForChatClient.ToList(),
Tools = chatOptions?.Tools as List<AITool> ?? chatOptions?.Tools?.ToList()
};
var invokingContext = new AIContextProvider.InvokingContext(this, typedSession, aiContext);
aiContext = await aiContextProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
if (aiContext.Tools is { Count: > 0 })
// Use the returned messages, tools and instructions directly since the provider accumulated them.
inputMessagesForChatClient = aiContext.Messages as List<ChatMessage> ?? aiContext.Messages?.ToList() ?? [];
if (chatOptions?.Tools is { Count: > 0 } || aiContext.Tools is { Count: > 0 })
{
chatOptions ??= new();
chatOptions.Tools ??= [];
foreach (AITool tool in aiContext.Tools)
{
chatOptions.Tools.Add(tool);
}
chatOptions.Tools = aiContext.Tools;
}
if (aiContext.Instructions is not null)
if (chatOptions?.Instructions is not null || aiContext.Instructions is not null)
{
chatOptions ??= new();
chatOptions.Instructions = string.IsNullOrWhiteSpace(chatOptions.Instructions) ? aiContext.Instructions : $"{chatOptions.Instructions}\n{aiContext.Instructions}";
chatOptions.Instructions = aiContext.Instructions;
}
}
}
@@ -727,7 +725,7 @@ public sealed partial class ChatClientAgent : AIAgent
chatOptions.ConversationId = typedSession.ConversationId;
}
return (typedSession, chatOptions, inputMessagesForProviders, inputMessagesForChatClient, continuationToken);
return (typedSession, chatOptions, inputMessagesForChatClient, continuationToken);
}
private void UpdateSessionConversationId(ChatClientAgentSession session, string? responseConversationId, CancellationToken cancellationToken)
@@ -141,6 +141,7 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
{
_ = Throw.IfNull(context);
var inputContext = context.AIContext;
var state = this.GetOrInitializeState(context.Session);
var searchScope = state?.SearchScope ?? new ChatHistoryMemoryProviderScope();
@@ -158,21 +159,26 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
description: this._toolDescription)
];
// Expose search tool for on-demand invocation by the model
return new AIContext { Tools = tools };
// Expose search tool for on-demand invocation by the model, accumulated with the input context
return new AIContext
{
Instructions = inputContext.Instructions,
Messages = inputContext.Messages,
Tools = (inputContext.Tools ?? []).Concat(tools).ToList()
};
}
try
{
// Get the text from the current request messages
var requestText = string.Join("\n", context.RequestMessages
var requestText = string.Join("\n", (inputContext.Messages ?? [])
.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External)
.Where(m => m != null && !string.IsNullOrWhiteSpace(m.Text))
.Select(m => m.Text));
if (string.IsNullOrWhiteSpace(requestText))
{
return new AIContext();
return inputContext;
}
// Search for relevant chat history
@@ -180,12 +186,20 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
if (string.IsNullOrWhiteSpace(contextText))
{
return new AIContext();
return inputContext;
}
return new AIContext
{
Messages = [new ChatMessage(ChatRole.User, contextText)]
Instructions = inputContext.Instructions,
Messages =
(inputContext.Messages ?? [])
.Concat(
[
new ChatMessage(ChatRole.User, contextText).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!)
])
.ToList(),
Tools = inputContext.Tools
};
}
catch (Exception ex)
@@ -201,7 +215,7 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable
this.SanitizeLogData(searchScope.UserId));
}
return new AIContext();
return inputContext;
}
}
@@ -89,10 +89,17 @@ public sealed class TextSearchProvider : AIContextProvider
/// <inheritdoc />
protected override async ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
var inputContext = context.AIContext;
if (this._searchTime != TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke)
{
// Expose the search tool for on-demand invocation.
return new AIContext { Tools = this._tools }; // No automatic message injection.
// Expose the search tool for on-demand invocation, accumulated with the input context.
return new AIContext
{
Instructions = inputContext.Instructions,
Messages = inputContext.Messages,
Tools = (inputContext.Tools ?? []).Concat(this._tools).ToList()
};
}
// Retrieve recent messages from the session state bag.
@@ -101,7 +108,7 @@ public sealed class TextSearchProvider : AIContextProvider
// Aggregate text from memory + current request messages.
var sbInput = new StringBuilder();
var requestMessagesText = context.RequestMessages
var requestMessagesText = (inputContext.Messages ?? [])
.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External)
.Where(x => !string.IsNullOrWhiteSpace(x?.Text)).Select(x => x.Text);
foreach (var messageText in recentMessagesText.Concat(requestMessagesText))
@@ -128,7 +135,7 @@ public sealed class TextSearchProvider : AIContextProvider
if (materialized.Count == 0)
{
return new AIContext();
return inputContext;
}
// Format search results
@@ -141,13 +148,21 @@ public sealed class TextSearchProvider : AIContextProvider
return new AIContext
{
Messages = [new ChatMessage(ChatRole.User, formatted) { AdditionalProperties = new AdditionalPropertiesDictionary() { ["IsTextSearchProviderOutput"] = true } }]
Instructions = inputContext.Instructions,
Messages =
(inputContext.Messages ?? [])
.Concat(
[
new ChatMessage(ChatRole.User, formatted).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!)
])
.ToList(),
Tools = inputContext.Tools
};
}
catch (Exception ex)
{
this._logger?.LogError(ex, "TextSearchProvider: Failed to search for data due to error");
return new AIContext();
return inputContext;
}
}
@@ -179,10 +194,7 @@ public sealed class TextSearchProvider : AIContextProvider
.Concat(context.ResponseMessages ?? [])
.Where(m =>
this._recentMessageRolesIncluded.Contains(m.Role) &&
!string.IsNullOrWhiteSpace(m.Text) &&
// Filter out any messages that were added by this class in InvokingAsync, since we don't want
// a feedback loop where previous search results are used to find new search results.
(m.AdditionalProperties == null || m.AdditionalProperties.TryGetValue("IsTextSearchProviderOutput", out bool isTextSearchProviderOutput) == false || !isTextSearchProviderOutput))
!string.IsNullOrWhiteSpace(m.Text))
.Select(m => m.Text);
// Combine existing messages with new messages, then take the most recent up to the limit.
@@ -3,7 +3,6 @@
using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
@@ -16,110 +15,6 @@ public class AIContextProviderTests
private static readonly AIAgent s_mockAgent = new Mock<AIAgent>().Object;
private static readonly AgentSession s_mockSession = new Mock<AgentSession>().Object;
#region InvokingAsync Message Stamping Tests
[Fact]
public async Task InvokingAsync_StampsMessagesWithSourceTypeAndSourceIdAsync()
{
// Arrange
var provider = new TestAIContextProviderWithMessages();
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]);
// Act
AIContext aiContext = await provider.InvokingAsync(context);
// Assert
Assert.NotNull(aiContext.Messages);
ChatMessage message = aiContext.Messages.Single();
Assert.NotNull(message.AdditionalProperties);
Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution));
var typedAttribution = Assert.IsType<AgentRequestMessageSourceAttribution>(attribution);
Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, typedAttribution.SourceType);
Assert.Equal(typeof(TestAIContextProviderWithMessages).FullName, typedAttribution.SourceId);
}
[Fact]
public async Task InvokingAsync_WithCustomSourceId_StampsMessagesWithCustomSourceIdAsync()
{
// Arrange
const string CustomSourceId = "CustomContextSource";
var provider = new TestAIContextProviderWithCustomSource(CustomSourceId);
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]);
// Act
AIContext aiContext = await provider.InvokingAsync(context);
// Assert
Assert.NotNull(aiContext.Messages);
ChatMessage message = aiContext.Messages.Single();
Assert.NotNull(message.AdditionalProperties);
Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution));
var typedAttribution = Assert.IsType<AgentRequestMessageSourceAttribution>(attribution);
Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, typedAttribution.SourceType);
Assert.Equal(CustomSourceId, typedAttribution.SourceId);
}
[Fact]
public async Task InvokingAsync_DoesNotReStampAlreadyStampedMessagesAsync()
{
// Arrange
var provider = new TestAIContextProviderWithPreStampedMessages();
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]);
// Act
AIContext aiContext = await provider.InvokingAsync(context);
// Assert
Assert.NotNull(aiContext.Messages);
ChatMessage message = aiContext.Messages.Single();
Assert.NotNull(message.AdditionalProperties);
Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution));
var typedAttribution = Assert.IsType<AgentRequestMessageSourceAttribution>(attribution);
Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, typedAttribution.SourceType);
Assert.Equal(typeof(TestAIContextProviderWithPreStampedMessages).FullName, typedAttribution.SourceId);
}
[Fact]
public async Task InvokingAsync_StampsMultipleMessagesAsync()
{
// Arrange
var provider = new TestAIContextProviderWithMultipleMessages();
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]);
// Act
AIContext aiContext = await provider.InvokingAsync(context);
// Assert
Assert.NotNull(aiContext.Messages);
List<ChatMessage> messageList = aiContext.Messages.ToList();
Assert.Equal(3, messageList.Count);
foreach (ChatMessage message in messageList)
{
Assert.NotNull(message.AdditionalProperties);
Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution));
var typedAttribution = Assert.IsType<AgentRequestMessageSourceAttribution>(attribution);
Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, typedAttribution.SourceType);
Assert.Equal(typeof(TestAIContextProviderWithMultipleMessages).FullName, typedAttribution.SourceId);
}
}
[Fact]
public async Task InvokingAsync_WithNullMessages_ReturnsContextWithoutStampingAsync()
{
// Arrange
var provider = new TestAIContextProvider();
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]);
// Act
AIContext aiContext = await provider.InvokingAsync(context);
// Assert
Assert.Null(aiContext.Messages);
}
#endregion
#region Basic Tests
[Fact]
@@ -271,39 +166,33 @@ public class AIContextProviderTests
#region InvokingContext Tests
[Fact]
public void InvokingContext_RequestMessages_SetterThrowsForNull()
public void InvokingContext_Constructor_ThrowsForNullAIContext()
{
// Arrange
var messages = new ReadOnlyCollection<ChatMessage>([new(ChatRole.User, "Hello")]);
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, messages);
// Act & Assert
Assert.Throws<ArgumentNullException>(() => context.RequestMessages = null!);
Assert.Throws<ArgumentNullException>(() => new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, null!));
}
[Fact]
public void InvokingContext_RequestMessages_SetterRoundtrips()
public void InvokingContext_AIContext_ConstructorValueRoundtrips()
{
// Arrange
var initialMessages = new ReadOnlyCollection<ChatMessage>([new(ChatRole.User, "Hello")]);
var newMessages = new List<ChatMessage> { new(ChatRole.User, "New message") };
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, initialMessages);
var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] };
// Act
context.RequestMessages = newMessages;
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, aiContext);
// Assert
Assert.Same(newMessages, context.RequestMessages);
Assert.Same(aiContext, context.AIContext);
}
[Fact]
public void InvokingContext_Agent_ReturnsConstructorValue()
{
// Arrange
var messages = new ReadOnlyCollection<ChatMessage>([new(ChatRole.User, "Hello")]);
var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] };
// Act
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, messages);
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, aiContext);
// Assert
Assert.Same(s_mockAgent, context.Agent);
@@ -313,10 +202,10 @@ public class AIContextProviderTests
public void InvokingContext_Session_ReturnsConstructorValue()
{
// Arrange
var messages = new ReadOnlyCollection<ChatMessage>([new(ChatRole.User, "Hello")]);
var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] };
// Act
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, messages);
var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, aiContext);
// Assert
Assert.Same(s_mockSession, context.Session);
@@ -326,10 +215,10 @@ public class AIContextProviderTests
public void InvokingContext_Session_CanBeNull()
{
// Arrange
var messages = new ReadOnlyCollection<ChatMessage>([new(ChatRole.User, "Hello")]);
var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] };
// Act
var context = new AIContextProvider.InvokingContext(s_mockAgent, null, messages);
var context = new AIContextProvider.InvokingContext(s_mockAgent, null, aiContext);
// Assert
Assert.Null(context.Session);
@@ -339,10 +228,10 @@ public class AIContextProviderTests
public void InvokingContext_Constructor_ThrowsForNullAgent()
{
// Arrange
var messages = new ReadOnlyCollection<ChatMessage>([new(ChatRole.User, "Hello")]);
var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] };
// Act & Assert
Assert.Throws<ArgumentNullException>(() => new AIContextProvider.InvokingContext(null!, s_mockSession, messages));
Assert.Throws<ArgumentNullException>(() => new AIContextProvider.InvokingContext(null!, s_mockSession, aiContext));
}
#endregion
@@ -461,55 +350,4 @@ public class AIContextProviderTests
protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new(new AIContext());
}
private sealed class TestAIContextProviderWithMessages : AIContextProvider
{
protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new(new AIContext
{
Messages = [new ChatMessage(ChatRole.System, "Context Message")]
});
}
private sealed class TestAIContextProviderWithCustomSource : AIContextProvider
{
public TestAIContextProviderWithCustomSource(string sourceId) : base(sourceId)
{
}
protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new(new AIContext
{
Messages = [new ChatMessage(ChatRole.System, "Context Message")]
});
}
private sealed class TestAIContextProviderWithPreStampedMessages : AIContextProvider
{
protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
var message = new ChatMessage(ChatRole.System, "Pre-stamped Message");
message.AdditionalProperties = new AdditionalPropertiesDictionary
{
[AgentRequestMessageSourceAttribution.AdditionalPropertiesKey] = new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!)
};
return new(new AIContext
{
Messages = [message]
});
}
}
private sealed class TestAIContextProviderWithMultipleMessages : AIContextProvider
{
protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new(new AIContext
{
Messages = [
new ChatMessage(ChatRole.System, "Message 1"),
new ChatMessage(ChatRole.User, "Message 2"),
new ChatMessage(ChatRole.Assistant, "Message 3")
]
});
}
}
@@ -18,92 +18,6 @@ public class ChatHistoryProviderTests
private static readonly AIAgent s_mockAgent = new Mock<AIAgent>().Object;
private static readonly AgentSession s_mockSession = new Mock<AgentSession>().Object;
#region InvokingAsync Message Stamping Tests
[Fact]
public async Task InvokingAsync_StampsMessagesWithSourceTypeAndSourceIdAsync()
{
// Arrange
var provider = new TestChatHistoryProvider();
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]);
// Act
IEnumerable<ChatMessage> messages = await provider.InvokingAsync(context);
// Assert
ChatMessage message = messages.Single();
Assert.NotNull(message.AdditionalProperties);
Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution));
var typedAttribution = Assert.IsType<AgentRequestMessageSourceAttribution>(attribution);
Assert.Equal(AgentRequestMessageSourceType.ChatHistory, typedAttribution.SourceType);
Assert.Equal(typeof(TestChatHistoryProvider).FullName, typedAttribution.SourceId);
}
[Fact]
public async Task InvokingAsync_WithCustomSourceId_StampsMessagesWithCustomSourceIdAsync()
{
// Arrange
const string CustomSourceId = "CustomHistorySource";
var provider = new TestChatHistoryProviderWithCustomSource(CustomSourceId);
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]);
// Act
IEnumerable<ChatMessage> messages = await provider.InvokingAsync(context);
// Assert
ChatMessage message = messages.Single();
Assert.NotNull(message.AdditionalProperties);
Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution));
var typedAttribution = Assert.IsType<AgentRequestMessageSourceAttribution>(attribution);
Assert.Equal(AgentRequestMessageSourceType.ChatHistory, typedAttribution.SourceType);
Assert.Equal(CustomSourceId, typedAttribution.SourceId);
}
[Fact]
public async Task InvokingAsync_DoesNotReStampAlreadyStampedMessagesAsync()
{
// Arrange
var provider = new TestChatHistoryProviderWithPreStampedMessages();
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]);
// Act
IEnumerable<ChatMessage> messages = await provider.InvokingAsync(context);
// Assert
ChatMessage message = messages.Single();
Assert.NotNull(message.AdditionalProperties);
Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution));
var typedAttribution = Assert.IsType<AgentRequestMessageSourceAttribution>(attribution);
Assert.Equal(AgentRequestMessageSourceType.ChatHistory, typedAttribution.SourceType);
Assert.Equal(typeof(TestChatHistoryProviderWithPreStampedMessages).FullName, typedAttribution.SourceId);
}
[Fact]
public async Task InvokingAsync_StampsMultipleMessagesAsync()
{
// Arrange
var provider = new TestChatHistoryProviderWithMultipleMessages();
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]);
// Act
IEnumerable<ChatMessage> messages = await provider.InvokingAsync(context);
// Assert
List<ChatMessage> messageList = messages.ToList();
Assert.Equal(3, messageList.Count);
foreach (ChatMessage message in messageList)
{
Assert.NotNull(message.AdditionalProperties);
Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution));
var typedAttribution = Assert.IsType<AgentRequestMessageSourceAttribution>(attribution);
Assert.Equal(AgentRequestMessageSourceType.ChatHistory, typedAttribution.SourceType);
Assert.Equal(typeof(TestChatHistoryProviderWithMultipleMessages).FullName, typedAttribution.SourceId);
}
}
#endregion
#region GetService Method Tests
[Fact]
@@ -371,49 +285,7 @@ public class ChatHistoryProviderTests
private sealed class TestChatHistoryProvider : ChatHistoryProvider
{
protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new([new ChatMessage(ChatRole.User, "Test Message")]);
protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
=> default;
}
private sealed class TestChatHistoryProviderWithCustomSource : ChatHistoryProvider
{
public TestChatHistoryProviderWithCustomSource(string sourceId) : base(sourceId)
{
}
protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new([new ChatMessage(ChatRole.User, "Test Message")]);
protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
=> default;
}
private sealed class TestChatHistoryProviderWithPreStampedMessages : ChatHistoryProvider
{
protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
var message = new ChatMessage(ChatRole.User, "Pre-stamped Message");
message.AdditionalProperties = new AdditionalPropertiesDictionary
{
[AgentRequestMessageSourceAttribution.AdditionalPropertiesKey] = new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)
};
return new([message]);
}
protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
=> default;
}
private sealed class TestChatHistoryProviderWithMultipleMessages : ChatHistoryProvider
{
protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new([
new ChatMessage(ChatRole.User, "Message 1"),
new ChatMessage(ChatRole.Assistant, "Message 2"),
new ChatMessage(ChatRole.User, "Message 3")
]);
=> new(new ChatMessage[] { new(ChatRole.User, "Test Message") }.Concat(context.RequestMessages));
protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
=> default;
@@ -350,7 +350,7 @@ public sealed class ChatMessageExtensionsTests
ChatMessage message = new(ChatRole.User, "Hello");
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.External, "TestSourceId");
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.External, "TestSourceId");
// Assert
Assert.NotSame(message, result);
@@ -368,7 +368,7 @@ public sealed class ChatMessageExtensionsTests
};
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.AIContextProvider, "ProviderSourceId");
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "ProviderSourceId");
// Assert
Assert.NotSame(message, result);
@@ -389,7 +389,7 @@ public sealed class ChatMessageExtensionsTests
};
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.ChatHistory, "HistoryId");
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "HistoryId");
// Assert
Assert.Same(message, result);
@@ -408,7 +408,7 @@ public sealed class ChatMessageExtensionsTests
};
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.AIContextProvider, "SourceId");
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "SourceId");
// Assert
Assert.NotSame(message, result);
@@ -429,7 +429,7 @@ public sealed class ChatMessageExtensionsTests
};
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.External, "NewId");
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.External, "NewId");
// Assert
Assert.NotSame(message, result);
@@ -444,7 +444,7 @@ public sealed class ChatMessageExtensionsTests
ChatMessage message = new(ChatRole.User, "Hello");
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.ChatHistory);
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory);
// Assert
Assert.NotSame(message, result);
@@ -465,7 +465,7 @@ public sealed class ChatMessageExtensionsTests
};
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.External);
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.External);
// Assert
Assert.Same(message, result);
@@ -478,7 +478,7 @@ public sealed class ChatMessageExtensionsTests
ChatMessage message = new(ChatRole.User, "Hello");
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.AIContextProvider, "ProviderId");
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "ProviderId");
// Assert
Assert.Null(message.AdditionalProperties);
@@ -499,7 +499,7 @@ public sealed class ChatMessageExtensionsTests
};
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.External, "SourceId");
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.External, "SourceId");
// Assert
Assert.NotSame(message, result);
@@ -514,7 +514,7 @@ public sealed class ChatMessageExtensionsTests
ChatMessage message = new(ChatRole.Assistant, "Test content");
// Act
ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.ChatHistory, "HistoryId");
ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "HistoryId");
// Assert
Assert.Equal(ChatRole.Assistant, result.Role);
@@ -71,7 +71,7 @@ public class InMemoryChatHistoryProviderTests
var requestMessages = new List<ChatMessage>
{
new(ChatRole.User, "Hello"),
new(ChatRole.System, "additional context") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "TestSource") } } },
new(ChatRole.System, "additional context") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "TestSource") } } },
};
var responseMessages = new List<ChatMessage>
{
@@ -120,6 +120,11 @@ public class InMemoryChatHistoryProviderTests
var session = CreateMockSession();
// Arrange
var requestMessages = new List<ChatMessage>
{
new(ChatRole.User, "Hello"),
};
var provider = new InMemoryChatHistoryProvider();
provider.SetMessages(session,
[
@@ -127,13 +132,18 @@ public class InMemoryChatHistoryProviderTests
new ChatMessage(ChatRole.Assistant, "Test2")
]);
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []);
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, requestMessages);
var result = (await provider.InvokingAsync(context, CancellationToken.None)).ToList();
// Assert
Assert.Equal(2, result.Count);
Assert.Equal(3, result.Count);
Assert.Contains(result, m => m.Text == "Test1");
Assert.Contains(result, m => m.Text == "Test2");
Assert.Contains(result, m => m.Text == "Hello");
Assert.Equal(AgentRequestMessageSourceType.ChatHistory, result[0].GetAgentRequestMessageSourceType());
Assert.Equal(AgentRequestMessageSourceType.ChatHistory, result[1].GetAgentRequestMessageSourceType());
Assert.Equal(AgentRequestMessageSourceType.External, result[2].GetAgentRequestMessageSourceType());
}
[Fact]
@@ -3135,7 +3135,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
{
protected override ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
return new ValueTask<AIContext>(new AIContext());
return new ValueTask<AIContext>(context.AIContext);
}
}
@@ -3146,7 +3146,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
{
protected override ValueTask<IEnumerable<ChatMessage>> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
{
return new ValueTask<IEnumerable<ChatMessage>>(Array.Empty<ChatMessage>());
return new ValueTask<IEnumerable<ChatMessage>>(context.RequestMessages);
}
protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
@@ -387,6 +387,8 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
Assert.Single(messageList2);
Assert.Equal("Message for conversation 1", messageList1[0].Text);
Assert.Equal("Message for conversation 2", messageList2[0].Text);
Assert.Equal(AgentRequestMessageSourceType.ChatHistory, messageList1[0].GetAgentRequestMessageSourceType());
Assert.Equal(AgentRequestMessageSourceType.ChatHistory, messageList2[0].GetAgentRequestMessageSourceType());
}
#endregion
@@ -1,6 +1,7 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Threading;
@@ -52,14 +53,14 @@ public sealed class Mem0ProviderTests : IDisposable
var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope));
await sut.ClearStoredMemoriesAsync(mockSession);
var ctxBefore = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, [question]));
var ctxBefore = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List<ChatMessage> { question } }));
Assert.DoesNotContain("Caoimhe", ctxBefore.Messages?[0].Text ?? string.Empty);
// Act
await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, [input]));
var ctxAfterAdding = await GetContextWithRetryAsync(sut, mockSession, question);
await sut.ClearStoredMemoriesAsync(mockSession);
var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, [question]));
var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List<ChatMessage> { question } }));
// Assert
Assert.Contains("Caoimhe", ctxAfterAdding.Messages?[0].Text ?? string.Empty);
@@ -77,14 +78,14 @@ public sealed class Mem0ProviderTests : IDisposable
var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope));
await sut.ClearStoredMemoriesAsync(mockSession);
var ctxBefore = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, [question]));
var ctxBefore = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List<ChatMessage> { question } }));
Assert.DoesNotContain("Caoimhe", ctxBefore.Messages?[0].Text ?? string.Empty);
// Act
await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, [assistantIntro]));
var ctxAfterAdding = await GetContextWithRetryAsync(sut, mockSession, question);
await sut.ClearStoredMemoriesAsync(mockSession);
var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, [question]));
var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List<ChatMessage> { question } }));
// Assert
Assert.Contains("Caoimhe", ctxAfterAdding.Messages?[0].Text ?? string.Empty);
@@ -107,8 +108,8 @@ public sealed class Mem0ProviderTests : IDisposable
await sut1.ClearStoredMemoriesAsync(mockSession1);
await sut2.ClearStoredMemoriesAsync(mockSession2);
var ctxBefore1 = await sut1.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession1, [question]));
var ctxBefore2 = await sut2.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession2, [question]));
var ctxBefore1 = await sut1.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession1, new AIContext { Messages = new List<ChatMessage> { question } }));
var ctxBefore2 = await sut2.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession2, new AIContext { Messages = new List<ChatMessage> { question } }));
Assert.DoesNotContain("Caoimhe", ctxBefore1.Messages?[0].Text ?? string.Empty);
Assert.DoesNotContain("Caoimhe", ctxBefore2.Messages?[0].Text ?? string.Empty);
@@ -131,7 +132,7 @@ public sealed class Mem0ProviderTests : IDisposable
AIContext? ctx = null;
for (int i = 0; i < attempts; i++)
{
ctx = await provider.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, session, [question]), CancellationToken.None);
ctx = await provider.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, session, new AIContext { Messages = new List<ChatMessage> { question } }), CancellationToken.None);
var text = ctx.Messages?[0].Text;
if (!string.IsNullOrEmpty(text) && text.IndexOf("Caoimhe", StringComparison.OrdinalIgnoreCase) >= 0)
{
@@ -103,7 +103,7 @@ public sealed class Mem0ProviderTests : IDisposable
};
var mockSession = new TestAgentSession();
var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), options: new() { EnableSensitiveTelemetryData = true }, loggerFactory: this._loggerFactoryMock.Object);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, [new ChatMessage(ChatRole.User, "What is my name?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "What is my name?") } });
// Act
var aiContext = await sut.InvokingAsync(invokingContext);
@@ -118,9 +118,12 @@ public sealed class Mem0ProviderTests : IDisposable
Assert.Equal("What is my name?", doc.RootElement.GetProperty("query").GetString());
Assert.NotNull(aiContext.Messages);
var contextMessage = Assert.Single(aiContext.Messages);
Assert.Equal(2, aiContext.Messages.Count);
Assert.Equal(AgentRequestMessageSourceType.External, aiContext.Messages[0].GetAgentRequestMessageSourceType());
var contextMessage = aiContext.Messages[1];
Assert.Equal(ChatRole.User, contextMessage.Role);
Assert.Contains("Name is Caoimhe", contextMessage.Text);
Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, contextMessage.GetAgentRequestMessageSourceType());
this._loggerMock.Verify(
l => l.Log(
@@ -169,7 +172,7 @@ public sealed class Mem0ProviderTests : IDisposable
var mockSession = new TestAgentSession();
var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), options: options, loggerFactory: this._loggerFactoryMock.Object);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, [new ChatMessage(ChatRole.User, "Who am I?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Who am I?") } });
// Act
await sut.InvokingAsync(invokingContext, CancellationToken.None);
@@ -371,13 +374,14 @@ public sealed class Mem0ProviderTests : IDisposable
var storageScope = new Mem0ProviderScope { ApplicationId = "app" };
var mockSession = new TestAgentSession();
var provider = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), loggerFactory: this._loggerFactoryMock.Object);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, [new ChatMessage(ChatRole.User, "Q?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Q?") } });
// Act
var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
// Assert
Assert.Null(aiContext.Messages);
Assert.NotNull(aiContext.Messages);
Assert.Single(aiContext.Messages);
Assert.Null(aiContext.Tools);
this._loggerMock.Verify(
l => l.Log(
@@ -403,7 +407,7 @@ public sealed class Mem0ProviderTests : IDisposable
initializerCallCount++;
return new Mem0Provider.State(storageScope);
});
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, [new ChatMessage(ChatRole.User, "Q?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Q?") } });
// Act
await sut.InvokingAsync(invokingContext, CancellationToken.None);
@@ -422,7 +426,7 @@ public sealed class Mem0ProviderTests : IDisposable
var mockSession = new TestAgentSession();
const string CustomKey = "MyCustomKey";
var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), options: new() { StateKey = CustomKey });
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, [new ChatMessage(ChatRole.User, "Q?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Q?") } });
// Act
await sut.InvokingAsync(invokingContext, CancellationToken.None);
@@ -345,12 +345,13 @@ public partial class ChatClientAgentTests
mockProvider
.Protected()
.Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(new AIContext
{
Messages = aiContextProviderMessages,
Instructions = "context provider instructions",
Tools = [AIFunctionFactory.Create(() => { }, "context provider function")]
});
.Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) =>
new ValueTask<AIContext>(new AIContext
{
Messages = (ctx.AIContext.Messages ?? []).Concat(aiContextProviderMessages).ToList(),
Instructions = ctx.AIContext.Instructions + "\ncontext provider instructions",
Tools = (ctx.AIContext.Tools ?? []).Concat(new[] { AIFunctionFactory.Create(() => { }, "context provider function") }).ToList()
}));
mockProvider
.Protected()
.Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -415,10 +416,11 @@ public partial class ChatClientAgentTests
mockProvider
.Protected()
.Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(new AIContext
{
Messages = aiContextProviderMessages,
});
.Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) =>
new ValueTask<AIContext>(new AIContext
{
Messages = (ctx.AIContext.Messages ?? []).Concat(aiContextProviderMessages).ToList(),
}));
mockProvider
.Protected()
.Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -472,7 +474,13 @@ public partial class ChatClientAgentTests
mockProvider
.Protected()
.Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(new AIContext());
.Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) =>
new ValueTask<AIContext>(new AIContext
{
Instructions = ctx.AIContext.Instructions,
Messages = ctx.AIContext.Messages?.ToList(),
Tools = ctx.AIContext.Tools
}));
ChatClientAgent agent = new(mockService.Object, options: new() { AIContextProvider = mockProvider.Object, ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
@@ -1423,12 +1431,13 @@ public partial class ChatClientAgentTests
mockProvider
.Protected()
.Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(new AIContext
{
Messages = aiContextProviderMessages,
Instructions = "context provider instructions",
Tools = [AIFunctionFactory.Create(() => { }, "context provider function")]
});
.Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) =>
new ValueTask<AIContext>(new AIContext
{
Messages = (ctx.AIContext.Messages ?? []).Concat(aiContextProviderMessages).ToList(),
Instructions = ctx.AIContext.Instructions + "\ncontext provider instructions",
Tools = (ctx.AIContext.Tools ?? []).Concat(new[] { AIFunctionFactory.Create(() => { }, "context provider function") }).ToList()
}));
mockProvider
.Protected()
.Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -1501,10 +1510,11 @@ public partial class ChatClientAgentTests
mockProvider
.Protected()
.Setup<ValueTask<AIContext>>("InvokingCoreAsync", ItExpr.IsAny<AIContextProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(new AIContext
{
Messages = aiContextProviderMessages,
});
.Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) =>
new ValueTask<AIContext>(new AIContext
{
Messages = (ctx.AIContext.Messages ?? []).Concat(aiContextProviderMessages).ToList(),
}));
mockProvider
.Protected()
.Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<AIContextProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -189,7 +189,8 @@ public class ChatClientAgent_ChatHistoryManagementTests
mockChatHistoryProvider
.Protected()
.Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync([new ChatMessage(ChatRole.User, "Existing Chat History")]);
.Returns((ChatHistoryProvider.InvokingContext ctx, CancellationToken _) =>
new ValueTask<IEnumerable<ChatMessage>>(new List<ChatMessage> { new(ChatRole.User, "Existing Chat History") }.Concat(ctx.RequestMessages).ToList()));
mockChatHistoryProvider
.Protected()
.Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -221,7 +222,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
mockChatHistoryProvider
.Protected()
.Verify<ValueTask>("InvokedCoreAsync", Times.Once(),
ItExpr.Is<ChatHistoryProvider.InvokedContext>(x => x.RequestMessages.Count() == 1 && x.ResponseMessages!.Count() == 1),
ItExpr.Is<ChatHistoryProvider.InvokedContext>(x => x.RequestMessages.Count() == 2 && x.ResponseMessages!.Count() == 1),
ItExpr.IsAny<CancellationToken>());
}
@@ -240,6 +241,11 @@ public class ChatClientAgent_ChatHistoryManagementTests
It.IsAny<CancellationToken>())).Throws(new InvalidOperationException("Test Error"));
Mock<ChatHistoryProvider> mockChatHistoryProvider = new();
mockChatHistoryProvider
.Protected()
.Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.Returns((ChatHistoryProvider.InvokingContext ctx, CancellationToken _) =>
new ValueTask<IEnumerable<ChatMessage>>(ctx.RequestMessages.ToList()));
ChatClientAgent agent = new(mockService.Object, options: new()
{
@@ -309,7 +315,8 @@ public class ChatClientAgent_ChatHistoryManagementTests
mockOverrideChatHistoryProvider
.Protected()
.Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync([new ChatMessage(ChatRole.User, "Existing Chat History")]);
.Returns((ChatHistoryProvider.InvokingContext ctx, CancellationToken _) =>
new ValueTask<IEnumerable<ChatMessage>>(new List<ChatMessage> { new(ChatRole.User, "Existing Chat History") }.Concat(ctx.RequestMessages).ToList()));
mockOverrideChatHistoryProvider
.Protected()
.Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
@@ -355,7 +362,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
mockOverrideChatHistoryProvider
.Protected()
.Verify<ValueTask>("InvokedCoreAsync", Times.Once(),
ItExpr.Is<ChatHistoryProvider.InvokedContext>(x => x.RequestMessages.Count() == 1 && x.ResponseMessages!.Count() == 1),
ItExpr.Is<ChatHistoryProvider.InvokedContext>(x => x.RequestMessages.Count() == 2 && x.ResponseMessages!.Count() == 1),
ItExpr.IsAny<CancellationToken>());
mockAgentOptionsChatHistoryProvider
@@ -90,10 +90,14 @@ public sealed class TextSearchProviderTests
var invokingContext = new AIContextProvider.InvokingContext(
s_mockAgent,
new TestAgentSession(),
[
new ChatMessage(ChatRole.User, "Sample user question?"),
new ChatMessage(ChatRole.User, "Additional part")
]);
new AIContext
{
Messages = new List<ChatMessage>
{
new(ChatRole.User, "Sample user question?"),
new(ChatRole.User, "Additional part")
}
});
// Act
var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
@@ -102,9 +106,14 @@ public sealed class TextSearchProviderTests
Assert.Equal("Sample user question?\nAdditional part", capturedInput);
Assert.Null(aiContext.Instructions); // TextSearchProvider uses a user message for context injection.
Assert.NotNull(aiContext.Messages);
Assert.Single(aiContext.Messages!);
var message = aiContext.Messages!.Single();
Assert.Equal(3, aiContext.Messages!.Count); // 2 input messages + 1 search result message
Assert.Equal("Sample user question?", aiContext.Messages![0].Text);
Assert.Equal("Additional part", aiContext.Messages![1].Text);
Assert.Equal(AgentRequestMessageSourceType.External, aiContext.Messages![0].GetAgentRequestMessageSourceType());
Assert.Equal(AgentRequestMessageSourceType.External, aiContext.Messages![1].GetAgentRequestMessageSourceType());
var message = aiContext.Messages!.Last();
Assert.Equal(ChatRole.User, message.Role);
Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, message.GetAgentRequestMessageSourceType());
string text = message.Text!;
if (overrideContextPrompt is null)
@@ -165,13 +174,15 @@ public sealed class TextSearchProviderTests
FunctionToolDescription = overrideDescription
};
var provider = new TextSearchProvider(this.NoResultSearchAsync, options);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), [new ChatMessage(ChatRole.User, "Q?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Q?") } });
// Act
var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
// Assert
Assert.Null(aiContext.Messages); // No automatic injection.
Assert.NotNull(aiContext.Messages); // Input messages are preserved.
Assert.Single(aiContext.Messages!);
Assert.Equal("Q?", aiContext.Messages![0].Text);
Assert.NotNull(aiContext.Tools);
Assert.Single(aiContext.Tools);
var tool = aiContext.Tools.Single();
@@ -184,13 +195,15 @@ public sealed class TextSearchProviderTests
{
// Arrange
var provider = new TextSearchProvider(this.FailingSearchAsync, loggerFactory: this._loggerFactoryMock.Object);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), [new ChatMessage(ChatRole.User, "Q?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Q?") } });
// Act
var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
// Assert
Assert.Null(aiContext.Messages);
Assert.NotNull(aiContext.Messages); // Input messages are preserved on error.
Assert.Single(aiContext.Messages!);
Assert.Equal("Q?", aiContext.Messages![0].Text);
Assert.Null(aiContext.Tools);
this._loggerMock.Verify(
l => l.Log(
@@ -277,15 +290,16 @@ public sealed class TextSearchProviderTests
ContextFormatter = r => $"Custom formatted context with {r.Count} results."
};
var provider = new TextSearchProvider(SearchDelegateAsync, options);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), [new ChatMessage(ChatRole.User, "Q?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Q?") } });
// Act
var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
// Assert
Assert.NotNull(aiContext.Messages);
Assert.Single(aiContext.Messages!);
Assert.Equal("Custom formatted context with 2 results.", aiContext.Messages![0].Text);
Assert.Equal(2, aiContext.Messages!.Count); // 1 input message + 1 formatted result message
Assert.Equal("Q?", aiContext.Messages![0].Text);
Assert.Equal("Custom formatted context with 2 results.", aiContext.Messages![1].Text);
}
[Fact]
@@ -311,15 +325,16 @@ public sealed class TextSearchProviderTests
ContextFormatter = r => string.Join(",", r.Select(x => ((RawPayload)x.RawRepresentation!).Id))
};
var provider = new TextSearchProvider(SearchDelegateAsync, options);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), [new ChatMessage(ChatRole.User, "Q?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Q?") } });
// Act
var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
// Assert
Assert.NotNull(aiContext.Messages);
Assert.Single(aiContext.Messages!);
Assert.Equal("R1,R2", aiContext.Messages![0].Text);
Assert.Equal(2, aiContext.Messages!.Count); // 1 input message + 1 formatted result message
Assert.Equal("Q?", aiContext.Messages![0].Text);
Assert.Equal("R1,R2", aiContext.Messages![1].Text);
}
[Fact]
@@ -328,13 +343,15 @@ public sealed class TextSearchProviderTests
// Arrange
var options = new TextSearchProviderOptions { SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke };
var provider = new TextSearchProvider(this.NoResultSearchAsync, options);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), [new ChatMessage(ChatRole.User, "Q?")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Q?") } });
// Act
var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
// Assert
Assert.Null(aiContext.Messages);
Assert.NotNull(aiContext.Messages); // Input messages are preserved when no results found.
Assert.Single(aiContext.Messages!);
Assert.Equal("Q?", aiContext.Messages![0].Text);
Assert.Null(aiContext.Instructions);
Assert.Null(aiContext.Tools);
}
@@ -373,9 +390,7 @@ public sealed class TextSearchProviderTests
var invokingContext = new AIContextProvider.InvokingContext(
s_mockAgent,
session,
[
new ChatMessage(ChatRole.User, "E")
]);
new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "E") } });
// Act
await provider.InvokingAsync(invokingContext, CancellationToken.None);
@@ -416,9 +431,7 @@ public sealed class TextSearchProviderTests
var invokingContext = new AIContextProvider.InvokingContext(
s_mockAgent,
session,
[
new ChatMessage(ChatRole.User, "E")
]);
new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "E") } });
// Act
await provider.InvokingAsync(invokingContext, CancellationToken.None);
@@ -465,7 +478,7 @@ public sealed class TextSearchProviderTests
new ChatMessage(ChatRole.User, "E"),
]));
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "F")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, session, new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "F") } });
// Act
await provider.InvokingAsync(invokingContext, CancellationToken.None);
@@ -506,9 +519,7 @@ public sealed class TextSearchProviderTests
var invokingContext = new AIContextProvider.InvokingContext(
s_mockAgent,
session,
[
new ChatMessage(ChatRole.User, "Question?") // Current request message always appended.
]);
new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "Question?") } }); // Current request message always appended.
// Act
await provider.InvokingAsync(invokingContext, CancellationToken.None);
@@ -589,8 +600,7 @@ public sealed class TextSearchProviderTests
SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke,
RecentMessageMemoryLimit = 4
});
var emptyMessages = Array.Empty<ChatMessage>();
await newProvider.InvokingAsync(new(s_mockAgent, restoredSession, emptyMessages), CancellationToken.None); // Trigger search to read memory.
await newProvider.InvokingAsync(new(s_mockAgent, restoredSession, new AIContext()), CancellationToken.None); // Trigger search to read memory.
// Assert
Assert.NotNull(capturedInput);
@@ -616,8 +626,7 @@ public sealed class TextSearchProviderTests
SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke,
RecentMessageMemoryLimit = 3
});
var emptyMessages = Array.Empty<ChatMessage>();
await provider.InvokingAsync(new(s_mockAgent, session, emptyMessages), CancellationToken.None);
await provider.InvokingAsync(new(s_mockAgent, session, new AIContext()), CancellationToken.None);
// Assert
Assert.NotNull(capturedInput);
@@ -381,10 +381,10 @@ public class ChatHistoryMemoryProviderTests
options: providerOptions);
var requestMsg = new ChatMessage(ChatRole.User, "requesting relevant history");
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), [requestMsg]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List<ChatMessage> { requestMsg } });
// Act
await provider.InvokingAsync(invokingContext, CancellationToken.None);
var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None);
// Assert
this._vectorStoreCollectionMock.Verify(
@@ -394,6 +394,11 @@ public class ChatHistoryMemoryProviderTests
It.IsAny<VectorSearchOptions<Dictionary<string, object?>>>(),
It.IsAny<CancellationToken>()),
Times.Once);
Assert.NotNull(aiContext.Messages);
Assert.Equal(2, aiContext.Messages.Count);
Assert.Equal(AgentRequestMessageSourceType.External, aiContext.Messages[0].GetAgentRequestMessageSourceType());
Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, aiContext.Messages[1].GetAgentRequestMessageSourceType());
}
[Fact]
@@ -437,7 +442,7 @@ public class ChatHistoryMemoryProviderTests
options: providerOptions);
var requestMsg = new ChatMessage(ChatRole.User, "requesting relevant history");
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), [requestMsg]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List<ChatMessage> { requestMsg } });
// Act
await provider.InvokingAsync(invokingContext, CancellationToken.None);
@@ -500,7 +505,7 @@ public class ChatHistoryMemoryProviderTests
options: options,
loggerFactory: this._loggerFactoryMock.Object);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), [new ChatMessage(ChatRole.User, "requesting relevant history")]);
var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List<ChatMessage> { new(ChatRole.User, "requesting relevant history") } });
// Act
await provider.InvokingAsync(invokingContext, CancellationToken.None);