autogen/dotnet/test/AutoGen.Tests/MiddlewareAgentTest.cs
Griffin Bassman 9af6883fbe
fix: dotnet test CI and standardize test categories (#5286)
Co-authored-by: Ryan Sweet <rysweet@microsoft.com>
2025-02-03 11:49:08 -05:00

107 lines
4.2 KiB
C#

// Copyright (c) Microsoft Corporation. All rights reserved.
// MiddlewareAgentTest.cs
using System.Linq;
using System.Threading.Tasks;
using FluentAssertions;
using Xunit;
namespace AutoGen.Tests;
[Trait("Category", "UnitV1")]
public class MiddlewareAgentTest
{
[Fact]
public async Task MiddlewareAgentUseTestAsync()
{
IAgent echoAgent = new EchoAgent("echo");
var middlewareAgent = new MiddlewareAgent(echoAgent);
// no middleware added
// the reply should be the same as the original agent
middlewareAgent.Name.Should().Be("echo");
var reply = await middlewareAgent.SendAsync("hello");
reply.GetContent().Should().Be("hello");
middlewareAgent.Use(async (messages, options, agent, ct) =>
{
var lastMessage = messages.Last() as TextMessage;
lastMessage!.Content = $"[middleware 0] {lastMessage.Content}";
return await agent.GenerateReplyAsync(messages, options, ct);
});
reply = await middlewareAgent.SendAsync("hello");
reply.GetContent().Should().Be("[middleware 0] hello");
middlewareAgent.Use(async (messages, options, agent, ct) =>
{
var lastMessage = messages.Last() as TextMessage;
lastMessage!.Content = $"[middleware 1] {lastMessage.Content}";
return await agent.GenerateReplyAsync(messages, options, ct);
});
// when multiple middleware are added, they will be executed in LIFO order
reply = await middlewareAgent.SendAsync("hello");
reply.GetContent().Should().Be("[middleware 0] [middleware 1] hello");
// test short cut
// short cut middleware will not call next middleware
middlewareAgent.Use(async (messages, options, next, ct) =>
{
var lastMessage = messages.Last() as TextMessage;
lastMessage!.Content = $"[middleware shortcut] {lastMessage.Content}";
return lastMessage;
});
reply = await middlewareAgent.SendAsync("hello");
reply.GetContent().Should().Be("[middleware shortcut] hello");
}
[Fact]
public async Task RegisterMiddlewareTestAsync()
{
var echoAgent = new EchoAgent("echo");
// RegisterMiddleware will return a new agent and keep the original agent unchanged
var middlewareAgent = echoAgent.RegisterMiddleware(async (messages, options, agent, ct) =>
{
var lastMessage = messages.Last() as TextMessage;
lastMessage!.Content = $"[middleware 0] {lastMessage.Content}";
return await agent.GenerateReplyAsync(messages, options, ct);
});
middlewareAgent.Should().BeOfType<MiddlewareAgent<EchoAgent>>();
middlewareAgent.Middlewares.Count().Should().Be(1);
var reply = await middlewareAgent.SendAsync("hello");
reply.GetContent().Should().Be("[middleware 0] hello");
reply = await echoAgent.SendAsync("hello");
reply.GetContent().Should().Be("hello");
// when multiple middleware are added, they will be executed in LIFO order
middlewareAgent = middlewareAgent.RegisterMiddleware(async (messages, options, agent, ct) =>
{
var lastMessage = messages.Last() as TextMessage;
lastMessage!.Content = $"[middleware 1] {lastMessage.Content}";
return await agent.GenerateReplyAsync(messages, options, ct);
});
middlewareAgent.Middlewares.Count().Should().Be(2);
reply = await middlewareAgent.SendAsync("hello");
reply.GetContent().Should().Be("[middleware 0] [middleware 1] hello");
// test short cut
// short cut middleware will not call next middleware
middlewareAgent = middlewareAgent.RegisterMiddleware(async (messages, options, agent, ct) =>
{
var lastMessage = messages.Last() as TextMessage;
lastMessage!.Content = $"[middleware shortcut] {lastMessage.Content}";
return lastMessage;
});
reply = await middlewareAgent.SendAsync("hello");
reply.GetContent().Should().Be("[middleware shortcut] hello");
middlewareAgent.Middlewares.Count().Should().Be(3);
}
}