meeting-assistant/MeetingAssistant.Tests/LiteLlmScreenshotOcrClientTests.cs

using System.Net;
using System.Net.Http.Headers;
using System.Text;
using System.Text.Json;
using System.Drawing;
using System.Drawing.Imaging;
using MeetingAssistant.Screenshots;
using Microsoft.Extensions.Logging.Abstractions;

#pragma warning disable CA1416

namespace MeetingAssistant.Tests;

public sealed class LiteLlmScreenshotOcrClientTests
{
    [Fact]
    public async Task ExtractUsesAgentStreamingTransportByDefault()
    {
        var screenshotPath = await CreateScreenshotAsync([1, 2, 3]);
        var handler = new RecordingHandler(
            CreateStreamedTextResponse("Streamed OCR text"),
            "text/event-stream");
        var client = new LiteLlmScreenshotOcrClient(
            () => handler,
            NullLogger<LiteLlmScreenshotOcrClient>.Instance);
        var options = new MeetingAssistantOptions
        {
            Agent =
            {
                Endpoint = "https://summary.local",
                Model = "vision-model",
                Key = "agent-key",
                UseStreaming = true
            }
        };

        var result = await client.ExtractAsync(
            screenshotPath,
            "Extract screenshot.",
            options,
            CancellationToken.None);

        Assert.Equal("Streamed OCR text", result.Text);
        using var payload = JsonDocument.Parse(handler.RequestBody!);
        Assert.True(payload.RootElement.GetProperty("stream").GetBoolean());
        var message = Assert.Single(payload.RootElement.GetProperty("input").EnumerateArray());
        Assert.Equal("message", message.GetProperty("type").GetString());
        var content = message.GetProperty("content").EnumerateArray().ToArray();
        Assert.Equal("input_text", content[0].GetProperty("type").GetString());
        Assert.Equal("Extract screenshot.", content[0].GetProperty("text").GetString());
        Assert.Equal("input_image", content[1].GetProperty("type").GetString());
        Assert.Equal("data:image/png;base64,AQID", content[1].GetProperty("image_url").GetString());
    }

    [Fact]
    public async Task ExtractUsesAgentEndpointAndModelWhenOcrEndpointAndModelAreBlank()
    {
        var screenshotPath = await CreateScreenshotAsync([1, 2, 3]);
        var handler = new RecordingHandler(CreateNonStreamingTextResponse("Visible slide text"));
        var client = new LiteLlmScreenshotOcrClient(
            () => handler,
            NullLogger<LiteLlmScreenshotOcrClient>.Instance);
        var options = new MeetingAssistantOptions
        {
            Agent =
            {
                Endpoint = "https://summary.local",
                Model = "summary-model",
                Key = "agent-key",
                UseStreaming = false
            },
            Screenshots =
            {
                Ocr =
                {
                    Endpoint = "",
                    Model = "",
                    Key = "ocr-key"
                }
            }
        };

        var result = await client.ExtractAsync(
            screenshotPath,
            "Extract screenshot.",
            options,
            CancellationToken.None);

        Assert.Equal("Visible slide text", result.Text);
        Assert.Null(result.Crop);
        Assert.Equal(new Uri("https://summary.local/v1/responses"), handler.RequestUri);
        Assert.Equal("Bearer", handler.Authorization?.Scheme);
        Assert.Equal("ocr-key", handler.Authorization?.Parameter);
        using var payload = JsonDocument.Parse(handler.RequestBody!);
        Assert.False(payload.RootElement.GetProperty("stream").GetBoolean());
        Assert.Equal("summary-model", payload.RootElement.GetProperty("model").GetString());
        var content = payload.RootElement
            .GetProperty("input")[0]
            .GetProperty("content");
        Assert.Equal("Extract screenshot.", content[0].GetProperty("text").GetString());
        Assert.Equal("data:image/png;base64,AQID", content[1].GetProperty("image_url").GetString());
    }

