// 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.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); } }