agent-framework/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestReplayAgent.cs

// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;

namespace Microsoft.Agents.AI.Workflows.UnitTests;

public class TestReplayAgent(List<List<ChatMessage>> messages, string? id = null, string? name = null) : AIAgent
{
    public TestReplayAgent(List<ChatMessage> messages, string? id = null, string? name = null) : this([messages ?? []], id, name)
    { }

    public TestReplayAgent(string? id = null, string? name = null) : this([[]], id, name)
    { }

    protected override string? IdCore => id;
    public override string? Name => name;

    public static List<ChatMessage> ToChatMessages(params string[] messages)
    {
        List<ChatMessage> result = messages.Select(ToMessage).ToList();

        static ChatMessage ToMessage(string text)
        {
            if (string.IsNullOrEmpty(text))
            {
                return new ChatMessage(ChatRole.Assistant, "") { MessageId = "" };
            }

            string[] splits = text.Split(' ');
            for (int i = 0; i < splits.Length - 1; i++)
            {
                splits[i] += ' ';
            }

            List<AIContent> contents = splits.Select<string, AIContent>(text => new TextContent(text) { RawRepresentation = text }).ToList();
            return new(ChatRole.Assistant, contents)
            {
                MessageId = Guid.NewGuid().ToString("N"),
                RawRepresentation = text,
                CreatedAt = DateTime.UtcNow,
            };
        }

        return result;
    }

    protected override ValueTask<AgentSession> CreateSessionCoreAsync(CancellationToken cancellationToken = default)
        => new(new ReplayAgentSession());

    protected override ValueTask<AgentSession> DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
        => new(new ReplayAgentSession());

    protected override ValueTask<JsonElement> SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
        => default;

    public static TestReplayAgent FromStrings(params string[] messages) =>
        new(ToChatMessages(messages));

    public List<List<ChatMessage>> Messages { get; } = Validate(messages) ?? [];

    protected override Task<AgentResponse> RunCoreAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
        => this.RunStreamingAsync(messages, session, options, cancellationToken).ToAgentResponseAsync(cancellationToken);

    public int Turn { get; set; }

    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        string responseId = Guid.NewGuid().ToString("N");

        if (this.Turn < this.Messages.Count)
        {
            foreach (ChatMessage message in this.Messages[this.Turn++])
            {
                foreach (AIContent content in message.Contents)
                {
                    yield return new AgentResponseUpdate()
                    {
                        AgentId = this.Id,
                        AuthorName = this.Name,
                        MessageId = message.MessageId,
                        ResponseId = responseId,
                        Contents = [content],
                        Role = message.Role,
                    };
                }
            }
        }
    }

    private static List<List<ChatMessage>>? Validate(List<List<ChatMessage>>? candidateMessages)
    {
        string? lastMessageId = null;

        if (candidateMessages != null)
        {
            foreach (List<ChatMessage> candidateMessagesTurn in candidateMessages)
            {
                foreach (ChatMessage message in candidateMessagesTurn)
                {
                    if (lastMessageId is null || lastMessageId != message.MessageId)
                    {
                        lastMessageId = message.MessageId;
                    }
                    else if (lastMessageId == message.MessageId)
                    {
                        throw new ArgumentException("Duplicate consecutive message ids");
                    }
                }
            }
        }

        return candidateMessages;
    }

    private sealed class ReplayAgentSession() : AgentSession();
}