    [Fact]
    public async Task ExtractUsesScreenshotOcrEndpointAndModelWhenConfigured()
    {
        var screenshotPath = await CreateScreenshotAsync([4, 5, 6]);
        var handler = new RecordingHandler(CreateNonStreamingTextResponse("OCR result"));
        var client = new LiteLlmScreenshotOcrClient(
            () => handler,
            NullLogger<LiteLlmScreenshotOcrClient>.Instance);
        var options = new MeetingAssistantOptions
        {
            Agent =
            {
                Endpoint = "https://summary.local",
                Model = "summary-model",
                Key = "agent-key",
                UseStreaming = false
            },
            Screenshots =
            {
                Ocr =
                {
                    Endpoint = "https://vision.local/v1",
                    Model = "vision-model",
                    Key = "ocr-key"
                }
            }
        };

        var result = await client.ExtractAsync(
            screenshotPath,
            "Extract screenshot.",
            options,
            CancellationToken.None);

        Assert.Equal("OCR result", result.Text);
        Assert.Equal(new Uri("https://vision.local/v1/responses"), handler.RequestUri);
        using var payload = JsonDocument.Parse(handler.RequestBody!);
        Assert.Equal("vision-model", payload.RootElement.GetProperty("model").GetString());
    }

    [Fact]
    public async Task ExtractParsesCropMetadataAndOmitsMetadataFromReturnedText()
    {
        var screenshotPath = await CreateScreenshotAsync(CreatePngBytes(8, 6));
        var handler = new RecordingHandler(CreateNonStreamingTextResponse(
            """
            Slide text

            ```json
            { "crop": { "x": 1, "y": 2, "width": 3, "height": 4 } }
            ```
            """));
        var client = new LiteLlmScreenshotOcrClient(
            () => handler,
            NullLogger<LiteLlmScreenshotOcrClient>.Instance);
        var options = new MeetingAssistantOptions
        {
            Agent =
            {
                Key = "agent-key",
                UseStreaming = false
            }
        };

        var result = await client.ExtractAsync(
            screenshotPath,
            "Extract screenshot.",
            options,
            CancellationToken.None);

        Assert.Equal("Slide text", result.Text);
        Assert.Equal(new ScreenshotCropCoordinates(1, 2, 3, 4), result.Crop);
        using var payload = JsonDocument.Parse(handler.RequestBody!);
        Assert.Contains(
            "Original screenshot dimensions: 8x6 pixels.",
            payload.RootElement.GetProperty("input")[0].GetProperty("content")[0].GetProperty("text").GetString(),
            StringComparison.Ordinal);
    }

    [Fact]
    public async Task ExtractParsesAttendeeMetadataAndOmitsMetadataFromReturnedText()
    {
        var screenshotPath = await CreateScreenshotAsync([1, 2, 3]);
        var handler = new RecordingHandler(CreateNonStreamingTextResponse(
            """
            Visible participant tiles: Ada and Grace.

            ```json
            { "crop": null, "attendees": ["Ada Lovelace", "Grace Hopper"] }
            ```
            """));
        var client = new LiteLlmScreenshotOcrClient(
            () => handler,
            NullLogger<LiteLlmScreenshotOcrClient>.Instance);
        var options = new MeetingAssistantOptions
        {
            Agent =
            {
                Key = "agent-key",
                UseStreaming = false
            }
        };

        var result = await client.ExtractAsync(
            screenshotPath,
            "Extract screenshot.",
            options,
            CancellationToken.None);

        Assert.Equal("Visible participant tiles: Ada and Grace.", result.Text);
        Assert.Equal(["Ada Lovelace", "Grace Hopper"], result.Attendees);
    }

