chore: import upstream snapshot with attribution

dotnet/samples/Demos/VectorStoreRAG/DataLoader.cs

@@ -0,0 +1,170 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Net;
+using Microsoft.Extensions.VectorData;
+using Microsoft.SemanticKernel;
+using Microsoft.SemanticKernel.ChatCompletion;
+using UglyToad.PdfPig;
+using UglyToad.PdfPig.Content;
+using UglyToad.PdfPig.DocumentLayoutAnalysis.PageSegmenter;
+
+namespace VectorStoreRAG;
+
+/// <summary>
+/// Class that loads text from a PDF file into a vector store.
+/// </summary>
+/// <typeparam name="TKey">The type of the data model key.</typeparam>
+/// <param name="uniqueKeyGenerator">A function to generate unique keys with.</param>
+/// <param name="vectorStoreRecordCollection">The collection to load the data into.</param>
+/// <param name="chatCompletionService">The chat completion service to use for generating text from images.</param>
+internal sealed class DataLoader<TKey>(
+    UniqueKeyGenerator<TKey> uniqueKeyGenerator,
+    VectorStoreCollection<TKey, TextSnippet<TKey>> vectorStoreRecordCollection,
+    IChatCompletionService chatCompletionService) : IDataLoader where TKey : notnull
+{
+    /// <inheritdoc/>
+    public async Task LoadPdf(string pdfPath, int batchSize, int betweenBatchDelayInMs, CancellationToken cancellationToken)
+    {
+        // Create the collection if it doesn't exist.
+        await vectorStoreRecordCollection.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false);
+
+        // Load the text and images from the PDF file and split them into batches.
+        var sections = LoadTextAndImages(pdfPath, cancellationToken);
+        var batches = sections.Chunk(batchSize);
+
+        // Process each batch of content items.
+        foreach (var batch in batches)
+        {
+            // Convert any images to text.
+            var textContentTasks = batch.Select(async content =>
+            {
+                if (content.Text != null)
+                {
+                    return content;
+                }
+
+                var textFromImage = await ConvertImageToTextWithRetryAsync(
+                    chatCompletionService,
+                    content.Image!.Value,
+                    cancellationToken).ConfigureAwait(false);
+                return new RawContent { Text = textFromImage, PageNumber = content.PageNumber };
+            });
+            var textContent = await Task.WhenAll(textContentTasks).ConfigureAwait(false);
+
+            // Map each paragraph to a TextSnippet.
+            var records = textContent.Select(content => new TextSnippet<TKey>
+            {
+                Key = uniqueKeyGenerator.GenerateKey(),
+                // The vector store will automatically generate the embedding for this text.
+                // See the TextEmbedding field on the TextSnippet class.
+                Text = content.Text,
+                ReferenceDescription = $"{new FileInfo(pdfPath).Name}#page={content.PageNumber}",
+                ReferenceLink = $"{new Uri(new FileInfo(pdfPath).FullName).AbsoluteUri}#page={content.PageNumber}",
+            });
+
+            // Upsert the records into the vector store.
+            await vectorStoreRecordCollection.UpsertAsync(records, cancellationToken: cancellationToken).ConfigureAwait(false);
+
+            await Task.Delay(betweenBatchDelayInMs, cancellationToken).ConfigureAwait(false);
+        }
+    }
+
+    /// <summary>
+    /// Read the text and images from each page in the provided PDF file.
+    /// </summary>
+    /// <param name="pdfPath">The pdf file to read the text and images from.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
+    /// <returns>The text and images from the pdf file, plus the page number that each is on.</returns>
+    private static IEnumerable<RawContent> LoadTextAndImages(string pdfPath, CancellationToken cancellationToken)
+    {
+        using (PdfDocument document = PdfDocument.Open(pdfPath))
+        {
+            foreach (Page page in document.GetPages())
+            {
+                if (cancellationToken.IsCancellationRequested)
+                {
+                    break;
+                }
+
+                foreach (var image in page.GetImages())
+                {
+                    if (image.TryGetPng(out var png))
+                    {
+                        yield return new RawContent { Image = png, PageNumber = page.Number };
+                    }
+                    else
+                    {
+                        Console.WriteLine($"Unsupported image format on page {page.Number}");
+                    }
+                }
+
+                var blocks = DefaultPageSegmenter.Instance.GetBlocks(page.GetWords());
+                foreach (var block in blocks)
+                {
+                    if (cancellationToken.IsCancellationRequested)
+                    {
+                        break;
+                    }
+
+                    yield return new RawContent { Text = block.Text, PageNumber = page.Number };
+                }
+            }
+        }
+    }
+
+    /// <summary>
+    /// Add a simple retry mechanism to image to text.
+    /// </summary>
+    /// <param name="chatCompletionService">The chat completion service to use for generating text from images.</param>
+    /// <param name="imageBytes">The image to generate the text for.</param>
+    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
+    /// <returns>The generated text.</returns>
+    private static async Task<string> ConvertImageToTextWithRetryAsync(
+        IChatCompletionService chatCompletionService,
+        ReadOnlyMemory<byte> imageBytes,
+        CancellationToken cancellationToken)
+    {
+        var tries = 0;
+
+        while (true)
+        {
+            try
+            {
+                var chatHistory = new ChatHistory();
+                chatHistory.AddUserMessage([
+                    new TextContent("What’s in this image?"),
+                    new ImageContent(imageBytes, "image/png"),
+                ]);
+                var result = await chatCompletionService.GetChatMessageContentsAsync(chatHistory, cancellationToken: cancellationToken).ConfigureAwait(false);
+                return string.Join("\n", result.Select(x => x.Content));
+            }
+            catch (HttpOperationException ex) when (ex.StatusCode == HttpStatusCode.TooManyRequests)
+            {
+                tries++;
+
+                if (tries < 3)
+                {
+                    Console.WriteLine($"Failed to generate text from image. Error: {ex}");
+                    Console.WriteLine("Retrying text to image conversion...");
+                    await Task.Delay(10_000, cancellationToken).ConfigureAwait(false);
+                }
+                else
+                {
+                    throw;
+                }
+            }
+        }
+    }
+
+    /// <summary>
+    /// Private model for returning the content items from a PDF file.
+    /// </summary>
+    private sealed class RawContent
+    {
+        public string? Text { get; init; }
+
+        public ReadOnlyMemory<byte>? Image { get; init; }
+
+        public int PageNumber { get; init; }
+    }
+}