[MEDI] Add TokenCount property to IngestionChunk

src/Libraries/Microsoft.Extensions.DataIngestion.Abstractions/IngestionChunk.cs

@@ -22,18 +22,27 @@ public sealed class IngestionChunk<T>
     /// </summary>
     /// <param name="content">The content of the chunk.</param>
     /// <param name="document">The document from which this chunk was extracted.</param>
+    /// <param name="tokenCount">The number of tokens used to represent the chunk.</param>
     /// <param name="context">Additional context for the chunk.</param>
     /// <exception cref="ArgumentNullException">
     /// <paramref name="content"/> or <paramref name="document"/> is <see langword="null"/>.
     /// </exception>
     /// <exception cref="ArgumentException">
     /// <paramref name="content"/> is a string that is empty or contains only white-space characters.
     /// </exception>
-    public IngestionChunk(T content, IngestionDocument document, string? context = null)
+    /// <exception cref="ArgumentOutOfRangeException">
+    /// <paramref name="tokenCount"/> is negative.
+    /// </exception>
+    public IngestionChunk(T content, IngestionDocument document, int tokenCount, string? context = null)
     {
         if (typeof(T) == typeof(string))
         {
             Content = (T)(object)Throw.IfNullOrEmpty((string)(object)content!);
+
+            if (tokenCount == 0)
+            {
+                Throw.ArgumentOutOfRangeException(nameof(tokenCount), "Token count must be greater than zero for non-empty content.");
+            }
         }
         else
         {
@@ -42,6 +51,7 @@ public IngestionChunk(T content, IngestionDocument document, string? context = n
 
         Document = Throw.IfNull(document);
         Context = context;
+        TokenCount = Throw.IfLessThan(tokenCount, 0);
     }
 
     /// <summary>
@@ -59,6 +69,11 @@ public IngestionChunk(T content, IngestionDocument document, string? context = n
     /// </summary>
     public string? Context { get; }
 
+    /// <summary>
+    /// Gets the number of tokens used to represent the chunk.
+    /// </summary>
+    public int TokenCount { get; }
+
     /// <summary>
     /// Gets a value indicating whether this chunk has metadata.
     /// </summary>

src/Libraries/Microsoft.Extensions.DataIngestion/Chunkers/DocumentTokenChunker.cs

@@ -63,7 +63,7 @@ public override async IAsyncEnumerable<IngestionChunk<string>> ProcessAsync(Inge
                         text: contentToProcess.Span,
                         maxTokenCount: _maxTokensPerChunk - stringBuilderTokenCount,
                         out string? _,
-                        out int _,
+                        out int addedTokenCount,
                         considerNormalization: false);
 
                     unsafe
@@ -73,6 +73,7 @@ public override async IAsyncEnumerable<IngestionChunk<string>> ProcessAsync(Inge
                             _ = stringBuilder.Append(ptr, index);
                         }
                     }
+                    stringBuilderTokenCount += addedTokenCount;
                     yield return FinalizeChunk();
 
                     contentToProcess = contentToProcess.Slice(index);
@@ -94,6 +95,7 @@ IngestionChunk<string> FinalizeChunk()
                 IngestionChunk<string> chunk = new IngestionChunk<string>(
                     content: stringBuilder.ToString(),
                     document: document,
+                    tokenCount: stringBuilderTokenCount,
                     context: string.Empty);
                 _ = stringBuilder.Clear();
                 stringBuilderTokenCount = 0;

src/Libraries/Microsoft.Extensions.DataIngestion/Chunkers/ElementsChunker.cs

@@ -196,7 +196,9 @@ internal IEnumerable<IngestionChunk<string>> Process(IngestionDocument document,
 
         if (totalTokenCount > contextTokenCount)
         {
-            chunks.Add(new(_currentChunk.ToString(), document, context));
+            string chunkContent = _currentChunk.ToString();
+            int chunkTokenCount = CountTokens(chunkContent.AsSpan());
+            chunks.Add(new(chunkContent, document, chunkTokenCount, context));
         }
 
         _currentChunk.Clear();
@@ -205,7 +207,9 @@ internal IEnumerable<IngestionChunk<string>> Process(IngestionDocument document,
 
         void Commit()
         {
-            chunks.Add(new(_currentChunk.ToString(), document, context));
+            string chunkContent = _currentChunk.ToString();
+            int chunkTokenCount = CountTokens(chunkContent.AsSpan());
+            chunks.Add(new(chunkContent, document, chunkTokenCount, context));
 
             // We keep the context in the current chunk as it's the same for all elements.
             _currentChunk.Remove(

test/Libraries/Microsoft.Extensions.DataIngestion.Tests/Chunkers/HeaderChunkerTests.cs

@@ -270,6 +270,35 @@ private static IngestionDocument CreateDocumentWithLargeTable()
         }
     }
 
+    [Fact]
+    public async Task VerifyTokenCountIsSet()
+    {
+        IngestionDocument doc = new("tokenCountTest");
+        doc.Sections.Add(new()
+        {
+            Elements =
+            {
+                new IngestionDocumentHeader("Header 1") { Level = 1 },
+                new IngestionDocumentParagraph("This is a test paragraph with some content."),
+                new IngestionDocumentParagraph("This is another paragraph with more content.")
+            }
+        });
+
+        var tokenizer = TiktokenTokenizer.CreateForModel("gpt-4");
+        HeaderChunker chunker = new(new(tokenizer));
+        IReadOnlyList<IngestionChunk<string>> chunks = await chunker.ProcessAsync(doc).ToListAsync();
+
+        foreach (var chunk in chunks)
+        {
+            // Verify that TokenCount property is set
+            Assert.True(chunk.TokenCount > 0);
+
+            // Verify that TokenCount matches actual token count of content
+            int actualTokenCount = tokenizer.CountTokens(chunk.Content, considerNormalization: false);
+            Assert.Equal(actualTokenCount, chunk.TokenCount);
+        }
+    }
+
     // We need plenty of more tests here, especially for edge cases:
     // - sentence splitting
     // - markdown splitting (e.g. lists, code blocks etc.)

test/Libraries/Microsoft.Extensions.DataIngestion.Tests/Chunkers/NoOverlapTokenChunkerTests.cs

@@ -60,5 +60,37 @@ public async Task ManyChunks()
 
             Assert.Equal(text, string.Join("", chunks.Select(c => c.Content)));
         }
+
+        [Fact]
+        public async Task VerifyTokenCount()
+        {
+            string text = string.Join(" ", Enumerable.Repeat("word", 600)); // each word is 1 token
+            IngestionDocument doc = new IngestionDocument("verifyTokenCountDoc");
+            doc.Sections.Add(new IngestionDocumentSection
+            {
+                Elements =
+                {
+                    new IngestionDocumentParagraph(text)
+                }
+            });
+
+            var tokenizer = TiktokenTokenizer.CreateForModel("gpt-4o");
+            IngestionChunker<string> chunker = new DocumentTokenChunker(new(tokenizer) { MaxTokensPerChunk = 512, OverlapTokens = 0 });
+            IReadOnlyList<IngestionChunk<string>> chunks = await chunker.ProcessAsync(doc).ToListAsync();
+
+            Assert.Equal(2, chunks.Count);
+            foreach (var chunk in chunks)
+            {
+                // Verify that TokenCount property is set
+                Assert.True(chunk.TokenCount > 0);
+
+                // Verify that TokenCount matches actual token count of content
+                int actualTokenCount = tokenizer.CountTokens(chunk.Content, considerNormalization: false);
+                Assert.Equal(actualTokenCount, chunk.TokenCount);
+
+                // Verify that TokenCount does not exceed max tokens per chunk
+                Assert.True(chunk.TokenCount <= 512);
+            }
+        }
     }
 }

test/Libraries/Microsoft.Extensions.DataIngestion.Tests/IngestionChunkTests.cs

@@ -0,0 +1,53 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using Microsoft.ML.Tokenizers;
+using Xunit;
+
+namespace Microsoft.Extensions.DataIngestion.Tests;
+
+public class IngestionChunkTests
+{
+    [Fact]
+    public void Constructor_SetsTokenCountProperty()
+    {
+        IngestionDocument document = new("test");
+        IngestionChunk<string> chunk = new("test content", document, 42);
+
+        Assert.Equal(42, chunk.TokenCount);
+    }
+
+    [Fact]
+    public void Constructor_ThrowsWhenTokenCountIsNegative()
+    {
+        IngestionDocument document = new("test");
+
+        ArgumentOutOfRangeException exception = Assert.Throws<ArgumentOutOfRangeException>(
+            () => new IngestionChunk<string>("test content", document, -1));
+
+        Assert.Equal("tokenCount", exception.ParamName);
+    }
+
+    [Fact]
+    public void Constructor_ThrowsWhenTokenCountIsZeroForNonEmptyContent()
+    {
+        IngestionDocument document = new("test");
+
+        ArgumentOutOfRangeException exception = Assert.Throws<ArgumentOutOfRangeException>(
+            () => new IngestionChunk<string>("test content", document, 0));
+
+        Assert.Equal("tokenCount", exception.ParamName);
+    }
+}
+
+public static class TestHelpers
+{
+    private static readonly Tokenizer s_tokenizer = TiktokenTokenizer.CreateForModel("gpt-4o");
+
+    public static IngestionChunk<string> CreateChunk(string content, IngestionDocument document)
+    {
+        int tokenCount = s_tokenizer.CountTokens(content, considerNormalization: false);
+        return new IngestionChunk<string>(content, document, tokenCount);
+    }
+}

test/Libraries/Microsoft.Extensions.DataIngestion.Tests/IngestionPipelineTests.cs

@@ -180,7 +180,8 @@ public override IAsyncEnumerable<IngestionChunk<DataContent>> ProcessAsync(Inges
                     .OfType<IngestionDocumentImage>()
                     .Select(image => new IngestionChunk<DataContent>(
                         content: new(image.Content.GetValueOrDefault(), image.MediaType!),
-                        document: document))
+                        document: document,
+                        tokenCount: 0))
                     .ToAsyncEnumerable();
     }
 