    [Fact]
    public async Task ExtractIgnoresMalformedAttendeesMetadataAndStillParsesCrop()
    {
        var screenshotPath = await CreateScreenshotAsync(CreatePngBytes(8, 6));
        var handler = new RecordingHandler(CreateNonStreamingTextResponse(
            """
            Slide text

            ```json
            { "crop": { "x": 1, "y": 2, "width": 3, "height": 4 }, "attendees": "Ada" }
            ```
            """));
        var client = new LiteLlmScreenshotOcrClient(
            () => handler,
            NullLogger<LiteLlmScreenshotOcrClient>.Instance);
        var options = new MeetingAssistantOptions
        {
            Agent =
            {
                Key = "agent-key",
                UseStreaming = false
            }
        };

        var result = await client.ExtractAsync(
            screenshotPath,
            "Extract screenshot.",
            options,
            CancellationToken.None);

        Assert.Equal("Slide text", result.Text);
        Assert.Equal(new ScreenshotCropCoordinates(1, 2, 3, 4), result.Crop);
        Assert.Empty(result.Attendees);
    }

    private sealed class RecordingHandler : HttpMessageHandler
    {
        private readonly string responseBody;
        private readonly string mediaType;

        public RecordingHandler(string responseBody, string mediaType = "application/json")
        {
            this.responseBody = responseBody;
            this.mediaType = mediaType;
        }

        public Uri? RequestUri { get; private set; }

        public AuthenticationHeaderValue? Authorization { get; private set; }

        public string? RequestBody { get; private set; }

        protected override async Task<HttpResponseMessage> SendAsync(
            HttpRequestMessage request,
            CancellationToken cancellationToken)
        {
            RequestUri = request.RequestUri;
            Authorization = request.Headers.Authorization;
            RequestBody = request.Content is null
                ? null
                : await request.Content.ReadAsStringAsync(cancellationToken);
            return new HttpResponseMessage(HttpStatusCode.OK)
            {
                Content = new StringContent(responseBody, Encoding.UTF8, mediaType)
            };
        }
    }

    private static byte[] CreatePngBytes(int width, int height)
    {
        using var bitmap = new Bitmap(width, height);
        using (var graphics = Graphics.FromImage(bitmap))
        {
            graphics.Clear(Color.White);
        }

        using var stream = new MemoryStream();
        bitmap.Save(stream, ImageFormat.Png);
        return stream.ToArray();
    }

    private static string CreateNonStreamingTextResponse(string text)
    {
        return JsonSerializer.Serialize(new
        {
            id = "resp_ocr",
            created_at = 1779147100,
            model = "vision-model",
            @object = "response",
            output = new[]
            {
                new
                {
                    id = "msg_ocr",
                    type = "message",
                    status = "completed",
                    content = new[]
                    {
                        new
                        {
                            type = "output_text",
                            annotations = Array.Empty<object>(),
                            text
                        }
                    },
                    role = "assistant"
                }
            },
            parallel_tool_calls = true,
            status = "completed",
            store = false
        });
    }

    private static string CreateStreamedTextResponse(string text)
    {
        var delta = JsonSerializer.Serialize(new
        {
            type = "response.output_text.delta",
            item_id = "msg_ocr",
            output_index = 0,
            content_index = 0,
            delta = text
        });
        return
            $"data: {delta}{Environment.NewLine}{Environment.NewLine}" +
            """data: {"type":"response.completed","response":{"id":"resp_ocr","created_at":1779147100,"model":"vision-model","object":"response","output":[],"parallel_tool_calls":true,"status":"completed","store":false}}""" +
            $"{Environment.NewLine}{Environment.NewLine}data: [DONE]{Environment.NewLine}{Environment.NewLine}";
    }

    private static async Task<string> CreateScreenshotAsync(byte[] bytes)
    {
        var screenshotPath = Path.Combine(Path.GetTempPath(), "meeting-assistant-tests", Guid.NewGuid().ToString("N") + ".png");
        Directory.CreateDirectory(Path.GetDirectoryName(screenshotPath)!);
        await File.WriteAllBytesAsync(screenshotPath, bytes);
        return screenshotPath;
    }
}

#pragma warning restore CA1416