Skip to content

Commit

Permalink
.Net Agents - Fix Aggregator Streaming for Nested Mode (#9669)
Browse files Browse the repository at this point in the history
### Motivation and Context
<!-- Thank you for your contribution to the semantic-kernel repo!
Please help reviewers and future users, providing the following
information:
  1. Why is this change required?
  2. What problem does it solve?
  3. What scenario does it contribute to?
  4. If it fixes an open issue, please link to the issue here.
-->

Fixes: #8677

Aggretator agent not yielding content for streamed response when
Mode=Nested.

### Description
<!-- Describe your changes, the overall approach, the underlying design.
These notes will help understanding how your code works. Thanks! -->

Analyzed and addressed behavior and added integration tests.

### Contribution Checklist
<!-- Before submitting this PR, please make sure: -->

- [X] The code builds clean without any errors or warnings
- [X] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [X] All unit tests pass, and I have added new tests where possible
- [X] I didn't break anyone 😄
  • Loading branch information
crickman authored Nov 12, 2024
1 parent 63d1dc7 commit 4cd7f07
Show file tree
Hide file tree
Showing 2 changed files with 192 additions and 6 deletions.
14 changes: 8 additions & 6 deletions dotnet/src/Agents/Abstractions/AggregatorChannel.cs
Original file line number Diff line number Diff line change
Expand Up @@ -54,29 +54,31 @@ protected internal override IAsyncEnumerable<ChatMessageContent> GetHistoryAsync
/// <inheritdoc/>
protected internal override async IAsyncEnumerable<StreamingChatMessageContent> InvokeStreamingAsync(AggregatorAgent agent, IList<ChatMessageContent> messages, [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
int messageCount = await this._chat.GetChatMessagesAsync(cancellationToken).CountAsync(cancellationToken).ConfigureAwait(false);
int initialCount = await this._chat.GetChatMessagesAsync(cancellationToken).CountAsync(cancellationToken).ConfigureAwait(false);

if (agent.Mode == AggregatorMode.Flat)
await foreach (StreamingChatMessageContent message in this._chat.InvokeStreamingAsync(cancellationToken).ConfigureAwait(false))
{
await foreach (StreamingChatMessageContent message in this._chat.InvokeStreamingAsync(cancellationToken).ConfigureAwait(false))
if (agent.Mode == AggregatorMode.Flat)
{
yield return message;
}
}

ChatMessageContent[] history = await this._chat.GetChatMessagesAsync(cancellationToken).ToArrayAsync(cancellationToken).ConfigureAwait(false);
if (history.Length > messageCount)
if (history.Length > initialCount)
{
if (agent.Mode == AggregatorMode.Flat)
{
for (int index = messageCount; index < messages.Count; ++index)
for (int index = history.Length - 1; index >= initialCount; --index)
{
messages.Add(history[index]);
}
}
else if (agent.Mode == AggregatorMode.Nested)
{
messages.Add(history[history.Length - 1]);
ChatMessageContent finalMessage = history[0]; // Order descending
yield return new StreamingChatMessageContent(finalMessage.Role, finalMessage.Content) { AuthorName = finalMessage.AuthorName };
messages.Add(finalMessage);
}
}
}
Expand Down
184 changes: 184 additions & 0 deletions dotnet/src/IntegrationTests/Agents/AggregatorAgentTests.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,184 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Azure.Identity;
using Microsoft.Extensions.Configuration;
using Microsoft.SemanticKernel;
using Microsoft.SemanticKernel.Agents;
using Microsoft.SemanticKernel.Agents.Chat;
using Microsoft.SemanticKernel.ChatCompletion;
using SemanticKernel.IntegrationTests.TestSettings;
using xRetry;
using Xunit;

namespace SemanticKernel.IntegrationTests.Agents;

#pragma warning disable xUnit1004 // Contains test methods used in manual verification. Disable warning for this file only.

public sealed class AggregatorAgentTests()
{
private readonly IKernelBuilder _kernelBuilder = Kernel.CreateBuilder();
private readonly IConfigurationRoot _configuration = new ConfigurationBuilder()
.AddJsonFile(path: "testsettings.json", optional: true, reloadOnChange: true)
.AddJsonFile(path: "testsettings.development.json", optional: true, reloadOnChange: true)
.AddEnvironmentVariables()
.AddUserSecrets<OpenAIAssistantAgentTests>()
.Build();

/// <summary>
/// Integration test for <see cref="AggregatorAgent"/> non-streamed nested response.
/// </summary>
[RetryFact(typeof(HttpOperationException))]
public async Task AggregatorAgentFlatResponseAsync()
{
// Arrange
AggregatorAgent aggregatorAgent = new(() => this.CreateChatProvider())
{
Mode = AggregatorMode.Flat,
};

AgentGroupChat chat = new();
chat.AddChatMessage(new ChatMessageContent(AuthorRole.User, "1"));

// Act
ChatMessageContent[] responses = await chat.InvokeAsync(aggregatorAgent).ToArrayAsync();

// Assert
ChatMessageContent[] innerHistory = await chat.GetChatMessagesAsync(aggregatorAgent).ToArrayAsync();
Assert.Equal(6, innerHistory.Length);
Assert.Equal(5, responses.Length);
Assert.NotNull(responses[4].Content);
AssertResponseContent(responses[4]);
}

/// <summary>
/// Integration test for <see cref="AggregatorAgent"/> non-streamed nested response.
/// </summary>
[RetryFact(typeof(HttpOperationException))]
public async Task AggregatorAgentNestedResponseAsync()
{
// Arrange
AggregatorAgent aggregatorAgent = new(() => this.CreateChatProvider())
{
Mode = AggregatorMode.Nested,
};

AgentGroupChat chat = new();
chat.AddChatMessage(new ChatMessageContent(AuthorRole.User, "1"));

// Act
ChatMessageContent[] responses = await chat.InvokeAsync(aggregatorAgent).ToArrayAsync();

// Assert
ChatMessageContent[] innerHistory = await chat.GetChatMessagesAsync(aggregatorAgent).ToArrayAsync();
Assert.Equal(6, innerHistory.Length);
Assert.Single(responses);
Assert.NotNull(responses[0].Content);
AssertResponseContent(responses[0]);
}

/// <summary>
/// Integration test for <see cref="AggregatorAgent"/> non-streamed response.
/// </summary>
[RetryFact(typeof(HttpOperationException))]
public async Task AggregatorAgentFlatStreamAsync()
{
// Arrange
AggregatorAgent aggregatorAgent = new(() => this.CreateChatProvider())
{
Mode = AggregatorMode.Flat,
};

AgentGroupChat chat = new();
chat.AddChatMessage(new ChatMessageContent(AuthorRole.User, "1"));

// Act
StreamingChatMessageContent[] streamedResponse = await chat.InvokeStreamingAsync(aggregatorAgent).ToArrayAsync();

// Assert
ChatMessageContent[] fullResponses = await chat.GetChatMessagesAsync().ToArrayAsync();
ChatMessageContent[] innerHistory = await chat.GetChatMessagesAsync(aggregatorAgent).ToArrayAsync();
Assert.NotEmpty(streamedResponse);
Assert.Equal(6, innerHistory.Length);
Assert.Equal(6, fullResponses.Length);
Assert.NotNull(fullResponses[0].Content);
AssertResponseContent(fullResponses[0]);
}

/// <summary>
/// Integration test for <see cref="AggregatorAgent"/> non-streamed response.
/// </summary>
[RetryFact(typeof(HttpOperationException))]
public async Task AggregatorAgentNestedStreamAsync()
{
// Arrange
AggregatorAgent aggregatorAgent = new(() => this.CreateChatProvider())
{
Mode = AggregatorMode.Nested,
};

AgentGroupChat chat = new();
chat.AddChatMessage(new ChatMessageContent(AuthorRole.User, "1"));

// Act
StreamingChatMessageContent[] streamedResponse = await chat.InvokeStreamingAsync(aggregatorAgent).ToArrayAsync();

// Assert
ChatMessageContent[] fullResponses = await chat.GetChatMessagesAsync().ToArrayAsync();
ChatMessageContent[] innerHistory = await chat.GetChatMessagesAsync(aggregatorAgent).ToArrayAsync();
Assert.NotEmpty(streamedResponse);
Assert.Equal(6, innerHistory.Length);
Assert.Equal(2, fullResponses.Length);
Assert.NotNull(fullResponses[0].Content);
AssertResponseContent(fullResponses[0]);
}

private static void AssertResponseContent(ChatMessageContent response)
{
// Counting is hard
Assert.True(
response.Content!.Contains("five", StringComparison.OrdinalIgnoreCase) ||
response.Content!.Contains("six", StringComparison.OrdinalIgnoreCase) ||
response.Content!.Contains("seven", StringComparison.OrdinalIgnoreCase) ||
response.Content!.Contains("eight", StringComparison.OrdinalIgnoreCase),
$"Content: {response}");
}

private AgentGroupChat CreateChatProvider()
{
// Arrange
AzureOpenAIConfiguration configuration = this._configuration.GetSection("AzureOpenAI").Get<AzureOpenAIConfiguration>()!;

this._kernelBuilder.AddAzureOpenAIChatCompletion(
configuration.ChatDeploymentName!,
configuration.Endpoint,
new AzureCliCredential());

Kernel kernel = this._kernelBuilder.Build();

ChatCompletionAgent agent =
new()
{
Kernel = kernel,
Instructions = "Your job is to count. Always add one to the previous number and respond using the english word for that number, without explanation.",
};

return new AgentGroupChat(agent)
{
ExecutionSettings = new()
{
TerminationStrategy = new CountTerminationStrategy(5)
}
};
}

private sealed class CountTerminationStrategy(int maximumResponseCount) : TerminationStrategy
{
// Terminate when the assistant has responded N times.
protected override Task<bool> ShouldAgentTerminateAsync(Agent agent, IReadOnlyList<ChatMessageContent> history, CancellationToken cancellationToken)
=> Task.FromResult(history.Count(message => message.Role == AuthorRole.Assistant) >= maximumResponseCount);
}
}

0 comments on commit 4cd7f07

Please sign in to comment.