[MEDI] Add TokenCount property to IngestionChunk

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

@@ -22,14 +22,18 @@ 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))
         {
@@ -42,6 +46,7 @@ public IngestionChunk(T content, IngestionDocument document, string? context = n
 
         Document = Throw.IfNull(document);
         Context = context;
+        TokenCount = Throw.IfLessThanOrEqual(tokenCount, 0);
     }
 
     /// <summary>
@@ -59,6 +64,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

@@ -15,6 +15,7 @@ internal sealed class ElementsChunker
 {
     private readonly Tokenizer _tokenizer;
     private readonly int _maxTokensPerChunk;
+    private readonly int _newLineTokenCount;
     private readonly StringBuilder _currentChunk;
 
     internal ElementsChunker(IngestionChunkerOptions options)
@@ -23,11 +24,47 @@ internal ElementsChunker(IngestionChunkerOptions options)
 
         _tokenizer = options.Tokenizer;
         _maxTokensPerChunk = options.MaxTokensPerChunk;
+        _newLineTokenCount = options.Tokenizer.CountTokens(Environment.NewLine, considerNormalization: false);
 
         // Token count != character count, but StringBuilder will grow as needed.
         _currentChunk = new(capacity: _maxTokensPerChunk);
     }
 
+    private static void AddMarkdownTableRow(IngestionDocumentTable table, int rowIndex, ref ValueStringBuilder vsb)
+    {
+        for (int columnIndex = 0; columnIndex < table.Cells.GetLength(1); columnIndex++)
+        {
+            vsb.Append('|');
+            vsb.Append(' ');
+            string? cellContent = table.Cells[rowIndex, columnIndex] switch
+            {
+                null => null,
+                IngestionDocumentImage img => img.AlternativeText ?? img.Text,
+                IngestionDocumentElement other => other.GetMarkdown()
+            };
+            vsb.Append(cellContent);
+            vsb.Append(' ');
+        }
+
+        vsb.Append('|');
+        vsb.Append(Environment.NewLine);
+    }
+
+    private static void AddMarkdownTableSeparatorRow(int columnCount, ref ValueStringBuilder vsb)
+    {
+        const int DashCount = 3; // The dash count does not need to match the header length.
+        for (int columnIndex = 0; columnIndex < columnCount; columnIndex++)
+        {
+            vsb.Append('|');
+            vsb.Append(' ');
+            vsb.Append('-', DashCount);
+            vsb.Append(' ');
+        }
+
+        vsb.Append('|');
+        vsb.Append(Environment.NewLine);
+    }
+
     // Goals:
     // 1. Create chunks that do not exceed _maxTokensPerChunk when tokenized.
     // 2. Maintain context in each chunk.
@@ -70,8 +107,8 @@ internal IEnumerable<IngestionChunk<string>> Process(IngestionDocument document,
             int elementTokenCount = CountTokens(semanticContent.AsSpan());
             if (elementTokenCount + totalTokenCount <= _maxTokensPerChunk)
             {
+                AppendNewLineAndSpan(_currentChunk, semanticContent.AsSpan(), ref totalTokenCount);
                 totalTokenCount += elementTokenCount;
-                AppendNewLineAndSpan(_currentChunk, semanticContent.AsSpan());
             }
             else if (element is IngestionDocumentTable table)
             {
@@ -113,7 +150,7 @@ internal IEnumerable<IngestionChunk<string>> Process(IngestionDocument document,
                             // We append the table as long as it's not just the header.
                             if (rowIndex != 1)
                             {
-                                AppendNewLineAndSpan(_currentChunk, tableBuilder.AsSpan(0, tableLength - Environment.NewLine.Length));
+                                AppendNewLineAndSpan(_currentChunk, tableBuilder.AsSpan(0, tableLength - Environment.NewLine.Length), ref totalTokenCount);
                             }
 
                             // And commit the table we built so far.
@@ -137,7 +174,7 @@ internal IEnumerable<IngestionChunk<string>> Process(IngestionDocument document,
                         totalTokenCount += lastRowTokens;
                     }
 
-                    AppendNewLineAndSpan(_currentChunk, tableBuilder.AsSpan(0, tableLength - Environment.NewLine.Length));
+                    AppendNewLineAndSpan(_currentChunk, tableBuilder.AsSpan(0, tableLength - Environment.NewLine.Length), ref totalTokenCount);
                 }
                 finally
                 {
@@ -172,7 +209,7 @@ internal IEnumerable<IngestionChunk<string>> Process(IngestionDocument document,
 
                         totalTokenCount += tokenCount;
                         ReadOnlySpan<char> spanToAppend = remainingContent.Slice(0, index);
-                        AppendNewLineAndSpan(_currentChunk, spanToAppend);
+                        AppendNewLineAndSpan(_currentChunk, spanToAppend, ref totalTokenCount);
                         remainingContent = remainingContent.Slice(index);
                     }
                     else if (totalTokenCount == contextTokenCount)
@@ -196,7 +233,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 +244,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(
@@ -218,12 +259,13 @@ static void ThrowTokenCountExceeded()
             => throw new InvalidOperationException("Can't fit in the current chunk. Consider increasing max tokens per chunk.");
     }
 
-    private static void AppendNewLineAndSpan(StringBuilder stringBuilder, ReadOnlySpan<char> chars)
+    private void AppendNewLineAndSpan(StringBuilder stringBuilder, ReadOnlySpan<char> chars, ref int totalTokenCount)
     {
         // Don't start an empty chunk (no context provided) with a new line.
         if (stringBuilder.Length > 0)
         {
             stringBuilder.AppendLine();
+            totalTokenCount += _newLineTokenCount;
         }
 
 #if NET
@@ -233,41 +275,6 @@ private static void AppendNewLineAndSpan(StringBuilder stringBuilder, ReadOnlySp
 #endif
     }
 
-    private static void AddMarkdownTableRow(IngestionDocumentTable table, int rowIndex, ref ValueStringBuilder vsb)
-    {
-        for (int columnIndex = 0; columnIndex < table.Cells.GetLength(1); columnIndex++)
-        {
-            vsb.Append('|');
-            vsb.Append(' ');
-            string? cellContent = table.Cells[rowIndex, columnIndex] switch
-            {
-                null => null,
-                IngestionDocumentImage img => img.AlternativeText ?? img.Text,
-                IngestionDocumentElement other => other.GetMarkdown()
-            };
-            vsb.Append(cellContent);
-            vsb.Append(' ');
-        }
-
-        vsb.Append('|');
-        vsb.Append(Environment.NewLine);
-    }
-
-    private static void AddMarkdownTableSeparatorRow(int columnCount, ref ValueStringBuilder vsb)
-    {
-        const int DashCount = 3; // The dash count does not need to match the header length.
-        for (int columnIndex = 0; columnIndex < columnCount; columnIndex++)
-        {
-            vsb.Append('|');
-            vsb.Append(' ');
-            vsb.Append('-', DashCount);
-            vsb.Append(' ');
-        }
-
-        vsb.Append('|');
-        vsb.Append(Environment.NewLine);
-    }
-
     private int CountTokens(ReadOnlySpan<char> input)
         => _tokenizer.CountTokens(input, considerNormalization: false);
 }

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.")
+            }
+        });
+
+        Tokenizer tokenizer = TiktokenTokenizer.CreateForModel("gpt-4");
+        HeaderChunker chunker = new(new(tokenizer));
+        IReadOnlyList<IngestionChunk<string>> chunks = await chunker.ProcessAsync(doc).ToListAsync();
+
+        foreach (IngestionChunk<string> chunk in chunks)
+        {
+            // Verify that TokenCount property is set and greater than zero
+            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)
+                }
+            });
+
+            Tokenizer tokenizer = TiktokenTokenizer.CreateForModel("gpt-4o");
+            IngestionChunker<string> chunker = CreateDocumentChunker(maxTokensPerChunk: 512, overlapTokens: 0);
+            IReadOnlyList<IngestionChunk<string>> chunks = await chunker.ProcessAsync(doc).ToListAsync();
+
+            Assert.Equal(2, chunks.Count);
+            foreach (IngestionChunk<string> 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,41 @@
+// 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 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_ThrowsWhenTokenCountIsZero()
+    {
+        IngestionDocument document = new("test");
+
+        ArgumentOutOfRangeException exception = Assert.Throws<ArgumentOutOfRangeException>(
+            () => new IngestionChunk<string>("test content", document, 0));
+
+        Assert.Equal("tokenCount", exception.ParamName);
+    }
+}

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: 123)) // made up number as we currently don't have the ability to easily count exact tokens
                     .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. " +
+        TestChunkFactory.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)." +
+        TestChunkFactory.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),
+        TestChunkFactory.CreateChunk("This text does not belong to any category.", _document),
     ];
 }