diff --git a/.github/workflows/medi.yml b/.github/workflows/medi.yml new file mode 100644 index 0000000..b8e774d --- /dev/null +++ b/.github/workflows/medi.yml @@ -0,0 +1,64 @@ +name: MEDI + +on: + push: + branches: [main] + paths: + - 'MEDI/**' + - 'Shared/**' + - 'eng/**' + - 'Directory.Build.props' + - 'Directory.Packages.props' + - 'global.json' + - '.github/workflows/medi.yml' + pull_request: + branches: [main] + paths: + - 'MEDI/**' + - 'Shared/**' + - 'eng/**' + - 'Directory.Build.props' + - 'Directory.Packages.props' + - 'global.json' + - '.github/workflows/medi.yml' + +env: + DOTNET_NOLOGO: true + DOTNET_CLI_TELEMETRY_OPTOUT: true + +jobs: + build: + name: Build + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Setup .NET + uses: actions/setup-dotnet@v4 + with: + global-json-file: global.json + + - name: Build + run: dotnet build MEDI/MEDI.slnf + + unit-tests: + name: 'Unit: ${{ matrix.package }} (${{ matrix.os }})' + needs: build + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest] + package: + - PdfPig.OnnxLayoutAnalysis + - PdfPig.DataIngestion + steps: + - uses: actions/checkout@v4 + + - name: Setup .NET + uses: actions/setup-dotnet@v4 + with: + global-json-file: global.json + + - name: Test + run: dotnet test MEDI/test/${{ matrix.package }}.UnitTests diff --git a/CommunityToolkit.AI.slnx b/CommunityToolkit.AI.slnx index 79732c7..8948b10 100644 --- a/CommunityToolkit.AI.slnx +++ b/CommunityToolkit.AI.slnx @@ -6,6 +6,15 @@ + + + + + + + + + diff --git a/Directory.Packages.props b/Directory.Packages.props index 7199d1e..4e4afed 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -22,6 +22,21 @@ + + + + + + + + + + + + + + + @@ -36,6 +51,8 @@ + + diff --git a/MEDI/MEDI.slnf b/MEDI/MEDI.slnf new file mode 100644 index 0000000..042c63f --- /dev/null +++ b/MEDI/MEDI.slnf @@ -0,0 +1,15 @@ +{ + "solution": { + "path": "../CommunityToolkit.AI.slnx", + "projects": + [ + "MEDI/src/PdfPig.OnnxLayoutAnalysis/PdfPig.OnnxLayoutAnalysis.csproj", + "MEDI/src/PdfPig.DataIngestion/PdfPig.DataIngestion.csproj", + "MEDI/src/DataRetrieval/DataRetrieval.csproj", + "MEDI/src/DataRetrieval.OnnxReranker/DataRetrieval.OnnxReranker.csproj", + + "MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/PdfPig.OnnxLayoutAnalysis.UnitTests.csproj", + "MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPig.DataIngestion.UnitTests.csproj" + ] + } +} diff --git a/MEDI/src/DataIngestion/ContextualChunkEnricher.cs b/MEDI/src/DataIngestion/ContextualChunkEnricher.cs new file mode 100644 index 0000000..ec1cc9f --- /dev/null +++ b/MEDI/src/DataIngestion/ContextualChunkEnricher.cs @@ -0,0 +1,52 @@ +using System.Runtime.CompilerServices; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; + +namespace CommunityToolkit.DataIngestion; + +/// +/// An that enriches text chunks with contextual +/// summaries for improved RAG retrieval. +/// +public sealed class ContextualChunkEnricher : IngestionChunkProcessor +{ + private readonly IChatClient _chatClient; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client used to generate contextual summaries. + public ContextualChunkEnricher(IChatClient chatClient) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + } + + /// + public override async IAsyncEnumerable> ProcessAsync( + IAsyncEnumerable> chunks, + [EnumeratorCancellation] CancellationToken ct = default) + { + await foreach (var chunk in chunks.WithCancellation(ct)) + { + var prompt = GetPromptForChunk(chunk); + var messages = new[] { new ChatMessage(ChatRole.User, prompt) }; + var response = await _chatClient.GetResponseAsync(messages, cancellationToken: ct); + + chunk.Metadata[MetadataKeys.ContextualSummary] = response.Text; + yield return chunk; + } + } + + private static string GetPromptForChunk(IngestionChunk chunk) + { + var elementType = chunk.Metadata.TryGetValue("element_type", out var t) ? t as string : null; + var instruction = elementType switch + { + "table" => "Summarize this table for search retrieval. Describe what data it contains, including key metrics, column headers, and notable values. Output only the summary sentence, nothing else.", + "picture" or "caption" => "Summarize what this figure or image refers to for search retrieval. Output only the summary sentence, nothing else.", + _ => "Provide a single concise sentence summarizing the following text for use in search retrieval. Output only the summary sentence, nothing else." + }; + + return instruction + "\n\n" + chunk.Content; + } +} diff --git a/MEDI/src/DataIngestion/DataIngestion.csproj b/MEDI/src/DataIngestion/DataIngestion.csproj new file mode 100644 index 0000000..698edd1 --- /dev/null +++ b/MEDI/src/DataIngestion/DataIngestion.csproj @@ -0,0 +1,25 @@ + + + 1.0.0-preview.1 + CommunityToolkit.DataIngestion + $(AssemblyName) + net10.0;net8.0 + + LLM-powered chunk processors for Microsoft.Extensions.DataIngestion pipelines + LLM-powered ingestion chunk processors including contextual chunk enrichment, entity extraction, topic classification, hypothetical query generation (reverse HyDE), and RAPTOR-style tree indexing for advanced RAG applications. + true + + + + + + + + + + + + + + + diff --git a/MEDI/src/DataIngestion/DataIngestionServiceCollectionExtensions.cs b/MEDI/src/DataIngestion/DataIngestionServiceCollectionExtensions.cs new file mode 100644 index 0000000..847f86e --- /dev/null +++ b/MEDI/src/DataIngestion/DataIngestionServiceCollectionExtensions.cs @@ -0,0 +1,39 @@ +using Microsoft.Extensions.DataIngestion; +using Microsoft.Extensions.DependencyInjection; + +namespace CommunityToolkit.DataIngestion; + +/// +/// Extension methods for registering ingestion pipeline builders with dependency injection. +/// +public static class DataIngestionServiceCollectionExtensions +{ + /// + /// Registers an singleton for composing + /// ingestion pipelines with LLM-powered chunk processors. + /// + /// + /// + /// builder.Services.AddIngestionPipeline<string>() + /// .UseEntityExtraction() + /// .UseTopicClassification(o => o.Taxonomy = ["web", "data", "security"]) + /// .UseHypotheticalQueries(o => o.QuestionsPerChunk = 5) + /// .UseContextualChunkEnrichment() + /// .UseTreeIndex(); + /// + /// + public static IngestionPipelineBuilder AddIngestionPipeline( + this IServiceCollection services) + { + var pipelineBuilder = new IngestionPipelineBuilder(); + services.AddSingleton(pipelineBuilder); + return pipelineBuilder; + } + + /// + /// Registers an with string content type (the common case). + /// + public static IngestionPipelineBuilder AddIngestionPipeline( + this IServiceCollection services) + => services.AddIngestionPipeline(); +} diff --git a/MEDI/src/DataIngestion/EntityExtractionProcessor.cs b/MEDI/src/DataIngestion/EntityExtractionProcessor.cs new file mode 100644 index 0000000..8f2bbd4 --- /dev/null +++ b/MEDI/src/DataIngestion/EntityExtractionProcessor.cs @@ -0,0 +1,105 @@ +using System.Runtime.CompilerServices; +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Logging.Abstractions; + +namespace CommunityToolkit.DataIngestion; + +/// +/// An that uses an LLM to extract named entities +/// (people, organizations, technologies, versions) from each chunk and stores them as +/// metadata for downstream filtered vector search. +/// +public sealed class EntityExtractionProcessor : IngestionChunkProcessor +{ + private readonly IChatClient _chatClient; + private readonly ILogger _logger; + + private static readonly string[] EntityKeys = + [MetadataKeys.EntitiesPeople, MetadataKeys.EntitiesOrganizations, MetadataKeys.EntitiesTechnologies, MetadataKeys.EntitiesVersions]; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client used for entity extraction. + /// Optional logger for diagnostics. + public EntityExtractionProcessor(IChatClient chatClient, ILogger? logger = null) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + _logger = logger ?? NullLogger.Instance; + } + + /// + public override async IAsyncEnumerable> ProcessAsync( + IAsyncEnumerable> chunks, + [EnumeratorCancellation] CancellationToken ct = default) + { + await foreach (var chunk in chunks.WithCancellation(ct)) + { + try + { + var prompt = $$""" + Extract named entities from the following text. + Return ONLY valid JSON matching: {"people": ["name"], "organizations": ["org"], "technologies": ["tech"], "versions": ["v1"]} + Use empty arrays [] when none found for a category. + + Text: + {{chunk.Content}} + """; + + var options = new ChatOptions + { + MaxOutputTokens = 300, + ResponseFormat = ChatResponseFormat.Json + }; + + var response = await _chatClient.GetResponseAsync(prompt, options, ct); + var entities = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options); + + if (entities is not null) + { + chunk.Metadata[MetadataKeys.EntitiesPeople] = string.Join(", ", entities.People ?? []); + chunk.Metadata[MetadataKeys.EntitiesOrganizations] = string.Join(", ", entities.Organizations ?? []); + chunk.Metadata[MetadataKeys.EntitiesTechnologies] = string.Join(", ", entities.Technologies ?? []); + chunk.Metadata[MetadataKeys.EntitiesVersions] = string.Join(", ", entities.Versions ?? []); + } + else + { + SetEmptyEntities(chunk); + } + } + catch (Exception ex) + { + _logger.LogWarning(ex, "Entity extraction failed for chunk; using empty defaults."); + SetEmptyEntities(chunk); + } + + yield return chunk; + } + } + + private static void SetEmptyEntities(IngestionChunk chunk) + { + foreach (var key in EntityKeys) + chunk.Metadata[key] = ""; + } + + private sealed class EntityResponse + { + [JsonPropertyName("people")] + public string[]? People { get; set; } + + [JsonPropertyName("organizations")] + public string[]? Organizations { get; set; } + + [JsonPropertyName("technologies")] + public string[]? Technologies { get; set; } + + [JsonPropertyName("versions")] + public string[]? Versions { get; set; } + } +} diff --git a/MEDI/src/DataIngestion/HypotheticalQueryProcessor.cs b/MEDI/src/DataIngestion/HypotheticalQueryProcessor.cs new file mode 100644 index 0000000..35bde79 --- /dev/null +++ b/MEDI/src/DataIngestion/HypotheticalQueryProcessor.cs @@ -0,0 +1,113 @@ +using System.Runtime.CompilerServices; +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Logging.Abstractions; + +namespace CommunityToolkit.DataIngestion; + +/// +/// An that generates hypothetical questions each chunk +/// could answer, storing them as additional retrieval vectors in the same collection. +/// This is "reverse HyDE" — bridging the query-document gap at ingestion time. +/// +/// +/// For each original chunk, yields: +/// +/// The original chunk (with chunk_type=original metadata) +/// N question chunks (with chunk_type=hypothetical_query and parent_chunk_id) +/// +/// Question chunks contain the question text as their Content — the embedding model creates +/// vectors in "question space," directly matching user queries. +/// +public sealed class HypotheticalQueryProcessor : IngestionChunkProcessor +{ + private readonly IChatClient _chatClient; + private readonly int _questionsPerChunk; + private readonly ILogger _logger; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client used for question generation. + /// Number of hypothetical questions to generate per chunk (default: 3). + /// Optional logger for diagnostics. + public HypotheticalQueryProcessor(IChatClient chatClient, int questionsPerChunk = 3, ILogger? logger = null) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + _questionsPerChunk = questionsPerChunk; + _logger = logger ?? NullLogger.Instance; + } + + /// + public override async IAsyncEnumerable> ProcessAsync( + IAsyncEnumerable> chunks, + [EnumeratorCancellation] CancellationToken ct = default) + { + await foreach (var chunk in chunks.WithCancellation(ct)) + { + var chunkId = chunk.Content.GetHashCode().ToString(); + + // Yield original chunk first (with consistent metadata keys) + chunk.Metadata[MetadataKeys.ChunkType] = MetadataKeys.ChunkTypeOriginal; + chunk.Metadata[MetadataKeys.ParentChunkId] = ""; + chunk.Metadata[MetadataKeys.HypotheticalQuestions] = ""; + yield return chunk; + + // Generate hypothetical questions (best-effort) + var questionChunks = new List>(); + try + { + var questions = await GenerateQuestionsAsync(chunk.Content, ct); + chunk.Metadata[MetadataKeys.HypotheticalQuestions] = string.Join(" | ", questions); + + foreach (var question in questions) + { + var qChunk = new IngestionChunk(question, chunk.Document, chunk.Context); + qChunk.Metadata[MetadataKeys.ChunkType] = MetadataKeys.ChunkTypeHypotheticalQuery; + qChunk.Metadata[MetadataKeys.ParentChunkId] = chunkId; + qChunk.Metadata[MetadataKeys.HypotheticalQuestions] = ""; + questionChunks.Add(qChunk); + } + } + catch (Exception ex) + { + _logger.LogWarning(ex, "Hypothetical question generation failed for chunk; continuing with original only."); + } + + foreach (var qChunk in questionChunks) + yield return qChunk; + } + } + + private async Task> GenerateQuestionsAsync(string chunkContent, CancellationToken ct) + { + var prompt = $"Generate exactly {_questionsPerChunk} questions that the following " + + "text passage answers.\n" + + "Return ONLY valid JSON matching: {\"questions\": [\"question1?\", \"question2?\", \"question3?\"]}" + + $"\n\nPassage:\n{chunkContent}"; + + var response = await _chatClient.GetResponseAsync(prompt, + new ChatOptions + { + MaxOutputTokens = 200, + ResponseFormat = ChatResponseFormat.Json + }, ct); + + var result = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options); + + return (result?.Questions ?? []) + .Where(q => !string.IsNullOrWhiteSpace(q) && q.Length > 10) + .Take(_questionsPerChunk) + .ToList(); + } + + private sealed class QuestionsResponse + { + [JsonPropertyName("questions")] + public List Questions { get; set; } = []; + } +} diff --git a/MEDI/src/DataIngestion/IngestionPipelineBuilder.cs b/MEDI/src/DataIngestion/IngestionPipelineBuilder.cs new file mode 100644 index 0000000..b297b56 --- /dev/null +++ b/MEDI/src/DataIngestion/IngestionPipelineBuilder.cs @@ -0,0 +1,149 @@ +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; + +namespace CommunityToolkit.DataIngestion; + +/// +/// Fluent builder for composing instances with +/// document and chunk processors. Registered via +/// . +/// +/// The chunk content type (typically string). +public class IngestionPipelineBuilder +{ + internal List> DocumentProcessorFactories { get; } = []; + internal List>> ChunkProcessorFactories { get; } = []; + + /// Adds a document processor created by the given factory. + public IngestionPipelineBuilder UseDocumentProcessor( + Func factory) + { + DocumentProcessorFactories.Add(factory); + return this; + } + + /// Adds a document processor resolved from DI. + public IngestionPipelineBuilder UseDocumentProcessor() + where TProcessor : IngestionDocumentProcessor + { + DocumentProcessorFactories.Add(sp => + ActivatorUtilities.CreateInstance(sp)); + return this; + } + + /// Adds a chunk processor created by the given factory. + public IngestionPipelineBuilder UseChunkProcessor( + Func> factory) + { + ChunkProcessorFactories.Add(factory); + return this; + } + + /// Adds a chunk processor resolved from DI. + public IngestionPipelineBuilder UseChunkProcessor() + where TProcessor : IngestionChunkProcessor + { + ChunkProcessorFactories.Add(sp => + ActivatorUtilities.CreateInstance(sp)); + return this; + } + + /// Adds entity extraction (people, organizations, technologies, versions) to chunks. + public IngestionPipelineBuilder UseEntityExtraction() + { + ChunkProcessorFactories.Add(sp => + (IngestionChunkProcessor)(object)new EntityExtractionProcessor( + sp.GetRequiredService(), + sp.GetService>())); + return this; + } + + /// Adds topic classification with a configurable taxonomy. + public IngestionPipelineBuilder UseTopicClassification( + Action? configure = null) + { + ChunkProcessorFactories.Add(sp => + { + var options = new TopicClassificationOptions(); + configure?.Invoke(options); + return (IngestionChunkProcessor)(object)new TopicClassificationProcessor( + sp.GetRequiredService(), + options.Taxonomy, + sp.GetService>()); + }); + return this; + } + + /// Adds hypothetical query generation for reverse-HyDE chunk enrichment. + public IngestionPipelineBuilder UseHypotheticalQueries( + Action? configure = null) + { + ChunkProcessorFactories.Add(sp => + { + var options = new HypotheticalQueryOptions(); + configure?.Invoke(options); + return (IngestionChunkProcessor)(object)new HypotheticalQueryProcessor( + sp.GetRequiredService(), + options.QuestionsPerChunk, + sp.GetService>()); + }); + return this; + } + + /// Adds contextual chunk enrichment for improved retrieval. + public IngestionPipelineBuilder UseContextualChunkEnrichment() + { + ChunkProcessorFactories.Add(sp => + (IngestionChunkProcessor)(object)new ContextualChunkEnricher( + sp.GetRequiredService())); + return this; + } + + /// Adds RAPTOR-style tree index generation (leaf → branch → root summaries). + public IngestionPipelineBuilder UseTreeIndex() + { + ChunkProcessorFactories.Add(sp => + (IngestionChunkProcessor)(object)new TreeIndexProcessor( + sp.GetRequiredService(), + sp.GetService>())); + return this; + } + + /// + /// Creates a configured with all registered processors. + /// The caller is responsible for disposing the returned pipeline. + /// + public IngestionPipeline Build( + IServiceProvider serviceProvider, + IngestionDocumentReader reader, + IngestionChunker chunker, + IngestionChunkWriter writer, + ILoggerFactory? loggerFactory = null) + { + var pipeline = new IngestionPipeline(reader, chunker, writer, loggerFactory: loggerFactory); + + foreach (var factory in DocumentProcessorFactories) + pipeline.DocumentProcessors.Add(factory(serviceProvider)); + + foreach (var factory in ChunkProcessorFactories) + pipeline.ChunkProcessors.Add(factory(serviceProvider)); + + return pipeline; + } +} + +/// Options for . +public class TopicClassificationOptions +{ + /// Valid topic labels for classification. + public string[] Taxonomy { get; set; } = ["web", "data", "performance", "security", "architecture"]; +} + +/// Options for . +public class HypotheticalQueryOptions +{ + /// Number of questions to generate per chunk. + public int QuestionsPerChunk { get; set; } = 3; +} diff --git a/MEDI/src/DataIngestion/JsonDefaults.cs b/MEDI/src/DataIngestion/JsonDefaults.cs new file mode 100644 index 0000000..fb74687 --- /dev/null +++ b/MEDI/src/DataIngestion/JsonDefaults.cs @@ -0,0 +1,13 @@ +using System.Text.Json; + +namespace CommunityToolkit.DataIngestion; + +/// Shared JSON serialization defaults for LLM response parsing. +internal static class JsonDefaults +{ + public static readonly JsonSerializerOptions Options = new() + { + PropertyNameCaseInsensitive = true, + PropertyNamingPolicy = JsonNamingPolicy.CamelCase + }; +} diff --git a/MEDI/src/DataIngestion/MetadataKeys.cs b/MEDI/src/DataIngestion/MetadataKeys.cs new file mode 100644 index 0000000..d862857 --- /dev/null +++ b/MEDI/src/DataIngestion/MetadataKeys.cs @@ -0,0 +1,55 @@ +namespace CommunityToolkit.DataIngestion; + +/// +/// Well-known metadata keys written by the ingestion chunk processors in this package. +/// +public static class MetadataKeys +{ + /// Comma-separated list of people entities extracted from the chunk. + public const string EntitiesPeople = "entities_people"; + + /// Comma-separated list of organization entities extracted from the chunk. + public const string EntitiesOrganizations = "entities_organizations"; + + /// Comma-separated list of technology entities extracted from the chunk. + public const string EntitiesTechnologies = "entities_technologies"; + + /// Comma-separated list of version strings extracted from the chunk. + public const string EntitiesVersions = "entities_versions"; + + /// Primary topic label assigned by topic classification. + public const string TopicPrimary = "topic_primary"; + + /// Comma-separated secondary topic labels assigned by topic classification. + public const string TopicSecondary = "topic_secondary"; + + /// The type of chunk (original, hypothetical_query, branch_summary, root_summary). + public const string ChunkType = "chunk_type"; + + /// The identifier of the parent chunk (for hypothetical query chunks). + public const string ParentChunkId = "parent_chunk_id"; + + /// Pipe-delimited hypothetical questions generated for a chunk. + public const string HypotheticalQuestions = "hypothetical_questions"; + + /// Contextual summary generated for a chunk. + public const string ContextualSummary = "contextual_summary"; + + /// Tree hierarchy level (0 = leaf, 1 = branch, 2 = root). + public const string Level = "level"; + + /// The identifier of the parent node in the tree index hierarchy. + public const string ParentId = "parent_id"; + + /// Chunk type value for original (leaf) content chunks. + public const string ChunkTypeOriginal = "original"; + + /// Chunk type value for hypothetical query expansion chunks. + public const string ChunkTypeHypotheticalQuery = "hypothetical_query"; + + /// Chunk type value for document-level (branch) summary chunks. + public const string ChunkTypeBranchSummary = "branch_summary"; + + /// Chunk type value for corpus-level (root) summary chunks. + public const string ChunkTypeRootSummary = "root_summary"; +} diff --git a/MEDI/src/DataIngestion/TopicClassificationProcessor.cs b/MEDI/src/DataIngestion/TopicClassificationProcessor.cs new file mode 100644 index 0000000..feefa5e --- /dev/null +++ b/MEDI/src/DataIngestion/TopicClassificationProcessor.cs @@ -0,0 +1,99 @@ +using System.Runtime.CompilerServices; +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Logging.Abstractions; + +namespace CommunityToolkit.DataIngestion; + +/// +/// An that classifies each chunk into primary and +/// optional secondary topics from a configurable taxonomy. Stores topic_primary +/// (single label) and topic_secondary (comma-separated) as chunk metadata. +/// +public sealed class TopicClassificationProcessor : IngestionChunkProcessor +{ + private readonly IChatClient _chatClient; + private readonly string[] _taxonomy; + private readonly ILogger _logger; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client used for topic classification. + /// Valid topic labels (e.g., "web", "data", "security", "performance", "architecture"). + /// Optional logger for diagnostics. + public TopicClassificationProcessor(IChatClient chatClient, string[] taxonomy, ILogger? logger = null) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + _taxonomy = taxonomy ?? throw new ArgumentNullException(nameof(taxonomy)); + _logger = logger ?? NullLogger.Instance; + } + + /// + public override async IAsyncEnumerable> ProcessAsync( + IAsyncEnumerable> chunks, + [EnumeratorCancellation] CancellationToken ct = default) + { + await foreach (var chunk in chunks.WithCancellation(ct)) + { + string primary = "unknown"; + string secondary = ""; + + try + { + var prompt = BuildPrompt(chunk.Content); + var options = new ChatOptions + { + MaxOutputTokens = 100, + ResponseFormat = ChatResponseFormat.Json + }; + + var response = await _chatClient.GetResponseAsync(prompt, options, ct); + var parsed = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options) + ?? new TopicResponse(); + + if (_taxonomy.Contains(parsed.Primary)) + { + primary = parsed.Primary; + var validSecondary = (parsed.Secondary ?? []) + .Where(t => _taxonomy.Contains(t) && t != primary) + .Distinct(); + secondary = string.Join(",", validSecondary); + } + } + catch (Exception ex) + { + _logger.LogWarning(ex, "Topic classification failed for chunk; defaulting to 'unknown'."); + } + + chunk.Metadata[MetadataKeys.TopicPrimary] = primary; + chunk.Metadata[MetadataKeys.TopicSecondary] = secondary; + + yield return chunk; + } + } + + private string BuildPrompt(string text) => $$""" + Classify this text into topics from: [{{string.Join(", ", _taxonomy.Select(t => $"\"{t}\""))}}]. + + Return ONLY valid JSON (no extra text): + {"primary": "topic", "secondary": ["topic2", "topic3"]} + + If the text fits only one topic, return an empty secondary array. + + Text: {{text}} + """; + + private sealed class TopicResponse + { + [JsonPropertyName("primary")] + public string Primary { get; set; } = ""; + + [JsonPropertyName("secondary")] + public List Secondary { get; set; } = []; + } +} diff --git a/MEDI/src/DataIngestion/TreeIndexProcessor.cs b/MEDI/src/DataIngestion/TreeIndexProcessor.cs new file mode 100644 index 0000000..af5e7b6 --- /dev/null +++ b/MEDI/src/DataIngestion/TreeIndexProcessor.cs @@ -0,0 +1,140 @@ +using System.Runtime.CompilerServices; +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Logging.Abstractions; + +namespace CommunityToolkit.DataIngestion; + +/// +/// A RAPTOR-style tree index processor that generates summary nodes at multiple levels. +/// +/// +/// For each document's chunks, generates: +/// +/// Level 1 (Branch): LLM summary of all chunks within the same document +/// Level 2 (Root): LLM summary of all branch summaries (corpus overview) +/// +/// Summary nodes are injected into the same chunk stream with level metadata. +/// The vector store writer stores them in the same collection as leaf chunks. +/// Standard VectorStoreCollection.SearchAsync naturally returns a mix of leaf +/// chunks (specific detail) and summary nodes (broader context). +/// +public sealed class TreeIndexProcessor : IngestionChunkProcessor +{ + private readonly IChatClient _chatClient; + private readonly ILogger _logger; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client used for summarization. + /// Optional logger for diagnostics. + public TreeIndexProcessor(IChatClient chatClient, ILogger? logger = null) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + _logger = logger ?? NullLogger.Instance; + } + + /// + public override async IAsyncEnumerable> ProcessAsync( + IAsyncEnumerable> chunks, + [EnumeratorCancellation] CancellationToken ct = default) + { + // Collect all chunks, grouping by source document for summarization + var chunksByDoc = new Dictionary>>(); + var allChunks = new List>(); + + await foreach (var chunk in chunks.WithCancellation(ct)) + { + // Mark as leaf node + chunk.Metadata[MetadataKeys.Level] = 0; + chunk.Metadata[MetadataKeys.ChunkType] = MetadataKeys.ChunkTypeOriginal; + chunk.Metadata[MetadataKeys.ParentId] = ""; + + allChunks.Add(chunk); + + var docId = SanitizeDocId(chunk.Document.Identifier); + if (!chunksByDoc.TryGetValue(docId, out var list)) + chunksByDoc[docId] = list = []; + list.Add(chunk); + } + + // Yield all original leaf chunks first + foreach (var chunk in allChunks) + yield return chunk; + + // Level 1 (Branch): Generate document-level summaries + var branchSummaries = new List<(string DocId, string Summary)>(); + foreach (var (docId, docChunks) in chunksByDoc) + { + var combinedText = string.Join("\n\n", docChunks.Select(c => c.Content)); + var summary = await SummarizeAsync( + $"Summarize the following text in 2-3 concise sentences. Capture key concepts and technologies.\n\n{combinedText}", ct); + + branchSummaries.Add((docId, summary)); + + // Update leaf chunks to reference their branch parent + foreach (var leaf in docChunks) + leaf.Metadata[MetadataKeys.ParentId] = $"branch-{docId}"; + + // Create and yield branch summary chunk + var branchChunk = new IngestionChunk(summary, docChunks[0].Document, docChunks[0].Context); + branchChunk.Metadata[MetadataKeys.Level] = 1; + branchChunk.Metadata[MetadataKeys.ChunkType] = MetadataKeys.ChunkTypeBranchSummary; + branchChunk.Metadata[MetadataKeys.ParentId] = "root"; + yield return branchChunk; + } + + // Level 2 (Root): Generate corpus-level summary + if (branchSummaries.Count > 0) + { + var allBranchText = string.Join("\n\n", + branchSummaries.Select(b => $"[{b.DocId}]: {b.Summary}")); + + var rootSummary = await SummarizeAsync( + "Write a single 2-3 sentence overview of the entire corpus:\n\n" + allBranchText, ct); + + var rootChunk = new IngestionChunk(rootSummary, allChunks[0].Document, allChunks[0].Context); + rootChunk.Metadata[MetadataKeys.Level] = 2; + rootChunk.Metadata[MetadataKeys.ChunkType] = MetadataKeys.ChunkTypeRootSummary; + rootChunk.Metadata[MetadataKeys.ParentId] = ""; + yield return rootChunk; + } + } + + private async Task SummarizeAsync(string prompt, CancellationToken ct) + { + var fullPrompt = prompt + "\nReturn ONLY valid JSON: {\"summary\": \"your summary here\"}"; + + try + { + var response = await _chatClient.GetResponseAsync(fullPrompt, + new ChatOptions + { + MaxOutputTokens = 200, + ResponseFormat = ChatResponseFormat.Json + }, ct); + + var result = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options); + return result?.Summary ?? "Summary unavailable"; + } + catch (Exception ex) + { + _logger.LogWarning(ex, "Tree index summarization failed; using placeholder."); + return "Summary unavailable"; + } + } + + private static string SanitizeDocId(string identifier) => + Path.GetFileNameWithoutExtension(identifier).Replace(" ", "-").ToLowerInvariant(); + + private sealed class SummaryResponse + { + [JsonPropertyName("summary")] + public string Summary { get; set; } = ""; + } +} diff --git a/MEDI/src/DataRetrieval.OnnxReranker/CrossEncoderReranker.cs b/MEDI/src/DataRetrieval.OnnxReranker/CrossEncoderReranker.cs new file mode 100644 index 0000000..547c5cb --- /dev/null +++ b/MEDI/src/DataRetrieval.OnnxReranker/CrossEncoderReranker.cs @@ -0,0 +1,156 @@ +using CommunityToolkit.DataRetrieval.OnnxReranker.Internal; +using Microsoft.Extensions.DataRetrieval; +using Microsoft.ML.Tokenizers; + +namespace CommunityToolkit.DataRetrieval.OnnxReranker; + +/// +/// ONNX cross-encoder reranker. Scores query-document pairs using a local cross-encoder +/// model (e.g., ms-marco-MiniLM-L-6-v2, BGE-reranker) for fast, high-quality reranking. +/// +/// Pipeline: text-pair tokenization → ONNX inference → sigmoid normalization. +/// +/// Dependencies: Microsoft.ML.OnnxRuntime.Managed + Microsoft.ML.Tokenizers (no Microsoft.ML). +/// +/// +/// Lazy initialization: the ONNX session and tokenizer are loaded on first use. +/// Call to release native ONNX resources. +/// +public class CrossEncoderReranker : RetrievalResultProcessor, IReranker, IDisposable +{ + private readonly CrossEncoderRerankerOptions _options; + private readonly object _initLock = new(); + + private OnnxModelSession? _session; + private TextPairTokenizer? _pairTokenizer; + private bool _disposed; + + /// + /// Initializes a new instance of the class. + /// + /// Reranker configuration options. + public CrossEncoderReranker(CrossEncoderRerankerOptions options) + { + _options = options ?? throw new ArgumentNullException(nameof(options)); + + if (string.IsNullOrWhiteSpace(options.ModelPath)) + throw new ArgumentException("ModelPath is required.", nameof(options)); + if (string.IsNullOrWhiteSpace(options.TokenizerPath)) + throw new ArgumentException("TokenizerPath is required.", nameof(options)); + } + + /// + public override async Task ProcessAsync( + RetrievalResults results, RetrievalQuery query, CancellationToken cancellationToken = default) + { + if (results.Chunks.Count <= 1) + return results; + + var reranked = await RerankAsync( + query.Text, + results.Chunks.ToList(), + cancellationToken); + + results.Chunks = reranked.Take(_options.MaxResults).ToList(); + results.Metadata["reranked"] = true; + results.Metadata["reranker"] = "CrossEncoder"; + results.Metadata["reranked_count"] = results.Chunks.Count; + return results; + } + + /// + public Task> RerankAsync( + string query, + IReadOnlyList chunks, + CancellationToken cancellationToken = default) + { + if (chunks.Count == 0) + return Task.FromResult>([]); + + return Task.Run(() => + { + EnsureInitialized(); + cancellationToken.ThrowIfCancellationRequested(); + + var documents = chunks.Select(r => r.Content).ToList(); + var scores = ScoreDocuments(query, documents); + + var reranked = chunks + .Select((chunk, i) => + { + chunk.Score = scores[i]; + return chunk; + }) + .OrderByDescending(c => c.Score) + .ToList(); + + return (IReadOnlyList)reranked; + }, cancellationToken); + } + + private float[] ScoreDocuments(string query, IList documents) + { + var allScores = new List(documents.Count); + int batchSize = _options.BatchSize; + + for (int start = 0; start < documents.Count; start += batchSize) + { + int count = Math.Min(batchSize, documents.Count - start); + var batchQueries = Enumerable.Repeat(query, count).ToList(); + var batchDocs = documents.Skip(start).Take(count).ToList(); + + var tokenized = _pairTokenizer!.TokenizePairs(batchQueries, batchDocs); + var rawOutputs = _session!.Score(tokenized); + + for (int i = 0; i < rawOutputs.Length; i++) + { + float logit = rawOutputs[i][0]; + allScores.Add(Sigmoid(logit)); + } + } + + return [.. allScores]; + } + + private void EnsureInitialized() + { + ObjectDisposedException.ThrowIf(_disposed, this); + + if (_session != null && _pairTokenizer != null) + return; + + lock (_initLock) + { + if (_session != null && _pairTokenizer != null) + return; + + var tokenizer = TokenizerFactory.LoadTokenizer(_options.TokenizerPath); + var (bosId, sepId, doubleSep) = TokenizerFactory.DiscoverSpecialTokens(_options.TokenizerPath); + + if (bosId == null || sepId == null) + throw new InvalidOperationException( + "Could not determine BOS/CLS and SEP token IDs for text-pair tokenization. " + + "Ensure the tokenizer directory contains a tokenizer_config.json with cls_token/sep_token " + + "or a recognized tokenizer_class."); + + _pairTokenizer = new TextPairTokenizer( + tokenizer, _options.MaxTokenLength, bosId.Value, sepId.Value, doubleSep); + + _session = new OnnxModelSession( + _options.ModelPath, _options.GpuDeviceId, _options.FallbackToCpu, _options.PreferredOutputNames); + } + } + + private static float Sigmoid(float x) => 1.0f / (1.0f + MathF.Exp(-x)); + + /// + public void Dispose() + { + if (_disposed) + return; + + _disposed = true; + _session?.Dispose(); + GC.SuppressFinalize(this); + } +} diff --git a/MEDI/src/DataRetrieval.OnnxReranker/CrossEncoderRerankerOptions.cs b/MEDI/src/DataRetrieval.OnnxReranker/CrossEncoderRerankerOptions.cs new file mode 100644 index 0000000..cdf38d9 --- /dev/null +++ b/MEDI/src/DataRetrieval.OnnxReranker/CrossEncoderRerankerOptions.cs @@ -0,0 +1,43 @@ +namespace CommunityToolkit.DataRetrieval.OnnxReranker; + +/// +/// Configuration for the ONNX cross-encoder reranker. +/// +public sealed class CrossEncoderRerankerOptions +{ + /// Path to the ONNX cross-encoder model file (.onnx). + public required string ModelPath { get; set; } + + /// + /// Path to tokenizer artifacts — either a directory containing tokenizer_config.json + /// (or vocab.txt, tokenizer.model, etc.) or a direct path to a vocab file. + /// + public required string TokenizerPath { get; set; } + + /// Maximum token length for the combined query + document pair. Default: 512. + public int MaxTokenLength { get; set; } = 512; + + /// Batch size for ONNX inference. Default: 32. + public int BatchSize { get; set; } = 32; + + /// + /// Optional GPU device ID. Null = CPU only. + /// Requires the consuming application to reference Microsoft.ML.OnnxRuntime.Gpu. + /// + public int? GpuDeviceId { get; set; } + + /// + /// If true and GPU initialization fails, fall back to CPU instead of throwing. + /// Default: false. + /// + public bool FallbackToCpu { get; set; } + + /// Maximum number of results to return after reranking. Default: 5. + public int MaxResults { get; set; } = 5; + + /// + /// Preferred ONNX output tensor names, searched in order. + /// Default: ["logits", "output"] — appropriate for cross-encoder models. + /// + public string[] PreferredOutputNames { get; set; } = ["logits", "output"]; +} diff --git a/MEDI/src/DataRetrieval.OnnxReranker/DataRetrieval.OnnxReranker.csproj b/MEDI/src/DataRetrieval.OnnxReranker/DataRetrieval.OnnxReranker.csproj new file mode 100644 index 0000000..75a0021 --- /dev/null +++ b/MEDI/src/DataRetrieval.OnnxReranker/DataRetrieval.OnnxReranker.csproj @@ -0,0 +1,24 @@ + + + 1.0.0-preview.1 + CommunityToolkit.DataRetrieval.OnnxReranker + $(AssemblyName) + net10.0;net8.0 + + ONNX cross-encoder re-ranking for DataRetrieval pipelines + ONNX Runtime cross-encoder re-ranker implementing IReranker for high-quality, low-latency passage re-ranking without LLM API calls. + true + + + + + + + + + + + + + + diff --git a/MEDI/src/DataRetrieval.OnnxReranker/Internal/OnnxModelSession.cs b/MEDI/src/DataRetrieval.OnnxReranker/Internal/OnnxModelSession.cs new file mode 100644 index 0000000..b0c454b --- /dev/null +++ b/MEDI/src/DataRetrieval.OnnxReranker/Internal/OnnxModelSession.cs @@ -0,0 +1,186 @@ +using Microsoft.ML.OnnxRuntime; + +namespace CommunityToolkit.DataRetrieval.OnnxReranker.Internal; + +/// +/// Discovered ONNX model tensor metadata. +/// +internal sealed record OnnxModelMetadata( + string InputIdsName, + string AttentionMaskName, + string? TokenTypeIdsName, + string OutputTensorName, + int HiddenDim, + int OutputRank); + +/// +/// Manages an ONNX InferenceSession for cross-encoder scoring. +/// Handles session creation, metadata discovery, batch inference, and disposal. +/// +internal sealed class OnnxModelSession : IDisposable +{ + private readonly InferenceSession _session; + private readonly OnnxModelMetadata _metadata; + + public OnnxModelMetadata Metadata => _metadata; + + public OnnxModelSession(string modelPath, int? gpuDeviceId, bool fallbackToCpu, string[]? preferredOutputNames) + { + if (!File.Exists(modelPath)) + throw new FileNotFoundException($"ONNX model not found: {modelPath}"); + + _session = CreateSession(modelPath, gpuDeviceId, fallbackToCpu); + _metadata = DiscoverMetadata(_session, preferredOutputNames); + } + + /// + /// Runs ONNX inference on a tokenized batch. Returns raw logits per item. + /// + public float[][] Score(TokenizedBatch batch) + { + int totalRows = batch.Count; + int seqLen = batch.SequenceLength; + var allOutputs = new List(totalRows); + + int batchSize = totalRows; + for (int start = 0; start < totalRows; start += batchSize) + { + int count = Math.Min(batchSize, totalRows - start); + var batchOutputs = RunOnnxBatch( + batch.TokenIds, batch.AttentionMasks, batch.TokenTypeIds, + start, count, seqLen); + allOutputs.AddRange(batchOutputs); + } + + return [.. allOutputs]; + } + + private float[][] RunOnnxBatch( + long[][] tokenIds, long[][] attentionMasks, long[][] tokenTypeIds, + int startIdx, int batchSize, int seqLen) + { + var idsArray = new long[batchSize * seqLen]; + var maskArray = new long[batchSize * seqLen]; + var typeIdsArray = _metadata.TokenTypeIdsName != null ? new long[batchSize * seqLen] : null; + + for (int b = 0; b < batchSize; b++) + { + Array.Copy(tokenIds[startIdx + b], 0, idsArray, b * seqLen, seqLen); + Array.Copy(attentionMasks[startIdx + b], 0, maskArray, b * seqLen, seqLen); + if (typeIdsArray != null) + Array.Copy(tokenTypeIds[startIdx + b], 0, typeIdsArray, b * seqLen, seqLen); + } + + var inputs = new Dictionary + { + [_metadata.InputIdsName] = OrtValue.CreateTensorValueFromMemory(idsArray, [batchSize, seqLen]), + [_metadata.AttentionMaskName] = OrtValue.CreateTensorValueFromMemory(maskArray, [batchSize, seqLen]) + }; + + if (_metadata.TokenTypeIdsName != null && typeIdsArray != null) + inputs[_metadata.TokenTypeIdsName] = OrtValue.CreateTensorValueFromMemory(typeIdsArray, [batchSize, seqLen]); + + try + { + using var results = _session.Run(new RunOptions(), inputs, [_metadata.OutputTensorName]); + var output = results[0]; + var outputSpan = output.GetTensorDataAsSpan(); + + var batchOutputs = new float[batchSize][]; + + if (_metadata.OutputRank == 2) + { + for (int b = 0; b < batchSize; b++) + batchOutputs[b] = outputSpan.Slice(b * _metadata.HiddenDim, _metadata.HiddenDim).ToArray(); + } + else + { + int rowSize = seqLen * _metadata.HiddenDim; + for (int b = 0; b < batchSize; b++) + batchOutputs[b] = outputSpan.Slice(b * rowSize, rowSize).ToArray(); + } + + return batchOutputs; + } + finally + { + foreach (var ortValue in inputs.Values) + ortValue.Dispose(); + } + } + + private static InferenceSession CreateSession(string modelPath, int? gpuDeviceId, bool fallbackToCpu) + { + var options = new SessionOptions(); + + if (gpuDeviceId.HasValue) + { + try + { + options.AppendExecutionProvider_CUDA(gpuDeviceId.Value); + } + catch (Exception) when (fallbackToCpu) + { + // CUDA libraries not available — fall back to CPU + } + } + + try + { + return new InferenceSession(modelPath, options); + } + catch (OnnxRuntimeException) when (fallbackToCpu) + { + return new InferenceSession(modelPath, new SessionOptions()); + } + } + + private static OnnxModelMetadata DiscoverMetadata(InferenceSession session, string[]? preferredOutputNames) + { + var inputMeta = session.InputMetadata; + var outputMeta = session.OutputMetadata; + + string inputIdsName = FindTensorName(inputMeta, ["input_ids"], inputMeta.Keys.First()); + string attentionMaskName = FindTensorName(inputMeta, ["attention_mask"], "attention_mask"); + string? tokenTypeIdsName = TryFindTensorName(inputMeta, ["token_type_ids"]); + + string outputName; + + if (preferredOutputNames != null) + { + var preferred = TryFindTensorName(outputMeta, preferredOutputNames); + outputName = preferred ?? outputMeta.Keys.First(); + } + else + { + var pooledName = TryFindTensorName(outputMeta, ["logits", "output", "sentence_embedding", "pooler_output"]); + outputName = pooledName ?? outputMeta.Keys.First(); + } + + var dims = outputMeta[outputName].Dimensions; + int outputRank = dims.Length; + int lastDim = (int)dims.Last(); + int hiddenDim = lastDim > 0 ? lastDim : 1; + + return new OnnxModelMetadata(inputIdsName, attentionMaskName, tokenTypeIdsName, outputName, hiddenDim, outputRank); + } + + private static string FindTensorName( + IReadOnlyDictionary metadata, string[] candidates, string fallback) + { + return TryFindTensorName(metadata, candidates) ?? fallback; + } + + private static string? TryFindTensorName( + IReadOnlyDictionary metadata, string[] candidates) + { + foreach (var candidate in candidates) + { + if (metadata.ContainsKey(candidate)) + return candidate; + } + return null; + } + + public void Dispose() => _session.Dispose(); +} diff --git a/MEDI/src/DataRetrieval.OnnxReranker/Internal/TextPairTokenizer.cs b/MEDI/src/DataRetrieval.OnnxReranker/Internal/TextPairTokenizer.cs new file mode 100644 index 0000000..75908e0 --- /dev/null +++ b/MEDI/src/DataRetrieval.OnnxReranker/Internal/TextPairTokenizer.cs @@ -0,0 +1,110 @@ +using Microsoft.ML.Tokenizers; + +namespace CommunityToolkit.DataRetrieval.OnnxReranker.Internal; + +/// +/// Batch of tokenized text pairs for ONNX cross-encoder inference. +/// +internal sealed class TokenizedBatch +{ + public long[][] TokenIds { get; } + public long[][] AttentionMasks { get; } + public long[][] TokenTypeIds { get; } + public int SequenceLength { get; } + public int Count => TokenIds.Length; + + public TokenizedBatch(long[][] tokenIds, long[][] attentionMasks, long[][] tokenTypeIds, int sequenceLength) + { + TokenIds = tokenIds; + AttentionMasks = attentionMasks; + TokenTypeIds = tokenTypeIds; + SequenceLength = sequenceLength; + } +} + +/// +/// Tokenizes query-document pairs for cross-encoder models. +/// Produces [CLS] query_tokens [SEP] document_tokens [SEP] with proper token_type_ids. +/// +internal sealed class TextPairTokenizer +{ + private readonly Tokenizer _tokenizer; + private readonly int _maxTokenLength; + private readonly int _bosTokenId; + private readonly int _sepTokenId; + private readonly bool _doubleSeparator; + + public TextPairTokenizer(Tokenizer tokenizer, int maxTokenLength, int bosTokenId, int sepTokenId, bool doubleSeparator = false) + { + _tokenizer = tokenizer ?? throw new ArgumentNullException(nameof(tokenizer)); + _maxTokenLength = maxTokenLength; + _bosTokenId = bosTokenId; + _sepTokenId = sepTokenId; + _doubleSeparator = doubleSeparator; + } + + /// + /// Tokenizes a batch of query-document pairs for cross-encoder scoring. + /// + public TokenizedBatch TokenizePairs(IReadOnlyList queries, IReadOnlyList documents) + { + if (queries.Count != documents.Count) + throw new ArgumentException("queries and documents must have the same length."); + + int batchSize = queries.Count; + var allTokenIds = new long[batchSize][]; + var allAttentionMasks = new long[batchSize][]; + var allTokenTypeIds = new long[batchSize][]; + + for (int i = 0; i < batchSize; i++) + { + var tokenIds = new long[_maxTokenLength]; + var attentionMask = new long[_maxTokenLength]; + var tokenTypeIds = new long[_maxTokenLength]; + + TokenizeSinglePair(queries[i], documents[i], tokenIds, attentionMask, tokenTypeIds); + + allTokenIds[i] = tokenIds; + allAttentionMasks[i] = attentionMask; + allTokenTypeIds[i] = tokenTypeIds; + } + + return new TokenizedBatch(allTokenIds, allAttentionMasks, allTokenTypeIds, _maxTokenLength); + } + + private void TokenizeSinglePair( + string textA, string textB, + long[] tokenIds, long[] attentionMask, long[] tokenTypeIds) + { + var encodedA = _tokenizer.EncodeToTokens(textA, out _); + var encodedB = _tokenizer.EncodeToTokens(textB, out _); + + var combined = new List(_maxTokenLength); + combined.Add(_bosTokenId); + + for (int j = 0; j < encodedA.Count; j++) + combined.Add(encodedA[j].Id); + + combined.Add(_sepTokenId); + + if (_doubleSeparator) + combined.Add(_sepTokenId); + + int firstSepIdx = combined.Count - 1; + + for (int j = 0; j < encodedB.Count; j++) + combined.Add(encodedB[j].Id); + + combined.Add(_sepTokenId); + + if (combined.Count > _maxTokenLength) + combined.RemoveRange(_maxTokenLength, combined.Count - _maxTokenLength); + + for (int s = 0; s < combined.Count; s++) + { + tokenIds[s] = combined[s]; + attentionMask[s] = 1; + tokenTypeIds[s] = s <= firstSepIdx ? 0 : 1; + } + } +} diff --git a/MEDI/src/DataRetrieval.OnnxReranker/Internal/TokenizerFactory.cs b/MEDI/src/DataRetrieval.OnnxReranker/Internal/TokenizerFactory.cs new file mode 100644 index 0000000..3682ab8 --- /dev/null +++ b/MEDI/src/DataRetrieval.OnnxReranker/Internal/TokenizerFactory.cs @@ -0,0 +1,248 @@ +using System.Text.Json; +using Microsoft.ML.Tokenizers; + +namespace CommunityToolkit.DataRetrieval.OnnxReranker.Internal; + +/// +/// Loads a HuggingFace-compatible tokenizer from a directory or file path. +/// Supports BERT/WordPiece, SentencePiece, and BPE tokenizer formats. +/// +internal static class TokenizerFactory +{ + /// + /// Loads a tokenizer from the given path. Accepts a directory (with tokenizer_config.json + /// or known vocab files) or a direct path to a config/vocab file. + /// + public static Tokenizer LoadTokenizer(string path) + { + if (Directory.Exists(path)) + return LoadFromDirectory(path); + + var fileName = Path.GetFileName(path).ToLowerInvariant(); + if (fileName == "tokenizer_config.json") + return LoadFromConfig(path); + + if (fileName == "tokenizer.json") + return LoadFromHuggingFaceTokenizerJson(path); + + return LoadFromVocabFile(path); + } + + /// + /// Extracts BOS (CLS) and SEP token IDs from tokenizer_config.json. + /// Also determines whether the model uses double separators (RoBERTa-family). + /// + public static (int? BosTokenId, int? SepTokenId, bool DoubleSeparator) DiscoverSpecialTokens(string tokenizerPath) + { + string? configPath = null; + + if (Directory.Exists(tokenizerPath)) + { + var candidate = Path.Combine(tokenizerPath, "tokenizer_config.json"); + if (File.Exists(candidate)) + configPath = candidate; + } + else if (Path.GetFileName(tokenizerPath).Equals("tokenizer_config.json", StringComparison.OrdinalIgnoreCase)) + { + configPath = tokenizerPath; + } + + if (configPath == null) + return (null, null, false); + + var json = File.ReadAllText(configPath); + using var doc = JsonDocument.Parse(json); + var root = doc.RootElement; + + int? bosId = null; + int? sepId = null; + + var clsTokenStr = root.TryGetProperty("cls_token", out var cls) ? cls.GetString() : null; + var sepTokenStr = root.TryGetProperty("sep_token", out var sep) ? sep.GetString() : null; + + if (clsTokenStr != null && sepTokenStr != null + && root.TryGetProperty("added_tokens_decoder", out var decoder)) + { + foreach (var entry in decoder.EnumerateObject()) + { + if (!int.TryParse(entry.Name, out int tokenId)) continue; + var content = entry.Value.TryGetProperty("content", out var c) ? c.GetString() : null; + if (content == null) continue; + + if (content == clsTokenStr && bosId == null) + bosId = tokenId; + if (content == sepTokenStr && sepId == null) + sepId = tokenId; + } + } + + if (bosId == null || sepId == null) + { + var tokenizerClass = root.TryGetProperty("tokenizer_class", out var tc) ? tc.GetString() ?? "" : ""; + if (tokenizerClass.EndsWith("Fast", StringComparison.Ordinal)) + tokenizerClass = tokenizerClass[..^4]; + + (int? defaultBos, int? defaultSep) = tokenizerClass switch + { + "BertTokenizer" or "DistilBertTokenizer" => ((int?)101, (int?)102), + "RobertaTokenizer" or "GPT2Tokenizer" => ((int?)0, (int?)2), + "DebertaTokenizer" or "DebertaV2Tokenizer" => ((int?)1, (int?)2), + "XLMRobertaTokenizer" => ((int?)0, (int?)2), + _ => ((int?)null, (int?)null) + }; + bosId ??= defaultBos; + sepId ??= defaultSep; + } + + var tokClass = root.TryGetProperty("tokenizer_class", out var t) ? t.GetString() ?? "" : ""; + if (tokClass.EndsWith("Fast", StringComparison.Ordinal)) + tokClass = tokClass[..^4]; + bool doubleSep = tokClass is "RobertaTokenizer" or "GPT2Tokenizer" or "XLMRobertaTokenizer"; + + return (bosId, sepId, doubleSep); + } + + private static Tokenizer LoadFromDirectory(string directory) + { + var configPath = Path.Combine(directory, "tokenizer_config.json"); + if (File.Exists(configPath)) + return LoadFromConfig(configPath); + + var vocabTxt = Path.Combine(directory, "vocab.txt"); + if (File.Exists(vocabTxt)) + return LoadFromVocabFile(vocabTxt); + + var spModel = Path.Combine(directory, "tokenizer.model"); + if (File.Exists(spModel)) + return LoadFromVocabFile(spModel); + + var spBpeModel = Path.Combine(directory, "sentencepiece.bpe.model"); + if (File.Exists(spBpeModel)) + return LoadFromVocabFile(spBpeModel); + + var spmModel = Path.Combine(directory, "spm.model"); + if (File.Exists(spmModel)) + return LoadFromVocabFile(spmModel); + + var tokenizerJson = Path.Combine(directory, "tokenizer.json"); + if (File.Exists(tokenizerJson)) + return LoadFromHuggingFaceTokenizerJson(tokenizerJson); + + throw new FileNotFoundException( + $"No tokenizer_config.json or known vocab file found in '{directory}'. " + + $"Expected one of: tokenizer_config.json, vocab.txt, tokenizer.model, sentencepiece.bpe.model, spm.model, tokenizer.json."); + } + + private static Tokenizer LoadFromConfig(string configPath) + { + var directory = Path.GetDirectoryName(configPath) + ?? throw new ArgumentException($"Cannot determine directory for config: {configPath}"); + + var json = File.ReadAllText(configPath); + using var doc = JsonDocument.Parse(json); + var root = doc.RootElement; + + var tokenizerClass = root.TryGetProperty("tokenizer_class", out var tcls) + ? tcls.GetString() ?? "" + : ""; + + if (tokenizerClass.EndsWith("Fast", StringComparison.Ordinal)) + tokenizerClass = tokenizerClass[..^4]; + + return tokenizerClass switch + { + "BertTokenizer" => LoadBertFromConfig(directory, root), + "DistilBertTokenizer" => LoadBertFromConfig(directory, root), + "XLMRobertaTokenizer" => LoadSentencePieceFromDirectory(directory), + "LlamaTokenizer" => LoadSentencePieceFromDirectory(directory), + "CamembertTokenizer" => LoadSentencePieceFromDirectory(directory), + "T5Tokenizer" => LoadSentencePieceFromDirectory(directory), + "AlbertTokenizer" => LoadSentencePieceFromDirectory(directory), + "DebertaTokenizer" => LoadSentencePieceFromDirectory(directory), + "DebertaV2Tokenizer" => LoadSentencePieceFromDirectory(directory), + "GPT2Tokenizer" => LoadBpeFromDirectory(directory), + "RobertaTokenizer" => LoadBpeFromDirectory(directory), + _ when !string.IsNullOrEmpty(tokenizerClass) => throw new NotSupportedException( + $"Unsupported tokenizer_class '{tokenizerClass}' in {configPath}. " + + $"Supported: BertTokenizer, DebertaV2Tokenizer, XLMRobertaTokenizer, LlamaTokenizer, GPT2Tokenizer, RobertaTokenizer. " + + $"Use a tokenizer_config.json with a supported tokenizer_class."), + _ => throw new InvalidOperationException( + $"No tokenizer_class found in {configPath}. Cannot auto-detect tokenizer type.") + }; + } + + private static Tokenizer LoadBertFromConfig(string directory, JsonElement config) + { + var vocabPath = Path.Combine(directory, "vocab.txt"); + if (!File.Exists(vocabPath)) + throw new FileNotFoundException( + $"BERT tokenizer requires vocab.txt in '{directory}'."); + + var lowerCase = config.TryGetProperty("do_lower_case", out var lc) && lc.GetBoolean(); + + using var stream = File.OpenRead(vocabPath); + return BertTokenizer.Create(stream, new BertOptions { LowerCaseBeforeTokenization = lowerCase }); + } + + private static Tokenizer LoadSentencePieceFromDirectory(string directory) + { + var candidates = new[] { "sentencepiece.bpe.model", "tokenizer.model", "spiece.model", "spm.model" }; + foreach (var candidate in candidates) + { + var spPath = Path.Combine(directory, candidate); + if (File.Exists(spPath)) + { + using var stream = File.OpenRead(spPath); + return SentencePieceTokenizer.Create(stream); + } + } + + throw new FileNotFoundException( + $"SentencePiece tokenizer requires one of [{string.Join(", ", candidates)}] in '{directory}'."); + } + + private static Tokenizer LoadBpeFromDirectory(string directory) + { + var vocabJson = Path.Combine(directory, "vocab.json"); + var mergesPath = Path.Combine(directory, "merges.txt"); + + if (!File.Exists(vocabJson)) + throw new FileNotFoundException($"BPE tokenizer requires vocab.json in '{directory}'."); + if (!File.Exists(mergesPath)) + throw new FileNotFoundException($"BPE tokenizer requires merges.txt in '{directory}'."); + + using var vocabStream = File.OpenRead(vocabJson); + using var mergesStream = File.OpenRead(mergesPath); + return CodeGenTokenizer.Create(vocabStream, mergesStream); + } + + private static Tokenizer LoadFromHuggingFaceTokenizerJson(string path) + { + var directory = Path.GetDirectoryName(path); + if (directory != null) + { + var vocabTxt = Path.Combine(directory, "vocab.txt"); + if (File.Exists(vocabTxt)) + return LoadFromVocabFile(vocabTxt); + } + + throw new NotSupportedException( + $"tokenizer.json without a tokenizer_config.json is not supported for cross-encoder reranking. " + + $"Please provide a directory with tokenizer_config.json and vocab.txt."); + } + + private static Tokenizer LoadFromVocabFile(string path) + { + var ext = Path.GetExtension(path).ToLowerInvariant(); + using var stream = File.OpenRead(path); + return ext switch + { + ".txt" => BertTokenizer.Create(stream), + ".model" => SentencePieceTokenizer.Create(stream), + _ => throw new NotSupportedException( + $"Unsupported vocab file extension '{ext}'. " + + $"Use .txt for BERT/WordPiece, .model for SentencePiece, " + + $"or point to a directory with tokenizer_config.json.") + }; + } +} diff --git a/MEDI/src/DataRetrieval/AdaptiveRouter.cs b/MEDI/src/DataRetrieval/AdaptiveRouter.cs new file mode 100644 index 0000000..376dfcc --- /dev/null +++ b/MEDI/src/DataRetrieval/AdaptiveRouter.cs @@ -0,0 +1,86 @@ +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataRetrieval; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// Adaptive RAG router — classifies queries and routes to the appropriate search paradigm. +/// +/// +/// Routes between: +/// +/// Vector: specific factual queries → standard vector search +/// TreeTraversal: broad/thematic queries → top-down tree traversal +/// EntityFiltered: entity-specific queries → vector search with entity metadata filter +/// +/// +public class AdaptiveRouter : RetrievalQueryProcessor +{ + private readonly IChatClient _chatClient; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client for LLM inference. + public AdaptiveRouter(IChatClient chatClient) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + } + + /// + public override async Task ProcessAsync( + RetrievalQuery query, CancellationToken cancellationToken = default) + { + var prompt = $$""" + Classify this search query into one of these categories: + - "vector": specific factual question about a particular topic, API, or concept + - "tree": broad/thematic question requiring multi-document overview or summary + - "entity": question about a specific named entity (person, organization, technology) + + Query: "{{query.Text}}" + + Return ONLY valid JSON: {"paradigm": "vector", "reasoning": "one sentence"} + """; + + string paradigm; + try + { + var response = await _chatClient.GetResponseAsync(prompt, + new ChatOptions + { + MaxOutputTokens = 100, + ResponseFormat = ChatResponseFormat.Json + }, cancellationToken); + + var result = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options); + + paradigm = result?.Paradigm switch + { + "vector" or "tree" or "entity" => result.Paradigm, + _ => "vector" + }; + + if (result?.Reasoning is not null) + query.Metadata["router_reasoning"] = result.Reasoning; + } + catch + { + paradigm = "vector"; + } + + query.Metadata["search_paradigm"] = paradigm; + return query; + } + + private sealed class RouterResponse + { + [JsonPropertyName("paradigm")] + public string Paradigm { get; set; } = "vector"; + + [JsonPropertyName("reasoning")] + public string? Reasoning { get; set; } + } +} diff --git a/MEDI/src/DataRetrieval/CragValidator.cs b/MEDI/src/DataRetrieval/CragValidator.cs new file mode 100644 index 0000000..60e5d50 --- /dev/null +++ b/MEDI/src/DataRetrieval/CragValidator.cs @@ -0,0 +1,119 @@ +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataRetrieval; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// Corrective RAG (CRAG) validator — scores retrieval quality and routes through +/// one of three paths: +/// +/// Correct (≥4): use results as-is +/// Ambiguous (2-3): flag for follow-up, keep results with warning +/// Incorrect (<2): clear chunks, set low_confidence flag +/// +/// +public class CragValidator : RetrievalResultProcessor +{ + private readonly IChatClient _chatClient; + + /// Number of top chunks to evaluate. + public int EvaluateTopN { get; init; } = 3; + + /// Maximum preview length per passage (characters). + public int PreviewLength { get; init; } = 300; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client for LLM inference. + public CragValidator(IChatClient chatClient) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + } + + /// + public override async Task ProcessAsync( + RetrievalResults results, RetrievalQuery query, CancellationToken cancellationToken = default) + { + if (results.Chunks.Count == 0) + { + results.Metadata["crag_score"] = 0; + results.Metadata["crag_path"] = "Incorrect"; + results.Metadata["low_confidence"] = true; + return results; + } + + var topChunks = results.Chunks.Take(EvaluateTopN).ToList(); + var passagesText = string.Join("\n---\n", + topChunks.Select(c => c.Content.Length > PreviewLength + ? c.Content[..PreviewLength] + : c.Content)); + + var prompt = $$""" + Rate how well these passages answer the query on a scale of 1-5: + 5 = directly answers with specific details + 4 = highly relevant, mostly answers + 3 = somewhat related but incomplete + 2 = tangentially related + 1 = completely different topic + + Query: {{query.Text}} + Passages: + {{passagesText}} + + Return ONLY valid JSON: {"score": 5, "reasoning": "explanation"} + """; + + var options = new ChatOptions + { + MaxOutputTokens = 200, + ResponseFormat = ChatResponseFormat.Json + }; + + int score; + try + { + var response = await _chatClient.GetResponseAsync(prompt, options, cancellationToken); + var result = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options); + score = result?.Score is >= 1 and <= 5 ? result.Score : 3; + if (result?.Reasoning is not null) + results.Metadata["crag_reasoning"] = result.Reasoning; + } + catch + { + score = 3; + } + + results.Metadata["crag_score"] = score; + + if (score >= 4) + { + results.Metadata["crag_path"] = "Correct"; + } + else if (score >= 2) + { + results.Metadata["crag_path"] = "Ambiguous"; + results.Metadata["needs_followup"] = true; + } + else + { + results.Metadata["crag_path"] = "Incorrect"; + results.Metadata["low_confidence"] = true; + results.Chunks = []; + } + + return results; + } + + private sealed class CragResponse + { + [JsonPropertyName("score")] + public int Score { get; set; } + + [JsonPropertyName("reasoning")] + public string? Reasoning { get; set; } + } +} diff --git a/MEDI/src/DataRetrieval/DataRetrieval.csproj b/MEDI/src/DataRetrieval/DataRetrieval.csproj new file mode 100644 index 0000000..fdcbab0 --- /dev/null +++ b/MEDI/src/DataRetrieval/DataRetrieval.csproj @@ -0,0 +1,26 @@ + + + 1.0.0-preview.1 + CommunityToolkit.DataRetrieval + $(AssemblyName) + net10.0;net8.0 + + Advanced retrieval processors for Microsoft.Extensions.DataRetrieval pipelines + LLM-powered retrieval pipeline processors including multi-query expansion, HyDE, adaptive routing, CRAG validation, LLM re-ranking, Self-RAG, and Speculative RAG for advanced RAG applications. + true + + + + + + + + + + + + + + + + diff --git a/MEDI/src/DataRetrieval/DataRetrievalServiceCollectionExtensions.cs b/MEDI/src/DataRetrieval/DataRetrievalServiceCollectionExtensions.cs new file mode 100644 index 0000000..6756e86 --- /dev/null +++ b/MEDI/src/DataRetrieval/DataRetrievalServiceCollectionExtensions.cs @@ -0,0 +1,72 @@ +using Microsoft.Extensions.DataRetrieval; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; +using RetrievalPipeline = Microsoft.Extensions.DataRetrieval.RetrievalPipeline; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// Extension methods for registering retrieval pipeline builders with dependency injection. +/// +public static class DataRetrievalServiceCollectionExtensions +{ + /// + /// Registers a singleton and returns a + /// for composing processors. + /// + /// + /// + /// builder.Services.AddRetrievalPipeline() + /// .UseQueryExpansion(o => o.VariantCount = 5) + /// .UseLlmReranking() + /// .UseCrag(); + /// + /// + public static RetrievalPipelineBuilder AddRetrievalPipeline( + this IServiceCollection services) + { + var pipelineBuilder = new RetrievalPipelineBuilder { Services = services }; + + services.AddSingleton(sp => + { + var loggerFactory = sp.GetService(); + var pipeline = new RetrievalPipeline(loggerFactory: loggerFactory); + + foreach (var factory in pipelineBuilder.QueryProcessorFactories) + pipeline.QueryProcessors.Add(factory(sp)); + + foreach (var factory in pipelineBuilder.ResultProcessorFactories) + pipeline.ResultProcessors.Add(factory(sp)); + + return pipeline; + }); + + return pipelineBuilder; + } + + /// + /// Registers a singleton using a custom factory and returns a + /// for composing processors on top. + /// + public static RetrievalPipelineBuilder AddRetrievalPipeline( + this IServiceCollection services, + Func pipelineFactory) + { + var pipelineBuilder = new RetrievalPipelineBuilder { Services = services }; + + services.AddSingleton(sp => + { + var pipeline = pipelineFactory(sp); + + foreach (var factory in pipelineBuilder.QueryProcessorFactories) + pipeline.QueryProcessors.Add(factory(sp)); + + foreach (var factory in pipelineBuilder.ResultProcessorFactories) + pipeline.ResultProcessors.Add(factory(sp)); + + return pipeline; + }); + + return pipelineBuilder; + } +} diff --git a/MEDI/src/DataRetrieval/HydeQueryTransformer.cs b/MEDI/src/DataRetrieval/HydeQueryTransformer.cs new file mode 100644 index 0000000..4b0829e --- /dev/null +++ b/MEDI/src/DataRetrieval/HydeQueryTransformer.cs @@ -0,0 +1,58 @@ +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataRetrieval; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// Hypothetical Document Embedding (HyDE) — generates a hypothetical answer to the query, +/// then searches using that answer's embedding instead of the raw query. +/// This bridges the query-document semantic gap. +/// +public class HydeQueryTransformer : RetrievalQueryProcessor +{ + private readonly IChatClient _chatClient; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client for LLM inference. + public HydeQueryTransformer(IChatClient chatClient) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + } + + /// + public override async Task ProcessAsync( + RetrievalQuery query, CancellationToken cancellationToken = default) + { + var prompt = $""" + Write a short, factual paragraph (3-4 sentences) that would directly answer this question. + Write as if you are a documentation page. Do NOT include any preamble or explanation. + + Question: {query.Text} + """; + + var options = new ChatOptions + { + MaxOutputTokens = 250 + }; + + string hypothetical; + try + { + var response = await _chatClient.GetResponseAsync(prompt, options, cancellationToken); + hypothetical = (response.Text ?? "").Trim(); + } + catch + { + return query; + } + + if (string.IsNullOrWhiteSpace(hypothetical)) + return query; + + query.Variants = [hypothetical]; + query.Metadata["hyde_hypothetical"] = hypothetical; + return query; + } +} diff --git a/MEDI/src/DataRetrieval/JsonDefaults.cs b/MEDI/src/DataRetrieval/JsonDefaults.cs new file mode 100644 index 0000000..9402378 --- /dev/null +++ b/MEDI/src/DataRetrieval/JsonDefaults.cs @@ -0,0 +1,13 @@ +using System.Text.Json; + +namespace CommunityToolkit.DataRetrieval; + +/// Shared JSON serialization defaults for LLM response parsing. +internal static class JsonDefaults +{ + public static readonly JsonSerializerOptions Options = new() + { + PropertyNameCaseInsensitive = true, + PropertyNamingPolicy = JsonNamingPolicy.CamelCase + }; +} diff --git a/MEDI/src/DataRetrieval/LlmReranker.cs b/MEDI/src/DataRetrieval/LlmReranker.cs new file mode 100644 index 0000000..33febf3 --- /dev/null +++ b/MEDI/src/DataRetrieval/LlmReranker.cs @@ -0,0 +1,108 @@ +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataRetrieval; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// LLM-based reranker — asks the model to rank passages by relevance to the query. +/// Implements for pipeline integration. +/// +public class LlmReranker : RetrievalResultProcessor +{ + private readonly IChatClient _chatClient; + + /// Maximum number of results to return after reranking. + public int MaxResults { get; init; } = 5; + + /// Maximum candidate passages to send to the LLM (controls token cost). + public int MaxCandidates { get; init; } = 8; + + /// Maximum preview length per passage (characters). + public int PreviewLength { get; init; } = 200; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client for LLM inference. + public LlmReranker(IChatClient chatClient) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + } + + /// + public override async Task ProcessAsync( + RetrievalResults results, RetrievalQuery query, CancellationToken cancellationToken = default) + { + if (results.Chunks.Count <= 2) + return results; + + var candidates = results.Chunks.Take(MaxCandidates).ToList(); + var rankedIndices = await GetRankedIndicesAsync( + query.Text, candidates, cancellationToken); + + var reranked = rankedIndices + .Where(idx => idx >= 1 && idx <= candidates.Count) + .Distinct() + .Take(MaxResults) + .Select(idx => candidates[idx - 1]) + .ToList(); + + if (reranked.Count == 0) + reranked = candidates.Take(MaxResults).ToList(); + + results.Chunks = reranked; + results.Metadata["reranked"] = true; + results.Metadata["reranked_count"] = reranked.Count; + return results; + } + + private async Task GetRankedIndicesAsync( + string query, IList candidates, CancellationToken ct) + { + var passagesBlock = string.Join("\n", + candidates.Select((c, i) => + { + var preview = c.Content.Length > PreviewLength + ? c.Content[..PreviewLength] + "..." + : c.Content; + return $"[{i + 1}] {preview}"; + })); + + var prompt = $$""" + Rank these passages by relevance to the query. Return the passage numbers + in order from most to least relevant. + + Query: {{query}} + Passages: + {{passagesBlock}} + + Return ONLY valid JSON: {"rankedIndices": [3, 1, 5, 2, 4]} + """; + + var options = new ChatOptions + { + MaxOutputTokens = 100, + ResponseFormat = ChatResponseFormat.Json + }; + + try + { + var response = await _chatClient.GetResponseAsync(prompt, options, ct); + var result = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options); + return result?.RankedIndices ?? []; + } + catch + { + return []; + } + } + + private sealed class RerankingResponse + { + [JsonPropertyName("rankedIndices")] + public int[] RankedIndices { get; set; } = []; + } +} diff --git a/MEDI/src/DataRetrieval/MultiQueryExpander.cs b/MEDI/src/DataRetrieval/MultiQueryExpander.cs new file mode 100644 index 0000000..124db6f --- /dev/null +++ b/MEDI/src/DataRetrieval/MultiQueryExpander.cs @@ -0,0 +1,75 @@ +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataRetrieval; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// Generates multiple query variants via LLM to improve recall. +/// The pipeline's RRF (Reciprocal Rank Fusion) merge deduplicates results from all variants. +/// +public class MultiQueryExpander : RetrievalQueryProcessor +{ + private readonly IChatClient _chatClient; + + /// Number of additional query variants to generate. + public int VariantCount { get; init; } = 3; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client for LLM inference. + public MultiQueryExpander(IChatClient chatClient) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + } + + /// + public override async Task ProcessAsync( + RetrievalQuery query, CancellationToken cancellationToken = default) + { + var prompt = $$""" + Given this search query, generate exactly {{VariantCount}} alternative phrasings that capture + different angles or terminology. + + Query: {{query.Text}} + + Return ONLY valid JSON: {"variants": ["query1", "query2", "query3"]} + """; + + var options = new ChatOptions + { + MaxOutputTokens = 200, + ResponseFormat = ChatResponseFormat.Json + }; + + List variants; + try + { + var response = await _chatClient.GetResponseAsync(prompt, options, cancellationToken); + var result = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options); + + variants = (result?.Variants ?? []) + .Where(v => !string.IsNullOrWhiteSpace(v) && v.Length > 10) + .Take(VariantCount) + .ToList(); + } + catch + { + variants = []; + } + + var allVariants = new List { query.Text }; + allVariants.AddRange(variants); + query.Variants = allVariants; + return query; + } + + private sealed class QueryExpansionResponse + { + [JsonPropertyName("variants")] + public string[] Variants { get; set; } = []; + } +} diff --git a/MEDI/src/DataRetrieval/RetrievalPipelineBuilder.cs b/MEDI/src/DataRetrieval/RetrievalPipelineBuilder.cs new file mode 100644 index 0000000..1333cc4 --- /dev/null +++ b/MEDI/src/DataRetrieval/RetrievalPipelineBuilder.cs @@ -0,0 +1,185 @@ +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataRetrieval; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.VectorData; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// Fluent builder for composing a with query and result processors. +/// Returned by . +/// +public class RetrievalPipelineBuilder +{ + internal List> QueryProcessorFactories { get; } = []; + internal List> ResultProcessorFactories { get; } = []; + internal IServiceCollection? Services { get; set; } + + /// Adds a query processor resolved from DI via . + public RetrievalPipelineBuilder UseQueryProcessor() where T : RetrievalQueryProcessor + { + QueryProcessorFactories.Add(sp => ActivatorUtilities.CreateInstance(sp)); + return this; + } + + /// Adds adaptive query routing that classifies queries and selects the best search paradigm. + public RetrievalPipelineBuilder UseAdaptiveRouting() + { + QueryProcessorFactories.Add(sp => + new AdaptiveRouter(sp.GetRequiredService())); + return this; + } + + /// Adds multi-query expansion with Reciprocal Rank Fusion. + public RetrievalPipelineBuilder UseQueryExpansion(Action? configure = null) + { + QueryProcessorFactories.Add(sp => + { + var options = new QueryExpansionOptions(); + configure?.Invoke(options); + return new MultiQueryExpander(sp.GetRequiredService()) + { + VariantCount = options.VariantCount + }; + }); + return this; + } + + /// Adds HyDE (Hypothetical Document Embeddings) query transformation. + public RetrievalPipelineBuilder UseHyDE() + { + QueryProcessorFactories.Add(sp => + new HydeQueryTransformer(sp.GetRequiredService())); + return this; + } + + /// Adds RAPTOR-style tree traversal for hierarchical search. + public RetrievalPipelineBuilder UseTreeSearch(Action? configure = null) + { + QueryProcessorFactories.Add(_ => + { + var options = new TreeSearchOptions(); + configure?.Invoke(options); + return new TreeSearchRetriever + { + ResultsPerLevel = options.ResultsPerLevel + }; + }); + return this; + } + + /// Adds a result processor resolved from DI via . + public RetrievalPipelineBuilder UseResultProcessor() where T : RetrievalResultProcessor + { + ResultProcessorFactories.Add(sp => ActivatorUtilities.CreateInstance(sp)); + return this; + } + + /// Adds LLM-based reranking of search results. + public RetrievalPipelineBuilder UseLlmReranking(Action? configure = null) + { + ResultProcessorFactories.Add(sp => + { + var options = new LlmRerankingOptions(); + configure?.Invoke(options); + return new LlmReranker(sp.GetRequiredService()) + { + MaxResults = options.MaxResults, + MaxCandidates = options.MaxCandidates, + PreviewLength = options.PreviewLength + }; + }); + return this; + } + + /// Adds CRAG quality gate that routes results based on relevance confidence. + public RetrievalPipelineBuilder UseCrag(Action? configure = null) + { + ResultProcessorFactories.Add(sp => + { + var options = new CragOptions(); + configure?.Invoke(options); + return new CragValidator(sp.GetRequiredService()) + { + EvaluateTopN = options.EvaluateTopN, + PreviewLength = options.PreviewLength + }; + }); + return this; + } + + /// + /// Registers an singleton that binds the configured pipeline + /// to a specific vector store collection. This is a terminal method — call it after + /// composing all processors. + /// + /// The vector store key type. + /// The vector store record type. + /// Factory to resolve the vector store collection from DI. + /// Optional function to extract text content from a record. + /// + /// + /// builder.Services.AddRetrievalPipeline() + /// .UseLlmReranking() + /// .AsRetriever<string, Article>( + /// sp => sp.GetRequiredService<VectorStoreCollection<string, Article>>(), + /// record => record.Content); + /// + /// + public void AsRetriever( + Func> collectionFactory, + Func? contentSelector = null) + where TKey : notnull + where TRecord : class + { + if (Services is null) + { + throw new InvalidOperationException( + "AsRetriever can only be called on a builder returned by AddRetrievalPipeline."); + } + + Services.AddSingleton(sp => + { + var pipeline = sp.GetRequiredService(); + var collection = collectionFactory(sp); + return pipeline.AsRetriever(collection, contentSelector); + }); + } +} + +/// Options for multi-query expansion. +public class QueryExpansionOptions +{ + /// Number of alternative query variants to generate. + public int VariantCount { get; set; } = 3; +} + +/// Options for tree search retrieval. +public class TreeSearchOptions +{ + /// Number of results per tree level. + public int ResultsPerLevel { get; set; } = 3; +} + +/// Options for LLM-based reranking. +public class LlmRerankingOptions +{ + /// Maximum results to return. + public int MaxResults { get; set; } = 5; + + /// Maximum candidates to evaluate. + public int MaxCandidates { get; set; } = 8; + + /// Maximum preview length per passage. + public int PreviewLength { get; set; } = 200; +} + +/// Options for CRAG validation. +public class CragOptions +{ + /// Number of top chunks to evaluate. + public int EvaluateTopN { get; set; } = 3; + + /// Maximum preview length per passage. + public int PreviewLength { get; set; } = 300; +} diff --git a/MEDI/src/DataRetrieval/SelfRagOrchestrator.cs b/MEDI/src/DataRetrieval/SelfRagOrchestrator.cs new file mode 100644 index 0000000..687e2d5 --- /dev/null +++ b/MEDI/src/DataRetrieval/SelfRagOrchestrator.cs @@ -0,0 +1,231 @@ +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.VectorData; +using Microsoft.Extensions.DataRetrieval; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// Self-RAG orchestrator — adaptive retrieval with self-evaluation. +/// Flow: decide retrieval → generate → self-critique → accept or retry with forced retrieval. +/// +/// +/// For compliance, medical, or legal domains where answer quality matters more than speed. +/// +public class SelfRagOrchestrator +{ + private readonly IChatClient _chatClient; + private readonly ILogger? _logger; + + /// Minimum average score (relevance + faithfulness) to accept an answer. + public double AcceptanceThreshold { get; init; } = 3.0; + + /// Maximum retry attempts before accepting best-effort. + public int MaxRetries { get; init; } = 1; + + /// + /// Initializes a new instance of the class. + /// + /// The chat client for LLM inference. + /// Optional logger factory. + public SelfRagOrchestrator(IChatClient chatClient, ILoggerFactory? loggerFactory = null) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + _logger = loggerFactory?.CreateLogger(); + } + + /// + /// Runs the Self-RAG flow for a query against a vector store collection. + /// + public async Task GenerateAsync( + VectorStoreCollection collection, + string query, + Func contentSelector, + int topK = 3, + CancellationToken cancellationToken = default) + where TKey : notnull + where TRecord : class + { + var (needsRetrieval, reasoning) = await DecideRetrievalAsync(query, cancellationToken); + _logger?.LogDebug("Self-RAG retrieval decision: {Decision} — {Reasoning}", needsRetrieval, reasoning); + + List? retrievedChunks = null; + + if (needsRetrieval) + { + retrievedChunks = await RetrieveChunksAsync(collection, query, contentSelector, topK, cancellationToken); + } + + var answer = await GenerateAnswerAsync(query, retrievedChunks, cancellationToken); + + var (relevance, faithfulness, critique) = await SelfCritiqueAsync( + query, answer, retrievedChunks, cancellationToken); + var avgScore = (relevance + faithfulness) / 2.0; + + _logger?.LogDebug("Self-RAG critique: relevance={Relevance}, faithfulness={Faithfulness}, avg={Avg}", + relevance, faithfulness, avgScore); + + if (avgScore < AcceptanceThreshold && MaxRetries > 0) + { + _logger?.LogInformation("Self-RAG score {Score} below threshold {Threshold} — retrying with forced retrieval", + avgScore, AcceptanceThreshold); + + retrievedChunks = await RetrieveChunksAsync(collection, query, contentSelector, topK, cancellationToken); + answer = await GenerateAnswerAsync(query, retrievedChunks, cancellationToken); + (relevance, faithfulness, critique) = await SelfCritiqueAsync( + query, answer, retrievedChunks, cancellationToken); + avgScore = (relevance + faithfulness) / 2.0; + } + + return new SelfRagResult + { + Answer = answer, + NeedsRetrieval = needsRetrieval, + RetrievalReasoning = reasoning, + Relevance = relevance, + Faithfulness = faithfulness, + AverageScore = avgScore, + Critique = critique, + RetrievedChunks = retrievedChunks ?? [] + }; + } + + private async Task<(bool NeedsRetrieval, string Reasoning)> DecideRetrievalAsync( + string query, CancellationToken ct) + { + var prompt = "You are an AI assistant deciding whether external retrieval is needed to answer a query.\n\n" + + "Rules:\n" + + "- If the query asks about SPECIFIC technical details, APIs, or framework features → needsRetrieval = true\n" + + "- If the query asks about GENERAL concepts that are widely known → needsRetrieval = false\n" + + "- If borderline → needsRetrieval = true\n\n" + + $"Query: \"{query}\"\n\n" + + "Return ONLY valid JSON: {\"needsRetrieval\": true, \"reasoning\": \"one sentence\"}"; + + try + { + var response = await _chatClient.GetResponseAsync(prompt, + new ChatOptions { MaxOutputTokens = 200, ResponseFormat = ChatResponseFormat.Json }, ct); + var result = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options); + return (result?.NeedsRetrieval ?? true, result?.Reasoning ?? ""); + } + catch + { + return (true, "defaulting to retrieval on failure"); + } + } + + private static async Task> RetrieveChunksAsync( + VectorStoreCollection collection, + string query, Func contentSelector, + int topK, CancellationToken ct) + where TKey : notnull + where TRecord : class + { + var results = new List(); + await foreach (var hit in collection.SearchAsync(query, top: topK, cancellationToken: ct)) + { + if (hit.Record is not null) + results.Add(contentSelector(hit.Record)); + } + return results; + } + + private async Task GenerateAnswerAsync( + string query, List? chunks, CancellationToken ct) + { + string prompt; + if (chunks is { Count: > 0 }) + { + var context = string.Join("\n---\n", chunks); + prompt = $"Using ONLY the following passages, answer the question concisely.\n\nPassages:\n{context}\n\nQuestion: {query}\n\nAnswer:"; + } + else + { + prompt = $"Answer the following question concisely from your general knowledge.\n\nQuestion: {query}\n\nAnswer:"; + } + + var response = await _chatClient.GetResponseAsync(prompt, + new ChatOptions { MaxOutputTokens = 500 }, ct); + return response.Text ?? ""; + } + + private async Task<(int Relevance, int Faithfulness, string Critique)> SelfCritiqueAsync( + string query, string answer, List? chunks, CancellationToken ct) + { + var contextSection = chunks is { Count: > 0 } + ? "Retrieved passages:\n" + string.Join("\n---\n", chunks) + "\n\n" + : ""; + + var prompt = "Score the following answer from 1 (worst) to 5 (best).\n\n" + + $"Question: \"{query}\"\n\n{contextSection}" + + $"Answer: \"{answer}\"\n\n" + + "Dimensions:\n- relevance: Does it address the question?\n- faithfulness: Is it supported by the passages?\n\n" + + "Return ONLY valid JSON: {\"relevance\": 4, \"faithfulness\": 5, \"critique\": \"summary\"}"; + + try + { + var response = await _chatClient.GetResponseAsync(prompt, + new ChatOptions { MaxOutputTokens = 200, ResponseFormat = ChatResponseFormat.Json }, ct); + var result = JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options); + + return ( + Math.Clamp(result?.Relevance ?? 3, 1, 5), + Math.Clamp(result?.Faithfulness ?? 3, 1, 5), + result?.Critique ?? ""); + } + catch + { + return (3, 3, "critique failed"); + } + } + + private sealed class RetrievalDecisionResponse + { + [JsonPropertyName("needsRetrieval")] + public bool NeedsRetrieval { get; set; } + [JsonPropertyName("reasoning")] + public string Reasoning { get; set; } = ""; + } + + private sealed class CritiqueResponse + { + [JsonPropertyName("relevance")] + public int Relevance { get; set; } + [JsonPropertyName("faithfulness")] + public int Faithfulness { get; set; } + [JsonPropertyName("critique")] + public string Critique { get; set; } = ""; + } +} + +/// Result of the Self-RAG flow. +public sealed class SelfRagResult +{ + /// The generated answer. + public required string Answer { get; init; } + + /// Whether the system decided retrieval was needed. + public bool NeedsRetrieval { get; init; } + + /// LLM reasoning for the retrieval decision. + public string RetrievalReasoning { get; init; } = ""; + + /// Relevance score (1-5). + public int Relevance { get; init; } + + /// Faithfulness score (1-5). + public int Faithfulness { get; init; } + + /// Average of relevance and faithfulness. + public double AverageScore { get; init; } + + /// Self-critique summary. + public string Critique { get; init; } = ""; + + /// The retrieved chunks used for generation. + public List RetrievedChunks { get; init; } = []; +} diff --git a/MEDI/src/DataRetrieval/SpeculativeRagOrchestrator.cs b/MEDI/src/DataRetrieval/SpeculativeRagOrchestrator.cs new file mode 100644 index 0000000..51d1a0e --- /dev/null +++ b/MEDI/src/DataRetrieval/SpeculativeRagOrchestrator.cs @@ -0,0 +1,219 @@ +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.VectorData; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// Speculative RAG orchestrator — parallel drafting with a small model + verification by a large model. +/// Leverages .NET's for true concurrent inference. +/// +/// +/// Uses keyed DI: register a "drafter" (small/fast model) and "verifier" (large/accurate model). +/// .NET's async/await + Task.WhenAll enables genuine concurrent LLM calls without GIL limitations. +/// +public class SpeculativeRagOrchestrator +{ + private readonly IChatClient _drafterClient; + private readonly IChatClient _verifierClient; + private readonly ILogger? _logger; + + /// Number of parallel drafts to generate. + public int DraftCount { get; init; } = 3; + + /// Maximum tokens for each draft response. + public int DraftMaxTokens { get; init; } = 500; + + /// Maximum tokens for the verification response. + public int VerifyMaxTokens { get; init; } = 300; + + /// + /// Initializes a new instance of the class. + /// + /// Small/fast model for parallel drafting. + /// Large/accurate model for verification. + /// Optional logger factory. + public SpeculativeRagOrchestrator( + IChatClient drafterClient, + IChatClient verifierClient, + ILoggerFactory? loggerFactory = null) + { + _drafterClient = drafterClient ?? throw new ArgumentNullException(nameof(drafterClient)); + _verifierClient = verifierClient ?? throw new ArgumentNullException(nameof(verifierClient)); + _logger = loggerFactory?.CreateLogger(); + } + + /// + /// Runs the Speculative RAG flow: retrieve → parallel draft → verify → select best. + /// + public async Task GenerateAsync( + VectorStoreCollection collection, + string query, + Func contentSelector, + int topK = 5, + CancellationToken cancellationToken = default) + where TKey : notnull + where TRecord : class + { + var retrievedChunks = new List(); + await foreach (var hit in collection.SearchAsync(query, top: topK, cancellationToken: cancellationToken)) + { + if (hit.Record is not null) + retrievedChunks.Add(contentSelector(hit.Record)); + } + + _logger?.LogDebug("Speculative RAG: retrieved {Count} chunks", retrievedChunks.Count); + + var subsets = CreateChunkSubsets(retrievedChunks, DraftCount); + + var draftTasks = subsets.Select(async (subset, idx) => + { + var prompt = $"Based on these passages, answer the question briefly:\n\n" + + string.Join("\n---\n", subset) + + $"\n\nQuestion: {query}"; + + var options = new ChatOptions { MaxOutputTokens = DraftMaxTokens }; + var response = await _drafterClient.GetResponseAsync(prompt, options, cancellationToken); + return new Draft + { + Index = idx + 1, + Text = (response.Text ?? "").Trim(), + ChunkSubset = subset + }; + }).ToArray(); + + var startTime = Environment.TickCount64; + var drafts = await Task.WhenAll(draftTasks); + var parallelMs = Environment.TickCount64 - startTime; + + _logger?.LogInformation( + "Speculative RAG: {Count} drafts generated in {Ms}ms (parallel)", + drafts.Length, parallelMs); + + var verification = await VerifyDraftsAsync(query, drafts, cancellationToken); + + var bestDraft = drafts.FirstOrDefault(d => d.Index == verification.BestDraftIndex) + ?? drafts.First(); + + return new SpeculativeRagResult + { + Answer = bestDraft.Text, + BestDraftIndex = verification.BestDraftIndex, + Confidence = verification.Confidence, + VerificationReasoning = verification.Reasoning, + Drafts = drafts.Select(d => new DraftResult + { + Index = d.Index, + Text = d.Text + }).ToList(), + ParallelDraftMs = parallelMs, + RetrievedChunks = retrievedChunks + }; + } + + private static List CreateChunkSubsets(List chunks, int subsetCount) + { + if (chunks.Count == 0) + return Enumerable.Range(0, subsetCount).Select(_ => Array.Empty()).ToList(); + + var subsets = new List(); + var chunkCount = chunks.Count; + + for (int i = 0; i < subsetCount; i++) + { + var start = (int)Math.Round((double)i * chunkCount / subsetCount); + var end = Math.Min(start + Math.Max(2, chunkCount / subsetCount + 1), chunkCount); + subsets.Add(chunks.Skip(start).Take(end - start).ToArray()); + } + + return subsets; + } + + private async Task VerifyDraftsAsync( + string query, Draft[] drafts, CancellationToken ct) + { + var draftSummaries = drafts.Select(d => + { + var text = d.Text.Length > 300 ? d.Text[..300] + "..." : d.Text; + return $"Draft {d.Index}: {text}"; + }); + + var prompt = $"Question: \"{query}\"\n\n" + + string.Join("\n\n", draftSummaries) + + "\n\nEvaluate which draft is most accurate and complete. " + + "Return ONLY valid JSON: {\"bestDraftIndex\": 1, \"confidence\": 0.85, \"reasoning\": \"one sentence\"}"; + + try + { + var response = await _verifierClient.GetResponseAsync(prompt, + new ChatOptions + { + MaxOutputTokens = VerifyMaxTokens, + ResponseFormat = ChatResponseFormat.Json + }, ct); + + return JsonSerializer.Deserialize( + response.Text ?? "{}", JsonDefaults.Options) ?? new VerificationResponse(); + } + catch + { + return new VerificationResponse { BestDraftIndex = 1, Confidence = 0.5, Reasoning = "verification failed" }; + } + } + + private sealed class Draft + { + public int Index { get; init; } + public string Text { get; init; } = ""; + public string[] ChunkSubset { get; init; } = []; + } + + private sealed class VerificationResponse + { + [JsonPropertyName("bestDraftIndex")] + public int BestDraftIndex { get; set; } = 1; + + [JsonPropertyName("confidence")] + public double Confidence { get; set; } + + [JsonPropertyName("reasoning")] + public string Reasoning { get; set; } = ""; + } +} + +/// Result of the Speculative RAG flow. +public sealed class SpeculativeRagResult +{ + /// The selected best answer. + public required string Answer { get; init; } + + /// Index (1-based) of the selected draft. + public int BestDraftIndex { get; init; } + + /// Verifier confidence (0-1). + public double Confidence { get; init; } + + /// Verifier reasoning. + public string VerificationReasoning { get; init; } = ""; + + /// All generated drafts. + public List Drafts { get; init; } = []; + + /// Time (ms) for parallel draft generation. + public long ParallelDraftMs { get; init; } + + /// Retrieved chunks used for drafting. + public List RetrievedChunks { get; init; } = []; +} + +/// Individual draft from the Speculative RAG drafting phase. +public sealed class DraftResult +{ + /// Draft index (1-based). + public int Index { get; init; } + + /// Draft text. + public string Text { get; init; } = ""; +} diff --git a/MEDI/src/DataRetrieval/TreeSearchRetriever.cs b/MEDI/src/DataRetrieval/TreeSearchRetriever.cs new file mode 100644 index 0000000..305b586 --- /dev/null +++ b/MEDI/src/DataRetrieval/TreeSearchRetriever.cs @@ -0,0 +1,70 @@ +using Microsoft.Extensions.DataRetrieval; +using Microsoft.Extensions.VectorData; + +namespace CommunityToolkit.DataRetrieval; + +/// +/// Top-down tree traversal retriever — searches root → branch → leaf levels. +/// Alternative search paradigm to flat vector search, best for exploratory or thematic queries. +/// +/// +/// Requires chunks to have a Level metadata field (set by TreeIndexProcessor). +/// Traversal: find best root summary → find best branch under it → return leaf chunks under that branch. +/// +public class TreeSearchRetriever : RetrievalQueryProcessor +{ + /// Number of results to retrieve at each level. + public int ResultsPerLevel { get; init; } = 3; + + /// + public override Task ProcessAsync( + RetrievalQuery query, CancellationToken cancellationToken = default) + { + query.Metadata["search_paradigm"] = "TreeTraversal"; + query.Metadata["results_per_level"] = ResultsPerLevel; + return Task.FromResult(query); + } + + /// + /// Performs top-down tree traversal search: root → branches → leaves. + /// + public static async Task> TraverseAsync( + VectorStoreCollection collection, + string query, + Func contentSelector, + Func levelSelector, + Func parentIdSelector, + Func idSelector, + int topK = 5, + CancellationToken cancellationToken = default) + where TKey : notnull + where TRecord : class + { + var results = new List(); + + var allResults = new List<(TRecord Record, double Score, int Level)>(); + await foreach (var hit in collection.SearchAsync(query, top: topK * 3, cancellationToken: cancellationToken)) + { + if (hit.Record is not null) + { + var level = levelSelector(hit.Record); + allResults.Add((hit.Record, hit.Score ?? 0, level)); + } + } + + var roots = allResults.Where(r => r.Level == 2).OrderByDescending(r => r.Score).Take(topK); + var branches = allResults.Where(r => r.Level == 1).OrderByDescending(r => r.Score).Take(topK); + var leaves = allResults.Where(r => r.Level == 0).OrderByDescending(r => r.Score).Take(topK); + + foreach (var item in roots.Concat(branches).Concat(leaves)) + { + var chunk = new RetrievalChunk(contentSelector(item.Record), item.Score); + chunk.Record["level"] = item.Level; + chunk.Record["id"] = idSelector(item.Record); + chunk.Record["parent_id"] = parentIdSelector(item.Record); + results.Add(chunk); + } + + return results.OrderByDescending(r => r.Score).Take(topK).ToList(); + } +} diff --git a/MEDI/src/Directory.Build.props b/MEDI/src/Directory.Build.props new file mode 100644 index 0000000..a2ca3cb --- /dev/null +++ b/MEDI/src/Directory.Build.props @@ -0,0 +1,13 @@ + + + + + + true + + + + + + + diff --git a/MEDI/src/PdfPig.DataIngestion/MetadataAwareSectionChunker.cs b/MEDI/src/PdfPig.DataIngestion/MetadataAwareSectionChunker.cs new file mode 100644 index 0000000..8f3647d --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/MetadataAwareSectionChunker.cs @@ -0,0 +1,120 @@ +using System; +using System.Collections.Generic; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.DataIngestion; +using Microsoft.Extensions.DataIngestion.Chunkers; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; + +/// +/// A chunker that wraps and propagates element metadata +/// (such as "element_type") to the resulting chunks. This works around the fact that +/// the built-in MEDI chunkers do not copy element metadata to chunks. +/// +/// +/// After the inner chunker produces chunks, this class matches each chunk's content +/// back to its source elements via substring containment and copies metadata from +/// the matching elements to the chunk. For keys with conflicting values across +/// multiple matching elements, the value from the first match wins. +/// +public sealed class MetadataAwareSectionChunker : IngestionChunker +{ + private readonly SectionChunker _inner; + + /// + /// Creates a new . + /// + /// The chunker options (tokenizer, max tokens, overlap). + public MetadataAwareSectionChunker(IngestionChunkerOptions options) + { + _inner = new SectionChunker(options); + } + + /// + public override async IAsyncEnumerable> ProcessAsync( + IngestionDocument document, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + var elementIndex = BuildElementIndex(document); + + await foreach (var chunk in _inner.ProcessAsync(document, cancellationToken).ConfigureAwait(false)) + { + ResolveMetadata(chunk, elementIndex); + yield return chunk; + } + } + + /// + /// Builds an index of element text snippets to their metadata. + /// Only includes elements that actually have metadata set. + /// + private static List<(string Text, IDictionary Metadata)> BuildElementIndex( + IngestionDocument document) + { + var index = new List<(string Text, IDictionary Metadata)>(); + + foreach (var section in document.Sections) + { + foreach (var element in section.Elements) + { + if (!element.HasMetadata) + { + continue; + } + + // Use the element's text content for matching. + // SectionChunker concatenates element GetMarkdown() values with newlines, + // so we match against GetMarkdown() for reliable substring containment. + var text = element.GetMarkdown(); + if (string.IsNullOrWhiteSpace(text)) + { + continue; + } + + index.Add((text, element.Metadata)); + } + } + + return index; + } + + /// + /// Matches a chunk's content back to source elements and copies their metadata. + /// For keys that appear in multiple matching elements, the first match's value wins. + /// + private static void ResolveMetadata( + IngestionChunk chunk, + List<(string Text, IDictionary Metadata)> elementIndex) + { + var content = chunk.Content; + if (string.IsNullOrEmpty(content)) + { + return; + } + + foreach (var (text, metadata) in elementIndex) + { + if (!content.Contains(text, StringComparison.Ordinal)) + { + continue; + } + + foreach (var kvp in metadata) + { + // Skip bounding box metadata — it's per-element, not meaningful at chunk level + if (kvp.Key.StartsWith("BoundingBox.", StringComparison.Ordinal)) + { + continue; + } + + // First match wins for each key + if (kvp.Value is not null && !chunk.Metadata.ContainsKey(kvp.Key)) + { + chunk.Metadata[kvp.Key] = kvp.Value; + } + } + } + } +} diff --git a/MEDI/src/PdfPig.DataIngestion/PageImageRenderer.cs b/MEDI/src/PdfPig.DataIngestion/PageImageRenderer.cs new file mode 100644 index 0000000..830fa91 --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/PageImageRenderer.cs @@ -0,0 +1,144 @@ +using System; +using System.Collections.Generic; +using SkiaSharp; +using UglyToad.PdfPig.Content; +using UglyToad.PdfPig.Core; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; + +/// +/// Renders PDF pages to PNG images using SkiaSharp for use with vision LLMs. +/// Unlike DocumentLayoutAnalysis.Onnx.PageImageRenderer which renders word bounding +/// boxes as black rectangles for ML model input, this renderer draws actual text characters +/// to produce human/LLM-readable page images. +/// +public static class PageImageRenderer +{ + private const float DefaultFontScale = 12f; + private const float MinFontSize = 2f; + private const float PdfDpi = 72f; + private const int PngQuality = 90; + + /// + /// Renders an entire PDF page to a PNG byte array. + /// + /// The PDF page to render. + /// Target resolution in dots per inch. Defaults to 150. + /// A PNG-encoded byte array of the rendered page. + /// is . + public static byte[] RenderPage(Page page, int dpi = 150) + { + ArgumentNullException.ThrowIfNull(page); + + float scale = dpi / PdfDpi; + int pixelWidth = Math.Max(1, (int)(page.Width * scale)); + int pixelHeight = Math.Max(1, (int)(page.Height * scale)); + + using var bitmap = new SKBitmap(pixelWidth, pixelHeight, SKColorType.Rgba8888, SKAlphaType.Premul); + using var canvas = new SKCanvas(bitmap); + canvas.Clear(SKColors.White); + + using var paint = new SKPaint { Color = SKColors.Black, IsAntialias = true }; + + var fontCache = new Dictionary(); + try + { + foreach (var letter in page.Letters) + { + float x = (float)letter.Location.X * scale; + float y = (float)(page.Height - letter.Location.Y) * scale; + float fontSize = (float)letter.PointSize * scale; + if (fontSize < MinFontSize) + { + fontSize = DefaultFontScale * scale; + } + + if (!fontCache.TryGetValue(fontSize, out var font)) + { + font = new SKFont { Size = fontSize }; + fontCache[fontSize] = font; + } + + canvas.DrawText(letter.Value, x, y, SKTextAlign.Left, font, paint); + } + } + finally + { + foreach (var font in fontCache.Values) + { + font.Dispose(); + } + } + + using var image = SKImage.FromBitmap(bitmap); + using var data = image.Encode(SKEncodedImageFormat.Png, PngQuality); + return data.ToArray(); + } + + /// + /// Renders a specific region of a PDF page to a PNG byte array. + /// Useful for sending only a table or element region to a vision LLM. + /// + /// The PDF page containing the region. + /// The bounding box of the region to render, in PDF coordinates. + /// Target resolution in dots per inch. Defaults to 150. + /// A PNG-encoded byte array of the rendered region. + /// is . + public static byte[] RenderRegion(Page page, PdfRectangle region, int dpi = 150) + { + ArgumentNullException.ThrowIfNull(page); + + float scale = dpi / PdfDpi; + int pixelWidth = Math.Max(1, (int)(region.Width * scale)); + int pixelHeight = Math.Max(1, (int)(region.Height * scale)); + + using var bitmap = new SKBitmap(pixelWidth, pixelHeight, SKColorType.Rgba8888, SKAlphaType.Premul); + using var canvas = new SKCanvas(bitmap); + canvas.Clear(SKColors.White); + + using var paint = new SKPaint { Color = SKColors.Black, IsAntialias = true }; + + var fontCache = new Dictionary(); + try + { + foreach (var letter in page.Letters) + { + double lx = letter.Location.X; + double ly = letter.Location.Y; + + if (lx < region.Left || lx > region.Right || ly < region.Bottom || ly > region.Top) + { + continue; + } + + // Translate so region's bottom-left maps to image (0,0) + float x = (float)(lx - region.Left) * scale; + float y = (float)(region.Top - ly) * scale; + float fontSize = (float)letter.PointSize * scale; + if (fontSize < MinFontSize) + { + fontSize = DefaultFontScale * scale; + } + + if (!fontCache.TryGetValue(fontSize, out var font)) + { + font = new SKFont { Size = fontSize }; + fontCache[fontSize] = font; + } + + canvas.DrawText(letter.Value, x, y, SKTextAlign.Left, font, paint); + } + } + finally + { + foreach (var font in fontCache.Values) + { + font.Dispose(); + } + } + + using var image = SKImage.FromBitmap(bitmap); + using var data = image.Encode(SKEncodedImageFormat.Png, PngQuality); + return data.ToArray(); + } +} diff --git a/MEDI/src/PdfPig.DataIngestion/PdfPig.DataIngestion.csproj b/MEDI/src/PdfPig.DataIngestion/PdfPig.DataIngestion.csproj new file mode 100644 index 0000000..26f6e6b --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/PdfPig.DataIngestion.csproj @@ -0,0 +1,29 @@ + + + 1.0.0-preview.1 + CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion + $(AssemblyName) + net10.0;net8.0 + + PdfPig PDF reader for Microsoft.Extensions.DataIngestion pipelines + PDF document reader, vision LLM enrichers (OCR, table extraction, contextual chunking), and metadata-aware chunking for MEDI ingestion pipelines using PdfPig. + true + + + + + + + + + + + + + + + + + + + diff --git a/MEDI/src/PdfPig.DataIngestion/PdfPigDataIngestionServiceCollectionExtensions.cs b/MEDI/src/PdfPig.DataIngestion/PdfPigDataIngestionServiceCollectionExtensions.cs new file mode 100644 index 0000000..78ee2fb --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/PdfPigDataIngestionServiceCollectionExtensions.cs @@ -0,0 +1,40 @@ +using System; +using CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; +using Microsoft.Extensions.DependencyInjection.Extensions; +using Microsoft.Extensions.Options; + +namespace Microsoft.Extensions.DependencyInjection; + +/// +/// Extension methods for configuring PdfPig data ingestion services. +/// +public static class PdfPigDataIngestionServiceCollectionExtensions +{ + /// + /// Adds a to the service collection. + /// + /// The service collection. + /// Optional delegate to configure reader options. + /// The service collection for chaining. + /// is . + public static IServiceCollection AddPdfPigReader( + this IServiceCollection services, + Action? configure = null) + { + ArgumentNullException.ThrowIfNull(services); + + services.AddOptions(); + + if (configure is not null) + { + services.Configure(configure); + } + + services.TryAddSingleton(); + + services.TryAddEnumerable( + ServiceDescriptor.Singleton, PdfPigReaderOptionsValidator>()); + + return services; + } +} diff --git a/MEDI/src/PdfPig.DataIngestion/PdfPigReader.cs b/MEDI/src/PdfPig.DataIngestion/PdfPigReader.cs new file mode 100644 index 0000000..a0e0a1c --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/PdfPigReader.cs @@ -0,0 +1,159 @@ +using System; +using System.IO; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.DataIngestion; +using Microsoft.Extensions.Options; +using UglyToad.PdfPig; +using UglyToad.PdfPig.DocumentLayoutAnalysis; +using UglyToad.PdfPig.DocumentLayoutAnalysis.PageSegmenter; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; + +/// +/// Reads PDF documents using PdfPig and converts them to MEDI format. +/// Supports pluggable page segmentation via and configurable +/// reading modes via . +/// +public sealed class PdfPigReader : IngestionDocumentReader +{ + private readonly IPageSegmenter segmenter; + private readonly PdfReadingMode mode; + private readonly int renderDpi; + private readonly Func? elementTypeResolver; + + /// + /// Creates a new using dependency-injected options. + /// + /// The configured reader options. + /// + /// Page segmenter for layout analysis. Defaults to if . + /// + /// + /// Optional delegate that resolves the element type label from a . + /// + public PdfPigReader( + IOptions options, + IPageSegmenter? segmenter = null, + Func? elementTypeResolver = null) + { + ArgumentNullException.ThrowIfNull(options); + var opts = options.Value; + this.segmenter = segmenter ?? DefaultPageSegmenter.Instance; + this.mode = opts.Mode; + this.renderDpi = opts.RenderDpi; + this.elementTypeResolver = elementTypeResolver; + } + + /// + /// Creates a new . + /// + /// + /// Page segmenter for layout analysis. Defaults to if . + /// Ignored when is . + /// + /// + /// Controls the text extraction strategy. Defaults to . + /// + /// + /// The DPI to use when rendering page images. Applies to + /// and modes. Defaults to 150. + /// + /// + /// Optional delegate that resolves the element type label from a . + /// When provided, the resolved type is stored in element metadata as "element_type". + /// This decouples the reader from specific segmenter implementations (e.g. ONNX). + /// + public PdfPigReader( + IPageSegmenter? segmenter = null, + PdfReadingMode mode = PdfReadingMode.TextOnly, + int renderDpi = 150, + Func? elementTypeResolver = null) + { + this.segmenter = segmenter ?? DefaultPageSegmenter.Instance; + this.mode = mode; + this.renderDpi = renderDpi; + this.elementTypeResolver = elementTypeResolver; + } + + /// + public override Task ReadAsync( + Stream source, string identifier, string mediaType, CancellationToken cancellationToken = default) + { + using var pdfDocument = PdfDocument.Open(source); + var document = new IngestionDocument(identifier); + + for (var i = 1; i <= pdfDocument.NumberOfPages; i++) + { + cancellationToken.ThrowIfCancellationRequested(); + + var page = pdfDocument.GetPage(i); + + var section = new IngestionDocumentSection + { + PageNumber = i + }; + + var renderImages = mode is PdfReadingMode.Hybrid or PdfReadingMode.VisionOnly; + + if (renderImages) + { + var imageBytes = PageImageRenderer.RenderPage(page, renderDpi); + section.Metadata["page_image"] = imageBytes; + section.Metadata["page_width"] = page.Width; + section.Metadata["page_height"] = page.Height; + } + + if (mode is not PdfReadingMode.VisionOnly) + { + var words = page.GetWords(); + var blocks = segmenter.GetBlocks(words); + + foreach (var block in blocks) + { + if (string.IsNullOrEmpty(block.Text)) + { + continue; + } + + var paragraph = new IngestionDocumentParagraph(block.Text) + { + Text = block.Text, + PageNumber = i + }; + + var bbox = block.BoundingBox; + paragraph.Metadata["BoundingBox.Left"] = bbox.Left; + paragraph.Metadata["BoundingBox.Bottom"] = bbox.Bottom; + paragraph.Metadata["BoundingBox.Right"] = bbox.Right; + paragraph.Metadata["BoundingBox.Top"] = bbox.Top; + + var elementType = elementTypeResolver?.Invoke(block); + if (elementType is not null) + { + paragraph.Metadata["element_type"] = elementType; + } + + section.Elements.Add(paragraph); + } + } + + // For scanned/image-only pages (or VisionOnly mode) with no elements, + // create a placeholder so VisionOcrEnricher can process the page image. + if (section.Elements.Count == 0 && renderImages) + { + var placeholder = new IngestionDocumentParagraph("[scanned-page]") + { + Text = string.Empty, + PageNumber = i + }; + placeholder.Metadata["placeholder"] = true; + section.Elements.Add(placeholder); + } + + document.Sections.Add(section); + } + + return Task.FromResult(document); + } +} diff --git a/MEDI/src/PdfPig.DataIngestion/PdfPigReaderOptions.cs b/MEDI/src/PdfPig.DataIngestion/PdfPigReaderOptions.cs new file mode 100644 index 0000000..5fc5973 --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/PdfPigReaderOptions.cs @@ -0,0 +1,18 @@ +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; + +/// +/// Options for configuring . +/// +public record PdfPigReaderOptions +{ + /// + /// Controls the text extraction strategy. Defaults to . + /// + public PdfReadingMode Mode { get; set; } = PdfReadingMode.TextOnly; + + /// + /// The DPI to use when rendering page images. Applies to + /// and modes. Defaults to 150. + /// + public int RenderDpi { get; set; } = 150; +} diff --git a/MEDI/src/PdfPig.DataIngestion/PdfPigReaderOptionsValidator.cs b/MEDI/src/PdfPig.DataIngestion/PdfPigReaderOptionsValidator.cs new file mode 100644 index 0000000..14b7fd3 --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/PdfPigReaderOptionsValidator.cs @@ -0,0 +1,20 @@ +using Microsoft.Extensions.Options; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; + +/// +/// Validates at startup. +/// +internal sealed class PdfPigReaderOptionsValidator : IValidateOptions +{ + /// + public ValidateOptionsResult Validate(string? name, PdfPigReaderOptions options) + { + if (options.RenderDpi <= 0) + { + return ValidateOptionsResult.Fail("RenderDpi must be a positive integer."); + } + + return ValidateOptionsResult.Success; + } +} diff --git a/MEDI/src/PdfPig.DataIngestion/PdfReadingMode.cs b/MEDI/src/PdfPig.DataIngestion/PdfReadingMode.cs new file mode 100644 index 0000000..ac9ad57 --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/PdfReadingMode.cs @@ -0,0 +1,28 @@ +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; + +/// +/// Controls how extracts content from PDF pages. +/// +public enum PdfReadingMode +{ + /// + /// Native text extraction with page segmentation. No page images are rendered. + /// This is the fastest mode with zero external dependencies. + /// Scanned pages with no extractable text produce empty sections. + /// + TextOnly, + + /// + /// Native text extraction with page segmentation and page image rendering. + /// Scanned pages that produce no text receive placeholder elements for + /// downstream VLM-based OCR via VisionOcrEnricher. + /// + Hybrid, + + /// + /// Skip native text extraction entirely. Every page is rendered as an image + /// with a placeholder element. Requires a downstream VisionOcrEnricher + /// with a vision-capable LLM to populate the text content. + /// + VisionOnly +} diff --git a/MEDI/src/PdfPig.DataIngestion/Processors/ContextualChunkEnricher.cs b/MEDI/src/PdfPig.DataIngestion/Processors/ContextualChunkEnricher.cs new file mode 100644 index 0000000..df27e7a --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/Processors/ContextualChunkEnricher.cs @@ -0,0 +1,73 @@ +using System; +using System.Collections.Generic; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.Processors; + +/// +/// Enriches text chunks with contextual summaries for improved RAG retrieval. +/// Uses to generate a brief summary of each chunk. +/// +public sealed class ContextualChunkEnricher : IngestionChunkProcessor +{ + private readonly IChatClient chatClient; + + /// + /// Creates a new . + /// + /// The chat client used to generate contextual summaries. + /// is . + public ContextualChunkEnricher(IChatClient chatClient) + { + this.chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + } + + /// + public override async IAsyncEnumerable> ProcessAsync( + IAsyncEnumerable> chunks, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + await foreach (var chunk in chunks.WithCancellation(cancellationToken).ConfigureAwait(false)) + { + var prompt = GetPromptForChunk(chunk); + + var messages = new[] + { + new ChatMessage(ChatRole.User, prompt) + }; + + var response = await chatClient.GetResponseAsync( + messages, + cancellationToken: cancellationToken).ConfigureAwait(false); + + chunk.Metadata["contextual_summary"] = response.Text; + + yield return chunk; + } + } + + private static string GetPromptForChunk(IngestionChunk chunk) + { + var elementType = chunk.Metadata.TryGetValue("element_type", out var t) ? t as string : null; + + var instruction = elementType switch + { + "table" => + "Summarize this table for search retrieval. Describe what data it contains, " + + "including key metrics, column headers, and notable values. " + + "Output only the summary sentence, nothing else.", + "picture" or "caption" => + "Summarize what this figure or image refers to for search retrieval. " + + "Output only the summary sentence, nothing else.", + _ => + "Provide a single concise sentence summarizing the following text for use " + + "in search retrieval. Output only the summary sentence, nothing else." + }; + + return instruction + "\n\n" + chunk.Content; + } +} diff --git a/MEDI/src/PdfPig.DataIngestion/Processors/VisionOcrEnricher.cs b/MEDI/src/PdfPig.DataIngestion/Processors/VisionOcrEnricher.cs new file mode 100644 index 0000000..d05411c --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/Processors/VisionOcrEnricher.cs @@ -0,0 +1,123 @@ +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.Processors; + +/// +/// Enriches document elements that have minimal or no text content by performing +/// vision LLM-based OCR, such as for scanned pages or image-heavy regions. +/// +public sealed class VisionOcrEnricher : IngestionDocumentProcessor +{ + private readonly IChatClient chatClient; + + /// + /// Creates a new . + /// + /// The chat client used to interact with a vision-capable LLM. + /// is . + public VisionOcrEnricher(IChatClient chatClient) + { + this.chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + } + + /// + public override async Task ProcessAsync( + IngestionDocument document, CancellationToken cancellationToken = default) + { + // Iterate by index so we can replace elements in-place. + // This is necessary because the MEDI chunker uses GetMarkdown() (which returns + // the immutable constructor parameter) rather than element.Text. Setting Text alone + // would leave stale placeholder content in GetMarkdown(), causing the chunker to + // produce chunks with placeholder text instead of the OCR'd content. + foreach (var section in document.Sections) + { + for (int i = 0; i < section.Elements.Count; i++) + { + var element = section.Elements[i]; + cancellationToken.ThrowIfCancellationRequested(); + + if (!string.IsNullOrWhiteSpace(element.Text)) + { + continue; + } + + // Try to get page image for vision-based OCR + byte[]? imageBytes = null; + if (element.PageNumber is int pageNumber) + { + foreach (var s in document.Sections) + { + if (s.PageNumber == pageNumber && + s.HasMetadata && + s.Metadata.TryGetValue("page_image", out var imageObj)) + { + imageBytes = imageObj as byte[]; + break; + } + } + } + + ChatMessage[] messages; + if (imageBytes is not null) + { + // Vision approach: send actual page image + messages = new[] + { + new ChatMessage(ChatRole.System, + "You are a precise OCR engine. Extract all visible text from the provided image exactly as it appears. " + + "Preserve line breaks and formatting. Output only the extracted text, no commentary."), + new ChatMessage(ChatRole.User, (IList)new AIContent[] + { + new DataContent(imageBytes, "image/png"), + new TextContent("Extract all text from this image.") + }) + }; + } + else + { + // Fallback: text-based approach when no image available + messages = new[] + { + new ChatMessage(ChatRole.User, + "You are an OCR engine. Extract all visible text from the following content. " + + "Return only the extracted text, preserving the original layout as much as possible.\n\n" + + (element.Text ?? string.Empty)) + }; + } + + var response = await chatClient.GetResponseAsync( + messages, + cancellationToken: cancellationToken).ConfigureAwait(false); + + if (!string.IsNullOrWhiteSpace(response.Text)) + { + // Replace the element with a new one whose markdown (constructor param) + // contains the OCR text, so downstream chunkers see it via GetMarkdown(). + var replacement = new IngestionDocumentParagraph(response.Text) + { + Text = response.Text, + PageNumber = element.PageNumber + }; + + if (element.HasMetadata) + { + foreach (var kv in element.Metadata) + { + replacement.Metadata[kv.Key] = kv.Value; + } + } + + replacement.Metadata["ocr_source"] = "vision_llm"; + section.Elements[i] = replacement; + } + } + } + + return document; + } +} diff --git a/MEDI/src/PdfPig.DataIngestion/Processors/VisionTableEnricher.cs b/MEDI/src/PdfPig.DataIngestion/Processors/VisionTableEnricher.cs new file mode 100644 index 0000000..e5e6b2e --- /dev/null +++ b/MEDI/src/PdfPig.DataIngestion/Processors/VisionTableEnricher.cs @@ -0,0 +1,95 @@ +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.Processors; + +/// +/// Enriches table elements in a document by sending their content +/// to a vision-capable LLM via to extract markdown table content. +/// +public sealed class VisionTableEnricher : IngestionDocumentProcessor +{ + private readonly IChatClient chatClient; + + /// + /// Creates a new . + /// + /// The chat client used to interact with a vision-capable LLM. + /// is . + public VisionTableEnricher(IChatClient chatClient) + { + this.chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + } + + /// + public override async Task ProcessAsync( + IngestionDocument document, CancellationToken cancellationToken = default) + { + foreach (var element in document.EnumerateContent()) + { + if (element is not IngestionDocumentTable table) + { + continue; + } + + cancellationToken.ThrowIfCancellationRequested(); + + // Try to get page image for vision-based extraction + byte[]? imageBytes = null; + if (element.PageNumber is int pageNumber) + { + foreach (var section in document.Sections) + { + if (section.PageNumber == pageNumber && + section.HasMetadata && + section.Metadata.TryGetValue("page_image", out var imageObj)) + { + imageBytes = imageObj as byte[]; + break; + } + } + } + + ChatMessage[] messages; + if (imageBytes is not null) + { + // Vision approach: send actual page image + messages = new[] + { + new ChatMessage(ChatRole.System, + "You are a table extraction engine. Extract the table from the image as a markdown table. " + + "Use | as column separators. Include a header separator (| --- | --- |). " + + "Output ONLY the markdown table, nothing else."), + new ChatMessage(ChatRole.User, (IList)new AIContent[] + { + new DataContent(imageBytes, "image/png"), + new TextContent("Extract the table from this image as a markdown table.") + }) + }; + } + else + { + // Fallback: text-based extraction + messages = new[] + { + new ChatMessage(ChatRole.User, + "Extract the following table content into a well-formatted markdown table. " + + "Only output the markdown table, no other text.\n\n" + + (table.Text ?? string.Empty)) + }; + } + + var response = await chatClient.GetResponseAsync( + messages, + cancellationToken: cancellationToken).ConfigureAwait(false); + + table.Metadata["enriched_markdown_table"] = response.Text; + } + + return document; + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/AnnotatedTextBlock.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/AnnotatedTextBlock.cs new file mode 100644 index 0000000..48063fb --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/AnnotatedTextBlock.cs @@ -0,0 +1,43 @@ +using System.Collections.Generic; +using UglyToad.PdfPig.Content; +using UglyToad.PdfPig.DocumentLayoutAnalysis; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; + +/// +/// A that carries the layout detection label and confidence +/// from an ONNX model. Downstream code that only knows about +/// sees a regular block; callers that know about this subclass can read +/// and . +/// +public sealed class AnnotatedTextBlock : TextBlock +{ + /// + /// The element type label detected by the ONNX layout model + /// (e.g. "table", "picture", "section_header", "text"). + /// + public string Label { get; } + + /// + /// The model's confidence score for this detection, between 0 and 1. + /// + public float Confidence { get; } + + /// + /// Create a new . + /// + /// The text lines in this block. + /// The layout detection label. + /// The detection confidence score. + /// The separator used between lines. + public AnnotatedTextBlock( + IReadOnlyList lines, + string label, + float confidence, + string separator = "\n") + : base(lines, separator) + { + Label = label; + Confidence = confidence; + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/DetectionPostprocessing.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/DetectionPostprocessing.cs new file mode 100644 index 0000000..3c2c876 --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/DetectionPostprocessing.cs @@ -0,0 +1,213 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using UglyToad.PdfPig.Core; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; + +/// +/// Static utility methods for postprocessing ONNX detection model outputs. +/// +public static class DetectionPostprocessing +{ + /// + /// Convert center-x, center-y, width, height (image coordinates, top-left origin) + /// to a (bottom-left origin, Y-up). + /// + /// Center X in image coordinates. + /// Center Y in image coordinates. + /// Width. + /// Height. + /// Page width for coordinate mapping. + /// Page height for coordinate mapping. + /// A PDF rectangle in PDF coordinate space. + public static PdfRectangle CxCyWhToRect(float cx, float cy, float w, float h, float pageWidth, float pageHeight) + { + float x1 = cx - w / 2f; + float y1 = cy - h / 2f; + float x2 = cx + w / 2f; + float y2 = cy + h / 2f; + + return XyxyToRect(x1, y1, x2, y2, pageWidth, pageHeight); + } + + /// + /// Convert x1, y1, x2, y2 (image coordinates, top-left origin, Y-down) + /// to a (bottom-left origin, Y-up). + /// + /// Left X in image coordinates. + /// Top Y in image coordinates. + /// Right X in image coordinates. + /// Bottom Y in image coordinates. + /// Page width for coordinate mapping. + /// Page height for coordinate mapping. + /// A PDF rectangle in PDF coordinate space. + public static PdfRectangle XyxyToRect(float x1, float y1, float x2, float y2, float pageWidth, float pageHeight) + { + // Image Y-axis is top-down; PDF Y-axis is bottom-up. + // Flip Y: pdfY = pageHeight - imageY + double pdfLeft = x1; + double pdfRight = x2; + double pdfBottom = pageHeight - y2; + double pdfTop = pageHeight - y1; + + return new PdfRectangle(pdfLeft, pdfBottom, pdfRight, pdfTop); + } + + /// + /// Apply Non-Maximum Suppression to remove overlapping detections. + /// Detections are processed in descending confidence order; a detection + /// is suppressed if its IoU with any already-kept detection exceeds the threshold. + /// + /// Input detections. + /// IoU threshold above which a detection is suppressed. + /// Filtered detections after NMS. + public static IReadOnlyList ApplyNms(IReadOnlyList detections, float iouThreshold = 0.45f) + { + ArgumentNullException.ThrowIfNull(detections); + + if (detections.Count <= 1) + { + return detections; + } + + var sorted = detections.OrderByDescending(d => d.Confidence).ToList(); + var kept = new List(); + var suppressed = new bool[sorted.Count]; + + for (int i = 0; i < sorted.Count; i++) + { + if (suppressed[i]) + { + continue; + } + + kept.Add(sorted[i]); + + for (int j = i + 1; j < sorted.Count; j++) + { + if (suppressed[j]) + { + continue; + } + + if (ComputeIoU(sorted[i].BoundingBox, sorted[j].BoundingBox) > iouThreshold) + { + suppressed[j] = true; + } + } + } + + return kept; + } + + /// + /// Scale detection bounding boxes from model coordinate space to PDF page coordinate space. + /// Handles optional letterbox padding removal. + /// + /// Detections in model coordinates. + /// Model input width. + /// Model input height. + /// Target page width (PDF units). + /// Target page height (PDF units). + /// Scale factor used during letterboxing, or null if not letterboxed. + /// Horizontal padding offset from letterboxing. + /// Vertical padding offset from letterboxing. + /// Detections with bounding boxes in page coordinate space. + public static IReadOnlyList ScaleToPage( + IReadOnlyList detections, + int modelWidth, + int modelHeight, + double pageWidth, + double pageHeight, + float? letterboxScale = null, + int padX = 0, + int padY = 0) + { + ArgumentNullException.ThrowIfNull(detections); + + var result = new List(detections.Count); + + foreach (var det in detections) + { + var box = det.BoundingBox; + + double left = box.Left; + double right = box.Right; + double bottom = box.Bottom; + double top = box.Top; + + if (letterboxScale.HasValue) + { + float s = letterboxScale.Value; + left = (left - padX) / s; + right = (right - padX) / s; + bottom = (bottom - padY) / s; + top = (top - padY) / s; + } + + // Scale from image pixel space to page space + double scaleX = pageWidth / modelWidth; + double scaleY = pageHeight / modelHeight; + + if (letterboxScale.HasValue) + { + // When letterboxed, coordinates were already unpadded/unscaled to original image size. + // Now scale from original image size to page size. + // The original image size = modelWidth / letterboxScale (approximately) + // but we already divided by letterboxScale, so coords are in original image space. + // We still need to map from original image dims to page dims. + // Since the original image was rendered from the page, they share the same aspect ratio. + scaleX = pageWidth / (modelWidth / letterboxScale.Value - padX * 2.0 / letterboxScale.Value); + scaleY = pageHeight / (modelHeight / letterboxScale.Value - padY * 2.0 / letterboxScale.Value); + } + + left *= scaleX; + right *= scaleX; + bottom *= scaleY; + top *= scaleY; + + // Clamp to page bounds + left = Math.Max(0, Math.Min(left, pageWidth)); + right = Math.Max(0, Math.Min(right, pageWidth)); + bottom = Math.Max(0, Math.Min(bottom, pageHeight)); + top = Math.Max(0, Math.Min(top, pageHeight)); + + var newBox = new PdfRectangle(left, bottom, right, top); + result.Add(det with { BoundingBox = newBox }); + } + + return result; + } + + /// + /// Compute the Intersection over Union (IoU) of two rectangles. + /// + /// First rectangle. + /// Second rectangle. + /// IoU value in [0, 1]. + public static float ComputeIoU(PdfRectangle a, PdfRectangle b) + { + double interLeft = Math.Max(a.Left, b.Left); + double interRight = Math.Min(a.Right, b.Right); + double interBottom = Math.Max(a.Bottom, b.Bottom); + double interTop = Math.Min(a.Top, b.Top); + + if (interLeft >= interRight || interBottom >= interTop) + { + return 0f; + } + + double interArea = (interRight - interLeft) * (interTop - interBottom); + double areaA = a.Width * a.Height; + double areaB = b.Width * b.Height; + double unionArea = areaA + areaB - interArea; + + if (unionArea <= 0) + { + return 0f; + } + + return (float)(interArea / unionArea); + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/ILayoutDetectionModel.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/ILayoutDetectionModel.cs new file mode 100644 index 0000000..f212885 --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/ILayoutDetectionModel.cs @@ -0,0 +1,41 @@ +using Microsoft.ML.OnnxRuntime; +using SkiaSharp; +using System; +using System.Collections.Generic; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; + +/// +/// Strategy interface encapsulating model-specific preprocessing and postprocessing +/// for ONNX-based layout detection models. +/// +public interface ILayoutDetectionModel : IDisposable +{ + /// + /// Convert a page image to the model's expected input tensor(s). + /// + /// The rendered page image. + /// Original page image width in pixels. + /// Original page image height in pixels. + /// Named ONNX values ready for inference. + IReadOnlyList Preprocess(SKBitmap pageImage, int originalWidth, int originalHeight); + + /// + /// Parse model output into layout detections. + /// + /// The raw ONNX inference results. + /// Original page image width in pixels. + /// Original page image height in pixels. + /// Detected layout elements. + IReadOnlyList Postprocess(IDisposableReadOnlyCollection results, int originalWidth, int originalHeight); + + /// + /// Mapping from class ID to human-readable label name. + /// + IReadOnlyDictionary LabelMapping { get; } + + /// + /// Path to the ONNX model file on disk. + /// + string ModelPath { get; } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/ImagePreprocessing.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/ImagePreprocessing.cs new file mode 100644 index 0000000..68bd16b --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/ImagePreprocessing.cs @@ -0,0 +1,174 @@ +using Microsoft.ML.OnnxRuntime.Tensors; +using SkiaSharp; +using System; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; + +/// +/// Static utility methods for image preprocessing before ONNX model inference. +/// +public static class ImagePreprocessing +{ + /// + /// Resize an image maintaining aspect ratio and pad to the target size (letterboxing). + /// + /// Source image. + /// Target width. + /// Target height. + /// Color to use for padding. + /// The scale factor applied to the image. + /// Horizontal padding offset in pixels. + /// Vertical padding offset in pixels. + /// A new letterboxed bitmap. + public static SKBitmap Letterbox(SKBitmap image, int targetW, int targetH, SKColor padColor, out float scale, out int padX, out int padY) + { + ArgumentNullException.ThrowIfNull(image); + + float scaleX = (float)targetW / image.Width; + float scaleY = (float)targetH / image.Height; + scale = Math.Min(scaleX, scaleY); + + int newW = (int)(image.Width * scale); + int newH = (int)(image.Height * scale); + padX = (targetW - newW) / 2; + padY = (targetH - newH) / 2; + + var info = new SKImageInfo(targetW, targetH, SKColorType.Rgba8888, SKAlphaType.Premul); + var result = new SKBitmap(info); + + using var canvas = new SKCanvas(result); + canvas.Clear(padColor); + + var destRect = SKRect.Create(padX, padY, newW, newH); + using var skImage = SKImage.FromBitmap(image); + canvas.DrawImage(skImage, destRect, new SKSamplingOptions(SKCubicResampler.Mitchell)); + + return result; + } + + /// + /// Resize an image to exact dimensions, stretching if necessary. + /// + /// Source image. + /// Target width. + /// Target height. + /// A new resized bitmap. + public static SKBitmap ResizeExact(SKBitmap image, int targetW, int targetH) + { + ArgumentNullException.ThrowIfNull(image); + + var info = new SKImageInfo(targetW, targetH, SKColorType.Rgba8888, SKAlphaType.Premul); + var result = new SKBitmap(info); + + using var canvas = new SKCanvas(result); + var destRect = SKRect.Create(0, 0, targetW, targetH); + using var skImage = SKImage.FromBitmap(image); + canvas.DrawImage(skImage, destRect, new SKSamplingOptions(SKCubicResampler.Mitchell)); + + return result; + } + + /// + /// Extract a CHW byte tensor [1, 3, H, W] from an image. + /// Channel order is RGB. + /// + /// Source image (must be RGBA8888 or compatible). + /// A dense tensor with shape [1, 3, H, W]. + public static DenseTensor ToChwUint8(SKBitmap image) + { + ArgumentNullException.ThrowIfNull(image); + + int w = image.Width; + int h = image.Height; + var tensor = new DenseTensor([1, 3, h, w]); + + for (int y = 0; y < h; y++) + { + for (int x = 0; x < w; x++) + { + var pixel = image.GetPixel(x, y); + tensor[0, 0, y, x] = pixel.Red; + tensor[0, 1, y, x] = pixel.Green; + tensor[0, 2, y, x] = pixel.Blue; + } + } + + return tensor; + } + + /// + /// Extract a CHW float tensor [1, 3, H, W] from an image. + /// Values are normalized to [0, 1] range. Channel order is RGB. + /// + /// Source image (must be RGBA8888 or compatible). + /// A dense tensor with shape [1, 3, H, W]. + public static DenseTensor ToChwFloat(SKBitmap image) + { + ArgumentNullException.ThrowIfNull(image); + + int w = image.Width; + int h = image.Height; + var tensor = new DenseTensor([1, 3, h, w]); + + for (int y = 0; y < h; y++) + { + for (int x = 0; x < w; x++) + { + var pixel = image.GetPixel(x, y); + tensor[0, 0, y, x] = pixel.Red / 255f; + tensor[0, 1, y, x] = pixel.Green / 255f; + tensor[0, 2, y, x] = pixel.Blue / 255f; + } + } + + return tensor; + } + + /// + /// Apply ImageNet normalization in-place: mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]. + /// + /// The CHW float buffer (length must be 3 * width * height). + /// Image width. + /// Image height. + public static void NormalizeImageNet(Span chw, int width, int height) + { + ReadOnlySpan mean = [0.485f, 0.456f, 0.406f]; + ReadOnlySpan std = [0.229f, 0.224f, 0.225f]; + Normalize(chw, width, height, mean, std); + } + + /// + /// Apply custom per-channel normalization in-place: (value - mean) / std. + /// + /// The CHW float buffer (length must be 3 * width * height). + /// Image width. + /// Image height. + /// Per-channel mean values (length 3). + /// Per-channel standard deviation values (length 3). + public static void Normalize(Span chw, int width, int height, ReadOnlySpan mean, ReadOnlySpan std) + { + if (mean.Length != 3) + { + throw new ArgumentException("Mean must have exactly 3 elements.", nameof(mean)); + } + + if (std.Length != 3) + { + throw new ArgumentException("Std must have exactly 3 elements.", nameof(std)); + } + + int channelSize = width * height; + + for (int c = 0; c < 3; c++) + { + int offset = c * channelSize; + float m = mean[c]; + float s = std[c]; + + for (int i = 0; i < channelSize; i++) + { + chw[offset + i] = (chw[offset + i] - m) / s; + } + } + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/LayoutDetection.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/LayoutDetection.cs new file mode 100644 index 0000000..376f581 --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/LayoutDetection.cs @@ -0,0 +1,12 @@ +using UglyToad.PdfPig.Core; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; + +/// +/// A single detected layout element from an ONNX model. +/// +public record LayoutDetection( + PdfRectangle BoundingBox, + string Label, + int ClassId, + float Confidence); diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/Models/ConfigurableLayoutModel.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/Models/ConfigurableLayoutModel.cs new file mode 100644 index 0000000..eaa58ab --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/Models/ConfigurableLayoutModel.cs @@ -0,0 +1,468 @@ +using Microsoft.ML.OnnxRuntime; +using Microsoft.ML.OnnxRuntime.Tensors; +using SkiaSharp; +using System; +using System.Collections.Generic; +using System.Linq; +using UglyToad.PdfPig.Core; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.Models; + +/// +/// A configuration-driven layout detection model implementation. +/// Uses to determine preprocessing and postprocessing behavior. +/// +/// +/// +/// Thread safety: This type is not thread-safe. The method +/// stores letterbox state in instance fields that reads. +/// Concurrent Preprocess/Postprocess call pairs will produce incorrect results. +/// Use a separate instance per thread, or synchronize access externally. +/// +/// +public sealed class ConfigurableLayoutModel : ILayoutDetectionModel +{ + private readonly LayoutModelOptions _options; + private bool _disposed; + + // Letterbox state preserved between pre- and post-processing + private float _letterboxScale; + private int _padX; + private int _padY; + private bool _wasLetterboxed; + + /// + /// Create a new . + /// + /// Path to the ONNX model file. + /// Model configuration options. + public ConfigurableLayoutModel(string modelPath, LayoutModelOptions options) + { + ModelPath = modelPath ?? throw new ArgumentNullException(nameof(modelPath)); + _options = options ?? throw new ArgumentNullException(nameof(options)); + } + + /// + public string ModelPath { get; } + + /// + public IReadOnlyDictionary LabelMapping => + _options.ClassLabels ?? new Dictionary(); + + /// + public IReadOnlyList Preprocess(SKBitmap pageImage, int originalWidth, int originalHeight) + { + ArgumentNullException.ThrowIfNull(pageImage); + + SKBitmap resized; + _wasLetterboxed = false; + + switch (_options.Resize) + { + case ResizeMode.Letterbox: + resized = ImagePreprocessing.Letterbox( + pageImage, + _options.InputWidth, + _options.InputHeight, + SKColors.Gray, + out _letterboxScale, + out _padX, + out _padY); + _wasLetterboxed = true; + break; + + case ResizeMode.AspectPreserve: + float scaleX = (float)_options.InputWidth / pageImage.Width; + float scaleY = (float)_options.InputHeight / pageImage.Height; + float scale = Math.Min(scaleX, scaleY); + int newW = (int)(pageImage.Width * scale); + int newH = (int)(pageImage.Height * scale); + resized = ImagePreprocessing.ResizeExact(pageImage, newW, newH); + break; + + default: // Exact + resized = ImagePreprocessing.ResizeExact(pageImage, _options.InputWidth, _options.InputHeight); + break; + } + + try + { + return _options.PixelFormat switch + { + PixelFormat.Float32Chw => CreateFloatInput(resized), + _ => CreateUint8Input(resized) + }; + } + finally + { + resized.Dispose(); + } + } + + private static List CreateUint8Input(SKBitmap image) + { + var tensor = ImagePreprocessing.ToChwUint8(image); + return new List + { + NamedOnnxValue.CreateFromTensor("images", tensor) + }; + } + + private List CreateFloatInput(SKBitmap image) + { + var tensor = ImagePreprocessing.ToChwFloat(image); + + if (_options.Normalization is not null) + { + var buffer = tensor.Buffer.Span; + ImagePreprocessing.Normalize( + buffer, + image.Width, + image.Height, + _options.Normalization.Mean.AsSpan(), + _options.Normalization.Std.AsSpan()); + } + + return new List + { + NamedOnnxValue.CreateFromTensor("images", tensor) + }; + } + + /// + public IReadOnlyList Postprocess( + IDisposableReadOnlyCollection results, + int originalWidth, + int originalHeight) + { + ArgumentNullException.ThrowIfNull(results); + + var detections = ParseOutputTensor(results, originalWidth, originalHeight); + + // Apply confidence filter + detections = detections + .Where(d => d.Confidence >= _options.ConfidenceThreshold) + .ToList(); + + // Apply NMS if required + if (_options.RequiresNms && detections.Count > 1) + { + detections = DetectionPostprocessing.ApplyNms(detections, _options.NmsIouThreshold).ToList(); + } + + // Scale to page coordinates if letterboxed + if (_wasLetterboxed) + { + detections = DetectionPostprocessing.ScaleToPage( + detections, + _options.InputWidth, + _options.InputHeight, + originalWidth, + originalHeight, + _letterboxScale, + _padX, + _padY).ToList(); + } + + return detections; + } + + private List ParseOutputTensor( + IDisposableReadOnlyCollection results, + int originalWidth, + int originalHeight) + { + var detections = new List(); + + // Try to find a single output tensor (common YOLO-style: [1, num_classes+4, num_detections]) + DisposableNamedOnnxValue? outputValue = null; + DisposableNamedOnnxValue? labelsValue = null; + DisposableNamedOnnxValue? boxesValue = null; + DisposableNamedOnnxValue? scoresValue = null; + + foreach (var result in results) + { + var name = result.Name; + if (name == "labels" || name == "pred_labels") + { + labelsValue = result; + } + else if (name == "boxes" || name == "pred_boxes") + { + boxesValue = result; + } + else if (name == "scores" || name == "pred_scores") + { + scoresValue = result; + } + else if (name == "output" || name == "output0") + { + outputValue = result; + } + } + + // If we have separate labels/boxes/scores outputs + if (labelsValue is not null && boxesValue is not null && scoresValue is not null) + { + return ParseSeparateOutputs(labelsValue, boxesValue, scoresValue, originalWidth, originalHeight); + } + + // Otherwise try a combined output tensor + if (outputValue?.Value is Tensor combinedTensor) + { + return ParseCombinedOutput(combinedTensor, originalWidth, originalHeight); + } + + // Fallback: try the first float tensor output + foreach (var result in results) + { + if (result.Value is Tensor tensor && tensor.Dimensions.Length >= 2) + { + return ParseCombinedOutput(tensor, originalWidth, originalHeight); + } + } + + return detections; + } + + private List ParseSeparateOutputs( + DisposableNamedOnnxValue labelsValue, + DisposableNamedOnnxValue boxesValue, + DisposableNamedOnnxValue scoresValue, + int originalWidth, + int originalHeight) + { + var detections = new List(); + + int[]? labels = ExtractIntArray(labelsValue); + float[]? scores = ExtractFloatArray(scoresValue); + float[]? boxes = ExtractBoxArray(boxesValue); + + if (labels is null || scores is null || boxes is null) + { + return detections; + } + + int count = labels.Length; + + for (int i = 0; i < count; i++) + { + int classId = labels[i]; + float confidence = scores[i]; + + float bx0 = boxes[i * 4 + 0]; + float bx1 = boxes[i * 4 + 1]; + float bx2 = boxes[i * 4 + 2]; + float bx3 = boxes[i * 4 + 3]; + + var rect = ConvertBbox(bx0, bx1, bx2, bx3, originalWidth, originalHeight); + string label = GetLabelName(classId); + + detections.Add(new LayoutDetection(rect, label, classId, confidence)); + } + + return detections; + } + + private List ParseCombinedOutput( + Tensor tensor, + int originalWidth, + int originalHeight) + { + var detections = new List(); + var dims = tensor.Dimensions; + + if (dims.Length == 3) + { + // Shape: [1, numDetections, 4+numClasses] or [1, 4+numClasses, numDetections] + int dim1 = dims[1]; + int dim2 = dims[2]; + + bool transposed = dim2 > dim1 && dim1 > 5; + + int numDetections; + int numValues; + + if (transposed) + { + // [1, 4+numClasses, numDetections] — need to transpose + numDetections = dim2; + numValues = dim1; + } + else + { + // [1, numDetections, 4+numClasses] + numDetections = dim1; + numValues = dim2; + } + + int numClasses = numValues - 4; + if (numClasses <= 0) + { + return detections; + } + + for (int i = 0; i < numDetections; i++) + { + float bx0, bx1, bx2, bx3; + int bestClass = 0; + float bestScore = float.MinValue; + + if (transposed) + { + bx0 = tensor[0, 0, i]; + bx1 = tensor[0, 1, i]; + bx2 = tensor[0, 2, i]; + bx3 = tensor[0, 3, i]; + + for (int c = 0; c < numClasses; c++) + { + float score = tensor[0, 4 + c, i]; + if (score > bestScore) + { + bestScore = score; + bestClass = c; + } + } + } + else + { + bx0 = tensor[0, i, 0]; + bx1 = tensor[0, i, 1]; + bx2 = tensor[0, i, 2]; + bx3 = tensor[0, i, 3]; + + for (int c = 0; c < numClasses; c++) + { + float score = tensor[0, i, 4 + c]; + if (score > bestScore) + { + bestScore = score; + bestClass = c; + } + } + } + + var rect = ConvertBbox(bx0, bx1, bx2, bx3, originalWidth, originalHeight); + string label = GetLabelName(bestClass); + + detections.Add(new LayoutDetection(rect, label, bestClass, bestScore)); + } + } + + return detections; + } + + private PdfRectangle ConvertBbox(float v0, float v1, float v2, float v3, int imageWidth, int imageHeight) + { + return _options.OutputBboxFormat switch + { + BboxFormat.CxCyWh => DetectionPostprocessing.CxCyWhToRect(v0, v1, v2, v3, imageWidth, imageHeight), + BboxFormat.Xywh => DetectionPostprocessing.XyxyToRect(v0, v1, v0 + v2, v1 + v3, imageWidth, imageHeight), + _ => DetectionPostprocessing.XyxyToRect(v0, v1, v2, v3, imageWidth, imageHeight) + }; + } + + private string GetLabelName(int classId) + { + if (_options.ClassLabels is not null && _options.ClassLabels.TryGetValue(classId, out var name)) + { + return name; + } + + return $"class_{classId}"; + } + + private static int[]? ExtractIntArray(DisposableNamedOnnxValue value) + { + if (value.Value is Tensor longTensor) + { + var dims = longTensor.Dimensions; + int count = dims.Length > 1 ? dims[1] : dims[0]; + var arr = new int[count]; + + for (int i = 0; i < count; i++) + { + arr[i] = dims.Length > 1 ? (int)longTensor[0, i] : (int)longTensor[i]; + } + + return arr; + } + + if (value.Value is Tensor intTensor) + { + var dims = intTensor.Dimensions; + int count = dims.Length > 1 ? dims[1] : dims[0]; + var arr = new int[count]; + + for (int i = 0; i < count; i++) + { + arr[i] = dims.Length > 1 ? intTensor[0, i] : intTensor[i]; + } + + return arr; + } + + return null; + } + + private static float[]? ExtractFloatArray(DisposableNamedOnnxValue value) + { + if (value.Value is Tensor floatTensor) + { + var dims = floatTensor.Dimensions; + int count = dims.Length > 1 ? dims[1] : dims[0]; + var arr = new float[count]; + + for (int i = 0; i < count; i++) + { + arr[i] = dims.Length > 1 ? floatTensor[0, i] : floatTensor[i]; + } + + return arr; + } + + return null; + } + + private static float[]? ExtractBoxArray(DisposableNamedOnnxValue value) + { + if (value.Value is Tensor floatTensor) + { + var dims = floatTensor.Dimensions; + int count = dims.Length > 1 ? dims[1] : dims[0]; + var arr = new float[count * 4]; + + for (int i = 0; i < count; i++) + { + if (dims.Length == 3) + { + arr[i * 4 + 0] = floatTensor[0, i, 0]; + arr[i * 4 + 1] = floatTensor[0, i, 1]; + arr[i * 4 + 2] = floatTensor[0, i, 2]; + arr[i * 4 + 3] = floatTensor[0, i, 3]; + } + else if (dims.Length == 2) + { + arr[i * 4 + 0] = floatTensor[i, 0]; + arr[i * 4 + 1] = floatTensor[i, 1]; + arr[i * 4 + 2] = floatTensor[i, 2]; + arr[i * 4 + 3] = floatTensor[i, 3]; + } + } + + return arr; + } + + return null; + } + + /// + public void Dispose() + { + if (!_disposed) + { + _disposed = true; + } + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/Models/LayoutModelOptions.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/Models/LayoutModelOptions.cs new file mode 100644 index 0000000..6ef3e1f --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/Models/LayoutModelOptions.cs @@ -0,0 +1,125 @@ +using System.Collections.Generic; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.Models; + +/// +/// Configuration options for a generic layout detection model, +/// driving the . +/// +public record LayoutModelOptions +{ + /// + /// Expected model input width in pixels. + /// + public int InputWidth { get; init; } = 640; + + /// + /// Expected model input height in pixels. + /// + public int InputHeight { get; init; } = 640; + + /// + /// How the image should be resized before feeding to the model. + /// + public ResizeMode Resize { get; init; } = ResizeMode.Exact; + + /// + /// Pixel format and data type for the input tensor. + /// + public PixelFormat PixelFormat { get; init; } = PixelFormat.Uint8Chw; + + /// + /// Optional per-channel normalization to apply after converting to float. + /// + public ImageNormalization? Normalization { get; init; } + + /// + /// Minimum confidence threshold for detections. + /// + public float ConfidenceThreshold { get; init; } = 0.3f; + + /// + /// IoU threshold for Non-Maximum Suppression. + /// + public float NmsIouThreshold { get; init; } = 0.45f; + + /// + /// Whether NMS should be applied to the output detections. + /// + public bool RequiresNms { get; init; } + + /// + /// Format of the bounding box coordinates in the model output. + /// + public BboxFormat OutputBboxFormat { get; init; } = BboxFormat.CxCyWh; + + /// + /// Mapping from class ID to human-readable label. If null, numeric labels are used. + /// + public IReadOnlyDictionary? ClassLabels { get; init; } +} + +/// +/// How the input image is resized before inference. +/// +public enum ResizeMode +{ + /// Stretch to exact target dimensions. + Exact, + + /// Preserve aspect ratio and pad with a solid color. + Letterbox, + + /// Preserve aspect ratio without padding (model must accept variable sizes). + AspectPreserve +} + +/// +/// Pixel format and data type for the model input tensor. +/// +public enum PixelFormat +{ + /// Unsigned 8-bit integer, CHW layout [1, 3, H, W]. + Uint8Chw, + + /// 32-bit float, CHW layout [1, 3, H, W], values in [0, 1]. + Float32Chw +} + +/// +/// Format of bounding box coordinates in the model output. +/// +public enum BboxFormat +{ + /// Center-X, Center-Y, Width, Height. + CxCyWh, + + /// Top-left X, top-left Y, bottom-right X, bottom-right Y. + Xyxy, + + /// Top-left X, top-left Y, Width, Height. + Xywh +} + +/// +/// Per-channel normalization parameters: (value - mean) / std. +/// +/// Per-channel mean values (RGB order, length 3). +/// Per-channel standard deviation values (RGB order, length 3). +public record ImageNormalization(float[] Mean, float[] Std); + +/// +/// Well-known normalization presets for common model architectures. +/// +public static class WellKnownNormalizations +{ + /// + /// ImageNet normalization: mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]. + /// + public static readonly ImageNormalization ImageNet = new([0.485f, 0.456f, 0.406f], [0.229f, 0.224f, 0.225f]); + + /// + /// Simple [0, 255] to [0, 1] normalization (mean=0, std=1/255). + /// + public static readonly ImageNormalization ZeroToOne = new([0f, 0f, 0f], [1f / 255f, 1f / 255f, 1f / 255f]); +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/Models/RtDetrLayoutModel.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/Models/RtDetrLayoutModel.cs new file mode 100644 index 0000000..faeac02 --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/Models/RtDetrLayoutModel.cs @@ -0,0 +1,265 @@ +using Microsoft.ML.OnnxRuntime; +using Microsoft.ML.OnnxRuntime.Tensors; +using SkiaSharp; +using System; +using System.Collections.Generic; +using System.Linq; +using UglyToad.PdfPig.Core; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.Models; + +/// +/// Layout detection model implementation for the Docling Layout Heron RT-DETR v2 model. +/// +/// +/// +/// RT-DETR (Real-Time DEtection TRansformer) uses built-in Hungarian matching, +/// so no external NMS is needed. +/// +/// +/// Input: Resize to 640×640 exact, uint8 CHW tensor (normalization is baked into the ONNX graph). +/// Also provides an orig_target_sizes int64 tensor with [height, width]. +/// +/// +/// Output: Handles both labels/boxes/scores and +/// pred_labels/pred_boxes/pred_scores naming conventions. +/// +/// +public sealed class RtDetrLayoutModel : ILayoutDetectionModel +{ + private const int ModelInputWidth = 640; + private const int ModelInputHeight = 640; + + private static readonly Dictionary DefaultLabelMapping = new Dictionary + { + [0] = "caption", + [1] = "footnote", + [2] = "formula", + [3] = "list_item", + [4] = "page_footer", + [5] = "page_header", + [6] = "picture", + [7] = "section_header", + [8] = "table", + [9] = "text", + [10] = "title", + [11] = "document_index", + [12] = "code", + [13] = "checkbox_selected", + [14] = "checkbox_unselected", + [15] = "form", + [16] = "key_value_region" + }; + + private bool _disposed; + + /// + /// Create a new . + /// + /// Path to the RT-DETR ONNX model file. + public RtDetrLayoutModel(string modelPath) + { + ModelPath = modelPath ?? throw new ArgumentNullException(nameof(modelPath)); + } + + /// + public string ModelPath { get; } + + /// + public IReadOnlyDictionary LabelMapping => DefaultLabelMapping; + + /// + public IReadOnlyList Preprocess(SKBitmap pageImage, int originalWidth, int originalHeight) + { + ArgumentNullException.ThrowIfNull(pageImage); + + // Resize to model input dimensions (exact, no letterbox) + using var resized = ImagePreprocessing.ResizeExact(pageImage, ModelInputWidth, ModelInputHeight); + + // Convert to CHW uint8 tensor — normalization is baked into the ONNX graph + var imageTensor = ImagePreprocessing.ToChwUint8(resized); + + // Create orig_target_sizes tensor [1, 2] with [height, width] + var origSizesTensor = new DenseTensor([1, 2]); + origSizesTensor[0, 0] = originalHeight; + origSizesTensor[0, 1] = originalWidth; + + return new List + { + NamedOnnxValue.CreateFromTensor("images", imageTensor), + NamedOnnxValue.CreateFromTensor("orig_target_sizes", origSizesTensor) + }; + } + + /// + public IReadOnlyList Postprocess( + IDisposableReadOnlyCollection results, + int originalWidth, + int originalHeight) + { + ArgumentNullException.ThrowIfNull(results); + + // Try both naming conventions + var labelsValue = TryGetOutput(results, "labels") ?? TryGetOutput(results, "pred_labels"); + var boxesValue = TryGetOutput(results, "boxes") ?? TryGetOutput(results, "pred_boxes"); + var scoresValue = TryGetOutput(results, "scores") ?? TryGetOutput(results, "pred_scores"); + + if (labelsValue is null || boxesValue is null || scoresValue is null) + { + return Array.Empty(); + } + + var labels = ExtractLabels(labelsValue); + var scores = ExtractScores(scoresValue); + var boxes = ExtractBoxes(boxesValue); + + if (labels is null || scores is null || boxes is null) + { + return Array.Empty(); + } + + int count = labels.Length; + var detections = new List(count); + + for (int i = 0; i < count; i++) + { + int classId = labels[i]; + float confidence = scores[i]; + + // Boxes are in absolute pixel coordinates (post orig_target_sizes scaling) + float x1 = boxes[i * 4 + 0]; + float y1 = boxes[i * 4 + 1]; + float x2 = boxes[i * 4 + 2]; + float y2 = boxes[i * 4 + 3]; + + string label = DefaultLabelMapping.TryGetValue(classId, out var name) + ? name + : $"class_{classId}"; + + // Convert from image coords (top-left origin) to PDF coords (bottom-left origin) + var rect = DetectionPostprocessing.XyxyToRect(x1, y1, x2, y2, originalWidth, originalHeight); + + detections.Add(new LayoutDetection(rect, label, classId, confidence)); + } + + return detections; + } + + private static DisposableNamedOnnxValue? TryGetOutput( + IDisposableReadOnlyCollection results, + string name) + { + foreach (var result in results) + { + if (string.Equals(result.Name, name, StringComparison.Ordinal)) + { + return result; + } + } + + return null; + } + + private static int[]? ExtractLabels(DisposableNamedOnnxValue value) + { + // Try int64 tensor first (most common) + if (value.Value is Tensor longTensor) + { + var dims = longTensor.Dimensions; + int count = dims.Length > 1 ? dims[1] : dims[0]; + var labels = new int[count]; + + for (int i = 0; i < count; i++) + { + labels[i] = dims.Length > 1 + ? (int)longTensor[0, i] + : (int)longTensor[i]; + } + + return labels; + } + + // Try int32 tensor + if (value.Value is Tensor intTensor) + { + var dims = intTensor.Dimensions; + int count = dims.Length > 1 ? dims[1] : dims[0]; + var labels = new int[count]; + + for (int i = 0; i < count; i++) + { + labels[i] = dims.Length > 1 + ? intTensor[0, i] + : intTensor[i]; + } + + return labels; + } + + return null; + } + + private static float[]? ExtractScores(DisposableNamedOnnxValue value) + { + if (value.Value is Tensor floatTensor) + { + var dims = floatTensor.Dimensions; + int count = dims.Length > 1 ? dims[1] : dims[0]; + var scores = new float[count]; + + for (int i = 0; i < count; i++) + { + scores[i] = dims.Length > 1 + ? floatTensor[0, i] + : floatTensor[i]; + } + + return scores; + } + + return null; + } + + private static float[]? ExtractBoxes(DisposableNamedOnnxValue value) + { + if (value.Value is Tensor floatTensor) + { + var dims = floatTensor.Dimensions; + int count = dims.Length > 1 ? dims[1] : dims[0]; + var boxes = new float[count * 4]; + + for (int i = 0; i < count; i++) + { + if (dims.Length == 3) + { + // Shape: [batch, num_detections, 4] + boxes[i * 4 + 0] = floatTensor[0, i, 0]; + boxes[i * 4 + 1] = floatTensor[0, i, 1]; + boxes[i * 4 + 2] = floatTensor[0, i, 2]; + boxes[i * 4 + 3] = floatTensor[0, i, 3]; + } + else if (dims.Length == 2) + { + // Shape: [num_detections, 4] + boxes[i * 4 + 0] = floatTensor[i, 0]; + boxes[i * 4 + 1] = floatTensor[i, 1]; + boxes[i * 4 + 2] = floatTensor[i, 2]; + boxes[i * 4 + 3] = floatTensor[i, 3]; + } + } + + return boxes; + } + + return null; + } + + /// + public void Dispose() + { + if (!_disposed) + { + _disposed = true; + } + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxPageSegmenter.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxPageSegmenter.cs new file mode 100644 index 0000000..8565a16 --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxPageSegmenter.cs @@ -0,0 +1,293 @@ +using Microsoft.Extensions.Options; +using SkiaSharp; +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using UglyToad.PdfPig.Content; +using UglyToad.PdfPig.Core; +using UglyToad.PdfPig.DocumentLayoutAnalysis; +using UglyToad.PdfPig.DocumentLayoutAnalysis.PageSegmenter; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; + +/// +/// Page segmenter that uses an ONNX layout detection model to identify +/// document regions and assign words to detected blocks. +/// +public class OnnxPageSegmenter : IPageSegmenter, IDisposable, IAsyncDisposable +{ + private readonly ILayoutDetectionModel _model; + private readonly Microsoft.ML.OnnxRuntime.InferenceSession _session; + private readonly float _confidenceThreshold; + private readonly int _renderDpi; + private bool _disposed; + + /// + /// Create a new using dependency-injected options. + /// + /// The layout detection model to use. + /// The configured options. + public OnnxPageSegmenter(ILayoutDetectionModel model, IOptions options) + : this(model, GetOptionsValue(options)) + { + } + + /// + /// Create a new . + /// + /// The layout detection model to use. + /// Optional segmenter configuration. + public OnnxPageSegmenter(ILayoutDetectionModel model, OnnxSegmenterOptions? options = null) + { + _model = model ?? throw new ArgumentNullException(nameof(model)); + _confidenceThreshold = options?.ConfidenceThreshold ?? 0.3f; + _renderDpi = options?.RenderDpi ?? 150; + var sessionOpts = options?.SessionOptions ?? new Microsoft.ML.OnnxRuntime.SessionOptions(); + _session = new Microsoft.ML.OnnxRuntime.InferenceSession(model.ModelPath, sessionOpts); + } + + /// + /// Get text blocks by running ONNX layout detection and assigning words to detected regions. + /// + /// The page's words to generate text blocks for. + /// A list of text blocks from this approach. + public IReadOnlyList GetBlocks(IEnumerable words) + { + ObjectDisposedException.ThrowIf(_disposed, this); + var wordList = words?.ToList() ?? throw new ArgumentNullException(nameof(words)); + if (wordList.Count == 0) + { + return Array.Empty(); + } + + // 1. Compute page bounds from words + double minX = double.MaxValue, minY = double.MaxValue; + double maxX = double.MinValue, maxY = double.MinValue; + foreach (var word in wordList) + { + var bb = word.BoundingBox; + minX = Math.Min(minX, bb.Left); + minY = Math.Min(minY, bb.Bottom); + maxX = Math.Max(maxX, bb.Right); + maxY = Math.Max(maxY, bb.Top); + } + + double pageWidth = maxX - minX; + double pageHeight = maxY - minY; + + if (pageWidth <= 0 || pageHeight <= 0) + { + return [new TextBlock(CreateSingleLine(wordList))]; + } + + // 2. Render page image using PageImageRenderer + using SKBitmap pageImage = PageImageRenderer.RenderWords(wordList, pageWidth, pageHeight, _renderDpi); + int imageWidth = pageImage.Width; + int imageHeight = pageImage.Height; + + // 3. Preprocess with model + var inputs = _model.Preprocess(pageImage, imageWidth, imageHeight); + + // 4. Run inference + using var results = _session.Run(inputs); + + // 5. Postprocess with model + var detections = _model.Postprocess(results, imageWidth, imageHeight); + + // 6. Filter by confidence + var filtered = new List(); + foreach (var det in detections) + { + if (det.Confidence >= _confidenceThreshold) + { + filtered.Add(det); + } + } + + if (filtered.Count == 0) + { + return [new TextBlock(CreateSingleLine(wordList))]; + } + + // 7. Map detections to TextBlocks via bbox overlap + return MapDetectionsToBlocks(filtered, wordList, pageWidth, pageHeight, minX, minY); + } + + private static List MapDetectionsToBlocks( + List detections, + List words, + double pageWidth, + double pageHeight, + double offsetX, + double offsetY) + { + var wordAssigned = new bool[words.Count]; + var blocks = new List(); + + foreach (var detection in detections) + { + var capturedWords = new List(); + + for (int i = 0; i < words.Count; i++) + { + if (wordAssigned[i]) + { + continue; + } + + if (HasOverlap(detection.BoundingBox, words[i].BoundingBox)) + { + capturedWords.Add(words[i]); + wordAssigned[i] = true; + } + } + + if (capturedWords.Count > 0) + { + var lines = GroupWordsIntoLines(capturedWords); + blocks.Add(new AnnotatedTextBlock(lines, detection.Label, detection.Confidence)); + } + } + + // 8. Handle uncaptured words — group into a fallback block + var uncaptured = new List(); + for (int i = 0; i < words.Count; i++) + { + if (!wordAssigned[i]) + { + uncaptured.Add(words[i]); + } + } + + if (uncaptured.Count > 0) + { + var lines = GroupWordsIntoLines(uncaptured); + blocks.Add(new TextBlock(lines)); + } + + return blocks; + } + + private static bool HasOverlap(PdfRectangle a, PdfRectangle b) + { + double interLeft = Math.Max(a.Left, b.Left); + double interRight = Math.Min(a.Right, b.Right); + double interBottom = Math.Max(a.Bottom, b.Bottom); + double interTop = Math.Min(a.Top, b.Top); + + return interLeft < interRight && interBottom < interTop; + } + + private static IReadOnlyList GroupWordsIntoLines(List words) + { + if (words.Count == 0) + { + return Array.Empty(); + } + + // Sort words by vertical position (top to bottom), then left to right + words.Sort((a, b) => + { + double aY = a.BoundingBox.Top; + double bY = b.BoundingBox.Top; + + // Group by Y-proximity: if vertical centers are close, treat as same line + double aCenter = (a.BoundingBox.Top + a.BoundingBox.Bottom) / 2.0; + double bCenter = (b.BoundingBox.Top + b.BoundingBox.Bottom) / 2.0; + double aHeight = a.BoundingBox.Height; + double bHeight = b.BoundingBox.Height; + double tolerance = Math.Min(aHeight, bHeight) * 0.5; + + if (Math.Abs(aCenter - bCenter) <= tolerance) + { + return a.BoundingBox.Left.CompareTo(b.BoundingBox.Left); + } + + // Higher Y = higher on page in PDF coords, so sort descending + return bY.CompareTo(aY); + }); + + var lines = new List(); + var currentLineWords = new List { words[0] }; + + for (int i = 1; i < words.Count; i++) + { + var prev = currentLineWords[^1]; + var curr = words[i]; + + double prevCenter = (prev.BoundingBox.Top + prev.BoundingBox.Bottom) / 2.0; + double currCenter = (curr.BoundingBox.Top + curr.BoundingBox.Bottom) / 2.0; + double tolerance = Math.Min(prev.BoundingBox.Height, curr.BoundingBox.Height) * 0.5; + + if (Math.Abs(prevCenter - currCenter) <= tolerance) + { + currentLineWords.Add(curr); + } + else + { + // Sort current line words left to right before creating line + currentLineWords.Sort((a, b) => a.BoundingBox.Left.CompareTo(b.BoundingBox.Left)); + lines.Add(new TextLine(currentLineWords.ToList())); + currentLineWords.Clear(); + currentLineWords.Add(curr); + } + } + + if (currentLineWords.Count > 0) + { + currentLineWords.Sort((a, b) => a.BoundingBox.Left.CompareTo(b.BoundingBox.Left)); + lines.Add(new TextLine(currentLineWords.ToList())); + } + + return lines; + } + + private static IReadOnlyList CreateSingleLine(List words) + { + return [new TextLine(words)]; + } + + private static OnnxSegmenterOptions GetOptionsValue(IOptions options) + { + ArgumentNullException.ThrowIfNull(options); + return options.Value; + } + + /// + /// Dispose resources held by this segmenter. + /// + public void Dispose() + { + Dispose(true); + GC.SuppressFinalize(this); + } + + /// + /// Asynchronously dispose resources held by this segmenter. + /// + public ValueTask DisposeAsync() + { + Dispose(true); + GC.SuppressFinalize(this); + return ValueTask.CompletedTask; + } + + /// + /// Dispose managed and unmanaged resources. + /// + /// Whether managed resources should be disposed. + protected virtual void Dispose(bool disposing) + { + if (!_disposed) + { + if (disposing) + { + _session.Dispose(); + _model.Dispose(); + } + + _disposed = true; + } + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxPageSegmenterServiceCollectionExtensions.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxPageSegmenterServiceCollectionExtensions.cs new file mode 100644 index 0000000..fd78e2a --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxPageSegmenterServiceCollectionExtensions.cs @@ -0,0 +1,50 @@ +using System; +using System.Diagnostics.CodeAnalysis; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.Models; +using Microsoft.Extensions.DependencyInjection.Extensions; +using Microsoft.Extensions.Options; +using UglyToad.PdfPig.DocumentLayoutAnalysis.PageSegmenter; + +namespace Microsoft.Extensions.DependencyInjection; + +/// +/// Extension methods for configuring ONNX-based page segmentation services. +/// +public static class OnnxPageSegmenterServiceCollectionExtensions +{ + /// + /// Adds an and its dependencies to the service collection. + /// + /// + /// The implementation to use (e.g. + /// or ). + /// + /// The service collection. + /// Optional delegate to configure . + /// The service collection for chaining. + /// is . + public static IServiceCollection AddOnnxPageSegmenter<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors)] TModel>( + this IServiceCollection services, + Action? configure = null) + where TModel : class, ILayoutDetectionModel + { + ArgumentNullException.ThrowIfNull(services); + + services.AddOptions(); + + if (configure is not null) + { + services.Configure(configure); + } + + services.TryAddSingleton(); + services.TryAddSingleton(); + services.TryAddSingleton(sp => sp.GetRequiredService()); + + services.TryAddEnumerable( + ServiceDescriptor.Singleton, OnnxSegmenterOptionsValidator>()); + + return services; + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxSegmenterOptions.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxSegmenterOptions.cs new file mode 100644 index 0000000..620cb7d --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxSegmenterOptions.cs @@ -0,0 +1,24 @@ +using Microsoft.ML.OnnxRuntime; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; + +/// +/// Options for configuring the . +/// +public record OnnxSegmenterOptions +{ + /// + /// Minimum confidence threshold for detections (0.0 to 1.0). + /// + public float ConfidenceThreshold { get; set; } = 0.3f; + + /// + /// ONNX Runtime session options. Use to configure GPU, thread count, etc. + /// + public SessionOptions? SessionOptions { get; set; } + + /// + /// DPI for rendering the page image. Higher values improve accuracy but are slower. + /// + public int RenderDpi { get; set; } = 150; +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxSegmenterOptionsValidator.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxSegmenterOptionsValidator.cs new file mode 100644 index 0000000..041df1c --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/OnnxSegmenterOptionsValidator.cs @@ -0,0 +1,25 @@ +using Microsoft.Extensions.Options; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; + +/// +/// Validates at startup to catch configuration errors early. +/// +internal sealed class OnnxSegmenterOptionsValidator : IValidateOptions +{ + /// + public ValidateOptionsResult Validate(string? name, OnnxSegmenterOptions options) + { + if (options.ConfidenceThreshold is < 0f or > 1f) + { + return ValidateOptionsResult.Fail("ConfidenceThreshold must be between 0.0 and 1.0."); + } + + if (options.RenderDpi <= 0) + { + return ValidateOptionsResult.Fail("RenderDpi must be a positive integer."); + } + + return ValidateOptionsResult.Success; + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/PageImageRenderer.cs b/MEDI/src/PdfPig.OnnxLayoutAnalysis/PageImageRenderer.cs new file mode 100644 index 0000000..55e8d02 --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/PageImageRenderer.cs @@ -0,0 +1,77 @@ +using SkiaSharp; +using System; +using System.Collections.Generic; +using UglyToad.PdfPig.Content; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; + +/// +/// Renders word bounding boxes to an for use as +/// input to ONNX layout detection models. +/// +public static class PageImageRenderer +{ + /// + /// Render word bounding boxes as filled rectangles on a white background. + /// PDF coordinates (bottom-left origin, Y-up) are converted to image coordinates + /// (top-left origin, Y-down). + /// + /// The words whose bounding boxes to render. + /// The page width in PDF units. + /// The page height in PDF units. + /// Rendering DPI. Higher values produce larger images with more detail. + /// A new bitmap with word bounding boxes rendered. + public static SKBitmap RenderWords(IReadOnlyList words, double pageWidth, double pageHeight, int dpi = 150) + { + ArgumentNullException.ThrowIfNull(words); + + // Scale from PDF points (72 dpi) to target DPI + double scale = dpi / 72.0; + int imageWidth = Math.Max(1, (int)(pageWidth * scale)); + int imageHeight = Math.Max(1, (int)(pageHeight * scale)); + + var info = new SKImageInfo(imageWidth, imageHeight, SKColorType.Rgba8888, SKAlphaType.Premul); + var bitmap = new SKBitmap(info); + + using var canvas = new SKCanvas(bitmap); + canvas.Clear(SKColors.White); + + using var paint = new SKPaint(); + paint.IsAntialias = false; + paint.Color = SKColors.Black; + paint.Style = SKPaintStyle.Fill; + + // Compute the bounds offset so we render relative to (0, 0) + double minX = double.MaxValue; + double minY = double.MaxValue; + foreach (var word in words) + { + minX = Math.Min(minX, word.BoundingBox.Left); + minY = Math.Min(minY, word.BoundingBox.Bottom); + } + + foreach (var word in words) + { + var bb = word.BoundingBox; + + // Translate to origin + double left = (bb.Left - minX) * scale; + double right = (bb.Right - minX) * scale; + double pdfBottom = (bb.Bottom - minY) * scale; + double pdfTop = (bb.Top - minY) * scale; + + // Convert PDF Y (bottom-up) to image Y (top-down) + float imgLeft = (float)left; + float imgTop = (float)(imageHeight - pdfTop); + float imgRight = (float)right; + float imgBottom = (float)(imageHeight - pdfBottom); + + if (imgRight > imgLeft && imgBottom > imgTop) + { + canvas.DrawRect(new SKRect(imgLeft, imgTop, imgRight, imgBottom), paint); + } + } + + return bitmap; + } +} diff --git a/MEDI/src/PdfPig.OnnxLayoutAnalysis/PdfPig.OnnxLayoutAnalysis.csproj b/MEDI/src/PdfPig.OnnxLayoutAnalysis/PdfPig.OnnxLayoutAnalysis.csproj new file mode 100644 index 0000000..182f186 --- /dev/null +++ b/MEDI/src/PdfPig.OnnxLayoutAnalysis/PdfPig.OnnxLayoutAnalysis.csproj @@ -0,0 +1,27 @@ + + + 1.0.0-preview.1 + CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis + $(AssemblyName) + net10.0;net8.0 + + ONNX-based layout detection for PdfPig document processing + ONNX Runtime layout detection models (RT-DETR, YOLO, configurable) that implement PdfPig's IPageSegmenter for intelligent document layout analysis. + true + + + + + + + + + + + + + + + + + diff --git a/MEDI/test/DataIngestion.UnitTests/DataIngestion.UnitTests.csproj b/MEDI/test/DataIngestion.UnitTests/DataIngestion.UnitTests.csproj new file mode 100644 index 0000000..26cc3c7 --- /dev/null +++ b/MEDI/test/DataIngestion.UnitTests/DataIngestion.UnitTests.csproj @@ -0,0 +1,34 @@ + + + + CommunityToolkit.DataIngestion.UnitTests + CommunityToolkit.DataIngestion.UnitTests + net10.0 + true + false + enable + enable + + + + + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + + + + + + + + + + diff --git a/MEDI/test/DataIngestion.UnitTests/Ingestion/ContextualChunkEnricherTests.cs b/MEDI/test/DataIngestion.UnitTests/Ingestion/ContextualChunkEnricherTests.cs new file mode 100644 index 0000000..8bb6180 --- /dev/null +++ b/MEDI/test/DataIngestion.UnitTests/Ingestion/ContextualChunkEnricherTests.cs @@ -0,0 +1,157 @@ +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; +using CommunityToolkit.DataIngestion.UnitTests.Utils; + +namespace CommunityToolkit.DataIngestion.UnitTests.Ingestion; + +public class ContextualChunkEnricherTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", + () => new ContextualChunkEnricher(null!)); + } + + [Fact] + public async Task ProcessAsync_AddsContextualSummaryMetadata() + { + var summaryText = "This chunk discusses PDF text extraction."; + using var client = CreateClientReturning(summaryText); + var processor = new ContextualChunkEnricher(client); + var chunks = CreateChunks("Some text about PDF extraction and analysis."); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + Assert.Single(results); + Assert.True(results[0].Metadata.ContainsKey(MetadataKeys.ContextualSummary)); + Assert.Equal(summaryText, results[0].Metadata[MetadataKeys.ContextualSummary]); + } + + [Fact] + public async Task ProcessAsync_MultipleChunks_AllGetSummaries() + { + var summaryText = "Summary text."; + using var client = CreateClientReturning(summaryText); + var processor = new ContextualChunkEnricher(client); + var chunks = CreateChunks("First chunk content.", "Second chunk content."); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + Assert.Equal(2, results.Count); + Assert.All(results, r => + { + Assert.True(r.Metadata.ContainsKey(MetadataKeys.ContextualSummary)); + Assert.Equal(summaryText, r.Metadata[MetadataKeys.ContextualSummary]); + }); + } + + [Fact] + public async Task ProcessAsync_EmptyChunks_NoResults() + { + using var client = CreateClientReturning("summary"); + var processor = new ContextualChunkEnricher(client); + + var results = await CollectAsync(processor.ProcessAsync(CreateChunks())); + + Assert.Empty(results); + } + + [Fact] + public async Task ProcessAsync_CancellationRequested_Throws() + { + using var client = new TestChatClient + { + GetResponseAsyncCallback = (_, _, ct) => + { + ct.ThrowIfCancellationRequested(); + return Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "summary"))); + } + }; + var processor = new ContextualChunkEnricher(client); + + using var cts = new CancellationTokenSource(); + cts.Cancel(); + + await Assert.ThrowsAnyAsync(async () => + { + await foreach (var _ in processor.ProcessAsync(CreateChunks("Some content."), cts.Token)) + { + } + }); + } + + [Fact] + public async Task ProcessAsync_PreservesChunkContentDocumentAndInstance() + { + using var client = CreateClientReturning("A summary."); + var processor = new ContextualChunkEnricher(client); + var document = new IngestionDocument("test-doc"); + var originalContent = "Original chunk content stays the same."; + var chunk = new IngestionChunk(originalContent, document); + + var results = await CollectAsync(processor.ProcessAsync(CreateChunkInstances(chunk))); + + var result = Assert.Single(results); + Assert.Same(chunk, result); + Assert.Same(document, result.Document); + Assert.Equal(originalContent, result.Content); + } + + [Theory] + [InlineData("table", "Summarize this table for search retrieval. Describe what data it contains, including key metrics, column headers, and notable values. Output only the summary sentence, nothing else.")] + [InlineData("picture", "Summarize what this figure or image refers to for search retrieval. Output only the summary sentence, nothing else.")] + [InlineData("caption", "Summarize what this figure or image refers to for search retrieval. Output only the summary sentence, nothing else.")] + [InlineData(null, "Provide a single concise sentence summarizing the following text for use in search retrieval. Output only the summary sentence, nothing else.")] + public async Task ProcessAsync_UsesElementTypeSpecificPrompt(string? elementType, string expectedInstruction) + { + string? prompt = null; + using var client = new TestChatClient + { + GetResponseAsyncCallback = (messages, _, _) => + { + prompt = Assert.Single(messages).Text; + return Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "summary"))); + } + }; + var processor = new ContextualChunkEnricher(client); + var document = new IngestionDocument("test-doc"); + var chunk = new IngestionChunk("Chunk content.", document); + if (elementType is not null) + chunk.Metadata["element_type"] = elementType; + + await CollectAsync(processor.ProcessAsync(CreateChunkInstances(chunk))); + + Assert.Equal(expectedInstruction + "\n\nChunk content.", prompt); + } + + private static TestChatClient CreateClientReturning(string response) + => new() + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, response))) + }; + + private static async IAsyncEnumerable> CreateChunks(params string[] contents) + { + await Task.CompletedTask; + var doc = new IngestionDocument("test-doc"); + foreach (var content in contents) + yield return new IngestionChunk(content, doc); + } + + private static async IAsyncEnumerable> CreateChunkInstances(params IngestionChunk[] chunks) + { + await Task.CompletedTask; + foreach (var chunk in chunks) + yield return chunk; + } + + private static async Task>> CollectAsync(IAsyncEnumerable> source) + { + var list = new List>(); + await foreach (var item in source) + list.Add(item); + return list; + } +} diff --git a/MEDI/test/DataIngestion.UnitTests/Ingestion/EntityExtractionProcessorTests.cs b/MEDI/test/DataIngestion.UnitTests/Ingestion/EntityExtractionProcessorTests.cs new file mode 100644 index 0000000..a303ab5 --- /dev/null +++ b/MEDI/test/DataIngestion.UnitTests/Ingestion/EntityExtractionProcessorTests.cs @@ -0,0 +1,69 @@ +using Microsoft.Extensions.DataIngestion; +using CommunityToolkit.DataIngestion.UnitTests.Utils; + +namespace CommunityToolkit.DataIngestion.UnitTests.Ingestion; + +public class EntityExtractionProcessorTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", + () => new EntityExtractionProcessor(null!)); + } + + [Fact] + public async Task ProcessAsync_ExtractsEntitiesIntoMetadata() + { + using var client = TestChatClient.WithJsonResponse(""" + { + "people": ["John Doe", "Jane Smith"], + "organizations": ["Microsoft", "OpenAI"], + "technologies": [".NET", "Azure"], + "versions": ["8.0", "10.0"] + } + """); + var processor = new EntityExtractionProcessor(client); + var chunks = CreateChunks("John Doe from Microsoft wrote about .NET 8.0"); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + Assert.Single(results); + Assert.Equal("John Doe, Jane Smith", results[0].Metadata[MetadataKeys.EntitiesPeople]); + Assert.Equal("Microsoft, OpenAI", results[0].Metadata[MetadataKeys.EntitiesOrganizations]); + Assert.Equal(".NET, Azure", results[0].Metadata[MetadataKeys.EntitiesTechnologies]); + Assert.Equal("8.0, 10.0", results[0].Metadata[MetadataKeys.EntitiesVersions]); + } + + [Fact] + public async Task ProcessAsync_LlmThrows_SetsEmptyDefaults() + { + using var client = TestChatClient.WithException(new InvalidOperationException("timeout")); + var processor = new EntityExtractionProcessor(client); + var chunks = CreateChunks("some text"); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + Assert.Single(results); + Assert.Equal("", results[0].Metadata[MetadataKeys.EntitiesPeople]); + Assert.Equal("", results[0].Metadata[MetadataKeys.EntitiesOrganizations]); + Assert.Equal("", results[0].Metadata[MetadataKeys.EntitiesTechnologies]); + Assert.Equal("", results[0].Metadata[MetadataKeys.EntitiesVersions]); + } + + private static async IAsyncEnumerable> CreateChunks(params string[] contents) + { + await Task.CompletedTask; + var doc = new IngestionDocument("test-doc"); + foreach (var content in contents) + yield return new IngestionChunk(content, doc); + } + + private static async Task>> CollectAsync(IAsyncEnumerable> source) + { + var list = new List>(); + await foreach (var item in source) + list.Add(item); + return list; + } +} diff --git a/MEDI/test/DataIngestion.UnitTests/Ingestion/HypotheticalQueryProcessorTests.cs b/MEDI/test/DataIngestion.UnitTests/Ingestion/HypotheticalQueryProcessorTests.cs new file mode 100644 index 0000000..ddab2bc --- /dev/null +++ b/MEDI/test/DataIngestion.UnitTests/Ingestion/HypotheticalQueryProcessorTests.cs @@ -0,0 +1,64 @@ +using Microsoft.Extensions.DataIngestion; +using CommunityToolkit.DataIngestion.UnitTests.Utils; + +namespace CommunityToolkit.DataIngestion.UnitTests.Ingestion; + +public class HypotheticalQueryProcessorTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", + () => new HypotheticalQueryProcessor(null!)); + } + + [Fact] + public async Task ProcessAsync_GeneratesQuestionsAndYieldsOriginalFirst() + { + using var client = TestChatClient.WithJsonResponse( + """{"questions": ["What is dependency injection?", "How to register services?", "What is IServiceProvider?"]}"""); + var processor = new HypotheticalQueryProcessor(client, questionsPerChunk: 3); + var chunks = CreateChunks("DI in .NET uses IServiceCollection."); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + // Original chunk first, then hypothetical query chunks + Assert.True(results.Count >= 2); + Assert.Equal(MetadataKeys.ChunkTypeOriginal, results[0].Metadata[MetadataKeys.ChunkType]); + + var queryChunks = results.Where(r => + (string)r.Metadata[MetadataKeys.ChunkType]! == MetadataKeys.ChunkTypeHypotheticalQuery).ToList(); + Assert.True(queryChunks.Count > 0); + Assert.All(queryChunks, q => Assert.NotEmpty(q.Content)); + } + + [Fact] + public async Task ProcessAsync_LlmThrows_YieldsOriginalOnly() + { + using var client = TestChatClient.WithException(new InvalidOperationException("timeout")); + var processor = new HypotheticalQueryProcessor(client); + var chunks = CreateChunks("some content"); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + Assert.Single(results); + Assert.Equal(MetadataKeys.ChunkTypeOriginal, results[0].Metadata[MetadataKeys.ChunkType]); + Assert.Equal("some content", results[0].Content); + } + + private static async IAsyncEnumerable> CreateChunks(params string[] contents) + { + await Task.CompletedTask; + var doc = new IngestionDocument("test-doc"); + foreach (var content in contents) + yield return new IngestionChunk(content, doc); + } + + private static async Task>> CollectAsync(IAsyncEnumerable> source) + { + var list = new List>(); + await foreach (var item in source) + list.Add(item); + return list; + } +} diff --git a/MEDI/test/DataIngestion.UnitTests/Ingestion/TopicClassificationProcessorTests.cs b/MEDI/test/DataIngestion.UnitTests/Ingestion/TopicClassificationProcessorTests.cs new file mode 100644 index 0000000..3ef6b27 --- /dev/null +++ b/MEDI/test/DataIngestion.UnitTests/Ingestion/TopicClassificationProcessorTests.cs @@ -0,0 +1,83 @@ +using Microsoft.Extensions.DataIngestion; +using CommunityToolkit.DataIngestion.UnitTests.Utils; + +namespace CommunityToolkit.DataIngestion.UnitTests.Ingestion; + +public class TopicClassificationProcessorTests +{ + private static readonly string[] Taxonomy = ["web", "data", "security", "performance", "architecture"]; + + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", + () => new TopicClassificationProcessor(null!, Taxonomy)); + } + + [Fact] + public void Constructor_ThrowsOnNullTaxonomy() + { + using var client = TestChatClient.WithJsonResponse("{}"); + Assert.Throws("taxonomy", + () => new TopicClassificationProcessor(client, null!)); + } + + [Fact] + public async Task ProcessAsync_ValidTopic_SetsMetadata() + { + using var client = TestChatClient.WithJsonResponse( + """{"primary": "security", "secondary": ["architecture", "web"]}"""); + var processor = new TopicClassificationProcessor(client, Taxonomy); + var chunks = CreateChunks("OAuth2 tokens need proper architecture"); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + Assert.Single(results); + Assert.Equal("security", results[0].Metadata[MetadataKeys.TopicPrimary]); + Assert.Contains("architecture", (string)results[0].Metadata[MetadataKeys.TopicSecondary]!); + } + + [Fact] + public async Task ProcessAsync_InvalidPrimary_DefaultsToUnknown() + { + using var client = TestChatClient.WithJsonResponse( + """{"primary": "not_in_taxonomy", "secondary": []}"""); + var processor = new TopicClassificationProcessor(client, Taxonomy); + var chunks = CreateChunks("some random text"); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + Assert.Single(results); + Assert.Equal("unknown", results[0].Metadata[MetadataKeys.TopicPrimary]); + } + + [Fact] + public async Task ProcessAsync_LlmThrows_DefaultsToUnknown() + { + using var client = TestChatClient.WithException(new InvalidOperationException("fail")); + var processor = new TopicClassificationProcessor(client, Taxonomy); + var chunks = CreateChunks("text"); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + Assert.Single(results); + Assert.Equal("unknown", results[0].Metadata[MetadataKeys.TopicPrimary]); + Assert.Equal("", results[0].Metadata[MetadataKeys.TopicSecondary]); + } + + private static async IAsyncEnumerable> CreateChunks(params string[] contents) + { + await Task.CompletedTask; + var doc = new IngestionDocument("test-doc"); + foreach (var content in contents) + yield return new IngestionChunk(content, doc); + } + + private static async Task>> CollectAsync(IAsyncEnumerable> source) + { + var list = new List>(); + await foreach (var item in source) + list.Add(item); + return list; + } +} diff --git a/MEDI/test/DataIngestion.UnitTests/Ingestion/TreeIndexProcessorTests.cs b/MEDI/test/DataIngestion.UnitTests/Ingestion/TreeIndexProcessorTests.cs new file mode 100644 index 0000000..1215602 --- /dev/null +++ b/MEDI/test/DataIngestion.UnitTests/Ingestion/TreeIndexProcessorTests.cs @@ -0,0 +1,82 @@ +using Microsoft.Extensions.DataIngestion; +using CommunityToolkit.DataIngestion.UnitTests.Utils; + +namespace CommunityToolkit.DataIngestion.UnitTests.Ingestion; + +public class TreeIndexProcessorTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", + () => new TreeIndexProcessor(null!)); + } + + [Fact] + public async Task ProcessAsync_MarksLeafChunksWithLevel0() + { + using var client = TestChatClient.WithJsonResponse("Branch summary text"); + var processor = new TreeIndexProcessor(client); + var chunks = CreateChunks("Chunk 1 content", "Chunk 2 content"); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + var leaves = results.Where(r => (int)r.Metadata[MetadataKeys.Level]! == 0).ToList(); + Assert.Equal(2, leaves.Count); + Assert.All(leaves, l => Assert.Equal(MetadataKeys.ChunkTypeOriginal, l.Metadata[MetadataKeys.ChunkType])); + } + + [Fact] + public async Task ProcessAsync_GeneratesBranchAndRootSummaries() + { + int callCount = 0; + using var client = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + callCount++; + var text = callCount == 1 + ? "Branch summary of the document." + : "Root summary of the entire corpus."; + return Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, text))); + } + }; + var processor = new TreeIndexProcessor(client); + var chunks = CreateChunks("Chunk 1", "Chunk 2", "Chunk 3"); + + var results = await CollectAsync(processor.ProcessAsync(chunks)); + + // Should have leaf chunks + at least 1 branch + 1 root summary + var branches = results.Where(r => + r.Metadata.ContainsKey(MetadataKeys.ChunkType) && + (string)r.Metadata[MetadataKeys.ChunkType]! == MetadataKeys.ChunkTypeBranchSummary).ToList(); + var roots = results.Where(r => + r.Metadata.ContainsKey(MetadataKeys.ChunkType) && + (string)r.Metadata[MetadataKeys.ChunkType]! == MetadataKeys.ChunkTypeRootSummary).ToList(); + + Assert.True(branches.Count >= 1, "Should generate at least one branch summary"); + Assert.True(roots.Count >= 1, "Should generate at least one root summary"); + + // Branch = level 1, Root = level 2 + Assert.All(branches, b => Assert.Equal(1, b.Metadata[MetadataKeys.Level])); + Assert.All(roots, r => Assert.Equal(2, r.Metadata[MetadataKeys.Level])); + } + + private static async IAsyncEnumerable> CreateChunks(params string[] contents) + { + await Task.CompletedTask; + var doc = new IngestionDocument("test-doc"); + foreach (var content in contents) + yield return new IngestionChunk(content, doc); + } + + private static async Task>> CollectAsync(IAsyncEnumerable> source) + { + var list = new List>(); + await foreach (var item in source) + list.Add(item); + return list; + } +} diff --git a/MEDI/test/DataIngestion.UnitTests/Utils/TestChatClient.cs b/MEDI/test/DataIngestion.UnitTests/Utils/TestChatClient.cs new file mode 100644 index 0000000..5e151f4 --- /dev/null +++ b/MEDI/test/DataIngestion.UnitTests/Utils/TestChatClient.cs @@ -0,0 +1,50 @@ +using Microsoft.Extensions.AI; + +namespace CommunityToolkit.DataIngestion.UnitTests.Utils; + +/// +/// A fake IChatClient for deterministic unit testing (no real LLM calls). +/// +public sealed class TestChatClient : IChatClient, IDisposable +{ + public Func, ChatOptions?, CancellationToken, Task>? + GetResponseAsyncCallback { get; set; } + + public Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + if (GetResponseAsyncCallback is not null) + return GetResponseAsyncCallback(messages, options, cancellationToken); + + return Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, ""))); + } + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() { } + + /// Creates a TestChatClient that always returns the given JSON string. + public static TestChatClient WithJsonResponse(string json) + => new() + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, json))) + }; + + /// Creates a TestChatClient that always throws the given exception. + public static TestChatClient WithException(Exception ex) + => new() + { + GetResponseAsyncCallback = (_, _, _) => throw ex + }; +} diff --git a/MEDI/test/DataRetrieval.UnitTests/DataRetrieval.UnitTests.csproj b/MEDI/test/DataRetrieval.UnitTests/DataRetrieval.UnitTests.csproj new file mode 100644 index 0000000..7cde736 --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/DataRetrieval.UnitTests.csproj @@ -0,0 +1,34 @@ + + + + CommunityToolkit.DataRetrieval.UnitTests + CommunityToolkit.DataRetrieval.UnitTests + net10.0 + true + false + enable + enable + + + + + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + + + + + + + + + + diff --git a/MEDI/test/DataRetrieval.UnitTests/Retrieval/AdaptiveRouterTests.cs b/MEDI/test/DataRetrieval.UnitTests/Retrieval/AdaptiveRouterTests.cs new file mode 100644 index 0000000..ecb855f --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Retrieval/AdaptiveRouterTests.cs @@ -0,0 +1,79 @@ +using Microsoft.Extensions.DataRetrieval; +using CommunityToolkit.DataRetrieval.UnitTests.Utils; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Retrieval; + +public class AdaptiveRouterTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", () => new AdaptiveRouter(null!)); + } + + [Theory] + [InlineData("vector")] + [InlineData("tree")] + [InlineData("entity")] + public async Task ProcessAsync_ValidParadigm_SetsMetadata(string paradigm) + { + using var client = TestChatClient.WithJsonResponse( + $$"""{"paradigm": "{{paradigm}}", "reasoning": "test reasoning"}"""); + var router = new AdaptiveRouter(client); + var query = new RetrievalQuery("test query"); + + var output = await router.ProcessAsync(query); + + Assert.Equal(paradigm, output.Metadata["search_paradigm"]); + } + + [Fact] + public async Task ProcessAsync_InvalidParadigm_DefaultsToVector() + { + using var client = TestChatClient.WithJsonResponse( + """{"paradigm": "unknown_paradigm", "reasoning": "bad"}"""); + var router = new AdaptiveRouter(client); + var query = new RetrievalQuery("test query"); + + var output = await router.ProcessAsync(query); + + Assert.Equal("vector", output.Metadata["search_paradigm"]); + } + + [Fact] + public async Task ProcessAsync_LlmThrows_DefaultsToVector() + { + using var client = TestChatClient.WithException(new InvalidOperationException("fail")); + var router = new AdaptiveRouter(client); + var query = new RetrievalQuery("test query"); + + var output = await router.ProcessAsync(query); + + Assert.Equal("vector", output.Metadata["search_paradigm"]); + } + + [Fact] + public async Task ProcessAsync_StoresReasoningWhenPresent() + { + using var client = TestChatClient.WithJsonResponse( + """{"paradigm": "tree", "reasoning": "broad question about architecture"}"""); + var router = new AdaptiveRouter(client); + var query = new RetrievalQuery("what are the design patterns used?"); + + var output = await router.ProcessAsync(query); + + Assert.Equal("broad question about architecture", output.Metadata["router_reasoning"]); + } + + [Fact] + public async Task ProcessAsync_NullReasoning_NoMetadataKey() + { + using var client = TestChatClient.WithJsonResponse("""{"paradigm": "vector"}"""); + var router = new AdaptiveRouter(client); + var query = new RetrievalQuery("test"); + + var output = await router.ProcessAsync(query); + + Assert.Equal("vector", output.Metadata["search_paradigm"]); + } +} diff --git a/MEDI/test/DataRetrieval.UnitTests/Retrieval/CragValidatorTests.cs b/MEDI/test/DataRetrieval.UnitTests/Retrieval/CragValidatorTests.cs new file mode 100644 index 0000000..975f59b --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Retrieval/CragValidatorTests.cs @@ -0,0 +1,122 @@ +using Microsoft.Extensions.DataRetrieval; +using CommunityToolkit.DataRetrieval.UnitTests.Utils; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Retrieval; + +public class CragValidatorTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", () => new CragValidator(null!)); + } + + [Fact] + public async Task ProcessAsync_EmptyChunks_ReturnsIncorrectPath() + { + using var client = TestChatClient.WithJsonResponse("""{"score": 5}"""); + var validator = new CragValidator(client); + var query = new RetrievalQuery("test"); + var results = new RetrievalResults { Chunks = [] }; + + var output = await validator.ProcessAsync(results, query); + + Assert.Equal(0, output.Metadata["crag_score"]); + Assert.Equal("Incorrect", output.Metadata["crag_path"]); + Assert.True((bool)output.Metadata["low_confidence"]!); + } + + [Fact] + public async Task ProcessAsync_HighScore_ReturnsCorrectPath() + { + using var client = TestChatClient.WithJsonResponse("""{"score": 5, "reasoning": "great match"}"""); + var validator = new CragValidator(client); + var (results, query) = CreateResultsWithChunks(3); + + var output = await validator.ProcessAsync(results, query); + + Assert.Equal(5, output.Metadata["crag_score"]); + Assert.Equal("Correct", output.Metadata["crag_path"]); + Assert.False(output.Metadata.ContainsKey("low_confidence")); + } + + [Fact] + public async Task ProcessAsync_MidScore_ReturnsAmbiguousPath() + { + using var client = TestChatClient.WithJsonResponse("""{"score": 3, "reasoning": "partial match"}"""); + var validator = new CragValidator(client); + var (results, query) = CreateResultsWithChunks(3); + + var output = await validator.ProcessAsync(results, query); + + Assert.Equal(3, output.Metadata["crag_score"]); + Assert.Equal("Ambiguous", output.Metadata["crag_path"]); + Assert.True((bool)output.Metadata["needs_followup"]!); + } + + [Fact] + public async Task ProcessAsync_LowScore_ReturnsIncorrectAndClearsChunks() + { + using var client = TestChatClient.WithJsonResponse("""{"score": 1, "reasoning": "off topic"}"""); + var validator = new CragValidator(client); + var (results, query) = CreateResultsWithChunks(3); + + var output = await validator.ProcessAsync(results, query); + + Assert.Equal(1, output.Metadata["crag_score"]); + Assert.Equal("Incorrect", output.Metadata["crag_path"]); + Assert.True((bool)output.Metadata["low_confidence"]!); + Assert.Empty(output.Chunks); + } + + [Fact] + public async Task ProcessAsync_ScoreOutOfRange_ClampedToDefault() + { + using var client = TestChatClient.WithJsonResponse("""{"score": 99}"""); + var validator = new CragValidator(client); + var (results, query) = CreateResultsWithChunks(3); + + var output = await validator.ProcessAsync(results, query); + + // Score 99 is not in [1,5] → defaults to 3 (Ambiguous) + Assert.Equal(3, output.Metadata["crag_score"]); + Assert.Equal("Ambiguous", output.Metadata["crag_path"]); + } + + [Fact] + public async Task ProcessAsync_LlmThrows_DefaultsToAmbiguous() + { + using var client = TestChatClient.WithException(new InvalidOperationException("fail")); + var validator = new CragValidator(client); + var (results, query) = CreateResultsWithChunks(3); + + var output = await validator.ProcessAsync(results, query); + + Assert.Equal(3, output.Metadata["crag_score"]); + Assert.Equal("Ambiguous", output.Metadata["crag_path"]); + } + + [Fact] + public async Task ProcessAsync_StoresReasoningInMetadata() + { + using var client = TestChatClient.WithJsonResponse("""{"score": 4, "reasoning": "relevant results"}"""); + var validator = new CragValidator(client); + var (results, query) = CreateResultsWithChunks(3); + + var output = await validator.ProcessAsync(results, query); + + Assert.Equal("relevant results", output.Metadata["crag_reasoning"]); + } + + private static (RetrievalResults results, RetrievalQuery query) CreateResultsWithChunks(int count) + { + var query = new RetrievalQuery("test query"); + var results = new RetrievalResults + { + Chunks = Enumerable.Range(0, count) + .Select(i => new RetrievalChunk($"Chunk {i} content", 0.9 - i * 0.1)) + .ToList() + }; + return (results, query); + } +} diff --git a/MEDI/test/DataRetrieval.UnitTests/Retrieval/HydeQueryTransformerTests.cs b/MEDI/test/DataRetrieval.UnitTests/Retrieval/HydeQueryTransformerTests.cs new file mode 100644 index 0000000..4a2502b --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Retrieval/HydeQueryTransformerTests.cs @@ -0,0 +1,87 @@ +using Microsoft.Extensions.DataRetrieval; +using CommunityToolkit.DataRetrieval.UnitTests.Utils; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Retrieval; + +public class HydeQueryTransformerTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", () => new HydeQueryTransformer(null!)); + } + + [Fact] + public async Task ProcessAsync_ReplacesVariantsWithHypothetical() + { + using var client = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, + "Dependency injection in .NET uses the IServiceCollection to register services."))) + }; + var transformer = new HydeQueryTransformer(client); + var query = new RetrievalQuery("How does DI work in .NET?"); + + var output = await transformer.ProcessAsync(query); + + Assert.Single(output.Variants); + Assert.Contains("Dependency injection", output.Variants[0]); + Assert.Equal("How does DI work in .NET?", output.Text); // Original preserved + } + + [Fact] + public async Task ProcessAsync_StoresHydeMetadata() + { + using var client = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, + "A hypothetical answer about the topic."))) + }; + var transformer = new HydeQueryTransformer(client); + var query = new RetrievalQuery("some question"); + + var output = await transformer.ProcessAsync(query); + + Assert.True(output.Metadata.ContainsKey("hyde_hypothetical")); + Assert.Equal("A hypothetical answer about the topic.", output.Metadata["hyde_hypothetical"]); + } + + [Fact] + public async Task ProcessAsync_EmptyResponse_ReturnsOriginal() + { + using var client = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, " "))) + }; + var transformer = new HydeQueryTransformer(client); + var query = new RetrievalQuery("my question"); + + var output = await transformer.ProcessAsync(query); + + Assert.Single(output.Variants); + Assert.Equal("my question", output.Variants[0]); // Unchanged + } + + [Fact] + public async Task ProcessAsync_LlmThrows_ReturnsOriginal() + { + using var client = TestChatClient.WithException(new InvalidOperationException("timeout")); + var transformer = new HydeQueryTransformer(client); + var query = new RetrievalQuery("my question"); + + var output = await transformer.ProcessAsync(query); + + Assert.Single(output.Variants); + Assert.Equal("my question", output.Variants[0]); + Assert.False(output.Metadata.ContainsKey("hyde_hypothetical")); + } +} diff --git a/MEDI/test/DataRetrieval.UnitTests/Retrieval/LlmRerankerTests.cs b/MEDI/test/DataRetrieval.UnitTests/Retrieval/LlmRerankerTests.cs new file mode 100644 index 0000000..76232a3 --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Retrieval/LlmRerankerTests.cs @@ -0,0 +1,118 @@ +using Microsoft.Extensions.DataRetrieval; +using CommunityToolkit.DataRetrieval.UnitTests.Utils; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Retrieval; + +public class LlmRerankerTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", () => new LlmReranker(null!)); + } + + [Fact] + public async Task ProcessAsync_TwoOrFewerChunks_ReturnsUnchanged() + { + using var client = TestChatClient.WithJsonResponse("""{"rankedIndices": [2, 1]}"""); + var reranker = new LlmReranker(client); + var (results, query) = CreateResults("chunk1", "chunk2"); + + var output = await reranker.ProcessAsync(results, query); + + Assert.Equal(2, output.Chunks.Count); + Assert.Equal("chunk1", output.Chunks[0].Content); + } + + [Fact] + public async Task ProcessAsync_ReranksBasedOnLlmResponse() + { + using var client = TestChatClient.WithJsonResponse("""{"rankedIndices": [3, 1, 2]}"""); + var reranker = new LlmReranker(client); + var (results, query) = CreateResults("A", "B", "C"); + + var output = await reranker.ProcessAsync(results, query); + + Assert.Equal("C", output.Chunks[0].Content); + Assert.Equal("A", output.Chunks[1].Content); + Assert.Equal("B", output.Chunks[2].Content); + } + + [Fact] + public async Task ProcessAsync_FiltersOutOfRangeIndices() + { + using var client = TestChatClient.WithJsonResponse("""{"rankedIndices": [99, 1, -1, 2]}"""); + var reranker = new LlmReranker(client); + var (results, query) = CreateResults("A", "B", "C"); + + var output = await reranker.ProcessAsync(results, query); + + Assert.Equal(2, output.Chunks.Count); + Assert.Equal("A", output.Chunks[0].Content); + Assert.Equal("B", output.Chunks[1].Content); + } + + [Fact] + public async Task ProcessAsync_DeduplicatesIndices() + { + using var client = TestChatClient.WithJsonResponse("""{"rankedIndices": [2, 2, 2, 1]}"""); + var reranker = new LlmReranker(client); + var (results, query) = CreateResults("A", "B", "C"); + + var output = await reranker.ProcessAsync(results, query); + + Assert.Equal(2, output.Chunks.Count); + Assert.Equal("B", output.Chunks[0].Content); + Assert.Equal("A", output.Chunks[1].Content); + } + + [Fact] + public async Task ProcessAsync_FallsBackToTopCandidatesOnInvalidJson() + { + using var client = TestChatClient.WithJsonResponse("not valid json"); + var reranker = new LlmReranker(client); + var (results, query) = CreateResults("A", "B", "C"); + + var output = await reranker.ProcessAsync(results, query); + + // Fallback: top MaxResults candidates returned as-is + Assert.True(output.Chunks.Count > 0); + Assert.Equal("A", output.Chunks[0].Content); + } + + [Fact] + public async Task ProcessAsync_FallsBackToTopCandidatesOnException() + { + using var client = TestChatClient.WithException(new InvalidOperationException("LLM down")); + var reranker = new LlmReranker(client); + var (results, query) = CreateResults("A", "B", "C"); + + var output = await reranker.ProcessAsync(results, query); + + Assert.True(output.Chunks.Count > 0); + Assert.Equal("A", output.Chunks[0].Content); + } + + [Fact] + public async Task ProcessAsync_SetsRerankedMetadata() + { + using var client = TestChatClient.WithJsonResponse("""{"rankedIndices": [1, 2, 3]}"""); + var reranker = new LlmReranker(client); + var (results, query) = CreateResults("A", "B", "C"); + + var output = await reranker.ProcessAsync(results, query); + + Assert.True((bool)output.Metadata["reranked"]!); + Assert.Equal(3, (int)output.Metadata["reranked_count"]!); + } + + private static (RetrievalResults results, RetrievalQuery query) CreateResults(params string[] contents) + { + var query = new RetrievalQuery("test query"); + var results = new RetrievalResults + { + Chunks = contents.Select((c, i) => new RetrievalChunk(c, 1.0 - i * 0.1)).ToList() + }; + return (results, query); + } +} diff --git a/MEDI/test/DataRetrieval.UnitTests/Retrieval/MultiQueryExpanderTests.cs b/MEDI/test/DataRetrieval.UnitTests/Retrieval/MultiQueryExpanderTests.cs new file mode 100644 index 0000000..62cd6e0 --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Retrieval/MultiQueryExpanderTests.cs @@ -0,0 +1,82 @@ +using Microsoft.Extensions.DataRetrieval; +using CommunityToolkit.DataRetrieval.UnitTests.Utils; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Retrieval; + +public class MultiQueryExpanderTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", () => new MultiQueryExpander(null!)); + } + + [Fact] + public async Task ProcessAsync_AddsVariantsAfterOriginal() + { + using var client = TestChatClient.WithJsonResponse( + """{"variants": ["how to configure RAG", "RAG pipeline setup guide", "setting up retrieval augmented generation"]}"""); + var expander = new MultiQueryExpander(client); + var query = new RetrievalQuery("how to set up RAG"); + + var output = await expander.ProcessAsync(query); + + Assert.Equal("how to set up RAG", output.Variants[0]); // Original first + Assert.True(output.Variants.Count >= 2); // At least original + 1 variant + } + + [Fact] + public async Task ProcessAsync_FiltersShortVariants() + { + using var client = TestChatClient.WithJsonResponse( + """{"variants": ["short", "", "this is a valid query variant about RAG"]}"""); + var expander = new MultiQueryExpander(client); + var query = new RetrievalQuery("original query text"); + + var output = await expander.ProcessAsync(query); + + // "short" (5 chars) and "" are filtered out, only the valid variant remains + Assert.Equal(2, output.Variants.Count); // original + 1 valid + Assert.DoesNotContain("short", output.Variants); + } + + [Fact] + public async Task ProcessAsync_OriginalAlwaysFirst() + { + using var client = TestChatClient.WithJsonResponse( + """{"variants": ["variant one that is long enough", "variant two that is long enough"]}"""); + var expander = new MultiQueryExpander(client); + var query = new RetrievalQuery("my original question"); + + var output = await expander.ProcessAsync(query); + + Assert.Equal("my original question", output.Variants[0]); + } + + [Fact] + public async Task ProcessAsync_LlmThrows_ReturnsOriginalOnly() + { + using var client = TestChatClient.WithException(new InvalidOperationException("fail")); + var expander = new MultiQueryExpander(client); + var query = new RetrievalQuery("my question"); + + var output = await expander.ProcessAsync(query); + + Assert.Single(output.Variants); + Assert.Equal("my question", output.Variants[0]); + } + + [Fact] + public async Task ProcessAsync_LimitsToVariantCount() + { + using var client = TestChatClient.WithJsonResponse( + """{"variants": ["variant one is long enough", "variant two is long enough", "variant three is long enough", "variant four is long enough", "variant five is long enough"]}"""); + var expander = new MultiQueryExpander(client) { VariantCount = 2 }; + var query = new RetrievalQuery("test"); + + var output = await expander.ProcessAsync(query); + + // 1 original + at most 2 variants + Assert.True(output.Variants.Count <= 3); + } +} diff --git a/MEDI/test/DataRetrieval.UnitTests/Retrieval/SelfRagOrchestratorTests.cs b/MEDI/test/DataRetrieval.UnitTests/Retrieval/SelfRagOrchestratorTests.cs new file mode 100644 index 0000000..552da32 --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Retrieval/SelfRagOrchestratorTests.cs @@ -0,0 +1,147 @@ +using Microsoft.Extensions.DataRetrieval; +using CommunityToolkit.DataRetrieval.UnitTests.Utils; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Retrieval; + +public class SelfRagOrchestratorTests +{ + [Fact] + public void Constructor_ThrowsOnNullChatClient() + { + Assert.Throws("chatClient", () => new SelfRagOrchestrator(null!)); + } + + [Fact] + public async Task GenerateAsync_NeedsRetrieval_RetrievesChunks() + { + int callCount = 0; + using var client = new TestChatClient + { + GetResponseAsyncCallback = (messages, options, ct) => + { + callCount++; + string json = callCount switch + { + 1 => """{"needsRetrieval": true, "reasoning": "technical question"}""", + 2 => "The answer is based on retrieved context.", + 3 => """{"relevance": 5, "faithfulness": 5, "critique": "excellent"}""", + _ => "{}" + }; + return Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, json))); + } + }; + + var collection = new TestVectorStoreCollection(); + collection.AddSearchResult("relevant chunk", 0.9); + + var orchestrator = new SelfRagOrchestrator(client); + var result = await orchestrator.GenerateAsync( + collection, "How does DI work?", + r => (string)r["content"]!); + + Assert.True(result.NeedsRetrieval); + Assert.NotEmpty(result.RetrievedChunks); + Assert.NotEmpty(result.Answer); + } + + [Fact] + public async Task GenerateAsync_NoRetrieval_GeneratesFromKnowledge() + { + int callCount = 0; + using var client = new TestChatClient + { + GetResponseAsyncCallback = (messages, options, ct) => + { + callCount++; + string json = callCount switch + { + 1 => """{"needsRetrieval": false, "reasoning": "general knowledge"}""", + 2 => "The answer from general knowledge.", + 3 => """{"relevance": 4, "faithfulness": 4, "critique": "good"}""", + _ => "{}" + }; + return Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, json))); + } + }; + + var collection = new TestVectorStoreCollection(); + var orchestrator = new SelfRagOrchestrator(client); + var result = await orchestrator.GenerateAsync( + collection, "What is water?", + r => (string)r["content"]!); + + Assert.False(result.NeedsRetrieval); + Assert.Empty(result.RetrievedChunks); + } + + [Fact] + public async Task GenerateAsync_LowScore_RetriesWithForcedRetrieval() + { + int callCount = 0; + using var client = new TestChatClient + { + GetResponseAsyncCallback = (messages, options, ct) => + { + callCount++; + string json = callCount switch + { + 1 => """{"needsRetrieval": false, "reasoning": "try without"}""", + 2 => "Weak answer.", + 3 => """{"relevance": 1, "faithfulness": 1, "critique": "poor"}""", + 4 => "Better answer with context.", + 5 => """{"relevance": 4, "faithfulness": 4, "critique": "improved"}""", + _ => "{}" + }; + return Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, json))); + } + }; + + var collection = new TestVectorStoreCollection(); + collection.AddSearchResult("helpful context", 0.9); + + var orchestrator = new SelfRagOrchestrator(client) { AcceptanceThreshold = 3.0 }; + var result = await orchestrator.GenerateAsync( + collection, "complex question", + r => (string)r["content"]!); + + Assert.True(callCount >= 4, "Should have retried after low score"); + } + + [Fact] + public async Task GenerateAsync_MaxRetriesZero_NoRetryEvenIfLowScore() + { + int callCount = 0; + using var client = new TestChatClient + { + GetResponseAsyncCallback = (messages, options, ct) => + { + callCount++; + string json = callCount switch + { + 1 => """{"needsRetrieval": false, "reasoning": "try"}""", + 2 => "Bad answer.", + 3 => """{"relevance": 1, "faithfulness": 1, "critique": "terrible"}""", + _ => "{}" + }; + return Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, json))); + } + }; + + var collection = new TestVectorStoreCollection(); + var orchestrator = new SelfRagOrchestrator(client) { MaxRetries = 0 }; + var result = await orchestrator.GenerateAsync( + collection, "test", + r => (string)r["content"]!); + + Assert.Equal(3, callCount); // No retry calls + Assert.Equal(1.0, result.AverageScore); + } +} diff --git a/MEDI/test/DataRetrieval.UnitTests/Retrieval/SpeculativeRagOrchestratorTests.cs b/MEDI/test/DataRetrieval.UnitTests/Retrieval/SpeculativeRagOrchestratorTests.cs new file mode 100644 index 0000000..602e7c9 --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Retrieval/SpeculativeRagOrchestratorTests.cs @@ -0,0 +1,110 @@ +using Microsoft.Extensions.DataRetrieval; +using CommunityToolkit.DataRetrieval.UnitTests.Utils; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Retrieval; + +public class SpeculativeRagOrchestratorTests +{ + [Fact] + public void Constructor_ThrowsOnNullDrafterClient() + { + using var verifier = TestChatClient.WithJsonResponse("{}"); + Assert.Throws("drafterClient", + () => new SpeculativeRagOrchestrator(null!, verifier)); + } + + [Fact] + public void Constructor_ThrowsOnNullVerifierClient() + { + using var drafter = TestChatClient.WithJsonResponse("{}"); + Assert.Throws("verifierClient", + () => new SpeculativeRagOrchestrator(drafter, null!)); + } + + [Fact] + public async Task GenerateAsync_ProducesDraftsInParallel() + { + using var drafter = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, + "Draft answer about the topic."))) + }; + using var verifier = TestChatClient.WithJsonResponse( + """{"bestDraftIndex": 1, "confidence": 0.9, "reasoning": "draft 1 is best"}"""); + + var collection = new TestVectorStoreCollection(); + collection.AddSearchResult("chunk1", 0.9); + collection.AddSearchResult("chunk2", 0.8); + + var orchestrator = new SpeculativeRagOrchestrator(drafter, verifier) + { + DraftCount = 3 + }; + var result = await orchestrator.GenerateAsync( + collection, "test query", + r => (string)r["content"]!); + + Assert.Equal(3, result.Drafts.Count); + Assert.NotEmpty(result.Answer); + Assert.True(result.ParallelDraftMs >= 0); + } + + [Fact] + public async Task GenerateAsync_SelectsBestDraft() + { + int draftIndex = 0; + using var drafter = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + var idx = Interlocked.Increment(ref draftIndex); + return Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, + $"Draft {idx} answer."))); + } + }; + using var verifier = TestChatClient.WithJsonResponse( + """{"bestDraftIndex": 2, "confidence": 0.85, "reasoning": "draft 2 most accurate"}"""); + + var collection = new TestVectorStoreCollection(); + collection.AddSearchResult("context", 0.9); + + var orchestrator = new SpeculativeRagOrchestrator(drafter, verifier) { DraftCount = 3 }; + var result = await orchestrator.GenerateAsync( + collection, "query", + r => (string)r["content"]!); + + Assert.Equal(2, result.BestDraftIndex); + Assert.Equal(0.85, result.Confidence); + } + + [Fact] + public async Task GenerateAsync_VerificationFails_FallsBackToFirstDraft() + { + using var drafter = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new Microsoft.Extensions.AI.ChatResponse( + new Microsoft.Extensions.AI.ChatMessage( + Microsoft.Extensions.AI.ChatRole.Assistant, + "Fallback draft."))) + }; + using var verifier = TestChatClient.WithException(new InvalidOperationException("verification failed")); + + var collection = new TestVectorStoreCollection(); + collection.AddSearchResult("context", 0.9); + + var orchestrator = new SpeculativeRagOrchestrator(drafter, verifier) { DraftCount = 2 }; + var result = await orchestrator.GenerateAsync( + collection, "query", + r => (string)r["content"]!); + + Assert.Equal(1, result.BestDraftIndex); + Assert.Equal(0.5, result.Confidence); + Assert.Equal("verification failed", result.VerificationReasoning); + } +} diff --git a/MEDI/test/DataRetrieval.UnitTests/Retrieval/TreeSearchRetrieverTests.cs b/MEDI/test/DataRetrieval.UnitTests/Retrieval/TreeSearchRetrieverTests.cs new file mode 100644 index 0000000..2c802bc --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Retrieval/TreeSearchRetrieverTests.cs @@ -0,0 +1,56 @@ +using Microsoft.Extensions.DataRetrieval; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Retrieval; + +public class TreeSearchRetrieverTests +{ + [Fact] + public async Task ProcessAsync_SetsTreeTraversalMetadata() + { + var retriever = new TreeSearchRetriever(); + var query = new RetrievalQuery("broad question about architecture"); + + var output = await retriever.ProcessAsync(query); + + Assert.Equal("TreeTraversal", output.Metadata["search_paradigm"]); + } + + [Fact] + public async Task ProcessAsync_SetsResultsPerLevel() + { + var retriever = new TreeSearchRetriever { ResultsPerLevel = 5 }; + var query = new RetrievalQuery("test"); + + var output = await retriever.ProcessAsync(query); + + Assert.Equal(5, output.Metadata["results_per_level"]); + } + + [Fact] + public async Task TraverseAsync_GroupsByLevelAndReturnsOrdered() + { + var collection = new Utils.TestVectorStoreCollection(); + collection.AddSearchResult("root summary", 0.7, + new Dictionary { ["level"] = 2, ["id"] = "root", ["parent_id"] = "" }); + collection.AddSearchResult("branch summary", 0.8, + new Dictionary { ["level"] = 1, ["id"] = "branch-1", ["parent_id"] = "root" }); + collection.AddSearchResult("leaf content", 0.9, + new Dictionary { ["level"] = 0, ["id"] = "leaf-1", ["parent_id"] = "branch-1" }); + + var results = await TreeSearchRetriever.TraverseAsync( + collection, + "test query", + contentSelector: r => (string)r["content"]!, + levelSelector: r => (int)r["level"]!, + parentIdSelector: r => r["parent_id"] as string, + idSelector: r => r["id"] as string, + topK: 5); + + Assert.True(results.Count > 0); + // Results should be ordered by score descending + for (int i = 1; i < results.Count; i++) + { + Assert.True(results[i - 1].Score >= results[i].Score); + } + } +} diff --git a/MEDI/test/DataRetrieval.UnitTests/Utils/TestChatClient.cs b/MEDI/test/DataRetrieval.UnitTests/Utils/TestChatClient.cs new file mode 100644 index 0000000..5cde1ae --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Utils/TestChatClient.cs @@ -0,0 +1,45 @@ +using Microsoft.Extensions.AI; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Utils; + +/// +/// Callback-based fake for unit testing. +/// +public sealed class TestChatClient : IChatClient +{ + public IServiceProvider? Services { get; set; } + + public Func, ChatOptions?, CancellationToken, Task>? + GetResponseAsyncCallback { get; set; } + + public Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + => GetResponseAsyncCallback?.Invoke(messages, options, cancellationToken) + ?? Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "default"))); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + => throw new NotImplementedException("Streaming not used in tests."); + + public object? GetService(Type serviceType, object? serviceKey = null) + => serviceType is not null && serviceKey is null && serviceType.IsInstanceOfType(this) ? this : null; + + public void Dispose() { } + + /// Creates a TestChatClient that returns the given JSON string as response text. + public static TestChatClient WithJsonResponse(string json) => new() + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, json))) + }; + + /// Creates a TestChatClient that throws the given exception. + public static TestChatClient WithException(Exception ex) => new() + { + GetResponseAsyncCallback = (_, _, _) => throw ex + }; +} diff --git a/MEDI/test/DataRetrieval.UnitTests/Utils/TestVectorStoreCollection.cs b/MEDI/test/DataRetrieval.UnitTests/Utils/TestVectorStoreCollection.cs new file mode 100644 index 0000000..a59fec7 --- /dev/null +++ b/MEDI/test/DataRetrieval.UnitTests/Utils/TestVectorStoreCollection.cs @@ -0,0 +1,90 @@ +using System.Linq.Expressions; +using System.Runtime.CompilerServices; +using Microsoft.Extensions.VectorData; + +namespace CommunityToolkit.DataRetrieval.UnitTests.Utils; + +/// +/// In-memory fake for testing +/// components that depend on vector search. +/// +public sealed class TestVectorStoreCollection : VectorStoreCollection> +{ + private readonly List>> _searchResults = []; + + public override string Name => "test-collection"; + + /// Configures search results to return. + public void SetSearchResults(params VectorSearchResult>[] results) + { + _searchResults.Clear(); + _searchResults.AddRange(results); + } + + /// Adds a search result with the given content and score. + public void AddSearchResult(string content, double score, Dictionary? extraProps = null) + { + var record = new Dictionary + { + ["content"] = content, + ["score"] = score + }; + + if (extraProps is not null) + { + foreach (var kv in extraProps) + record[kv.Key] = kv.Value; + } + + _searchResults.Add(new(record, score)); + } + + public override async IAsyncEnumerable>> SearchAsync( + TInput query, + int top = 5, + VectorSearchOptions>? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + await Task.CompletedTask; + foreach (var result in _searchResults.Take(top)) + { + cancellationToken.ThrowIfCancellationRequested(); + yield return result; + } + } + + public override Task CollectionExistsAsync(CancellationToken cancellationToken = default) + => Task.FromResult(true); + + public override Task EnsureCollectionExistsAsync(CancellationToken cancellationToken = default) + => Task.CompletedTask; + + public override Task EnsureCollectionDeletedAsync(CancellationToken cancellationToken = default) + => Task.CompletedTask; + + public override Task?> GetAsync( + string key, + RecordRetrievalOptions? options = null, + CancellationToken cancellationToken = default) + => Task.FromResult?>(null); + + public override Task DeleteAsync(string key, CancellationToken cancellationToken = default) + => Task.CompletedTask; + + public override Task UpsertAsync(Dictionary record, CancellationToken cancellationToken = default) + => Task.FromResult("test-key"); + + public override Task> UpsertAsync( + IEnumerable> records, + CancellationToken cancellationToken = default) + => Task.FromResult>(["test-key"]); + + public override IAsyncEnumerable> GetAsync( + Expression, bool>> filter, + int top = 100, + FilteredRecordRetrievalOptions>? options = null, + CancellationToken cancellationToken = default) + => AsyncEnumerable.Empty>(); + + public override object? GetService(Type serviceType, object? serviceKey = null) => null; +} diff --git a/MEDI/test/Directory.Build.props b/MEDI/test/Directory.Build.props new file mode 100644 index 0000000..2cf1445 --- /dev/null +++ b/MEDI/test/Directory.Build.props @@ -0,0 +1,16 @@ + + + + + + $(NoWarn);CA1515 + $(NoWarn);CA1707 + $(NoWarn);CA1716 + $(NoWarn);CA1720 + $(NoWarn);CA1819 + $(NoWarn);CA1861 + $(NoWarn);CA2007;VSTHRD111 + $(NoWarn);CS1591 + + + diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/ContextualChunkEnricherTests.cs b/MEDI/test/PdfPig.DataIngestion.UnitTests/ContextualChunkEnricherTests.cs new file mode 100644 index 0000000..c5a1110 --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/ContextualChunkEnricherTests.cs @@ -0,0 +1,197 @@ +#if NET8_0_OR_GREATER +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; +using CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.Processors; +using Xunit; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests; + +public class ContextualChunkEnricherTests +{ + private class TestChatClient : IChatClient + { + private readonly string _response; + + public TestChatClient(string response) => _response = response; + + public Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + var msg = new ChatMessage(ChatRole.Assistant, _response); + return Task.FromResult(new ChatResponse(msg)); + } + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + => throw new NotImplementedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() { } + } + + [Fact] + public void Constructor_NullChatClient_ThrowsArgumentNullException() + { + Assert.Throws(() => new ContextualChunkEnricher(null!)); + } + + [Fact] + public async Task ProcessAsync_AddsContextualSummaryMetadata() + { + var summaryText = "This chunk discusses PDF text extraction."; + var client = new TestChatClient(summaryText); + var enricher = new ContextualChunkEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var chunk = new IngestionChunk("Some text about PDF extraction and analysis.", doc, "Page 1"); + + var results = new List>(); + await foreach (var c in enricher.ProcessAsync(ToAsyncEnumerable(chunk))) + { + results.Add(c); + } + + Assert.Single(results); + Assert.True(results[0].Metadata.ContainsKey("contextual_summary")); + Assert.Equal(summaryText, results[0].Metadata["contextual_summary"]); + } + + [Fact] + public async Task ProcessAsync_MultipleChunks_AllGetSummaries() + { + var summaryText = "Summary text."; + var client = new TestChatClient(summaryText); + var enricher = new ContextualChunkEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var chunk1 = new IngestionChunk("First chunk content.", doc, "Page 1"); + var chunk2 = new IngestionChunk("Second chunk content.", doc, "Page 2"); + + var results = new List>(); + await foreach (var c in enricher.ProcessAsync(ToAsyncEnumerable(chunk1, chunk2))) + { + results.Add(c); + } + + Assert.Equal(2, results.Count); + Assert.All(results, r => + { + Assert.True(r.Metadata.ContainsKey("contextual_summary")); + Assert.Equal(summaryText, r.Metadata["contextual_summary"]); + }); + } + + [Fact] + public async Task ProcessAsync_EmptyChunks_NoResults() + { + var client = new TestChatClient("summary"); + var enricher = new ContextualChunkEnricher(client); + + var results = new List>(); + await foreach (var c in enricher.ProcessAsync(ToAsyncEnumerable>())) + { + results.Add(c); + } + + Assert.Empty(results); + } + + [Fact] + public async Task ProcessAsync_CancellationRequested_Throws() + { + var client = new TestChatClient("summary"); + var enricher = new ContextualChunkEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var chunk = new IngestionChunk("Some content.", doc, "Page 1"); + + using var cts = new CancellationTokenSource(); + cts.Cancel(); + + await Assert.ThrowsAnyAsync(async () => + { + await foreach (var _ in enricher.ProcessAsync(ToAsyncEnumerable(chunk), cts.Token)) + { + } + }); + } + + [Fact] + public async Task ProcessAsync_PreservesChunkContent() + { + var client = new TestChatClient("A summary."); + var enricher = new ContextualChunkEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var originalContent = "Original chunk content stays the same."; + var chunk = new IngestionChunk(originalContent, doc, "Page 1"); + + var results = new List>(); + await foreach (var c in enricher.ProcessAsync(ToAsyncEnumerable(chunk))) + { + results.Add(c); + } + + Assert.Single(results); + Assert.Equal(originalContent, results[0].Content); + } + + [Fact] + public async Task ProcessAsync_PreservesChunkDocumentReference() + { + var client = new TestChatClient("A summary."); + var enricher = new ContextualChunkEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var chunk = new IngestionChunk("Content.", doc, "Page 1"); + + var results = new List>(); + await foreach (var c in enricher.ProcessAsync(ToAsyncEnumerable(chunk))) + { + results.Add(c); + } + + Assert.Single(results); + Assert.Same(doc, results[0].Document); + } + + [Fact] + public async Task ProcessAsync_YieldsSameChunkInstances() + { + var client = new TestChatClient("Summary."); + var enricher = new ContextualChunkEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var chunk = new IngestionChunk("Content.", doc, "Page 1"); + + var results = new List>(); + await foreach (var c in enricher.ProcessAsync(ToAsyncEnumerable(chunk))) + { + results.Add(c); + } + + Assert.Single(results); + Assert.Same(chunk, results[0]); + } + + private static async IAsyncEnumerable ToAsyncEnumerable(params T[] items) + { + foreach (var item in items) + { + await Task.Yield(); + yield return item; + } + } +} +#endif diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/Documents/cat-genetics.pdf b/MEDI/test/PdfPig.DataIngestion.UnitTests/Documents/cat-genetics.pdf new file mode 100644 index 0000000..90c7afe Binary files /dev/null and b/MEDI/test/PdfPig.DataIngestion.UnitTests/Documents/cat-genetics.pdf differ diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/Documents/data.pdf b/MEDI/test/PdfPig.DataIngestion.UnitTests/Documents/data.pdf new file mode 100644 index 0000000..8bdda7f Binary files /dev/null and b/MEDI/test/PdfPig.DataIngestion.UnitTests/Documents/data.pdf differ diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/Documents/inherited_mediabox.pdf b/MEDI/test/PdfPig.DataIngestion.UnitTests/Documents/inherited_mediabox.pdf new file mode 100644 index 0000000..2d4468c Binary files /dev/null and b/MEDI/test/PdfPig.DataIngestion.UnitTests/Documents/inherited_mediabox.pdf differ diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/IntegrationHelpers.cs b/MEDI/test/PdfPig.DataIngestion.UnitTests/IntegrationHelpers.cs new file mode 100644 index 0000000..310985b --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/IntegrationHelpers.cs @@ -0,0 +1,33 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests; + +/// +/// Helper to locate test PDF documents for integration tests. +/// Looks in a "Documents" folder relative to the test output directory. +/// +internal static class IntegrationHelpers +{ + private static readonly string DocumentsFolder = Path.Combine( + AppDomain.CurrentDomain.BaseDirectory, "Documents"); + + public static string GetDocumentPath(string name) + { + // Try with .pdf extension first + var path = Path.Combine(DocumentsFolder, name + ".pdf"); + if (File.Exists(path)) + { + return path; + } + + // Try without extension + path = Path.Combine(DocumentsFolder, name); + if (File.Exists(path)) + { + return path; + } + + // Return the .pdf path (caller's test will fail with a clear "file not found" message) + return Path.Combine(DocumentsFolder, name + ".pdf"); + } +} +#endif diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/MetadataAwareSectionChunkerTests.cs b/MEDI/test/PdfPig.DataIngestion.UnitTests/MetadataAwareSectionChunkerTests.cs new file mode 100644 index 0000000..a9040c5 --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/MetadataAwareSectionChunkerTests.cs @@ -0,0 +1,179 @@ +#if NET8_0_OR_GREATER +using System.Collections.Generic; +using System.Threading.Tasks; +using Microsoft.Extensions.DataIngestion; +using Microsoft.Extensions.DataIngestion.Chunkers; +using Microsoft.ML.Tokenizers; +using CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; +using Xunit; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests; + +public class MetadataAwareSectionChunkerTests +{ + private static MetadataAwareSectionChunker CreateChunker() + { + var tokenizer = TiktokenTokenizer.CreateForModel("gpt-4o"); + var options = new IngestionChunkerOptions(tokenizer) { MaxTokensPerChunk = 10000 }; + return new MetadataAwareSectionChunker(options); + } + + private static async Task>> CollectChunksAsync( + MetadataAwareSectionChunker chunker, + IngestionDocument doc) + { + var chunks = new List>(); + await foreach (var chunk in chunker.ProcessAsync(doc)) + { + chunks.Add(chunk); + } + return chunks; + } + + [Fact] + public async Task ProcessAsync_PropagatesElementTypeMetadata() + { + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + var paragraph = new IngestionDocumentParagraph("Some text") { Text = "Some text" }; + paragraph.Metadata["element_type"] = "text"; + section.Elements.Add(paragraph); + doc.Sections.Add(section); + + var chunker = CreateChunker(); + var chunks = await CollectChunksAsync(chunker, doc); + + Assert.Single(chunks); + Assert.True(chunks[0].Metadata.ContainsKey("element_type")); + Assert.Equal("text", chunks[0].Metadata["element_type"]); + } + + [Fact] + public async Task ProcessAsync_SkipsBoundingBoxMetadata() + { + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + var paragraph = new IngestionDocumentParagraph("Heading text") { Text = "Heading text" }; + paragraph.Metadata["BoundingBox.Left"] = 10.0; + paragraph.Metadata["element_type"] = "heading"; + section.Elements.Add(paragraph); + doc.Sections.Add(section); + + var chunker = CreateChunker(); + var chunks = await CollectChunksAsync(chunker, doc); + + Assert.Single(chunks); + Assert.True(chunks[0].Metadata.ContainsKey("element_type")); + Assert.Equal("heading", chunks[0].Metadata["element_type"]); + Assert.False(chunks[0].Metadata.ContainsKey("BoundingBox.Left")); + } + + [Fact] + public async Task ProcessAsync_FirstMatchWinsForConflictingKeys() + { + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + + var paragraph1 = new IngestionDocumentParagraph("First paragraph") { Text = "First paragraph" }; + paragraph1.Metadata["element_type"] = "text"; + section.Elements.Add(paragraph1); + + var paragraph2 = new IngestionDocumentParagraph("Second paragraph") { Text = "Second paragraph" }; + paragraph2.Metadata["element_type"] = "heading"; + section.Elements.Add(paragraph2); + + doc.Sections.Add(section); + + var chunker = CreateChunker(); + var chunks = await CollectChunksAsync(chunker, doc); + + Assert.Single(chunks); + Assert.Equal("text", chunks[0].Metadata["element_type"]); + } + + [Fact] + public async Task ProcessAsync_NullMetadataValuesSkipped() + { + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + var paragraph = new IngestionDocumentParagraph("Null meta text") { Text = "Null meta text" }; + paragraph.Metadata["element_type"] = null; + section.Elements.Add(paragraph); + doc.Sections.Add(section); + + var chunker = CreateChunker(); + var chunks = await CollectChunksAsync(chunker, doc); + + Assert.Single(chunks); + Assert.False(chunks[0].Metadata.ContainsKey("element_type")); + } + + [Fact] + public async Task ProcessAsync_ElementWithoutMetadata_Skipped() + { + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + var paragraph = new IngestionDocumentParagraph("Plain text") { Text = "Plain text" }; + section.Elements.Add(paragraph); + doc.Sections.Add(section); + + var chunker = CreateChunker(); + var chunks = await CollectChunksAsync(chunker, doc); + + Assert.Single(chunks); + Assert.Empty(chunks[0].Metadata); + } + + [Fact] + public async Task ProcessAsync_EmptyTextElement_Skipped() + { + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + + var emptyParagraph = new IngestionDocumentParagraph(" ") { Text = "" }; + emptyParagraph.Metadata["element_type"] = "empty"; + section.Elements.Add(emptyParagraph); + + var visibleParagraph = new IngestionDocumentParagraph("Visible text") { Text = "Visible text" }; + section.Elements.Add(visibleParagraph); + + doc.Sections.Add(section); + + var chunker = CreateChunker(); + var chunks = await CollectChunksAsync(chunker, doc); + + Assert.Single(chunks); + Assert.False(chunks[0].Metadata.ContainsKey("element_type")); + } + + [Fact] + public async Task ProcessAsync_SubstringMatchIsCaseSensitive() + { + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + var paragraph = new IngestionDocumentParagraph("Hello World") { Text = "Hello World" }; + paragraph.Metadata["element_type"] = "greeting"; + section.Elements.Add(paragraph); + doc.Sections.Add(section); + + var chunker = CreateChunker(); + var chunks = await CollectChunksAsync(chunker, doc); + + Assert.Single(chunks); + Assert.Contains("Hello World", chunks[0].Content); + Assert.True(chunks[0].Metadata.ContainsKey("element_type")); + Assert.Equal("greeting", chunks[0].Metadata["element_type"]); + } + + [Fact] + public async Task ProcessAsync_EmptyDocument_ReturnsNoChunks() + { + var doc = new IngestionDocument("test.pdf"); + + var chunker = CreateChunker(); + var chunks = await CollectChunksAsync(chunker, doc); + + Assert.Empty(chunks); + } +} +#endif diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/PageImageRendererTests.cs b/MEDI/test/PdfPig.DataIngestion.UnitTests/PageImageRendererTests.cs new file mode 100644 index 0000000..c5b05bf --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/PageImageRendererTests.cs @@ -0,0 +1,92 @@ +#if NET8_0_OR_GREATER +using System; +using System.IO; +using UglyToad.PdfPig; +using UglyToad.PdfPig.Core; +using CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; +using Xunit; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests; + +public class PageImageRendererTests +{ + [Fact] + public void RenderPage_NullPage_ThrowsArgumentNullException() + { + Assert.Throws(() => PageImageRenderer.RenderPage(null!)); + } + + [Fact] + public void RenderRegion_NullPage_ThrowsArgumentNullException() + { + var region = new PdfRectangle(0, 0, 100, 100); + Assert.Throws(() => PageImageRenderer.RenderRegion(null!, region)); + } + + [Fact] + public void RenderPage_ReturnsValidPng() + { + var path = IntegrationHelpers.GetDocumentPath("data"); + using var pdfDoc = PdfDocument.Open(path); + var page = pdfDoc.GetPage(1); + + var result = PageImageRenderer.RenderPage(page); + + Assert.NotNull(result); + Assert.True(result.Length > 4, "PNG output should have more than 4 bytes."); + // PNG magic bytes: 0x89 0x50 0x4E 0x47 + Assert.Equal(0x89, result[0]); + Assert.Equal(0x50, result[1]); + Assert.Equal(0x4E, result[2]); + Assert.Equal(0x47, result[3]); + } + + [Fact] + public void RenderPage_DifferentDpiProducesDifferentSizedOutput() + { + var path = IntegrationHelpers.GetDocumentPath("data"); + using var pdfDoc = PdfDocument.Open(path); + var page = pdfDoc.GetPage(1); + + var lowDpi = PageImageRenderer.RenderPage(page, dpi: 72); + var highDpi = PageImageRenderer.RenderPage(page, dpi: 300); + + // Higher DPI should produce a larger image + Assert.True(highDpi.Length > lowDpi.Length, + $"300 DPI image ({highDpi.Length} bytes) should be larger than 72 DPI image ({lowDpi.Length} bytes)."); + } + + [Fact] + public void RenderRegion_ReturnsValidPng() + { + var path = IntegrationHelpers.GetDocumentPath("data"); + using var pdfDoc = PdfDocument.Open(path); + var page = pdfDoc.GetPage(1); + + var region = new PdfRectangle(0, 0, page.Width, page.Height); + var result = PageImageRenderer.RenderRegion(page, region); + + Assert.NotNull(result); + Assert.True(result.Length > 4); + Assert.Equal(0x89, result[0]); + Assert.Equal(0x50, result[1]); + Assert.Equal(0x4E, result[2]); + Assert.Equal(0x47, result[3]); + } + + [Fact] + public void RenderRegion_SmallRegion_ProducesSmallerOutputThanFullPage() + { + var path = IntegrationHelpers.GetDocumentPath("data"); + using var pdfDoc = PdfDocument.Open(path); + var page = pdfDoc.GetPage(1); + + var fullPage = PageImageRenderer.RenderPage(page, dpi: 150); + var smallRegion = new PdfRectangle(0, 0, page.Width / 4, page.Height / 4); + var regionImage = PageImageRenderer.RenderRegion(page, smallRegion, dpi: 150); + + Assert.True(regionImage.Length < fullPage.Length, + $"Region image ({regionImage.Length} bytes) should be smaller than full page ({fullPage.Length} bytes)."); + } +} +#endif diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPig.DataIngestion.UnitTests.csproj b/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPig.DataIngestion.UnitTests.csproj new file mode 100644 index 0000000..f31cb2f --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPig.DataIngestion.UnitTests.csproj @@ -0,0 +1,41 @@ + + + + CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests + CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests + net10.0 + true + false + + + + + + + + + + + + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + + + + PreserveNewest + + + + + + + + diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPigReaderDiTests.cs b/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPigReaderDiTests.cs new file mode 100644 index 0000000..7976286 --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPigReaderDiTests.cs @@ -0,0 +1,118 @@ +#if NET8_0_OR_GREATER +using System; +using System.IO; +using System.Linq; +using System.Threading.Tasks; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Options; +using UglyToad.PdfPig.Content; +using CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; +using Xunit; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests; + +public class PdfPigReaderDiTests +{ + [Fact] + public void AddPdfPigReader_RegistersPdfPigReader() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(); + + Assert.Contains(services, d => d.ServiceType == typeof(PdfPigReader)); + } + + [Fact] + public void AddPdfPigReader_ConfigureAppliesOptions() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(o => + { + o.Mode = PdfReadingMode.Hybrid; + o.RenderDpi = 300; + }); + var provider = services.BuildServiceProvider(); + + var options = provider.GetRequiredService>(); + var opts = options.Value; + + Assert.Equal(PdfReadingMode.Hybrid, opts.Mode); + Assert.Equal(300, opts.RenderDpi); + } + + [Fact] + public void AddPdfPigReader_DefaultOptionsWithoutConfigure() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(); + var provider = services.BuildServiceProvider(); + + var options = provider.GetRequiredService>(); + var opts = options.Value; + + Assert.Equal(PdfReadingMode.TextOnly, opts.Mode); + Assert.Equal(150, opts.RenderDpi); + } + + [Fact] + public void AddPdfPigReader_InvalidDpi_ThrowsOnResolve() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(o => o.RenderDpi = 0); + var provider = services.BuildServiceProvider(); + + var options = provider.GetRequiredService>(); + var ex = Assert.Throws(() => options.Value); + + Assert.Contains("RenderDpi", ex.Message); + } + + [Fact] + public void IOptionsCtor_NullOptions_Throws() + { + Assert.Throws( + () => new PdfPigReader((IOptions)null!)); + } + + [Fact] + public async Task IOptionsCtor_UsesOptionsValues() + { + var options = Options.Create(new PdfPigReaderOptions + { + Mode = PdfReadingMode.Hybrid, + RenderDpi = 72 + }); + var reader = new PdfPigReader(options); + var pdfBytes = CreateBlankPagePdf(); + + using var stream = new MemoryStream(pdfBytes); + var doc = await reader.ReadAsync(stream, "blank.pdf", "application/pdf"); + + Assert.Single(doc.Sections); + var section = doc.Sections[0]; + Assert.True(section.Metadata.ContainsKey("page_image"), + "Hybrid mode should render page images."); + } + + [Fact] + public void AddPdfPigReader_IdempotentRegistration() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(); + services.AddPdfPigReader(); + + var readerDescriptors = services + .Where(d => d.ServiceType == typeof(PdfPigReader)) + .ToList(); + + Assert.Single(readerDescriptors); + } + + private static byte[] CreateBlankPagePdf() + { + using var builder = new UglyToad.PdfPig.Writer.PdfDocumentBuilder(); + builder.AddPage(PageSize.A4); + return builder.Build(); + } +} +#endif diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPigReaderOptionsValidatorTests.cs b/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPigReaderOptionsValidatorTests.cs new file mode 100644 index 0000000..7162349 --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPigReaderOptionsValidatorTests.cs @@ -0,0 +1,82 @@ +#if NET8_0_OR_GREATER +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Options; +using CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; +using Xunit; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests; + +public class PdfPigReaderOptionsValidatorTests +{ + [Fact] + public void DefaultOptions_Succeeds() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(); + var provider = services.BuildServiceProvider(); + + var options = provider.GetRequiredService>(); + var opts = options.Value; + + Assert.Equal(PdfReadingMode.TextOnly, opts.Mode); + Assert.Equal(150, opts.RenderDpi); + } + + [Fact] + public void RenderDpi_Zero_Fails() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(o => o.RenderDpi = 0); + var provider = services.BuildServiceProvider(); + + var options = provider.GetRequiredService>(); + var ex = Assert.Throws(() => options.Value); + + Assert.Contains("RenderDpi", ex.Message); + } + + [Fact] + public void RenderDpi_Negative_Fails() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(o => o.RenderDpi = -1); + var provider = services.BuildServiceProvider(); + + var options = provider.GetRequiredService>(); + var ex = Assert.Throws(() => options.Value); + + Assert.Contains("RenderDpi", ex.Message); + } + + [Fact] + public void RenderDpi_One_Succeeds() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(o => o.RenderDpi = 1); + var provider = services.BuildServiceProvider(); + + var options = provider.GetRequiredService>(); + var opts = options.Value; + + Assert.Equal(1, opts.RenderDpi); + } + + [Fact] + public void CustomValidOptions_Succeeds() + { + var services = new ServiceCollection(); + services.AddPdfPigReader(o => + { + o.Mode = PdfReadingMode.Hybrid; + o.RenderDpi = 300; + }); + var provider = services.BuildServiceProvider(); + + var options = provider.GetRequiredService>(); + var opts = options.Value; + + Assert.Equal(PdfReadingMode.Hybrid, opts.Mode); + Assert.Equal(300, opts.RenderDpi); + } +} +#endif diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPigReaderTests.cs b/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPigReaderTests.cs new file mode 100644 index 0000000..e95b1aa --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/PdfPigReaderTests.cs @@ -0,0 +1,499 @@ +#if NET8_0_OR_GREATER +using System; +using System.IO; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.DataIngestion; +using UglyToad.PdfPig; +using UglyToad.PdfPig.Content; +using UglyToad.PdfPig.DocumentLayoutAnalysis.PageSegmenter; +using CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion; +using Xunit; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests; + +public class PdfPigReaderTests +{ + #region Default (TextOnly) mode + + [Fact] + public async Task ReadAsync_DefaultMode_ReturnsDocumentWithSections() + { + var reader = new PdfPigReader(); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + Assert.NotNull(doc); + Assert.Equal("data.pdf", doc.Identifier); + Assert.NotEmpty(doc.Sections); + } + + [Fact] + public async Task ReadAsync_DefaultMode_SectionCountMatchesPageCount() + { + var reader = new PdfPigReader(); + var path = IntegrationHelpers.GetDocumentPath("cat-genetics"); + + int expectedPages; + using (var pdfDoc = PdfDocument.Open(path)) + { + expectedPages = pdfDoc.NumberOfPages; + } + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "cat-genetics.pdf", "application/pdf"); + + Assert.Equal(expectedPages, doc.Sections.Count); + } + + [Fact] + public async Task ReadAsync_DefaultMode_SectionsContainCorrectPageNumbers() + { + var reader = new PdfPigReader(); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + for (int i = 0; i < doc.Sections.Count; i++) + { + Assert.Equal(i + 1, doc.Sections[i].PageNumber); + } + } + + [Fact] + public async Task ReadAsync_DefaultMode_ParagraphsContainBoundingBoxMetadata() + { + var reader = new PdfPigReader(); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + var elementsWithMetadata = doc.EnumerateContent() + .Where(e => e.HasMetadata) + .ToList(); + + Assert.NotEmpty(elementsWithMetadata); + + var first = elementsWithMetadata.First(); + Assert.True(first.Metadata.ContainsKey("BoundingBox.Left")); + Assert.True(first.Metadata.ContainsKey("BoundingBox.Bottom")); + Assert.True(first.Metadata.ContainsKey("BoundingBox.Right")); + Assert.True(first.Metadata.ContainsKey("BoundingBox.Top")); + } + + [Fact] + public async Task ReadAsync_DefaultMode_BoundingBoxValuesAreNumeric() + { + var reader = new PdfPigReader(); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + var element = doc.EnumerateContent().First(e => e.HasMetadata); + + Assert.IsType(element.Metadata["BoundingBox.Left"]); + Assert.IsType(element.Metadata["BoundingBox.Bottom"]); + Assert.IsType(element.Metadata["BoundingBox.Right"]); + Assert.IsType(element.Metadata["BoundingBox.Top"]); + } + + [Fact] + public async Task ReadAsync_DefaultMode_ParagraphsHaveNonEmptyText() + { + var reader = new PdfPigReader(); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + var allElements = doc.EnumerateContent().ToList(); + + foreach (var element in allElements) + { + Assert.False(string.IsNullOrEmpty(element.Text), + "All paragraphs should have non-empty text (empty blocks are skipped)."); + } + } + + [Fact] + public async Task ReadAsync_DefaultMode_ParagraphPageNumbersMatchSectionPageNumbers() + { + var reader = new PdfPigReader(); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + foreach (var section in doc.Sections) + { + foreach (var element in section.Elements) + { + Assert.Equal(section.PageNumber, element.PageNumber); + } + } + } + + [Fact] + public async Task ReadAsync_WithExplicitSegmenter_ReturnsStructuredSections() + { + var reader = new PdfPigReader(RecursiveXYCut.Instance); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + Assert.NotNull(doc); + Assert.NotEmpty(doc.Sections); + } + + [Fact] + public async Task ReadAsync_MinimalPdf_ReturnsDocument() + { + var reader = new PdfPigReader(); + var path = IntegrationHelpers.GetDocumentPath("inherited_mediabox"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "inherited_mediabox.pdf", "application/pdf"); + + Assert.NotNull(doc); + Assert.Equal("inherited_mediabox.pdf", doc.Identifier); + } + + [Fact] + public async Task ReadAsync_CancellationAlreadyCancelled_ThrowsOperationCanceledException() + { + var reader = new PdfPigReader(); + var path = IntegrationHelpers.GetDocumentPath("data"); + using var cts = new CancellationTokenSource(); + cts.Cancel(); + + using var stream = File.OpenRead(path); + + await Assert.ThrowsAnyAsync( + () => reader.ReadAsync(stream, "data.pdf", "application/pdf", cts.Token)); + } + + #endregion + + #region TextOnly mode + + [Fact] + public async Task ReadAsync_TextOnly_DoesNotStorePageImage() + { + var reader = new PdfPigReader(mode: PdfReadingMode.TextOnly); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + foreach (var section in doc.Sections) + { + Assert.False(section.Metadata.ContainsKey("page_image"), + $"Section for page {section.PageNumber} should NOT contain page_image in TextOnly mode."); + Assert.False(section.Metadata.ContainsKey("page_width")); + Assert.False(section.Metadata.ContainsKey("page_height")); + } + } + + [Fact] + public async Task ReadAsync_TextOnly_BlankPage_NoPlaceholder() + { + var reader = new PdfPigReader(mode: PdfReadingMode.TextOnly); + var pdfBytes = CreateBlankPagePdf(); + + using var stream = new MemoryStream(pdfBytes); + var doc = await reader.ReadAsync(stream, "blank.pdf", "application/pdf"); + + Assert.Single(doc.Sections); + Assert.Empty(doc.Sections[0].Elements); + } + + #endregion + + #region Hybrid mode + + [Fact] + public async Task ReadAsync_Hybrid_StoresPageImageInSectionMetadata() + { + var reader = new PdfPigReader(mode: PdfReadingMode.Hybrid); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + foreach (var section in doc.Sections) + { + Assert.True(section.Metadata.ContainsKey("page_image"), + $"Section for page {section.PageNumber} should contain page_image in Hybrid mode."); + } + } + + [Fact] + public async Task ReadAsync_Hybrid_PageImageIsValidPng() + { + var reader = new PdfPigReader(mode: PdfReadingMode.Hybrid); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + var section = doc.Sections[0]; + var imageBytes = section.Metadata["page_image"] as byte[]; + + Assert.NotNull(imageBytes); + Assert.True(imageBytes.Length > 4); + // PNG magic bytes + Assert.Equal(0x89, imageBytes[0]); + Assert.Equal(0x50, imageBytes[1]); + Assert.Equal(0x4E, imageBytes[2]); + Assert.Equal(0x47, imageBytes[3]); + } + + [Fact] + public async Task ReadAsync_Hybrid_StoresPageDimensions() + { + var reader = new PdfPigReader(mode: PdfReadingMode.Hybrid); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + var section = doc.Sections[0]; + Assert.True(section.Metadata.ContainsKey("page_width")); + Assert.True(section.Metadata.ContainsKey("page_height")); + Assert.IsType(section.Metadata["page_width"]); + Assert.IsType(section.Metadata["page_height"]); + } + + [Fact] + public async Task ReadAsync_Hybrid_CustomDpi_ProducesValidDocument() + { + var reader = new PdfPigReader(mode: PdfReadingMode.Hybrid, renderDpi: 72); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + Assert.NotNull(doc); + Assert.NotEmpty(doc.Sections); + + var imageBytes = doc.Sections[0].Metadata["page_image"] as byte[]; + Assert.NotNull(imageBytes); + Assert.Equal(0x89, imageBytes[0]); + Assert.Equal(0x50, imageBytes[1]); + } + + [Fact] + public async Task ReadAsync_Hybrid_WithAllParameters_ProducesValidDocument() + { + var reader = new PdfPigReader( + segmenter: RecursiveXYCut.Instance, + mode: PdfReadingMode.Hybrid, + renderDpi: 200); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + Assert.NotNull(doc); + Assert.NotEmpty(doc.Sections); + Assert.True(doc.Sections[0].Metadata.ContainsKey("page_image")); + } + + [Fact] + public async Task ReadAsync_Hybrid_BlankPage_CreatesPlaceholderElement() + { + var reader = new PdfPigReader(mode: PdfReadingMode.Hybrid); + var pdfBytes = CreateBlankPagePdf(); + + using var stream = new MemoryStream(pdfBytes); + var doc = await reader.ReadAsync(stream, "blank.pdf", "application/pdf"); + + Assert.Single(doc.Sections); + var section = doc.Sections[0]; + Assert.Single(section.Elements); + + var element = section.Elements[0]; + Assert.Equal(string.Empty, element.Text); + } + + [Fact] + public async Task ReadAsync_Hybrid_BlankPage_PlaceholderHasCorrectPageNumber() + { + var reader = new PdfPigReader(mode: PdfReadingMode.Hybrid); + var pdfBytes = CreateBlankPagePdf(); + + using var stream = new MemoryStream(pdfBytes); + var doc = await reader.ReadAsync(stream, "blank.pdf", "application/pdf"); + + var element = doc.Sections[0].Elements[0]; + Assert.Equal(1, element.PageNumber); + } + + [Fact] + public async Task ReadAsync_Hybrid_BlankPage_PlaceholderHasPlaceholderMetadata() + { + var reader = new PdfPigReader(mode: PdfReadingMode.Hybrid); + var pdfBytes = CreateBlankPagePdf(); + + using var stream = new MemoryStream(pdfBytes); + var doc = await reader.ReadAsync(stream, "blank.pdf", "application/pdf"); + + var element = doc.Sections[0].Elements[0]; + Assert.True(element.HasMetadata); + Assert.True(element.Metadata.ContainsKey("placeholder")); + Assert.Equal(true, element.Metadata["placeholder"]); + } + + [Fact] + public async Task ReadAsync_Hybrid_BlankPage_SectionStillHasPageImage() + { + var reader = new PdfPigReader(mode: PdfReadingMode.Hybrid); + var pdfBytes = CreateBlankPagePdf(); + + using var stream = new MemoryStream(pdfBytes); + var doc = await reader.ReadAsync(stream, "blank.pdf", "application/pdf"); + + var section = doc.Sections[0]; + Assert.True(section.Metadata.ContainsKey("page_image")); + var imageBytes = section.Metadata["page_image"] as byte[]; + Assert.NotNull(imageBytes); + Assert.True(imageBytes.Length > 0); + } + + [Fact] + public async Task ReadAsync_Hybrid_MixedPdf_OnlyBlankPagesGetPlaceholders() + { + var reader = new PdfPigReader(mode: PdfReadingMode.Hybrid); + var pdfBytes = CreateMixedPdf(); + + using var stream = new MemoryStream(pdfBytes); + var doc = await reader.ReadAsync(stream, "mixed.pdf", "application/pdf"); + + Assert.Equal(2, doc.Sections.Count); + + // Page 1 has text — should have normal elements, no placeholder + var page1 = doc.Sections[0]; + Assert.NotEmpty(page1.Elements); + foreach (var el in page1.Elements) + { + Assert.False(string.IsNullOrEmpty(el.Text)); + Assert.False(el.Metadata.ContainsKey("placeholder")); + } + + // Page 2 is blank — should have single placeholder + var page2 = doc.Sections[1]; + Assert.Single(page2.Elements); + Assert.Equal(string.Empty, page2.Elements[0].Text); + Assert.True(page2.Elements[0].Metadata.ContainsKey("placeholder")); + Assert.Equal(2, page2.Elements[0].PageNumber); + } + + #endregion + + #region VisionOnly mode + + [Fact] + public async Task ReadAsync_VisionOnly_EveryPageGetsPlaceholder() + { + var reader = new PdfPigReader(mode: PdfReadingMode.VisionOnly); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + Assert.NotEmpty(doc.Sections); + + foreach (var section in doc.Sections) + { + Assert.Single(section.Elements); + var element = section.Elements[0]; + Assert.Equal(string.Empty, element.Text); + Assert.True(element.Metadata.ContainsKey("placeholder")); + Assert.Equal(true, element.Metadata["placeholder"]); + } + } + + [Fact] + public async Task ReadAsync_VisionOnly_AllSectionsHavePageImage() + { + var reader = new PdfPigReader(mode: PdfReadingMode.VisionOnly); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + foreach (var section in doc.Sections) + { + Assert.True(section.Metadata.ContainsKey("page_image")); + var imageBytes = section.Metadata["page_image"] as byte[]; + Assert.NotNull(imageBytes); + Assert.True(imageBytes.Length > 0); + } + } + + [Fact] + public async Task ReadAsync_VisionOnly_PlaceholderPageNumbersMatchSections() + { + var reader = new PdfPigReader(mode: PdfReadingMode.VisionOnly); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + for (int i = 0; i < doc.Sections.Count; i++) + { + Assert.Equal(i + 1, doc.Sections[i].Elements[0].PageNumber); + } + } + + [Fact] + public async Task ReadAsync_VisionOnly_SkipsTextExtraction() + { + // VisionOnly should produce placeholders even for PDFs with text + var reader = new PdfPigReader(mode: PdfReadingMode.VisionOnly); + var path = IntegrationHelpers.GetDocumentPath("data"); + + using var stream = File.OpenRead(path); + var doc = await reader.ReadAsync(stream, "data.pdf", "application/pdf"); + + // All elements should be placeholders with empty text — no native text extraction + foreach (var element in doc.EnumerateContent()) + { + Assert.Equal(string.Empty, element.Text); + Assert.True(element.Metadata.ContainsKey("placeholder")); + } + } + + #endregion + + #region Helpers + + private static byte[] CreateBlankPagePdf() + { + using var builder = new UglyToad.PdfPig.Writer.PdfDocumentBuilder(); + builder.AddPage(PageSize.A4); + return builder.Build(); + } + + private static byte[] CreateMixedPdf() + { + using var builder = new UglyToad.PdfPig.Writer.PdfDocumentBuilder(); + // Page 1: has text + var page1 = builder.AddPage(PageSize.A4); + var font = builder.AddStandard14Font(UglyToad.PdfPig.Fonts.Standard14Fonts.Standard14Font.Helvetica); + page1.AddText("Hello World", 12, new UglyToad.PdfPig.Core.PdfPoint(72, 720), font); + // Page 2: blank (no text) + builder.AddPage(PageSize.A4); + return builder.Build(); + } + + #endregion +} +#endif diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/VisionOcrEnricherTests.cs b/MEDI/test/PdfPig.DataIngestion.UnitTests/VisionOcrEnricherTests.cs new file mode 100644 index 0000000..5fcf72c --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/VisionOcrEnricherTests.cs @@ -0,0 +1,292 @@ +#if NET8_0_OR_GREATER +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; +using CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.Processors; +using Xunit; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests; + +public class VisionOcrEnricherTests +{ + private class TestChatClient : IChatClient + { + private readonly string _response; + + public List LastMessages { get; private set; } = new(); + + public TestChatClient(string response) => _response = response; + + public Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + LastMessages = messages.ToList(); + var msg = new ChatMessage(ChatRole.Assistant, _response); + return Task.FromResult(new ChatResponse(msg)); + } + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + => throw new NotImplementedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() { } + } + + [Fact] + public void Constructor_NullChatClient_ThrowsArgumentNullException() + { + Assert.Throws(() => new VisionOcrEnricher(null!)); + } + + [Fact] + public async Task ProcessAsync_EmptyTextElement_GetsOcrText() + { + var ocrText = "Extracted OCR text from image"; + var client = new TestChatClient(ocrText); + var fallback = new VisionOcrEnricher(client); + + var doc = CreateDocumentWithEmptyTextElement(); + + var result = await fallback.ProcessAsync(doc); + + var element = result.EnumerateContent().First(); + Assert.Equal(ocrText, element.Text); + } + + [Fact] + public async Task ProcessAsync_EmptyTextElement_SetsOcrSourceMetadata() + { + var client = new TestChatClient("OCR result"); + var fallback = new VisionOcrEnricher(client); + + var doc = CreateDocumentWithEmptyTextElement(); + + var result = await fallback.ProcessAsync(doc); + + var element = result.EnumerateContent().First(); + Assert.True(element.Metadata.ContainsKey("ocr_source")); + Assert.Equal("vision_llm", element.Metadata["ocr_source"]); + } + + [Fact] + public async Task ProcessAsync_ElementWithText_NotModified() + { + var client = new TestChatClient("Should not replace"); + var fallback = new VisionOcrEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + var paragraph = new IngestionDocumentParagraph("Existing content") { Text = "Existing content" }; + section.Elements.Add(paragraph); + doc.Sections.Add(section); + + var result = await fallback.ProcessAsync(doc); + + var element = result.EnumerateContent().First(); + Assert.Equal("Existing content", element.Text); + Assert.False(element.Metadata.ContainsKey("ocr_source")); + } + + [Fact] + public async Task ProcessAsync_MixedElements_OnlyEmptyOnesEnriched() + { + var ocrText = "Extracted text"; + var client = new TestChatClient(ocrText); + var fallback = new VisionOcrEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + + var withText = new IngestionDocumentParagraph("Has content") { Text = "Has content" }; + var withoutText = new IngestionDocumentParagraph("empty region") { Text = "" }; + + section.Elements.Add(withText); + section.Elements.Add(withoutText); + doc.Sections.Add(section); + + var result = await fallback.ProcessAsync(doc); + var elements = result.EnumerateContent().ToList(); + + Assert.Equal("Has content", elements[0].Text); + Assert.False(elements[0].Metadata.ContainsKey("ocr_source")); + + Assert.Equal(ocrText, elements[1].Text); + Assert.True(elements[1].Metadata.ContainsKey("ocr_source")); + } + + [Fact] + public async Task ProcessAsync_EmptyDocument_NoException() + { + var client = new TestChatClient("response"); + var fallback = new VisionOcrEnricher(client); + + var doc = new IngestionDocument("empty.pdf"); + + var result = await fallback.ProcessAsync(doc); + + Assert.NotNull(result); + } + + [Fact] + public async Task ProcessAsync_LlmReturnsWhitespace_ElementNotModified() + { + var client = new TestChatClient(" "); + var fallback = new VisionOcrEnricher(client); + + var doc = CreateDocumentWithEmptyTextElement(); + + var result = await fallback.ProcessAsync(doc); + + var element = result.EnumerateContent().First(); + Assert.True(string.IsNullOrWhiteSpace(element.Text)); + Assert.False(element.Metadata.ContainsKey("ocr_source")); + } + + [Fact] + public async Task ProcessAsync_CancellationRequested_Throws() + { + var client = new TestChatClient("response"); + var fallback = new VisionOcrEnricher(client); + + var doc = CreateDocumentWithEmptyTextElement(); + using var cts = new CancellationTokenSource(); + cts.Cancel(); + + await Assert.ThrowsAnyAsync( + () => fallback.ProcessAsync(doc, cts.Token)); + } + + [Fact] + public async Task ProcessAsync_ReturnsSameDocumentInstance() + { + var client = new TestChatClient("ocr text"); + var fallback = new VisionOcrEnricher(client); + + var doc = CreateDocumentWithEmptyTextElement(); + + var result = await fallback.ProcessAsync(doc); + + Assert.Same(doc, result); + } + + [Fact] + public async Task ProcessAsync_NullTextElement_GetsOcrText() + { + var ocrText = "OCR from null text"; + var client = new TestChatClient(ocrText); + var fallback = new VisionOcrEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + // Text defaults to null when not explicitly set + var paragraph = new IngestionDocumentParagraph("image placeholder"); + section.Elements.Add(paragraph); + doc.Sections.Add(section); + + var result = await fallback.ProcessAsync(doc); + + var element = result.EnumerateContent().First(); + // If Text was null (IsNullOrWhiteSpace), it should get OCR'd + if (string.IsNullOrWhiteSpace(element.Text) == false) + { + // Text was set by OCR + Assert.Equal(ocrText, element.Text); + } + } + + [Fact] + public async Task ProcessAsync_WithPageImage_SendsDataContentToLlm() + { + var ocrText = "Vision OCR result"; + var client = new TestChatClient(ocrText); + var fallback = new VisionOcrEnricher(client); + + var fakeImageBytes = new byte[] { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A }; + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + section.Metadata["page_image"] = fakeImageBytes; + var emptyParagraph = new IngestionDocumentParagraph("image region") { Text = "", PageNumber = 1 }; + section.Elements.Add(emptyParagraph); + doc.Sections.Add(section); + + await fallback.ProcessAsync(doc); + + var userMsg = client.LastMessages.Last(m => m.Role == ChatRole.User); + Assert.Contains(userMsg.Contents, c => c is DataContent dc && dc.MediaType == "image/png"); + } + + [Fact] + public async Task ProcessAsync_WithoutPageImage_SendsTextOnlyToLlm() + { + var ocrText = "Text fallback OCR result"; + var client = new TestChatClient(ocrText); + var fallback = new VisionOcrEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + // No page_image metadata + var emptyParagraph = new IngestionDocumentParagraph("image region") { Text = "", PageNumber = 1 }; + section.Elements.Add(emptyParagraph); + doc.Sections.Add(section); + + await fallback.ProcessAsync(doc); + + var userMsg = client.LastMessages.Last(m => m.Role == ChatRole.User); + Assert.DoesNotContain(userMsg.Contents, c => c is DataContent); + } + + [Fact] + public async Task ProcessAsync_PlaceholderElement_WithPageImage_PerformsVisionOcr() + { + var ocrText = "Text extracted from scanned page image"; + var client = new TestChatClient(ocrText); + var enricher = new VisionOcrEnricher(client); + + // Simulate what PdfPigReader produces for a scanned page + var fakeImageBytes = new byte[] { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A }; + var doc = new IngestionDocument("scanned.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + section.Metadata["page_image"] = fakeImageBytes; + + var placeholder = new IngestionDocumentParagraph("[scanned-page]") + { + Text = string.Empty, + PageNumber = 1 + }; + placeholder.Metadata["placeholder"] = true; + section.Elements.Add(placeholder); + doc.Sections.Add(section); + + var result = await enricher.ProcessAsync(doc); + + var element = result.EnumerateContent().First(); + Assert.Equal(ocrText, element.Text); + Assert.Equal("vision_llm", element.Metadata["ocr_source"]); + + // Verify vision approach was used (DataContent with image) + var userMsg = client.LastMessages.Last(m => m.Role == ChatRole.User); + Assert.Contains(userMsg.Contents, c => c is DataContent dc && dc.MediaType == "image/png"); + } + + private static IngestionDocument CreateDocumentWithEmptyTextElement() + { + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + var emptyParagraph = new IngestionDocumentParagraph("image region") { Text = "" }; + section.Elements.Add(emptyParagraph); + doc.Sections.Add(section); + return doc; + } +} +#endif diff --git a/MEDI/test/PdfPig.DataIngestion.UnitTests/VisionTableEnricherTests.cs b/MEDI/test/PdfPig.DataIngestion.UnitTests/VisionTableEnricherTests.cs new file mode 100644 index 0000000..8d3ec70 --- /dev/null +++ b/MEDI/test/PdfPig.DataIngestion.UnitTests/VisionTableEnricherTests.cs @@ -0,0 +1,242 @@ +#if NET8_0_OR_GREATER +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DataIngestion; +using CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.Processors; +using Xunit; + +namespace CommunityToolkit.DocumentProcessing.PdfPig.DataIngestion.UnitTests; + +public class VisionTableEnricherTests +{ + private class TestChatClient : IChatClient + { + private readonly string _response; + + public List LastMessages { get; private set; } = new(); + + public TestChatClient(string response) => _response = response; + + public Task GetResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + LastMessages = messages.ToList(); + var msg = new ChatMessage(ChatRole.Assistant, _response); + return Task.FromResult(new ChatResponse(msg)); + } + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, + ChatOptions? options = null, + CancellationToken cancellationToken = default) + => throw new NotImplementedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() { } + } + + [Fact] + public void Constructor_NullChatClient_ThrowsArgumentNullException() + { + Assert.Throws(() => new VisionTableEnricher(null!)); + } + + [Fact] + public async Task ProcessAsync_TableElement_GetsEnrichedWithMarkdown() + { + var expectedMarkdown = "| H1 | H2 |\n|---|---|\n| V1 | V2 |"; + var client = new TestChatClient(expectedMarkdown); + var enricher = new VisionTableEnricher(client); + + var doc = CreateDocumentWithTable(); + + var result = await enricher.ProcessAsync(doc); + + var table = result.EnumerateContent().OfType().First(); + Assert.True(table.HasMetadata); + Assert.True(table.Metadata.ContainsKey("enriched_markdown_table")); + Assert.Equal(expectedMarkdown, table.Metadata["enriched_markdown_table"]); + } + + [Fact] + public async Task ProcessAsync_NonTableElement_NotModified() + { + var client = new TestChatClient("should not appear"); + var enricher = new VisionTableEnricher(client); + + var doc = CreateDocumentWithTable(); + + var result = await enricher.ProcessAsync(doc); + + var paragraphs = result.EnumerateContent() + .OfType() + .ToList(); + + foreach (var p in paragraphs) + { + Assert.False(p.Metadata.ContainsKey("enriched_markdown_table")); + } + } + + [Fact] + public async Task ProcessAsync_EmptyDocument_NoException() + { + var client = new TestChatClient("response"); + var enricher = new VisionTableEnricher(client); + + var doc = new IngestionDocument("empty.pdf"); + + var result = await enricher.ProcessAsync(doc); + + Assert.NotNull(result); + Assert.Empty(result.Sections); + } + + [Fact] + public async Task ProcessAsync_DocumentWithNoTables_NoEnrichment() + { + var client = new TestChatClient("response"); + var enricher = new VisionTableEnricher(client); + + var doc = new IngestionDocument("no-tables.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + section.Elements.Add(new IngestionDocumentParagraph("Regular text") { Text = "Regular text" }); + doc.Sections.Add(section); + + var result = await enricher.ProcessAsync(doc); + + var elements = result.EnumerateContent().ToList(); + foreach (var e in elements) + { + Assert.False(e.Metadata.ContainsKey("enriched_markdown_table")); + } + } + + [Fact] + public async Task ProcessAsync_MultipleTables_AllEnriched() + { + var expectedMarkdown = "| A | B |"; + var client = new TestChatClient(expectedMarkdown); + var enricher = new VisionTableEnricher(client); + + var doc = new IngestionDocument("multi-table.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + + section.Elements.Add(CreateTable("| T1C1 | T1C2 |")); + section.Elements.Add(CreateTable("| T2C1 | T2C2 |")); + doc.Sections.Add(section); + + var result = await enricher.ProcessAsync(doc); + + var tables = result.EnumerateContent().OfType().ToList(); + Assert.Equal(2, tables.Count); + Assert.All(tables, t => + { + Assert.True(t.Metadata.ContainsKey("enriched_markdown_table")); + Assert.Equal(expectedMarkdown, t.Metadata["enriched_markdown_table"]); + }); + } + + [Fact] + public async Task ProcessAsync_CancellationRequested_Throws() + { + var client = new TestChatClient("response"); + var enricher = new VisionTableEnricher(client); + + var doc = CreateDocumentWithTable(); + using var cts = new CancellationTokenSource(); + cts.Cancel(); + + await Assert.ThrowsAnyAsync( + () => enricher.ProcessAsync(doc, cts.Token)); + } + + [Fact] + public async Task ProcessAsync_ReturnsSameDocumentInstance() + { + var client = new TestChatClient("markdown"); + var enricher = new VisionTableEnricher(client); + + var doc = CreateDocumentWithTable(); + + var result = await enricher.ProcessAsync(doc); + + Assert.Same(doc, result); + } + + [Fact] + public async Task ProcessAsync_WithPageImage_SendsDataContentToLlm() + { + var expectedMarkdown = "| A | B |"; + var client = new TestChatClient(expectedMarkdown); + var enricher = new VisionTableEnricher(client); + + var fakeImageBytes = new byte[] { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A }; + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + section.Metadata["page_image"] = fakeImageBytes; + + var table = CreateTable("| raw | data |"); + table.PageNumber = 1; + section.Elements.Add(table); + doc.Sections.Add(section); + + await enricher.ProcessAsync(doc); + + var userMsg = client.LastMessages.Last(m => m.Role == ChatRole.User); + Assert.Contains(userMsg.Contents, c => c is DataContent dc && dc.MediaType == "image/png"); + } + + [Fact] + public async Task ProcessAsync_WithoutPageImage_SendsTextOnlyToLlm() + { + var expectedMarkdown = "| A | B |"; + var client = new TestChatClient(expectedMarkdown); + var enricher = new VisionTableEnricher(client); + + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + // No page_image metadata + + var table = CreateTable("| raw | data |"); + table.PageNumber = 1; + section.Elements.Add(table); + doc.Sections.Add(section); + + await enricher.ProcessAsync(doc); + + var userMsg = client.LastMessages.Last(m => m.Role == ChatRole.User); + Assert.DoesNotContain(userMsg.Contents, c => c is DataContent); + } + + private static IngestionDocument CreateDocumentWithTable() + { + var doc = new IngestionDocument("test.pdf"); + var section = new IngestionDocumentSection { PageNumber = 1 }; + + section.Elements.Add(CreateTable("| H1 | H2 |\n|---|---|\n| V1 | V2 |")); + + var paragraph = new IngestionDocumentParagraph("Regular text") { Text = "Regular text" }; + section.Elements.Add(paragraph); + + doc.Sections.Add(section); + return doc; + } + + private static IngestionDocumentTable CreateTable(string markdown) + { + var cells = new IngestionDocumentElement[1, 2]; + cells[0, 0] = new IngestionDocumentParagraph("C1") { Text = "C1" }; + cells[0, 1] = new IngestionDocumentParagraph("C2") { Text = "C2" }; + return new IngestionDocumentTable(markdown, cells); + } +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/AnnotatedTextBlockTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/AnnotatedTextBlockTests.cs new file mode 100644 index 0000000..efd9b23 --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/AnnotatedTextBlockTests.cs @@ -0,0 +1,61 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using UglyToad.PdfPig.Content; +using UglyToad.PdfPig.Core; +using UglyToad.PdfPig.DocumentLayoutAnalysis; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; +using UglyToad.PdfPig.PdfFonts; +using Xunit; + +public class AnnotatedTextBlockTests +{ + [Fact] + public void Constructor_SetsLabelAndConfidence() + { + var word = CreateWord(new PdfRectangle(10, 10, 50, 25)); + var line = new TextLine(new[] { word }); + var block = new AnnotatedTextBlock(new[] { line }, "table", 0.95f); + + Assert.Equal("table", block.Label); + Assert.Equal(0.95f, block.Confidence, 4); + } + + [Fact] + public void TextProperty_ReturnsLineText() + { + var word = CreateWord(new PdfRectangle(10, 10, 50, 25)); + var line = new TextLine(new[] { word }); + var block = new AnnotatedTextBlock(new[] { line }, "text", 0.9f); + + Assert.Equal("a", block.Text); + } + + [Fact] + public void DefaultSeparator_IsNewline() + { + var w1 = CreateWord(new PdfRectangle(10, 100, 50, 120)); + var w2 = CreateWord(new PdfRectangle(10, 50, 50, 70)); + var line1 = new TextLine(new[] { w1 }); + var line2 = new TextLine(new[] { w2 }); + var block = new AnnotatedTextBlock(new[] { line1, line2 }, "text", 0.8f); + + Assert.Equal("a\na", block.Text); + } + + private static Word CreateWord(PdfRectangle boundingBox) + { + var letter = new Letter( + "a", + boundingBox, + boundingBox, + boundingBox.BottomLeft, + boundingBox.BottomRight, + 10, 1, + (FontDetails)null!, + TextRenderingMode.NeitherClip, + null!, null!, + 0, 0); + return new Word(new[] { letter }); + } +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/ConfigurableLayoutModelTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/ConfigurableLayoutModelTests.cs new file mode 100644 index 0000000..1c9448d --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/ConfigurableLayoutModelTests.cs @@ -0,0 +1,82 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using System; +using System.Collections.Generic; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.Models; +using Xunit; + +public class ConfigurableLayoutModelTests +{ + #region Constructor guard clauses + + [Fact] + public void Constructor_NullModelPath_ThrowsArgumentNullException() + { + var ex = Assert.Throws( + () => new ConfigurableLayoutModel(null!, new LayoutModelOptions())); + + Assert.Equal("modelPath", ex.ParamName); + } + + [Fact] + public void Constructor_NullOptions_ThrowsArgumentNullException() + { + var ex = Assert.Throws( + () => new ConfigurableLayoutModel("model.onnx", null!)); + + Assert.Equal("options", ex.ParamName); + } + + #endregion + + #region Properties + + [Fact] + public void ModelPath_ReturnsConstructorValue() + { + var model = new ConfigurableLayoutModel("path/to/model.onnx", new LayoutModelOptions()); + + Assert.Equal("path/to/model.onnx", model.ModelPath); + } + + [Fact] + public void LabelMapping_WithClassLabels_ReturnsSameLabels() + { + var labels = new Dictionary { [0] = "text", [1] = "table" }; + var options = new LayoutModelOptions { ClassLabels = labels }; + var model = new ConfigurableLayoutModel("model.onnx", options); + + var mapping = model.LabelMapping; + + Assert.Equal(2, mapping.Count); + Assert.Equal("text", mapping[0]); + Assert.Equal("table", mapping[1]); + } + + [Fact] + public void LabelMapping_WithoutClassLabels_ReturnsEmptyDictionary() + { + var model = new ConfigurableLayoutModel("model.onnx", new LayoutModelOptions()); + + var mapping = model.LabelMapping; + + Assert.NotNull(mapping); + Assert.Empty(mapping); + } + + #endregion + + #region Dispose + + [Fact] + public void Dispose_CalledTwice_DoesNotThrow() + { + var model = new ConfigurableLayoutModel("model.onnx", new LayoutModelOptions()); + + model.Dispose(); + model.Dispose(); + } + + #endregion +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/DetectionPostprocessingTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/DetectionPostprocessingTests.cs new file mode 100644 index 0000000..a3e4d59 --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/DetectionPostprocessingTests.cs @@ -0,0 +1,314 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using UglyToad.PdfPig.Core; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; +using Xunit; + +public class DetectionPostprocessingTests +{ + #region CxCyWhToRect + + [Fact] + public void CxCyWhToRect_CenteredBox() + { + // Center at (320, 240) in image coords, 100×50 box on a 640×480 page + var rect = DetectionPostprocessing.CxCyWhToRect(320, 240, 100, 50, 640, 480); + + // Image coords: x1=270, y1=215, x2=370, y2=265 + // PDF Y-flip: pdfBottom = 480 - 265 = 215, pdfTop = 480 - 215 = 265 + Assert.Equal(270, rect.Left, 1); + Assert.Equal(370, rect.Right, 1); + Assert.Equal(215, rect.Bottom, 1); + Assert.Equal(265, rect.Top, 1); + } + + [Fact] + public void CxCyWhToRect_TopLeftCornerBox() + { + // Box centered at (25, 25) with size 50×50 on 100×100 page + var rect = DetectionPostprocessing.CxCyWhToRect(25, 25, 50, 50, 100, 100); + + // Image coords: x1=0, y1=0, x2=50, y2=50 + // PDF Y-flip: pdfBottom = 100-50=50, pdfTop = 100-0=100 + Assert.Equal(0, rect.Left, 1); + Assert.Equal(50, rect.Right, 1); + Assert.Equal(50, rect.Bottom, 1); + Assert.Equal(100, rect.Top, 1); + } + + [Fact] + public void CxCyWhToRect_BottomRightCorner() + { + // Box centered at (75, 75) with size 50×50 on 100×100 page + var rect = DetectionPostprocessing.CxCyWhToRect(75, 75, 50, 50, 100, 100); + + // Image coords: x1=50, y1=50, x2=100, y2=100 + // PDF Y-flip: pdfBottom = 100-100=0, pdfTop = 100-50=50 + Assert.Equal(50, rect.Left, 1); + Assert.Equal(100, rect.Right, 1); + Assert.Equal(0, rect.Bottom, 1); + Assert.Equal(50, rect.Top, 1); + } + + #endregion + + #region XyxyToRect + + [Fact] + public void XyxyToRect_FlipsYAxis() + { + // Image coords: top-left (10, 20), bottom-right (100, 80) on 200×300 page + var rect = DetectionPostprocessing.XyxyToRect(10, 20, 100, 80, 200, 300); + + Assert.Equal(10, rect.Left, 1); + Assert.Equal(100, rect.Right, 1); + // PDF Y-flip: pdfBottom = 300 - 80 = 220, pdfTop = 300 - 20 = 280 + Assert.Equal(220, rect.Bottom, 1); + Assert.Equal(280, rect.Top, 1); + } + + [Fact] + public void XyxyToRect_FullPage() + { + // Image covers full page 640×480 + var rect = DetectionPostprocessing.XyxyToRect(0, 0, 640, 480, 640, 480); + + Assert.Equal(0, rect.Left, 1); + Assert.Equal(640, rect.Right, 1); + Assert.Equal(0, rect.Bottom, 1); + Assert.Equal(480, rect.Top, 1); + } + + [Fact] + public void XyxyToRect_SmallBox() + { + var rect = DetectionPostprocessing.XyxyToRect(0, 0, 10, 10, 1000, 1000); + + Assert.Equal(0, rect.Left, 1); + Assert.Equal(10, rect.Right, 1); + Assert.Equal(990, rect.Bottom, 1); + Assert.Equal(1000, rect.Top, 1); + } + + #endregion + + #region ComputeIoU + + [Fact] + public void ComputeIoU_IdenticalBoxes_ReturnsOne() + { + var a = new PdfRectangle(0, 0, 100, 100); + var b = new PdfRectangle(0, 0, 100, 100); + + float iou = DetectionPostprocessing.ComputeIoU(a, b); + + Assert.Equal(1f, iou, 4); + } + + [Fact] + public void ComputeIoU_NonOverlapping_ReturnsZero() + { + var a = new PdfRectangle(0, 0, 50, 50); + var b = new PdfRectangle(100, 100, 200, 200); + + float iou = DetectionPostprocessing.ComputeIoU(a, b); + + Assert.Equal(0f, iou, 4); + } + + [Fact] + public void ComputeIoU_EdgeAdjacent_ReturnsZero() + { + // Two boxes sharing an edge (no overlap area) + var a = new PdfRectangle(0, 0, 50, 50); + var b = new PdfRectangle(50, 0, 100, 50); + + float iou = DetectionPostprocessing.ComputeIoU(a, b); + + Assert.Equal(0f, iou, 4); + } + + [Fact] + public void ComputeIoU_HalfOverlap() + { + // Two 100×100 boxes overlapping by 50 in X + var a = new PdfRectangle(0, 0, 100, 100); + var b = new PdfRectangle(50, 0, 150, 100); + + float iou = DetectionPostprocessing.ComputeIoU(a, b); + + // Intersection: 50×100 = 5000 + // Union: 10000 + 10000 - 5000 = 15000 + // IoU = 5000/15000 = 1/3 + Assert.Equal(1f / 3f, iou, 3); + } + + [Fact] + public void ComputeIoU_ContainedBox() + { + // Small box fully inside large box + var large = new PdfRectangle(0, 0, 100, 100); + var small = new PdfRectangle(25, 25, 75, 75); + + float iou = DetectionPostprocessing.ComputeIoU(large, small); + + // Intersection = 50×50 = 2500 + // Union = 10000 + 2500 - 2500 = 10000 + // IoU = 2500/10000 = 0.25 + Assert.Equal(0.25f, iou, 3); + } + + [Fact] + public void ComputeIoU_SymmetricProperty() + { + var a = new PdfRectangle(10, 10, 60, 80); + var b = new PdfRectangle(30, 20, 90, 70); + + Assert.Equal( + DetectionPostprocessing.ComputeIoU(a, b), + DetectionPostprocessing.ComputeIoU(b, a), + 4); + } + + #endregion + + #region ApplyNms + + [Fact] + public void ApplyNms_EmptyList_ReturnsEmpty() + { + var detections = new List(); + var result = DetectionPostprocessing.ApplyNms(detections); + + Assert.Empty(result); + } + + [Fact] + public void ApplyNms_SingleDetection_ReturnsSame() + { + var detection = new LayoutDetection(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.9f); + var detections = new List { detection }; + + var result = DetectionPostprocessing.ApplyNms(detections); + + Assert.Single(result); + Assert.Equal(detection, result[0]); + } + + [Fact] + public void ApplyNms_NonOverlapping_AllPreserved() + { + var detections = new List + { + new(new PdfRectangle(0, 0, 50, 50), "text", 0, 0.9f), + new(new PdfRectangle(200, 200, 300, 300), "text", 0, 0.8f), + new(new PdfRectangle(400, 400, 500, 500), "text", 0, 0.7f) + }; + + var result = DetectionPostprocessing.ApplyNms(detections); + + Assert.Equal(3, result.Count); + } + + [Fact] + public void ApplyNms_IdenticalBoxes_KeepsHighestConfidence() + { + var detections = new List + { + new(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.5f), + new(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.9f), + new(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.7f) + }; + + var result = DetectionPostprocessing.ApplyNms(detections, 0.45f); + + Assert.Single(result); + Assert.Equal(0.9f, result[0].Confidence, 4); + } + + [Fact] + public void ApplyNms_HighOverlap_Suppressed() + { + // Two boxes with ~70% overlap + var detections = new List + { + new(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.9f), + new(new PdfRectangle(20, 20, 120, 120), "text", 0, 0.7f) + }; + + var result = DetectionPostprocessing.ApplyNms(detections, 0.3f); + + // IoU of these boxes is significant → lower-confidence one suppressed + Assert.Single(result); + Assert.Equal(0.9f, result[0].Confidence, 4); + } + + [Fact] + public void ApplyNms_HighThreshold_KeepsBoth() + { + // Same boxes, but with very high IoU threshold → both kept + var detections = new List + { + new(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.9f), + new(new PdfRectangle(20, 20, 120, 120), "text", 0, 0.7f) + }; + + var result = DetectionPostprocessing.ApplyNms(detections, 0.99f); + + Assert.Equal(2, result.Count); + } + + [Fact] + public void ApplyNms_NullDetections_Throws() + { + Assert.Throws(() => + DetectionPostprocessing.ApplyNms(null!)); + } + + [Fact] + public void ApplyNms_MixedOverlap_CorrectBehavior() + { + // Three detections: A and B overlap heavily, C is separate + var detections = new List + { + new(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.9f), // A + new(new PdfRectangle(10, 10, 110, 110), "text", 0, 0.8f), // B (overlaps A) + new(new PdfRectangle(500, 500, 600, 600), "text", 0, 0.7f) // C (no overlap) + }; + + var result = DetectionPostprocessing.ApplyNms(detections, 0.3f); + + Assert.Equal(2, result.Count); + Assert.Equal(0.9f, result[0].Confidence, 4); + Assert.Equal(0.7f, result[1].Confidence, 4); + } + + #endregion + + #region ScaleToPage + + [Fact] + public void ScaleToPage_SimpleScaling() + { + var det = new LayoutDetection(new PdfRectangle(0, 0, 320, 240), "text", 0, 0.9f); + var detections = new List { det }; + + var result = DetectionPostprocessing.ScaleToPage(detections, 640, 480, 1280, 960); + + // Scale 2x in both dimensions + Assert.Equal(0, result[0].BoundingBox.Left, 1); + Assert.Equal(640, result[0].BoundingBox.Right, 1); + Assert.Equal(0, result[0].BoundingBox.Bottom, 1); + Assert.Equal(480, result[0].BoundingBox.Top, 1); + } + + [Fact] + public void ScaleToPage_NullDetections_Throws() + { + Assert.Throws(() => + DetectionPostprocessing.ScaleToPage(null!, 640, 640, 100, 100)); + } + + #endregion +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/ImagePreprocessingTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/ImagePreprocessingTests.cs new file mode 100644 index 0000000..b4b9e05 --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/ImagePreprocessingTests.cs @@ -0,0 +1,391 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using SkiaSharp; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; +using Xunit; + +public class ImagePreprocessingTests +{ + #region ResizeExact + + [Fact] + public void ResizeExact_ProducesCorrectDimensions() + { + using var src = CreateSolidBitmap(100, 200, SKColors.Red); + using var result = ImagePreprocessing.ResizeExact(src, 50, 75); + + Assert.Equal(50, result.Width); + Assert.Equal(75, result.Height); + } + + [Fact] + public void ResizeExact_SquareImage() + { + using var src = CreateSolidBitmap(64, 64, SKColors.Blue); + using var result = ImagePreprocessing.ResizeExact(src, 640, 640); + + Assert.Equal(640, result.Width); + Assert.Equal(640, result.Height); + } + + [Fact] + public void ResizeExact_1x1Image() + { + using var src = CreateSolidBitmap(1, 1, SKColors.Green); + using var result = ImagePreprocessing.ResizeExact(src, 10, 10); + + Assert.Equal(10, result.Width); + Assert.Equal(10, result.Height); + } + + [Fact] + public void ResizeExact_Upscale() + { + using var src = CreateSolidBitmap(10, 10, SKColors.White); + using var result = ImagePreprocessing.ResizeExact(src, 640, 480); + + Assert.Equal(640, result.Width); + Assert.Equal(480, result.Height); + } + + [Fact] + public void ResizeExact_NullImage_Throws() + { + Assert.Throws(() => ImagePreprocessing.ResizeExact(null!, 10, 10)); + } + + #endregion + + #region Letterbox + + [Fact] + public void Letterbox_ProducesTargetDimensions() + { + using var src = CreateSolidBitmap(100, 200, SKColors.Red); + using var result = ImagePreprocessing.Letterbox(src, 640, 640, SKColors.Gray, out _, out _, out _); + + Assert.Equal(640, result.Width); + Assert.Equal(640, result.Height); + } + + [Fact] + public void Letterbox_PreservesAspectRatio_TallImage() + { + // Tall image (100×200): scale limited by height → 640/200=3.2 + using var src = CreateSolidBitmap(100, 200, SKColors.Red); + using var result = ImagePreprocessing.Letterbox(src, 640, 640, SKColors.Gray, out float scale, out int padX, out int padY); + + Assert.Equal(640 / 200f, scale, 3); + Assert.True(padX > 0, "Tall image should have horizontal padding"); + Assert.Equal(0, padY); + } + + [Fact] + public void Letterbox_PreservesAspectRatio_WideImage() + { + // Wide image (200×100): scale limited by width → 640/200=3.2 + using var src = CreateSolidBitmap(200, 100, SKColors.Blue); + using var result = ImagePreprocessing.Letterbox(src, 640, 640, SKColors.Gray, out float scale, out int padX, out int padY); + + Assert.Equal(640 / 200f, scale, 3); + Assert.Equal(0, padX); + Assert.True(padY > 0, "Wide image should have vertical padding"); + } + + [Fact] + public void Letterbox_SquareImage_NoPadding() + { + using var src = CreateSolidBitmap(100, 100, SKColors.Green); + using var result = ImagePreprocessing.Letterbox(src, 640, 640, SKColors.Gray, out float scale, out int padX, out int padY); + + Assert.Equal(640 / 100f, scale, 3); + Assert.Equal(0, padX); + Assert.Equal(0, padY); + } + + [Fact] + public void Letterbox_PaddingColor_Applied() + { + using var src = CreateSolidBitmap(100, 200, SKColors.Red); + using var result = ImagePreprocessing.Letterbox(src, 640, 640, SKColors.Gray, out _, out int padX, out _); + + // The top-left corner should be in the pad area (if padX > 0) + if (padX > 0) + { + var pixel = result.GetPixel(0, 0); + // Pad color should be gray + Assert.Equal(SKColors.Gray.Red, pixel.Red); + Assert.Equal(SKColors.Gray.Green, pixel.Green); + Assert.Equal(SKColors.Gray.Blue, pixel.Blue); + } + } + + [Fact] + public void Letterbox_NullImage_Throws() + { + Assert.Throws(() => + ImagePreprocessing.Letterbox(null!, 640, 640, SKColors.Gray, out _, out _, out _)); + } + + #endregion + + #region ToChwUint8 + + [Fact] + public void ToChwUint8_CorrectShape() + { + using var src = CreateSolidBitmap(4, 3, SKColors.Red); + var tensor = ImagePreprocessing.ToChwUint8(src); + + Assert.Equal(4, tensor.Dimensions.Length); // [1, 3, H, W] + Assert.Equal(1, tensor.Dimensions[0]); + Assert.Equal(3, tensor.Dimensions[1]); + Assert.Equal(3, tensor.Dimensions[2]); // height + Assert.Equal(4, tensor.Dimensions[3]); // width + } + + [Fact] + public void ToChwUint8_RedImage_CorrectChannelValues() + { + using var src = CreateSolidBitmap(2, 2, SKColors.Red); + var tensor = ImagePreprocessing.ToChwUint8(src); + + // Red channel = 255, Green = 0, Blue = 0 + Assert.Equal(255, tensor[0, 0, 0, 0]); // R + Assert.Equal(0, tensor[0, 1, 0, 0]); // G + Assert.Equal(0, tensor[0, 2, 0, 0]); // B + } + + [Fact] + public void ToChwUint8_WhiteImage_AllChannels255() + { + using var src = CreateSolidBitmap(2, 2, SKColors.White); + var tensor = ImagePreprocessing.ToChwUint8(src); + + for (int c = 0; c < 3; c++) + { + for (int y = 0; y < 2; y++) + { + for (int x = 0; x < 2; x++) + { + Assert.Equal(255, tensor[0, c, y, x]); + } + } + } + } + + [Fact] + public void ToChwUint8_1x1Image() + { + using var src = CreateSolidBitmap(1, 1, new SKColor(10, 20, 30)); + var tensor = ImagePreprocessing.ToChwUint8(src); + + Assert.Equal(1, tensor.Dimensions[2]); + Assert.Equal(1, tensor.Dimensions[3]); + Assert.Equal(10, tensor[0, 0, 0, 0]); + Assert.Equal(20, tensor[0, 1, 0, 0]); + Assert.Equal(30, tensor[0, 2, 0, 0]); + } + + [Fact] + public void ToChwUint8_NullImage_Throws() + { + Assert.Throws(() => ImagePreprocessing.ToChwUint8(null!)); + } + + #endregion + + #region ToChwFloat + + [Fact] + public void ToChwFloat_CorrectShape() + { + using var src = CreateSolidBitmap(5, 7, SKColors.Blue); + var tensor = ImagePreprocessing.ToChwFloat(src); + + Assert.Equal(4, tensor.Dimensions.Length); + Assert.Equal(1, tensor.Dimensions[0]); + Assert.Equal(3, tensor.Dimensions[1]); + Assert.Equal(7, tensor.Dimensions[2]); // height + Assert.Equal(5, tensor.Dimensions[3]); // width + } + + [Fact] + public void ToChwFloat_ValuesInZeroOneRange() + { + using var src = CreateSolidBitmap(3, 3, new SKColor(128, 64, 255)); + var tensor = ImagePreprocessing.ToChwFloat(src); + + for (int c = 0; c < 3; c++) + { + for (int y = 0; y < 3; y++) + { + for (int x = 0; x < 3; x++) + { + float val = tensor[0, c, y, x]; + Assert.InRange(val, 0f, 1f); + } + } + } + } + + [Fact] + public void ToChwFloat_WhiteImage_AllOnes() + { + using var src = CreateSolidBitmap(2, 2, SKColors.White); + var tensor = ImagePreprocessing.ToChwFloat(src); + + for (int c = 0; c < 3; c++) + { + Assert.Equal(1f, tensor[0, c, 0, 0], 4); + } + } + + [Fact] + public void ToChwFloat_BlackImage_AllZeros() + { + using var src = CreateSolidBitmap(2, 2, SKColors.Black); + var tensor = ImagePreprocessing.ToChwFloat(src); + + for (int c = 0; c < 3; c++) + { + Assert.Equal(0f, tensor[0, c, 0, 0], 4); + } + } + + [Fact] + public void ToChwFloat_SpecificColor_CorrectNormalization() + { + using var src = CreateSolidBitmap(1, 1, new SKColor(128, 0, 255)); + var tensor = ImagePreprocessing.ToChwFloat(src); + + Assert.Equal(128f / 255f, tensor[0, 0, 0, 0], 3); // R + Assert.Equal(0f, tensor[0, 1, 0, 0], 3); // G + Assert.Equal(1f, tensor[0, 2, 0, 0], 3); // B + } + + [Fact] + public void ToChwFloat_NullImage_Throws() + { + Assert.Throws(() => ImagePreprocessing.ToChwFloat(null!)); + } + + #endregion + + #region NormalizeImageNet + + [Fact] + public void NormalizeImageNet_AppliesCorrectTransformation() + { + int w = 2, h = 2; + int channelSize = w * h; + var chw = new float[3 * channelSize]; + + // Fill with 0.5 for all channels + Array.Fill(chw, 0.5f); + + ImagePreprocessing.NormalizeImageNet(chw, w, h); + + // Expected: (0.5 - mean) / std per channel + float[] mean = [0.485f, 0.456f, 0.406f]; + float[] std = [0.229f, 0.224f, 0.225f]; + + for (int c = 0; c < 3; c++) + { + float expected = (0.5f - mean[c]) / std[c]; + for (int i = 0; i < channelSize; i++) + { + Assert.Equal(expected, chw[c * channelSize + i], 4); + } + } + } + + [Fact] + public void NormalizeImageNet_ZeroInput_NegativeMeanDividedByStd() + { + int w = 1, h = 1; + var chw = new float[3]; + + ImagePreprocessing.NormalizeImageNet(chw, w, h); + + Assert.Equal(-0.485f / 0.229f, chw[0], 4); + Assert.Equal(-0.456f / 0.224f, chw[1], 4); + Assert.Equal(-0.406f / 0.225f, chw[2], 4); + } + + #endregion + + #region Normalize (custom) + + [Fact] + public void Normalize_CustomMeanStd() + { + int w = 2, h = 1; + var chw = new float[6]; // 3 channels × 2 pixels + Array.Fill(chw, 1.0f); + + float[] mean = [0.5f, 0.5f, 0.5f]; + float[] std = [0.5f, 0.5f, 0.5f]; + ImagePreprocessing.Normalize(chw, w, h, mean, std); + + // (1.0 - 0.5) / 0.5 = 1.0 + for (int i = 0; i < 6; i++) + { + Assert.Equal(1.0f, chw[i], 4); + } + } + + [Fact] + public void Normalize_IdentityTransform() + { + int w = 1, h = 1; + var chw = new float[] { 0.3f, 0.6f, 0.9f }; + + float[] mean = [0f, 0f, 0f]; + float[] std = [1f, 1f, 1f]; + ImagePreprocessing.Normalize(chw, w, h, mean, std); + + // (x - 0) / 1 = x, unchanged + Assert.Equal(0.3f, chw[0], 4); + Assert.Equal(0.6f, chw[1], 4); + Assert.Equal(0.9f, chw[2], 4); + } + + [Fact] + public void Normalize_InvalidMeanLength_Throws() + { + var chw = new float[3]; + float[] badMean = [0.5f, 0.5f]; // length 2 instead of 3 + float[] std = [1f, 1f, 1f]; + + Assert.Throws(() => + ImagePreprocessing.Normalize(chw, 1, 1, badMean, std)); + } + + [Fact] + public void Normalize_InvalidStdLength_Throws() + { + var chw = new float[3]; + float[] mean = [0.5f, 0.5f, 0.5f]; + float[] badStd = [1f]; // length 1 instead of 3 + + Assert.Throws(() => + ImagePreprocessing.Normalize(chw, 1, 1, mean, badStd)); + } + + #endregion + + #region Helpers + + private static SKBitmap CreateSolidBitmap(int width, int height, SKColor color) + { + var info = new SKImageInfo(width, height, SKColorType.Rgba8888, SKAlphaType.Premul); + var bitmap = new SKBitmap(info); + using var canvas = new SKCanvas(bitmap); + canvas.Clear(color); + return bitmap; + } + + #endregion +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/LayoutModelOptionsTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/LayoutModelOptionsTests.cs new file mode 100644 index 0000000..90bb4a2 --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/LayoutModelOptionsTests.cs @@ -0,0 +1,241 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.Models; +using Xunit; + +public class LayoutModelOptionsTests +{ + #region Default values + + [Fact] + public void DefaultOptions_InputWidth() + { + var options = new LayoutModelOptions(); + Assert.Equal(640, options.InputWidth); + } + + [Fact] + public void DefaultOptions_InputHeight() + { + var options = new LayoutModelOptions(); + Assert.Equal(640, options.InputHeight); + } + + [Fact] + public void DefaultOptions_ResizeMode() + { + var options = new LayoutModelOptions(); + Assert.Equal(ResizeMode.Exact, options.Resize); + } + + [Fact] + public void DefaultOptions_PixelFormat() + { + var options = new LayoutModelOptions(); + Assert.Equal(PixelFormat.Uint8Chw, options.PixelFormat); + } + + [Fact] + public void DefaultOptions_NormalizationIsNull() + { + var options = new LayoutModelOptions(); + Assert.Null(options.Normalization); + } + + [Fact] + public void DefaultOptions_ConfidenceThreshold() + { + var options = new LayoutModelOptions(); + Assert.Equal(0.3f, options.ConfidenceThreshold, 4); + } + + [Fact] + public void DefaultOptions_NmsIouThreshold() + { + var options = new LayoutModelOptions(); + Assert.Equal(0.45f, options.NmsIouThreshold, 4); + } + + [Fact] + public void DefaultOptions_RequiresNmsFalse() + { + var options = new LayoutModelOptions(); + Assert.False(options.RequiresNms); + } + + [Fact] + public void DefaultOptions_OutputBboxFormat() + { + var options = new LayoutModelOptions(); + Assert.Equal(BboxFormat.CxCyWh, options.OutputBboxFormat); + } + + [Fact] + public void DefaultOptions_ClassLabelsNull() + { + var options = new LayoutModelOptions(); + Assert.Null(options.ClassLabels); + } + + #endregion + + #region Custom init values + + [Fact] + public void Options_CustomValues() + { + var labels = new Dictionary { [0] = "text", [1] = "image" }; + var options = new LayoutModelOptions + { + InputWidth = 1024, + InputHeight = 768, + Resize = ResizeMode.Letterbox, + PixelFormat = PixelFormat.Float32Chw, + Normalization = WellKnownNormalizations.ImageNet, + ConfidenceThreshold = 0.5f, + NmsIouThreshold = 0.6f, + RequiresNms = true, + OutputBboxFormat = BboxFormat.Xyxy, + ClassLabels = labels + }; + + Assert.Equal(1024, options.InputWidth); + Assert.Equal(768, options.InputHeight); + Assert.Equal(ResizeMode.Letterbox, options.Resize); + Assert.Equal(PixelFormat.Float32Chw, options.PixelFormat); + Assert.NotNull(options.Normalization); + Assert.Equal(0.5f, options.ConfidenceThreshold, 4); + Assert.Equal(0.6f, options.NmsIouThreshold, 4); + Assert.True(options.RequiresNms); + Assert.Equal(BboxFormat.Xyxy, options.OutputBboxFormat); + Assert.Equal(2, options.ClassLabels!.Count); + } + + #endregion + + #region WellKnownNormalizations + + [Fact] + public void ImageNet_HasCorrectMean() + { + var norm = WellKnownNormalizations.ImageNet; + + Assert.Equal(3, norm.Mean.Length); + Assert.Equal(0.485f, norm.Mean[0], 4); + Assert.Equal(0.456f, norm.Mean[1], 4); + Assert.Equal(0.406f, norm.Mean[2], 4); + } + + [Fact] + public void ImageNet_HasCorrectStd() + { + var norm = WellKnownNormalizations.ImageNet; + + Assert.Equal(3, norm.Std.Length); + Assert.Equal(0.229f, norm.Std[0], 4); + Assert.Equal(0.224f, norm.Std[1], 4); + Assert.Equal(0.225f, norm.Std[2], 4); + } + + [Fact] + public void ZeroToOne_HasCorrectMean() + { + var norm = WellKnownNormalizations.ZeroToOne; + + Assert.Equal(3, norm.Mean.Length); + Assert.Equal(0f, norm.Mean[0], 4); + Assert.Equal(0f, norm.Mean[1], 4); + Assert.Equal(0f, norm.Mean[2], 4); + } + + [Fact] + public void ZeroToOne_HasCorrectStd() + { + var norm = WellKnownNormalizations.ZeroToOne; + + Assert.Equal(3, norm.Std.Length); + float expected = 1f / 255f; + Assert.Equal(expected, norm.Std[0], 6); + Assert.Equal(expected, norm.Std[1], 6); + Assert.Equal(expected, norm.Std[2], 6); + } + + #endregion + + #region Enums + + [Theory] + [InlineData(ResizeMode.Exact)] + [InlineData(ResizeMode.Letterbox)] + [InlineData(ResizeMode.AspectPreserve)] + public void ResizeMode_AllValuesExist(ResizeMode mode) + { + Assert.True(Enum.IsDefined(mode)); + } + + [Theory] + [InlineData(PixelFormat.Uint8Chw)] + [InlineData(PixelFormat.Float32Chw)] + public void PixelFormat_AllValuesExist(PixelFormat format) + { + Assert.True(Enum.IsDefined(format)); + } + + [Theory] + [InlineData(BboxFormat.CxCyWh)] + [InlineData(BboxFormat.Xyxy)] + [InlineData(BboxFormat.Xywh)] + public void BboxFormat_AllValuesExist(BboxFormat format) + { + Assert.True(Enum.IsDefined(format)); + } + + [Fact] + public void ResizeMode_HasThreeValues() + { + Assert.Equal(3, Enum.GetValues().Length); + } + + [Fact] + public void PixelFormat_HasTwoValues() + { + Assert.Equal(2, Enum.GetValues().Length); + } + + [Fact] + public void BboxFormat_HasThreeValues() + { + Assert.Equal(3, Enum.GetValues().Length); + } + + #endregion + + #region ImageNormalization record + + [Fact] + public void ImageNormalization_StoresMeanAndStd() + { + float[] mean = [0.1f, 0.2f, 0.3f]; + float[] std = [0.4f, 0.5f, 0.6f]; + + var norm = new ImageNormalization(mean, std); + + Assert.Same(mean, norm.Mean); + Assert.Same(std, norm.Std); + } + + [Fact] + public void ImageNormalization_RecordEquality() + { + float[] mean = [0.5f, 0.5f, 0.5f]; + float[] std = [1f, 1f, 1f]; + + var a = new ImageNormalization(mean, std); + var b = new ImageNormalization(mean, std); + + Assert.Equal(a, b); + } + + #endregion +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/OnnxPageSegmenterDiTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/OnnxPageSegmenterDiTests.cs new file mode 100644 index 0000000..60e663b --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/OnnxPageSegmenterDiTests.cs @@ -0,0 +1,141 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Options; +using Microsoft.ML.OnnxRuntime; +using SkiaSharp; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; +using UglyToad.PdfPig.DocumentLayoutAnalysis.PageSegmenter; +using Xunit; + +public class OnnxPageSegmenterDiTests +{ + #region Helpers + + private sealed class StubModel : ILayoutDetectionModel + { + public string ModelPath => "stub.onnx"; + + public IReadOnlyDictionary LabelMapping { get; } = new Dictionary(); + + public IReadOnlyList Preprocess(SKBitmap p, int w, int h) + { + return Array.Empty(); + } + + public IReadOnlyList Postprocess( + IDisposableReadOnlyCollection r, int w, int h) + { + return Array.Empty(); + } + + public void Dispose() { } + } + + #endregion + + #region Service registration + + [Fact] + public void AddOnnxPageSegmenter_RegistersILayoutDetectionModel() + { + var services = new ServiceCollection(); + services.AddOnnxPageSegmenter(); + var provider = services.BuildServiceProvider(); + + var model = provider.GetService(); + + Assert.NotNull(model); + Assert.IsType(model); + } + + [Fact] + public void AddOnnxPageSegmenter_RegistersIPageSegmenter() + { + var services = new ServiceCollection(); + services.AddOnnxPageSegmenter(); + + var descriptor = services.FirstOrDefault(d => d.ServiceType == typeof(IPageSegmenter)); + + Assert.NotNull(descriptor); + Assert.Equal(ServiceLifetime.Singleton, descriptor.Lifetime); + } + + [Fact] + public void AddOnnxPageSegmenter_IdempotentRegistration() + { + var services = new ServiceCollection(); + services.AddOnnxPageSegmenter(); + services.AddOnnxPageSegmenter(); + + int modelCount = services.Count(d => d.ServiceType == typeof(ILayoutDetectionModel)); + int segmenterCount = services.Count(d => d.ServiceType == typeof(OnnxPageSegmenter)); + int pageSegmenterCount = services.Count(d => d.ServiceType == typeof(IPageSegmenter)); + + Assert.Equal(1, modelCount); + Assert.Equal(1, segmenterCount); + Assert.Equal(1, pageSegmenterCount); + } + + #endregion + + #region Options configuration + + [Fact] + public void AddOnnxPageSegmenter_ConfigureAppliesOptions() + { + var services = new ServiceCollection(); + services.AddOnnxPageSegmenter(o => + { + o.ConfidenceThreshold = 0.8f; + o.RenderDpi = 300; + }); + var provider = services.BuildServiceProvider(); + + var opts = provider.GetRequiredService>().Value; + + Assert.Equal(0.8f, opts.ConfidenceThreshold, 4); + Assert.Equal(300, opts.RenderDpi); + } + + [Fact] + public void AddOnnxPageSegmenter_DefaultOptionsWithoutConfigure() + { + var services = new ServiceCollection(); + services.AddOnnxPageSegmenter(); + var provider = services.BuildServiceProvider(); + + var opts = provider.GetRequiredService>().Value; + + Assert.Equal(0.3f, opts.ConfidenceThreshold, 4); + Assert.Equal(150, opts.RenderDpi); + Assert.Null(opts.SessionOptions); + } + + [Fact] + public void AddOnnxPageSegmenter_InvalidOptions_ThrowsOnResolve() + { + var services = new ServiceCollection(); + services.AddOnnxPageSegmenter(o => o.ConfidenceThreshold = -1f); + var provider = services.BuildServiceProvider(); + var options = provider.GetRequiredService>(); + + Assert.Throws(() => options.Value); + } + + #endregion + + #region Constructor null checks + + [Fact] + public void IOptionsCtor_NullOptions_Throws() + { + var model = new StubModel(); + + Assert.Throws( + () => new OnnxPageSegmenter(model, (IOptions)null!)); + } + + #endregion +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/OnnxPageSegmenterTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/OnnxPageSegmenterTests.cs new file mode 100644 index 0000000..a89c88c --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/OnnxPageSegmenterTests.cs @@ -0,0 +1,311 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using Microsoft.ML.OnnxRuntime; +using SkiaSharp; +using UglyToad.PdfPig.Content; +using UglyToad.PdfPig.Core; +using UglyToad.PdfPig.DocumentLayoutAnalysis; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; +using UglyToad.PdfPig.PdfFonts; +using Xunit; + +public class OnnxPageSegmenterTests +{ + #region Constructor + + [Fact] + public void Constructor_NullModel_Throws() + { + Assert.Throws(() => new OnnxPageSegmenter(null!)); + } + + #endregion + + #region IDisposable + + [Fact] + public void Dispose_CanBeCalledOnStubModel() + { + // RtDetrLayoutModel doesn't create a session in its constructor, + // so we can construct it with a fake path for disposal testing. + var model = new StubLayoutModel(); + // OnnxPageSegmenter creates an InferenceSession in its constructor + // which requires a valid model file. We test dispose via the stub model directly. + model.Dispose(); + + // Calling Dispose again should not throw + model.Dispose(); + } + + [Fact] + public void StubModel_ImplementsIDisposable() + { + var model = new StubLayoutModel(); + Assert.IsAssignableFrom(model); + model.Dispose(); + } + + #endregion + + #region GetBlocks - empty words + + [SkippableFact] + public void GetBlocks_EmptyWords_ReturnsEmpty() + { + // This test requires a real ONNX model file + string modelPath = GetTestModelPath(); + Skip.IfNot(File.Exists(modelPath), "ONNX model file not found; skipping integration test."); + + using var model = new StubLayoutModelWithFile(modelPath); + using var segmenter = new OnnxPageSegmenter(model); + + var result = segmenter.GetBlocks(new List()); + Assert.Empty(result); + } + + #endregion + + #region Detection-to-TextBlock mapping logic (via DetectionPostprocessing + synthetic data) + + [Fact] + public void DetectionOverlap_WordInsideDetection_Captured() + { + // Verify the overlap logic: a word fully inside a detection box + var detBox = new PdfRectangle(0, 0, 100, 100); + var wordBox = new PdfRectangle(10, 10, 50, 50); + + // Verify overlap exists (using ComputeIoU as proxy) + float iou = DetectionPostprocessing.ComputeIoU(detBox, wordBox); + Assert.True(iou > 0, "Word inside detection should overlap"); + } + + [Fact] + public void DetectionOverlap_WordOutsideDetection_NotCaptured() + { + var detBox = new PdfRectangle(0, 0, 100, 100); + var wordBox = new PdfRectangle(200, 200, 250, 250); + + float iou = DetectionPostprocessing.ComputeIoU(detBox, wordBox); + Assert.Equal(0f, iou, 4); + } + + [Fact] + public void ConfidenceFiltering_BelowThreshold_NotIncluded() + { + var detections = new List + { + new(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.1f), // below 0.3 threshold + new(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.5f), // above threshold + }; + + // Simulate filtering as OnnxPageSegmenter does + float threshold = 0.3f; + var filtered = detections.Where(d => d.Confidence >= threshold).ToList(); + + Assert.Single(filtered); + Assert.Equal(0.5f, filtered[0].Confidence, 4); + } + + [Fact] + public void ConfidenceFiltering_AllBelowThreshold_ReturnsNone() + { + var detections = new List + { + new(new PdfRectangle(0, 0, 100, 100), "text", 0, 0.1f), + new(new PdfRectangle(50, 50, 150, 150), "table", 1, 0.2f), + }; + + float threshold = 0.3f; + var filtered = detections.Where(d => d.Confidence >= threshold).ToList(); + + Assert.Empty(filtered); + } + + #endregion + + #region Synthetic Word creation and TextBlock composition + + [Fact] + public void SyntheticWord_CanBeCreated() + { + var word = CreateWord(new PdfRectangle(10, 10, 50, 25)); + Assert.NotNull(word); + Assert.Equal("a", word.Text); + } + + [Fact] + public void SyntheticTextBlock_FromMultipleWords() + { + var w1 = CreateWord(new PdfRectangle(10, 10, 50, 25)); + var w2 = CreateWord(new PdfRectangle(60, 10, 100, 25)); + + var line = new TextLine(new[] { w1, w2 }); + var block = new TextBlock(new[] { line }); + + Assert.Equal(2, block.TextLines[0].Words.Count); + } + + [Fact] + public void WordsOnDifferentLines_GroupedCorrectly() + { + // Two words on different vertical positions + var w1 = CreateWord(new PdfRectangle(10, 100, 50, 120)); + var w2 = CreateWord(new PdfRectangle(10, 50, 50, 70)); + + var line1 = new TextLine(new[] { w1 }); + var line2 = new TextLine(new[] { w2 }); + var block = new TextBlock(new[] { line1, line2 }); + + Assert.Equal(2, block.TextLines.Count); + } + + #endregion + + #region LayoutDetection record + + [Fact] + public void LayoutDetection_RecordProperties() + { + var box = new PdfRectangle(10, 20, 100, 80); + var det = new LayoutDetection(box, "table", 5, 0.95f); + + Assert.Equal(box, det.BoundingBox); + Assert.Equal("table", det.Label); + Assert.Equal(5, det.ClassId); + Assert.Equal(0.95f, det.Confidence, 4); + } + + [Fact] + public void LayoutDetection_WithExpression() + { + var det = new LayoutDetection(new PdfRectangle(0, 0, 10, 10), "text", 0, 0.5f); + var updated = det with { Confidence = 0.9f }; + + Assert.Equal(0.9f, updated.Confidence, 4); + Assert.Equal("text", updated.Label); + } + + #endregion + + #region OnnxSegmenterOptions + + [Fact] + public void OnnxSegmenterOptions_DefaultValues() + { + var opts = new OnnxSegmenterOptions(); + Assert.Equal(0.3f, opts.ConfidenceThreshold, 4); + Assert.Equal(150, opts.RenderDpi); + Assert.Null(opts.SessionOptions); + } + + [Fact] + public void OnnxSegmenterOptions_CustomValues() + { + var sessionOpts = new SessionOptions(); + var opts = new OnnxSegmenterOptions + { + ConfidenceThreshold = 0.7f, + RenderDpi = 300, + SessionOptions = sessionOpts + }; + + Assert.Equal(0.7f, opts.ConfidenceThreshold, 4); + Assert.Equal(300, opts.RenderDpi); + Assert.Same(sessionOpts, opts.SessionOptions); + sessionOpts.Dispose(); + } + + #endregion + + #region Helpers + + private static Word CreateWord(PdfRectangle boundingBox) + { + var letter = new Letter( + "a", + boundingBox, + boundingBox, + boundingBox.BottomLeft, + boundingBox.BottomRight, + 10, 1, + (FontDetails)null!, + TextRenderingMode.NeitherClip, + null!, null!, + 0, 0); + return new Word(new[] { letter }); + } + + private static string GetTestModelPath() + { + // Look for model in common locations + var candidates = new[] + { + Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "models", "rtdetr.onnx"), + Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "rtdetr.onnx"), + }; + return candidates.FirstOrDefault(File.Exists) ?? candidates[0]; + } + + /// + /// A minimal stub implementing ILayoutDetectionModel for unit testing without an ONNX model file. + /// + private sealed class StubLayoutModel : ILayoutDetectionModel + { + public string ModelPath => "stub_model.onnx"; + + public IReadOnlyDictionary LabelMapping { get; } = new Dictionary + { + [0] = "text", + [1] = "table" + }; + + public IReadOnlyList Preprocess(SKBitmap pageImage, int originalWidth, int originalHeight) + { + return Array.Empty(); + } + + public IReadOnlyList Postprocess( + IDisposableReadOnlyCollection results, + int originalWidth, int originalHeight) + { + return Array.Empty(); + } + + public void Dispose() { } + } + + /// + /// A stub layout model that accepts an actual model file path for integration tests. + /// + private sealed class StubLayoutModelWithFile : ILayoutDetectionModel + { + public StubLayoutModelWithFile(string modelPath) + { + ModelPath = modelPath; + } + + public string ModelPath { get; } + + public IReadOnlyDictionary LabelMapping { get; } = new Dictionary + { + [0] = "text" + }; + + public IReadOnlyList Preprocess(SKBitmap pageImage, int originalWidth, int originalHeight) + { + return Array.Empty(); + } + + public IReadOnlyList Postprocess( + IDisposableReadOnlyCollection results, + int originalWidth, int originalHeight) + { + return Array.Empty(); + } + + public void Dispose() { } + } + + #endregion +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/OnnxSegmenterOptionsValidatorTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/OnnxSegmenterOptionsValidatorTests.cs new file mode 100644 index 0000000..3e4146a --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/OnnxSegmenterOptionsValidatorTests.cs @@ -0,0 +1,129 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Options; +using Microsoft.ML.OnnxRuntime; +using SkiaSharp; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; +using Xunit; + +public class OnnxSegmenterOptionsValidatorTests +{ + #region Helpers + + private static IOptions BuildOptions(Action? configure = null) + { + var services = new ServiceCollection(); + services.AddOnnxPageSegmenter(configure); + var provider = services.BuildServiceProvider(); + return provider.GetRequiredService>(); + } + + private sealed class StubModel : ILayoutDetectionModel + { + public string ModelPath => "stub.onnx"; + + public IReadOnlyDictionary LabelMapping { get; } = new Dictionary(); + + public IReadOnlyList Preprocess(SKBitmap p, int w, int h) + { + return Array.Empty(); + } + + public IReadOnlyList Postprocess( + IDisposableReadOnlyCollection r, int w, int h) + { + return Array.Empty(); + } + + public void Dispose() { } + } + + #endregion + + #region ConfidenceThreshold + + [Fact] + public void DefaultOptions_Succeeds() + { + var options = BuildOptions(); + + var opts = options.Value; + + Assert.Equal(0.3f, opts.ConfidenceThreshold, 4); + Assert.Equal(150, opts.RenderDpi); + } + + [Fact] + public void ConfidenceThreshold_Negative_Fails() + { + var options = BuildOptions(o => o.ConfidenceThreshold = -0.1f); + + var ex = Assert.Throws(() => options.Value); + Assert.Contains("ConfidenceThreshold", ex.Message); + } + + [Fact] + public void ConfidenceThreshold_AboveOne_Fails() + { + var options = BuildOptions(o => o.ConfidenceThreshold = 1.1f); + + var ex = Assert.Throws(() => options.Value); + Assert.Contains("ConfidenceThreshold", ex.Message); + } + + [Fact] + public void ConfidenceThreshold_ExactlyZero_Succeeds() + { + var options = BuildOptions(o => o.ConfidenceThreshold = 0f); + + var opts = options.Value; + + Assert.Equal(0f, opts.ConfidenceThreshold, 4); + } + + [Fact] + public void ConfidenceThreshold_ExactlyOne_Succeeds() + { + var options = BuildOptions(o => o.ConfidenceThreshold = 1f); + + var opts = options.Value; + + Assert.Equal(1f, opts.ConfidenceThreshold, 4); + } + + #endregion + + #region RenderDpi + + [Fact] + public void RenderDpi_Zero_Fails() + { + var options = BuildOptions(o => o.RenderDpi = 0); + + var ex = Assert.Throws(() => options.Value); + Assert.Contains("RenderDpi", ex.Message); + } + + [Fact] + public void RenderDpi_Negative_Fails() + { + var options = BuildOptions(o => o.RenderDpi = -1); + + var ex = Assert.Throws(() => options.Value); + Assert.Contains("RenderDpi", ex.Message); + } + + [Fact] + public void RenderDpi_One_Succeeds() + { + var options = BuildOptions(o => o.RenderDpi = 1); + + var opts = options.Value; + + Assert.Equal(1, opts.RenderDpi); + } + + #endregion +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/PageImageRendererTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/PageImageRendererTests.cs new file mode 100644 index 0000000..0c7c864 --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/PageImageRendererTests.cs @@ -0,0 +1,94 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using System.Collections.Generic; +using SkiaSharp; +using UglyToad.PdfPig.Content; +using UglyToad.PdfPig.Core; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; +using UglyToad.PdfPig.PdfFonts; +using Xunit; + +public class PageImageRendererTests +{ + [Fact] + public void RenderWords_NullWords_ThrowsArgumentNullException() + { + Assert.Throws(() => + PageImageRenderer.RenderWords(null!, 595, 842)); + } + + [Fact] + public void RenderWords_DefaultDpi_CorrectBitmapDimensions() + { + using var bitmap = PageImageRenderer.RenderWords( + new List(), 595, 842); + + Assert.Equal(1239, bitmap.Width); + Assert.Equal(1754, bitmap.Height); + } + + [Fact] + public void RenderWords_CustomDpi_ScalesDimensions() + { + using var bitmap = PageImageRenderer.RenderWords( + new List(), 595, 842, dpi: 72); + + Assert.Equal(595, bitmap.Width); + Assert.Equal(842, bitmap.Height); + } + + [Fact] + public void RenderWords_EmptyWordsList_ReturnsWhiteBitmap() + { + using var bitmap = PageImageRenderer.RenderWords( + new List(), 595, 842); + + Assert.NotNull(bitmap); + Assert.Equal(1239, bitmap.Width); + Assert.Equal(1754, bitmap.Height); + Assert.Equal(SKColors.White, bitmap.GetPixel(0, 0)); + } + + [Fact] + public void RenderWords_WithWords_RendersNonWhitePixels() + { + var words = new List + { + CreateWord(new PdfRectangle(10, 10, 100, 30)) + }; + + using var bitmap = PageImageRenderer.RenderWords(words, 200, 200, dpi: 72); + + // The word occupies a region — at least one pixel should be non-white (black fill) + bool hasNonWhite = false; + for (int y = 0; y < bitmap.Height && !hasNonWhite; y++) + { + for (int x = 0; x < bitmap.Width && !hasNonWhite; x++) + { + if (bitmap.GetPixel(x, y) != SKColors.White) + { + hasNonWhite = true; + } + } + } + + Assert.True(hasNonWhite, "Expected at least one non-white pixel from rendered word bounding box."); + } + + private static Word CreateWord(PdfRectangle boundingBox) + { + var letter = new Letter( + "a", + boundingBox, + boundingBox, + boundingBox.BottomLeft, + boundingBox.BottomRight, + 10, 1, + (FontDetails)null!, + TextRenderingMode.NeitherClip, + null!, null!, + 0, 0); + return new Word(new[] { letter }); + } +} +#endif diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/PdfPig.OnnxLayoutAnalysis.UnitTests.csproj b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/PdfPig.OnnxLayoutAnalysis.UnitTests.csproj new file mode 100644 index 0000000..9512caa --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/PdfPig.OnnxLayoutAnalysis.UnitTests.csproj @@ -0,0 +1,35 @@ + + + + CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests + CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests + net10.0 + true + false + + + + + + + + + + + + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + + + + + + diff --git a/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/RtDetrLayoutModelTests.cs b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/RtDetrLayoutModelTests.cs new file mode 100644 index 0000000..e708dfd --- /dev/null +++ b/MEDI/test/PdfPig.OnnxLayoutAnalysis.UnitTests/RtDetrLayoutModelTests.cs @@ -0,0 +1,154 @@ +#if NET8_0_OR_GREATER +namespace CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.UnitTests; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis; +using CommunityToolkit.DocumentProcessing.PdfPig.OnnxLayoutAnalysis.Models; +using Xunit; + +public class RtDetrLayoutModelTests +{ + #region Constructor + + [Fact] + public void Constructor_StoresModelPath() + { + using var model = new RtDetrLayoutModel("test_model.onnx"); + Assert.Equal("test_model.onnx", model.ModelPath); + } + + [Fact] + public void Constructor_NullModelPath_Throws() + { + Assert.Throws(() => new RtDetrLayoutModel(null!)); + } + + #endregion + + #region LabelMapping + + [Fact] + public void LabelMapping_Has17Entries() + { + using var model = new RtDetrLayoutModel("test.onnx"); + Assert.Equal(17, model.LabelMapping.Count); + } + + [Fact] + public void LabelMapping_ContainsExpectedLabels() + { + using var model = new RtDetrLayoutModel("test.onnx"); + var mapping = model.LabelMapping; + + Assert.Equal("caption", mapping[0]); + Assert.Equal("footnote", mapping[1]); + Assert.Equal("formula", mapping[2]); + Assert.Equal("list_item", mapping[3]); + Assert.Equal("page_footer", mapping[4]); + Assert.Equal("page_header", mapping[5]); + Assert.Equal("picture", mapping[6]); + Assert.Equal("section_header", mapping[7]); + Assert.Equal("table", mapping[8]); + Assert.Equal("text", mapping[9]); + Assert.Equal("title", mapping[10]); + Assert.Equal("document_index", mapping[11]); + Assert.Equal("code", mapping[12]); + Assert.Equal("checkbox_selected", mapping[13]); + Assert.Equal("checkbox_unselected", mapping[14]); + Assert.Equal("form", mapping[15]); + Assert.Equal("key_value_region", mapping[16]); + } + + [Fact] + public void LabelMapping_ContiguousKeys_0To16() + { + using var model = new RtDetrLayoutModel("test.onnx"); + var mapping = model.LabelMapping; + + for (int i = 0; i <= 16; i++) + { + Assert.True(mapping.ContainsKey(i), $"Missing key {i}"); + } + } + + [Fact] + public void LabelMapping_AllValuesNonEmpty() + { + using var model = new RtDetrLayoutModel("test.onnx"); + + foreach (var kvp in model.LabelMapping) + { + Assert.False(string.IsNullOrWhiteSpace(kvp.Value), + $"Label for key {kvp.Key} should not be empty"); + } + } + + #endregion + + #region ILayoutDetectionModel interface + + [Fact] + public void ImplementsILayoutDetectionModel() + { + using var model = new RtDetrLayoutModel("test.onnx"); + Assert.IsAssignableFrom(model); + } + + [Fact] + public void ImplementsIDisposable() + { + using var model = new RtDetrLayoutModel("test.onnx"); + Assert.IsAssignableFrom(model); + } + + [Fact] + public void Dispose_CanBeCalledMultipleTimes() + { + var model = new RtDetrLayoutModel("test.onnx"); + model.Dispose(); + model.Dispose(); // Should not throw + } + + #endregion + + #region Preprocess (requires model file) + + [SkippableFact] + public void Preprocess_NullImage_Throws() + { + string modelPath = FindModelPath(); + Skip.IfNot(File.Exists(modelPath), "ONNX model file not found; skipping."); + + using var model = new RtDetrLayoutModel(modelPath); + Assert.Throws(() => model.Preprocess(null!, 100, 100)); + } + + #endregion + + #region Postprocess (requires model file) + + [SkippableFact] + public void Postprocess_NullResults_Throws() + { + string modelPath = FindModelPath(); + Skip.IfNot(File.Exists(modelPath), "ONNX model file not found; skipping."); + + using var model = new RtDetrLayoutModel(modelPath); + Assert.Throws(() => model.Postprocess(null!, 640, 640)); + } + + #endregion + + #region Helpers + + private static string FindModelPath() + { + var candidates = new[] + { + Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "models", "rtdetr.onnx"), + Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "rtdetr.onnx"), + }; + return candidates.FirstOrDefault(File.Exists) ?? candidates[0]; + } + + #endregion +} +#endif diff --git a/nuget.config b/nuget.config new file mode 100644 index 0000000..039ca38 --- /dev/null +++ b/nuget.config @@ -0,0 +1,20 @@ + + + + + + + + + + + + + + + + + + + +