test/Libraries/Microsoft.Extensions.DataIngestion.Tests/Processors/ClassificationEnricherTests.cs

@@ -7,6 +7,7 @@
 using System.Text.Json;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DataIngestion.Tests;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Testing;
 using Xunit;
@@ -119,11 +120,11 @@ public async Task FailureDoesNotStopTheProcessing()
 
     private static List<IngestionChunk<string>> CreateChunks() =>
     [
-        new(".NET developers need to integrate and interact with a growing variety of artificial intelligence (AI) services in their apps. " +
+        TestHelpers.CreateChunk(".NET developers need to integrate and interact with a growing variety of artificial intelligence (AI) services in their apps. " +
             "The Microsoft.Extensions.AI libraries provide a unified approach for representing generative AI components, and enable seamless" +
             " integration and interoperability with various AI services.", _document),
-        new ("Rabbits are small mammals in the family Leporidae of the order Lagomorpha (along with the hare and the pika)." +
+        TestHelpers.CreateChunk("Rabbits are small mammals in the family Leporidae of the order Lagomorpha (along with the hare and the pika)." +
             "They are herbivorous animals and are known for their long ears, large hind legs, and short fluffy tails.", _document),
-        new("This text does not belong to any category.", _document),
+        TestHelpers.CreateChunk("This text does not belong to any category.", _document),
     ];
 }

test/Libraries/Microsoft.Extensions.DataIngestion.Tests/Processors/KeywordEnricherTests.cs

@@ -7,6 +7,7 @@
 using System.Text.Json;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DataIngestion.Tests;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Testing;
 using Xunit;
@@ -119,7 +120,7 @@ public async Task FailureDoesNotStopTheProcessing()
 
     private static List<IngestionChunk<string>> CreateChunks() =>
     [
-        new("The Microsoft.Extensions.AI libraries provide a unified approach for representing generative AI components", _document),
-        new("Rabbits are great pets. They are friendly and make excellent companions.", _document)
+        TestHelpers.CreateChunk("The Microsoft.Extensions.AI libraries provide a unified approach for representing generative AI components", _document),
+        TestHelpers.CreateChunk("Rabbits are great pets. They are friendly and make excellent companions.", _document)
     ];
 }

