FunctionInvokingChatClient reorder messages when using FunctionApprovalResponseContent #7156

using System.Text.Json;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.Logging;

#pragma warning disable MEAI001

namespace AiResearch;

public class ChatService(IChatClient rawClient, ILogger<ChatService> logger, ILoggerFactory loggerFactory)
{
    public async Task RunAsync()
    {
        var sink = new MessageSink();
        var client = CreateChatClient(sink);
        foreach (var (testCase, method) in GetMethods())
        {
            await TestMethod(client, sink, testCase, method);
        }
    }


    private IChatClient CreateChatClient(MessageSink messageSink)
    {
        // build in reverse order. logging must wrap original logging to see real messages and not that in function invoking
        return rawClient.AsBuilder()
            .UseFunctionInvocation(loggerFactory)
            .Use(c => new SinkingChatClient(c, loggerFactory.CreateLogger<SinkingChatClient>(), messageSink))
            .Build();
    }


    private async Task TestMethod(IChatClient chatClient, MessageSink messageSink, string testCase, Action<IList<ChatMessage>> method)
    {
        using var scope = logger.BeginScope("{TestCase}", testCase);
        logger.LogInformation("Begin {TestCase}", testCase);

        var chatOptions = new ChatOptions
        {
            Tools =
            [
                new ApprovalRequiredAIFunction(AIFunctionFactory.Create(ConvertString)),
            ],
        };

        try
        {
            string prompt = """
                                You are test engineer that need to test AI Functions.

                                As first test case, call ConvertString function with random input.

                                Respect user response if tool was rejected.
                                """;
            List<ChatMessage> messages = [new ChatMessage(ChatRole.System, prompt)];

            ChatResponse response;
            bool firstCall = true;
            while ((response = await chatClient.GetResponseAsync(messages, chatOptions)).FinishReason != ChatFinishReason.Stop)
            {
                messages.AddRange(response.Messages);
                if (response.FinishReason == ChatFinishReason.ToolCalls)
                {
                    if (firstCall)
                    {
                        method(messages);
                        firstCall = false;
                    }
                    else
                    {
                        messages.Add(new ChatMessage(ChatRole.User,
                            [response.Messages.Last().Contents.OfType<FunctionApprovalRequestContent>().First().CreateResponse(true)]));

                        messages.Add(new ChatMessage(ChatRole.User, "this is result message"));
                    }
                }
            }

            logger.LogInformation("Prompt processed {Log}", JsonSerializer.Serialize(messageSink.ChatMessages, AIJsonUtilities.DefaultOptions));
        }
        catch (Exception ex)
        {
            logger.LogError(ex, "Error occurred while getting response.");
        }
    }

    private string ConvertString(string value) => $"Converted: {value}";

