Skip to content
Open
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,9 @@ internal static class ProcessFile
internal const string ActivityName = "ProcessFile";
internal const string FilePathTagName = "rag.file.path";
}

internal static class ProcessDocument
{
internal const string ActivityName = "ProcessDocument";
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ namespace Microsoft.Extensions.DataIngestion;
/// </summary>
public sealed class IngestionPipeline : IDisposable
{
private readonly IngestionDocumentReader _reader;
private readonly IngestionChunker _chunker;
private readonly IngestionChunkWriter _writer;
private readonly ActivitySource _activitySource;
Expand All @@ -32,19 +31,16 @@ public sealed class IngestionPipeline : IDisposable
/// <summary>
/// Initializes a new instance of the <see cref="IngestionPipeline"/> class.
/// </summary>
/// <param name="reader">The reader for ingestion documents.</param>
/// <param name="chunker">The chunker to split documents into chunks.</param>
/// <param name="writer">The writer for processing chunks.</param>
/// <param name="options">The options for the ingestion pipeline.</param>
/// <param name="loggerFactory">The logger factory for creating loggers.</param>
public IngestionPipeline(
IngestionDocumentReader reader,
IngestionChunker chunker,
IngestionChunkWriter writer,
IngestionPipelineOptions? options = default,
ILoggerFactory? loggerFactory = default)
{
_reader = Throw.IfNull(reader);
_chunker = Throw.IfNull(chunker);
_writer = Throw.IfNull(writer);
_activitySource = new((options ?? new()).ActivitySourceName);
Expand All @@ -68,17 +64,36 @@ public void Dispose()
/// </summary>
public IList<IngestionChunkProcessor> ChunkProcessors { get; } = [];

/// <summary>
/// Processes the specified document through the pipeline.
/// </summary>
/// <param name="document">The document to process.</param>
/// <param name="cancellationToken">The cancellation token for the operation.</param>
/// <returns>A task representing the asynchronous operation, returning the processed document.</returns>
public async Task<IngestionDocument> ProcessAsync(IngestionDocument document, CancellationToken cancellationToken = default)
{
Throw.IfNull(document);

using (Activity? activity = _activitySource.StartActivity(ProcessDocument.ActivityName))
{
activity?.SetTag(ProcessSource.DocumentIdTagName, document.Identifier);
return await IngestAsync(document, activity, cancellationToken).ConfigureAwait(false);
}
}

/// <summary>
/// Processes all files in the specified directory that match the given search pattern and option.
/// </summary>
/// <param name="reader">The reader for ingestion documents.</param>
/// <param name="directory">The directory to process.</param>
/// <param name="searchPattern">The search pattern for file selection.</param>
/// <param name="searchOption">The search option for directory traversal.</param>
/// <param name="cancellationToken">The cancellation token for the operation.</param>
/// <returns>A task representing the asynchronous operation.</returns>
public async IAsyncEnumerable<IngestionResult> ProcessAsync(DirectoryInfo directory, string searchPattern = "*.*",
public async IAsyncEnumerable<IngestionResult> ProcessAsync(IngestionDocumentReader reader, DirectoryInfo directory, string searchPattern = "*.*",
SearchOption searchOption = SearchOption.TopDirectoryOnly, [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
Throw.IfNull(reader);
Throw.IfNull(directory);
Throw.IfNullOrEmpty(searchPattern);
Throw.IfOutOfRange((int)searchOption, (int)SearchOption.TopDirectoryOnly, (int)SearchOption.AllDirectories);
Expand All @@ -90,7 +105,7 @@ public async IAsyncEnumerable<IngestionResult> ProcessAsync(DirectoryInfo direct
.SetTag(ProcessDirectory.SearchOptionTagName, searchOption.ToString());
_logger?.ProcessingDirectory(directory.FullName, searchPattern, searchOption);

await foreach (var ingestionResult in ProcessAsync(directory.EnumerateFiles(searchPattern, searchOption), rootActivity, cancellationToken).ConfigureAwait(false))
await foreach (IngestionResult ingestionResult in ProcessAsync(reader, directory.EnumerateFiles(searchPattern, searchOption), rootActivity, cancellationToken).ConfigureAwait(false))
{
yield return ingestionResult;
}
Expand All @@ -100,16 +115,18 @@ public async IAsyncEnumerable<IngestionResult> ProcessAsync(DirectoryInfo direct
/// <summary>
/// Processes the specified files.
/// </summary>
/// <param name="reader">The reader for ingestion documents.</param>
/// <param name="files">The collection of files to process.</param>
/// <param name="cancellationToken">The cancellation token for the operation.</param>
/// <returns>A task representing the asynchronous operation.</returns>
public async IAsyncEnumerable<IngestionResult> ProcessAsync(IEnumerable<FileInfo> files, [EnumeratorCancellation] CancellationToken cancellationToken = default)
public async IAsyncEnumerable<IngestionResult> ProcessAsync(IngestionDocumentReader reader, IEnumerable<FileInfo> files, [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
Throw.IfNull(reader);
Throw.IfNull(files);

using (Activity? rootActivity = _activitySource.StartActivity(ProcessFiles.ActivityName))
{
await foreach (var ingestionResult in ProcessAsync(files, rootActivity, cancellationToken).ConfigureAwait(false))
await foreach (IngestionResult ingestionResult in ProcessAsync(reader, files, rootActivity, cancellationToken).ConfigureAwait(false))
{
yield return ingestionResult;
}
Expand All @@ -124,7 +141,7 @@ private static void TraceException(Activity? activity, Exception ex)
.SetStatus(ActivityStatusCode.Error, ex.Message);
}

private async IAsyncEnumerable<IngestionResult> ProcessAsync(IEnumerable<FileInfo> files, Activity? rootActivity,
private async IAsyncEnumerable<IngestionResult> ProcessAsync(IngestionDocumentReader reader, IEnumerable<FileInfo> files, Activity? rootActivity,
[EnumeratorCancellation] CancellationToken cancellationToken)
{
#if NET
Expand All @@ -142,13 +159,13 @@ private async IAsyncEnumerable<IngestionResult> ProcessAsync(IEnumerable<FileInf
using (Activity? processFileActivity = _activitySource.StartActivity(ProcessFile.ActivityName, ActivityKind.Internal, parentContext: rootActivity?.Context ?? default))
{
processFileActivity?.SetTag(ProcessFile.FilePathTagName, fileInfo.FullName);
_logger?.ReadingFile(fileInfo.FullName, GetShortName(_reader));
_logger?.ReadingFile(fileInfo.FullName, GetShortName(reader));

IngestionDocument? document = null;
Exception? failure = null;
try
{
document = await _reader.ReadAsync(fileInfo, cancellationToken).ConfigureAwait(false);
document = await reader.ReadAsync(fileInfo, cancellationToken).ConfigureAwait(false);

processFileActivity?.SetTag(ProcessSource.DocumentIdTagName, document.Identifier);
_logger?.ReadDocument(document.Identifier);
Expand Down Expand Up @@ -180,7 +197,7 @@ private async Task<IngestionDocument> IngestAsync(IngestionDocument document, Ac
}

IAsyncEnumerable<IngestionChunk> chunks = _chunker.ProcessAsync(document, cancellationToken);
foreach (var processor in ChunkProcessors)
foreach (IngestionChunkProcessor processor in ChunkProcessors)
{
chunks = processor.ProcessAsync(chunks, cancellationToken);
}
Expand Down
32 changes: 32 additions & 0 deletions src/Libraries/Microsoft.Extensions.DataIngestion/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,38 @@ using VectorStoreWriter<IngestionChunkVectorRecord> writer = new(collection);
await writer.WriteAsync(chunks);
```

## Using the IngestionPipeline

### Processing documents from the file system

To process documents from the file system, create an `IngestionPipeline` and pass a reader to the `ProcessAsync` method:

```csharp
using IngestionPipeline pipeline = new(chunker, writer);

IngestionDocumentReader reader = new MarkdownReader();
await foreach (IngestionResult result in pipeline.ProcessAsync(reader, directory, "*.md"))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you show a code sample demonstrating error handling when processing the async enumerable under the proposed model?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It could be sth like this:

IngestionDocumentReader reader = new MarkdownReader();
await foreach (IngestionResult result in pipeline.ProcessAsync(reader, directory, "*.md"))
{
    if (!result.Succeeded)
    {
         logger.Error($"Failed to process'{result.DocumentId}'.", result.Exception);
    }
}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But that doesn't process the full stream E2E right? What if I wanted to skip over the ingestion result that failed?

{
Console.WriteLine($"Processed '{result.DocumentId}'. Succeeded: {result.Succeeded}");
}
```

### Processing documents without a reader

The `IngestionPipeline` can also process documents that are already in memory, without requiring a reader:

```csharp
using IngestionPipeline pipeline = new(chunker, writer);

IngestionDocument document = new("my-document-id");
IngestionDocumentSection section = new("Main");
section.Elements.Add(new IngestionDocumentHeader("# Introduction") { Level = 1 });
section.Elements.Add(new IngestionDocumentParagraph("This is the content of my document."));
document.Sections.Add(section);

IngestionDocument processedDocument = await pipeline.ProcessAsync(document);
```

### Custom metadata

To store custom metadata alongside each chunk, create a type derived from `IngestionChunkVectorRecord` with additional properties, and a `VectorStoreWriter` subclass that overrides `SetMetadata`:
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
using Microsoft.Extensions.AI;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DataIngestion;
using Microsoft.Extensions.DataIngestion.Chunkers;
using Microsoft.Extensions.VectorData;
Expand All @@ -19,13 +19,13 @@ public async Task IngestDataAsync(DirectoryInfo directory, string searchPattern)
IncrementalIngestion = false,
});

DocumentReader reader = new(directory);
using var pipeline = new IngestionPipeline(
reader: new DocumentReader(directory),
chunker: new SemanticSimilarityChunker(embeddingGenerator, new(TiktokenTokenizer.CreateForModel("gpt-4o"))),
writer: writer,
loggerFactory: loggerFactory);

await foreach (var result in pipeline.ProcessAsync(directory, searchPattern))
await foreach (var result in pipeline.ProcessAsync(reader, directory, searchPattern))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the reader already encapsulates the directory, why does it also need to be passed in ProcessAsync?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Encapsulating the directory is specific only to this particular reader implementation:

internal sealed class DocumentReader(DirectoryInfo rootDirectory) : IngestionDocumentReader
{
private readonly MarkdownReader _markdownReader = new();
#if (IsAspire)
private readonly MarkItDownMcpReader _pdfReader = new(mcpServerUri: GetMarkItDownMcpServerUrl());
#else
private readonly PdfPigReader _pdfReader = new();
#endif
public override Task<IngestionDocument> ReadAsync(FileInfo source, string identifier, string? mediaType = null, CancellationToken cancellationToken = default)
{
if (Path.IsPathFullyQualified(identifier))
{
// Normalize the identifier to its relative path
identifier = Path.GetRelativePath(rootDirectory.FullName, identifier);

Other readers don't do that.

{
logger.LogInformation("Completed processing '{id}'. Succeeded: '{succeeded}'.", result.DocumentId, result.Succeeded);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -90,8 +90,8 @@ public async Task CanProcessDocuments()
"chunks", TestEmbeddingGenerator<AIContent>.DimensionCount);
using VectorStoreWriter<IngestionChunkVectorRecord> vectorStoreWriter = new(collection);

using IngestionPipeline pipeline = new(CreateReader(), CreateChunker(), vectorStoreWriter);
List<IngestionResult> ingestionResults = await pipeline.ProcessAsync(_sampleFiles).ToListAsync();
using IngestionPipeline pipeline = new(CreateChunker(), vectorStoreWriter);
List<IngestionResult> ingestionResults = await pipeline.ProcessAsync(CreateReader(), _sampleFiles).ToListAsync();

Assert.Equal(_sampleFiles.Count, ingestionResults.Count);
AssertAllIngestionsSucceeded(ingestionResults);
Expand Down Expand Up @@ -127,10 +127,11 @@ public async Task CanProcessDocumentsInDirectory()
"chunks-dir", TestEmbeddingGenerator<AIContent>.DimensionCount);
using VectorStoreWriter<IngestionChunkVectorRecord> vectorStoreWriter = new(collection);

using IngestionPipeline pipeline = new(CreateReader(), CreateChunker(), vectorStoreWriter);
IngestionDocumentReader reader = CreateReader();
using IngestionPipeline pipeline = new(CreateChunker(), vectorStoreWriter);

DirectoryInfo directory = new("TestFiles");
List<IngestionResult> ingestionResults = await pipeline.ProcessAsync(directory, "*.md").ToListAsync();
List<IngestionResult> ingestionResults = await pipeline.ProcessAsync(reader, directory, "*.md").ToListAsync();
Assert.Equal(directory.EnumerateFiles("*.md").Count(), ingestionResults.Count);
AssertAllIngestionsSucceeded(ingestionResults);

Expand Down Expand Up @@ -164,10 +165,11 @@ public async Task ChunksCanBeMoreThanJustText()
VectorStoreCollection<Guid, IngestionChunkVectorRecord> collection = testVectorStore.GetIngestionRecordCollection<IngestionChunkVectorRecord>(
"chunks-img", TestEmbeddingGenerator<AIContent>.DimensionCount);
using VectorStoreWriter<IngestionChunkVectorRecord> vectorStoreWriter = new(collection);
using IngestionPipeline pipeline = new(CreateReader(), new ImageChunker(), vectorStoreWriter);
IngestionDocumentReader reader = CreateReader();
using IngestionPipeline pipeline = new(new ImageChunker(), vectorStoreWriter);

Assert.False(embeddingGenerator.WasCalled);
List<IngestionResult> ingestionResults = await pipeline.ProcessAsync(_sampleFiles).ToListAsync();
List<IngestionResult> ingestionResults = await pipeline.ProcessAsync(reader, _sampleFiles).ToListAsync();
AssertAllIngestionsSucceeded(ingestionResults);

List<IngestionChunkVectorRecord> retrieved = await vectorStoreWriter.VectorStoreCollection
Expand Down Expand Up @@ -248,8 +250,9 @@ public async Task PipelineWorksWithEmbeddingGenerator()
"chunks-aicontent", TestEmbeddingGenerator<AIContent>.DimensionCount);
using VectorStoreWriter<IngestionChunkVectorRecord> vectorStoreWriter = new(collection);

using IngestionPipeline pipeline = new(CreateReader(), CreateChunker(), vectorStoreWriter);
List<IngestionResult> ingestionResults = await pipeline.ProcessAsync(_sampleFiles).ToListAsync();
IngestionDocumentReader reader = CreateReader();
using IngestionPipeline pipeline = new(CreateChunker(), vectorStoreWriter);
List<IngestionResult> ingestionResults = await pipeline.ProcessAsync(reader, _sampleFiles).ToListAsync();

AssertAllIngestionsSucceeded(ingestionResults);
Assert.True(embeddingGenerator.WasCalled, "The embedding generator should have been called.");
Expand Down Expand Up @@ -326,10 +329,10 @@ public async Task SingleFailureDoesNotTearDownEntirePipeline()
"chunks-fail", TestEmbeddingGenerator<AIContent>.DimensionCount);
using VectorStoreWriter<IngestionChunkVectorRecord> vectorStoreWriter = new(collection);

using IngestionPipeline pipeline = new(failingForFirstReader, CreateChunker(), vectorStoreWriter);
using IngestionPipeline pipeline = new(CreateChunker(), vectorStoreWriter);

await Verify(pipeline.ProcessAsync(_sampleFiles));
await Verify(pipeline.ProcessAsync(_sampleDirectory));
await Verify(pipeline.ProcessAsync(failingForFirstReader, _sampleFiles));
await Verify(pipeline.ProcessAsync(failingForFirstReader, _sampleDirectory));

async Task Verify(IAsyncEnumerable<IngestionResult> results)
{
Expand All @@ -346,6 +349,49 @@ async Task Verify(IAsyncEnumerable<IngestionResult> results)
}
}

[Fact]
public async Task CanProcessDocumentWithoutReader()
{
List<Activity> activities = [];
using TracerProvider tracerProvider = CreateTraceProvider(activities);

TestEmbeddingGenerator<AIContent> embeddingGenerator = new();
using InMemoryVectorStore testVectorStore = new(new() { EmbeddingGenerator = embeddingGenerator });

VectorStoreCollection<Guid, IngestionChunkVectorRecord> collection = testVectorStore.GetIngestionRecordCollection<IngestionChunkVectorRecord>(
"chunks-doc", TestEmbeddingGenerator<AIContent>.DimensionCount);
using VectorStoreWriter<IngestionChunkVectorRecord> vectorStoreWriter = new(collection);

using IngestionPipeline pipeline = new(CreateChunker(), vectorStoreWriter);

IngestionDocument document = new("test-document-id");
IngestionDocumentSection section = new("Main");
section.Elements.Add(new IngestionDocumentHeader("# Introduction") { Level = 1 });
section.Elements.Add(new IngestionDocumentParagraph("This is the content of a manually created document."));
document.Sections.Add(section);

IngestionDocument processedDocument = await pipeline.ProcessAsync(document);

Assert.NotNull(processedDocument);
Assert.Equal("test-document-id", processedDocument.Identifier);
Assert.True(embeddingGenerator.WasCalled, "Embedding generator should have been called.");

List<IngestionChunkVectorRecord> retrieved = await vectorStoreWriter.VectorStoreCollection
.GetAsync(record => record.DocumentId == "test-document-id", top: 100)
.ToListAsync();

Assert.NotEmpty(retrieved);
for (int i = 0; i < retrieved.Count; i++)
{
Assert.NotEqual(Guid.Empty, retrieved[i].Key);
Assert.NotEmpty(retrieved[i].SerializedContent!);
Assert.Equal("test-document-id", retrieved[i].DocumentId);
}

Assert.NotEmpty(activities);
Assert.Contains(activities, a => a.OperationName == "ProcessDocument");
}

private static IngestionDocumentReader CreateReader() => new MarkdownReader();

private static IngestionChunker CreateChunker() => new HeaderChunker(new(TiktokenTokenizer.CreateForModel("gpt-4")));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,13 @@ public async Task IngestDataAsync(DirectoryInfo directory, string searchPattern)
IncrementalIngestion = false,
});

DocumentReader reader = new(directory);
using var pipeline = new IngestionPipeline(
reader: new DocumentReader(directory),
chunker: new SemanticSimilarityChunker(embeddingGenerator, new(TiktokenTokenizer.CreateForModel("gpt-4o"))),
writer: writer,
loggerFactory: loggerFactory);

await foreach (var result in pipeline.ProcessAsync(directory, searchPattern))
await foreach (var result in pipeline.ProcessAsync(reader, directory, searchPattern))
{
logger.LogInformation("Completed processing '{id}'. Succeeded: '{succeeded}'.", result.DocumentId, result.Succeeded);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,13 @@ public async Task IngestDataAsync(DirectoryInfo directory, string searchPattern)
IncrementalIngestion = false,
});

DocumentReader reader = new(directory);
using var pipeline = new IngestionPipeline(
reader: new DocumentReader(directory),
chunker: new SemanticSimilarityChunker(embeddingGenerator, new(TiktokenTokenizer.CreateForModel("gpt-4o"))),
writer: writer,
loggerFactory: loggerFactory);

await foreach (var result in pipeline.ProcessAsync(directory, searchPattern))
await foreach (var result in pipeline.ProcessAsync(reader, directory, searchPattern))
{
logger.LogInformation("Completed processing '{id}'. Succeeded: '{succeeded}'.", result.DocumentId, result.Succeeded);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,13 @@ public async Task IngestDataAsync(DirectoryInfo directory, string searchPattern)
IncrementalIngestion = false,
});

DocumentReader reader = new(directory);
using var pipeline = new IngestionPipeline(
reader: new DocumentReader(directory),
chunker: new SemanticSimilarityChunker(embeddingGenerator, new(TiktokenTokenizer.CreateForModel("gpt-4o"))),
writer: writer,
loggerFactory: loggerFactory);

await foreach (var result in pipeline.ProcessAsync(directory, searchPattern))
await foreach (var result in pipeline.ProcessAsync(reader, directory, searchPattern))
{
logger.LogInformation("Completed processing '{id}'. Succeeded: '{succeeded}'.", result.DocumentId, result.Succeeded);
}
Expand Down
Loading
Loading