test/Libraries/Microsoft.Extensions.DataIngestion.Tests/Processors/SentimentEnricherTests.cs

@@ -7,6 +7,7 @@
 using System.Text.Json;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DataIngestion.Tests;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Testing;
 using Xunit;
@@ -106,9 +107,9 @@ public async Task FailureDoesNotStopTheProcessing()
 
     private static List<IngestionChunk<string>> CreateChunks() =>
     [
-        new("I love programming! It's so much fun and rewarding.", _document),
-        new("I hate bugs. They are so frustrating and time-consuming.", _document),
-        new("The weather is okay, not too bad but not great either.", _document),
-        new("I hate you. I am sorry, I actually don't. I am not sure myself what my feelings are.", _document)
+        TestHelpers.CreateChunk("I love programming! It's so much fun and rewarding.", _document),
+        TestHelpers.CreateChunk("I hate bugs. They are so frustrating and time-consuming.", _document),
+        TestHelpers.CreateChunk("The weather is okay, not too bad but not great either.", _document),
+        TestHelpers.CreateChunk("I hate you. I am sorry, I actually don't. I am not sure myself what my feelings are.", _document)
     ];
 }

test/Libraries/Microsoft.Extensions.DataIngestion.Tests/Processors/SummaryEnricherTests.cs

