llm-pii-poc/tests/PiiRedaction.Core.Tests/TestSupport/CapturingChatClient.cs

using System.Runtime.CompilerServices;
using Microsoft.Extensions.AI;

namespace PiiRedaction.Core.Tests.TestSupport;

public sealed class CapturingChatClient : IChatClient
{
    public string? LastUserMessage { get; private set; }

    public Task<ChatResponse> GetResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        LastUserMessage = messages.LastOrDefault(message => message.Role == ChatRole.User)?.Text;

        var response = new ChatResponse(new ChatMessage(
            ChatRole.Assistant,
            $"Captured {LastUserMessage?.Length ?? 0} chars."));

        return Task.FromResult(response);
    }

    public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        var response = await GetResponseAsync(messages, options, cancellationToken).ConfigureAwait(false);
        yield return new ChatResponseUpdate(ChatRole.Assistant, response.Messages.Last().Text);
    }

    public object? GetService(Type serviceType, object? serviceKey = null) => null;

    public void Dispose()
    {
    }
}