diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/Program.cs index f5f567dbd2..09cd540378 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/Program.cs @@ -105,12 +105,27 @@ static string GetDateTime() var contextProviderAgent = originalAgent .AsBuilder() - .Use([new DateTimeContextProvider()]) + .UseAIContextProviders(new DateTimeContextProvider()) .Build(); var contextResponse = await contextProviderAgent.RunAsync("Is it almost time for lunch?"); Console.WriteLine($"Context-enriched response: {contextResponse}"); +// AIContextProvider at the chat client level. Unlike the agent-level MessageAIContextProvider, +// this operates within the IChatClient pipeline and can also enrich tools and instructions. +// It must be used within the context of a running AIAgent (uses AIAgent.CurrentRunContext). +// In this case we are attaching an AIContextProvider that only adds messages. +Console.WriteLine("\n\n=== Example 6: AIContextProvider on chat client pipeline ==="); + +var chatClientProviderAgent = azureOpenAIClient.AsIChatClient() + .AsBuilder() + .UseAIContextProviders(new DateTimeContextProvider()) + .BuildAIAgent( + instructions: "You are an AI assistant that helps people find information."); + +var chatClientContextResponse = await chatClientProviderAgent.RunAsync("Is it almost time for lunch?"); +Console.WriteLine($"Chat client context-enriched response: {chatClientContextResponse}"); + // Function invocation middleware that logs before and after function calls. async ValueTask FunctionCallMiddleware(AIAgent agent, FunctionInvocationContext context, Func> next, CancellationToken cancellationToken) { @@ -278,7 +293,7 @@ async Task PerRequestChatClientMiddleware(IEnumerable /// /// A that injects the current date and time into the agent's context. /// This is a simple example of how to use a MessageAIContextProvider to enrich agent messages -/// via the extension method. +/// via the extension method. /// internal sealed class DateTimeContextProvider : MessageAIContextProvider { diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/README.md b/dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/README.md index d7193b6982..53a59c20a7 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/README.md +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/README.md @@ -15,6 +15,7 @@ This sample demonstrates how to add middleware to intercept: 6. Per‑request function pipeline with approval 7. Combining agent‑level and per‑request middleware 8. MessageAIContextProvider middleware via `AIAgentBuilder.Use(...)` for injecting additional context messages +9. AIContextProvider middleware via `ChatClientBuilder.Use(...)` for enriching messages, tools, and instructions at the chat client level ## Function Invocation Middleware diff --git a/dotnet/src/Microsoft.Agents.AI/AIAgentBuilder.cs b/dotnet/src/Microsoft.Agents.AI/AIAgentBuilder.cs index 52293db621..a2b1cd23a7 100644 --- a/dotnet/src/Microsoft.Agents.AI/AIAgentBuilder.cs +++ b/dotnet/src/Microsoft.Agents.AI/AIAgentBuilder.cs @@ -172,7 +172,7 @@ public AIAgentBuilder Use( /// context enrichment, not just agents that natively support instances. /// /// - public AIAgentBuilder Use(MessageAIContextProvider[] providers) + public AIAgentBuilder UseAIContextProviders(params MessageAIContextProvider[] providers) { return this.Use((innerAgent, _) => new MessageAIContextProviderAgent(innerAgent, providers)); } diff --git a/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClient.cs b/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClient.cs new file mode 100644 index 0000000000..305abe0465 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClient.cs @@ -0,0 +1,215 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI; + +/// +/// A delegating chat client that enriches input messages, tools, and instructions by invoking a pipeline of +/// instances before delegating to the inner chat client, and notifies those +/// providers after the inner client completes. +/// +/// +/// +/// This chat client must be used within the context of a running . It retrieves the current +/// agent and session from , which is set automatically when an agent's +/// or +/// method is called. +/// An is thrown if no run context is available. +/// +/// +internal sealed class AIContextProviderChatClient : DelegatingChatClient +{ + private readonly IReadOnlyList _providers; + + /// + /// Initializes a new instance of the class. + /// + /// The underlying chat client that will handle the core operations. + /// The AI context providers to invoke before and after the inner chat client. + public AIContextProviderChatClient(IChatClient innerClient, IReadOnlyList providers) + : base(innerClient) + { + Throw.IfNull(providers); + + if (providers.Count == 0) + { + Throw.ArgumentException(nameof(providers), "At least one AIContextProvider must be provided."); + } + + this._providers = providers; + } + + /// + public override async Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + var runContext = GetRequiredRunContext(); + var (enrichedMessages, enrichedOptions) = await this.InvokeProvidersAsync(runContext, messages, options, cancellationToken).ConfigureAwait(false); + + ChatResponse response; + try + { + response = await base.GetResponseAsync(enrichedMessages, enrichedOptions, cancellationToken).ConfigureAwait(false); + } + catch (Exception ex) + { + await this.NotifyProvidersOfFailureAsync(runContext, enrichedMessages, ex, cancellationToken).ConfigureAwait(false); + throw; + } + + await this.NotifyProvidersOfSuccessAsync(runContext, enrichedMessages, response.Messages, cancellationToken).ConfigureAwait(false); + + return response; + } + + /// + public override async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + var runContext = GetRequiredRunContext(); + var (enrichedMessages, enrichedOptions) = await this.InvokeProvidersAsync(runContext, messages, options, cancellationToken).ConfigureAwait(false); + + List responseUpdates = []; + + IAsyncEnumerator enumerator; + try + { + enumerator = base.GetStreamingResponseAsync(enrichedMessages, enrichedOptions, cancellationToken).GetAsyncEnumerator(cancellationToken); + } + catch (Exception ex) + { + await this.NotifyProvidersOfFailureAsync(runContext, enrichedMessages, ex, cancellationToken).ConfigureAwait(false); + throw; + } + + bool hasUpdates; + try + { + hasUpdates = await enumerator.MoveNextAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + await this.NotifyProvidersOfFailureAsync(runContext, enrichedMessages, ex, cancellationToken).ConfigureAwait(false); + throw; + } + + while (hasUpdates) + { + var update = enumerator.Current; + responseUpdates.Add(update); + yield return update; + + try + { + hasUpdates = await enumerator.MoveNextAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + await this.NotifyProvidersOfFailureAsync(runContext, enrichedMessages, ex, cancellationToken).ConfigureAwait(false); + throw; + } + } + + var chatResponse = responseUpdates.ToChatResponse(); + await this.NotifyProvidersOfSuccessAsync(runContext, enrichedMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false); + } + + /// + /// Gets the current , throwing if not available. + /// + private static AgentRunContext GetRequiredRunContext() + { + return AIAgent.CurrentRunContext + ?? throw new InvalidOperationException( + $"{nameof(AIContextProviderChatClient)} can only be used within the context of a running AIAgent. " + + "Ensure that the chat client is being invoked as part of an AIAgent.RunAsync or AIAgent.RunStreamingAsync call."); + } + + /// + /// Invokes each provider's in sequence, + /// accumulating context (messages, tools, instructions) from each. + /// + private async Task<(IEnumerable Messages, ChatOptions? Options)> InvokeProvidersAsync( + AgentRunContext runContext, + IEnumerable messages, + ChatOptions? options, + CancellationToken cancellationToken) + { + var aiContext = new AIContext + { + Instructions = options?.Instructions, + Messages = messages, + Tools = options?.Tools + }; + + foreach (var provider in this._providers) + { + var invokingContext = new AIContextProvider.InvokingContext(runContext.Agent, runContext.Session, aiContext); + aiContext = await provider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false); + } + + // Materialize the accumulated context back into messages and options. + var enrichedMessages = aiContext.Messages ?? []; + + var tools = aiContext.Tools as IList ?? aiContext.Tools?.ToList(); + if (options?.Tools is { Count: > 0 } || tools is { Count: > 0 }) + { + options ??= new(); + options.Tools = tools; + } + + if (options?.Instructions is not null || aiContext.Instructions is not null) + { + options ??= new(); + options.Instructions = aiContext.Instructions; + } + + return (enrichedMessages, options); + } + + /// + /// Notifies each provider of a successful invocation. + /// + private async Task NotifyProvidersOfSuccessAsync( + AgentRunContext runContext, + IEnumerable requestMessages, + IEnumerable responseMessages, + CancellationToken cancellationToken) + { + var invokedContext = new AIContextProvider.InvokedContext(runContext.Agent, runContext.Session, requestMessages, responseMessages); + + foreach (var provider in this._providers) + { + await provider.InvokedAsync(invokedContext, cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Notifies each provider of a failed invocation. + /// + private async Task NotifyProvidersOfFailureAsync( + AgentRunContext runContext, + IEnumerable requestMessages, + Exception exception, + CancellationToken cancellationToken) + { + var invokedContext = new AIContextProvider.InvokedContext(runContext.Agent, runContext.Session, requestMessages, exception); + + foreach (var provider in this._providers) + { + await provider.InvokedAsync(invokedContext, cancellationToken).ConfigureAwait(false); + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClientBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClientBuilderExtensions.cs new file mode 100644 index 0000000000..1152c9843e --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClientBuilderExtensions.cs @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Agents.AI; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// +/// Provides extension methods for adding support to instances. +/// +public static class AIContextProviderChatClientBuilderExtensions +{ + /// + /// Adds one or more instances to the chat client pipeline, enabling context enrichment + /// (messages, tools, and instructions) for any . + /// + /// The to which the providers will be added. + /// + /// The instances to invoke before and after each chat client call. + /// Providers are called in sequence, with each receiving the accumulated context from the previous provider. + /// + /// The with the providers added, enabling method chaining. + /// or is . + /// is empty. + /// + /// + /// This method wraps the inner chat client with a decorator that calls each provider's + /// in sequence before the inner client is called, + /// and calls on each provider after the inner client completes. + /// + /// + /// The chat client must be used within the context of a running . The agent and session + /// are retrieved from . An + /// is thrown at invocation time if no run context is available. + /// + /// + public static ChatClientBuilder UseAIContextProviders(this ChatClientBuilder builder, params AIContextProvider[] providers) + { + _ = Throw.IfNull(builder); + + return builder.Use(innerClient => new AIContextProviderChatClient(innerClient, providers)); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI/MessageAIContextProviderAgent.cs b/dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/MessageAIContextProviderAgent.cs similarity index 100% rename from dotnet/src/Microsoft.Agents.AI/MessageAIContextProviderAgent.cs rename to dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/MessageAIContextProviderAgent.cs diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs new file mode 100644 index 0000000000..51eb4be3ab --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs @@ -0,0 +1,430 @@ +// 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.UnitTests; + +/// +/// Unit tests for the class and +/// the builder extension. +/// +public class AIContextProviderChatClientTests +{ + private static readonly AgentSession s_mockSession = new Mock().Object; + + #region Constructor Tests + + [Fact] + public void Constructor_NullInnerClient_ThrowsArgumentNullException() + { + // Arrange + var provider = new TestAIContextProvider("key1"); + + // Act & Assert + Assert.Throws(() => new AIContextProviderChatClient(null!, [provider])); + } + + [Fact] + public void Constructor_NullProviders_ThrowsArgumentNullException() + { + // Arrange + var innerClient = new Mock().Object; + + // Act & Assert + Assert.Throws(() => new AIContextProviderChatClient(innerClient, null!)); + } + + [Fact] + public void Constructor_EmptyProviders_ThrowsArgumentException() + { + // Arrange + var innerClient = new Mock().Object; + + // Act & Assert + Assert.Throws(() => new AIContextProviderChatClient(innerClient, [])); + } + + #endregion + + #region GetResponseAsync Tests + + [Fact] + public async Task GetResponseAsync_NoRunContext_ThrowsInvalidOperationExceptionAsync() + { + // Arrange + var innerClient = new Mock(); + var provider = new TestAIContextProvider("key1"); + var chatClient = new AIContextProviderChatClient(innerClient.Object, [provider]); + + // Act & Assert — no AIAgent.CurrentRunContext is set + await Assert.ThrowsAsync( + () => chatClient.GetResponseAsync([new ChatMessage(ChatRole.User, "Hello")])); + } + + [Fact] + public async Task GetResponseAsync_SingleProvider_EnrichesMessagesAsync() + { + // Arrange + IEnumerable? capturedMessages = null; + var innerClient = CreateMockChatClient( + onGetResponse: (messages, _, _) => + { + capturedMessages = messages; + return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")])); + }); + + var provider = new TestAIContextProvider("key1", provideMessages: [new ChatMessage(ChatRole.System, "Extra context")]); + var chatClient = new AIContextProviderChatClient(innerClient, [provider]); + + // Act — run through an agent so CurrentRunContext is set + await RunWithAgentContextAsync(chatClient); + + // Assert + Assert.NotNull(capturedMessages); + var messageList = capturedMessages!.ToList(); + Assert.Equal(2, messageList.Count); + Assert.Equal("Hello", messageList[0].Text); + Assert.Contains("Extra context", messageList[1].Text); + } + + [Fact] + public async Task GetResponseAsync_MultipleProviders_CalledInSequenceAsync() + { + // Arrange + IEnumerable? capturedMessages = null; + var innerClient = CreateMockChatClient( + onGetResponse: (messages, _, _) => + { + capturedMessages = messages; + return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")])); + }); + + var provider1 = new TestAIContextProvider("key1", provideMessages: [new ChatMessage(ChatRole.System, "From P1")]); + var provider2 = new TestAIContextProvider("key2", provideMessages: [new ChatMessage(ChatRole.System, "From P2")]); + var chatClient = new AIContextProviderChatClient(innerClient, [provider1, provider2]); + + // Act + await RunWithAgentContextAsync(chatClient); + + // Assert + Assert.NotNull(capturedMessages); + var messageList = capturedMessages!.ToList(); + Assert.Equal(3, messageList.Count); + } + + [Fact] + public async Task GetResponseAsync_Provider_EnrichesToolsAndInstructionsAsync() + { + // Arrange + ChatOptions? capturedOptions = null; + var innerClient = CreateMockChatClient( + onGetResponse: (_, options, _) => + { + capturedOptions = options; + return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")])); + }); + + var provider = new TestAIContextProvider("key1", provideInstructions: "Extra instructions", provideTools: [new TestAITool()]); + var chatClient = new AIContextProviderChatClient(innerClient, [provider]); + + // Act + await RunWithAgentContextAsync(chatClient); + + // Assert + Assert.NotNull(capturedOptions); + Assert.Equal("Extra instructions", capturedOptions!.Instructions); + Assert.Single(capturedOptions.Tools!); + } + + [Fact] + public async Task GetResponseAsync_OnSuccess_InvokedAsyncCalledAsync() + { + // Arrange + var innerClient = CreateMockChatClient( + onGetResponse: (_, _, _) => Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")]))); + + var provider = new TestAIContextProvider("key1"); + var chatClient = new AIContextProviderChatClient(innerClient, [provider]); + + // Act + await RunWithAgentContextAsync(chatClient); + + // Assert + Assert.True(provider.InvokedAsyncCalled); + Assert.Null(provider.LastInvokedContext!.InvokeException); + Assert.NotNull(provider.LastInvokedContext.ResponseMessages); + } + + [Fact] + public async Task GetResponseAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync() + { + // Arrange + var expectedException = new InvalidOperationException("Chat failed"); + var innerClient = CreateMockChatClient( + onGetResponse: (_, _, _) => throw expectedException); + + var provider = new TestAIContextProvider("key1"); + var chatClient = new AIContextProviderChatClient(innerClient, [provider]); + + // Act & Assert + await Assert.ThrowsAsync(() => RunWithAgentContextAsync(chatClient)); + + Assert.True(provider.InvokedAsyncCalled); + Assert.Same(expectedException, provider.LastInvokedContext!.InvokeException); + } + + #endregion + + #region GetStreamingResponseAsync Tests + + [Fact] + public async Task GetStreamingResponseAsync_SingleProvider_EnrichesAndStreamsAsync() + { + // Arrange + IEnumerable? capturedMessages = null; + var innerClient = CreateMockStreamingChatClient( + onGetStreamingResponse: (messages, _, _) => + { + capturedMessages = messages; + return ToAsyncEnumerableAsync( + new ChatResponseUpdate(ChatRole.Assistant, "Part1"), + new ChatResponseUpdate(ChatRole.Assistant, "Part2")); + }); + + var provider = new TestAIContextProvider("key1", provideMessages: [new ChatMessage(ChatRole.System, "Extra context")]); + var chatClient = new AIContextProviderChatClient(innerClient, [provider]); + + // Act + var updates = new List(); + await RunStreamingWithAgentContextAsync(chatClient, updates); + + // Assert + Assert.Equal(2, updates.Count); + Assert.NotNull(capturedMessages); + Assert.Equal(2, capturedMessages!.ToList().Count); + } + + [Fact] + public async Task GetStreamingResponseAsync_OnSuccess_InvokedAsyncCalledAsync() + { + // Arrange + var innerClient = CreateMockStreamingChatClient( + onGetStreamingResponse: (_, _, _) => ToAsyncEnumerableAsync( + new ChatResponseUpdate(ChatRole.Assistant, "Response"))); + + var provider = new TestAIContextProvider("key1"); + var chatClient = new AIContextProviderChatClient(innerClient, [provider]); + + // Act + await RunStreamingWithAgentContextAsync(chatClient, []); + + // Assert + Assert.True(provider.InvokedAsyncCalled); + Assert.Null(provider.LastInvokedContext!.InvokeException); + } + + [Fact] + public async Task GetStreamingResponseAsync_OnFailure_InvokedAsyncCalledWithExceptionAsync() + { + // Arrange + var expectedException = new InvalidOperationException("Stream failed"); + var innerClient = CreateMockStreamingChatClient( + onGetStreamingResponse: (_, _, _) => throw expectedException); + + var provider = new TestAIContextProvider("key1"); + var chatClient = new AIContextProviderChatClient(innerClient, [provider]); + + // Act & Assert + await Assert.ThrowsAsync( + () => RunStreamingWithAgentContextAsync(chatClient, [])); + + Assert.True(provider.InvokedAsyncCalled); + Assert.Same(expectedException, provider.LastInvokedContext!.InvokeException); + } + + #endregion + + #region Builder Extension Tests + + [Fact] + public void UseExtension_NullBuilder_ThrowsArgumentNullException() + { + // Arrange + var provider = new TestAIContextProvider("key1"); + + // Act & Assert + Assert.Throws(() => + AIContextProviderChatClientBuilderExtensions.UseAIContextProviders(null!, provider)); + } + + [Fact] + public async Task UseExtension_CreatesWorkingPipelineAsync() + { + // Arrange + IEnumerable? capturedMessages = null; + var innerClient = CreateMockChatClient( + onGetResponse: (messages, _, _) => + { + capturedMessages = messages; + return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Response")])); + }); + + var provider = new TestAIContextProvider("key1", provideMessages: [new ChatMessage(ChatRole.System, "Pipeline context")]); + + var pipeline = new ChatClientBuilder(innerClient) + .UseAIContextProviders(provider) + .Build(); + + // Act — wrap in an agent to set CurrentRunContext + var agent = new TestAIAgent + { + RunAsyncFunc = async (messages, session, options, ct) => + { + var response = await pipeline.GetResponseAsync(messages, cancellationToken: ct); + return new AgentResponse(response); + } + }; + + await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession); + + // Assert + Assert.NotNull(capturedMessages); + var messageList = capturedMessages!.ToList(); + Assert.Equal(2, messageList.Count); + } + + #endregion + + #region Helpers + + /// + /// Runs a chat client within an agent context so that AIAgent.CurrentRunContext is set. + /// + private static async Task RunWithAgentContextAsync(AIContextProviderChatClient chatClient) + { + var agent = new TestAIAgent + { + RunAsyncFunc = async (messages, session, options, ct) => + { + var response = await chatClient.GetResponseAsync(messages, cancellationToken: ct); + return new AgentResponse(response); + } + }; + + await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession); + } + + /// + /// Runs a streaming chat client within an agent context so that AIAgent.CurrentRunContext is set. + /// + private static async Task RunStreamingWithAgentContextAsync(AIContextProviderChatClient chatClient, List updates) + { + var agent = new TestAIAgent + { + RunAsyncFunc = async (messages, session, options, ct) => + { + await foreach (var update in chatClient.GetStreamingResponseAsync(messages, cancellationToken: ct)) + { + updates.Add(update); + } + + return new AgentResponse([new ChatMessage(ChatRole.Assistant, "done")]); + } + }; + + await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], s_mockSession); + } + + private static IChatClient CreateMockChatClient( + Func, ChatOptions?, CancellationToken, Task> onGetResponse) + { + var mock = new Mock(); + mock.Setup(c => c.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns((IEnumerable m, ChatOptions? o, CancellationToken ct) => onGetResponse(m, o, ct)); + return mock.Object; + } + + private static IChatClient CreateMockStreamingChatClient( + Func, ChatOptions?, CancellationToken, IAsyncEnumerable> onGetStreamingResponse) + { + var mock = new Mock(); + mock.Setup(c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns((IEnumerable m, ChatOptions? o, CancellationToken ct) => onGetStreamingResponse(m, o, ct)); + return mock.Object; + } + + private static async IAsyncEnumerable ToAsyncEnumerableAsync(params ChatResponseUpdate[] updates) + { + foreach (var update in updates) + { + yield return update; + } + + await Task.CompletedTask; + } + + /// + /// A test AIContextProvider that provides configurable messages, tools, and instructions. + /// + private sealed class TestAIContextProvider : AIContextProvider + { + private readonly string _stateKey; + private readonly IEnumerable _provideMessages; + private readonly string? _provideInstructions; + private readonly IEnumerable? _provideTools; + + public bool InvokedAsyncCalled { get; private set; } + + public InvokedContext? LastInvokedContext { get; private set; } + + public override string StateKey => this._stateKey; + + public TestAIContextProvider( + string stateKey, + IEnumerable? provideMessages = null, + string? provideInstructions = null, + IEnumerable? provideTools = null) + { + this._stateKey = stateKey; + this._provideMessages = provideMessages ?? []; + this._provideInstructions = provideInstructions; + this._provideTools = provideTools; + } + + protected override ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + return new ValueTask(new AIContext + { + Messages = this._provideMessages, + Instructions = this._provideInstructions, + Tools = this._provideTools, + }); + } + + protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + this.InvokedAsyncCalled = true; + this.LastInvokedContext = context; + return default; + } + } + + /// + /// A minimal AITool for testing. + /// + private sealed class TestAITool : AITool; + + #endregion +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/MessageAIContextProviderAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs similarity index 98% rename from dotnet/tests/Microsoft.Agents.AI.UnitTests/MessageAIContextProviderAgentTests.cs rename to dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs index ed2f82321d..3ba170a8d5 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/MessageAIContextProviderAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs @@ -12,7 +12,7 @@ namespace Microsoft.Agents.AI.UnitTests; /// /// Unit tests for the class and -/// the builder extension. +/// the builder extension. /// public class MessageAIContextProviderAgentTests { @@ -355,7 +355,7 @@ public async Task UseExtension_CreatesWorkingPipelineAsync() }); var pipeline = new AIAgentBuilder(innerAgent) - .Use([provider]) + .UseAIContextProviders([provider]) .Build(); // Act @@ -385,7 +385,7 @@ public async Task UseExtension_MultipleProviders_AllAppliedAsync() }); var pipeline = new AIAgentBuilder(innerAgent) - .Use([provider1, provider2]) + .UseAIContextProviders([provider1, provider2]) .Build(); // Act