@@ -7,6 +7,7 @@
 using System.Text.Json;
 using System.Threading.Tasks;
 using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DataIngestion.Tests;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Testing;
 using Xunit;
@@ -103,7 +104,7 @@ public async Task FailureDoesNotStopTheProcessing()
 
     private static List<IngestionChunk<string>> CreateChunks() =>
     [
-        new("I love programming! It's so much fun and rewarding.", _document),
-        new("I hate bugs. They are so frustrating and time-consuming.", _document)
+        TestHelpers.CreateChunk("I love programming! It's so much fun and rewarding.", _document),
+        TestHelpers.CreateChunk("I hate bugs. They are so frustrating and time-consuming.", _document)
     ];
 }

test/Libraries/Microsoft.Extensions.DataIngestion.Tests/Writers/VectorStoreWriterTests.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading.Tasks;
+using Microsoft.Extensions.DataIngestion.Tests;
 using Microsoft.Extensions.VectorData;
 using Xunit;
 
@@ -24,19 +25,13 @@ public async Task CanGenerateDynamicSchema()
             dimensionCount: TestEmbeddingGenerator<string>.DimensionCount);
 
         IngestionDocument document = new(documentId);
-        List<IngestionChunk<string>> chunks =
-        [
-            new("some content", document)
-            {
-                Metadata =
-                {
-                    { "key1", "value1" },
-                    { "key2", 123 },
-                    { "key3", true },
-                    { "key4", 123.45 },
-                }
-            }
-        ];
+        IngestionChunk<string> chunk = TestHelpers.CreateChunk("some content", document);
+        chunk.Metadata["key1"] = "value1";
+        chunk.Metadata["key2"] = 123;
+        chunk.Metadata["key3"] = true;
+        chunk.Metadata["key4"] = 123.45;
+
+        List<IngestionChunk<string>> chunks = [chunk];
 
         Assert.False(testEmbeddingGenerator.WasCalled);
         await writer.WriteAsync(chunks.ToAsyncEnumerable());
@@ -73,17 +68,12 @@ public async Task DoesSupportIncrementalIngestion()
             });
 
         IngestionDocument document = new(documentId);
-        List<IngestionChunk<string>> chunks =
-        [
-            new("first chunk", document)
-            {
-                Metadata =
-                {
-                    { "key1", "value1" }
-                }
-            },
-            new("second chunk", document)
-        ];
+        IngestionChunk<string> chunk1 = TestHelpers.CreateChunk("first chunk", document);
+        chunk1.Metadata["key1"] = "value1";
+
+        IngestionChunk<string> chunk2 = TestHelpers.CreateChunk("second chunk", document);
+
+        List<IngestionChunk<string>> chunks = [chunk1, chunk2];
 
         await writer.WriteAsync(chunks.ToAsyncEnumerable());
 
@@ -93,16 +83,10 @@ public async Task DoesSupportIncrementalIngestion()
         Assert.Equal(chunks.Count, recordCount);
 
         // Now we will do an incremental ingestion that updates the chunk(s).
-        List<IngestionChunk<string>> updatedChunks =
-        [
-            new("different content", document)
-            {
-                Metadata =
-                {
-                    { "key1", "value2" },
-                }
-            }
-        ];
+        IngestionChunk<string> updatedChunk = TestHelpers.CreateChunk("different content", document);
+        updatedChunk.Metadata["key1"] = "value2";
+
+        List<IngestionChunk<string>> updatedChunks = [updatedChunk];
 
         await writer.WriteAsync(updatedChunks.ToAsyncEnumerable());