From d761c92a5207599f4df94656c88fd451694239a7 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Thu, 12 Jun 2025 07:18:47 +0100 Subject: [PATCH] .Net ChatClientAgent Streaming API Impl. (#69) * Add Streaming API * Removing InstructionsRole * Updating thread notification strategy * Fix net472 failing * Address typo --- .../ChatCompletion/ChatClientAgent.cs | 169 +++++--- .../ChatClientAgentExtensions.cs | 39 +- .../ChatClientAgentLogMessages.cs | 12 +- .../ChatCompletion/ChatClientAgentTests.cs | 171 ++------ .../ChatClientAgentThreadTests.cs | 402 ++++++++++++++++++ 5 files changed, 587 insertions(+), 206 deletions(-) diff --git a/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgent.cs index 199069d46a..119cbfcddb 100644 --- a/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgent.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -19,6 +20,7 @@ public sealed class ChatClientAgent : Agent { private readonly ChatClientAgentOptions? _agentOptions; private readonly ILogger _logger; + private readonly Type _chatClientType; /// /// Initializes a new instance of the class. @@ -30,27 +32,17 @@ public sealed class ChatClientAgent : Agent { Throw.IfNull(chatClient); + this._chatClientType = chatClient.GetType(); this.ChatClient = chatClient.AsAgentInvokingChatClient(); this._agentOptions = options; this._logger = (loggerFactory ?? chatClient.GetService() ?? NullLoggerFactory.Instance).CreateLogger(); } /// - /// The chat client. + /// The underlying chat client used by the agent to invoke chat completions. /// public IChatClient ChatClient { get; } - /// - /// Gets the role used for agent instructions. Defaults to "system". - /// - /// - /// Certain versions of "O*" series (deep reasoning) models require the instructions - /// to be provided as "developer" role. Other versions support neither role and - /// an agent targeting such a model cannot provide instructions. Agent functionality - /// will be dictated entirely by the provided plugins. - /// - public ChatRole InstructionsRole { get; set; } = ChatRole.System; - /// public override string Id => this._agentOptions?.Id ?? base.Id; @@ -72,6 +64,110 @@ public sealed class ChatClientAgent : Agent { Throw.IfNull(messages); + (ChatClientAgentThread chatClientThread, ChatOptions? chatOptions, List threadMessages) = + await this.PrepareThreadAndMessagesAsync(thread, messages, options, cancellationToken).ConfigureAwait(false); + + var agentName = this.GetAgentName(); + + this._logger.LogAgentChatClientInvokingAgent(nameof(RunAsync), this.Id, agentName, this._chatClientType); + + ChatResponse chatResponse = await this.ChatClient.GetResponseAsync(threadMessages, chatOptions, cancellationToken).ConfigureAwait(false); + + this._logger.LogAgentChatClientInvokedAgent(nameof(RunAsync), this.Id, agentName, this._chatClientType, messages.Count); + + // Only notify the thread of new messages if the chatResponse was successful to avoid inconsistent messages state in the thread. + await this.NotifyThreadOfNewMessagesAsync(chatClientThread, messages, cancellationToken).ConfigureAwait(false); + + // Ensure that the author name is set for each message in the response. + foreach (ChatMessage chatResponseMessage in chatResponse.Messages) + { + chatResponseMessage.AuthorName ??= agentName; + } + + // Convert the chat response messages to a valid IReadOnlyCollection for notification signatures below. + var chatResponseMessages = chatResponse.Messages as IReadOnlyCollection ?? chatResponse.Messages.ToArray(); + + await this.NotifyThreadOfNewMessagesAsync(chatClientThread, chatResponseMessages, cancellationToken).ConfigureAwait(false); + if (options?.OnIntermediateMessages is not null) + { + await options.OnIntermediateMessages(chatResponseMessages).ConfigureAwait(false); + } + + return chatResponse; + } + + /// + public override async IAsyncEnumerable RunStreamingAsync( + IReadOnlyCollection messages, + AgentThread? thread = null, + AgentRunOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + Throw.IfNull(messages); + + (ChatClientAgentThread chatClientThread, ChatOptions? chatOptions, List threadMessages) = + await this.PrepareThreadAndMessagesAsync(thread, messages, options, cancellationToken).ConfigureAwait(false); + + int messageCount = threadMessages.Count; + var agentName = this.GetAgentName(); + + this._logger.LogAgentChatClientInvokingAgent(nameof(RunStreamingAsync), this.Id, agentName, this._chatClientType); + + // Using the enumerator to ensure we consider the case where no updates are returned for notification. + var responseUpdatesEnumerator = this.ChatClient.GetStreamingResponseAsync(threadMessages, chatOptions, cancellationToken).GetAsyncEnumerator(cancellationToken); + + this._logger.LogAgentChatClientInvokedStreamingAgent(nameof(RunStreamingAsync), this.Id, agentName, this._chatClientType); + + List responseUpdates = []; + + // Ensure we start the streaming request + var hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); + + // To avoid inconsistent state we only notify the thread of the input messages if no error occurs after the initial request. + await this.NotifyThreadOfNewMessagesAsync(chatClientThread, messages, cancellationToken).ConfigureAwait(false); + + while (hasUpdates) + { + var update = responseUpdatesEnumerator.Current; + if (update is not null) + { + responseUpdates.Add(update); + update.AuthorName ??= agentName; + yield return update; + } + + hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false); + } + + var chatResponse = responseUpdates.ToChatResponse(); + var chatResponseMessages = chatResponse.Messages as IReadOnlyCollection ?? chatResponse.Messages.ToArray(); + + await this.NotifyThreadOfNewMessagesAsync(chatClientThread, chatResponseMessages, cancellationToken).ConfigureAwait(false); + if (options?.OnIntermediateMessages is not null) + { + await options.OnIntermediateMessages(chatResponseMessages).ConfigureAwait(false); + } + } + + /// + public override AgentThread GetNewThread() => new ChatClientAgentThread(); + + #region Private + + /// + /// Prepares the thread, chat options, and messages for agent execution. + /// + /// The conversation thread to use or create. + /// The input messages to use. + /// Optional parameters for agent invocation. + /// The cancellation token. + /// A tuple containing the thread, chat options, and thread messages. + private async Task<(ChatClientAgentThread thread, ChatOptions? chatOptions, List threadMessages)> PrepareThreadAndMessagesAsync( + AgentThread? thread, + IReadOnlyCollection inputMessages, + AgentRunOptions? options, + CancellationToken cancellationToken) + { // Retrieve chat options from the provided AgentRunOptions if available. ChatOptions? chatOptions = (options as ChatClientAgentRunOptions)?.ChatOptions; @@ -87,65 +183,28 @@ public sealed class ChatClientAgent : Agent } } - // Append to the existing thread messages the messages that were passed in to this call. - threadMessages.AddRange(messages); - // Update the messages with agent instructions. this.UpdateThreadMessagesWithAgentInstructions(threadMessages, options); - var agentName = this.Name ?? "UnnamedAgent"; - Type serviceType = this.ChatClient.GetType(); + // Add the input messages to the end of thread messages. + threadMessages.AddRange(inputMessages); - this._logger.LogAgentChatClientInvokingAgent(nameof(RunAsync), this.Id, agentName, serviceType); - - ChatResponse chatResponse = await this.ChatClient.GetResponseAsync(threadMessages, chatOptions, cancellationToken).ConfigureAwait(false); - - this._logger.LogAgentChatClientInvokedAgent(nameof(RunAsync), this.Id, agentName, serviceType, messages.Count); - - // Only notify the thread of new messages if the chatResponse was successful to avoid inconsistent messages state in the thread. - await this.NotifyThreadOfNewMessagesAsync(chatClientThread, messages, cancellationToken).ConfigureAwait(false); - - // Ensure that the author name is set for each message in the response. - foreach (ChatMessage chatResponseMessage in chatResponse.Messages) - { - chatResponseMessage.AuthorName ??= agentName; - } - - // Convert the chat response messages to a valid IReadOnlyCollection for notification signatures below. - var chatResponseMessages = chatResponse.Messages.ToArray(); - - await this.NotifyThreadOfNewMessagesAsync(chatClientThread, chatResponseMessages, cancellationToken).ConfigureAwait(false); - if (options?.OnIntermediateMessages is not null) - { - await options.OnIntermediateMessages(chatResponseMessages).ConfigureAwait(false); - } - - return chatResponse; + return (chatClientThread, chatOptions, threadMessages); } - /// - public override IAsyncEnumerable RunStreamingAsync(IReadOnlyCollection messages, AgentThread? thread = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) - { - throw new System.NotImplementedException(); - } - - /// - public override AgentThread GetNewThread() => new ChatClientAgentThread(); - - #region Private - private void UpdateThreadMessagesWithAgentInstructions(List threadMessages, AgentRunOptions? options) { if (!string.IsNullOrWhiteSpace(options?.AdditionalInstructions)) { - threadMessages.Insert(0, new(this.InstructionsRole, options?.AdditionalInstructions) { AuthorName = this.Name }); + threadMessages.Insert(0, new(ChatRole.System, options?.AdditionalInstructions) { AuthorName = this.Name }); } if (!string.IsNullOrWhiteSpace(this.Instructions)) { - threadMessages.Insert(0, new(this.InstructionsRole, this.Instructions) { AuthorName = this.Name }); + threadMessages.Insert(0, new(ChatRole.System, this.Instructions) { AuthorName = this.Name }); } } + private string GetAgentName() => this.Name ?? "UnnamedAgent"; #endregion } diff --git a/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgentExtensions.cs b/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgentExtensions.cs index a61be61921..32dfc3ff7e 100644 --- a/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgentExtensions.cs +++ b/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgentExtensions.cs @@ -14,26 +14,49 @@ namespace Microsoft.Agents; public static class ChatClientAgentExtensions { /// - /// Allow running a chat client agent with a configuration. + /// Run the agent with the provided message and arguments. /// /// Target agent to run. - /// Messages to send to the agent. - /// Optional thread to use for the agent. - /// Optional agent run options. + /// The messages to pass to the agent. + /// The conversation thread to continue with this invocation. If not provided, creates a new thread. The thread will be mutated with the provided messages and agent reponse. + /// Optional parameters for agent invocation. /// Optional chat options. - /// Optional cancellation token. - /// A task representing the asynchronous operation, with the chat response. + /// The to monitor for cancellation requests. The default is . + /// A containing the list of items. public static Task RunAsync( this ChatClientAgent agent, IReadOnlyCollection messages, AgentThread? thread = null, - AgentRunOptions? agentOptions = null, + AgentRunOptions? agentRunOptions = null, ChatOptions? chatOptions = null, CancellationToken cancellationToken = default) { Throw.IfNull(agent); Throw.IfNull(messages); - return agent.RunAsync(messages, thread, new ChatClientAgentRunOptions(agentOptions, chatOptions), cancellationToken); + return agent.RunAsync(messages, thread, new ChatClientAgentRunOptions(agentRunOptions, chatOptions), cancellationToken); + } + + /// + /// Run the agent with the provided message and arguments. + /// + /// Target agent to run. + /// The messages to pass to the agent. + /// The conversation thread to continue with this invocation. If not provided, creates a new thread. The thread will be mutated with the provided messages and agent reponse. + /// Optional parameters for agent invocation. + /// Optional chat options. + /// The to monitor for cancellation requests. The default is . + public static IAsyncEnumerable RunStreamingAsync( + this ChatClientAgent agent, + IReadOnlyCollection messages, + AgentThread? thread = null, + AgentRunOptions? agentRunOptions = null, + ChatOptions? chatOptions = null, + CancellationToken cancellationToken = default) + { + Throw.IfNull(agent); + Throw.IfNull(messages); + + return agent.RunStreamingAsync(messages, thread, new ChatClientAgentRunOptions(agentRunOptions, chatOptions), cancellationToken); } } diff --git a/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgentLogMessages.cs b/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgentLogMessages.cs index 87347dcebb..a54be742de 100644 --- a/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgentLogMessages.cs +++ b/dotnet/src/Microsoft.Agents/ChatCompletion/ChatClientAgentLogMessages.cs @@ -23,13 +23,13 @@ internal static partial class ChatClientAgentLogMessages [LoggerMessage( EventId = 0, Level = LogLevel.Debug, - Message = "[{MethodName}] Agent {AgentId}/{AgentName} Invoking service {ServiceType}.")] + Message = "[{MethodName}] Agent {AgentId}/{AgentName} Invoking client {ClientType}.")] public static partial void LogAgentChatClientInvokingAgent( this ILogger logger, string methodName, string agentId, string agentName, - Type serviceType); + Type clientType); /// /// Logs invoked agent (complete). @@ -37,13 +37,13 @@ internal static partial class ChatClientAgentLogMessages [LoggerMessage( EventId = 0, Level = LogLevel.Information, - Message = "[{MethodName}] Agent {AgentId}/{AgentName} Invoked service {ServiceType} with message count: {MessageCount}.")] + Message = "[{MethodName}] Agent {AgentId}/{AgentName} Invoked client {ClientType} with message count: {MessageCount}.")] public static partial void LogAgentChatClientInvokedAgent( this ILogger logger, string methodName, string agentId, string agentName, - Type serviceType, + Type clientType, int messageCount); /// @@ -52,11 +52,11 @@ internal static partial class ChatClientAgentLogMessages [LoggerMessage( EventId = 0, Level = LogLevel.Information, - Message = "[{MethodName}] Agent {AgentId}/{AgentName} Invoked service {ServiceType}.")] + Message = "[{MethodName}] Agent {AgentId}/{AgentName} Invoked client {ClientType}.")] public static partial void LogAgentChatClientInvokedStreamingAgent( this ILogger logger, string methodName, string agentId, string agentName, - Type serviceType); + Type clientType); } diff --git a/dotnet/tests/Microsoft.Agents.UnitTests/ChatCompletion/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.UnitTests/ChatCompletion/ChatClientAgentTests.cs index 9e6a136180..94126ae57a 100644 --- a/dotnet/tests/Microsoft.Agents.UnitTests/ChatCompletion/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.UnitTests/ChatCompletion/ChatClientAgentTests.cs @@ -16,7 +16,7 @@ public class ChatClientAgentTests /// Verify the invocation and response of . /// [Fact] - public void VerifyChatCompletionAgentDefinition() + public void VerifyChatClientAgentDefinition() { // Arrange var chatClient = new Mock().Object; @@ -38,42 +38,6 @@ public class ChatClientAgentTests Assert.Equal("test instructions", agent.Instructions); Assert.NotNull(agent.ChatClient); Assert.Equal("AgentInvokingChatClient", agent.ChatClient.GetType().Name); - Assert.Equal(ChatRole.System, agent.InstructionsRole); - } - - /// - /// Verify the invocation and response of . - /// - [Fact] - public async Task VerifyChatCompletionAgentInvocationAsync() - { - // Arrange - Mock mockService = new(); - mockService.Setup( - s => s.GetResponseAsync( - It.IsAny>(), - It.IsAny(), - It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "what?")])); - - ChatClientAgent agent = - new(mockService.Object, new() - { - Instructions = "test instructions" - }); - - // Act - ChatResponse result = await agent.RunAsync([]); - - // Assert - Assert.Single(result.Messages); - - mockService.Verify( - x => - x.GetResponseAsync( - It.IsAny>(), - It.IsAny(), - It.IsAny()), - Times.Once); } /// @@ -119,47 +83,6 @@ public class ChatClientAgentTests }); } - /// - /// Verify the streaming invocation and response of . - /// - [Fact(Skip = "Not implemented yet")] - public async Task VerifyChatClientAgentStreamingAsync() - { - // Arrange - ChatResponseUpdate[] returnUpdates = - [ - new ChatResponseUpdate(role: ChatRole.Assistant, content: "wh"), - new ChatResponseUpdate(role: null, content: "at?"), - ]; - - Mock mockService = new(); - mockService.Setup( - s => s.GetStreamingResponseAsync( - It.IsAny>(), - It.IsAny(), - It.IsAny())).Returns(returnUpdates.ToAsyncEnumerable()); - - ChatClientAgent agent = - new(mockService.Object, new() - { - Instructions = "test instructions" - }); - - // Act - ChatResponseUpdate[] result = await agent.RunStreamingAsync([]).ToArrayAsync(); - - // Assert - Assert.Equal(2, result.Length); - - mockService.Verify( - x => - x.GetStreamingResponseAsync( - It.IsAny>(), - It.IsAny(), - It.IsAny()), - Times.Once); - } - /// /// Verify that RunAsync throws ArgumentNullException when messages parameter is null. /// @@ -607,77 +530,51 @@ public class ChatClientAgentTests Assert.Null(agent.Instructions); } - /// - /// Verify that InstructionsRole property has default value of System. - /// - [Fact] - public void InstructionsRoleHasDefaultValueOfSystem() - { - // Arrange - var chatClient = new Mock().Object; - ChatClientAgent agent = new(chatClient, new()); - - // Act & Assert - Assert.Equal(ChatRole.System, agent.InstructionsRole); - } - - /// - /// Verify that InstructionsRole property can be set to custom values. - /// - [Fact] - public void InstructionsRoleCanBeSetToCustomValue() - { - // Arrange - var chatClient = new Mock().Object; - ChatClientAgent agent = new(chatClient, new()); - - // Act - agent.InstructionsRole = ChatRole.User; - - // Assert - Assert.Equal(ChatRole.User, agent.InstructionsRole); - } - #endregion #region RunStreamingAsync Tests /// - /// Verify that RunStreamingAsync throws NotImplementedException. + /// Verify the streaming invocation and response of . /// [Fact] - public void RunStreamingAsyncThrowsNotImplementedException() + public async Task VerifyChatClientAgentStreamingAsync() { // Arrange - var chatClient = new Mock().Object; - ChatClientAgent agent = new(chatClient, new() { Instructions = "test instructions" }); + ChatResponseUpdate[] returnUpdates = + [ + new ChatResponseUpdate(role: ChatRole.Assistant, content: "wh"), + new ChatResponseUpdate(role: ChatRole.Assistant, content: "at?"), + ]; - // Act & Assert - Assert.Throws(() => - { - var result = agent.RunStreamingAsync([new(ChatRole.User, "test")]); - // Force enumeration to trigger the exception - result.GetAsyncEnumerator(); - }); - } + Mock mockService = new(); + mockService.Setup( + s => s.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())).Returns(returnUpdates.ToAsyncEnumerable()); - /// - /// Verify that RunStreamingAsync with string message throws NotImplementedException. - /// - [Fact] - public void RunStreamingAsyncWithStringMessageThrowsNotImplementedException() - { - // Arrange - var chatClient = new Mock().Object; - ChatClientAgent agent = new(chatClient, new() { Instructions = "test instructions" }); + ChatClientAgent agent = + new(mockService.Object, new() + { + Instructions = "test instructions" + }); - // Act & Assert - Assert.Throws(() => - { - var result = agent.RunStreamingAsync("test message"); - // Force enumeration to trigger the exception - result.GetAsyncEnumerator(); - }); + // Act + ChatResponseUpdate[] result = await agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hello")]).ToArrayAsync(); + + // Assert + Assert.Equal(2, result.Length); + Assert.Equal("wh", result[0].Text); + Assert.Equal("at?", result[1].Text); + + mockService.Verify( + x => + x.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny()), + Times.Once); } #endregion diff --git a/dotnet/tests/Microsoft.Agents.UnitTests/ChatCompletion/ChatClientAgentThreadTests.cs b/dotnet/tests/Microsoft.Agents.UnitTests/ChatCompletion/ChatClientAgentThreadTests.cs index 3893f63f99..5dbd512086 100644 --- a/dotnet/tests/Microsoft.Agents.UnitTests/ChatCompletion/ChatClientAgentThreadTests.cs +++ b/dotnet/tests/Microsoft.Agents.UnitTests/ChatCompletion/ChatClientAgentThreadTests.cs @@ -1,11 +1,15 @@ // 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; +#pragma warning disable CS0162 // Unreachable code detected + namespace Microsoft.Agents.UnitTests.ChatCompletion; public class ChatClientAgentThreadTests @@ -265,4 +269,402 @@ public class ChatClientAgentThreadTests } #endregion + + #region RunStreamingAsync Thread Notification Tests + + /// + /// Verify that thread is notified of both input and response messages when invoking the streaming API with RunStreamingAsync. + /// + [Fact] + public async Task VerifyThreadNotificationDuringStreamingAsync() + { + // Arrange + var userMessage = new ChatMessage(ChatRole.User, "Hello, streaming!"); + var assistantMessage = new ChatMessage(ChatRole.Assistant, "Hi there, streaming response!"); + + // Create streaming response updates + ChatResponseUpdate[] returnUpdates = + [ + new ChatResponseUpdate(role: ChatRole.Assistant, content: "Hi there, "), + new ChatResponseUpdate(role: null, content: "streaming response!"), + ]; + + var mockChatClient = new Mock(); + mockChatClient.Setup( + c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(returnUpdates.ToAsyncEnumerable()); + + // Create ChatClientAgent with the mocked client + var agent = new ChatClientAgent(mockChatClient.Object, new() + { + Instructions = "You are a helpful assistant" + }); + + // Get a new thread from the agent + var thread = agent.GetNewThread(); + + // Act - Run the agent with streaming to populate the thread with messages + var streamingResults = new List(); + await foreach (var update in agent.RunStreamingAsync([userMessage], thread)) + { + streamingResults.Add(update); + } + + // Assert - Verify streaming worked + Assert.Equal(2, streamingResults.Count); + + // Retrieve messages from the thread to verify notification occurred + var messagesRetrievableThread = (IMessagesRetrievableThread)thread; + var retrievedMessages = new List(); + await foreach (var message in messagesRetrievableThread.GetMessagesAsync()) + { + retrievedMessages.Add(message); + } + + // Assert - Verify that the thread was notified and contains both user and assistant messages + Assert.NotEmpty(retrievedMessages); + Assert.Equal(2, retrievedMessages.Count); + Assert.Contains(retrievedMessages, m => m.Text == "Hello, streaming!" && m.Role == ChatRole.User); + Assert.Contains(retrievedMessages, m => m.Text == "Hi there, streaming response!" && m.Role == ChatRole.Assistant); + } + + /// + /// Verify that thread accumulates both input and response messages across multiple streaming calls. + /// + [Fact] + public async Task VerifyThreadAccumulatesMessagesAcrossMultipleStreamingCallsAsync() + { + // Arrange + var firstUserMessage = new ChatMessage(ChatRole.User, "First streaming message"); + var secondUserMessage = new ChatMessage(ChatRole.User, "Second streaming message"); + + // Create streaming response updates for first call + ChatResponseUpdate[] firstReturnUpdates = + [ + new ChatResponseUpdate(role: ChatRole.Assistant, content: "First "), + new ChatResponseUpdate(role: null, content: "response"), + ]; + + // Create streaming response updates for second call + ChatResponseUpdate[] secondReturnUpdates = + [ + new ChatResponseUpdate(role: ChatRole.Assistant, content: "Second "), + new ChatResponseUpdate(role: null, content: "response"), + ]; + + var mockChatClient = new Mock(); + mockChatClient.SetupSequence( + c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(firstReturnUpdates.ToAsyncEnumerable()) + .Returns(secondReturnUpdates.ToAsyncEnumerable()); + + var agent = new ChatClientAgent(mockChatClient.Object, new()); + var thread = agent.GetNewThread(); + + // Act - Make two streaming calls + var firstStreamingResults = new List(); + await foreach (var update in agent.RunStreamingAsync([firstUserMessage], thread)) + { + firstStreamingResults.Add(update); + } + + var secondStreamingResults = new List(); + await foreach (var update in agent.RunStreamingAsync([secondUserMessage], thread)) + { + secondStreamingResults.Add(update); + } + + // Assert - Verify both streaming calls worked + Assert.Equal(2, firstStreamingResults.Count); + Assert.Equal(2, secondStreamingResults.Count); + + // Retrieve all messages from the thread + var messagesRetrievableThread = (IMessagesRetrievableThread)thread; + var retrievedMessages = new List(); + await foreach (var message in messagesRetrievableThread.GetMessagesAsync()) + { + retrievedMessages.Add(message); + } + + // Assert - Verify that the thread contains all messages in order + Assert.Equal(4, retrievedMessages.Count); + Assert.Equal("First streaming message", retrievedMessages[0].Text); + Assert.Equal("First response", retrievedMessages[1].Text); + Assert.Equal("Second streaming message", retrievedMessages[2].Text); + Assert.Equal("Second response", retrievedMessages[3].Text); + } + + /// + /// Verify that thread notification works correctly when streaming with existing thread messages. + /// Both RunAsync and RunStreamingAsync should add both input and response messages to the thread. + /// + [Fact] + public async Task VerifyStreamingWithExistingThreadMessagesAsync() + { + // Arrange + var initialUserMessage = new ChatMessage(ChatRole.User, "Initial message"); + var initialAssistantMessage = new ChatMessage(ChatRole.Assistant, "Initial response"); + var newUserMessage = new ChatMessage(ChatRole.User, "New streaming message"); + + // Setup for initial non-streaming call + var mockChatClient = new Mock(); + mockChatClient.Setup( + c => c.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(new ChatResponse([initialAssistantMessage])); + + // Setup for streaming call + ChatResponseUpdate[] streamingUpdates = + [ + new ChatResponseUpdate(role: ChatRole.Assistant, content: "Streaming "), + new ChatResponseUpdate(role: null, content: "response"), + ]; + + mockChatClient.Setup( + c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(streamingUpdates.ToAsyncEnumerable()); + + var agent = new ChatClientAgent(mockChatClient.Object, new()); + var thread = agent.GetNewThread(); + + // Act - First, make a regular call to populate the thread + await agent.RunAsync([initialUserMessage], thread); + + // Then make a streaming call + var streamingResults = new List(); + await foreach (var update in agent.RunStreamingAsync([newUserMessage], thread)) + { + streamingResults.Add(update); + } + + // Assert - Verify streaming worked + Assert.Equal(2, streamingResults.Count); + + // Retrieve all messages from the thread + var messagesRetrievableThread = (IMessagesRetrievableThread)thread; + var retrievedMessages = new List(); + await foreach (var message in messagesRetrievableThread.GetMessagesAsync()) + { + retrievedMessages.Add(message); + } + + // Assert - Verify that the thread contains all messages including the new streaming ones + Assert.Equal(4, retrievedMessages.Count); + Assert.Equal("Initial message", retrievedMessages[0].Text); + Assert.Equal("Initial response", retrievedMessages[1].Text); + Assert.Equal("New streaming message", retrievedMessages[2].Text); + Assert.Equal("Streaming response", retrievedMessages[3].Text); + } + + /// + /// Verify that thread is notified of input messages even when zero streaming updates are received. + /// + [Fact] + public async Task VerifyThreadNotificationWithZeroStreamingUpdatesAsync() + { + // Arrange + var userMessage = new ChatMessage(ChatRole.User, "Hello with no response!"); + + // Create empty streaming response (no updates) + ChatResponseUpdate[] returnUpdates = []; + + var mockChatClient = new Mock(); + mockChatClient.Setup( + c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(returnUpdates.ToAsyncEnumerable()); + + var agent = new ChatClientAgent(mockChatClient.Object, new()); + var thread = agent.GetNewThread(); + + // Act - Run the agent with streaming that returns no updates + var streamingResults = new List(); + await foreach (var update in agent.RunStreamingAsync([userMessage], thread)) + { + streamingResults.Add(update); + } + + // Assert - Verify no streaming updates were received + Assert.Empty(streamingResults); + + // Retrieve messages from the thread to verify notification occurred + var messagesRetrievableThread = (IMessagesRetrievableThread)thread; + var retrievedMessages = new List(); + await foreach (var message in messagesRetrievableThread.GetMessagesAsync()) + { + retrievedMessages.Add(message); + } + + // Assert - Verify that the thread was notified of input messages even with zero updates + // The fallback mechanism should ensure input messages are added to the thread + Assert.Single(retrievedMessages); + Assert.Contains(retrievedMessages, m => m.Text == "Hello with no response!" && m.Role == ChatRole.User); + } + + /// + /// Verify that thread is notified of input messages only once even with multiple streaming updates. + /// + [Fact] + public async Task VerifyThreadNotificationWithMultipleStreamingUpdatesAsync() + { + // Arrange + var userMessage = new ChatMessage(ChatRole.User, "Hello with many updates!"); + + // Create multiple streaming response updates + ChatResponseUpdate[] returnUpdates = + [ + new ChatResponseUpdate(role: ChatRole.Assistant, content: "First "), + new ChatResponseUpdate(role: null, content: "update, "), + new ChatResponseUpdate(role: null, content: "second "), + new ChatResponseUpdate(role: null, content: "update, "), + new ChatResponseUpdate(role: null, content: "third "), + new ChatResponseUpdate(role: null, content: "update!"), + ]; + + var mockChatClient = new Mock(); + mockChatClient.Setup( + c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(returnUpdates.ToAsyncEnumerable()); + + var agent = new ChatClientAgent(mockChatClient.Object, new()); + var thread = agent.GetNewThread(); + + // Act - Run the agent with streaming that returns multiple updates + var streamingResults = new List(); + await foreach (var update in agent.RunStreamingAsync([userMessage], thread)) + { + streamingResults.Add(update); + } + + // Assert - Verify all streaming updates were received + Assert.Equal(6, streamingResults.Count); + + // Retrieve messages from the thread to verify notification occurred + var messagesRetrievableThread = (IMessagesRetrievableThread)thread; + var retrievedMessages = new List(); + await foreach (var message in messagesRetrievableThread.GetMessagesAsync()) + { + retrievedMessages.Add(message); + } + + // Assert - Verify that the thread contains both input and response messages + // Input message should be added only once despite multiple updates + Assert.Equal(2, retrievedMessages.Count); + Assert.Contains(retrievedMessages, m => m.Text == "Hello with many updates!" && m.Role == ChatRole.User); + Assert.Contains(retrievedMessages, m => m.Text == "First update, second update, third update!" && m.Role == ChatRole.Assistant); + } + + /// + /// Verify that thread is NOT notified of input messages when an exception occurs during streaming. + /// + [Fact] + public async Task VerifyThreadNotNotifiedWhenStreamingThrowsExceptionAsync() + { + // Arrange + var userMessage = new ChatMessage(ChatRole.User, "Hello that will fail!"); + + var mockChatClient = new Mock(); + mockChatClient.Setup( + c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Throws(new InvalidOperationException("Streaming failed")); + + var agent = new ChatClientAgent(mockChatClient.Object, new()); + var thread = agent.GetNewThread(); + + // Act & Assert - Verify that streaming throws an exception + await Assert.ThrowsAsync(async () => + { + await foreach (var update in agent.RunStreamingAsync([userMessage], thread)) + { + Assert.Fail("Should not yield updates."); + } + }); + + // Retrieve messages from the thread to verify NO notification occurred + var messagesRetrievableThread = (IMessagesRetrievableThread)thread; + var retrievedMessages = new List(); + await foreach (var message in messagesRetrievableThread.GetMessagesAsync()) + { + retrievedMessages.Add(message); + } + + // Assert - Verify that the thread was NOT notified of any messages due to the exception + // This ensures that failed operations don't leave the thread in an inconsistent state + Assert.Empty(retrievedMessages); + } + + /// + /// Verify that thread is NOT notified of input messages when an exception occurs after some streaming updates. + /// + [Fact] + public async Task VerifyThreadNotNotifiedWhenStreamingThrowsExceptionAfterUpdatesAsync() + { + // Arrange + var userMessage = new ChatMessage(ChatRole.User, "Hello that will partially fail!"); + + // Create an async enumerable that yields some updates then throws + static async IAsyncEnumerable GetUpdatesWithExceptionAsync() + { + await Task.CompletedTask; // Simulate async operation + throw new InvalidOperationException("Streaming failed after partial response"); + yield break; + } + + var mockChatClient = new Mock(); + mockChatClient.Setup( + c => c.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Returns(GetUpdatesWithExceptionAsync()); + + var agent = new ChatClientAgent(mockChatClient.Object, new()); + var thread = agent.GetNewThread(); + + // Act & Assert - Verify that streaming throws an exception after some updates + var streamingResults = new List(); + await Assert.ThrowsAsync(async () => + { + await foreach (var update in agent.RunStreamingAsync([userMessage], thread)) + { + streamingResults.Add(update); + } + }); + + // Verify that some updates were received before the exception + Assert.Empty(streamingResults); + + // Retrieve messages from the thread to verify NO notification occurred + var messagesRetrievableThread = (IMessagesRetrievableThread)thread; + var retrievedMessages = new List(); + await foreach (var message in messagesRetrievableThread.GetMessagesAsync()) + { + retrievedMessages.Add(message); + } + + // Assert - Verify that the thread was NOT notified of any messages due to the exception + // Even though some updates were received, the exception should prevent thread notification + Assert.Empty(retrievedMessages); + } + + #endregion }