    private static IList<(string, Action<IList<ChatMessage>>)> GetMethods()
    {
        const string userMessage = "Please use another input.";

        return
        [
            (nameof(Manual_AddUserMessage_AsExtraTextContent), Manual_AddUserMessage_AsExtraTextContent),
            (nameof(Manual_AddUserMessage_AsNextChatMessage), Manual_AddUserMessage_AsNextChatMessage),
            (nameof(Manual_AddUserMessage_AsToolResult), Manual_AddUserMessage_AsToolResult),
            (nameof(Default_AddUserMessageAsExtraTextContent_ResponseAsNextMessage), Default_AddUserMessageAsExtraTextContent_ResponseAsNextMessage),
            (nameof(Default_AddUserMessageAsNextMessage_ResponseAsNextMessage), Default_AddUserMessageAsNextMessage_ResponseAsNextMessage),
            (nameof(Default_AddUserMessageAsNextMessage_ResponseAsSameMessage), Default_AddUserMessageAsNextMessage_ResponseAsSameMessage),
            (nameof(Default_AddUserMessage_AsExtraTextContent_ResponseAsSameMessage), Default_AddUserMessage_AsExtraTextContent_ResponseAsSameMessage),
        ];


        static void Manual_AddUserMessage_AsToolResult(IList<ChatMessage> allMessages)
        {
            for (var i = 0; i < allMessages.Count; i++)
            {
                var chatMessage = allMessages[i];
                if (chatMessage.Contents.OfType<FunctionApprovalRequestContent>().FirstOrDefault() is { } approval)
                {
                    var index = chatMessage.Contents.IndexOf(approval);
                    chatMessage.Contents[index] = approval.FunctionCall;

                    allMessages.Insert(i + 1, new ChatMessage(ChatRole.Tool, [new FunctionResultContent(approval.FunctionCall.CallId, userMessage)] ));
                }
            }
        }

        static void Manual_AddUserMessage_AsExtraTextContent(IList<ChatMessage> allMessages)
        {
            for (var i = 0; i < allMessages.Count; i++)
            {
                var chatMessage = allMessages[i];
                if (chatMessage.Contents.OfType<FunctionApprovalRequestContent>().FirstOrDefault() is { } approval)
                {
                    var index = chatMessage.Contents.IndexOf(approval);
                    chatMessage.Contents[index] = approval.FunctionCall;

                    allMessages.Insert(i + 1,
                        new ChatMessage(ChatRole.Tool,
                            [
                                new FunctionResultContent(approval.FunctionCall.CallId, "Error: Tool call invocation was rejected by user."),
                                new TextContent(userMessage)
                            ]
                        )
                    );
                }
            }
        }

        static void Manual_AddUserMessage_AsNextChatMessage(IList<ChatMessage> allMessages)
        {
            for (var i = 0; i < allMessages.Count; i++)
            {
                var chatMessage = allMessages[i];
                if (chatMessage.Contents.OfType<FunctionApprovalRequestContent>().FirstOrDefault() is { } approval)
                {
                    var index = chatMessage.Contents.IndexOf(approval);
                    chatMessage.Contents[index] = approval.FunctionCall;

                    allMessages.Insert(i + 1,
                        new ChatMessage(ChatRole.Tool,
                            [
                                new FunctionResultContent(approval.FunctionCall.CallId, "Error: Tool call invocation was rejected by user."),
                            ]
                        )
                    );
                    allMessages.Insert(i + 2, new ChatMessage(ChatRole.User, userMessage));
                }
            }
        }

        static void Default_AddUserMessage_AsExtraTextContent_ResponseAsSameMessage(IList<ChatMessage> allMessages)
        {
            foreach (var chatMessage in allMessages)
            {
                if (chatMessage.Contents.OfType<FunctionApprovalRequestContent>().FirstOrDefault() is { } approval)
                {
                    var response = approval.CreateResponse(false);
                    var index = chatMessage.Contents.IndexOf(approval);
                    chatMessage.Contents[index] = response;
                    chatMessage.Contents.Insert(index + 1, new TextContent(userMessage));
                }
            }
        }

        static void Default_AddUserMessageAsExtraTextContent_ResponseAsNextMessage(IList<ChatMessage> allMessages)
        {
            for (var i = 0; i < allMessages.Count; i++)
            {
                var chatMessage = allMessages[i];
                if (chatMessage.Contents.OfType<FunctionApprovalRequestContent>().FirstOrDefault() is { } approval)
                {
                    allMessages.Insert(i + 1, new ChatMessage(ChatRole.User,
                    [
                        approval.CreateResponse(false),
                        new TextContent(userMessage),
                    ]));
                }
            }
        }

        static void Default_AddUserMessageAsNextMessage_ResponseAsSameMessage(IList<ChatMessage> allMessages)
        {
            for (var i = 0; i < allMessages.Count; i++)
            {
                var chatMessage = allMessages[i];
                if (chatMessage.Contents.OfType<FunctionApprovalRequestContent>().FirstOrDefault() is { } approval)
                {
                    chatMessage.Contents.Add(approval.CreateResponse(false));
                    allMessages.Insert(i + 1, new ChatMessage(ChatRole.User,
                    [
                        new TextContent(userMessage),
                    ]));
                }
            }
        }

        static void Default_AddUserMessageAsNextMessage_ResponseAsNextMessage(IList<ChatMessage> allMessages)
        {
            for (var i = 0; i < allMessages.Count; i++)
            {
                var chatMessage = allMessages[i];
                if (chatMessage.Contents.OfType<FunctionApprovalRequestContent>().FirstOrDefault() is { } approval)
                {
                    allMessages.Insert(i + 1, new ChatMessage(ChatRole.User,
                    [
                        approval.CreateResponse(false),
                    ]));
                    allMessages.Insert(i + 2, new ChatMessage(ChatRole.User,
                    [
                        new TextContent(userMessage),
                    ]));
                }
            }
        }
    }


    private class SinkingChatClient(IChatClient innerClient, ILogger<SinkingChatClient> logger, MessageSink messageSink) : DelegatingChatClient(innerClient)
    {
        public override async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null,
            CancellationToken cancellationToken = default)
        {
            var log = messages.ToList();
            var result = await base.GetResponseAsync(log.ToList(), options, cancellationToken);
            log.AddRange(result.Messages);

            messageSink.ChatMessages = log;

            return result;
        }
    }

    private class MessageSink
    {
        public IList<ChatMessage> ChatMessages { get; set; }
    }

}