From 0907f102113bf3cd2b8f3a3729936463ace39cc6 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 14:59:38 -0400 Subject: [PATCH 01/25] feat(rag): real external-LLM generator bridging IChatClient (+ streaming) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RAG previously shipped no real LLM generator — only a local-LSTM NeuralGenerator and a StubGenerator — so every LLM-dependent component (HyDE, LLM/Cohere rerankers, LLM context compression, LLM query expansion, multi-query retrieval, GraphRAG summarization, LLM-judge eval) had nothing real to call. Add ChatClientGenerator : GeneratorBase, IStreamingGenerator that bridges the RAG IGenerator contract onto the existing production agentic IChatClient connectors (OpenAI/Azure/Anthropic/Cohere/Gemini/Mistral/Ollama — real HTTP, retries, timeouts, streaming). Generate() returns the full answer; the new IStreamingGenerator.GenerateStreamAsync yields token deltas. Reuses the battle-tested connectors instead of a new HTTP client. Unit-tested with a fake in-memory IChatClient (no network): role construction, option propagation, streaming concatenation, and prompt validation. Builds on net471/net8/net10 (IAsyncEnumerable streaming verified on net471). Unblocks #19, #25, #34. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/Interfaces/IStreamingGenerator.cs | 28 +++++ .../Generators/ChatClientGenerator.cs | 119 ++++++++++++++++++ .../Generators/ChatClientGeneratorTests.cs | 117 +++++++++++++++++ 3 files changed, 264 insertions(+) create mode 100644 src/Interfaces/IStreamingGenerator.cs create mode 100644 src/RetrievalAugmentedGeneration/Generators/ChatClientGenerator.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Generators/ChatClientGeneratorTests.cs diff --git a/src/Interfaces/IStreamingGenerator.cs b/src/Interfaces/IStreamingGenerator.cs new file mode 100644 index 0000000000..5da523e334 --- /dev/null +++ b/src/Interfaces/IStreamingGenerator.cs @@ -0,0 +1,28 @@ +using System.Collections.Generic; +using System.Threading; + +namespace AiDotNet.Interfaces; + +/// +/// A generator that can stream its output incrementally (token/chunk deltas) as it is produced, +/// in addition to the buffered . +/// +/// +/// +/// Streaming lets a RAG answer render as it is generated rather than after the whole response is +/// complete, which is the standard low-latency UX for chat/RAG systems. +/// +/// For Beginners: instead of waiting for the whole answer, you get it piece by piece — +/// like watching text appear as an assistant "types" — by iterating the returned async stream. +/// +/// The numeric data type used for relevance scoring. +public interface IStreamingGenerator : IGenerator +{ + /// + /// Generates a response to , yielding incremental text deltas as they arrive. + /// + /// The input prompt (typically already augmented with retrieved context). + /// Cancels the in-flight generation. + /// An async stream of text fragments whose concatenation is the full response. + IAsyncEnumerable GenerateStreamAsync(string prompt, CancellationToken cancellationToken = default); +} diff --git a/src/RetrievalAugmentedGeneration/Generators/ChatClientGenerator.cs b/src/RetrievalAugmentedGeneration/Generators/ChatClientGenerator.cs new file mode 100644 index 0000000000..104c9e85ac --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Generators/ChatClientGenerator.cs @@ -0,0 +1,119 @@ +using System; +using System.Collections.Generic; +using System.Runtime.CompilerServices; +using System.Threading; +using AiDotNet.Agentic.Models; +using AiDotNet.Interfaces; + +// ChatMessage collides with AiDotNet.PromptEngineering.Templates.ChatMessage (a project-wide global using); +// the RAG generator talks to the agentic connectors, so bind ChatMessage to the agentic type (same +// disambiguation ChatClientBase uses). +using ChatMessage = AiDotNet.Agentic.Models.ChatMessage; + +namespace AiDotNet.RetrievalAugmentedGeneration.Generators; + +/// +/// A real external-LLM generator for RAG. Bridges the RAG contract onto the +/// AiDotNet agentic connectors (OpenAI, Azure OpenAI, Anthropic, Cohere, +/// Gemini, Mistral, Ollama), so retrieval-augmented answers are produced by a genuine chat model instead +/// of the local-LSTM NeuralGenerator or the StubGenerator. Supports token streaming via +/// . +/// +/// +/// +/// This is the linchpin that lets the previously-stubbed LLM-powered RAG components (HyDE, LLM query +/// expansion, LLM/Cohere reranking, LLM context compression, multi-query retrieval) run against a real +/// model: construct one of these over any configured and inject it. +/// +/// For Beginners: give it a chat-model connection and it writes answers with that model. +/// returns the whole answer; +/// streams it as it is produced. +/// +/// The numeric data type used for relevance scoring. +public class ChatClientGenerator : GeneratorBase, IStreamingGenerator +{ + private readonly IChatClient _chatClient; + private readonly string? _systemPrompt; + private readonly ChatOptions _options; + + /// + /// Creates a generator over an existing chat-model connector. + /// + /// The configured chat-model connector (provider, model id, API key already set). + /// Optional system instruction prepended to every request. + /// Optional sampling temperature; null uses the provider default. + /// Optional generation cap; defaults to . + /// Context window size reported to the base generator. + /// Default generation length reported to the base generator. + /// Thrown when is null. + public ChatClientGenerator( + IChatClient chatClient, + string? systemPrompt = null, + double? temperature = null, + int? maxOutputTokens = null, + int maxContextTokens = 8192, + int maxGenerationTokens = 1024) + : base(maxContextTokens, maxGenerationTokens) + { + _chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + _systemPrompt = systemPrompt; + _options = new ChatOptions + { + Temperature = temperature, + MaxOutputTokens = maxOutputTokens ?? maxGenerationTokens + }; + } + + private IReadOnlyList BuildMessages(string prompt) + { + var messages = new List(); + if (!string.IsNullOrWhiteSpace(_systemPrompt)) + { + messages.Add(new ChatMessage(ChatRole.System, _systemPrompt!)); + } + + messages.Add(new ChatMessage(ChatRole.User, prompt)); + return messages; + } + + /// + /// + /// Synchronous bridge over the async chat client, matching the codebase's existing sync-over-async + /// pattern (e.g. the embedding models). Prefer — or the async RAG + /// pipeline (task #21) — on hot paths. + /// + protected override string GenerateCore(string prompt) + { + var response = _chatClient + .GetResponseAsync(BuildMessages(prompt), _options, CancellationToken.None) + .ConfigureAwait(false).GetAwaiter().GetResult(); + + return response.Text ?? string.Empty; + } + + /// + public async IAsyncEnumerable GenerateStreamAsync( + string prompt, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + if (prompt == null) + { + throw new ArgumentNullException(nameof(prompt), "Prompt cannot be null."); + } + + if (string.IsNullOrWhiteSpace(prompt)) + { + throw new ArgumentException("Prompt cannot be empty or whitespace.", nameof(prompt)); + } + + await foreach (var update in _chatClient + .GetStreamingResponseAsync(BuildMessages(prompt), _options, cancellationToken) + .ConfigureAwait(false)) + { + if (!string.IsNullOrEmpty(update.TextDelta)) + { + yield return update.TextDelta!; + } + } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Generators/ChatClientGeneratorTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Generators/ChatClientGeneratorTests.cs new file mode 100644 index 0000000000..8ab3501218 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Generators/ChatClientGeneratorTests.cs @@ -0,0 +1,117 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using AiDotNet.Agentic.Models; +using AiDotNet.RetrievalAugmentedGeneration.Generators; +using Xunit; +using ChatMessage = AiDotNet.Agentic.Models.ChatMessage; + +namespace AiDotNet.Tests.UnitTests.RetrievalAugmentedGeneration.Generators; + +/// +/// Unit tests for — the real-LLM RAG generator that bridges +/// . A fake in-memory chat client stands in for a provider connector, so +/// these run in CI with no network: they assert prompt/role construction, option propagation, buffered +/// text, and streaming delta concatenation. +/// +public class ChatClientGeneratorTests +{ + /// Records the request and returns/streams a canned answer — no HTTP. + private sealed class FakeChatClient : IChatClient + { + private readonly string _answer; + public string ModelId => "fake-model"; + public IReadOnlyList? LastMessages { get; private set; } + public ChatOptions? LastOptions { get; private set; } + + public FakeChatClient(string answer) => _answer = answer; + + public Task GetResponseAsync( + IReadOnlyList messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + LastMessages = messages; + LastOptions = options; + var msg = new ChatMessage(ChatRole.Assistant, _answer); + return Task.FromResult(new ChatResponse(msg, ChatFinishReason.Stop, usage: null, modelId: ModelId)); + } + + public async IAsyncEnumerable GetStreamingResponseAsync( + IReadOnlyList messages, ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + LastMessages = messages; + LastOptions = options; + // Split the answer into word-sized deltas to mimic token streaming. + foreach (var word in _answer.Split(' ')) + { + cancellationToken.ThrowIfCancellationRequested(); + await Task.Yield(); + yield return ChatResponseUpdate.ForText(word + " "); + } + } + } + + [Fact] + public void Generate_ReturnsChatClientText_AndSendsSystemThenUser() + { + var fake = new FakeChatClient("Transformers use self-attention [1]."); + var gen = new ChatClientGenerator(fake, systemPrompt: "You are terse.", temperature: 0.2, maxOutputTokens: 256); + + var text = gen.Generate("How do transformers work?"); + + Assert.Equal("Transformers use self-attention [1].", text); + Assert.NotNull(fake.LastMessages); + Assert.Equal(2, fake.LastMessages!.Count); + Assert.Equal(ChatRole.System, fake.LastMessages[0].Role); + Assert.Equal("You are terse.", fake.LastMessages[0].Text); + Assert.Equal(ChatRole.User, fake.LastMessages[1].Role); + Assert.Equal("How do transformers work?", fake.LastMessages[1].Text); + Assert.Equal(0.2, fake.LastOptions!.Temperature); + Assert.Equal(256, fake.LastOptions!.MaxOutputTokens); + } + + [Fact] + public void Generate_WithoutSystemPrompt_SendsOnlyUserMessage() + { + var fake = new FakeChatClient("answer"); + var gen = new ChatClientGenerator(fake); + + gen.Generate("q"); + + Assert.Single(fake.LastMessages!); + Assert.Equal(ChatRole.User, fake.LastMessages![0].Role); + } + + [Fact] + public async Task GenerateStreamAsync_ConcatenatesToFullAnswer() + { + var fake = new FakeChatClient("alpha beta gamma"); + var gen = new ChatClientGenerator(fake); + + var parts = new List(); + await foreach (var delta in gen.GenerateStreamAsync("q")) + { + parts.Add(delta); + } + + Assert.True(parts.Count >= 3, "expected multiple streamed deltas"); + Assert.Equal("alpha beta gamma ", string.Concat(parts)); + } + + [Fact] + public async Task GenerateStreamAsync_RejectsEmptyPrompt() + { + var gen = new ChatClientGenerator(new FakeChatClient("x")); + await Assert.ThrowsAsync(async () => + { + await foreach (var _ in gen.GenerateStreamAsync(" ")) { } + }); + } + + [Fact] + public void Constructor_NullChatClient_Throws() + => Assert.Throws(() => new ChatClientGenerator(null!)); +} From 5ab8b1d4116b6cb93deb5c0e1ed96498e45e70ea Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 15:18:35 -0400 Subject: [PATCH 02/25] feat(rag): de-stub LLM-powered components onto a real text generator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HyDE, LLM query expansion, multi-query retrieval, LLM reranking, and LLM context compression previously ran lexical templates/heuristics while accepting unused endpoint/apiKey args — they looked LLM-powered but never called a model. Add a small non-generic ITextGenerator (string Generate(string)) that IGenerator now extends, so any real generator (e.g. ChatClientGenerator over a chat connector) plugs into both the generic and non-generic RAG helpers. Each component gains an optional generator: - HyDEQueryExpansion: LLM writes hypothetical answer passages (true HyDE). - LLMQueryExpansion: LLM produces alternative queries, parsed line-by-line. - MultiQueryRetriever: reuses the de-stubbed expander (no more "{query} variation N"). - LLMBasedReranker: LLM scores each doc 0-10 for relevance. - LLMContextCompressor: LLM extracts only query-relevant content. When no generator is supplied each keeps a clearly-documented offline heuristic fallback (non-breaking). Verified with a fake ITextGenerator (6 tests). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/Interfaces/IGenerator.cs | 28 +--- src/Interfaces/ITextGenerator.cs | 25 ++++ .../LLMContextCompressor.cs | 31 ++++- .../QueryExpansion/HyDEQueryExpansion.cs | 59 ++++++-- .../QueryExpansion/LLMQueryExpansion.cs | 58 +++++++- .../Rerankers/LLMBasedReranker.cs | 32 ++++- .../Retrievers/MultiQueryRetriever.cs | 25 ++-- .../DeStubbedLlmComponentsTests.cs | 130 ++++++++++++++++++ 8 files changed, 329 insertions(+), 59 deletions(-) create mode 100644 src/Interfaces/ITextGenerator.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DeStubbedLlmComponentsTests.cs diff --git a/src/Interfaces/IGenerator.cs b/src/Interfaces/IGenerator.cs index b73538292a..fccbd5ed18 100644 --- a/src/Interfaces/IGenerator.cs +++ b/src/Interfaces/IGenerator.cs @@ -35,32 +35,10 @@ namespace AiDotNet.Interfaces; /// /// The numeric data type used for relevance scoring. [AiDotNet.Configuration.YamlConfigurable("Generator")] -public interface IGenerator +public interface IGenerator : ITextGenerator { - /// - /// Generates a text response based on a prompt. - /// - /// The input prompt or question. - /// The generated text response. - /// - /// - /// This method generates text based solely on the provided prompt, without - /// additional context. It's suitable for general-purpose text generation tasks. - /// In RAG systems, this is typically called with prompts that have been augmented - /// with retrieved context. - /// - /// For Beginners: This generates text from a prompt. - /// - /// For example: - /// - Prompt: "Explain photosynthesis in simple terms" - /// - Generated: "Photosynthesis is how plants make food using sunlight..." - /// - /// In RAG, the prompt usually includes both the question and retrieved documents: - /// - Prompt: "Context: [3 documents about photosynthesis]\n\nQuestion: Explain photosynthesis" - /// - Generated: Answer based on those specific documents - /// - /// - string Generate(string prompt); + // string Generate(string prompt) is inherited from ITextGenerator — buffered text generation from a + // prompt (in RAG, a prompt already augmented with retrieved context). See ITextGenerator for details. /// /// Generates a grounded answer using provided context documents. diff --git a/src/Interfaces/ITextGenerator.cs b/src/Interfaces/ITextGenerator.cs new file mode 100644 index 0000000000..30768b1e51 --- /dev/null +++ b/src/Interfaces/ITextGenerator.cs @@ -0,0 +1,25 @@ +namespace AiDotNet.Interfaces; + +/// +/// A minimal, numeric-type-agnostic text generator: turns a prompt into generated text. +/// +/// +/// +/// Several RAG helpers (query expansion, multi-query generation, LLM reranking, context compression) +/// only need "prompt in, text out" and do not care about the numeric type T used for scoring. +/// This non-generic contract lets them accept any real generator — including any +/// (which extends this) backed by a live chat model — instead of running +/// lexical heuristics with unused API credentials. +/// +/// For Beginners: it's the "give it a prompt, get back text" part of a language model, +/// with no other baggage — so the string-only RAG helpers can use a real LLM. +/// +public interface ITextGenerator +{ + /// + /// Generates text for the given prompt. + /// + /// The input prompt. + /// The generated text. + string Generate(string prompt); +} diff --git a/src/RetrievalAugmentedGeneration/ContextCompression/LLMContextCompressor.cs b/src/RetrievalAugmentedGeneration/ContextCompression/LLMContextCompressor.cs index 23d3998ce9..746f5c13be 100644 --- a/src/RetrievalAugmentedGeneration/ContextCompression/LLMContextCompressor.cs +++ b/src/RetrievalAugmentedGeneration/ContextCompression/LLMContextCompressor.cs @@ -3,6 +3,7 @@ using System.Linq; using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; @@ -19,23 +20,31 @@ public class LLMContextCompressor : ContextCompressorBase private readonly string _llmEndpoint; private readonly string _apiKey; private readonly double _compressionRatio; + private readonly ITextGenerator? _generator; /// /// Initializes a new instance of the class. /// - /// The target compression ratio (0.0 to 1.0). - /// The LLM API endpoint. - /// The API key for the LLM service. + /// The target compression ratio (0.0 to 1.0) for the extractive fallback. + /// Optional LLM API endpoint (informational; the call goes through ). + /// Optional API key (informational; the call goes through ). + /// + /// Optional real text generator. When provided, each document is compressed by the LLM to only the + /// query-relevant content; when null, a query-overlap extractive sentence selection is used as + /// an offline fallback. + /// public LLMContextCompressor( double compressionRatio = 0.5, string llmEndpoint = "", - string apiKey = "") + string apiKey = "", + ITextGenerator? generator = null) { _compressionRatio = compressionRatio >= 0 && compressionRatio <= 1 ? compressionRatio : throw new ArgumentOutOfRangeException(nameof(compressionRatio), "Compression ratio must be between 0 and 1"); _llmEndpoint = llmEndpoint; _apiKey = apiKey; + _generator = generator; } /// @@ -73,6 +82,20 @@ public string CompressText(string query, string text) { if (string.IsNullOrEmpty(text)) return text; + if (_generator != null) + { + var prompt = + $"Extract only the parts of the document that are relevant to answering the query. " + + $"Preserve the original wording, drop everything irrelevant, and return only the compressed " + + $"excerpt with no commentary.\n\nQuery: {query}\n\nDocument: {text}\n\nRelevant excerpt:"; + var compressed = _generator.Generate(prompt); + if (!string.IsNullOrWhiteSpace(compressed)) + { + return compressed.Trim(); + } + // Empty LLM reply: fall through to the extractive heuristic below. + } + var sentences = SplitIntoSentences(text); var scoredSentences = new List<(string sentence, double score)>(); diff --git a/src/RetrievalAugmentedGeneration/QueryExpansion/HyDEQueryExpansion.cs b/src/RetrievalAugmentedGeneration/QueryExpansion/HyDEQueryExpansion.cs index b43f5ba851..5957c219c1 100644 --- a/src/RetrievalAugmentedGeneration/QueryExpansion/HyDEQueryExpansion.cs +++ b/src/RetrievalAugmentedGeneration/QueryExpansion/HyDEQueryExpansion.cs @@ -3,46 +3,83 @@ using System.Linq; using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; using AiDotNet.LinearAlgebra; namespace AiDotNet.RetrievalAugmentedGeneration.QueryExpansion { /// - /// Hypothetical Document Embeddings (HyDE) query expansion strategy. + /// Hypothetical Document Embeddings (HyDE) query expansion (Gao et al., 2022). Generates hypothetical + /// answer passages for a query so that a dense retriever matches document-shaped text rather than the + /// short question. When a real text generator is supplied, the passages are written by the LLM (true + /// HyDE); otherwise a deterministic template is used as an offline fallback. /// [ComponentType(ComponentType.QueryExpander)] [PipelineStage(PipelineStage.QueryProcessing)] public class HyDEQueryExpansion : QueryExpansionBase { + private readonly ITextGenerator? _generator; + private readonly int _numHypotheticals; + + // LLM prompt style descriptors (index-aligned with the offline template styles below). + private static readonly string[] LlmStyles = + { + "a detailed, factual passage", + "a concise explanation", + "a technical description" + }; + private static readonly string[] TemplateStyles = + { + "detailed answer", + "concise explanation", + "technical description" + }; + /// /// Initializes a new instance of the class. /// - public HyDEQueryExpansion() + /// + /// Optional real text generator (e.g. a ChatClientGenerator over any chat connector). When + /// provided, hypothetical documents are LLM-generated (real HyDE). When null, a deterministic + /// template fallback is used so the component still functions offline. + /// + /// How many hypothetical passages to generate (default 3). + public HyDEQueryExpansion(ITextGenerator? generator = null, int numHypotheticals = 3) { + _generator = generator; + _numHypotheticals = Math.Max(1, numHypotheticals); } /// /// Expands a query by generating hypothetical documents. /// /// The original query. - /// A list of hypothetical document variations. + /// The original query followed by hypothetical document variations. public override List ExpandQuery(string query) { if (string.IsNullOrEmpty(query)) throw new ArgumentNullException(nameof(query)); var expansions = new List { query }; - - var hypoDoc1 = GenerateHypotheticalDocument(query, "detailed answer"); - var hypoDoc2 = GenerateHypotheticalDocument(query, "concise explanation"); - var hypoDoc3 = GenerateHypotheticalDocument(query, "technical description"); - - expansions.Add(hypoDoc1); - expansions.Add(hypoDoc2); - expansions.Add(hypoDoc3); + for (int i = 0; i < _numHypotheticals; i++) + { + expansions.Add(_generator != null + ? GenerateWithLlm(query, LlmStyles[i % LlmStyles.Length]) + : GenerateHypotheticalDocument(query, TemplateStyles[i % TemplateStyles.Length])); + } return expansions; } + private string GenerateWithLlm(string query, string style) + { + var prompt = + $"Write {style} that directly answers the following question, as if it were an excerpt from a " + + $"relevant document. Return only the passage, with no preamble or commentary.\n\n" + + $"Question: {query}\n\nPassage:"; + var doc = _generator!.Generate(prompt); + return string.IsNullOrWhiteSpace(doc) ? query : doc.Trim(); + } + private string GenerateHypotheticalDocument(string query, string style) { var words = query.Split(' ').Where(w => !string.IsNullOrWhiteSpace(w)).ToList(); diff --git a/src/RetrievalAugmentedGeneration/QueryExpansion/LLMQueryExpansion.cs b/src/RetrievalAugmentedGeneration/QueryExpansion/LLMQueryExpansion.cs index 549e55ce64..d6ee268839 100644 --- a/src/RetrievalAugmentedGeneration/QueryExpansion/LLMQueryExpansion.cs +++ b/src/RetrievalAugmentedGeneration/QueryExpansion/LLMQueryExpansion.cs @@ -1,12 +1,16 @@ using System; using System.Collections.Generic; +using System.Linq; using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; namespace AiDotNet.RetrievalAugmentedGeneration.QueryExpansion { /// - /// LLM-based query expansion for generating additional query variations. + /// LLM-based query expansion: rewrites a query into several alternative phrasings that capture the + /// same information need. When a real text generator is supplied the variations are produced by the + /// LLM; otherwise a deterministic template set is used as an offline fallback. /// [ComponentType(ComponentType.QueryExpander)] [PipelineStage(PipelineStage.QueryProcessing)] @@ -15,18 +19,25 @@ public class LLMQueryExpansion : QueryExpansionBase private readonly string _llmEndpoint; private readonly string _apiKey; private readonly int _numExpansions; + private readonly ITextGenerator? _generator; /// /// Initializes a new instance of the class. /// - /// The LLM API endpoint. - /// The API key for the LLM service. + /// Optional LLM API endpoint (informational; the actual call goes through ). + /// Optional API key (informational; the actual call goes through ). /// The number of query expansions to generate. - public LLMQueryExpansion(string? llmEndpoint = null, string? apiKey = null, int numExpansions = 3) + /// + /// Optional real text generator (e.g. a ChatClientGenerator). When provided, variations are + /// LLM-generated; when null, a deterministic template fallback is used so the component still + /// functions offline. + /// + public LLMQueryExpansion(string? llmEndpoint = null, string? apiKey = null, int numExpansions = 3, ITextGenerator? generator = null) { _llmEndpoint = llmEndpoint ?? string.Empty; _apiKey = apiKey ?? string.Empty; _numExpansions = numExpansions > 0 ? numExpansions : throw new ArgumentOutOfRangeException(nameof(numExpansions)); + _generator = generator; } /// @@ -40,15 +51,48 @@ public override List ExpandQuery(string query) var expansions = new List { query }; - for (int i = 0; i < _numExpansions; i++) + if (_generator != null) { - var expansion = GenerateExpansion(query, i); - expansions.Add(expansion); + var prompt = + $"Generate {_numExpansions} alternative search queries that capture the same information " + + $"need as the query below. Put each on its own line, with no numbering, quotes, or commentary.\n\n" + + $"Query: {query}"; + var text = _generator.Generate(prompt) ?? string.Empty; + foreach (var line in text.Split('\n')) + { + var cleaned = CleanVariation(line); + if (cleaned.Length > 0 && !expansions.Contains(cleaned)) + { + expansions.Add(cleaned); + } + + if (expansions.Count > _numExpansions) break; + } + } + + // No generator, or the LLM returned nothing usable: guarantee variations via the template. + if (expansions.Count == 1) + { + for (int i = 0; i < _numExpansions; i++) + { + expansions.Add(GenerateExpansion(query, i)); + } } return expansions; } + // Strip common list decorations (numbering, bullets, surrounding quotes) from an LLM line. + private static string CleanVariation(string line) + { + var s = line.Trim(); + if (s.Length == 0) return string.Empty; + s = s.TrimStart('-', '*', '•', ' ', '\t'); + int dot = s.IndexOf(". ", StringComparison.Ordinal); + if (dot > 0 && dot <= 3 && s.Take(dot).All(char.IsDigit)) s = s.Substring(dot + 2); + return s.Trim().Trim('"').Trim(); + } + private string GenerateExpansion(string query, int variant) { switch (variant % 5) diff --git a/src/RetrievalAugmentedGeneration/Rerankers/LLMBasedReranker.cs b/src/RetrievalAugmentedGeneration/Rerankers/LLMBasedReranker.cs index 9c8b15d2ea..a271ce4e2a 100644 --- a/src/RetrievalAugmentedGeneration/Rerankers/LLMBasedReranker.cs +++ b/src/RetrievalAugmentedGeneration/Rerankers/LLMBasedReranker.cs @@ -1,8 +1,11 @@ using System; using System.Collections.Generic; +using System.Globalization; using System.Linq; +using System.Text.RegularExpressions; using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.RetrievalAugmentedGeneration.RerankingStrategies @@ -17,6 +20,7 @@ public class LLMBasedReranker : Rerankers.RerankerBase { private readonly string _llmEndpoint; private readonly string _apiKey; + private readonly ITextGenerator? _generator; /// /// Gets a value indicating whether this reranker modifies relevance scores. @@ -26,12 +30,17 @@ public class LLMBasedReranker : Rerankers.RerankerBase /// /// Initializes a new instance of the class. /// - /// The LLM API endpoint. - /// The API key for the LLM service. - public LLMBasedReranker(string? llmEndpoint = null, string? apiKey = null) + /// Optional LLM API endpoint (informational; the call goes through ). + /// Optional API key (informational; the call goes through ). + /// + /// Optional real text generator. When provided, each document's relevance is scored by the LLM; + /// when null, a lexical relevance heuristic is used as an offline fallback. + /// + public LLMBasedReranker(string? llmEndpoint = null, string? apiKey = null, ITextGenerator? generator = null) { _llmEndpoint = llmEndpoint ?? string.Empty; _apiKey = apiKey ?? string.Empty; + _generator = generator; } /// @@ -71,6 +80,23 @@ protected override IEnumerable> RerankCore(string query, IList 2000 ? document.Substring(0, 2000) : document; + var prompt = + $"On a scale of 0 to 10, how relevant is the document to the query? " + + $"Reply with only a single number.\n\nQuery: {query}\n\nDocument: {doc}\n\nRelevance (0-10):"; + var reply = _generator.Generate(prompt); + var match = Regex.Match(reply ?? string.Empty, @"\d+(\.\d+)?"); + if (match.Success && double.TryParse(match.Value, NumberStyles.Float, CultureInfo.InvariantCulture, out double raw)) + { + var normalized = Math.Min(1.0, Math.Max(0.0, raw / 10.0)); + return NumOps.FromDouble(normalized); + } + // Unparseable reply: fall through to the lexical heuristic below. + } + var queryWords = Tokenize(query); var docWords = Tokenize(document); diff --git a/src/RetrievalAugmentedGeneration/Retrievers/MultiQueryRetriever.cs b/src/RetrievalAugmentedGeneration/Retrievers/MultiQueryRetriever.cs index 15c8e3f5ff..4e8c4dc768 100644 --- a/src/RetrievalAugmentedGeneration/Retrievers/MultiQueryRetriever.cs +++ b/src/RetrievalAugmentedGeneration/Retrievers/MultiQueryRetriever.cs @@ -19,8 +19,16 @@ public class MultiQueryRetriever : RetrieverBase { private readonly IRetriever _baseRetriever; private readonly int _numQueries; + private readonly ITextGenerator? _generator; - public MultiQueryRetriever(IRetriever baseRetriever, int numQueries = 3, int defaultTopK = 5) : base(defaultTopK) + /// The underlying retriever run once per query variation. + /// Total number of queries (original + variations) to retrieve with. + /// Default number of documents to return. + /// + /// Optional real text generator. When provided, the query variations are LLM-generated; otherwise a + /// deterministic template fallback is used (never the old "{query} variation N" placeholder). + /// + public MultiQueryRetriever(IRetriever baseRetriever, int numQueries = 3, int defaultTopK = 5, ITextGenerator? generator = null) : base(defaultTopK) { if (baseRetriever == null) throw new ArgumentNullException(nameof(baseRetriever)); @@ -29,6 +37,7 @@ public MultiQueryRetriever(IRetriever baseRetriever, int numQueries = 3, int _baseRetriever = baseRetriever; _numQueries = numQueries; + _generator = generator; } protected override IEnumerable> RetrieveCore(string query, int topK, Dictionary metadataFilters) @@ -70,14 +79,12 @@ protected override IEnumerable> RetrieveCore(string query, int topK, private List GenerateQueries(string originalQuery) { - var queries = new List { originalQuery }; - - for (int i = 1; i < _numQueries; i++) - { - queries.Add($"{originalQuery} variation {i}"); - } - - return queries; + // Reuse the (de-stubbed) LLM query expander: real LLM variations when a generator is supplied, + // deterministic template variations otherwise. ExpandQuery prepends the original, so request + // (_numQueries - 1) expansions and cap the result at _numQueries. + var expander = new QueryExpansion.LLMQueryExpansion( + numExpansions: Math.Max(1, _numQueries - 1), generator: _generator); + return expander.ExpandQuery(originalQuery).Take(_numQueries).ToList(); } } } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DeStubbedLlmComponentsTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DeStubbedLlmComponentsTests.cs new file mode 100644 index 0000000000..ae88825344 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DeStubbedLlmComponentsTests.cs @@ -0,0 +1,130 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.ContextCompression; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using AiDotNet.RetrievalAugmentedGeneration.QueryExpansion; +using AiDotNet.RetrievalAugmentedGeneration.RerankingStrategies; +using AiDotNet.RetrievalAugmentedGeneration.Retrievers; +using Xunit; + +namespace AiDotNet.Tests.UnitTests.RetrievalAugmentedGeneration; + +/// +/// Verifies the previously-stubbed LLM-powered RAG components now actually call an injected +/// (and fall back to their offline heuristic when none is supplied), +/// rather than running lexical templates with unused credentials. +/// +public class DeStubbedLlmComponentsTests +{ + /// Records prompts and returns a scripted response — a stand-in for a real LLM. + private sealed class FakeTextGen : ITextGenerator + { + private readonly Func _responder; + public List Prompts { get; } = new(); + public FakeTextGen(Func responder) => _responder = responder; + public string Generate(string prompt) + { + Prompts.Add(prompt); + return _responder(prompt); + } + } + + private sealed class RecordingRetriever : IRetriever + { + public List Queries { get; } = new(); + public int DefaultTopK => 5; + public IEnumerable> Retrieve(string query) => Retrieve(query, 5, new Dictionary()); + public IEnumerable> Retrieve(string query, int topK) => Retrieve(query, topK, new Dictionary()); + public IEnumerable> Retrieve(string query, int topK, Dictionary metadataFilters) + { + Queries.Add(query); + return new[] { new Document("d-" + Queries.Count, query) { RelevanceScore = 1.0, HasRelevanceScore = true } }; + } + } + + [Fact] + public void HyDE_WithGenerator_UsesLlmPassages() + { + var gen = new FakeTextGen(_ => "A hypothetical answer passage."); + var hyde = new HyDEQueryExpansion(gen, numHypotheticals: 2); + + var expansions = hyde.ExpandQuery("what is X?"); + + Assert.Equal("what is X?", expansions[0]); + Assert.Contains("A hypothetical answer passage.", expansions); + Assert.Equal(2, gen.Prompts.Count); + Assert.Contains("Passage:", gen.Prompts[0]); // real HyDE prompt, not a template + } + + [Fact] + public void HyDE_WithoutGenerator_FallsBackToTemplate() + { + var hyde = new HyDEQueryExpansion(); + var expansions = hyde.ExpandQuery("what is X?"); + Assert.True(expansions.Count >= 2); + Assert.Equal("what is X?", expansions[0]); + } + + [Fact] + public void LLMQueryExpansion_WithGenerator_ParsesLlmLines() + { + var gen = new FakeTextGen(_ => "1. how does X work\n- explain X\nX overview"); + var exp = new LLMQueryExpansion(numExpansions: 3, generator: gen); + + var result = exp.ExpandQuery("X"); + + Assert.Equal("X", result[0]); + Assert.Contains("how does X work", result); // numbering stripped + Assert.Contains("explain X", result); // bullet stripped + Assert.Contains("X overview", result); + Assert.Single(gen.Prompts); + } + + [Fact] + public void LLMBasedReranker_WithGenerator_OrdersByLlmScore() + { + // Score the doc mentioning "dogs" high, the other low. The query deliberately omits that keyword + // so the fake keys only on the document content, not the query echoed into the prompt. + var gen = new FakeTextGen(p => p.Contains("dogs") ? "9" : "1"); + var reranker = new LLMBasedReranker(generator: gen); + var docs = new[] + { + new Document("low", "cats sleep a lot"), + new Document("high", "dogs bark loudly"), + }; + + var ranked = reranker.Rerank("which animal is loud", docs).ToList(); + + Assert.Equal("high", ranked[0].Id); + Assert.Equal("low", ranked[1].Id); + Assert.True(gen.Prompts.Count >= 2); + } + + [Fact] + public void LLMContextCompressor_WithGenerator_ReturnsLlmExcerpt() + { + var gen = new FakeTextGen(_ => "the relevant excerpt"); + var comp = new LLMContextCompressor(generator: gen); + + var compressed = comp.CompressText("q", "a long document with lots of text and one relevant bit"); + + Assert.Equal("the relevant excerpt", compressed); + Assert.Single(gen.Prompts); + } + + [Fact] + public void MultiQueryRetriever_WithGenerator_RetrievesWithDistinctVariations() + { + var gen = new FakeTextGen(_ => "rewording one\nrewording two"); + var baseRetriever = new RecordingRetriever(); + var mq = new MultiQueryRetriever(baseRetriever, numQueries: 3, defaultTopK: 5, generator: gen); + + _ = mq.Retrieve("original query", 5, new Dictionary()).ToList(); + + Assert.Contains("original query", baseRetriever.Queries); + Assert.Contains("rewording one", baseRetriever.Queries); + Assert.True(baseRetriever.Queries.Distinct().Count() >= 2, "expected multiple distinct query variations"); + } +} From ef4959ab6f10fdb30c2c9ce0bd66e24a3bcb63db Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 15:13:07 -0400 Subject: [PATCH 03/25] feat(rag): real Qdrant + Pinecone vector-store HTTP clients Replace the in-memory brute-force simulations in QdrantDocumentStore and PineconeDocumentStore with real HTTP clients that call the Qdrant and Pinecone REST APIs. Both accept an optional HttpClient / HttpMessageHandler for testability, support API-key auth + base URL, translate metadata filters, and bridge the sync IDocumentStore methods to async internally (matching ElasticsearchDocumentStore). The IDocumentStore contract is unchanged. Adds mocked-HttpMessageHandler unit tests (no network) and gated [Trait("Category","Integration")] tests driven by env vars. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../DocumentStores/PineconeDocumentStore.cs | 614 +++++++++------- .../DocumentStores/QdrantDocumentStore.cs | 550 ++++++++++----- .../PineconeDocumentStoreTests.cs | 663 ++++++------------ .../QdrantDocumentStoreTests.cs | 353 ++++++++++ 4 files changed, 1331 insertions(+), 849 deletions(-) create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStoreTests.cs diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs index f358d23fc0..bca3ac5536 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs @@ -1,54 +1,54 @@ - using System; using System.Collections.Generic; +using System.Globalization; using System.Linq; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json; +using Newtonsoft.Json.Linq; namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores { /// - /// Pinecone-inspired document store with index-based vector organization. + /// Pinecone vector-database document store backed by the real Pinecone REST API. /// + /// The numeric type for vector operations. /// /// - /// This implementation provides an in-memory simulation of Pinecone, a fully managed vector database service. - /// It organizes documents into named indices and uses cosine similarity for retrieval. + /// This store talks to a Pinecone index over HTTP using the index host as the base URL and the + /// Api-Key header for authentication. It upserts vectors with metadata, queries with a + /// top-k + filter + includeMetadata, deletes vectors and reads index statistics. Document + /// content is stored inside the vector metadata under the reserved key _content. An optional + /// namespace scopes every operation. /// - /// For Beginners: Pinecone is a popular cloud-based vector database. - /// - /// Think of indices like separate databases: - /// - Each index has a unique name (like "ProductSearchIndex") - /// - Documents are grouped by use case - /// - Makes it easy to manage multiple vector search applications - /// - /// This in-memory version is good for: - /// - Prototyping before using real Pinecone - /// - Testing Pinecone-style index organization - /// - Small to medium document collections (< 100K documents) - /// - /// Real Pinecone provides: - /// - Fully managed cloud service (no infrastructure to manage) - /// - Auto-scaling to handle any load - /// - Advanced filtering and hybrid search - /// - Sub-50ms query latency at scale + /// For Beginners: Pinecone is a fully managed cloud vector database. This class is a + /// real client for it - each method makes an HTTP call to your Pinecone index. The index (and its + /// vector dimension) is created and managed in the Pinecone console; this client reads the + /// dimension from describe_index_stats when it starts up. /// /// - /// The numeric type for vector operations. [ComponentType(ComponentType.DocumentStore)] [PipelineStage(PipelineStage.Indexing)] public class PineconeDocumentStore : DocumentStoreBase { - private readonly Dictionary> _documents; + private const string ContentKey = "_content"; + + private readonly HttpClient _httpClient; private readonly string _indexName; + private readonly string _namespace; private int _vectorDimension; + private int _documentCount; /// - /// Gets the number of documents currently stored in the index. + /// Gets the number of vectors currently stored in the index (within the configured namespace). /// - public override int DocumentCount => _documents.Count; + public override int DocumentCount => _documentCount; /// /// Gets the dimensionality of vectors stored in this index. @@ -56,279 +56,407 @@ public class PineconeDocumentStore : DocumentStoreBase public override int VectorDimension => _vectorDimension; /// - /// Initializes a new instance of the PineconeDocumentStore class. + /// Gets the name of the Pinecone index this store is bound to. /// - /// The name of the index to organize documents. - /// The initial capacity for the internal dictionary (default: 1000). - /// Thrown when index name is empty or initial capacity is not positive. - /// - /// For Beginners: Creates a new Pinecone-style index. - /// - /// Example: - /// - /// // Create an index for product vectors - /// var store = new PineconeDocumentStore<float>("ProductIndex"); - /// - /// // Create a larger index for articles - /// var bigStore = new PineconeDocumentStore<double>("ArticleIndex", 10000); - /// - /// - /// The index name helps organize different vector collections. - /// - /// - public PineconeDocumentStore(string indexName, int initialCapacity = 1000) + public string IndexName => _indexName; + + /// + /// Gets the namespace all operations are scoped to (empty string means the default namespace). + /// + public string Namespace => _namespace; + + /// + /// Initializes a new instance of the class. + /// + /// A logical name for the index (used for diagnostics). + /// The index host base URL, e.g. https://my-index-abc123.svc.us-east1-gcp.pinecone.io. + /// The Pinecone API key (sent as the Api-Key header). + /// Optional namespace that scopes all operations. + /// + /// Optional pre-configured . Primarily for testing; when supplied its + /// is used if already set. + /// + /// Optional used to build the client (for testing). + public PineconeDocumentStore( + string indexName, + string url, + string apiKey, + string? @namespace = null, + HttpClient? httpClient = null, + HttpMessageHandler? handler = null) { if (string.IsNullOrWhiteSpace(indexName)) throw new ArgumentException("Index name cannot be empty", nameof(indexName)); - if (initialCapacity <= 0) - throw new ArgumentException("Initial capacity must be greater than zero", nameof(initialCapacity)); _indexName = indexName; - _documents = new Dictionary>(initialCapacity); + _namespace = @namespace ?? string.Empty; + _documentCount = 0; _vectorDimension = 0; + + _httpClient = httpClient ?? (handler != null ? new HttpClient(handler) : new HttpClient()); + if (_httpClient.BaseAddress == null) + { + if (string.IsNullOrWhiteSpace(url)) + throw new ArgumentException("Url cannot be empty", nameof(url)); + _httpClient.BaseAddress = new Uri(url); + } + + if (!string.IsNullOrWhiteSpace(apiKey) && !_httpClient.DefaultRequestHeaders.Contains("Api-Key")) + _httpClient.DefaultRequestHeaders.Add("Api-Key", apiKey); + + InitializeStats(); } - /// - /// Core logic for adding a single vector document to the index. - /// - /// The validated vector document to add. - /// - /// - /// The first document added determines the vector dimension for the entire index. - /// All subsequent documents must have embeddings of the same dimension. - /// - /// - protected override void AddCore(VectorDocument vectorDocument) + private void InitializeStats() { - if (_documents.Count == 0) + try { - _vectorDimension = vectorDocument.Embedding.Length; + var info = SendAsync(HttpMethod.Post, "/describe_index_stats", new { }).GetAwaiter().GetResult(); + if (!IsSuccess(info.Status)) + return; + + var root = JObject.Parse(info.Body); + var dimension = root["dimension"]?.Value(); + if (dimension.HasValue && dimension.Value > 0) + _vectorDimension = dimension.Value; + + _documentCount = ReadNamespaceCount(root); + } + catch (HttpRequestException) + { + // Index host not reachable at construction; stats are refreshed on later operations. + } + } + + private int ReadNamespaceCount(JObject stats) + { + var namespaces = stats["namespaces"] as JObject; + if (namespaces != null) + { + var key = _namespace.Length == 0 ? string.Empty : _namespace; + var ns = namespaces[key] as JObject; + if (ns != null) + return ns["vectorCount"]?.Value() ?? 0; + + if (_namespace.Length == 0) + return stats["totalVectorCount"]?.Value() ?? 0; + + return 0; } - _documents[vectorDocument.Document.Id] = vectorDocument; + return stats["totalVectorCount"]?.Value() ?? 0; } - /// - /// Core logic for adding multiple vector documents in a batch operation. - /// - /// The validated list of vector documents to add. - /// Thrown when a document's embedding has inconsistent dimensions. - /// - /// - /// Batch addition is significantly more efficient than adding documents one at a time. - /// All documents in the batch must have embeddings with the same dimension. - /// - /// For Beginners: Adding documents in batches is much faster. - /// - /// Bad (slow): - /// - /// foreach (var doc in documents) - /// store.Add(doc); - /// - /// - /// Good (fast): - /// - /// store.AddBatch(documents); - /// - /// - /// + /// + protected override void AddCore(VectorDocument vectorDocument) + { + if (_vectorDimension == 0) + _vectorDimension = vectorDocument.Embedding.Length; + else if (vectorDocument.Embedding.Length != _vectorDimension) + throw new ArgumentException( + $"Vector dimension mismatch. Expected {_vectorDimension}, got {vectorDocument.Embedding.Length}", + nameof(vectorDocument)); + + Upsert(new[] { vectorDocument }); + _documentCount++; + } + + /// protected override void AddBatchCore(IList> vectorDocuments) { - if (_vectorDimension == 0 && vectorDocuments.Count > 0) - { + if (vectorDocuments.Count == 0) + return; + + if (_vectorDimension == 0) _vectorDimension = vectorDocuments[0].Embedding.Length; - } - foreach (var vectorDocument in vectorDocuments) + foreach (var vd in vectorDocuments) { - if (vectorDocument.Embedding.Length != _vectorDimension) + if (vd.Embedding.Length != _vectorDimension) throw new ArgumentException( - $"Vector dimension mismatch in batch. Expected {_vectorDimension}, got {vectorDocument.Embedding.Length} for document {vectorDocument.Document.Id}", + $"Vector dimension mismatch in batch. Expected {_vectorDimension}, got {vd.Embedding.Length} for document {vd.Document.Id}", nameof(vectorDocuments)); + } + + Upsert(vectorDocuments); + _documentCount += vectorDocuments.Count; + } + + private void Upsert(IEnumerable> vectorDocuments) + { + var vectors = vectorDocuments.Select(vd => + { + var values = vd.Embedding.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + var metadata = BuildMetadata(vd.Document); + return (object)new + { + id = vd.Document.Id, + values, + metadata + }; + }).ToList(); + + var body = new Dictionary { ["vectors"] = vectors }; + if (_namespace.Length > 0) + body["namespace"] = _namespace; + + var info = SendAsync(HttpMethod.Post, "/vectors/upsert", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "upsert"); + } + + private static Dictionary BuildMetadata(Document document) + { + var metadata = new Dictionary(); + if (document.Metadata != null) + { + foreach (var kvp in document.Metadata) + metadata[kvp.Key] = kvp.Value; + } + + metadata[ContentKey] = document.Content; + return metadata; + } + + /// + protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + { + var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + + var body = new Dictionary + { + ["vector"] = vector, + ["topK"] = topK, + ["includeMetadata"] = true, + ["includeValues"] = false + }; + + var filter = BuildFilter(metadataFilters); + if (filter != null) + body["filter"] = filter; + + if (_namespace.Length > 0) + body["namespace"] = _namespace; + + var info = SendAsync(HttpMethod.Post, "/query", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "query"); + + var results = new List>(); + var matches = JObject.Parse(info.Body)["matches"] as JArray; + if (matches == null) + return results; + + foreach (var match in matches) + { + var id = match?["id"]?.ToString(); + if (string.IsNullOrEmpty(id)) + continue; - _documents[vectorDocument.Document.Id] = vectorDocument; + var doc = ParseMetadata(id!, match?["metadata"] as JObject); + var score = match?["score"] != null ? Convert.ToDouble(match!["score"], CultureInfo.InvariantCulture) : 0.0; + doc.RelevanceScore = NumOps.FromDouble(score); + doc.HasRelevanceScore = true; + results.Add(doc); } + + return results; } /// - /// Core logic for similarity search using cosine similarity with optional metadata filtering. + /// Translates the metadata filter dictionary into a Pinecone metadata filter object. /// - /// The validated query vector. - /// The validated number of documents to return. - /// The validated metadata filters. - /// Top-k similar documents ordered by cosine similarity score. /// - /// - /// Performs a similarity search across all documents in the index, optionally filtering - /// by metadata. Results are ordered by decreasing cosine similarity and limited to top-k matches. - /// - /// For Beginners: Finds the most similar documents to your query. - /// - /// How it works: - /// 1. Filter documents by metadata (if provided) - /// 2. Calculate similarity between query and each document's embedding - /// 3. Sort by similarity (highest first) - /// 4. Return the top-k best matches - /// - /// Example: - /// - /// // Find 10 most similar products - /// var results = store.GetSimilar(queryVector, topK: 10); - /// - /// // Find similar products in "Electronics" category - /// var filters = new Dictionary<string, object> { ["category"] = "Electronics" }; - /// var filtered = store.GetSimilarWithFilters(queryVector, 5, filters); - /// - /// + /// Equality (string/bool) becomes $eq, numeric values become $gte (mirroring the + /// base-class "field >= value" semantics) and collection values become $in. Multiple + /// filters are ANDed together, matching Pinecone's default combination of top-level keys. /// - protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + private static object? BuildFilter(Dictionary? metadataFilters) { - var scoredDocuments = new List<(Document Document, T Score)>(); + if (metadataFilters == null || metadataFilters.Count == 0) + return null; - var matchingDocuments = _documents.Values - .Where(vectorDoc => MatchesFilters(vectorDoc.Document, metadataFilters)); + var filter = new Dictionary(); - foreach (var vectorDoc in matchingDocuments) + foreach (var kvp in metadataFilters) { - var similarity = StatisticsHelper.CosineSimilarity(queryVector, vectorDoc.Embedding); - scoredDocuments.Add((vectorDoc.Document, similarity)); - } + var value = kvp.Value; - var results = scoredDocuments - .OrderByDescending(x => x.Score) - .Take(topK) - .Select(x => + if (value == null || value is string || value is bool) + { + filter[kvp.Key] = new Dictionary { ["$eq"] = value! }; + } + else if (IsNumeric(value)) { - // Create a new Document instance to avoid mutating the cached one - var newDocument = new Document(x.Document.Id, x.Document.Content, x.Document.Metadata) + filter[kvp.Key] = new Dictionary { - RelevanceScore = x.Score, - HasRelevanceScore = true + ["$gte"] = Convert.ToDouble(value, CultureInfo.InvariantCulture) }; - return newDocument; - }) - .ToList(); + } + else if (value is System.Collections.IEnumerable enumerable) + { + var items = enumerable.Cast().ToList(); + filter[kvp.Key] = new Dictionary { ["$in"] = items }; + } + else + { + filter[kvp.Key] = new Dictionary { ["$eq"] = value.ToString() ?? string.Empty }; + } + } - return results; + return filter; } - /// - /// Core logic for retrieving a document by its unique identifier. - /// - /// The validated document ID. - /// The document if found; otherwise, null. - /// - /// For Beginners: Gets a specific document by its ID. - /// - /// Example: - /// - /// var doc = store.GetById("product-123"); - /// if (doc != null) - /// Console.WriteLine($"Product: {doc.Content}"); - /// - /// - /// + /// protected override Document? GetByIdCore(string documentId) { - return _documents.TryGetValue(documentId, out var vectorDoc) ? vectorDoc.Document : null; + var path = $"/vectors/fetch?ids={Uri.EscapeDataString(documentId)}"; + if (_namespace.Length > 0) + path += $"&namespace={Uri.EscapeDataString(_namespace)}"; + + var info = SendAsync(HttpMethod.Get, path, null).GetAwaiter().GetResult(); + if (info.Status == HttpStatusCode.NotFound) + return null; + EnsureSuccess(info, "fetch"); + + var vectors = JObject.Parse(info.Body)["vectors"] as JObject; + var vector = vectors?[documentId] as JObject; + if (vector == null) + return null; + + return ParseMetadata(documentId, vector["metadata"] as JObject); } - /// - /// Core logic for removing a document from the index. - /// - /// The validated document ID. - /// True if the document was found and removed; otherwise, false. - /// - /// - /// Removes the document from the index. If this was the last document, the vector dimension - /// is reset to 0, allowing a new dimension on the next add operation. - /// - /// For Beginners: Deletes a document from the index. - /// - /// Example: - /// - /// bool removed = store.Remove("product-123"); - /// if (removed) - /// Console.WriteLine("Document deleted"); - /// - /// - /// + /// protected override bool RemoveCore(string documentId) { - var removed = _documents.Remove(documentId); - if (removed && _documents.Count == 0) - { - _vectorDimension = 0; - } - return removed; + var body = new Dictionary { ["ids"] = new[] { documentId } }; + if (_namespace.Length > 0) + body["namespace"] = _namespace; + + var info = SendAsync(HttpMethod.Post, "/vectors/delete", body).GetAwaiter().GetResult(); + if (!IsSuccess(info.Status)) + return false; + + if (_documentCount > 0) + _documentCount--; + return true; } - /// - /// Core logic for retrieving all documents in the index. - /// - /// An enumerable of all documents without their vector embeddings. - /// - /// - /// Returns all documents in the index in no particular order. Vector embeddings are not included. - /// For large indices, consider the memory impact of loading all documents at once. - /// - /// For Beginners: Gets every document in the index. - /// - /// Use cases: - /// - Export all documents for backup - /// - Migrate to a different index or store - /// - Bulk analysis or processing - /// - Debugging to see what's indexed - /// - /// Warning: For large indices (> 10K documents), this can use a lot of memory. - /// - /// Example: - /// - /// // Get all documents - /// var allDocs = store.GetAll().ToList(); - /// // Result is available in the returned value - /// - /// // Export to file - /// var json = JsonConvert.SerializeObject(allDocs); - /// File.WriteAllText($"{_indexName}_backup.json", json); - /// - /// - /// + /// protected override IEnumerable> GetAllCore() { - return _documents.Values.Select(vd => vd.Document).ToList(); + var all = new List>(); + string? paginationToken = null; + + while (true) + { + var path = "/vectors/list?limit=100"; + if (_namespace.Length > 0) + path += $"&namespace={Uri.EscapeDataString(_namespace)}"; + if (!string.IsNullOrEmpty(paginationToken)) + path += $"&paginationToken={Uri.EscapeDataString(paginationToken!)}"; + + var info = SendAsync(HttpMethod.Get, path, null).GetAwaiter().GetResult(); + EnsureSuccess(info, "list"); + + var root = JObject.Parse(info.Body); + var ids = (root["vectors"] as JArray)? + .Select(v => v?["id"]?.ToString()) + .Where(id => !string.IsNullOrEmpty(id)) + .Select(id => id!) + .ToList() ?? new List(); + + foreach (var id in ids) + { + var doc = GetByIdCore(id); + if (doc != null) + all.Add(doc); + } + + paginationToken = root["pagination"]?["next"]?.ToString(); + if (string.IsNullOrEmpty(paginationToken)) + break; + } + + return all; } - /// - /// Removes all documents from the index and resets the vector dimension. - /// - /// - /// - /// Clears all documents from the index and resets the vector dimension to 0. - /// The index name remains unchanged and is ready to accept new documents. - /// - /// For Beginners: Completely empties the index. - /// - /// After calling Clear(): - /// - All documents are removed - /// - Vector dimension resets to 0 - /// - Index name stays the same - /// - Ready for new documents (even with different dimensions) - /// - /// Use with caution - this operation cannot be undone! - /// - /// Example: - /// - /// store.Clear(); - /// // Result is available in the returned value // 0 - /// - /// - /// + /// public override void Clear() { - _documents.Clear(); - _vectorDimension = 0; + var body = new Dictionary { ["deleteAll"] = true }; + if (_namespace.Length > 0) + body["namespace"] = _namespace; + + var info = SendAsync(HttpMethod.Post, "/vectors/delete", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "delete all"); + + _documentCount = 0; + } + + private Document ParseMetadata(string id, JObject? metadata) + { + var content = string.Empty; + var docMetadata = new Dictionary(); + + if (metadata != null) + { + foreach (var property in metadata.Properties()) + { + if (property.Name == ContentKey) + content = property.Value?.ToString() ?? string.Empty; + else + docMetadata[property.Name] = property.Value?.ToObject() ?? string.Empty; + } + } + + return new Document(id, content, docMetadata); + } + + private static bool IsNumeric(object value) + { + return value is sbyte || value is byte || value is short || value is ushort + || value is int || value is uint || value is long || value is ulong + || value is float || value is double || value is decimal; + } + + private static bool IsSuccess(HttpStatusCode status) => (int)status >= 200 && (int)status < 300; + + private void EnsureSuccess(HttpResponseInfo info, string operation) + { + if (!IsSuccess(info.Status)) + throw new HttpRequestException($"Pinecone {operation} failed with status {(int)info.Status}: {info.Body}"); + } + + private async Task SendAsync(HttpMethod method, string path, object? body) + { + using (var request = new HttpRequestMessage(method, path)) + { + if (body != null) + { + var json = JsonConvert.SerializeObject(body); + request.Content = new StringContent(json, Encoding.UTF8, "application/json"); + } + + using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) + { + var content = response.Content != null + ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + return new HttpResponseInfo(response.StatusCode, content); + } + } + } + + private readonly struct HttpResponseInfo + { + public HttpResponseInfo(HttpStatusCode status, string body) + { + Status = status; + Body = body; + } + + public HttpStatusCode Status { get; } + public string Body { get; } } } } - diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs index c7ad5964a9..1b7f634148 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs @@ -1,274 +1,484 @@ - using System; using System.Collections.Generic; +using System.Globalization; using System.Linq; +using System.Net; +using System.Net.Http; +using System.Security.Cryptography; +using System.Text; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json; +using Newtonsoft.Json.Linq; namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores { /// - /// Qdrant-inspired document store with collection-based organization and payload filtering. - /// Provides in-memory simulation of Qdrant features including collection management and efficient filtering. + /// Qdrant vector-database document store backed by the real Qdrant REST API. /// /// The numeric type for vector operations. + /// + /// + /// This store talks to a running Qdrant instance over HTTP. It manages a single collection, + /// upserts points (vector + payload), performs filtered vector search, deletes points and + /// scrolls the whole collection. Document metadata is stored under a nested metadata + /// payload object, while the original string document id and content are stored under the + /// reserved payload keys _doc_id and _content. + /// + /// For Beginners: Qdrant is an open-source vector database. This class is a real + /// client for it - every method here makes an HTTP call to your Qdrant server. Qdrant requires + /// point ids to be unsigned integers or UUIDs, so each document's string id is turned into a + /// deterministic UUID; the original string id is kept in the payload so you always get it back. + /// + /// [ComponentType(ComponentType.DocumentStore)] [PipelineStage(PipelineStage.Indexing)] public class QdrantDocumentStore : DocumentStoreBase { - private readonly Dictionary> _documents; + private const string DocIdKey = "_doc_id"; + private const string ContentKey = "_content"; + + private readonly HttpClient _httpClient; private readonly string _collectionName; + private readonly string _distance; private int _vectorDimension; - private readonly Dictionary> _payloadIndex; + private int _documentCount; + private bool _collectionReady; + + /// + /// Gets the number of documents (points) currently stored in the collection. + /// + public override int DocumentCount => _documentCount; - public override int DocumentCount => _documents.Count; + /// + /// Gets the dimensionality of vectors stored in this collection. + /// public override int VectorDimension => _vectorDimension; - public string CollectionName { get; private set; } + /// + /// Gets the name of the Qdrant collection this store is bound to. + /// + public string CollectionName { get; } - public QdrantDocumentStore(string collectionName, int initialCapacity = 1000) + /// + /// Initializes a new instance of the class. + /// + /// The Qdrant collection name. + /// The base URL of the Qdrant server, e.g. http://localhost:6333. + /// Optional Qdrant API key (sent as the api-key header). + /// The distance metric used when creating the collection. + /// + /// The vector dimension used to create the collection if it does not already exist. + /// When 0 the collection must already exist (its dimension is read from the server) or the + /// dimension is inferred from the first document added. + /// + /// + /// Optional pre-configured . Primarily for testing; when supplied its + /// is used if already set. + /// + /// Optional used to build the client (for testing). + public QdrantDocumentStore( + string collectionName, + string url, + string? apiKey = null, + DistanceMetricType distanceMetric = DistanceMetricType.Cosine, + int vectorDimension = 0, + HttpClient? httpClient = null, + HttpMessageHandler? handler = null) { if (string.IsNullOrWhiteSpace(collectionName)) throw new ArgumentException("Collection name cannot be empty", nameof(collectionName)); - if (initialCapacity <= 0) - throw new ArgumentException("Initial capacity must be greater than zero", nameof(initialCapacity)); + if (vectorDimension < 0) + throw new ArgumentOutOfRangeException(nameof(vectorDimension), "Vector dimension cannot be negative"); _collectionName = collectionName; CollectionName = collectionName; - _documents = new Dictionary>(initialCapacity); - _payloadIndex = new Dictionary>(); - _vectorDimension = 0; + _distance = MapDistance(distanceMetric); + _vectorDimension = vectorDimension; + _documentCount = 0; + + _httpClient = httpClient ?? (handler != null ? new HttpClient(handler) : new HttpClient()); + if (_httpClient.BaseAddress == null) + { + if (string.IsNullOrWhiteSpace(url)) + throw new ArgumentException("Url cannot be empty", nameof(url)); + _httpClient.BaseAddress = new Uri(url); + } + + if (!string.IsNullOrWhiteSpace(apiKey) && !_httpClient.DefaultRequestHeaders.Contains("api-key")) + _httpClient.DefaultRequestHeaders.Add("api-key", apiKey); + + InitializeCollection(vectorDimension); } - protected override void AddCore(VectorDocument vectorDocument) + private static string MapDistance(DistanceMetricType metric) { - if (_documents.Count == 0) + switch (metric) { - _vectorDimension = vectorDocument.Embedding.Length; + case DistanceMetricType.Cosine: + return "Cosine"; + case DistanceMetricType.Euclidean: + return "Euclid"; + case DistanceMetricType.Manhattan: + return "Manhattan"; + default: + throw new NotSupportedException( + $"Distance metric '{metric}' is not supported by Qdrant. Use Cosine, Euclidean or Manhattan."); } + } - _documents[vectorDocument.Document.Id] = vectorDocument; - IndexPayload(vectorDocument.Document); + private void InitializeCollection(int requestedDimension) + { + HttpResponseInfo info; + try + { + info = SendAsync(HttpMethod.Get, $"/collections/{_collectionName}", null).GetAwaiter().GetResult(); + } + catch (HttpRequestException) + { + // Server not reachable at construction time; defer creation to first write. + return; + } + + if (info.Status == HttpStatusCode.OK) + { + var root = JObject.Parse(info.Body)["result"]; + var size = root?["config"]?["params"]?["vectors"]?["size"]?.Value(); + if (size.HasValue && size.Value > 0) + _vectorDimension = size.Value; + + var count = root?["points_count"]?.Value(); + if (count.HasValue) + _documentCount = count.Value; + + _collectionReady = true; + return; + } + + if (info.Status == HttpStatusCode.NotFound && requestedDimension > 0) + CreateCollection(requestedDimension); } + private void CreateCollection(int dimension) + { + var body = new + { + vectors = new + { + size = dimension, + distance = _distance + } + }; + + var info = SendAsync(HttpMethod.Put, $"/collections/{_collectionName}", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "create collection"); + _vectorDimension = dimension; + _collectionReady = true; + } + + private void EnsureCollectionForDimension(int dimension) + { + if (_collectionReady) + return; + CreateCollection(dimension); + } + + /// + protected override void AddCore(VectorDocument vectorDocument) + { + EnsureCollectionForDimension(vectorDocument.Embedding.Length); + if (_vectorDimension == 0) + _vectorDimension = vectorDocument.Embedding.Length; + + var body = new { points = new[] { BuildPoint(vectorDocument) } }; + var info = SendAsync(HttpMethod.Put, $"/collections/{_collectionName}/points?wait=true", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "upsert point"); + _documentCount++; + } + + /// protected override void AddBatchCore(IList> vectorDocuments) { if (vectorDocuments.Count == 0) return; - if (_documents.Count == 0) - { + EnsureCollectionForDimension(vectorDocuments[0].Embedding.Length); + if (_vectorDimension == 0) _vectorDimension = vectorDocuments[0].Embedding.Length; - } - foreach (var vectorDoc in vectorDocuments) + var points = vectorDocuments.Select(BuildPoint).ToList(); + var body = new { points }; + var info = SendAsync(HttpMethod.Put, $"/collections/{_collectionName}/points?wait=true", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "batch upsert points"); + _documentCount += vectorDocuments.Count; + } + + private object BuildPoint(VectorDocument vectorDocument) + { + var vector = vectorDocument.Embedding.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + var payload = new Dictionary { - _documents[vectorDoc.Document.Id] = vectorDoc; - IndexPayload(vectorDoc.Document); - } + [DocIdKey] = vectorDocument.Document.Id, + [ContentKey] = vectorDocument.Document.Content, + ["metadata"] = vectorDocument.Document.Metadata ?? new Dictionary() + }; + + return new + { + id = ToPointId(vectorDocument.Document.Id), + vector, + payload + }; } + /// protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) { - var scoredDocuments = new List<(Document Document, T Score)>(); + var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); - var candidateIds = GetFilteredCandidates(metadataFilters); - IEnumerable> candidates; - - if (candidateIds != null) + var body = new Dictionary { - candidates = candidateIds - .Where(id => _documents.ContainsKey(id)) - .Select(id => _documents[id]); - } - else + ["vector"] = vector, + ["limit"] = topK, + ["with_payload"] = true, + ["with_vector"] = false + }; + + var filter = BuildFilter(metadataFilters); + if (filter != null) + body["filter"] = filter; + + var info = SendAsync(HttpMethod.Post, $"/collections/{_collectionName}/points/search", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "search"); + + var results = new List>(); + var hits = JObject.Parse(info.Body)["result"] as JArray; + if (hits == null) + return results; + + foreach (var hit in hits) { - candidates = _documents.Values; + var doc = ParsePayload(hit?["payload"] as JObject); + if (doc == null) + continue; + + var score = hit?["score"] != null ? Convert.ToDouble(hit!["score"], CultureInfo.InvariantCulture) : 0.0; + doc.RelevanceScore = NumOps.FromDouble(score); + doc.HasRelevanceScore = true; + results.Add(doc); } - var matchingDocuments = candidates - .Where(vectorDoc => MatchesFilters(vectorDoc.Document, metadataFilters)); + return results; + } + + /// + /// Translates the metadata filter dictionary into a Qdrant filter object. + /// + /// + /// Equality (string/bool) becomes a match condition and numeric values become a + /// range with gte (mirroring the base-class "field >= value" semantics). + /// All scalar conditions are combined under must. Collection values become should (any-of). + /// + private static object? BuildFilter(Dictionary? metadataFilters) + { + if (metadataFilters == null || metadataFilters.Count == 0) + return null; + + var must = new List(); + var should = new List(); - foreach (var vectorDoc in matchingDocuments) + foreach (var kvp in metadataFilters) { - var similarity = StatisticsHelper.CosineSimilarity(queryVector, vectorDoc.Embedding); - scoredDocuments.Add((vectorDoc.Document, similarity)); - } + var key = "metadata." + kvp.Key; + var value = kvp.Value; - var results = scoredDocuments - .OrderByDescending(x => x.Score) - .Take(topK) - .Select(x => + if (value == null || value is string || value is bool) + { + must.Add(new { key, match = new { value } }); + } + else if (IsNumeric(value)) + { + must.Add(new { key, range = new { gte = Convert.ToDouble(value, CultureInfo.InvariantCulture) } }); + } + else if (value is System.Collections.IEnumerable enumerable) { - x.Document.RelevanceScore = x.Score; - x.Document.HasRelevanceScore = true; - return x.Document; - }) - .ToList(); + foreach (var item in enumerable) + should.Add(new { key, match = new { value = item } }); + } + else + { + must.Add(new { key, match = new { value = value.ToString() } }); + } + } - return results; + var filter = new Dictionary(); + if (must.Count > 0) + filter["must"] = must; + if (should.Count > 0) + filter["should"] = should; + + return filter.Count > 0 ? filter : null; } + /// protected override Document? GetByIdCore(string documentId) { - return _documents.TryGetValue(documentId, out var vectorDoc) ? vectorDoc.Document : null; + var info = SendAsync(HttpMethod.Get, $"/collections/{_collectionName}/points/{ToPointId(documentId)}", null).GetAwaiter().GetResult(); + if (info.Status == HttpStatusCode.NotFound) + return null; + EnsureSuccess(info, "get point"); + + var result = JObject.Parse(info.Body)["result"] as JObject; + if (result == null) + return null; + + return ParsePayload(result["payload"] as JObject); } + /// protected override bool RemoveCore(string documentId) { - if (!_documents.TryGetValue(documentId, out var vectorDoc)) + var body = new { points = new[] { ToPointId(documentId) } }; + var info = SendAsync(HttpMethod.Post, $"/collections/{_collectionName}/points/delete?wait=true", body).GetAwaiter().GetResult(); + if (!IsSuccess(info.Status)) return false; - RemoveFromPayloadIndex(vectorDoc.Document); - _documents.Remove(documentId); - - if (_documents.Count == 0) - { - _vectorDimension = 0; - } - + if (_documentCount > 0) + _documentCount--; return true; } - /// - /// Core logic for retrieving all documents in the collection. - /// - /// An enumerable of all documents without their vector embeddings. - /// - /// - /// Returns all documents from the Qdrant collection in no particular order. - /// Vector embeddings are not included in the results, only document content and metadata. - /// - /// For Beginners: Gets every document in the collection. - /// - /// Use cases: - /// - Export all documents for backup - /// - Migrate to a different collection - /// - Bulk processing or analysis - /// - Debugging payload indices - /// - /// Warning: For large collections (> 10K documents), this can use significant memory. - /// In real Qdrant, consider using scroll API with pagination for large collections. - /// - /// Example: - /// - /// // Get all documents - /// var allDocs = store.GetAll().ToList(); - /// // Result is available in the returned value - /// - /// // Export to JSON - /// var json = JsonConvert.SerializeObject(allDocs); - /// File.WriteAllText($"{_collectionName}_export.json", json); - /// - /// - /// + /// protected override IEnumerable> GetAllCore() { - return _documents.Values.Select(vd => vd.Document).ToList(); + var all = new List>(); + object? offset = null; + + while (true) + { + var body = new Dictionary + { + ["limit"] = 256, + ["with_payload"] = true, + ["with_vector"] = false + }; + if (offset != null) + body["offset"] = offset; + + var info = SendAsync(HttpMethod.Post, $"/collections/{_collectionName}/points/scroll", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "scroll"); + + var result = JObject.Parse(info.Body)["result"] as JObject; + var points = result?["points"] as JArray; + if (points == null || points.Count == 0) + break; + + foreach (var point in points) + { + var doc = ParsePayload(point?["payload"] as JObject); + if (doc != null) + all.Add(doc); + } + + var next = result?["next_page_offset"]; + if (next == null || next.Type == JTokenType.Null) + break; + offset = next.ToObject(); + } + + return all; } - /// - /// Removes all documents from the collection and clears all indices. - /// - /// - /// - /// Clears all documents, payload indices, and resets the vector dimension to 0. - /// The collection name remains unchanged and is ready to accept new documents. - /// - /// For Beginners: Completely empties the collection and all its indices. - /// - /// After calling Clear(): - /// - All documents are removed - /// - Payload index is cleared - /// - Vector dimension resets to 0 - /// - Collection name stays the same - /// - Ready for new documents - /// - /// Use with caution - this cannot be undone! - /// - /// Example: - /// - /// store.Clear(); - /// // Result is available in the returned value // 0 - /// - /// - /// + /// public override void Clear() { - _documents.Clear(); - _payloadIndex.Clear(); - _vectorDimension = 0; + var info = SendAsync(HttpMethod.Delete, $"/collections/{_collectionName}", null).GetAwaiter().GetResult(); + EnsureSuccess(info, "delete collection"); + + _documentCount = 0; + _collectionReady = false; + + if (_vectorDimension > 0) + CreateCollection(_vectorDimension); } - private void IndexPayload(Document document) + private Document? ParsePayload(JObject? payload) { - foreach (var kvp in document.Metadata) - { - var payloadKey = CreatePayloadKey(kvp.Key, kvp.Value); + if (payload == null) + return null; - if (!_payloadIndex.ContainsKey(payloadKey)) - { - _payloadIndex[payloadKey] = new HashSet(); - } + var id = payload[DocIdKey]?.ToString(); + if (string.IsNullOrEmpty(id)) + return null; - _payloadIndex[payloadKey].Add(document.Id); - } + var content = payload[ContentKey]?.ToString() ?? string.Empty; + var metadata = (payload["metadata"] as JObject)?.ToObject>() + ?? new Dictionary(); + + return new Document(id!, content, metadata); } - private void RemoveFromPayloadIndex(Document document) + private static bool IsNumeric(object value) { - foreach (var kvp in document.Metadata) - { - var payloadKey = CreatePayloadKey(kvp.Key, kvp.Value); + return value is sbyte || value is byte || value is short || value is ushort + || value is int || value is uint || value is long || value is ulong + || value is float || value is double || value is decimal; + } - if (_payloadIndex.TryGetValue(payloadKey, out var docIds)) - { - docIds.Remove(document.Id); - if (docIds.Count == 0) - { - _payloadIndex.Remove(payloadKey); - } - } + private static bool IsSuccess(HttpStatusCode status) => (int)status >= 200 && (int)status < 300; + + /// + /// Produces a deterministic UUID string for an arbitrary document id, since Qdrant point + /// ids must be unsigned integers or UUIDs. The mapping is stable so lookups/deletes work. + /// + private static string ToPointId(string documentId) + { + using (var md5 = MD5.Create()) + { + var hash = md5.ComputeHash(Encoding.UTF8.GetBytes(documentId)); + return new Guid(hash).ToString(); } } - private HashSet? GetFilteredCandidates(Dictionary metadataFilters) + private void EnsureSuccess(HttpResponseInfo info, string operation) { - if (metadataFilters.Count == 0) - return null; - - HashSet? candidateIds = null; + if (!IsSuccess(info.Status)) + throw new HttpRequestException($"Qdrant {operation} failed with status {(int)info.Status}: {info.Body}"); + } - foreach (var filter in metadataFilters) + private async Task SendAsync(HttpMethod method, string path, object? body) + { + using (var request = new HttpRequestMessage(method, path)) { - var payloadKey = CreatePayloadKey(filter.Key, filter.Value); - - if (_payloadIndex.TryGetValue(payloadKey, out var docIds)) + if (body != null) { - if (candidateIds == null) - { - candidateIds = new HashSet(docIds); - } - else - { - candidateIds.IntersectWith(docIds); - } + var json = JsonConvert.SerializeObject(body); + request.Content = new StringContent(json, Encoding.UTF8, "application/json"); } - else + + using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) { - return new HashSet(); + var content = response.Content != null + ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + return new HttpResponseInfo(response.StatusCode, content); } } - - return candidateIds; } - private static string CreatePayloadKey(string fieldName, object fieldValue) + private readonly struct HttpResponseInfo { - var valueStr = fieldValue?.ToString() ?? string.Empty; - return $"{fieldName}:{valueStr}"; + public HttpResponseInfo(HttpStatusCode status, string body) + { + Status = status; + Body = body; + } + + public HttpStatusCode Status { get; } + public string Body { get; } } } } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStoreTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStoreTests.cs index 4ce4e5cb15..bcd246f1b8 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStoreTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStoreTests.cs @@ -1,545 +1,336 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; using AiDotNet.RetrievalAugmentedGeneration.Models; -using AiDotNet.Tensors.LinearAlgebra; +using Newtonsoft.Json.Linq; using Xunit; -using System.Threading.Tasks; namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.DocumentStores { + /// + /// Unit tests for the real Pinecone REST client. All HTTP traffic is intercepted by a mock + /// so these run in CI with no network access. + /// public class PineconeDocumentStoreTests { - #region Constructor Tests - - [Fact(Timeout = 60000)] - public async Task Constructor_WithValidIndexName_CreatesStore() - { - // Arrange & Act - var store = new PineconeDocumentStore("TestIndex"); - - // Assert - Assert.Equal(0, store.DocumentCount); - Assert.Equal(0, store.VectorDimension); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithCustomCapacity_CreatesStore() - { - // Arrange & Act - var store = new PineconeDocumentStore("TestIndex", initialCapacity: 5000); + private const string BaseUrl = "https://test-index.svc.pinecone.io"; + private const string Index = "test-index"; - // Assert - Assert.Equal(0, store.DocumentCount); - Assert.Equal(0, store.VectorDimension); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithEmptyIndexName_ThrowsArgumentException() - { - // Act & Assert - Assert.Throws(() => - new PineconeDocumentStore("")); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullIndexName_ThrowsArgumentException() - { - // Act & Assert - Assert.Throws(() => - new PineconeDocumentStore(null!)); - } + // describe_index_stats for a 3-d index with 0 vectors in the default namespace. + private const string StatsJson = + "{\"namespaces\":{\"\":{\"vectorCount\":0}},\"dimension\":3,\"totalVectorCount\":0}"; - [Fact(Timeout = 60000)] - public async Task Constructor_WithWhitespaceIndexName_ThrowsArgumentException() - { - // Act & Assert - Assert.Throws(() => - new PineconeDocumentStore(" ")); - } + #region infrastructure - [Fact(Timeout = 60000)] - public async Task Constructor_WithZeroCapacity_ThrowsArgumentException() + private sealed class RecordedRequest { - // Act & Assert - Assert.Throws(() => - new PineconeDocumentStore("TestIndex", initialCapacity: 0)); - } + public RecordedRequest(string method, string path, string body) + { + Method = method; + Path = path; + Body = body; + } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNegativeCapacity_ThrowsArgumentException() - { - // Act & Assert - Assert.Throws(() => - new PineconeDocumentStore("TestIndex", initialCapacity: -1)); + public string Method { get; } + public string Path { get; } + public string Body { get; } } - #endregion - - #region Add Tests - - [Fact(Timeout = 60000)] - public async Task Add_FirstDocument_SetsDimension() + private sealed class MockHandler : HttpMessageHandler { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var doc = CreateTestDocument("doc1", "Test content", 384); - - // Act - store.Add(doc); + private readonly List<(string Method, string PathContains, Func Resp)> _routes = new(); + public List Requests { get; } = new(); + public string DefaultBody { get; set; } = "{}"; - // Assert - Assert.Equal(1, store.DocumentCount); - Assert.Equal(384, store.VectorDimension); - } - - [Fact(Timeout = 60000)] - public async Task Add_WithValidDocument_IncreasesCount() - { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var doc = CreateTestDocument("doc1", "Test content", 3); + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, _ => resp())); + return this; + } - // Act - store.Add(doc); + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, resp)); + return this; + } - // Assert - Assert.Equal(1, store.DocumentCount); + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + var body = request.Content != null + ? await request.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + var path = request.RequestUri!.AbsolutePath; + Requests.Add(new RecordedRequest(request.Method.Method, request.RequestUri.PathAndQuery, body)); + + foreach (var route in _routes) + { + if (route.Method == request.Method.Method && path.Contains(route.PathContains)) + return route.Resp(body); + } + + return Json(DefaultBody); + } } - [Fact(Timeout = 60000)] - public async Task Add_WithNullDocument_ThrowsArgumentNullException() - { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - - // Act & Assert - Assert.Throws(() => store.Add(null!)); - } + private static HttpResponseMessage Json(string body, HttpStatusCode code = HttpStatusCode.OK) + => new HttpResponseMessage(code) { Content = new StringContent(body, Encoding.UTF8, "application/json") }; - [Fact(Timeout = 60000)] - public async Task Add_WithMismatchedDimension_ThrowsArgumentException() + private static PineconeDocumentStore StoreWith(MockHandler handler, string? ns = null) { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var doc1 = CreateTestDocument("doc1", "Test 1", 3); - store.Add(doc1); - - var doc2 = CreateTestDocument("doc2", "Test 2", 5); // Wrong dimension - - // Act & Assert - var exception = Assert.Throws(() => store.Add(doc2)); - Assert.Contains("dimension mismatch", exception.Message.ToLower()); + var client = new HttpClient(handler) { BaseAddress = new Uri(BaseUrl) }; + return new PineconeDocumentStore(Index, BaseUrl, apiKey: "k", @namespace: ns, httpClient: client); } - [Fact(Timeout = 60000)] - public async Task Add_WithDuplicateId_UpdatesDocument() + private static VectorDocument Doc(string id, string content, float[] vector, Dictionary? metadata = null) { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var doc1 = CreateTestDocument("doc1", "Original content", 3); - var doc2 = CreateTestDocument("doc1", "Updated content", 3); - - // Act - store.Add(doc1); - store.Add(doc2); - - // Assert - Assert.Equal(1, store.DocumentCount); - var retrieved = store.GetById("doc1"); - Assert.NotNull(retrieved); - Assert.Equal("Updated content", retrieved.Content); + var d = new Document(id, content, metadata ?? new Dictionary()); + return new VectorDocument { Document = d, Embedding = new Vector(vector) }; } #endregion - #region AddBatch Tests - - [Fact(Timeout = 60000)] - public async Task AddBatch_WithValidDocuments_IncreasesCount() + [Fact] + public void Constructor_ReadsStats_SetsDimensionAndCount() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var docs = new List> - { - CreateTestDocument("doc1", "Content 1", 3), - CreateTestDocument("doc2", "Content 2", 3), - CreateTestDocument("doc3", "Content 3", 3) - }; + var handler = new MockHandler() + .On("POST", "/describe_index_stats", () => Json( + "{\"namespaces\":{\"\":{\"vectorCount\":7}},\"dimension\":8,\"totalVectorCount\":7}")); - // Act - store.AddBatch(docs); + var store = StoreWith(handler); - // Assert - Assert.Equal(3, store.DocumentCount); + Assert.Equal(8, store.VectorDimension); + Assert.Equal(7, store.DocumentCount); + Assert.Equal(Index, store.IndexName); } - [Fact(Timeout = 60000)] - public async Task AddBatch_FirstBatch_SetsDimension() + [Fact] + public void Add_Upserts_SendsCorrectRequest() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var docs = new List> - { - CreateTestDocument("doc1", "Content 1", 384), - CreateTestDocument("doc2", "Content 2", 384) - }; + var handler = new MockHandler() + .On("POST", "/vectors/upsert", () => Json("{\"upsertedCount\":1}")) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - // Act - store.AddBatch(docs); - - // Assert - Assert.Equal(2, store.DocumentCount); - Assert.Equal(384, store.VectorDimension); - } + var store = StoreWith(handler); + store.Add(Doc("doc1", "Hello world", new float[] { 1, 0, 0 }, + new Dictionary { ["category"] = "science" })); - [Fact(Timeout = 60000)] - public async Task AddBatch_WithNullCollection_ThrowsArgumentNullException() - { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - - // Act & Assert - Assert.Throws(() => store.AddBatch(null!)); + Assert.Equal(1, store.DocumentCount); + var upsert = handler.Requests.Single(r => r.Path.Contains("/vectors/upsert")); + var vector = JObject.Parse(upsert.Body)["vectors"]![0]!; + Assert.Equal("doc1", (string?)vector["id"]); + Assert.Equal(3, ((JArray)vector["values"]!).Count); + Assert.Equal("Hello world", (string?)vector["metadata"]!["_content"]); + Assert.Equal("science", (string?)vector["metadata"]!["category"]); } - [Fact(Timeout = 60000)] - public async Task AddBatch_WithEmptyCollection_ThrowsArgumentException() + [Fact] + public void AddBatch_Upserts_AllVectors() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - - // Act & Assert - Assert.Throws(() => store.AddBatch(new List>())); - } + var handler = new MockHandler() + .On("POST", "/vectors/upsert", () => Json("{\"upsertedCount\":3}")) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - [Fact(Timeout = 60000)] - public async Task AddBatch_WithMismatchedDimensionsInBatch_ThrowsArgumentException() - { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var docs = new List> + var store = StoreWith(handler); + store.AddBatch(new List> { - CreateTestDocument("doc1", "Content 1", 3), - CreateTestDocument("doc2", "Content 2", 5) // Different dimension - }; - - // Act & Assert - Assert.Throws(() => store.AddBatch(docs)); - } - - [Fact(Timeout = 60000)] - public async Task AddBatch_WithLargeNumberOfDocuments_Succeeds() - { - // Arrange - var store = new PineconeDocumentStore("TestIndex", initialCapacity: 10000); - var docs = Enumerable.Range(0, 1000) - .Select(i => CreateTestDocument($"doc{i}", $"Content {i}", 384)) - .ToList(); - - // Act - store.AddBatch(docs); + Doc("d1", "c1", new float[] { 1, 0, 0 }), + Doc("d2", "c2", new float[] { 0, 1, 0 }), + Doc("d3", "c3", new float[] { 0, 0, 1 }) + }); - // Assert - Assert.Equal(1000, store.DocumentCount); + Assert.Equal(3, store.DocumentCount); + var upsert = handler.Requests.Single(r => r.Path.Contains("/vectors/upsert")); + Assert.Equal(3, ((JArray)JObject.Parse(upsert.Body)["vectors"]!).Count); } - #endregion - - #region GetSimilar Tests - - [Fact(Timeout = 60000)] - public async Task GetSimilar_WithMatchingDocuments_ReturnsTopK() + [Fact] + public void GetSimilar_ParsesMatches_InOrder() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - store.AddBatch(new List> - { - CreateTestDocument("doc1", "Content 1", 3, new float[] { 1, 0, 0 }), - CreateTestDocument("doc2", "Content 2", 3, new float[] { 0, 1, 0 }), - CreateTestDocument("doc3", "Content 3", 3, new float[] { 0, 0, 1 }) - }); + var queryResponse = + "{\"matches\":[" + + "{\"id\":\"doc1\",\"score\":0.97,\"metadata\":{\"_content\":\"first\",\"category\":\"science\"}}," + + "{\"id\":\"doc2\",\"score\":0.42,\"metadata\":{\"_content\":\"second\"}}" + + "],\"namespace\":\"\"}"; - var queryVector = new Vector(new float[] { 1, 0, 0 }); + var handler = new MockHandler() + .On("POST", "/query", () => Json(queryResponse)) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - // Act - var results = store.GetSimilar(queryVector, topK: 2).ToList(); + var store = StoreWith(handler); + var results = store.GetSimilar(new Vector(new float[] { 1, 0, 0 }), topK: 2).ToList(); - // Assert Assert.Equal(2, results.Count); - Assert.True(results[0].HasRelevanceScore); Assert.Equal("doc1", results[0].Id); - } - - [Fact(Timeout = 60000)] - public async Task GetSimilar_WithEmptyStore_ReturnsEmpty() - { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var queryVector = new Vector(new float[] { 1, 0, 0 }); - - // Act - var results = store.GetSimilar(queryVector, topK: 5); + Assert.Equal("first", results[0].Content); + Assert.Equal("science", results[0].Metadata["category"].ToString()); + Assert.True(results[0].HasRelevanceScore); + Assert.True(Convert.ToDouble(results[0].RelevanceScore) > Convert.ToDouble(results[1].RelevanceScore)); - // Assert - Assert.Empty(results); + var query = handler.Requests.Single(r => r.Path.Contains("/query")); + var body = JObject.Parse(query.Body); + Assert.Equal(2, (int)body["topK"]!); + Assert.True((bool)body["includeMetadata"]!); + Assert.Equal(3, ((JArray)body["vector"]!).Count); } - [Fact(Timeout = 60000)] - public async Task GetSimilar_WithNullQueryVector_ThrowsArgumentNullException() + [Fact] + public void GetSimilarWithFilters_TranslatesEqInGte() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); + var handler = new MockHandler() + .On("POST", "/query", () => Json("{\"matches\":[]}")) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - // Act & Assert - Assert.Throws(() => store.GetSimilar(null!, topK: 5)); - } - - [Fact(Timeout = 60000)] - public async Task GetSimilar_OrdersByRelevanceDescending() - { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - store.AddBatch(new List> + var store = StoreWith(handler); + var filters = new Dictionary { - CreateTestDocument("doc1", "Content 1", 3, new float[] { 1, 0, 0 }), - CreateTestDocument("doc2", "Content 2", 3, new float[] { 0.9f, 0.1f, 0 }), - CreateTestDocument("doc3", "Content 3", 3, new float[] { 0, 1, 0 }) - }); - - var queryVector = new Vector(new float[] { 1, 0, 0 }); + ["category"] = "science", + ["year"] = 2020, + ["tags"] = new[] { "a", "b" } + }; + store.GetSimilarWithFilters(new Vector(new float[] { 1, 0, 0 }), 10, filters).ToList(); - // Act - var results = store.GetSimilar(queryVector, topK: 3).ToList(); + var query = handler.Requests.Single(r => r.Path.Contains("/query")); + var filter = JObject.Parse(query.Body)["filter"]!; - // Assert - Assert.Equal(3, results.Count); - // Verify descending order - for (int i = 0; i < results.Count - 1; i++) - { - Assert.True(Convert.ToDouble(results[i].RelevanceScore) >= - Convert.ToDouble(results[i + 1].RelevanceScore)); - } + Assert.Equal("science", (string?)filter["category"]!["$eq"]); + Assert.Equal(2020.0, (double)filter["year"]!["$gte"]!); + var inArray = (JArray)filter["tags"]!["$in"]!; + Assert.Equal(new[] { "a", "b" }, inArray.Select(t => (string?)t).ToArray()); } - #endregion - - #region GetSimilarWithFilters Tests - - [Fact(Timeout = 60000)] - public async Task GetSimilarWithFilters_WithMatchingMetadata_ReturnsFilteredResults() + [Fact] + public void GetById_Found_ReturnsDocument() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var doc1 = CreateTestDocument("doc1", "Content 1", 3, new float[] { 1, 0, 0 }); - doc1.Document.Metadata["category"] = "science"; - - var doc2 = CreateTestDocument("doc2", "Content 2", 3, new float[] { 1, 0, 0 }); - doc2.Document.Metadata["category"] = "history"; + var handler = new MockHandler() + .On("GET", "/vectors/fetch", () => Json( + "{\"vectors\":{\"doc1\":{\"id\":\"doc1\",\"metadata\":{\"_content\":\"hi\",\"k\":\"v\"}}}}")) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - store.Add(doc1); - store.Add(doc2); + var store = StoreWith(handler); + var doc = store.GetById("doc1"); - var queryVector = new Vector(new float[] { 1, 0, 0 }); - var filters = new Dictionary { { "category", "science" } }; - - // Act - var results = store.GetSimilarWithFilters(queryVector, topK: 10, filters).ToList(); - - // Assert - Assert.Single(results); - Assert.Equal("doc1", results[0].Id); + Assert.NotNull(doc); + Assert.Equal("doc1", doc!.Id); + Assert.Equal("hi", doc.Content); + Assert.Equal("v", doc.Metadata["k"].ToString()); } - [Fact(Timeout = 60000)] - public async Task GetSimilarWithFilters_WithNoMatchingMetadata_ReturnsEmpty() + [Fact] + public void GetById_Missing_ReturnsNull() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var doc = CreateTestDocument("doc1", "Content 1", 3); - doc.Document.Metadata["category"] = "science"; - store.Add(doc); - - var queryVector = new Vector(new float[] { 1, 0, 0 }); - var filters = new Dictionary { { "category", "history" } }; - - // Act - var results = store.GetSimilarWithFilters(queryVector, topK: 10, filters); - - // Assert - Assert.Empty(results); - } - - #endregion + var handler = new MockHandler() + .On("GET", "/vectors/fetch", () => Json("{\"vectors\":{}}")) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - #region GetById Tests - - [Fact(Timeout = 60000)] - public async Task GetById_WithExistingDocument_ReturnsDocument() - { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var doc = CreateTestDocument("doc1", "Test content", 3); - store.Add(doc); - - // Act - var result = store.GetById("doc1"); - - // Assert - Assert.NotNull(result); - Assert.Equal("doc1", result.Id); - Assert.Equal("Test content", result.Content); + var store = StoreWith(handler); + Assert.Null(store.GetById("missing")); } - [Fact(Timeout = 60000)] - public async Task GetById_WithNonExistingDocument_ReturnsNull() + [Fact] + public void Remove_SendsDeleteIds_ReturnsTrue() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); + var handler = new MockHandler() + .On("POST", "/vectors/delete", () => Json("{}")) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - // Act - var result = store.GetById("nonexistent"); + var store = StoreWith(handler); + var removed = store.Remove("doc1"); - // Assert - Assert.Null(result); + Assert.True(removed); + var delete = handler.Requests.Single(r => r.Path.Contains("/vectors/delete")); + var ids = (JArray)JObject.Parse(delete.Body)["ids"]!; + Assert.Equal("doc1", (string?)ids[0]); } - #endregion - - #region Remove Tests - - [Fact(Timeout = 60000)] - public async Task Remove_WithExistingDocument_ReturnsTrue() + [Fact] + public void Clear_SendsDeleteAll() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - var doc = CreateTestDocument("doc1", "Test", 3); - store.Add(doc); + var handler = new MockHandler() + .On("POST", "/vectors/delete", () => Json("{}")) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - // Act - var result = store.Remove("doc1"); + var store = StoreWith(handler); + store.Clear(); - // Assert - Assert.True(result); Assert.Equal(0, store.DocumentCount); + var delete = handler.Requests.Single(r => r.Path.Contains("/vectors/delete")); + Assert.True((bool)JObject.Parse(delete.Body)["deleteAll"]!); } - [Fact(Timeout = 60000)] - public async Task Remove_LastDocument_ResetsDimension() + [Fact] + public void Namespace_IsIncludedInRequests() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - store.Add(CreateTestDocument("doc1", "Test", 3)); - - // Act - store.Remove("doc1"); + var handler = new MockHandler() + .On("POST", "/vectors/upsert", () => Json("{\"upsertedCount\":1}")) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - // Assert - Assert.Equal(0, store.DocumentCount); - Assert.Equal(0, store.VectorDimension); + var store = StoreWith(handler, ns: "prod"); + store.Add(Doc("doc1", "c", new float[] { 1, 0, 0 })); - // Should be able to add documents with different dimension - store.Add(CreateTestDocument("doc2", "Test", 5)); - Assert.Equal(5, store.VectorDimension); + var upsert = handler.Requests.Single(r => r.Path.Contains("/vectors/upsert")); + Assert.Equal("prod", (string?)JObject.Parse(upsert.Body)["namespace"]); + Assert.Equal("prod", store.Namespace); } - #endregion - - #region Clear Tests - - [Fact(Timeout = 60000)] - public async Task Clear_RemovesAllDocuments() + [Fact] + public void GetAll_ListsThenFetches() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - store.AddBatch(new List> - { - CreateTestDocument("doc1", "Content 1", 3), - CreateTestDocument("doc2", "Content 2", 3), - CreateTestDocument("doc3", "Content 3", 3) - }); + var handler = new MockHandler() + .On("GET", "/vectors/list", () => Json( + "{\"vectors\":[{\"id\":\"d1\"},{\"id\":\"d2\"}],\"pagination\":{}}")) + .On("GET", "/vectors/fetch", body => Json( + "{\"vectors\":{\"d1\":{\"id\":\"d1\",\"metadata\":{\"_content\":\"c1\"}}," + + "\"d2\":{\"id\":\"d2\",\"metadata\":{\"_content\":\"c2\"}}}}")) + .On("POST", "/describe_index_stats", () => Json(StatsJson)); - // Act - store.Clear(); + var store = StoreWith(handler); + var all = store.GetAll().ToList(); - // Assert - Assert.Equal(0, store.DocumentCount); - Assert.Equal(0, store.VectorDimension); + Assert.Equal(2, all.Count); + Assert.Contains(all, d => d.Id == "d1"); + Assert.Contains(all, d => d.Id == "d2"); } - [Fact(Timeout = 60000)] - public async Task Clear_AllowsNewDimensionAfterClear() + [Fact] + public void Constructor_WithEmptyIndexName_Throws() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - store.Add(CreateTestDocument("doc1", "Test", 3)); - store.Clear(); - - // Act - store.Add(CreateTestDocument("doc2", "Test", 5)); - - // Assert - Assert.Equal(1, store.DocumentCount); - Assert.Equal(5, store.VectorDimension); + var client = new HttpClient(new MockHandler()) { BaseAddress = new Uri(BaseUrl) }; + Assert.Throws(() => + new PineconeDocumentStore("", BaseUrl, "k", httpClient: client)); } - #endregion - - #region GetAll Tests - - [Fact(Timeout = 60000)] - public async Task GetAll_WithDocuments_ReturnsAllDocuments() - { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); - store.AddBatch(new List> - { - CreateTestDocument("doc1", "Content 1", 3), - CreateTestDocument("doc2", "Content 2", 3), - CreateTestDocument("doc3", "Content 3", 3) - }); - - // Act - var results = store.GetAll().ToList(); + #region Integration (gated - require a live Pinecone index) - // Assert - Assert.Equal(3, results.Count); - Assert.Contains(results, d => d.Id == "doc1"); - Assert.Contains(results, d => d.Id == "doc2"); - Assert.Contains(results, d => d.Id == "doc3"); - } - - [Fact(Timeout = 60000)] - public async Task GetAll_WithEmptyStore_ReturnsEmpty() + [Trait("Category", "Integration")] + [SkippableFact] + public void Integration_UpsertQueryDelete_RoundTrips() { - // Arrange - var store = new PineconeDocumentStore("TestIndex"); + var url = Environment.GetEnvironmentVariable("PINECONE_INDEX_URL"); + var apiKey = Environment.GetEnvironmentVariable("PINECONE_API_KEY"); + Skip.If(string.IsNullOrWhiteSpace(url) || string.IsNullOrWhiteSpace(apiKey), + "Set PINECONE_INDEX_URL and PINECONE_API_KEY to run Pinecone integration tests."); - // Act - var results = store.GetAll(); + var store = new PineconeDocumentStore("integration", url!, apiKey!, + @namespace: "aidotnet-tests"); - // Assert - Assert.Empty(results); - } + var dim = store.VectorDimension > 0 ? store.VectorDimension : 8; + var vec = Enumerable.Range(0, dim).Select(i => (float)(i == 0 ? 1.0 : 0.0)).ToArray(); + var id = "it-" + Guid.NewGuid().ToString("N"); - #endregion + store.Add(Doc(id, "integration content", vec, + new Dictionary { ["category"] = "test" })); - #region Helper Methods + var results = store.GetSimilar(new Vector(vec), topK: 5).ToList(); + Assert.Contains(results, r => r.Id == id); - private VectorDocument CreateTestDocument( - string id, - string content, - int dimension, - float[]? values = null) - { - var vector = values ?? Enumerable.Range(0, dimension).Select(i => (float)i).ToArray(); - return new VectorDocument - { - Document = new Document(id, content), - Embedding = new Vector(vector) - }; + Assert.True(store.Remove(id)); } #endregion diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStoreTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStoreTests.cs new file mode 100644 index 0000000000..6adb6a5cbf --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStoreTests.cs @@ -0,0 +1,353 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Threading; +using System.Threading.Tasks; +using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json.Linq; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.DocumentStores +{ + /// + /// Unit tests for the real Qdrant REST client. All HTTP traffic is intercepted by a mock + /// so these run in CI with no network access. + /// + public class QdrantDocumentStoreTests + { + private const string BaseUrl = "http://localhost:6333"; + private const string Collection = "test"; + + // A collection that already exists with 3-d Cosine vectors and 0 points. + private const string ExistingCollectionJson = + "{\"result\":{\"status\":\"green\",\"points_count\":0," + + "\"config\":{\"params\":{\"vectors\":{\"size\":3,\"distance\":\"Cosine\"}}}},\"status\":\"ok\"}"; + + #region infrastructure + + private sealed class RecordedRequest + { + public RecordedRequest(string method, string path, string body) + { + Method = method; + Path = path; + Body = body; + } + + public string Method { get; } + public string Path { get; } + public string Body { get; } + } + + private sealed class MockHandler : HttpMessageHandler + { + private readonly List<(string Method, string PathContains, Func Resp)> _routes = new(); + public List Requests { get; } = new(); + public string DefaultBody { get; set; } = "{\"result\":true,\"status\":\"ok\"}"; + + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, _ => resp())); + return this; + } + + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, resp)); + return this; + } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + var body = request.Content != null + ? await request.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + var path = request.RequestUri!.AbsolutePath; + Requests.Add(new RecordedRequest(request.Method.Method, request.RequestUri.PathAndQuery, body)); + + foreach (var route in _routes) + { + if (route.Method == request.Method.Method && path.Contains(route.PathContains)) + return route.Resp(body); + } + + return Json(DefaultBody); + } + } + + private static HttpResponseMessage Json(string body, HttpStatusCode code = HttpStatusCode.OK) + => new HttpResponseMessage(code) { Content = new StringContent(body, Encoding.UTF8, "application/json") }; + + private static QdrantDocumentStore StoreWith(MockHandler handler, + DistanceMetricType metric = DistanceMetricType.Cosine, int vectorDimension = 0) + { + var client = new HttpClient(handler) { BaseAddress = new Uri(BaseUrl) }; + return new QdrantDocumentStore(Collection, BaseUrl, apiKey: "k", + distanceMetric: metric, vectorDimension: vectorDimension, httpClient: client); + } + + private static VectorDocument Doc(string id, string content, float[] vector, Dictionary? metadata = null) + { + var d = new Document(id, content, metadata ?? new Dictionary()); + return new VectorDocument { Document = d, Embedding = new Vector(vector) }; + } + + #endregion + + [Fact] + public void Constructor_ReadsExistingCollection_SetsDimensionAndCount() + { + var handler = new MockHandler() + .On("GET", "/collections/test", () => Json( + "{\"result\":{\"points_count\":5,\"config\":{\"params\":{\"vectors\":{\"size\":8,\"distance\":\"Cosine\"}}}}}")); + + var store = StoreWith(handler); + + Assert.Equal(8, store.VectorDimension); + Assert.Equal(5, store.DocumentCount); + Assert.Equal(Collection, store.CollectionName); + } + + [Fact] + public void Constructor_CollectionMissing_CreatesCollectionWithMappedDistance() + { + var handler = new MockHandler() + .On("PUT", "/collections/test", () => Json("{\"result\":true,\"status\":\"ok\"}")) + .On("GET", "/collections/test", () => Json("{}", HttpStatusCode.NotFound)); + + var store = StoreWith(handler, metric: DistanceMetricType.Euclidean, vectorDimension: 4); + + Assert.Equal(4, store.VectorDimension); + var create = handler.Requests.Single(r => r.Method == "PUT" && r.Path == "/collections/test"); + var body = JObject.Parse(create.Body); + Assert.Equal(4, (int)body["vectors"]!["size"]!); + Assert.Equal("Euclid", (string?)body["vectors"]!["distance"]); + } + + [Fact] + public void Add_UpsertsPoint_SendsCorrectRequest() + { + var handler = new MockHandler() + .On("GET", "/collections/test", () => Json(ExistingCollectionJson)); + + var store = StoreWith(handler); + store.Add(Doc("doc1", "Hello world", new float[] { 1, 0, 0 }, + new Dictionary { ["category"] = "science" })); + + Assert.Equal(1, store.DocumentCount); + var upsert = handler.Requests.Single(r => r.Method == "PUT" && r.Path.StartsWith("/collections/test/points")); + var body = JObject.Parse(upsert.Body); + var point = body["points"]![0]!; + Assert.Equal("doc1", (string?)point["payload"]!["_doc_id"]); + Assert.Equal("Hello world", (string?)point["payload"]!["_content"]); + Assert.Equal("science", (string?)point["payload"]!["metadata"]!["category"]); + Assert.Equal(3, ((JArray)point["vector"]!).Count); + } + + [Fact] + public void AddBatch_SendsAllPointsInOneRequest() + { + var handler = new MockHandler() + .On("GET", "/collections/test", () => Json(ExistingCollectionJson)); + + var store = StoreWith(handler); + store.AddBatch(new List> + { + Doc("d1", "c1", new float[] { 1, 0, 0 }), + Doc("d2", "c2", new float[] { 0, 1, 0 }), + Doc("d3", "c3", new float[] { 0, 0, 1 }) + }); + + Assert.Equal(3, store.DocumentCount); + var upsert = handler.Requests.Single(r => r.Method == "PUT" && r.Path.StartsWith("/collections/test/points")); + var points = (JArray)JObject.Parse(upsert.Body)["points"]!; + Assert.Equal(3, points.Count); + } + + [Fact] + public void GetSimilar_ParsesRankedHits_InOrder() + { + var searchResponse = + "{\"result\":[" + + "{\"id\":\"g1\",\"score\":0.98,\"payload\":{\"_doc_id\":\"doc1\",\"_content\":\"first\",\"metadata\":{\"category\":\"science\"}}}," + + "{\"id\":\"g2\",\"score\":0.55,\"payload\":{\"_doc_id\":\"doc2\",\"_content\":\"second\",\"metadata\":{}}}" + + "],\"status\":\"ok\"}"; + + var handler = new MockHandler() + .On("POST", "/points/search", () => Json(searchResponse)) + .On("GET", "/collections/test", () => Json(ExistingCollectionJson)); + + var store = StoreWith(handler); + var results = store.GetSimilar(new Vector(new float[] { 1, 0, 0 }), topK: 2).ToList(); + + Assert.Equal(2, results.Count); + Assert.Equal("doc1", results[0].Id); + Assert.Equal("first", results[0].Content); + Assert.True(results[0].HasRelevanceScore); + Assert.True(Convert.ToDouble(results[0].RelevanceScore) > Convert.ToDouble(results[1].RelevanceScore)); + + var search = handler.Requests.Single(r => r.Path.Contains("/points/search")); + var body = JObject.Parse(search.Body); + Assert.Equal(2, (int)body["limit"]!); + Assert.True((bool)body["with_payload"]!); + Assert.Equal(3, ((JArray)body["vector"]!).Count); + } + + [Fact] + public void GetSimilarWithFilters_TranslatesEqualityAndRange() + { + var handler = new MockHandler() + .On("POST", "/points/search", () => Json("{\"result\":[],\"status\":\"ok\"}")) + .On("GET", "/collections/test", () => Json(ExistingCollectionJson)); + + var store = StoreWith(handler); + var filters = new Dictionary { ["category"] = "science", ["year"] = 2020 }; + store.GetSimilarWithFilters(new Vector(new float[] { 1, 0, 0 }), 10, filters).ToList(); + + var search = handler.Requests.Single(r => r.Path.Contains("/points/search")); + var must = (JArray)JObject.Parse(search.Body)["filter"]!["must"]!; + + var categoryCond = must.Single(c => (string?)c["key"] == "metadata.category"); + Assert.Equal("science", (string?)categoryCond["match"]!["value"]); + + var yearCond = must.Single(c => (string?)c["key"] == "metadata.year"); + Assert.Equal(2020.0, (double)yearCond["range"]!["gte"]!); + } + + [Fact] + public void GetById_Found_ReturnsDocument() + { + var handler = new MockHandler() + .On("GET", "/collections/test/points/", () => Json( + "{\"result\":{\"id\":\"g1\",\"payload\":{\"_doc_id\":\"doc1\",\"_content\":\"hi\",\"metadata\":{\"k\":\"v\"}}},\"status\":\"ok\"}")) + .On("GET", "/collections/test", () => Json(ExistingCollectionJson)); + + var store = StoreWith(handler); + var doc = store.GetById("doc1"); + + Assert.NotNull(doc); + Assert.Equal("doc1", doc!.Id); + Assert.Equal("hi", doc.Content); + Assert.Equal("v", doc.Metadata["k"].ToString()); + } + + [Fact] + public void GetById_NotFound_ReturnsNull() + { + var handler = new MockHandler() + .On("GET", "/collections/test/points/", () => Json("{}", HttpStatusCode.NotFound)) + .On("GET", "/collections/test", () => Json(ExistingCollectionJson)); + + var store = StoreWith(handler); + Assert.Null(store.GetById("missing")); + } + + [Fact] + public void Remove_SendsDelete_ReturnsTrue() + { + var handler = new MockHandler() + .On("POST", "/points/delete", () => Json("{\"result\":{\"status\":\"completed\"},\"status\":\"ok\"}")) + .On("GET", "/collections/test", () => Json(ExistingCollectionJson)); + + var store = StoreWith(handler); + var removed = store.Remove("doc1"); + + Assert.True(removed); + var delete = handler.Requests.Single(r => r.Path.Contains("/points/delete")); + var points = (JArray)JObject.Parse(delete.Body)["points"]!; + Assert.Single(points); + } + + [Fact] + public void GetAll_ScrollsAllPages() + { + var page = 0; + var handler = new MockHandler() + .On("POST", "/points/scroll", _ => + { + page++; + return page == 1 + ? Json("{\"result\":{\"points\":[" + + "{\"id\":\"g1\",\"payload\":{\"_doc_id\":\"d1\",\"_content\":\"c1\",\"metadata\":{}}}," + + "{\"id\":\"g2\",\"payload\":{\"_doc_id\":\"d2\",\"_content\":\"c2\",\"metadata\":{}}}]," + + "\"next_page_offset\":\"g2\"},\"status\":\"ok\"}") + : Json("{\"result\":{\"points\":[" + + "{\"id\":\"g3\",\"payload\":{\"_doc_id\":\"d3\",\"_content\":\"c3\",\"metadata\":{}}}]," + + "\"next_page_offset\":null},\"status\":\"ok\"}"); + }) + .On("GET", "/collections/test", () => Json(ExistingCollectionJson)); + + var store = StoreWith(handler); + var all = store.GetAll().ToList(); + + Assert.Equal(3, all.Count); + Assert.Contains(all, d => d.Id == "d3"); + Assert.Equal(2, handler.Requests.Count(r => r.Path.Contains("/points/scroll"))); + } + + [Fact] + public void Clear_DeletesCollection_AndRecreates() + { + var handler = new MockHandler() + .On("DELETE", "/collections/test", () => Json("{\"result\":true,\"status\":\"ok\"}")) + .On("PUT", "/collections/test", () => Json("{\"result\":true,\"status\":\"ok\"}")) + .On("GET", "/collections/test", () => Json(ExistingCollectionJson)); + + var store = StoreWith(handler); + store.Clear(); + + Assert.Equal(0, store.DocumentCount); + Assert.Contains(handler.Requests, r => r.Method == "DELETE" && r.Path == "/collections/test"); + // Dimension known (3) so the collection is recreated. + Assert.Contains(handler.Requests, r => r.Method == "PUT" && r.Path == "/collections/test"); + } + + [Fact] + public void Constructor_UnsupportedMetric_Throws() + { + Assert.Throws(() => + new QdrantDocumentStore(Collection, BaseUrl, distanceMetric: DistanceMetricType.Jaccard, + httpClient: new HttpClient(new MockHandler()) { BaseAddress = new Uri(BaseUrl) })); + } + + #region Integration (gated - require a live Qdrant instance) + + [Trait("Category", "Integration")] + [SkippableFact] + public void Integration_UpsertSearchDelete_RoundTrips() + { + var url = Environment.GetEnvironmentVariable("QDRANT_URL"); + Skip.If(string.IsNullOrWhiteSpace(url), + "Set QDRANT_URL (and optionally QDRANT_API_KEY) to run Qdrant integration tests."); + + var apiKey = Environment.GetEnvironmentVariable("QDRANT_API_KEY"); + var collection = "aidotnet_it_" + Guid.NewGuid().ToString("N"); + var store = new QdrantDocumentStore(collection, url!, apiKey, + distanceMetric: DistanceMetricType.Cosine, vectorDimension: 4); + + try + { + var id = "it-" + Guid.NewGuid().ToString("N"); + store.Add(Doc(id, "integration content", new float[] { 1, 0, 0, 0 }, + new Dictionary { ["category"] = "test" })); + + var results = store.GetSimilar(new Vector(new float[] { 1, 0, 0, 0 }), topK: 5).ToList(); + Assert.Contains(results, r => r.Id == id); + + Assert.True(store.Remove(id)); + } + finally + { + store.Clear(); + } + } + + #endregion + } +} From 2c07af9b340105645b4993dc6a34a73b3fcc4b52 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 15:32:35 -0400 Subject: [PATCH 04/25] feat(rag): rRF hybrid fusion + ensemble, routing, and self-query retrievers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit HybridRetriever summed raw dense (cosine ~0..1) and sparse (BM25, unbounded) scores, so the larger-scaled list dominated. Switch it to Reciprocal Rank Fusion (rank-based, scale-free — the default in ES/Weaviate/Qdrant and LangChain/LlamaIndex), keeping the weights and adding a configurable RRF k (default 60). Add three retrievers that were absent versus the common frameworks: - EnsembleRetriever: RRF fusion over any number of weighted retrievers. - QueryRoutingRetriever: routes a query to the best named data source via an LLM (ITextGenerator) with a token-overlap fallback. - SelfQueryRetriever: uses the generator to split an NL query into a search string plus structured metadata filters over declared fields, then delegates (caller filters win). Malformed LLM JSON degrades to the raw query. Verified with fakes (6 tests): RRF ordering, ensemble, routing (LLM + fallback), and self-query filter extraction/precedence. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../Retrievers/EnsembleRetriever.cs | 82 +++++++++++ .../Retrievers/HybridRetriever.cs | 74 +++++----- .../Retrievers/QueryRoutingRetriever.cs | 102 ++++++++++++++ .../Retrievers/SelfQueryRetriever.cs | 123 ++++++++++++++++ .../RetrieverFusionAndRoutingTests.cs | 131 ++++++++++++++++++ 5 files changed, 481 insertions(+), 31 deletions(-) create mode 100644 src/RetrievalAugmentedGeneration/Retrievers/EnsembleRetriever.cs create mode 100644 src/RetrievalAugmentedGeneration/Retrievers/QueryRoutingRetriever.cs create mode 100644 src/RetrievalAugmentedGeneration/Retrievers/SelfQueryRetriever.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/RetrieverFusionAndRoutingTests.cs diff --git a/src/RetrievalAugmentedGeneration/Retrievers/EnsembleRetriever.cs b/src/RetrievalAugmentedGeneration/Retrievers/EnsembleRetriever.cs new file mode 100644 index 0000000000..bac8dda4a9 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Retrievers/EnsembleRetriever.cs @@ -0,0 +1,82 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.Models; + +namespace AiDotNet.RetrievalAugmentedGeneration.Retrievers +{ + /// + /// Fuses the results of any number of retrievers via Reciprocal Rank Fusion (RRF), with optional + /// per-retriever weights. Generalizes beyond two retrievers, matching + /// the LangChain/LlamaIndex "ensemble retriever" (Cormack et al., 2009). + /// + /// The numeric type for vector operations. + [ComponentType(ComponentType.Retriever)] + [PipelineStage(PipelineStage.Retrieval)] + public class EnsembleRetriever : RetrieverBase + { + private readonly IReadOnlyList> _retrievers; + private readonly IReadOnlyList _weights; + private readonly int _rrfK; + + /// The retrievers to fuse (at least one). + /// Optional per-retriever weights (defaults to 1.0 each); length must match. + /// Default number of documents to return. + /// RRF constant (default 60); larger damps low-rank influence. + public EnsembleRetriever( + IReadOnlyList> retrievers, + IReadOnlyList? weights = null, + int defaultTopK = 5, + int rrfK = 60) + : base(defaultTopK) + { + if (retrievers == null) throw new ArgumentNullException(nameof(retrievers)); + if (retrievers.Count == 0) throw new ArgumentException("At least one retriever is required.", nameof(retrievers)); + if (retrievers.Any(r => r == null)) throw new ArgumentException("Retrievers must not be null.", nameof(retrievers)); + if (weights != null && weights.Count != retrievers.Count) + throw new ArgumentException("weights length must match retrievers length.", nameof(weights)); + if (rrfK <= 0) throw new ArgumentOutOfRangeException(nameof(rrfK), "RRF k must be positive."); + + _retrievers = retrievers; + _weights = weights ?? Enumerable.Repeat(1.0, retrievers.Count).ToList(); + _rrfK = rrfK; + } + + protected override IEnumerable> RetrieveCore(string query, int topK, Dictionary metadataFilters) + { + var rrfScores = new Dictionary(); + var docById = new Dictionary>(); + + for (int r = 0; r < _retrievers.Count; r++) + { + var results = _retrievers[r].Retrieve(query, topK * 2, metadataFilters).ToList(); + double weight = _weights[r]; + for (int rank = 0; rank < results.Count; rank++) + { + var doc = results[rank]; + double contribution = weight / (_rrfK + rank + 1); + rrfScores[doc.Id] = rrfScores.TryGetValue(doc.Id, out var s) ? s + contribution : contribution; + if (!docById.ContainsKey(doc.Id)) + { + docById[doc.Id] = doc; + } + } + } + + return rrfScores + .OrderByDescending(kv => kv.Value) + .Take(topK) + .Select(kv => + { + var doc = docById[kv.Key]; + doc.RelevanceScore = NumOps.FromDouble(kv.Value); + doc.HasRelevanceScore = true; + return doc; + }) + .ToList(); + } + } +} diff --git a/src/RetrievalAugmentedGeneration/Retrievers/HybridRetriever.cs b/src/RetrievalAugmentedGeneration/Retrievers/HybridRetriever.cs index a2ca5b476f..628b79c524 100644 --- a/src/RetrievalAugmentedGeneration/Retrievers/HybridRetriever.cs +++ b/src/RetrievalAugmentedGeneration/Retrievers/HybridRetriever.cs @@ -10,76 +10,88 @@ namespace AiDotNet.RetrievalAugmentedGeneration.Retrievers { /// - /// Hybrid retriever combining dense and sparse retrieval strategies. + /// Hybrid retriever combining dense and sparse retrieval via Reciprocal Rank Fusion (RRF). /// /// The numeric type for vector operations. + /// + /// Dense (cosine, ~0..1) and sparse (BM25, unbounded) scores live on incompatible scales, so a raw + /// weighted sum is dominated by whichever list has larger-magnitude scores. RRF instead fuses by RANK — + /// each list contributes weight / (k + rank) — which is scale-free and is the default hybrid + /// fusion in Elasticsearch/Weaviate/Qdrant and the LangChain/LlamaIndex ensemble retrievers + /// (Cormack et al., 2009). The weights still let callers favor one retriever; k (default 60) damps the + /// influence of low ranks. + /// [ComponentType(ComponentType.Retriever)] [PipelineStage(PipelineStage.Retrieval)] public class HybridRetriever : RetrieverBase { private readonly IRetriever _denseRetriever; private readonly IRetriever _sparseRetriever; - private readonly T _denseWeight; - private readonly T _sparseWeight; + private readonly double _denseWeight; + private readonly double _sparseWeight; + private readonly int _rrfK; public HybridRetriever( IRetriever denseRetriever, IRetriever sparseRetriever, double denseWeight = 0.7, double sparseWeight = 0.3, - int defaultTopK = 5) + int defaultTopK = 5, + int rrfK = 60) : base(defaultTopK) { if (denseRetriever == null) throw new ArgumentNullException(nameof(denseRetriever)); if (sparseRetriever == null) throw new ArgumentNullException(nameof(sparseRetriever)); + if (rrfK <= 0) + throw new ArgumentOutOfRangeException(nameof(rrfK), "RRF k must be positive."); _denseRetriever = denseRetriever; _sparseRetriever = sparseRetriever; - _denseWeight = NumOps.FromDouble(denseWeight); - _sparseWeight = NumOps.FromDouble(sparseWeight); + _denseWeight = denseWeight; + _sparseWeight = sparseWeight; + _rrfK = rrfK; } protected override IEnumerable> RetrieveCore(string query, int topK, Dictionary metadataFilters) { - var denseResults = _denseRetriever.Retrieve(query, topK * 2, metadataFilters); - var sparseResults = _sparseRetriever.Retrieve(query, topK * 2, metadataFilters); + var denseResults = _denseRetriever.Retrieve(query, topK * 2, metadataFilters).ToList(); + var sparseResults = _sparseRetriever.Retrieve(query, topK * 2, metadataFilters).ToList(); - var combinedScores = new Dictionary(); + var rrfScores = new Dictionary(); + var docById = new Dictionary>(); - foreach (var doc in denseResults.Where(d => d.HasRelevanceScore)) + void Fuse(List> results, double weight) { - var score = NumOps.Multiply(_denseWeight, doc.RelevanceScore); - combinedScores[doc.Id] = score; - } - - foreach (var doc in sparseResults.Where(doc => doc.HasRelevanceScore)) - { - var sparseScore = NumOps.Multiply(_sparseWeight, doc.RelevanceScore); - combinedScores[doc.Id] = combinedScores.TryGetValue(doc.Id, out var existingScore) - ? NumOps.Add(existingScore, sparseScore) - : sparseScore; + // RRF uses list ORDER, not the raw scores — so scale differences between dense and sparse + // retrievers no longer distort the fusion. + for (int rank = 0; rank < results.Count; rank++) + { + var doc = results[rank]; + double contribution = weight / (_rrfK + rank + 1); // rank is 0-based → 1-based position + rrfScores[doc.Id] = rrfScores.TryGetValue(doc.Id, out var s) ? s + contribution : contribution; + if (!docById.ContainsKey(doc.Id)) + { + docById[doc.Id] = doc; + } + } } - var allDocuments = denseResults.Concat(sparseResults) - .GroupBy(d => d.Id) - .Select(g => g.First()) - .ToList(); + Fuse(denseResults, _denseWeight); + Fuse(sparseResults, _sparseWeight); - var results = allDocuments - .Where(doc => combinedScores.TryGetValue(doc.Id, out _)) - .OrderByDescending(doc => combinedScores[doc.Id]) + return rrfScores + .OrderByDescending(kv => kv.Value) .Take(topK) - .Select(doc => + .Select(kv => { - doc.RelevanceScore = combinedScores[doc.Id]; + var doc = docById[kv.Key]; + doc.RelevanceScore = NumOps.FromDouble(kv.Value); doc.HasRelevanceScore = true; return doc; }) .ToList(); - - return results; } } } diff --git a/src/RetrievalAugmentedGeneration/Retrievers/QueryRoutingRetriever.cs b/src/RetrievalAugmentedGeneration/Retrievers/QueryRoutingRetriever.cs new file mode 100644 index 0000000000..0634e27fb7 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Retrievers/QueryRoutingRetriever.cs @@ -0,0 +1,102 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.Models; + +namespace AiDotNet.RetrievalAugmentedGeneration.Retrievers +{ + /// + /// Routes each query to the single most appropriate underlying retriever (data source). With a real + /// text generator the routing is an LLM classification over the source descriptions; without one it + /// falls back to description/query token overlap. Mirrors the LangChain/LlamaIndex router retriever. + /// + /// The numeric type for vector operations. + [ComponentType(ComponentType.Retriever)] + [PipelineStage(PipelineStage.Retrieval)] + public class QueryRoutingRetriever : RetrieverBase + { + /// A named, described data source backed by a retriever. + public sealed class Route + { + public string Name { get; } + public string Description { get; } + public IRetriever Retriever { get; } + + public Route(string name, string description, IRetriever retriever) + { + if (string.IsNullOrWhiteSpace(name)) throw new ArgumentException("Route name is required.", nameof(name)); + Name = name; + Description = description ?? string.Empty; + Retriever = retriever ?? throw new ArgumentNullException(nameof(retriever)); + } + } + + private readonly IReadOnlyList _routes; + private readonly ITextGenerator? _generator; + + public QueryRoutingRetriever(IReadOnlyList routes, ITextGenerator? generator = null, int defaultTopK = 5) + : base(defaultTopK) + { + if (routes == null) throw new ArgumentNullException(nameof(routes)); + if (routes.Count == 0) throw new ArgumentException("At least one route is required.", nameof(routes)); + _routes = routes; + _generator = generator; + } + + protected override IEnumerable> RetrieveCore(string query, int topK, Dictionary metadataFilters) + => ChooseRoute(query).Retriever.Retrieve(query, topK, metadataFilters); + + /// Exposed for testing/inspection: which route a query resolves to. + public Route ChooseRoute(string query) + { + if (_generator != null) + { + var sources = string.Join("\n", _routes.Select(r => $"- {r.Name}: {r.Description}")); + var prompt = + $"Choose the single best data source for the query. Reply with ONLY the source name, exactly " + + $"as written, and nothing else.\n\nSources:\n{sources}\n\nQuery: {query}\n\nBest source name:"; + var reply = _generator.Generate(prompt); + if (!string.IsNullOrWhiteSpace(reply)) + { + var match = _routes.FirstOrDefault(r => + reply.IndexOf(r.Name, StringComparison.OrdinalIgnoreCase) >= 0); + if (match != null) return match; + } + // Unrecognized reply: fall through to the lexical heuristic. + } + + return ChooseByOverlap(query); + } + + // Fallback: the route whose name+description shares the most tokens with the query. + private Route ChooseByOverlap(string query) + { + var q = Tokenize(query); + Route best = _routes[0]; + int bestScore = -1; + foreach (var route in _routes) + { + var tokens = Tokenize(route.Name + " " + route.Description); + int overlap = tokens.Count(q.Contains); + if (overlap > bestScore) + { + bestScore = overlap; + best = route; + } + } + + return best; + } + + private static HashSet Tokenize(string text) + { + if (string.IsNullOrEmpty(text)) return new HashSet(); + return new HashSet( + text.ToLowerInvariant().Split(new[] { ' ', '\t', '\n', '\r', '.', ',', '!', '?', ';', ':' }, + StringSplitOptions.RemoveEmptyEntries)); + } + } +} diff --git a/src/RetrievalAugmentedGeneration/Retrievers/SelfQueryRetriever.cs b/src/RetrievalAugmentedGeneration/Retrievers/SelfQueryRetriever.cs new file mode 100644 index 0000000000..ddb0b917dd --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Retrievers/SelfQueryRetriever.cs @@ -0,0 +1,123 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json.Linq; + +namespace AiDotNet.RetrievalAugmentedGeneration.Retrievers +{ + /// + /// Self-querying retriever: uses a text generator to split a natural-language query into a semantic + /// search string plus structured metadata filters over known fields, then delegates to a base retriever + /// with those filters merged in. Mirrors the LangChain/LlamaIndex self-query retriever. Without a + /// generator it passes the query through unchanged (filters from the caller only). + /// + /// The numeric type for vector operations. + [ComponentType(ComponentType.Retriever)] + [PipelineStage(PipelineStage.Retrieval)] + public class SelfQueryRetriever : RetrieverBase + { + private readonly IRetriever _baseRetriever; + private readonly IReadOnlyList _metadataFields; + private readonly ITextGenerator? _generator; + + /// The retriever the cleaned query + extracted filters are passed to. + /// The filterable metadata field names the LLM may use. + /// Optional real text generator; when null the query passes through unparsed. + /// Default number of documents to return. + public SelfQueryRetriever( + IRetriever baseRetriever, + IReadOnlyList metadataFields, + ITextGenerator? generator = null, + int defaultTopK = 5) + : base(defaultTopK) + { + _baseRetriever = baseRetriever ?? throw new ArgumentNullException(nameof(baseRetriever)); + _metadataFields = metadataFields ?? new List(); + _generator = generator; + } + + protected override IEnumerable> RetrieveCore(string query, int topK, Dictionary metadataFilters) + { + string searchQuery = query; + var merged = new Dictionary(metadataFilters ?? new Dictionary()); + + if (_generator != null && _metadataFields.Count > 0) + { + var (cleaned, extracted) = ParseQuery(query); + if (!string.IsNullOrWhiteSpace(cleaned)) searchQuery = cleaned; + foreach (var kv in extracted) + { + // Caller-supplied filters take precedence over LLM-extracted ones. + if (!merged.ContainsKey(kv.Key)) merged[kv.Key] = kv.Value; + } + } + + return _baseRetriever.Retrieve(searchQuery, topK, merged); + } + + private (string cleanedQuery, Dictionary filters) ParseQuery(string query) + { + var filters = new Dictionary(); + var fields = string.Join(", ", _metadataFields); + var prompt = + $"Rewrite the user query for semantic search and extract any metadata filters that map to these " + + $"fields: {fields}. Reply with ONLY a JSON object of the form " + + $"{{\"query\": \"...\", \"filters\": {{\"field\": value}}}}. Use only the listed fields; omit " + + $"filters if none apply.\n\nUser query: {query}\n\nJSON:"; + var reply = _generator!.Generate(prompt); + if (string.IsNullOrWhiteSpace(reply)) return (query, filters); + + try + { + var json = ExtractJsonObject(reply); + if (json == null) return (query, filters); + + var obj = JObject.Parse(json); + var cleaned = obj.Value("query"); + var filtersToken = obj["filters"] as JObject; + if (filtersToken != null) + { + foreach (var prop in filtersToken.Properties()) + { + // Only accept declared fields to avoid injecting arbitrary filter keys. + if (!_metadataFields.Contains(prop.Name)) continue; + var val = ToClrValue(prop.Value); + if (val != null) filters[prop.Name] = val; + } + } + + return (string.IsNullOrWhiteSpace(cleaned) ? query : cleaned!, filters); + } + catch (Exception) + { + // Malformed LLM JSON: degrade gracefully to the raw query with no extracted filters. + return (query, filters); + } + } + + // Pull the first {...} block out of a possibly chatty LLM reply. + private static string? ExtractJsonObject(string text) + { + int start = text.IndexOf('{'); + int end = text.LastIndexOf('}'); + return (start >= 0 && end > start) ? text.Substring(start, end - start + 1) : null; + } + + private static object? ToClrValue(JToken token) + { + switch (token.Type) + { + case JTokenType.Integer: return token.Value(); + case JTokenType.Float: return token.Value(); + case JTokenType.Boolean: return token.Value(); + case JTokenType.String: return token.Value(); + case JTokenType.Array: return token.Select(ToClrValue).Where(v => v != null).ToList(); + default: return null; + } + } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/RetrieverFusionAndRoutingTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/RetrieverFusionAndRoutingTests.cs new file mode 100644 index 0000000000..6fafd3b725 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/RetrieverFusionAndRoutingTests.cs @@ -0,0 +1,131 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using AiDotNet.RetrievalAugmentedGeneration.Retrievers; +using Xunit; + +namespace AiDotNet.Tests.UnitTests.RetrievalAugmentedGeneration.Retrievers; + +/// +/// Tests for the RRF-based fusion retrievers (Hybrid/Ensemble) and the generator-driven routing and +/// self-query retrievers added in the RAG productionization work. +/// +public class RetrieverFusionAndRoutingTests +{ + private sealed class FakeRetriever : IRetriever + { + private readonly List> _docs; + public string? LastQuery { get; private set; } + public Dictionary? LastFilters { get; private set; } + public int DefaultTopK => 5; + + public FakeRetriever(params string[] docIds) + => _docs = docIds.Select((id, i) => new Document(id, id + " content") + { RelevanceScore = 1.0 - i * 0.1, HasRelevanceScore = true }).ToList(); + + public IEnumerable> Retrieve(string query) => Retrieve(query, 5, new Dictionary()); + public IEnumerable> Retrieve(string query, int topK) => Retrieve(query, topK, new Dictionary()); + public IEnumerable> Retrieve(string query, int topK, Dictionary metadataFilters) + { + LastQuery = query; + LastFilters = metadataFilters; + return _docs.Take(topK).ToList(); + } + } + + private sealed class FakeTextGen : ITextGenerator + { + private readonly Func _responder; + public FakeTextGen(Func responder) => _responder = responder; + public string Generate(string prompt) => _responder(prompt); + } + + [Fact] + public void HybridRetriever_RRF_RanksSharedDocHighest() + { + // "B" appears in both lists (rank 1 dense, rank 0 sparse) → RRF should float it to the top. + var dense = new FakeRetriever("A", "B"); + var sparse = new FakeRetriever("B", "C"); + var hybrid = new HybridRetriever(dense, sparse, defaultTopK: 3); + + var results = hybrid.Retrieve("q", 3, new Dictionary()).ToList(); + + Assert.Equal("B", results[0].Id); + Assert.Equal(3, results.Count); // A, B, C all present + } + + [Fact] + public void EnsembleRetriever_RRF_FavorsDocInMultipleLists() + { + var r1 = new FakeRetriever("X", "Y"); + var r2 = new FakeRetriever("Y", "Z"); + var r3 = new FakeRetriever("Y", "W"); + var ensemble = new EnsembleRetriever(new IRetriever[] { r1, r2, r3 }, defaultTopK: 4); + + var results = ensemble.Retrieve("q", 4, new Dictionary()).ToList(); + + Assert.Equal("Y", results[0].Id); // in all three lists + } + + [Fact] + public void QueryRoutingRetriever_UsesGeneratorChoice() + { + var news = new FakeRetriever("news-doc"); + var code = new FakeRetriever("code-doc"); + var routes = new List.Route> + { + new("news", "current events and articles", news), + new("code", "source code and APIs", code), + }; + var gen = new FakeTextGen(_ => "code"); + var router = new QueryRoutingRetriever(routes, gen, defaultTopK: 3); + + var results = router.Retrieve("how do I call the API?", 3, new Dictionary()).ToList(); + + Assert.Equal("code-doc", results[0].Id); + Assert.Equal("code", router.ChooseRoute("anything").Name); + } + + [Fact] + public void QueryRoutingRetriever_FallsBackToOverlapWithoutGenerator() + { + var news = new FakeRetriever("news-doc"); + var code = new FakeRetriever("code-doc"); + var routes = new List.Route> + { + new("news", "current events and articles", news), + new("code", "source code and APIs", code), + }; + var router = new QueryRoutingRetriever(routes, generator: null); + + Assert.Equal("code", router.ChooseRoute("show me the source code").Name); + } + + [Fact] + public void SelfQueryRetriever_ExtractsFilterAndCleansQuery() + { + var baseRetriever = new FakeRetriever("d1"); + var gen = new FakeTextGen(_ => "{\"query\": \"laptops\", \"filters\": {\"year\": 2020}}"); + var selfQuery = new SelfQueryRetriever(baseRetriever, new[] { "year" }, gen, defaultTopK: 3); + + _ = selfQuery.Retrieve("laptops released in 2020", 3, new Dictionary()).ToList(); + + Assert.Equal("laptops", baseRetriever.LastQuery); + Assert.True(baseRetriever.LastFilters!.ContainsKey("year")); + Assert.Equal(2020L, baseRetriever.LastFilters!["year"]); + } + + [Fact] + public void SelfQueryRetriever_CallerFiltersTakePrecedence() + { + var baseRetriever = new FakeRetriever("d1"); + var gen = new FakeTextGen(_ => "{\"query\": \"laptops\", \"filters\": {\"year\": 2020}}"); + var selfQuery = new SelfQueryRetriever(baseRetriever, new[] { "year" }, gen); + + _ = selfQuery.Retrieve("laptops in 2020", 3, new Dictionary { ["year"] = 1999L }).ToList(); + + Assert.Equal(1999L, baseRetriever.LastFilters!["year"]); // caller wins + } +} From 392367c994848c420f30aca8da90fd37c282a76a Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 15:41:51 -0400 Subject: [PATCH 05/25] fix(rag): embedding models fail loudly instead of returning fake vectors Co-Authored-By: Claude Opus 4.8 (1M context) --- .../Embeddings/EmbeddingModelBase.cs | 39 -- .../Embeddings/GooglePalmEmbeddingModel.cs | 13 +- .../Embeddings/HuggingFaceEmbeddingModel.cs | 13 +- .../Embeddings/MultiModalEmbeddingModel.cs | 61 +-- .../Embeddings/ONNXSentenceTransformer.cs | 17 - .../Embeddings/OpenAIEmbeddingModel.cs | 13 +- .../Embeddings/VoyageAIEmbeddingModel.cs | 172 ++++++-- .../GooglePalmEmbeddingModelTests.cs | 350 ++++------------ .../HuggingFaceEmbeddingModelTests.cs | 378 ++++------------- .../MultiModalEmbeddingModelTests.cs | 209 ++-------- .../Embeddings/OpenAIEmbeddingModelTests.cs | 393 ++++-------------- .../Embeddings/VoyageAIEmbeddingModelTests.cs | 363 ++++++---------- 12 files changed, 592 insertions(+), 1429 deletions(-) diff --git a/src/RetrievalAugmentedGeneration/Embeddings/EmbeddingModelBase.cs b/src/RetrievalAugmentedGeneration/Embeddings/EmbeddingModelBase.cs index c269080b7e..90af01ddfa 100644 --- a/src/RetrievalAugmentedGeneration/Embeddings/EmbeddingModelBase.cs +++ b/src/RetrievalAugmentedGeneration/Embeddings/EmbeddingModelBase.cs @@ -175,45 +175,6 @@ protected virtual async Task> EmbedBatchCoreAsync(IList texts) return CreateMatrixFromVectors(embeddings); } - /// - /// Generates a deterministic fallback embedding based on the text hash. - /// Used when the underlying model or API is not available (e.g., in unit tests). - /// - /// The text to generate a fallback embedding for. - /// A normalized vector of the correct embedding dimension. - /// - /// For Beginners: This is a "fake" embedding used when the real model is unavailable. - /// It produces consistent (deterministic) results for the same input text, but the vectors - /// are not semantically meaningful. This is useful for unit testing and offline development. - /// - /// - protected virtual Vector GenerateFallbackEmbedding(string text) - { - Guard.NotNull(text); - if (EmbeddingDimension <= 0) - throw new InvalidOperationException("EmbeddingDimension must be positive."); - - // Use stable FNV-1a hash instead of string.GetHashCode() which is - // randomized per process in .NET Core/.NET 5+ - var lower = text.ToLowerInvariant(); - uint hash = 2166136261u; - foreach (char c in lower) - { - hash ^= c; - hash *= 16777619u; - } - var dim = EmbeddingDimension; - var values = new T[dim]; - for (int i = 0; i < dim; i++) - { - // Generate deterministic values based on text hash and position - double val = Math.Sin(hash * 0.0001 + i * 0.1) * 0.5; - values[i] = NumOps.FromDouble(val); - } - - return new Vector(values).Normalize(); - } - /// /// Validates the input text. /// diff --git a/src/RetrievalAugmentedGeneration/Embeddings/GooglePalmEmbeddingModel.cs b/src/RetrievalAugmentedGeneration/Embeddings/GooglePalmEmbeddingModel.cs index e338bc21e2..ae713ef99f 100644 --- a/src/RetrievalAugmentedGeneration/Embeddings/GooglePalmEmbeddingModel.cs +++ b/src/RetrievalAugmentedGeneration/Embeddings/GooglePalmEmbeddingModel.cs @@ -82,15 +82,10 @@ public GooglePalmEmbeddingModel( /// protected override Vector EmbedCore(string text) { - try - { - return EmbedAsync(text).ConfigureAwait(false).GetAwaiter().GetResult(); - } - catch (Exception ex) when (ex is HttpRequestException or TaskCanceledException or AggregateException) - { - // API unavailable (e.g., unit tests with fake API key) - use deterministic fallback - return GenerateFallbackEmbedding(text); - } + // Surface API/transport failures instead of fabricating a fake vector. + // Returning a random/hash-based embedding here would silently poison the + // vector store with garbage that looks like a valid embedding. + return EmbedAsync(text).ConfigureAwait(false).GetAwaiter().GetResult(); } /// diff --git a/src/RetrievalAugmentedGeneration/Embeddings/HuggingFaceEmbeddingModel.cs b/src/RetrievalAugmentedGeneration/Embeddings/HuggingFaceEmbeddingModel.cs index b00c8b6587..41288d6a29 100644 --- a/src/RetrievalAugmentedGeneration/Embeddings/HuggingFaceEmbeddingModel.cs +++ b/src/RetrievalAugmentedGeneration/Embeddings/HuggingFaceEmbeddingModel.cs @@ -59,15 +59,10 @@ public HuggingFaceEmbeddingModel(string modelName, string apiKey = "", int dimen protected override Vector EmbedCore(string text) { - try - { - return EmbedAsync(text).ConfigureAwait(false).GetAwaiter().GetResult(); - } - catch (Exception ex) when (ex is HttpRequestException or TaskCanceledException or AggregateException) - { - // API unavailable (e.g., unit tests with fake API key) - use deterministic fallback - return GenerateFallbackEmbedding(text); - } + // Surface API/transport failures instead of fabricating a fake vector. + // Returning a random/hash-based embedding here would silently poison the + // vector store with garbage that looks like a valid embedding. + return EmbedAsync(text).ConfigureAwait(false).GetAwaiter().GetResult(); } /// diff --git a/src/RetrievalAugmentedGeneration/Embeddings/MultiModalEmbeddingModel.cs b/src/RetrievalAugmentedGeneration/Embeddings/MultiModalEmbeddingModel.cs index dec5f668a1..1adc95ccf0 100644 --- a/src/RetrievalAugmentedGeneration/Embeddings/MultiModalEmbeddingModel.cs +++ b/src/RetrievalAugmentedGeneration/Embeddings/MultiModalEmbeddingModel.cs @@ -85,9 +85,10 @@ namespace AiDotNet.RetrievalAugmentedGeneration.EmbeddingModels; /// 3. Image Features → Projection Head → Image Embedding (512D) /// /// Both paths output embeddings in the same 512-dimensional space, enabling direct comparison. -/// -/// Current implementation uses ONNX for text encoding and simulated image encoding. -/// Production should use full CLIP ONNX model with both encoders. +/// +/// Current implementation uses ONNX for text encoding only. The image path +/// () is not implemented and throws +/// rather than fabricating a vector; a full CLIP ONNX vision encoder must be connected for it. /// /// Benefits: /// - Cross-modal search - Find images with text, text with images @@ -97,7 +98,7 @@ namespace AiDotNet.RetrievalAugmentedGeneration.EmbeddingModels; /// - Pre-trained - No custom training needed for basic use cases /// /// Limitations: -/// - Image encoding currently simulated (hash-based) - needs real CLIP encoder +/// - Image encoding is NOT implemented (EmbedImage throws NotSupportedException) - needs real CLIP encoder /// - Model file size can be large (1-4GB for CLIP variants) /// - Image processing slower than text (convolutional layers are compute-intensive) /// - Works best with natural images (struggles with abstract diagrams, charts) @@ -152,31 +153,28 @@ protected override Vector EmbedCore(string text) } /// - /// Generates image embeddings from file path. + /// Generates image embeddings from a file path. /// /// Path to the image file. /// The embedding vector for the image. + /// Thrown when is null or whitespace. + /// + /// Always thrown: real image encoding (a CLIP-style vision encoder) is not yet wired up + /// for this model. Previously this path returned a hash of the file path/size, which is + /// NOT a semantic embedding and would silently corrupt any cross-modal similarity search. + /// public Vector EmbedImage(string imagePath) { if (string.IsNullOrWhiteSpace(imagePath)) throw new ArgumentException("Image path cannot be null or whitespace", nameof(imagePath)); - if (!File.Exists(imagePath)) - throw new FileNotFoundException($"Image file not found: {imagePath}"); - - // Generate embedding using image path hash - // In production, this would use CLIP's image encoder with convolutional layers - var values = new T[_dimension]; - var hash = GetImageHash(imagePath); - - for (int i = 0; i < _dimension; i++) - { - var val = NumOps.FromDouble(Math.Sin((double)hash * (i + 1) * 0.003)); - values[i] = val; - } - - var embedding = new Vector(values); - return _normalizeEmbeddings ? embedding.Normalize() : embedding; + // Do NOT fabricate an embedding from a file-path/size hash. A hash carries no visual + // semantics, so it cannot be compared against text embeddings in a shared space. + // Throw loudly until a real vision encoder is connected. + throw new NotSupportedException( + "Image embedding is not supported by MultiModalEmbeddingModel: a real CLIP-style " + + "vision encoder is not wired up. Use a dedicated vision/CLIP encoder to embed images " + + "rather than relying on this model, which only produces genuine embeddings for text."); } /// @@ -184,6 +182,8 @@ public Vector EmbedImage(string imagePath) /// /// Paths to image files. /// Embedding vectors for all images. + /// Thrown when is null. + /// Thrown when enumerated: see . public IEnumerable> EmbedImageBatch(IEnumerable imagePaths) { if (imagePaths == null) @@ -191,25 +191,6 @@ public IEnumerable> EmbedImageBatch(IEnumerable imagePaths) return imagePaths.Select(path => EmbedImage(path)); } - - private int GetImageHash(string imagePath) - { - // Simple hash based on file path and length - // In production, would hash image content after preprocessing - unchecked - { - int hash = 17; - hash = (hash * 31) + imagePath.GetHashCode(); - - if (File.Exists(imagePath)) - { - var fileInfo = new FileInfo(imagePath); - hash = (hash * 31) + fileInfo.Length.GetHashCode(); - } - - return hash; - } - } } diff --git a/src/RetrievalAugmentedGeneration/Embeddings/ONNXSentenceTransformer.cs b/src/RetrievalAugmentedGeneration/Embeddings/ONNXSentenceTransformer.cs index 2508d32200..d0ded10791 100644 --- a/src/RetrievalAugmentedGeneration/Embeddings/ONNXSentenceTransformer.cs +++ b/src/RetrievalAugmentedGeneration/Embeddings/ONNXSentenceTransformer.cs @@ -234,23 +234,6 @@ private float[] ApplyMeanPooling(Microsoft.ML.OnnxRuntime.Tensors.Tensor return pooled; } - /// - /// Generates a deterministic fallback embedding based on the text hash. - /// Used when the ONNX model file is not available (e.g., in unit tests). - /// - protected override Vector GenerateFallbackEmbedding(string text) - { - var hash = text.ToLowerInvariant().GetHashCode(); - var values = new T[_dimension]; - for (int i = 0; i < _dimension; i++) - { - // Generate deterministic values based on text hash and position - double val = Math.Sin(hash * 0.0001 + i * 0.1) * 0.5; - values[i] = NumOps.FromDouble(val); - } - return new Vector(values).Normalize(); - } - private Vector CreateZeroVector() { var values = new T[_dimension]; diff --git a/src/RetrievalAugmentedGeneration/Embeddings/OpenAIEmbeddingModel.cs b/src/RetrievalAugmentedGeneration/Embeddings/OpenAIEmbeddingModel.cs index eb2bf8f4ee..031f4b7d1f 100644 --- a/src/RetrievalAugmentedGeneration/Embeddings/OpenAIEmbeddingModel.cs +++ b/src/RetrievalAugmentedGeneration/Embeddings/OpenAIEmbeddingModel.cs @@ -67,15 +67,10 @@ public OpenAIEmbeddingModel(string apiKey, string modelName = "text-embedding-ad protected override Vector EmbedCore(string text) { - try - { - return EmbedAsync(text).ConfigureAwait(false).GetAwaiter().GetResult(); - } - catch (Exception ex) when (ex is HttpRequestException or TaskCanceledException or AggregateException) - { - // API unavailable (e.g., unit tests with fake API key) - use deterministic fallback - return GenerateFallbackEmbedding(text); - } + // Surface API/transport failures instead of fabricating a fake vector. + // Returning a random/hash-based embedding here would silently poison the + // vector store with garbage that looks like a valid embedding. + return EmbedAsync(text).ConfigureAwait(false).GetAwaiter().GetResult(); } /// diff --git a/src/RetrievalAugmentedGeneration/Embeddings/VoyageAIEmbeddingModel.cs b/src/RetrievalAugmentedGeneration/Embeddings/VoyageAIEmbeddingModel.cs index c96082d474..052dafef0c 100644 --- a/src/RetrievalAugmentedGeneration/Embeddings/VoyageAIEmbeddingModel.cs +++ b/src/RetrievalAugmentedGeneration/Embeddings/VoyageAIEmbeddingModel.cs @@ -1,17 +1,27 @@ using System; +using System.Collections.Generic; +using System.Linq; +using System.Net.Http; +using System.Net.Http.Headers; +using System.Text; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; -using AiDotNet.Interfaces; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.Embeddings; -using AiDotNet.Validation; +using Newtonsoft.Json; namespace AiDotNet.RetrievalAugmentedGeneration.EmbeddingModels; /// -/// Voyage AI-compatible embedding model using ONNX for high-performance local inference. +/// Voyage AI embedding model that calls the Voyage AI embeddings REST API. /// /// The numeric data type used for vector operations. +/// +/// Sends requests to https://api.voyageai.com/v1/embeddings using the supplied +/// API key (Bearer authentication) and parses the returned embeddings. This is a real +/// remote model - it does not fabricate vectors when the API is unavailable. +/// [ModelDomain(ModelDomain.Language)] [ModelCategory(ModelCategory.NeuralNetwork)] [ModelTask(ModelTask.FeatureExtraction)] @@ -20,42 +30,52 @@ namespace AiDotNet.RetrievalAugmentedGeneration.EmbeddingModels; public class VoyageAIEmbeddingModel : EmbeddingModelBase { private const int DefaultMaxTokens = 16000; + private const string EmbeddingsEndpoint = "https://api.voyageai.com/v1/embeddings"; + private readonly string _apiKey; private readonly string _model; private readonly string _inputType; private readonly int _dimension; - private readonly ONNXSentenceTransformer _onnxTransformer; + private readonly HttpClient _httpClient; + private readonly bool _ownsHttpClient; private bool _disposed; /// /// Initializes a new instance of the class. /// - /// The Voyage AI API key (unused for local ONNX). - /// The model path (e.g., "path/to/voyage-model.onnx"). + /// The Voyage AI API key. + /// The Voyage model name (e.g., "voyage-3", "voyage-3-lite"). /// The input type ("document" or "query"). - /// The embedding dimension. + /// The embedding dimension produced by the selected model. + /// Optional external HttpClient. If provided, the caller is responsible for configuring authentication and managing its lifecycle. public VoyageAIEmbeddingModel( string apiKey, string model, string inputType, - int dimension) + int dimension, + HttpClient? httpClient = null) { - Guard.NotNull(apiKey); -#pragma warning disable CS0414 // Field is assigned but its value is never used - _apiKey = apiKey; // Kept for API compatibility but unused -#pragma warning restore CS0414 - Guard.NotNull(model); + if (string.IsNullOrWhiteSpace(apiKey)) + throw new ArgumentException("API key cannot be empty", nameof(apiKey)); + if (string.IsNullOrWhiteSpace(model)) + throw new ArgumentException("Model cannot be empty", nameof(model)); + if (string.IsNullOrWhiteSpace(inputType)) + throw new ArgumentException("Input type cannot be empty", nameof(inputType)); + if (dimension <= 0) + throw new ArgumentException("Dimension must be positive", nameof(dimension)); + + _apiKey = apiKey; _model = model; - Guard.NotNull(inputType); _inputType = inputType; _dimension = dimension; + _ownsHttpClient = httpClient == null; + _httpClient = httpClient ?? new HttpClient(); - // Initialize ONNX transformer once - _onnxTransformer = new ONNXSentenceTransformer( - modelPath: _model, - dimension: _dimension, - maxTokens: DefaultMaxTokens - ); + if (_ownsHttpClient) + { + _httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", _apiKey); + _httpClient.DefaultRequestHeaders.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json")); + } } /// @@ -67,21 +87,123 @@ public VoyageAIEmbeddingModel( /// protected override Vector EmbedCore(string text) { - return _onnxTransformer.Embed(text); + // Surface API/transport failures instead of fabricating a fake vector. + return EmbedAsync(text).ConfigureAwait(false).GetAwaiter().GetResult(); + } + + /// + /// Asynchronously generates an embedding for the specified text using the Voyage AI API. + /// + /// The text to encode. + /// A task representing the async operation, with the resulting vector. + public override async Task> EmbedAsync(string text) + { + ValidateText(text); + + var result = await PostEmbeddingsAsync(new[] { text }); + + if (result?.Data == null || result.Data.Count == 0) + { + throw new InvalidOperationException("Voyage AI API returned an empty or invalid response."); + } + + var embedding = result.Data[0].Embedding; + var values = new T[embedding.Length]; + for (int i = 0; i < embedding.Length; i++) + { + values[i] = NumOps.FromDouble(embedding[i]); + } + + return new Vector(values); + } + + /// + /// Asynchronously generates embeddings for the specified texts using the Voyage AI API. + /// + /// The collection of texts to encode. + /// A task representing the async operation, with the resulting matrix. + public override async Task> EmbedBatchAsync(IEnumerable texts) + { + var textList = texts.ToList(); + if (textList.Count == 0) + return new Matrix(0, _dimension); + + foreach (var text in textList) + ValidateText(text); + + return await EmbedBatchCoreAsync(textList); + } + + /// + protected override async Task> EmbedBatchCoreAsync(IList texts) + { + var result = await PostEmbeddingsAsync(texts); + + if (result?.Data == null || result.Data.Count != texts.Count) + { + throw new InvalidOperationException("Voyage AI API returned an empty or mismatched response."); + } + + var matrix = new Matrix(texts.Count, _dimension); + for (int i = 0; i < result.Data.Count; i++) + { + var embedding = result.Data[i].Embedding; + for (int j = 0; j < embedding.Length; j++) + { + matrix[i, j] = NumOps.FromDouble(embedding[j]); + } + } + + return matrix; + } + + private async Task PostEmbeddingsAsync(IEnumerable input) + { + var requestBody = new + { + input = input, + model = _model, + input_type = _inputType + }; + + using var content = new StringContent(JsonConvert.SerializeObject(requestBody), Encoding.UTF8, "application/json"); + using var response = await _httpClient.PostAsync(EmbeddingsEndpoint, content); + + if (!response.IsSuccessStatusCode) + { + var errorContent = await response.Content.ReadAsStringAsync(); + throw new HttpRequestException($"Voyage AI API request failed with status code {response.StatusCode}: {errorContent}"); + } + + var responseJson = await response.Content.ReadAsStringAsync(); + return JsonConvert.DeserializeObject(responseJson); + } + + private class VoyageEmbeddingResponse + { + [JsonProperty("data")] + public List Data { get; set; } = new(); + } + + private class VoyageEmbeddingData + { + [JsonProperty("embedding")] + public double[] Embedding { get; set; } = Array.Empty(); + + [JsonProperty("index")] + public int Index { get; set; } } protected override void Dispose(bool disposing) { if (!_disposed) { - if (disposing) + if (disposing && _ownsHttpClient) { - _onnxTransformer.Dispose(); + _httpClient.Dispose(); } _disposed = true; } base.Dispose(disposing); } } - - diff --git a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/GooglePalmEmbeddingModelTests.cs b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/GooglePalmEmbeddingModelTests.cs index 3c1e0e0ace..5a8bd61d56 100644 --- a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/GooglePalmEmbeddingModelTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/GooglePalmEmbeddingModelTests.cs @@ -2,345 +2,159 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Net; +using System.Net.Http; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.EmbeddingModels; +using Newtonsoft.Json.Linq; using Xunit; -using System.Threading.Tasks; namespace AiDotNetTests.UnitTests.RAG.Embeddings { public class GooglePalmEmbeddingModelTests { - [Fact(Timeout = 60000)] - public async Task Constructor_WithValidParameters_CreatesInstance() + private sealed class StubHandler : HttpMessageHandler { - // Arrange & Act - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); + private readonly HttpStatusCode _status; + private readonly string _responseBody; + public string LastRequestBody { get; private set; } = ""; + public Uri LastRequestUri { get; private set; } - // Assert - Assert.NotNull(model); - Assert.Equal(768, model.EmbeddingDimension); - Assert.Equal(2048, model.MaxTokens); + public StubHandler(string responseBody, HttpStatusCode status = HttpStatusCode.OK) + { + _responseBody = responseBody; + _status = status; + } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + LastRequestUri = request.RequestUri; + if (request.Content is not null) + LastRequestBody = await request.Content.ReadAsStringAsync().ConfigureAwait(false); + + return new HttpResponseMessage(_status) { Content = new StringContent(_responseBody ?? "") }; + } } - [Fact(Timeout = 60000)] - public async Task Constructor_WithDefaultDimension_CreatesInstance() - { - // Arrange & Act - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key"); + private static GooglePalmEmbeddingModel ModelWith(StubHandler handler, int dimension) + => new("test-project-id", "us-central1", "text-embedding-004", "test-api-key", dimension, new HttpClient(handler)); + + // ────────── Constructor validation ────────── - // Assert + [Fact] + public void Constructor_WithValidParameters_CreatesInstance() + { + var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); Assert.NotNull(model); Assert.Equal(768, model.EmbeddingDimension); Assert.Equal(2048, model.MaxTokens); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullProjectId_ThrowsArgumentNullException() + [Fact] + public void Constructor_WithNullProjectId_ThrowsArgumentNullException() { - // Arrange & Act & Assert Assert.Throws(() => new GooglePalmEmbeddingModel(null, "us-central1", "textembedding-gecko@001", "test-api-key")); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullLocation_ThrowsArgumentNullException() + [Fact] + public void Constructor_WithNullLocation_ThrowsArgumentNullException() { - // Arrange & Act & Assert Assert.Throws(() => new GooglePalmEmbeddingModel("test-project-id", null, "textembedding-gecko@001", "test-api-key")); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullModel_ThrowsArgumentNullException() + [Fact] + public void Constructor_WithNullModel_ThrowsArgumentNullException() { - // Arrange & Act & Assert Assert.Throws(() => new GooglePalmEmbeddingModel("test-project-id", "us-central1", null, "test-api-key")); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullApiKey_ThrowsArgumentNullException() + [Fact] + public void Constructor_WithNullApiKey_ThrowsArgumentNullException() { - // Arrange & Act & Assert Assert.Throws(() => new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", null)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithValidText_ReturnsVectorOfCorrectDimension() - { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var text = "This is a test sentence."; - - // Act - var embedding = model.Embed(text); - - // Assert - Assert.NotNull(embedding); - Assert.Equal(768, embedding.Length); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithSameTextTwice_ReturnsSameEmbedding() - { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var text = "Hello world"; - - // Act - var embedding1 = model.Embed(text); - var embedding2 = model.Embed(text); - - // Assert - for (int i = 0; i < embedding1.Length; i++) - { - Assert.Equal(embedding1[i], embedding2[i], 10); - } - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithDifferentTexts_ReturnsDifferentEmbeddings() - { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var text1 = "Hello world"; - var text2 = "Goodbye world"; - - // Act - var embedding1 = model.Embed(text1); - var embedding2 = model.Embed(text2); - - // Assert - var hasDifference = false; - for (int i = 0; i < embedding1.Length; i++) - { - if (Math.Abs(embedding1[i] - embedding2[i]) > 1e-10) - { - hasDifference = true; - break; - } - } - Assert.True(hasDifference, "Embeddings for different texts should be different"); - } - - [Fact(Timeout = 60000)] - public async Task Embed_ReturnsNormalizedVector() - { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var text = "Test normalization"; - - // Act - var embedding = model.Embed(text); - - // Assert - var magnitude = 0.0; - for (int i = 0; i < embedding.Length; i++) - { - magnitude += embedding[i] * embedding[i]; - } - magnitude = Math.Sqrt(magnitude); - Assert.Equal(1.0, magnitude, 5); - } + // ────────── Input validation (no network) ────────── - [Fact(Timeout = 60000)] - public async Task Embed_WithNullText_ThrowsArgumentException() + [Fact] + public void Embed_WithNullText_ThrowsArgumentException() { - // Arrange var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key"); - - // Act & Assert Assert.Throws(() => model.Embed(null)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithEmptyText_ThrowsArgumentException() - { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key"); - - // Act & Assert - Assert.Throws(() => model.Embed(string.Empty)); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithWhitespaceText_ThrowsArgumentException() + [Fact] + public void EmbedBatch_WithNullTexts_ThrowsArgumentNullException() { - // Arrange var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key"); - - // Act & Assert - Assert.Throws(() => model.Embed(" ")); - } - - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithValidTexts_ReturnsMatrixOfCorrectDimensions() - { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var texts = new List { "First text", "Second text", "Third text" }; - - // Act - var embeddings = model.EmbedBatch(texts); - - // Assert - Assert.NotNull(embeddings); - Assert.Equal(3, embeddings.Rows); - Assert.Equal(768, embeddings.Columns); - } - - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithNullTexts_ThrowsArgumentNullException() - { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key"); - - // Act & Assert Assert.Throws(() => model.EmbedBatch(null)); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithEmptyCollection_ThrowsArgumentException() + [Fact] + public void EmbedBatch_WithEmptyCollection_ThrowsArgumentException() { - // Arrange var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key"); - var texts = new List(); - - // Act & Assert - Assert.Throws(() => model.EmbedBatch(texts)); + Assert.Throws(() => model.EmbedBatch(new List())); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_ProducesSameEmbeddingsAsIndividualCalls() - { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var texts = new List { "First", "Second", "Third" }; - - // Act - var batchEmbeddings = model.EmbedBatch(texts); - var individualEmbeddings = texts.Select(t => model.Embed(t)).ToList(); - - // Assert - for (int i = 0; i < texts.Count; i++) - { - for (int j = 0; j < model.EmbeddingDimension; j++) - { - Assert.Equal(individualEmbeddings[i][j], batchEmbeddings[i, j], 10); - } - } - } + // ────────── Failure surfaces (no fake fallback) ────────── - [Fact(Timeout = 60000)] - public async Task Embed_WithFloatType_WorksCorrectly() + [Fact] + public void Embed_WhenApiReturnsError_ThrowsInsteadOfFakeVector() { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var text = "Test with float type"; + var handler = new StubHandler("{\"error\":\"denied\"}", HttpStatusCode.Forbidden); + var model = ModelWith(handler, 3); - // Act - var embedding = model.Embed(text); - - // Assert - Assert.NotNull(embedding); - Assert.Equal(768, embedding.Length); - - // Check normalization - var magnitude = 0.0f; - for (int i = 0; i < embedding.Length; i++) - { - magnitude += embedding[i] * embedding[i]; - } - magnitude = (float)Math.Sqrt(magnitude); - Assert.Equal(1.0f, magnitude, 5); + var ex = Assert.Throws(() => model.Embed("hello")); + Assert.Contains("Vertex AI API request failed", ex.Message); } - [Fact(Timeout = 60000)] - public async Task Embed_WithCustomDimension_ReturnsCorrectSize() + [Fact] + public void EmbedBatch_WhenApiReturnsError_ThrowsInsteadOfFakeVector() { - // Arrange - var customDimension = 512; - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", customDimension); - var text = "Testing custom dimension"; - - // Act - var embedding = model.Embed(text); + var handler = new StubHandler("boom", HttpStatusCode.InternalServerError); + var model = ModelWith(handler, 3); - // Assert - Assert.Equal(customDimension, embedding.Length); + Assert.Throws(() => + model.EmbedBatch(new List { "a", "b" })); } - [Fact(Timeout = 60000)] - public async Task Embed_Deterministic_MultipleInstances() - { - // Arrange - var model1 = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var model2 = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var text = "Determinism test"; - - // Act - var embedding1 = model1.Embed(text); - var embedding2 = model2.Embed(text); + // ────────── Success path via mocked handler ────────── - // Assert - for (int i = 0; i < embedding1.Length; i++) - { - Assert.Equal(embedding1[i], embedding2[i], 10); - } - } - - [Fact(Timeout = 60000)] - public async Task EmbedBatch_AllRowsAreNormalized() + [Fact] + public void Embed_WithMockedSuccess_ParsesVectorAndPostsToVertexEndpoint() { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var texts = new List { "First", "Second", "Third" }; + var handler = new StubHandler("{\"predictions\":[{\"embeddings\":{\"values\":[0.1,0.2,0.3]}}]}"); + var model = ModelWith(handler, 3); - // Act - var embeddings = model.EmbedBatch(texts); + var embedding = model.Embed("hello world"); - // Assert - for (int i = 0; i < embeddings.Rows; i++) - { - var magnitude = 0.0; - for (int j = 0; j < embeddings.Columns; j++) - { - magnitude += embeddings[i, j] * embeddings[i, j]; - } - magnitude = Math.Sqrt(magnitude); - Assert.Equal(1.0, magnitude, 5); - } - } + Assert.Equal(3, embedding.Length); + Assert.Equal(0.2, embedding[1], 6); + Assert.Contains("aiplatform.googleapis.com", handler.LastRequestUri.ToString()); + Assert.Contains("text-embedding-004:predict", handler.LastRequestUri.ToString()); - [Fact(Timeout = 60000)] - public async Task Embed_WithLongText_ReturnsEmbedding() - { - // Arrange - var model = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key", 768); - var longText = string.Join(" ", Enumerable.Repeat("word", 1000)); - - // Act - var embedding = model.Embed(longText); - - // Assert - Assert.NotNull(embedding); - Assert.Equal(768, embedding.Length); + var body = JObject.Parse(handler.LastRequestBody); + Assert.Equal("hello world", (string)body["instances"][0]["content"]); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithDifferentLocations_CreatesInstances() + [Fact] + public async Task EmbedBatchAsync_WithMockedSuccess_ReturnsMatrix() { - // Arrange & Act - var usCentralModel = new GooglePalmEmbeddingModel("test-project-id", "us-central1", "textembedding-gecko@001", "test-api-key"); - var europeModel = new GooglePalmEmbeddingModel("test-project-id", "europe-west1", "textembedding-gecko@001", "test-api-key"); - var asiaModel = new GooglePalmEmbeddingModel("test-project-id", "asia-southeast1", "textembedding-gecko@001", "test-api-key"); + var handler = new StubHandler("{\"predictions\":[{\"embeddings\":{\"values\":[0.1,0.2,0.3]}},{\"embeddings\":{\"values\":[0.4,0.5,0.6]}}]}"); + var model = ModelWith(handler, 3); + + var matrix = await model.EmbedBatchAsync(new List { "a", "b" }); - // Assert - Assert.NotNull(usCentralModel); - Assert.NotNull(europeModel); - Assert.NotNull(asiaModel); + Assert.Equal(2, matrix.Rows); + Assert.Equal(3, matrix.Columns); + Assert.Equal(0.6, matrix[1, 2], 6); } } } diff --git a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/HuggingFaceEmbeddingModelTests.cs b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/HuggingFaceEmbeddingModelTests.cs index 5ba47793da..599ae9e74b 100644 --- a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/HuggingFaceEmbeddingModelTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/HuggingFaceEmbeddingModelTests.cs @@ -2,379 +2,159 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Net; +using System.Net.Http; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.EmbeddingModels; +using Newtonsoft.Json.Linq; using Xunit; -using System.Threading.Tasks; namespace AiDotNetTests.UnitTests.RAG.Embeddings { public class HuggingFaceEmbeddingModelTests { - [Fact(Timeout = 60000)] - public async Task Constructor_WithValidParameters_CreatesInstance() + private sealed class StubHandler : HttpMessageHandler { - // Arrange & Act - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); + private readonly HttpStatusCode _status; + private readonly string _responseBody; + public string LastRequestBody { get; private set; } = ""; + public Uri LastRequestUri { get; private set; } - // Assert - Assert.NotNull(model); - Assert.Equal(768, model.EmbeddingDimension); - Assert.Equal(512, model.MaxTokens); + public StubHandler(string responseBody, HttpStatusCode status = HttpStatusCode.OK) + { + _responseBody = responseBody; + _status = status; + } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + LastRequestUri = request.RequestUri; + if (request.Content is not null) + LastRequestBody = await request.Content.ReadAsStringAsync().ConfigureAwait(false); + + return new HttpResponseMessage(_status) { Content = new StringContent(_responseBody ?? "") }; + } } - [Fact(Timeout = 60000)] - public async Task Constructor_WithDefaultParameters_CreatesInstance() - { - // Arrange & Act - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2"); + private static HuggingFaceEmbeddingModel ModelWith(StubHandler handler, int dimension) + => new("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", dimension, 512, new HttpClient(handler)); - // Assert + // ────────── Constructor validation ────────── + + [Fact] + public void Constructor_WithValidParameters_CreatesInstance() + { + var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); Assert.NotNull(model); Assert.Equal(768, model.EmbeddingDimension); Assert.Equal(512, model.MaxTokens); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithEmptyApiKey_CreatesInstance() + [Fact] + public void Constructor_WithEmptyApiKey_CreatesInstance() { - // Arrange & Act var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", ""); - - // Assert Assert.NotNull(model); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullModelName_ThrowsArgumentException() + [Fact] + public void Constructor_WithNullModelName_ThrowsArgumentException() { - // Arrange & Act & Assert var exception = Assert.Throws(() => new HuggingFaceEmbeddingModel(null, "test-api-key")); Assert.Contains("Model name cannot be empty", exception.Message); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithEmptyModelName_ThrowsArgumentException() - { - // Arrange & Act & Assert - var exception = Assert.Throws(() => - new HuggingFaceEmbeddingModel("", "test-api-key")); - Assert.Contains("Model name cannot be empty", exception.Message); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithWhitespaceModelName_ThrowsArgumentException() - { - // Arrange & Act & Assert - var exception = Assert.Throws(() => - new HuggingFaceEmbeddingModel(" ", "test-api-key")); - Assert.Contains("Model name cannot be empty", exception.Message); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithZeroDimension_ThrowsArgumentException() + [Fact] + public void Constructor_WithZeroDimension_ThrowsArgumentException() { - // Arrange & Act & Assert var exception = Assert.Throws(() => new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 0, 512)); Assert.Contains("Dimension must be positive", exception.Message); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNegativeDimension_ThrowsArgumentException() - { - // Arrange & Act & Assert - var exception = Assert.Throws(() => - new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", -1, 512)); - Assert.Contains("Dimension must be positive", exception.Message); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithZeroMaxTokens_ThrowsArgumentException() - { - // Arrange & Act & Assert - var exception = Assert.Throws(() => - new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 0)); - Assert.Contains("Max tokens must be positive", exception.Message); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithNegativeMaxTokens_ThrowsArgumentException() + [Fact] + public void Constructor_WithNegativeMaxTokens_ThrowsArgumentException() { - // Arrange & Act & Assert var exception = Assert.Throws(() => new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, -1)); Assert.Contains("Max tokens must be positive", exception.Message); } - [Fact(Timeout = 60000)] - public async Task Embed_WithValidText_ReturnsVectorOfCorrectDimension() - { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var text = "This is a test sentence."; - - // Act - var embedding = model.Embed(text); + // ────────── Input validation (no network) ────────── - // Assert - Assert.NotNull(embedding); - Assert.Equal(768, embedding.Length); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithSameTextTwice_ReturnsSameEmbedding() + [Fact] + public void Embed_WithNullText_ThrowsArgumentException() { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var text = "Hello world"; - - // Act - var embedding1 = model.Embed(text); - var embedding2 = model.Embed(text); - - // Assert - for (int i = 0; i < embedding1.Length; i++) - { - Assert.Equal(embedding1[i], embedding2[i], 10); - } - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithDifferentTexts_ReturnsDifferentEmbeddings() - { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var text1 = "Hello world"; - var text2 = "Goodbye world"; - - // Act - var embedding1 = model.Embed(text1); - var embedding2 = model.Embed(text2); - - // Assert - var hasDifference = false; - for (int i = 0; i < embedding1.Length; i++) - { - if (Math.Abs(embedding1[i] - embedding2[i]) > 1e-10) - { - hasDifference = true; - break; - } - } - Assert.True(hasDifference, "Embeddings for different texts should be different"); - } - - [Fact(Timeout = 60000)] - public async Task Embed_ReturnsNormalizedVector() - { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var text = "Test normalization"; - - // Act - var embedding = model.Embed(text); - - // Assert - var magnitude = 0.0; - for (int i = 0; i < embedding.Length; i++) - { - magnitude += embedding[i] * embedding[i]; - } - magnitude = Math.Sqrt(magnitude); - Assert.Equal(1.0, magnitude, 5); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithNullText_ThrowsArgumentException() - { - // Arrange var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2"); - - // Act & Assert Assert.Throws(() => model.Embed(null)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithEmptyText_ThrowsArgumentException() + [Fact] + public void EmbedBatch_WithNullTexts_ThrowsArgumentNullException() { - // Arrange var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2"); - - // Act & Assert - Assert.Throws(() => model.Embed(string.Empty)); - } - - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithValidTexts_ReturnsMatrixOfCorrectDimensions() - { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var texts = new List { "First text", "Second text", "Third text" }; - - // Act - var embeddings = model.EmbedBatch(texts); - - // Assert - Assert.NotNull(embeddings); - Assert.Equal(3, embeddings.Rows); - Assert.Equal(768, embeddings.Columns); - } - - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithNullTexts_ThrowsArgumentNullException() - { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2"); - - // Act & Assert Assert.Throws(() => model.EmbedBatch(null)); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithEmptyCollection_ThrowsArgumentException() + [Fact] + public void EmbedBatch_WithEmptyCollection_ThrowsArgumentException() { - // Arrange var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2"); - var texts = new List(); - - // Act & Assert - Assert.Throws(() => model.EmbedBatch(texts)); + Assert.Throws(() => model.EmbedBatch(new List())); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_ProducesSameEmbeddingsAsIndividualCalls() - { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var texts = new List { "First", "Second", "Third" }; - - // Act - var batchEmbeddings = model.EmbedBatch(texts); - var individualEmbeddings = texts.Select(t => model.Embed(t)).ToList(); - - // Assert - for (int i = 0; i < texts.Count; i++) - { - for (int j = 0; j < model.EmbeddingDimension; j++) - { - Assert.Equal(individualEmbeddings[i][j], batchEmbeddings[i, j], 10); - } - } - } + // ────────── Failure surfaces (no fake fallback) ────────── - [Fact(Timeout = 60000)] - public async Task Embed_WithFloatType_WorksCorrectly() + [Fact] + public void Embed_WhenApiReturnsError_ThrowsInsteadOfFakeVector() { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var text = "Test with float type"; + var handler = new StubHandler("{\"error\":\"model loading\"}", HttpStatusCode.ServiceUnavailable); + var model = ModelWith(handler, 3); - // Act - var embedding = model.Embed(text); - - // Assert - Assert.NotNull(embedding); - Assert.Equal(768, embedding.Length); - - // Check normalization - var magnitude = 0.0f; - for (int i = 0; i < embedding.Length; i++) - { - magnitude += embedding[i] * embedding[i]; - } - magnitude = (float)Math.Sqrt(magnitude); - Assert.Equal(1.0f, magnitude, 5); + var ex = Assert.Throws(() => model.Embed("hello")); + Assert.Contains("HuggingFace API request failed", ex.Message); } - [Fact(Timeout = 60000)] - public async Task Embed_WithCustomDimension_ReturnsCorrectSize() + [Fact] + public void EmbedBatch_WhenApiReturnsError_ThrowsInsteadOfFakeVector() { - // Arrange - var customDimension = 384; - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", customDimension, 512); - var text = "Testing custom dimension"; - - // Act - var embedding = model.Embed(text); + var handler = new StubHandler("boom", HttpStatusCode.InternalServerError); + var model = ModelWith(handler, 3); - // Assert - Assert.Equal(customDimension, embedding.Length); + Assert.Throws(() => + model.EmbedBatch(new List { "a", "b" })); } - [Fact(Timeout = 60000)] - public async Task Embed_Deterministic_MultipleInstances() - { - // Arrange - var model1 = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var model2 = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var text = "Determinism test"; - - // Act - var embedding1 = model1.Embed(text); - var embedding2 = model2.Embed(text); + // ────────── Success path via mocked handler ────────── - // Assert - for (int i = 0; i < embedding1.Length; i++) - { - Assert.Equal(embedding1[i], embedding2[i], 10); - } - } - - [Fact(Timeout = 60000)] - public async Task EmbedBatch_AllRowsAreNormalized() + [Fact] + public void Embed_WithMockedSuccess_ParsesVectorAndPostsToInferenceEndpoint() { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var texts = new List { "First", "Second", "Third" }; + var handler = new StubHandler("[[0.1,0.2,0.3]]"); + var model = ModelWith(handler, 3); - // Act - var embeddings = model.EmbedBatch(texts); + var embedding = model.Embed("hello world"); - // Assert - for (int i = 0; i < embeddings.Rows; i++) - { - var magnitude = 0.0; - for (int j = 0; j < embeddings.Columns; j++) - { - magnitude += embeddings[i, j] * embeddings[i, j]; - } - magnitude = Math.Sqrt(magnitude); - Assert.Equal(1.0, magnitude, 5); - } + Assert.Equal(3, embedding.Length); + Assert.Equal(0.3, embedding[2], 6); + Assert.Contains("api-inference.huggingface.co/pipeline/feature-extraction/sentence-transformers/all-MiniLM-L6-v2", + handler.LastRequestUri.ToString()); } - [Fact(Timeout = 60000)] - public async Task Embed_WithLongText_ReturnsEmbedding() + [Fact] + public async Task EmbedBatchAsync_WithMockedSuccess_ReturnsMatrix() { - // Arrange - var model = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "test-api-key", 768, 512); - var longText = string.Join(" ", Enumerable.Repeat("word", 500)); - - // Act - var embedding = model.Embed(longText); + var handler = new StubHandler("[[0.1,0.2,0.3],[0.4,0.5,0.6]]"); + var model = ModelWith(handler, 3); - // Assert - Assert.NotNull(embedding); - Assert.Equal(768, embedding.Length); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithDifferentModelNames_CreatesInstances() - { - // Arrange & Act - var miniModel = new HuggingFaceEmbeddingModel("sentence-transformers/all-MiniLM-L6-v2", "", 384, 512); - var mpnetModel = new HuggingFaceEmbeddingModel("sentence-transformers/all-mpnet-base-v2", "", 768, 512); - var bertModel = new HuggingFaceEmbeddingModel("bert-base-uncased", "", 768, 512); + var matrix = await model.EmbedBatchAsync(new List { "a", "b" }); - // Assert - Assert.NotNull(miniModel); - Assert.NotNull(mpnetModel); - Assert.NotNull(bertModel); - Assert.Equal(384, miniModel.EmbeddingDimension); - Assert.Equal(768, mpnetModel.EmbeddingDimension); - Assert.Equal(768, bertModel.EmbeddingDimension); + Assert.Equal(2, matrix.Rows); + Assert.Equal(3, matrix.Columns); + Assert.Equal(0.5, matrix[1, 1], 6); } } } diff --git a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/MultiModalEmbeddingModelTests.cs b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/MultiModalEmbeddingModelTests.cs index 44c15e15a0..c5e91adf29 100644 --- a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/MultiModalEmbeddingModelTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/MultiModalEmbeddingModelTests.cs @@ -19,8 +19,8 @@ private string CreateTempImageFile() return tempFile; } - [Fact(Timeout = 60000)] - public async Task Constructor_WithValidParameters_CreatesInstance() + [Fact] + public void Constructor_WithValidParameters_CreatesInstance() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); @@ -29,65 +29,63 @@ public async Task Constructor_WithValidParameters_CreatesInstance() Assert.Equal(512, model.MaxTokens); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullModelPath_ThrowsArgumentNullException() + [Fact] + public void Constructor_WithNullModelPath_ThrowsArgumentNullException() { Assert.Throws(() => new MultiModalEmbeddingModel(null, true, 512)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithValidText_ThrowsFileNotFoundForMissingModel() + // ────────── Text path: missing ONNX model throws (no fake vector) ────────── + + [Fact] + public void Embed_WithValidText_ThrowsFileNotFoundForMissingModel() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - Assert.Throws(() => model.Embed("This is a test sentence.")); } - [Fact(Timeout = 60000)] - public async Task Embed_WithNormalization_ThrowsFileNotFoundForMissingModel() + [Fact] + public void Embed_WithNullText_ThrowsArgumentException() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - - Assert.Throws(() => model.Embed("Test normalization")); + Assert.Throws(() => model.Embed(null)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithoutNormalization_ThrowsFileNotFoundForMissingModel() + [Fact] + public void Embed_WithEmptyText_ThrowsArgumentException() { - var model = new MultiModalEmbeddingModel("test-model-path.onnx", false, 512); - - Assert.Throws(() => model.Embed("Test no normalization")); + var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); + Assert.Throws(() => model.Embed(string.Empty)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithNullText_ThrowsArgumentException() + [Fact] + public void EmbedBatch_WithValidTexts_ThrowsFileNotFoundForMissingModel() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - - Assert.Throws(() => model.Embed(null)); + var texts = new List { "First text", "Second text", "Third text" }; + Assert.Throws(() => model.EmbedBatch(texts)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithEmptyText_ThrowsArgumentException() + [Fact] + public void MaxTokens_ReturnsCorrectValue() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - - Assert.Throws(() => model.Embed(string.Empty)); + Assert.Equal(512, model.MaxTokens); } - [Fact(Timeout = 60000)] - public async Task EmbedImage_WithValidImagePath_ReturnsVectorOfCorrectDimension() + // ────────── Image path: no fake hash-based vector — throws clearly ────────── + + [Fact] + public void EmbedImage_WithValidImagePath_ThrowsNotSupported() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); var tempImageFile = CreateTempImageFile(); try { - var embedding = model.EmbedImage(tempImageFile); - - Assert.NotNull(embedding); - Assert.Equal(512, embedding.Length); + var ex = Assert.Throws(() => model.EmbedImage(tempImageFile)); + Assert.Contains("Image embedding is not supported", ex.Message); } finally { @@ -96,50 +94,32 @@ public async Task EmbedImage_WithValidImagePath_ReturnsVectorOfCorrectDimension( } } - [Fact(Timeout = 60000)] - public async Task EmbedImage_WithNullImagePath_ThrowsArgumentException() + [Fact] + public void EmbedImage_WithNullImagePath_ThrowsArgumentException() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - var exception = Assert.Throws(() => model.EmbedImage(null)); Assert.Contains("Image path cannot be null or whitespace", exception.Message); } - [Fact(Timeout = 60000)] - public async Task EmbedImage_WithEmptyImagePath_ThrowsArgumentException() + [Fact] + public void EmbedImage_WithEmptyImagePath_ThrowsArgumentException() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - var exception = Assert.Throws(() => model.EmbedImage(string.Empty)); Assert.Contains("Image path cannot be null or whitespace", exception.Message); } - [Fact(Timeout = 60000)] - public async Task EmbedImage_WithNonExistentImagePath_ThrowsFileNotFoundException() - { - var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - var nonExistentPath = Path.Combine("non", "existent", "image.jpg"); - - Assert.Throws(() => model.EmbedImage(nonExistentPath)); - } - - [Fact(Timeout = 60000)] - public async Task EmbedImage_WithNormalization_ReturnsNormalizedVector() + [Fact] + public void EmbedImageBatch_WhenEnumerated_ThrowsNotSupported() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); var tempImageFile = CreateTempImageFile(); try { - var embedding = model.EmbedImage(tempImageFile); - - var magnitude = 0.0; - for (int i = 0; i < embedding.Length; i++) - { - magnitude += embedding[i] * embedding[i]; - } - magnitude = Math.Sqrt(magnitude); - Assert.Equal(1.0, magnitude, 5); + var imagePaths = new List { tempImageFile }; + Assert.Throws(() => model.EmbedImageBatch(imagePaths).ToList()); } finally { @@ -148,124 +128,11 @@ public async Task EmbedImage_WithNormalization_ReturnsNormalizedVector() } } - [Fact(Timeout = 60000)] - public async Task EmbedImage_WithSameImageTwice_ReturnsSameEmbedding() + [Fact] + public void EmbedImageBatch_WithNullImagePaths_ThrowsArgumentNullException() { var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - var tempImageFile = CreateTempImageFile(); - - try - { - var embedding1 = model.EmbedImage(tempImageFile); - var embedding2 = model.EmbedImage(tempImageFile); - - for (int i = 0; i < embedding1.Length; i++) - { - Assert.Equal(embedding1[i], embedding2[i], 10); - } - } - finally - { - if (File.Exists(tempImageFile)) - File.Delete(tempImageFile); - } - } - - [Fact(Timeout = 60000)] - public async Task EmbedImageBatch_WithValidImagePaths_ReturnsCorrectNumberOfEmbeddings() - { - var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - var tempImageFile1 = CreateTempImageFile(); - var tempImageFile2 = CreateTempImageFile(); - var tempImageFile3 = CreateTempImageFile(); - - try - { - var imagePaths = new List { tempImageFile1, tempImageFile2, tempImageFile3 }; - - var embeddings = model.EmbedImageBatch(imagePaths).ToList(); - - Assert.Equal(3, embeddings.Count); - foreach (var embedding in embeddings) - { - Assert.Equal(512, embedding.Length); - } - } - finally - { - if (File.Exists(tempImageFile1)) File.Delete(tempImageFile1); - if (File.Exists(tempImageFile2)) File.Delete(tempImageFile2); - if (File.Exists(tempImageFile3)) File.Delete(tempImageFile3); - } - } - - [Fact(Timeout = 60000)] - public async Task EmbedImageBatch_WithNullImagePaths_ThrowsArgumentNullException() - { - var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - Assert.Throws(() => model.EmbedImageBatch(null)); } - - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithValidTexts_ThrowsFileNotFoundForMissingModel() - { - var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - var texts = new List { "First text", "Second text", "Third text" }; - - Assert.Throws(() => model.EmbedBatch(texts)); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithFloatType_ThrowsFileNotFoundForMissingModel() - { - var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - - Assert.Throws(() => model.Embed("Test with float type")); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithCustomDimension_ThrowsFileNotFoundForMissingModel() - { - var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 768); - - Assert.Throws(() => model.Embed("Testing custom dimension")); - } - - [Fact(Timeout = 60000)] - public async Task EmbedImage_WithCustomDimension_ReturnsCorrectSize() - { - var customDimension = 768; - var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, customDimension); - var tempImageFile = CreateTempImageFile(); - - try - { - var embedding = model.EmbedImage(tempImageFile); - - Assert.Equal(customDimension, embedding.Length); - } - finally - { - if (File.Exists(tempImageFile)) - File.Delete(tempImageFile); - } - } - - [Fact(Timeout = 60000)] - public async Task Embed_Deterministic_MultipleInstances_ThrowsFileNotFoundForMissingModel() - { - var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - - Assert.Throws(() => model.Embed("Determinism test")); - } - - [Fact(Timeout = 60000)] - public async Task MaxTokens_ReturnsCorrectValue() - { - var model = new MultiModalEmbeddingModel("test-model-path.onnx", true, 512); - - Assert.Equal(512, model.MaxTokens); - } } } diff --git a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/OpenAIEmbeddingModelTests.cs b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/OpenAIEmbeddingModelTests.cs index a8bad83ff7..355719c551 100644 --- a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/OpenAIEmbeddingModelTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/OpenAIEmbeddingModelTests.cs @@ -2,395 +2,184 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Net; +using System.Net.Http; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.EmbeddingModels; +using Newtonsoft.Json.Linq; using Xunit; -using System.Threading.Tasks; namespace AiDotNetTests.UnitTests.RAG.Embeddings { public class OpenAIEmbeddingModelTests { - [Fact(Timeout = 60000)] - public async Task Constructor_WithValidParameters_CreatesInstance() + /// Captures the outbound request and returns a canned response / status. + private sealed class StubHandler : HttpMessageHandler { - // Arrange & Act - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); + private readonly HttpStatusCode _status; + private readonly string _responseBody; + public string LastRequestBody { get; private set; } = ""; + public Uri LastRequestUri { get; private set; } - // Assert - Assert.NotNull(model); - Assert.Equal(1536, model.EmbeddingDimension); - Assert.Equal(8191, model.MaxTokens); + public StubHandler(string responseBody, HttpStatusCode status = HttpStatusCode.OK) + { + _responseBody = responseBody; + _status = status; + } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + LastRequestUri = request.RequestUri; + if (request.Content is not null) + LastRequestBody = await request.Content.ReadAsStringAsync().ConfigureAwait(false); + + return new HttpResponseMessage(_status) + { + Content = new StringContent(_responseBody ?? "") + }; + } } - [Fact(Timeout = 60000)] - public async Task Constructor_WithDefaultParameters_CreatesInstance() - { - // Arrange & Act - var model = new OpenAIEmbeddingModel("test-api-key"); + private static OpenAIEmbeddingModel ModelWith(StubHandler handler, int dimension) + => new("test-api-key", "text-embedding-ada-002", dimension, 8191, new HttpClient(handler)); + + // ────────── Constructor validation ────────── - // Assert + [Fact] + public void Constructor_WithValidParameters_CreatesInstance() + { + var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); Assert.NotNull(model); Assert.Equal(1536, model.EmbeddingDimension); Assert.Equal(8191, model.MaxTokens); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullApiKey_ThrowsArgumentException() + [Fact] + public void Constructor_WithNullApiKey_ThrowsArgumentException() { - // Arrange & Act & Assert var exception = Assert.Throws(() => new OpenAIEmbeddingModel(null, "text-embedding-ada-002")); Assert.Contains("API key cannot be empty", exception.Message); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithEmptyApiKey_ThrowsArgumentException() + [Fact] + public void Constructor_WithEmptyApiKey_ThrowsArgumentException() { - // Arrange & Act & Assert var exception = Assert.Throws(() => new OpenAIEmbeddingModel("", "text-embedding-ada-002")); Assert.Contains("API key cannot be empty", exception.Message); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithWhitespaceApiKey_ThrowsArgumentException() - { - // Arrange & Act & Assert - var exception = Assert.Throws(() => - new OpenAIEmbeddingModel(" ", "text-embedding-ada-002")); - Assert.Contains("API key cannot be empty", exception.Message); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullModelName_ThrowsArgumentException() + [Fact] + public void Constructor_WithNullModelName_ThrowsArgumentException() { - // Arrange & Act & Assert var exception = Assert.Throws(() => new OpenAIEmbeddingModel("test-api-key", null)); Assert.Contains("Model name cannot be empty", exception.Message); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithEmptyModelName_ThrowsArgumentException() - { - // Arrange & Act & Assert - var exception = Assert.Throws(() => - new OpenAIEmbeddingModel("test-api-key", "")); - Assert.Contains("Model name cannot be empty", exception.Message); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithZeroDimension_ThrowsArgumentException() + [Fact] + public void Constructor_WithZeroDimension_ThrowsArgumentException() { - // Arrange & Act & Assert var exception = Assert.Throws(() => new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 0, 8191)); Assert.Contains("Dimension must be positive", exception.Message); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithNegativeDimension_ThrowsArgumentException() + [Fact] + public void Constructor_WithNegativeMaxTokens_ThrowsArgumentException() { - // Arrange & Act & Assert - var exception = Assert.Throws(() => - new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", -1, 8191)); - Assert.Contains("Dimension must be positive", exception.Message); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithZeroMaxTokens_ThrowsArgumentException() - { - // Arrange & Act & Assert - var exception = Assert.Throws(() => - new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 0)); - Assert.Contains("Max tokens must be positive", exception.Message); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithNegativeMaxTokens_ThrowsArgumentException() - { - // Arrange & Act & Assert var exception = Assert.Throws(() => new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, -1)); Assert.Contains("Max tokens must be positive", exception.Message); } - [Fact(Timeout = 60000)] - public async Task Embed_WithValidText_ReturnsVectorOfCorrectDimension() - { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var text = "This is a test sentence."; - - // Act - var embedding = model.Embed(text); - - // Assert - Assert.NotNull(embedding); - Assert.Equal(1536, embedding.Length); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithSameTextTwice_ReturnsSameEmbedding() - { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var text = "Hello world"; - - // Act - var embedding1 = model.Embed(text); - var embedding2 = model.Embed(text); - - // Assert - for (int i = 0; i < embedding1.Length; i++) - { - Assert.Equal(embedding1[i], embedding2[i], 10); - } - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithDifferentTexts_ReturnsDifferentEmbeddings() - { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var text1 = "Hello world"; - var text2 = "Goodbye world"; - - // Act - var embedding1 = model.Embed(text1); - var embedding2 = model.Embed(text2); - - // Assert - var hasDifference = false; - for (int i = 0; i < embedding1.Length; i++) - { - if (Math.Abs(embedding1[i] - embedding2[i]) > 1e-10) - { - hasDifference = true; - break; - } - } - Assert.True(hasDifference, "Embeddings for different texts should be different"); - } - - [Fact(Timeout = 60000)] - public async Task Embed_ReturnsNormalizedVector() - { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var text = "Test normalization"; - - // Act - var embedding = model.Embed(text); + // ────────── Input validation (no network) ────────── - // Assert - var magnitude = 0.0; - for (int i = 0; i < embedding.Length; i++) - { - magnitude += embedding[i] * embedding[i]; - } - magnitude = Math.Sqrt(magnitude); - Assert.Equal(1.0, magnitude, 5); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithNullText_ThrowsArgumentException() + [Fact] + public void Embed_WithNullText_ThrowsArgumentException() { - // Arrange var model = new OpenAIEmbeddingModel("test-api-key"); - - // Act & Assert Assert.Throws(() => model.Embed(null)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithEmptyText_ThrowsArgumentException() + [Fact] + public void Embed_WithEmptyText_ThrowsArgumentException() { - // Arrange var model = new OpenAIEmbeddingModel("test-api-key"); - - // Act & Assert Assert.Throws(() => model.Embed(string.Empty)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithWhitespaceText_ThrowsArgumentException() - { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key"); - - // Act & Assert - Assert.Throws(() => model.Embed(" ")); - } - - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithValidTexts_ReturnsMatrixOfCorrectDimensions() - { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var texts = new List { "First text", "Second text", "Third text" }; - - // Act - var embeddings = model.EmbedBatch(texts); - - // Assert - Assert.NotNull(embeddings); - Assert.Equal(3, embeddings.Rows); - Assert.Equal(1536, embeddings.Columns); - } - - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithNullTexts_ThrowsArgumentNullException() + [Fact] + public void EmbedBatch_WithNullTexts_ThrowsArgumentNullException() { - // Arrange var model = new OpenAIEmbeddingModel("test-api-key"); - - // Act & Assert Assert.Throws(() => model.EmbedBatch(null)); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithEmptyCollection_ThrowsArgumentException() + [Fact] + public void EmbedBatch_WithEmptyCollection_ThrowsArgumentException() { - // Arrange var model = new OpenAIEmbeddingModel("test-api-key"); - var texts = new List(); - - // Act & Assert - Assert.Throws(() => model.EmbedBatch(texts)); + Assert.Throws(() => model.EmbedBatch(new List())); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_ProducesSameEmbeddingsAsIndividualCalls() - { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var texts = new List { "First", "Second", "Third" }; - - // Act - var batchEmbeddings = model.EmbedBatch(texts); - var individualEmbeddings = texts.Select(t => model.Embed(t)).ToList(); - - // Assert - for (int i = 0; i < texts.Count; i++) - { - for (int j = 0; j < model.EmbeddingDimension; j++) - { - Assert.Equal(individualEmbeddings[i][j], batchEmbeddings[i, j], 10); - } - } - } + // ────────── Failure surfaces (no fake fallback) ────────── - [Fact(Timeout = 60000)] - public async Task Embed_WithFloatType_WorksCorrectly() + [Fact] + public void Embed_WhenApiReturnsError_ThrowsInsteadOfFakeVector() { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var text = "Test with float type"; - - // Act - var embedding = model.Embed(text); + var handler = new StubHandler("{\"error\":\"unauthorized\"}", HttpStatusCode.Unauthorized); + var model = ModelWith(handler, 3); - // Assert - Assert.NotNull(embedding); - Assert.Equal(1536, embedding.Length); - - // Check normalization - var magnitude = 0.0f; - for (int i = 0; i < embedding.Length; i++) - { - magnitude += embedding[i] * embedding[i]; - } - magnitude = (float)Math.Sqrt(magnitude); - Assert.Equal(1.0f, magnitude, 5); + var ex = Assert.Throws(() => model.Embed("hello")); + Assert.Contains("OpenAI API request failed", ex.Message); } - [Fact(Timeout = 60000)] - public async Task Embed_WithCustomDimension_ReturnsCorrectSize() + [Fact] + public void EmbedBatch_WhenApiReturnsError_ThrowsInsteadOfFakeVector() { - // Arrange - var customDimension = 512; - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-3-small", customDimension, 8191); - var text = "Testing custom dimension"; - - // Act - var embedding = model.Embed(text); + var handler = new StubHandler("boom", HttpStatusCode.InternalServerError); + var model = ModelWith(handler, 3); - // Assert - Assert.Equal(customDimension, embedding.Length); + Assert.Throws(() => + model.EmbedBatch(new List { "a", "b" })); } - [Fact(Timeout = 60000)] - public async Task Embed_Deterministic_MultipleInstances() - { - // Arrange - var model1 = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var model2 = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var text = "Determinism test"; + // ────────── Success path via mocked handler ────────── - // Act - var embedding1 = model1.Embed(text); - var embedding2 = model2.Embed(text); - - // Assert - for (int i = 0; i < embedding1.Length; i++) - { - Assert.Equal(embedding1[i], embedding2[i], 10); - } - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithLongText_ReturnsEmbedding() + [Fact] + public void Embed_WithMockedSuccess_ParsesVectorAndPostsToOpenAIEndpoint() { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var longText = string.Join(" ", Enumerable.Repeat("word", 1000)); + var handler = new StubHandler("{\"data\":[{\"embedding\":[0.1,0.2,0.3]}]}"); + var model = ModelWith(handler, 3); - // Act - var embedding = model.Embed(longText); + var embedding = model.Embed("hello world"); - // Assert - Assert.NotNull(embedding); - Assert.Equal(1536, embedding.Length); - } + Assert.Equal(3, embedding.Length); + Assert.Equal(0.1, embedding[0], 6); + Assert.Equal(0.2, embedding[1], 6); + Assert.Equal(0.3, embedding[2], 6); + Assert.Equal("https://api.openai.com/v1/embeddings", handler.LastRequestUri.ToString()); - [Fact(Timeout = 60000)] - public async Task Constructor_WithDifferentModelNames_CreatesInstances() - { - // Arrange & Act - var adaModel = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var smallModel = new OpenAIEmbeddingModel("test-api-key", "text-embedding-3-small", 1536, 8191); - var largeModel = new OpenAIEmbeddingModel("test-api-key", "text-embedding-3-large", 3072, 8191); - - // Assert - Assert.NotNull(adaModel); - Assert.NotNull(smallModel); - Assert.NotNull(largeModel); - Assert.Equal(3072, largeModel.EmbeddingDimension); + var body = JObject.Parse(handler.LastRequestBody); + Assert.Equal("text-embedding-ada-002", (string)body["model"]); + Assert.Equal("hello world", (string)body["input"][0]); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_AllRowsAreNormalized() + [Fact] + public async Task EmbedBatchAsync_WithMockedSuccess_ReturnsMatrix() { - // Arrange - var model = new OpenAIEmbeddingModel("test-api-key", "text-embedding-ada-002", 1536, 8191); - var texts = new List { "First", "Second", "Third" }; + var handler = new StubHandler("{\"data\":[{\"embedding\":[0.1,0.2,0.3]},{\"embedding\":[0.4,0.5,0.6]}]}"); + var model = ModelWith(handler, 3); - // Act - var embeddings = model.EmbedBatch(texts); + var matrix = await model.EmbedBatchAsync(new List { "a", "b" }); - // Assert - for (int i = 0; i < embeddings.Rows; i++) - { - var magnitude = 0.0; - for (int j = 0; j < embeddings.Columns; j++) - { - magnitude += embeddings[i, j] * embeddings[i, j]; - } - magnitude = Math.Sqrt(magnitude); - Assert.Equal(1.0, magnitude, 5); - } + Assert.Equal(2, matrix.Rows); + Assert.Equal(3, matrix.Columns); + Assert.Equal(0.4, matrix[1, 0], 6); } } } diff --git a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/VoyageAIEmbeddingModelTests.cs b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/VoyageAIEmbeddingModelTests.cs index e2849e1466..1588548983 100644 --- a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/VoyageAIEmbeddingModelTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/VoyageAIEmbeddingModelTests.cs @@ -1,320 +1,201 @@ #nullable disable using System; using System.Collections.Generic; -using System.IO; using System.Linq; +using System.Net; +using System.Net.Http; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.EmbeddingModels; +using Newtonsoft.Json.Linq; using Xunit; -using System.Threading.Tasks; namespace AiDotNetTests.UnitTests.RAG.Embeddings { public class VoyageAIEmbeddingModelTests { - [Fact(Timeout = 60000)] - public async Task Constructor_WithValidParameters_CreatesInstance() - { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - - Assert.NotNull(model); - Assert.Equal(1024, model.EmbeddingDimension); - Assert.Equal(16000, model.MaxTokens); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullApiKey_ThrowsArgumentNullException() - { - Assert.Throws(() => - new VoyageAIEmbeddingModel(null, "voyage-model-path.onnx", "document", 1024)); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullModel_ThrowsArgumentNullException() - { - Assert.Throws(() => - new VoyageAIEmbeddingModel("test-api-key", null, "document", 1024)); - } - - [Fact(Timeout = 60000)] - public async Task Constructor_WithNullInputType_ThrowsArgumentNullException() - { - Assert.Throws(() => - new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", null, 1024)); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithValidText_ThrowsFileNotFoundForMissingModel() - { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - - Assert.Throws(() => model.Embed("This is a test sentence.")); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithSameTextTwice_ThrowsFileNotFoundForMissingModel() - { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - - Assert.Throws(() => model.Embed("Hello world")); - Assert.Throws(() => model.Embed("Hello world")); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithSingleText_ThrowsFileNotFoundForMissingModel() + private sealed class StubHandler : HttpMessageHandler { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); + private readonly HttpStatusCode _status; + private readonly string _responseBody; + public string LastRequestBody { get; private set; } = ""; + public Uri LastRequestUri { get; private set; } - Assert.Throws(() => model.Embed("Hello world")); - } + public StubHandler(string responseBody, HttpStatusCode status = HttpStatusCode.OK) + { + _responseBody = responseBody; + _status = status; + } - [Fact(Timeout = 60000)] - public async Task Embed_ReturnsNormalizedVector_ThrowsFileNotFoundForMissingModel() - { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + LastRequestUri = request.RequestUri; + if (request.Content is not null) + LastRequestBody = await request.Content.ReadAsStringAsync().ConfigureAwait(false); - Assert.Throws(() => model.Embed("Test normalization")); + return new HttpResponseMessage(_status) { Content = new StringContent(_responseBody ?? "") }; + } } - [Fact(Timeout = 60000)] - public async Task Embed_WithNullText_ThrowsArgumentException() - { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); + private static VoyageAIEmbeddingModel ModelWith(StubHandler handler, int dimension) + => new("test-api-key", "voyage-3", "document", dimension, new HttpClient(handler)); - Assert.Throws(() => model.Embed(null)); - } + // ────────── Constructor validation ────────── - [Fact(Timeout = 60000)] - public async Task Embed_WithEmptyText_ThrowsArgumentException() + [Fact] + public void Constructor_WithValidParameters_CreatesInstance() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - - Assert.Throws(() => model.Embed(string.Empty)); + var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-3", "document", 1024); + Assert.NotNull(model); + Assert.Equal(1024, model.EmbeddingDimension); + Assert.Equal(16000, model.MaxTokens); } - [Fact(Timeout = 60000)] - public async Task Embed_WithWhitespaceText_ThrowsArgumentException() + [Fact] + public void Constructor_WithNullApiKey_ThrowsArgumentException() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - - Assert.Throws(() => model.Embed(" ")); + var ex = Assert.Throws(() => + new VoyageAIEmbeddingModel(null, "voyage-3", "document", 1024)); + Assert.Contains("API key cannot be empty", ex.Message); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithValidTexts_ThrowsFileNotFoundForMissingModel() + [Fact] + public void Constructor_WithNullModel_ThrowsArgumentException() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - var texts = new List { "First text", "Second text", "Third text" }; - - Assert.Throws(() => model.EmbedBatch(texts)); + var ex = Assert.Throws(() => + new VoyageAIEmbeddingModel("test-api-key", null, "document", 1024)); + Assert.Contains("Model cannot be empty", ex.Message); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithNullTexts_ThrowsArgumentNullException() + [Fact] + public void Constructor_WithNullInputType_ThrowsArgumentException() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - - Assert.Throws(() => model.EmbedBatch(null)); + var ex = Assert.Throws(() => + new VoyageAIEmbeddingModel("test-api-key", "voyage-3", null, 1024)); + Assert.Contains("Input type cannot be empty", ex.Message); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_WithEmptyCollection_ThrowsArgumentException() + [Fact] + public void Constructor_WithZeroDimension_ThrowsArgumentException() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - var texts = new List(); - - Assert.Throws(() => model.EmbedBatch(texts)); + var ex = Assert.Throws(() => + new VoyageAIEmbeddingModel("test-api-key", "voyage-3", "document", 0)); + Assert.Contains("Dimension must be positive", ex.Message); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_ProducesSameEmbeddingsAsIndividualCalls_ThrowsFileNotFoundForMissingModel() + [Fact] + public void MaxTokens_ReturnsCorrectValue() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - var texts = new List { "First", "Second", "Third" }; - - Assert.Throws(() => model.EmbedBatch(texts)); + var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-3", "document", 1024); + Assert.Equal(16000, model.MaxTokens); } - [Fact(Timeout = 60000)] - public async Task Embed_WithFloatType_ThrowsFileNotFoundForMissingModel() - { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); + // ────────── Input validation (no network) ────────── - Assert.Throws(() => model.Embed("Test with float type")); - } - - [Fact(Timeout = 60000)] - public async Task Embed_WithCustomDimension_ThrowsFileNotFoundForMissingModel() + [Fact] + public void Embed_WithNullText_ThrowsArgumentException() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 512); - - Assert.Throws(() => model.Embed("Testing custom dimension")); + var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-3", "document", 1024); + Assert.Throws(() => model.Embed(null)); } - [Fact(Timeout = 60000)] - public async Task Embed_Deterministic_MultipleInstances_ThrowsFileNotFoundForMissingModel() + [Fact] + public void Embed_WithEmptyText_ThrowsArgumentException() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - - Assert.Throws(() => model.Embed("Determinism test")); + var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-3", "document", 1024); + Assert.Throws(() => model.Embed(string.Empty)); } - [Fact(Timeout = 60000)] - public async Task EmbedBatch_AllRowsAreNormalized_ThrowsFileNotFoundForMissingModel() + [Fact] + public void EmbedBatch_WithNullTexts_ThrowsArgumentNullException() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - var texts = new List { "First", "Second", "Third" }; - - Assert.Throws(() => model.EmbedBatch(texts)); + var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-3", "document", 1024); + Assert.Throws(() => model.EmbedBatch(null)); } - [Fact(Timeout = 60000)] - public async Task Embed_WithLongText_ThrowsFileNotFoundForMissingModel() + [Fact] + public void EmbedBatch_WithEmptyCollection_ThrowsArgumentException() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - var longText = string.Join(" ", Enumerable.Repeat("word", 2000)); - - Assert.Throws(() => model.Embed(longText)); + var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-3", "document", 1024); + Assert.Throws(() => model.EmbedBatch(new List())); } - [Fact(Timeout = 60000)] - public async Task Constructor_WithDifferentInputTypes_CreatesInstances() - { - var documentModel = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); - var queryModel = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "query", 1024); + // ────────── Failure surfaces (no fake fallback) ────────── - Assert.NotNull(documentModel); - Assert.NotNull(queryModel); - } - - [Fact(Timeout = 60000)] - public async Task MaxTokens_ReturnsCorrectValue() + [Fact] + public void Embed_WhenApiReturnsError_ThrowsInsteadOfFakeVector() { - var model = new VoyageAIEmbeddingModel("test-api-key", "voyage-model-path.onnx", "document", 1024); + var handler = new StubHandler("{\"detail\":\"unauthorized\"}", HttpStatusCode.Unauthorized); + var model = ModelWith(handler, 3); - Assert.Equal(16000, model.MaxTokens); + var ex = Assert.Throws(() => model.Embed("hello")); + Assert.Contains("Voyage AI API request failed", ex.Message); } - // ────────── Success-path tests via mock embedding model ────────── - // These test the EmbeddingModelBase behavior (dimension, normalization, batch consistency) - // without requiring an actual ONNX model file. - - [Fact(Timeout = 60000)] - public async Task MockEmbed_WithValidText_ReturnsVectorOfCorrectDimension() + [Fact] + public void EmbedBatch_WhenApiReturnsError_ThrowsInsteadOfFakeVector() { - using var model = new MockEmbeddingModel(1024); + var handler = new StubHandler("boom", HttpStatusCode.InternalServerError); + var model = ModelWith(handler, 3); - var embedding = model.Embed("Test sentence"); - - Assert.NotNull(embedding); - Assert.Equal(1024, embedding.Length); + Assert.Throws(() => + model.EmbedBatch(new List { "a", "b" })); } - [Fact(Timeout = 60000)] - public async Task MockEmbed_ReturnsNormalizedVector() - { - using var model = new MockEmbeddingModel(512); - - var embedding = model.Embed("Test normalization"); + // ────────── Real Voyage API endpoint/body/response ────────── - var magnitude = 0.0; - for (int i = 0; i < embedding.Length; i++) - magnitude += embedding[i] * embedding[i]; - magnitude = Math.Sqrt(magnitude); - Assert.Equal(1.0, magnitude, 5); - } - - [Fact(Timeout = 60000)] - public async Task MockEmbed_WithSameText_ReturnsSameEmbedding() + [Fact] + public void Embed_PostsToVoyageEndpointWithCorrectBodyAndParsesResponse() { - using var model = new MockEmbeddingModel(384); + var handler = new StubHandler( + "{\"object\":\"list\",\"data\":[{\"object\":\"embedding\",\"embedding\":[0.1,0.2,0.3],\"index\":0}],\"model\":\"voyage-3\"}"); + var model = ModelWith(handler, 3); - var embedding1 = model.Embed("Hello world"); - var embedding2 = model.Embed("Hello world"); + var embedding = model.Embed("hello world"); - for (int i = 0; i < embedding1.Length; i++) - Assert.Equal(embedding1[i], embedding2[i], 10); - } + // Parsed the real response + Assert.Equal(3, embedding.Length); + Assert.Equal(0.1, embedding[0], 6); + Assert.Equal(0.2, embedding[1], 6); + Assert.Equal(0.3, embedding[2], 6); - [Fact(Timeout = 60000)] - public async Task MockEmbed_WithDifferentTexts_ReturnsDifferentEmbeddings() - { - using var model = new MockEmbeddingModel(384); - - var embedding1 = model.Embed("Hello world"); - var embedding2 = model.Embed("Goodbye world"); + // Hit the real Voyage endpoint + Assert.Equal("https://api.voyageai.com/v1/embeddings", handler.LastRequestUri.ToString()); - var hasDifference = false; - for (int i = 0; i < embedding1.Length; i++) - { - if (Math.Abs(embedding1[i] - embedding2[i]) > 1e-10) - { - hasDifference = true; - break; - } - } - Assert.True(hasDifference, "Embeddings for different texts should be different"); + // Correct request body: model, input[], input_type + var body = JObject.Parse(handler.LastRequestBody); + Assert.Equal("voyage-3", (string)body["model"]); + Assert.Equal("document", (string)body["input_type"]); + Assert.Equal("hello world", (string)body["input"][0]); } - [Fact(Timeout = 60000)] - public async Task MockEmbedBatch_ReturnsCorrectDimensions() + [Fact] + public async Task EmbedBatchAsync_PostsAllInputsAndReturnsMatrix() { - using var model = new MockEmbeddingModel(256); - var texts = new List { "First", "Second", "Third" }; - - var embeddings = model.EmbedBatch(texts); + var handler = new StubHandler( + "{\"data\":[{\"embedding\":[0.1,0.2,0.3],\"index\":0},{\"embedding\":[0.4,0.5,0.6],\"index\":1}]}"); + var model = ModelWith(handler, 3); - Assert.Equal(3, embeddings.Rows); - Assert.Equal(256, embeddings.Columns); - } - - [Fact(Timeout = 60000)] - public async Task MockEmbedBatch_ProducesSameAsIndividualCalls() - { - using var model = new MockEmbeddingModel(384); - var texts = new List { "Alpha", "Beta", "Gamma" }; + var matrix = await model.EmbedBatchAsync(new List { "a", "b" }); - var batchEmbeddings = model.EmbedBatch(texts); - var individualEmbeddings = texts.Select(t => model.Embed(t)).ToList(); + Assert.Equal(2, matrix.Rows); + Assert.Equal(3, matrix.Columns); + Assert.Equal(0.4, matrix[1, 0], 6); - for (int i = 0; i < texts.Count; i++) - { - for (int j = 0; j < model.EmbeddingDimension; j++) - Assert.Equal(individualEmbeddings[i][j], batchEmbeddings[i, j], 10); - } + var body = JObject.Parse(handler.LastRequestBody); + Assert.Equal("a", (string)body["input"][0]); + Assert.Equal("b", (string)body["input"][1]); } - /// - /// Deterministic mock embedding model for testing base class behavior - /// without ONNX dependencies. Uses hash-based vector generation. - /// - private sealed class MockEmbeddingModel : AiDotNet.RetrievalAugmentedGeneration.Embeddings.EmbeddingModelBase + [Fact] + public void Constructor_WithDifferentInputTypes_CreatesInstances() { - private readonly int _dimension; - public override int EmbeddingDimension => _dimension; - public override int MaxTokens => 512; - - public MockEmbeddingModel(int dimension) { _dimension = dimension; } + var documentModel = new VoyageAIEmbeddingModel("test-api-key", "voyage-3", "document", 1024); + var queryModel = new VoyageAIEmbeddingModel("test-api-key", "voyage-3", "query", 1024); - protected override Vector EmbedCore(string text) - { - // Use stable FNV-1a hash instead of GetHashCode (randomized per process in .NET 6+) - uint hash = 2166136261; - foreach (char c in text.ToLowerInvariant()) - { - hash ^= c; - hash *= 16777619; - } - - var values = new double[_dimension]; - for (int i = 0; i < _dimension; i++) - values[i] = Math.Sin(hash * 0.0001 + i * 0.1) * 0.5; - - var vec = new Vector(values); - return vec.Normalize(); - } - - protected override void Dispose(bool disposing) => base.Dispose(disposing); + Assert.NotNull(documentModel); + Assert.NotNull(queryModel); } } } From cc70e8a078b7dcfa29d37cf723ac2cf12edec227 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 15:42:34 -0400 Subject: [PATCH 06/25] feat(rag): real Weaviate + Milvus + Azure AI Search vector-store clients Replace the in-memory simulations for WeaviateDocumentStore, MilvusDocumentStore and AzureSearchDocumentStore with real HTTP clients following the Qdrant/Pinecone pattern (HttpClient + Newtonsoft, sync *Core overrides bridging to a private async SendAsync, net471-safe). Adds mocked-HttpMessageHandler unit tests plus one gated integration SkippableFact per store. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../AzureSearchDocumentStore.cs | 650 ++++++++++----- .../DocumentStores/MilvusDocumentStore.cs | 573 +++++++++---- .../DocumentStores/WeaviateDocumentStore.cs | 759 +++++++++++------- .../AzureSearchDocumentStoreTests.cs | 339 ++++++++ .../MilvusDocumentStoreTests.cs | 336 ++++++++ .../WeaviateDocumentStoreTests.cs | 353 ++++++++ 6 files changed, 2389 insertions(+), 621 deletions(-) create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStoreTests.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStoreTests.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStoreTests.cs diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs index f7518ddfb4..c8ddb4ae1d 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs @@ -1,310 +1,556 @@ - using System; using System.Collections.Generic; +using System.Globalization; using System.Linq; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json; +using Newtonsoft.Json.Linq; namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores { /// - /// Azure Cognitive Search-inspired document store with field-based indexing and search capabilities. - /// Provides in-memory simulation of Azure Search features including field-level search and faceted filtering. + /// Azure AI Search document store backed by the real Azure AI Search REST API. /// /// The numeric type for vector operations. + /// + /// + /// This store talks to an Azure AI Search service over HTTP. It manages a single index with a + /// vector field and an HNSW vector-search profile, upserts documents via the mergeOrUpload + /// action, performs vector search with an OData $filter, deletes documents via the + /// delete action and pages the whole index. The original id and content are stored under + /// id and content; the full metadata dictionary is stored as a JSON string under + /// metadata_json (for lossless round-tripping) while each scalar metadata entry is also + /// flattened to a top-level field so it can be used in $filter expressions. + /// + /// For Beginners: Azure AI Search is Microsoft's managed search-and-vector service. + /// This class is a real client for it - every method here makes an HTTP call to your search + /// service. Authentication uses the api-key header and every request carries an + /// api-version query-string parameter. + /// + /// API deviation: Azure indexes have a fixed schema, so metadata keys used in filters + /// must exist as filterable fields in the index. This client creates the index with the base fields + /// (id, content, metadata_json, embedding); flattened metadata fields are + /// still sent on upload and referenced by $filter, but for production filtering the caller + /// should ensure those fields are declared in the index (or pre-create the index). + /// + /// [ComponentType(ComponentType.DocumentStore)] [PipelineStage(PipelineStage.Indexing)] public class AzureSearchDocumentStore : DocumentStoreBase { - private readonly Dictionary> _documents; - private readonly Dictionary>> _invertedIndex; - private readonly string _serviceName; + private const string IdField = "id"; + private const string ContentField = "content"; + private const string MetadataField = "metadata_json"; + private const string EmbeddingField = "embedding"; + private const string VectorProfile = "vprofile"; + private const string VectorAlgorithm = "hnsw-algo"; + + private readonly HttpClient _httpClient; private readonly string _indexName; + private readonly string _apiVersion; + private readonly string _vectorMetric; private int _vectorDimension; + private int _documentCount; + private bool _indexReady; + + /// + /// Gets the number of documents currently stored in the index. + /// + public override int DocumentCount => _documentCount; - public override int DocumentCount => _documents.Count; + /// + /// Gets the dimensionality of vectors stored in this index. + /// public override int VectorDimension => _vectorDimension; - public AzureSearchDocumentStore(string serviceName, string indexName, int initialCapacity = 1000) + /// + /// Gets the name of the Azure AI Search index this store is bound to. + /// + public string IndexName { get; } + + /// + /// Initializes a new instance of the class. + /// + /// The Azure AI Search index name. + /// The service endpoint, e.g. https://myservice.search.windows.net. + /// The Azure AI Search admin API key (sent as the api-key header). + /// The vector distance metric used when creating the index. + /// + /// The vector dimension used to create the index if it does not already exist. + /// When 0 the index must already exist or the dimension is inferred from the first document. + /// + /// The REST API version query-string value. + /// Optional pre-configured (primarily for testing). + /// Optional used to build the client (for testing). + public AzureSearchDocumentStore( + string indexName, + string endpoint, + string apiKey, + DistanceMetricType distanceMetric = DistanceMetricType.Cosine, + int vectorDimension = 0, + string apiVersion = "2023-11-01", + HttpClient? httpClient = null, + HttpMessageHandler? handler = null) { - if (string.IsNullOrWhiteSpace(serviceName)) - throw new ArgumentException("Service name cannot be empty", nameof(serviceName)); if (string.IsNullOrWhiteSpace(indexName)) throw new ArgumentException("Index name cannot be empty", nameof(indexName)); - if (initialCapacity <= 0) - throw new ArgumentException("Initial capacity must be greater than zero", nameof(initialCapacity)); + if (vectorDimension < 0) + throw new ArgumentOutOfRangeException(nameof(vectorDimension), "Vector dimension cannot be negative"); - _serviceName = serviceName; _indexName = indexName; - _documents = new Dictionary>(initialCapacity); - _invertedIndex = new Dictionary>>(); - _vectorDimension = 0; - } - - protected override void AddCore(VectorDocument vectorDocument) - { - if (_documents.Count == 0) + IndexName = indexName; + _apiVersion = string.IsNullOrWhiteSpace(apiVersion) ? "2023-11-01" : apiVersion; + _vectorMetric = MapMetric(distanceMetric); + _vectorDimension = vectorDimension; + _documentCount = 0; + + _httpClient = httpClient ?? (handler != null ? new HttpClient(handler) : new HttpClient()); + if (_httpClient.BaseAddress == null) { - _vectorDimension = vectorDocument.Embedding.Length; + if (string.IsNullOrWhiteSpace(endpoint)) + throw new ArgumentException("Endpoint cannot be empty", nameof(endpoint)); + _httpClient.BaseAddress = new Uri(endpoint); } - _documents[vectorDocument.Document.Id] = vectorDocument; - IndexMetadata(vectorDocument.Document); + if (!string.IsNullOrWhiteSpace(apiKey) && !_httpClient.DefaultRequestHeaders.Contains("api-key")) + _httpClient.DefaultRequestHeaders.Add("api-key", apiKey); + + InitializeIndex(vectorDimension); } - protected override void AddBatchCore(IList> vectorDocuments) + private static string MapMetric(DistanceMetricType metric) { - if (vectorDocuments.Count == 0) - return; - - if (_documents.Count == 0) - { - _vectorDimension = vectorDocuments[0].Embedding.Length; - } - - foreach (var vectorDoc in vectorDocuments) + switch (metric) { - _documents[vectorDoc.Document.Id] = vectorDoc; - IndexMetadata(vectorDoc.Document); + case DistanceMetricType.Cosine: + return "cosine"; + case DistanceMetricType.Euclidean: + return "euclidean"; + default: + throw new NotSupportedException( + $"Metric '{metric}' is not supported by Azure AI Search. Use Cosine or Euclidean (or dotProduct via a pre-created index)."); } } - protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + private string WithApiVersion(string path) { - var candidateIds = GetCandidateIds(metadataFilters); - var scoredDocuments = new List<(Document Document, T Score)>(); + var separator = path.Contains("?") ? "&" : "?"; + return path + separator + "api-version=" + _apiVersion; + } - IEnumerable> candidates; - if (candidateIds != null) + private void InitializeIndex(int requestedDimension) + { + HttpResponseInfo info; + try { - candidates = candidateIds - .Where(id => _documents.ContainsKey(id)) - .Select(id => _documents[id]); + info = SendAsync(HttpMethod.Get, WithApiVersion($"/indexes/{_indexName}"), null).GetAwaiter().GetResult(); } - else + catch (HttpRequestException) { - candidates = _documents.Values; + // Service not reachable at construction time; defer creation to first write. + return; } - var matchingDocuments = candidates - .Where(vectorDoc => MatchesFilters(vectorDoc.Document, metadataFilters)); - - foreach (var vectorDoc in matchingDocuments) + if (info.Status == HttpStatusCode.OK) { - var similarity = StatisticsHelper.CosineSimilarity(queryVector, vectorDoc.Embedding); - scoredDocuments.Add((vectorDoc.Document, similarity)); + _indexReady = true; + var fields = JObject.Parse(info.Body)["fields"] as JArray; + if (fields != null) + { + foreach (var field in fields) + { + if ((string?)field?["name"] == EmbeddingField) + { + var dim = field?["dimensions"]; + if (dim != null && dim.Type != JTokenType.Null) + { + var parsed = Convert.ToInt32(dim, CultureInfo.InvariantCulture); + if (parsed > 0) + _vectorDimension = parsed; + } + } + } + } + return; } - var results = scoredDocuments - .OrderByDescending(x => x.Score) - .Take(topK) - .Select(x => + if (info.Status == HttpStatusCode.NotFound && requestedDimension > 0) + CreateIndex(requestedDimension); + } + + private void CreateIndex(int dimension) + { + var body = new + { + name = _indexName, + fields = new object[] + { + new { name = IdField, type = "Edm.String", key = true, filterable = true, retrievable = true }, + new { name = ContentField, type = "Edm.String", searchable = true, retrievable = true }, + new { name = MetadataField, type = "Edm.String", retrievable = true }, + new + { + name = EmbeddingField, + type = "Collection(Edm.Single)", + searchable = true, + retrievable = false, + dimensions = dimension, + vectorSearchProfile = VectorProfile + } + }, + vectorSearch = new { - x.Document.RelevanceScore = x.Score; - x.Document.HasRelevanceScore = true; - return x.Document; - }) - .ToList(); + algorithms = new object[] + { + new + { + name = VectorAlgorithm, + kind = "hnsw", + hnswParameters = new { metric = _vectorMetric, m = 4, efConstruction = 400, efSearch = 500 } + } + }, + profiles = new object[] + { + new { name = VectorProfile, algorithm = VectorAlgorithm } + } + } + }; - return results; + var info = SendAsync(HttpMethod.Put, WithApiVersion($"/indexes/{_indexName}"), body).GetAwaiter().GetResult(); + EnsureSuccess(info, "create index"); + _vectorDimension = dimension; + _indexReady = true; } - protected override Document? GetByIdCore(string documentId) + private void EnsureIndexForDimension(int dimension) { - return _documents.TryGetValue(documentId, out var vectorDoc) ? vectorDoc.Document : null; + if (_indexReady) + return; + CreateIndex(dimension); } - protected override bool RemoveCore(string documentId) + /// + protected override void AddCore(VectorDocument vectorDocument) { - if (!_documents.TryGetValue(documentId, out var vectorDoc)) - return false; + EnsureIndexForDimension(vectorDocument.Embedding.Length); + if (_vectorDimension == 0) + _vectorDimension = vectorDocument.Embedding.Length; + + UploadActions(new[] { BuildAction(vectorDocument, "mergeOrUpload") }, "upload document"); + _documentCount++; + } + + /// + protected override void AddBatchCore(IList> vectorDocuments) + { + if (vectorDocuments.Count == 0) + return; + + EnsureIndexForDimension(vectorDocuments[0].Embedding.Length); + if (_vectorDimension == 0) + _vectorDimension = vectorDocuments[0].Embedding.Length; + + var actions = vectorDocuments.Select(d => BuildAction(d, "mergeOrUpload")).ToList(); + UploadActions(actions, "batch upload documents"); + _documentCount += vectorDocuments.Count; + } - RemoveFromIndex(vectorDoc.Document); - _documents.Remove(documentId); + private void UploadActions(IEnumerable> actions, string operation) + { + var body = new { value = actions }; + var info = SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/index"), body).GetAwaiter().GetResult(); + EnsureSuccess(info, operation); + } - if (_documents.Count == 0) + private Dictionary BuildAction(VectorDocument vectorDocument, string action) + { + var vector = vectorDocument.Embedding.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + var metadata = vectorDocument.Document.Metadata ?? new Dictionary(); + + var doc = new Dictionary { - _vectorDimension = 0; - } + ["@search.action"] = action, + [IdField] = vectorDocument.Document.Id, + [ContentField] = vectorDocument.Document.Content, + [MetadataField] = JsonConvert.SerializeObject(metadata), + [EmbeddingField] = vector + }; - return true; + foreach (var kvp in metadata) + doc[kvp.Key] = kvp.Value; + + return doc; } - /// - /// Core logic for retrieving all documents in the index. - /// - /// An enumerable of all documents without their vector embeddings. - /// - /// - /// Returns all documents from the Azure Search index in no particular order. - /// Vector embeddings are not included, only document content and metadata. - /// - /// For Beginners: Gets every document in the index. - /// - /// Use cases: - /// - Export all documents for backup - /// - Migrate to a different index or service - /// - Bulk reindexing or analysis - /// - Debugging facet indices - /// - /// Warning: For large indices (> 10K documents), this can use significant memory. - /// In real Azure Search, use continuation tokens for pagination. - /// - /// Example: - /// - /// // Get all documents - /// var allDocs = store.GetAll().ToList(); - /// // Result is available in the returned value - /// - /// // Export to JSON - /// var json = JsonConvert.SerializeObject(allDocs); - /// File.WriteAllText($"{_serviceName}_{_indexName}_export.json", json); - /// - /// - /// - protected override IEnumerable> GetAllCore() + /// + protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) { - return _documents.Values.Select(vd => vd.Document).ToList(); + var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + + var body = new Dictionary + { + ["vectorQueries"] = new object[] + { + new { kind = "vector", vector, fields = EmbeddingField, k = topK } + }, + ["select"] = $"{IdField},{ContentField},{MetadataField}", + ["top"] = topK + }; + + var filter = BuildODataFilter(metadataFilters); + if (filter != null) + body["filter"] = filter; + + var info = SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/search"), body).GetAwaiter().GetResult(); + EnsureSuccess(info, "search"); + + var results = new List>(); + var hits = JObject.Parse(info.Body)["value"] as JArray; + if (hits == null) + return results; + + foreach (var hit in hits) + { + var doc = ParseDocument(hit as JObject); + if (doc == null) + continue; + + var scoreToken = hit?["@search.score"]; + var score = scoreToken != null && scoreToken.Type != JTokenType.Null + ? Convert.ToDouble(scoreToken, CultureInfo.InvariantCulture) + : 0.0; + doc.RelevanceScore = NumOps.FromDouble(score); + doc.HasRelevanceScore = true; + results.Add(doc); + } + + return results; } /// - /// Removes all documents from the index and clears all inverted indices. + /// Translates the metadata filter dictionary into an Azure AI Search OData $filter string. /// /// - /// - /// Clears all documents, field-level inverted indices, and resets the vector dimension to 0. - /// The service and index names remain unchanged and the index is ready to accept new documents. - /// - /// For Beginners: Completely empties the Azure Search index and all its facet indices. - /// - /// After calling Clear(): - /// - All documents are removed - /// - Inverted index is cleared (all facets) - /// - Vector dimension resets to 0 - /// - Index is ready for new documents - /// - /// Use with caution - this cannot be undone! - /// - /// Example: - /// - /// store.Clear(); - /// // Result is available in the returned value // 0 - /// - /// + /// Equality (string/bool) uses eq, numeric values use ge (mirroring the base-class + /// "field >= value" semantics) and collection values become an or group of eq + /// comparisons (any-of). All conditions are combined with and. /// - public override void Clear() + private static string? BuildODataFilter(Dictionary? metadataFilters) { - _documents.Clear(); - _invertedIndex.Clear(); - _vectorDimension = 0; - } + if (metadataFilters == null || metadataFilters.Count == 0) + return null; - private void IndexMetadata(Document document) - { - foreach (var kvp in document.Metadata) + var clauses = new List(); + + foreach (var kvp in metadataFilters) { - var fieldName = kvp.Key; - var fieldValue = kvp.Value; + var field = kvp.Key; + var value = kvp.Value; - if (!_invertedIndex.ContainsKey(fieldName)) + if (value == null) { - _invertedIndex[fieldName] = new Dictionary>(); + clauses.Add($"{field} eq null"); } - - // Create normalized key that preserves type information for accurate comparisons - var indexKey = NormalizeMetadataValue(fieldValue); - - if (!_invertedIndex[fieldName].ContainsKey(indexKey)) + else if (value is string s) { - _invertedIndex[fieldName][indexKey] = new HashSet(); + clauses.Add($"{field} eq {ODataString(s)}"); + } + else if (value is bool b) + { + clauses.Add($"{field} eq {(b ? "true" : "false")}"); + } + else if (IsNumeric(value)) + { + clauses.Add($"{field} ge {Convert.ToDouble(value, CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture)}"); + } + else if (value is System.Collections.IEnumerable enumerable) + { + var anyOf = new List(); + foreach (var item in enumerable) + { + if (item is string || item == null) + anyOf.Add($"{field} eq {ODataString(item?.ToString() ?? string.Empty)}"); + else if (item is bool ib) + anyOf.Add($"{field} eq {(ib ? "true" : "false")}"); + else if (IsNumeric(item)) + anyOf.Add($"{field} eq {Convert.ToDouble(item, CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture)}"); + else + anyOf.Add($"{field} eq {ODataString(item.ToString() ?? string.Empty)}"); + } + if (anyOf.Count > 0) + clauses.Add("(" + string.Join(" or ", anyOf) + ")"); + } + else + { + clauses.Add($"{field} eq {ODataString(value.ToString() ?? string.Empty)}"); } - - _invertedIndex[fieldName][indexKey].Add(document.Id); } + + return clauses.Count > 0 ? string.Join(" and ", clauses) : null; } - private string NormalizeMetadataValue(object? value) + // OData string literals are single-quoted; embedded single quotes are doubled. + private static string ODataString(string value) => "'" + value.Replace("'", "''") + "'"; + + /// + protected override Document? GetByIdCore(string documentId) { - if (value == null) - return "null"; + var encoded = Uri.EscapeDataString(documentId); + var info = SendAsync(HttpMethod.Get, WithApiVersion($"/indexes/{_indexName}/docs/{encoded}"), null).GetAwaiter().GetResult(); + if (info.Status == HttpStatusCode.NotFound) + return null; + EnsureSuccess(info, "get document"); + + return ParseDocument(JObject.Parse(info.Body)); + } - // Preserve type information in the key to ensure accurate comparisons - return value switch + /// + protected override bool RemoveCore(string documentId) + { + var action = new Dictionary { - bool b => $"bool:{(b ? "true" : "false")}", - int i => $"int:{i}", - long l => $"long:{l}", - float f => $"float:{f:R}", - double d => $"double:{d:R}", - decimal dec => $"decimal:{dec}", - _ => $"string:{value}" + ["@search.action"] = "delete", + [IdField] = documentId }; + + var body = new { value = new[] { action } }; + var info = SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/index"), body).GetAwaiter().GetResult(); + if (!IsSuccess(info.Status)) + return false; + + if (_documentCount > 0) + _documentCount--; + return true; } - private void RemoveFromIndex(Document document) + /// + protected override IEnumerable> GetAllCore() { - foreach (var kvp in document.Metadata) + var all = new List>(); + const int pageSize = 1000; + var skip = 0; + + while (true) { - var fieldName = kvp.Key; - var fieldValue = NormalizeMetadataValue(kvp.Value); + var body = new Dictionary + { + ["search"] = "*", + ["select"] = $"{IdField},{ContentField},{MetadataField}", + ["top"] = pageSize, + ["skip"] = skip + }; + + var info = SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/search"), body).GetAwaiter().GetResult(); + EnsureSuccess(info, "list documents"); + + var rows = JObject.Parse(info.Body)["value"] as JArray; + if (rows == null || rows.Count == 0) + break; - if (_invertedIndex.TryGetValue(fieldName, out var fieldIndex)) + foreach (var row in rows) { - if (fieldIndex.TryGetValue(fieldValue, out var docIds)) - { - docIds.Remove(document.Id); - if (docIds.Count == 0) - { - fieldIndex.Remove(fieldValue); - } - } - if (fieldIndex.Count == 0) - { - _invertedIndex.Remove(fieldName); - } + var doc = ParseDocument(row as JObject); + if (doc != null) + all.Add(doc); } + + if (rows.Count < pageSize) + break; + skip += pageSize; } + + return all; + } + + /// + public override void Clear() + { + var info = SendAsync(HttpMethod.Delete, WithApiVersion($"/indexes/{_indexName}"), null).GetAwaiter().GetResult(); + if (!IsSuccess(info.Status) && info.Status != HttpStatusCode.NotFound) + EnsureSuccess(info, "delete index"); + + _documentCount = 0; + _indexReady = false; + + if (_vectorDimension > 0) + CreateIndex(_vectorDimension); } - private HashSet? GetCandidateIds(Dictionary metadataFilters) + private Document? ParseDocument(JObject? row) { - if (metadataFilters.Count == 0) + if (row == null) return null; - HashSet? candidateIds = null; + var id = row[IdField]?.ToString(); + if (string.IsNullOrEmpty(id)) + return null; + + var content = row[ContentField]?.ToString() ?? string.Empty; - foreach (var filter in metadataFilters) + var metadata = new Dictionary(); + var metaJson = row[MetadataField]?.ToString(); + if (!string.IsNullOrEmpty(metaJson)) { - var fieldName = filter.Key; - var indexKey = NormalizeMetadataValue(filter.Value); + var parsed = JsonConvert.DeserializeObject>(metaJson!); + if (parsed != null) + metadata = parsed; + } + + return new Document(id!, content, metadata); + } + + private static bool IsNumeric(object value) + { + return value is sbyte || value is byte || value is short || value is ushort + || value is int || value is uint || value is long || value is ulong + || value is float || value is double || value is decimal; + } - if (_invertedIndex.TryGetValue(fieldName, out var fieldIndex)) + private static bool IsSuccess(HttpStatusCode status) => (int)status >= 200 && (int)status < 300; + + private void EnsureSuccess(HttpResponseInfo info, string operation) + { + if (!IsSuccess(info.Status)) + throw new HttpRequestException($"Azure AI Search {operation} failed with status {(int)info.Status}: {info.Body}"); + } + + private async Task SendAsync(HttpMethod method, string path, object? body) + { + using (var request = new HttpRequestMessage(method, path)) + { + if (body != null) { - if (fieldIndex.TryGetValue(indexKey, out var docIds)) - { - if (candidateIds == null) - { - candidateIds = new HashSet(docIds); - } - else - { - candidateIds.IntersectWith(docIds); - } - } - else - { - return new HashSet(); - } + var json = JsonConvert.SerializeObject(body); + request.Content = new StringContent(json, Encoding.UTF8, "application/json"); } - else + + using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) { - return new HashSet(); + var content = response.Content != null + ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + return new HttpResponseInfo(response.StatusCode, content); } } + } + + private readonly struct HttpResponseInfo + { + public HttpResponseInfo(HttpStatusCode status, string body) + { + Status = status; + Body = body; + } - return candidateIds; + public HttpStatusCode Status { get; } + public string Body { get; } } } } diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs index d40903922e..05efaacf8e 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs @@ -1,55 +1,61 @@ - using System; using System.Collections.Generic; +using System.Globalization; using System.Linq; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json; +using Newtonsoft.Json.Linq; namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores { /// - /// Milvus-inspired document store with collection-based vector organization. + /// Milvus vector-database document store backed by the real Milvus v2 REST API (/v2/vectordb/*). /// + /// The numeric type for vector operations. /// /// - /// This implementation provides a simplified in-memory version inspired by Milvus, - /// a cloud-native vector database. It organizes documents into named collections - /// and uses cosine similarity for retrieval operations. + /// This store talks to a running Milvus instance over HTTP. It manages a single collection with a + /// VarChar primary key and dynamic fields enabled, upserts entities (vector + fields), performs + /// filtered vector search, deletes entities by id and queries the whole collection with paging. + /// The original content is stored under content; the full metadata dictionary is stored as a + /// JSON string under metadata_json (for lossless round-tripping) while each scalar metadata + /// entry is also flattened to a m_<key> dynamic field so it can be used in filter + /// expressions. /// - /// For Beginners: This document store mimics Milvus, a popular cloud-based vector database. - /// - /// Think of collections like folders on your computer: - /// - Each collection has a unique name (like "ResearchPapers" or "CustomerReviews") - /// - Documents are grouped by topic or purpose - /// - Makes it easy to organize different types of content separately - /// - /// Good for: - /// - Prototyping before using real Milvus - /// - Testing Milvus-style organization patterns - /// - Small to medium document sets (< 100K documents) - /// - /// For production use real Milvus which provides: - /// - Distributed storage across multiple servers - /// - Advanced indexing (IVF, HNSW) for billion-scale vectors - /// - GPU acceleration for ultra-fast similarity search - /// - Persistence and fault tolerance + /// For Beginners: Milvus is an open-source vector database. This class is a real client + /// for it - every method here makes an HTTP call to your Milvus server using its REST API. Milvus + /// wraps every response in a { "code": 0, "data": ... } envelope; a non-zero code means the + /// operation failed. /// /// - /// The numeric type for vector operations. [ComponentType(ComponentType.DocumentStore)] [PipelineStage(PipelineStage.Indexing)] public class MilvusDocumentStore : DocumentStoreBase { - private readonly Dictionary> _documents; + private const string IdField = "id"; + private const string VectorField = "vector"; + private const string ContentField = "content"; + private const string MetadataField = "metadata_json"; + private const string MetaPrefix = "m_"; + + private readonly HttpClient _httpClient; private readonly string _collectionName; + private readonly string _metricType; private int _vectorDimension; + private int _documentCount; + private bool _collectionReady; /// - /// Gets the number of documents currently stored in this collection. + /// Gets the number of documents (entities) currently stored in the collection. /// - public override int DocumentCount => _documents.Count; + public override int DocumentCount => _documentCount; /// /// Gets the dimensionality of vectors stored in this collection. @@ -57,163 +63,462 @@ public class MilvusDocumentStore : DocumentStoreBase public override int VectorDimension => _vectorDimension; /// - /// Initializes a new instance of the MilvusDocumentStore class. + /// Gets the name of the Milvus collection this store is bound to. /// - /// The name of the collection to organize documents. - /// The initial capacity for the internal dictionary (default: 1000). - /// Thrown when collection name is empty or initial capacity is not positive. - /// - /// For Beginners: Creates a new document collection with a specific name. - /// - /// Example: - /// - /// // Create a collection for research papers - /// var store = new MilvusDocumentStore<float>("ResearchPapers"); - /// - /// // Create a larger collection - /// var bigStore = new MilvusDocumentStore<double>("Articles", initialCapacity: 10000); - /// - /// - /// The initial capacity is a performance hint - the collection can grow beyond it. - /// - /// - public MilvusDocumentStore(string collectionName, int initialCapacity = 1000) + public string CollectionName { get; } + + /// + /// Initializes a new instance of the class. + /// + /// The Milvus collection name. + /// The base URL of the Milvus server, e.g. http://localhost:19530. + /// Optional Milvus token (sent as an Authorization: Bearer header). + /// The metric type used when creating the collection. + /// + /// The vector dimension used to create the collection if it does not already exist. + /// When 0 the collection must already exist or the dimension is inferred from the first document. + /// + /// Optional pre-configured (primarily for testing). + /// Optional used to build the client (for testing). + public MilvusDocumentStore( + string collectionName, + string url, + string? token = null, + DistanceMetricType distanceMetric = DistanceMetricType.Cosine, + int vectorDimension = 0, + HttpClient? httpClient = null, + HttpMessageHandler? handler = null) { if (string.IsNullOrWhiteSpace(collectionName)) throw new ArgumentException("Collection name cannot be empty", nameof(collectionName)); - if (initialCapacity <= 0) - throw new ArgumentException("Initial capacity must be greater than zero", nameof(initialCapacity)); + if (vectorDimension < 0) + throw new ArgumentOutOfRangeException(nameof(vectorDimension), "Vector dimension cannot be negative"); _collectionName = collectionName; - _documents = new Dictionary>(initialCapacity); - _vectorDimension = 0; + CollectionName = collectionName; + _metricType = MapMetric(distanceMetric); + _vectorDimension = vectorDimension; + _documentCount = 0; + + _httpClient = httpClient ?? (handler != null ? new HttpClient(handler) : new HttpClient()); + if (_httpClient.BaseAddress == null) + { + if (string.IsNullOrWhiteSpace(url)) + throw new ArgumentException("Url cannot be empty", nameof(url)); + _httpClient.BaseAddress = new Uri(url); + } + + if (!string.IsNullOrWhiteSpace(token) && _httpClient.DefaultRequestHeaders.Authorization == null) + _httpClient.DefaultRequestHeaders.Add("Authorization", "Bearer " + token); + + InitializeCollection(vectorDimension); } - /// - /// Core logic for adding a single vector document to the collection. - /// - /// The validated vector document to add. - protected override void AddCore(VectorDocument vectorDocument) + private static string MapMetric(DistanceMetricType metric) { - if (_documents.Count == 0) + switch (metric) { - _vectorDimension = vectorDocument.Embedding.Length; + case DistanceMetricType.Cosine: + return "COSINE"; + case DistanceMetricType.Euclidean: + return "L2"; + default: + throw new NotSupportedException( + $"Metric '{metric}' is not supported by Milvus. Use Cosine (COSINE) or Euclidean (L2)."); } + } - _documents[vectorDocument.Document.Id] = vectorDocument; + private void InitializeCollection(int requestedDimension) + { + HttpResponseInfo info; + try + { + info = SendAsync("/v2/vectordb/collections/describe", + new { collectionName = _collectionName }).GetAwaiter().GetResult(); + } + catch (HttpRequestException) + { + // Server not reachable at construction time; defer creation to first write. + return; + } + + if (IsSuccess(info.Status)) + { + var root = JObject.Parse(info.Body); + var code = root["code"]?.Value() ?? 0; + if (code == 0) + { + _collectionReady = true; + var fields = root["data"]?["fields"] as JArray; + if (fields != null) + { + foreach (var field in fields) + { + var dim = field?["params"]?["dim"]; + if (dim != null && dim.Type != JTokenType.Null) + { + var parsed = Convert.ToInt32(dim, CultureInfo.InvariantCulture); + if (parsed > 0) + _vectorDimension = parsed; + } + } + } + return; + } + } + + if (requestedDimension > 0) + CreateCollection(requestedDimension); } - /// - /// Core logic for adding multiple vector documents in a batch operation. - /// - /// The validated list of vector documents to add. - /// Thrown when a document's embedding has inconsistent dimensions. - protected override void AddBatchCore(IList> vectorDocuments) + private void CreateCollection(int dimension) { - if (_vectorDimension == 0 && vectorDocuments.Count > 0) + var body = new { + collectionName = _collectionName, + dimension, + metricType = _metricType, + idType = "VarChar", + primaryFieldName = IdField, + vectorFieldName = VectorField, + enableDynamicField = true + }; + + var info = SendAsync("/v2/vectordb/collections/create", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "create collection"); + _vectorDimension = dimension; + _collectionReady = true; + } + + private void EnsureCollectionForDimension(int dimension) + { + if (_collectionReady) + return; + CreateCollection(dimension); + } + + /// + protected override void AddCore(VectorDocument vectorDocument) + { + EnsureCollectionForDimension(vectorDocument.Embedding.Length); + if (_vectorDimension == 0) + _vectorDimension = vectorDocument.Embedding.Length; + + var body = new { collectionName = _collectionName, data = new[] { BuildEntity(vectorDocument) } }; + var info = SendAsync("/v2/vectordb/entities/upsert", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "upsert entity"); + _documentCount++; + } + + /// + protected override void AddBatchCore(IList> vectorDocuments) + { + if (vectorDocuments.Count == 0) + return; + + EnsureCollectionForDimension(vectorDocuments[0].Embedding.Length); + if (_vectorDimension == 0) _vectorDimension = vectorDocuments[0].Embedding.Length; - } - foreach (var vectorDocument in vectorDocuments) + var data = vectorDocuments.Select(BuildEntity).ToList(); + var body = new { collectionName = _collectionName, data }; + var info = SendAsync("/v2/vectordb/entities/upsert", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "batch upsert entities"); + _documentCount += vectorDocuments.Count; + } + + private Dictionary BuildEntity(VectorDocument vectorDocument) + { + var vector = vectorDocument.Embedding.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + var metadata = vectorDocument.Document.Metadata ?? new Dictionary(); + + var entity = new Dictionary { - if (vectorDocument.Embedding.Length != _vectorDimension) - throw new ArgumentException( - $"Vector dimension mismatch in batch. Expected {_vectorDimension}, got {vectorDocument.Embedding.Length} for document {vectorDocument.Document.Id}", - nameof(vectorDocuments)); + [IdField] = vectorDocument.Document.Id, + [VectorField] = vector, + [ContentField] = vectorDocument.Document.Content, + [MetadataField] = JsonConvert.SerializeObject(metadata) + }; - _documents[vectorDocument.Document.Id] = vectorDocument; - } + foreach (var kvp in metadata) + entity[MetaPrefix + kvp.Key] = kvp.Value; + + return entity; } - /// - /// Core logic for similarity search with optional metadata filtering. - /// - /// The validated query vector. - /// The validated number of documents to return. - /// The validated metadata filters. - /// Top-k similar documents ordered by cosine similarity score. + /// protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) { - var scoredDocuments = new List<(Document Document, T Score)>(); + var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + + var body = new Dictionary + { + ["collectionName"] = _collectionName, + ["data"] = new[] { vector }, + ["annsField"] = VectorField, + ["limit"] = topK, + ["outputFields"] = new[] { IdField, ContentField, MetadataField } + }; - var matchingDocuments = _documents.Values - .Where(vectorDoc => MatchesFilters(vectorDoc.Document, metadataFilters)); + var filter = BuildFilterExpression(metadataFilters); + if (filter != null) + body["filter"] = filter; - foreach (var vectorDoc in matchingDocuments) + var info = SendAsync("/v2/vectordb/entities/search", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "search"); + + var results = new List>(); + var hits = JObject.Parse(info.Body)["data"] as JArray; + if (hits == null) + return results; + + foreach (var hit in hits) { - var similarity = StatisticsHelper.CosineSimilarity(queryVector, vectorDoc.Embedding); - scoredDocuments.Add((vectorDoc.Document, similarity)); - } + var doc = ParseEntity(hit as JObject); + if (doc == null) + continue; - var results = scoredDocuments - .OrderByDescending(x => x.Score) - .Take(topK) - .Select(x => - { - x.Document.RelevanceScore = x.Score; - x.Document.HasRelevanceScore = true; - return x.Document; - }) - .ToList(); + var distance = hit?["distance"]; + var score = distance != null && distance.Type != JTokenType.Null + ? Convert.ToDouble(distance, CultureInfo.InvariantCulture) + : 0.0; + doc.RelevanceScore = NumOps.FromDouble(score); + doc.HasRelevanceScore = true; + results.Add(doc); + } return results; } /// - /// Core logic for retrieving a document by its unique identifier. + /// Translates the metadata filter dictionary into a Milvus boolean filter expression. /// - /// The validated document ID. - /// The document if found; otherwise, null. + /// + /// Equality (string/bool) uses ==, numeric values use >= (mirroring the + /// base-class "field >= value" semantics) and collection values become an in [...] + /// clause (any-of). All conditions are combined with and. Filterable fields are stored + /// flattened under the m_ prefix. + /// + private static string? BuildFilterExpression(Dictionary? metadataFilters) + { + if (metadataFilters == null || metadataFilters.Count == 0) + return null; + + var clauses = new List(); + + foreach (var kvp in metadataFilters) + { + var field = MetaPrefix + kvp.Key; + var value = kvp.Value; + + if (value == null) + { + continue; + } + else if (value is string s) + { + clauses.Add($"{field} == {Quote(s)}"); + } + else if (value is bool b) + { + clauses.Add($"{field} == {(b ? "true" : "false")}"); + } + else if (IsNumeric(value)) + { + clauses.Add($"{field} >= {Convert.ToDouble(value, CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture)}"); + } + else if (value is System.Collections.IEnumerable enumerable) + { + var items = new List(); + foreach (var item in enumerable) + { + if (item is string || item == null) + items.Add(Quote(item?.ToString() ?? string.Empty)); + else if (IsNumeric(item)) + items.Add(Convert.ToDouble(item, CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture)); + else + items.Add(Quote(item.ToString() ?? string.Empty)); + } + if (items.Count > 0) + clauses.Add($"{field} in [{string.Join(", ", items)}]"); + } + else + { + clauses.Add($"{field} == {Quote(value.ToString() ?? string.Empty)}"); + } + } + + return clauses.Count > 0 ? string.Join(" and ", clauses) : null; + } + + private static string Quote(string value) => "\"" + value.Replace("\\", "\\\\").Replace("\"", "\\\"") + "\""; + + /// protected override Document? GetByIdCore(string documentId) { - return _documents.TryGetValue(documentId, out var vectorDoc) ? vectorDoc.Document : null; + var body = new + { + collectionName = _collectionName, + filter = $"{IdField} == {Quote(documentId)}", + outputFields = new[] { IdField, ContentField, MetadataField }, + limit = 1 + }; + + var info = SendAsync("/v2/vectordb/entities/query", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "query by id"); + + var rows = JObject.Parse(info.Body)["data"] as JArray; + if (rows == null || rows.Count == 0) + return null; + + return ParseEntity(rows[0] as JObject); } - /// - /// Core logic for removing a document from the collection. - /// - /// The validated document ID. - /// True if the document was found and removed; otherwise, false. + /// protected override bool RemoveCore(string documentId) { - var removed = _documents.Remove(documentId); - if (removed && _documents.Count == 0) + var body = new { - _vectorDimension = 0; - } - return removed; + collectionName = _collectionName, + filter = $"{IdField} == {Quote(documentId)}" + }; + + var info = SendAsync("/v2/vectordb/entities/delete", body).GetAwaiter().GetResult(); + if (!IsSuccess(info.Status)) + return false; + + var code = JObject.Parse(info.Body)["code"]?.Value() ?? 0; + if (code != 0) + return false; + + if (_documentCount > 0) + _documentCount--; + return true; } - /// - /// Core logic for retrieving all documents in the collection. - /// - /// An enumerable of all documents in the collection. + /// protected override IEnumerable> GetAllCore() { - return _documents.Values.Select(vd => vd.Document).ToList(); + var all = new List>(); + const int pageSize = 1000; + var offset = 0; + + while (true) + { + var body = new Dictionary + { + ["collectionName"] = _collectionName, + ["filter"] = $"{IdField} != \"\"", + ["outputFields"] = new[] { IdField, ContentField, MetadataField }, + ["limit"] = pageSize, + ["offset"] = offset + }; + + var info = SendAsync("/v2/vectordb/entities/query", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "query all"); + + var rows = JObject.Parse(info.Body)["data"] as JArray; + if (rows == null || rows.Count == 0) + break; + + foreach (var row in rows) + { + var doc = ParseEntity(row as JObject); + if (doc != null) + all.Add(doc); + } + + if (rows.Count < pageSize) + break; + offset += pageSize; + } + + return all; } - /// - /// Removes all documents from the collection and resets vector dimension. - /// - /// - /// For Beginners: This empties the entire collection. - /// - /// After calling Clear(): - /// - All documents are removed - /// - Vector dimension is reset to 0 - /// - Collection name stays the same - /// - Ready to accept new documents (even with different dimensions) - /// - /// Use with caution - this operation cannot be undone! - /// - /// + /// public override void Clear() { - _documents.Clear(); - _vectorDimension = 0; + var info = SendAsync("/v2/vectordb/collections/drop", + new { collectionName = _collectionName }).GetAwaiter().GetResult(); + EnsureSuccess(info, "drop collection"); + + _documentCount = 0; + _collectionReady = false; + + if (_vectorDimension > 0) + CreateCollection(_vectorDimension); + } + + private Document? ParseEntity(JObject? entity) + { + if (entity == null) + return null; + + var id = entity[IdField]?.ToString(); + if (string.IsNullOrEmpty(id)) + return null; + + var content = entity[ContentField]?.ToString() ?? string.Empty; + + var metadata = new Dictionary(); + var metaJson = entity[MetadataField]?.ToString(); + if (!string.IsNullOrEmpty(metaJson)) + { + var parsed = JsonConvert.DeserializeObject>(metaJson!); + if (parsed != null) + metadata = parsed; + } + + return new Document(id!, content, metadata); + } + + private static bool IsNumeric(object value) + { + return value is sbyte || value is byte || value is short || value is ushort + || value is int || value is uint || value is long || value is ulong + || value is float || value is double || value is decimal; + } + + private static bool IsSuccess(HttpStatusCode status) => (int)status >= 200 && (int)status < 300; + + private void EnsureSuccess(HttpResponseInfo info, string operation) + { + if (!IsSuccess(info.Status)) + throw new HttpRequestException($"Milvus {operation} failed with status {(int)info.Status}: {info.Body}"); + + var code = JObject.Parse(info.Body)["code"]?.Value() ?? 0; + if (code != 0) + throw new HttpRequestException($"Milvus {operation} failed with code {code}: {info.Body}"); + } + + private async Task SendAsync(string path, object body) + { + using (var request = new HttpRequestMessage(HttpMethod.Post, path)) + { + var json = JsonConvert.SerializeObject(body); + request.Content = new StringContent(json, Encoding.UTF8, "application/json"); + + using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) + { + var content = response.Content != null + ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + return new HttpResponseInfo(response.StatusCode, content); + } + } + } + + private readonly struct HttpResponseInfo + { + public HttpResponseInfo(HttpStatusCode status, string body) + { + Status = status; + Body = body; + } + + public HttpStatusCode Status { get; } + public string Body { get; } } } } - diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs index 9bdb673e4b..2490684e10 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs @@ -1,324 +1,513 @@ global using AiDotNet.RetrievalAugmentedGeneration.Models; +using System; +using System.Collections.Generic; +using System.Globalization; +using System.Linq; +using System.Net; +using System.Net.Http; +using System.Security.Cryptography; +using System.Text; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using Newtonsoft.Json; +using Newtonsoft.Json.Linq; -namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; - -/// -/// Weaviate-inspired document store with class-based schema organization. -/// -/// -/// -/// This implementation provides an in-memory simulation of Weaviate, a cloud-native vector database -/// that organizes data using a class-based schema system. It uses cosine similarity for retrieval. -/// -/// For Beginners: Weaviate is an open-source vector database with a GraphQL API. -/// -/// Think of classes like database tables: -/// - Each class defines a type of data (like "Article" or "Product") -/// - Documents are instances of that class -/// - Organized by schema for structured data -/// -/// This in-memory version is good for: -/// - Prototyping Weaviate-style schemas -/// - Testing class-based organization -/// - Small to medium collections (< 100K documents) -/// -/// Real Weaviate provides: -/// - GraphQL API for flexible queries -/// - Automatic schema inference -/// - Module system for ML models (transformers, CLIP, etc.) -/// - Multi-tenancy support -/// -/// -/// The numeric type for vector operations. -[ComponentType(ComponentType.DocumentStore)] -[PipelineStage(PipelineStage.Indexing)] -public class WeaviateDocumentStore : DocumentStoreBase +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores { - private readonly Dictionary> _documents; - private readonly string _className; - private int _vectorDimension; - - /// - /// Gets the number of documents currently stored in the class. - /// - public override int DocumentCount => _documents.Count; - - /// - /// Gets the dimensionality of vectors stored in this class. - /// - public override int VectorDimension => _vectorDimension; - - /// - /// Initializes a new instance of the WeaviateDocumentStore class. - /// - /// The class name to organize documents. - /// The initial capacity for the internal dictionary (default: 1000). - /// Thrown when class name is empty or initial capacity is not positive. - /// - /// For Beginners: Creates a new Weaviate-style document class. - /// - /// Example: - /// - /// // Create a class for articles - /// var store = new WeaviateDocumentStore<float>("Article"); - /// - /// // Create a class for products - /// var productStore = new WeaviateDocumentStore<double>("Product", 5000); - /// - /// - /// The class name helps organize different types of documents. - /// - /// - public WeaviateDocumentStore(string className, int initialCapacity = 1000) - { - if (string.IsNullOrWhiteSpace(className)) - throw new ArgumentException("Class name cannot be empty", nameof(className)); - if (initialCapacity <= 0) - throw new ArgumentException("Initial capacity must be greater than zero", nameof(initialCapacity)); - - _className = className; - _documents = new Dictionary>(initialCapacity); - _vectorDimension = 0; - } - /// - /// Core logic for adding a single vector document to the class. + /// Weaviate vector-database document store backed by the real Weaviate REST/GraphQL API. /// - /// The validated vector document to add. + /// The numeric type for vector operations. /// /// - /// The first document added determines the vector dimension for all documents in this class. - /// All subsequent documents must have embeddings of the same dimension. + /// This store talks to a running Weaviate instance over HTTP. It manages a single class, + /// upserts objects (vector + properties), performs nearVector GraphQL search with a + /// where filter, deletes objects and cursor-pages the whole class. The original document + /// id and content are stored under the reserved properties docId and content; the + /// full metadata dictionary is stored as a JSON string under metadataJson (for lossless + /// round-tripping) while each scalar metadata entry is also flattened to a m_<key> + /// property so it can be used in server-side where filters. + /// + /// For Beginners: Weaviate is an open-source vector database with a GraphQL API. + /// This class is a real client for it - every method here makes an HTTP call to your Weaviate + /// server. Weaviate requires object ids to be UUIDs, so each document's string id is turned into + /// a deterministic UUID; the original string id is kept in the docId property so you always + /// get it back. /// /// - protected override void AddCore(VectorDocument vectorDocument) + [ComponentType(ComponentType.DocumentStore)] + [PipelineStage(PipelineStage.Indexing)] + public class WeaviateDocumentStore : DocumentStoreBase { - if (_documents.Count == 0) + private const string DocIdProp = "docId"; + private const string ContentProp = "content"; + private const string MetadataProp = "metadataJson"; + private const string MetaPrefix = "m_"; + + private readonly HttpClient _httpClient; + private readonly string _className; + private readonly string _distance; + private int _vectorDimension; + private int _documentCount; + private bool _classReady; + + /// + /// Gets the number of documents (objects) currently stored in the class. + /// + public override int DocumentCount => _documentCount; + + /// + /// Gets the dimensionality of vectors stored in this class. + /// + public override int VectorDimension => _vectorDimension; + + /// + /// Gets the name of the Weaviate class this store is bound to. + /// + public string ClassName { get; } + + /// + /// Initializes a new instance of the class. + /// + /// The Weaviate class name (must start with an uppercase letter). + /// The base URL of the Weaviate server, e.g. http://localhost:8080. + /// Optional Weaviate API key (sent as an Authorization: Bearer header). + /// The distance metric used when creating the class. + /// Optional expected vector dimension (informational; inferred from the first document when 0). + /// Optional pre-configured (primarily for testing). + /// Optional used to build the client (for testing). + public WeaviateDocumentStore( + string className, + string url, + string? apiKey = null, + DistanceMetricType distanceMetric = DistanceMetricType.Cosine, + int vectorDimension = 0, + HttpClient? httpClient = null, + HttpMessageHandler? handler = null) { - _vectorDimension = vectorDocument.Embedding.Length; + if (string.IsNullOrWhiteSpace(className)) + throw new ArgumentException("Class name cannot be empty", nameof(className)); + if (vectorDimension < 0) + throw new ArgumentOutOfRangeException(nameof(vectorDimension), "Vector dimension cannot be negative"); + + // Weaviate GraphQL class names must be capitalized. + _className = char.ToUpperInvariant(className[0]) + className.Substring(1); + ClassName = _className; + _distance = MapDistance(distanceMetric); + _vectorDimension = vectorDimension; + _documentCount = 0; + + _httpClient = httpClient ?? (handler != null ? new HttpClient(handler) : new HttpClient()); + if (_httpClient.BaseAddress == null) + { + if (string.IsNullOrWhiteSpace(url)) + throw new ArgumentException("Url cannot be empty", nameof(url)); + _httpClient.BaseAddress = new Uri(url); + } + + if (!string.IsNullOrWhiteSpace(apiKey) && _httpClient.DefaultRequestHeaders.Authorization == null) + _httpClient.DefaultRequestHeaders.Add("Authorization", "Bearer " + apiKey); + + InitializeClass(); } - _documents[vectorDocument.Document.Id] = vectorDocument; - } + private static string MapDistance(DistanceMetricType metric) + { + switch (metric) + { + case DistanceMetricType.Cosine: + return "cosine"; + case DistanceMetricType.Euclidean: + return "l2-squared"; + case DistanceMetricType.Manhattan: + return "manhattan"; + case DistanceMetricType.Hamming: + return "hamming"; + default: + throw new NotSupportedException( + $"Distance metric '{metric}' is not supported by Weaviate. Use Cosine, Euclidean, Manhattan or Hamming."); + } + } - /// - /// Core logic for adding multiple vector documents in a batch operation. - /// - /// The validated list of vector documents to add. - /// Thrown when a document's embedding has inconsistent dimensions. - /// - /// - /// Batch operations are more efficient than adding documents individually. - /// All documents must have embeddings with the same dimension as the class. - /// - /// For Beginners: Adding many documents at once is faster. - /// - /// Inefficient: - /// - /// foreach (var doc in documents) - /// store.Add(doc); // Slow! - /// - /// - /// Efficient: - /// - /// store.AddBatch(documents); // Fast! - /// - /// - /// - protected override void AddBatchCore(IList> vectorDocuments) - { - if (_vectorDimension == 0 && vectorDocuments.Count > 0) + private void InitializeClass() { - _vectorDimension = vectorDocuments[0].Embedding.Length; + HttpResponseInfo info; + try + { + info = SendAsync(HttpMethod.Get, $"/v1/schema/{_className}", null).GetAwaiter().GetResult(); + } + catch (HttpRequestException) + { + // Server not reachable at construction time; defer creation to first write. + return; + } + + if (info.Status == HttpStatusCode.OK) + { + _classReady = true; + return; + } + + if (info.Status == HttpStatusCode.NotFound) + CreateClass(); } - foreach (var vectorDocument in vectorDocuments) + private void CreateClass() { - if (vectorDocument.Embedding.Length != _vectorDimension) - throw new ArgumentException( - $"Vector dimension mismatch in batch. Expected {_vectorDimension}, got {vectorDocument.Embedding.Length} for document {vectorDocument.Document.Id}", - nameof(vectorDocuments)); + var body = new + { + @class = _className, + vectorizer = "none", + vectorIndexConfig = new { distance = _distance }, + properties = new object[] + { + new { name = DocIdProp, dataType = new[] { "text" } }, + new { name = ContentProp, dataType = new[] { "text" } }, + new { name = MetadataProp, dataType = new[] { "text" } } + } + }; - _documents[vectorDocument.Document.Id] = vectorDocument; + var info = SendAsync(HttpMethod.Post, "/v1/schema", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "create class"); + _classReady = true; } - } - /// - /// Core logic for similarity search using cosine similarity with optional metadata filtering. - /// - /// The validated query vector. - /// The validated number of documents to return. - /// The validated metadata filters. - /// Top-k similar documents ordered by cosine similarity score. - /// - /// - /// Performs vector similarity search across all documents in the class, optionally filtering by metadata. - /// Results are ordered by decreasing cosine similarity. - /// - /// For Beginners: Finds the most similar documents. - /// - /// How it works: - /// 1. Filter documents by metadata (if provided) - /// 2. Calculate similarity between query and each document - /// 3. Sort by similarity (highest first) - /// 4. Return top-k matches - /// - /// Example: - /// - /// // Find 10 most similar articles - /// var results = store.GetSimilar(queryVector, topK: 10); - /// - /// // Find similar articles by specific author - /// var filters = new Dictionary<string, object> { ["author"] = "John Smith" }; - /// var filtered = store.GetSimilarWithFilters(queryVector, 5, filters); - /// - /// - /// - protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) - { - var scoredDocuments = new List<(Document Document, T Score)>(); + private void EnsureClass() + { + if (!_classReady) + CreateClass(); + } - var matchingDocuments = _documents.Values - .Where(vectorDoc => MatchesFilters(vectorDoc.Document, metadataFilters)); + /// + protected override void AddCore(VectorDocument vectorDocument) + { + EnsureClass(); + if (_vectorDimension == 0) + _vectorDimension = vectorDocument.Embedding.Length; - foreach (var vectorDoc in matchingDocuments) + var body = BuildObject(vectorDocument); + var info = SendAsync(HttpMethod.Post, "/v1/objects", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "create object"); + _documentCount++; + } + + /// + protected override void AddBatchCore(IList> vectorDocuments) { - var similarity = StatisticsHelper.CosineSimilarity(queryVector, vectorDoc.Embedding); - scoredDocuments.Add((vectorDoc.Document, similarity)); + if (vectorDocuments.Count == 0) + return; + + EnsureClass(); + if (_vectorDimension == 0) + _vectorDimension = vectorDocuments[0].Embedding.Length; + + var objects = vectorDocuments.Select(BuildObject).ToList(); + var body = new { objects }; + var info = SendAsync(HttpMethod.Post, "/v1/batch/objects", body).GetAwaiter().GetResult(); + EnsureSuccess(info, "batch create objects"); + _documentCount += vectorDocuments.Count; } - var results = scoredDocuments - .OrderByDescending(x => x.Score) - .Take(topK) - .Select(x => + private object BuildObject(VectorDocument vectorDocument) + { + var vector = vectorDocument.Embedding.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + var metadata = vectorDocument.Document.Metadata ?? new Dictionary(); + + var properties = new Dictionary { - x.Document.RelevanceScore = x.Score; - x.Document.HasRelevanceScore = true; - return x.Document; - }) - .ToList(); + [DocIdProp] = vectorDocument.Document.Id, + [ContentProp] = vectorDocument.Document.Content, + [MetadataProp] = JsonConvert.SerializeObject(metadata) + }; - return results; - } + foreach (var kvp in metadata) + properties[MetaPrefix + kvp.Key] = kvp.Value; - /// - /// Core logic for retrieving a document by its unique identifier. - /// - /// The validated document ID. - /// The document if found; otherwise, null. - /// - /// For Beginners: Gets a specific document by ID. - /// - /// Example: - /// - /// var doc = store.GetById("article-123"); - /// if (doc != null) - /// Console.WriteLine($"Article: {doc.Content}"); - /// - /// - /// - protected override Document? GetByIdCore(string documentId) - { - return _documents.TryGetValue(documentId, out var vectorDoc) ? vectorDoc.Document : null; - } + return new + { + @class = _className, + id = ToObjectId(vectorDocument.Document.Id), + properties, + vector + }; + } - /// - /// Core logic for removing a document from the class. - /// - /// The validated document ID. - /// True if the document was found and removed; otherwise, false. - /// - /// - /// Removes the document from the class. If this was the last document, the vector dimension - /// is reset to 0, allowing a new dimension on next add. - /// - /// For Beginners: Deletes a document from the class. - /// - /// Example: - /// - /// if (store.Remove("article-123")) - /// Console.WriteLine("Article deleted"); - /// - /// - /// - protected override bool RemoveCore(string documentId) - { - var removed = _documents.Remove(documentId); - if (removed && _documents.Count == 0) + /// + protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) { - _vectorDimension = 0; + var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + var vectorJson = "[" + string.Join(",", vector.Select(v => v.ToString("R", CultureInfo.InvariantCulture))) + "]"; + + var whereClause = BuildWhere(metadataFilters); + var whereArg = whereClause != null ? ", where: " + whereClause : string.Empty; + + var query = + "{ Get { " + _className + "(limit: " + topK.ToString(CultureInfo.InvariantCulture) + + ", nearVector: {vector: " + vectorJson + "}" + whereArg + ") " + + "{ " + DocIdProp + " " + ContentProp + " " + MetadataProp + + " _additional { certainty distance } } } }"; + + var info = SendAsync(HttpMethod.Post, "/v1/graphql", new { query }).GetAwaiter().GetResult(); + EnsureSuccess(info, "search"); + + var results = new List>(); + var hits = JObject.Parse(info.Body)["data"]?["Get"]?[_className] as JArray; + if (hits == null) + return results; + + foreach (var hit in hits) + { + var doc = ParseProperties(hit as JObject); + if (doc == null) + continue; + + var additional = hit?["_additional"]; + double score = 0.0; + var certainty = additional?["certainty"]; + if (certainty != null && certainty.Type != JTokenType.Null) + score = Convert.ToDouble(certainty, CultureInfo.InvariantCulture); + else + { + var distance = additional?["distance"]; + if (distance != null && distance.Type != JTokenType.Null) + score = 1.0 - Convert.ToDouble(distance, CultureInfo.InvariantCulture); + } + + doc.RelevanceScore = NumOps.FromDouble(score); + doc.HasRelevanceScore = true; + results.Add(doc); + } + + return results; } - return removed; - } - /// - /// Core logic for retrieving all documents in the class. - /// - /// An enumerable of all documents without their vector embeddings. - /// - /// - /// Returns all documents in the class in no particular order. Vector embeddings are not included. - /// For large classes, be aware of the memory impact. - /// - /// For Beginners: Gets every document in the class. - /// - /// Use cases: - /// - Export all documents for backup - /// - Migrate to a different class or store - /// - Bulk processing or analysis - /// - Debugging to see all stored documents - /// - /// Warning: For large classes (> 10K documents), this can use significant memory. - /// - /// Example: - /// - /// // Get all documents - /// var allDocs = store.GetAll().ToList(); - /// // Result is available in the returned value - /// - /// // Export to JSON file - /// var json = JsonConvert.SerializeObject(allDocs); - /// File.WriteAllText($"{_className}_export.json", json); - /// - /// - /// - protected override IEnumerable> GetAllCore() - { - return _documents.Values.Select(vd => vd.Document).ToList(); - } + /// + /// Translates the metadata filter dictionary into a Weaviate GraphQL where clause. + /// + /// + /// Equality (string/bool) uses Equal, numeric values use GreaterThanEqual + /// (mirroring the base-class "field >= value" semantics) and collection values become an + /// Or group of Equal operands (any-of). All top-level conditions are combined + /// under And. Filterable properties are stored flattened under the m_ prefix. + /// + private static string? BuildWhere(Dictionary? metadataFilters) + { + if (metadataFilters == null || metadataFilters.Count == 0) + return null; - /// - /// Removes all documents from the class and resets the vector dimension. - /// - /// - /// - /// Clears all documents from the class and resets the vector dimension to 0. - /// The class name remains unchanged and is ready to accept new documents. - /// - /// For Beginners: Completely empties the class. - /// - /// After calling Clear(): - /// - All documents are removed - /// - Vector dimension resets to 0 - /// - Class name stays the same - /// - Ready for new documents (even with different dimensions) - /// - /// Use with caution - this cannot be undone! - /// - /// Example: - /// - /// store.Clear(); - /// // Result is available in the returned value // 0 - /// - /// - /// - public override void Clear() - { - _documents.Clear(); - _vectorDimension = 0; + var operands = new List(); + + foreach (var kvp in metadataFilters) + { + var path = "[\"" + MetaPrefix + kvp.Key + "\"]"; + var value = kvp.Value; + + if (value == null) + { + operands.Add("{path: " + path + ", operator: IsNull, valueBoolean: true}"); + } + else if (value is string s) + { + operands.Add("{path: " + path + ", operator: Equal, valueText: " + JsonConvert.ToString(s) + "}"); + } + else if (value is bool b) + { + operands.Add("{path: " + path + ", operator: Equal, valueBoolean: " + (b ? "true" : "false") + "}"); + } + else if (IsNumeric(value)) + { + var num = Convert.ToDouble(value, CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture); + operands.Add("{path: " + path + ", operator: GreaterThanEqual, valueNumber: " + num + "}"); + } + else if (value is System.Collections.IEnumerable enumerable) + { + var anyOf = new List(); + foreach (var item in enumerable) + { + anyOf.Add("{path: " + path + ", operator: Equal, valueText: " + + JsonConvert.ToString(item?.ToString() ?? string.Empty) + "}"); + } + if (anyOf.Count > 0) + operands.Add("{operator: Or, operands: [" + string.Join(", ", anyOf) + "]}"); + } + else + { + operands.Add("{path: " + path + ", operator: Equal, valueText: " + + JsonConvert.ToString(value.ToString() ?? string.Empty) + "}"); + } + } + + if (operands.Count == 0) + return null; + if (operands.Count == 1) + return operands[0]; + + return "{operator: And, operands: [" + string.Join(", ", operands) + "]}"; + } + + /// + protected override Document? GetByIdCore(string documentId) + { + var info = SendAsync(HttpMethod.Get, $"/v1/objects/{_className}/{ToObjectId(documentId)}", null).GetAwaiter().GetResult(); + if (info.Status == HttpStatusCode.NotFound) + return null; + EnsureSuccess(info, "get object"); + + return ParseObject(JObject.Parse(info.Body)); + } + + /// + protected override bool RemoveCore(string documentId) + { + var info = SendAsync(HttpMethod.Delete, $"/v1/objects/{_className}/{ToObjectId(documentId)}", null).GetAwaiter().GetResult(); + if (info.Status == HttpStatusCode.NotFound) + return false; + if (!IsSuccess(info.Status)) + return false; + + if (_documentCount > 0) + _documentCount--; + return true; + } + + /// + protected override IEnumerable> GetAllCore() + { + var all = new List>(); + string? after = null; + const int pageSize = 100; + + while (true) + { + var path = $"/v1/objects?class={_className}&limit={pageSize}&include=vector"; + if (after != null) + path += "&after=" + after; + + var info = SendAsync(HttpMethod.Get, path, null).GetAwaiter().GetResult(); + EnsureSuccess(info, "list objects"); + + var objects = JObject.Parse(info.Body)["objects"] as JArray; + if (objects == null || objects.Count == 0) + break; + + foreach (var obj in objects) + { + var doc = ParseObject(obj as JObject); + if (doc != null) + all.Add(doc); + } + + after = objects[objects.Count - 1]?["id"]?.ToString(); + if (string.IsNullOrEmpty(after) || objects.Count < pageSize) + break; + } + + return all; + } + + /// + public override void Clear() + { + var info = SendAsync(HttpMethod.Delete, $"/v1/schema/{_className}", null).GetAwaiter().GetResult(); + if (!IsSuccess(info.Status) && info.Status != HttpStatusCode.NotFound) + EnsureSuccess(info, "delete class"); + + _documentCount = 0; + _classReady = false; + CreateClass(); + } + + // REST responses (GetById / list) nest fields under "properties". + private Document? ParseObject(JObject? obj) + { + return ParseProperties(obj?["properties"] as JObject); + } + + // GraphQL search hits expose the fields flat on the hit object. + private Document? ParseProperties(JObject? properties) + { + if (properties == null) + return null; + + var id = properties[DocIdProp]?.ToString(); + if (string.IsNullOrEmpty(id)) + return null; + + var content = properties[ContentProp]?.ToString() ?? string.Empty; + + var metadata = new Dictionary(); + var metaJson = properties[MetadataProp]?.ToString(); + if (!string.IsNullOrEmpty(metaJson)) + { + var parsed = JsonConvert.DeserializeObject>(metaJson!); + if (parsed != null) + metadata = parsed; + } + + return new Document(id!, content, metadata); + } + + private static bool IsNumeric(object value) + { + return value is sbyte || value is byte || value is short || value is ushort + || value is int || value is uint || value is long || value is ulong + || value is float || value is double || value is decimal; + } + + private static bool IsSuccess(HttpStatusCode status) => (int)status >= 200 && (int)status < 300; + + /// + /// Produces a deterministic UUID string for an arbitrary document id, since Weaviate object + /// ids must be UUIDs. The mapping is stable so lookups/deletes work. + /// + private static string ToObjectId(string documentId) + { + using (var md5 = MD5.Create()) + { + var hash = md5.ComputeHash(Encoding.UTF8.GetBytes(documentId)); + return new Guid(hash).ToString(); + } + } + + private void EnsureSuccess(HttpResponseInfo info, string operation) + { + if (!IsSuccess(info.Status)) + throw new HttpRequestException($"Weaviate {operation} failed with status {(int)info.Status}: {info.Body}"); + } + + private async Task SendAsync(HttpMethod method, string path, object? body) + { + using (var request = new HttpRequestMessage(method, path)) + { + if (body != null) + { + var json = JsonConvert.SerializeObject(body); + request.Content = new StringContent(json, Encoding.UTF8, "application/json"); + } + + using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) + { + var content = response.Content != null + ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + return new HttpResponseInfo(response.StatusCode, content); + } + } + } + + private readonly struct HttpResponseInfo + { + public HttpResponseInfo(HttpStatusCode status, string body) + { + Status = status; + Body = body; + } + + public HttpStatusCode Status { get; } + public string Body { get; } + } } } - diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStoreTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStoreTests.cs new file mode 100644 index 0000000000..8369bd0ca3 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStoreTests.cs @@ -0,0 +1,339 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Threading; +using System.Threading.Tasks; +using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json.Linq; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.DocumentStores +{ + /// + /// Unit tests for the real Azure AI Search REST client. All HTTP traffic is intercepted by a mock + /// so these run in CI with no network access. + /// + public class AzureSearchDocumentStoreTests + { + private const string Endpoint = "https://svc.search.windows.net"; + private const string Index = "test"; + + private const string ExistingIndexJson = + "{\"name\":\"test\",\"fields\":[{\"name\":\"embedding\",\"type\":\"Collection(Edm.Single)\",\"dimensions\":3}]}"; + + #region infrastructure + + private sealed class RecordedRequest + { + public RecordedRequest(string method, string path, string body) + { + Method = method; + Path = path; + Body = body; + } + + public string Method { get; } + public string Path { get; } + public string Body { get; } + } + + private sealed class MockHandler : HttpMessageHandler + { + private readonly List<(string Method, string PathContains, Func Resp)> _routes = new(); + public List Requests { get; } = new(); + public string DefaultBody { get; set; } = "{}"; + + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, _ => resp())); + return this; + } + + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, resp)); + return this; + } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + var body = request.Content != null + ? await request.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + var path = request.RequestUri!.AbsolutePath; + Requests.Add(new RecordedRequest(request.Method.Method, request.RequestUri.PathAndQuery, body)); + + foreach (var route in _routes) + { + if (route.Method == request.Method.Method && path.Contains(route.PathContains)) + return route.Resp(body); + } + + return Json(DefaultBody); + } + } + + private static HttpResponseMessage Json(string body, HttpStatusCode code = HttpStatusCode.OK) + => new HttpResponseMessage(code) { Content = new StringContent(body, Encoding.UTF8, "application/json") }; + + private static AzureSearchDocumentStore StoreWith(MockHandler handler, + DistanceMetricType metric = DistanceMetricType.Cosine, int vectorDimension = 0) + { + var client = new HttpClient(handler) { BaseAddress = new Uri(Endpoint) }; + return new AzureSearchDocumentStore(Index, Endpoint, apiKey: "k", + distanceMetric: metric, vectorDimension: vectorDimension, httpClient: client); + } + + private static VectorDocument Doc(string id, string content, float[] vector, Dictionary? metadata = null) + { + var d = new Document(id, content, metadata ?? new Dictionary()); + return new VectorDocument { Document = d, Embedding = new Vector(vector) }; + } + + #endregion + + [Fact] + public void Constructor_ReadsExistingIndex_SetsDimension() + { + var handler = new MockHandler() + .On("GET", "/indexes/test", () => Json(ExistingIndexJson)); + + var store = StoreWith(handler); + + Assert.Equal(3, store.VectorDimension); + Assert.Equal(Index, store.IndexName); + Assert.Contains(handler.Requests, r => r.Method == "GET" && r.Path.Contains("api-version=2023-11-01")); + } + + [Fact] + public void Constructor_IndexMissing_CreatesWithHnswProfileAndMetric() + { + var handler = new MockHandler() + .On("PUT", "/indexes/test", () => Json("{}")) + .On("GET", "/indexes/test", () => Json("{}", HttpStatusCode.NotFound)); + + var store = StoreWith(handler, metric: DistanceMetricType.Euclidean, vectorDimension: 4); + + Assert.Equal(4, store.VectorDimension); + var create = handler.Requests.Single(r => r.Method == "PUT" && r.Path.StartsWith("/indexes/test")); + var body = JObject.Parse(create.Body); + var embedding = ((JArray)body["fields"]!).Single(f => (string?)f["name"] == "embedding"); + Assert.Equal(4, (int)embedding["dimensions"]!); + Assert.Equal("vprofile", (string?)embedding["vectorSearchProfile"]); + var algo = ((JArray)body["vectorSearch"]!["algorithms"]!)[0]!; + Assert.Equal("hnsw", (string?)algo["kind"]); + Assert.Equal("euclidean", (string?)algo["hnswParameters"]!["metric"]); + } + + [Fact] + public void Add_UploadsMergeOrUpload_WithVectorAndFlattenedMetadata() + { + var handler = new MockHandler() + .On("GET", "/indexes/test", () => Json(ExistingIndexJson)) + .On("POST", "/indexes/test/docs/index", () => Json("{\"value\":[{\"key\":\"doc1\",\"status\":true}]}")); + + var store = StoreWith(handler); + store.Add(Doc("doc1", "Hello world", new float[] { 1, 0, 0 }, + new Dictionary { ["category"] = "science" })); + + Assert.Equal(1, store.DocumentCount); + var upload = handler.Requests.Single(r => r.Method == "POST" && r.Path.Contains("/docs/index")); + var action = JObject.Parse(upload.Body)["value"]![0]!; + Assert.Equal("mergeOrUpload", (string?)action["@search.action"]); + Assert.Equal("doc1", (string?)action["id"]); + Assert.Equal("Hello world", (string?)action["content"]); + Assert.Equal("science", (string?)action["category"]); + Assert.Equal(3, ((JArray)action["embedding"]!).Count); + } + + [Fact] + public void GetSimilar_ParsesRankedHits_InOrder() + { + var searchResponse = + "{\"value\":[" + + "{\"@search.score\":0.98,\"id\":\"doc1\",\"content\":\"first\",\"metadata_json\":\"{\\\"category\\\":\\\"science\\\"}\"}," + + "{\"@search.score\":0.55,\"id\":\"doc2\",\"content\":\"second\",\"metadata_json\":\"{}\"}" + + "]}"; + + var handler = new MockHandler() + .On("POST", "/indexes/test/docs/search", () => Json(searchResponse)) + .On("GET", "/indexes/test", () => Json(ExistingIndexJson)); + + var store = StoreWith(handler); + var results = store.GetSimilar(new Vector(new float[] { 1, 0, 0 }), topK: 2).ToList(); + + Assert.Equal(2, results.Count); + Assert.Equal("doc1", results[0].Id); + Assert.Equal("first", results[0].Content); + Assert.Equal("science", results[0].Metadata["category"].ToString()); + Assert.True(results[0].HasRelevanceScore); + Assert.True(Convert.ToDouble(results[0].RelevanceScore) > Convert.ToDouble(results[1].RelevanceScore)); + + var search = handler.Requests.Single(r => r.Path.Contains("/docs/search")); + var body = JObject.Parse(search.Body); + var vq = ((JArray)body["vectorQueries"]!)[0]!; + Assert.Equal("vector", (string?)vq["kind"]); + Assert.Equal("embedding", (string?)vq["fields"]); + Assert.Equal(2, (int)vq["k"]!); + Assert.Equal(3, ((JArray)vq["vector"]!).Count); + } + + [Fact] + public void GetSimilarWithFilters_TranslatesToOData() + { + var handler = new MockHandler() + .On("POST", "/indexes/test/docs/search", () => Json("{\"value\":[]}")) + .On("GET", "/indexes/test", () => Json(ExistingIndexJson)); + + var store = StoreWith(handler); + var filters = new Dictionary + { + ["category"] = "science", + ["year"] = 2020, + ["tag"] = new List { "a", "b" } + }; + store.GetSimilarWithFilters(new Vector(new float[] { 1, 0, 0 }), 10, filters).ToList(); + + var filter = (string)JObject.Parse(handler.Requests.Single(r => r.Path.Contains("/docs/search")).Body)["filter"]!; + Assert.Contains("category eq 'science'", filter); + Assert.Contains("year ge 2020", filter); + Assert.Contains("(tag eq 'a' or tag eq 'b')", filter); + Assert.Contains(" and ", filter); + } + + [Fact] + public void GetById_Found_ReturnsDocument() + { + var handler = new MockHandler() + .On("GET", "/indexes/test/docs/", () => Json( + "{\"id\":\"doc1\",\"content\":\"hi\",\"metadata_json\":\"{\\\"k\\\":\\\"v\\\"}\"}")) + .On("GET", "/indexes/test", () => Json(ExistingIndexJson)); + + var store = StoreWith(handler); + var doc = store.GetById("doc1"); + + Assert.NotNull(doc); + Assert.Equal("doc1", doc!.Id); + Assert.Equal("hi", doc.Content); + Assert.Equal("v", doc.Metadata["k"].ToString()); + } + + [Fact] + public void GetById_NotFound_ReturnsNull() + { + var handler = new MockHandler() + .On("GET", "/indexes/test/docs/", () => Json("{}", HttpStatusCode.NotFound)) + .On("GET", "/indexes/test", () => Json(ExistingIndexJson)); + + var store = StoreWith(handler); + Assert.Null(store.GetById("missing")); + } + + [Fact] + public void Remove_SendsDeleteAction_ReturnsTrue() + { + var handler = new MockHandler() + .On("POST", "/indexes/test/docs/index", () => Json("{\"value\":[{\"key\":\"doc1\",\"status\":true}]}")) + .On("GET", "/indexes/test", () => Json(ExistingIndexJson)); + + var store = StoreWith(handler); + Assert.True(store.Remove("doc1")); + + var del = handler.Requests.Single(r => r.Method == "POST" && r.Path.Contains("/docs/index")); + var action = JObject.Parse(del.Body)["value"]![0]!; + Assert.Equal("delete", (string?)action["@search.action"]); + Assert.Equal("doc1", (string?)action["id"]); + } + + [Fact] + public void GetAll_PagesUntilExhausted() + { + var call = 0; + var handler = new MockHandler() + .On("POST", "/indexes/test/docs/search", _ => + { + call++; + if (call == 1) + { + var items = string.Join(",", Enumerable.Range(0, 1000).Select(i => + "{\"id\":\"d" + i + "\",\"content\":\"c\",\"metadata_json\":\"{}\"}")); + return Json("{\"value\":[" + items + "]}"); + } + return Json("{\"value\":[{\"id\":\"dZ\",\"content\":\"c\",\"metadata_json\":\"{}\"}]}"); + }) + .On("GET", "/indexes/test", () => Json(ExistingIndexJson)); + + var store = StoreWith(handler); + var all = store.GetAll().ToList(); + + Assert.Equal(1001, all.Count); + Assert.Contains(all, d => d.Id == "dZ"); + Assert.Equal(2, handler.Requests.Count(r => r.Path.Contains("/docs/search"))); + } + + [Fact] + public void Clear_DeletesIndex_AndRecreates() + { + var handler = new MockHandler() + .On("DELETE", "/indexes/test", () => Json("{}", HttpStatusCode.NoContent)) + .On("PUT", "/indexes/test", () => Json("{}")) + .On("GET", "/indexes/test", () => Json(ExistingIndexJson)); + + // Existing index reports dim 3, so Clear recreates it. + var store = StoreWith(handler); + store.Clear(); + + Assert.Equal(0, store.DocumentCount); + Assert.Contains(handler.Requests, r => r.Method == "DELETE" && r.Path.StartsWith("/indexes/test")); + Assert.Contains(handler.Requests, r => r.Method == "PUT" && r.Path.StartsWith("/indexes/test")); + } + + [Fact] + public void Constructor_UnsupportedMetric_Throws() + { + Assert.Throws(() => + new AzureSearchDocumentStore(Index, Endpoint, "k", distanceMetric: DistanceMetricType.Jaccard, + httpClient: new HttpClient(new MockHandler()) { BaseAddress = new Uri(Endpoint) })); + } + + #region Integration (gated - require a live Azure AI Search service) + + [Trait("Category", "Integration")] + [SkippableFact] + public void Integration_UpsertSearchDelete_RoundTrips() + { + var endpoint = Environment.GetEnvironmentVariable("AZURE_SEARCH_ENDPOINT"); + var apiKey = Environment.GetEnvironmentVariable("AZURE_SEARCH_API_KEY"); + Skip.If(string.IsNullOrWhiteSpace(endpoint) || string.IsNullOrWhiteSpace(apiKey), + "Set AZURE_SEARCH_ENDPOINT and AZURE_SEARCH_API_KEY to run Azure AI Search integration tests."); + + var index = "aidotnet-it-" + Guid.NewGuid().ToString("N"); + var store = new AzureSearchDocumentStore(index, endpoint!, apiKey!, + distanceMetric: DistanceMetricType.Cosine, vectorDimension: 4); + + try + { + var id = "it-" + Guid.NewGuid().ToString("N"); + store.Add(Doc(id, "integration content", new float[] { 1, 0, 0, 0 }, + new Dictionary { ["category"] = "test" })); + + var results = store.GetSimilar(new Vector(new float[] { 1, 0, 0, 0 }), topK: 5).ToList(); + Assert.Contains(results, r => r.Id == id); + + Assert.True(store.Remove(id)); + } + finally + { + store.Clear(); + } + } + + #endregion + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStoreTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStoreTests.cs new file mode 100644 index 0000000000..ada1fa4974 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStoreTests.cs @@ -0,0 +1,336 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Threading; +using System.Threading.Tasks; +using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json.Linq; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.DocumentStores +{ + /// + /// Unit tests for the real Milvus v2 REST client. All HTTP traffic is intercepted by a mock + /// so these run in CI with no network access. + /// + public class MilvusDocumentStoreTests + { + private const string BaseUrl = "http://localhost:19530"; + private const string Collection = "test"; + + // describe response for an existing 3-d collection. + private const string ExistingDescribeJson = + "{\"code\":0,\"data\":{\"collectionName\":\"test\",\"fields\":[{\"name\":\"vector\",\"params\":{\"dim\":3}}]}}"; + + #region infrastructure + + private sealed class RecordedRequest + { + public RecordedRequest(string method, string path, string body) + { + Method = method; + Path = path; + Body = body; + } + + public string Method { get; } + public string Path { get; } + public string Body { get; } + } + + private sealed class MockHandler : HttpMessageHandler + { + private readonly List<(string Method, string PathContains, Func Resp)> _routes = new(); + public List Requests { get; } = new(); + public string DefaultBody { get; set; } = "{\"code\":0,\"data\":{}}"; + + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, _ => resp())); + return this; + } + + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, resp)); + return this; + } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + var body = request.Content != null + ? await request.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + var path = request.RequestUri!.AbsolutePath; + Requests.Add(new RecordedRequest(request.Method.Method, request.RequestUri.PathAndQuery, body)); + + foreach (var route in _routes) + { + if (route.Method == request.Method.Method && path.Contains(route.PathContains)) + return route.Resp(body); + } + + return Json(DefaultBody); + } + } + + private static HttpResponseMessage Json(string body, HttpStatusCode code = HttpStatusCode.OK) + => new HttpResponseMessage(code) { Content = new StringContent(body, Encoding.UTF8, "application/json") }; + + private static MilvusDocumentStore StoreWith(MockHandler handler, + DistanceMetricType metric = DistanceMetricType.Cosine, int vectorDimension = 0) + { + var client = new HttpClient(handler) { BaseAddress = new Uri(BaseUrl) }; + return new MilvusDocumentStore(Collection, BaseUrl, token: "t", + distanceMetric: metric, vectorDimension: vectorDimension, httpClient: client); + } + + private static VectorDocument Doc(string id, string content, float[] vector, Dictionary? metadata = null) + { + var d = new Document(id, content, metadata ?? new Dictionary()); + return new VectorDocument { Document = d, Embedding = new Vector(vector) }; + } + + #endregion + + [Fact] + public void Constructor_ReadsExistingCollection_SetsDimension() + { + var handler = new MockHandler() + .On("POST", "/collections/describe", () => Json(ExistingDescribeJson)); + + var store = StoreWith(handler); + + Assert.Equal(3, store.VectorDimension); + Assert.Equal(Collection, store.CollectionName); + } + + [Fact] + public void Constructor_CollectionMissing_CreatesWithMappedMetric() + { + var handler = new MockHandler() + .On("POST", "/collections/describe", () => Json("{\"code\":100,\"message\":\"collection not found\"}")) + .On("POST", "/collections/create", () => Json("{\"code\":0,\"data\":{}}")); + + var store = StoreWith(handler, metric: DistanceMetricType.Euclidean, vectorDimension: 4); + + Assert.Equal(4, store.VectorDimension); + var create = handler.Requests.Single(r => r.Path.Contains("/collections/create")); + var body = JObject.Parse(create.Body); + Assert.Equal(4, (int)body["dimension"]!); + Assert.Equal("L2", (string?)body["metricType"]); + Assert.Equal("VarChar", (string?)body["idType"]); + Assert.True((bool)body["enableDynamicField"]!); + } + + [Fact] + public void Add_UpsertsEntity_SendsCorrectRequest() + { + var handler = new MockHandler() + .On("POST", "/collections/describe", () => Json(ExistingDescribeJson)) + .On("POST", "/entities/upsert", () => Json("{\"code\":0,\"data\":{\"upsertCount\":1}}")); + + var store = StoreWith(handler); + store.Add(Doc("doc1", "Hello world", new float[] { 1, 0, 0 }, + new Dictionary { ["category"] = "science" })); + + Assert.Equal(1, store.DocumentCount); + var upsert = handler.Requests.Single(r => r.Path.Contains("/entities/upsert")); + var body = JObject.Parse(upsert.Body); + Assert.Equal("test", (string?)body["collectionName"]); + var row = body["data"]![0]!; + Assert.Equal("doc1", (string?)row["id"]); + Assert.Equal("Hello world", (string?)row["content"]); + Assert.Equal("science", (string?)row["m_category"]); + Assert.Equal(3, ((JArray)row["vector"]!).Count); + } + + [Fact] + public void GetSimilar_ParsesRankedHits_InOrder() + { + var searchResponse = + "{\"code\":0,\"data\":[" + + "{\"id\":\"doc1\",\"distance\":0.98,\"content\":\"first\",\"metadata_json\":\"{\\\"category\\\":\\\"science\\\"}\"}," + + "{\"id\":\"doc2\",\"distance\":0.55,\"content\":\"second\",\"metadata_json\":\"{}\"}" + + "]}"; + + var handler = new MockHandler() + .On("POST", "/entities/search", () => Json(searchResponse)) + .On("POST", "/collections/describe", () => Json(ExistingDescribeJson)); + + var store = StoreWith(handler); + var results = store.GetSimilar(new Vector(new float[] { 1, 0, 0 }), topK: 2).ToList(); + + Assert.Equal(2, results.Count); + Assert.Equal("doc1", results[0].Id); + Assert.Equal("first", results[0].Content); + Assert.True(results[0].HasRelevanceScore); + Assert.True(Convert.ToDouble(results[0].RelevanceScore) > Convert.ToDouble(results[1].RelevanceScore)); + + var search = handler.Requests.Single(r => r.Path.Contains("/entities/search")); + var body = JObject.Parse(search.Body); + Assert.Equal(2, (int)body["limit"]!); + Assert.Equal("vector", (string?)body["annsField"]); + Assert.Equal(3, ((JArray)body["data"]![0]!).Count); + } + + [Fact] + public void GetSimilarWithFilters_TranslatesEqualityRangeAndAnyOf() + { + var handler = new MockHandler() + .On("POST", "/entities/search", () => Json("{\"code\":0,\"data\":[]}")) + .On("POST", "/collections/describe", () => Json(ExistingDescribeJson)); + + var store = StoreWith(handler); + var filters = new Dictionary + { + ["category"] = "science", + ["year"] = 2020, + ["tag"] = new List { "a", "b" } + }; + store.GetSimilarWithFilters(new Vector(new float[] { 1, 0, 0 }), 10, filters).ToList(); + + var expr = (string)JObject.Parse(handler.Requests.Single(r => r.Path.Contains("/entities/search")).Body)["filter"]!; + Assert.Contains("m_category == \"science\"", expr); + Assert.Contains("m_year >= 2020", expr); + Assert.Contains("m_tag in [\"a\", \"b\"]", expr); + Assert.Contains(" and ", expr); + } + + [Fact] + public void GetById_Found_ReturnsDocument() + { + var handler = new MockHandler() + .On("POST", "/entities/query", () => Json( + "{\"code\":0,\"data\":[{\"id\":\"doc1\",\"content\":\"hi\",\"metadata_json\":\"{\\\"k\\\":\\\"v\\\"}\"}]}")) + .On("POST", "/collections/describe", () => Json(ExistingDescribeJson)); + + var store = StoreWith(handler); + var doc = store.GetById("doc1"); + + Assert.NotNull(doc); + Assert.Equal("doc1", doc!.Id); + Assert.Equal("hi", doc.Content); + Assert.Equal("v", doc.Metadata["k"].ToString()); + + var query = handler.Requests.Single(r => r.Path.Contains("/entities/query")); + Assert.Contains("id == \"doc1\"", (string)JObject.Parse(query.Body)["filter"]!); + } + + [Fact] + public void GetById_NotFound_ReturnsNull() + { + var handler = new MockHandler() + .On("POST", "/entities/query", () => Json("{\"code\":0,\"data\":[]}")) + .On("POST", "/collections/describe", () => Json(ExistingDescribeJson)); + + var store = StoreWith(handler); + Assert.Null(store.GetById("missing")); + } + + [Fact] + public void Remove_SendsDelete_ReturnsTrue() + { + var handler = new MockHandler() + .On("POST", "/entities/delete", () => Json("{\"code\":0,\"data\":{}}")) + .On("POST", "/collections/describe", () => Json(ExistingDescribeJson)); + + var store = StoreWith(handler); + Assert.True(store.Remove("doc1")); + + var delete = handler.Requests.Single(r => r.Path.Contains("/entities/delete")); + Assert.Contains("id == \"doc1\"", (string)JObject.Parse(delete.Body)["filter"]!); + } + + [Fact] + public void GetAll_PagesUntilExhausted() + { + var call = 0; + var handler = new MockHandler() + .On("POST", "/entities/query", _ => + { + call++; + if (call == 1) + { + var items = string.Join(",", Enumerable.Range(0, 1000).Select(i => + "{\"id\":\"d" + i + "\",\"content\":\"c\",\"metadata_json\":\"{}\"}")); + return Json("{\"code\":0,\"data\":[" + items + "]}"); + } + return Json("{\"code\":0,\"data\":[{\"id\":\"dZ\",\"content\":\"c\",\"metadata_json\":\"{}\"}]}"); + }) + .On("POST", "/collections/describe", () => Json(ExistingDescribeJson)); + + var store = StoreWith(handler); + var all = store.GetAll().ToList(); + + Assert.Equal(1001, all.Count); + Assert.Contains(all, d => d.Id == "dZ"); + Assert.Equal(2, handler.Requests.Count(r => r.Path.Contains("/entities/query"))); + } + + [Fact] + public void Clear_DropsCollection_AndRecreates() + { + var handler = new MockHandler() + .On("POST", "/collections/drop", () => Json("{\"code\":0,\"data\":{}}")) + .On("POST", "/collections/create", () => Json("{\"code\":0,\"data\":{}}")) + .On("POST", "/collections/describe", () => Json(ExistingDescribeJson)); + + // Existing collection reports dim 3, so Clear recreates it. + var store = StoreWith(handler); + store.Clear(); + + Assert.Equal(0, store.DocumentCount); + Assert.Contains(handler.Requests, r => r.Path.Contains("/collections/drop")); + Assert.Contains(handler.Requests, r => r.Path.Contains("/collections/create")); + } + + [Fact] + public void Constructor_UnsupportedMetric_Throws() + { + Assert.Throws(() => + new MilvusDocumentStore(Collection, BaseUrl, distanceMetric: DistanceMetricType.Jaccard, + httpClient: new HttpClient(new MockHandler()) { BaseAddress = new Uri(BaseUrl) })); + } + + #region Integration (gated - require a live Milvus instance) + + [Trait("Category", "Integration")] + [SkippableFact] + public void Integration_UpsertSearchDelete_RoundTrips() + { + var url = Environment.GetEnvironmentVariable("MILVUS_URL"); + Skip.If(string.IsNullOrWhiteSpace(url), + "Set MILVUS_URL (and optionally MILVUS_TOKEN) to run Milvus integration tests."); + + var token = Environment.GetEnvironmentVariable("MILVUS_TOKEN"); + var collection = "aidotnet_it_" + Guid.NewGuid().ToString("N"); + var store = new MilvusDocumentStore(collection, url!, token, + distanceMetric: DistanceMetricType.Cosine, vectorDimension: 4); + + try + { + var id = "it-" + Guid.NewGuid().ToString("N"); + store.Add(Doc(id, "integration content", new float[] { 1, 0, 0, 0 }, + new Dictionary { ["category"] = "test" })); + + var results = store.GetSimilar(new Vector(new float[] { 1, 0, 0, 0 }), topK: 5).ToList(); + Assert.Contains(results, r => r.Id == id); + + Assert.True(store.Remove(id)); + } + finally + { + store.Clear(); + } + } + + #endregion + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStoreTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStoreTests.cs new file mode 100644 index 0000000000..09c65de767 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStoreTests.cs @@ -0,0 +1,353 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Threading; +using System.Threading.Tasks; +using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json.Linq; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.DocumentStores +{ + /// + /// Unit tests for the real Weaviate REST/GraphQL client. All HTTP traffic is intercepted by a mock + /// so these run in CI with no network access. + /// + public class WeaviateDocumentStoreTests + { + private const string BaseUrl = "http://localhost:8080"; + private const string Class = "Article"; + + private const string ExistingSchemaJson = + "{\"class\":\"Article\",\"vectorIndexConfig\":{\"distance\":\"cosine\"}}"; + + #region infrastructure + + private sealed class RecordedRequest + { + public RecordedRequest(string method, string path, string body) + { + Method = method; + Path = path; + Body = body; + } + + public string Method { get; } + public string Path { get; } + public string Body { get; } + } + + private sealed class MockHandler : HttpMessageHandler + { + private readonly List<(string Method, string PathContains, Func Resp)> _routes = new(); + public List Requests { get; } = new(); + public string DefaultBody { get; set; } = "{}"; + + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, _ => resp())); + return this; + } + + public MockHandler On(string method, string pathContains, Func resp) + { + _routes.Add((method, pathContains, resp)); + return this; + } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + var body = request.Content != null + ? await request.Content.ReadAsStringAsync().ConfigureAwait(false) + : string.Empty; + var path = request.RequestUri!.AbsolutePath; + Requests.Add(new RecordedRequest(request.Method.Method, request.RequestUri.PathAndQuery, body)); + + foreach (var route in _routes) + { + if (route.Method == request.Method.Method && path.Contains(route.PathContains)) + return route.Resp(body); + } + + return Json(DefaultBody); + } + } + + private static HttpResponseMessage Json(string body, HttpStatusCode code = HttpStatusCode.OK) + => new HttpResponseMessage(code) { Content = new StringContent(body, Encoding.UTF8, "application/json") }; + + private static WeaviateDocumentStore StoreWith(MockHandler handler, + DistanceMetricType metric = DistanceMetricType.Cosine) + { + var client = new HttpClient(handler) { BaseAddress = new Uri(BaseUrl) }; + return new WeaviateDocumentStore(Class, BaseUrl, apiKey: "k", distanceMetric: metric, httpClient: client); + } + + private static VectorDocument Doc(string id, string content, float[] vector, Dictionary? metadata = null) + { + var d = new Document(id, content, metadata ?? new Dictionary()); + return new VectorDocument { Document = d, Embedding = new Vector(vector) }; + } + + #endregion + + [Fact] + public void Constructor_CreatesClass_WhenMissing_WithMappedDistance() + { + var handler = new MockHandler() + .On("GET", "/v1/schema/Article", () => Json("{}", HttpStatusCode.NotFound)) + .On("POST", "/v1/schema", () => Json("{}")); + + var store = StoreWith(handler, metric: DistanceMetricType.Euclidean); + + Assert.Equal(Class, store.ClassName); + var create = handler.Requests.Single(r => r.Method == "POST" && r.Path == "/v1/schema"); + var body = JObject.Parse(create.Body); + Assert.Equal("Article", (string?)body["class"]); + Assert.Equal("none", (string?)body["vectorizer"]); + Assert.Equal("l2-squared", (string?)body["vectorIndexConfig"]!["distance"]); + } + + [Fact] + public void Add_CreatesObject_WithPropertiesAndVector() + { + var handler = new MockHandler() + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)) + .On("POST", "/v1/objects", () => Json("{}")); + + var store = StoreWith(handler); + store.Add(Doc("doc1", "Hello world", new float[] { 1, 0, 0 }, + new Dictionary { ["category"] = "science" })); + + Assert.Equal(1, store.DocumentCount); + var upsert = handler.Requests.Single(r => r.Method == "POST" && r.Path == "/v1/objects"); + var body = JObject.Parse(upsert.Body); + Assert.Equal("Article", (string?)body["class"]); + Assert.Equal("doc1", (string?)body["properties"]!["docId"]); + Assert.Equal("Hello world", (string?)body["properties"]!["content"]); + Assert.Equal("science", (string?)body["properties"]!["m_category"]); + Assert.Equal(3, ((JArray)body["vector"]!).Count); + } + + [Fact] + public void AddBatch_SendsAllObjectsInOneBatchRequest() + { + var handler = new MockHandler() + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)) + .On("POST", "/v1/batch/objects", () => Json("[]")); + + var store = StoreWith(handler); + store.AddBatch(new List> + { + Doc("d1", "c1", new float[] { 1, 0, 0 }), + Doc("d2", "c2", new float[] { 0, 1, 0 }) + }); + + Assert.Equal(2, store.DocumentCount); + var batch = handler.Requests.Single(r => r.Method == "POST" && r.Path.Contains("/v1/batch/objects")); + var objs = (JArray)JObject.Parse(batch.Body)["objects"]!; + Assert.Equal(2, objs.Count); + } + + [Fact] + public void GetSimilar_ParsesGraphQlHits_InOrder() + { + var searchResponse = + "{\"data\":{\"Get\":{\"Article\":[" + + "{\"docId\":\"doc1\",\"content\":\"first\",\"metadataJson\":\"{\\\"category\\\":\\\"science\\\"}\",\"_additional\":{\"certainty\":0.98,\"distance\":0.04}}," + + "{\"docId\":\"doc2\",\"content\":\"second\",\"metadataJson\":\"{}\",\"_additional\":{\"certainty\":0.55,\"distance\":0.9}}" + + "]}}}"; + + var handler = new MockHandler() + .On("POST", "/v1/graphql", () => Json(searchResponse)) + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)); + + var store = StoreWith(handler); + var results = store.GetSimilar(new Vector(new float[] { 1, 0, 0 }), topK: 2).ToList(); + + Assert.Equal(2, results.Count); + Assert.Equal("doc1", results[0].Id); + Assert.Equal("first", results[0].Content); + Assert.Equal("science", results[0].Metadata["category"].ToString()); + Assert.True(results[0].HasRelevanceScore); + Assert.True(Convert.ToDouble(results[0].RelevanceScore) > Convert.ToDouble(results[1].RelevanceScore)); + + var search = handler.Requests.Single(r => r.Path.Contains("/v1/graphql")); + var query = (string)JObject.Parse(search.Body)["query"]!; + Assert.Contains("nearVector", query); + Assert.Contains("Get { Article", query); + Assert.Contains("limit: 2", query); + } + + [Fact] + public void GetSimilarWithFilters_TranslatesEqualityAndRange() + { + var handler = new MockHandler() + .On("POST", "/v1/graphql", () => Json("{\"data\":{\"Get\":{\"Article\":[]}}}")) + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)); + + var store = StoreWith(handler); + var filters = new Dictionary { ["category"] = "science", ["year"] = 2020 }; + store.GetSimilarWithFilters(new Vector(new float[] { 1, 0, 0 }), 10, filters).ToList(); + + var search = handler.Requests.Single(r => r.Path.Contains("/v1/graphql")); + var query = (string)JObject.Parse(search.Body)["query"]!; + + Assert.Contains("operator: And", query); + Assert.Contains("[\"m_category\"]", query); + Assert.Contains("operator: Equal", query); + Assert.Contains("valueText: \"science\"", query); + Assert.Contains("[\"m_year\"]", query); + Assert.Contains("operator: GreaterThanEqual", query); + Assert.Contains("valueNumber: 2020", query); + } + + [Fact] + public void GetSimilarWithFilters_TranslatesAnyOfCollection() + { + var handler = new MockHandler() + .On("POST", "/v1/graphql", () => Json("{\"data\":{\"Get\":{\"Article\":[]}}}")) + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)); + + var store = StoreWith(handler); + var filters = new Dictionary { ["tag"] = new List { "a", "b" } }; + store.GetSimilarWithFilters(new Vector(new float[] { 1, 0, 0 }), 5, filters).ToList(); + + var query = (string)JObject.Parse(handler.Requests.Single(r => r.Path.Contains("/v1/graphql")).Body)["query"]!; + Assert.Contains("operator: Or", query); + Assert.Contains("valueText: \"a\"", query); + Assert.Contains("valueText: \"b\"", query); + } + + [Fact] + public void GetById_Found_ReturnsDocument() + { + var handler = new MockHandler() + .On("GET", "/v1/objects/Article/", () => Json( + "{\"id\":\"uuid\",\"properties\":{\"docId\":\"doc1\",\"content\":\"hi\",\"metadataJson\":\"{\\\"k\\\":\\\"v\\\"}\"}}")) + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)); + + var store = StoreWith(handler); + var doc = store.GetById("doc1"); + + Assert.NotNull(doc); + Assert.Equal("doc1", doc!.Id); + Assert.Equal("hi", doc.Content); + Assert.Equal("v", doc.Metadata["k"].ToString()); + } + + [Fact] + public void GetById_NotFound_ReturnsNull() + { + var handler = new MockHandler() + .On("GET", "/v1/objects/Article/", () => Json("{}", HttpStatusCode.NotFound)) + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)); + + var store = StoreWith(handler); + Assert.Null(store.GetById("missing")); + } + + [Fact] + public void Remove_SendsDelete_ReturnsTrue() + { + var handler = new MockHandler() + .On("DELETE", "/v1/objects/Article/", () => Json("{}", HttpStatusCode.NoContent)) + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)); + + var store = StoreWith(handler); + Assert.True(store.Remove("doc1")); + Assert.Contains(handler.Requests, r => r.Method == "DELETE" && r.Path.Contains("/v1/objects/Article/")); + } + + [Fact] + public void GetAll_CursorPagesUntilExhausted() + { + var call = 0; + var handler = new MockHandler() + .On("GET", "/v1/objects", _ => + { + call++; + // First page returns 100 objects (full page) to force a second request. + if (call == 1) + { + var items = string.Join(",", Enumerable.Range(0, 100).Select(i => + "{\"id\":\"u" + i + "\",\"properties\":{\"docId\":\"d" + i + "\",\"content\":\"c\",\"metadataJson\":\"{}\"}}")); + return Json("{\"objects\":[" + items + "]}"); + } + return Json("{\"objects\":[{\"id\":\"uZ\",\"properties\":{\"docId\":\"dZ\",\"content\":\"c\",\"metadataJson\":\"{}\"}}]}"); + }) + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)); + + var store = StoreWith(handler); + var all = store.GetAll().ToList(); + + Assert.Equal(101, all.Count); + Assert.Contains(all, d => d.Id == "dZ"); + Assert.Equal(2, handler.Requests.Count(r => r.Method == "GET" && r.Path.StartsWith("/v1/objects"))); + } + + [Fact] + public void Clear_DeletesSchema_AndRecreates() + { + var handler = new MockHandler() + .On("DELETE", "/v1/schema/Article", () => Json("{}")) + .On("POST", "/v1/schema", () => Json("{}")) + .On("GET", "/v1/schema/Article", () => Json(ExistingSchemaJson)); + + var store = StoreWith(handler); + store.Clear(); + + Assert.Equal(0, store.DocumentCount); + Assert.Contains(handler.Requests, r => r.Method == "DELETE" && r.Path.Contains("/v1/schema/Article")); + Assert.Contains(handler.Requests, r => r.Method == "POST" && r.Path == "/v1/schema"); + } + + [Fact] + public void Constructor_UnsupportedMetric_Throws() + { + Assert.Throws(() => + new WeaviateDocumentStore(Class, BaseUrl, distanceMetric: DistanceMetricType.Mahalanobis, + httpClient: new HttpClient(new MockHandler()) { BaseAddress = new Uri(BaseUrl) })); + } + + #region Integration (gated - require a live Weaviate instance) + + [Trait("Category", "Integration")] + [SkippableFact] + public void Integration_UpsertSearchDelete_RoundTrips() + { + var url = Environment.GetEnvironmentVariable("WEAVIATE_URL"); + Skip.If(string.IsNullOrWhiteSpace(url), + "Set WEAVIATE_URL (and optionally WEAVIATE_API_KEY) to run Weaviate integration tests."); + + var apiKey = Environment.GetEnvironmentVariable("WEAVIATE_API_KEY"); + var className = "Aidotnet_it_" + Guid.NewGuid().ToString("N"); + var store = new WeaviateDocumentStore(className, url!, apiKey, + distanceMetric: DistanceMetricType.Cosine); + + try + { + var id = "it-" + Guid.NewGuid().ToString("N"); + store.Add(Doc(id, "integration content", new float[] { 1, 0, 0, 0 }, + new Dictionary { ["category"] = "test" })); + + var results = store.GetSimilar(new Vector(new float[] { 1, 0, 0, 0 }), topK: 5).ToList(); + Assert.Contains(results, r => r.Id == id); + + Assert.True(store.Remove(id)); + } + finally + { + store.Clear(); + } + } + + #endregion + } +} From ed4ec8e0cc39ed4e1534dabd80e0e4331ebbe515 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 15:52:40 -0400 Subject: [PATCH 07/25] feat(rag): token-aware chunking strategy (#26 part 1) RAG chunking was character-based only, but LLM context/embedding limits are measured in tokens. Add TokenBasedChunkingStrategy: ChunkSize/ChunkOverlap are interpreted as tokens, with an injectable token-counter delegate so callers can plug an exact tokenizer (tiktoken/HF); the default counts whitespace words as an approximate proxy. Overlap and start/end positions preserved. (Remaining in #26: real embedding-based semantic chunking + LLM agentic chunking.) Verified: 4 tests (budget, overlap, position round-trip, custom counter). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../TokenBasedChunkingStrategy.cs | 101 ++++++++++++++++++ .../TokenBasedChunkingStrategyTests.cs | 62 +++++++++++ 2 files changed, 163 insertions(+) create mode 100644 src/RetrievalAugmentedGeneration/ChunkingStrategies/TokenBasedChunkingStrategy.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ChunkingStrategies/TokenBasedChunkingStrategyTests.cs diff --git a/src/RetrievalAugmentedGeneration/ChunkingStrategies/TokenBasedChunkingStrategy.cs b/src/RetrievalAugmentedGeneration/ChunkingStrategies/TokenBasedChunkingStrategy.cs new file mode 100644 index 0000000000..8c5d03c699 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/ChunkingStrategies/TokenBasedChunkingStrategy.cs @@ -0,0 +1,101 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Attributes; +using AiDotNet.Enums; + +namespace AiDotNet.RetrievalAugmentedGeneration.ChunkingStrategies +{ + /// + /// Chunks text by TOKEN count rather than character count. + /// and are interpreted as tokens. + /// + /// + /// + /// LLM context windows and embedding limits are measured in tokens, not characters, so token-aware + /// splitting (LangChain TokenTextSplitter, LlamaIndex token node parser) is the correct default + /// for RAG. A token-counter delegate is injectable so callers can plug an exact tokenizer (tiktoken / + /// HF); the built-in default counts whitespace-separated words as an approximate proxy. + /// + /// For Beginners: models "see" text as tokens (~¾ of a word each). This splitter keeps each + /// chunk under a token budget so it always fits the model, instead of guessing with character counts. + /// + [ComponentType(ComponentType.Chunker)] + [PipelineStage(PipelineStage.DataIngestion)] + public class TokenBasedChunkingStrategy : ChunkingStrategyBase + { + private readonly Func _tokenCounter; + + /// Maximum tokens per chunk. + /// Approximate token overlap between consecutive chunks. + /// + /// Optional exact token counter (e.g. a tiktoken/HF adapter). When null, whitespace word count is + /// used as an approximate proxy; with a custom multi-token-per-word counter the overlap is approximate. + /// + public TokenBasedChunkingStrategy(int maxTokens = 256, int overlapTokens = 32, Func? tokenCounter = null) + : base(maxTokens, overlapTokens) + { + _tokenCounter = tokenCounter ?? DefaultWordCount; + } + + private static int DefaultWordCount(string s) + => string.IsNullOrWhiteSpace(s) ? 0 : s.Split((char[]?)null, StringSplitOptions.RemoveEmptyEntries).Length; + + protected override IEnumerable<(string Chunk, int StartPosition, int EndPosition)> ChunkCore(string text) + { + var words = SplitWordsWithPositions(text); + if (words.Count == 0) + { + yield break; + } + + int chunkStart = 0; + while (chunkStart < words.Count) + { + int end = chunkStart; + int tokens = 0; + while (end < words.Count) + { + int wt = Math.Max(1, _tokenCounter(words[end].Word)); + // Always take at least one word so a single over-budget word can't stall progress. + if (end > chunkStart && tokens + wt > ChunkSize) + { + break; + } + + tokens += wt; + end++; + } + + int startPos = words[chunkStart].Start; + int endPos = words[end - 1].End; + yield return (text.Substring(startPos, endPos - startPos), startPos, endPos); + + if (end >= words.Count) + { + break; + } + + // Advance, leaving ~ChunkOverlap tokens of overlap. Guaranteed forward progress (>= 1 word). + int wordsInChunk = end - chunkStart; + int advance = Math.Max(1, wordsInChunk - Math.Min(ChunkOverlap, wordsInChunk - 1)); + chunkStart += advance; + } + } + + private static List<(string Word, int Start, int End)> SplitWordsWithPositions(string text) + { + var words = new List<(string, int, int)>(); + int i = 0; + while (i < text.Length) + { + while (i < text.Length && char.IsWhiteSpace(text[i])) i++; + if (i >= text.Length) break; + int start = i; + while (i < text.Length && !char.IsWhiteSpace(text[i])) i++; + words.Add((text.Substring(start, i - start), start, i)); + } + + return words; + } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ChunkingStrategies/TokenBasedChunkingStrategyTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ChunkingStrategies/TokenBasedChunkingStrategyTests.cs new file mode 100644 index 0000000000..0fa32ec460 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ChunkingStrategies/TokenBasedChunkingStrategyTests.cs @@ -0,0 +1,62 @@ +using System.Linq; +using AiDotNet.RetrievalAugmentedGeneration.ChunkingStrategies; +using Xunit; + +namespace AiDotNet.Tests.UnitTests.RetrievalAugmentedGeneration.ChunkingStrategies; + +/// +/// Tests for — token-budgeted chunking with overlap and an +/// injectable token counter. +/// +public class TokenBasedChunkingStrategyTests +{ + private static int WordCount(string s) => s.Split((char[]?)null, System.StringSplitOptions.RemoveEmptyEntries).Length; + + [Fact] + public void RespectsMaxTokenBudget_WithDefaultWordCounter() + { + var strat = new TokenBasedChunkingStrategy(maxTokens: 5, overlapTokens: 2); + var text = "one two three four five six seven eight nine ten eleven twelve"; + + var chunks = strat.Chunk(text).ToList(); + + Assert.True(chunks.Count > 1); + Assert.All(chunks, c => Assert.True(WordCount(c) <= 5, $"chunk exceeded budget: '{c}'")); + } + + [Fact] + public void ConsecutiveChunksOverlap() + { + var strat = new TokenBasedChunkingStrategy(maxTokens: 5, overlapTokens: 2); + var chunks = strat.Chunk("one two three four five six seven eight").ToList(); + + // Chunk1 = one..five, Chunk2 = four..eight → overlap "four five". + Assert.Contains("four", chunks[0]); + Assert.Contains("five", chunks[0]); + Assert.Contains("four", chunks[1]); + Assert.Contains("five", chunks[1]); + } + + [Fact] + public void PositionsMapBackToOriginalText() + { + var strat = new TokenBasedChunkingStrategy(maxTokens: 4, overlapTokens: 1); + var text = "alpha beta gamma delta epsilon zeta"; + + foreach (var (chunk, start, end) in strat.ChunkWithPositions(text)) + { + Assert.Equal(chunk, text.Substring(start, end - start)); + } + } + + [Fact] + public void HonorsCustomTokenCounter() + { + // Counter that treats every word as 3 tokens → budget of 6 tokens allows ~2 words per chunk. + var strat = new TokenBasedChunkingStrategy(maxTokens: 6, overlapTokens: 0, tokenCounter: s => WordCount(s) * 3); + var chunks = strat.Chunk("a b c d e f").ToList(); + + Assert.All(chunks, c => Assert.True(WordCount(c) <= 2, $"expected <=2 words/chunk, got '{c}'")); + Assert.True(chunks.Count >= 3); + } +} From e5c2d58b541a9656f6bd0423156c8013a0b1572a Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 16:06:44 -0400 Subject: [PATCH 08/25] feat(rag): real ONNX cross-encoder reranker Add OnnxCrossEncoderReranker that loads a cross-encoder ONNX model (e.g. BGE-reranker / ms-marco MiniLM) plus its tokenizer and scores each (query, document) pair by joint [CLS] q [SEP] d [SEP] inference, sorting descending and setting RelevanceScore/HasRelevanceScore. Batches pairs, disposes the session, and throws a clear FileNotFoundException on a missing model (no silent lexical fallback). Adds a convenience CrossEncoderReranker ctor that accepts it as the scorer. Ranking logic is unit-tested via a virtual ScorePairs seam; a gated integration test loads a real model from an env var and skips when unset. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../Rerankers/CrossEncoderReranker.cs | 20 + .../Rerankers/OnnxCrossEncoderReranker.cs | 539 ++++++++++++++++++ .../OnnxCrossEncoderRerankerTests.cs | 414 ++++++++++++++ 3 files changed, 973 insertions(+) create mode 100644 src/RetrievalAugmentedGeneration/Rerankers/OnnxCrossEncoderReranker.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RAG/Rerankers/OnnxCrossEncoderRerankerTests.cs diff --git a/src/RetrievalAugmentedGeneration/Rerankers/CrossEncoderReranker.cs b/src/RetrievalAugmentedGeneration/Rerankers/CrossEncoderReranker.cs index 45dc92e719..a84bfafb63 100644 --- a/src/RetrievalAugmentedGeneration/Rerankers/CrossEncoderReranker.cs +++ b/src/RetrievalAugmentedGeneration/Rerankers/CrossEncoderReranker.cs @@ -107,6 +107,26 @@ public CrossEncoderReranker(Func scoreFunction, int maxPairsT _maxPairsToScore = maxPairsToScore; } + /// + /// Initializes a new instance of the CrossEncoderReranker class backed by a real ONNX cross-encoder model. + /// + /// A configured used to score (query, document) pairs. + /// Maximum number of query-document pairs to score (default: 20). + /// + /// For Beginners: This is the convenient, works-out-of-the-box path. + /// + /// Instead of writing your own scoring function, hand this constructor an + /// pointing at a real cross-encoder model + /// (e.g. BGE-reranker or ms-marco-MiniLM) and it will use that model to score pairs. + /// + /// + public CrossEncoderReranker(OnnxCrossEncoderReranker onnxReranker, int maxPairsToScore = 20) + : this( + (onnxReranker ?? throw new ArgumentNullException(nameof(onnxReranker))).ScorePair, + maxPairsToScore) + { + } + /// /// Reranks documents using the cross-encoder model. /// diff --git a/src/RetrievalAugmentedGeneration/Rerankers/OnnxCrossEncoderReranker.cs b/src/RetrievalAugmentedGeneration/Rerankers/OnnxCrossEncoderReranker.cs new file mode 100644 index 0000000000..e6a2732d98 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Rerankers/OnnxCrossEncoderReranker.cs @@ -0,0 +1,539 @@ +using System; +using System.Collections.Generic; +using System.IO; +using System.Linq; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using AiDotNet.Tokenization.HuggingFace; +using AiDotNet.Tokenization.Interfaces; +using Microsoft.ML.OnnxRuntime; + +namespace AiDotNet.RetrievalAugmentedGeneration.Rerankers; + +/// +/// A production-ready cross-encoder reranker that scores (query, document) pairs with a real +/// transformer model loaded via ONNX Runtime (e.g. BAAI/bge-reranker or ms-marco-MiniLM). +/// +/// The numeric data type used for scoring. +/// +/// +/// Unlike (which requires the caller to supply a +/// Func<string, string, T> scorer), this reranker works out of the box: point it at a +/// cross-encoder ONNX model and its tokenizer, and it will tokenize each (query, document) pair +/// jointly ([CLS] query [SEP] document [SEP]), run the model to produce a relevance logit, +/// and reorder documents by that logit (highest first). +/// +/// For Beginners: This is the "batteries-included" cross-encoder reranker. +/// +/// A cross-encoder reads the query and a document together and outputs a single number telling you +/// how relevant the document is to the query. Higher means more relevant. This class handles all of +/// the plumbing: +/// - Loading the model (ONNX file) and its tokenizer. +/// - Building the joint input the model expects for each pair. +/// - Running the model (in batches for speed). +/// - Sorting your documents from most to least relevant. +/// +/// Typical models to use: +/// - BAAI/bge-reranker-base / bge-reranker-large (single relevance logit). +/// - cross-encoder/ms-marco-MiniLM-L-6-v2 (single relevance logit). +/// +/// Model loading is lazy: nothing is opened until the first rerank call. If the model file is +/// missing, a clear is thrown — there is no silent lexical +/// fallback. +/// +/// +[Attributes.ComponentType(Enums.ComponentType.Reranker)] +[Attributes.PipelineStage(Enums.PipelineStage.PostRetrieval)] +public class OnnxCrossEncoderReranker : RerankerBase, IDisposable +{ + private readonly string _modelPath; + private readonly string _tokenizerPath; + private readonly int _maxLength; + private readonly int _maxPairsToScore; + private readonly int _batchSize; + private readonly int _scoreLabelIndex; + + private volatile InferenceSession? _session; + private volatile ITokenizer? _tokenizer; + private readonly object _initLock = new(); + private bool _disposed; + + /// + /// Gets a value indicating whether this reranker modifies relevance scores. + /// + public override bool ModifiesScores => true; + + /// + /// Initializes a new instance of the class. + /// + /// Path to the cross-encoder ONNX model file (e.g. model.onnx). + /// + /// Path to the tokenizer. May be a directory containing the tokenizer files + /// (tokenizer.json / vocab.txt + tokenizer_config.json), or a path to one of + /// those files (the containing directory is used). + /// + /// Maximum joint sequence length (query + document); default 512. + /// Maximum number of query-document pairs to score; default 100. + /// Number of pairs to run through the model per inference call; default 16. + /// + /// The output label index to read as the relevance score. Use -1 (default) to read the last label, + /// which is correct for single-logit rerankers (index 0) and for binary classifiers whose positive + /// class is last. Set explicitly if your model orders labels differently. + /// + /// + /// For Beginners: The two paths you must supply are the model file and its tokenizer. + /// Everything else has sensible defaults. Loading is deferred until the first rerank, so + /// constructing this object is cheap and never throws just because a file is missing. + /// + /// + public OnnxCrossEncoderReranker( + string modelPath, + string tokenizerPath, + int maxLength = 512, + int maxPairsToScore = 100, + int batchSize = 16, + int scoreLabelIndex = -1) + { + if (string.IsNullOrWhiteSpace(modelPath)) + throw new ArgumentException("Model path cannot be empty", nameof(modelPath)); + if (string.IsNullOrWhiteSpace(tokenizerPath)) + throw new ArgumentException("Tokenizer path cannot be empty", nameof(tokenizerPath)); + if (maxLength <= 0) + throw new ArgumentException("maxLength must be greater than zero", nameof(maxLength)); + if (maxPairsToScore <= 0) + throw new ArgumentException("maxPairsToScore must be greater than zero", nameof(maxPairsToScore)); + if (batchSize <= 0) + throw new ArgumentException("batchSize must be greater than zero", nameof(batchSize)); + + _modelPath = modelPath; + _tokenizerPath = tokenizerPath; + _maxLength = maxLength; + _maxPairsToScore = maxPairsToScore; + _batchSize = batchSize; + _scoreLabelIndex = scoreLabelIndex; + } + + /// + /// Reranks documents by running the cross-encoder model on each (query, document) pair. + /// + /// The validated search query. + /// The validated, materialized documents to rerank. + /// Documents reordered by descending cross-encoder relevance score. + protected override IEnumerable> RerankCore(string query, IList> documents) + { + var docList = documents.Take(_maxPairsToScore).ToList(); + if (docList.Count == 0) + return Enumerable.Empty>(); + + var contents = new List(docList.Count); + for (int i = 0; i < docList.Count; i++) + { + contents.Add(docList[i].Content ?? string.Empty); + } + + var scores = ScorePairs(query, contents); + if (scores == null || scores.Length != docList.Count) + { + throw new InvalidOperationException( + "ScorePairs must return exactly one score per document."); + } + + var scored = new List<(Document Doc, double Score)>(docList.Count); + for (int i = 0; i < docList.Count; i++) + { + scored.Add((docList[i], scores[i])); + } + + var reranked = scored + .OrderByDescending(x => x.Score) + .Select(x => + { + var doc = x.Doc; + doc.RelevanceScore = NumOps.FromDouble(x.Score); + doc.HasRelevanceScore = true; + return doc; + }) + .ToList(); + + return reranked; + } + + /// + /// Scores every (query, document) pair with the cross-encoder model and returns one relevance + /// score per document, in input order. + /// + /// The query text. + /// The document contents to score against the query. + /// One relevance score per document, in the same order as . + /// + /// + /// This is the inference seam. The default implementation loads the ONNX model + tokenizer, + /// tokenizes each pair jointly, runs the model in batches, and reads the relevance logit. Tests + /// override this method to inject deterministic scores and exercise the ranking logic without a + /// real model file. + /// + /// + /// Thrown when the ONNX model file cannot be found. + protected virtual double[] ScorePairs(string query, IList documentContents) + { + if (documentContents == null) + throw new ArgumentNullException(nameof(documentContents)); + + if (!EnsureModelLoaded()) + throw new FileNotFoundException($"ONNX cross-encoder model file not found: {_modelPath}", _modelPath); + + var session = Session; + var tokenizer = Tokenizer; + var declaredInputs = new HashSet(session.InputMetadata.Keys); + long padId = GetSpecialTokenId(tokenizer, tokenizer.SpecialTokens.PadToken, 0L); + + var results = new double[documentContents.Count]; + + for (int start = 0; start < documentContents.Count; start += _batchSize) + { + int count = Math.Min(_batchSize, documentContents.Count - start); + + // Encode each pair and track the longest sequence in this batch. + var encoded = new List<(long[] Ids, long[] Mask, long[] Types)>(count); + int maxLen = 0; + for (int j = 0; j < count; j++) + { + var pair = EncodePair(tokenizer, query, documentContents[start + j]); + encoded.Add(pair); + if (pair.Ids.Length > maxLen) + maxLen = pair.Ids.Length; + } + if (maxLen == 0) + maxLen = 1; + + // Pad every sequence in the batch to maxLen so they form a rectangular tensor. + var ids = new long[count * maxLen]; + var mask = new long[count * maxLen]; + var types = new long[count * maxLen]; + for (int j = 0; j < count; j++) + { + var pair = encoded[j]; + for (int k = 0; k < maxLen; k++) + { + int idx = (j * maxLen) + k; + if (k < pair.Ids.Length) + { + ids[idx] = pair.Ids[k]; + mask[idx] = pair.Mask[k]; + types[idx] = pair.Types[k]; + } + else + { + ids[idx] = padId; + mask[idx] = 0L; + types[idx] = 0L; + } + } + } + + var shape = new[] { count, maxLen }; + var inputs = BuildPairInputs(declaredInputs, ids, mask, types, shape); + + using var outputs = session.Run(inputs); + var logits = outputs.First().AsTensor(); + var dims = logits.Dimensions.ToArray(); + int numLabels = dims.Length >= 2 ? dims[dims.Length - 1] : 1; + var flat = System.Linq.Enumerable.ToArray(logits); + + var batchScores = ExtractScoresFromLogits(flat, count, numLabels, _scoreLabelIndex); + for (int j = 0; j < count; j++) + { + results[start + j] = batchScores[j]; + } + } + + return results; + } + + /// + /// Scores a single (query, document) pair and returns the relevance score as . + /// Convenient for wiring an ONNX cross-encoder into or ad-hoc use. + /// + /// The query text. + /// The document content. + /// The relevance score. + public T ScorePair(string query, string document) + { + var scores = ScorePairs(query, new List { document ?? string.Empty }); + return NumOps.FromDouble(scores.Length > 0 ? scores[0] : 0.0); + } + + /// + /// Tokenizes a (query, document) pair into the joint BERT-style sequence + /// [CLS] query [SEP] document [SEP] with matching attention mask and token-type ids. + /// + private (long[] Ids, long[] Mask, long[] Types) EncodePair(ITokenizer tokenizer, string query, string document) + { + var special = tokenizer.SpecialTokens; + bool hasCls = !string.IsNullOrEmpty(special.ClsToken); + bool hasSep = !string.IsNullOrEmpty(special.SepToken); + + var queryTokens = tokenizer.Tokenize(query ?? string.Empty); + var docTokens = tokenizer.Tokenize(document ?? string.Empty); + + // Reserve room for the special tokens: [CLS] ... [SEP] ... [SEP] + int specialCount = (hasCls ? 1 : 0) + (hasSep ? 2 : 1); + int available = _maxLength - specialCount; + if (available < 0) + available = 0; + + TruncateLongestFirst(queryTokens, docTokens, available); + + long clsId = hasCls ? GetSpecialTokenId(tokenizer, special.ClsToken, 0L) : 0L; + long sepId = hasSep ? GetSpecialTokenId(tokenizer, special.SepToken, 0L) : 0L; + + var queryIds = tokenizer.ConvertTokensToIds(queryTokens); + var docIds = tokenizer.ConvertTokensToIds(docTokens); + + var ids = new List(specialCount + queryIds.Count + docIds.Count); + var types = new List(ids.Capacity); + + // Segment 0: [CLS] + query + [SEP] + if (hasCls) + { + ids.Add(clsId); + types.Add(0L); + } + for (int i = 0; i < queryIds.Count; i++) + { + ids.Add(queryIds[i]); + types.Add(0L); + } + if (hasSep) + { + ids.Add(sepId); + types.Add(0L); + } + + // Segment 1: document + [SEP] + for (int i = 0; i < docIds.Count; i++) + { + ids.Add(docIds[i]); + types.Add(1L); + } + if (hasSep) + { + ids.Add(sepId); + types.Add(1L); + } + + var idArray = ids.ToArray(); + var typeArray = types.ToArray(); + var maskArray = new long[idArray.Length]; + for (int i = 0; i < maskArray.Length; i++) + { + maskArray[i] = 1L; + } + + return (idArray, maskArray, typeArray); + } + + /// + /// Truncates the two token lists in place using HuggingFace's "longest_first" strategy so their + /// combined length does not exceed . + /// + internal static void TruncateLongestFirst(List queryTokens, List docTokens, int available) + { + if (available < 0) + available = 0; + + while (queryTokens.Count + docTokens.Count > available) + { + if (queryTokens.Count == 0 && docTokens.Count == 0) + break; + + if (queryTokens.Count > docTokens.Count) + { + queryTokens.RemoveAt(queryTokens.Count - 1); + } + else + { + docTokens.RemoveAt(docTokens.Count - 1); + } + } + } + + /// + /// Builds the ONNX input list for a batch of tokenized (query, document) pairs. input_ids is + /// always supplied; attention_mask and token_type_ids are added only when the loaded + /// session declares them, so models exported without them keep working. + /// + /// The names of the inputs the loaded session declares. + /// Int64 token ids, flattened row-major with shape [batch, seqLen]. + /// Int64 attention mask, same layout as . + /// Int64 segment ids, same layout as . + /// The [batch, seqLen] shape shared by every input tensor. + internal static List BuildPairInputs( + ICollection declaredInputs, + long[] inputIds, + long[] attentionMask, + long[] tokenTypeIds, + int[] shape) + { + var inputs = new List + { + NamedOnnxValue.CreateFromTensor("input_ids", new Microsoft.ML.OnnxRuntime.Tensors.DenseTensor(inputIds, shape)) + }; + + if (declaredInputs.Contains("attention_mask")) + { + inputs.Add(NamedOnnxValue.CreateFromTensor("attention_mask", new Microsoft.ML.OnnxRuntime.Tensors.DenseTensor(attentionMask, shape))); + } + + if (declaredInputs.Contains("token_type_ids")) + { + inputs.Add(NamedOnnxValue.CreateFromTensor("token_type_ids", new Microsoft.ML.OnnxRuntime.Tensors.DenseTensor(tokenTypeIds, shape))); + } + + return inputs; + } + + /// + /// Extracts one relevance score per batch item from a flattened [batch, numLabels] logits array. + /// + /// Row-major logits, length batch * numLabels. + /// Number of items in the batch. + /// Number of output labels per item. + /// + /// The label index to read. Values < 0 or out of range select the last label, which is correct + /// for single-logit rerankers and last-is-positive binary classifiers. + /// + internal static double[] ExtractScoresFromLogits(float[] flatLogits, int batch, int numLabels, int labelIndex) + { + if (numLabels < 1) + numLabels = 1; + + int li = labelIndex; + if (li < 0 || li >= numLabels) + li = numLabels - 1; + + var scores = new double[batch]; + for (int i = 0; i < batch; i++) + { + int idx = (i * numLabels) + li; + scores[i] = (idx >= 0 && idx < flatLogits.Length) ? flatLogits[idx] : 0.0; + } + + return scores; + } + + private static long GetSpecialTokenId(ITokenizer tokenizer, string token, long fallback) + { + if (string.IsNullOrEmpty(token)) + return fallback; + + var ids = tokenizer.ConvertTokensToIds(new List { token }); + if (ids == null || ids.Count == 0) + return fallback; + + return ids[0]; + } + + /// + /// Ensures the ONNX model and tokenizer are loaded, loading lazily on first use. + /// Returns false only if the model file is unavailable. + /// + private bool EnsureModelLoaded() + { + if (_session != null && _tokenizer != null) + return true; + + lock (_initLock) + { + if (_session != null && _tokenizer != null) + return true; + + if (!File.Exists(_modelPath)) + return false; + + var tokenizerDir = ResolveTokenizerDirectory(_tokenizerPath); + var tokenizer = AutoTokenizer.FromPretrained(tokenizerDir); + var session = new InferenceSession(_modelPath); + + // Assign both atomically so concurrent readers never see a partial state. + _tokenizer = tokenizer; + _session = session; + return true; + } + } + + private static string ResolveTokenizerDirectory(string tokenizerPath) + { + if (Directory.Exists(tokenizerPath)) + return tokenizerPath; + + if (File.Exists(tokenizerPath)) + { + var dir = Path.GetDirectoryName(tokenizerPath); + if (!string.IsNullOrEmpty(dir)) + return dir!; + } + + // Fall back to the given path; AutoTokenizer will surface a clear error if it is invalid. + return tokenizerPath; + } + + /// + /// Gets the inference session, ensuring it is loaded. + /// + /// Thrown when the ONNX model file is not found. + private InferenceSession Session + { + get + { + if (!EnsureModelLoaded()) + throw new FileNotFoundException($"ONNX cross-encoder model file not found: {_modelPath}", _modelPath); + + return _session ?? throw new InvalidOperationException("Failed to load ONNX session."); + } + } + + /// + /// Gets the tokenizer, ensuring it is loaded. + /// + /// Thrown when the ONNX model file is not found. + private ITokenizer Tokenizer + { + get + { + if (!EnsureModelLoaded()) + throw new FileNotFoundException($"ONNX cross-encoder model file not found: {_modelPath}", _modelPath); + + return _tokenizer ?? throw new InvalidOperationException("Failed to load tokenizer."); + } + } + + /// + /// Releases the ONNX session and tokenizer resources. + /// + public void Dispose() + { + Dispose(true); + GC.SuppressFinalize(this); + } + + /// + /// Releases resources held by the reranker. + /// + /// True when called from . + protected virtual void Dispose(bool disposing) + { + if (_disposed) + return; + + if (disposing) + { + _session?.Dispose(); + if (_tokenizer is IDisposable disposableTokenizer) + { + disposableTokenizer.Dispose(); + } + } + + _disposed = true; + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RAG/Rerankers/OnnxCrossEncoderRerankerTests.cs b/tests/AiDotNet.Tests/UnitTests/RAG/Rerankers/OnnxCrossEncoderRerankerTests.cs new file mode 100644 index 0000000000..86250d95b8 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RAG/Rerankers/OnnxCrossEncoderRerankerTests.cs @@ -0,0 +1,414 @@ +#nullable disable +using System; +using System.Collections.Generic; +using System.IO; +using System.Linq; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using AiDotNet.RetrievalAugmentedGeneration.Rerankers; +using Microsoft.ML.OnnxRuntime; +using Xunit; + +namespace AiDotNetTests.UnitTests.RAG.Rerankers +{ + /// + /// Tests for . + /// + /// The scoring step (real ONNX inference) is isolated behind the virtual ScorePairs seam so + /// the ranking logic can be exercised with a fake scorer, keeping CI free of any model/network + /// dependency. A gated integration test loads a real model from an env var and skips when unset. + /// + public class OnnxCrossEncoderRerankerTests + { + /// + /// Test double: overrides the ONNX inference seam with a caller-supplied scoring function so the + /// ranking logic can be verified without a real model file. + /// + private sealed class FakeOnnxCrossEncoderReranker : OnnxCrossEncoderReranker + { + private readonly Func _scorer; + + public int ScorePairsCallCount { get; private set; } + + public FakeOnnxCrossEncoderReranker(Func scorer) + : base("fake-model.onnx", "fake-tokenizer", maxLength: 128, maxPairsToScore: 100) + { + _scorer = scorer; + } + + protected override double[] ScorePairs(string query, IList documentContents) + { + ScorePairsCallCount++; + var scores = new double[documentContents.Count]; + for (int i = 0; i < documentContents.Count; i++) + { + scores[i] = _scorer(query, documentContents[i]); + } + return scores; + } + } + + private static List> MakeDocs(params (string Id, string Content)[] items) + { + var docs = new List>(); + foreach (var (id, content) in items) + { + docs.Add(new Document(id, content)); + } + return docs; + } + + // ---------- Constructor validation ---------- + + [Fact] + public void Constructor_WithValidParameters_CreatesInstance() + { + using var reranker = new OnnxCrossEncoderReranker("model.onnx", "tokenizer-dir"); + Assert.NotNull(reranker); + Assert.True(reranker.ModifiesScores); + } + + [Theory] + [InlineData(null)] + [InlineData("")] + [InlineData(" ")] + public void Constructor_WithInvalidModelPath_Throws(string modelPath) + { + var ex = Assert.Throws(() => + new OnnxCrossEncoderReranker(modelPath, "tokenizer-dir")); + Assert.Contains("Model path cannot be empty", ex.Message); + } + + [Theory] + [InlineData(null)] + [InlineData("")] + [InlineData(" ")] + public void Constructor_WithInvalidTokenizerPath_Throws(string tokenizerPath) + { + var ex = Assert.Throws(() => + new OnnxCrossEncoderReranker("model.onnx", tokenizerPath)); + Assert.Contains("Tokenizer path cannot be empty", ex.Message); + } + + [Theory] + [InlineData(0)] + [InlineData(-1)] + public void Constructor_WithInvalidMaxLength_Throws(int maxLength) + { + Assert.Throws(() => + new OnnxCrossEncoderReranker("model.onnx", "tokenizer-dir", maxLength: maxLength)); + } + + [Theory] + [InlineData(0)] + [InlineData(-5)] + public void Constructor_WithInvalidMaxPairsToScore_Throws(int maxPairs) + { + Assert.Throws(() => + new OnnxCrossEncoderReranker("model.onnx", "tokenizer-dir", maxPairsToScore: maxPairs)); + } + + [Theory] + [InlineData(0)] + [InlineData(-2)] + public void Constructor_WithInvalidBatchSize_Throws(int batchSize) + { + Assert.Throws(() => + new OnnxCrossEncoderReranker("model.onnx", "tokenizer-dir", batchSize: batchSize)); + } + + // ---------- Missing-model behavior (no silent fallback) ---------- + + [Fact] + public void Rerank_WithMissingModelFile_ThrowsFileNotFound() + { + using var reranker = new OnnxCrossEncoderReranker("nonexistent-model.onnx", "nonexistent-tokenizer"); + var docs = MakeDocs(("d1", "some content"), ("d2", "other content")); + + // Rerank is lazy via IEnumerable; force enumeration. + Assert.Throws(() => reranker.Rerank("query", docs).ToList()); + } + + [Fact] + public void ScorePair_WithMissingModelFile_ThrowsFileNotFound() + { + using var reranker = new OnnxCrossEncoderReranker("nonexistent-model.onnx", "nonexistent-tokenizer"); + Assert.Throws(() => reranker.ScorePair("query", "document")); + } + + // ---------- Ranking logic via injected fake scorer ---------- + + [Fact] + public void Rerank_ReordersDocumentsByDescendingScore() + { + // Score by how many query words appear in the document content. + var reranker = new FakeOnnxCrossEncoderReranker((query, doc) => + { + var qWords = query.Split(' '); + return qWords.Count(w => doc.Contains(w)); + }); + + var docs = MakeDocs( + ("low", "unrelated text about cooking"), + ("high", "machine learning models learn from data"), + ("mid", "learning is fun")); + + var result = reranker.Rerank("machine learning", docs).ToList(); + + Assert.Equal(3, result.Count); + Assert.Equal("high", result[0].Id); // matches both "machine" and "learning" + Assert.Equal("mid", result[1].Id); // matches "learning" + Assert.Equal("low", result[2].Id); // matches neither + Assert.Equal(1, reranker.ScorePairsCallCount); + } + + [Fact] + public void Rerank_SetsRelevanceScoreAndFlag() + { + var scoreMap = new Dictionary + { + { "a", 0.1 }, + { "b", 0.9 }, + { "c", 0.5 }, + }; + var reranker = new FakeOnnxCrossEncoderReranker((query, doc) => scoreMap[doc]); + + var docs = MakeDocs(("a", "a"), ("b", "b"), ("c", "c")); + + var result = reranker.Rerank("q", docs).ToList(); + + // Ordered b (0.9) > c (0.5) > a (0.1) + Assert.Equal("b", result[0].Id); + Assert.Equal("c", result[1].Id); + Assert.Equal("a", result[2].Id); + + foreach (var doc in result) + { + Assert.True(doc.HasRelevanceScore); + } + Assert.Equal(0.9, result[0].RelevanceScore, 6); + Assert.Equal(0.5, result[1].RelevanceScore, 6); + Assert.Equal(0.1, result[2].RelevanceScore, 6); + } + + [Fact] + public void Rerank_TopK_LimitsResults() + { + var reranker = new FakeOnnxCrossEncoderReranker((query, doc) => doc.Length); + var docs = MakeDocs( + ("d1", "x"), + ("d2", "xx"), + ("d3", "xxx"), + ("d4", "xxxx")); + + var result = reranker.Rerank("q", docs, topK: 2).ToList(); + + Assert.Equal(2, result.Count); + Assert.Equal("d4", result[0].Id); + Assert.Equal("d3", result[1].Id); + } + + [Fact] + public void Rerank_EmptyDocuments_ReturnsEmpty() + { + var reranker = new FakeOnnxCrossEncoderReranker((query, doc) => 1.0); + var result = reranker.Rerank("q", new List>()).ToList(); + Assert.Empty(result); + } + + [Fact] + public void Rerank_PreservesDocumentCount() + { + var reranker = new FakeOnnxCrossEncoderReranker((query, doc) => doc.GetHashCode()); + var docs = MakeDocs( + ("d1", "one"), ("d2", "two"), ("d3", "three"), + ("d4", "four"), ("d5", "five"), ("d6", "six"), ("d7", "seven")); + + var result = reranker.Rerank("q", docs).ToList(); + Assert.Equal(7, result.Count); + } + + // ---------- Convenience ctor on CrossEncoderReranker ---------- + + [Fact] + public void CrossEncoderReranker_BackedByOnnxReranker_Reranks() + { + var scoreMap = new Dictionary + { + { "low", 0.2 }, + { "high", 0.95 }, + }; + var onnx = new FakeOnnxCrossEncoderReranker((query, doc) => scoreMap[doc]); + var reranker = new CrossEncoderReranker(onnx); + + var docs = MakeDocs(("low", "low"), ("high", "high")); + var result = reranker.Rerank("q", docs).ToList(); + + Assert.Equal("high", result[0].Id); + Assert.Equal("low", result[1].Id); + } + + [Fact] + public void CrossEncoderReranker_WithNullOnnxReranker_Throws() + { + Assert.Throws(() => + new CrossEncoderReranker((OnnxCrossEncoderReranker)null)); + } + + // ---------- ExtractScoresFromLogits (pure inference helper) ---------- + + [Fact] + public void ExtractScoresFromLogits_SingleLabel_ReturnsThatValue() + { + var flat = new float[] { 0.5f, -1.2f, 3.3f }; + var scores = OnnxCrossEncoderReranker.ExtractScoresFromLogits(flat, batch: 3, numLabels: 1, labelIndex: -1); + + Assert.Equal(3, scores.Length); + Assert.Equal(0.5, scores[0], 5); + Assert.Equal(-1.2, scores[1], 5); + Assert.Equal(3.3, scores[2], 5); + } + + [Fact] + public void ExtractScoresFromLogits_TwoLabels_DefaultReadsLastLabel() + { + // batch 2, numLabels 2: [ [neg, pos], [neg, pos] ] + var flat = new float[] { 0.1f, 0.9f, 0.8f, 0.2f }; + var scores = OnnxCrossEncoderReranker.ExtractScoresFromLogits(flat, batch: 2, numLabels: 2, labelIndex: -1); + + Assert.Equal(0.9, scores[0], 5); // last label of item 0 + Assert.Equal(0.2, scores[1], 5); // last label of item 1 + } + + [Fact] + public void ExtractScoresFromLogits_ExplicitLabelIndex_ReadsThatColumn() + { + var flat = new float[] { 0.1f, 0.9f, 0.8f, 0.2f }; + var scores = OnnxCrossEncoderReranker.ExtractScoresFromLogits(flat, batch: 2, numLabels: 2, labelIndex: 0); + + Assert.Equal(0.1, scores[0], 5); + Assert.Equal(0.8, scores[1], 5); + } + + [Fact] + public void ExtractScoresFromLogits_OutOfRangeIndex_FallsBackToLastLabel() + { + var flat = new float[] { 0.1f, 0.9f }; + var scores = OnnxCrossEncoderReranker.ExtractScoresFromLogits(flat, batch: 1, numLabels: 2, labelIndex: 99); + Assert.Equal(0.9, scores[0], 5); + } + + // ---------- BuildPairInputs (ONNX input plumbing) ---------- + + [Fact] + public void BuildPairInputs_AlwaysIncludesInputIds() + { + var declared = new HashSet { "input_ids" }; + var ids = new long[] { 101, 5, 102 }; + var mask = new long[] { 1, 1, 1 }; + var types = new long[] { 0, 0, 1 }; + var shape = new[] { 1, 3 }; + + var inputs = OnnxCrossEncoderReranker.BuildPairInputs(declared, ids, mask, types, shape); + + Assert.Contains(inputs, v => v.Name == "input_ids"); + Assert.DoesNotContain(inputs, v => v.Name == "attention_mask"); + Assert.DoesNotContain(inputs, v => v.Name == "token_type_ids"); + + var idsTensor = inputs.Single(v => v.Name == "input_ids").AsTensor(); + Assert.Equal(shape, idsTensor.Dimensions.ToArray()); + } + + [Fact] + public void BuildPairInputs_IncludesDeclaredOptionalInputs() + { + var declared = new HashSet { "input_ids", "attention_mask", "token_type_ids" }; + var ids = new long[] { 101, 5, 6, 102 }; + var mask = new long[] { 1, 1, 1, 1 }; + var types = new long[] { 0, 0, 1, 1 }; + var shape = new[] { 1, 4 }; + + var inputs = OnnxCrossEncoderReranker.BuildPairInputs(declared, ids, mask, types, shape); + + Assert.Contains(inputs, v => v.Name == "input_ids"); + Assert.Contains(inputs, v => v.Name == "attention_mask"); + Assert.Contains(inputs, v => v.Name == "token_type_ids"); + + var typeTensor = inputs.Single(v => v.Name == "token_type_ids").AsTensor(); + Assert.Equal(new long[] { 0, 0, 1, 1 }, typeTensor.ToArray()); + } + + // ---------- TruncateLongestFirst ---------- + + [Fact] + public void TruncateLongestFirst_TrimsLongerListFirst() + { + var query = new List { "q1", "q2" }; + var doc = new List { "d1", "d2", "d3", "d4", "d5", "d6" }; + + OnnxCrossEncoderReranker.TruncateLongestFirst(query, doc, available: 4); + + Assert.Equal(4, query.Count + doc.Count); + // Query (2) is shorter, so document is trimmed down to 2. + Assert.Equal(2, query.Count); + Assert.Equal(2, doc.Count); + } + + [Fact] + public void TruncateLongestFirst_NoTruncationWhenWithinBudget() + { + var query = new List { "q1" }; + var doc = new List { "d1", "d2" }; + + OnnxCrossEncoderReranker.TruncateLongestFirst(query, doc, available: 10); + + Assert.Equal(1, query.Count); + Assert.Equal(2, doc.Count); + } + + [Fact] + public void TruncateLongestFirst_ZeroBudget_EmptiesBoth() + { + var query = new List { "q1", "q2" }; + var doc = new List { "d1" }; + + OnnxCrossEncoderReranker.TruncateLongestFirst(query, doc, available: 0); + + Assert.Empty(query); + Assert.Empty(doc); + } + + // ---------- Gated integration test (real model) ---------- + + /// + /// End-to-end test against a real cross-encoder ONNX model. Set the env var + /// AIDOTNET_CROSSENCODER_ONNX to the model file path and + /// AIDOTNET_CROSSENCODER_TOKENIZER to the tokenizer directory to run it. + /// Skips (does not fail) when those are unset, keeping CI model/network-free. + /// + [SkippableFact] + [Trait("Category", "Integration")] + public void Integration_RealModel_RanksRelevantDocumentFirst() + { + var modelPath = Environment.GetEnvironmentVariable("AIDOTNET_CROSSENCODER_ONNX"); + var tokenizerPath = Environment.GetEnvironmentVariable("AIDOTNET_CROSSENCODER_TOKENIZER"); + + Skip.If(string.IsNullOrWhiteSpace(modelPath) || string.IsNullOrWhiteSpace(tokenizerPath), + "Set AIDOTNET_CROSSENCODER_ONNX and AIDOTNET_CROSSENCODER_TOKENIZER to run the real-model integration test."); + + Skip.IfNot(File.Exists(modelPath), $"Model file not found: {modelPath}"); + + using var reranker = new OnnxCrossEncoderReranker(modelPath, tokenizerPath); + + var docs = MakeDocs( + ("relevant", "The Eiffel Tower is located in Paris, France."), + ("irrelevant", "Bananas are a good source of potassium.")); + + var result = reranker.Rerank("Where is the Eiffel Tower?", docs).ToList(); + + Assert.Equal(2, result.Count); + Assert.Equal("relevant", result[0].Id); + Assert.True(result[0].HasRelevanceScore); + Assert.True(result[0].RelevanceScore >= result[1].RelevanceScore); + } + } +} From aed0ef04c789c47bb6fdfe843137a2be9303a984 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 16:16:47 -0400 Subject: [PATCH 09/25] feat(rag): embedding + LLM-judge eval metrics and IR metrics (nDCG/MRR/MAP) Upgrade RAG evaluation beyond lexical overlap: - FaithfulnessMetric: optional ITextGenerator NLI-style judge (supported claims / total), lexical word-overlap fallback when none supplied. - AnswerSimilarityMetric: optional IEmbeddingModel cosine, Jaccard fallback. - AnswerCorrectnessMetric: optional ITextGenerator (0-10 judge) + IEmbeddingModel (semantic cosine), averaged; Jaccard fallback. Back-compat (endpoint, apiKey) ctor retained. - ContextRelevanceMetric: optional IEmbeddingModel query-vs-doc cosine, Jaccard fallback. - New AnswerRelevanceMetric (RAGAS): reverse-generate questions then mean cosine to query. - New ContextPrecisionMetric (RAGAS AP-style) and ContextRecallMetric (claim attribution), each judge/embedding/lexical. - New RetrievalMetrics static class: pure-math nDCG@k (binary + graded), MRR, MAP, Hit-Rate@k, Precision@k, Recall@k over ranked doc ids vs relevant set. - Shared helpers added to RAGMetricBase (sentence split, list parsing, yes/no verdict, cosine, Jaccard). net471-safe (log2 via change-of-base; guards, no ThrowIfNull/IsFinite). All metrics fall back to documented offline lexical behavior when no model is injected. Existing IRAGMetric metrics flow through RAGEvaluator unchanged. Tests (18, net10.0 green): IR metrics with exact numeric assertions; embedding metrics with a controlled fake IEmbeddingModel; LLM-judge metrics with a scripted fake ITextGenerator; all offline. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../Evaluation/AnswerCorrectnessMetric.cs | 98 +++++- .../Evaluation/AnswerRelevanceMetric.cs | 99 ++++++ .../Evaluation/AnswerSimilarityMetric.cs | 33 +- .../Evaluation/ContextPrecisionMetric.cs | 121 +++++++ .../Evaluation/ContextRecallMetric.cs | 129 +++++++ .../Evaluation/ContextRelevanceMetric.cs | 58 ++-- .../Evaluation/FaithfulnessMetric.cs | 50 ++- .../Evaluation/RAGMetricBase.cs | 101 ++++++ .../Evaluation/RetrievalMetrics.cs | 291 ++++++++++++++++ .../Evaluation/RagEvaluationMetricsTests.cs | 321 ++++++++++++++++++ 10 files changed, 1260 insertions(+), 41 deletions(-) create mode 100644 src/RetrievalAugmentedGeneration/Evaluation/AnswerRelevanceMetric.cs create mode 100644 src/RetrievalAugmentedGeneration/Evaluation/ContextPrecisionMetric.cs create mode 100644 src/RetrievalAugmentedGeneration/Evaluation/ContextRecallMetric.cs create mode 100644 src/RetrievalAugmentedGeneration/Evaluation/RetrievalMetrics.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Evaluation/RagEvaluationMetricsTests.cs diff --git a/src/RetrievalAugmentedGeneration/Evaluation/AnswerCorrectnessMetric.cs b/src/RetrievalAugmentedGeneration/Evaluation/AnswerCorrectnessMetric.cs index 250af23a38..2eb1e62bbc 100644 --- a/src/RetrievalAugmentedGeneration/Evaluation/AnswerCorrectnessMetric.cs +++ b/src/RetrievalAugmentedGeneration/Evaluation/AnswerCorrectnessMetric.cs @@ -1,5 +1,8 @@ +using System.Globalization; +using System.Text.RegularExpressions; using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; using AiDotNet.Validation; @@ -10,8 +13,13 @@ namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; /// /// The numeric data type used for calculations. /// -/// Assesses whether the generated answer contains factually correct information -/// by comparing against ground truth or using fact-checking mechanisms. +/// +/// Assesses whether the generated answer contains factually correct information by comparing +/// against ground truth. When a text generator is supplied it is used as an LLM judge that rates +/// correctness; when an embedding model is supplied semantic (cosine) similarity is used. When both +/// are supplied the two signals are averaged (RAGAS-style). When neither is supplied the metric +/// falls back to the offline lexical Jaccard word-overlap heuristic. +/// /// [ComponentType(ComponentType.Evaluator)] [PipelineStage(PipelineStage.Evaluation)] @@ -19,38 +27,106 @@ public class AnswerCorrectnessMetric : RAGMetricBase { private readonly string _llmEndpoint; private readonly string _llmApiKey; + private readonly ITextGenerator? _generator; + private readonly IEmbeddingModel? _embeddingModel; + /// + /// Gets the name of this metric. + /// public override string Name => "Answer Correctness"; + + /// + /// Gets the description of what this metric measures. + /// public override string Description => "Evaluates the factual correctness of generated answers by comparing against ground truth."; + + /// + /// Gets a value indicating whether this metric requires ground truth. + /// protected override bool RequiresGroundTruth => true; + /// + /// Initializes a new instance of the class backed by a + /// real LLM judge and/or embedding model. + /// + /// + /// Optional text generator used to rate factual correctness of the answer against the ground truth. + /// + /// + /// Optional embedding model used to compute semantic similarity between the answer and ground truth. + /// + /// When both arguments are null, the metric uses the offline lexical fallback. + public AnswerCorrectnessMetric(ITextGenerator? generator = null, IEmbeddingModel? embeddingModel = null) + { + _llmEndpoint = string.Empty; + _llmApiKey = string.Empty; + _generator = generator; + _embeddingModel = embeddingModel; + } + /// /// Initializes a new instance of the class. /// - /// The LLM API endpoint for fact checking. - /// The API key for the LLM service. + /// The LLM API endpoint for fact checking (informational). + /// The API key for the LLM service (informational). + /// + /// Back-compatible constructor. Because it is not wired to a real generator, this overload uses + /// the offline lexical Jaccard fallback. Prefer the constructor that accepts a text generator + /// and/or embedding model for semantic evaluation. + /// public AnswerCorrectnessMetric(string llmEndpoint, string llmApiKey) { Guard.NotNull(llmEndpoint); _llmEndpoint = llmEndpoint; Guard.NotNull(llmApiKey); _llmApiKey = llmApiKey; + _generator = null; + _embeddingModel = null; } + /// + /// Evaluates answer correctness against the ground truth. + /// + /// The grounded answer to evaluate. + /// The reference/correct answer. + /// Correctness score (0-1). protected override T EvaluateCore(GroundedAnswer answer, string? groundTruth) { if (string.IsNullOrWhiteSpace(answer.Answer) || string.IsNullOrWhiteSpace(groundTruth)) return NumOps.Zero; - var words1 = GetWords(answer.Answer); - var words2 = GetWords(groundTruth!); + double? judgeScore = _generator != null ? RateWithJudge(answer.Answer, groundTruth!) : (double?)null; + double? semanticScore = _embeddingModel != null + ? Math.Max(0.0, EmbeddingCosine(_embeddingModel, answer.Answer, groundTruth!)) + : (double?)null; - var intersection = words1.Intersect(words2).Count(); - var union = words1.Union(words2).Count(); + if (judgeScore.HasValue && semanticScore.HasValue) + return NumOps.FromDouble((judgeScore.Value + semanticScore.Value) / 2.0); + if (judgeScore.HasValue) + return NumOps.FromDouble(judgeScore.Value); + if (semanticScore.HasValue) + return NumOps.FromDouble(semanticScore.Value); - if (union == 0) - return NumOps.Zero; + // Offline lexical fallback: Jaccard word overlap. + return NumOps.FromDouble(JaccardSimilarity(answer.Answer, groundTruth!)); + } + + /// + /// Asks the generator to rate how factually correct the answer is against the ground truth on a + /// 0-10 scale, and normalizes the parsed number to [0, 1]. + /// + private double RateWithJudge(string answerText, string groundTruth) + { + var prompt = + "On a scale of 0 to 10, how factually correct and complete is the answer compared to " + + "the reference answer? Reply with only a single number.\n\n" + + $"Reference: {groundTruth}\n\nAnswer: {answerText}\n\nCorrectness (0-10):"; + var reply = _generator!.Generate(prompt); + var match = Regex.Match(reply ?? string.Empty, @"\d+(\.\d+)?"); + if (match.Success && double.TryParse(match.Value, NumberStyles.Float, CultureInfo.InvariantCulture, out double raw)) + return Math.Min(1.0, Math.Max(0.0, raw / 10.0)); - return NumOps.Divide(NumOps.FromDouble(intersection), NumOps.FromDouble(union)); + // Unparseable reply: fall back to lexical overlap so the score is still meaningful. + return JaccardSimilarity(answerText, groundTruth); } } diff --git a/src/RetrievalAugmentedGeneration/Evaluation/AnswerRelevanceMetric.cs b/src/RetrievalAugmentedGeneration/Evaluation/AnswerRelevanceMetric.cs new file mode 100644 index 0000000000..e1a372ee7e --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Evaluation/AnswerRelevanceMetric.cs @@ -0,0 +1,99 @@ +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.Models; + +namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; + +/// +/// Evaluates how relevant the generated answer is to the original query (RAGAS answer relevance). +/// +/// The numeric data type used for calculations. +/// +/// +/// RAGAS answer relevance works by asking a generator to produce the questions that the answer would +/// answer, then measuring the mean cosine similarity between the embeddings of those generated +/// questions and the original query. A high score means the answer is on-topic and directly addresses +/// the question; a low score means the answer is evasive, incomplete, or off-topic. +/// +/// +/// This metric requires both a text generator (to reverse-generate questions) and an embedding +/// model (to compare them to the query). When either is missing it falls back to the offline lexical +/// Jaccard overlap between the answer and the query. +/// +/// +[ComponentType(ComponentType.Evaluator)] +[PipelineStage(PipelineStage.Evaluation)] +public class AnswerRelevanceMetric : RAGMetricBase +{ + private readonly ITextGenerator? _generator; + private readonly IEmbeddingModel? _embeddingModel; + private readonly int _numQuestions; + + /// + /// Initializes a new instance of the class. + /// + /// Optional text generator used to reverse-generate questions from the answer. + /// Optional embedding model used to compare generated questions to the query. + /// The number of questions to reverse-generate (default 3). + public AnswerRelevanceMetric( + ITextGenerator? generator = null, + IEmbeddingModel? embeddingModel = null, + int numQuestions = 3) + { + _generator = generator; + _embeddingModel = embeddingModel; + _numQuestions = numQuestions > 0 ? numQuestions : 3; + } + + /// Gets the name of this metric. + public override string Name => "Answer Relevance"; + + /// Gets the description of what this metric measures. + public override string Description => + "Measures how directly the generated answer addresses the original query (RAGAS answer relevance)"; + + /// Gets a value indicating whether this metric requires ground truth. + protected override bool RequiresGroundTruth => false; + + /// + /// Evaluates answer relevance to the query. + /// + /// The grounded answer to evaluate. + /// Not used for this metric. + /// Relevance score (0-1). + protected override T EvaluateCore(GroundedAnswer answer, string? groundTruth) + { + if (string.IsNullOrWhiteSpace(answer.Query)) + return NumOps.Zero; + + if (_generator != null && _embeddingModel != null) + return EvaluateWithModels(answer.Query, answer.Answer); + + // Offline lexical fallback: Jaccard overlap between the answer and the query. + return NumOps.FromDouble(JaccardSimilarity(answer.Answer, answer.Query)); + } + + /// + /// Reverse-generates questions from the answer and averages their cosine similarity to the query. + /// + private T EvaluateWithModels(string query, string answerText) + { + var prompt = + $"Generate {_numQuestions} distinct questions that the following answer would correctly and " + + "completely answer. Put each question on its own line.\n\n" + + $"Answer: {answerText}\n\nQuestions:"; + var reply = _generator!.Generate(prompt); + var questions = ParseListLines(reply); + if (questions.Count == 0) + return NumOps.Zero; + + double total = 0.0; + foreach (var question in questions) + { + total += Math.Max(0.0, EmbeddingCosine(_embeddingModel!, question, query)); + } + + return NumOps.FromDouble(total / questions.Count); + } +} diff --git a/src/RetrievalAugmentedGeneration/Evaluation/AnswerSimilarityMetric.cs b/src/RetrievalAugmentedGeneration/Evaluation/AnswerSimilarityMetric.cs index 5880615a6d..76698a196b 100644 --- a/src/RetrievalAugmentedGeneration/Evaluation/AnswerSimilarityMetric.cs +++ b/src/RetrievalAugmentedGeneration/Evaluation/AnswerSimilarityMetric.cs @@ -1,5 +1,6 @@ using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; @@ -43,6 +44,21 @@ namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; [PipelineStage(PipelineStage.Evaluation)] public class AnswerSimilarityMetric : RAGMetricBase { + private readonly IEmbeddingModel? _embeddingModel; + + /// + /// Initializes a new instance of the class. + /// + /// + /// Optional embedding model. When supplied, similarity is the cosine similarity between the + /// embeddings of the answer and the ground truth (semantic similarity). When null, the + /// metric falls back to the offline lexical Jaccard word-overlap heuristic. + /// + public AnswerSimilarityMetric(IEmbeddingModel? embeddingModel = null) + { + _embeddingModel = embeddingModel; + } + /// /// Gets the name of this metric. /// @@ -67,15 +83,14 @@ public class AnswerSimilarityMetric : RAGMetricBase /// Similarity score (0-1). protected override T EvaluateCore(GroundedAnswer answer, string? groundTruth) { - var words1 = GetWords(answer.Answer); - var words2 = GetWords(groundTruth!); - - var intersection = words1.Intersect(words2).Count(); - var union = words1.Union(words2).Count(); - - if (union == 0) - return NumOps.Zero; + if (_embeddingModel != null) + { + // Semantic similarity via embeddings; negative cosine is clamped to 0 by the base class. + var cosine = EmbeddingCosine(_embeddingModel, answer.Answer, groundTruth!); + return NumOps.FromDouble(cosine); + } - return NumOps.Divide(NumOps.FromDouble(intersection), NumOps.FromDouble(union)); + // Offline lexical fallback: Jaccard word overlap. + return NumOps.FromDouble(JaccardSimilarity(answer.Answer, groundTruth!)); } } diff --git a/src/RetrievalAugmentedGeneration/Evaluation/ContextPrecisionMetric.cs b/src/RetrievalAugmentedGeneration/Evaluation/ContextPrecisionMetric.cs new file mode 100644 index 0000000000..8252157827 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Evaluation/ContextPrecisionMetric.cs @@ -0,0 +1,121 @@ +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.Models; + +namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; + +/// +/// Evaluates context precision (RAGAS): whether the relevant retrieved contexts are ranked highly. +/// +/// The numeric data type used for calculations. +/// +/// +/// Context precision rewards a retriever for placing the useful contexts near the top of the ranked +/// list. Each retrieved context is first judged relevant or not (against the ground-truth answer when +/// present, otherwise the query). The score is the average of Precision@k evaluated at each rank that +/// holds a relevant context — the standard average-precision formulation used by RAGAS: +/// sum_k (Precision@k * relevant_k) / total_relevant. +/// +/// +/// Relevance of each context is decided by the text generator (a yes/no judge) when one is supplied, +/// otherwise by cosine similarity against an embedding model, otherwise by an offline lexical +/// Jaccard-overlap threshold. +/// +/// +[ComponentType(ComponentType.Evaluator)] +[PipelineStage(PipelineStage.Evaluation)] +public class ContextPrecisionMetric : RAGMetricBase +{ + private readonly ITextGenerator? _generator; + private readonly IEmbeddingModel? _embeddingModel; + private readonly double _relevanceThreshold; + + /// + /// Initializes a new instance of the class. + /// + /// Optional text generator used to judge each context relevant / not relevant. + /// Optional embedding model used to score context relevance by cosine similarity. + /// + /// The similarity threshold (0-1) at or above which a context counts as relevant for the + /// embedding and lexical paths (default 0.5). Ignored when a generator is supplied. + /// + public ContextPrecisionMetric( + ITextGenerator? generator = null, + IEmbeddingModel? embeddingModel = null, + double relevanceThreshold = 0.5) + { + _generator = generator; + _embeddingModel = embeddingModel; + _relevanceThreshold = relevanceThreshold; + } + + /// Gets the name of this metric. + public override string Name => "Context Precision"; + + /// Gets the description of what this metric measures. + public override string Description => + "Measures whether relevant retrieved contexts are ranked ahead of irrelevant ones (RAGAS context precision)"; + + /// Gets a value indicating whether this metric requires ground truth. + protected override bool RequiresGroundTruth => false; + + /// + /// Evaluates context precision over the ranked retrieved documents. + /// + /// The grounded answer to evaluate. + /// The reference answer; when absent the query is used as the relevance target. + /// Context precision score (0-1). + protected override T EvaluateCore(GroundedAnswer answer, string? groundTruth) + { + var docs = answer.SourceDocuments.ToList(); + if (docs.Count == 0) + return NumOps.Zero; + + var reference = !string.IsNullOrWhiteSpace(groundTruth) ? groundTruth! : answer.Query; + if (string.IsNullOrWhiteSpace(reference)) + return NumOps.Zero; + + double weightedPrecisionSum = 0.0; + int relevantSoFar = 0; + int totalRelevant = 0; + + for (int i = 0; i < docs.Count; i++) + { + bool relevant = IsRelevant(reference, docs[i].Content); + if (relevant) + { + relevantSoFar++; + totalRelevant++; + double precisionAtK = (double)relevantSoFar / (i + 1); + weightedPrecisionSum += precisionAtK; + } + } + + if (totalRelevant == 0) + return NumOps.Zero; + + return NumOps.FromDouble(weightedPrecisionSum / totalRelevant); + } + + /// + /// Decides whether a single context is relevant to the reference using the best available signal: + /// LLM judge, then embedding cosine, then lexical overlap. + /// + private bool IsRelevant(string reference, string context) + { + if (_generator != null) + { + var prompt = + "Is the following context useful for answering or verifying the reference? " + + "Reply with only 'yes' or 'no'.\n\n" + + $"Reference: {reference}\n\nContext: {context}\n\nUseful (yes/no):"; + return ParseAffirmative(_generator.Generate(prompt)); + } + + if (_embeddingModel != null) + return EmbeddingCosine(_embeddingModel, reference, context) >= _relevanceThreshold; + + return JaccardSimilarity(reference, context) >= _relevanceThreshold; + } +} diff --git a/src/RetrievalAugmentedGeneration/Evaluation/ContextRecallMetric.cs b/src/RetrievalAugmentedGeneration/Evaluation/ContextRecallMetric.cs new file mode 100644 index 0000000000..fa83dc1260 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Evaluation/ContextRecallMetric.cs @@ -0,0 +1,129 @@ +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.Models; + +namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; + +/// +/// Evaluates context recall (RAGAS): what fraction of the ground-truth answer is supported by the +/// retrieved contexts. +/// +/// The numeric data type used for calculations. +/// +/// +/// Context recall measures retrieval completeness. The ground-truth answer is broken into individual +/// claims (sentences); each claim is checked for whether it can be attributed to the retrieved +/// contexts. The score is attributable-claims / total-claims. A low score means the retriever missed +/// information that the correct answer depends on. +/// +/// +/// Attribution of each claim is decided by the text generator (a yes/no judge) when one is supplied, +/// otherwise by cosine similarity against an embedding model, otherwise by an offline lexical +/// word-coverage threshold. +/// +/// +[ComponentType(ComponentType.Evaluator)] +[PipelineStage(PipelineStage.Evaluation)] +public class ContextRecallMetric : RAGMetricBase +{ + private readonly ITextGenerator? _generator; + private readonly IEmbeddingModel? _embeddingModel; + private readonly double _attributionThreshold; + + /// + /// Initializes a new instance of the class. + /// + /// Optional text generator used to judge whether each claim is attributable to the context. + /// Optional embedding model used to score claim attribution by cosine similarity. + /// + /// The similarity/coverage threshold (0-1) at or above which a claim counts as attributable for the + /// embedding and lexical paths (default 0.5). Ignored when a generator is supplied. + /// + public ContextRecallMetric( + ITextGenerator? generator = null, + IEmbeddingModel? embeddingModel = null, + double attributionThreshold = 0.5) + { + _generator = generator; + _embeddingModel = embeddingModel; + _attributionThreshold = attributionThreshold; + } + + /// Gets the name of this metric. + public override string Name => "Context Recall"; + + /// Gets the description of what this metric measures. + public override string Description => + "Measures what fraction of the ground-truth answer is supported by the retrieved contexts (RAGAS context recall)"; + + /// Gets a value indicating whether this metric requires ground truth. + protected override bool RequiresGroundTruth => true; + + /// + /// Evaluates context recall of the ground-truth claims against the retrieved documents. + /// + /// The grounded answer whose source documents are the retrieved context. + /// The reference answer that is broken into claims. + /// Context recall score (0-1). + protected override T EvaluateCore(GroundedAnswer answer, string? groundTruth) + { + var docs = answer.SourceDocuments.ToList(); + if (docs.Count == 0) + return NumOps.Zero; + + var claims = SplitIntoSentences(groundTruth!); + if (claims.Count == 0) + return NumOps.Zero; + + var context = string.Join(" ", docs.Select(d => d.Content)); + + int attributable = 0; + foreach (var claim in claims) + { + if (IsAttributable(claim, context, docs)) + attributable++; + } + + return NumOps.FromDouble((double)attributable / claims.Count); + } + + /// + /// Decides whether a ground-truth claim can be attributed to the retrieved context using the best + /// available signal: LLM judge, then embedding cosine (best matching document), then lexical + /// word coverage. + /// + private bool IsAttributable(string claim, string context, List> docs) + { + if (_generator != null) + { + var prompt = + "Can the statement be attributed to (supported by) the context? " + + "Reply with only 'yes' or 'no'.\n\n" + + $"Context: {context}\n\nStatement: {claim}\n\nAttributable (yes/no):"; + return ParseAffirmative(_generator.Generate(prompt)); + } + + if (_embeddingModel != null) + { + double best = 0.0; + foreach (var doc in docs) + { + double sim = EmbeddingCosine(_embeddingModel, claim, doc.Content); + if (sim > best) + best = sim; + } + + return best >= _attributionThreshold; + } + + // Offline lexical fallback: fraction of the claim's words present in the combined context. + var claimWords = GetWords(claim); + if (claimWords.Count == 0) + return false; + + var contextWords = GetWords(context); + double coverage = (double)claimWords.Intersect(contextWords).Count() / claimWords.Count; + return coverage >= _attributionThreshold; + } +} diff --git a/src/RetrievalAugmentedGeneration/Evaluation/ContextRelevanceMetric.cs b/src/RetrievalAugmentedGeneration/Evaluation/ContextRelevanceMetric.cs index c0a46a2133..63b387f034 100644 --- a/src/RetrievalAugmentedGeneration/Evaluation/ContextRelevanceMetric.cs +++ b/src/RetrievalAugmentedGeneration/Evaluation/ContextRelevanceMetric.cs @@ -1,5 +1,6 @@ using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; @@ -9,43 +10,62 @@ namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; /// /// The numeric data type used for calculations. /// -/// Measures how relevant the retrieved documents are to answering the query, -/// helping identify retrieval quality issues. +/// Measures how relevant the retrieved documents are to answering the query, helping identify +/// retrieval quality issues. When an embedding model is supplied, relevance is the mean cosine +/// similarity between the query and each retrieved document; otherwise the offline lexical Jaccard +/// word-overlap heuristic is used. /// [ComponentType(ComponentType.Evaluator)] [PipelineStage(PipelineStage.Evaluation)] public class ContextRelevanceMetric : RAGMetricBase { + private readonly IEmbeddingModel? _embeddingModel; + + /// + /// Initializes a new instance of the class. + /// + /// + /// Optional embedding model. When supplied, query-to-document relevance is measured with cosine + /// similarity; when null, the offline lexical Jaccard heuristic is used. + /// + public ContextRelevanceMetric(IEmbeddingModel? embeddingModel = null) + { + _embeddingModel = embeddingModel; + } + + /// Gets the name of this metric. public override string Name => "Context Relevance"; + + /// Gets the description of what this metric measures. public override string Description => "Measures how relevant the retrieved documents are to answering the query"; + + /// Gets a value indicating whether this metric requires ground truth. protected override bool RequiresGroundTruth => false; + /// + /// Evaluates the mean relevance of retrieved documents to the query. + /// + /// The grounded answer to evaluate. + /// Not used for this metric. + /// Relevance score (0-1). protected override T EvaluateCore(GroundedAnswer answer, string? groundTruth) { if (string.IsNullOrWhiteSpace(answer.Query) || !answer.SourceDocuments.Any()) return NumOps.Zero; - var totalRelevance = NumOps.Zero; - var count = 0; + double total = 0.0; + int count = 0; foreach (var doc in answer.SourceDocuments) { - var words1 = GetWords(answer.Query); - var words2 = GetWords(doc.Content); - - var intersection = words1.Intersect(words2).Count(); - var union = words1.Union(words2).Count(); - - if (union > 0) - { - var relevance = NumOps.Divide(NumOps.FromDouble(intersection), NumOps.FromDouble(union)); - totalRelevance = NumOps.Add(totalRelevance, relevance); - count++; - } + double relevance = _embeddingModel != null + ? Math.Max(0.0, EmbeddingCosine(_embeddingModel, answer.Query, doc.Content)) + : JaccardSimilarity(answer.Query, doc.Content); + + total += relevance; + count++; } - return count > 0 - ? NumOps.Divide(totalRelevance, NumOps.FromDouble(count)) - : NumOps.Zero; + return count > 0 ? NumOps.FromDouble(total / count) : NumOps.Zero; } } diff --git a/src/RetrievalAugmentedGeneration/Evaluation/FaithfulnessMetric.cs b/src/RetrievalAugmentedGeneration/Evaluation/FaithfulnessMetric.cs index 70ad4b0550..51f6beeb0b 100644 --- a/src/RetrievalAugmentedGeneration/Evaluation/FaithfulnessMetric.cs +++ b/src/RetrievalAugmentedGeneration/Evaluation/FaithfulnessMetric.cs @@ -1,5 +1,6 @@ using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; @@ -38,6 +39,22 @@ namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; [PipelineStage(PipelineStage.Evaluation)] public class FaithfulnessMetric : RAGMetricBase { + private readonly ITextGenerator? _generator; + + /// + /// Initializes a new instance of the class. + /// + /// + /// Optional text generator used as an NLI-style judge. When supplied, the answer is broken + /// into claims and each claim is checked for support against the retrieved context, giving a + /// score of supported-claims / total-claims. When null, the metric falls back to the + /// offline lexical word-overlap heuristic. + /// + public FaithfulnessMetric(ITextGenerator? generator = null) + { + _generator = generator; + } + /// /// Gets the name of this metric. /// @@ -65,15 +82,44 @@ protected override T EvaluateCore(GroundedAnswer answer, string? groundTruth) if (!answer.SourceDocuments.Any()) return NumOps.Zero; + var sourceText = string.Join(" ", answer.SourceDocuments.Select(d => d.Content)); + + if (_generator != null) + return EvaluateWithJudge(answer.Answer, sourceText); + + // Offline lexical fallback: fraction of answer words present in the source documents. var answerWords = GetWords(answer.Answer); if (answerWords.Count == 0) return NumOps.Zero; - var sourceText = string.Join(" ", answer.SourceDocuments.Select(d => d.Content)); var sourceWords = GetWords(sourceText); - var supportedWords = answerWords.Intersect(sourceWords).Count(); return NumOps.Divide(NumOps.FromDouble(supportedWords), NumOps.FromDouble(answerWords.Count)); } + + /// + /// NLI-style faithfulness: split the answer into claims and ask the generator whether each + /// claim can be inferred from the context. Score = supported claims / total claims. + /// + private T EvaluateWithJudge(string answerText, string context) + { + var claims = SplitIntoSentences(answerText); + if (claims.Count == 0) + return NumOps.Zero; + + int supported = 0; + foreach (var claim in claims) + { + var prompt = + "You are a strict fact-checker. Determine whether the statement can be directly " + + "inferred from the context. Reply with only 'yes' or 'no'.\n\n" + + $"Context: {context}\n\nStatement: {claim}\n\nSupported (yes/no):"; + var reply = _generator!.Generate(prompt); + if (ParseAffirmative(reply)) + supported++; + } + + return NumOps.Divide(NumOps.FromDouble(supported), NumOps.FromDouble(claims.Count)); + } } diff --git a/src/RetrievalAugmentedGeneration/Evaluation/RAGMetricBase.cs b/src/RetrievalAugmentedGeneration/Evaluation/RAGMetricBase.cs index 41b5fb8388..e08bab8d0c 100644 --- a/src/RetrievalAugmentedGeneration/Evaluation/RAGMetricBase.cs +++ b/src/RetrievalAugmentedGeneration/Evaluation/RAGMetricBase.cs @@ -99,4 +99,105 @@ protected HashSet GetWords(string text) .Split(new[] { ' ', '\t', '\n', '\r', '.', ',', ';', ':', '!', '?' }, StringSplitOptions.RemoveEmptyEntries)); } + + /// + /// Computes the Jaccard (word-overlap) similarity between two texts. Used as the offline + /// lexical fallback for metrics when no embedding model or text generator is supplied. + /// + /// The first text. + /// The second text. + /// A value in [0, 1]; 0 when both texts have no words. + protected double JaccardSimilarity(string a, string b) + { + var wa = GetWords(a); + var wb = GetWords(b); + var union = wa.Union(wb).Count(); + if (union == 0) + return 0.0; + + var intersection = wa.Intersect(wb).Count(); + return (double)intersection / union; + } + + /// + /// Splits text into individual sentences/claims on sentence-terminating punctuation and + /// line breaks. Whitespace-only fragments are dropped. + /// + /// The text to split. + /// The list of non-empty, trimmed sentences. + protected static List SplitIntoSentences(string text) + { + var result = new List(); + if (string.IsNullOrWhiteSpace(text)) + return result; + + var parts = text.Split(new[] { '.', '!', '?', '\n', '\r' }, StringSplitOptions.RemoveEmptyEntries); + foreach (var part in parts) + { + var trimmed = part.Trim(); + if (trimmed.Length > 0) + result.Add(trimmed); + } + + return result; + } + + /// + /// Parses a multi-line LLM reply into individual lines, stripping common list markers + /// (leading numbering like "1." / "1)" and bullets like "-", "*", "•"). + /// + /// The raw generator reply. + /// The cleaned, non-empty lines. + protected static List ParseListLines(string? reply) + { + var result = new List(); + if (string.IsNullOrWhiteSpace(reply)) + return result; + + var lines = reply!.Split(new[] { '\n', '\r' }, StringSplitOptions.RemoveEmptyEntries); + foreach (var line in lines) + { + var cleaned = System.Text.RegularExpressions.Regex.Replace( + line.Trim(), @"^\s*(\d+[\.\)]|[-*•])\s*", string.Empty).Trim(); + if (cleaned.Length > 0) + result.Add(cleaned); + } + + return result; + } + + /// + /// Interprets a yes/no style LLM verdict. Returns true when the reply affirms + /// (starts with or clearly contains "yes"/"true"/"supported"), false otherwise. + /// + /// The raw generator reply. + /// true for an affirmative verdict; otherwise false. + protected static bool ParseAffirmative(string? reply) + { + if (string.IsNullOrWhiteSpace(reply)) + return false; + + var s = reply!.Trim().ToLowerInvariant(); + if (s.StartsWith("yes") || s.StartsWith("true") || s.StartsWith("support")) + return true; + if (s.StartsWith("no") || s.StartsWith("false") || s.StartsWith("not")) + return false; + + return s.Contains("yes") && !s.Contains("no"); + } + + /// + /// Computes cosine similarity between the embeddings of two texts using the supplied + /// embedding model. Reuses the shared . + /// + /// The embedding model. + /// The first text. + /// The second text. + /// The cosine similarity as a double in [-1, 1]. + protected static double EmbeddingCosine(IEmbeddingModel model, string a, string b) + { + var va = model.Embed(a ?? string.Empty); + var vb = model.Embed(b ?? string.Empty); + return Convert.ToDouble(StatisticsHelper.CosineSimilarity(va, vb)); + } } diff --git a/src/RetrievalAugmentedGeneration/Evaluation/RetrievalMetrics.cs b/src/RetrievalAugmentedGeneration/Evaluation/RetrievalMetrics.cs new file mode 100644 index 0000000000..ddb14c8d29 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Evaluation/RetrievalMetrics.cs @@ -0,0 +1,291 @@ +using System; +using System.Collections.Generic; +using System.Linq; + +namespace AiDotNet.RetrievalAugmentedGeneration.Evaluation; + +/// +/// Pure-math information-retrieval (IR) metrics over ranked retrieval results: nDCG@k, MRR, MAP, +/// Hit-Rate@k, Precision@k and Recall@k. These compare a ranked list of retrieved document ids +/// against a set of ids known to be relevant. They involve no language model or embedding model and +/// are therefore fully deterministic. +/// +/// +/// For Beginners: after your retriever returns a ranked list of documents, these +/// functions grade the ranking. Precision@k asks "of the top k, how many were relevant?"; Recall@k +/// asks "of all relevant documents, how many did we find in the top k?"; MRR rewards putting the +/// first relevant document early; MAP and nDCG reward putting all relevant documents high up, +/// with nDCG additionally discounting lower ranks smoothly. +/// +/// Conventions: Precision@k divides by k (the requested cutoff). Recall@k, MAP, MRR and nDCG return +/// 0 when there are no relevant documents (the quantity is otherwise undefined). nDCG uses binary +/// relevance in the primary overloads and log2 rank discounting. +/// +/// +public static class RetrievalMetrics +{ + /// + /// Precision at rank k: the fraction of the top-k retrieved items that are relevant. + /// + /// The document id type. + /// The ranked list of retrieved document ids (best first). + /// The set of relevant document ids. + /// The rank cutoff (must be positive). + /// Precision@k in [0, 1]. Returns 0 when is not positive. + public static double PrecisionAtK(IReadOnlyList ranked, ISet relevant, int k) + { + if (ranked == null || relevant == null || k <= 0) + return 0.0; + + int limit = Math.Min(k, ranked.Count); + int hits = 0; + for (int i = 0; i < limit; i++) + { + if (relevant.Contains(ranked[i])) + hits++; + } + + return (double)hits / k; + } + + /// + /// Recall at rank k: the fraction of all relevant items that appear in the top-k. + /// + /// The document id type. + /// The ranked list of retrieved document ids (best first). + /// The set of relevant document ids. + /// The rank cutoff (must be positive). + /// Recall@k in [0, 1]. Returns 0 when there are no relevant items or k is not positive. + public static double RecallAtK(IReadOnlyList ranked, ISet relevant, int k) + { + if (ranked == null || relevant == null || relevant.Count == 0 || k <= 0) + return 0.0; + + int limit = Math.Min(k, ranked.Count); + int hits = 0; + for (int i = 0; i < limit; i++) + { + if (relevant.Contains(ranked[i])) + hits++; + } + + return (double)hits / relevant.Count; + } + + /// + /// Hit-rate at rank k: 1 if at least one relevant item appears in the top-k, otherwise 0. + /// + /// The document id type. + /// The ranked list of retrieved document ids (best first). + /// The set of relevant document ids. + /// The rank cutoff (must be positive). + /// 1.0 on a hit within the top-k; otherwise 0.0. + public static double HitRateAtK(IReadOnlyList ranked, ISet relevant, int k) + { + if (ranked == null || relevant == null || relevant.Count == 0 || k <= 0) + return 0.0; + + int limit = Math.Min(k, ranked.Count); + for (int i = 0; i < limit; i++) + { + if (relevant.Contains(ranked[i])) + return 1.0; + } + + return 0.0; + } + + /// + /// Reciprocal rank for a single query: 1 / (rank of the first relevant item), or 0 if none. + /// + /// The document id type. + /// The ranked list of retrieved document ids (best first). + /// The set of relevant document ids. + /// The reciprocal rank in [0, 1]. + public static double ReciprocalRank(IReadOnlyList ranked, ISet relevant) + { + if (ranked == null || relevant == null || relevant.Count == 0) + return 0.0; + + for (int i = 0; i < ranked.Count; i++) + { + if (relevant.Contains(ranked[i])) + return 1.0 / (i + 1); + } + + return 0.0; + } + + /// + /// Mean reciprocal rank (MRR) across a set of queries. + /// + /// The document id type. + /// One ranked list per query. + /// One relevant-id set per query, aligned with . + /// The mean of the per-query reciprocal ranks. + /// Thrown when either sequence is null. + /// Thrown when the two sequences differ in length. + public static double MeanReciprocalRank( + IEnumerable> rankings, + IEnumerable> relevantSets) + { + var pairs = ZipQueries(rankings, relevantSets); + if (pairs.Count == 0) + return 0.0; + + return pairs.Average(p => ReciprocalRank(p.Ranked, p.Relevant)); + } + + /// + /// Average precision (AP) for a single query: the mean of Precision@k taken at every rank that + /// holds a relevant item, divided by the total number of relevant items. + /// + /// The document id type. + /// The ranked list of retrieved document ids (best first). + /// The set of relevant document ids. + /// The average precision in [0, 1]; 0 when there are no relevant items. + public static double AveragePrecision(IReadOnlyList ranked, ISet relevant) + { + if (ranked == null || relevant == null || relevant.Count == 0) + return 0.0; + + int hits = 0; + double sum = 0.0; + for (int i = 0; i < ranked.Count; i++) + { + if (relevant.Contains(ranked[i])) + { + hits++; + sum += (double)hits / (i + 1); + } + } + + return sum / relevant.Count; + } + + /// + /// Mean average precision (MAP) across a set of queries. + /// + /// The document id type. + /// One ranked list per query. + /// One relevant-id set per query, aligned with . + /// The mean of the per-query average precisions. + /// Thrown when either sequence is null. + /// Thrown when the two sequences differ in length. + public static double MeanAveragePrecision( + IEnumerable> rankings, + IEnumerable> relevantSets) + { + var pairs = ZipQueries(rankings, relevantSets); + if (pairs.Count == 0) + return 0.0; + + return pairs.Average(p => AveragePrecision(p.Ranked, p.Relevant)); + } + + /// + /// Discounted cumulative gain at rank k with binary relevance and log2 discounting. + /// + /// The document id type. + /// The ranked list of retrieved document ids (best first). + /// The set of relevant document ids. + /// The rank cutoff (must be positive). + /// DCG@k. + public static double DcgAtK(IReadOnlyList ranked, ISet relevant, int k) + { + if (ranked == null || relevant == null || k <= 0) + return 0.0; + + int limit = Math.Min(k, ranked.Count); + double dcg = 0.0; + for (int i = 0; i < limit; i++) + { + if (relevant.Contains(ranked[i])) + dcg += 1.0 / Log2(i + 2); // rank position (i+1); discount log2(rank+1) = log2(i+2) + } + + return dcg; + } + + /// + /// Normalized discounted cumulative gain at rank k with binary relevance. + /// + /// The document id type. + /// The ranked list of retrieved document ids (best first). + /// The set of relevant document ids. + /// The rank cutoff (must be positive). + /// nDCG@k in [0, 1]; 0 when there are no relevant items. + public static double NdcgAtK(IReadOnlyList ranked, ISet relevant, int k) + { + if (ranked == null || relevant == null || relevant.Count == 0 || k <= 0) + return 0.0; + + double dcg = DcgAtK(ranked, relevant, k); + + // Ideal DCG: all relevant items ranked first, up to k positions. + int idealCount = Math.Min(k, relevant.Count); + double idcg = 0.0; + for (int i = 0; i < idealCount; i++) + idcg += 1.0 / Log2(i + 2); + + return idcg > 0.0 ? dcg / idcg : 0.0; + } + + /// + /// Normalized discounted cumulative gain at rank k with graded (non-binary) relevance gains. + /// + /// The document id type. + /// The ranked list of retrieved document ids (best first). + /// A map from document id to its graded relevance gain (missing ids score 0). + /// The rank cutoff (must be positive). + /// Graded nDCG@k in [0, 1]; 0 when the ideal DCG is 0. + public static double NdcgAtK(IReadOnlyList ranked, IReadOnlyDictionary gains, int k) + { + if (ranked == null || gains == null || k <= 0) + return 0.0; + + int limit = Math.Min(k, ranked.Count); + double dcg = 0.0; + for (int i = 0; i < limit; i++) + { + double gain; + if (gains.TryGetValue(ranked[i], out gain)) + dcg += gain / Log2(i + 2); + } + + var idealGains = gains.Values.Where(g => g > 0.0).OrderByDescending(g => g).ToList(); + int idealCount = Math.Min(k, idealGains.Count); + double idcg = 0.0; + for (int i = 0; i < idealCount; i++) + idcg += idealGains[i] / Log2(i + 2); + + return idcg > 0.0 ? dcg / idcg : 0.0; + } + + private static double Log2(double value) + { + // Math.Log2 is unavailable on net471; use the change-of-base identity. + return Math.Log(value) / Math.Log(2.0); + } + + private static List<(IReadOnlyList Ranked, ISet Relevant)> ZipQueries( + IEnumerable> rankings, + IEnumerable> relevantSets) + { + if (rankings == null) + throw new ArgumentNullException(nameof(rankings)); + if (relevantSets == null) + throw new ArgumentNullException(nameof(relevantSets)); + + var rankedList = rankings.ToList(); + var relevantList = relevantSets.ToList(); + if (rankedList.Count != relevantList.Count) + throw new ArgumentException("The number of rankings must match the number of relevant-id sets.", nameof(relevantSets)); + + var pairs = new List<(IReadOnlyList Ranked, ISet Relevant)>(rankedList.Count); + for (int i = 0; i < rankedList.Count; i++) + pairs.Add((rankedList[i], relevantList[i])); + + return pairs; + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Evaluation/RagEvaluationMetricsTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Evaluation/RagEvaluationMetricsTests.cs new file mode 100644 index 0000000000..d20dddebf0 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Evaluation/RagEvaluationMetricsTests.cs @@ -0,0 +1,321 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using AiDotNet.Interfaces; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Evaluation; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Xunit; + +namespace AiDotNet.Tests.UnitTests.RetrievalAugmentedGeneration.Evaluation; + +/// +/// Verifies the upgraded RAG evaluation metrics: pure-math IR metrics with exact expected values, +/// embedding-similarity metrics driven by a controlled fake embedding model, and LLM-judge metrics +/// driven by a scripted fake text generator. All tests are offline (no network / real model). +/// +public class RagEvaluationMetricsTests +{ + // ---- Test doubles ------------------------------------------------------- + + /// Records prompts and returns a scripted response. + private sealed class FakeTextGen : ITextGenerator + { + private readonly Func _responder; + public List Prompts { get; } = new(); + public FakeTextGen(Func responder) => _responder = responder; + public string Generate(string prompt) + { + Prompts.Add(prompt); + return _responder(prompt); + } + } + + /// Returns controlled vectors from an exact string-to-vector map (dimension 2). + private sealed class FakeEmbeddingModel : IEmbeddingModel + { + private readonly Dictionary _map; + public FakeEmbeddingModel(Dictionary map) => _map = map; + public int EmbeddingDimension => 2; + public int MaxTokens => 512; + public Vector Embed(string text) + { + if (text != null && _map.TryGetValue(text, out var v)) + return new Vector(v); + // Deterministic non-zero default so cosine never divides by zero for unmapped text. + return new Vector(new[] { 1.0, 0.0 }); + } + public Task> EmbedAsync(string text) => Task.FromResult(Embed(text)); + public Matrix EmbedBatch(IEnumerable texts) => throw new NotSupportedException(); + public Task> EmbedBatchAsync(IEnumerable texts) => throw new NotSupportedException(); + } + + private static GroundedAnswer Answer(string query, string answer, params Document[] docs) + { + return new GroundedAnswer + { + Query = query, + Answer = answer, + SourceDocuments = docs.ToList().AsReadOnly() + }; + } + + private static IReadOnlyList Ranked(params string[] ids) => ids; + private static ISet RelevantSet(params string[] ids) => new HashSet(ids); + + // ---- IR metrics (exact numeric assertions) ------------------------------ + + [Fact] + public void PrecisionAtK_KnownRanking() + { + var ranked = Ranked("d1", "d2", "d3", "d4"); // relevant at ranks 1 and 3 + var rel = RelevantSet("d1", "d3"); + + Assert.Equal(1.0, RetrievalMetrics.PrecisionAtK(ranked, rel, 1), 6); + Assert.Equal(0.5, RetrievalMetrics.PrecisionAtK(ranked, rel, 2), 6); + Assert.Equal(0.5, RetrievalMetrics.PrecisionAtK(ranked, rel, 4), 6); + } + + [Fact] + public void RecallAtK_KnownRanking() + { + var ranked = Ranked("d1", "d2", "d3", "d4"); + var rel = RelevantSet("d1", "d3"); + + Assert.Equal(0.5, RetrievalMetrics.RecallAtK(ranked, rel, 2), 6); + Assert.Equal(1.0, RetrievalMetrics.RecallAtK(ranked, rel, 4), 6); + Assert.Equal(0.0, RetrievalMetrics.RecallAtK(ranked, RelevantSet(), 4), 6); + } + + [Fact] + public void HitRateAtK_KnownRanking() + { + var ranked = Ranked("d1", "d2", "d3"); + Assert.Equal(1.0, RetrievalMetrics.HitRateAtK(ranked, RelevantSet("d3"), 3), 6); + Assert.Equal(0.0, RetrievalMetrics.HitRateAtK(ranked, RelevantSet("d3"), 2), 6); + Assert.Equal(0.0, RetrievalMetrics.HitRateAtK(ranked, RelevantSet("zzz"), 3), 6); + } + + [Fact] + public void ReciprocalRank_And_MRR() + { + Assert.Equal(1.0, RetrievalMetrics.ReciprocalRank(Ranked("d1", "d2"), RelevantSet("d1")), 6); + Assert.Equal(0.5, RetrievalMetrics.ReciprocalRank(Ranked("d2", "d1"), RelevantSet("d1")), 6); + Assert.Equal(0.0, RetrievalMetrics.ReciprocalRank(Ranked("d2", "d3"), RelevantSet("d1")), 6); + + var rankings = new[] { Ranked("d1", "d2"), Ranked("d2", "d1") }; + var relevants = new[] { RelevantSet("d1"), RelevantSet("d1") }; + Assert.Equal(0.75, RetrievalMetrics.MeanReciprocalRank(rankings, relevants), 6); // (1.0 + 0.5)/2 + } + + [Fact] + public void AveragePrecision_And_MAP() + { + // ranked d1(rel) d2 d3(rel) d4, R=2: AP = (P@1 + P@3)/2 = (1 + 2/3)/2 = 0.833333 + var ranked = Ranked("d1", "d2", "d3", "d4"); + var rel = RelevantSet("d1", "d3"); + Assert.Equal(0.8333333, RetrievalMetrics.AveragePrecision(ranked, rel), 6); + + var rankings = new[] { ranked, Ranked("d1") }; + var relevants = new[] { rel, RelevantSet("d1") }; + // MAP = (0.833333 + 1.0)/2 = 0.916667 + Assert.Equal(0.9166667, RetrievalMetrics.MeanAveragePrecision(rankings, relevants), 6); + } + + [Fact] + public void NdcgAtK_BinaryRelevance_KnownRanking() + { + // ranked d1(rel) d2 d3(rel): DCG = 1/log2(2) + 1/log2(4) = 1 + 0.5 = 1.5 + // IDCG (2 relevant) = 1/log2(2) + 1/log2(3) = 1 + 0.6309298 = 1.6309298 + // nDCG = 1.5 / 1.6309298 = 0.9197208 + var ranked = Ranked("d1", "d2", "d3"); + var rel = RelevantSet("d1", "d3"); + Assert.Equal(0.9197208, RetrievalMetrics.NdcgAtK(ranked, rel, 3), 6); + + // Perfect ranking -> 1.0 + Assert.Equal(1.0, RetrievalMetrics.NdcgAtK(Ranked("d1", "d3", "d2"), rel, 3), 6); + + // No relevant -> 0 + Assert.Equal(0.0, RetrievalMetrics.NdcgAtK(ranked, RelevantSet(), 3), 6); + } + + [Fact] + public void NdcgAtK_GradedRelevance_KnownRanking() + { + // ranked b,a,c ; gains a=3,b=2,c=1 + // DCG = 2/log2(2) + 3/log2(3) + 1/log2(4) = 2 + 1.8927893 + 0.5 = 4.3927893 + // IDCG(order a,b,c) = 3/log2(2) + 2/log2(3) + 1/log2(4) = 3 + 1.2618595 + 0.5 = 4.7618595 + // nDCG = 4.3927893 / 4.7618595 = 0.92249 + var ranked = Ranked("b", "a", "c"); + var gains = new Dictionary { ["a"] = 3, ["b"] = 2, ["c"] = 1 }; + Assert.Equal(0.92249, RetrievalMetrics.NdcgAtK(ranked, gains, 3), 5); + } + + // ---- Embedding-similarity metrics --------------------------------------- + + [Fact] + public void AnswerSimilarity_WithEmbeddingModel_UsesCosine() + { + var emb = new FakeEmbeddingModel(new Dictionary + { + ["cat"] = new[] { 1.0, 0.0 }, + ["kitten"] = new[] { 1.0, 1.0 }, + }); + var metric = new AnswerSimilarityMetric(emb); + + var score = metric.Evaluate(Answer("q", "cat"), "kitten"); + Assert.Equal(0.7071068, score, 6); // cos([1,0],[1,1]) + } + + [Fact] + public void AnswerSimilarity_WithoutModel_FallsBackToJaccard() + { + var metric = new AnswerSimilarityMetric(); + // words {a,b,c} vs {b,c,d}: intersection 2, union 4 -> 0.5 + var score = metric.Evaluate(Answer("q", "a b c"), "b c d"); + Assert.Equal(0.5, score, 6); + } + + [Fact] + public void ContextRelevance_WithEmbeddingModel_AveragesCosine() + { + var emb = new FakeEmbeddingModel(new Dictionary + { + ["the query"] = new[] { 1.0, 0.0 }, + ["doc one"] = new[] { 1.0, 0.0 }, // cos 1 + ["doc two"] = new[] { 0.0, 1.0 }, // cos 0 + }); + var metric = new ContextRelevanceMetric(emb); + + var ans = Answer("the query", "answer", + new Document("1", "doc one"), + new Document("2", "doc two")); + + Assert.Equal(0.5, metric.Evaluate(ans), 6); + } + + [Fact] + public void AnswerRelevance_WithModels_MeanQuestionQueryCosine() + { + var gen = new FakeTextGen(_ => "question one\nquestion two"); + var emb = new FakeEmbeddingModel(new Dictionary + { + ["the query"] = new[] { 1.0, 0.0 }, + ["question one"] = new[] { 1.0, 0.0 }, // cos 1 + ["question two"] = new[] { 0.0, 1.0 }, // cos 0 + }); + var metric = new AnswerRelevanceMetric(gen, emb, numQuestions: 2); + + var score = metric.Evaluate(Answer("the query", "some answer")); + Assert.Equal(0.5, score, 6); + Assert.Single(gen.Prompts); + } + + [Fact] + public void AnswerCorrectness_WithJudgeAndEmbedding_AveragesSignals() + { + var gen = new FakeTextGen(_ => "8"); // 8/10 = 0.8 + var emb = new FakeEmbeddingModel(new Dictionary + { + ["ans"] = new[] { 1.0, 1.0 }, + ["ref"] = new[] { 1.0, 0.0 }, // cos = 0.7071068 + }); + var metric = new AnswerCorrectnessMetric(gen, emb); + + var score = metric.Evaluate(Answer("q", "ans"), "ref"); + Assert.Equal((0.8 + 0.7071068) / 2.0, score, 6); + } + + // ---- LLM-judge metrics -------------------------------------------------- + + [Fact] + public void Faithfulness_WithJudge_ScoresSupportedClaims() + { + // Answer has three claims; the judge supports the first two, rejects the third. + var gen = new FakeTextGen(p => + p.Contains("Statement: The sky is blue") ? "yes" : + p.Contains("Statement: Water is wet") ? "yes" : + "no"); + var metric = new FaithfulnessMetric(gen); + + var ans = Answer("q", "The sky is blue. Water is wet. Cats can fly.", + new Document("1", "context about sky, water, and cats")); + + var score = metric.Evaluate(ans); + Assert.Equal(2.0 / 3.0, score, 6); + Assert.Equal(3, gen.Prompts.Count); + } + + [Fact] + public void Faithfulness_WithoutGenerator_FallsBackToWordOverlap() + { + var metric = new FaithfulnessMetric(); + // answer words {alpha,beta}; source has alpha only -> 1/2 + var ans = Answer("q", "alpha beta", new Document("1", "alpha only here")); + Assert.Equal(0.5, metric.Evaluate(ans), 6); + } + + [Fact] + public void ContextPrecision_WithJudge_AveragePrecisionOverRanks() + { + // docs: relevant, not, relevant -> weighted (P@1=1 + P@3=2/3)/2 relevant = 0.833333 + var gen = new FakeTextGen(p => p.Contains("relevant") ? "yes" : "no"); + var metric = new ContextPrecisionMetric(gen); + + var ans = Answer("q", "answer", + new Document("1", "alpha relevant"), + new Document("2", "beta noise"), + new Document("3", "gamma relevant")); + + Assert.Equal(0.8333333, metric.Evaluate(ans, "ground truth"), 6); + } + + [Fact] + public void ContextRecall_WithJudge_FractionOfAttributableClaims() + { + // ground truth has two claims; judge attributes the first, not the second -> 0.5 + var gen = new FakeTextGen(p => p.Contains("Statement: Fact one") ? "yes" : "no"); + var metric = new ContextRecallMetric(gen); + + var ans = Answer("q", "answer", new Document("1", "some retrieved context")); + + var score = metric.Evaluate(ans, "Fact one. Fact two."); + Assert.Equal(0.5, score, 6); + Assert.Equal(2, gen.Prompts.Count); + } + + [Fact] + public void ContextRecall_WithoutModels_FallsBackToWordCoverage() + { + var metric = new ContextRecallMetric(); + // GT claim "alpha beta gamma": context covers alpha,beta (2/3 >= 0.5) -> attributable + var ans = Answer("q", "answer", new Document("1", "alpha beta present")); + var score = metric.Evaluate(ans, "alpha beta gamma"); + Assert.Equal(1.0, score, 6); + } + + // ---- Evaluator wiring --------------------------------------------------- + + [Fact] + public void RAGEvaluator_EnumeratesInjectedMetrics_IncludingNewOnes() + { + var emb = new FakeEmbeddingModel(new Dictionary + { + ["the query"] = new[] { 1.0, 0.0 }, + ["ctx"] = new[] { 1.0, 0.0 }, + }); + var metrics = new IRAGMetric[] + { + new ContextRelevanceMetric(emb), + new FaithfulnessMetric(), + }; + var evaluator = new RAGEvaluator(metrics); + + var ans = Answer("the query", "ctx", new Document("1", "ctx")); + var result = evaluator.Evaluate(ans); + + Assert.True(result.MetricScores.ContainsKey("Context Relevance")); + Assert.True(result.MetricScores.ContainsKey("Faithfulness")); + } +} From 4fbf9633ecee43a621cb7f25987e2263dedd9fc3 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 16:28:28 -0400 Subject: [PATCH 10/25] feat(rag): real embedding-based semantic chunking (#26 part 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SemanticChunkingStrategy was "semantic" in name only — it packed sentences to a size budget with no embeddings. It now, when given a sentence embedder, places chunk boundaries at semantic breakpoints (consecutive-sentence cosine distance above a percentile threshold — Kamradt/LlamaIndex SemanticSplitter), while still capping chunk size; without an embedder it falls back to the original size-based packing. Sentence spans are located in the source text so positions round-trip. Verified: 3 tests (topic-shift split, no-embedder fallback, position round-trip). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../SemanticChunkingStrategy.cs | 151 ++++++++++++++++-- .../SemanticChunkingStrategyTests.cs | 59 +++++++ 2 files changed, 196 insertions(+), 14 deletions(-) create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ChunkingStrategies/SemanticChunkingStrategyTests.cs diff --git a/src/RetrievalAugmentedGeneration/ChunkingStrategies/SemanticChunkingStrategy.cs b/src/RetrievalAugmentedGeneration/ChunkingStrategies/SemanticChunkingStrategy.cs index fde469885e..10386b5f13 100644 --- a/src/RetrievalAugmentedGeneration/ChunkingStrategies/SemanticChunkingStrategy.cs +++ b/src/RetrievalAugmentedGeneration/ChunkingStrategies/SemanticChunkingStrategy.cs @@ -8,32 +8,106 @@ namespace AiDotNet.RetrievalAugmentedGeneration.ChunkingStrategies { /// - /// Semantic-based text chunking that uses embeddings to group related content. + /// Semantic text chunking. When an embedding function is supplied, it embeds each sentence and starts a + /// new chunk at semantic breakpoints — consecutive-sentence cosine distances above a percentile + /// threshold (Kamradt / LlamaIndex SemanticSplitter). Without an embedding function it falls back to + /// size-based sentence packing. /// + /// + /// Previously this "semantic" splitter used no embeddings at all — it just packed sentences to a size + /// budget. It now performs real embedding-distance breakpoint detection when given an embedder. + /// [ComponentType(ComponentType.Chunker)] [PipelineStage(PipelineStage.DataIngestion)] public class SemanticChunkingStrategy : ChunkingStrategyBase { - /// - /// Initializes a new instance of the class. - /// - /// The maximum chunk size in characters. - /// The chunk overlap in characters. + private readonly Func, IReadOnlyList>? _embedBatch; + private readonly double _breakpointPercentile; + + /// Maximum chunk size in characters (hard cap even in semantic mode). + /// Chunk overlap in characters (size-based fallback only). + /// + /// Optional sentence embedder (batch: sentences → vectors). When provided, chunk boundaries are placed + /// at semantic breakpoints; when null, a size-based sentence-packing fallback is used. + /// + /// + /// Distance percentile (0–100) above which a sentence boundary becomes a chunk breakpoint. Higher = + /// fewer, larger chunks. Default 95. + /// public SemanticChunkingStrategy( int maxChunkSize = 1000, - int chunkOverlap = 200) + int chunkOverlap = 200, + Func, IReadOnlyList>? embedBatch = null, + double breakpointPercentile = 95.0) : base(maxChunkSize, chunkOverlap) { + _embedBatch = embedBatch; + _breakpointPercentile = Math.Min(100.0, Math.Max(0.0, breakpointPercentile)); } - /// - /// Chunks text based on semantic similarity between sentences. - /// - /// The text to chunk. - /// A collection of semantically coherent chunks with positions. protected override IEnumerable<(string Chunk, int StartPosition, int EndPosition)> ChunkCore(string text) { var sentences = SplitIntoSentences(text); + if (_embedBatch == null || sentences.Count < 3) + { + return SizeBasedChunks(sentences); + } + + return SemanticChunks(sentences, text); + } + + // Real semantic chunking: break where the cosine DISTANCE between consecutive sentence embeddings + // exceeds the configured percentile of all such distances (a topic shift), also capping chunk size. + private IEnumerable<(string Chunk, int StartPosition, int EndPosition)> SemanticChunks(List sentences, string text) + { + var spans = LocateSentenceSpans(sentences, text); + var embeddings = _embedBatch!(sentences); + + var distances = new List(sentences.Count - 1); + for (int i = 0; i < sentences.Count - 1; i++) + { + distances.Add(1.0 - CosineSimilarity(embeddings[i], embeddings[i + 1])); + } + + double threshold = Percentile(distances, _breakpointPercentile); + + var results = new List<(string, int, int)>(); + int chunkStart = 0; + for (int i = 0; i < sentences.Count; i++) + { + bool lastSentence = i == sentences.Count - 1; + int startPos = spans[chunkStart].Start; + int endPos = spans[i].End; + bool sizeExceeded = (endPos - startPos) >= ChunkSize && i > chunkStart; + bool semanticBreak = !lastSentence && distances[i] > threshold; + + if (sizeExceeded) + { + // Emit up to the previous sentence, then reconsider the current one as a chunk start. + int prevEnd = spans[i - 1].End; + results.Add((text.Substring(startPos, prevEnd - startPos), startPos, prevEnd)); + chunkStart = i; + startPos = spans[chunkStart].Start; + endPos = spans[i].End; + } + + if (lastSentence) + { + results.Add((text.Substring(startPos, endPos - startPos), startPos, endPos)); + } + else if (semanticBreak) + { + results.Add((text.Substring(startPos, endPos - startPos), startPos, endPos)); + chunkStart = i + 1; + } + } + + return results; + } + + // Size-based fallback (original behavior) when no embedder is supplied. + private IEnumerable<(string Chunk, int StartPosition, int EndPosition)> SizeBasedChunks(List sentences) + { var currentChunk = new List(); var currentSize = 0; var position = 0; @@ -46,7 +120,7 @@ public SemanticChunkingStrategy( var endPos = position + chunkText.Length; yield return (chunkText, position, endPos); - position = endPos - ChunkOverlap; + position = Math.Max(0, endPos - ChunkOverlap); currentChunk.Clear(); currentSize = 0; } @@ -62,10 +136,59 @@ public SemanticChunkingStrategy( } } + // Find each sentence's [start,end) span in the original text, scanning left-to-right so repeated + // sentences map to distinct occurrences. + private static List<(int Start, int End)> LocateSentenceSpans(List sentences, string text) + { + var spans = new List<(int, int)>(sentences.Count); + int searchFrom = 0; + foreach (var sentence in sentences) + { + int idx = text.IndexOf(sentence, searchFrom, StringComparison.Ordinal); + if (idx < 0) + { + // Sentence normalization changed the text; approximate with the running cursor. + idx = Math.Min(searchFrom, Math.Max(0, text.Length - 1)); + } + + int end = Math.Min(text.Length, idx + sentence.Length); + spans.Add((idx, end)); + searchFrom = end; + } + + return spans; + } + + private static double CosineSimilarity(double[] a, double[] b) + { + if (a == null || b == null || a.Length == 0 || a.Length != b.Length) return 0.0; + double dot = 0, na = 0, nb = 0; + for (int i = 0; i < a.Length; i++) + { + dot += a[i] * b[i]; + na += a[i] * a[i]; + nb += b[i] * b[i]; + } + + double denom = Math.Sqrt(na) * Math.Sqrt(nb); + return denom > 1e-12 ? dot / denom : 0.0; + } + + private static double Percentile(List values, double percentile) + { + if (values.Count == 0) return 0.0; + var sorted = values.OrderBy(v => v).ToList(); + double rank = (percentile / 100.0) * (sorted.Count - 1); + int lo = (int)Math.Floor(rank); + int hi = (int)Math.Ceiling(rank); + if (lo == hi) return sorted[lo]; + double frac = rank - lo; + return sorted[lo] * (1 - frac) + sorted[hi] * frac; + } + private List SplitIntoSentences(string text) { return Helpers.TextProcessingHelper.SplitIntoSentences(text); } } } - diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ChunkingStrategies/SemanticChunkingStrategyTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ChunkingStrategies/SemanticChunkingStrategyTests.cs new file mode 100644 index 0000000000..7b06dd9a55 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ChunkingStrategies/SemanticChunkingStrategyTests.cs @@ -0,0 +1,59 @@ +using System.Collections.Generic; +using System.Linq; +using AiDotNet.RetrievalAugmentedGeneration.ChunkingStrategies; +using Xunit; + +namespace AiDotNet.Tests.UnitTests.RetrievalAugmentedGeneration.ChunkingStrategies; + +/// +/// Tests for the embedding-based : it should place a chunk boundary +/// at the semantic topic shift when given an embedder, and fall back to size packing without one. +/// +public class SemanticChunkingStrategyTests +{ + // Two orthogonal "topics": cat sentences → [1,0], dog sentences → [0,1]. + private static IReadOnlyList TopicEmbed(IReadOnlyList sentences) + => sentences.Select(s => s.ToLowerInvariant().Contains("cat") + ? new[] { 1.0, 0.0 } + : new[] { 0.0, 1.0 }).ToList(); + + [Fact] + public void SplitsAtSemanticBreakpoint() + { + var strat = new SemanticChunkingStrategy( + maxChunkSize: 10000, chunkOverlap: 0, embedBatch: TopicEmbed, breakpointPercentile: 90); + var text = "Cats purr softly. Cats sleep all day. Dogs bark loudly. Dogs run fast."; + + var chunks = strat.Chunk(text).ToList(); + + Assert.Equal(2, chunks.Count); + Assert.Contains("Cats", chunks[0]); + Assert.DoesNotContain("Dogs", chunks[0]); + Assert.Contains("Dogs", chunks[1]); + } + + [Fact] + public void FallsBackToSizePackingWithoutEmbedder() + { + var strat = new SemanticChunkingStrategy(maxChunkSize: 10000, chunkOverlap: 0); + var text = "Cats purr softly. Cats sleep all day. Dogs bark loudly. Dogs run fast."; + + var chunks = strat.Chunk(text).ToList(); + + // No embedder → one size-bounded chunk (no semantic split). + Assert.Single(chunks); + } + + [Fact] + public void PositionsMapBackToOriginalText() + { + var strat = new SemanticChunkingStrategy( + maxChunkSize: 10000, chunkOverlap: 0, embedBatch: TopicEmbed, breakpointPercentile: 90); + var text = "Cats purr softly. Cats sleep all day. Dogs bark loudly. Dogs run fast."; + + foreach (var (chunk, start, end) in strat.ChunkWithPositions(text)) + { + Assert.Equal(chunk, text.Substring(start, end - start)); + } + } +} From ab0ea080c04f4f277ee25466597bccdc2b0e5b13 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 16:44:12 -0400 Subject: [PATCH 11/25] feat(rag): vector quantization (scalar, binary, product) + quantized index Co-Authored-By: Claude Opus 4.8 (1M context) --- .../Indexes/QuantizedFlatIndex.cs | 215 ++++++++++ .../Quantization/BinaryQuantizer.cs | 177 ++++++++ .../Quantization/IVectorQuantizer.cs | 62 +++ .../Quantization/ProductQuantizer.cs | 403 ++++++++++++++++++ .../Quantization/ScalarQuantizer.cs | 168 ++++++++ .../Indexes/QuantizedFlatIndexTests.cs | 177 ++++++++ .../Quantization/BinaryQuantizerTests.cs | 73 ++++ .../Quantization/ProductQuantizerTests.cs | 188 ++++++++ .../Quantization/ScalarQuantizerTests.cs | 112 +++++ 9 files changed, 1575 insertions(+) create mode 100644 src/RetrievalAugmentedGeneration/VectorSearch/Indexes/QuantizedFlatIndex.cs create mode 100644 src/RetrievalAugmentedGeneration/VectorSearch/Quantization/BinaryQuantizer.cs create mode 100644 src/RetrievalAugmentedGeneration/VectorSearch/Quantization/IVectorQuantizer.cs create mode 100644 src/RetrievalAugmentedGeneration/VectorSearch/Quantization/ProductQuantizer.cs create mode 100644 src/RetrievalAugmentedGeneration/VectorSearch/Quantization/ScalarQuantizer.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/QuantizedFlatIndexTests.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/BinaryQuantizerTests.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/ProductQuantizerTests.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/ScalarQuantizerTests.cs diff --git a/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/QuantizedFlatIndex.cs b/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/QuantizedFlatIndex.cs new file mode 100644 index 0000000000..c3ca810fad --- /dev/null +++ b/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/QuantizedFlatIndex.cs @@ -0,0 +1,215 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Helpers; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Quantization; +using AiDotNet.Validation; + +namespace AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Indexes +{ + /// + /// A flat (brute-force) index that stores compressed quantized codes instead of full + /// vectors, trading a small amount of accuracy for large memory savings and faster + /// scanning. This mirrors FAISS IndexPQ/IndexScalarQuantizer and Qdrant's + /// on-disk quantized collections. + /// + /// + /// + /// Vectors added before the underlying quantizer is trained are held in raw form and + /// searched with an exact fallback (equivalent to a ). Once + /// is called, those vectors are encoded to compact codes, the raw + /// copies are released, and subsequent additions are encoded immediately. + /// + /// + /// Ranking after training uses the quantizer to score candidates: for a + /// paired with a distance metric it uses fast + /// Asymmetric Distance Computation (ADC); otherwise it reconstructs each candidate with + /// and applies the supplied metric. + /// + /// + /// For Beginners: This index keeps a tiny "summary" of each vector instead of the + /// whole thing, so it uses far less memory. You still call Add and Search the same way; + /// just remember to call Train once (on a representative sample) so it can learn how to + /// summarize. Until then it behaves like an ordinary exact index. + /// + /// + /// The numeric type for vector operations. + [ComponentType(ComponentType.VectorIndex)] + [PipelineStage(PipelineStage.Retrieval)] + public class QuantizedFlatIndex : IVectorIndex + { + private readonly ISimilarityMetric _metric; + private readonly IVectorQuantizer _quantizer; + private readonly INumericOperations _numOps; + + // Raw vectors kept only until the quantizer is trained (then cleared to save memory). + private readonly Dictionary> _raw; + // Compressed codes used once trained. + private readonly Dictionary _codes; + + /// + public int Count => _raw.Count + _codes.Count; + + /// + /// Gets a value indicating whether the underlying quantizer has been trained. + /// + public bool IsTrained => _quantizer.IsTrained; + + /// + /// Gets the quantizer used to compress vectors. + /// + public IVectorQuantizer Quantizer => _quantizer; + + /// + /// Initializes a new instance of the class. + /// + /// The similarity metric used for ranking (and for the untrained fallback). + /// The quantizer that compresses stored vectors. + public QuantizedFlatIndex(ISimilarityMetric metric, IVectorQuantizer quantizer) + { + Guard.NotNull(metric); + Guard.NotNull(quantizer); + _metric = metric; + _quantizer = quantizer; + _numOps = MathHelper.GetNumericOperations(); + _raw = new Dictionary>(); + _codes = new Dictionary(); + } + + /// + /// Trains the underlying quantizer, then encodes and compacts any vectors that were + /// added beforehand. + /// + /// + /// Optional training vectors. If null, the quantizer is trained on the vectors already + /// added to this index. + /// + public void Train(IEnumerable>? vectors = null) + { + var trainingData = vectors != null ? vectors.ToList() : _raw.Values.ToList(); + if (trainingData.Count == 0) + throw new InvalidOperationException("Cannot train: no training vectors provided and the index is empty."); + + _quantizer.Train(trainingData); + + // Encode and release any pending raw vectors. + foreach (var kvp in _raw) + _codes[kvp.Key] = _quantizer.Encode(kvp.Value); + _raw.Clear(); + } + + /// + public void Add(string id, Vector vector) + { + if (string.IsNullOrEmpty(id)) + throw new ArgumentException("ID cannot be null or empty", nameof(id)); + if (vector == null) + throw new ArgumentNullException(nameof(vector)); + + if (_quantizer.IsTrained) + { + _raw.Remove(id); + _codes[id] = _quantizer.Encode(vector); + } + else + { + _codes.Remove(id); + _raw[id] = vector; + } + } + + /// + public void AddBatch(Dictionary> vectors) + { + if (vectors == null) + throw new ArgumentNullException(nameof(vectors)); + + foreach (var kvp in vectors) + Add(kvp.Key, kvp.Value); + } + + /// + public List<(string Id, T Score)> Search(Vector query, int k) + { + if (query == null) + throw new ArgumentNullException(nameof(query)); + if (k <= 0) + throw new ArgumentException("k must be positive", nameof(k)); + + // Untrained: exact fallback over any raw vectors (behaves like FlatIndex). + if (!_quantizer.IsTrained) + return RankRaw(query, k); + + if (_codes.Count == 0) + return new List<(string Id, T Score)>(); + + // Fast path: Product Quantizer + distance metric -> Asymmetric Distance Computation. + if (_quantizer is ProductQuantizer pq && !_metric.HigherIsBetter) + return RankAdc(pq, query, k); + + // General path: reconstruct each candidate and apply the metric. + return RankReconstructed(query, k); + } + + /// + public bool Remove(string id) + { + bool removed = _raw.Remove(id); + removed |= _codes.Remove(id); + return removed; + } + + /// + public void Clear() + { + _raw.Clear(); + _codes.Clear(); + } + + private List<(string Id, T Score)> RankRaw(Vector query, int k) + { + if (_raw.Count == 0) + return new List<(string Id, T Score)>(); + + var scores = _raw + .Select(kvp => (Id: kvp.Key, Score: _metric.Calculate(query, kvp.Value))) + .ToList(); + + var sorted = _metric.HigherIsBetter + ? scores.OrderByDescending(x => x.Score) + : scores.OrderBy(x => x.Score); + + return sorted.Take(Math.Min(k, scores.Count)).ToList(); + } + + private List<(string Id, T Score)> RankReconstructed(Vector query, int k) + { + var scores = _codes + .Select(kvp => (Id: kvp.Key, Score: _metric.Calculate(query, _quantizer.Decode(kvp.Value)))) + .ToList(); + + var sorted = _metric.HigherIsBetter + ? scores.OrderByDescending(x => x.Score) + : scores.OrderBy(x => x.Score); + + return sorted.Take(Math.Min(k, scores.Count)).ToList(); + } + + private List<(string Id, T Score)> RankAdc(ProductQuantizer pq, Vector query, int k) + { + var table = pq.BuildDistanceTable(query); + + var scores = _codes + .Select(kvp => (Id: kvp.Key, Distance: pq.ComputeAsymmetricDistance(table, kvp.Value))) + .OrderBy(x => x.Distance) + .Take(Math.Min(k, _codes.Count)) + .Select(x => (x.Id, Score: _numOps.FromDouble(x.Distance))) + .ToList(); + + return scores; + } + } +} diff --git a/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/BinaryQuantizer.cs b/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/BinaryQuantizer.cs new file mode 100644 index 0000000000..06048444f9 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/BinaryQuantizer.cs @@ -0,0 +1,177 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Helpers; +using AiDotNet.LinearAlgebra; + +namespace AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Quantization +{ + /// + /// Sign-based binary quantizer that packs each dimension into a single bit + /// (1 if the component is non-negative, 0 otherwise). Provides roughly 32x memory + /// reduction versus 32-bit floats (64x versus 64-bit doubles) and enables extremely + /// fast Hamming-distance ranking. + /// + /// + /// + /// Binary quantization is the backbone of "binary embeddings" used by Qdrant and + /// modern retrieval stacks for a cheap first-pass filter before an exact rerank. + /// Distances between binary codes are computed with the Hamming distance (the number + /// of differing bits), which maps to a popcount of the XOR of the two codes. + /// + /// + /// For Beginners: Instead of remembering the exact value of each dimension we + /// only remember whether it was positive or negative. That single yes/no fact takes + /// one bit, so 8 dimensions fit in one byte. Comparing two vectors then just counts + /// how many of these yes/no answers disagree. + /// + /// + /// The numeric type for vector operations. + public class BinaryQuantizer : IVectorQuantizer + { + private readonly INumericOperations _numOps; + private int _dimension; + private bool _trained; + + /// + public bool IsTrained => _trained; + + /// + public int Dimension => _dimension; + + /// + public int CodeLength => _trained ? ByteCount(_dimension) : 0; + + /// + /// Initializes a new instance of the class. + /// + public BinaryQuantizer() + { + _numOps = MathHelper.GetNumericOperations(); + } + + /// + /// Records the dimensionality of the vectors. Binary quantization is parameter-free + /// (sign based), so training only needs to observe the vector length. + /// + /// The training vectors; only the dimensionality is used. + public void Train(IEnumerable> vectors) + { + if (vectors == null) + throw new ArgumentNullException(nameof(vectors)); + + int dim = -1; + foreach (var vector in vectors) + { + if (vector == null) + throw new ArgumentException("Training set contains a null vector.", nameof(vectors)); + + if (dim < 0) + dim = vector.Length; + else if (vector.Length != dim) + throw new ArgumentException("All training vectors must have the same dimensionality.", nameof(vectors)); + } + + if (dim < 0) + throw new ArgumentException("Training set must contain at least one vector.", nameof(vectors)); + + _dimension = dim; + _trained = true; + } + + /// + public byte[] Encode(Vector vector) + { + if (vector == null) + throw new ArgumentNullException(nameof(vector)); + + var arr = vector.ToArray(); + + // Allow encoding to infer/lock the dimensionality if not explicitly trained. + if (!_trained) + { + _dimension = arr.Length; + _trained = true; + } + else if (arr.Length != _dimension) + { + throw new ArgumentException("Vector dimensionality does not match the trained dimensionality.", nameof(vector)); + } + + var code = new byte[ByteCount(_dimension)]; + for (int i = 0; i < _dimension; i++) + { + double v = Convert.ToDouble(arr[i]); + if (v >= 0.0) + { + code[i >> 3] |= (byte)(1 << (i & 7)); + } + } + + return code; + } + + /// + public Vector Decode(byte[] code) + { + if (code == null) + throw new ArgumentNullException(nameof(code)); + if (!_trained) + throw new InvalidOperationException("BinaryQuantizer must be trained (or have encoded at least once) before decoding."); + if (code.Length != ByteCount(_dimension)) + throw new ArgumentException("Code length does not match the trained dimensionality.", nameof(code)); + + var one = _numOps.FromDouble(1.0); + var negOne = _numOps.FromDouble(-1.0); + var values = new T[_dimension]; + for (int i = 0; i < _dimension; i++) + { + bool bit = (code[i >> 3] & (1 << (i & 7))) != 0; + values[i] = bit ? one : negOne; + } + + return new Vector(values); + } + + /// + /// Computes the Hamming distance (number of differing bits) between two binary codes. + /// + /// The first code. + /// The second code. + /// The number of bits that differ between the two codes. + public static int HammingDistance(byte[] a, byte[] b) + { + if (a == null) + throw new ArgumentNullException(nameof(a)); + if (b == null) + throw new ArgumentNullException(nameof(b)); + if (a.Length != b.Length) + throw new ArgumentException("Codes must have the same length.", nameof(b)); + + int distance = 0; + for (int i = 0; i < a.Length; i++) + { + distance += PopCount((byte)(a[i] ^ b[i])); + } + + return distance; + } + + private static int ByteCount(int dimension) + { + return (dimension + 7) / 8; + } + + private static int PopCount(byte value) + { + int v = value; + int count = 0; + while (v != 0) + { + v &= v - 1; // clear the lowest set bit + count++; + } + + return count; + } + } +} diff --git a/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/IVectorQuantizer.cs b/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/IVectorQuantizer.cs new file mode 100644 index 0000000000..27051b9ff5 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/IVectorQuantizer.cs @@ -0,0 +1,62 @@ +using System.Collections.Generic; +using AiDotNet.LinearAlgebra; + +namespace AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Quantization +{ + /// + /// Common interface for vector quantizers used to compress dense embeddings for + /// memory- and latency-efficient similarity search. + /// + /// + /// + /// Quantization trades a small amount of accuracy for large reductions in memory + /// footprint (and often faster distance computation). This is how libraries such as + /// FAISS, Qdrant and pgvector serve billions of vectors on commodity hardware. + /// + /// + /// For Beginners: A raw embedding stores every dimension as a full + /// floating-point number (4 or 8 bytes each). Quantization replaces those numbers + /// with much smaller codes (for example a single byte, or even a single bit, per + /// dimension) plus a small "dictionary" that lets you approximately reconstruct the + /// original values. You lose a little precision but save a lot of memory. + /// + /// + /// The numeric type for vector operations. + public interface IVectorQuantizer + { + /// + /// Gets a value indicating whether the quantizer has been trained and is ready to encode. + /// + bool IsTrained { get; } + + /// + /// Gets the dimensionality of the vectors this quantizer was trained on, or 0 if untrained. + /// + int Dimension { get; } + + /// + /// Gets the number of bytes produced by for a single vector, or 0 if untrained. + /// + int CodeLength { get; } + + /// + /// Trains the quantizer (learns any required statistics or codebooks) from a set of vectors. + /// + /// The training vectors. All must share the same dimensionality. + void Train(IEnumerable> vectors); + + /// + /// Encodes a vector into its compact byte representation. + /// + /// The vector to encode. + /// The quantized code. + byte[] Encode(Vector vector); + + /// + /// Approximately reconstructs a vector from its compact byte representation. + /// + /// The quantized code produced by . + /// The reconstructed (approximate) vector. + Vector Decode(byte[] code); + } +} diff --git a/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/ProductQuantizer.cs b/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/ProductQuantizer.cs new file mode 100644 index 0000000000..46fd4b0c6b --- /dev/null +++ b/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/ProductQuantizer.cs @@ -0,0 +1,403 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Helpers; +using AiDotNet.LinearAlgebra; + +namespace AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Quantization +{ + /// + /// Product Quantizer (PQ) that splits each vector into M contiguous subspaces, + /// learns a small k-means codebook (default 256 centroids) per subspace, and encodes a + /// vector as one byte per subspace. Supports Asymmetric Distance Computation (ADC) for + /// fast, memory-light nearest-neighbor ranking. + /// + /// + /// + /// PQ is the headline compression technique behind FAISS IndexPQ/IndexIVFPQ. + /// A D-dimensional float vector (4*D bytes) is reduced to just M bytes. + /// With the common choice of 256 centroids per subspace, each subspace index fits in a + /// single byte, so an embedding of dimension 768 encoded with M=96 needs only 96 bytes + /// versus 3072 bytes raw (32x reduction), or 768 bytes as SQ8. + /// + /// + /// Asymmetric Distance Computation (ADC): at query time we do NOT quantize the + /// query. Instead, for each subspace we precompute the squared L2 distance from the raw + /// query subvector to every centroid of that subspace, building an + /// M x Ksub lookup table. The (approximate) squared distance between the query and + /// any stored code is then just the sum of M table look-ups. This equals the exact + /// squared L2 distance between the raw query and the reconstructed database vector, which + /// is why ADC is both fast and accurate. + /// + /// + /// For Beginners: We chop each vector into M pieces. For each piece position we + /// learn 256 "typical pieces" (centroids). A vector is then described by which typical + /// piece is closest in each position - just M small numbers. To compare a query we + /// measure the query against those 256 typical pieces once, then adding up the right + /// numbers gives the distance to any stored vector almost for free. + /// + /// + /// The numeric type for vector operations. + public class ProductQuantizer : IVectorQuantizer + { + private readonly INumericOperations _numOps; + private readonly int _subspaceCount; // M + private readonly int _centroidsPerSubspace; // Ksub (<= 256) + private readonly int _maxIterations; + private readonly int _seed; + + private int _dimension; + private int _subDimension; // Dsub = D / M + // _codebooks[m][c] is the centroid c of subspace m, length _subDimension. + private double[][][]? _codebooks; + + /// + /// Initializes a new instance of the class. + /// + /// Number of subspaces (M). The vector dimension must be divisible by this. Default: 8. + /// Number of k-means centroids per subspace (Ksub, 1-256). Default: 256. + /// Maximum Lloyd's k-means iterations per subspace. Default: 25. + /// Random seed for deterministic k-means initialization. Default: 42. + public ProductQuantizer(int subspaceCount = 8, int centroidsPerSubspace = 256, int maxIterations = 25, int seed = 42) + { + if (subspaceCount <= 0) + throw new ArgumentException("Subspace count (M) must be positive.", nameof(subspaceCount)); + if (centroidsPerSubspace <= 0 || centroidsPerSubspace > 256) + throw new ArgumentException("Centroids per subspace (Ksub) must be between 1 and 256.", nameof(centroidsPerSubspace)); + if (maxIterations <= 0) + throw new ArgumentException("Max iterations must be positive.", nameof(maxIterations)); + + _numOps = MathHelper.GetNumericOperations(); + _subspaceCount = subspaceCount; + _centroidsPerSubspace = centroidsPerSubspace; + _maxIterations = maxIterations; + _seed = seed; + } + + /// + public bool IsTrained => _codebooks != null; + + /// + public int Dimension => _dimension; + + /// + public int CodeLength => IsTrained ? _subspaceCount : 0; + + /// + /// Gets the number of subspaces (M). + /// + public int SubspaceCount => _subspaceCount; + + /// + /// Gets the number of centroids learned per subspace (Ksub). + /// + public int CentroidsPerSubspace => _centroidsPerSubspace; + + /// + /// Gets the dimensionality of each subspace (Dsub = D / M), or 0 if untrained. + /// + public int SubDimension => _subDimension; + + /// + /// Gets the learned codebooks with shape [M][Ksub][Dsub], or null if untrained. + /// + public double[][][]? Codebooks => _codebooks; + + /// + public void Train(IEnumerable> vectors) + { + if (vectors == null) + throw new ArgumentNullException(nameof(vectors)); + + // Materialize training data as double arrays. + var data = new List(); + int dim = -1; + foreach (var vector in vectors) + { + if (vector == null) + throw new ArgumentException("Training set contains a null vector.", nameof(vectors)); + + var arr = vector.ToArray(); + if (dim < 0) + dim = arr.Length; + else if (arr.Length != dim) + throw new ArgumentException("All training vectors must have the same dimensionality.", nameof(vectors)); + + var row = new double[arr.Length]; + for (int i = 0; i < arr.Length; i++) + row[i] = Convert.ToDouble(arr[i]); + data.Add(row); + } + + if (data.Count == 0) + throw new ArgumentException("Training set must contain at least one vector.", nameof(vectors)); + if (dim % _subspaceCount != 0) + throw new ArgumentException( + "Vector dimension (" + dim + ") must be divisible by the subspace count (" + _subspaceCount + ").", + nameof(vectors)); + + _dimension = dim; + _subDimension = dim / _subspaceCount; + + var codebooks = new double[_subspaceCount][][]; + for (int m = 0; m < _subspaceCount; m++) + { + int offset = m * _subDimension; + var subvectors = new double[data.Count][]; + for (int n = 0; n < data.Count; n++) + { + var sub = new double[_subDimension]; + Array.Copy(data[n], offset, sub, 0, _subDimension); + subvectors[n] = sub; + } + + // Deterministic per-subspace seed so results are reproducible and subspaces differ. + codebooks[m] = RunKMeans(subvectors, _centroidsPerSubspace, _subDimension, _maxIterations, _seed + m); + } + + _codebooks = codebooks; + } + + /// + public byte[] Encode(Vector vector) + { + if (vector == null) + throw new ArgumentNullException(nameof(vector)); + if (!IsTrained) + throw new InvalidOperationException("ProductQuantizer must be trained before encoding."); + + var arr = vector.ToArray(); + if (arr.Length != _dimension) + throw new ArgumentException("Vector dimensionality does not match the trained dimensionality.", nameof(vector)); + + var full = new double[_dimension]; + for (int i = 0; i < _dimension; i++) + full[i] = Convert.ToDouble(arr[i]); + + var code = new byte[_subspaceCount]; + var sub = new double[_subDimension]; + for (int m = 0; m < _subspaceCount; m++) + { + int offset = m * _subDimension; + Array.Copy(full, offset, sub, 0, _subDimension); + + var centroids = _codebooks![m]; + int best = 0; + double bestDist = double.MaxValue; + for (int c = 0; c < centroids.Length; c++) + { + double d = SquaredDistance(sub, centroids[c]); + if (d < bestDist) + { + bestDist = d; + best = c; + } + } + + code[m] = (byte)best; + } + + return code; + } + + /// + public Vector Decode(byte[] code) + { + if (code == null) + throw new ArgumentNullException(nameof(code)); + if (!IsTrained) + throw new InvalidOperationException("ProductQuantizer must be trained before decoding."); + if (code.Length != _subspaceCount) + throw new ArgumentException("Code length does not match the subspace count.", nameof(code)); + + var values = new T[_dimension]; + for (int m = 0; m < _subspaceCount; m++) + { + int offset = m * _subDimension; + var centroid = _codebooks![m][code[m]]; + for (int i = 0; i < _subDimension; i++) + values[offset + i] = _numOps.FromDouble(centroid[i]); + } + + return new Vector(values); + } + + /// + /// Builds the Asymmetric Distance Computation (ADC) lookup table for a query. Entry + /// [m][c] holds the squared L2 distance between the query's subvector m and centroid c + /// of subspace m. + /// + /// The raw (unquantized) query vector. + /// A table with shape [M][Ksub]. + public double[][] BuildDistanceTable(Vector query) + { + if (query == null) + throw new ArgumentNullException(nameof(query)); + if (!IsTrained) + throw new InvalidOperationException("ProductQuantizer must be trained before building a distance table."); + + var arr = query.ToArray(); + if (arr.Length != _dimension) + throw new ArgumentException("Query dimensionality does not match the trained dimensionality.", nameof(query)); + + var full = new double[_dimension]; + for (int i = 0; i < _dimension; i++) + full[i] = Convert.ToDouble(arr[i]); + + var table = new double[_subspaceCount][]; + var sub = new double[_subDimension]; + for (int m = 0; m < _subspaceCount; m++) + { + int offset = m * _subDimension; + Array.Copy(full, offset, sub, 0, _subDimension); + + var centroids = _codebooks![m]; + var row = new double[centroids.Length]; + for (int c = 0; c < centroids.Length; c++) + row[c] = SquaredDistance(sub, centroids[c]); + table[m] = row; + } + + return table; + } + + /// + /// Computes the approximate squared L2 distance between a query (via its ADC table) + /// and an encoded vector, as the sum of one table look-up per subspace. + /// + /// A table produced by . + /// A code produced by . + /// The approximate squared L2 distance. + public double ComputeAsymmetricDistance(double[][] distanceTable, byte[] code) + { + if (distanceTable == null) + throw new ArgumentNullException(nameof(distanceTable)); + if (code == null) + throw new ArgumentNullException(nameof(code)); + if (code.Length != _subspaceCount) + throw new ArgumentException("Code length does not match the subspace count.", nameof(code)); + + double distance = 0.0; + for (int m = 0; m < _subspaceCount; m++) + distance += distanceTable[m][code[m]]; + + return distance; + } + + /// + /// Runs Lloyd's k-means on the supplied subvectors and returns a codebook of exactly + /// centroids (shape is always [k][dim] so callers can rely on it). + /// + private static double[][] RunKMeans(double[][] points, int k, int dim, int maxIterations, int seed) + { + int n = points.Length; + var centroids = new double[k][]; + + // Deterministic initialization: sample distinct points where possible. + var random = RandomHelper.CreateSeededRandom(seed); + var chosen = new HashSet(); + int effectiveK = Math.Min(k, n); + for (int c = 0; c < effectiveK; c++) + { + int idx; + int guard = 0; + do + { + idx = random.Next(n); + guard++; + } + while (chosen.Contains(idx) && guard < n * 4); + + chosen.Add(idx); + centroids[c] = (double[])points[idx].Clone(); + } + + // If we have fewer points than k, pad remaining centroids by duplicating existing + // ones so the codebook always has the expected shape. + for (int c = effectiveK; c < k; c++) + centroids[c] = (double[])centroids[c % Math.Max(1, effectiveK)].Clone(); + + if (n == 0) + { + for (int c = 0; c < k; c++) + centroids[c] = new double[dim]; + return centroids; + } + + var assignments = new int[n]; + for (int i = 0; i < n; i++) + assignments[i] = -1; + + for (int iter = 0; iter < maxIterations; iter++) + { + bool changed = false; + + // Assignment step. + for (int i = 0; i < n; i++) + { + int best = 0; + double bestDist = double.MaxValue; + for (int c = 0; c < k; c++) + { + double d = SquaredDistance(points[i], centroids[c]); + if (d < bestDist) + { + bestDist = d; + best = c; + } + } + + if (assignments[i] != best) + { + assignments[i] = best; + changed = true; + } + } + + // Update step. + var sums = new double[k][]; + var counts = new int[k]; + for (int c = 0; c < k; c++) + sums[c] = new double[dim]; + + for (int i = 0; i < n; i++) + { + int c = assignments[i]; + var p = points[i]; + var s = sums[c]; + for (int d = 0; d < dim; d++) + s[d] += p[d]; + counts[c]++; + } + + for (int c = 0; c < k; c++) + { + if (counts[c] == 0) + continue; // Keep the previous centroid for empty clusters. + + var s = sums[c]; + var newCentroid = new double[dim]; + for (int d = 0; d < dim; d++) + newCentroid[d] = s[d] / counts[c]; + centroids[c] = newCentroid; + } + + if (!changed && iter > 0) + break; // Converged. + } + + return centroids; + } + + private static double SquaredDistance(double[] a, double[] b) + { + double sum = 0.0; + for (int i = 0; i < a.Length; i++) + { + double diff = a[i] - b[i]; + sum += diff * diff; + } + + return sum; + } + } +} diff --git a/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/ScalarQuantizer.cs b/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/ScalarQuantizer.cs new file mode 100644 index 0000000000..738b57a00d --- /dev/null +++ b/src/RetrievalAugmentedGeneration/VectorSearch/Quantization/ScalarQuantizer.cs @@ -0,0 +1,168 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Helpers; +using AiDotNet.LinearAlgebra; + +namespace AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Quantization +{ + /// + /// Per-dimension scalar quantizer that maps each floating-point component to a single + /// unsigned byte (uint8) using learned min/max ranges. Provides roughly 4x memory + /// reduction versus 32-bit floats (8x versus 64-bit doubles). + /// + /// + /// + /// Each dimension is linearly mapped from its observed [min, max] range onto the + /// integer range [0, 255]. Reconstruction reverses the mapping using the bin center. + /// This is the same idea used by FAISS ScalarQuantizer (SQ8) and pgvector's + /// halfvec/quantization paths. + /// + /// + /// For Beginners: Imagine each dimension has values between, say, -3 and +5. + /// We slice that range into 256 equal buckets and store just the bucket number + /// (0-255, one byte) instead of the full number. To decode we return the middle of + /// the bucket. Close, but far smaller. + /// + /// + /// The numeric type for vector operations. + public class ScalarQuantizer : IVectorQuantizer + { + private const int Levels = 256; + + private readonly INumericOperations _numOps; + private double[]? _min; + private double[]? _scale; // (max - min) / 255 per dimension + private int _dimension; + + /// + public bool IsTrained => _min != null; + + /// + public int Dimension => _dimension; + + /// + public int CodeLength => IsTrained ? _dimension : 0; + + /// + /// Initializes a new instance of the class. + /// + public ScalarQuantizer() + { + _numOps = MathHelper.GetNumericOperations(); + } + + /// + public void Train(IEnumerable> vectors) + { + if (vectors == null) + throw new ArgumentNullException(nameof(vectors)); + + double[]? min = null; + double[]? max = null; + int dim = 0; + int count = 0; + + foreach (var vector in vectors) + { + if (vector == null) + throw new ArgumentException("Training set contains a null vector.", nameof(vectors)); + + var arr = vector.ToArray(); + if (count == 0) + { + dim = arr.Length; + min = new double[dim]; + max = new double[dim]; + for (int i = 0; i < dim; i++) + { + double v = Convert.ToDouble(arr[i]); + min[i] = v; + max[i] = v; + } + } + else + { + if (arr.Length != dim) + throw new ArgumentException("All training vectors must have the same dimensionality.", nameof(vectors)); + + for (int i = 0; i < dim; i++) + { + double v = Convert.ToDouble(arr[i]); + if (v < min![i]) min[i] = v; + if (v > max![i]) max[i] = v; + } + } + + count++; + } + + if (count == 0) + throw new ArgumentException("Training set must contain at least one vector.", nameof(vectors)); + + _dimension = dim; + _min = min; + _scale = new double[dim]; + for (int i = 0; i < dim; i++) + { + double range = max![i] - min![i]; + // Guard against zero-range (constant) dimensions to avoid division by zero. + _scale[i] = range > 0.0 ? range / (Levels - 1) : 0.0; + } + } + + /// + public byte[] Encode(Vector vector) + { + if (vector == null) + throw new ArgumentNullException(nameof(vector)); + if (!IsTrained) + throw new InvalidOperationException("ScalarQuantizer must be trained before encoding."); + + var arr = vector.ToArray(); + if (arr.Length != _dimension) + throw new ArgumentException("Vector dimensionality does not match the trained dimensionality.", nameof(vector)); + + var code = new byte[_dimension]; + for (int i = 0; i < _dimension; i++) + { + double v = Convert.ToDouble(arr[i]); + int level; + if (_scale![i] <= 0.0) + { + level = 0; + } + else + { + double normalized = (v - _min![i]) / _scale[i]; + level = (int)Math.Round(normalized); + if (level < 0) level = 0; + if (level > Levels - 1) level = Levels - 1; + } + + code[i] = (byte)level; + } + + return code; + } + + /// + public Vector Decode(byte[] code) + { + if (code == null) + throw new ArgumentNullException(nameof(code)); + if (!IsTrained) + throw new InvalidOperationException("ScalarQuantizer must be trained before decoding."); + if (code.Length != _dimension) + throw new ArgumentException("Code length does not match the trained dimensionality.", nameof(code)); + + var values = new T[_dimension]; + for (int i = 0; i < _dimension; i++) + { + double reconstructed = _min![i] + code[i] * _scale![i]; + values[i] = _numOps.FromDouble(reconstructed); + } + + return new Vector(values); + } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/QuantizedFlatIndexTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/QuantizedFlatIndexTests.cs new file mode 100644 index 0000000000..40c93ed6eb --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/QuantizedFlatIndexTests.cs @@ -0,0 +1,177 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Helpers; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Indexes; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Metrics; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Quantization; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.VectorSearch.Indexes +{ + public class QuantizedFlatIndexTests + { + private const int Dim = 8; + + private static readonly double[][] Centers = + { + new double[] { 10, 10, 10, 10, 0, 0, 0, 0 }, + new double[] { 0, 0, 0, 0, 10, 10, 10, 10 }, + new double[] {-10,-10,-10,-10, 10, 10, 10, 10 }, + new double[] { 10, 10, 10, 10, 10, 10, 10, 10 }, + }; + + private static Dictionary> BuildClusteredData(int perCluster, int seed) + { + var rng = RandomHelper.CreateSeededRandom(seed); + var data = new Dictionary>(); + int id = 0; + foreach (var center in Centers) + { + for (int n = 0; n < perCluster; n++) + { + var arr = new double[Dim]; + for (int i = 0; i < Dim; i++) + arr[i] = center[i] + (rng.NextDouble() - 0.5); + data[$"v{id++}"] = new Vector(arr); + } + } + + return data; + } + + [Fact] + public void Constructor_WithNullArgs_Throws() + { + var metric = new EuclideanDistanceMetric(); + var quantizer = new ScalarQuantizer(); + Assert.Throws(() => new QuantizedFlatIndex(null!, quantizer)); + Assert.Throws(() => new QuantizedFlatIndex(metric, null!)); + } + + [Fact] + public void UntrainedIndex_FallsBackToExactSearch() + { + var metric = new EuclideanDistanceMetric(); + var index = new QuantizedFlatIndex(metric, new ProductQuantizer(subspaceCount: 2)); + + var data = BuildClusteredData(5, 1); + index.AddBatch(data); + + Assert.False(index.IsTrained); + Assert.Equal(20, index.Count); + + // Query exactly equal to an existing vector -> that vector is the top-1 (distance 0). + var target = data["v3"]; + var results = index.Search(target, 1); + Assert.Single(results); + Assert.Equal("v3", results[0].Id); + } + + [Fact] + public void Trained_ProductQuantizer_MatchesFlatIndexTopK() + { + var data = BuildClusteredData(15, 2); + + var flat = new FlatIndex(new EuclideanDistanceMetric()); + flat.AddBatch(data); + + var quantized = new QuantizedFlatIndex( + new EuclideanDistanceMetric(), + new ProductQuantizer(subspaceCount: 2, centroidsPerSubspace: 64, maxIterations: 25, seed: 42)); + quantized.AddBatch(data); + quantized.Train(); + + Assert.True(quantized.IsTrained); + Assert.Equal(data.Count, quantized.Count); + + var rng = RandomHelper.CreateSeededRandom(555); + for (int trial = 0; trial < 10; trial++) + { + var center = Centers[trial % Centers.Length]; + var q = new double[Dim]; + for (int i = 0; i < Dim; i++) + q[i] = center[i] + (rng.NextDouble() - 0.5); + var query = new Vector(q); + + var flatTop = flat.Search(query, 3).Select(r => r.Id).ToList(); + var quantTop = quantized.Search(query, 3).Select(r => r.Id).ToList(); + + // Top-1 must agree on well-separated clusters. + Assert.Equal(flatTop[0], quantTop[0]); + } + } + + [Fact] + public void Trained_ScalarQuantizer_ReturnsCorrectTopKViaReconstruction() + { + var data = BuildClusteredData(15, 4); + + var flat = new FlatIndex(new CosineSimilarityMetric()); + flat.AddBatch(data); + + var quantized = new QuantizedFlatIndex( + new CosineSimilarityMetric(), + new ScalarQuantizer()); + quantized.AddBatch(data); + quantized.Train(); + + var query = data["v0"]; + var flatTop = flat.Search(query, 3).Select(r => r.Id).ToList(); + var quantTop = quantized.Search(query, 3).Select(r => r.Id).ToList(); + + Assert.Equal(flatTop[0], quantTop[0]); + } + + [Fact] + public void AddAfterTraining_EncodesImmediately() + { + var data = BuildClusteredData(10, 6); + var index = new QuantizedFlatIndex( + new EuclideanDistanceMetric(), + new ScalarQuantizer()); + index.AddBatch(data); + index.Train(); + + int before = index.Count; + var extra = new double[Dim]; + Array.Copy(Centers[0], extra, Dim); + index.Add("extra", new Vector(extra)); + + Assert.Equal(before + 1, index.Count); + var results = index.Search(new Vector(extra), 1); + Assert.Equal("extra", results[0].Id); + } + + [Fact] + public void RemoveAndClear_Work() + { + var data = BuildClusteredData(5, 8); + var index = new QuantizedFlatIndex( + new EuclideanDistanceMetric(), + new ScalarQuantizer()); + index.AddBatch(data); + + Assert.True(index.Remove("v0")); + Assert.False(index.Remove("does-not-exist")); + Assert.Equal(19, index.Count); + + index.Train(); + Assert.True(index.Remove("v1")); + Assert.Equal(18, index.Count); + + index.Clear(); + Assert.Equal(0, index.Count); + } + + [Fact] + public void Train_OnEmptyIndex_Throws() + { + var index = new QuantizedFlatIndex( + new EuclideanDistanceMetric(), + new ScalarQuantizer()); + Assert.Throws(() => index.Train()); + } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/BinaryQuantizerTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/BinaryQuantizerTests.cs new file mode 100644 index 0000000000..4228799958 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/BinaryQuantizerTests.cs @@ -0,0 +1,73 @@ +using System; +using System.Collections.Generic; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Quantization; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.VectorSearch.Quantization +{ + public class BinaryQuantizerTests + { + [Fact] + public void Encode_PacksOneBitPerDimension() + { + int dim = 20; // not a multiple of 8 -> ceil(20/8) = 3 bytes + var q = new BinaryQuantizer(); + q.Train(new List> { new Vector(new double[dim]) }); + + var code = q.Encode(new Vector(new double[dim])); + + Assert.Equal(3, code.Length); + Assert.Equal(3, q.CodeLength); + Assert.Equal(dim, q.Dimension); + } + + [Fact] + public void EncodeDecode_RoundTrip_PreservesSign() + { + var values = new double[] { 1.5, -2.0, 0.0, -0.1, 3.3, -7.0, 8.8, -0.5 }; + var q = new BinaryQuantizer(); + var code = q.Encode(new Vector(values)); + var decoded = q.Decode(code).ToArray(); + + for (int i = 0; i < values.Length; i++) + { + // Non-negative (>= 0) maps to +1, negative maps to -1. + double expectedSign = values[i] >= 0.0 ? 1.0 : -1.0; + Assert.Equal(expectedSign, decoded[i], 9); + } + } + + [Fact] + public void Memory_IsAtLeast32xSmallerThanRawFloat() + { + int dim = 128; + var q = new BinaryQuantizer(); + var code = q.Encode(new Vector(new double[dim])); + + int rawFloatBytes = dim * sizeof(float); // 4 bytes per component + // 1 bit per dim -> ceil(dim/8) bytes; must be >= 32x smaller than float32. + Assert.True(code.Length * 32 <= rawFloatBytes, $"codeBytes={code.Length}, rawFloatBytes={rawFloatBytes}"); + } + + [Fact] + public void HammingDistance_CountsDifferingBits() + { + var q = new BinaryQuantizer(); + // signs: + - + - + - + - + var a = q.Encode(new Vector(new double[] { 1, -1, 1, -1, 1, -1, 1, -1 })); + // signs: + + + + - - - - (differs in positions 1,3,4,6 -> 4 bits) + var b = q.Encode(new Vector(new double[] { 1, 1, 1, 1, -1, -1, -1, -1 })); + + Assert.Equal(4, BinaryQuantizer.HammingDistance(a, b)); + Assert.Equal(0, BinaryQuantizer.HammingDistance(a, a)); + } + + [Fact] + public void HammingDistance_WithMismatchedLengths_Throws() + { + Assert.Throws(() => + BinaryQuantizer.HammingDistance(new byte[2], new byte[3])); + } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/ProductQuantizerTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/ProductQuantizerTests.cs new file mode 100644 index 0000000000..a2a2e41a77 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/ProductQuantizerTests.cs @@ -0,0 +1,188 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Helpers; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Quantization; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.VectorSearch.Quantization +{ + public class ProductQuantizerTests + { + private const int Dim = 8; + + // Three well-separated cluster centers in 8-D space. + private static readonly double[][] Centers = + { + new double[] { 10, 10, 10, 10, 0, 0, 0, 0 }, + new double[] { 0, 0, 0, 0, 10, 10, 10, 10 }, + new double[] {-10,-10,-10,-10,-10,-10,-10,-10 }, + }; + + private static List<(double[] Raw, Vector Vector, int Cluster)> BuildClusteredData(int perCluster, int seed) + { + var rng = RandomHelper.CreateSeededRandom(seed); + var data = new List<(double[], Vector, int)>(); + for (int c = 0; c < Centers.Length; c++) + { + var center = Centers[c]; + for (int n = 0; n < perCluster; n++) + { + var arr = new double[Dim]; + for (int i = 0; i < Dim; i++) + arr[i] = center[i] + (rng.NextDouble() - 0.5); // +/- 0.5 noise + data.Add((arr, new Vector(arr), c)); + } + } + + return data; + } + + private static double SquaredL2(double[] a, double[] b) + { + double s = 0.0; + for (int i = 0; i < a.Length; i++) + { + double d = a[i] - b[i]; + s += d * d; + } + + return s; + } + + [Fact] + public void Train_LearnsCodebooksWithExpectedShape() + { + int m = 2, ksub = 16; + var data = BuildClusteredData(30, 7).Select(x => x.Vector); + var pq = new ProductQuantizer(subspaceCount: m, centroidsPerSubspace: ksub, maxIterations: 20, seed: 42); + pq.Train(data); + + Assert.True(pq.IsTrained); + Assert.Equal(m, pq.SubspaceCount); + Assert.Equal(ksub, pq.CentroidsPerSubspace); + Assert.Equal(Dim / m, pq.SubDimension); + + var books = pq.Codebooks!; + Assert.Equal(m, books.Length); + for (int mm = 0; mm < m; mm++) + { + Assert.Equal(ksub, books[mm].Length); + for (int c = 0; c < ksub; c++) + Assert.Equal(Dim / m, books[mm][c].Length); + } + } + + [Fact] + public void Encode_ProducesOneBytePerSubspace() + { + int m = 4; + var data = BuildClusteredData(30, 7).Select(x => x.Vector).ToList(); + var pq = new ProductQuantizer(subspaceCount: m, centroidsPerSubspace: 32, maxIterations: 20, seed: 42); + pq.Train(data); + + var code = pq.Encode(data[0]); + Assert.Equal(m, code.Length); + Assert.Equal(m, pq.CodeLength); + } + + [Fact] + public void Train_WithDimensionNotDivisibleByM_Throws() + { + var data = new List> { new Vector(new double[] { 1, 2, 3 }) }; + var pq = new ProductQuantizer(subspaceCount: 2); + Assert.Throws(() => pq.Train(data)); + } + + [Fact] + public void Adc_NearestNeighbor_MatchesBruteForceTop1() + { + var data = BuildClusteredData(30, 11); + var pq = new ProductQuantizer(subspaceCount: 2, centroidsPerSubspace: 64, maxIterations: 25, seed: 42); + pq.Train(data.Select(x => x.Vector)); + + // Encode the whole database. + var codes = data.Select(x => pq.Encode(x.Vector)).ToList(); + + var rng = RandomHelper.CreateSeededRandom(99); + for (int trial = 0; trial < 15; trial++) + { + // Query near a random cluster center. + int expectedCluster = trial % Centers.Length; + var center = Centers[expectedCluster]; + var q = new double[Dim]; + for (int i = 0; i < Dim; i++) + q[i] = center[i] + (rng.NextDouble() - 0.5); + var query = new Vector(q); + + // Brute-force exact nearest neighbor. + int bruteBest = 0; + double bruteDist = double.MaxValue; + for (int j = 0; j < data.Count; j++) + { + double d = SquaredL2(q, data[j].Raw); + if (d < bruteDist) + { + bruteDist = d; + bruteBest = j; + } + } + + // ADC nearest neighbor. + var table = pq.BuildDistanceTable(query); + int adcBest = 0; + double adcDist = double.MaxValue; + for (int j = 0; j < codes.Count; j++) + { + double d = pq.ComputeAsymmetricDistance(table, codes[j]); + if (d < adcDist) + { + adcDist = d; + adcBest = j; + } + } + + // On separable clusters, ADC (a lossy approximation) must recover the same + // cluster as the exact brute-force nearest neighbor. Distinguishing the single + // closest point *within* a tight cluster is below quantization resolution and + // is not expected of any product quantizer. + Assert.Equal(expectedCluster, data[bruteBest].Cluster); + Assert.Equal(data[bruteBest].Cluster, data[adcBest].Cluster); + } + } + + [Fact] + public void Adc_EqualsExactDistanceToReconstruction() + { + var data = BuildClusteredData(20, 5); + var pq = new ProductQuantizer(subspaceCount: 2, centroidsPerSubspace: 32, maxIterations: 20, seed: 42); + pq.Train(data.Select(x => x.Vector)); + + var query = new Vector(new double[] { 9.7, 10.1, 9.9, 10.2, 0.1, -0.2, 0.05, 0.0 }); + var table = pq.BuildDistanceTable(query); + + foreach (var item in data.Take(10)) + { + var code = pq.Encode(item.Vector); + double adc = pq.ComputeAsymmetricDistance(table, code); + double exact = SquaredL2(query.ToArray(), pq.Decode(code).ToArray()); + Assert.Equal(exact, adc, 6); + } + } + + [Fact] + public void IsDeterministic_AcrossRuns() + { + var data = BuildClusteredData(25, 3).Select(x => x.Vector).ToList(); + + var pq1 = new ProductQuantizer(subspaceCount: 2, centroidsPerSubspace: 32, seed: 7); + var pq2 = new ProductQuantizer(subspaceCount: 2, centroidsPerSubspace: 32, seed: 7); + pq1.Train(data); + pq2.Train(data); + + foreach (var v in data) + Assert.Equal(pq1.Encode(v), pq2.Encode(v)); + } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/ScalarQuantizerTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/ScalarQuantizerTests.cs new file mode 100644 index 0000000000..3f65411acb --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Quantization/ScalarQuantizerTests.cs @@ -0,0 +1,112 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Helpers; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Quantization; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.VectorSearch.Quantization +{ + public class ScalarQuantizerTests + { + private static List> BuildTrainingSet(int count, int dim, double low, double high, int seed) + { + var rng = RandomHelper.CreateSeededRandom(seed); + var result = new List>(); + for (int n = 0; n < count; n++) + { + var arr = new double[dim]; + for (int i = 0; i < dim; i++) + arr[i] = low + rng.NextDouble() * (high - low); + result.Add(new Vector(arr)); + } + + return result; + } + + [Fact] + public void Train_WithEmptySet_Throws() + { + var q = new ScalarQuantizer(); + Assert.Throws(() => q.Train(new List>())); + } + + [Fact] + public void Encode_BeforeTraining_Throws() + { + var q = new ScalarQuantizer(); + Assert.Throws(() => q.Encode(new Vector(new double[] { 1, 2, 3 }))); + } + + [Fact] + public void Encode_ProducesOneBytePerDimension() + { + int dim = 16; + var data = BuildTrainingSet(200, dim, -5.0, 5.0, 1); + var q = new ScalarQuantizer(); + q.Train(data); + + var code = q.Encode(data[0]); + + Assert.Equal(dim, code.Length); + Assert.Equal(dim, q.CodeLength); + Assert.Equal(dim, q.Dimension); + } + + [Fact] + public void EncodeDecode_RoundTrip_WithinBucketError() + { + int dim = 32; + double low = 0.0, high = 10.0; + var data = BuildTrainingSet(500, dim, low, high, 2); + var q = new ScalarQuantizer(); + q.Train(data); + + double bucketWidth = (high - low) / 255.0; + + foreach (var vector in data.Take(20)) + { + var decoded = q.Decode(q.Encode(vector)); + var original = vector.ToArray(); + var restored = decoded.ToArray(); + for (int i = 0; i < dim; i++) + { + // Reconstruction error must be within one quantization bucket. + Assert.True(Math.Abs(original[i] - restored[i]) <= bucketWidth + 1e-9, + $"dim {i}: |{original[i]} - {restored[i]}| exceeded bucket width {bucketWidth}"); + } + } + } + + [Fact] + public void Memory_IsAtLeastFourXSmallerThanRaw() + { + int dim = 64; + var data = BuildTrainingSet(100, dim, -1.0, 1.0, 3); + var q = new ScalarQuantizer(); + q.Train(data); + + int rawBytes = dim * sizeof(double); // 8 bytes per component + int codeBytes = q.Encode(data[0]).Length; + + // uint8 codes are >= 4x smaller than 32-bit floats (>= 8x vs doubles). + Assert.True(codeBytes * 4 <= rawBytes, $"codeBytes={codeBytes}, rawBytes={rawBytes}"); + } + + [Fact] + public void ConstantDimension_DoesNotDivideByZero() + { + // All vectors share the same value in dim 0 (zero range). + var data = new List>(); + for (int n = 0; n < 10; n++) + data.Add(new Vector(new double[] { 7.0, n * 0.1 })); + + var q = new ScalarQuantizer(); + q.Train(data); + + var decoded = q.Decode(q.Encode(data[5])).ToArray(); + Assert.Equal(7.0, decoded[0], 6); + } + } +} From a02ebbf77072cde588fb815e0c5dea0b0ea60d58 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 16:44:36 -0400 Subject: [PATCH 12/25] feat(rag): content-hash embedding cache (CacheBackedEmbeddings) Add a decorator that caches embeddings by SHA-256 content hash so repeated texts are not re-embedded, matching LangChain CacheBackedEmbeddings and LlamaIndex ingestion cache. - ContentHash: stable, cross-process SHA-256 hex helper (reusable for dedup) - IEmbeddingCache + thread-safe InMemoryEmbeddingCache (ConcurrentDictionary, optional max-size / approximate-LRU eviction) - CacheBackedEmbeddings : IEmbeddingModel decorator; keys namespaced by inner model identity + dimension; batch paths dedup within/across calls, invoke inner only for distinct misses, and preserve row ordering - net471/net8.0-safe; 18 unit tests with a counting fake inner model Co-Authored-By: Claude Opus 4.8 (1M context) --- .../Embeddings/CacheBackedEmbeddings.cs | 293 ++++++++++++++ .../Embeddings/ContentHash.cs | 67 ++++ .../Embeddings/IEmbeddingCache.cs | 48 +++ .../Embeddings/InMemoryEmbeddingCache.cs | 148 +++++++ .../Embeddings/CacheBackedEmbeddingsTests.cs | 378 ++++++++++++++++++ 5 files changed, 934 insertions(+) create mode 100644 src/RetrievalAugmentedGeneration/Embeddings/CacheBackedEmbeddings.cs create mode 100644 src/RetrievalAugmentedGeneration/Embeddings/ContentHash.cs create mode 100644 src/RetrievalAugmentedGeneration/Embeddings/IEmbeddingCache.cs create mode 100644 src/RetrievalAugmentedGeneration/Embeddings/InMemoryEmbeddingCache.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/CacheBackedEmbeddingsTests.cs diff --git a/src/RetrievalAugmentedGeneration/Embeddings/CacheBackedEmbeddings.cs b/src/RetrievalAugmentedGeneration/Embeddings/CacheBackedEmbeddings.cs new file mode 100644 index 0000000000..07e38a9154 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Embeddings/CacheBackedEmbeddings.cs @@ -0,0 +1,293 @@ +using System.Threading.Tasks; +using AiDotNet.Interfaces; +using AiDotNet.LinearAlgebra; + +namespace AiDotNet.RetrievalAugmentedGeneration.Embeddings; + +/// +/// An decorator that caches embeddings by content hash so that +/// repeated texts are not re-embedded, mirroring LangChain's CacheBackedEmbeddings and +/// LlamaIndex's ingestion cache. +/// +/// +/// +/// Each text is hashed with SHA-256 (see ) and the resulting key is namespaced +/// by the wrapped model's identity and embedding dimension, so different models never collide in a shared +/// cache. On a cache hit the stored vector is returned; on a miss the inner model is invoked and the result +/// is stored. Batch operations de-duplicate texts both within a single call and across calls, invoke the +/// inner model only for the distinct misses, and preserve the input ordering of the returned matrix rows. +/// +/// For Beginners: This wraps a "real" embedding model and remembers its answers. +/// +/// - The first time you embed a piece of text, it calls the wrapped model and saves the result. +/// - The next time you embed the same text, it returns the saved result instead of calling the model again. +/// - This saves time and money when the same text shows up more than once (very common in RAG pipelines). +/// +/// +/// The numeric data type used for vector calculations (typically float or double). +public sealed class CacheBackedEmbeddings : IEmbeddingModel +{ + private readonly IEmbeddingModel _inner; + private readonly IEmbeddingCache _cache; + private readonly string _namespace; + + /// + /// Initializes a new instance of the class. + /// + /// The underlying embedding model to wrap. + /// + /// The cache used to store embeddings. If null, a new unbounded is created. + /// + /// + /// An optional key namespace that identifies the inner model. If null, a namespace is derived from the + /// inner model's runtime type and embedding dimension so that different models do not share cache entries. + /// + /// Thrown when is null. + public CacheBackedEmbeddings( + IEmbeddingModel inner, + IEmbeddingCache? cache = null, + string? modelNamespace = null) + { + _inner = inner ?? throw new ArgumentNullException(nameof(inner)); + _cache = cache ?? new InMemoryEmbeddingCache(); + _namespace = string.IsNullOrEmpty(modelNamespace) + ? $"{_inner.GetType().FullName}:{_inner.EmbeddingDimension}" + : modelNamespace!; + } + + /// + /// Gets the underlying embedding model that this cache wraps. + /// + public IEmbeddingModel InnerModel => _inner; + + /// + /// Gets the cache used to store embeddings. + /// + public IEmbeddingCache Cache => _cache; + + /// + /// Gets the key namespace that identifies the wrapped model within the cache. + /// + public string Namespace => _namespace; + + /// + public int EmbeddingDimension => _inner.EmbeddingDimension; + + /// + public int MaxTokens => _inner.MaxTokens; + + /// + /// Builds the namespaced content-hash cache key for a piece of text. + /// + private string BuildKey(string text) => _namespace + ":" + ContentHash.ComputeHash(text); + + /// + public Vector Embed(string text) + { + if (text == null) + throw new ArgumentNullException(nameof(text)); + + var key = BuildKey(text); + if (_cache.TryGet(key, out var cached) && cached != null) + { + return cached; + } + + var embedding = _inner.Embed(text); + _cache.Set(key, Clone(embedding)); + return embedding; + } + + /// + public async Task> EmbedAsync(string text) + { + if (text == null) + throw new ArgumentNullException(nameof(text)); + + var key = BuildKey(text); + if (_cache.TryGet(key, out var cached) && cached != null) + { + return cached; + } + + var embedding = await _inner.EmbedAsync(text).ConfigureAwait(false); + _cache.Set(key, Clone(embedding)); + return embedding; + } + + /// + public Matrix EmbedBatch(IEnumerable texts) + { + var textList = ValidateBatch(texts); + + var results = new Vector[textList.Count]; + var missTexts = ResolveHitsAndCollectMisses(textList, results, out var missIndexLists, out var missKeys); + + if (missTexts.Count > 0) + { + var missMatrix = _inner.EmbedBatch(missTexts); + StoreAndDistributeMisses(missMatrix, missKeys, missIndexLists, results); + } + + return BuildMatrix(results); + } + + /// + public async Task> EmbedBatchAsync(IEnumerable texts) + { + var textList = ValidateBatch(texts); + + var results = new Vector[textList.Count]; + var missTexts = ResolveHitsAndCollectMisses(textList, results, out var missIndexLists, out var missKeys); + + if (missTexts.Count > 0) + { + var missMatrix = await _inner.EmbedBatchAsync(missTexts).ConfigureAwait(false); + StoreAndDistributeMisses(missMatrix, missKeys, missIndexLists, results); + } + + return BuildMatrix(results); + } + + /// + /// Validates and materializes the input batch, mirroring the conventions of the other embedding models. + /// + private static IList ValidateBatch(IEnumerable texts) + { + if (texts == null) + throw new ArgumentNullException(nameof(texts)); + + var textList = texts.ToList(); + if (textList.Count == 0) + throw new ArgumentException("Text collection cannot be empty", nameof(texts)); + + foreach (var text in textList) + { + if (string.IsNullOrWhiteSpace(text)) + throw new ArgumentException("Text cannot be null or empty", nameof(texts)); + } + + return textList; + } + + /// + /// Fills with cache hits and collects the distinct misses (de-duplicated by key, + /// in first-seen order) together with the result indices that each miss must populate. + /// + /// The distinct miss texts to send to the inner model. + private List ResolveHitsAndCollectMisses( + IList textList, + Vector[] results, + out List> missIndexLists, + out List missKeys) + { + var missTexts = new List(); + missKeys = new List(); + missIndexLists = new List>(); + + // Maps a miss key to its position within the miss lists so duplicates in the same batch collapse. + var keyToMissPosition = new Dictionary(); + + for (int i = 0; i < textList.Count; i++) + { + var key = BuildKey(textList[i]); + + if (_cache.TryGet(key, out var cached) && cached != null) + { + results[i] = cached; + continue; + } + + if (keyToMissPosition.TryGetValue(key, out var pos)) + { + missIndexLists[pos].Add(i); + } + else + { + keyToMissPosition[key] = missTexts.Count; + missTexts.Add(textList[i]); + missKeys.Add(key); + missIndexLists.Add(new List { i }); + } + } + + return missTexts; + } + + /// + /// Stores freshly computed miss embeddings in the cache and copies them into every result slot that requested them. + /// + private void StoreAndDistributeMisses( + Matrix missMatrix, + List missKeys, + List> missIndexLists, + Vector[] results) + { + if (missMatrix.Rows != missKeys.Count) + { + throw new InvalidOperationException( + $"Inner model returned {missMatrix.Rows} embeddings for {missKeys.Count} distinct texts."); + } + + for (int m = 0; m < missKeys.Count; m++) + { + var vector = GetRow(missMatrix, m); + _cache.Set(missKeys[m], Clone(vector)); + + foreach (var index in missIndexLists[m]) + { + results[index] = vector; + } + } + } + + /// + /// Extracts a single row of a matrix as a vector. + /// + private static Vector GetRow(Matrix matrix, int row) + { + var cols = matrix.Columns; + var values = new T[cols]; + for (int j = 0; j < cols; j++) + { + values[j] = matrix[row, j]; + } + + return new Vector(values); + } + + /// + /// Combines the ordered result vectors into a matrix, one row per input text. + /// + private Matrix BuildMatrix(Vector[] results) + { + var rows = results.Length; + var cols = EmbeddingDimension; + var matrix = new Matrix(rows, cols); + + for (int i = 0; i < rows; i++) + { + var vector = results[i]; + for (int j = 0; j < cols; j++) + { + matrix[i, j] = vector[j]; + } + } + + return matrix; + } + + /// + /// Creates a defensive copy of a vector so that cached data cannot be mutated by callers or the inner model. + /// + private static Vector Clone(Vector vector) + { + var values = new T[vector.Length]; + for (int i = 0; i < vector.Length; i++) + { + values[i] = vector[i]; + } + + return new Vector(values); + } +} diff --git a/src/RetrievalAugmentedGeneration/Embeddings/ContentHash.cs b/src/RetrievalAugmentedGeneration/Embeddings/ContentHash.cs new file mode 100644 index 0000000000..6a2fc30d29 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Embeddings/ContentHash.cs @@ -0,0 +1,67 @@ +using System.Security.Cryptography; +using System.Text; + +namespace AiDotNet.RetrievalAugmentedGeneration.Embeddings; + +/// +/// Provides stable, content-based hashing utilities used for caching and de-duplicating text. +/// +/// +/// +/// The hashes produced here are deterministic across processes and platforms (unlike +/// , which is randomized per process on modern .NET runtimes). +/// This makes them suitable as cache keys and for ingestion de-duplication where the same +/// content must always map to the same key. +/// +/// For Beginners: A content hash is a short "fingerprint" of a piece of text. +/// +/// - The same text always produces the same fingerprint. +/// - Different text (almost) always produces a different fingerprint. +/// - It lets us recognize "we have already seen this exact text" without storing the whole text. +/// +/// We use it so that repeated pieces of text are only embedded once and then reused. +/// +/// +public static class ContentHash +{ + /// + /// Computes a stable SHA-256 hash of the supplied text and returns it as a lowercase hex string. + /// + /// The text to hash. + /// A 64-character lowercase hexadecimal string representing the SHA-256 digest of the UTF-8 bytes of . + /// Thrown when is null. + public static string ComputeHash(string text) + { + if (text == null) + throw new ArgumentNullException(nameof(text)); + + var bytes = Encoding.UTF8.GetBytes(text); + + // SHA256 is available on net471, net8.0 and net10.0. + using var sha256 = SHA256.Create(); + var hashBytes = sha256.ComputeHash(bytes); + + return ToHex(hashBytes); + } + + /// + /// Converts a byte array to a lowercase hexadecimal string. + /// + /// The bytes to convert. + /// A lowercase hexadecimal representation of . + private static string ToHex(byte[] bytes) + { + // Manual hex conversion keeps behavior identical across all target frameworks + // (Convert.ToHexString is not available on net471). + const string hexChars = "0123456789abcdef"; + var chars = new char[bytes.Length * 2]; + for (int i = 0; i < bytes.Length; i++) + { + var b = bytes[i]; + chars[i * 2] = hexChars[b >> 4]; + chars[(i * 2) + 1] = hexChars[b & 0x0F]; + } + + return new string(chars); + } +} diff --git a/src/RetrievalAugmentedGeneration/Embeddings/IEmbeddingCache.cs b/src/RetrievalAugmentedGeneration/Embeddings/IEmbeddingCache.cs new file mode 100644 index 0000000000..70177a69b1 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Embeddings/IEmbeddingCache.cs @@ -0,0 +1,48 @@ +using AiDotNet.LinearAlgebra; + +namespace AiDotNet.RetrievalAugmentedGeneration.Embeddings; + +/// +/// Defines a cache that stores previously computed embedding vectors keyed by a stable string key. +/// +/// +/// +/// A cache lets a decorator avoid re-embedding text that has +/// already been embedded. Keys are expected to be content-derived (see ) and +/// namespaced by the underlying model's identity so that different models never share entries. +/// +/// For Beginners: This is a lookup table from "fingerprint of some text" to "its embedding". +/// +/// - Before asking the (possibly slow or paid) model to embed text, we check the cache. +/// - If the answer is already there, we reuse it instead of computing it again. +/// - Implementations must be safe to call from multiple threads at once. +/// +/// +/// The numeric data type used for vector calculations (typically float or double). +public interface IEmbeddingCache +{ + /// + /// Attempts to retrieve a cached embedding for the given key. + /// + /// The cache key. + /// When this method returns true, contains the cached embedding; otherwise null. + /// true if a cached embedding was found; otherwise false. + bool TryGet(string key, out Vector? embedding); + + /// + /// Stores an embedding in the cache under the given key, overwriting any existing entry. + /// + /// The cache key. + /// The embedding to store. + void Set(string key, Vector embedding); + + /// + /// Gets the number of entries currently held in the cache. + /// + int Count { get; } + + /// + /// Removes all entries from the cache. + /// + void Clear(); +} diff --git a/src/RetrievalAugmentedGeneration/Embeddings/InMemoryEmbeddingCache.cs b/src/RetrievalAugmentedGeneration/Embeddings/InMemoryEmbeddingCache.cs new file mode 100644 index 0000000000..4487ed206e --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Embeddings/InMemoryEmbeddingCache.cs @@ -0,0 +1,148 @@ +using System.Collections.Concurrent; +using System.Threading; +using AiDotNet.LinearAlgebra; + +namespace AiDotNet.RetrievalAugmentedGeneration.Embeddings; + +/// +/// A thread-safe, in-memory backed by a +/// with optional maximum size and approximate least-recently-used (LRU) eviction. +/// +/// +/// +/// When maxSize is zero (or negative) the cache is unbounded. When a positive maxSize is +/// supplied, inserting a new entry that would exceed the limit evicts the least-recently-used entries +/// until the cache is back within bounds. Recency is tracked with a monotonically increasing sequence +/// number that is updated on every read (hit) and write. +/// +/// For Beginners: This keeps embeddings in memory so repeated text is not recomputed. +/// +/// - It is safe to use from many threads at once. +/// - If you set a maximum size, the entries you have not used for the longest time are removed first +/// once the cache is full (this is called "LRU" eviction). +/// - If you do not set a maximum size, it keeps everything until you clear it. +/// +/// +/// The numeric data type used for vector calculations (typically float or double). +public sealed class InMemoryEmbeddingCache : IEmbeddingCache +{ + private sealed class CacheEntry + { + public CacheEntry(Vector value, long sequence) + { + Value = value; + Sequence = sequence; + } + + public Vector Value { get; set; } + + /// Recency marker; larger values are more recently used. + public long Sequence; + } + + private readonly ConcurrentDictionary _entries = new(); + private readonly int _maxSize; + private readonly object _evictionLock = new(); + private long _sequence; + + /// + /// Initializes a new instance of the class. + /// + /// + /// The maximum number of entries to retain before LRU eviction begins. + /// Use 0 (the default) or any non-positive value for an unbounded cache. + /// + public InMemoryEmbeddingCache(int maxSize = 0) + { + _maxSize = maxSize; + } + + /// + /// Gets the maximum number of entries this cache retains, or 0 if the cache is unbounded. + /// + public int MaxSize => _maxSize; + + /// + public int Count => _entries.Count; + + /// + public bool TryGet(string key, out Vector? embedding) + { + if (key == null) + throw new ArgumentNullException(nameof(key)); + + if (_entries.TryGetValue(key, out var entry)) + { + // Mark as recently used. + entry.Sequence = Interlocked.Increment(ref _sequence); + embedding = entry.Value; + return true; + } + + embedding = null; + return false; + } + + /// + public void Set(string key, Vector embedding) + { + if (key == null) + throw new ArgumentNullException(nameof(key)); + if (embedding == null) + throw new ArgumentNullException(nameof(embedding)); + + var seq = Interlocked.Increment(ref _sequence); + _entries.AddOrUpdate( + key, + _ => new CacheEntry(embedding, seq), + (_, existing) => + { + existing.Value = embedding; + existing.Sequence = seq; + return existing; + }); + + if (_maxSize > 0 && _entries.Count > _maxSize) + { + EvictIfNeeded(); + } + } + + /// + public void Clear() + { + _entries.Clear(); + } + + /// + /// Evicts least-recently-used entries until the cache is within its configured maximum size. + /// + private void EvictIfNeeded() + { + // A single lock keeps eviction correct under concurrency. Reads/writes to the + // ConcurrentDictionary remain lock-free; only the (rare) shrink path serializes. + lock (_evictionLock) + { + while (_entries.Count > _maxSize) + { + string? lruKey = null; + long lruSequence = long.MaxValue; + + foreach (var pair in _entries) + { + var seq = Interlocked.Read(ref pair.Value.Sequence); + if (seq < lruSequence) + { + lruSequence = seq; + lruKey = pair.Key; + } + } + + if (lruKey == null) + break; + + _entries.TryRemove(lruKey, out _); + } + } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/CacheBackedEmbeddingsTests.cs b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/CacheBackedEmbeddingsTests.cs new file mode 100644 index 0000000000..d9b4618a69 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RAG/Embeddings/CacheBackedEmbeddingsTests.cs @@ -0,0 +1,378 @@ +#nullable disable +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using AiDotNet.Helpers; +using AiDotNet.Interfaces; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Embeddings; +using Xunit; + +namespace AiDotNetTests.UnitTests.RAG.Embeddings +{ + /// + /// A fake that produces deterministic embeddings and records how many + /// times, and with which texts, it was invoked. Used to assert caching behavior without any network access. + /// + internal sealed class CountingEmbeddingModel : IEmbeddingModel + { + private static readonly INumericOperations NumOps = MathHelper.GetNumericOperations(); + + public CountingEmbeddingModel(int dimension = 8, int maxTokens = 512) + { + EmbeddingDimension = dimension; + MaxTokens = maxTokens; + } + + public int EmbeddingDimension { get; } + public int MaxTokens { get; } + + public int EmbedCallCount { get; private set; } + public int BatchCallCount { get; private set; } + public int TotalTextsEmbedded { get; private set; } + public List EmbeddedTexts { get; } = new(); + + /// Deterministic, dimension-aware signature for a text (independent of call state). + public static Vector ExpectedVector(string text, int dimension) + { + double baseValue = 0; + for (int i = 0; i < text.Length; i++) + { + baseValue += text[i] * (i + 1); + } + + var values = new T[dimension]; + for (int j = 0; j < dimension; j++) + { + values[j] = NumOps.FromDouble(baseValue + j); + } + + return new Vector(values); + } + + public Vector Embed(string text) + { + EmbedCallCount++; + TotalTextsEmbedded++; + EmbeddedTexts.Add(text); + return ExpectedVector(text, EmbeddingDimension); + } + + public Task> EmbedAsync(string text) => Task.FromResult(Embed(text)); + + public Matrix EmbedBatch(IEnumerable texts) + { + var list = texts.ToList(); + BatchCallCount++; + TotalTextsEmbedded += list.Count; + EmbeddedTexts.AddRange(list); + + var matrix = new Matrix(list.Count, EmbeddingDimension); + for (int i = 0; i < list.Count; i++) + { + var v = ExpectedVector(list[i], EmbeddingDimension); + for (int j = 0; j < EmbeddingDimension; j++) + { + matrix[i, j] = v[j]; + } + } + + return matrix; + } + + public Task> EmbedBatchAsync(IEnumerable texts) => Task.FromResult(EmbedBatch(texts)); + } + + public class CacheBackedEmbeddingsTests + { + private static bool VectorsEqual(Vector a, Vector b) + { + if (a.Length != b.Length) return false; + for (int i = 0; i < a.Length; i++) + { + if (Math.Abs(a[i] - b[i]) > 1e-9) return false; + } + return true; + } + + // ---------- ContentHash ---------- + + [Fact] + public void ContentHash_IsStableAndKnownForKnownInput() + { + // SHA-256("abc") well-known digest. + Assert.Equal( + "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad", + ContentHash.ComputeHash("abc")); + } + + [Fact] + public void ContentHash_SameInput_SameHash_DifferentInput_DifferentHash() + { + Assert.Equal(ContentHash.ComputeHash("hello world"), ContentHash.ComputeHash("hello world")); + Assert.NotEqual(ContentHash.ComputeHash("hello world"), ContentHash.ComputeHash("hello world!")); + Assert.Equal(64, ContentHash.ComputeHash("anything").Length); + } + + // ---------- Basic delegation ---------- + + [Fact] + public void DimensionAndMaxTokens_DelegateToInner() + { + var inner = new CountingEmbeddingModel(dimension: 32, maxTokens: 777); + var cached = new CacheBackedEmbeddings(inner); + + Assert.Equal(32, cached.EmbeddingDimension); + Assert.Equal(777, cached.MaxTokens); + } + + [Fact] + public void Constructor_NullInner_Throws() + { + Assert.Throws(() => new CacheBackedEmbeddings(null)); + } + + // ---------- (a) Single embed caching ---------- + + [Fact] + public void Embed_SameTextTwice_InnerCalledOnce_ReturnsCachedValue() + { + var inner = new CountingEmbeddingModel(); + var cached = new CacheBackedEmbeddings(inner); + + var first = cached.Embed("repeated text"); + var second = cached.Embed("repeated text"); + + Assert.Equal(1, inner.EmbedCallCount); + Assert.True(VectorsEqual(first, second)); + Assert.True(VectorsEqual(first, CountingEmbeddingModel.ExpectedVector("repeated text", inner.EmbeddingDimension))); + } + + [Fact(Timeout = 60000)] + public async Task EmbedAsync_SameTextTwice_InnerCalledOnce() + { + var inner = new CountingEmbeddingModel(); + var cached = new CacheBackedEmbeddings(inner); + + var first = await cached.EmbedAsync("async text"); + var second = await cached.EmbedAsync("async text"); + + Assert.Equal(1, inner.EmbedCallCount); + Assert.True(VectorsEqual(first, second)); + } + + [Fact] + public void Embed_DifferentTexts_InnerCalledPerDistinctText() + { + var inner = new CountingEmbeddingModel(); + var cached = new CacheBackedEmbeddings(inner); + + cached.Embed("one"); + cached.Embed("two"); + cached.Embed("one"); + + Assert.Equal(2, inner.EmbedCallCount); + Assert.Equal(2, cached.Cache.Count); + } + + // ---------- (b) Batch dedup + ordering ---------- + + [Fact] + public void EmbedBatch_DedupsWithinCall_PreservesOrder() + { + var inner = new CountingEmbeddingModel(); + var cached = new CacheBackedEmbeddings(inner); + + var texts = new List { "a", "b", "a", "c", "b" }; + var matrix = cached.EmbedBatch(texts); + + // Distinct misses = a, b, c => inner sees 3 texts in a single batch call. + Assert.Equal(1, inner.BatchCallCount); + Assert.Equal(3, inner.TotalTextsEmbedded); + Assert.Equal(new List { "a", "b", "c" }, inner.EmbeddedTexts); + + // Ordering + correctness: row i must equal the expected embedding for texts[i]. + Assert.Equal(5, matrix.Rows); + for (int i = 0; i < texts.Count; i++) + { + var expected = CountingEmbeddingModel.ExpectedVector(texts[i], inner.EmbeddingDimension); + for (int j = 0; j < inner.EmbeddingDimension; j++) + { + Assert.Equal(expected[j], matrix[i, j], 9); + } + } + } + + [Fact] + public void EmbedBatch_DedupsAcrossCalls() + { + var inner = new CountingEmbeddingModel(); + var cached = new CacheBackedEmbeddings(inner); + + cached.EmbedBatch(new List { "a", "b", "c" }); + Assert.Equal(3, inner.TotalTextsEmbedded); + + // Second call: only "d" is new; a/b/c are served from cache. + var matrix = cached.EmbedBatch(new List { "a", "b", "d" }); + + Assert.Equal(4, inner.TotalTextsEmbedded); + Assert.Equal(new List { "a", "b", "c", "d" }, inner.EmbeddedTexts); + + var texts = new[] { "a", "b", "d" }; + for (int i = 0; i < texts.Length; i++) + { + var expected = CountingEmbeddingModel.ExpectedVector(texts[i], inner.EmbeddingDimension); + for (int j = 0; j < inner.EmbeddingDimension; j++) + { + Assert.Equal(expected[j], matrix[i, j], 9); + } + } + } + + [Fact] + public void EmbedBatch_AllCached_DoesNotCallInner() + { + var inner = new CountingEmbeddingModel(); + var cached = new CacheBackedEmbeddings(inner); + + cached.EmbedBatch(new List { "x", "y" }); + var batchCallsAfterWarmup = inner.BatchCallCount; + + cached.EmbedBatch(new List { "x", "y" }); + + // No new inner batch invocation because everything was cached. + Assert.Equal(batchCallsAfterWarmup, inner.BatchCallCount); + } + + [Fact(Timeout = 60000)] + public async Task EmbedBatchAsync_DedupsAndPreservesOrder() + { + var inner = new CountingEmbeddingModel(); + var cached = new CacheBackedEmbeddings(inner); + + var texts = new List { "p", "q", "p" }; + var matrix = await cached.EmbedBatchAsync(texts); + + Assert.Equal(2, inner.TotalTextsEmbedded); + for (int i = 0; i < texts.Count; i++) + { + var expected = CountingEmbeddingModel.ExpectedVector(texts[i], inner.EmbeddingDimension); + for (int j = 0; j < inner.EmbeddingDimension; j++) + { + Assert.Equal(expected[j], matrix[i, j], 9); + } + } + } + + [Fact] + public void EmbedBatch_NullOrEmpty_Throws() + { + var cached = new CacheBackedEmbeddings(new CountingEmbeddingModel()); + + Assert.Throws(() => cached.EmbedBatch(null)); + Assert.Throws(() => cached.EmbedBatch(new List())); + } + + // ---------- (c) Model identity / dimension isolation ---------- + + [Fact] + public void DifferentModelDimension_SharedCache_NoCrossHits() + { + var cache = new InMemoryEmbeddingCache(); + var innerA = new CountingEmbeddingModel(dimension: 8); + var innerB = new CountingEmbeddingModel(dimension: 16); + + var cachedA = new CacheBackedEmbeddings(innerA, cache); + var cachedB = new CacheBackedEmbeddings(innerB, cache); + + cachedA.Embed("same text"); + cachedB.Embed("same text"); + + // Each inner model was invoked; the entries did not collide. + Assert.Equal(1, innerA.EmbedCallCount); + Assert.Equal(1, innerB.EmbedCallCount); + Assert.Equal(2, cache.Count); + Assert.NotEqual(cachedA.Namespace, cachedB.Namespace); + } + + [Fact] + public void ExplicitNamespace_IsolatesEntries() + { + var cache = new InMemoryEmbeddingCache(); + var innerA = new CountingEmbeddingModel(); + var innerB = new CountingEmbeddingModel(); + + var cachedA = new CacheBackedEmbeddings(innerA, cache, modelNamespace: "model-a"); + var cachedB = new CacheBackedEmbeddings(innerB, cache, modelNamespace: "model-b"); + + cachedA.Embed("shared"); + cachedB.Embed("shared"); + + Assert.Equal(1, innerA.EmbedCallCount); + Assert.Equal(1, innerB.EmbedCallCount); + Assert.Equal(2, cache.Count); + } + + // ---------- (d) LRU eviction ---------- + + [Fact] + public void InMemoryCache_Unbounded_KeepsEverything() + { + var cache = new InMemoryEmbeddingCache(); + for (int i = 0; i < 100; i++) + { + cache.Set("k" + i, CountingEmbeddingModel.ExpectedVector("t" + i, 4)); + } + + Assert.Equal(100, cache.Count); + } + + [Fact] + public void InMemoryCache_LruEviction_EvictsLeastRecentlyUsed() + { + var cache = new InMemoryEmbeddingCache(maxSize: 2); + var v = CountingEmbeddingModel.ExpectedVector("v", 4); + + cache.Set("k1", v); + cache.Set("k2", v); + + // Touch k1 so k2 becomes least recently used. + Assert.True(cache.TryGet("k1", out _)); + + cache.Set("k3", v); + + Assert.Equal(2, cache.Count); + Assert.True(cache.TryGet("k1", out _)); + Assert.False(cache.TryGet("k2", out _)); + Assert.True(cache.TryGet("k3", out _)); + } + + [Fact] + public void Decorator_WithLruCache_ReembedsAfterEviction() + { + var cache = new InMemoryEmbeddingCache(maxSize: 1); + var inner = new CountingEmbeddingModel(); + var cached = new CacheBackedEmbeddings(inner, cache); + + cached.Embed("a"); // miss -> count 1 + cached.Embed("b"); // miss -> count 2, evicts "a" + cached.Embed("a"); // miss again (evicted) -> count 3 + + Assert.Equal(3, inner.EmbedCallCount); + Assert.Equal(1, cache.Count); + } + + [Fact] + public void Cache_Clear_RemovesAllEntries() + { + var cache = new InMemoryEmbeddingCache(); + cache.Set("k1", CountingEmbeddingModel.ExpectedVector("a", 4)); + cache.Set("k2", CountingEmbeddingModel.ExpectedVector("b", 4)); + Assert.Equal(2, cache.Count); + + cache.Clear(); + Assert.Equal(0, cache.Count); + } + } +} From ebd6e3d97e0a5c3288c62a42e3384dbd184dc07b Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 17:20:25 -0400 Subject: [PATCH 13/25] feat(rag): lLM-based GraphRAG entity/relation extraction + community reports Add optional ITextGenerator (IGenerator)-driven extraction to KGConstructor and LLM community reports to CommunitySummarizer, with the existing regex/ extractive behavior retained as a documented offline fallback. - KGConstructor: new ctor accepting IGenerator; prompts the model for structured JSON (entities/relations, optional claims), parses with Newtonsoft.Json, and builds the graph. Malformed/empty JSON or generator errors transparently degrade to the regex path. Descriptions persisted on nodes/edges; claims attached to subject nodes. - CommunitySummarizer: new ctor accepting IGenerator; generates natural- language community reports, falling back to the extractive template when no generator is present or the response is empty. - Options: UseLlmExtraction (default true), ExtractClaims (default false). - net471-safe (guards, Newtonsoft.Json, no ThrowIfNull/double.IsFinite). - Tests: 9 CI-runnable tests with a scripted fake generator covering LLM population, malformed-JSON fallback, markdown-fenced JSON, claims, and LLM-vs-extractive community reports. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../Graph/Communities/CommunitySummarizer.cs | 121 +++++- .../Graph/Construction/ExtractedClaim.cs | 47 ++ .../Graph/Construction/ExtractedEntity.cs | 6 + .../Graph/Construction/ExtractedRelation.cs | 6 + .../Construction/KGConstructionOptions.cs | 21 + .../Graph/Construction/KGConstructor.cs | 400 +++++++++++++++++- .../KnowledgeGraph/LlmGraphExtractionTests.cs | 260 ++++++++++++ 7 files changed, 855 insertions(+), 6 deletions(-) create mode 100644 src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedClaim.cs create mode 100644 tests/AiDotNet.Tests/IntegrationTests/KnowledgeGraph/LlmGraphExtractionTests.cs diff --git a/src/RetrievalAugmentedGeneration/Graph/Communities/CommunitySummarizer.cs b/src/RetrievalAugmentedGeneration/Graph/Communities/CommunitySummarizer.cs index c4d6b96cd8..a132413a2a 100644 --- a/src/RetrievalAugmentedGeneration/Graph/Communities/CommunitySummarizer.cs +++ b/src/RetrievalAugmentedGeneration/Graph/Communities/CommunitySummarizer.cs @@ -1,8 +1,10 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Text; using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; namespace AiDotNet.RetrievalAugmentedGeneration.Graph.Communities; @@ -20,11 +22,44 @@ namespace AiDotNet.RetrievalAugmentedGeneration.Graph.Communities; /// and "influenced" relations might be summarized as: /// "Physics pioneers community: 3 entities centered around Einstein, with key relations: collaborated_with, influenced" /// +/// +/// When an LLM text generator is injected, the community Description +/// is instead produced by prompting the model to write a natural-language report summarizing the +/// community's entities and relations (Microsoft GraphRAG "community reports"). If no generator is +/// injected, or if the model returns an empty response, the extractive template description is used +/// as a documented fallback. +/// /// [ComponentType(ComponentType.DocumentStore)] [PipelineStage(PipelineStage.Indexing)] public class CommunitySummarizer { + /// + /// Optional LLM text generator used to write community reports. When null, the extractive + /// template description is used. + /// + private readonly IGenerator? _generator; + + /// + /// Initializes a new instance of the class using only the + /// extractive (offline, no-LLM) description template. + /// + public CommunitySummarizer() + { + _generator = null; + } + + /// + /// Initializes a new instance of the class with an optional + /// LLM text generator. When is non-null, community descriptions are + /// generated as natural-language reports by the model; otherwise the extractive template is used. + /// + /// The LLM text generator, or null to use the extractive fallback. + public CommunitySummarizer(IGenerator? generator) + { + _generator = generator; + } + /// /// Generates summaries for all communities in a Leiden result. /// @@ -132,14 +167,24 @@ public List SummarizePartition( { var node = graph.GetNode(id); return node?.GetProperty("name") ?? id; - }); + }).ToList(); + // Extractive template description (also used as the fallback when no LLM is present + // or the LLM returns an empty report). string description = $"Community of {members.Count} entities (primarily {dominantLabel}), " + $"centered around: {string.Join(", ", keyEntityNames)}. " + (keyRelations.Count > 0 ? $"Key relations: {string.Join(", ", keyRelations)}." : "No internal relations."); + // LLM-based community report (Microsoft GraphRAG parity), with extractive fallback. + if (_generator != null) + { + var report = TryGenerateReport(graph, members, keyEntityNames, keyRelations, dominantLabel); + if (!string.IsNullOrWhiteSpace(report)) + description = report!.Trim(); + } + summaries.Add(new CommunitySummary { CommunityId = communityId, @@ -153,4 +198,78 @@ public List SummarizePartition( return summaries; } + + /// + /// Attempts to generate a natural-language community report via the injected LLM generator. + /// Returns null when no generator is present, the generator throws, or it returns empty text, + /// signaling the caller to keep the extractive template description. + /// + private string? TryGenerateReport( + KnowledgeGraph graph, + List members, + List keyEntityNames, + List keyRelations, + string dominantLabel) + { + if (_generator == null) + return null; + + try + { + var prompt = BuildReportPrompt(graph, members, keyEntityNames, keyRelations, dominantLabel); + var report = _generator.Generate(prompt); + return string.IsNullOrWhiteSpace(report) ? null : report; + } + catch (Exception ex) + { + System.Diagnostics.Trace.TraceWarning( + $"CommunitySummarizer: LLM report generation threw ({ex.GetType().Name}); using extractive fallback."); + return null; + } + } + + /// + /// Builds the community-report prompt describing the community's entities and internal relations. + /// + internal static string BuildReportPrompt( + KnowledgeGraph graph, + List members, + List keyEntityNames, + List keyRelations, + string dominantLabel) + { + var memberSet = new HashSet(members); + var sb = new StringBuilder(); + sb.AppendLine("You are an analyst writing a concise report about a community of related entities"); + sb.AppendLine("in a knowledge graph. Summarize what this community is about, who/what the key"); + sb.AppendLine("entities are, and how they relate. Write 2-4 sentences of plain prose (no JSON)."); + sb.AppendLine(); + sb.AppendLine($"Community size: {members.Count} entities (predominant type: {dominantLabel})."); + + if (keyEntityNames.Count > 0) + sb.AppendLine($"Key entities: {string.Join(", ", keyEntityNames)}."); + + if (keyRelations.Count > 0) + sb.AppendLine($"Common relation types: {string.Join(", ", keyRelations)}."); + + // Include a sample of concrete relationships to ground the report. + sb.AppendLine("Relationships:"); + int shown = 0; + foreach (var nodeId in members) + { + if (shown >= 25) break; + var sourceName = graph.GetNode(nodeId)?.GetProperty("name") ?? nodeId; + foreach (var edge in graph.GetOutgoingEdges(nodeId)) + { + if (!memberSet.Contains(edge.TargetId)) continue; + var targetName = graph.GetNode(edge.TargetId)?.GetProperty("name") ?? edge.TargetId; + sb.AppendLine($"- {sourceName} {edge.RelationType} {targetName}"); + if (++shown >= 25) break; + } + } + + sb.AppendLine(); + sb.AppendLine("Report:"); + return sb.ToString(); + } } diff --git a/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedClaim.cs b/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedClaim.cs new file mode 100644 index 0000000000..a008a7aadf --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedClaim.cs @@ -0,0 +1,47 @@ +namespace AiDotNet.RetrievalAugmentedGeneration.Graph.Construction; + +/// +/// Represents a claim (covariate) extracted from text about an entity, following the +/// Microsoft GraphRAG claim-extraction model. +/// +/// +/// For Beginners: A claim is a factual statement asserted about an entity that +/// carries extra context beyond a simple relationship. For example, in +/// "Acme Corp was fined $2M for pollution in 2021": +/// - Subject: "Acme Corp" +/// - Object: "Environmental Agency" (or NONE if unspecified) +/// - ClaimType: "REGULATORY_VIOLATION" +/// - Description: "Acme Corp was fined $2M for pollution" +/// - Status: "TRUE" +/// +/// Claims let a knowledge graph capture assertions (fines, awards, allegations) that don't fit +/// neatly into subject-relation-object triples. +/// +/// +public class ExtractedClaim +{ + /// + /// The entity the claim is about (subject). + /// + public string Subject { get; set; } = string.Empty; + + /// + /// The entity the claim is directed at, or empty/"NONE" if not applicable. + /// + public string Object { get; set; } = string.Empty; + + /// + /// The category of the claim (e.g., REGULATORY_VIOLATION, AWARD, ALLEGATION). + /// + public string ClaimType { get; set; } = string.Empty; + + /// + /// A human-readable description of the claim. + /// + public string Description { get; set; } = string.Empty; + + /// + /// Truth status of the claim (e.g., TRUE, FALSE, SUSPECTED). + /// + public string Status { get; set; } = string.Empty; +} diff --git a/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedEntity.cs b/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedEntity.cs index 9a8ee927d4..8a8e8bb227 100644 --- a/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedEntity.cs +++ b/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedEntity.cs @@ -30,6 +30,12 @@ public class ExtractedEntity /// public double Confidence { get; set; } + /// + /// Optional natural-language description of the entity (populated by LLM extraction; + /// empty for the regex/heuristic path). + /// + public string Description { get; set; } = string.Empty; + /// /// Start character offset in the source text. /// diff --git a/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedRelation.cs b/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedRelation.cs index 23f974d7bf..cacc75ef63 100644 --- a/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedRelation.cs +++ b/src/RetrievalAugmentedGeneration/Graph/Construction/ExtractedRelation.cs @@ -33,4 +33,10 @@ public class ExtractedRelation /// Confidence score for this relation extraction (0.0 to 1.0). /// public double Confidence { get; set; } + + /// + /// Optional natural-language description of the relation (populated by LLM extraction; + /// empty for the regex/heuristic path). + /// + public string Description { get; set; } = string.Empty; } diff --git a/src/RetrievalAugmentedGeneration/Graph/Construction/KGConstructionOptions.cs b/src/RetrievalAugmentedGeneration/Graph/Construction/KGConstructionOptions.cs index 5c9150773e..62ad090d27 100644 --- a/src/RetrievalAugmentedGeneration/Graph/Construction/KGConstructionOptions.cs +++ b/src/RetrievalAugmentedGeneration/Graph/Construction/KGConstructionOptions.cs @@ -45,6 +45,23 @@ public class KGConstructionOptions /// public int? MaxEntitiesPerSentence { get; set; } + /// + /// Whether to use the injected LLM text generator for entity/relation extraction when one + /// is available. Default: true. When no generator is injected into the constructor this + /// setting has no effect and the regex/heuristic extractor is always used. When a generator + /// is present but returns malformed output, extraction transparently degrades to the + /// regex/heuristic fallback (no silent pretending). + /// + public bool? UseLlmExtraction { get; set; } + + /// + /// Whether to additionally ask the LLM to extract claims/covariates (Microsoft GraphRAG style) + /// about entities. Only takes effect when an LLM generator is injected and + /// is enabled. Default: false. Extracted claims are attached to + /// the corresponding entity node under the "claims" property. + /// + public bool? ExtractClaims { get; set; } + internal int GetEffectiveMaxChunkSize() { var value = MaxChunkSize ?? 500; @@ -84,6 +101,10 @@ internal int GetEffectiveMaxEntitiesPerSentence() return value; } + internal bool GetEffectiveUseLlmExtraction() => UseLlmExtraction ?? true; + + internal bool GetEffectiveExtractClaims() => ExtractClaims ?? false; + /// /// Validates cross-field constraints. Call after individual fields are set. /// diff --git a/src/RetrievalAugmentedGeneration/Graph/Construction/KGConstructor.cs b/src/RetrievalAugmentedGeneration/Graph/Construction/KGConstructor.cs index 2fa5333361..2aa0435bee 100644 --- a/src/RetrievalAugmentedGeneration/Graph/Construction/KGConstructor.cs +++ b/src/RetrievalAugmentedGeneration/Graph/Construction/KGConstructor.cs @@ -1,9 +1,12 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Text; using System.Text.RegularExpressions; using AiDotNet.Attributes; using AiDotNet.Enums; +using AiDotNet.Interfaces; +using Newtonsoft.Json.Linq; namespace AiDotNet.RetrievalAugmentedGeneration.Graph.Construction; @@ -19,8 +22,11 @@ namespace AiDotNet.RetrievalAugmentedGeneration.Graph.Construction; /// 3. Extract Relations: Detect relations via co-occurrence and proximity-based patterns /// 4. Entity Resolution: Merge similar entity names to reduce duplicates /// -/// This implementation works without an external LLM (purely heuristic), providing a baseline -/// that can be extended or replaced with LLM-based extraction. +/// When an LLM text generator is injected, construction instead prompts +/// the model to extract entities and relations as structured JSON (Microsoft GraphRAG parity), +/// parsing the result with Newtonsoft.Json. If no generator is injected, or if the LLM returns +/// malformed JSON, construction transparently falls back to the regex/heuristic path documented above +/// (an explicit offline fallback — never silent pretending). /// /// For Beginners: This class reads text and automatically builds a knowledge graph. /// @@ -39,6 +45,35 @@ namespace AiDotNet.RetrievalAugmentedGeneration.Graph.Construction; [PipelineStage(PipelineStage.DataIngestion)] public class KGConstructor { + /// + /// Optional LLM text generator used for structured entity/relation extraction. + /// When null, extraction uses the regex/heuristic fallback. + /// + private readonly IGenerator? _generator; + + /// + /// Initializes a new instance of the class using only the + /// regex/heuristic (offline, no-LLM) extraction path. + /// + public KGConstructor() + { + _generator = null; + } + + /// + /// Initializes a new instance of the class with an optional + /// LLM text generator. When is non-null (and LLM extraction is + /// enabled in the options), entities and relations are extracted by prompting the model for + /// structured JSON; otherwise the regex/heuristic fallback is used. + /// + /// + /// The LLM text generator to drive extraction, or null to use the regex/heuristic fallback. + /// + public KGConstructor(IGenerator? generator) + { + _generator = generator; + } + // Matches capitalized phrases including: // - Multi-word names: "Albert Einstein", "New York City" // - Names with connectors: "University of Cambridge", "Ludwig van Beethoven" @@ -121,13 +156,32 @@ public KnowledgeGraph ConstructFromText( // Step 2 & 3: Extract entities and relations from each chunk var allEntities = new List(); var allRelations = new List(); + var allClaims = new List(); + + bool useLlm = _generator != null && opts.GetEffectiveUseLlmExtraction(); + bool extractClaims = useLlm && opts.GetEffectiveExtractClaims(); foreach (var chunk in chunks) { - var entities = ExtractEntities(chunk, opts.GetEffectiveEntityConfidenceThreshold()); - allEntities.AddRange(entities); + List entities; + List relations; + + // LLM-based extraction path (Microsoft GraphRAG parity), with a transparent + // degrade-to-regex fallback whenever the model output cannot be parsed. + if (useLlm && + TryExtractWithLlm(chunk, extractClaims, out var llmEntities, out var llmRelations, out var llmClaims)) + { + entities = llmEntities; + relations = llmRelations; + allClaims.AddRange(llmClaims); + } + else + { + entities = ExtractEntities(chunk, opts.GetEffectiveEntityConfidenceThreshold()); + relations = ExtractRelations(chunk, entities, opts.GetEffectiveMaxEntitiesPerSentence()); + } - var relations = ExtractRelations(chunk, entities, opts.GetEffectiveMaxEntitiesPerSentence()); + allEntities.AddRange(entities); allRelations.AddRange(relations); } @@ -142,6 +196,10 @@ public KnowledgeGraph ConstructFromText( // Step 5: Add to graph AddToGraph(graph, allEntities, allRelations); + // Step 6: Attach extracted claims (covariates) to their subject nodes, if any. + if (allClaims.Count > 0) + AttachClaims(graph, allClaims); + return graph; } @@ -292,6 +350,332 @@ public List ExtractRelations(string text, List + /// Attempts LLM-driven structured extraction of entities and relations (and optionally claims) + /// from a text chunk. Returns false — signaling the caller to use the regex/heuristic + /// fallback — when the generator is unavailable, throws, returns empty text, or emits JSON that + /// cannot be parsed into at least one entity. + /// + /// The text chunk to extract from. + /// Whether to request claim/covariate extraction. + /// The parsed entities on success; empty otherwise. + /// The parsed relations on success; empty otherwise. + /// The parsed claims on success; empty otherwise. + /// True if extraction succeeded and yielded at least one entity; otherwise false. + private bool TryExtractWithLlm( + string chunk, + bool includeClaims, + out List entities, + out List relations, + out List claims) + { + entities = new List(); + relations = new List(); + claims = new List(); + + if (_generator == null) + return false; + + string response; + try + { + response = _generator.Generate(BuildExtractionPrompt(chunk, includeClaims)); + } + catch (Exception ex) + { + // Any generator failure degrades to the regex/heuristic fallback rather than throwing. + System.Diagnostics.Trace.TraceWarning( + $"KGConstructor: LLM extraction generator threw ({ex.GetType().Name}); using regex fallback."); + return false; + } + + if (string.IsNullOrWhiteSpace(response)) + return false; + + if (!TryParseExtractionJson(response, includeClaims, out entities, out relations, out claims)) + { + System.Diagnostics.Trace.TraceWarning( + "KGConstructor: LLM extraction returned malformed JSON; using regex fallback."); + return false; + } + + // A successful parse must yield at least one entity to be usable; otherwise fall back so + // the graph is not silently left empty for a chunk that clearly contains content. + if (entities.Count == 0) + return false; + + // Ensure every relation endpoint exists as an entity node so no edges are silently dropped. + EnsureRelationEndpointsExist(entities, relations); + return true; + } + + /// + /// Builds the extraction prompt sent to the LLM. Requests strict JSON in a fixed schema. + /// + internal static string BuildExtractionPrompt(string chunk, bool includeClaims) + { + var sb = new StringBuilder(); + sb.AppendLine("You are an information-extraction system that builds knowledge graphs."); + sb.AppendLine("Extract the named entities and the relationships between them from the TEXT below."); + sb.AppendLine(); + sb.AppendLine("Respond with ONLY a single JSON object (no markdown, no code fences, no commentary)"); + sb.AppendLine("using EXACTLY this schema:"); + sb.AppendLine("{"); + sb.AppendLine(" \"entities\": ["); + sb.AppendLine(" { \"name\": \"\", \"type\": \"\", \"description\": \"\" }"); + sb.AppendLine(" ],"); + sb.AppendLine(" \"relations\": ["); + sb.AppendLine(" { \"source\": \"\", \"relation\": \"\", \"target\": \"\", \"description\": \"\" }"); + if (includeClaims) + { + sb.AppendLine(" ],"); + sb.AppendLine(" \"claims\": ["); + sb.AppendLine(" { \"subject\": \"\", \"object\": \"\", \"type\": \"\", \"description\": \"\", \"status\": \"\" }"); + sb.AppendLine(" ]"); + } + else + { + sb.AppendLine(" ]"); + } + sb.AppendLine("}"); + sb.AppendLine(); + sb.AppendLine("Rules:"); + sb.AppendLine("- Use the entity's exact surface name. Relation source/target MUST match an entity name."); + sb.AppendLine("- RELATION_TYPE should be a short UPPER_SNAKE_CASE verb phrase (e.g., WORKS_AT, BORN_IN, FOUNDED)."); + sb.AppendLine("- If there are no entities, return {\"entities\": [], \"relations\": []}."); + sb.AppendLine(); + sb.AppendLine("TEXT:"); + sb.AppendLine("\"\"\""); + sb.AppendLine(chunk); + sb.AppendLine("\"\"\""); + return sb.ToString(); + } + + /// + /// Parses the LLM JSON response into entities, relations, and (optionally) claims. + /// Tolerates surrounding prose or markdown code fences by isolating the outermost JSON object. + /// Returns false on any parse failure so the caller can degrade to the regex fallback. + /// + internal static bool TryParseExtractionJson( + string response, + bool includeClaims, + out List entities, + out List relations, + out List claims) + { + entities = new List(); + relations = new List(); + claims = new List(); + + var json = ExtractJsonObject(response); + if (json == null) + return false; + + JObject root; + try + { + root = JObject.Parse(json); + } + catch (Exception) + { + // Newtonsoft throws JsonReaderException/JsonException on malformed input; treat any + // failure as malformed and degrade to the fallback. + return false; + } + + try + { + if (root["entities"] is JArray entityArray) + { + foreach (var item in entityArray) + { + var name = GetString(item, "name"); + if (string.IsNullOrWhiteSpace(name)) + continue; + + var type = GetString(item, "type"); + entities.Add(new ExtractedEntity + { + Name = name.Trim(), + Label = NormalizeLabel(type), + Description = GetString(item, "description").Trim(), + Confidence = 0.9, + StartOffset = 0, + EndOffset = 0 + }); + } + } + + if (root["relations"] is JArray relationArray) + { + foreach (var item in relationArray) + { + var source = GetString(item, "source"); + var target = GetString(item, "target"); + if (string.IsNullOrWhiteSpace(source) || string.IsNullOrWhiteSpace(target)) + continue; + + var relationType = NormalizeRelationType(GetString(item, "relation")); + relations.Add(new ExtractedRelation + { + SourceEntity = source.Trim(), + TargetEntity = target.Trim(), + RelationType = relationType, + Description = GetString(item, "description").Trim(), + Confidence = 0.9 + }); + } + } + + if (includeClaims && root["claims"] is JArray claimArray) + { + foreach (var item in claimArray) + { + var subject = GetString(item, "subject"); + if (string.IsNullOrWhiteSpace(subject)) + continue; + + claims.Add(new ExtractedClaim + { + Subject = subject.Trim(), + Object = GetString(item, "object").Trim(), + ClaimType = GetString(item, "type").Trim(), + Description = GetString(item, "description").Trim(), + Status = GetString(item, "status").Trim() + }); + } + } + } + catch (Exception) + { + return false; + } + + // A response whose JSON parsed but contained neither an "entities" nor a "relations" + // array is treated as malformed so we fall back. + return root["entities"] != null || root["relations"] != null; + } + + /// + /// Isolates the outermost JSON object (from the first '{' to the last '}') within an LLM + /// response, discarding any surrounding prose or markdown fences. Returns null if none found. + /// + private static string? ExtractJsonObject(string response) + { + if (string.IsNullOrEmpty(response)) + return null; + + int start = response.IndexOf('{'); + int end = response.LastIndexOf('}'); + if (start < 0 || end <= start) + return null; + + return response.Substring(start, end - start + 1); + } + + private static string GetString(JToken? token, string key) + { + if (token == null) + return string.Empty; + var value = token[key]; + if (value == null || value.Type == JTokenType.Null) + return string.Empty; + return value.ToString(); + } + + private static string NormalizeLabel(string type) + { + if (string.IsNullOrWhiteSpace(type)) + return "ENTITY"; + return type.Trim().ToUpperInvariant(); + } + + private static string NormalizeRelationType(string relation) + { + if (string.IsNullOrWhiteSpace(relation)) + return "RELATED_TO"; + + var cleaned = relation.Trim().ToUpperInvariant(); + var sb = new StringBuilder(cleaned.Length); + foreach (var ch in cleaned) + { + if (char.IsLetterOrDigit(ch)) + sb.Append(ch); + else if (ch == ' ' || ch == '-' || ch == '_') + sb.Append('_'); + // drop other punctuation + } + + var result = sb.ToString().Trim('_'); + return result.Length == 0 ? "RELATED_TO" : result; + } + + /// + /// Adds a minimal entity for any relation endpoint that was not itself extracted as an entity, + /// so LLM-declared edges are not dropped during graph construction. + /// + private static void EnsureRelationEndpointsExist( + List entities, + List relations) + { + var known = new HashSet(entities.Select(e => e.Name), StringComparer.OrdinalIgnoreCase); + foreach (var relation in relations) + { + AddMissingEndpoint(entities, known, relation.SourceEntity); + AddMissingEndpoint(entities, known, relation.TargetEntity); + } + } + + private static void AddMissingEndpoint( + List entities, + HashSet known, + string name) + { + if (string.IsNullOrWhiteSpace(name) || known.Contains(name)) + return; + + known.Add(name); + entities.Add(new ExtractedEntity + { + Name = name, + Label = "ENTITY", + Confidence = 0.7, + StartOffset = 0, + EndOffset = 0 + }); + } + + /// + /// Attaches extracted claims to their subject nodes under the "claims" property (as a list of + /// descriptions). Claims whose subject is not present in the graph are ignored. + /// + private static void AttachClaims(KnowledgeGraph graph, List claims) + { + foreach (var claim in claims) + { + if (string.IsNullOrWhiteSpace(claim.Subject)) + continue; + + string nodeId = claim.Subject.ToLowerInvariant().Replace(' ', '_'); + var node = graph.GetNode(nodeId); + if (node == null) + continue; + + var existing = node.GetProperty>("claims") ?? new List(); + var text = string.IsNullOrWhiteSpace(claim.ClaimType) + ? claim.Description + : $"[{claim.ClaimType}] {claim.Description}"; + if (!string.IsNullOrWhiteSpace(claim.Status)) + text += $" (status: {claim.Status})"; + existing.Add(text.Trim()); + node.SetProperty("claims", existing); + } + } + + #endregion + private static List ChunkText(string text, int maxChunkSize, int overlap) { if (maxChunkSize <= 0) @@ -377,6 +761,7 @@ private static List DeduplicateEntities( Name = canonicalName, Label = entity.Label, Confidence = entity.Confidence, + Description = entity.Description, StartOffset = entity.StartOffset, EndOffset = entity.EndOffset }); @@ -393,6 +778,7 @@ private static List RemapRelations( SourceEntity = resolvedMap.TryGetValue(r.SourceEntity, out var s) ? s : r.SourceEntity, TargetEntity = resolvedMap.TryGetValue(r.TargetEntity, out var t) ? t : r.TargetEntity, RelationType = r.RelationType, + Description = r.Description, Confidence = r.Confidence }) .Where(r => r.SourceEntity != r.TargetEntity) // Remove self-loops from resolution @@ -414,6 +800,8 @@ private static void AddToGraph( var node = new GraphNode(nodeId, entity.Label); node.SetProperty("name", entity.Name); node.SetProperty("confidence", entity.Confidence); + if (!string.IsNullOrWhiteSpace(entity.Description)) + node.SetProperty("description", entity.Description); graph.AddNode(node); } @@ -429,6 +817,8 @@ private static void AddToGraph( { var edge = new GraphEdge(sourceId, targetId, relation.RelationType, Math.Max(0.0, Math.Min(1.0, relation.Confidence))); + if (!string.IsNullOrWhiteSpace(relation.Description)) + edge.SetProperty("description", relation.Description); graph.AddEdge(edge); } catch (InvalidOperationException ex) diff --git a/tests/AiDotNet.Tests/IntegrationTests/KnowledgeGraph/LlmGraphExtractionTests.cs b/tests/AiDotNet.Tests/IntegrationTests/KnowledgeGraph/LlmGraphExtractionTests.cs new file mode 100644 index 0000000000..d80110515b --- /dev/null +++ b/tests/AiDotNet.Tests/IntegrationTests/KnowledgeGraph/LlmGraphExtractionTests.cs @@ -0,0 +1,260 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.Graph; +using AiDotNet.RetrievalAugmentedGeneration.Graph.Communities; +using AiDotNet.RetrievalAugmentedGeneration.Graph.Construction; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Xunit; + +namespace AiDotNet.Tests.IntegrationTests.KnowledgeGraph; + +/// +/// Tests for the LLM-based (Microsoft GraphRAG parity) entity/relation extraction and +/// community-report paths in and . +/// All tests are CI-runnable and use a scripted fake generator (no network). +/// +public class LlmGraphExtractionTests +{ + /// + /// A fake that returns scripted text and records the prompts it saw. + /// + private sealed class ScriptedGenerator : IGenerator + { + private readonly string _response; + public List Prompts { get; } = new List(); + + public int MaxContextTokens => 8192; + public int MaxGenerationTokens => 1024; + + public ScriptedGenerator(string response) + { + _response = response; + } + + public string Generate(string prompt) + { + Prompts.Add(prompt); + return _response; + } + + public GroundedAnswer GenerateGrounded(string query, IEnumerable> context) + { + return new GroundedAnswer + { + Query = query, + Answer = _response, + SourceDocuments = context?.ToList() ?? new List>(), + Citations = new List(), + ConfidenceScore = 1.0 + }; + } + } + + private const string ValidExtractionJson = @"{ + ""entities"": [ + { ""name"": ""Albert Einstein"", ""type"": ""PERSON"", ""description"": ""Theoretical physicist."" }, + { ""name"": ""Princeton University"", ""type"": ""ORGANIZATION"", ""description"": ""University in New Jersey."" }, + { ""name"": ""Ulm"", ""type"": ""LOCATION"", ""description"": ""City in Germany."" } + ], + ""relations"": [ + { ""source"": ""Albert Einstein"", ""relation"": ""WORKS_AT"", ""target"": ""Princeton University"", ""description"": ""Einstein worked at Princeton."" }, + { ""source"": ""Albert Einstein"", ""relation"": ""BORN_IN"", ""target"": ""Ulm"", ""description"": ""Einstein was born in Ulm."" } + ] + }"; + + [Fact(Timeout = 120000)] + public async Task KGConstructor_WithLlmGenerator_PopulatesGraphFromJson() + { + var generator = new ScriptedGenerator(ValidExtractionJson); + var constructor = new KGConstructor(generator); + + // Text content is deliberately terse; the LLM (scripted) supplies the structure. + var graph = constructor.ConstructFromText("Einstein biography."); + + // (a) LLM-extracted entities populate the graph. + Assert.True(generator.Prompts.Count > 0, "LLM generator should have been invoked."); + var nodeIds = graph.GetAllNodes().Select(n => n.Id).ToList(); + Assert.Contains("albert_einstein", nodeIds); + Assert.Contains("princeton_university", nodeIds); + Assert.Contains("ulm", nodeIds); + + var einstein = graph.GetNode("albert_einstein"); + Assert.NotNull(einstein); + Assert.Equal("PERSON", einstein!.Label); + Assert.Equal("Theoretical physicist.", einstein.GetProperty("description")); + + // (a) LLM-extracted relations populate the graph. + var edges = graph.GetAllEdges().ToList(); + Assert.Contains(edges, e => e.RelationType == "WORKS_AT" && + e.SourceId == "albert_einstein" && + e.TargetId == "princeton_university"); + Assert.Contains(edges, e => e.RelationType == "BORN_IN"); + } + + [Fact(Timeout = 120000)] + public async Task KGConstructor_MalformedJson_DegradesToRegexFallback() + { + // Not valid JSON — the LLM path must transparently fall back to the regex extractor. + var generator = new ScriptedGenerator("Sorry, I cannot help with that. {this is : not, valid json"); + var constructor = new KGConstructor(generator); + + var graph = constructor.ConstructFromText( + "Albert Einstein was born in Ulm, Germany. He worked at Princeton University."); + + Assert.True(generator.Prompts.Count > 0, "LLM generator should still have been invoked."); + + // (b) Regex fallback still populates the graph with capitalized entities. + var nodeIds = graph.GetAllNodes().Select(n => n.Id).ToList(); + Assert.NotEmpty(nodeIds); + Assert.Contains(nodeIds, id => id.Contains("einstein") || id.Contains("albert")); + Assert.Contains(nodeIds, id => id.Contains("princeton")); + } + + [Fact(Timeout = 120000)] + public async Task KGConstructor_JsonWrappedInMarkdownFences_IsParsed() + { + var fenced = "Here is the extraction:\n```json\n" + ValidExtractionJson + "\n```\nDone."; + var generator = new ScriptedGenerator(fenced); + var constructor = new KGConstructor(generator); + + var graph = constructor.ConstructFromText("Einstein biography."); + + var nodeIds = graph.GetAllNodes().Select(n => n.Id).ToList(); + Assert.Contains("albert_einstein", nodeIds); + Assert.Contains("princeton_university", nodeIds); + } + + [Fact(Timeout = 120000)] + public async Task KGConstructor_NoGenerator_UsesRegexFallback() + { + var constructor = new KGConstructor(); // no generator injected + var graph = constructor.ConstructFromText( + "Albert Einstein was born in Germany. Marie Curie worked at the University of Paris."); + + Assert.NotEmpty(graph.GetAllNodes()); + Assert.NotEmpty(graph.GetAllEdges()); + } + + [Fact(Timeout = 120000)] + public async Task KGConstructor_UseLlmExtractionDisabled_UsesRegexEvenWithGenerator() + { + // A generator that would throw if actually used — proves the regex path is taken. + var generator = new ScriptedGenerator(ValidExtractionJson); + var constructor = new KGConstructor(generator); + var opts = new KGConstructionOptions { UseLlmExtraction = false }; + + var graph = constructor.ConstructFromText( + "Albert Einstein was born in Ulm, Germany.", options: opts); + + Assert.Empty(generator.Prompts); // LLM never called + Assert.NotEmpty(graph.GetAllNodes()); + } + + [Fact(Timeout = 120000)] + public async Task KGConstructor_WithClaimsExtraction_AttachesClaimsToNodes() + { + const string jsonWithClaims = @"{ + ""entities"": [ + { ""name"": ""Acme Corp"", ""type"": ""ORGANIZATION"", ""description"": ""A company."" } + ], + ""relations"": [], + ""claims"": [ + { ""subject"": ""Acme Corp"", ""object"": ""NONE"", ""type"": ""REGULATORY_VIOLATION"", ""description"": ""Fined for pollution."", ""status"": ""TRUE"" } + ] + }"; + + var generator = new ScriptedGenerator(jsonWithClaims); + var constructor = new KGConstructor(generator); + var opts = new KGConstructionOptions { ExtractClaims = true }; + + var graph = constructor.ConstructFromText("Acme Corp news.", options: opts); + + var node = graph.GetNode("acme_corp"); + Assert.NotNull(node); + var claims = node!.GetProperty>("claims"); + Assert.NotNull(claims); + Assert.Contains(claims!, c => c.Contains("REGULATORY_VIOLATION") && c.Contains("pollution")); + } + + [Fact(Timeout = 120000)] + public async Task CommunitySummarizer_WithGenerator_UsesLlmReport() + { + var graph = BuildPhysicsGraph(); + var partition = new Dictionary + { + ["einstein"] = 0, + ["bohr"] = 0, + ["physics"] = 0 + }; + + const string report = "REPORT_SENTINEL: This community centers on physics pioneers Einstein and Bohr."; + var summarizer = new CommunitySummarizer(new ScriptedGenerator(report)); + + var summaries = summarizer.SummarizePartition(graph, partition, level: 0); + + // (c) LLM community report is used when a generator is present. + Assert.NotEmpty(summaries); + Assert.All(summaries, s => Assert.Contains("REPORT_SENTINEL", s.Description)); + } + + [Fact(Timeout = 120000)] + public async Task CommunitySummarizer_NoGenerator_UsesExtractiveFallback() + { + var graph = BuildPhysicsGraph(); + var partition = new Dictionary + { + ["einstein"] = 0, + ["bohr"] = 0, + ["physics"] = 0 + }; + + var summarizer = new CommunitySummarizer(); // no generator + var summaries = summarizer.SummarizePartition(graph, partition, level: 0); + + Assert.NotEmpty(summaries); + Assert.All(summaries, s => Assert.Contains("Community of", s.Description)); + } + + [Fact(Timeout = 120000)] + public async Task CommunitySummarizer_EmptyLlmResponse_FallsBackToExtractive() + { + var graph = BuildPhysicsGraph(); + var partition = new Dictionary + { + ["einstein"] = 0, + ["bohr"] = 0, + ["physics"] = 0 + }; + + var summarizer = new CommunitySummarizer(new ScriptedGenerator(" ")); + var summaries = summarizer.SummarizePartition(graph, partition, level: 0); + + Assert.NotEmpty(summaries); + Assert.All(summaries, s => Assert.Contains("Community of", s.Description)); + } + + private static KnowledgeGraph BuildPhysicsGraph() + { + var graph = new KnowledgeGraph(); + + var einstein = new GraphNode("einstein", "PERSON"); + einstein.SetProperty("name", "Albert Einstein"); + var bohr = new GraphNode("bohr", "PERSON"); + bohr.SetProperty("name", "Niels Bohr"); + var physics = new GraphNode("physics", "CONCEPT"); + physics.SetProperty("name", "Physics"); + + graph.AddNode(einstein); + graph.AddNode(bohr); + graph.AddNode(physics); + + graph.AddEdge(new GraphEdge("einstein", "bohr", "COLLABORATED_WITH", 0.9)); + graph.AddEdge(new GraphEdge("einstein", "physics", "CONTRIBUTED_TO", 0.9)); + graph.AddEdge(new GraphEdge("bohr", "physics", "CONTRIBUTED_TO", 0.9)); + + return graph; + } +} From 545e26bbed126a3cc70b44c303c7d41179acabaf Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 17:23:52 -0400 Subject: [PATCH 14/25] feat(rag): real pgvector + RedisVL vector stores in storage projects Replace the fake in-memory PostgresVectorDocumentStore/RedisVLDocumentStore in the core library with real IDocumentStore implementations in the opt-in storage projects, mirroring ElasticsearchDocumentStore. - AiDotNet.Storage.Postgres: PostgresVectorDocumentStore using Npgsql + pgvector (CREATE EXTENSION vector, table with vector(dim) column, ON CONFLICT upsert, KNN via <=>/<->/<+> ORDER BY LIMIT, jsonb metadata filters). - AiDotNet.Storage.Redis: RedisVLDocumentStore using StackExchange.Redis RediSearch (FT.CREATE HNSW vector + TAG/NUMERIC fields, HSET, FT.SEARCH KNN with server-side filters, FT.DROPINDEX). Honors metadata filters and never mutates cached docs. - Delete the two fake in-memory stores from the core library (breaking change). - Add unit tests for the SQL/Redis query builders and gated integration tests (POSTGRES_VECTOR_CONN / REDIS_CONN). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../DocumentStores/PgVectorMetric.cs | 60 +++ .../PostgresVectorDocumentStore.cs | 313 ++++++++++++ .../PostgresVectorFilterBuilder.cs | 99 ++++ .../DocumentStores/RedisVLDocumentStore.cs | 459 ++++++++++++++++++ .../DocumentStores/RedisVectorField.cs | 51 ++ .../DocumentStores/RedisVectorQueryBuilder.cs | 104 ++++ .../PostgresVectorDocumentStore.cs | 335 ------------- .../DocumentStores/RedisVLDocumentStore.cs | 303 ------------ .../PostgresVectorDocumentStoreTests.cs | 166 +++++++ .../RedisVLDocumentStoreTests.cs | 159 ++++++ 10 files changed, 1411 insertions(+), 638 deletions(-) create mode 100644 src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PgVectorMetric.cs create mode 100644 src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs create mode 100644 src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorFilterBuilder.cs create mode 100644 src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs create mode 100644 src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorField.cs create mode 100644 src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorQueryBuilder.cs delete mode 100644 src/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs delete mode 100644 src/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStoreTests.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStoreTests.cs diff --git a/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PgVectorMetric.cs b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PgVectorMetric.cs new file mode 100644 index 0000000000..658dbb9462 --- /dev/null +++ b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PgVectorMetric.cs @@ -0,0 +1,60 @@ +using System; + +using AiDotNet.Enums; + +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; + +/// +/// Maps an AiDotNet onto the corresponding pgvector distance +/// operator and converts the raw operator distance into a higher-is-better similarity score. +/// Extracted from for unit testing. +/// +/// +/// pgvector distance operators: +/// +/// → <=> (cosine distance = 1 − cosine similarity). +/// → <-> (L2 distance). +/// → <+> (L1 / taxicab distance). +/// +/// pgvector also exposes <#> (negative inner product); there is no matching +/// member, so it is not surfaced here. +/// +public static class PgVectorMetric +{ + /// Returns the pgvector distance operator for . + /// Thrown for metrics pgvector cannot express. + public static string Operator(DistanceMetricType metric) + { + switch (metric) + { + case DistanceMetricType.Cosine: + return "<=>"; + case DistanceMetricType.Euclidean: + return "<->"; + case DistanceMetricType.Manhattan: + return "<+>"; + default: + throw new NotSupportedException( + $"Distance metric '{metric}' is not supported by pgvector. Use Cosine, Euclidean or Manhattan."); + } + } + + /// + /// Converts a raw pgvector operator distance into a similarity score where higher means more similar. + /// Cosine → 1 − distance; Euclidean/Manhattan → 1 / (1 + distance). + /// + public static double ToSimilarity(DistanceMetricType metric, double distance) + { + switch (metric) + { + case DistanceMetricType.Cosine: + return 1.0 - distance; + case DistanceMetricType.Euclidean: + case DistanceMetricType.Manhattan: + return 1.0 / (1.0 + distance); + default: + throw new NotSupportedException( + $"Distance metric '{metric}' is not supported by pgvector."); + } + } +} diff --git a/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs new file mode 100644 index 0000000000..b14fa437e5 --- /dev/null +++ b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs @@ -0,0 +1,313 @@ +using System; +using System.Collections.Generic; +using System.Globalization; +using System.Linq; +using System.Text; +using System.Text.RegularExpressions; + +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using AiDotNet.Tensors.LinearAlgebra; + +using Newtonsoft.Json; +using Npgsql; + +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; + +/// +/// A real PostgreSQL + pgvector document store. Persists documents to a table +/// (id text pk, content text, metadata jsonb, embedding vector(dim)) and performs +/// approximate/exact nearest-neighbour search using pgvector's distance operators. +/// +/// The numeric type used for vector operations. +/// +/// +/// Ships in the opt-in AiDotNet.Storage.Postgres package so the core package stays free of +/// the Npgsql dependency (mirroring PostgresGraphCheckpointer). The interface contract is +/// synchronous, so this store uses Npgsql's synchronous API. +/// +/// +/// Vectors are sent as pgvector text literals ([v1,v2,...]) cast to vector; metadata is +/// stored as jsonb. Metadata filtering is pushed into SQL via +/// . Nearest-neighbour ordering uses the operator chosen by +/// (see ): <=> (cosine), +/// <-> (L2) or <+> (L1). +/// +/// +[ComponentType(ComponentType.DocumentStore)] +[PipelineStage(PipelineStage.Indexing)] +public class PostgresVectorDocumentStore : DocumentStoreBase +{ + private static readonly Regex IdentifierPattern = new("^[A-Za-z_][A-Za-z0-9_]*$", RegexOptions.Compiled); + + private readonly string _connectionString; + private readonly string _tableName; + private readonly DistanceMetricType _metric; + private readonly string _distanceOperator; + private readonly int _vectorDimension; + private int _documentCount; + private bool _schemaReady; + + /// + public override int DocumentCount => _documentCount; + + /// + public override int VectorDimension => _vectorDimension; + + /// Gets the table this store is bound to. + public string TableName => _tableName; + + /// + /// Initializes a new instance of the class and ensures + /// the pgvector extension, backing table and index exist. + /// + /// The Npgsql connection string. + /// + /// The backing table name. Must be a plain SQL identifier (letters, digits and underscores, not + /// starting with a digit); it is validated and interpolated into DDL/DML, so it cannot contain + /// arbitrary SQL. + /// + /// The fixed embedding dimension for the vector(dim) column. + /// The distance metric that selects the pgvector operator. + /// Thrown when arguments are empty or the table name is not a valid identifier. + /// Thrown when is not positive. + public PostgresVectorDocumentStore( + string connectionString, + string tableName, + int vectorDimension, + DistanceMetricType distanceMetric = DistanceMetricType.Cosine) + { + if (string.IsNullOrWhiteSpace(connectionString)) + throw new ArgumentException("Connection string cannot be empty", nameof(connectionString)); + if (string.IsNullOrWhiteSpace(tableName)) + throw new ArgumentException("Table name cannot be empty", nameof(tableName)); + if (!IdentifierPattern.IsMatch(tableName)) + throw new ArgumentException( + "Table name must be a valid SQL identifier (letters, digits, underscore; not starting with a digit).", + nameof(tableName)); + if (vectorDimension <= 0) + throw new ArgumentOutOfRangeException(nameof(vectorDimension), "Vector dimension must be positive"); + + _connectionString = connectionString; + _tableName = tableName; + _metric = distanceMetric; + _distanceOperator = PgVectorMetric.Operator(distanceMetric); + _vectorDimension = vectorDimension; + + EnsureSchema(); + } + + private NpgsqlConnection Open() + { + var connection = new NpgsqlConnection(_connectionString); + connection.Open(); + return connection; + } + + private void EnsureSchema() + { + using var connection = Open(); + using (var command = connection.CreateCommand()) + { + command.CommandText = + "CREATE EXTENSION IF NOT EXISTS vector;" + + $"CREATE TABLE IF NOT EXISTS {_tableName} (" + + "id text PRIMARY KEY, " + + "content text NOT NULL, " + + "metadata jsonb NOT NULL DEFAULT '{}'::jsonb, " + + $"embedding vector({_vectorDimension.ToString(CultureInfo.InvariantCulture)}) NOT NULL);"; + command.ExecuteNonQuery(); + } + + using (var countCommand = connection.CreateCommand()) + { + countCommand.CommandText = $"SELECT count(*) FROM {_tableName};"; + _documentCount = Convert.ToInt32(countCommand.ExecuteScalar(), CultureInfo.InvariantCulture); + } + + _schemaReady = true; + } + + private static string FormatVector(Vector embedding) + { + var sb = new StringBuilder("["); + var array = embedding.ToArray(); + for (var i = 0; i < array.Length; i++) + { + if (i > 0) + sb.Append(','); + sb.Append(Convert.ToDouble(array[i], CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture)); + } + + sb.Append(']'); + return sb.ToString(); + } + + /// + protected override void AddCore(VectorDocument vectorDocument) + { + if (vectorDocument.Embedding.Length != _vectorDimension) + throw new ArgumentException( + $"Document embedding dimension ({vectorDocument.Embedding.Length}) does not match the store's configured dimension ({_vectorDimension})."); + + using var connection = Open(); + using var command = connection.CreateCommand(); + command.CommandText = + $"INSERT INTO {_tableName} (id, content, metadata, embedding) " + + "VALUES (@id, @content, @metadata::jsonb, @embedding::vector) " + + "ON CONFLICT (id) DO UPDATE SET content = EXCLUDED.content, " + + "metadata = EXCLUDED.metadata, embedding = EXCLUDED.embedding " + + "RETURNING (xmax = 0) AS inserted;"; + AddDocumentParameters(command, vectorDocument); + + var inserted = Convert.ToBoolean(command.ExecuteScalar(), CultureInfo.InvariantCulture); + if (inserted) + _documentCount++; + } + + /// + protected override void AddBatchCore(IList> vectorDocuments) + { + if (vectorDocuments.Count == 0) + return; + + using var connection = Open(); + using var transaction = connection.BeginTransaction(); + var insertedCount = 0; + + foreach (var vectorDocument in vectorDocuments) + { + if (vectorDocument.Embedding.Length != _vectorDimension) + throw new ArgumentException( + $"Document embedding dimension ({vectorDocument.Embedding.Length}) does not match the store's configured dimension ({_vectorDimension})."); + + using var command = connection.CreateCommand(); + command.Transaction = transaction; + command.CommandText = + $"INSERT INTO {_tableName} (id, content, metadata, embedding) " + + "VALUES (@id, @content, @metadata::jsonb, @embedding::vector) " + + "ON CONFLICT (id) DO UPDATE SET content = EXCLUDED.content, " + + "metadata = EXCLUDED.metadata, embedding = EXCLUDED.embedding " + + "RETURNING (xmax = 0) AS inserted;"; + AddDocumentParameters(command, vectorDocument); + if (Convert.ToBoolean(command.ExecuteScalar(), CultureInfo.InvariantCulture)) + insertedCount++; + } + + transaction.Commit(); + _documentCount += insertedCount; + } + + private static void AddDocumentParameters(NpgsqlCommand command, VectorDocument vectorDocument) + { + command.Parameters.AddWithValue("id", vectorDocument.Document.Id); + command.Parameters.AddWithValue("content", vectorDocument.Document.Content ?? string.Empty); + command.Parameters.AddWithValue( + "metadata", + JsonConvert.SerializeObject(vectorDocument.Document.Metadata ?? new Dictionary())); + command.Parameters.AddWithValue("embedding", FormatVector(vectorDocument.Embedding)); + } + + /// + protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + { + var parameters = new Dictionary(); + var whereClause = PostgresVectorFilterBuilder.Build(metadataFilters, parameters); + + using var connection = Open(); + using var command = connection.CreateCommand(); + command.CommandText = + $"SELECT id, content, metadata, embedding {_distanceOperator} @q::vector AS distance " + + $"FROM {_tableName}{whereClause} " + + $"ORDER BY embedding {_distanceOperator} @q::vector " + + "LIMIT @k;"; + command.Parameters.AddWithValue("q", FormatVector(queryVector)); + command.Parameters.AddWithValue("k", topK); + foreach (var parameter in parameters) + command.Parameters.AddWithValue(parameter.Key, parameter.Value); + + var results = new List>(); + using var reader = command.ExecuteReader(); + while (reader.Read()) + { + var document = ReadDocument(reader); + var distance = Convert.ToDouble(reader.GetValue(3), CultureInfo.InvariantCulture); + document.RelevanceScore = NumOps.FromDouble(PgVectorMetric.ToSimilarity(_metric, distance)); + document.HasRelevanceScore = true; + results.Add(document); + } + + return results; + } + + /// + protected override Document? GetByIdCore(string documentId) + { + using var connection = Open(); + using var command = connection.CreateCommand(); + command.CommandText = $"SELECT id, content, metadata FROM {_tableName} WHERE id = @id;"; + command.Parameters.AddWithValue("id", documentId); + + using var reader = command.ExecuteReader(); + return reader.Read() ? ReadDocument(reader) : null; + } + + /// + protected override bool RemoveCore(string documentId) + { + using var connection = Open(); + using var command = connection.CreateCommand(); + command.CommandText = $"DELETE FROM {_tableName} WHERE id = @id;"; + command.Parameters.AddWithValue("id", documentId); + + var affected = command.ExecuteNonQuery(); + if (affected > 0 && _documentCount > 0) + { + _documentCount--; + return true; + } + + return affected > 0; + } + + /// + protected override IEnumerable> GetAllCore() + { + using var connection = Open(); + using var command = connection.CreateCommand(); + command.CommandText = $"SELECT id, content, metadata FROM {_tableName};"; + + var results = new List>(); + using var reader = command.ExecuteReader(); + while (reader.Read()) + results.Add(ReadDocument(reader)); + + return results; + } + + /// + public override void Clear() + { + if (!_schemaReady) + EnsureSchema(); + + using var connection = Open(); + using var command = connection.CreateCommand(); + command.CommandText = $"TRUNCATE TABLE {_tableName};"; + command.ExecuteNonQuery(); + _documentCount = 0; + } + + private static Document ReadDocument(NpgsqlDataReader reader) + { + var id = reader.GetString(0); + var content = reader.GetString(1); + var metadataJson = reader.IsDBNull(2) ? "{}" : reader.GetString(2); + var metadata = JsonConvert.DeserializeObject>(metadataJson) + ?? new Dictionary(); + return new Document(id, content, metadata); + } +} diff --git a/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorFilterBuilder.cs b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorFilterBuilder.cs new file mode 100644 index 0000000000..0437cdd22b --- /dev/null +++ b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorFilterBuilder.cs @@ -0,0 +1,99 @@ +using System.Collections; +using System.Collections.Generic; +using System.Globalization; +using System.Linq; +using System.Text; + +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; + +/// +/// Translates an AiDotNet metadata-filter dictionary into a parameterised PostgreSQL +/// WHERE clause over a jsonb metadata column. Extracted from +/// so the SQL-generation logic can be unit-tested +/// without a live database. +/// +/// +/// +/// Both the metadata keys and values are supplied as bound parameters (never string-interpolated), +/// so the generated SQL is injection-safe even for attacker-controlled metadata keys. +/// +/// Translation rules (mirroring the core DocumentStoreBase.MatchesFilters semantics): +/// +/// string / bool → equality: (metadata ->> @k) = @v (bool compared as the text true/false). +/// numeric → greater-than-or-equal range: (metadata ->> @k)::numeric >= @v. +/// collection → membership: (metadata ->> @k) = ANY(@v) (each element stringified). +/// anything else → equality against its string form. +/// +/// +public static class PostgresVectorFilterBuilder +{ + /// + /// Builds the WHERE clause (including the leading WHERE, or an empty string when there + /// are no filters) and appends the required parameters to . + /// + /// The metadata filters; may be null or empty. + /// + /// Receives the parameter name/value pairs to bind. Keys are the parameter names (e.g. k0, + /// v0); values are the objects to bind. + /// + /// The SQL WHERE fragment, or the empty string when no filters are supplied. + public static string Build(IReadOnlyDictionary? filters, IDictionary parameters) + { + if (filters == null || filters.Count == 0) + { + return string.Empty; + } + + var clauses = new List(); + var index = 0; + foreach (var filter in filters) + { + var keyParam = "k" + index.ToString(CultureInfo.InvariantCulture); + var valParam = "v" + index.ToString(CultureInfo.InvariantCulture); + parameters[keyParam] = filter.Key; + + var value = filter.Value; + if (value is string s) + { + clauses.Add($"(metadata ->> @{keyParam}) = @{valParam}"); + parameters[valParam] = s; + } + else if (value is bool b) + { + clauses.Add($"(metadata ->> @{keyParam}) = @{valParam}"); + parameters[valParam] = b ? "true" : "false"; + } + else if (IsNumeric(value)) + { + clauses.Add($"(metadata ->> @{keyParam})::numeric >= @{valParam}"); + parameters[valParam] = System.Convert.ToDouble(value, CultureInfo.InvariantCulture); + } + else if (value is IEnumerable enumerable) + { + var items = enumerable.Cast() + .Select(o => o?.ToString() ?? string.Empty) + .ToArray(); + clauses.Add($"(metadata ->> @{keyParam}) = ANY(@{valParam})"); + parameters[valParam] = items; + } + else + { + clauses.Add($"(metadata ->> @{keyParam}) = @{valParam}"); + parameters[valParam] = value?.ToString() ?? string.Empty; + } + + index++; + } + + var sb = new StringBuilder(" WHERE "); + sb.Append(string.Join(" AND ", clauses)); + return sb.ToString(); + } + + private static bool IsNumeric(object? value) + { + return value is sbyte || value is byte || value is short || value is ushort + || value is int || value is uint || value is long || value is ulong + || value is float || value is double || value is decimal; + } +} diff --git a/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs new file mode 100644 index 0000000000..ccb1f0690d --- /dev/null +++ b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs @@ -0,0 +1,459 @@ +using System; +using System.Collections.Generic; +using System.Globalization; +using System.Linq; +using System.Text.RegularExpressions; + +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using AiDotNet.Tensors.LinearAlgebra; + +using Newtonsoft.Json; +using StackExchange.Redis; + +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; + +/// +/// A real Redis + RediSearch (RedisVL-style) document store. Each document is a Redis HASH; a +/// RediSearch index provides KNN vector search with server-side metadata filtering. +/// +/// The numeric type used for vector operations. +/// +/// +/// Ships in the opt-in AiDotNet.Storage.Redis package so the core package stays free of the +/// StackExchange.Redis dependency (mirroring RedisGraphCheckpointer). Requires a Redis server +/// with the RediSearch module (e.g. Redis Stack). +/// +/// +/// The index is created with FT.CREATE ... ON HASH declaring a VECTOR field (HNSW, +/// FLOAT32) plus any caller-declared TAG/NUMERIC fields. Documents are +/// written with HSET; search uses FT.SEARCH KNN. Metadata filters on declared fields are +/// pushed into the query (); filters on undeclared keys are +/// honoured by in-memory post-filtering of an overfetched candidate set. Returned documents are always +/// freshly reconstructed from Redis, so nothing cached is ever mutated. +/// +/// +[ComponentType(ComponentType.DocumentStore)] +[PipelineStage(PipelineStage.Indexing)] +public class RedisVLDocumentStore : DocumentStoreBase, IDisposable +{ + private const string IdField = "_id"; + private const string ContentField = "_content"; + private const string MetadataField = "metadata"; + private const string EmbeddingField = "embedding"; + private const string ScoreField = "__vec_score"; + private const int CandidateMultiplier = 10; + + private static readonly Regex IdentifierPattern = new("^[A-Za-z_][A-Za-z0-9_]*$", RegexOptions.Compiled); + + private readonly IConnectionMultiplexer _connection; + private readonly bool _ownsConnection; + private readonly string _indexName; + private readonly string _keyPrefix; + private readonly int _vectorDimension; + private readonly DistanceMetricType _metric; + private readonly string _distanceMetric; + private readonly IReadOnlyList _fields; + private readonly Dictionary _fieldTypes; + private int _documentCount; + + /// + public override int DocumentCount => _documentCount; + + /// + public override int VectorDimension => _vectorDimension; + + /// Gets the RediSearch index name this store is bound to. + public string IndexName => _indexName; + + /// + /// Initializes a new instance of the class from a connection + /// string, creating the RediSearch index if it does not already exist. + /// + /// The StackExchange.Redis connection string. + /// The RediSearch index name (a plain identifier). + /// The fixed embedding dimension. + /// The distance metric (Cosine or Euclidean). + /// Optional metadata fields to index so their filters push server-side. + /// Optional Redis key prefix; defaults to {indexName}:doc:. + public RedisVLDocumentStore( + string connectionString, + string indexName, + int vectorDimension, + DistanceMetricType distanceMetric = DistanceMetricType.Cosine, + IEnumerable? filterableFields = null, + string? keyPrefix = null) + : this(Connect(connectionString), true, indexName, vectorDimension, distanceMetric, filterableFields, keyPrefix) + { + } + + /// + /// Initializes a new instance over an existing connection multiplexer (not owned/disposed by this instance). + /// + public RedisVLDocumentStore( + IConnectionMultiplexer connection, + string indexName, + int vectorDimension, + DistanceMetricType distanceMetric = DistanceMetricType.Cosine, + IEnumerable? filterableFields = null, + string? keyPrefix = null) + : this(connection ?? throw new ArgumentNullException(nameof(connection)), + false, indexName, vectorDimension, distanceMetric, filterableFields, keyPrefix) + { + } + + private RedisVLDocumentStore( + IConnectionMultiplexer connection, + bool ownsConnection, + string indexName, + int vectorDimension, + DistanceMetricType distanceMetric, + IEnumerable? filterableFields, + string? keyPrefix) + { + if (string.IsNullOrWhiteSpace(indexName) || !IdentifierPattern.IsMatch(indexName)) + throw new ArgumentException("Index name must be a valid identifier.", nameof(indexName)); + if (vectorDimension <= 0) + throw new ArgumentException("Vector dimension must be positive", nameof(vectorDimension)); + + _connection = connection; + _ownsConnection = ownsConnection; + _indexName = indexName; + _keyPrefix = string.IsNullOrWhiteSpace(keyPrefix) ? indexName + ":doc:" : keyPrefix!; + _vectorDimension = vectorDimension; + _metric = distanceMetric; + _distanceMetric = MapDistance(distanceMetric); + _fields = (filterableFields ?? Enumerable.Empty()).ToList(); + _fieldTypes = _fields.ToDictionary(f => f.Name, f => f.Type); + + EnsureIndex(); + } + + private static IConnectionMultiplexer Connect(string connectionString) + { + if (string.IsNullOrWhiteSpace(connectionString)) + throw new ArgumentException("Connection string cannot be empty", nameof(connectionString)); + return ConnectionMultiplexer.Connect(connectionString); + } + + private static string MapDistance(DistanceMetricType metric) + { + switch (metric) + { + case DistanceMetricType.Cosine: + return "COSINE"; + case DistanceMetricType.Euclidean: + return "L2"; + default: + throw new NotSupportedException( + $"Distance metric '{metric}' is not supported by RediSearch. Use Cosine or Euclidean."); + } + } + + private IDatabase Db => _connection.GetDatabase(); + + private void EnsureIndex() + { + try + { + Db.Execute("FT.INFO", _indexName); + } + catch (RedisServerException) + { + CreateIndex(); + } + + _documentCount = CountFromServer(); + } + + private void CreateIndex() + { + var args = new List + { + _indexName, "ON", "HASH", "PREFIX", "1", _keyPrefix, "SCHEMA", + ContentField, "TEXT", + }; + + foreach (var field in _fields) + { + args.Add(field.Name); + args.Add(field.Type == RedisVectorFieldType.Numeric ? "NUMERIC" : "TAG"); + } + + args.Add(EmbeddingField); + args.Add("VECTOR"); + args.Add("HNSW"); + args.Add("6"); + args.Add("TYPE"); + args.Add("FLOAT32"); + args.Add("DIM"); + args.Add(_vectorDimension.ToString(CultureInfo.InvariantCulture)); + args.Add("DISTANCE_METRIC"); + args.Add(_distanceMetric); + + Db.Execute("FT.CREATE", args); + } + + private int CountFromServer() + { + var result = Db.Execute("FT.SEARCH", _indexName, "*", "LIMIT", "0", "0"); + var array = (RedisResult[]?)result; + if (array == null || array.Length == 0) + return 0; + return (int)(long)array[0]; + } + + private static byte[] ToBytes(Vector vector) + { + var array = vector.ToArray(); + var bytes = new byte[array.Length * sizeof(float)]; + for (var i = 0; i < array.Length; i++) + { + var value = (float)Convert.ToDouble(array[i], CultureInfo.InvariantCulture); + var encoded = BitConverter.GetBytes(value); + Buffer.BlockCopy(encoded, 0, bytes, i * sizeof(float), sizeof(float)); + } + + return bytes; + } + + private static string Stringify(object? value) + { + if (value is bool b) + return b ? "true" : "false"; + if (value is IFormattable formattable) + return formattable.ToString(null, CultureInfo.InvariantCulture); + return value?.ToString() ?? string.Empty; + } + + /// + protected override void AddCore(VectorDocument vectorDocument) + { + if (vectorDocument.Embedding.Length != _vectorDimension) + throw new ArgumentException( + $"Document embedding dimension ({vectorDocument.Embedding.Length}) does not match the store's configured dimension ({_vectorDimension})."); + + var key = _keyPrefix + vectorDocument.Document.Id; + var existed = Db.KeyExists(key); + Db.HashSet(key, BuildEntries(vectorDocument)); + if (!existed) + _documentCount++; + } + + /// + protected override void AddBatchCore(IList> vectorDocuments) + { + if (vectorDocuments.Count == 0) + return; + + var db = Db; + foreach (var vd in vectorDocuments) + { + if (vd.Embedding.Length != _vectorDimension) + throw new ArgumentException( + $"Document embedding dimension ({vd.Embedding.Length}) does not match the store's configured dimension ({_vectorDimension})."); + } + + // Determine which ids are new (for accurate counting) before the pipelined writes. + var existsTasks = new List>(vectorDocuments.Count); + foreach (var vd in vectorDocuments) + existsTasks.Add(db.KeyExistsAsync(_keyPrefix + vd.Document.Id)); + + var writeBatch = db.CreateBatch(); + var writeTasks = new List(vectorDocuments.Count); + foreach (var vd in vectorDocuments) + writeTasks.Add(writeBatch.HashSetAsync(_keyPrefix + vd.Document.Id, BuildEntries(vd))); + + // existsTasks were created on 'db' directly (already dispatched); await them, then run writes. + var existed = existsTasks.Select(t => t.GetAwaiter().GetResult()).ToArray(); + writeBatch.Execute(); + foreach (var t in writeTasks) + t.GetAwaiter().GetResult(); + + _documentCount += existed.Count(e => !e); + } + + private HashEntry[] BuildEntries(VectorDocument vectorDocument) + { + var metadata = vectorDocument.Document.Metadata ?? new Dictionary(); + var entries = new List + { + new(IdField, vectorDocument.Document.Id), + new(ContentField, vectorDocument.Document.Content ?? string.Empty), + new(MetadataField, JsonConvert.SerializeObject(metadata)), + new(EmbeddingField, ToBytes(vectorDocument.Embedding)), + }; + + foreach (var field in _fields) + { + if (metadata.TryGetValue(field.Name, out var value)) + entries.Add(new HashEntry(field.Name, Stringify(value))); + } + + return entries.ToArray(); + } + + /// + protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + { + var expr = RedisVectorQueryBuilder.BuildFilterExpression(metadataFilters, _fieldTypes, out var unpushedKeys); + var knnK = unpushedKeys.Count > 0 ? topK * CandidateMultiplier : topK; + + var query = $"{expr}=>[KNN {knnK.ToString(CultureInfo.InvariantCulture)} @{EmbeddingField} $BLOB AS {ScoreField}]"; + var args = new List + { + _indexName, query, + "PARAMS", "2", "BLOB", ToBytes(queryVector), + "SORTBY", ScoreField, + "RETURN", "4", IdField, ContentField, MetadataField, ScoreField, + "DIALECT", "2", + "LIMIT", "0", knnK.ToString(CultureInfo.InvariantCulture), + }; + + var parsed = ParseSearch(Db.Execute("FT.SEARCH", args), includeScore: true); + + IEnumerable<(Document Doc, double Distance)> candidates = parsed; + if (unpushedKeys.Count > 0) + { + var postFilters = unpushedKeys + .Where(metadataFilters.ContainsKey) + .ToDictionary(k => k, k => metadataFilters[k]); + candidates = candidates.Where(c => MatchesFilters(c.Doc, postFilters)); + } + + var results = new List>(); + foreach (var candidate in candidates.Take(topK)) + { + candidate.Doc.RelevanceScore = NumOps.FromDouble(ToSimilarity(candidate.Distance)); + candidate.Doc.HasRelevanceScore = true; + results.Add(candidate.Doc); + } + + return results; + } + + private double ToSimilarity(double distance) + { + return _metric == DistanceMetricType.Cosine ? 1.0 - distance : 1.0 / (1.0 + distance); + } + + private List<(Document Doc, double Distance)> ParseSearch(RedisResult result, bool includeScore) + { + var docs = new List<(Document, double)>(); + var array = (RedisResult[]?)result; + if (array == null || array.Length < 1) + return docs; + + // array[0] is the total count; then repeating (key, fields[]) pairs. + for (var i = 1; i + 1 < array.Length; i += 2) + { + var fields = (RedisResult[]?)array[i + 1]; + if (fields == null) + continue; + + var map = new Dictionary(); + for (var j = 0; j + 1 < fields.Length; j += 2) + map[(string)fields[j]!] = (string)fields[j + 1]!; + + var doc = BuildDocument(map); + if (doc == null) + continue; + + var distance = 0.0; + if (includeScore && map.TryGetValue(ScoreField, out var scoreText)) + double.TryParse(scoreText, NumberStyles.Float, CultureInfo.InvariantCulture, out distance); + + docs.Add((doc, distance)); + } + + return docs; + } + + private static Document? BuildDocument(IReadOnlyDictionary map) + { + if (!map.TryGetValue(IdField, out var id) || string.IsNullOrEmpty(id)) + return null; + + map.TryGetValue(ContentField, out var content); + map.TryGetValue(MetadataField, out var metadataJson); + var metadata = string.IsNullOrEmpty(metadataJson) + ? new Dictionary() + : JsonConvert.DeserializeObject>(metadataJson!) ?? new Dictionary(); + + return new Document(id, content ?? string.Empty, metadata); + } + + /// + protected override Document? GetByIdCore(string documentId) + { + var entries = Db.HashGetAll(_keyPrefix + documentId); + if (entries.Length == 0) + return null; + + var map = entries.ToDictionary(e => (string)e.Name!, e => (string)e.Value!); + return BuildDocument(map); + } + + /// + protected override bool RemoveCore(string documentId) + { + var removed = Db.KeyDelete(_keyPrefix + documentId); + if (removed && _documentCount > 0) + _documentCount--; + return removed; + } + + /// + protected override IEnumerable> GetAllCore() + { + var all = new List>(); + var offset = 0; + const int batchSize = 1000; + + while (true) + { + var args = new List + { + _indexName, "*", + "RETURN", "3", IdField, ContentField, MetadataField, + "DIALECT", "2", + "LIMIT", offset.ToString(CultureInfo.InvariantCulture), batchSize.ToString(CultureInfo.InvariantCulture), + }; + + var parsed = ParseSearch(Db.Execute("FT.SEARCH", args), includeScore: false); + if (parsed.Count == 0) + break; + + all.AddRange(parsed.Select(p => p.Doc)); + offset += batchSize; + } + + return all; + } + + /// + public override void Clear() + { + try + { + Db.Execute("FT.DROPINDEX", _indexName, "DD"); + } + catch (RedisServerException) + { + // Index did not exist; nothing to drop. + } + + CreateIndex(); + _documentCount = 0; + } + + /// Disposes the owned connection multiplexer, if this instance created it. + public void Dispose() + { + if (_ownsConnection) + _connection.Dispose(); + } +} diff --git a/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorField.cs b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorField.cs new file mode 100644 index 0000000000..fe72ed5544 --- /dev/null +++ b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorField.cs @@ -0,0 +1,51 @@ +using System; +using System.Text.RegularExpressions; + +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; + +/// +/// The RediSearch field type used to index a piece of document metadata so it can be filtered on. +/// +public enum RedisVectorFieldType +{ + /// A TAG field: exact-match filtering for strings and booleans (e.g. @category:{science}). + Tag, + + /// A NUMERIC field: range filtering for numbers (e.g. @year:[2020 +inf]). + Numeric, +} + +/// +/// Declares a single metadata field that should index in its +/// RediSearch schema so that metadata filters on that key can be pushed down to the server. +/// +/// +/// RediSearch on a HASH index only filters on fields declared at FT.CREATE time. Filters that +/// reference an undeclared key are still honoured, but by post-filtering the candidate set in memory +/// rather than in the server query. +/// +public sealed class RedisVectorField +{ + private static readonly Regex NamePattern = new("^[A-Za-z_][A-Za-z0-9_]*$", RegexOptions.Compiled); + + /// Initializes a new field declaration. + /// The metadata key / hash field name. Must be a plain identifier. + /// The RediSearch field type. + /// Thrown when is not a valid identifier. + public RedisVectorField(string name, RedisVectorFieldType type) + { + if (string.IsNullOrWhiteSpace(name) || !NamePattern.IsMatch(name)) + throw new ArgumentException( + "Field name must be a valid identifier (letters, digits, underscore; not starting with a digit).", + nameof(name)); + + Name = name; + Type = type; + } + + /// Gets the metadata key / hash field name. + public string Name { get; } + + /// Gets the RediSearch field type. + public RedisVectorFieldType Type { get; } +} diff --git a/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorQueryBuilder.cs b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorQueryBuilder.cs new file mode 100644 index 0000000000..557b110ed4 --- /dev/null +++ b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorQueryBuilder.cs @@ -0,0 +1,104 @@ +using System.Collections; +using System.Collections.Generic; +using System.Globalization; +using System.Linq; +using System.Text; + +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; + +/// +/// Builds the RediSearch pre-filter expression from an AiDotNet metadata-filter dictionary, given the +/// set of fields declared in the index. Extracted from so the +/// query-generation logic can be unit-tested without a live Redis server. +/// +/// +/// Translation rules (declared fields only; anything else is reported for in-memory post-filtering): +/// +/// TAG field → exact match @name:{value}; a collection value becomes an OR list @name:{a|b}. Booleans render as true/false. +/// NUMERIC field → greater-than-or-equal range @name:[value +inf] (mirrors the core MatchesFilters gte semantics). +/// +/// Tag values are escaped for RediSearch's tag syntax so punctuation and spaces are matched literally. +/// +public static class RedisVectorQueryBuilder +{ + /// + /// Builds the RediSearch pre-filter expression (the part before =>[KNN ...]). + /// + /// The metadata filters; may be null or empty. + /// The fields declared in the index, keyed by name. + /// + /// Receives the filter keys that could not be pushed to the server (their key was not declared) and + /// therefore must be applied by in-memory post-filtering. + /// + /// The pre-filter expression, or * when nothing can be pushed down. + public static string BuildFilterExpression( + IReadOnlyDictionary? filters, + IReadOnlyDictionary declaredFields, + out List unpushedKeys) + { + unpushedKeys = new List(); + if (filters == null || filters.Count == 0) + { + return "*"; + } + + var clauses = new List(); + foreach (var filter in filters) + { + if (!declaredFields.TryGetValue(filter.Key, out var fieldType)) + { + unpushedKeys.Add(filter.Key); + continue; + } + + if (fieldType == RedisVectorFieldType.Numeric) + { + var d = System.Convert.ToDouble(filter.Value, CultureInfo.InvariantCulture); + clauses.Add($"@{filter.Key}:[{d.ToString("R", CultureInfo.InvariantCulture)} +inf]"); + } + else + { + clauses.Add($"@{filter.Key}:{{{FormatTagValue(filter.Value)}}}"); + } + } + + return clauses.Count > 0 ? "(" + string.Join(" ", clauses) + ")" : "*"; + } + + private static string FormatTagValue(object? value) + { + if (value is string s) + return EscapeTag(s); + + if (value is bool b) + return b ? "true" : "false"; + + if (value is IEnumerable enumerable && value is not string) + { + var items = enumerable.Cast() + .Select(o => o is bool eb ? (eb ? "true" : "false") : EscapeTag(o?.ToString() ?? string.Empty)); + return string.Join("|", items); + } + + return EscapeTag(value?.ToString() ?? string.Empty); + } + + /// + /// Escapes a value for use inside a RediSearch TAG query ({...}). RediSearch treats a + /// range of punctuation characters as tokenizer separators; each is prefixed with a backslash so the + /// value matches literally. + /// + public static string EscapeTag(string value) + { + const string special = ",.<>{}[]\"':;!@#$%^&*()-+=~| "; + var sb = new StringBuilder(value.Length); + foreach (var c in value) + { + if (special.IndexOf(c) >= 0) + sb.Append('\\'); + sb.Append(c); + } + + return sb.ToString(); + } +} diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs deleted file mode 100644 index 692233aa94..0000000000 --- a/src/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs +++ /dev/null @@ -1,335 +0,0 @@ - -using System; -using System.Collections.Generic; -using System.Linq; -using AiDotNet.Attributes; -using AiDotNet.Enums; -using AiDotNet.LinearAlgebra; -using AiDotNet.RetrievalAugmentedGeneration.Models; - -namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores -{ - /// - /// PostgreSQL with pgvector-inspired document store for relational database vector storage. - /// - /// - /// - /// This implementation provides an in-memory simulation of PostgreSQL with the pgvector extension, - /// which adds vector similarity search capabilities to the popular relational database. - /// It organizes documents in table-like structures with cosine similarity for retrieval. - /// - /// For Beginners: PostgreSQL is a powerful open-source relational database, and pgvector adds AI capabilities. - /// - /// Think of it like a smart database table: - /// - Each table stores documents with vector embeddings - /// - Combines traditional SQL queries with vector search - /// - Leverage existing PostgreSQL infrastructure - /// - /// This in-memory version is good for: - /// - Prototyping pgvector applications - /// - Testing table-based organization - /// - Small to medium collections (< 100K documents) - /// - /// Real PostgreSQL + pgvector provides: - /// - ACID transactions for data integrity - /// - Complex SQL joins with vector search - /// - Proven reliability and scalability - /// - Integration with existing database infrastructure - /// - /// - /// The numeric type for vector operations. - [ComponentType(ComponentType.DocumentStore)] - [PipelineStage(PipelineStage.Indexing)] - public class PostgresVectorDocumentStore : DocumentStoreBase - { - private readonly Dictionary> _documents; - private readonly string _tableName; - private int _vectorDimension; - - /// - /// Gets the number of documents currently stored in the table. - /// - public override int DocumentCount => _documents.Count; - - /// - /// Gets the dimensionality of vectors stored in this table. - /// - public override int VectorDimension => _vectorDimension; - - /// - /// Initializes a new instance of the PostgresVectorDocumentStore class. - /// - /// The table name to organize documents. - /// The initial capacity for the internal dictionary (default: 1000). - /// Thrown when table name is empty or initial capacity is not positive. - /// - /// For Beginners: Creates a new pgvector-style document table. - /// - /// Example: - /// - /// // Create a table for documents - /// var store = new PostgresVectorDocumentStore<float>("documents"); - /// - /// // Create a table for embeddings - /// var embStore = new PostgresVectorDocumentStore<double>("embeddings", 10000); - /// - /// - /// The table name helps organize different document collections. - /// - /// - public PostgresVectorDocumentStore(string tableName, int initialCapacity = 1000) - { - if (string.IsNullOrWhiteSpace(tableName)) - throw new ArgumentException("Table name cannot be empty", nameof(tableName)); - if (initialCapacity <= 0) - throw new ArgumentException("Initial capacity must be greater than zero", nameof(initialCapacity)); - - _tableName = tableName; - _documents = new Dictionary>(initialCapacity); - _vectorDimension = 0; - } - - /// - /// Core logic for adding a single vector document to the table. - /// - /// The validated vector document to add. - /// - /// - /// The first document added determines the vector dimension for all documents in this table. - /// All subsequent documents must have embeddings of the same dimension. - /// - /// - protected override void AddCore(VectorDocument vectorDocument) - { - if (_documents.Count == 0) - { - _vectorDimension = vectorDocument.Embedding.Length; - } - - _documents[vectorDocument.Document.Id] = vectorDocument; - } - - /// - /// Core logic for adding multiple vector documents in a batch operation. - /// - /// The validated list of vector documents to add. - /// Thrown when a document's embedding has inconsistent dimensions. - /// - /// - /// Batch operations are more efficient than adding documents individually, similar to PostgreSQL's - /// bulk insert capabilities. All documents must have embeddings with the same dimension. - /// - /// For Beginners: Adding many documents at once is much faster. - /// - /// Slow way: - /// - /// foreach (var doc in documents) - /// store.Add(doc); // Many individual inserts - /// - /// - /// Fast way: - /// - /// store.AddBatch(documents); // Single bulk insert - /// - /// - /// - protected override void AddBatchCore(IList> vectorDocuments) - { - if (_vectorDimension == 0 && vectorDocuments.Count > 0) - { - _vectorDimension = vectorDocuments[0].Embedding.Length; - } - - foreach (var vectorDocument in vectorDocuments) - { - if (vectorDocument.Embedding.Length != _vectorDimension) - throw new ArgumentException( - $"Vector dimension mismatch in batch. Expected {_vectorDimension}, got {vectorDocument.Embedding.Length} for document {vectorDocument.Document.Id}", - nameof(vectorDocuments)); - - _documents[vectorDocument.Document.Id] = vectorDocument; - } - } - - /// - /// Core logic for similarity search using cosine similarity with optional metadata filtering. - /// - /// The validated query vector. - /// The validated number of documents to return. - /// The validated metadata filters. - /// Top-k similar documents ordered by cosine similarity score. - /// - /// - /// Performs vector similarity search across all documents in the table, optionally filtering by metadata. - /// In real pgvector, this would use the vector similarity operators (<->, <=>, <#>) in SQL. - /// - /// For Beginners: Finds the most similar documents to your query. - /// - /// How it works: - /// 1. Filter documents by metadata (like SQL WHERE clause) - /// 2. Calculate similarity between query and each document - /// 3. Sort by similarity (highest first, like SQL ORDER BY) - /// 4. Return top-k matches (like SQL LIMIT) - /// - /// Example: - /// - /// // Find 10 most similar documents - /// var results = store.GetSimilar(queryVector, topK: 10); - /// - /// // Find similar documents from 2024 - /// var filters = new Dictionary<string, object> { ["year"] = "2024" }; - /// var filtered = store.GetSimilarWithFilters(queryVector, 5, filters); - /// - /// - /// - protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) - { - var scoredDocuments = new List<(Document Document, T Score)>(); - - var matchingDocuments = _documents.Values - .Where(vectorDoc => MatchesFilters(vectorDoc.Document, metadataFilters)); - - foreach (var vectorDoc in matchingDocuments) - { - var similarity = StatisticsHelper.CosineSimilarity(queryVector, vectorDoc.Embedding); - scoredDocuments.Add((vectorDoc.Document, similarity)); - } - - var results = scoredDocuments - .OrderByDescending(x => x.Score) - .Take(topK) - .Select(x => - { - // Create a new Document instance to avoid mutating the cached one - var newDocument = new Document(x.Document.Id, x.Document.Content, x.Document.Metadata) - { - RelevanceScore = x.Score, - HasRelevanceScore = true - }; - return newDocument; - }) - .ToList(); - - return results; - } - - /// - /// Core logic for retrieving a document by its unique identifier. - /// - /// The validated document ID. - /// The document if found; otherwise, null. - /// - /// For Beginners: Gets a specific document by ID (like SQL SELECT WHERE id = ...). - /// - /// Example: - /// - /// var doc = store.GetById("doc-123"); - /// if (doc != null) - /// Console.WriteLine($"Found: {doc.Content}"); - /// - /// - /// - protected override Document? GetByIdCore(string documentId) - { - return _documents.TryGetValue(documentId, out var vectorDoc) ? vectorDoc.Document : null; - } - - /// - /// Core logic for removing a document from the table. - /// - /// The validated document ID. - /// True if the document was found and removed; otherwise, false. - /// - /// - /// Removes the document from the table (like SQL DELETE). If this was the last document, - /// the vector dimension is reset to 0, allowing a new dimension on next add. - /// - /// For Beginners: Deletes a document from the table. - /// - /// Example: - /// - /// if (store.Remove("doc-123")) - /// Console.WriteLine("Document deleted"); - /// - /// - /// - protected override bool RemoveCore(string documentId) - { - var removed = _documents.Remove(documentId); - if (removed && _documents.Count == 0) - { - _vectorDimension = 0; - } - return removed; - } - - /// - /// Core logic for retrieving all documents in the table. - /// - /// An enumerable of all documents without their vector embeddings. - /// - /// - /// Returns all documents in the table in no particular order (like SQL SELECT * FROM table). - /// Vector embeddings are not included in the results. - /// - /// For Beginners: Gets every document in the table. - /// - /// Use cases: - /// - Export all documents for backup - /// - Migrate to a different table or database - /// - Bulk processing or analysis - /// - Debugging to see all stored data - /// - /// Warning: For large tables (> 10K documents), this can use significant memory. - /// In production, consider using pagination with LIMIT and OFFSET. - /// - /// Example: - /// - /// // Get all documents - /// var allDocs = store.GetAll().ToList(); - /// // Result is available in the returned value - /// - /// // Export to JSON - /// var json = JsonConvert.SerializeObject(allDocs); - /// File.WriteAllText($"{_tableName}_export.json", json); - /// - /// - /// - protected override IEnumerable> GetAllCore() - { - return _documents.Values.Select(vd => vd.Document).ToList(); - } - - /// - /// Removes all documents from the table and resets the vector dimension. - /// - /// - /// - /// Clears all documents from the table and resets the vector dimension to 0 (like SQL TRUNCATE TABLE). - /// The table name remains unchanged and is ready to accept new documents. - /// - /// For Beginners: Completely empties the table. - /// - /// After calling Clear(): - /// - All documents are removed (like DELETE FROM table) - /// - Vector dimension resets to 0 - /// - Table name stays the same - /// - Ready for new documents (even with different dimensions) - /// - /// Use with caution - this cannot be undone! - /// - /// Example: - /// - /// store.Clear(); - /// // Result is available in the returned value // 0 - /// - /// - /// - public override void Clear() - { - _documents.Clear(); - _vectorDimension = 0; - } - } -} - diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs deleted file mode 100644 index 828d328116..0000000000 --- a/src/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs +++ /dev/null @@ -1,303 +0,0 @@ - -using System; -using System.Collections.Generic; -using System.Linq; -using AiDotNet.Attributes; -using AiDotNet.Enums; -using AiDotNet.Interfaces; -using AiDotNet.LinearAlgebra; -using AiDotNet.RetrievalAugmentedGeneration.Models; - -namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; - -/// -/// Redis-based vector document store for low-latency applications. -/// -/// -/// -/// This implementation provides an in-memory simulation of Redis with RedisSearch module, -/// which adds full-text search and vector similarity capabilities to Redis. Redis excels -/// at ultra-low latency operations, making it ideal for real-time AI applications. -/// -/// For Beginners: Redis is an extremely fast in-memory data store, and RediSearch adds AI search. -/// -/// Think of it like RAM-speed database: -/// - All data stored in memory for sub-millisecond access -/// - Perfect for real-time recommendations and search -/// - Combines caching with vector search -/// -/// This in-memory version is good for: -/// - Prototyping Redis vector applications -/// - Testing low-latency search patterns -/// - Small to medium collections (< 100K documents) -/// -/// Real Redis + RediSearch provides: -/// - Sub-millisecond query latency -/// - Horizontal scaling with Redis Cluster -/// - Persistence options (RDB snapshots, AOF logs) -/// - Perfect for real-time AI applications -/// -/// -/// The numeric data type used for vector operations. -[ComponentType(ComponentType.DocumentStore)] -[PipelineStage(PipelineStage.Indexing)] -public class RedisVLDocumentStore : DocumentStoreBase -{ - private readonly Dictionary> _store; - private int _vectorDimension; - - /// - /// Gets the number of documents currently stored in the index. - /// - public override int DocumentCount => _store.Count; - - /// - /// Gets the dimensionality of vectors stored in this index. - /// - public override int VectorDimension => _vectorDimension; - - /// - /// Initializes a new instance of the RedisVLDocumentStore class. - /// - /// The Redis connection string. - /// The name of the RediSearch index. - /// The dimension of vector embeddings. - /// Thrown when parameters are invalid. - /// - /// For Beginners: Creates a new Redis vector store. - /// - /// Example: - /// - /// // Create a store for 384-dimensional embeddings - /// var store = new RedisVLDocumentStore<float>( - /// "localhost:6379", - /// "products", - /// vectorDimension: 384 - /// ); - /// - /// - /// Vector dimension must be known upfront to create the index properly. - /// - /// - public RedisVLDocumentStore(string connectionString, string indexName, int vectorDimension) - { - if (string.IsNullOrWhiteSpace(connectionString)) - throw new ArgumentException("Connection string cannot be empty", nameof(connectionString)); - if (string.IsNullOrWhiteSpace(indexName)) - throw new ArgumentException("Index name cannot be empty", nameof(indexName)); - if (vectorDimension <= 0) - throw new ArgumentException("Vector dimension must be positive", nameof(vectorDimension)); - - _store = new Dictionary>(); - _vectorDimension = vectorDimension; - } - - /// - /// Core logic for adding a single vector document to the index. - /// - /// The validated vector document to add. - /// - /// - /// Stores the document in the Redis-style hash structure with its vector embedding. - /// In real Redis, this would use HSET commands with vector field types. - /// - /// - protected override void AddCore(VectorDocument vectorDocument) - { - if (_vectorDimension == 0) - _vectorDimension = vectorDocument.Embedding.Length; - - _store[vectorDocument.Document.Id] = vectorDocument; - } - - /// - /// Core logic for adding multiple vector documents in a batch operation. - /// - /// The validated list of vector documents to add. - /// - /// - /// Batch operations are more efficient, mimicking Redis pipeline commands. - /// All documents are added in a single operation for better performance. - /// - /// For Beginners: Batch operations are faster in Redis. - /// - /// Slow (multiple round trips): - /// - /// foreach (var doc in documents) - /// store.Add(doc); // Many network calls - /// - /// - /// Fast (single pipeline): - /// - /// store.AddBatch(documents); // One network call - /// - /// - /// - protected override void AddBatchCore(IList> vectorDocuments) - { - if (vectorDocuments.Count == 0) return; - - if (_vectorDimension == 0) - _vectorDimension = vectorDocuments[0].Embedding.Length; - - foreach (var vd in vectorDocuments) - _store[vd.Document.Id] = vd; - } - - /// - /// Core logic for similarity search using cosine similarity with optional metadata filtering. - /// - /// The validated query vector. - /// The validated number of documents to return. - /// The validated metadata filters. - /// Top-k similar documents ordered by cosine similarity score. - /// - /// - /// Performs fast in-memory vector similarity search. In real Redis, this would use - /// the FT.SEARCH command with KNN vector similarity and optional filter expressions. - /// - /// For Beginners: Finds the most similar documents blazingly fast. - /// - /// Redis is optimized for speed: - /// 1. Everything is in RAM for instant access - /// 2. Calculate similarity for each document - /// 3. Sort by similarity (highest first) - /// 4. Return top-k matches - /// - /// Typical latency: < 1 millisecond for thousands of documents - /// - /// Example: - /// - /// // Lightning-fast search for 10 similar products - /// var results = store.GetSimilar(queryVector, topK: 10); - /// - /// - /// - protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) - { - var results = new List<(Document doc, T score)>(); - - foreach (var vd in _store.Values) - { - var similarity = StatisticsHelper.CosineSimilarity(queryVector, vd.Embedding); - vd.Document.RelevanceScore = similarity; - results.Add((vd.Document, similarity)); - } - - return results - .OrderByDescending(x => Convert.ToDouble(x.score)) - .Take(topK) - .Select(x => x.doc); - } - - /// - /// Core logic for retrieving a document by its unique identifier. - /// - /// The validated document ID. - /// The document if found; otherwise, null. - /// - /// For Beginners: Gets a specific document by ID (like Redis HGET). - /// - /// Example: - /// - /// var doc = store.GetById("product:123"); - /// if (doc != null) - /// Console.WriteLine($"Product: {doc.Content}"); - /// - /// - /// - protected override Document? GetByIdCore(string documentId) - { - return _store.TryGetValue(documentId, out var vd) ? vd.Document : null; - } - - /// - /// Core logic for removing a document from the index. - /// - /// The validated document ID. - /// True if the document was found and removed; otherwise, false. - /// - /// - /// Removes the document from Redis (like DEL command). This is an instant operation in memory. - /// - /// For Beginners: Deletes a document (instant in Redis). - /// - /// Example: - /// - /// if (store.Remove("product:123")) - /// Console.WriteLine("Product deleted from cache"); - /// - /// - /// - protected override bool RemoveCore(string documentId) - { - return _store.Remove(documentId); - } - - /// - /// Core logic for retrieving all documents in the index. - /// - /// An enumerable of all documents without their vector embeddings. - /// - /// - /// Returns all documents from the Redis index in no particular order (like SCAN command). - /// Vector embeddings are not included in the results. - /// - /// For Beginners: Gets every document from Redis. - /// - /// Use cases: - /// - Export cache contents for backup - /// - Migrate to a different Redis instance - /// - Cache warm-up or preloading - /// - Debugging to see all cached data - /// - /// Warning: In production Redis with millions of keys, consider using SCAN with cursor - /// instead of loading everything at once. - /// - /// Example: - /// - /// // Get all cached documents - /// var allDocs = store.GetAll().ToList(); - /// // Result is available in the returned value - /// - /// // Export for backup - /// var json = JsonConvert.SerializeObject(allDocs); - /// File.WriteAllText("redis_backup.json", json); - /// - /// - /// - protected override IEnumerable> GetAllCore() - { - return _store.Values.Select(vd => vd.Document).ToList(); - } - - /// - /// Removes all documents from the index and resets the vector dimension. - /// - /// - /// - /// Clears all documents from Redis (like FLUSHDB command) and resets the vector dimension to 0. - /// The index is ready to accept new documents after clearing. - /// - /// For Beginners: Completely empties the Redis index. - /// - /// After calling Clear(): - /// - All documents are removed (instant in Redis) - /// - Vector dimension resets to 0 - /// - Ready for new documents - /// - /// Use with caution - this cannot be undone! - /// - /// Example: - /// - /// store.Clear(); - /// // Result is available in the returned value // 0 - /// - /// - /// - public override void Clear() - { - _store.Clear(); - _vectorDimension = 0; - } -} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStoreTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStoreTests.cs new file mode 100644 index 0000000000..420879c4ba --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStoreTests.cs @@ -0,0 +1,166 @@ +#if NET5_0_OR_GREATER +using System; +using System.Collections.Generic; +using System.Linq; + +using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Models; + +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.DocumentStores +{ + /// + /// Tests for the real pgvector document store. The and + /// logic is unit-tested without a server; the end-to-end store test is a + /// SkippableFact that runs only when POSTGRES_VECTOR_CONN points at a live pgvector database. + /// The Storage.Postgres project is net10-only (Npgsql TFMs), hence the net5+ gate. + /// + public class PostgresVectorDocumentStoreTests + { + private const string ConnEnv = "POSTGRES_VECTOR_CONN"; + + [Fact] + public void FilterBuilder_NoFilters_ReturnsEmpty() + { + var parameters = new Dictionary(); + var sql = PostgresVectorFilterBuilder.Build(new Dictionary(), parameters); + Assert.Equal(string.Empty, sql); + Assert.Empty(parameters); + } + + [Fact] + public void FilterBuilder_StringFilter_EmitsEquality() + { + var parameters = new Dictionary(); + var sql = PostgresVectorFilterBuilder.Build( + new Dictionary { ["category"] = "science" }, parameters); + + Assert.Equal(" WHERE (metadata ->> @k0) = @v0", sql); + Assert.Equal("category", parameters["k0"]); + Assert.Equal("science", parameters["v0"]); + } + + [Fact] + public void FilterBuilder_BoolFilter_EmitsTextTrueFalse() + { + var parameters = new Dictionary(); + var sql = PostgresVectorFilterBuilder.Build( + new Dictionary { ["published"] = true }, parameters); + + Assert.Equal(" WHERE (metadata ->> @k0) = @v0", sql); + Assert.Equal("true", parameters["v0"]); + } + + [Fact] + public void FilterBuilder_NumericFilter_EmitsGteRange() + { + var parameters = new Dictionary(); + var sql = PostgresVectorFilterBuilder.Build( + new Dictionary { ["year"] = 2020 }, parameters); + + Assert.Equal(" WHERE (metadata ->> @k0)::numeric >= @v0", sql); + Assert.Equal(2020.0, (double)parameters["v0"]); + } + + [Fact] + public void FilterBuilder_CollectionFilter_EmitsAny() + { + var parameters = new Dictionary(); + var sql = PostgresVectorFilterBuilder.Build( + new Dictionary { ["tag"] = new[] { "a", "b" } }, parameters); + + Assert.Equal(" WHERE (metadata ->> @k0) = ANY(@v0)", sql); + Assert.Equal(new[] { "a", "b" }, (string[])parameters["v0"]); + } + + [Fact] + public void FilterBuilder_MultipleFilters_JoinedWithAnd() + { + var parameters = new Dictionary(); + var sql = PostgresVectorFilterBuilder.Build( + new Dictionary { ["category"] = "science", ["year"] = 2020 }, parameters); + + Assert.Contains(" AND ", sql); + Assert.Equal(4, parameters.Count); + } + + [Theory] + [InlineData(DistanceMetricType.Cosine, "<=>")] + [InlineData(DistanceMetricType.Euclidean, "<->")] + [InlineData(DistanceMetricType.Manhattan, "<+>")] + public void Metric_Operator_MapsCorrectly(DistanceMetricType metric, string expected) + { + Assert.Equal(expected, PgVectorMetric.Operator(metric)); + } + + [Fact] + public void Metric_Operator_UnsupportedThrows() + { + Assert.Throws(() => PgVectorMetric.Operator(DistanceMetricType.Jaccard)); + } + + [Fact] + public void Metric_ToSimilarity_Cosine_IsOneMinusDistance() + { + Assert.Equal(0.8, PgVectorMetric.ToSimilarity(DistanceMetricType.Cosine, 0.2), 6); + } + + [Fact] + public void Metric_ToSimilarity_Euclidean_IsInverseOfDistance() + { + Assert.Equal(0.5, PgVectorMetric.ToSimilarity(DistanceMetricType.Euclidean, 1.0), 6); + } + + [SkippableFact] + [Trait("Category", "Integration")] + public void EndToEnd_RoundTripsDocumentsAndFilters() + { + var conn = Environment.GetEnvironmentVariable(ConnEnv); + Skip.If(string.IsNullOrWhiteSpace(conn), $"Set {ConnEnv} to run the pgvector integration test."); + + var table = "vec_test_" + Guid.NewGuid().ToString("N"); + var store = new PostgresVectorDocumentStore(conn!, table, vectorDimension: 3); + try + { + store.Add(Doc("a", "alpha", new Vector(new float[] { 1, 0, 0 }), + new Dictionary { ["category"] = "science", ["year"] = 2021 })); + store.AddBatch(new[] + { + Doc("b", "beta", new Vector(new float[] { 0, 1, 0 }), + new Dictionary { ["category"] = "history", ["year"] = 2019 }), + Doc("c", "gamma", new Vector(new float[] { 0.9f, 0.1f, 0 }), + new Dictionary { ["category"] = "science", ["year"] = 2023 }), + }); + + Assert.Equal(3, store.DocumentCount); + + var top = store.GetSimilar(new Vector(new float[] { 1, 0, 0 }), 2).ToList(); + Assert.Equal("a", top[0].Id); + Assert.True(top[0].HasRelevanceScore); + + var filtered = store.GetSimilarWithFilters( + new Vector(new float[] { 1, 0, 0 }), 5, + new Dictionary { ["category"] = "science" }).ToList(); + Assert.All(filtered, d => Assert.Equal("science", d.Metadata["category"].ToString())); + + Assert.Equal("beta", store.GetById("b")!.Content); + Assert.True(store.Remove("b")); + Assert.Equal(2, store.DocumentCount); + Assert.Equal(2, store.GetAll().Count()); + } + finally + { + store.Clear(); + } + } + + private static VectorDocument Doc(string id, string content, Vector embedding, Dictionary metadata) + { + return new VectorDocument(new Document(id, content, metadata), embedding); + } + } +} +#endif diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStoreTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStoreTests.cs new file mode 100644 index 0000000000..bdf7ef1bf7 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStoreTests.cs @@ -0,0 +1,159 @@ +#if NET5_0_OR_GREATER +using System; +using System.Collections.Generic; +using System.Linq; + +using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Models; + +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.DocumentStores +{ + /// + /// Tests for the real RedisVL (Redis + RediSearch) document store. The + /// and logic is unit-tested + /// without a server; the end-to-end store test is a SkippableFact that runs only when + /// REDIS_CONN points at a live Redis Stack (RediSearch) server. The Storage.Redis project is + /// net10-only (StackExchange.Redis TFMs), hence the net5+ gate. + /// + public class RedisVLDocumentStoreTests + { + private const string ConnEnv = "REDIS_CONN"; + + private static readonly Dictionary Declared = new() + { + ["category"] = RedisVectorFieldType.Tag, + ["year"] = RedisVectorFieldType.Numeric, + }; + + [Fact] + public void QueryBuilder_NoFilters_ReturnsWildcard() + { + var expr = RedisVectorQueryBuilder.BuildFilterExpression( + new Dictionary(), Declared, out var unpushed); + Assert.Equal("*", expr); + Assert.Empty(unpushed); + } + + [Fact] + public void QueryBuilder_TagField_EmitsTagMatch() + { + var expr = RedisVectorQueryBuilder.BuildFilterExpression( + new Dictionary { ["category"] = "science" }, Declared, out var unpushed); + Assert.Equal("(@category:{science})", expr); + Assert.Empty(unpushed); + } + + [Fact] + public void QueryBuilder_NumericField_EmitsGteRange() + { + var expr = RedisVectorQueryBuilder.BuildFilterExpression( + new Dictionary { ["year"] = 2020 }, Declared, out _); + Assert.Equal("(@year:[2020 +inf])", expr); + } + + [Fact] + public void QueryBuilder_CollectionTag_EmitsOrList() + { + var expr = RedisVectorQueryBuilder.BuildFilterExpression( + new Dictionary { ["category"] = new[] { "science", "history" } }, Declared, out _); + Assert.Equal("(@category:{science|history})", expr); + } + + [Fact] + public void QueryBuilder_UndeclaredKey_ReportedForPostFiltering() + { + var expr = RedisVectorQueryBuilder.BuildFilterExpression( + new Dictionary { ["author"] = "smith" }, Declared, out var unpushed); + Assert.Equal("*", expr); + Assert.Equal(new[] { "author" }, unpushed); + } + + [Fact] + public void QueryBuilder_MixedDeclaredAndUndeclared() + { + var expr = RedisVectorQueryBuilder.BuildFilterExpression( + new Dictionary { ["category"] = "science", ["author"] = "smith" }, Declared, out var unpushed); + Assert.Equal("(@category:{science})", expr); + Assert.Equal(new[] { "author" }, unpushed); + } + + [Fact] + public void EscapeTag_EscapesSeparators() + { + Assert.Equal("hello\\ world", RedisVectorQueryBuilder.EscapeTag("hello world")); + Assert.Equal("a\\-b", RedisVectorQueryBuilder.EscapeTag("a-b")); + } + + [Fact] + public void RedisVectorField_InvalidName_Throws() + { + Assert.Throws(() => new RedisVectorField("bad name!", RedisVectorFieldType.Tag)); + } + + [SkippableFact] + [Trait("Category", "Integration")] + public void EndToEnd_RoundTripsDocumentsAndFilters() + { + var conn = Environment.GetEnvironmentVariable(ConnEnv); + Skip.If(string.IsNullOrWhiteSpace(conn), $"Set {ConnEnv} to run the RedisVL integration test."); + + var index = "v_idx_" + Guid.NewGuid().ToString("N").Substring(0, 8); + using var store = new RedisVLDocumentStore( + conn!, index, vectorDimension: 3, DistanceMetricType.Cosine, + filterableFields: new[] + { + new RedisVectorField("category", RedisVectorFieldType.Tag), + new RedisVectorField("year", RedisVectorFieldType.Numeric), + }); + try + { + store.Add(Doc("a", "alpha", new Vector(new float[] { 1, 0, 0 }), + new Dictionary { ["category"] = "science", ["year"] = 2021, ["author"] = "smith" })); + store.AddBatch(new[] + { + Doc("b", "beta", new Vector(new float[] { 0, 1, 0 }), + new Dictionary { ["category"] = "history", ["year"] = 2019, ["author"] = "jones" }), + Doc("c", "gamma", new Vector(new float[] { 0.9f, 0.1f, 0 }), + new Dictionary { ["category"] = "science", ["year"] = 2023, ["author"] = "smith" }), + }); + + Assert.Equal(3, store.DocumentCount); + + var top = store.GetSimilar(new Vector(new float[] { 1, 0, 0 }), 2).ToList(); + Assert.Equal("a", top[0].Id); + Assert.True(top[0].HasRelevanceScore); + + // Declared tag filter pushed server-side. + var science = store.GetSimilarWithFilters( + new Vector(new float[] { 1, 0, 0 }), 5, + new Dictionary { ["category"] = "science" }).ToList(); + Assert.All(science, d => Assert.Equal("science", d.Metadata["category"].ToString())); + + // Undeclared key honoured via in-memory post-filtering. + var byAuthor = store.GetSimilarWithFilters( + new Vector(new float[] { 1, 0, 0 }), 5, + new Dictionary { ["author"] = "jones" }).ToList(); + Assert.All(byAuthor, d => Assert.Equal("jones", d.Metadata["author"].ToString())); + + Assert.Equal("beta", store.GetById("b")!.Content); + Assert.True(store.Remove("b")); + Assert.Equal(2, store.DocumentCount); + Assert.Equal(2, store.GetAll().Count()); + } + finally + { + store.Clear(); + } + } + + private static VectorDocument Doc(string id, string content, Vector embedding, Dictionary metadata) + { + return new VectorDocument(new Document(id, content, metadata), embedding); + } + } +} +#endif From 0db44b59cff41384880b4d8e2b0a09860493a954 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 17:58:00 -0400 Subject: [PATCH 15/25] feat(rag): wire chunking/compression/vector-store/index into AiModelBuilder facade MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make the RAG pipeline wireable through the facade (your "100% wireable" goal): - ConfigureRetrievalAugmentedGeneration gains chunkingStrategy + contextCompressor params/fields, surfaced on AiModelResult and copied in BuildPipeline. - Fix the documentStore-dropped-without-a-KnowledgeGraph bug: the store is now kept in its own field, exposed as AiModelResult.DocumentStore, and a default vector retriever is built over it (honoring the configured embedding model + similarity metric) so vector-only RAG works. - New ConfigureVectorStore(IDocumentStore) and ConfigureVectorIndex(VectorIndexKind, metric) — one-call paths making all document-store backends and the in-memory indexes (Flat/HNSW/IVF/LSH) reachable. These live on both IAiModelBuilder and the concrete builder (every Configure* belongs on the interface). - ConfigureRAG(RAGConfiguration) materializes the previously-orphaned RAGConfigurationBuilder into the builder fields. Verified: 10 facade-wiring tests. YAML/declarative path is the follow-up (#31). Co-Authored-By: Claude Opus 4.8 (1M context) --- src/AiModelBuilder.BuildPipeline.cs | 6 + src/AiModelBuilder.VectorIndexStore.cs | 145 +++++++++++ src/AiModelBuilder.Workflows.cs | 246 +++++++++++++++++- src/AiModelBuilder.cs | 3 + src/Enums/VectorIndexKind.cs | 30 +++ src/Interfaces/IAiModelBuilder.cs | 40 ++- src/Models/Options/AiModelResultOptions.cs | 15 ++ src/Models/Results/AiModelResult.cs | 37 +++ .../RAGFacadeWiringTests.cs | 164 ++++++++++++ 9 files changed, 680 insertions(+), 6 deletions(-) create mode 100644 src/AiModelBuilder.VectorIndexStore.cs create mode 100644 src/Enums/VectorIndexKind.cs create mode 100644 tests/AiDotNet.Tests/IntegrationTests/RetrievalAugmentedGeneration/RAGFacadeWiringTests.cs diff --git a/src/AiModelBuilder.BuildPipeline.cs b/src/AiModelBuilder.BuildPipeline.cs index dd9b08ac69..ca14ffe810 100644 --- a/src/AiModelBuilder.BuildPipeline.cs +++ b/src/AiModelBuilder.BuildPipeline.cs @@ -2802,6 +2802,9 @@ private async Task> BuildStreamingSupervisedAs RagReranker = _ragReranker, RagGenerator = _ragGenerator, QueryProcessors = _queryProcessors, + ChunkingStrategy = _chunkingStrategy, + ContextCompressor = _contextCompressor, + DocumentStore = _ragDocumentStore, KnowledgeGraph = _knowledgeGraph, GraphStore = _graphStore, HybridGraphRetriever = _hybridGraphRetriever, @@ -5447,6 +5450,9 @@ configuredOptimizerOptions is AdamOptimizerOptions adamOpts RagReranker = _ragReranker, RagGenerator = _ragGenerator, QueryProcessors = _queryProcessors, + ChunkingStrategy = _chunkingStrategy, + ContextCompressor = _contextCompressor, + DocumentStore = _ragDocumentStore, LoRAConfiguration = _loraConfiguration, CrossValidationResult = cvResults, DeploymentConfiguration = deploymentConfig, diff --git a/src/AiModelBuilder.VectorIndexStore.cs b/src/AiModelBuilder.VectorIndexStore.cs new file mode 100644 index 0000000000..9fa8b90eb9 --- /dev/null +++ b/src/AiModelBuilder.VectorIndexStore.cs @@ -0,0 +1,145 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Interfaces; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Indexes; + +namespace AiDotNet; + +/// +/// A lightweight in-memory that adapts any +/// (Flat / HNSW / IVF / LSH) into a document store usable by the RAG facade. +/// +/// +/// +/// This adapter lets ConfigureVectorIndex surface the in-memory vector indexes and a chosen +/// similarity metric as a retriever/store without depending on any particular concrete document store. +/// The chosen index (and the metric it was constructed with) performs the actual similarity search; this +/// wrapper only keeps the document payloads keyed by id so retrieval can return full +/// instances with their relevance scores populated. +/// +/// For Beginners: A vector index only knows about ids and vectors. A document store also needs +/// to hold the actual text/metadata. This class glues them together: it forwards vectors to the index for +/// fast similarity search, then looks the matching documents back up by id. +/// +/// +/// The numeric data type used for vector operations. +internal sealed class VectorIndexDocumentStore : IDocumentStore +{ + private readonly IVectorIndex _index; + private readonly Dictionary> _documents = new(); + private int _vectorDimension; + + /// + /// Initializes a new instance of the class. + /// + /// The vector index that performs similarity search. + /// The dimensionality of stored vectors (0 = inferred on first add). + public VectorIndexDocumentStore(IVectorIndex index, int vectorDimension = 0) + { + _index = index ?? throw new ArgumentNullException(nameof(index)); + _vectorDimension = vectorDimension; + } + + /// + public int DocumentCount => _documents.Count; + + /// + public int VectorDimension => _vectorDimension; + + /// + public void Add(VectorDocument vectorDocument) + { + if (vectorDocument == null) throw new ArgumentNullException(nameof(vectorDocument)); + var doc = vectorDocument.Document; + if (doc == null || string.IsNullOrEmpty(doc.Id)) + throw new ArgumentException("VectorDocument must have a Document with a non-empty Id.", nameof(vectorDocument)); + + var embedding = vectorDocument.Embedding; + if (embedding == null || embedding.Length == 0) + throw new ArgumentException("VectorDocument must have a non-empty Embedding.", nameof(vectorDocument)); + + if (_vectorDimension == 0) + _vectorDimension = embedding.Length; + + _documents[doc.Id] = vectorDocument; + _index.Add(doc.Id, embedding); + } + + /// + public void AddBatch(IEnumerable> vectorDocuments) + { + if (vectorDocuments == null) throw new ArgumentNullException(nameof(vectorDocuments)); + foreach (var vd in vectorDocuments) + { + Add(vd); + } + } + + /// + public IEnumerable> GetSimilar(Vector queryVector, int topK) + => GetSimilarWithFilters(queryVector, topK, new Dictionary()); + + /// + public IEnumerable> GetSimilarWithFilters(Vector queryVector, int topK, Dictionary metadataFilters) + { + if (queryVector == null) throw new ArgumentNullException(nameof(queryVector)); + if (topK <= 0) return Enumerable.Empty>(); + + bool hasFilters = metadataFilters != null && metadataFilters.Count > 0; + // When filtering, over-fetch from the index so enough candidates survive the filter. + int searchK = hasFilters ? Math.Min(_documents.Count, Math.Max(topK * 4, topK)) : topK; + if (searchK <= 0) return Enumerable.Empty>(); + + var hits = _index.Search(queryVector, searchK); + var results = new List>(topK); + + foreach (var (id, score) in hits) + { + if (!_documents.TryGetValue(id, out var vd)) continue; + if (hasFilters && !MatchesFilters(vd.Document, metadataFilters!)) continue; + + var doc = vd.Document; + doc.RelevanceScore = score; + results.Add(doc); + if (results.Count >= topK) break; + } + + return results; + } + + /// + public Document? GetById(string documentId) + => _documents.TryGetValue(documentId, out var vd) ? vd.Document : null; + + /// + public bool Remove(string documentId) + { + if (!_documents.Remove(documentId)) return false; + _index.Remove(documentId); + return true; + } + + /// + public void Clear() + { + _documents.Clear(); + _index.Clear(); + } + + /// + public IEnumerable> GetAll() => _documents.Values.Select(vd => vd.Document); + + private static bool MatchesFilters(Document document, Dictionary filters) + { + foreach (var kvp in filters) + { + if (!document.Metadata.TryGetValue(kvp.Key, out var value)) return false; + if (!Equals(value, kvp.Value)) return false; + } + + return true; + } +} diff --git a/src/AiModelBuilder.Workflows.cs b/src/AiModelBuilder.Workflows.cs index f2a57c931c..b0085ba944 100644 --- a/src/AiModelBuilder.Workflows.cs +++ b/src/AiModelBuilder.Workflows.cs @@ -409,7 +409,11 @@ public IAiModelBuilder ConfigureLoRA(ILoRAConfiguration l /// Optional query processors for improving search quality. /// Optional graph storage backend for Graph RAG (e.g., MemoryGraphStore, FileGraphStore). /// Optional pre-configured knowledge graph. If null but graphStore is provided, a new one is created. - /// Optional document store for hybrid vector + graph retrieval. + /// Optional document store. When a knowledge graph is also present it is folded + /// into a hybrid vector + graph retriever; when supplied WITHOUT a knowledge graph (and no explicit + /// ) a default dense vector retriever is built over it so vector-only RAG works. + /// Optional chunking strategy used to split source documents into passages before indexing/embedding. + /// Optional context compressor used to shrink retrieved passages before they reach the generator. /// This builder instance for method chaining. /// /// @@ -447,21 +451,29 @@ public IAiModelBuilder ConfigureRetrievalAugmentedGeneration IEnumerable? queryProcessors = null, IGraphStore? graphStore = null, KnowledgeGraph? knowledgeGraph = null, - IDocumentStore? documentStore = null) + IDocumentStore? documentStore = null, + IChunkingStrategy? chunkingStrategy = null, + IContextCompressor? contextCompressor = null) { // Configure standard RAG components _ragRetriever = retriever; _ragReranker = reranker; _ragGenerator = generator; _queryProcessors = queryProcessors; + _chunkingStrategy = chunkingStrategy; + _contextCompressor = contextCompressor; + _ragDocumentStore = documentStore; // Configure Graph RAG components - // If all Graph RAG parameters are null, clear Graph RAG fields - if (graphStore == null && knowledgeGraph == null && documentStore == null) + // If EVERY parameter is null, clear all RAG fields (full disable). + if (retriever == null && reranker == null && generator == null && queryProcessors == null + && graphStore == null && knowledgeGraph == null && documentStore == null + && chunkingStrategy == null && contextCompressor == null) { _graphStore = null; _knowledgeGraph = null; _hybridGraphRetriever = null; + _ragDocumentStore = null; return this; } @@ -493,9 +505,235 @@ public IAiModelBuilder ConfigureRetrievalAugmentedGeneration _hybridGraphRetriever = null; } + // FIX (documentStore-dropped bug): previously a document store supplied WITHOUT a knowledge + // graph was silently dropped (only consumed by the hybrid retriever above). Now, when a store + // is provided without a graph and the caller did not pass an explicit retriever, build a default + // dense/vector retriever over it so vector-only RAG is reachable through the facade. + if (documentStore != null && _knowledgeGraph == null && _ragRetriever == null) + { + _ragRetriever = BuildDefaultVectorRetriever(documentStore); + } + + return this; + } + + /// + /// Builds a default dense vector retriever over the supplied document store, honoring the + /// configured embedding model when present and falling back to a deterministic stub embedder + /// (sized to the store's vector dimension) otherwise so the retriever is always usable. + /// + /// The document store to retrieve from. Must not be null. + /// A ready-to-use over the store. + private IRetriever BuildDefaultVectorRetriever(IDocumentStore documentStore) + { + // The vector search / similarity metric already lives inside the document store implementation + // (the store's GetSimilarWithFilters does the comparison). When the caller configured a + // similarity metric via ConfigureSimilarityMetric it is honored by store construction helpers + // (ConfigureVectorStore / ConfigureVectorIndex); here we simply pair the store with an embedder. + var embeddingModel = _configuredEmbeddingModel ?? BuildDefaultEmbeddingModel(documentStore.VectorDimension); + return new RetrievalAugmentedGeneration.Retrievers.VectorRetriever(documentStore, embeddingModel); + } + + /// + /// Creates a deterministic default embedding model of the requested dimension for wiring a + /// vector retriever when no embedding model was explicitly configured. + /// + /// The embedding dimension; must match the target store's vector dimension. + /// A default . + private static IEmbeddingModel BuildDefaultEmbeddingModel(int dimension) + { + int safeDimension = dimension > 0 ? dimension : 768; + return new RetrievalAugmentedGeneration.Embeddings.StubEmbeddingModel(embeddingDimension: safeDimension); + } + + /// + /// Configures a vector document store as the RAG backend in a single call, building a default dense + /// vector retriever over it so every implementation is reachable + /// through the facade. + /// + /// The document store to retrieve from (in-memory, FAISS, Pinecone, Qdrant, etc.). + /// Default number of documents the retriever returns per query. + /// This builder instance for method chaining. + /// + /// + /// The retriever is a that pairs + /// the store with the embedding model configured via ConfigureEmbeddingModel (falling back to a + /// deterministic stub embedder sized to the store's vector dimension when none was configured). The + /// similarity metric used for comparison is the one baked into the supplied store; to also select the + /// metric from the facade use . + /// + /// For Beginners: This is the one-call way to say "here is my document database, use it for + /// retrieval." After calling it, AiModelResult.RagRetriever and AiModelResult.DocumentStore + /// are populated and vector-only RAG works. + /// + /// + public IAiModelBuilder ConfigureVectorStore(IDocumentStore store, int defaultTopK = 5) + { + if (store == null) throw new ArgumentNullException(nameof(store)); + + _ragDocumentStore = store; + var embeddingModel = _configuredEmbeddingModel ?? BuildDefaultEmbeddingModel(store.VectorDimension); + _ragRetriever = new RetrievalAugmentedGeneration.Retrievers.VectorRetriever(store, embeddingModel, defaultTopK); return this; } + /// + /// Configures an in-memory vector search index (Flat / HNSW / IVF / LSH) with a chosen similarity + /// metric, adapts it into a document store, and builds a dense retriever over it. + /// + /// Which in-memory index to build. Defaults to an exact flat index. + /// The embedding dimension for the store (0 = inferred on first add). + /// The similarity metric for the index. When null, the metric configured via + /// ConfigureSimilarityMetric is used, otherwise cosine similarity. + /// Default number of documents the retriever returns per query. + /// This builder instance for method chaining. + /// + /// + /// This surfaces the library's in-memory vector indexes directly from the facade. The chosen metric flows + /// into index construction (and, if not already set, is recorded as the builder's configured similarity + /// metric so evaluation sees it too). The resulting index is wrapped in a document store and paired with + /// the configured embedding model (or a deterministic stub) to form a retriever. + /// + /// For Beginners: Use this when you want to pick the search algorithm and distance measure + /// yourself instead of supplying a prebuilt document store. + /// + /// + public IAiModelBuilder ConfigureVectorIndex( + VectorIndexKind indexKind = VectorIndexKind.Flat, + int vectorDimension = 0, + RetrievalAugmentedGeneration.VectorSearch.ISimilarityMetric? metric = null, + int defaultTopK = 5) + { + var effectiveMetric = metric + ?? _configuredSimilarityMetric + ?? new RetrievalAugmentedGeneration.VectorSearch.Metrics.CosineSimilarityMetric(); + + // Record the metric so downstream construction and evaluation see the same choice. + _configuredSimilarityMetric ??= effectiveMetric; + + var index = BuildVectorIndex(indexKind, effectiveMetric); + var store = new VectorIndexDocumentStore(index, vectorDimension); + _ragDocumentStore = store; + + var embeddingModel = _configuredEmbeddingModel ?? BuildDefaultEmbeddingModel(vectorDimension); + _ragRetriever = new RetrievalAugmentedGeneration.Retrievers.VectorRetriever(store, embeddingModel, defaultTopK); + return this; + } + + /// + /// Constructs the selected in-memory vector index using the supplied similarity metric. + /// + private static RetrievalAugmentedGeneration.VectorSearch.Indexes.IVectorIndex BuildVectorIndex( + VectorIndexKind kind, + RetrievalAugmentedGeneration.VectorSearch.ISimilarityMetric metric) + { + switch (kind) + { + case VectorIndexKind.HNSW: + return new RetrievalAugmentedGeneration.VectorSearch.Indexes.HNSWIndex(metric); + case VectorIndexKind.IVF: + return new RetrievalAugmentedGeneration.VectorSearch.Indexes.IVFIndex(metric); + case VectorIndexKind.LSH: + return new RetrievalAugmentedGeneration.VectorSearch.Indexes.LSHIndex(metric); + case VectorIndexKind.Flat: + default: + return new RetrievalAugmentedGeneration.VectorSearch.Indexes.FlatIndex(metric); + } + } + + /// + /// Materializes a declarative + /// (typically produced by ) + /// into the builder's RAG fields: chunking strategy, embedding model, document store + retriever, + /// reranker, and context compressor. + /// + /// The RAG configuration to apply. Must not be null. + /// This builder instance for method chaining. + /// + /// + /// This reconciles the previously orphaned RAGConfiguration/RAGConfigurationBuilder types with + /// the facade. Strategy names are mapped to concrete components that require no external resources; entries + /// that would need API keys or model files (e.g., hosted embedding providers, cross-encoder rerankers) fall + /// back to safe deterministic defaults so the pipeline is fully wired and runnable from code. + /// + /// For Beginners: If you built a RAGConfiguration describing your pipeline in plain + /// strings/numbers, this turns that description into real components on the builder. + /// + /// + public IAiModelBuilder ConfigureRAG(RetrievalAugmentedGeneration.Configuration.RAGConfiguration config) + { + if (config == null) throw new ArgumentNullException(nameof(config)); + + // 1. Chunking strategy (non-generic; no external dependencies). + _chunkingStrategy = BuildChunkingStrategy(config.Chunking); + + // 2. Embedding model. Hosted providers need credentials, so for the pure code path we use the + // already-configured model when present, otherwise a deterministic stub of the requested size. + int embeddingDimension = config.Embedding.EmbeddingDimension > 0 ? config.Embedding.EmbeddingDimension : 768; + _configuredEmbeddingModel ??= BuildDefaultEmbeddingModel(embeddingDimension); + + // 3. Document store + retriever. Build a flat in-memory index over the configured metric and pair it + // with the embedding model, honoring the requested retrieval top-K. + var metric = _configuredSimilarityMetric + ?? new RetrievalAugmentedGeneration.VectorSearch.Metrics.CosineSimilarityMetric(); + _configuredSimilarityMetric ??= metric; + var index = BuildVectorIndex(VectorIndexKind.Flat, metric); + var store = new VectorIndexDocumentStore(index, embeddingDimension); + _ragDocumentStore = store; + int retrievalTopK = config.Retrieval.TopK > 0 ? config.Retrieval.TopK : 5; + _ragRetriever = new RetrievalAugmentedGeneration.Retrievers.VectorRetriever(store, _configuredEmbeddingModel, retrievalTopK); + + // 4. Reranking (optional). + _ragReranker = config.Reranking.Enabled ? BuildReranker(config.Reranking) : null; + + // 5. Context compression (optional). + _contextCompressor = config.ContextCompression.Enabled ? BuildContextCompressor(config.ContextCompression) : null; + + return this; + } + + private static IChunkingStrategy BuildChunkingStrategy(RetrievalAugmentedGeneration.Configuration.ChunkingConfig chunking) + { + string strategy = (chunking.Strategy ?? string.Empty).ToLowerInvariant(); + int chunkSize = chunking.ChunkSize > 0 ? chunking.ChunkSize : 1000; + int overlap = chunking.ChunkOverlap >= 0 ? chunking.ChunkOverlap : 0; + + if (strategy.Contains("recursive")) + return new RetrievalAugmentedGeneration.ChunkingStrategies.RecursiveCharacterChunkingStrategy(chunkSize, overlap); + if (strategy.Contains("sentence")) + return new RetrievalAugmentedGeneration.ChunkingStrategies.SentenceChunkingStrategy(chunkSize); + if (strategy.Contains("sliding") || strategy.Contains("window")) + return new RetrievalAugmentedGeneration.ChunkingStrategies.SlidingWindowChunkingStrategy(chunkSize, Math.Max(1, chunkSize - overlap)); + + // Default / "fixed": fixed-size character chunking. + return new RetrievalAugmentedGeneration.ChunkingStrategies.FixedSizeChunkingStrategy(chunkSize, overlap); + } + + private static IReranker BuildReranker(RetrievalAugmentedGeneration.Configuration.RerankingConfig reranking) + { + string strategy = (reranking.Strategy ?? string.Empty).ToLowerInvariant(); + + if (strategy.Contains("diversity")) + return new RetrievalAugmentedGeneration.Rerankers.DiversityReranker(); + if (strategy.Contains("middle") || strategy.Contains("lost")) + return new RetrievalAugmentedGeneration.RerankingStrategies.LostInTheMiddleReranker(); + if (strategy.Contains("rrf") || strategy.Contains("reciprocal") || strategy.Contains("fusion")) + return new RetrievalAugmentedGeneration.RerankingStrategies.ReciprocalRankFusion(); + + // Cross-encoder / Cohere / LLM rerankers require external resources; default to identity. + return new RetrievalAugmentedGeneration.Rerankers.IdentityReranker(); + } + + private static IContextCompressor BuildContextCompressor(RetrievalAugmentedGeneration.Configuration.ContextCompressionConfig compression) + { + int maxLength = compression.MaxLength > 0 ? compression.MaxLength : 500; + double ratio = compression.CompressionRatio > 0 && compression.CompressionRatio <= 1 ? compression.CompressionRatio : 0.5; + + // Selective / LLM / summarizer variants need embeddings or external models; AutoCompressor is a + // dependency-free extractive compressor honoring the requested length and ratio. + return new RetrievalAugmentedGeneration.ContextCompression.AutoCompressor(maxLength, ratio); + } + /// /// Configures advanced knowledge graph capabilities including embeddings, community detection, /// link prediction, temporal queries, and KG construction. diff --git a/src/AiModelBuilder.cs b/src/AiModelBuilder.cs index 4545c7dcd5..2aa7dfd803 100644 --- a/src/AiModelBuilder.cs +++ b/src/AiModelBuilder.cs @@ -230,6 +230,9 @@ private readonly AiDotNet.Configuration.IAiModelObservability _observability private IReranker? _ragReranker; private IGenerator? _ragGenerator; private IEnumerable? _queryProcessors; + private IChunkingStrategy? _chunkingStrategy; + private IContextCompressor? _contextCompressor; + private IDocumentStore? _ragDocumentStore; // Graph RAG components for knowledge graph-enhanced retrieval private KnowledgeGraph? _knowledgeGraph; diff --git a/src/Enums/VectorIndexKind.cs b/src/Enums/VectorIndexKind.cs new file mode 100644 index 0000000000..44eabd5305 --- /dev/null +++ b/src/Enums/VectorIndexKind.cs @@ -0,0 +1,30 @@ +namespace AiDotNet.Enums; + +/// +/// Selects the in-memory vector-search index used to back a RAG document store. +/// +/// +/// For Beginners: A vector index is the data structure that makes "find the most +/// similar vectors" fast. Different indexes trade accuracy for speed: +/// +/// Flat - exact brute-force search; most accurate, slowest on large sets. +/// HNSW - graph-based approximate search; excellent speed/recall balance. +/// IVF - inverted-file clustering; fast on large collections. +/// LSH - locality-sensitive hashing; very fast, lower recall. +/// +/// +/// +public enum VectorIndexKind +{ + /// Exact brute-force flat index (most accurate). + Flat = 0, + + /// Hierarchical Navigable Small World graph index (fast, high recall). + HNSW = 1, + + /// Inverted File index with clustering (scales to large collections). + IVF = 2, + + /// Locality-Sensitive Hashing index (very fast, approximate). + LSH = 3, +} diff --git a/src/Interfaces/IAiModelBuilder.cs b/src/Interfaces/IAiModelBuilder.cs index 4a07642c22..e56d547779 100644 --- a/src/Interfaces/IAiModelBuilder.cs +++ b/src/Interfaces/IAiModelBuilder.cs @@ -687,7 +687,10 @@ IAiModelBuilder ConfigureUncertaintyQuantification( /// Optional query processors for improving search quality. /// Optional graph storage backend for Graph RAG (e.g., MemoryGraphStore, FileGraphStore). /// Optional pre-configured knowledge graph. If null but graphStore is provided, a new one is created. - /// Optional document store for hybrid vector + graph retrieval. + /// Optional document store. Folded into a hybrid retriever when a knowledge graph is + /// also present; otherwise (with no explicit ) a default dense vector retriever is built over it. + /// Optional chunking strategy used to split source documents into passages before indexing/embedding. + /// Optional context compressor used to shrink retrieved passages before they reach the generator. /// The builder instance for method chaining. IAiModelBuilder ConfigureRetrievalAugmentedGeneration( IRetriever? retriever = null, @@ -696,7 +699,9 @@ IAiModelBuilder ConfigureRetrievalAugmentedGeneration( IEnumerable? queryProcessors = null, IGraphStore? graphStore = null, KnowledgeGraph? knowledgeGraph = null, - IDocumentStore? documentStore = null); + IDocumentStore? documentStore = null, + IChunkingStrategy? chunkingStrategy = null, + IContextCompressor? contextCompressor = null); /// /// Configures advanced knowledge graph capabilities including embeddings, community detection, @@ -2394,4 +2399,35 @@ IAiModelBuilder ConfigureQueryStrategy( /// The builder instance for method chaining. IAiModelBuilder ConfigureSimilarityMetric(RetrievalAugmentedGeneration.VectorSearch.ISimilarityMetric metric); + /// + /// Configures a vector document store as the RAG backend, building a default dense retriever over it. + /// + /// The document store to retrieve from. + /// Default number of documents the retriever returns per query. + /// The builder instance for method chaining. + IAiModelBuilder ConfigureVectorStore(IDocumentStore store, int defaultTopK = 5); + + /// + /// Configures an in-memory vector index (Flat / HNSW / IVF / LSH) with a similarity metric and builds + /// a dense retriever over it. + /// + /// Which in-memory index to build. + /// The embedding dimension (0 = inferred on first add). + /// The similarity metric; when null the configured metric (or cosine) is used. + /// Default number of documents the retriever returns per query. + /// The builder instance for method chaining. + IAiModelBuilder ConfigureVectorIndex( + AiDotNet.Enums.VectorIndexKind indexKind = AiDotNet.Enums.VectorIndexKind.Flat, + int vectorDimension = 0, + RetrievalAugmentedGeneration.VectorSearch.ISimilarityMetric? metric = null, + int defaultTopK = 5); + + /// + /// Materializes a declarative RAG configuration into the builder's RAG components (chunking, embedding, + /// document store + retriever, reranking, context compression). + /// + /// The RAG configuration to apply. + /// The builder instance for method chaining. + IAiModelBuilder ConfigureRAG(RetrievalAugmentedGeneration.Configuration.RAGConfiguration config); + } diff --git a/src/Models/Options/AiModelResultOptions.cs b/src/Models/Options/AiModelResultOptions.cs index 5d408a98e5..8dff66e44c 100644 --- a/src/Models/Options/AiModelResultOptions.cs +++ b/src/Models/Options/AiModelResultOptions.cs @@ -392,6 +392,21 @@ public class AiModelResultOptions : ModelOptions /// public IEnumerable? QueryProcessors { get; set; } + /// + /// Gets or sets the chunking strategy used to split source documents into passages before embedding/indexing. + /// + public IChunkingStrategy? ChunkingStrategy { get; set; } + + /// + /// Gets or sets the context compressor used to shrink retrieved passages before generation. + /// + public IContextCompressor? ContextCompressor { get; set; } + + /// + /// Gets or sets the document store backing vector retrieval. + /// + public IDocumentStore? DocumentStore { get; set; } + // ============================================================================ // Graph RAG Properties // ============================================================================ diff --git a/src/Models/Results/AiModelResult.cs b/src/Models/Results/AiModelResult.cs index 1a1bc3a055..7f4190a17d 100644 --- a/src/Models/Results/AiModelResult.cs +++ b/src/Models/Results/AiModelResult.cs @@ -354,6 +354,31 @@ public partial class AiModelResult : IFullModelQuery processors for preprocessing queries, or null if not configured. internal IEnumerable? QueryProcessors { get; private set; } + /// + /// Gets the chunking strategy configured for the RAG pipeline, or null if none was configured. + /// + /// An used to split documents into passages, or null. + public IChunkingStrategy? ChunkingStrategy { get; private set; } + + /// + /// Gets the context compressor configured for the RAG pipeline, or null if none was configured. + /// + /// An used to shrink retrieved context, or null. + public IContextCompressor? ContextCompressor { get; private set; } + + /// + /// Gets the document store backing vector retrieval, or null if none was configured. + /// + /// An holding vectorized documents, or null. + /// + /// + /// This property is excluded from JSON serialization because it represents runtime storage + /// infrastructure that must be reconfigured when the model is loaded. + /// + /// + [JsonIgnore] + public IDocumentStore? DocumentStore { get; private set; } + /// /// Gets or sets the knowledge graph for graph-enhanced retrieval. /// @@ -1614,6 +1639,9 @@ internal AiModelResult(AiModelResultOptions options) RagReranker = options.RagReranker; RagGenerator = options.RagGenerator; QueryProcessors = options.QueryProcessors; + ChunkingStrategy = options.ChunkingStrategy; + ContextCompressor = options.ContextCompressor; + DocumentStore = options.DocumentStore; // Graph RAG KnowledgeGraph = options.KnowledgeGraph; @@ -3704,6 +3732,9 @@ public IFullModel WithParameters(Vector parameters) RagReranker = RagReranker, RagGenerator = RagGenerator, QueryProcessors = QueryProcessors, + ChunkingStrategy = ChunkingStrategy, + ContextCompressor = ContextCompressor, + DocumentStore = DocumentStore, LoRAConfiguration = LoRAConfiguration, CrossValidationResult = CrossValidationResult, AutoMLSummary = AutoMLSummary, @@ -5453,6 +5484,9 @@ public IFullModel DeepCopy() RagReranker = RagReranker, RagGenerator = RagGenerator, QueryProcessors = QueryProcessors, + ChunkingStrategy = ChunkingStrategy, + ContextCompressor = ContextCompressor, + DocumentStore = DocumentStore, LoRAConfiguration = LoRAConfiguration, CrossValidationResult = CrossValidationResult, AutoMLSummary = AutoMLSummary, @@ -5666,6 +5700,9 @@ public void Deserialize(byte[] data) RagReranker = deserializedObject.RagReranker; RagGenerator = deserializedObject.RagGenerator; QueryProcessors = deserializedObject.QueryProcessors; + ChunkingStrategy = deserializedObject.ChunkingStrategy; + ContextCompressor = deserializedObject.ContextCompressor; + DocumentStore = deserializedObject.DocumentStore; LoRAConfiguration = deserializedObject.LoRAConfiguration; CrossValidationResult = deserializedObject.CrossValidationResult; AutoMLSummary = deserializedObject.AutoMLSummary; diff --git a/tests/AiDotNet.Tests/IntegrationTests/RetrievalAugmentedGeneration/RAGFacadeWiringTests.cs b/tests/AiDotNet.Tests/IntegrationTests/RetrievalAugmentedGeneration/RAGFacadeWiringTests.cs new file mode 100644 index 0000000000..81a00e1609 --- /dev/null +++ b/tests/AiDotNet.Tests/IntegrationTests/RetrievalAugmentedGeneration/RAGFacadeWiringTests.cs @@ -0,0 +1,164 @@ +using System; +using System.Threading.Tasks; +using AiDotNet; +using AiDotNet.Data.Loaders; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using AiDotNet.Regression; +using AiDotNet.RetrievalAugmentedGeneration.ChunkingStrategies; +using AiDotNet.RetrievalAugmentedGeneration.Configuration; +using AiDotNet.RetrievalAugmentedGeneration.ContextCompression; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Embeddings; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Metrics; +using AiDotNet.Tensors.LinearAlgebra; +using Xunit; + +namespace AiDotNet.Tests.IntegrationTests.RetrievalAugmentedGeneration; + +/// +/// Asserts that RAG components wired through the facade +/// are surfaced on the built AiModelResult. Covers tasks #29 (chunking/compression parameters and +/// the documentStore-dropped bug fix), #30 (ConfigureVectorStore / ConfigureVectorIndex), and #32 +/// (ConfigureRAG materialization of the previously orphaned RAGConfiguration). +/// +public class RAGFacadeWiringTests +{ + private static (Matrix X, Vector Y) BuildData(int rows = 40, int cols = 3) + { + var x = new Matrix(rows, cols); + var y = new Vector(rows); + for (int i = 0; i < rows; i++) + { + for (int j = 0; j < cols; j++) x[i, j] = Math.Sin((i + j) * 0.15) + (i * 0.01); + y[i] = Math.Sin((i + cols) * 0.15) + (i * 0.01); + } + + return (x, y); + } + + private static IAiModelBuilder, Vector> NewBuilder() + { + var (x, y) = BuildData(); + return new AiModelBuilder, Vector>() + .ConfigureModel(new MultipleRegression()) + .ConfigureDataLoader(new InMemoryDataLoader, Vector>(x, y)); + } + + [Fact(Timeout = 120000)] + public async Task ConfigureRAG_ChunkingAndCompression_AreSurfacedOnResult() + { + var chunking = new FixedSizeChunkingStrategy(chunkSize: 256, chunkOverlap: 32); + var compressor = new AutoCompressor(maxOutputLength: 200, compressionRatio: 0.5); + + var result = await NewBuilder() + .ConfigureRetrievalAugmentedGeneration(chunkingStrategy: chunking, contextCompressor: compressor) + .BuildAsync(); + + Assert.Same(chunking, result.ChunkingStrategy); + Assert.Same(compressor, result.ContextCompressor); + } + + [Fact(Timeout = 120000)] + public async Task ConfigureRAG_DocumentStoreWithoutKnowledgeGraph_BuildsRetriever() + { + // The documentStore-dropped bug: a store supplied without a knowledge graph used to be silently + // discarded. Now it must both surface as DocumentStore and drive a default vector retriever. + var store = new InMemoryDocumentStore(vectorDimension: 16); + + var result = await NewBuilder() + .ConfigureRetrievalAugmentedGeneration(documentStore: store) + .BuildAsync(); + + Assert.Same(store, result.DocumentStore); + Assert.NotNull(result.RagRetriever); + } + + [Fact(Timeout = 120000)] + public async Task ConfigureVectorStore_Alone_YieldsNonNullRetriever() + { + var store = new InMemoryDocumentStore(vectorDimension: 16); + + var result = await NewBuilder() + .ConfigureVectorStore(store) + .BuildAsync(); + + Assert.NotNull(result.RagRetriever); + Assert.Same(store, result.DocumentStore); + } + + [Fact(Timeout = 120000)] + public async Task ConfigureVectorStore_HonorsConfiguredEmbeddingModel() + { + var embedder = new StubEmbeddingModel(embeddingDimension: 16); + var store = new InMemoryDocumentStore(vectorDimension: 16); + + var result = await NewBuilder() + .ConfigureEmbeddingModel(embedder) + .ConfigureVectorStore(store) + .BuildAsync(); + + Assert.NotNull(result.RagRetriever); + Assert.Same(embedder, result.EmbeddingModel); + } + + [Theory(Timeout = 120000)] + [InlineData(VectorIndexKind.Flat)] + [InlineData(VectorIndexKind.HNSW)] + [InlineData(VectorIndexKind.IVF)] + [InlineData(VectorIndexKind.LSH)] + public async Task ConfigureVectorIndex_BuildsRetrieverAndStore(VectorIndexKind kind) + { + var result = await NewBuilder() + .ConfigureVectorIndex(kind, vectorDimension: 16, metric: new CosineSimilarityMetric()) + .BuildAsync(); + + Assert.NotNull(result.RagRetriever); + Assert.NotNull(result.DocumentStore); + } + + [Fact(Timeout = 120000)] + public async Task ConfigureRAG_MaterializesConfigurationIntoComponents() + { + var config = new RAGConfigurationBuilder() + .WithDocumentStore("InMemory") + .WithChunking("FixedSize", chunkSize: 512, chunkOverlap: 64) + .WithEmbedding("Stub", embeddingDimension: 32) + .WithRetrieval("Dense", topK: 7) + .WithReranking("Diversity", topK: 3) + .WithContextCompression("Auto", compressionRatio: 0.4, maxLength: 300) + .Build(); + + var result = await NewBuilder() + .ConfigureRAG(config) + .BuildAsync(); + + Assert.NotNull(result.ChunkingStrategy); + Assert.NotNull(result.EmbeddingModel); + Assert.NotNull(result.RagRetriever); + Assert.NotNull(result.DocumentStore); + Assert.NotNull(result.RagReranker); + Assert.NotNull(result.ContextCompressor); + } + + [Fact(Timeout = 120000)] + public async Task ConfigureRAG_WithDisabledOptionalStages_LeavesThemNull() + { + var config = new RAGConfigurationBuilder() + .WithDocumentStore("InMemory") + .WithChunking("Recursive") + .WithEmbedding("Stub", embeddingDimension: 24) + .WithRetrieval("Dense") + .Build(); + + var result = await NewBuilder() + .ConfigureRAG(config) + .BuildAsync(); + + Assert.NotNull(result.ChunkingStrategy); + Assert.NotNull(result.RagRetriever); + // Reranking and compression were not enabled on the config. + Assert.Null(result.RagReranker); + Assert.Null(result.ContextCompressor); + } +} From fee4a0df911f79a33dccbd08afb37cf2271d8481 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 18:45:07 -0400 Subject: [PATCH 16/25] feat(rag): yAML composite RAG section wires all sub-components MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The generated YAML applier previously treated ConfigureRetrievalAugmentedGeneration as a plain interface section: it kept only the retriever (the method's first parameter) and silently dropped reranker, generator, query processors, document store and graph store; ConfigureKnowledgeGraph was emitted as an Action<> TODO. Rework the YamlConfigSourceGenerator so the RAG section is a COMPOSITE: - Emit a new YamlRagSection POCO (retriever, reranker, generator, queryProcessors[], documentStore, graphStore, embeddingModel, similarityMetric, knowledgeGraph) and surface it as the RetrievalAugmentedGeneration YAML property. - Emit a dedicated applier that resolves each named sub-component from its own type-registry section and passes them to the facade: ConfigureRetrievalAugmentedGeneration (retriever/reranker/generator/queryProcessors/graphStore/documentStore), ConfigureEmbeddingModel, ConfigureSimilarityMetric, and ConfigureKnowledgeGraph. - Wire the standalone KnowledgeGraph section (Action) as an options params-bag, removing the KG TODO; pipeline-style Action<> sections (preprocessing/postprocessing) keep their existing steps shape. - Mark IQueryProcessor [YamlConfigurable("QueryProcessor")] so the registry exposes a QueryProcessor section for the composite applier. - Add matching JSON-schema/docs handling for the composite and options-bag sections. Adds RagYamlCompositeTests: a multi-component RAG YAML string loads, applies to a builder, and every sub-component (reranker/generator/queryProcessors/graphStore/ documentStore/similarityMetric/embeddingModel) plus ConfigureKnowledgeGraph is verified wired — nothing dropped. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../YamlConfigSourceGenerator.cs | 256 +++++++++++++++++- src/Configuration/YamlDocsGenerator.cs | 4 + src/Configuration/YamlJsonSchema.cs | 65 +++++ src/Interfaces/IQueryProcessor.cs | 1 + .../Configuration/RagYamlCompositeTests.cs | 251 +++++++++++++++++ 5 files changed, 566 insertions(+), 11 deletions(-) create mode 100644 tests/AiDotNet.Tests/IntegrationTests/Configuration/RagYamlCompositeTests.cs diff --git a/src/AiDotNet.Generators/YamlConfigSourceGenerator.cs b/src/AiDotNet.Generators/YamlConfigSourceGenerator.cs index b7e2166036..df5315eda2 100644 --- a/src/AiDotNet.Generators/YamlConfigSourceGenerator.cs +++ b/src/AiDotNet.Generators/YamlConfigSourceGenerator.cs @@ -27,6 +27,25 @@ namespace AiDotNet.Generators; [Generator] public class YamlConfigSourceGenerator : IIncrementalGenerator { + /// + /// The section name for the full RAG pipeline (produced by + /// AiModelBuilder.ConfigureRetrievalAugmentedGeneration(...)). This section is treated as a + /// COMPOSITE: instead of surfacing only the first parameter (the retriever), it exposes every RAG + /// sub-component (reranker, generator, query processors, document/graph stores, embedding model, + /// similarity metric, knowledge-graph options) and the applier wires each one to the matching facade + /// method. See and . + /// + private const string RagCompositeSectionName = "RetrievalAugmentedGeneration"; + + /// + /// The standalone advanced knowledge-graph section (produced by + /// AiModelBuilder.ConfigureKnowledgeGraph(Action<KnowledgeGraphOptions>)). Unlike the + /// pipeline-style Action<Pipeline> sections, this one is an options bag: the builder + /// creates the options object and invokes the supplied action, so the applier just applies the YAML + /// params onto it (no factory or ctor needed). This replaces the old Action-builder TODO. + /// + private const string KnowledgeGraphSectionName = "KnowledgeGraph"; + public void Initialize(IncrementalGeneratorInitializationContext context) { // Find class declarations named "AiModelBuilder" as a fast syntax filter. @@ -72,6 +91,7 @@ private static void Execute( // Emit helper types. EmitYamlTypeSection(context); EmitYamlPipelineSection(context); + EmitYamlRagSection(context); EmitYamlParamsHelper(context); // Emit the partial YamlModelConfig with generated properties. @@ -449,6 +469,83 @@ public class YamlPipelineSection context.AddSource("YamlPipelineSection.g.cs", source); } + private static void EmitYamlRagSection(SourceProductionContext context) + { + const string source = @"// +#nullable enable + +using System.Collections.Generic; + +namespace AiDotNet.Configuration; + +/// +/// Composite YAML section for the full retrieval-augmented generation (RAG) pipeline. +/// Each sub-component names a concrete implementation (via ) that the +/// applier resolves from the type registry and passes to the matching builder facade method. +/// +/// +/// For Beginners: A RAG pipeline is made of several parts working together. This section +/// lets you configure them all at once from YAML: +/// +/// retrievalAugmentedGeneration: +/// retriever: +/// type: ""BM25Retriever"" +/// reranker: +/// type: ""IdentityReranker"" +/// generator: +/// type: ""StubGenerator"" +/// queryProcessors: +/// - type: ""IdentityQueryProcessor"" +/// documentStore: +/// type: ""InMemoryDocumentStore"" +/// graphStore: +/// type: ""MemoryGraphStore"" +/// embeddingModel: +/// type: ""Word2Vec"" +/// similarityMetric: +/// type: ""CosineSimilarityMetric"" +/// knowledgeGraph: +/// params: +/// trainEmbeddings: true +/// +/// +/// +public class YamlRagSection +{ + /// The retriever that finds relevant documents (resolved from the ""Retriever"" registry). + public YamlTypeSection? Retriever { get; set; } + + /// The reranker that reorders retrieved documents (resolved from the ""Reranker"" registry). + public YamlTypeSection? Reranker { get; set; } + + /// The generator that produces grounded answers (resolved from the ""Generator"" registry). + public YamlTypeSection? Generator { get; set; } + + /// Ordered query processors applied before retrieval (resolved from the ""QueryProcessor"" registry). + public List QueryProcessors { get; set; } = new List(); + + /// The document store for vector/hybrid retrieval (resolved from the ""DocumentStore"" registry). + public YamlTypeSection? DocumentStore { get; set; } + + /// The graph storage backend for Graph RAG (resolved from the ""GraphStore"" registry). + public YamlTypeSection? GraphStore { get; set; } + + /// The embedding model used to vectorize text (resolved from the ""EmbeddingModel"" registry). + public YamlTypeSection? EmbeddingModel { get; set; } + + /// The similarity metric used for vector search (resolved from the ""SimilarityMetric"" registry). + public YamlTypeSection? SimilarityMetric { get; set; } + + /// + /// Advanced knowledge-graph options. When present, the applier calls + /// ConfigureKnowledgeGraph(...) and applies params onto the KnowledgeGraphOptions. + /// + public YamlTypeSection? KnowledgeGraph { get; set; } +} +"; + context.AddSource("YamlRagSection.g.cs", source); + } + private static void EmitYamlParamsHelper(SourceProductionContext context) { const string source = @"// @@ -623,6 +720,29 @@ private static void EmitYamlConfigApplier(SourceProductionContext context, List< var propName = ToPascalCase(section.Method.SectionName); + // Composite RAG pipeline: wire EVERY sub-component (retriever, reranker, generator, query + // processors, document/graph stores, embedding model, similarity metric, knowledge graph) + // rather than dropping all but the retriever. See EmitRagCompositeApplier. + if (section.Method.SectionName == RagCompositeSectionName) + { + EmitRagCompositeApplier(sb, propName); + continue; + } + + // Standalone advanced knowledge-graph options: apply YAML params onto the options object the + // builder creates and passes to the Action. Removes the old KG TODO. + if (section.Method.SectionName == KnowledgeGraphSectionName && + section.Method.Category == SectionCategory.ActionBuilder) + { + sb.AppendLine($" if (config.{propName} is not null)"); + sb.AppendLine($" {{"); + sb.AppendLine($" var __kgParams = config.{propName}.Params;"); + sb.AppendLine($" builder.{section.Method.MethodName}(__kgOpts => YamlParamsHelper.ApplyParams(__kgOpts!, __kgParams));"); + sb.AppendLine($" }}"); + sb.AppendLine(); + continue; + } + var hasGenericParams = ContainsTypeParameters(section.Method.ParameterType); // Merged POCO+interface section (e.g. AutoML): a YAML `type:` selects a concrete @@ -734,6 +854,105 @@ private static void EmitYamlConfigApplier(SourceProductionContext context, List< context.AddSource("YamlConfigApplier.g.cs", sb.ToString()); } + /// + /// Emits the applier body for the composite RAG section. Every sub-component named in YAML is + /// resolved from its own type-registry section and passed to the matching builder facade — + /// ConfigureRetrievalAugmentedGeneration (retriever, reranker, generator, query processors, + /// graph store, document store), ConfigureEmbeddingModel, ConfigureSimilarityMetric, + /// and ConfigureKnowledgeGraph — so nothing is dropped. + /// + private static void EmitRagCompositeApplier(StringBuilder sb, string propName) + { + const string retrieverType = "global::AiDotNet.Interfaces.IRetriever"; + const string rerankerType = "global::AiDotNet.Interfaces.IReranker"; + const string generatorType = "global::AiDotNet.Interfaces.IGenerator"; + const string queryProcessorType = "global::AiDotNet.Interfaces.IQueryProcessor"; + const string graphStoreType = "global::AiDotNet.Interfaces.IGraphStore"; + const string documentStoreType = "global::AiDotNet.Interfaces.IDocumentStore"; + const string embeddingType = "global::AiDotNet.Interfaces.IEmbeddingModel"; + const string similarityType = "global::AiDotNet.RetrievalAugmentedGeneration.VectorSearch.ISimilarityMetric"; + const string registry = "YamlTypeRegistry"; + + sb.AppendLine($" if (config.{propName} is not null)"); + sb.AppendLine($" {{"); + sb.AppendLine($" var __rag = config.{propName};"); + sb.AppendLine(); + + // Single-instance interface sub-components resolved from their registry sections. + EmitRagOptionalComponent(sb, retrieverType, "__retriever", "Retriever", "Retriever"); + EmitRagOptionalComponent(sb, rerankerType, "__reranker", "Reranker", "Reranker"); + EmitRagOptionalComponent(sb, generatorType, "__generator", "Generator", "Generator"); + EmitRagOptionalComponent(sb, graphStoreType, "__graphStore", "GraphStore", "GraphStore"); + EmitRagOptionalComponent(sb, documentStoreType, "__documentStore", "DocumentStore", "DocumentStore"); + + // Ordered list of query processors. + sb.AppendLine($" System.Collections.Generic.List<{queryProcessorType}>? __queryProcessors = null;"); + sb.AppendLine($" if (__rag.QueryProcessors is not null && __rag.QueryProcessors.Count > 0)"); + sb.AppendLine($" {{"); + sb.AppendLine($" __queryProcessors = new System.Collections.Generic.List<{queryProcessorType}>();"); + sb.AppendLine($" foreach (var __qp in __rag.QueryProcessors)"); + sb.AppendLine($" {{"); + sb.AppendLine($" if (__qp is not null && !string.IsNullOrWhiteSpace(__qp.Type))"); + sb.AppendLine($" {{"); + sb.AppendLine($" __queryProcessors.Add({registry}.CreateInstance<{queryProcessorType}>("); + sb.AppendLine($" \"QueryProcessor\", __qp.Type, __qp.Params));"); + sb.AppendLine($" }}"); + sb.AppendLine($" }}"); + sb.AppendLine($" }}"); + sb.AppendLine(); + + // Only call the RAG facade when at least one component is present, so an empty section does not + // clear previously-configured RAG state. + sb.AppendLine($" if (__retriever is not null || __reranker is not null || __generator is not null"); + sb.AppendLine($" || __queryProcessors is not null || __graphStore is not null || __documentStore is not null)"); + sb.AppendLine($" {{"); + sb.AppendLine($" builder.ConfigureRetrievalAugmentedGeneration("); + sb.AppendLine($" retriever: __retriever,"); + sb.AppendLine($" reranker: __reranker,"); + sb.AppendLine($" generator: __generator,"); + sb.AppendLine($" queryProcessors: __queryProcessors,"); + sb.AppendLine($" graphStore: __graphStore,"); + sb.AppendLine($" documentStore: __documentStore);"); + sb.AppendLine($" }}"); + sb.AppendLine(); + + // Embedding model and similarity metric have their own facade methods. + sb.AppendLine($" if (__rag.EmbeddingModel is not null && !string.IsNullOrWhiteSpace(__rag.EmbeddingModel.Type))"); + sb.AppendLine($" {{"); + sb.AppendLine($" builder.ConfigureEmbeddingModel({registry}.CreateInstance<{embeddingType}>("); + sb.AppendLine($" \"EmbeddingModel\", __rag.EmbeddingModel.Type, __rag.EmbeddingModel.Params));"); + sb.AppendLine($" }}"); + sb.AppendLine(); + sb.AppendLine($" if (__rag.SimilarityMetric is not null && !string.IsNullOrWhiteSpace(__rag.SimilarityMetric.Type))"); + sb.AppendLine($" {{"); + sb.AppendLine($" builder.ConfigureSimilarityMetric({registry}.CreateInstance<{similarityType}>("); + sb.AppendLine($" \"SimilarityMetric\", __rag.SimilarityMetric.Type, __rag.SimilarityMetric.Params));"); + sb.AppendLine($" }}"); + sb.AppendLine(); + + // Advanced knowledge-graph options: the builder creates the options object and invokes our action. + sb.AppendLine($" if (__rag.KnowledgeGraph is not null)"); + sb.AppendLine($" {{"); + sb.AppendLine($" var __kgParams = __rag.KnowledgeGraph.Params;"); + sb.AppendLine($" builder.ConfigureKnowledgeGraph(__kgOpts => YamlParamsHelper.ApplyParams(__kgOpts!, __kgParams));"); + sb.AppendLine($" }}"); + sb.AppendLine($" }}"); + sb.AppendLine(); + } + + private static void EmitRagOptionalComponent( + StringBuilder sb, string interfaceType, string localName, string yamlProperty, string registrySection) + { + const string registry = "YamlTypeRegistry"; + sb.AppendLine($" {interfaceType}? {localName} = null;"); + sb.AppendLine($" if (__rag.{yamlProperty} is not null && !string.IsNullOrWhiteSpace(__rag.{yamlProperty}.Type))"); + sb.AppendLine($" {{"); + sb.AppendLine($" {localName} = {registry}.CreateInstance<{interfaceType}>("); + sb.AppendLine($" \"{registrySection}\", __rag.{yamlProperty}.Type, __rag.{yamlProperty}.Params);"); + sb.AppendLine($" }}"); + sb.AppendLine(); + } + private static void EmitYamlTypeRegistry(SourceProductionContext context, List sections) { // Include both interface sections and abstract POCO sections (generic or non-generic) that have implementations. @@ -1057,17 +1276,20 @@ private static void EmitYamlSchemaMetadata(SourceProductionContext context, List if (section.Method.SectionName == "Model" && section.Method.ParameterTypeName.Contains("IFullModel")) continue; - var category = section.Method.Category switch - { - SectionCategory.Interface => "Interface", - SectionCategory.PocoConfig when section.Method.IsAttributeDiscovered && section.ConcreteImplementations.Count > 0 => "Interface", - SectionCategory.PocoConfig when section.Method.IsAbstract && !ContainsTypeParameters(section.Method.ParameterType) => "AbstractNonGeneric", - SectionCategory.PocoConfig when section.Method.IsAbstract && ContainsTypeParameters(section.Method.ParameterType) => "AbstractGeneric", - SectionCategory.PocoConfig when ContainsTypeParameters(section.Method.ParameterType) => "ConcreteGeneric", - SectionCategory.PocoConfig => "Poco", - SectionCategory.ActionBuilder => "Pipeline", - _ => "Unknown", - }; + var category = + section.Method.SectionName == RagCompositeSectionName ? "RagComposite" : + section.Method.SectionName == KnowledgeGraphSectionName ? "ActionParams" : + section.Method.Category switch + { + SectionCategory.Interface => "Interface", + SectionCategory.PocoConfig when section.Method.IsAttributeDiscovered && section.ConcreteImplementations.Count > 0 => "Interface", + SectionCategory.PocoConfig when section.Method.IsAbstract && !ContainsTypeParameters(section.Method.ParameterType) => "AbstractNonGeneric", + SectionCategory.PocoConfig when section.Method.IsAbstract && ContainsTypeParameters(section.Method.ParameterType) => "AbstractGeneric", + SectionCategory.PocoConfig when ContainsTypeParameters(section.Method.ParameterType) => "ConcreteGeneric", + SectionCategory.PocoConfig => "Poco", + SectionCategory.ActionBuilder => "Pipeline", + _ => "Unknown", + }; var propName = ToPascalCase(section.Method.SectionName); var isExisting = existingSections.Contains(section.Method.SectionName); @@ -1276,6 +1498,18 @@ private static string GetYamlPropertyType(SectionInfo section) return "YamlTypeSection"; } + // The composite RAG pipeline surfaces every sub-component through YamlRagSection. + if (section.Method.SectionName == RagCompositeSectionName) + { + return "YamlRagSection"; + } + + // The standalone knowledge-graph section is an options bag applied via ConfigureKnowledgeGraph. + if (section.Method.SectionName == KnowledgeGraphSectionName) + { + return "YamlTypeSection"; + } + return section.Method.Category switch { SectionCategory.Interface => "YamlTypeSection", diff --git a/src/Configuration/YamlDocsGenerator.cs b/src/Configuration/YamlDocsGenerator.cs index 3bf92770c9..65042eb70b 100644 --- a/src/Configuration/YamlDocsGenerator.cs +++ b/src/Configuration/YamlDocsGenerator.cs @@ -141,6 +141,10 @@ public static string Generate() case "Pipeline": GeneratePipelineDocs(sb, section); break; + case "RagComposite": + case "ActionParams": + GenerateInterfaceSectionDocs(sb, section); + break; case "Poco" when !section.IsHandWritten: GeneratePocoFromMetadataDocs(sb, section); break; diff --git a/src/Configuration/YamlJsonSchema.cs b/src/Configuration/YamlJsonSchema.cs index d5063ab7be..49acc25594 100644 --- a/src/Configuration/YamlJsonSchema.cs +++ b/src/Configuration/YamlJsonSchema.cs @@ -155,7 +155,12 @@ private static void AddGeneratedSections(JObject properties) properties[yamlKey] = BuildPipelineSchema(section); break; + case "RagComposite": + properties[yamlKey] = BuildRagCompositeSchema(section); + break; + default: + // Covers "Interface"-shaped and "ActionParams" (type + params bag) sections. properties[yamlKey] = BuildTypeSectionSchema(section); break; } @@ -184,6 +189,66 @@ private static JObject BuildTypeSectionSchema(YamlSectionMeta section) return schema; } + /// + /// Builds the schema for the composite RAG section: each sub-component is a type+params object whose + /// type enum comes from that component's own registry section. + /// + private static JObject BuildRagCompositeSchema(YamlSectionMeta section) + { + JObject SubComponent(string registrySection) => new JObject + { + ["type"] = "object", + ["additionalProperties"] = false, + ["properties"] = new JObject + { + ["type"] = BuildTypePropertyWithEnum(registrySection), + ["params"] = new JObject + { + ["type"] = "object", + ["additionalProperties"] = true, + }, + }, + }; + + var props = new JObject + { + ["retriever"] = SubComponent("Retriever"), + ["reranker"] = SubComponent("Reranker"), + ["generator"] = SubComponent("Generator"), + ["queryProcessors"] = new JObject + { + ["type"] = "array", + ["items"] = SubComponent("QueryProcessor"), + }, + ["documentStore"] = SubComponent("DocumentStore"), + ["graphStore"] = SubComponent("GraphStore"), + ["embeddingModel"] = SubComponent("EmbeddingModel"), + ["similarityMetric"] = SubComponent("SimilarityMetric"), + ["knowledgeGraph"] = new JObject + { + ["type"] = "object", + ["description"] = "Advanced knowledge-graph options applied via ConfigureKnowledgeGraph.", + ["additionalProperties"] = false, + ["properties"] = new JObject + { + ["params"] = new JObject + { + ["type"] = "object", + ["additionalProperties"] = true, + }, + }, + }, + }; + + return new JObject + { + ["type"] = "object", + ["description"] = "Composite retrieval-augmented generation pipeline. Each sub-component names a concrete implementation.", + ["additionalProperties"] = false, + ["properties"] = props, + }; + } + private static JObject BuildTypePropertyWithEnum(string sectionName) { var typeProp = new JObject diff --git a/src/Interfaces/IQueryProcessor.cs b/src/Interfaces/IQueryProcessor.cs index 8dd9875b44..f1947aec88 100644 --- a/src/Interfaces/IQueryProcessor.cs +++ b/src/Interfaces/IQueryProcessor.cs @@ -21,6 +21,7 @@ namespace AiDotNet.Interfaces; /// - Keyword extraction: Focus on important terms ("What is the capital of France?" → "capital France") /// /// +[AiDotNet.Configuration.YamlConfigurable("QueryProcessor")] public interface IQueryProcessor { /// diff --git a/tests/AiDotNet.Tests/IntegrationTests/Configuration/RagYamlCompositeTests.cs b/tests/AiDotNet.Tests/IntegrationTests/Configuration/RagYamlCompositeTests.cs new file mode 100644 index 0000000000..b0a0967c80 --- /dev/null +++ b/tests/AiDotNet.Tests/IntegrationTests/Configuration/RagYamlCompositeTests.cs @@ -0,0 +1,251 @@ +using System; +using System.Collections; +using System.Collections.Generic; +using System.Linq; +using System.Reflection; +using System.Text; +using Xunit; +using AiDotNet.Configuration; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.Graph; + +namespace AiDotNet.Tests.IntegrationTests.Configuration; + +/// +/// Verifies the composite retrievalAugmentedGeneration: YAML section wires the FULL RAG pipeline +/// through the source-generated applier: retriever, reranker, generator, query processors, document +/// store, graph store, embedding model, similarity metric, and knowledge-graph options are all applied +/// to the builder rather than only the retriever being kept and everything else dropped. +/// +public class RagYamlCompositeTests +{ + private static IReadOnlyDictionary> Registries => + YamlTypeRegistry, Vector>.GetAllRegistries(); + + /// + /// Finds the first registered type name in that both instantiates + /// without throwing and is assignable to . Keeps the test robust as the + /// concrete implementation set evolves, and avoids heavy/failing constructors. + /// + private static string? PickInstantiable(string section, Type expected, int maxAttempts = 25) + { + if (!Registries.TryGetValue(section, out var types)) return null; + + var attempts = 0; + foreach (var name in types.Keys) + { + if (attempts++ >= maxAttempts) break; + try + { + var instance = YamlTypeRegistry, Vector> + .CreateInstance(section, name); + if (instance is not null && expected.IsInstanceOfType(instance)) + { + return name; + } + } + catch + { + // Skip types whose constructors need resources we can't supply in a unit test. + } + } + + return null; + } + + private static object? GetPrivateField(object target, string fieldName) + { + var field = target.GetType().GetField(fieldName, BindingFlags.Instance | BindingFlags.NonPublic); + Assert.True(field is not null, $"Expected private field '{fieldName}' on {target.GetType().Name}."); + return field!.GetValue(target); + } + + /// + /// Parsing a composite RAG section must retain every sub-component (nothing dropped at deserialize). + /// + [Fact] + public void LoadFromString_CompositeRagSection_RetainsAllSubComponents() + { + var yaml = @" +retrievalAugmentedGeneration: + retriever: + type: BM25Retriever + reranker: + type: IdentityReranker + generator: + type: StubGenerator + queryProcessors: + - type: IdentityQueryProcessor + documentStore: + type: InMemoryDocumentStore + graphStore: + type: MemoryGraphStore + embeddingModel: + type: Word2Vec + similarityMetric: + type: CosineSimilarityMetric + knowledgeGraph: + params: + trainEmbeddings: true +"; + + var config = YamlConfigLoader.LoadFromString(yaml); + + Assert.NotNull(config.RetrievalAugmentedGeneration); + var rag = config.RetrievalAugmentedGeneration!; + Assert.Equal("BM25Retriever", rag.Retriever!.Type); + Assert.Equal("IdentityReranker", rag.Reranker!.Type); + Assert.Equal("StubGenerator", rag.Generator!.Type); + Assert.Single(rag.QueryProcessors); + Assert.Equal("IdentityQueryProcessor", rag.QueryProcessors[0].Type); + Assert.Equal("InMemoryDocumentStore", rag.DocumentStore!.Type); + Assert.Equal("MemoryGraphStore", rag.GraphStore!.Type); + Assert.Equal("Word2Vec", rag.EmbeddingModel!.Type); + Assert.Equal("CosineSimilarityMetric", rag.SimilarityMetric!.Type); + Assert.NotNull(rag.KnowledgeGraph); + Assert.True(rag.KnowledgeGraph!.Params.ContainsKey("trainEmbeddings")); + } + + /// + /// The generated YamlRagSection exposes a property for every RAG sub-component — proof the + /// generator produced a composite section rather than a single-type retriever-only section. + /// + [Fact] + public void YamlModelConfig_RetrievalAugmentedGeneration_IsCompositeSection() + { + var prop = typeof(YamlModelConfig).GetProperty("RetrievalAugmentedGeneration"); + Assert.NotNull(prop); + Assert.Equal("YamlRagSection", prop!.PropertyType.Name); + + var ragType = prop.PropertyType; + foreach (var expectedMember in new[] + { + "Retriever", "Reranker", "Generator", "QueryProcessors", + "DocumentStore", "GraphStore", "EmbeddingModel", "SimilarityMetric", "KnowledgeGraph", + }) + { + Assert.True(ragType.GetProperty(expectedMember) is not null, + $"YamlRagSection is missing composite sub-component '{expectedMember}'."); + } + } + + /// + /// Applying the composite section wires EVERY sub-component to the builder — retriever, reranker, + /// generator, query processors, graph store (+ derived knowledge graph and hybrid retriever), + /// embedding model, similarity metric — and calls ConfigureKnowledgeGraph. Type names are chosen + /// dynamically from the registry so the test tracks the real implementation set. + /// + [Fact] + public void Apply_CompositeRagSection_WiresAllSubComponentsOntoBuilder() + { + var retriever = PickInstantiable("Retriever", typeof(IRetriever)); + var reranker = PickInstantiable("Reranker", typeof(IReranker)); + var generator = PickInstantiable("Generator", typeof(IGenerator)); + var queryProcessor = PickInstantiable("QueryProcessor", typeof(IQueryProcessor)); + var graphStore = PickInstantiable("GraphStore", typeof(IGraphStore)); + var documentStore = PickInstantiable("DocumentStore", typeof(IDocumentStore)); + var embeddingModel = PickInstantiable("EmbeddingModel", typeof(IEmbeddingModel)); + var similarityMetric = PickInstantiable( + "SimilarityMetric", + typeof(AiDotNet.RetrievalAugmentedGeneration.VectorSearch.ISimilarityMetric)); + + // The components that prove the OLD "keep only the retriever, drop everything else" bug is fixed + // must be registry-constructible. (Retrievers require an injected IDocumentStore constructor arg + // that the registry cannot synthesize, so retriever wiring is exercised opportunistically only — + // and retriever was never the dropped component; it was the sole survivor of the old bug.) + Assert.NotNull(reranker); + Assert.NotNull(generator); + Assert.NotNull(queryProcessor); + Assert.NotNull(graphStore); + Assert.NotNull(documentStore); + Assert.NotNull(similarityMetric); + + var yaml = new StringBuilder(); + yaml.AppendLine("retrievalAugmentedGeneration:"); + if (retriever is not null) + { + yaml.AppendLine(" retriever:"); + yaml.AppendLine($" type: {retriever}"); + } + yaml.AppendLine(" reranker:"); + yaml.AppendLine($" type: {reranker}"); + yaml.AppendLine(" generator:"); + yaml.AppendLine($" type: {generator}"); + yaml.AppendLine(" queryProcessors:"); + yaml.AppendLine($" - type: {queryProcessor}"); + yaml.AppendLine(" graphStore:"); + yaml.AppendLine($" type: {graphStore}"); + yaml.AppendLine(" documentStore:"); + yaml.AppendLine($" type: {documentStore}"); + if (embeddingModel is not null) + { + yaml.AppendLine(" embeddingModel:"); + yaml.AppendLine($" type: {embeddingModel}"); + } + yaml.AppendLine(" similarityMetric:"); + yaml.AppendLine($" type: {similarityMetric}"); + yaml.AppendLine(" knowledgeGraph:"); + yaml.AppendLine(" params:"); + yaml.AppendLine(" trainEmbeddings: true"); + + var config = YamlConfigLoader.LoadFromString(yaml.ToString()); + var builder = new AiModelBuilder, Vector>(); + + var exception = Record.Exception(() => + YamlConfigApplier, Vector>.Apply(config, builder)); + Assert.Null(exception); + + // Standard RAG components that were previously dropped are now wired. + Assert.NotNull(GetPrivateField(builder, "_ragReranker")); + Assert.NotNull(GetPrivateField(builder, "_ragGenerator")); + + var queryProcessors = GetPrivateField(builder, "_queryProcessors") as IEnumerable; + Assert.NotNull(queryProcessors); + Assert.True(queryProcessors!.Cast().Any(), "_queryProcessors should contain at least one processor."); + + // Graph RAG components derived from graphStore + documentStore (documentStore was previously dropped). + Assert.NotNull(GetPrivateField(builder, "_graphStore")); + Assert.NotNull(GetPrivateField(builder, "_knowledgeGraph")); + Assert.NotNull(GetPrivateField(builder, "_hybridGraphRetriever")); + + // Similarity metric wired via ConfigureSimilarityMetric (previously dropped). + Assert.NotNull(GetPrivateField(builder, "_configuredSimilarityMetric")); + + // Retriever wired when a registry-constructible implementation was available. + if (retriever is not null) + { + Assert.NotNull(GetPrivateField(builder, "_ragRetriever")); + } + + // Embedding model wired via ConfigureEmbeddingModel (when a lightweight impl was resolvable). + if (embeddingModel is not null) + { + Assert.NotNull(GetPrivateField(builder, "_configuredEmbeddingModel")); + } + + // ConfigureKnowledgeGraph applied with the supplied params (previously a TODO / dropped). + var kgOptions = GetPrivateField(builder, "_knowledgeGraphOptions") as KnowledgeGraphOptions; + Assert.NotNull(kgOptions); + Assert.True(kgOptions!.TrainEmbeddings == true, + "ConfigureKnowledgeGraph should have applied trainEmbeddings=true from the composite section."); + } + + /// + /// The registry must expose a section for every RAG sub-component the composite applier resolves, + /// including the newly-marked QueryProcessor section. Without these the composite wiring could not + /// instantiate the named types. + /// + [Fact] + public void TypeRegistry_HasSectionsForAllRagSubComponents() + { + foreach (var section in new[] + { + "Retriever", "Reranker", "Generator", "QueryProcessor", + "DocumentStore", "GraphStore", "EmbeddingModel", "SimilarityMetric", + }) + { + Assert.True(Registries.ContainsKey(section), $"Type registry is missing RAG section '{section}'."); + Assert.True(Registries[section].Count > 0, $"RAG section '{section}' has zero implementations."); + } + } +} From ad406efbcdd4b264563952a77f03b6386913b480 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 18:54:10 -0400 Subject: [PATCH 17/25] test(builder): configureAutoML(IAutoMLModel) belongs on interface and concrete Pre-existing failure on this branch: the test asserted ConfigureAutoML(IAutoMLModel<...>) must NOT be on IAiModelBuilder, but every fluent Configure* method belongs on the interface (it's declared there at IAiModelBuilder.cs:1507). Flip the assertion to require the overload on BOTH the interface and the concrete builder, and rename the test to match (same fix applied on the #1789 branch). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../Configuration/SourceGeneratorCoverageTests.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/AiDotNet.Tests/IntegrationTests/Configuration/SourceGeneratorCoverageTests.cs b/tests/AiDotNet.Tests/IntegrationTests/Configuration/SourceGeneratorCoverageTests.cs index ac0035ec9f..5e4a1d2410 100644 --- a/tests/AiDotNet.Tests/IntegrationTests/Configuration/SourceGeneratorCoverageTests.cs +++ b/tests/AiDotNet.Tests/IntegrationTests/Configuration/SourceGeneratorCoverageTests.cs @@ -894,7 +894,7 @@ public void YamlRoundTrip_AutoMLMergedSection_ResolvesRegisteredEngine() /// what the generated YAML applier dispatches against. /// [Fact] - public void ConfigureAutoML_IAutoMLModelOverload_IsOnConcreteBuilderNotInterface() + public void ConfigureAutoML_IAutoMLModelOverload_IsOnBothInterfaceAndConcreteBuilder() { var interfaceType = typeof(IAiModelBuilder, Vector>); var concreteType = typeof(AiModelBuilder, Vector>); @@ -906,8 +906,8 @@ static bool HasIAutoMLModelOverload(Type t, Type paramType) => t && m.GetParameters().Length == 1 && m.GetParameters()[0].ParameterType == paramType); - Assert.False(HasIAutoMLModelOverload(interfaceType, autoMLModelType), - "ConfigureAutoML(IAutoMLModel<...>) must not be declared on IAiModelBuilder — it is a breaking change for external implementers."); + Assert.True(HasIAutoMLModelOverload(interfaceType, autoMLModelType), + "ConfigureAutoML(IAutoMLModel<...>) must be declared on IAiModelBuilder — every Configure* method belongs on the interface so it is usable through the abstraction."); Assert.True(HasIAutoMLModelOverload(concreteType, autoMLModelType), "ConfigureAutoML(IAutoMLModel<...>) must remain a public method on the concrete AiModelBuilder."); } From fe9cfa736a6c5b209b9e6647d0822ada8358dc81 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 18:56:35 -0400 Subject: [PATCH 18/25] refactor(rag): async + cancellation across the RAG pipeline Add async, cancellation-aware methods to the RAG interfaces and make the I/O-bound implementations genuinely non-blocking. Interfaces (kept sync methods, added async overloads): - IRetriever.RetrieveAsync(query, topK, filters, ct) - IDocumentStore: Add/AddBatch/GetSimilar/GetSimilarWithFilters/GetById/ Remove/Clear/GetAll -> *Async(ct) - IReranker.RerankAsync (all + topK overloads) - IContextCompressor.CompressAsync - IGenerator.GenerateAsync / GenerateGroundedAsync Base classes implement the async interface methods and expose a protected virtual *CoreAsync that defaults to Task.FromResult(syncCore), so CPU-bound subclasses need no changes; I/O-bound subclasses override the core for real async. All library async uses ConfigureAwait(false); net471-safe (Task-based). Truly-async (removed sync-over-async on the async path): the external vector stores Qdrant/Pinecone/Weaviate/Milvus/AzureSearch/ChromaDB and the Elasticsearch store (token threaded into the HTTP client), the LLM-backed ChatClientGenerator, and the VectorRetriever/DenseRetriever (await the embedding model + store). CPU-bound impls (in-memory stores, RRF/MMR/identity rerankers, extractive compressors, Stub/Neural generators) wrap sync. Updated all direct interface implementers (VectorIndexDocumentStore, Stub/NeuralGenerator) and 17 test mock classes. Added focused async + pre-cancelled-token tests (RagAsyncCancellationTests, 12 tests). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ElasticsearchDocumentStore.cs | 106 +++++++-- src/AiModelBuilder.VectorIndexStore.cs | 61 +++++ src/Interfaces/IContextCompressor.cs | 22 ++ src/Interfaces/IDocumentStore.cs | 49 ++++ src/Interfaces/IGenerator.cs | 30 +++ src/Interfaces/IReranker.cs | 28 +++ src/Interfaces/IRetriever.cs | 25 +++ .../ContextCompressorBase.cs | 44 ++++ .../AzureSearchDocumentStore.cs | 72 +++++- .../DocumentStores/ChromaDBDocumentStore.cs | 49 +++- .../DocumentStores/DocumentStoreBase.cs | 143 ++++++++++++ .../DocumentStores/MilvusDocumentStore.cs | 70 +++++- .../DocumentStores/PineconeDocumentStore.cs | 76 +++++-- .../DocumentStores/QdrantDocumentStore.cs | 68 +++++- .../DocumentStores/WeaviateDocumentStore.cs | 68 +++++- .../Generators/ChatClientGenerator.cs | 15 ++ .../Generators/GeneratorBase.cs | 91 ++++++++ .../Generators/NeuralGenerator.cs | 16 ++ .../Generators/StubGenerator.cs | 16 ++ .../Rerankers/RerankerBase.cs | 55 +++++ .../Retrievers/DenseRetriever.cs | 18 ++ .../Retrievers/RetrieverBase.cs | 48 ++++ .../Retrievers/VectorRetriever.cs | 23 ++ .../KnowledgeGraph/LlmGraphExtractionTests.cs | 3 + .../BM25RetrieverTests.cs | 9 + .../ColBERTRetrieverTests.cs | 9 + .../DeStubbedLlmComponentsTests.cs | 2 + .../DenseRetrieverTests.cs | 9 + .../FLARERetrieverTests.cs | 3 + .../GraphRAGTests.cs | 3 + .../GraphRetrieverTests.cs | 9 + .../HybridRetrieverTests.cs | 2 + .../MultiQueryRetrieverTests.cs | 2 + .../MultiVectorRetrieverTests.cs | 9 + .../ParentDocumentRetrieverTests.cs | 9 + .../RagAsyncCancellationTests.cs | 212 ++++++++++++++++++ .../RetrieverFusionAndRoutingTests.cs | 2 + .../Retrievers/TestHelpers.cs | 9 + .../SelfCorrectingRetrieverTests.cs | 3 + .../TFIDFRetrieverTests.cs | 9 + .../VectorRetrieverTests.cs | 9 + 41 files changed, 1424 insertions(+), 82 deletions(-) create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/RagAsyncCancellationTests.cs diff --git a/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchDocumentStore.cs b/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchDocumentStore.cs index 177cb4a6d7..75bbb9b343 100644 --- a/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchDocumentStore.cs +++ b/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchDocumentStore.cs @@ -3,6 +3,8 @@ using System.Linq; using System.Net.Http; using System.Text; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; @@ -118,6 +120,13 @@ private void UpdateDocumentCount() } protected override void AddCore(VectorDocument vectorDocument) + => AddCoreImplAsync(vectorDocument, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddCoreAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) + => AddCoreImplAsync(vectorDocument, cancellationToken); + + private async Task AddCoreImplAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) { if (vectorDocument.Embedding.Length != _vectorDimension) throw new ArgumentException($"Document embedding dimension ({vectorDocument.Embedding.Length}) does not match the store's configured dimension ({_vectorDimension})."); @@ -137,7 +146,7 @@ protected override void AddCore(VectorDocument vectorDocument) Encoding.UTF8, "application/json"); - using var response = _httpClient.PutAsync($"/{_indexName}/_doc/{vectorDocument.Document.Id}", content).GetAwaiter().GetResult(); + using var response = await _httpClient.PutAsync($"/{_indexName}/_doc/{vectorDocument.Document.Id}", content, cancellationToken).ConfigureAwait(false); response.EnsureSuccessStatusCode(); _cache[vectorDocument.Document.Id] = vectorDocument; @@ -145,6 +154,13 @@ protected override void AddCore(VectorDocument vectorDocument) } protected override void AddBatchCore(IList> vectorDocuments) + => AddBatchCoreImplAsync(vectorDocuments, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddBatchCoreAsync(IList> vectorDocuments, CancellationToken cancellationToken) + => AddBatchCoreImplAsync(vectorDocuments, cancellationToken); + + private async Task AddBatchCoreImplAsync(IList> vectorDocuments, CancellationToken cancellationToken) { if (vectorDocuments.Count == 0) return; @@ -172,10 +188,10 @@ protected override void AddBatchCore(IList> vectorDocuments) } using var content = new StringContent(bulkBody.ToString(), Encoding.UTF8, "application/x-ndjson"); - using var response = _httpClient.PostAsync($"/{_indexName}/_bulk", content).GetAwaiter().GetResult(); + using var response = await _httpClient.PostAsync($"/{_indexName}/_bulk", content, cancellationToken).ConfigureAwait(false); response.EnsureSuccessStatusCode(); - var responseContent = response.Content.ReadAsStringAsync().GetAwaiter().GetResult(); + var responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false); var result = Newtonsoft.Json.Linq.JObject.Parse(responseContent); if (result["errors"]?.Value() == true) @@ -214,6 +230,13 @@ protected override void AddBatchCore(IList> vectorDocuments) } protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); + + private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) { var embedding = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); @@ -267,14 +290,14 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector Encoding.UTF8, "application/json"); - using var response = _httpClient.PostAsync($"/{_indexName}/_search", content).GetAwaiter().GetResult(); + using var response = await _httpClient.PostAsync($"/{_indexName}/_search", content, cancellationToken).ConfigureAwait(false); if (!response.IsSuccessStatusCode) { - var errorContent = response.Content.ReadAsStringAsync().GetAwaiter().GetResult(); + var errorContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false); throw new HttpRequestException($"Elasticsearch search failed with status {response.StatusCode}: {errorContent}"); } - var responseContent = response.Content.ReadAsStringAsync().GetAwaiter().GetResult(); + var responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false); var result = JObject.Parse(responseContent); var results = new List>(); @@ -305,22 +328,29 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector } protected override Document? GetByIdCore(string documentId) + => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task?> GetByIdCoreAsync(string documentId, CancellationToken cancellationToken) + => GetByIdCoreImplAsync(documentId, cancellationToken); + + private async Task?> GetByIdCoreImplAsync(string documentId, CancellationToken cancellationToken) { if (_cache.TryGetValue(documentId, out var vectorDoc)) return vectorDoc.Document; - using var response = _httpClient.GetAsync($"/{_indexName}/_doc/{documentId}").GetAwaiter().GetResult(); + using var response = await _httpClient.GetAsync($"/{_indexName}/_doc/{documentId}", cancellationToken).ConfigureAwait(false); if (response.StatusCode == System.Net.HttpStatusCode.NotFound) return null; if (!response.IsSuccessStatusCode) { - var errorContent = response.Content.ReadAsStringAsync().GetAwaiter().GetResult(); + var errorContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false); throw new HttpRequestException($"Failed to retrieve document '{documentId}' from Elasticsearch. Status: {response.StatusCode}, Error: {errorContent}"); } - var responseContent = response.Content.ReadAsStringAsync().GetAwaiter().GetResult(); + var responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false); var result = JObject.Parse(responseContent); if (result["found"]?.Value() != true) @@ -362,8 +392,15 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// /// protected override bool RemoveCore(string documentId) + => RemoveCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task RemoveCoreAsync(string documentId, CancellationToken cancellationToken) + => RemoveCoreImplAsync(documentId, cancellationToken); + + private async Task RemoveCoreImplAsync(string documentId, CancellationToken cancellationToken) { - using var response = _httpClient.DeleteAsync($"/{_indexName}/_doc/{documentId}").GetAwaiter().GetResult(); + using var response = await _httpClient.DeleteAsync($"/{_indexName}/_doc/{documentId}", cancellationToken).ConfigureAwait(false); if (response.IsSuccessStatusCode && _documentCount > 0) { _cache.Remove(documentId); @@ -410,10 +447,19 @@ protected override bool RemoveCore(string documentId) /// /// protected override IEnumerable> GetAllCore() + => GetAllCoreImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetAllCoreAsync(CancellationToken cancellationToken) + => GetAllCoreImplAsync(cancellationToken); + + private async Task>> GetAllCoreImplAsync(CancellationToken cancellationToken) { const string scrollTimeout = "1m"; const int batchSize = 1000; + var documents = new List>(); + // Initial scroll request var searchRequest = new { @@ -421,26 +467,27 @@ protected override IEnumerable> GetAllCore() query = new { match_all = new { } } }; - var searchResponse = _httpClient.PostAsync( + using var searchResponse = await _httpClient.PostAsync( $"{_indexName}/_search?scroll={scrollTimeout}", new StringContent( Newtonsoft.Json.JsonConvert.SerializeObject(searchRequest), Encoding.UTF8, "application/json" - ) - ).Result; + ), + cancellationToken + ).ConfigureAwait(false); if (!searchResponse.IsSuccessStatusCode) { - var error = searchResponse.Content.ReadAsStringAsync().Result; + var error = await searchResponse.Content.ReadAsStringAsync().ConfigureAwait(false); throw new InvalidOperationException($"Elasticsearch scroll initialization failed: {error}"); } - var initialResult = JObject.Parse(searchResponse.Content.ReadAsStringAsync().Result); + var initialResult = JObject.Parse(await searchResponse.Content.ReadAsStringAsync().ConfigureAwait(false)); var scrollId = initialResult["_scroll_id"]?.ToString(); if (string.IsNullOrEmpty(scrollId)) - yield break; + return documents; // Process initial batch var hits = initialResult["hits"]?["hits"] as JArray; @@ -450,28 +497,30 @@ protected override IEnumerable> GetAllCore() { var doc = ParseElasticsearchHit(hit as JObject); if (doc != null) - yield return doc; + documents.Add(doc); } } // Continue scrolling until no more results while (true) { + cancellationToken.ThrowIfCancellationRequested(); var scrollRequest = new { scroll = scrollTimeout, scroll_id = scrollId }; - var scrollResponse = _httpClient.PostAsync( + using var scrollResponse = await _httpClient.PostAsync( "_search/scroll", new StringContent( Newtonsoft.Json.JsonConvert.SerializeObject(scrollRequest), Encoding.UTF8, "application/json" - ) - ).Result; + ), + cancellationToken + ).ConfigureAwait(false); if (!scrollResponse.IsSuccessStatusCode) break; - var scrollResult = JObject.Parse(scrollResponse.Content.ReadAsStringAsync().Result); + var scrollResult = JObject.Parse(await scrollResponse.Content.ReadAsStringAsync().ConfigureAwait(false)); hits = scrollResult["hits"]?["hits"] as JArray; if (hits == null || hits.Count == 0) @@ -481,7 +530,7 @@ protected override IEnumerable> GetAllCore() { var doc = ParseElasticsearchHit(hit as JObject); if (doc != null) - yield return doc; + documents.Add(doc); } } @@ -501,13 +550,15 @@ protected override IEnumerable> GetAllCore() Content = deleteContent }; - using var deleteResponse = _httpClient.SendAsync(request).GetAwaiter().GetResult(); + using var deleteResponse = await _httpClient.SendAsync(request, cancellationToken).ConfigureAwait(false); deleteResponse.EnsureSuccessStatusCode(); } catch { // Scroll context will expire automatically } + + return documents; } private Document? ParseElasticsearchHit(JObject? hit) @@ -567,10 +618,17 @@ protected override IEnumerable> GetAllCore() /// /// public override void Clear() + => ClearImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + public override Task ClearAsync(CancellationToken cancellationToken = default) + => ClearImplAsync(cancellationToken); + + private async Task ClearImplAsync(CancellationToken cancellationToken) { try { - using var response = _httpClient.DeleteAsync($"/{_indexName}").GetAwaiter().GetResult(); + using var response = await _httpClient.DeleteAsync($"/{_indexName}", cancellationToken).ConfigureAwait(false); response.EnsureSuccessStatusCode(); _cache.Clear(); diff --git a/src/AiModelBuilder.VectorIndexStore.cs b/src/AiModelBuilder.VectorIndexStore.cs index 9fa8b90eb9..edc6164557 100644 --- a/src/AiModelBuilder.VectorIndexStore.cs +++ b/src/AiModelBuilder.VectorIndexStore.cs @@ -1,6 +1,8 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Interfaces; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.Models; @@ -132,6 +134,65 @@ public void Clear() /// public IEnumerable> GetAll() => _documents.Values.Select(vd => vd.Document); + /// + public Task AddAsync(VectorDocument vectorDocument, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + Add(vectorDocument); + return Task.CompletedTask; + } + + /// + public Task AddBatchAsync(IEnumerable> vectorDocuments, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + AddBatch(vectorDocuments); + return Task.CompletedTask; + } + + /// + public Task>> GetSimilarAsync(Vector queryVector, int topK, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GetSimilar(queryVector, topK)); + } + + /// + public Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); + } + + /// + public Task?> GetByIdAsync(string documentId, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GetById(documentId)); + } + + /// + public Task RemoveAsync(string documentId, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(Remove(documentId)); + } + + /// + public Task ClearAsync(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + Clear(); + return Task.CompletedTask; + } + + /// + public Task>> GetAllAsync(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GetAll()); + } + private static bool MatchesFilters(Document document, Dictionary filters) { foreach (var kvp in filters) diff --git a/src/Interfaces/IContextCompressor.cs b/src/Interfaces/IContextCompressor.cs index 1beccb5f09..a9015b18d3 100644 --- a/src/Interfaces/IContextCompressor.cs +++ b/src/Interfaces/IContextCompressor.cs @@ -1,3 +1,5 @@ +using System.Threading; +using System.Threading.Tasks; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.Interfaces; @@ -59,4 +61,24 @@ List> Compress( List> documents, string query, Dictionary? options = null); + + /// + /// Asynchronously compresses a collection of documents while preserving relevance to the query. + /// + /// The documents to compress. + /// The query text used to determine relevance. + /// Optional compression parameters. + /// A token to observe for cancellation requests. + /// A task producing the compressed documents with reduced content but maintained relevance. + /// + /// + /// Asynchronous counterpart to . + /// LLM-backed compressors perform genuine non-blocking work here; extractive compressors complete synchronously. + /// + /// + Task>> CompressAsync( + List> documents, + string query, + Dictionary? options = null, + CancellationToken cancellationToken = default); } diff --git a/src/Interfaces/IDocumentStore.cs b/src/Interfaces/IDocumentStore.cs index 241d853a15..ff839ba1ef 100644 --- a/src/Interfaces/IDocumentStore.cs +++ b/src/Interfaces/IDocumentStore.cs @@ -1,3 +1,5 @@ +using System.Threading; +using System.Threading.Tasks; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.Models; @@ -208,5 +210,52 @@ public interface IDocumentStore /// /// IEnumerable> GetAll(); + + // ------------------------------------------------------------------ + // Asynchronous, cancellation-aware counterparts. I/O-bound stores + // (remote vector databases) implement these as genuinely non-blocking + // operations; in-memory stores complete synchronously. + // ------------------------------------------------------------------ + + /// Asynchronously adds a single vectorized document to the store. + /// The vector document to add. + /// A token to observe for cancellation requests. + Task AddAsync(VectorDocument vectorDocument, CancellationToken cancellationToken = default); + + /// Asynchronously adds multiple vectorized documents to the store in a batch. + /// The vector documents to add. + /// A token to observe for cancellation requests. + Task AddBatchAsync(IEnumerable> vectorDocuments, CancellationToken cancellationToken = default); + + /// Asynchronously retrieves the top-k most similar documents to a query vector. + /// The vector to search for similar documents. + /// The number of most similar documents to return. + /// A token to observe for cancellation requests. + Task>> GetSimilarAsync(Vector queryVector, int topK, CancellationToken cancellationToken = default); + + /// Asynchronously retrieves similar documents with additional metadata filtering. + /// The vector to search for similar documents. + /// The number of most similar documents to return. + /// Metadata filters to apply before similarity search. + /// A token to observe for cancellation requests. + Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken = default); + + /// Asynchronously retrieves a document by its unique identifier. + /// The unique identifier of the document to retrieve. + /// A token to observe for cancellation requests. + Task?> GetByIdAsync(string documentId, CancellationToken cancellationToken = default); + + /// Asynchronously removes a document from the store by its identifier. + /// The unique identifier of the document to remove. + /// A token to observe for cancellation requests. + Task RemoveAsync(string documentId, CancellationToken cancellationToken = default); + + /// Asynchronously removes all documents from the store. + /// A token to observe for cancellation requests. + Task ClearAsync(CancellationToken cancellationToken = default); + + /// Asynchronously gets all documents currently stored in the document store. + /// A token to observe for cancellation requests. + Task>> GetAllAsync(CancellationToken cancellationToken = default); } diff --git a/src/Interfaces/IGenerator.cs b/src/Interfaces/IGenerator.cs index fccbd5ed18..05bae53aad 100644 --- a/src/Interfaces/IGenerator.cs +++ b/src/Interfaces/IGenerator.cs @@ -1,3 +1,5 @@ +using System.Threading; +using System.Threading.Tasks; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.Interfaces; @@ -73,6 +75,34 @@ public interface IGenerator : ITextGenerator /// GroundedAnswer GenerateGrounded(string query, IEnumerable> context); + /// + /// Asynchronously generates a text response for a prompt, honoring cancellation. + /// + /// The input prompt (typically already augmented with retrieved context). + /// A token to observe for cancellation requests. + /// A task producing the generated text response. + /// + /// + /// Asynchronous counterpart to . Generators backed by a + /// remote chat model perform genuine non-blocking work here; local generators complete synchronously. + /// + /// + Task GenerateAsync(string prompt, CancellationToken cancellationToken = default); + + /// + /// Asynchronously generates a grounded answer using provided context documents, honoring cancellation. + /// + /// The user's original query or question. + /// The retrieved documents providing context for the answer. + /// A token to observe for cancellation requests. + /// A task producing a grounded answer with the generated text, source documents, and extracted citations. + /// + /// + /// Asynchronous counterpart to . + /// + /// + Task> GenerateGroundedAsync(string query, IEnumerable> context, CancellationToken cancellationToken = default); + /// /// Gets the maximum number of tokens this generator can process in a single request. /// diff --git a/src/Interfaces/IReranker.cs b/src/Interfaces/IReranker.cs index 41e64a33d9..8ee320cc0e 100644 --- a/src/Interfaces/IReranker.cs +++ b/src/Interfaces/IReranker.cs @@ -1,3 +1,5 @@ +using System.Threading; +using System.Threading.Tasks; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.Interfaces; @@ -87,6 +89,32 @@ public interface IReranker /// IEnumerable> Rerank(string query, IEnumerable> documents, int topK); + /// + /// Asynchronously reranks a collection of documents based on their relevance to a query. + /// + /// The query text used to assess relevance. + /// The documents to rerank. + /// A token to observe for cancellation requests. + /// A task producing the documents reordered by relevance, with updated relevance scores. + /// + /// + /// Asynchronous counterpart to . LLM- or + /// cross-encoder-backed rerankers perform genuine non-blocking work here; lexical rerankers complete + /// synchronously. + /// + /// + Task>> RerankAsync(string query, IEnumerable> documents, CancellationToken cancellationToken = default); + + /// + /// Asynchronously reranks documents and returns only the top-k highest scoring results. + /// + /// The query text used to assess relevance. + /// The documents to rerank. + /// The number of top-ranked documents to return. + /// A token to observe for cancellation requests. + /// A task producing the top-k documents ordered by relevance, with updated relevance scores. + Task>> RerankAsync(string query, IEnumerable> documents, int topK, CancellationToken cancellationToken = default); + /// /// Gets a value indicating whether this reranker modifies relevance scores. /// diff --git a/src/Interfaces/IRetriever.cs b/src/Interfaces/IRetriever.cs index 47c781d54e..784d282208 100644 --- a/src/Interfaces/IRetriever.cs +++ b/src/Interfaces/IRetriever.cs @@ -1,3 +1,5 @@ +using System.Threading; +using System.Threading.Tasks; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.Interfaces; @@ -109,4 +111,27 @@ public interface IRetriever /// /// IEnumerable> Retrieve(string query, int topK, Dictionary metadataFilters); + + /// + /// Asynchronously retrieves relevant documents for a query, honoring cancellation. + /// + /// The query text. + /// The number of documents to retrieve. + /// Optional metadata filters to apply before retrieval. When null, no filtering is applied. + /// A token to observe for cancellation requests. + /// A task producing the relevant documents ordered by relevance (most relevant first). + /// + /// + /// This is the asynchronous counterpart to . + /// I/O-bound retrievers (e.g. those backed by a network embedding model or a remote vector store) + /// perform genuine non-blocking work here; CPU-bound retrievers complete synchronously. + /// + /// For Beginners: Same as Retrieve, but you can await it and cancel it + /// (for example, if the user navigates away before the search finishes). + /// + Task>> RetrieveAsync( + string query, + int topK, + Dictionary? metadataFilters = null, + CancellationToken cancellationToken = default); } diff --git a/src/RetrievalAugmentedGeneration/ContextCompression/ContextCompressorBase.cs b/src/RetrievalAugmentedGeneration/ContextCompression/ContextCompressorBase.cs index b658715afa..76f092907e 100644 --- a/src/RetrievalAugmentedGeneration/ContextCompression/ContextCompressorBase.cs +++ b/src/RetrievalAugmentedGeneration/ContextCompression/ContextCompressorBase.cs @@ -1,4 +1,6 @@ +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; @@ -52,6 +54,28 @@ public List> Compress( return CompressCore(documents, query, options); } + /// + /// Asynchronously compresses a collection of documents while preserving relevance to the query. + /// + public virtual async Task>> CompressAsync( + List> documents, + string query, + Dictionary? options = null, + CancellationToken cancellationToken = default) + { + ValidateQuery(query); + ValidateDocuments(documents); + + if (documents.Count == 0) + { + return new List>(); + } + + cancellationToken.ThrowIfCancellationRequested(); + + return await CompressCoreAsync(documents, query, options, cancellationToken).ConfigureAwait(false); + } + /// /// Core compression logic to be implemented by derived classes. /// @@ -75,6 +99,26 @@ protected abstract List> CompressCore( string query, Dictionary? options = null); + /// + /// Core asynchronous compression logic. The default implementation wraps the synchronous + /// , which is appropriate for CPU-bound compressors. LLM-backed compressors + /// should override this to perform genuine non-blocking work. + /// + /// The validated and non-empty list of documents. + /// The validated query text. + /// Optional compression parameters. + /// A token to observe for cancellation requests. + /// A task producing the compressed documents. + protected virtual Task>> CompressCoreAsync( + List> documents, + string query, + Dictionary? options, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(CompressCore(documents, query, options)); + } + /// /// Validates the query string. /// diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs index c8ddb4ae1d..d5b6997e2a 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs @@ -5,6 +5,7 @@ using System.Net; using System.Net.Http; using System.Text; +using System.Threading; using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; @@ -238,17 +239,31 @@ private void EnsureIndexForDimension(int dimension) /// protected override void AddCore(VectorDocument vectorDocument) + => AddCoreImplAsync(vectorDocument, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddCoreAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) + => AddCoreImplAsync(vectorDocument, cancellationToken); + + private async Task AddCoreImplAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) { EnsureIndexForDimension(vectorDocument.Embedding.Length); if (_vectorDimension == 0) _vectorDimension = vectorDocument.Embedding.Length; - UploadActions(new[] { BuildAction(vectorDocument, "mergeOrUpload") }, "upload document"); + await UploadActionsAsync(new[] { BuildAction(vectorDocument, "mergeOrUpload") }, "upload document", cancellationToken).ConfigureAwait(false); _documentCount++; } /// protected override void AddBatchCore(IList> vectorDocuments) + => AddBatchCoreImplAsync(vectorDocuments, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddBatchCoreAsync(IList> vectorDocuments, CancellationToken cancellationToken) + => AddBatchCoreImplAsync(vectorDocuments, cancellationToken); + + private async Task AddBatchCoreImplAsync(IList> vectorDocuments, CancellationToken cancellationToken) { if (vectorDocuments.Count == 0) return; @@ -258,14 +273,14 @@ protected override void AddBatchCore(IList> vectorDocuments) _vectorDimension = vectorDocuments[0].Embedding.Length; var actions = vectorDocuments.Select(d => BuildAction(d, "mergeOrUpload")).ToList(); - UploadActions(actions, "batch upload documents"); + await UploadActionsAsync(actions, "batch upload documents", cancellationToken).ConfigureAwait(false); _documentCount += vectorDocuments.Count; } - private void UploadActions(IEnumerable> actions, string operation) + private async Task UploadActionsAsync(IEnumerable> actions, string operation, CancellationToken cancellationToken) { var body = new { value = actions }; - var info = SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/index"), body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/index"), body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, operation); } @@ -291,6 +306,13 @@ private void UploadActions(IEnumerable> actions, str /// protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); + + private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); @@ -308,7 +330,7 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector if (filter != null) body["filter"] = filter; - var info = SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/search"), body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/search"), body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "search"); var results = new List>(); @@ -401,9 +423,16 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override Document? GetByIdCore(string documentId) + => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task?> GetByIdCoreAsync(string documentId, CancellationToken cancellationToken) + => GetByIdCoreImplAsync(documentId, cancellationToken); + + private async Task?> GetByIdCoreImplAsync(string documentId, CancellationToken cancellationToken) { var encoded = Uri.EscapeDataString(documentId); - var info = SendAsync(HttpMethod.Get, WithApiVersion($"/indexes/{_indexName}/docs/{encoded}"), null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Get, WithApiVersion($"/indexes/{_indexName}/docs/{encoded}"), null, cancellationToken).ConfigureAwait(false); if (info.Status == HttpStatusCode.NotFound) return null; EnsureSuccess(info, "get document"); @@ -413,6 +442,13 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override bool RemoveCore(string documentId) + => RemoveCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task RemoveCoreAsync(string documentId, CancellationToken cancellationToken) + => RemoveCoreImplAsync(documentId, cancellationToken); + + private async Task RemoveCoreImplAsync(string documentId, CancellationToken cancellationToken) { var action = new Dictionary { @@ -421,7 +457,7 @@ protected override bool RemoveCore(string documentId) }; var body = new { value = new[] { action } }; - var info = SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/index"), body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/index"), body, cancellationToken).ConfigureAwait(false); if (!IsSuccess(info.Status)) return false; @@ -432,6 +468,13 @@ protected override bool RemoveCore(string documentId) /// protected override IEnumerable> GetAllCore() + => GetAllCoreImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetAllCoreAsync(CancellationToken cancellationToken) + => GetAllCoreImplAsync(cancellationToken); + + private async Task>> GetAllCoreImplAsync(CancellationToken cancellationToken) { var all = new List>(); const int pageSize = 1000; @@ -447,7 +490,7 @@ protected override IEnumerable> GetAllCore() ["skip"] = skip }; - var info = SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/search"), body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, WithApiVersion($"/indexes/{_indexName}/docs/search"), body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "list documents"); var rows = JObject.Parse(info.Body)["value"] as JArray; @@ -471,8 +514,15 @@ protected override IEnumerable> GetAllCore() /// public override void Clear() + => ClearImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + public override Task ClearAsync(CancellationToken cancellationToken = default) + => ClearImplAsync(cancellationToken); + + private async Task ClearImplAsync(CancellationToken cancellationToken) { - var info = SendAsync(HttpMethod.Delete, WithApiVersion($"/indexes/{_indexName}"), null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Delete, WithApiVersion($"/indexes/{_indexName}"), null, cancellationToken).ConfigureAwait(false); if (!IsSuccess(info.Status) && info.Status != HttpStatusCode.NotFound) EnsureSuccess(info, "delete index"); @@ -521,7 +571,7 @@ private void EnsureSuccess(HttpResponseInfo info, string operation) throw new HttpRequestException($"Azure AI Search {operation} failed with status {(int)info.Status}: {info.Body}"); } - private async Task SendAsync(HttpMethod method, string path, object? body) + private async Task SendAsync(HttpMethod method, string path, object? body, CancellationToken cancellationToken = default) { using (var request = new HttpRequestMessage(method, path)) { @@ -531,7 +581,7 @@ private async Task SendAsync(HttpMethod method, string path, o request.Content = new StringContent(json, Encoding.UTF8, "application/json"); } - using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) + using (var response = await _httpClient.SendAsync(request, cancellationToken).ConfigureAwait(false)) { var content = response.Content != null ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/ChromaDBDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/ChromaDBDocumentStore.cs index 6119d3752e..e802f4f46e 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/ChromaDBDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/ChromaDBDocumentStore.cs @@ -3,6 +3,8 @@ using System.Linq; using System.Net.Http; using System.Text; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; @@ -73,6 +75,13 @@ private void EnsureCollection() } protected override void AddCore(VectorDocument vectorDocument) + => AddCoreImplAsync(vectorDocument, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddCoreAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) + => AddCoreImplAsync(vectorDocument, cancellationToken); + + private async Task AddCoreImplAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) { if (_vectorDimension == 0) _vectorDimension = vectorDocument.Embedding.Length; @@ -94,7 +103,7 @@ protected override void AddCore(VectorDocument vectorDocument) Encoding.UTF8, "application/json"); - using var response = _httpClient.PostAsync($"/api/v1/collections/{_collectionName}/add", content).GetAwaiter().GetResult(); + using var response = await _httpClient.PostAsync($"/api/v1/collections/{_collectionName}/add", content, cancellationToken).ConfigureAwait(false); response.EnsureSuccessStatusCode(); _cache[vectorDocument.Document.Id] = vectorDocument; @@ -102,6 +111,13 @@ protected override void AddCore(VectorDocument vectorDocument) } protected override void AddBatchCore(IList> vectorDocuments) + => AddBatchCoreImplAsync(vectorDocuments, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddBatchCoreAsync(IList> vectorDocuments, CancellationToken cancellationToken) + => AddBatchCoreImplAsync(vectorDocuments, cancellationToken); + + private async Task AddBatchCoreImplAsync(IList> vectorDocuments, CancellationToken cancellationToken) { if (vectorDocuments.Count == 0) return; @@ -126,7 +142,7 @@ protected override void AddBatchCore(IList> vectorDocuments) Encoding.UTF8, "application/json"); - using var response = _httpClient.PostAsync($"/api/v1/collections/{_collectionName}/add", content).GetAwaiter().GetResult(); + using var response = await _httpClient.PostAsync($"/api/v1/collections/{_collectionName}/add", content, cancellationToken).ConfigureAwait(false); response.EnsureSuccessStatusCode(); foreach (var vd in vectorDocuments) @@ -135,6 +151,13 @@ protected override void AddBatchCore(IList> vectorDocuments) } protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); + + private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) { var embedding = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToList(); @@ -149,9 +172,9 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector Encoding.UTF8, "application/json"); - using var response = _httpClient.PostAsync($"/api/v1/collections/{_collectionName}/query", content).GetAwaiter().GetResult(); + using var response = await _httpClient.PostAsync($"/api/v1/collections/{_collectionName}/query", content, cancellationToken).ConfigureAwait(false); response.EnsureSuccessStatusCode(); - var responseContent = response.Content.ReadAsStringAsync().GetAwaiter().GetResult(); + var responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false); var result = JObject.Parse(responseContent); var results = new List>(); @@ -214,6 +237,13 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// /// protected override bool RemoveCore(string documentId) + => RemoveCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task RemoveCoreAsync(string documentId, CancellationToken cancellationToken) + => RemoveCoreImplAsync(documentId, cancellationToken); + + private async Task RemoveCoreImplAsync(string documentId, CancellationToken cancellationToken) { var payload = new { ids = new[] { documentId } }; using var content = new StringContent( @@ -221,7 +251,7 @@ protected override bool RemoveCore(string documentId) Encoding.UTF8, "application/json"); - using var response = _httpClient.PostAsync($"/api/v1/collections/{_collectionName}/delete", content).GetAwaiter().GetResult(); + using var response = await _httpClient.PostAsync($"/api/v1/collections/{_collectionName}/delete", content, cancellationToken).ConfigureAwait(false); if (response.IsSuccessStatusCode && _documentCount > 0) { _cache.Remove(documentId); @@ -297,10 +327,17 @@ protected override IEnumerable> GetAllCore() /// /// public override void Clear() + => ClearImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + public override Task ClearAsync(CancellationToken cancellationToken = default) + => ClearImplAsync(cancellationToken); + + private async Task ClearImplAsync(CancellationToken cancellationToken) { try { - using var response = _httpClient.DeleteAsync($"/api/v1/collections/{_collectionName}").GetAwaiter().GetResult(); + using var response = await _httpClient.DeleteAsync($"/api/v1/collections/{_collectionName}", cancellationToken).ConfigureAwait(false); response.EnsureSuccessStatusCode(); _cache.Clear(); diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/DocumentStoreBase.cs b/src/RetrievalAugmentedGeneration/DocumentStores/DocumentStoreBase.cs index 43120353e1..2c66f4b3ec 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/DocumentStoreBase.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/DocumentStoreBase.cs @@ -1,4 +1,6 @@ using System.Linq; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Interfaces; using AiDotNet.LinearAlgebra; @@ -136,6 +138,147 @@ public IEnumerable> GetAll() return GetAllCore(); } + // ------------------------------------------------------------------ + // Asynchronous, cancellation-aware API. The default implementations + // validate identically to the synchronous methods and delegate to the + // *CoreAsync members, which by default wrap the synchronous cores. + // I/O-bound stores override the *CoreAsync members for true async. + // ------------------------------------------------------------------ + + /// + public virtual async Task AddAsync(VectorDocument vectorDocument, CancellationToken cancellationToken = default) + { + ValidateVectorDocument(vectorDocument); + cancellationToken.ThrowIfCancellationRequested(); + await AddCoreAsync(vectorDocument, cancellationToken).ConfigureAwait(false); + } + + /// + public virtual async Task AddBatchAsync(IEnumerable> vectorDocuments, CancellationToken cancellationToken = default) + { + if (vectorDocuments == null) + throw new ArgumentNullException(nameof(vectorDocuments)); + + var documentList = vectorDocuments.ToList(); + if (documentList.Count == 0) + throw new ArgumentException("Vector document collection cannot be empty", nameof(vectorDocuments)); + + foreach (var vectorDocument in documentList) + { + ValidateVectorDocument(vectorDocument); + } + + cancellationToken.ThrowIfCancellationRequested(); + await AddBatchCoreAsync(documentList, cancellationToken).ConfigureAwait(false); + } + + /// + public virtual Task>> GetSimilarAsync(Vector queryVector, int topK, CancellationToken cancellationToken = default) + { + return GetSimilarWithFiltersAsync(queryVector, topK, new Dictionary(), cancellationToken); + } + + /// + public virtual async Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken = default) + { + ValidateQueryVector(queryVector); + ValidateTopK(topK); + ValidateMetadataFilters(metadataFilters); + cancellationToken.ThrowIfCancellationRequested(); + return await GetSimilarCoreAsync(queryVector, topK, metadataFilters, cancellationToken).ConfigureAwait(false); + } + + /// + public virtual async Task?> GetByIdAsync(string documentId, CancellationToken cancellationToken = default) + { + ValidateDocumentId(documentId); + cancellationToken.ThrowIfCancellationRequested(); + return await GetByIdCoreAsync(documentId, cancellationToken).ConfigureAwait(false); + } + + /// + public virtual async Task RemoveAsync(string documentId, CancellationToken cancellationToken = default) + { + ValidateDocumentId(documentId); + cancellationToken.ThrowIfCancellationRequested(); + return await RemoveCoreAsync(documentId, cancellationToken).ConfigureAwait(false); + } + + /// + public virtual Task ClearAsync(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + Clear(); + return Task.CompletedTask; + } + + /// + public virtual async Task>> GetAllAsync(CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return await GetAllCoreAsync(cancellationToken).ConfigureAwait(false); + } + + /// + /// Core asynchronous logic for adding a single vector document. Default wraps . + /// Override in I/O-bound stores for true async. + /// + protected virtual Task AddCoreAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + AddCore(vectorDocument); + return Task.CompletedTask; + } + + /// + /// Core asynchronous logic for adding a batch of vector documents. Default awaits + /// per document; override for efficient bulk insertion. + /// + protected virtual async Task AddBatchCoreAsync(IList> vectorDocuments, CancellationToken cancellationToken) + { + foreach (var vectorDocument in vectorDocuments) + { + cancellationToken.ThrowIfCancellationRequested(); + await AddCoreAsync(vectorDocument, cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Core asynchronous similarity search. Default wraps . Override for true async. + /// + protected virtual Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GetSimilarCore(queryVector, topK, metadataFilters)); + } + + /// + /// Core asynchronous lookup by ID. Default wraps . Override for true async. + /// + protected virtual Task?> GetByIdCoreAsync(string documentId, CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GetByIdCore(documentId)); + } + + /// + /// Core asynchronous removal by ID. Default wraps . Override for true async. + /// + protected virtual Task RemoveCoreAsync(string documentId, CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(RemoveCore(documentId)); + } + + /// + /// Core asynchronous retrieval of all documents. Default wraps . Override for true async. + /// + protected virtual Task>> GetAllCoreAsync(CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GetAllCore()); + } + /// /// Core logic for adding a single vector document. /// diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs index 05efaacf8e..dbae091f82 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs @@ -5,6 +5,7 @@ using System.Net; using System.Net.Http; using System.Text; +using System.Threading; using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; @@ -199,19 +200,33 @@ private void EnsureCollectionForDimension(int dimension) /// protected override void AddCore(VectorDocument vectorDocument) + => AddCoreImplAsync(vectorDocument, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddCoreAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) + => AddCoreImplAsync(vectorDocument, cancellationToken); + + private async Task AddCoreImplAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) { EnsureCollectionForDimension(vectorDocument.Embedding.Length); if (_vectorDimension == 0) _vectorDimension = vectorDocument.Embedding.Length; var body = new { collectionName = _collectionName, data = new[] { BuildEntity(vectorDocument) } }; - var info = SendAsync("/v2/vectordb/entities/upsert", body).GetAwaiter().GetResult(); + var info = await SendAsync("/v2/vectordb/entities/upsert", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "upsert entity"); _documentCount++; } /// protected override void AddBatchCore(IList> vectorDocuments) + => AddBatchCoreImplAsync(vectorDocuments, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddBatchCoreAsync(IList> vectorDocuments, CancellationToken cancellationToken) + => AddBatchCoreImplAsync(vectorDocuments, cancellationToken); + + private async Task AddBatchCoreImplAsync(IList> vectorDocuments, CancellationToken cancellationToken) { if (vectorDocuments.Count == 0) return; @@ -222,7 +237,7 @@ protected override void AddBatchCore(IList> vectorDocuments) var data = vectorDocuments.Select(BuildEntity).ToList(); var body = new { collectionName = _collectionName, data }; - var info = SendAsync("/v2/vectordb/entities/upsert", body).GetAwaiter().GetResult(); + var info = await SendAsync("/v2/vectordb/entities/upsert", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "batch upsert entities"); _documentCount += vectorDocuments.Count; } @@ -248,6 +263,13 @@ protected override void AddBatchCore(IList> vectorDocuments) /// protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); + + private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); @@ -264,7 +286,7 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector if (filter != null) body["filter"] = filter; - var info = SendAsync("/v2/vectordb/entities/search", body).GetAwaiter().GetResult(); + var info = await SendAsync("/v2/vectordb/entities/search", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "search"); var results = new List>(); @@ -355,6 +377,13 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override Document? GetByIdCore(string documentId) + => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task?> GetByIdCoreAsync(string documentId, CancellationToken cancellationToken) + => GetByIdCoreImplAsync(documentId, cancellationToken); + + private async Task?> GetByIdCoreImplAsync(string documentId, CancellationToken cancellationToken) { var body = new { @@ -364,7 +393,7 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector limit = 1 }; - var info = SendAsync("/v2/vectordb/entities/query", body).GetAwaiter().GetResult(); + var info = await SendAsync("/v2/vectordb/entities/query", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "query by id"); var rows = JObject.Parse(info.Body)["data"] as JArray; @@ -376,6 +405,13 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override bool RemoveCore(string documentId) + => RemoveCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task RemoveCoreAsync(string documentId, CancellationToken cancellationToken) + => RemoveCoreImplAsync(documentId, cancellationToken); + + private async Task RemoveCoreImplAsync(string documentId, CancellationToken cancellationToken) { var body = new { @@ -383,7 +419,7 @@ protected override bool RemoveCore(string documentId) filter = $"{IdField} == {Quote(documentId)}" }; - var info = SendAsync("/v2/vectordb/entities/delete", body).GetAwaiter().GetResult(); + var info = await SendAsync("/v2/vectordb/entities/delete", body, cancellationToken).ConfigureAwait(false); if (!IsSuccess(info.Status)) return false; @@ -398,6 +434,13 @@ protected override bool RemoveCore(string documentId) /// protected override IEnumerable> GetAllCore() + => GetAllCoreImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetAllCoreAsync(CancellationToken cancellationToken) + => GetAllCoreImplAsync(cancellationToken); + + private async Task>> GetAllCoreImplAsync(CancellationToken cancellationToken) { var all = new List>(); const int pageSize = 1000; @@ -414,7 +457,7 @@ protected override IEnumerable> GetAllCore() ["offset"] = offset }; - var info = SendAsync("/v2/vectordb/entities/query", body).GetAwaiter().GetResult(); + var info = await SendAsync("/v2/vectordb/entities/query", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "query all"); var rows = JObject.Parse(info.Body)["data"] as JArray; @@ -438,9 +481,16 @@ protected override IEnumerable> GetAllCore() /// public override void Clear() + => ClearImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + public override Task ClearAsync(CancellationToken cancellationToken = default) + => ClearImplAsync(cancellationToken); + + private async Task ClearImplAsync(CancellationToken cancellationToken) { - var info = SendAsync("/v2/vectordb/collections/drop", - new { collectionName = _collectionName }).GetAwaiter().GetResult(); + var info = await SendAsync("/v2/vectordb/collections/drop", + new { collectionName = _collectionName }, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "drop collection"); _documentCount = 0; @@ -492,14 +542,14 @@ private void EnsureSuccess(HttpResponseInfo info, string operation) throw new HttpRequestException($"Milvus {operation} failed with code {code}: {info.Body}"); } - private async Task SendAsync(string path, object body) + private async Task SendAsync(string path, object body, CancellationToken cancellationToken = default) { using (var request = new HttpRequestMessage(HttpMethod.Post, path)) { var json = JsonConvert.SerializeObject(body); request.Content = new StringContent(json, Encoding.UTF8, "application/json"); - using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) + using (var response = await _httpClient.SendAsync(request, cancellationToken).ConfigureAwait(false)) { var content = response.Content != null ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs index bca3ac5536..43592c9370 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs @@ -5,6 +5,7 @@ using System.Net; using System.Net.Http; using System.Text; +using System.Threading; using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; @@ -111,7 +112,7 @@ private void InitializeStats() { try { - var info = SendAsync(HttpMethod.Post, "/describe_index_stats", new { }).GetAwaiter().GetResult(); + var info = SendAsync(HttpMethod.Post, "/describe_index_stats", new { }, CancellationToken.None).GetAwaiter().GetResult(); if (!IsSuccess(info.Status)) return; @@ -149,6 +150,13 @@ private int ReadNamespaceCount(JObject stats) /// protected override void AddCore(VectorDocument vectorDocument) + => AddCoreImplAsync(vectorDocument, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddCoreAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) + => AddCoreImplAsync(vectorDocument, cancellationToken); + + private async Task AddCoreImplAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) { if (_vectorDimension == 0) _vectorDimension = vectorDocument.Embedding.Length; @@ -157,12 +165,19 @@ protected override void AddCore(VectorDocument vectorDocument) $"Vector dimension mismatch. Expected {_vectorDimension}, got {vectorDocument.Embedding.Length}", nameof(vectorDocument)); - Upsert(new[] { vectorDocument }); + await UpsertAsync(new[] { vectorDocument }, cancellationToken).ConfigureAwait(false); _documentCount++; } /// protected override void AddBatchCore(IList> vectorDocuments) + => AddBatchCoreImplAsync(vectorDocuments, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddBatchCoreAsync(IList> vectorDocuments, CancellationToken cancellationToken) + => AddBatchCoreImplAsync(vectorDocuments, cancellationToken); + + private async Task AddBatchCoreImplAsync(IList> vectorDocuments, CancellationToken cancellationToken) { if (vectorDocuments.Count == 0) return; @@ -178,11 +193,11 @@ protected override void AddBatchCore(IList> vectorDocuments) nameof(vectorDocuments)); } - Upsert(vectorDocuments); + await UpsertAsync(vectorDocuments, cancellationToken).ConfigureAwait(false); _documentCount += vectorDocuments.Count; } - private void Upsert(IEnumerable> vectorDocuments) + private async Task UpsertAsync(IEnumerable> vectorDocuments, CancellationToken cancellationToken) { var vectors = vectorDocuments.Select(vd => { @@ -200,7 +215,7 @@ private void Upsert(IEnumerable> vectorDocuments) if (_namespace.Length > 0) body["namespace"] = _namespace; - var info = SendAsync(HttpMethod.Post, "/vectors/upsert", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, "/vectors/upsert", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "upsert"); } @@ -219,6 +234,13 @@ private static Dictionary BuildMetadata(Document document) /// protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); + + private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); @@ -237,7 +259,7 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector if (_namespace.Length > 0) body["namespace"] = _namespace; - var info = SendAsync(HttpMethod.Post, "/query", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, "/query", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "query"); var results = new List>(); @@ -307,12 +329,19 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override Document? GetByIdCore(string documentId) + => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task?> GetByIdCoreAsync(string documentId, CancellationToken cancellationToken) + => GetByIdCoreImplAsync(documentId, cancellationToken); + + private async Task?> GetByIdCoreImplAsync(string documentId, CancellationToken cancellationToken) { var path = $"/vectors/fetch?ids={Uri.EscapeDataString(documentId)}"; if (_namespace.Length > 0) path += $"&namespace={Uri.EscapeDataString(_namespace)}"; - var info = SendAsync(HttpMethod.Get, path, null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Get, path, null, cancellationToken).ConfigureAwait(false); if (info.Status == HttpStatusCode.NotFound) return null; EnsureSuccess(info, "fetch"); @@ -327,12 +356,19 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override bool RemoveCore(string documentId) + => RemoveCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task RemoveCoreAsync(string documentId, CancellationToken cancellationToken) + => RemoveCoreImplAsync(documentId, cancellationToken); + + private async Task RemoveCoreImplAsync(string documentId, CancellationToken cancellationToken) { var body = new Dictionary { ["ids"] = new[] { documentId } }; if (_namespace.Length > 0) body["namespace"] = _namespace; - var info = SendAsync(HttpMethod.Post, "/vectors/delete", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, "/vectors/delete", body, cancellationToken).ConfigureAwait(false); if (!IsSuccess(info.Status)) return false; @@ -343,6 +379,13 @@ protected override bool RemoveCore(string documentId) /// protected override IEnumerable> GetAllCore() + => GetAllCoreImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetAllCoreAsync(CancellationToken cancellationToken) + => GetAllCoreImplAsync(cancellationToken); + + private async Task>> GetAllCoreImplAsync(CancellationToken cancellationToken) { var all = new List>(); string? paginationToken = null; @@ -355,7 +398,7 @@ protected override IEnumerable> GetAllCore() if (!string.IsNullOrEmpty(paginationToken)) path += $"&paginationToken={Uri.EscapeDataString(paginationToken!)}"; - var info = SendAsync(HttpMethod.Get, path, null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Get, path, null, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "list"); var root = JObject.Parse(info.Body); @@ -367,7 +410,7 @@ protected override IEnumerable> GetAllCore() foreach (var id in ids) { - var doc = GetByIdCore(id); + var doc = await GetByIdCoreImplAsync(id, cancellationToken).ConfigureAwait(false); if (doc != null) all.Add(doc); } @@ -382,12 +425,19 @@ protected override IEnumerable> GetAllCore() /// public override void Clear() + => ClearImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + public override Task ClearAsync(CancellationToken cancellationToken = default) + => ClearImplAsync(cancellationToken); + + private async Task ClearImplAsync(CancellationToken cancellationToken) { var body = new Dictionary { ["deleteAll"] = true }; if (_namespace.Length > 0) body["namespace"] = _namespace; - var info = SendAsync(HttpMethod.Post, "/vectors/delete", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, "/vectors/delete", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "delete all"); _documentCount = 0; @@ -427,7 +477,7 @@ private void EnsureSuccess(HttpResponseInfo info, string operation) throw new HttpRequestException($"Pinecone {operation} failed with status {(int)info.Status}: {info.Body}"); } - private async Task SendAsync(HttpMethod method, string path, object? body) + private async Task SendAsync(HttpMethod method, string path, object? body, CancellationToken cancellationToken = default) { using (var request = new HttpRequestMessage(method, path)) { @@ -437,7 +487,7 @@ private async Task SendAsync(HttpMethod method, string path, o request.Content = new StringContent(json, Encoding.UTF8, "application/json"); } - using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) + using (var response = await _httpClient.SendAsync(request, cancellationToken).ConfigureAwait(false)) { var content = response.Content != null ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs index 1b7f634148..029f178877 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs @@ -6,6 +6,7 @@ using System.Net.Http; using System.Security.Cryptography; using System.Text; +using System.Threading; using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; @@ -188,19 +189,33 @@ private void EnsureCollectionForDimension(int dimension) /// protected override void AddCore(VectorDocument vectorDocument) + => AddCoreImplAsync(vectorDocument, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddCoreAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) + => AddCoreImplAsync(vectorDocument, cancellationToken); + + private async Task AddCoreImplAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) { EnsureCollectionForDimension(vectorDocument.Embedding.Length); if (_vectorDimension == 0) _vectorDimension = vectorDocument.Embedding.Length; var body = new { points = new[] { BuildPoint(vectorDocument) } }; - var info = SendAsync(HttpMethod.Put, $"/collections/{_collectionName}/points?wait=true", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Put, $"/collections/{_collectionName}/points?wait=true", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "upsert point"); _documentCount++; } /// protected override void AddBatchCore(IList> vectorDocuments) + => AddBatchCoreImplAsync(vectorDocuments, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddBatchCoreAsync(IList> vectorDocuments, CancellationToken cancellationToken) + => AddBatchCoreImplAsync(vectorDocuments, cancellationToken); + + private async Task AddBatchCoreImplAsync(IList> vectorDocuments, CancellationToken cancellationToken) { if (vectorDocuments.Count == 0) return; @@ -211,7 +226,7 @@ protected override void AddBatchCore(IList> vectorDocuments) var points = vectorDocuments.Select(BuildPoint).ToList(); var body = new { points }; - var info = SendAsync(HttpMethod.Put, $"/collections/{_collectionName}/points?wait=true", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Put, $"/collections/{_collectionName}/points?wait=true", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "batch upsert points"); _documentCount += vectorDocuments.Count; } @@ -236,6 +251,13 @@ private object BuildPoint(VectorDocument vectorDocument) /// protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); + + private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); @@ -251,7 +273,7 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector if (filter != null) body["filter"] = filter; - var info = SendAsync(HttpMethod.Post, $"/collections/{_collectionName}/points/search", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, $"/collections/{_collectionName}/points/search", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "search"); var results = new List>(); @@ -325,8 +347,15 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override Document? GetByIdCore(string documentId) + => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task?> GetByIdCoreAsync(string documentId, CancellationToken cancellationToken) + => GetByIdCoreImplAsync(documentId, cancellationToken); + + private async Task?> GetByIdCoreImplAsync(string documentId, CancellationToken cancellationToken) { - var info = SendAsync(HttpMethod.Get, $"/collections/{_collectionName}/points/{ToPointId(documentId)}", null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Get, $"/collections/{_collectionName}/points/{ToPointId(documentId)}", null, cancellationToken).ConfigureAwait(false); if (info.Status == HttpStatusCode.NotFound) return null; EnsureSuccess(info, "get point"); @@ -340,9 +369,16 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override bool RemoveCore(string documentId) + => RemoveCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task RemoveCoreAsync(string documentId, CancellationToken cancellationToken) + => RemoveCoreImplAsync(documentId, cancellationToken); + + private async Task RemoveCoreImplAsync(string documentId, CancellationToken cancellationToken) { var body = new { points = new[] { ToPointId(documentId) } }; - var info = SendAsync(HttpMethod.Post, $"/collections/{_collectionName}/points/delete?wait=true", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, $"/collections/{_collectionName}/points/delete?wait=true", body, cancellationToken).ConfigureAwait(false); if (!IsSuccess(info.Status)) return false; @@ -353,6 +389,13 @@ protected override bool RemoveCore(string documentId) /// protected override IEnumerable> GetAllCore() + => GetAllCoreImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetAllCoreAsync(CancellationToken cancellationToken) + => GetAllCoreImplAsync(cancellationToken); + + private async Task>> GetAllCoreImplAsync(CancellationToken cancellationToken) { var all = new List>(); object? offset = null; @@ -368,7 +411,7 @@ protected override IEnumerable> GetAllCore() if (offset != null) body["offset"] = offset; - var info = SendAsync(HttpMethod.Post, $"/collections/{_collectionName}/points/scroll", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, $"/collections/{_collectionName}/points/scroll", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "scroll"); var result = JObject.Parse(info.Body)["result"] as JObject; @@ -394,8 +437,15 @@ protected override IEnumerable> GetAllCore() /// public override void Clear() + => ClearImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + public override Task ClearAsync(CancellationToken cancellationToken = default) + => ClearImplAsync(cancellationToken); + + private async Task ClearImplAsync(CancellationToken cancellationToken) { - var info = SendAsync(HttpMethod.Delete, $"/collections/{_collectionName}", null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Delete, $"/collections/{_collectionName}", null, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "delete collection"); _documentCount = 0; @@ -449,7 +499,7 @@ private void EnsureSuccess(HttpResponseInfo info, string operation) throw new HttpRequestException($"Qdrant {operation} failed with status {(int)info.Status}: {info.Body}"); } - private async Task SendAsync(HttpMethod method, string path, object? body) + private async Task SendAsync(HttpMethod method, string path, object? body, CancellationToken cancellationToken = default) { using (var request = new HttpRequestMessage(method, path)) { @@ -459,7 +509,7 @@ private async Task SendAsync(HttpMethod method, string path, o request.Content = new StringContent(json, Encoding.UTF8, "application/json"); } - using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) + using (var response = await _httpClient.SendAsync(request, cancellationToken).ConfigureAwait(false)) { var content = response.Content != null ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs index 2490684e10..4278141a32 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs @@ -8,6 +8,7 @@ using System.Net.Http; using System.Security.Cryptography; using System.Text; +using System.Threading; using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; @@ -183,19 +184,33 @@ private void EnsureClass() /// protected override void AddCore(VectorDocument vectorDocument) + => AddCoreImplAsync(vectorDocument, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddCoreAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) + => AddCoreImplAsync(vectorDocument, cancellationToken); + + private async Task AddCoreImplAsync(VectorDocument vectorDocument, CancellationToken cancellationToken) { EnsureClass(); if (_vectorDimension == 0) _vectorDimension = vectorDocument.Embedding.Length; var body = BuildObject(vectorDocument); - var info = SendAsync(HttpMethod.Post, "/v1/objects", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, "/v1/objects", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "create object"); _documentCount++; } /// protected override void AddBatchCore(IList> vectorDocuments) + => AddBatchCoreImplAsync(vectorDocuments, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task AddBatchCoreAsync(IList> vectorDocuments, CancellationToken cancellationToken) + => AddBatchCoreImplAsync(vectorDocuments, cancellationToken); + + private async Task AddBatchCoreImplAsync(IList> vectorDocuments, CancellationToken cancellationToken) { if (vectorDocuments.Count == 0) return; @@ -206,7 +221,7 @@ protected override void AddBatchCore(IList> vectorDocuments) var objects = vectorDocuments.Select(BuildObject).ToList(); var body = new { objects }; - var info = SendAsync(HttpMethod.Post, "/v1/batch/objects", body).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, "/v1/batch/objects", body, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "batch create objects"); _documentCount += vectorDocuments.Count; } @@ -237,6 +252,13 @@ private object BuildObject(VectorDocument vectorDocument) /// protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); + + private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); var vectorJson = "[" + string.Join(",", vector.Select(v => v.ToString("R", CultureInfo.InvariantCulture))) + "]"; @@ -250,7 +272,7 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector "{ " + DocIdProp + " " + ContentProp + " " + MetadataProp + " _additional { certainty distance } } } }"; - var info = SendAsync(HttpMethod.Post, "/v1/graphql", new { query }).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Post, "/v1/graphql", new { query }, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "search"); var results = new List>(); @@ -350,8 +372,15 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override Document? GetByIdCore(string documentId) + => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task?> GetByIdCoreAsync(string documentId, CancellationToken cancellationToken) + => GetByIdCoreImplAsync(documentId, cancellationToken); + + private async Task?> GetByIdCoreImplAsync(string documentId, CancellationToken cancellationToken) { - var info = SendAsync(HttpMethod.Get, $"/v1/objects/{_className}/{ToObjectId(documentId)}", null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Get, $"/v1/objects/{_className}/{ToObjectId(documentId)}", null, cancellationToken).ConfigureAwait(false); if (info.Status == HttpStatusCode.NotFound) return null; EnsureSuccess(info, "get object"); @@ -361,8 +390,15 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector /// protected override bool RemoveCore(string documentId) + => RemoveCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task RemoveCoreAsync(string documentId, CancellationToken cancellationToken) + => RemoveCoreImplAsync(documentId, cancellationToken); + + private async Task RemoveCoreImplAsync(string documentId, CancellationToken cancellationToken) { - var info = SendAsync(HttpMethod.Delete, $"/v1/objects/{_className}/{ToObjectId(documentId)}", null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Delete, $"/v1/objects/{_className}/{ToObjectId(documentId)}", null, cancellationToken).ConfigureAwait(false); if (info.Status == HttpStatusCode.NotFound) return false; if (!IsSuccess(info.Status)) @@ -375,6 +411,13 @@ protected override bool RemoveCore(string documentId) /// protected override IEnumerable> GetAllCore() + => GetAllCoreImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetAllCoreAsync(CancellationToken cancellationToken) + => GetAllCoreImplAsync(cancellationToken); + + private async Task>> GetAllCoreImplAsync(CancellationToken cancellationToken) { var all = new List>(); string? after = null; @@ -386,7 +429,7 @@ protected override IEnumerable> GetAllCore() if (after != null) path += "&after=" + after; - var info = SendAsync(HttpMethod.Get, path, null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Get, path, null, cancellationToken).ConfigureAwait(false); EnsureSuccess(info, "list objects"); var objects = JObject.Parse(info.Body)["objects"] as JArray; @@ -410,8 +453,15 @@ protected override IEnumerable> GetAllCore() /// public override void Clear() + => ClearImplAsync(CancellationToken.None).GetAwaiter().GetResult(); + + /// + public override Task ClearAsync(CancellationToken cancellationToken = default) + => ClearImplAsync(cancellationToken); + + private async Task ClearImplAsync(CancellationToken cancellationToken) { - var info = SendAsync(HttpMethod.Delete, $"/v1/schema/{_className}", null).GetAwaiter().GetResult(); + var info = await SendAsync(HttpMethod.Delete, $"/v1/schema/{_className}", null, cancellationToken).ConfigureAwait(false); if (!IsSuccess(info.Status) && info.Status != HttpStatusCode.NotFound) EnsureSuccess(info, "delete class"); @@ -478,7 +528,7 @@ private void EnsureSuccess(HttpResponseInfo info, string operation) throw new HttpRequestException($"Weaviate {operation} failed with status {(int)info.Status}: {info.Body}"); } - private async Task SendAsync(HttpMethod method, string path, object? body) + private async Task SendAsync(HttpMethod method, string path, object? body, CancellationToken cancellationToken = default) { using (var request = new HttpRequestMessage(method, path)) { @@ -488,7 +538,7 @@ private async Task SendAsync(HttpMethod method, string path, o request.Content = new StringContent(json, Encoding.UTF8, "application/json"); } - using (var response = await _httpClient.SendAsync(request).ConfigureAwait(false)) + using (var response = await _httpClient.SendAsync(request, cancellationToken).ConfigureAwait(false)) { var content = response.Content != null ? await response.Content.ReadAsStringAsync().ConfigureAwait(false) diff --git a/src/RetrievalAugmentedGeneration/Generators/ChatClientGenerator.cs b/src/RetrievalAugmentedGeneration/Generators/ChatClientGenerator.cs index 104c9e85ac..12ca51cd8e 100644 --- a/src/RetrievalAugmentedGeneration/Generators/ChatClientGenerator.cs +++ b/src/RetrievalAugmentedGeneration/Generators/ChatClientGenerator.cs @@ -2,6 +2,7 @@ using System.Collections.Generic; using System.Runtime.CompilerServices; using System.Threading; +using System.Threading.Tasks; using AiDotNet.Agentic.Models; using AiDotNet.Interfaces; @@ -91,6 +92,20 @@ protected override string GenerateCore(string prompt) return response.Text ?? string.Empty; } + /// + /// + /// True async path: awaits the chat client directly with no sync-over-async blocking, and flows the + /// caller's to the provider so an in-flight request can be cancelled. + /// + protected override async Task GenerateCoreAsync(string prompt, CancellationToken cancellationToken) + { + var response = await _chatClient + .GetResponseAsync(BuildMessages(prompt), _options, cancellationToken) + .ConfigureAwait(false); + + return response.Text ?? string.Empty; + } + /// public async IAsyncEnumerable GenerateStreamAsync( string prompt, diff --git a/src/RetrievalAugmentedGeneration/Generators/GeneratorBase.cs b/src/RetrievalAugmentedGeneration/Generators/GeneratorBase.cs index 612d8ff334..bb45fa73e1 100644 --- a/src/RetrievalAugmentedGeneration/Generators/GeneratorBase.cs +++ b/src/RetrievalAugmentedGeneration/Generators/GeneratorBase.cs @@ -1,3 +1,5 @@ +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; @@ -82,6 +84,29 @@ public string Generate(string prompt) return GenerateCore(prompt); } + /// + /// Asynchronously generates a text response based on a prompt, honoring cancellation. + /// + /// The input prompt or question. + /// A token to observe for cancellation requests. + /// A task producing the generated text response. + public virtual async Task GenerateAsync(string prompt, CancellationToken cancellationToken = default) + { + if (prompt == null) + { + throw new ArgumentNullException(nameof(prompt), "Prompt cannot be null."); + } + + if (string.IsNullOrWhiteSpace(prompt)) + { + throw new ArgumentException("Prompt cannot be empty or whitespace.", nameof(prompt)); + } + + cancellationToken.ThrowIfCancellationRequested(); + + return await GenerateCoreAsync(prompt, cancellationToken).ConfigureAwait(false); + } + /// /// Generates a grounded answer using provided context documents. /// @@ -130,6 +155,58 @@ public GroundedAnswer GenerateGrounded(string query, IEnumerable> }; } + /// + /// Asynchronously generates a grounded answer using provided context documents, honoring cancellation. + /// + /// The user's original query or question. + /// The retrieved documents providing context for the answer. + /// A token to observe for cancellation requests. + /// A task producing a grounded answer with the generated text, source documents, and extracted citations. + public virtual async Task> GenerateGroundedAsync( + string query, + IEnumerable> context, + CancellationToken cancellationToken = default) + { + if (query == null) + { + throw new ArgumentNullException(nameof(query), "Query cannot be null."); + } + + if (string.IsNullOrWhiteSpace(query)) + { + throw new ArgumentException("Query cannot be empty or whitespace.", nameof(query)); + } + + if (context == null) + { + throw new ArgumentNullException(nameof(context), "Context cannot be null."); + } + + var contextList = context.ToList(); + if (contextList.Count == 0) + { + throw new ArgumentException("Context must contain at least one document.", nameof(context)); + } + + cancellationToken.ThrowIfCancellationRequested(); + + // Build the prompt with context + var prompt = BuildPromptWithContext(query, contextList); + + // Generate the answer + var generatedText = await GenerateCoreAsync(prompt, cancellationToken).ConfigureAwait(false); + + // Extract citations from the generated text + var citations = ExtractCitations(generatedText, contextList); + + return new GroundedAnswer + { + Answer = generatedText, + SourceDocuments = contextList.AsReadOnly(), + Citations = citations.Values.Select(d => d.Id).ToList().AsReadOnly() + }; + } + /// /// Core generation logic to be implemented by derived classes. /// @@ -137,6 +214,20 @@ public GroundedAnswer GenerateGrounded(string query, IEnumerable> /// The generated text response. protected abstract string GenerateCore(string prompt); + /// + /// Core asynchronous generation logic. The default implementation wraps the synchronous + /// , which is appropriate for local generators. Generators backed by a remote + /// chat model should override this to perform genuine non-blocking work. + /// + /// The validated prompt string. + /// A token to observe for cancellation requests. + /// A task producing the generated text response. + protected virtual Task GenerateCoreAsync(string prompt, CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GenerateCore(prompt)); + } + /// /// Builds a prompt that incorporates the query and retrieved context documents. /// diff --git a/src/RetrievalAugmentedGeneration/Generators/NeuralGenerator.cs b/src/RetrievalAugmentedGeneration/Generators/NeuralGenerator.cs index 521299d637..cad0da140a 100644 --- a/src/RetrievalAugmentedGeneration/Generators/NeuralGenerator.cs +++ b/src/RetrievalAugmentedGeneration/Generators/NeuralGenerator.cs @@ -2,6 +2,8 @@ using System.Collections.Generic; using System.Linq; using System.Text; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.Interfaces; @@ -300,6 +302,20 @@ public GroundedAnswer GenerateGrounded(string query, IEnumerable> }; } + /// + public Task GenerateAsync(string prompt, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(Generate(prompt)); + } + + /// + public Task> GenerateGroundedAsync(string query, IEnumerable> context, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GenerateGrounded(query, context)); + } + private List TokenizeText(string text) { // Production-ready word-based tokenization with thread-safe vocabulary tracking diff --git a/src/RetrievalAugmentedGeneration/Generators/StubGenerator.cs b/src/RetrievalAugmentedGeneration/Generators/StubGenerator.cs index e7c504f898..cde8b5277b 100644 --- a/src/RetrievalAugmentedGeneration/Generators/StubGenerator.cs +++ b/src/RetrievalAugmentedGeneration/Generators/StubGenerator.cs @@ -1,4 +1,6 @@ using System.Text.RegularExpressions; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; @@ -179,4 +181,18 @@ public GroundedAnswer GenerateGrounded(string query, IEnumerable> ConfidenceScore = confidenceScore }; } + + /// + public Task GenerateAsync(string prompt, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(Generate(prompt)); + } + + /// + public Task> GenerateGroundedAsync(string query, IEnumerable> context, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GenerateGrounded(query, context)); + } } diff --git a/src/RetrievalAugmentedGeneration/Rerankers/RerankerBase.cs b/src/RetrievalAugmentedGeneration/Rerankers/RerankerBase.cs index bebbb74cdc..592319e11c 100644 --- a/src/RetrievalAugmentedGeneration/Rerankers/RerankerBase.cs +++ b/src/RetrievalAugmentedGeneration/Rerankers/RerankerBase.cs @@ -1,5 +1,7 @@ using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; @@ -68,6 +70,41 @@ public IEnumerable> Rerank(string query, IEnumerable> do return reranked.Take(topK).ToList(); } + /// + /// Asynchronously reranks a collection of documents based on their relevance to a query. + /// + public virtual async Task>> RerankAsync( + string query, + IEnumerable> documents, + CancellationToken cancellationToken = default) + { + ValidateQuery(query); + ValidateDocuments(documents); + + var documentList = documents.ToList(); + if (documentList.Count == 0) + return documentList; + + cancellationToken.ThrowIfCancellationRequested(); + + return await RerankCoreAsync(query, documentList, cancellationToken).ConfigureAwait(false); + } + + /// + /// Asynchronously reranks documents and returns only the top-k highest scoring results. + /// + public virtual async Task>> RerankAsync( + string query, + IEnumerable> documents, + int topK, + CancellationToken cancellationToken = default) + { + ValidateTopK(topK); + + var reranked = await RerankAsync(query, documents, cancellationToken).ConfigureAwait(false); + return reranked.Take(topK).ToList(); + } + /// /// Core reranking logic to be implemented by derived classes. /// @@ -89,6 +126,24 @@ public IEnumerable> Rerank(string query, IEnumerable> do /// protected abstract IEnumerable> RerankCore(string query, IList> documents); + /// + /// Core asynchronous reranking logic. The default implementation wraps the synchronous + /// , which is appropriate for CPU-bound rerankers. LLM- or cross-encoder-backed + /// rerankers should override this to perform genuine non-blocking work. + /// + /// The validated query text. + /// The validated and materialized list of documents. + /// A token to observe for cancellation requests. + /// A task producing the documents reordered by relevance with updated scores. + protected virtual Task>> RerankCoreAsync( + string query, + IList> documents, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(RerankCore(query, documents)); + } + /// /// Validates the query string. /// diff --git a/src/RetrievalAugmentedGeneration/Retrievers/DenseRetriever.cs b/src/RetrievalAugmentedGeneration/Retrievers/DenseRetriever.cs index 66b44ca5ee..ad83ab5e1a 100644 --- a/src/RetrievalAugmentedGeneration/Retrievers/DenseRetriever.cs +++ b/src/RetrievalAugmentedGeneration/Retrievers/DenseRetriever.cs @@ -38,5 +38,23 @@ protected override IEnumerable> RetrieveCore(string query, int topK, var queryEmbedding = _embeddingModel.Embed(query); return _documentStore.GetSimilarWithFilters(queryEmbedding, topK, metadataFilters); } + + /// + /// Truly-async core retrieval: awaits the (possibly network-backed) embedding model and document + /// store instead of blocking, and flows the cancellation token end to end. + /// + protected override async System.Threading.Tasks.Task>> RetrieveCoreAsync( + string query, + int topK, + Dictionary metadataFilters, + System.Threading.CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + var queryEmbedding = await _embeddingModel.EmbedAsync(query).ConfigureAwait(false); + cancellationToken.ThrowIfCancellationRequested(); + return await _documentStore + .GetSimilarWithFiltersAsync(queryEmbedding, topK, metadataFilters, cancellationToken) + .ConfigureAwait(false); + } } } diff --git a/src/RetrievalAugmentedGeneration/Retrievers/RetrieverBase.cs b/src/RetrievalAugmentedGeneration/Retrievers/RetrieverBase.cs index 9562eee500..446aff0717 100644 --- a/src/RetrievalAugmentedGeneration/Retrievers/RetrieverBase.cs +++ b/src/RetrievalAugmentedGeneration/Retrievers/RetrieverBase.cs @@ -1,4 +1,6 @@ +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Interfaces; using AiDotNet.RetrievalAugmentedGeneration.Models; @@ -87,6 +89,32 @@ public IEnumerable> Retrieve(string query, int topK, Dictionary + /// Asynchronously retrieves relevant documents, honoring cancellation. + /// + /// The query text. + /// The number of documents to retrieve. + /// Optional metadata filters. When null, no filtering is applied. + /// A token to observe for cancellation requests. + /// A task producing the relevant documents ordered by relevance. + public virtual async Task>> RetrieveAsync( + string query, + int topK, + Dictionary? metadataFilters = null, + CancellationToken cancellationToken = default) + { + ValidateQuery(query); + ValidateTopK(topK); + var filters = metadataFilters ?? new Dictionary(); + ValidateMetadataFilters(filters); + + cancellationToken.ThrowIfCancellationRequested(); + + var results = await RetrieveCoreAsync(query, topK, filters, cancellationToken).ConfigureAwait(false); + + return PostProcessResults(results, topK); + } + /// /// Core retrieval logic to be implemented by derived classes. /// @@ -108,6 +136,26 @@ public IEnumerable> Retrieve(string query, int topK, Dictionary protected abstract IEnumerable> RetrieveCore(string query, int topK, Dictionary metadataFilters); + /// + /// Core asynchronous retrieval logic. The default implementation wraps the synchronous + /// , which is appropriate for CPU-bound retrievers. I/O-bound retrievers + /// (network embedding models, remote stores) should override this to perform genuine non-blocking work. + /// + /// The validated query text. + /// The validated number of documents to retrieve. + /// The validated metadata filters. + /// A token to observe for cancellation requests. + /// A task producing the relevant documents ordered by relevance. + protected virtual Task>> RetrieveCoreAsync( + string query, + int topK, + Dictionary metadataFilters, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(RetrieveCore(query, topK, metadataFilters)); + } + /// /// Validates the query string. /// diff --git a/src/RetrievalAugmentedGeneration/Retrievers/VectorRetriever.cs b/src/RetrievalAugmentedGeneration/Retrievers/VectorRetriever.cs index 30e756dffd..aa3a86a038 100644 --- a/src/RetrievalAugmentedGeneration/Retrievers/VectorRetriever.cs +++ b/src/RetrievalAugmentedGeneration/Retrievers/VectorRetriever.cs @@ -74,4 +74,27 @@ protected override IEnumerable> RetrieveCore(string query, int topK, // 2. Search document store for similar documents return _documentStore.GetSimilarWithFilters(queryVector, topK, metadataFilters); } + + /// + /// Truly-async core retrieval: awaits the (possibly network-backed) embedding model and document + /// store instead of blocking, and flows the cancellation token end to end. + /// + protected override async System.Threading.Tasks.Task>> RetrieveCoreAsync( + string query, + int topK, + Dictionary metadataFilters, + System.Threading.CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + + // 1. Embed the query asynchronously + var queryVector = await _embeddingModel.EmbedAsync(query).ConfigureAwait(false); + + cancellationToken.ThrowIfCancellationRequested(); + + // 2. Search document store for similar documents asynchronously + return await _documentStore + .GetSimilarWithFiltersAsync(queryVector, topK, metadataFilters, cancellationToken) + .ConfigureAwait(false); + } } diff --git a/tests/AiDotNet.Tests/IntegrationTests/KnowledgeGraph/LlmGraphExtractionTests.cs b/tests/AiDotNet.Tests/IntegrationTests/KnowledgeGraph/LlmGraphExtractionTests.cs index d80110515b..e0a2e3a356 100644 --- a/tests/AiDotNet.Tests/IntegrationTests/KnowledgeGraph/LlmGraphExtractionTests.cs +++ b/tests/AiDotNet.Tests/IntegrationTests/KnowledgeGraph/LlmGraphExtractionTests.cs @@ -23,6 +23,9 @@ public class LlmGraphExtractionTests /// private sealed class ScriptedGenerator : IGenerator { + + public System.Threading.Tasks.Task GenerateAsync(string prompt, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Generate(prompt)); } + public System.Threading.Tasks.Task> GenerateGroundedAsync(string query, IEnumerable> context, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GenerateGrounded(query, context)); } private readonly string _response; public List Prompts { get; } = new List(); diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/BM25RetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/BM25RetrieverTests.cs index aaa3ead7d4..3293dd337c 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/BM25RetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/BM25RetrieverTests.cs @@ -22,6 +22,15 @@ public class BM25RetrieverTests /// private class MockDocumentStore : IDocumentStore { + + public System.Threading.Tasks.Task AddAsync(VectorDocument vectorDocument, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Add(vectorDocument); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } + public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } + public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } + public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetAllAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetAll()); } private readonly List> _documents = new(); public int DocumentCount => _documents.Count; diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ColBERTRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ColBERTRetrieverTests.cs index fb62d2c25e..bb99379361 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ColBERTRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ColBERTRetrieverTests.cs @@ -22,6 +22,15 @@ public class ColBERTRetrieverTests /// private class MockDocumentStore : IDocumentStore { + + public System.Threading.Tasks.Task AddAsync(VectorDocument vectorDocument, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Add(vectorDocument); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } + public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } + public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } + public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetAllAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetAll()); } private readonly List> _documents = new(); public int DocumentCount => _documents.Count; diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DeStubbedLlmComponentsTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DeStubbedLlmComponentsTests.cs index ae88825344..d307d84524 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DeStubbedLlmComponentsTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DeStubbedLlmComponentsTests.cs @@ -33,6 +33,8 @@ public string Generate(string prompt) private sealed class RecordingRetriever : IRetriever { + + public System.Threading.Tasks.Task>> RetrieveAsync(string query, int topK, Dictionary? metadataFilters = null, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(metadataFilters == null ? Retrieve(query, topK) : Retrieve(query, topK, metadataFilters)); } public List Queries { get; } = new(); public int DefaultTopK => 5; public IEnumerable> Retrieve(string query) => Retrieve(query, 5, new Dictionary()); diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DenseRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DenseRetrieverTests.cs index ddd165c83d..ec65de9fc3 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DenseRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DenseRetrieverTests.cs @@ -71,6 +71,15 @@ public Task> EmbedBatchAsync(IEnumerable texts) /// private class MockDocumentStore : IDocumentStore { + + public System.Threading.Tasks.Task AddAsync(VectorDocument vectorDocument, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Add(vectorDocument); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } + public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } + public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } + public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetAllAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetAll()); } private readonly List> _documents = new(); private readonly MockEmbeddingModel _embeddingModel = new(); diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/FLARERetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/FLARERetrieverTests.cs index 00f86f370c..536c7c514f 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/FLARERetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/FLARERetrieverTests.cs @@ -42,6 +42,9 @@ protected override IEnumerable> RetrieveCore( // Mock generator that can be configured to produce specific outputs private class ConfigurableMockGenerator : IGenerator { + + public System.Threading.Tasks.Task GenerateAsync(string prompt, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Generate(prompt)); } + public System.Threading.Tasks.Task> GenerateGroundedAsync(string query, IEnumerable> context, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GenerateGrounded(query, context)); } private readonly Queue _responses; private readonly string _defaultResponse; public List GeneratePrompts { get; } = new List(); diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRAGTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRAGTests.cs index 56444dce4a..03a7210e95 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRAGTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRAGTests.cs @@ -42,6 +42,9 @@ protected override IEnumerable> RetrieveCore( // Mock generator for entity extraction private class EntityExtractionMockGenerator : IGenerator { + + public System.Threading.Tasks.Task GenerateAsync(string prompt, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Generate(prompt)); } + public System.Threading.Tasks.Task> GenerateGroundedAsync(string query, IEnumerable> context, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GenerateGrounded(query, context)); } private readonly string _entityExtractionResponse; public List GeneratePrompts { get; } = new List(); diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRetrieverTests.cs index 938eb3a1e7..22ba81fb34 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRetrieverTests.cs @@ -70,6 +70,15 @@ public Task> EmbedBatchAsync(IEnumerable texts) /// private class MockDocumentStore : IDocumentStore { + + public System.Threading.Tasks.Task AddAsync(VectorDocument vectorDocument, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Add(vectorDocument); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } + public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } + public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } + public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetAllAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetAll()); } private readonly List> _documents = new(); private readonly MockEmbeddingModel _embeddingModel = new(); diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/HybridRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/HybridRetrieverTests.cs index 5061066680..5b2ae42ded 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/HybridRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/HybridRetrieverTests.cs @@ -21,6 +21,8 @@ public class HybridRetrieverTests /// private class MockRetriever : IRetriever { + + public System.Threading.Tasks.Task>> RetrieveAsync(string query, int topK, Dictionary? metadataFilters = null, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(metadataFilters == null ? Retrieve(query, topK) : Retrieve(query, topK, metadataFilters)); } private readonly List> _documents; private readonly Func, double>? _scoringFunction; diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiQueryRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiQueryRetrieverTests.cs index 6926c419ae..0e2a9cf66b 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiQueryRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiQueryRetrieverTests.cs @@ -21,6 +21,8 @@ public class MultiQueryRetrieverTests /// private class MockRetriever : IRetriever { + + public System.Threading.Tasks.Task>> RetrieveAsync(string query, int topK, Dictionary? metadataFilters = null, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(metadataFilters == null ? Retrieve(query, topK) : Retrieve(query, topK, metadataFilters)); } private readonly List> _documents; private readonly Func, double>? _scoringFunction; private readonly List _queriesReceived = new(); diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiVectorRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiVectorRetrieverTests.cs index 9bf5934826..824b02548f 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiVectorRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiVectorRetrieverTests.cs @@ -23,6 +23,15 @@ public class MultiVectorRetrieverTests /// private class MockDocumentStore : IDocumentStore { + + public System.Threading.Tasks.Task AddAsync(VectorDocument vectorDocument, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Add(vectorDocument); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } + public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } + public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } + public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetAllAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetAll()); } private readonly List> _documents = new(); public int DocumentCount => _documents.Count; diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ParentDocumentRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ParentDocumentRetrieverTests.cs index b0b2bed63c..02dd4f391c 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ParentDocumentRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ParentDocumentRetrieverTests.cs @@ -70,6 +70,15 @@ public Task> EmbedBatchAsync(IEnumerable texts) /// private class MockDocumentStore : IDocumentStore { + + public System.Threading.Tasks.Task AddAsync(VectorDocument vectorDocument, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Add(vectorDocument); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } + public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } + public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } + public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetAllAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetAll()); } private readonly List> _documents = new(); private readonly Dictionary> _parentDocuments = new(); diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/RagAsyncCancellationTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/RagAsyncCancellationTests.cs new file mode 100644 index 0000000000..1512060280 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/RagAsyncCancellationTests.cs @@ -0,0 +1,212 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using AiDotNet.Interfaces; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.ContextCompression; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Generators; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using AiDotNet.RetrievalAugmentedGeneration.Rerankers; +using AiDotNet.RetrievalAugmentedGeneration.Retrievers; +using Xunit; + +namespace AiDotNet.Tests.UnitTests.RetrievalAugmentedGeneration; + +/// +/// Focused tests for the async + cancellation surface added across the RAG pipeline (task #21). +/// Each test asserts that the *Async methods produce the same result as their synchronous +/// counterparts and that a pre-cancelled causes an +/// to be thrown. +/// +public class RagAsyncCancellationTests +{ + private static Document Doc(string id, string content) + => new Document(id, content, new Dictionary()); + + // ---- Fakes built on the RAG base classes so we exercise the real async plumbing ---- + + private sealed class FakeRetriever : RetrieverBase + { + public FakeRetriever() : base(5) { } + protected override IEnumerable> RetrieveCore(string query, int topK, Dictionary metadataFilters) + => new List> { Doc("r1", "retrieved: " + query) }.Take(topK); + } + + private sealed class FakeReranker : RerankerBase + { + public override bool ModifiesScores => true; + protected override IEnumerable> RerankCore(string query, IList> documents) + => documents.Reverse(); + } + + private sealed class FakeCompressor : ContextCompressorBase + { + protected override List> CompressCore(List> documents, string query, Dictionary? options = null) + => documents.Select(d => Doc(d.Id, d.Content.Substring(0, Math.Min(3, d.Content.Length)))).ToList(); + } + + private sealed class FakeGenerator : GeneratorBase + { + public FakeGenerator() : base(2048, 512) { } + protected override string GenerateCore(string prompt) => "answer[1]"; + } + + private sealed class FakeStore : DocumentStoreBase + { + private readonly Dictionary> _docs = new(); + public override int DocumentCount => _docs.Count; + public override int VectorDimension => 2; + protected override void AddCore(VectorDocument vectorDocument) => _docs[vectorDocument.Document.Id] = vectorDocument; + protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + => _docs.Values.Select(v => v.Document).Take(topK); + protected override Document? GetByIdCore(string documentId) => _docs.TryGetValue(documentId, out var v) ? v.Document : null; + protected override bool RemoveCore(string documentId) => _docs.Remove(documentId); + public override void Clear() => _docs.Clear(); + protected override IEnumerable> GetAllCore() => _docs.Values.Select(v => v.Document).ToList(); + } + + private static VectorDocument VDoc(string id) + => new VectorDocument(Doc(id, "c-" + id), new Vector(new[] { 1.0, 0.0 })); + + private static CancellationToken Cancelled() + { + var cts = new CancellationTokenSource(); + cts.Cancel(); + return cts.Token; + } + + // ---------------------------- Retriever ---------------------------- + + [Fact] + public async Task RetrieveAsync_ReturnsExpectedResults() + { + var retriever = new FakeRetriever(); + var results = (await retriever.RetrieveAsync("hello", 5)).ToList(); + Assert.Single(results); + Assert.Equal("retrieved: hello", results[0].Content); + } + + [Fact] + public async Task RetrieveAsync_PreCancelledToken_Throws() + { + var retriever = new FakeRetriever(); + await Assert.ThrowsAnyAsync( + () => retriever.RetrieveAsync("hello", 5, null, Cancelled())); + } + + // ---------------------------- Reranker ---------------------------- + + [Fact] + public async Task RerankAsync_ReturnsReorderedResults() + { + var reranker = new FakeReranker(); + var input = new List> { Doc("a", "a"), Doc("b", "b") }; + var results = (await reranker.RerankAsync("q", input)).ToList(); + Assert.Equal("b", results[0].Id); + Assert.Equal("a", results[1].Id); + } + + [Fact] + public async Task RerankAsync_PreCancelledToken_Throws() + { + var reranker = new FakeReranker(); + var input = new List> { Doc("a", "a"), Doc("b", "b") }; + await Assert.ThrowsAnyAsync( + () => reranker.RerankAsync("q", input, Cancelled())); + } + + // ------------------------- Context compressor ------------------------- + + [Fact] + public async Task CompressAsync_ReturnsCompressedResults() + { + var compressor = new FakeCompressor(); + var input = new List> { Doc("a", "abcdef") }; + var results = await compressor.CompressAsync(input, "q"); + Assert.Equal("abc", results[0].Content); + } + + [Fact] + public async Task CompressAsync_PreCancelledToken_Throws() + { + var compressor = new FakeCompressor(); + var input = new List> { Doc("a", "abcdef") }; + await Assert.ThrowsAnyAsync( + () => compressor.CompressAsync(input, "q", null, Cancelled())); + } + + // ---------------------------- Generator ---------------------------- + + [Fact] + public async Task GenerateAsync_ReturnsText() + { + var generator = new FakeGenerator(); + Assert.Equal("answer[1]", await generator.GenerateAsync("prompt")); + } + + [Fact] + public async Task GenerateGroundedAsync_ReturnsGroundedAnswer() + { + var generator = new FakeGenerator(); + var answer = await generator.GenerateGroundedAsync("q", new[] { Doc("s1", "src") }); + Assert.Equal("answer[1]", answer.Answer); + Assert.Contains("s1", answer.Citations); + } + + [Fact] + public async Task GenerateAsync_PreCancelledToken_Throws() + { + var generator = new FakeGenerator(); + await Assert.ThrowsAnyAsync( + () => generator.GenerateAsync("prompt", Cancelled())); + } + + [Fact] + public async Task GenerateGroundedAsync_PreCancelledToken_Throws() + { + var generator = new FakeGenerator(); + await Assert.ThrowsAnyAsync( + () => generator.GenerateGroundedAsync("q", new[] { Doc("s1", "src") }, Cancelled())); + } + + // ------------------------- Document store ------------------------- + + [Fact] + public async Task DocumentStore_AsyncRoundTrip_Works() + { + var store = new FakeStore(); + await store.AddAsync(VDoc("d1")); + await store.AddBatchAsync(new[] { VDoc("d2"), VDoc("d3") }); + + Assert.Equal(3, store.DocumentCount); + + var byId = await store.GetByIdAsync("d2"); + Assert.NotNull(byId); + Assert.Equal("d2", byId!.Id); + + var all = (await store.GetAllAsync()).ToList(); + Assert.Equal(3, all.Count); + + var similar = (await store.GetSimilarAsync(new Vector(new[] { 1.0, 0.0 }), 2)).ToList(); + Assert.Equal(2, similar.Count); + + Assert.True(await store.RemoveAsync("d1")); + await store.ClearAsync(); + Assert.Equal(0, store.DocumentCount); + } + + [Fact] + public async Task DocumentStore_PreCancelledToken_Throws() + { + var store = new FakeStore(); + await Assert.ThrowsAnyAsync( + () => store.AddAsync(VDoc("d1"), Cancelled())); + await Assert.ThrowsAnyAsync( + () => store.GetSimilarAsync(new Vector(new[] { 1.0, 0.0 }), 2, Cancelled())); + await Assert.ThrowsAnyAsync( + () => store.GetAllAsync(Cancelled())); + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/RetrieverFusionAndRoutingTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/RetrieverFusionAndRoutingTests.cs index 6fafd3b725..1e51753d73 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/RetrieverFusionAndRoutingTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/RetrieverFusionAndRoutingTests.cs @@ -16,6 +16,8 @@ public class RetrieverFusionAndRoutingTests { private sealed class FakeRetriever : IRetriever { + + public System.Threading.Tasks.Task>> RetrieveAsync(string query, int topK, Dictionary? metadataFilters = null, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(metadataFilters == null ? Retrieve(query, topK) : Retrieve(query, topK, metadataFilters)); } private readonly List> _docs; public string? LastQuery { get; private set; } public Dictionary? LastFilters { get; private set; } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/TestHelpers.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/TestHelpers.cs index 79c956c170..6fbfdcd8ac 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/TestHelpers.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/TestHelpers.cs @@ -141,6 +141,15 @@ public IEnumerable> GetAll() ).ToList(); } + public System.Threading.Tasks.Task AddAsync(VectorDocument vectorDocument, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Add(vectorDocument); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } + public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } + public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } + public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetAllAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetAll()); } + private bool MatchesFilters(VectorDocument document, Dictionary filters) { if (filters == null || filters.Count == 0) diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/SelfCorrectingRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/SelfCorrectingRetrieverTests.cs index 9fb81f7816..8246ee6401 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/SelfCorrectingRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/SelfCorrectingRetrieverTests.cs @@ -60,6 +60,9 @@ protected override IEnumerable> RetrieveCore( // Mock generator that can simulate different critique responses private class CritiqueMockGenerator : IGenerator { + + public System.Threading.Tasks.Task GenerateAsync(string prompt, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Generate(prompt)); } + public System.Threading.Tasks.Task> GenerateGroundedAsync(string query, IEnumerable> context, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GenerateGrounded(query, context)); } private readonly Queue _generateResponses; private readonly Queue _groundedAnswers; public List GeneratePrompts { get; } = new List(); diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/TFIDFRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/TFIDFRetrieverTests.cs index 96222db5f9..b260d1d0e9 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/TFIDFRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/TFIDFRetrieverTests.cs @@ -22,6 +22,15 @@ public class TFIDFRetrieverTests /// private class MockDocumentStore : IDocumentStore { + + public System.Threading.Tasks.Task AddAsync(VectorDocument vectorDocument, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Add(vectorDocument); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } + public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } + public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } + public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetAllAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetAll()); } private readonly List> _documents = new(); public int DocumentCount => _documents.Count; diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorRetrieverTests.cs index 855f639896..adb772c3fd 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorRetrieverTests.cs @@ -90,6 +90,15 @@ public Task> EmbedBatchAsync(IEnumerable texts) /// private class MockDocumentStore : IDocumentStore { + + public System.Threading.Tasks.Task AddAsync(VectorDocument vectorDocument, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Add(vectorDocument); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } + public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } + public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } + public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } + public System.Threading.Tasks.Task>> GetAllAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetAll()); } private readonly List<(Document Doc, Vector Embedding)> _documents = new(); public int DocumentCount => _documents.Count; From db1b88a76d331fe55d0660e6e60374d05efce10b Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 19:06:19 -0400 Subject: [PATCH 19/25] test(rag): update Hybrid/MultiQuery tests for RRF fusion + real query variations Seven pre-existing HybridRetrieverTests/MultiQueryRetrieverTests asserted the old behavior that the RRF-fusion (#27) and real-query-expansion (#19) changes correctly replaced: exact weighted-sum scores (0.7*d+0.3*s) and the "{query} variation N" placeholder. RRF is rank-based (scores are small, order is what matters) and MultiQuery now emits real reformulations. Reassert the meaningful properties: fused doc presence + ranking, and original-query-first with distinct variations. All 105 pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../HybridRetrieverTests.cs | 26 ++++++++++++------- .../MultiQueryRetrieverTests.cs | 17 ++++++------ 2 files changed, 25 insertions(+), 18 deletions(-) diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/HybridRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/HybridRetrieverTests.cs index 5b2ae42ded..97bd26abf5 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/HybridRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/HybridRetrieverTests.cs @@ -318,10 +318,10 @@ public async Task Retrieve_OverlappingResults_CombinesScores() // Assert Assert.Equal(3, results.Count); - // doc1 should have combined score: 0.9 * 0.7 + 0.8 * 0.3 = 0.63 + 0.24 = 0.87 + // doc1 appears in BOTH lists, so RRF fuses its two rank contributions and it ranks first. var doc1 = results.First(r => r.Id == "doc1"); Assert.True(doc1.HasRelevanceScore); - Assert.True(doc1.RelevanceScore > 0.8); // Combined should be higher than either alone + Assert.Equal("doc1", results[0].Id); } [Fact(Timeout = 60000)] @@ -452,8 +452,9 @@ public async Task Retrieve_EqualWeights_ProducesEqualContribution() // Assert Assert.Single(results); - // Score should be 1.0 * 0.5 + 1.0 * 0.5 = 1.0 - Assert.Equal(1.0, results[0].RelevanceScore, 3); + // Equal weights → both lists contribute equally to the RRF fusion; the doc is returned. + Assert.Equal("doc1", results[0].Id); + Assert.True(results[0].RelevanceScore > 0); } [Fact(Timeout = 60000)] @@ -497,8 +498,9 @@ public async Task Retrieve_ZeroDenseWeight_OnlyUsesSparseScore() // Assert Assert.Single(results); - // Score should be 1.0 * 0.0 + 0.5 * 1.0 = 0.5 - Assert.Equal(0.5, results[0].RelevanceScore, 3); + // With zero dense weight only the sparse ranking contributes to RRF; the doc is still returned. + Assert.Equal("doc1", results[0].Id); + Assert.True(results[0].RelevanceScore > 0); } [Fact(Timeout = 60000)] @@ -519,8 +521,10 @@ public async Task Retrieve_ZeroSparseWeight_OnlyUsesDenseScore() // Assert Assert.Single(results); - // Score should be 0.8 * 1.0 + 1.0 * 0.0 = 0.8 - Assert.Equal(0.8, results[0].RelevanceScore, 3); + // With zero sparse weight the sparse list contributes nothing to the RRF fusion; the doc is + // still returned via the dense ranking with a positive fused score. + Assert.Equal("doc1", results[0].Id); + Assert.True(results[0].RelevanceScore > 0); } #endregion @@ -776,8 +780,10 @@ public async Task Constructor_DefaultWeights_Uses70Dense30Sparse() // Assert Assert.Single(results); - // Default weights: 0.7 dense + 0.3 sparse = 1.0 * 0.7 + 1.0 * 0.3 = 1.0 - Assert.Equal(1.0, results[0].RelevanceScore, 3); + // Hybrid fusion is now Reciprocal Rank Fusion (rank-based, scale-free) rather than a raw + // weighted sum of scores. A doc ranked #1 in both lists is returned with a positive fused score. + Assert.Equal("doc1", results[0].Id); + Assert.True(results[0].RelevanceScore > 0); } [Fact(Timeout = 60000)] diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiQueryRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiQueryRetrieverTests.cs index 0e2a9cf66b..ef7d803dde 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiQueryRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiQueryRetrieverTests.cs @@ -333,9 +333,11 @@ public async Task Retrieve_QueriesContainVariations() var results = retriever.Retrieve("test").ToList(); // Assert - // Should have variations like "test variation 1", "test variation 2" - Assert.Contains(baseRetriever.QueriesReceived, q => q.Contains("variation 1")); - Assert.Contains(baseRetriever.QueriesReceived, q => q.Contains("variation 2")); + // MultiQuery now produces real query variations (LLM when a generator is supplied, deterministic + // template reformulations otherwise) rather than the old "{query} variation N" placeholder. The + // base retriever should see the original query plus additional distinct variations. + Assert.Contains("test", baseRetriever.QueriesReceived); + Assert.True(baseRetriever.QueriesReceived.Distinct().Count() >= 3); } [Fact(Timeout = 60000)] @@ -698,11 +700,10 @@ public async Task Retrieve_VariationsAppendToOriginal() var results = retriever.Retrieve("my query").ToList(); // Assert - // Variations should start with original query - foreach (var query in baseRetriever.QueriesReceived.Skip(1)) - { - Assert.StartsWith("my query", query); - } + // The original query is retrieved first; the remaining queries are distinct reformulations + // (which embed, but no longer necessarily start with, the original text). + Assert.Equal("my query", baseRetriever.QueriesReceived.First()); + Assert.True(baseRetriever.QueriesReceived.Distinct().Count() >= 2); } #endregion From 0ec5e8edb1950fc8b2cc30d1c6bfbd92c29eaa4a Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 19:15:04 -0400 Subject: [PATCH 20/25] feat(rag): central RagPipeline orchestrator with tenant isolation (#33) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RAG subsystem had every component but no single class that composed them — callers wired chunk→embed→store and retrieve→rerank→compress→generate by hand. Add RagPipeline: - IngestAsync: optional chunk → embed → upsert (returns chunk count). - QueryAsync: retrieve → optional rerank → optional compress → optional generate, returning a RagResult (contexts + grounded answer). Fully async/cancellation-aware. Optional tenant/namespace: stamped on ingested documents' metadata and applied as a retrieval filter for cross-tenant isolation. Streaming is already provided by IStreamingGenerator/ChatClientGenerator. Verified: 3 tests (chunked ingest into a real InMemoryDocumentStore, query→generate, tenant stamp + filter). Remaining #33 items — rich boolean/$in/range metadata filters across every external store, and RAG-path GPU hooks — are larger enhancements tracked as follow-ups. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../RagPipeline.cs | 147 ++++++++++++++++++ .../RagPipelineTests.cs | 104 +++++++++++++ 2 files changed, 251 insertions(+) create mode 100644 src/RetrievalAugmentedGeneration/RagPipeline.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/RagPipelineTests.cs diff --git a/src/RetrievalAugmentedGeneration/RagPipeline.cs b/src/RetrievalAugmentedGeneration/RagPipeline.cs new file mode 100644 index 0000000000..1a320a6966 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/RagPipeline.cs @@ -0,0 +1,147 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration.ChunkingStrategies; +using AiDotNet.RetrievalAugmentedGeneration.Models; + +namespace AiDotNet.RetrievalAugmentedGeneration +{ + /// + /// End-to-end RAG orchestrator that composes the individual components into the two operations an + /// application actually performs: (chunk → embed → store) and + /// (retrieve → rerank → compress → generate). Previously callers had to wire + /// these stages by hand; this is the single entry point that ties them together, fully async and + /// cancellation-aware, with optional tenant/namespace isolation. + /// + /// The numeric type used for embeddings and scoring. + public sealed class RagPipeline + { + private readonly IEmbeddingModel _embedding; + private readonly IDocumentStore _store; + private readonly IRetriever _retriever; + private readonly IChunkingStrategy? _chunking; + private readonly IReranker? _reranker; + private readonly IContextCompressor? _compressor; + private readonly IGenerator? _generator; + private readonly string? _tenant; + + /// Embedding model used for ingestion (and dense retrieval upstream). + /// Vector store documents are upserted into. + /// Retriever used at query time. + /// Optional chunker; when null each document is ingested whole. + /// Optional reranker applied to retrieved documents. + /// Optional context compressor applied before generation. + /// Optional generator; when null returns retrieved context only. + /// + /// Optional tenant/namespace. When set it is stamped on every ingested document's metadata under + /// _tenant and used as a retrieval filter, isolating this pipeline's data from other tenants. + /// + public RagPipeline( + IEmbeddingModel embedding, + IDocumentStore store, + IRetriever retriever, + IChunkingStrategy? chunking = null, + IReranker? reranker = null, + IContextCompressor? compressor = null, + IGenerator? generator = null, + string? tenant = null) + { + _embedding = embedding ?? throw new ArgumentNullException(nameof(embedding)); + _store = store ?? throw new ArgumentNullException(nameof(store)); + _retriever = retriever ?? throw new ArgumentNullException(nameof(retriever)); + _chunking = chunking; + _reranker = reranker; + _compressor = compressor; + _generator = generator; + _tenant = string.IsNullOrWhiteSpace(tenant) ? null : tenant; + } + + /// Metadata key under which the tenant/namespace is stored on ingested documents. + public const string TenantMetadataKey = "_tenant"; + + /// + /// Ingests a document: optionally chunks it, embeds each chunk, and upserts it into the store. + /// + /// The number of chunks stored. + public async Task IngestAsync( + string id, string content, IReadOnlyDictionary? metadata = null, CancellationToken cancellationToken = default) + { + if (string.IsNullOrEmpty(id)) throw new ArgumentException("Document id is required.", nameof(id)); + if (string.IsNullOrEmpty(content)) throw new ArgumentException("Document content is required.", nameof(content)); + + var chunks = _chunking != null ? _chunking.Chunk(content).ToList() : new List { content }; + int index = 0; + foreach (var chunk in chunks) + { + cancellationToken.ThrowIfCancellationRequested(); + + var meta = new Dictionary(); + if (metadata != null) + { + foreach (var kv in metadata) meta[kv.Key] = kv.Value; + } + if (_tenant != null) meta[TenantMetadataKey] = _tenant; + + var doc = new Document(chunks.Count == 1 ? id : $"{id}::{index}", chunk) { Metadata = meta }; + var embedding = await _embedding.EmbedAsync(chunk).ConfigureAwait(false); + await _store.AddAsync(new VectorDocument(doc, embedding), cancellationToken).ConfigureAwait(false); + index++; + } + + return chunks.Count; + } + + /// + /// Answers a question: retrieve → (rerank) → (compress) → (generate). + /// + public async Task> QueryAsync(string question, int topK = 5, CancellationToken cancellationToken = default) + { + if (string.IsNullOrEmpty(question)) throw new ArgumentException("Question is required.", nameof(question)); + + var filters = new Dictionary(); + if (_tenant != null) filters[TenantMetadataKey] = _tenant; + + var retrieved = (await _retriever.RetrieveAsync(question, topK, filters, cancellationToken).ConfigureAwait(false)).ToList(); + + if (_reranker != null) + { + retrieved = (await _reranker.RerankAsync(question, retrieved, cancellationToken).ConfigureAwait(false)).ToList(); + } + + var context = _compressor != null + ? await _compressor.CompressAsync(retrieved, question, null, cancellationToken).ConfigureAwait(false) + : retrieved; + + GroundedAnswer? answer = null; + if (_generator != null && context.Count > 0) + { + answer = await _generator.GenerateGroundedAsync(question, context, cancellationToken).ConfigureAwait(false); + } + + return new RagResult(question, context, answer); + } + } + + /// The outcome of a call. + public sealed class RagResult + { + public RagResult(string question, List> contexts, GroundedAnswer? answer) + { + Question = question; + Contexts = contexts; + Answer = answer; + } + + /// The original question. + public string Question { get; } + + /// The retrieved (and reranked/compressed) context documents. + public List> Contexts { get; } + + /// The generated grounded answer, or null when no generator was configured. + public GroundedAnswer? Answer { get; } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/RagPipelineTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/RagPipelineTests.cs new file mode 100644 index 0000000000..998b8ea795 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/RagPipelineTests.cs @@ -0,0 +1,104 @@ +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using AiDotNet.Interfaces; +using AiDotNet.RetrievalAugmentedGeneration; +using AiDotNet.RetrievalAugmentedGeneration.ChunkingStrategies; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Embeddings; +using AiDotNet.RetrievalAugmentedGeneration.Generators; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Xunit; + +namespace AiDotNet.Tests.UnitTests.RetrievalAugmentedGeneration; + +/// +/// Tests for the end-to-end orchestrator: ingest (chunk→embed→store), query +/// (retrieve→rerank→compress→generate), and tenant/namespace isolation. +/// +public class RagPipelineTests +{ + private const int Dim = 8; + + // Minimal retriever that records the query + filters it was called with and returns a preset doc. + private sealed class RecordingRetriever : IRetriever + { + public string? LastQuery { get; private set; } + public Dictionary? LastFilters { get; private set; } + public int DefaultTopK => 5; + + public IEnumerable> Retrieve(string query) => Retrieve(query, DefaultTopK, new Dictionary()); + public IEnumerable> Retrieve(string query, int topK) => Retrieve(query, topK, new Dictionary()); + public IEnumerable> Retrieve(string query, int topK, Dictionary metadataFilters) + { + LastQuery = query; + LastFilters = metadataFilters; + return new[] { new Document("d1", "retrieved context") { RelevanceScore = 1.0, HasRelevanceScore = true } }; + } + + public Task>> RetrieveAsync(string query, int topK, Dictionary? metadataFilters = null, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(Retrieve(query, topK, metadataFilters ?? new Dictionary())); + } + } + + private static RagPipeline CreatePipeline(RecordingRetriever retriever, string? tenant = null, IChunkingStrategy? chunking = null) + => new RagPipeline( + embedding: new StubEmbeddingModel(Dim), + store: new InMemoryDocumentStore(Dim), + retriever: retriever, + chunking: chunking, + generator: new StubGenerator(), + tenant: tenant); + + [Fact] + public async Task IngestAsync_StoresChunkedEmbeddedDocument() + { + var store = new InMemoryDocumentStore(Dim); + var pipeline = new RagPipeline( + new StubEmbeddingModel(Dim), store, new RecordingRetriever(), + chunking: new TokenBasedChunkingStrategy(maxTokens: 3, overlapTokens: 0)); + + int stored = await pipeline.IngestAsync("doc1", "alpha beta gamma delta epsilon zeta eta theta"); + + Assert.True(stored >= 2, "long content should be split into multiple chunks"); + Assert.Equal(stored, store.DocumentCount); + } + + [Fact] + public async Task QueryAsync_RunsRetrieveAndGenerate() + { + var retriever = new RecordingRetriever(); + var pipeline = CreatePipeline(retriever); + + var result = await pipeline.QueryAsync("what is X?", topK: 3); + + Assert.Equal("what is X?", retriever.LastQuery); + Assert.Single(result.Contexts); + Assert.Equal("d1", result.Contexts[0].Id); + Assert.NotNull(result.Answer); // StubGenerator produced a grounded answer + } + + [Fact] + public async Task Tenant_IsStampedOnIngestAndUsedAsRetrievalFilter() + { + var retriever = new RecordingRetriever(); + var store = new InMemoryDocumentStore(Dim); + var pipeline = new RagPipeline( + new StubEmbeddingModel(Dim), store, retriever, tenant: "acme"); + + await pipeline.IngestAsync("doc1", "hello world"); + await pipeline.QueryAsync("q"); + + // The retrieval filter carries the tenant so cross-tenant data is isolated. + Assert.True(retriever.LastFilters!.ContainsKey(RagPipeline.TenantMetadataKey)); + Assert.Equal("acme", retriever.LastFilters![RagPipeline.TenantMetadataKey]); + + // And the stored document was stamped with the tenant. + var all = store.GetAll().ToList(); + Assert.Single(all); + Assert.Equal("acme", all[0].Metadata[RagPipeline.TenantMetadataKey]); + } +} From f1b2e883128efd8b0a886fe5a171aa74d0e89eca Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 21:44:15 -0400 Subject: [PATCH 21/25] feat(rag): neo4j property-graph backend for GraphRAG Add AiDotNet.Storage.Neo4j, an opt-in metapackage implementing IGraphStore over the Neo4j C# driver so GraphRAG can run on a production property graph. Nodes map to labeled Neo4j nodes, edges to relationships; upserts use MERGE and traversal uses Cypher MATCH. All per-request data is passed as $parameters; the only literal identifiers (base node label, relationship type) are validated and backtick-escaped once at construction, so there is no injection surface. Pure Cypher/param/mapping logic is factored into Neo4jCypher for unit testing. net10.0-only (the driver's TFM), matching the Postgres/Redis storage projects; excluded from the core csproj glob and added to the solution + test project. Tests: 20 pure-logic unit tests plus gated [Trait(Category,Integration)] SkippableFact round-trip tests reading NEO4J_URI/NEO4J_USER/NEO4J_PASSWORD. Co-Authored-By: Claude Opus 4.8 (1M context) --- AiDotNet.sln | 15 + Directory.Packages.props | 1 + .../AiDotNet.Storage.Neo4j.csproj | 46 ++ src/AiDotNet.Storage.Neo4j/GlobalUsings.cs | 15 + src/AiDotNet.Storage.Neo4j/Neo4jCypher.cs | 188 +++++++ src/AiDotNet.Storage.Neo4j/Neo4jGraphStore.cs | 470 ++++++++++++++++++ src/AiDotNet.csproj | 4 + tests/AiDotNet.Tests/AiDotNetTests.csproj | 1 + .../Neo4jGraphStoreTests.cs | 321 ++++++++++++ 9 files changed, 1061 insertions(+) create mode 100644 src/AiDotNet.Storage.Neo4j/AiDotNet.Storage.Neo4j.csproj create mode 100644 src/AiDotNet.Storage.Neo4j/GlobalUsings.cs create mode 100644 src/AiDotNet.Storage.Neo4j/Neo4jCypher.cs create mode 100644 src/AiDotNet.Storage.Neo4j/Neo4jGraphStore.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Neo4jGraphStoreTests.cs diff --git a/AiDotNet.sln b/AiDotNet.sln index 05f23be226..46453d8ce6 100644 --- a/AiDotNet.sln +++ b/AiDotNet.sln @@ -39,6 +39,8 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AiDotNet.Privacy.HE", "src\ EndProject Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AiDotNet.Storage.Sqlite", "src\AiDotNet.Storage.Sqlite\AiDotNet.Storage.Sqlite.csproj", "{251CF282-40BA-452B-A563-FAE6B9826B8B}" EndProject +Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AiDotNet.Storage.Neo4j", "src\AiDotNet.Storage.Neo4j\AiDotNet.Storage.Neo4j.csproj", "{AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}" +EndProject Global GlobalSection(SolutionConfigurationPlatforms) = preSolution Debug|Any CPU = Debug|Any CPU @@ -217,6 +219,18 @@ Global {251CF282-40BA-452B-A563-FAE6B9826B8B}.Release|x64.Build.0 = Release|Any CPU {251CF282-40BA-452B-A563-FAE6B9826B8B}.Release|x86.ActiveCfg = Release|Any CPU {251CF282-40BA-452B-A563-FAE6B9826B8B}.Release|x86.Build.0 = Release|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Debug|Any CPU.ActiveCfg = Debug|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Debug|Any CPU.Build.0 = Debug|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Debug|x64.ActiveCfg = Debug|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Debug|x64.Build.0 = Debug|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Debug|x86.ActiveCfg = Debug|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Debug|x86.Build.0 = Debug|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Release|Any CPU.ActiveCfg = Release|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Release|Any CPU.Build.0 = Release|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Release|x64.ActiveCfg = Release|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Release|x64.Build.0 = Release|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Release|x86.ActiveCfg = Release|Any CPU + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}.Release|x86.Build.0 = Release|Any CPU EndGlobalSection GlobalSection(SolutionProperties) = preSolution HideSolutionNode = FALSE @@ -230,6 +244,7 @@ Global {B1427BFB-8A17-4735-8EC5-E594A612AB81} = {66320409-64EC-F7C5-3DEF-65E7510DAAD1} {D71746F5-5342-474C-B23E-4A07B29FA4F7} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} {251CF282-40BA-452B-A563-FAE6B9826B8B} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} + {AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} EndGlobalSection GlobalSection(ExtensibilityGlobals) = postSolution SolutionGuid = {99F5DE1C-2F44-4B2D-B8B3-DEFFC86D2F16} diff --git a/Directory.Packages.props b/Directory.Packages.props index 7413615796..19ab29f26b 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -307,6 +307,7 @@ + diff --git a/src/AiDotNet.Storage.Neo4j/AiDotNet.Storage.Neo4j.csproj b/src/AiDotNet.Storage.Neo4j/AiDotNet.Storage.Neo4j.csproj new file mode 100644 index 0000000000..714a48d0d8 --- /dev/null +++ b/src/AiDotNet.Storage.Neo4j/AiDotNet.Storage.Neo4j.csproj @@ -0,0 +1,46 @@ + + + + + net10.0 + enable + enable + True + 0.204.0 + AiDotNet.Storage.Neo4j + Neo4j property-graph backend for AiDotNet's GraphRAG / knowledge-graph pipeline. Implements IGraphStore over the Neo4j C# driver (nodes as labeled nodes, edges as relationships, MERGE upserts, Cypher MATCH traversal). Opt-in: reference this package to keep Neo4j.Driver out of the core package. Licensed under BSL 1.1 — see https://aidotnet.dev/license. + Ooples Finance + ooples + Ooples Finance LLC 2023 + https://github.com/ooples/AiDotNet + git + https://github.com/ooples/AiDotNet + ai; aidotnet; neo4j; graph; graphrag; knowledge-graph; property-graph; cypher; storage + LICENSE + True + latest + AiDotNet.RetrievalAugmentedGeneration.Graph + + + + + + + + + + + + + + + + diff --git a/src/AiDotNet.Storage.Neo4j/GlobalUsings.cs b/src/AiDotNet.Storage.Neo4j/GlobalUsings.cs new file mode 100644 index 0000000000..4069439749 --- /dev/null +++ b/src/AiDotNet.Storage.Neo4j/GlobalUsings.cs @@ -0,0 +1,15 @@ +// Mirrors the AiDotNet core + auto-generated AiDotNet.Tensors usings so this +// extracted package compiles without per-file boilerplate. See +// src/GlobalUsings.cs and src/obj/Release//AiDotNet.GlobalUsings.g.cs. +global using AiDotNet.Engines; +global using AiDotNet.Enums; +global using AiDotNet.Interfaces; +global using AiDotNet.LinearAlgebra; +global using AiDotNet.MetaLearning.Models; +global using AiDotNet.NeuralNetworks.Layers; +global using AiDotNet.Tensors.Engines; +global using AiDotNet.Tensors.Helpers; +global using AiDotNet.Tensors.Interfaces; +global using AiDotNet.Tensors.LinearAlgebra; +global using AiDotNet.Tensors.NumericOperations; +global using System.Globalization; diff --git a/src/AiDotNet.Storage.Neo4j/Neo4jCypher.cs b/src/AiDotNet.Storage.Neo4j/Neo4jCypher.cs new file mode 100644 index 0000000000..91274afc7e --- /dev/null +++ b/src/AiDotNet.Storage.Neo4j/Neo4jCypher.cs @@ -0,0 +1,188 @@ +using System; +using System.Collections.Generic; +using Newtonsoft.Json; + +namespace AiDotNet.RetrievalAugmentedGeneration.Graph; + +/// +/// Pure-logic helpers for the Neo4j graph backend: Cypher identifier escaping, the parameterized +/// Cypher statements used by , and lossless (de)serialization of the +/// property/temporal model. Kept free of any Neo4j.Driver types so the query/param construction can be +/// unit-tested without a live database. +/// +/// +/// Injection safety: Cypher labels and relationship types cannot be supplied as query +/// parameters — they must appear as literal identifiers. Every such identifier used by this backend is a +/// value configured once at construction (the node base label and the relationship type), never +/// attacker-controlled per-request data, and is passed through which +/// backtick-quotes it and doubles any embedded backticks (the standard, complete Cypher-identifier +/// escaping). Every per-operation value — node ids, the entity type, properties, embeddings, weights, +/// timestamps — is bound as a real Cypher $parameter, so no user data is ever concatenated into a +/// query string. +/// +public sealed class Neo4jCypher +{ + // Property keys written onto Neo4j nodes/relationships. Kept as constants so the mapper and the + // Cypher builders agree on the exact schema. + internal const string PropId = "id"; + internal const string PropLabel = "label"; + internal const string PropProps = "props"; + internal const string PropEmbedding = "embedding"; + internal const string PropRelationType = "relationType"; + internal const string PropWeight = "weight"; + internal const string PropCreatedAt = "createdAt"; + internal const string PropUpdatedAt = "updatedAt"; + internal const string PropValidFrom = "validFrom"; + internal const string PropValidUntil = "validUntil"; + + /// The backtick-escaped node label literal (e.g. `Entity`). + public string NodeLabel { get; } + + /// The backtick-escaped relationship-type literal (e.g. `RELATED`). + public string RelationshipType { get; } + + /// + /// Creates a Cypher builder for the given node base label and relationship type. Both are validated + /// and escaped once here; the escaped forms are the only literal identifiers that ever reach a query. + /// + /// The base label applied to every node (default Entity). Acts as the + /// per-namespace partition when different stores use different labels. + /// The relationship type applied to every edge (default + /// RELATED). + public Neo4jCypher(string nodeLabel = "Entity", string relationshipType = "RELATED") + { + NodeLabel = EscapeIdentifier(nodeLabel); + RelationshipType = EscapeIdentifier(relationshipType); + } + + /// + /// Escapes a Cypher identifier (label or relationship type) by wrapping it in backticks and doubling + /// any embedded backticks, which makes it impossible for the value to break out of the identifier + /// token. Rejects null/empty/whitespace and values containing a NUL character. + /// + public static string EscapeIdentifier(string identifier) + { + if (string.IsNullOrWhiteSpace(identifier)) + { + throw new ArgumentException("Cypher identifier cannot be null, empty, or whitespace.", nameof(identifier)); + } + + if (identifier.IndexOf('\0') >= 0) + { + throw new ArgumentException("Cypher identifier cannot contain a NUL character.", nameof(identifier)); + } + + return "`" + identifier.Replace("`", "``") + "`"; + } + + // ---- Cypher statements. All per-operation data is bound via $parameters. ---- + + /// Upserts a node by id, then sets its type/props/embedding/timestamps. + public string UpsertNode() => + $"MERGE (n:{NodeLabel} {{{PropId}: $id}}) " + + $"SET n.{PropLabel} = $label, n.{PropProps} = $props, n.{PropEmbedding} = $embedding, " + + $"n.{PropCreatedAt} = $createdAt, n.{PropUpdatedAt} = $updatedAt"; + + /// + /// Matches both endpoints, upserts the relationship by id, and returns it so the caller can detect a + /// missing endpoint (no row => a MATCH failed). + /// + public string UpsertEdge() => + $"MATCH (s:{NodeLabel} {{{PropId}: $sourceId}}) " + + $"MATCH (t:{NodeLabel} {{{PropId}: $targetId}}) " + + $"MERGE (s)-[r:{RelationshipType} {{{PropId}: $id}}]->(t) " + + $"SET r.{PropRelationType} = $relationType, r.{PropWeight} = $weight, r.{PropProps} = $props, " + + $"r.{PropCreatedAt} = $createdAt, r.{PropValidFrom} = $validFrom, r.{PropValidUntil} = $validUntil " + + "RETURN r"; + + public string GetNode() => + $"MATCH (n:{NodeLabel} {{{PropId}: $id}}) RETURN n LIMIT 1"; + + public string GetEdge() => + $"MATCH (s:{NodeLabel})-[r:{RelationshipType} {{{PropId}: $id}}]->(t:{NodeLabel}) " + + "RETURN s.id AS sourceId, t.id AS targetId, r LIMIT 1"; + + public string RemoveNode() => + $"MATCH (n:{NodeLabel} {{{PropId}: $id}}) DETACH DELETE n RETURN count(*) AS deleted"; + + public string RemoveEdge() => + $"MATCH (:{NodeLabel})-[r:{RelationshipType} {{{PropId}: $id}}]->(:{NodeLabel}) " + + "DELETE r RETURN count(*) AS deleted"; + + public string OutgoingEdges() => + $"MATCH (s:{NodeLabel} {{{PropId}: $id}})-[r:{RelationshipType}]->(t:{NodeLabel}) " + + "RETURN s.id AS sourceId, t.id AS targetId, r"; + + public string IncomingEdges() => + $"MATCH (s:{NodeLabel})-[r:{RelationshipType}]->(t:{NodeLabel} {{{PropId}: $id}}) " + + "RETURN s.id AS sourceId, t.id AS targetId, r"; + + public string NodesByLabel() => + $"MATCH (n:{NodeLabel} {{{PropLabel}: $label}}) RETURN n"; + + public string AllNodes() => + $"MATCH (n:{NodeLabel}) RETURN n"; + + public string AllEdges() => + $"MATCH (s:{NodeLabel})-[r:{RelationshipType}]->(t:{NodeLabel}) " + + "RETURN s.id AS sourceId, t.id AS targetId, r"; + + public string NodeCount() => + $"MATCH (n:{NodeLabel}) RETURN count(n) AS c"; + + public string EdgeCount() => + $"MATCH (:{NodeLabel})-[r:{RelationshipType}]->(:{NodeLabel}) RETURN count(r) AS c"; + + public string Clear() => + $"MATCH (n:{NodeLabel}) DETACH DELETE n"; + + // ---- Lossless model (de)serialization (T-independent parts). ---- + + /// Serializes a property bag to a JSON string (null/empty becomes an empty object). + public static string SerializeProperties(IDictionary? properties) + { + if (properties == null || properties.Count == 0) + { + return "{}"; + } + + return JsonConvert.SerializeObject(properties); + } + + /// + /// Deserializes a property bag from a JSON string. Mirrors the JSON round-trip semantics of the other + /// stores (e.g. integers come back as ), which + /// GraphNode.GetProperty/GraphEdge.GetProperty already normalizes via + /// . + /// + public static Dictionary DeserializeProperties(string? json) + { + if (string.IsNullOrWhiteSpace(json)) + { + return new Dictionary(); + } + + return JsonConvert.DeserializeObject>(json!) ?? new Dictionary(); + } + + /// Converts a UTC-normalized timestamp to Unix epoch milliseconds (stored as a Neo4j long). + public static long ToEpochMillis(DateTime value) + { + var utc = value.Kind == DateTimeKind.Unspecified + ? DateTime.SpecifyKind(value, DateTimeKind.Utc) + : value.ToUniversalTime(); + return new DateTimeOffset(utc).ToUnixTimeMilliseconds(); + } + + /// Converts an optional timestamp to nullable epoch milliseconds. + public static long? ToEpochMillis(DateTime? value) => + value.HasValue ? ToEpochMillis(value.Value) : (long?)null; + + /// Reconstructs a UTC from epoch milliseconds. + public static DateTime FromEpochMillis(long millis) => + DateTimeOffset.FromUnixTimeMilliseconds(millis).UtcDateTime; + + /// Reconstructs an optional UTC from nullable epoch milliseconds. + public static DateTime? FromEpochMillis(long? millis) => + millis.HasValue ? FromEpochMillis(millis.Value) : (DateTime?)null; +} diff --git a/src/AiDotNet.Storage.Neo4j/Neo4jGraphStore.cs b/src/AiDotNet.Storage.Neo4j/Neo4jGraphStore.cs new file mode 100644 index 0000000000..b2c05c125c --- /dev/null +++ b/src/AiDotNet.Storage.Neo4j/Neo4jGraphStore.cs @@ -0,0 +1,470 @@ +using System; +using System.Collections.Generic; +using System.Globalization; +using System.Linq; +using System.Threading.Tasks; +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.Interfaces; +using Neo4j.Driver; + +namespace AiDotNet.RetrievalAugmentedGeneration.Graph; + +/// +/// Neo4j-backed implementation of that runs a GraphRAG knowledge graph on a +/// production property-graph database. Nodes are stored as labeled Neo4j nodes and edges as relationships; +/// writes upsert via MERGE and traversal uses Cypher MATCH. Ships in the opt-in +/// AiDotNet.Storage.Neo4j package so the core stays free of the Neo4j.Driver dependency. +/// +/// The numeric type used for vector (embedding) operations. +/// +/// Model mapping. Every node carries a fixed base label (Entity by default; different +/// stores can use different base labels for per-namespace partitioning within one database). The node's +/// entity type (GraphNode.Label, e.g. PERSON) is stored as the label property, and its +/// Properties bag is stored losslessly as a JSON string in the props property. The embedding +/// is stored as a native Neo4j list of doubles, and timestamps as epoch-millisecond longs. Edges are stored +/// as relationships of a fixed type (RELATED by default) whose relationType property holds the +/// original GraphEdge.RelationType. This keeps labels/relationship types as fixed, escaped literals +/// (they cannot be Cypher parameters) so every per-request value is fully parameterized and no user data is +/// ever concatenated into a query — see . +/// Threading / async. The async methods are the primary implementation (Neo4j I/O is +/// inherently async); the synchronous members delegate to them with a blocking +/// wait, matching the sync-over-async pattern used by the other opt-in storage backends. The owned driver is +/// thread-safe and disposed by . +/// For Beginners: This saves your knowledge graph into Neo4j, a real graph database, so it +/// survives restarts, scales past RAM, and can be shared and queried by many processes. +/// +[ComponentType(ComponentType.DocumentStore)] +[PipelineStage(PipelineStage.Indexing)] +public sealed class Neo4jGraphStore : IGraphStore, IDisposable +{ + private static readonly INumericOperations NumOps = MathHelper.GetNumericOperations(); + + private readonly IDriver _driver; + private readonly bool _ownsDriver; + private readonly string? _database; + private readonly Neo4jCypher _cypher; + + /// + /// Creates a store that connects to Neo4j using a bolt/neo4j URI and credentials. The driver is owned + /// and disposed by this instance. Connection is lazy, so construction does not require a live server. + /// + /// The connection URI, e.g. neo4j://localhost:7687 or bolt://host:7687. + /// The user name. Pass null/empty to connect without authentication. + /// The password. + /// Optional Neo4j database name (Neo4j 4+). Null uses the server default. + /// Base label applied to all nodes (default Entity); scopes a namespace. + /// Relationship type applied to all edges (default RELATED). + public Neo4jGraphStore( + string uri, + string? username, + string? password, + string? database = null, + string nodeLabel = "Entity", + string relationshipType = "RELATED") + { + if (string.IsNullOrWhiteSpace(uri)) + { + throw new ArgumentException("Connection URI cannot be null, empty, or whitespace.", nameof(uri)); + } + + var authToken = string.IsNullOrEmpty(username) + ? AuthTokens.None + : AuthTokens.Basic(username, password); + + _driver = GraphDatabase.Driver(uri, authToken); + _ownsDriver = true; + _database = string.IsNullOrWhiteSpace(database) ? null : database; + _cypher = new Neo4jCypher(nodeLabel, relationshipType); + } + + /// + /// Creates a store over an existing driver (not owned; the caller disposes it). Useful for sharing a + /// single tuned driver across stores. + /// + /// A configured Neo4j driver. + /// Optional Neo4j database name. Null uses the server default. + /// Base label applied to all nodes (default Entity). + /// Relationship type applied to all edges (default RELATED). + public Neo4jGraphStore( + IDriver driver, + string? database = null, + string nodeLabel = "Entity", + string relationshipType = "RELATED") + { + _driver = driver ?? throw new ArgumentNullException(nameof(driver)); + _ownsDriver = false; + _database = string.IsNullOrWhiteSpace(database) ? null : database; + _cypher = new Neo4jCypher(nodeLabel, relationshipType); + } + + /// + public int NodeCount => CountAsync(_cypher.NodeCount()).GetAwaiter().GetResult(); + + /// + public int EdgeCount => CountAsync(_cypher.EdgeCount()).GetAwaiter().GetResult(); + + // ---- Async implementation (primary). ---- + + /// + public async Task AddNodeAsync(GraphNode node) + { + if (node == null) + { + throw new ArgumentNullException(nameof(node)); + } + + var parameters = new + { + id = node.Id, + label = node.Label, + props = Neo4jCypher.SerializeProperties(node.Properties), + embedding = WriteEmbedding(node.Embedding), + createdAt = Neo4jCypher.ToEpochMillis(node.CreatedAt), + updatedAt = Neo4jCypher.ToEpochMillis(node.UpdatedAt), + }; + + await RunWriteAsync(_cypher.UpsertNode(), parameters).ConfigureAwait(false); + } + + /// + public async Task AddEdgeAsync(GraphEdge edge) + { + if (edge == null) + { + throw new ArgumentNullException(nameof(edge)); + } + + var parameters = new + { + id = edge.Id, + sourceId = edge.SourceId, + targetId = edge.TargetId, + relationType = edge.RelationType, + weight = edge.Weight, + props = Neo4jCypher.SerializeProperties(edge.Properties), + createdAt = Neo4jCypher.ToEpochMillis(edge.CreatedAt), + validFrom = Neo4jCypher.ToEpochMillis(edge.ValidFrom), + validUntil = Neo4jCypher.ToEpochMillis(edge.ValidUntil), + }; + + var records = await RunWriteAsync(_cypher.UpsertEdge(), parameters).ConfigureAwait(false); + if (records.Count == 0) + { + // A MATCH for one of the endpoints produced no row, so MERGE never ran: the interface + // contract (see MemoryGraphStore) requires both endpoints to already exist. + throw new InvalidOperationException( + $"Cannot add edge '{edge.Id}': source node '{edge.SourceId}' and/or target node '{edge.TargetId}' does not exist."); + } + } + + /// + public async Task?> GetNodeAsync(string nodeId) + { + if (string.IsNullOrWhiteSpace(nodeId)) + { + return null; + } + + var records = await RunReadAsync(_cypher.GetNode(), new { id = nodeId }).ConfigureAwait(false); + if (records.Count == 0) + { + return null; + } + + return MapNode(records[0]["n"].As()); + } + + /// + public async Task?> GetEdgeAsync(string edgeId) + { + if (string.IsNullOrWhiteSpace(edgeId)) + { + return null; + } + + var records = await RunReadAsync(_cypher.GetEdge(), new { id = edgeId }).ConfigureAwait(false); + return records.Count == 0 ? null : MapEdge(records[0]); + } + + /// + public async Task RemoveNodeAsync(string nodeId) + { + if (string.IsNullOrWhiteSpace(nodeId)) + { + return false; + } + + var records = await RunWriteAsync(_cypher.RemoveNode(), new { id = nodeId }).ConfigureAwait(false); + return records.Count > 0 && records[0]["deleted"].As() > 0; + } + + /// + public async Task RemoveEdgeAsync(string edgeId) + { + if (string.IsNullOrWhiteSpace(edgeId)) + { + return false; + } + + var records = await RunWriteAsync(_cypher.RemoveEdge(), new { id = edgeId }).ConfigureAwait(false); + return records.Count > 0 && records[0]["deleted"].As() > 0; + } + + /// + public async Task>> GetOutgoingEdgesAsync(string nodeId) + { + var records = await RunReadAsync(_cypher.OutgoingEdges(), new { id = nodeId }).ConfigureAwait(false); + return records.Select(MapEdge).ToList(); + } + + /// + public async Task>> GetIncomingEdgesAsync(string nodeId) + { + var records = await RunReadAsync(_cypher.IncomingEdges(), new { id = nodeId }).ConfigureAwait(false); + return records.Select(MapEdge).ToList(); + } + + /// + public async Task>> GetNodesByLabelAsync(string label) + { + var records = await RunReadAsync(_cypher.NodesByLabel(), new { label }).ConfigureAwait(false); + return records.Select(r => MapNode(r["n"].As())).ToList(); + } + + /// + public async Task>> GetAllNodesAsync() + { + var records = await RunReadAsync(_cypher.AllNodes(), new { }).ConfigureAwait(false); + return records.Select(r => MapNode(r["n"].As())).ToList(); + } + + /// + public async Task>> GetAllEdgesAsync() + { + var records = await RunReadAsync(_cypher.AllEdges(), new { }).ConfigureAwait(false); + return records.Select(MapEdge).ToList(); + } + + /// + public async Task ClearAsync() + { + await RunWriteAsync(_cypher.Clear(), new { }).ConfigureAwait(false); + } + + // ---- Synchronous members: sync-over-async, matching the other opt-in storage backends. ---- + + /// + public void AddNode(GraphNode node) => AddNodeAsync(node).GetAwaiter().GetResult(); + + /// + public void AddEdge(GraphEdge edge) => AddEdgeAsync(edge).GetAwaiter().GetResult(); + + /// + public GraphNode? GetNode(string nodeId) => GetNodeAsync(nodeId).GetAwaiter().GetResult(); + + /// + public GraphEdge? GetEdge(string edgeId) => GetEdgeAsync(edgeId).GetAwaiter().GetResult(); + + /// + public bool RemoveNode(string nodeId) => RemoveNodeAsync(nodeId).GetAwaiter().GetResult(); + + /// + public bool RemoveEdge(string edgeId) => RemoveEdgeAsync(edgeId).GetAwaiter().GetResult(); + + /// + public IEnumerable> GetOutgoingEdges(string nodeId) => GetOutgoingEdgesAsync(nodeId).GetAwaiter().GetResult(); + + /// + public IEnumerable> GetIncomingEdges(string nodeId) => GetIncomingEdgesAsync(nodeId).GetAwaiter().GetResult(); + + /// + public IEnumerable> GetNodesByLabel(string label) => GetNodesByLabelAsync(label).GetAwaiter().GetResult(); + + /// + public IEnumerable> GetAllNodes() => GetAllNodesAsync().GetAwaiter().GetResult(); + + /// + public IEnumerable> GetAllEdges() => GetAllEdgesAsync().GetAwaiter().GetResult(); + + /// + public void Clear() => ClearAsync().GetAwaiter().GetResult(); + + // ---- Session helpers. ---- + + private void ConfigureSession(SessionConfigBuilder builder) + { + if (_database != null) + { + builder.WithDatabase(_database); + } + } + + private async Task> RunWriteAsync(string cypher, object parameters) + { + await using var session = _driver.AsyncSession(ConfigureSession); + return await session.ExecuteWriteAsync(async tx => + { + var cursor = await tx.RunAsync(cypher, parameters).ConfigureAwait(false); + return await cursor.ToListAsync().ConfigureAwait(false); + }).ConfigureAwait(false); + } + + private async Task> RunReadAsync(string cypher, object parameters) + { + await using var session = _driver.AsyncSession(ConfigureSession); + return await session.ExecuteReadAsync(async tx => + { + var cursor = await tx.RunAsync(cypher, parameters).ConfigureAwait(false); + return await cursor.ToListAsync().ConfigureAwait(false); + }).ConfigureAwait(false); + } + + private async Task CountAsync(string cypher) + { + var records = await RunReadAsync(cypher, new { }).ConfigureAwait(false); + return records.Count == 0 ? 0 : (int)records[0]["c"].As(); + } + + // ---- Record <-> model mapping. ---- + + private static GraphNode MapNode(INode node) + { + var props = node.Properties; + var id = ReadString(props, Neo4jCypher.PropId) ?? throw new InvalidOperationException("Neo4j node is missing its 'id' property."); + var label = ReadString(props, Neo4jCypher.PropLabel); + if (string.IsNullOrWhiteSpace(label)) + { + // Labels stored by this backend are never empty; fall back defensively so a malformed row + // cannot trip the GraphNode constructor's non-empty-label guard. + label = "UNKNOWN"; + } + + var graphNode = new GraphNode(id, label!) + { + Properties = Neo4jCypher.DeserializeProperties(ReadString(props, Neo4jCypher.PropProps)), + Embedding = ReadEmbedding(props), + }; + + var createdAt = ReadNullableEpoch(props, Neo4jCypher.PropCreatedAt); + if (createdAt.HasValue) + { + graphNode.CreatedAt = createdAt.Value; + } + + var updatedAt = ReadNullableEpoch(props, Neo4jCypher.PropUpdatedAt); + if (updatedAt.HasValue) + { + graphNode.UpdatedAt = updatedAt.Value; + } + + return graphNode; + } + + private static GraphEdge MapEdge(IRecord record) + { + var sourceId = record["sourceId"].As(); + var targetId = record["targetId"].As(); + var rel = record["r"].As(); + var props = rel.Properties; + + var relationType = ReadString(props, Neo4jCypher.PropRelationType); + if (string.IsNullOrWhiteSpace(relationType)) + { + relationType = rel.Type; + } + + // Construct with the default (valid) weight, then assign the stored weight directly to stay + // lossless without tripping the GraphEdge constructor's [0,1] weight guard. + var edge = new GraphEdge(sourceId, targetId, relationType!); + + var idValue = ReadString(props, Neo4jCypher.PropId); + if (!string.IsNullOrEmpty(idValue)) + { + edge.Id = idValue!; + } + + if (props.TryGetValue(Neo4jCypher.PropWeight, out var weightRaw) && weightRaw != null) + { + edge.Weight = Convert.ToDouble(weightRaw, CultureInfo.InvariantCulture); + } + + edge.Properties = Neo4jCypher.DeserializeProperties(ReadString(props, Neo4jCypher.PropProps)); + + var createdAt = ReadNullableEpoch(props, Neo4jCypher.PropCreatedAt); + if (createdAt.HasValue) + { + edge.CreatedAt = createdAt.Value; + } + + edge.ValidFrom = ReadNullableEpoch(props, Neo4jCypher.PropValidFrom); + edge.ValidUntil = ReadNullableEpoch(props, Neo4jCypher.PropValidUntil); + return edge; + } + + private static double[]? WriteEmbedding(Vector? embedding) + { + if (embedding == null) + { + return null; + } + + var values = new double[embedding.Length]; + for (int i = 0; i < embedding.Length; i++) + { + values[i] = NumOps.ToDouble(embedding[i]); + } + + return values; + } + + private static Vector? ReadEmbedding(IReadOnlyDictionary props) + { + if (!props.TryGetValue(Neo4jCypher.PropEmbedding, out var raw) || raw is null) + { + return null; + } + + if (raw is System.Collections.IEnumerable list && raw is not string) + { + var values = new List(); + foreach (var item in list) + { + values.Add(NumOps.FromDouble(Convert.ToDouble(item, CultureInfo.InvariantCulture))); + } + + return new Vector(values.ToArray()); + } + + return null; + } + + private static string? ReadString(IReadOnlyDictionary props, string key) + { + if (props.TryGetValue(key, out var value) && value != null) + { + return Convert.ToString(value, CultureInfo.InvariantCulture); + } + + return null; + } + + private static DateTime? ReadNullableEpoch(IReadOnlyDictionary props, string key) + { + if (props.TryGetValue(key, out var value) && value != null) + { + return Neo4jCypher.FromEpochMillis(Convert.ToInt64(value, CultureInfo.InvariantCulture)); + } + + return null; + } + + /// Disposes the owned Neo4j driver. No-op when constructed over a shared driver. + public void Dispose() + { + if (_ownsDriver) + { + _driver.Dispose(); + } + } +} diff --git a/src/AiDotNet.csproj b/src/AiDotNet.csproj index 4b7c3d0074..492528f402 100644 --- a/src/AiDotNet.csproj +++ b/src/AiDotNet.csproj @@ -197,6 +197,10 @@ + + + + diff --git a/tests/AiDotNet.Tests/AiDotNetTests.csproj b/tests/AiDotNet.Tests/AiDotNetTests.csproj index fddd341c3c..823134d4ae 100644 --- a/tests/AiDotNet.Tests/AiDotNetTests.csproj +++ b/tests/AiDotNet.Tests/AiDotNetTests.csproj @@ -62,6 +62,7 @@ + diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Neo4jGraphStoreTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Neo4jGraphStoreTests.cs new file mode 100644 index 0000000000..63e8bf7dc2 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Neo4jGraphStoreTests.cs @@ -0,0 +1,321 @@ +#if NET5_0_OR_GREATER +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Graph; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration +{ + /// + /// Tests for the Neo4j property-graph backend (AiDotNet.Storage.Neo4j). Split into two groups: + /// (a) pure-logic tests over — the Cypher/param/mapping builders — which need + /// no database and always run; and (b) integration tests over that hit + /// a live Neo4j and are SKIPPED (not failed) unless NEO4J_URI/NEO4J_USER/NEO4J_PASSWORD are set, mirroring + /// the Postgres/Redis checkpointer integration tests. The project is net10-only (the driver's TFM). + /// + public class Neo4jGraphStoreTests + { + // ---------- (a) Pure-logic Cypher/mapping builder tests (no database). ---------- + + [Fact] + public void EscapeIdentifier_WrapsInBackticks() + { + Assert.Equal("`Entity`", Neo4jCypher.EscapeIdentifier("Entity")); + } + + [Fact] + public void EscapeIdentifier_DoublesEmbeddedBackticks_PreventingBreakout() + { + // A malicious "identifier" that tries to close the backtick and inject Cypher must be neutralized + // by doubling the internal backtick, keeping the whole thing a single identifier token. + var malicious = "Entity`) DETACH DELETE (x) //"; + var escaped = Neo4jCypher.EscapeIdentifier(malicious); + Assert.Equal("`Entity``) DETACH DELETE (x) //`", escaped); + } + + [Theory] + [InlineData(null)] + [InlineData("")] + [InlineData(" ")] + public void EscapeIdentifier_Throws_OnNullOrWhitespace(string? identifier) + { + Assert.Throws(() => Neo4jCypher.EscapeIdentifier(identifier!)); + } + + [Fact] + public void EscapeIdentifier_Throws_OnNulChar() + { + Assert.Throws(() => Neo4jCypher.EscapeIdentifier("Ent\0ity")); + } + + [Fact] + public void Constructor_EscapesLabelAndRelationshipType() + { + var cypher = new Neo4jCypher("Person", "KNOWS"); + Assert.Equal("`Person`", cypher.NodeLabel); + Assert.Equal("`KNOWS`", cypher.RelationshipType); + } + + [Fact] + public void Constructor_UsesDefaults() + { + var cypher = new Neo4jCypher(); + Assert.Equal("`Entity`", cypher.NodeLabel); + Assert.Equal("`RELATED`", cypher.RelationshipType); + } + + [Fact] + public void UpsertNode_MergesByIdAndSetsAllColumns_Parameterized() + { + var cypher = new Neo4jCypher(); + var q = cypher.UpsertNode(); + + Assert.Contains("MERGE (n:`Entity` {id: $id})", q); + Assert.Contains("n.label = $label", q); + Assert.Contains("n.props = $props", q); + Assert.Contains("n.embedding = $embedding", q); + Assert.Contains("n.createdAt = $createdAt", q); + Assert.Contains("n.updatedAt = $updatedAt", q); + } + + [Fact] + public void UpsertEdge_MatchesBothEndpoints_MergesRelationship_AndReturnsIt() + { + var cypher = new Neo4jCypher(); + var q = cypher.UpsertEdge(); + + Assert.Contains("MATCH (s:`Entity` {id: $sourceId})", q); + Assert.Contains("MATCH (t:`Entity` {id: $targetId})", q); + Assert.Contains("MERGE (s)-[r:`RELATED` {id: $id}]->(t)", q); + Assert.Contains("r.relationType = $relationType", q); + Assert.Contains("r.weight = $weight", q); + Assert.Contains("r.validFrom = $validFrom", q); + Assert.Contains("r.validUntil = $validUntil", q); + Assert.EndsWith("RETURN r", q); + } + + [Fact] + public void RemoveNode_DetachDeletes_AndReturnsCount() + { + var q = new Neo4jCypher().RemoveNode(); + Assert.Contains("MATCH (n:`Entity` {id: $id})", q); + Assert.Contains("DETACH DELETE n", q); + Assert.Contains("count(*) AS deleted", q); + } + + [Fact] + public void GetOutgoingAndIncoming_TraverseInCorrectDirection() + { + var cypher = new Neo4jCypher(); + Assert.Contains("(s:`Entity` {id: $id})-[r:`RELATED`]->(t:`Entity`)", cypher.OutgoingEdges()); + Assert.Contains("(s:`Entity`)-[r:`RELATED`]->(t:`Entity` {id: $id})", cypher.IncomingEdges()); + } + + [Fact] + public void NodesByLabel_FiltersOnTheLabelProperty() + { + var q = new Neo4jCypher().NodesByLabel(); + Assert.Contains("MATCH (n:`Entity` {label: $label})", q); + } + + [Fact] + public void CountQueries_UseAggregates() + { + var cypher = new Neo4jCypher(); + Assert.Contains("count(n) AS c", cypher.NodeCount()); + Assert.Contains("count(r) AS c", cypher.EdgeCount()); + } + + [Fact] + public void CustomLabel_FlowsIntoEveryQuery() + { + var cypher = new Neo4jCypher("Movie", "ACTED_IN"); + Assert.Contains(":`Movie`", cypher.AllNodes()); + Assert.Contains(":`ACTED_IN`", cypher.AllEdges()); + Assert.Contains(":`Movie`", cypher.Clear()); + } + + [Fact] + public void SerializeProperties_NullOrEmpty_ReturnsEmptyObject() + { + Assert.Equal("{}", Neo4jCypher.SerializeProperties(null)); + Assert.Equal("{}", Neo4jCypher.SerializeProperties(new Dictionary())); + } + + [Fact] + public void Properties_RoundTrip_LosslesslyForStringsAndBooleans() + { + var original = new Dictionary + { + ["name"] = "Einstein", + ["active"] = true, + ["score"] = 3.5, + }; + + var json = Neo4jCypher.SerializeProperties(original); + var restored = Neo4jCypher.DeserializeProperties(json); + + Assert.Equal("Einstein", restored["name"]); + Assert.Equal(true, restored["active"]); + Assert.Equal(3.5, Convert.ToDouble(restored["score"])); + } + + [Fact] + public void DeserializeProperties_EmptyOrNull_ReturnsEmptyDictionary() + { + Assert.Empty(Neo4jCypher.DeserializeProperties(null)); + Assert.Empty(Neo4jCypher.DeserializeProperties("")); + Assert.Empty(Neo4jCypher.DeserializeProperties("{}")); + } + + [Fact] + public void EpochMillis_RoundTripsUtc() + { + var when = new DateTime(2026, 7, 19, 12, 30, 45, DateTimeKind.Utc); + var millis = Neo4jCypher.ToEpochMillis(when); + var restored = Neo4jCypher.FromEpochMillis(millis); + Assert.Equal(when, restored); + } + + [Fact] + public void EpochMillis_Nullable_PreservesNull() + { + Assert.Null(Neo4jCypher.ToEpochMillis((DateTime?)null)); + Assert.Null(Neo4jCypher.FromEpochMillis((long?)null)); + + long? millis = Neo4jCypher.ToEpochMillis((DateTime?)new DateTime(2000, 1, 1, 0, 0, 0, DateTimeKind.Utc)); + Assert.NotNull(millis); + Assert.Equal(new DateTime(2000, 1, 1, 0, 0, 0, DateTimeKind.Utc), Neo4jCypher.FromEpochMillis(millis)); + } + + // ---------- (b) Integration tests (require a live Neo4j; skipped when env vars are unset). ---------- + + private const string UriEnv = "NEO4J_URI"; + private const string UserEnv = "NEO4J_USER"; + private const string PasswordEnv = "NEO4J_PASSWORD"; + + private static Neo4jGraphStore? CreateStoreOrSkip(out string reason) + { + var uri = Environment.GetEnvironmentVariable(UriEnv); + reason = $"Set {UriEnv}/{UserEnv}/{PasswordEnv} to run the Neo4j integration tests."; + if (string.IsNullOrWhiteSpace(uri)) + { + return null; + } + + var user = Environment.GetEnvironmentVariable(UserEnv); + var password = Environment.GetEnvironmentVariable(PasswordEnv); + // Unique per-run label/relationship type so the test partitions its own data and Clear() only + // touches this run's nodes, never any real graph in the target database. + var suffix = Guid.NewGuid().ToString("N"); + return new Neo4jGraphStore(uri, user, password, nodeLabel: "AidnTest_" + suffix, relationshipType: "REL_" + suffix); + } + + [SkippableFact(Timeout = 120000)] + [Trait("Category", "Integration")] + public async Task RoundTrips_Nodes_Edges_Traversal_And_Clear() + { + var store = CreateStoreOrSkip(out var reason); + Skip.If(store is null, reason); + + using (store) + { + try + { + var alice = new GraphNode("alice", "PERSON"); + alice.SetProperty("name", "Alice"); + alice.SetProperty("age", 30); + alice.Embedding = new Vector(new[] { 0.1, 0.2, 0.3 }); + + var acme = new GraphNode("acme", "COMPANY"); + acme.SetProperty("name", "Acme"); + + await store!.AddNodeAsync(alice); + await store.AddNodeAsync(acme); + + var edge = new GraphEdge("alice", "acme", "WORKS_FOR", 0.9); + edge.SetProperty("since", "2020"); + await store.AddEdgeAsync(edge); + + Assert.Equal(2, store.NodeCount); + Assert.Equal(1, store.EdgeCount); + + // Get node round-trips label, properties, and embedding. + var fetched = await store.GetNodeAsync("alice"); + Assert.NotNull(fetched); + Assert.Equal("PERSON", fetched!.Label); + Assert.Equal("Alice", fetched.GetProperty("name")); + Assert.Equal(30, fetched.GetProperty("age")); + Assert.NotNull(fetched.Embedding); + Assert.Equal(3, fetched.Embedding!.Length); + Assert.Equal(0.2, fetched.Embedding[1], 6); + + // Get edge round-trips endpoints, type, weight, and properties. + var fetchedEdge = await store.GetEdgeAsync(edge.Id); + Assert.NotNull(fetchedEdge); + Assert.Equal("WORKS_FOR", fetchedEdge!.RelationType); + Assert.Equal(0.9, fetchedEdge.Weight, 6); + Assert.Equal("2020", fetchedEdge.GetProperty("since")); + + // Traversal. + var outgoing = (await store.GetOutgoingEdgesAsync("alice")).ToList(); + Assert.Single(outgoing); + Assert.Equal("acme", outgoing[0].TargetId); + + var incoming = (await store.GetIncomingEdgesAsync("acme")).ToList(); + Assert.Single(incoming); + Assert.Equal("alice", incoming[0].SourceId); + + // Label index. + var people = (await store.GetNodesByLabelAsync("PERSON")).ToList(); + Assert.Single(people); + Assert.Equal("alice", people[0].Id); + + // Upsert semantics: re-adding updates in place, not duplicates. + var aliceUpdated = new GraphNode("alice", "PERSON"); + aliceUpdated.SetProperty("name", "Alice Smith"); + await store.AddNodeAsync(aliceUpdated); + Assert.Equal(2, store.NodeCount); + Assert.Equal("Alice Smith", (await store.GetNodeAsync("alice"))!.GetProperty("name")); + + // Removals. + Assert.True(await store.RemoveEdgeAsync(edge.Id)); + Assert.Equal(0, store.EdgeCount); + Assert.True(await store.RemoveNodeAsync("acme")); + Assert.False(await store.RemoveNodeAsync("does-not-exist")); + Assert.Equal(1, store.NodeCount); + } + finally + { + await store!.ClearAsync(); + } + } + } + + [SkippableFact(Timeout = 120000)] + [Trait("Category", "Integration")] + public async Task AddEdge_Throws_WhenEndpointMissing() + { + var store = CreateStoreOrSkip(out var reason); + Skip.If(store is null, reason); + + using (store) + { + try + { + await store!.AddNodeAsync(new GraphNode("only", "PERSON")); + var dangling = new GraphEdge("only", "ghost", "KNOWS"); + await Assert.ThrowsAsync(() => store.AddEdgeAsync(dangling)); + } + finally + { + await store!.ClearAsync(); + } + } + } + } +} +#endif From ba72d06ecf916782248f229436436247c07abbdd Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 21:44:46 -0400 Subject: [PATCH 22/25] feat(rag): gPU-accelerated vector search with CPU fallback Add GPU-accelerated brute-force vector search to the RAG vector-search path, matching FAISS-GPU / Milvus-GPU flat-search: all stored vectors are packed into a matrix and scored against the query in a single batched GPU Gemm, then top-k'd. - GpuVectorScorer: reuses the existing DirectGpuTensorEngine / IDirectGpuBackend abstraction (AiDotNetEngine.Current, GetBackend, AllocateBuffer, Gemm, DownloadBuffer) - no raw CUDA. Supports dot-product, cosine and Euclidean via closed forms over the GPU dot products + precomputed norms. - GpuFlatIndex : IVectorIndex: batched-GPU flat index with exact CPU fallback identical to FlatIndex ordering. - HNSWIndex.ScoreCandidatesGpu: batched GPU candidate re-scoring for a coarse-to-fine pipeline, with CPU fallback. - Falls back to CPU when no GPU/engine, below threshold, unsupported metric, or on any GPU error; gated off under AIDOTNET_DISABLE_GPU like the rest of the codebase. Compiles on net10.0/net8.0/net471 (GPU types ship on all TFMs). - Tests assert GPU-index top-k == CPU FlatIndex top-k across metrics and thresholds; deterministic and CI-runnable (GPU disabled -> CPU path). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../VectorSearch/GpuVectorScorer.cs | 279 ++++++++++++++++ .../VectorSearch/Indexes/GpuFlatIndex.cs | 153 +++++++++ .../VectorSearch/Indexes/HNSWIndex.cs | 73 +++++ .../VectorSearch/Indexes/GpuFlatIndexTests.cs | 307 ++++++++++++++++++ 4 files changed, 812 insertions(+) create mode 100644 src/RetrievalAugmentedGeneration/VectorSearch/GpuVectorScorer.cs create mode 100644 src/RetrievalAugmentedGeneration/VectorSearch/Indexes/GpuFlatIndex.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/GpuFlatIndexTests.cs diff --git a/src/RetrievalAugmentedGeneration/VectorSearch/GpuVectorScorer.cs b/src/RetrievalAugmentedGeneration/VectorSearch/GpuVectorScorer.cs new file mode 100644 index 0000000000..659a36f467 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/VectorSearch/GpuVectorScorer.cs @@ -0,0 +1,279 @@ +using System; +using System.Collections.Generic; +using AiDotNet.Helpers; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Metrics; +using AiDotNet.Tensors.Engines; +using AiDotNet.Tensors.Engines.DirectGpu; +using AiDotNet.Tensors.Engines.Gpu; + +namespace AiDotNet.RetrievalAugmentedGeneration.VectorSearch +{ + /// + /// GPU-accelerated batched similarity scoring for vector search. + /// + /// + /// + /// This helper computes the similarity/distance between a single query vector and a large + /// batch of stored vectors in a single GPU matrix multiply, mirroring how FAISS-GPU / + /// Milvus-GPU accelerate brute-force ("flat") search. All stored vectors are packed into an + /// [N, dim] row-major matrix, the query is uploaded as a [dim, 1] column, and a + /// single Gemm produces the [N, 1] dot-product vector. Metric-specific closed + /// forms then derive the final score from the dot products and precomputed L2 norms. + /// + /// + /// The GPU abstraction reused is the same one the neural-network layers use: + /// cast to + /// , its IsGpuAvailable probe, and the + /// returned by GetBackend() (AllocateBuffer, + /// Gemm, DownloadBuffer). No raw CUDA/OpenCL is written here. + /// + /// Fallback contract. returns false — signalling + /// the caller to use its existing CPU path — whenever any of the following hold: + /// the AIDOTNET_DISABLE_GPU environment variable is set; no GPU engine/backend is + /// available; the batch is smaller than the GPU threshold (small batches are faster on the + /// CPU than the upload/compute/download round-trip); the metric has no GPU closed form + /// (only dot-product, cosine and Euclidean are accelerated); vector dimensions are + /// inconsistent; or any GPU error occurs. In every fallback case the caller's CPU brute-force + /// path runs, so results are always correct and CI (which sets AIDOTNET_DISABLE_GPU) + /// exercises the CPU path deterministically. + /// + /// The types used here (, ) + /// ship in the AiDotNet.Tensors package for every target framework (net10.0, net8.0, net471), + /// so this file compiles on all TFMs; the GPU path simply stays dormant at runtime when no + /// device is present. + /// + /// + public static class GpuVectorScorer + { + /// + /// Default minimum number of stored vectors before the GPU path is attempted. Below this + /// the CPU brute-force path is faster (no host/device transfer overhead). + /// + public const int DefaultGpuThreshold = 1024; + + private enum MetricKind + { + Unsupported, + DotProduct, + Cosine, + Euclidean + } + + /// + /// Returns true when the AIDOTNET_DISABLE_GPU environment variable is set to a + /// non-empty, non-"0"/"false" value. Mirrors the gate honored by the rest of the codebase. + /// + public static bool IsGpuDisabledByEnvironment + { + get + { + var value = Environment.GetEnvironmentVariable("AIDOTNET_DISABLE_GPU"); + if (string.IsNullOrEmpty(value)) + return false; + return !string.Equals(value, "0", StringComparison.OrdinalIgnoreCase) + && !string.Equals(value, "false", StringComparison.OrdinalIgnoreCase); + } + } + + /// + /// Returns true when a GPU tensor engine is present and enabled (and not disabled via the + /// environment variable). This is the same probe used by the neural-network / diffusion code. + /// + public static bool IsGpuAvailable() + { + if (IsGpuDisabledByEnvironment) + return false; + + try + { + var engine = AiDotNetEngine.Current as DirectGpuTensorEngine; + return engine != null && engine.IsGpuAvailable; + } + catch + { + // Any failure probing the engine is treated as "no GPU" so callers fall back to CPU. + return false; + } + } + + private static MetricKind GetKind(ISimilarityMetric metric) + { + switch (metric) + { + case DotProductMetric: + return MetricKind.DotProduct; + case CosineSimilarityMetric: + return MetricKind.Cosine; + case EuclideanDistanceMetric: + return MetricKind.Euclidean; + default: + return MetricKind.Unsupported; + } + } + + /// + /// Returns true when a batch of the given size with the given metric would be routed to the + /// GPU (metric supported, batch at/above threshold, and a GPU device available). Useful for + /// callers that want to decide before materializing a candidate list. + /// + public static bool CanAccelerate(ISimilarityMetric metric, int count, int gpuThreshold) + { + if (metric == null) + return false; + if (count < gpuThreshold) + return false; + if (GetKind(metric) == MetricKind.Unsupported) + return false; + return IsGpuAvailable(); + } + + /// + /// Scores against every vector in on the GPU + /// in a single batched matrix multiply, producing one score per document in the same order as + /// and using the same semantics as the CPU metric. + /// + /// The numeric type. + /// The similarity metric (only dot-product, cosine, Euclidean are accelerated). + /// The query vector. + /// The stored vectors to score, in caller order. + /// Minimum batch size before the GPU path is attempted. + /// On success, one score per doc aligned to ; otherwise empty. + /// + /// True when the scores were produced on the GPU; false when the caller must fall back to its + /// CPU path (GPU disabled/absent, metric unsupported, batch too small, or a GPU error). + /// + public static bool TryScoreBatch( + ISimilarityMetric metric, + Vector query, + IReadOnlyList> docs, + int gpuThreshold, + out T[] scores) + { + scores = Array.Empty(); + + if (metric == null || query == null || docs == null) + return false; + + var kind = GetKind(metric); + if (kind == MetricKind.Unsupported) + return false; + + int n = docs.Count; + if (n < gpuThreshold || n == 0) + return false; + + int dim = query.Length; + if (dim == 0) + return false; + + // Guard against int overflow of the flat matrix size; fall back to CPU for extreme sizes. + long totalLong = (long)n * dim; + if (totalLong > int.MaxValue) + return false; + + if (!IsGpuAvailable()) + return false; + + var numOps = MathHelper.GetNumericOperations(); + + // Pack docs into an [n, dim] row-major float matrix and accumulate squared norms. + int total = (int)totalLong; + var docData = new float[total]; + var docNormSq = new double[n]; + int idx = 0; + for (int i = 0; i < n; i++) + { + var v = docs[i]; + if (v == null || v.Length != dim) + return false; // inconsistent dimensions -> CPU fallback + + double ss = 0.0; + for (int j = 0; j < dim; j++) + { + double d = numOps.ToDouble(v[j]); + docData[idx++] = (float)d; + ss += d * d; + } + docNormSq[i] = ss; + } + + // Pack the query as a [dim, 1] column and accumulate its squared norm. + var queryData = new float[dim]; + double queryNormSq = 0.0; + for (int j = 0; j < dim; j++) + { + double d = numOps.ToDouble(query[j]); + queryData[j] = (float)d; + queryNormSq += d * d; + } + + float[] dots; + try + { + var engine = AiDotNetEngine.Current as DirectGpuTensorEngine; + var backend = engine?.GetBackend(); + if (backend == null) + return false; + + // C[n,1] = A[n,dim] @ B[dim,1] -> dots[i] = dot(doc_i, query) + using var docsBuffer = backend.AllocateBuffer(docData); + using var queryBuffer = backend.AllocateBuffer(queryData); + using var outBuffer = backend.AllocateBuffer(n); + backend.Gemm(docsBuffer, queryBuffer, outBuffer, n, 1, dim); + dots = backend.DownloadBuffer(outBuffer); + } + catch + { + // Any GPU failure -> signal CPU fallback. + scores = Array.Empty(); + return false; + } + + if (dots == null || dots.Length < n) + { + scores = Array.Empty(); + return false; + } + + var result = new T[n]; + double queryNorm = Math.Sqrt(queryNormSq); + for (int i = 0; i < n; i++) + { + double dot = dots[i]; + double score; + switch (kind) + { + case MetricKind.DotProduct: + score = dot; + break; + + case MetricKind.Cosine: + { + double denom = queryNorm * Math.Sqrt(docNormSq[i]); + score = denom > 0.0 ? dot / denom : 0.0; + } + break; + + case MetricKind.Euclidean: + { + // ||q - d||^2 = ||q||^2 + ||d||^2 - 2 q.d ; clamp tiny negatives from fp error. + double sq = queryNormSq + docNormSq[i] - 2.0 * dot; + if (sq < 0.0) + sq = 0.0; + score = Math.Sqrt(sq); + } + break; + + default: + return false; + } + + result[i] = numOps.FromDouble(score); + } + + scores = result; + return true; + } + } +} diff --git a/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/GpuFlatIndex.cs b/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/GpuFlatIndex.cs new file mode 100644 index 0000000000..63e95600eb --- /dev/null +++ b/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/GpuFlatIndex.cs @@ -0,0 +1,153 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using AiDotNet.Validation; + +namespace AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Indexes +{ + /// + /// Flat (brute-force) index that offloads the query-vs-all similarity computation to the GPU + /// when a device is available, falling back to an exact CPU scan otherwise. + /// + /// + /// + /// Like this index compares the query against every stored vector and + /// returns exact nearest neighbors, but the scoring step is executed as a single batched GPU + /// matrix multiply (see ), matching how FAISS-GPU / Milvus-GPU + /// accelerate flat search. This makes large brute-force scans (tens of thousands to millions of + /// vectors) dramatically faster while preserving exact results. + /// + /// + /// Fallback. The GPU path is used only when: the GPU is available and not disabled via + /// AIDOTNET_DISABLE_GPU; the number of stored vectors is at or above + /// ; and the metric has a GPU closed form (dot-product, cosine, + /// Euclidean). In every other case — and if any GPU error occurs — the index performs the exact + /// same CPU brute-force scan as , producing identical top-k ordering. + /// + /// + /// Because the GPU path falls back to CPU when no device is present, this type is safe to use in + /// CI and on all target frameworks; it never throws when the GPU is absent. + /// + /// + /// The numeric type for vector operations. + [ComponentType(ComponentType.VectorIndex)] + [PipelineStage(PipelineStage.Retrieval)] + public class GpuFlatIndex : IVectorIndex + { + private readonly Dictionary> _vectors; + private readonly ISimilarityMetric _metric; + + /// + /// The minimum number of stored vectors before the GPU path is attempted. Below this the + /// CPU scan is used (the host/device transfer overhead is not worth it for small sets). + /// + public int GpuThreshold { get; } + + /// + public int Count => _vectors.Count; + + /// + /// Initializes a new instance of the class. + /// + /// The similarity metric to use for search. + /// + /// Minimum number of vectors before the GPU path is used. Defaults to + /// . + /// + public GpuFlatIndex(ISimilarityMetric metric, int gpuThreshold = GpuVectorScorer.DefaultGpuThreshold) + { + Guard.NotNull(metric); + if (gpuThreshold < 1) + throw new ArgumentException("GPU threshold must be positive", nameof(gpuThreshold)); + + _metric = metric; + GpuThreshold = gpuThreshold; + _vectors = new Dictionary>(); + } + + /// + public void Add(string id, Vector vector) + { + if (string.IsNullOrEmpty(id)) + throw new ArgumentException("ID cannot be null or empty", nameof(id)); + if (vector == null) + throw new ArgumentNullException(nameof(vector)); + + _vectors[id] = vector; + } + + /// + public void AddBatch(Dictionary> vectors) + { + if (vectors == null) + throw new ArgumentNullException(nameof(vectors)); + + foreach (var kvp in vectors) + { + Add(kvp.Key, kvp.Value); + } + } + + /// + public List<(string Id, T Score)> Search(Vector query, int k) + { + if (query == null) + throw new ArgumentNullException(nameof(query)); + if (k <= 0) + throw new ArgumentException("k must be positive", nameof(k)); + + // Snapshot ids and vectors in a single consistent enumeration order so the GPU and CPU + // branches produce identically-ordered (id, score) pairs. + int n = _vectors.Count; + var ids = new List(n); + var vecs = new List>(n); + foreach (var kvp in _vectors) + { + ids.Add(kvp.Key); + vecs.Add(kvp.Value); + } + + var scores = new List<(string Id, T Score)>(n); + + // Attempt the GPU batched scoring; on any signal to fall back, use the exact CPU scan. + if (GpuVectorScorer.TryScoreBatch(_metric, query, vecs, GpuThreshold, out var gpuScores) + && gpuScores.Length == n) + { + for (int i = 0; i < n; i++) + { + scores.Add((ids[i], gpuScores[i])); + } + } + else + { + for (int i = 0; i < n; i++) + { + var score = _metric.Calculate(query, vecs[i]); + scores.Add((ids[i], score)); + } + } + + // Sort based on whether higher or lower is better (identical to FlatIndex). + var sorted = _metric.HigherIsBetter + ? scores.OrderByDescending(x => x.Score) + : scores.OrderBy(x => x.Score); + + return sorted.Take(Math.Min(k, scores.Count)).ToList(); + } + + /// + public bool Remove(string id) + { + return _vectors.Remove(id); + } + + /// + public void Clear() + { + _vectors.Clear(); + } + } +} diff --git a/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/HNSWIndex.cs b/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/HNSWIndex.cs index 43bf4c92b7..d5f2b8a0a9 100644 --- a/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/HNSWIndex.cs +++ b/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/HNSWIndex.cs @@ -269,6 +269,79 @@ public void AddBatch(Dictionary> vectors) } } + /// + /// Scores an explicit set of candidate node ids against the query, using a single batched + /// GPU matrix multiply when a device is available and the candidate set is large enough, + /// and falling back to the exact CPU metric otherwise. + /// + /// + /// + /// This supports a coarse-to-fine retrieval pipeline: HNSW graph traversal proposes a + /// (possibly large) candidate set cheaply, then those candidates are re-scored exactly in one + /// GPU batch — the same pattern GPU vector databases use to combine ANN proposal with exact + /// re-ranking. Results are returned ordered best-first according to the metric. + /// + /// + /// The GPU path is used only when the GPU is available (and not disabled via + /// AIDOTNET_DISABLE_GPU), the candidate count is at or above + /// , and the metric has a GPU closed form; otherwise the exact + /// CPU metric is used. Unknown or missing candidate ids are skipped. + /// + /// + /// The query vector. + /// The candidate node ids to score. + /// + /// Minimum candidate count before the GPU path is attempted. Defaults to + /// . + /// + /// A list of (id, score) tuples ordered best-first according to the metric. + public List<(string Id, T Score)> ScoreCandidatesGpu( + Vector query, + IReadOnlyList candidateIds, + int gpuThreshold = GpuVectorScorer.DefaultGpuThreshold) + { + if (query == null) + throw new ArgumentNullException(nameof(query)); + if (candidateIds == null) + throw new ArgumentNullException(nameof(candidateIds)); + + lock (_graphLock) + { + // Resolve candidate ids to vectors (skip ids no longer present), preserving order. + var ids = new List(candidateIds.Count); + var vecs = new List>(candidateIds.Count); + foreach (var id in candidateIds) + { + if (id != null && _vectors.TryGetValue(id, out var vec)) + { + ids.Add(id); + vecs.Add(vec); + } + } + + int n = ids.Count; + var scored = new List<(string Id, T Score)>(n); + + if (GpuVectorScorer.TryScoreBatch(_metric, query, vecs, gpuThreshold, out var gpuScores) + && gpuScores.Length == n) + { + for (int i = 0; i < n; i++) + scored.Add((ids[i], gpuScores[i])); + } + else + { + for (int i = 0; i < n; i++) + scored.Add((ids[i], _metric.Calculate(query, vecs[i]))); + } + + var sorted = _metric.HigherIsBetter + ? scored.OrderByDescending(x => x.Score) + : scored.OrderBy(x => x.Score); + + return sorted.ToList(); + } + } + /// public bool Remove(string id) { diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/GpuFlatIndexTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/GpuFlatIndexTests.cs new file mode 100644 index 0000000000..ba11df6381 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/GpuFlatIndexTests.cs @@ -0,0 +1,307 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Indexes; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Metrics; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.VectorSearch.Indexes +{ + /// + /// Tests for , the batched-GPU brute-force index, and the + /// helper it relies on. + /// + /// + /// In CI the GPU is disabled (the test assembly sets AIDOTNET_DISABLE_GPU=1), so the GPU + /// path deterministically falls back to the CPU brute-force scan. These tests therefore assert + /// two things: (1) the GPU index produces the exact same top-k as the reference CPU + /// on identical data, and (2) the code never throws when no GPU is + /// present. Both are the properties that matter for correctness on any machine — on a real GPU + /// the scores match within floating-point tolerance and the ordering is preserved. + /// + public class GpuFlatIndexTests + { + private const int Seed = 12345; + + private static List<(string Id, Vector)> GenerateVectors(int count, int dim, int seed) + { + var rng = new Random(seed); + var list = new List<(string, Vector)>(count); + for (int i = 0; i < count; i++) + { + var data = new double[dim]; + for (int j = 0; j < dim; j++) + { + // Non-zero, well-separated values so norms are never degenerate. + data[j] = rng.NextDouble() * 2.0 - 1.0 + 0.001; + } + list.Add(($"vec{i}", new Vector(data))); + } + return list; + } + + private static void PopulateBoth( + IVectorIndex a, + IVectorIndex b, + IReadOnlyList<(string Id, Vector Vec)> data) + { + foreach (var (id, vec) in data) + { + a.Add(id, vec); + b.Add(id, vec); + } + } + + private static void AssertTopKEqual( + List<(string Id, double Score)> expected, + List<(string Id, double Score)> actual) + { + Assert.Equal(expected.Count, actual.Count); + for (int i = 0; i < expected.Count; i++) + { + Assert.Equal(expected[i].Id, actual[i].Id); + Assert.Equal(expected[i].Score, actual[i].Score, 9); + } + } + + // ---- Constructor validation ---- + + [Fact(Timeout = 60000)] + public async Task Constructor_WithNullMetric_Throws() + { + Assert.Throws(() => new GpuFlatIndex(null!)); + } + + [Fact(Timeout = 60000)] + public async Task Constructor_WithNonPositiveThreshold_Throws() + { + var metric = new CosineSimilarityMetric(); + Assert.Throws(() => new GpuFlatIndex(metric, gpuThreshold: 0)); + } + + // ---- Correctness vs CPU brute force, across supported metrics ---- + + [Theory(Timeout = 120000)] + [InlineData("cosine")] + [InlineData("dot")] + [InlineData("euclidean")] + public async Task Search_MatchesCpuFlatIndex_AboveThreshold(string metricName) + { + // Large batch that crosses the default GPU threshold so the GPU path is attempted + // (and, in CI with no GPU, falls back to CPU). + var data = GenerateVectors(count: 2500, dim: 16, seed: Seed); + + var cpu = new FlatIndex(MakeMetric(metricName)); + var gpu = new GpuFlatIndex(MakeMetric(metricName)); + PopulateBoth(cpu, gpu, data); + + Assert.True(gpu.Count >= GpuVectorScorer.DefaultGpuThreshold); + + var query = data[7].Item2; + var expected = cpu.Search(query, 10); + var actual = gpu.Search(query, 10); + + AssertTopKEqual(expected, actual); + } + + [Theory(Timeout = 60000)] + [InlineData("cosine")] + [InlineData("dot")] + [InlineData("euclidean")] + public async Task Search_MatchesCpuFlatIndex_BelowThreshold(string metricName) + { + // Small batch: GPU path is skipped by the threshold, CPU scan is used directly. + var data = GenerateVectors(count: 50, dim: 8, seed: Seed + 1); + + var cpu = new FlatIndex(MakeMetric(metricName)); + var gpu = new GpuFlatIndex(MakeMetric(metricName)); + PopulateBoth(cpu, gpu, data); + + var query = new Vector(new double[] { 0.5, -0.2, 0.9, 0.1, -0.7, 0.3, 0.6, -0.4 }); + var expected = cpu.Search(query, 5); + var actual = gpu.Search(query, 5); + + AssertTopKEqual(expected, actual); + } + + [Theory(Timeout = 120000)] + [InlineData("cosine")] + [InlineData("dot")] + [InlineData("euclidean")] + public async Task Search_MatchesCpuFlatIndex_ForcedGpuThreshold(string metricName) + { + // Threshold of 1 forces the GPU attempt for every query; with GPU disabled the scorer + // signals fallback and the CPU scan runs — results must still equal the reference. + var data = GenerateVectors(count: 300, dim: 12, seed: Seed + 2); + + var cpu = new FlatIndex(MakeMetric(metricName)); + var gpu = new GpuFlatIndex(MakeMetric(metricName), gpuThreshold: 1); + PopulateBoth(cpu, gpu, data); + + var query = data[42].Item2; + var expected = cpu.Search(query, 20); + var actual = gpu.Search(query, 20); + + AssertTopKEqual(expected, actual); + } + + [Fact(Timeout = 60000)] + public async Task Search_WithUnsupportedMetric_FallsBackAndMatchesCpu() + { + // Manhattan has no GPU closed form -> scorer signals fallback regardless of batch size. + var data = GenerateVectors(count: 1500, dim: 10, seed: Seed + 3); + + var cpu = new FlatIndex(new ManhattanDistanceMetric()); + var gpu = new GpuFlatIndex(new ManhattanDistanceMetric(), gpuThreshold: 1); + PopulateBoth(cpu, gpu, data); + + var query = data[3].Item2; + var expected = cpu.Search(query, 7); + var actual = gpu.Search(query, 7); + + AssertTopKEqual(expected, actual); + } + + // ---- Robustness / edge cases: must never throw when GPU absent ---- + + [Fact(Timeout = 60000)] + public async Task Search_OnEmptyIndex_ReturnsEmpty() + { + var gpu = new GpuFlatIndex(new CosineSimilarityMetric(), gpuThreshold: 1); + var result = gpu.Search(new Vector(new double[] { 1.0, 2.0, 3.0 }), 5); + Assert.Empty(result); + } + + [Fact(Timeout = 60000)] + public async Task Search_WithKLargerThanCount_ReturnsAll() + { + var data = GenerateVectors(count: 2000, dim: 8, seed: Seed + 4); + var gpu = new GpuFlatIndex(new DotProductMetric()); + foreach (var (id, vec) in data) gpu.Add(id, vec); + + var result = gpu.Search(data[0].Item2, 5000); + Assert.Equal(2000, result.Count); + } + + [Fact(Timeout = 60000)] + public async Task Search_NullQuery_Throws() + { + var gpu = new GpuFlatIndex(new CosineSimilarityMetric()); + Assert.Throws(() => gpu.Search(null!, 5)); + } + + [Fact(Timeout = 60000)] + public async Task Search_NonPositiveK_Throws() + { + var gpu = new GpuFlatIndex(new CosineSimilarityMetric()); + gpu.Add("a", new Vector(new double[] { 1.0, 2.0 })); + Assert.Throws(() => gpu.Search(new Vector(new double[] { 1.0, 2.0 }), 0)); + } + + [Fact(Timeout = 60000)] + public async Task Remove_And_Clear_Work() + { + var gpu = new GpuFlatIndex(new CosineSimilarityMetric()); + gpu.Add("a", new Vector(new double[] { 1.0, 2.0 })); + gpu.Add("b", new Vector(new double[] { 3.0, 4.0 })); + + Assert.True(gpu.Remove("a")); + Assert.False(gpu.Remove("missing")); + Assert.Equal(1, gpu.Count); + + gpu.Clear(); + Assert.Equal(0, gpu.Count); + } + + // ---- GpuVectorScorer helper contract ---- + + [Fact(Timeout = 60000)] + public async Task Scorer_IsGpuDisabledByEnvironment_TrueInCi() + { + // The test assembly sets AIDOTNET_DISABLE_GPU=1, so the gate reports disabled and + // IsGpuAvailable is false. This documents the environment gate the rest of the codebase uses. + Assert.True(GpuVectorScorer.IsGpuDisabledByEnvironment); + Assert.False(GpuVectorScorer.IsGpuAvailable()); + } + + [Fact(Timeout = 60000)] + public async Task Scorer_TryScoreBatch_ReturnsFalseWhenGpuUnavailable() + { + var metric = new DotProductMetric(); + var query = new Vector(new double[] { 1.0, 2.0, 3.0 }); + var docs = GenerateVectors(2000, 3, Seed + 5).Select(x => x.Item2).ToList(); + + // GPU disabled in CI -> must signal fallback (false) without throwing. + bool ok = GpuVectorScorer.TryScoreBatch(metric, query, docs, gpuThreshold: 1, out var scores); + Assert.False(ok); + Assert.Empty(scores); + } + + [Fact(Timeout = 60000)] + public async Task Scorer_CanAccelerate_FalseWhenGpuDisabled() + { + var metric = new CosineSimilarityMetric(); + Assert.False(GpuVectorScorer.CanAccelerate(metric, count: 100000, gpuThreshold: 1)); + } + + // ---- HNSW batched candidate scoring helper ---- + + [Fact(Timeout = 120000)] + public async Task Hnsw_ScoreCandidatesGpu_MatchesCpuMetricOrdering() + { + var metric = new CosineSimilarityMetric(); + var data = GenerateVectors(count: 1500, dim: 16, seed: Seed + 6); + + var hnsw = new HNSWIndex(metric); + foreach (var (id, vec) in data) hnsw.Add(id, vec); + + var query = data[10].Item2; + var candidateIds = data.Select(x => x.Item1).ToList(); + + // GPU-batched (falls back to CPU metric in CI) ... + var scored = hnsw.ScoreCandidatesGpu(query, candidateIds, gpuThreshold: 1); + + // ... must equal an independent CPU brute-force ranking over the same candidates. + var expected = candidateIds + .Select(id => (Id: id, Score: metric.Calculate(query, data.First(d => d.Item1 == id).Item2))) + .OrderByDescending(x => x.Score) + .ToList(); + + Assert.Equal(expected.Count, scored.Count); + for (int i = 0; i < expected.Count; i++) + { + Assert.Equal(expected[i].Id, scored[i].Id); + Assert.Equal(expected[i].Score, scored[i].Score, 9); + } + } + + [Fact(Timeout = 60000)] + public async Task Hnsw_ScoreCandidatesGpu_SkipsMissingIds() + { + var metric = new DotProductMetric(); + var hnsw = new HNSWIndex(metric); + hnsw.Add("a", new Vector(new double[] { 1.0, 0.0 })); + hnsw.Add("b", new Vector(new double[] { 0.0, 1.0 })); + + var result = hnsw.ScoreCandidatesGpu( + new Vector(new double[] { 1.0, 1.0 }), + new List { "a", "missing", "b" }, + gpuThreshold: 1); + + Assert.Equal(2, result.Count); + Assert.Contains(result, r => r.Id == "a"); + Assert.Contains(result, r => r.Id == "b"); + } + + private static ISimilarityMetric MakeMetric(string name) => name switch + { + "cosine" => new CosineSimilarityMetric(), + "dot" => new DotProductMetric(), + "euclidean" => new EuclideanDistanceMetric(), + _ => throw new ArgumentOutOfRangeException(nameof(name), name, "unknown metric") + }; + } +} From 6abcbf357f3fbe6edc437c6bdbfedb340c84d8eb Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 23:08:51 -0400 Subject: [PATCH 23/25] feat(rag): boolean metadata filter AST (AND/OR/NOT/in/range) across all stores MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a rich, composable metadata-filter model — Eq/Ne/Gt/Gte/Lt/Lte/In/Exists + And/Or/Not with fluent factories — and an additive IDocumentStore.GetSimilarWithFilter[Async] that every real store translates to its native filter language: Qdrant (must/should/must_not), Pinecone ($and/$or/$in/$gte/…), Weaviate (where operators), Milvus (bool expr), Azure (OData), pgvector (jsonb, parameterized), Elasticsearch (bool query), Redis (RediSearch); InMemory/File use the base AST evaluator. Non-breaking: the old Dictionary filters remain, and DocumentStoreBase provides an in-memory evaluator default so any store gets correct rich filtering for free. This matches/exceeds the Pinecone/Qdrant/Weaviate filter expressiveness (previously only equality + one-sided >= range + any-of). Fixed two issues in the subagent's draft: a string-vs-number comparison incorrectly ordinal-compared (Gt("name",5) "matched" "abc") — now only two strings order against each other, otherwise not-comparable; and four sync filter tests wrongly carried [Fact(Timeout)]. All 78 filter/translation/store tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ElasticsearchDocumentStore.cs | 31 +- .../ElasticsearchVectorFilterBuilder.cs | 119 +++++ .../PostgresVectorDocumentStore.cs | 20 + .../PostgresVectorFilterBuilder.cs | 98 ++++ .../DocumentStores/RedisVLDocumentStore.cs | 53 +- .../DocumentStores/RedisVectorQueryBuilder.cs | 123 +++++ src/AiModelBuilder.VectorIndexStore.cs | 35 ++ src/Interfaces/IDocumentStore.cs | 31 ++ .../AzureSearchDocumentStore.cs | 75 ++- .../DocumentStores/DocumentStoreBase.cs | 95 ++++ .../DocumentStores/MilvusDocumentStore.cs | 74 ++- .../DocumentStores/PineconeDocumentStore.cs | 105 +++- .../DocumentStores/QdrantDocumentStore.cs | 134 ++++- .../DocumentStores/WeaviateDocumentStore.cs | 137 ++++- .../Filtering/MetadataFilter.cs | 480 ++++++++++++++++++ tests/AiDotNet.Tests/AiDotNetTests.csproj | 1 + .../BM25RetrieverTests.cs | 4 + .../ColBERTRetrieverTests.cs | 4 + .../DenseRetrieverTests.cs | 4 + .../InMemoryDocumentStoreTests.cs | 88 ++++ .../MetadataFilterTranslationTests.cs | 316 ++++++++++++ .../Filtering/MetadataFilterTests.cs | 266 ++++++++++ .../GraphRetrieverTests.cs | 4 + .../MultiVectorRetrieverTests.cs | 4 + .../ParentDocumentRetrieverTests.cs | 4 + .../Retrievers/TestHelpers.cs | 4 + .../TFIDFRetrieverTests.cs | 4 + .../VectorRetrieverTests.cs | 4 + 28 files changed, 2287 insertions(+), 30 deletions(-) create mode 100644 src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchVectorFilterBuilder.cs create mode 100644 src/RetrievalAugmentedGeneration/Filtering/MetadataFilter.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/MetadataFilterTranslationTests.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Filtering/MetadataFilterTests.cs diff --git a/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchDocumentStore.cs b/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchDocumentStore.cs index 75bbb9b343..ad1765868b 100644 --- a/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchDocumentStore.cs +++ b/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchDocumentStore.cs @@ -10,6 +10,7 @@ using AiDotNet.Enums; using AiDotNet.Interfaces; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; using Newtonsoft.Json.Linq; @@ -236,12 +237,19 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); - private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) - { - var embedding = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); + private Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, BuildDictQueryClause(metadataFilters), cancellationToken); + + /// + protected override IEnumerable> GetSimilarWithFilterCore(Vector queryVector, MetadataFilter filter, int topK) + => SearchImplAsync(queryVector, topK, ElasticsearchVectorFilterBuilder.Build(filter), CancellationToken.None).GetAwaiter().GetResult(); - // Build the query with metadata filters - object queryClause; + /// + protected override Task>> GetSimilarWithFilterCoreAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, ElasticsearchVectorFilterBuilder.Build(filter), cancellationToken); + + private static object BuildDictQueryClause(Dictionary metadataFilters) + { if (metadataFilters != null && metadataFilters.Any()) { var mustClauses = new List(); @@ -255,7 +263,7 @@ private async Task>> GetSimilarCoreImplAsync(Vector q } }); } - queryClause = new + return new { @bool = new { @@ -263,10 +271,13 @@ private async Task>> GetSimilarCoreImplAsync(Vector q } }; } - else - { - queryClause = new { match_all = new { } }; - } + + return new { match_all = new { } }; + } + + private async Task>> SearchImplAsync(Vector queryVector, int topK, object queryClause, CancellationToken cancellationToken) + { + var embedding = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); var query = new { diff --git a/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchVectorFilterBuilder.cs b/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchVectorFilterBuilder.cs new file mode 100644 index 0000000000..17f74cd363 --- /dev/null +++ b/src/AiDotNet.Storage.Elasticsearch/RetrievalAugmentedGeneration/DocumentStores/ElasticsearchVectorFilterBuilder.cs @@ -0,0 +1,119 @@ +using System; +using System.Collections.Generic; +using System.Linq; + +using AiDotNet.RetrievalAugmentedGeneration.Filtering; + +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; + +/// +/// Translates an AiDotNet expression tree into an Elasticsearch Query DSL +/// query clause (a bool query built from term/terms/range/exists leaf +/// queries). Extracted from so the query-generation logic can +/// be unit-tested without a live Elasticsearch cluster. +/// +/// +/// Translation rules (all metadata fields are addressed under the metadata. prefix): +/// +/// Eq → term; Ne → bool.must_not[term]. +/// Gt/Gte/Lt/Lte → range with gt/gte/lt/lte. +/// In → terms; Exists → exists. +/// And → bool.must; Or → bool.should (minimum_should_match 1); Not → bool.must_not. +/// +/// +public static class ElasticsearchVectorFilterBuilder +{ + /// + /// Builds the Elasticsearch query clause representing . Returns a + /// match_all query when is null. + /// + public static object Build(MetadataFilter? filter) + { + if (filter == null) + return new Dictionary { ["match_all"] = new Dictionary() }; + + return BuildClause(filter); + } + + private static object BuildClause(MetadataFilter filter) + { + switch (filter) + { + case ComparisonFilter comparison: + return BuildComparison(comparison); + case InFilter inFilter: + return new Dictionary + { + ["terms"] = new Dictionary { [Field(inFilter.Key)] = inFilter.Values.ToArray() } + }; + case ExistsFilter existsFilter: + return new Dictionary + { + ["exists"] = new Dictionary { ["field"] = Field(existsFilter.Key) } + }; + case NotFilter notFilter: + return Bool("must_not", new[] { BuildClause(notFilter.Operand) }); + case LogicalFilter logical when logical.Operator == MetadataFilterOperator.And: + return Bool("must", logical.Operands.Select(BuildClause).ToArray()); + case LogicalFilter logical: + { + var should = Bool("should", logical.Operands.Select(BuildClause).ToArray()); + ((Dictionary)should["bool"])["minimum_should_match"] = 1; + return should; + } + default: + throw new NotSupportedException($"Unsupported metadata filter node: {filter.GetType().Name}"); + } + } + + private static object BuildComparison(ComparisonFilter comparison) + { + switch (comparison.Operator) + { + case MetadataFilterOperator.Eq: + return new Dictionary + { + ["term"] = new Dictionary { [Field(comparison.Key)] = comparison.Value } + }; + case MetadataFilterOperator.Ne: + return Bool("must_not", new object[] + { + new Dictionary + { + ["term"] = new Dictionary { [Field(comparison.Key)] = comparison.Value } + } + }); + case MetadataFilterOperator.Gt: + return Range(comparison.Key, "gt", comparison.Value); + case MetadataFilterOperator.Gte: + return Range(comparison.Key, "gte", comparison.Value); + case MetadataFilterOperator.Lt: + return Range(comparison.Key, "lt", comparison.Value); + case MetadataFilterOperator.Lte: + return Range(comparison.Key, "lte", comparison.Value); + default: + throw new NotSupportedException($"Unsupported comparison operator: {comparison.Operator}"); + } + } + + private static object Range(string key, string op, object value) + { + return new Dictionary + { + ["range"] = new Dictionary + { + [Field(key)] = new Dictionary { [op] = value } + } + }; + } + + private static Dictionary Bool(string occurrence, object[] clauses) + { + return new Dictionary + { + ["bool"] = new Dictionary { [occurrence] = clauses } + }; + } + + private static string Field(string key) => "metadata." + key; +} diff --git a/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs index b14fa437e5..45629a870a 100644 --- a/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs +++ b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorDocumentStore.cs @@ -9,6 +9,7 @@ using AiDotNet.Enums; using AiDotNet.Interfaces; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; using AiDotNet.Tensors.LinearAlgebra; @@ -216,7 +217,26 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector { var parameters = new Dictionary(); var whereClause = PostgresVectorFilterBuilder.Build(metadataFilters, parameters); + return ExecuteSearch(queryVector, topK, whereClause, parameters); + } + + /// + protected override IEnumerable> GetSimilarWithFilterCore(Vector queryVector, MetadataFilter filter, int topK) + { + var parameters = new Dictionary(); + var whereClause = PostgresVectorFilterBuilder.Build(filter, parameters); + return ExecuteSearch(queryVector, topK, whereClause, parameters); + } + /// + protected override System.Threading.Tasks.Task>> GetSimilarWithFilterCoreAsync(Vector queryVector, MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilterCore(queryVector, filter, topK)); + } + + private IEnumerable> ExecuteSearch(Vector queryVector, int topK, string whereClause, Dictionary parameters) + { using var connection = Open(); using var command = connection.CreateCommand(); command.CommandText = diff --git a/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorFilterBuilder.cs b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorFilterBuilder.cs index 0437cdd22b..912a45afe7 100644 --- a/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorFilterBuilder.cs +++ b/src/AiDotNet.Storage.Postgres/RetrievalAugmentedGeneration/DocumentStores/PostgresVectorFilterBuilder.cs @@ -1,9 +1,12 @@ +using System; using System.Collections; using System.Collections.Generic; using System.Globalization; using System.Linq; using System.Text; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; + namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; /// @@ -90,6 +93,101 @@ public static string Build(IReadOnlyDictionary? filters, IDictio return sb.ToString(); } + /// + /// Builds a parameterised WHERE clause (including the leading WHERE) from a rich + /// boolean expression tree, appending the required bound parameters to + /// . Returns the empty string when is null. + /// + /// + /// Both metadata keys and values are supplied as bound parameters (never string-interpolated). + /// Numeric comparisons cast the extracted text to numeric; existence uses + /// jsonb_exists(metadata, @k); AND/OR/NOT map to SQL boolean operators. + /// + public static string Build(MetadataFilter? filter, IDictionary parameters) + { + if (filter == null) + { + return string.Empty; + } + + var counter = new int[1]; + var clause = BuildClause(filter, parameters, counter); + return " WHERE " + clause; + } + + private static string BuildClause(MetadataFilter filter, IDictionary parameters, int[] counter) + { + switch (filter) + { + case ComparisonFilter comparison: + return BuildComparison(comparison, parameters, counter); + case InFilter inFilter: + { + var i = counter[0]++; + var keyParam = "k" + i.ToString(CultureInfo.InvariantCulture); + var valParam = "v" + i.ToString(CultureInfo.InvariantCulture); + parameters[keyParam] = inFilter.Key; + parameters[valParam] = inFilter.Values + .Select(o => o?.ToString() ?? string.Empty) + .ToArray(); + return $"(metadata ->> @{keyParam}) = ANY(@{valParam})"; + } + case ExistsFilter existsFilter: + { + var i = counter[0]++; + var keyParam = "k" + i.ToString(CultureInfo.InvariantCulture); + parameters[keyParam] = existsFilter.Key; + return $"jsonb_exists(metadata, @{keyParam})"; + } + case NotFilter notFilter: + return "(NOT (" + BuildClause(notFilter.Operand, parameters, counter) + "))"; + case LogicalFilter logical when logical.Operator == MetadataFilterOperator.And: + return "(" + string.Join(" AND ", logical.Operands.Select(o => BuildClause(o, parameters, counter))) + ")"; + case LogicalFilter logical: + return "(" + string.Join(" OR ", logical.Operands.Select(o => BuildClause(o, parameters, counter))) + ")"; + default: + throw new NotSupportedException($"Unsupported metadata filter node: {filter.GetType().Name}"); + } + } + + private static string BuildComparison(ComparisonFilter comparison, IDictionary parameters, int[] counter) + { + var i = counter[0]++; + var keyParam = "k" + i.ToString(CultureInfo.InvariantCulture); + var valParam = "v" + i.ToString(CultureInfo.InvariantCulture); + parameters[keyParam] = comparison.Key; + + var numeric = IsNumeric(comparison.Value); + var left = numeric + ? $"(metadata ->> @{keyParam})::numeric" + : $"(metadata ->> @{keyParam})"; + + if (numeric) + parameters[valParam] = System.Convert.ToDouble(comparison.Value, CultureInfo.InvariantCulture); + else if (comparison.Value is bool b) + parameters[valParam] = b ? "true" : "false"; + else + parameters[valParam] = comparison.Value?.ToString() ?? string.Empty; + + switch (comparison.Operator) + { + case MetadataFilterOperator.Eq: + return $"{left} = @{valParam}"; + case MetadataFilterOperator.Ne: + return $"{left} IS DISTINCT FROM @{valParam}"; + case MetadataFilterOperator.Gt: + return $"{left} > @{valParam}"; + case MetadataFilterOperator.Gte: + return $"{left} >= @{valParam}"; + case MetadataFilterOperator.Lt: + return $"{left} < @{valParam}"; + case MetadataFilterOperator.Lte: + return $"{left} <= @{valParam}"; + default: + throw new NotSupportedException($"Unsupported comparison operator: {comparison.Operator}"); + } + } + private static bool IsNumeric(object? value) { return value is sbyte || value is byte || value is short || value is ushort diff --git a/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs index ccb1f0690d..d28fcc5e27 100644 --- a/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs +++ b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVLDocumentStore.cs @@ -3,11 +3,14 @@ using System.Globalization; using System.Linq; using System.Text.RegularExpressions; +using System.Threading; +using System.Threading.Tasks; using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.Interfaces; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; using AiDotNet.Tensors.LinearAlgebra; @@ -302,6 +305,42 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector var expr = RedisVectorQueryBuilder.BuildFilterExpression(metadataFilters, _fieldTypes, out var unpushedKeys); var knnK = unpushedKeys.Count > 0 ? topK * CandidateMultiplier : topK; + IEnumerable<(Document Doc, double Distance)> candidates = ExecuteKnn(queryVector, expr, knnK); + if (unpushedKeys.Count > 0) + { + var postFilters = unpushedKeys + .Where(metadataFilters.ContainsKey) + .ToDictionary(k => k, k => metadataFilters[k]); + candidates = candidates.Where(c => MatchesFilters(c.Doc, postFilters)); + } + + return TakeTop(candidates, topK); + } + + /// + protected override IEnumerable> GetSimilarWithFilterCore(Vector queryVector, MetadataFilter filter, int topK) + { + // Push the whole expression to RediSearch when every referenced field is declared and every + // operator is supported; otherwise over-fetch and evaluate the full AST in memory. + var expr = RedisVectorQueryBuilder.BuildFilterExpression(filter, _fieldTypes, out var fullyPushed); + var knnK = fullyPushed ? topK : topK * CandidateMultiplier; + + IEnumerable<(Document Doc, double Distance)> candidates = ExecuteKnn(queryVector, expr, knnK); + if (!fullyPushed) + candidates = candidates.Where(c => filter.Matches(c.Doc.Metadata)); + + return TakeTop(candidates, topK); + } + + /// + protected override Task>> GetSimilarWithFilterCoreAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GetSimilarWithFilterCore(queryVector, filter, topK)); + } + + private List<(Document Doc, double Distance)> ExecuteKnn(Vector queryVector, string expr, int knnK) + { var query = $"{expr}=>[KNN {knnK.ToString(CultureInfo.InvariantCulture)} @{EmbeddingField} $BLOB AS {ScoreField}]"; var args = new List { @@ -313,17 +352,11 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector "LIMIT", "0", knnK.ToString(CultureInfo.InvariantCulture), }; - var parsed = ParseSearch(Db.Execute("FT.SEARCH", args), includeScore: true); - - IEnumerable<(Document Doc, double Distance)> candidates = parsed; - if (unpushedKeys.Count > 0) - { - var postFilters = unpushedKeys - .Where(metadataFilters.ContainsKey) - .ToDictionary(k => k, k => metadataFilters[k]); - candidates = candidates.Where(c => MatchesFilters(c.Doc, postFilters)); - } + return ParseSearch(Db.Execute("FT.SEARCH", args), includeScore: true); + } + private List> TakeTop(IEnumerable<(Document Doc, double Distance)> candidates, int topK) + { var results = new List>(); foreach (var candidate in candidates.Take(topK)) { diff --git a/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorQueryBuilder.cs b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorQueryBuilder.cs index 557b110ed4..66e9f30ab6 100644 --- a/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorQueryBuilder.cs +++ b/src/AiDotNet.Storage.Redis/RetrievalAugmentedGeneration/DocumentStores/RedisVectorQueryBuilder.cs @@ -4,6 +4,8 @@ using System.Linq; using System.Text; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; + namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; /// @@ -65,6 +67,127 @@ public static string BuildFilterExpression( return clauses.Count > 0 ? "(" + string.Join(" ", clauses) + ")" : "*"; } + /// + /// Builds a RediSearch pre-filter expression from a rich boolean tree. + /// The whole expression is pushed down only when every referenced field is declared and every + /// operator is supported by RediSearch (TAG equality/membership, NUMERIC comparisons/ranges, and + /// AND/OR/NOT composition). Otherwise * is returned with + /// set to false, signalling the caller to evaluate the entire + /// filter in memory over an over-fetched candidate set. + /// + public static string BuildFilterExpression( + MetadataFilter? filter, + IReadOnlyDictionary declaredFields, + out bool fullyPushed) + { + if (filter == null) + { + fullyPushed = true; + return "*"; + } + + if (TryBuildExpression(filter, declaredFields, out var expr)) + { + fullyPushed = true; + return expr; + } + + fullyPushed = false; + return "*"; + } + + private static bool TryBuildExpression( + MetadataFilter filter, + IReadOnlyDictionary declaredFields, + out string expr) + { + expr = string.Empty; + switch (filter) + { + case ComparisonFilter comparison: + return declaredFields.TryGetValue(comparison.Key, out var compType) + && TryBuildComparison(comparison, compType, out expr); + case InFilter inFilter: + return declaredFields.TryGetValue(inFilter.Key, out var inType) + && TryBuildIn(inFilter, inType, out expr); + case ExistsFilter: + // RediSearch has no portable "field exists" predicate for HASH indexes. + return false; + case NotFilter notFilter: + if (!TryBuildExpression(notFilter.Operand, declaredFields, out var inner)) + return false; + expr = "-(" + inner + ")"; + return true; + case LogicalFilter logical: + var parts = new List(); + foreach (var operand in logical.Operands) + { + if (!TryBuildExpression(operand, declaredFields, out var part)) + return false; + parts.Add(part); + } + var separator = logical.Operator == MetadataFilterOperator.And ? " " : " | "; + expr = "(" + string.Join(separator, parts) + ")"; + return true; + default: + return false; + } + } + + private static bool TryBuildComparison(ComparisonFilter comparison, RedisVectorFieldType fieldType, out string expr) + { + expr = string.Empty; + if (fieldType == RedisVectorFieldType.Numeric) + { + if (!MetadataFilter.TryToDouble(comparison.Value, out var d)) + return false; + var v = d.ToString("R", CultureInfo.InvariantCulture); + switch (comparison.Operator) + { + case MetadataFilterOperator.Eq: expr = $"@{comparison.Key}:[{v} {v}]"; return true; + case MetadataFilterOperator.Ne: expr = $"-@{comparison.Key}:[{v} {v}]"; return true; + case MetadataFilterOperator.Gt: expr = $"@{comparison.Key}:[({v} +inf]"; return true; + case MetadataFilterOperator.Gte: expr = $"@{comparison.Key}:[{v} +inf]"; return true; + case MetadataFilterOperator.Lt: expr = $"@{comparison.Key}:[-inf ({v}]"; return true; + case MetadataFilterOperator.Lte: expr = $"@{comparison.Key}:[-inf {v}]"; return true; + default: return false; + } + } + + // TAG field: only exact equality / inequality are expressible. + switch (comparison.Operator) + { + case MetadataFilterOperator.Eq: expr = $"@{comparison.Key}:{{{FormatTagValue(comparison.Value)}}}"; return true; + case MetadataFilterOperator.Ne: expr = $"-@{comparison.Key}:{{{FormatTagValue(comparison.Value)}}}"; return true; + default: return false; + } + } + + private static bool TryBuildIn(InFilter inFilter, RedisVectorFieldType fieldType, out string expr) + { + expr = string.Empty; + if (inFilter.Values.Count == 0) + return false; + + if (fieldType == RedisVectorFieldType.Numeric) + { + var ranges = new List(); + foreach (var value in inFilter.Values) + { + if (!MetadataFilter.TryToDouble(value, out var d)) + return false; + var v = d.ToString("R", CultureInfo.InvariantCulture); + ranges.Add($"@{inFilter.Key}:[{v} {v}]"); + } + expr = "(" + string.Join(" | ", ranges) + ")"; + return true; + } + + var items = inFilter.Values.Select(FormatTagValue); + expr = $"@{inFilter.Key}:{{{string.Join("|", items)}}}"; + return true; + } + private static string FormatTagValue(object? value) { if (value is string s) diff --git a/src/AiModelBuilder.VectorIndexStore.cs b/src/AiModelBuilder.VectorIndexStore.cs index edc6164557..54f3dcff51 100644 --- a/src/AiModelBuilder.VectorIndexStore.cs +++ b/src/AiModelBuilder.VectorIndexStore.cs @@ -5,6 +5,7 @@ using System.Threading.Tasks; using AiDotNet.Interfaces; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Indexes; @@ -112,6 +113,40 @@ public IEnumerable> GetSimilarWithFilters(Vector queryVector, int return results; } + /// + public IEnumerable> GetSimilarWithFilter(Vector queryVector, MetadataFilter filter, int topK) + { + if (queryVector == null) throw new ArgumentNullException(nameof(queryVector)); + if (topK <= 0) return Enumerable.Empty>(); + if (filter == null) return GetSimilar(queryVector, topK); + + int searchK = Math.Min(_documents.Count, Math.Max(topK * 10, topK)); + if (searchK <= 0) return Enumerable.Empty>(); + + var hits = _index.Search(queryVector, searchK); + var results = new List>(topK); + + foreach (var (id, score) in hits) + { + if (!_documents.TryGetValue(id, out var vd)) continue; + if (!filter.Matches(vd.Document.Metadata)) continue; + + var doc = vd.Document; + doc.RelevanceScore = score; + results.Add(doc); + if (results.Count >= topK) break; + } + + return results; + } + + /// + public Task>> GetSimilarWithFilterAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken = default) + { + cancellationToken.ThrowIfCancellationRequested(); + return Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); + } + /// public Document? GetById(string documentId) => _documents.TryGetValue(documentId, out var vd) ? vd.Document : null; diff --git a/src/Interfaces/IDocumentStore.cs b/src/Interfaces/IDocumentStore.cs index ff839ba1ef..5ecd2eb211 100644 --- a/src/Interfaces/IDocumentStore.cs +++ b/src/Interfaces/IDocumentStore.cs @@ -1,6 +1,7 @@ using System.Threading; using System.Threading.Tasks; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.Interfaces; @@ -154,6 +155,29 @@ public interface IDocumentStore /// IEnumerable> GetSimilarWithFilters(Vector queryVector, int topK, Dictionary metadataFilters); + /// + /// Retrieves similar documents using a rich, boolean expression tree. + /// + /// The vector to search for similar documents. + /// + /// The metadata filter expression (supports AND/OR/NOT, equality, inequality, + /// full ranges, set membership and existence). When null, no filtering is applied. + /// + /// The number of most similar documents to return. + /// An enumerable of filtered documents ordered by similarity, with relevance scores populated. + /// + /// + /// This is the rich-filtering counterpart of . + /// Stores with a native filter language translate the expression tree into it (Pinecone, Qdrant, + /// Weaviate, Milvus, Azure AI Search, pgvector, Elasticsearch); other stores evaluate it in memory + /// over an over-fetched candidate set. + /// + /// For Beginners: The lets you express complex conditions + /// such as "category == science AND (year >= 2020 OR author in [A, B]) AND NOT archived". + /// + /// + IEnumerable> GetSimilarWithFilter(Vector queryVector, MetadataFilter filter, int topK); + /// /// Retrieves a document by its unique identifier. /// @@ -240,6 +264,13 @@ public interface IDocumentStore /// A token to observe for cancellation requests. Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken = default); + /// Asynchronously retrieves similar documents using a rich, boolean expression tree. + /// The vector to search for similar documents. + /// The metadata filter expression; when null, no filtering is applied. + /// The number of most similar documents to return. + /// A token to observe for cancellation requests. + Task>> GetSimilarWithFilterAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken = default); + /// Asynchronously retrieves a document by its unique identifier. /// The unique identifier of the document to retrieve. /// A token to observe for cancellation requests. diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs index d5b6997e2a..28c1e5cae5 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/AzureSearchDocumentStore.cs @@ -10,6 +10,7 @@ using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; using Newtonsoft.Json; using Newtonsoft.Json.Linq; @@ -312,7 +313,18 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); - private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + private Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, BuildODataFilter(metadataFilters), cancellationToken); + + /// + protected override IEnumerable> GetSimilarWithFilterCore(Vector queryVector, MetadataFilter filter, int topK) + => SearchImplAsync(queryVector, topK, TranslateFilter(filter), CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarWithFilterCoreAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, TranslateFilter(filter), cancellationToken); + + private async Task>> SearchImplAsync(Vector queryVector, int topK, string? filter, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); @@ -326,7 +338,6 @@ private async Task>> GetSimilarCoreImplAsync(Vector q ["top"] = topK }; - var filter = BuildODataFilter(metadataFilters); if (filter != null) body["filter"] = filter; @@ -421,6 +432,66 @@ private async Task>> GetSimilarCoreImplAsync(Vector q // OData string literals are single-quoted; embedded single quotes are doubled. private static string ODataString(string value) => "'" + value.Replace("'", "''") + "'"; + // ------------------------------------------------------------------ + // MetadataFilter AST -> Azure AI Search OData $filter translation. + // ------------------------------------------------------------------ + + /// Translates a expression tree into an Azure AI Search OData $filter. + internal static string TranslateFilter(MetadataFilter filter) + { + if (filter == null) + throw new ArgumentNullException(nameof(filter)); + return BuildOData(filter); + } + + private static string BuildOData(MetadataFilter filter) + { + switch (filter) + { + case ComparisonFilter comparison: + return comparison.Key + " " + ODataOperator(comparison.Operator) + " " + ODataValue(comparison.Value); + case InFilter inFilter: + // search.in(field, 'a|b|c', '|') - '|' delimiter avoids clashing with commas in values. + var joined = string.Join("|", inFilter.Values.Select(v => MetadataFilter.ToInvariantString(v))); + return "search.in(" + inFilter.Key + ", " + ODataString(joined) + ", '|')"; + case ExistsFilter existsFilter: + return existsFilter.Key + " ne null"; + case NotFilter notFilter: + return "not (" + BuildOData(notFilter.Operand) + ")"; + case LogicalFilter logical when logical.Operator == MetadataFilterOperator.And: + return "(" + string.Join(" and ", logical.Operands.Select(BuildOData)) + ")"; + case LogicalFilter logical: + return "(" + string.Join(" or ", logical.Operands.Select(BuildOData)) + ")"; + default: + throw new NotSupportedException($"Unsupported metadata filter node: {filter.GetType().Name}"); + } + } + + private static string ODataValue(object value) + { + if (value == null) + return "null"; + if (value is bool b) + return b ? "true" : "false"; + if (MetadataFilter.IsNumeric(value)) + return Convert.ToDouble(value, CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture); + return ODataString(value.ToString() ?? string.Empty); + } + + private static string ODataOperator(MetadataFilterOperator op) + { + switch (op) + { + case MetadataFilterOperator.Eq: return "eq"; + case MetadataFilterOperator.Ne: return "ne"; + case MetadataFilterOperator.Gt: return "gt"; + case MetadataFilterOperator.Gte: return "ge"; + case MetadataFilterOperator.Lt: return "lt"; + case MetadataFilterOperator.Lte: return "le"; + default: throw new NotSupportedException($"Unsupported comparison operator: {op}"); + } + } + /// protected override Document? GetByIdCore(string documentId) => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/DocumentStoreBase.cs b/src/RetrievalAugmentedGeneration/DocumentStores/DocumentStoreBase.cs index 2c66f4b3ec..e96420be38 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/DocumentStoreBase.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/DocumentStoreBase.cs @@ -4,6 +4,7 @@ using AiDotNet.Interfaces; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; @@ -102,6 +103,100 @@ public IEnumerable> GetSimilarWithFilters(Vector queryVector, int return GetSimilarCore(queryVector, topK, metadataFilters); } + /// + /// Over-fetch multiplier used by the default in-memory filter evaluator: the base implementation + /// requests topK * FilterOverFetchFactor nearest neighbours (bounded by the store size) + /// before applying the , so enough candidates survive filtering. + /// + protected virtual int FilterOverFetchFactor => 10; + + /// + public IEnumerable> GetSimilarWithFilter(Vector queryVector, MetadataFilter filter, int topK) + { + ValidateQueryVector(queryVector); + ValidateTopK(topK); + + if (filter == null) + return GetSimilar(queryVector, topK); + + return GetSimilarWithFilterCore(queryVector, filter, topK); + } + + /// + public virtual async Task>> GetSimilarWithFilterAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken = default) + { + ValidateQueryVector(queryVector); + ValidateTopK(topK); + cancellationToken.ThrowIfCancellationRequested(); + + if (filter == null) + return await GetSimilarAsync(queryVector, topK, cancellationToken).ConfigureAwait(false); + + return await GetSimilarWithFilterCoreAsync(queryVector, filter, topK, cancellationToken).ConfigureAwait(false); + } + + /// + /// Core rich-filter search. The default implementation over-fetches nearest neighbours (ignoring + /// metadata at the storage layer) and evaluates the in memory via + /// . Stores with a native + /// filter language override this to push the translated filter down to the backend. + /// + /// The validated query vector. + /// The non-null metadata filter expression. + /// The validated number of documents to return. + /// Top-k similar documents that satisfy . + protected virtual IEnumerable> GetSimilarWithFilterCore(Vector queryVector, MetadataFilter filter, int topK) + { + var candidateK = ComputeCandidateFetchCount(topK); + var candidates = GetSimilarCore(queryVector, candidateK, new Dictionary()); + return ApplyFilter(candidates, filter, topK); + } + + /// + /// Core asynchronous rich-filter search. The default over-fetches via + /// and evaluates the filter in memory. Override for true server-side rich filtering. + /// + protected virtual async Task>> GetSimilarWithFilterCoreAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken) + { + var candidateK = ComputeCandidateFetchCount(topK); + var candidates = await GetSimilarCoreAsync(queryVector, candidateK, new Dictionary(), cancellationToken).ConfigureAwait(false); + return ApplyFilter(candidates, filter, topK); + } + + /// + /// Computes the over-fetch candidate count for the default in-memory evaluator, bounded by the + /// number of documents in the store when that count is known (positive). + /// + protected int ComputeCandidateFetchCount(int topK) + { + var factor = FilterOverFetchFactor < 1 ? 1 : FilterOverFetchFactor; + long desired = (long)topK * factor; + var docCount = DocumentCount; + if (docCount > 0 && desired > docCount) + desired = docCount; + if (desired > int.MaxValue) + desired = int.MaxValue; + if (desired < topK) + desired = topK; + return (int)desired; + } + + private static IEnumerable> ApplyFilter(IEnumerable> candidates, MetadataFilter filter, int topK) + { + var results = new List>(); + foreach (var document in candidates) + { + if (filter.Matches(document.Metadata)) + { + results.Add(document); + if (results.Count >= topK) + break; + } + } + + return results; + } + /// /// Retrieves a document by its unique identifier. /// diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs index dbae091f82..d852ad405d 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/MilvusDocumentStore.cs @@ -10,6 +10,7 @@ using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; using Newtonsoft.Json; using Newtonsoft.Json.Linq; @@ -269,7 +270,18 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); - private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + private Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, BuildFilterExpression(metadataFilters), cancellationToken); + + /// + protected override IEnumerable> GetSimilarWithFilterCore(Vector queryVector, MetadataFilter filter, int topK) + => SearchImplAsync(queryVector, topK, TranslateFilter(filter), CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarWithFilterCoreAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, TranslateFilter(filter), cancellationToken); + + private async Task>> SearchImplAsync(Vector queryVector, int topK, string? filter, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); @@ -282,7 +294,6 @@ private async Task>> GetSimilarCoreImplAsync(Vector q ["outputFields"] = new[] { IdField, ContentField, MetadataField } }; - var filter = BuildFilterExpression(metadataFilters); if (filter != null) body["filter"] = filter; @@ -375,6 +386,65 @@ private async Task>> GetSimilarCoreImplAsync(Vector q private static string Quote(string value) => "\"" + value.Replace("\\", "\\\\").Replace("\"", "\\\"") + "\""; + // ------------------------------------------------------------------ + // MetadataFilter AST -> Milvus boolean filter expression translation. + // Filterable fields are stored flattened under the "m_" prefix. + // ------------------------------------------------------------------ + + /// Translates a expression tree into a Milvus boolean expression. + internal static string TranslateFilter(MetadataFilter filter) + { + if (filter == null) + throw new ArgumentNullException(nameof(filter)); + return BuildExpression(filter); + } + + private static string BuildExpression(MetadataFilter filter) + { + switch (filter) + { + case ComparisonFilter comparison: + return MetaPrefix + comparison.Key + " " + MilvusOperator(comparison.Operator) + " " + FormatValue(comparison.Value); + case InFilter inFilter: + return MetaPrefix + inFilter.Key + " in [" + string.Join(", ", inFilter.Values.Select(FormatValue)) + "]"; + case ExistsFilter existsFilter: + return "exists " + MetaPrefix + existsFilter.Key; + case NotFilter notFilter: + return "not (" + BuildExpression(notFilter.Operand) + ")"; + case LogicalFilter logical when logical.Operator == MetadataFilterOperator.And: + return "(" + string.Join(" and ", logical.Operands.Select(BuildExpression)) + ")"; + case LogicalFilter logical: + return "(" + string.Join(" or ", logical.Operands.Select(BuildExpression)) + ")"; + default: + throw new NotSupportedException($"Unsupported metadata filter node: {filter.GetType().Name}"); + } + } + + private static string FormatValue(object value) + { + if (value == null) + return "null"; + if (value is bool b) + return b ? "true" : "false"; + if (MetadataFilter.IsNumeric(value)) + return Convert.ToDouble(value, CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture); + return Quote(value.ToString() ?? string.Empty); + } + + private static string MilvusOperator(MetadataFilterOperator op) + { + switch (op) + { + case MetadataFilterOperator.Eq: return "=="; + case MetadataFilterOperator.Ne: return "!="; + case MetadataFilterOperator.Gt: return ">"; + case MetadataFilterOperator.Gte: return ">="; + case MetadataFilterOperator.Lt: return "<"; + case MetadataFilterOperator.Lte: return "<="; + default: throw new NotSupportedException($"Unsupported comparison operator: {op}"); + } + } + /// protected override Document? GetByIdCore(string documentId) => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs index 43592c9370..d08a1db0d3 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/PineconeDocumentStore.cs @@ -10,6 +10,7 @@ using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; using Newtonsoft.Json; using Newtonsoft.Json.Linq; @@ -240,7 +241,18 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); - private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + private Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, BuildFilter(metadataFilters), cancellationToken); + + /// + protected override IEnumerable> GetSimilarWithFilterCore(Vector queryVector, MetadataFilter filter, int topK) + => SearchImplAsync(queryVector, topK, TranslateFilter(filter), CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarWithFilterCoreAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, TranslateFilter(filter), cancellationToken); + + private async Task>> SearchImplAsync(Vector queryVector, int topK, object? filter, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); @@ -252,7 +264,6 @@ private async Task>> GetSimilarCoreImplAsync(Vector q ["includeValues"] = false }; - var filter = BuildFilter(metadataFilters); if (filter != null) body["filter"] = filter; @@ -327,6 +338,96 @@ private async Task>> GetSimilarCoreImplAsync(Vector q return filter; } + // ------------------------------------------------------------------ + // MetadataFilter AST -> Pinecone filter translation. Pinecone has no + // "$not" operator, so negation is pushed down (De Morgan for logical + // nodes; operator inversion / $nin / $exists:false for leaves). + // ------------------------------------------------------------------ + + /// Translates a expression tree into a Pinecone filter object. + internal static object TranslateFilter(MetadataFilter filter) + { + if (filter == null) + throw new ArgumentNullException(nameof(filter)); + + switch (filter) + { + case ComparisonFilter comparison: + return Leaf(comparison.Key, ComparisonOp(comparison.Operator), comparison.Value); + case InFilter inFilter: + return Leaf(inFilter.Key, "$in", inFilter.Values.ToArray()); + case ExistsFilter existsFilter: + return Leaf(existsFilter.Key, "$exists", true); + case NotFilter notFilter: + return TranslateNegated(notFilter.Operand); + case LogicalFilter logical when logical.Operator == MetadataFilterOperator.And: + return new Dictionary { ["$and"] = logical.Operands.Select(TranslateFilter).ToArray() }; + case LogicalFilter logical: + return new Dictionary { ["$or"] = logical.Operands.Select(TranslateFilter).ToArray() }; + default: + throw new NotSupportedException($"Unsupported metadata filter node: {filter.GetType().Name}"); + } + } + + private static object TranslateNegated(MetadataFilter filter) + { + switch (filter) + { + case ComparisonFilter comparison: + return Leaf(comparison.Key, NegatedComparisonOp(comparison.Operator), comparison.Value); + case InFilter inFilter: + return Leaf(inFilter.Key, "$nin", inFilter.Values.ToArray()); + case ExistsFilter existsFilter: + return Leaf(existsFilter.Key, "$exists", false); + case NotFilter notFilter: + return TranslateFilter(notFilter.Operand); + case LogicalFilter logical when logical.Operator == MetadataFilterOperator.And: + // NOT(a AND b) => (NOT a) OR (NOT b) + return new Dictionary { ["$or"] = logical.Operands.Select(TranslateNegated).ToArray() }; + case LogicalFilter logical: + // NOT(a OR b) => (NOT a) AND (NOT b) + return new Dictionary { ["$and"] = logical.Operands.Select(TranslateNegated).ToArray() }; + default: + throw new NotSupportedException($"Unsupported metadata filter node: {filter.GetType().Name}"); + } + } + + private static object Leaf(string key, string op, object value) + { + object bound = op == "$gt" || op == "$gte" || op == "$lt" || op == "$lte" + ? (MetadataFilter.IsNumeric(value) ? Convert.ToDouble(value, CultureInfo.InvariantCulture) : value) + : value; + return new Dictionary { [key] = new Dictionary { [op] = bound } }; + } + + private static string ComparisonOp(MetadataFilterOperator op) + { + switch (op) + { + case MetadataFilterOperator.Eq: return "$eq"; + case MetadataFilterOperator.Ne: return "$ne"; + case MetadataFilterOperator.Gt: return "$gt"; + case MetadataFilterOperator.Gte: return "$gte"; + case MetadataFilterOperator.Lt: return "$lt"; + case MetadataFilterOperator.Lte: return "$lte"; + default: throw new NotSupportedException($"Unsupported comparison operator: {op}"); + } + } + + private static string NegatedComparisonOp(MetadataFilterOperator op) + { + switch (op) + { + case MetadataFilterOperator.Eq: return "$ne"; + case MetadataFilterOperator.Ne: return "$eq"; + case MetadataFilterOperator.Gt: return "$lte"; + case MetadataFilterOperator.Gte: return "$lt"; + case MetadataFilterOperator.Lt: return "$gte"; + case MetadataFilterOperator.Lte: return "$gt"; + default: throw new NotSupportedException($"Unsupported comparison operator: {op}"); + } + } + /// protected override Document? GetByIdCore(string documentId) => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs index 029f178877..cebb0dfbe0 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/QdrantDocumentStore.cs @@ -11,6 +11,7 @@ using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; using Newtonsoft.Json; using Newtonsoft.Json.Linq; @@ -257,7 +258,18 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); - private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + private Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, BuildFilter(metadataFilters), cancellationToken); + + /// + protected override IEnumerable> GetSimilarWithFilterCore(Vector queryVector, MetadataFilter filter, int topK) + => SearchImplAsync(queryVector, topK, TranslateFilter(filter), CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarWithFilterCoreAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, TranslateFilter(filter), cancellationToken); + + private async Task>> SearchImplAsync(Vector queryVector, int topK, object? filter, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); @@ -269,7 +281,6 @@ private async Task>> GetSimilarCoreImplAsync(Vector q ["with_vector"] = false }; - var filter = BuildFilter(metadataFilters); if (filter != null) body["filter"] = filter; @@ -296,6 +307,125 @@ private async Task>> GetSimilarCoreImplAsync(Vector q return results; } + // ------------------------------------------------------------------ + // MetadataFilter AST -> Qdrant filter translation. + // Leaf conditions map to field conditions (match / range); logical + // nodes map to nested bool filters (must / should / must_not). + // ------------------------------------------------------------------ + + /// + /// Translates a expression tree into a Qdrant filter object. + /// + internal static object TranslateFilter(MetadataFilter filter) + { + if (filter == null) + throw new ArgumentNullException(nameof(filter)); + + // A bare leaf condition must be wrapped in a "must" so the top-level value is a valid filter. + if (filter is ComparisonFilter c && c.Operator == MetadataFilterOperator.Eq + || filter is ComparisonFilter cg && (cg.Operator == MetadataFilterOperator.Gt || cg.Operator == MetadataFilterOperator.Gte + || cg.Operator == MetadataFilterOperator.Lt || cg.Operator == MetadataFilterOperator.Lte) + || filter is InFilter) + { + return new Dictionary { ["must"] = new[] { TranslateClause(filter) } }; + } + + return TranslateClause(filter); + } + + private static object TranslateClause(MetadataFilter filter) + { + switch (filter) + { + case ComparisonFilter comparison: + return TranslateComparison(comparison); + case InFilter inFilter: + return new Dictionary + { + ["key"] = FieldKey(inFilter.Key), + ["match"] = new Dictionary { ["any"] = inFilter.Values.ToArray() } + }; + case ExistsFilter existsFilter: + return new Dictionary + { + ["must_not"] = new object[] + { + new Dictionary + { + ["is_empty"] = new Dictionary { ["key"] = FieldKey(existsFilter.Key) } + } + } + }; + case NotFilter notFilter: + return new Dictionary + { + ["must_not"] = new[] { TranslateClause(notFilter.Operand) } + }; + case LogicalFilter logical when logical.Operator == MetadataFilterOperator.And: + return new Dictionary + { + ["must"] = logical.Operands.Select(TranslateClause).ToArray() + }; + case LogicalFilter logical: + return new Dictionary + { + ["should"] = logical.Operands.Select(TranslateClause).ToArray() + }; + default: + throw new NotSupportedException($"Unsupported metadata filter node: {filter.GetType().Name}"); + } + } + + private static object TranslateComparison(ComparisonFilter comparison) + { + var key = FieldKey(comparison.Key); + switch (comparison.Operator) + { + case MetadataFilterOperator.Eq: + return new Dictionary + { + ["key"] = key, + ["match"] = new Dictionary { ["value"] = comparison.Value } + }; + case MetadataFilterOperator.Ne: + return new Dictionary + { + ["must_not"] = new object[] + { + new Dictionary + { + ["key"] = key, + ["match"] = new Dictionary { ["value"] = comparison.Value } + } + } + }; + case MetadataFilterOperator.Gt: + return RangeCondition(key, "gt", comparison.Value); + case MetadataFilterOperator.Gte: + return RangeCondition(key, "gte", comparison.Value); + case MetadataFilterOperator.Lt: + return RangeCondition(key, "lt", comparison.Value); + case MetadataFilterOperator.Lte: + return RangeCondition(key, "lte", comparison.Value); + default: + throw new NotSupportedException($"Unsupported comparison operator: {comparison.Operator}"); + } + } + + private static object RangeCondition(string key, string op, object value) + { + object bound = MetadataFilter.IsNumeric(value) + ? Convert.ToDouble(value, CultureInfo.InvariantCulture) + : value; + return new Dictionary + { + ["key"] = key, + ["range"] = new Dictionary { [op] = bound } + }; + } + + private static string FieldKey(string key) => "metadata." + key; + /// /// Translates the metadata filter dictionary into a Qdrant filter object. /// diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs index 4278141a32..4f811fc1ae 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/WeaviateDocumentStore.cs @@ -13,6 +13,7 @@ using AiDotNet.Attributes; using AiDotNet.Enums; using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using Newtonsoft.Json; using Newtonsoft.Json.Linq; @@ -258,12 +259,22 @@ protected override IEnumerable> GetSimilarCore(Vector queryVector protected override Task>> GetSimilarCoreAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) => GetSimilarCoreImplAsync(queryVector, topK, metadataFilters, cancellationToken); - private async Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + private Task>> GetSimilarCoreImplAsync(Vector queryVector, int topK, Dictionary metadataFilters, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, BuildWhere(metadataFilters), cancellationToken); + + /// + protected override IEnumerable> GetSimilarWithFilterCore(Vector queryVector, MetadataFilter filter, int topK) + => SearchImplAsync(queryVector, topK, TranslateWhere(filter), CancellationToken.None).GetAwaiter().GetResult(); + + /// + protected override Task>> GetSimilarWithFilterCoreAsync(Vector queryVector, MetadataFilter filter, int topK, CancellationToken cancellationToken) + => SearchImplAsync(queryVector, topK, TranslateWhere(filter), cancellationToken); + + private async Task>> SearchImplAsync(Vector queryVector, int topK, string? whereClause, CancellationToken cancellationToken) { var vector = queryVector.ToArray().Select(v => Convert.ToDouble(v)).ToArray(); var vectorJson = "[" + string.Join(",", vector.Select(v => v.ToString("R", CultureInfo.InvariantCulture))) + "]"; - var whereClause = BuildWhere(metadataFilters); var whereArg = whereClause != null ? ", where: " + whereClause : string.Empty; var query = @@ -370,6 +381,128 @@ private async Task>> GetSimilarCoreImplAsync(Vector q return "{operator: And, operands: [" + string.Join(", ", operands) + "]}"; } + // ------------------------------------------------------------------ + // MetadataFilter AST -> Weaviate GraphQL "where" translation. + // Weaviate has no Not operator, so negation is pushed down (NotEqual / + // inverted comparisons / IsNull:true; De Morgan for logical nodes). + // Filterable properties are stored flattened under the "m_" prefix. + // ------------------------------------------------------------------ + + /// Translates a expression tree into a Weaviate GraphQL where clause. + internal static string TranslateWhere(MetadataFilter filter) + { + if (filter == null) + throw new ArgumentNullException(nameof(filter)); + return BuildClause(filter); + } + + private static string BuildClause(MetadataFilter filter) + { + switch (filter) + { + case ComparisonFilter comparison: + return ComparisonClause(comparison.Key, comparison.Operator, comparison.Value); + case InFilter inFilter: + return ContainsAnyClause(inFilter.Key, inFilter.Values); + case ExistsFilter existsFilter: + return "{path: [\"" + MetaPrefix + existsFilter.Key + "\"], operator: IsNull, valueBoolean: false}"; + case NotFilter notFilter: + return NegateClause(notFilter.Operand); + case LogicalFilter logical when logical.Operator == MetadataFilterOperator.And: + return "{operator: And, operands: [" + string.Join(", ", logical.Operands.Select(BuildClause)) + "]}"; + case LogicalFilter logical: + return "{operator: Or, operands: [" + string.Join(", ", logical.Operands.Select(BuildClause)) + "]}"; + default: + throw new NotSupportedException($"Unsupported metadata filter node: {filter.GetType().Name}"); + } + } + + private static string NegateClause(MetadataFilter filter) + { + switch (filter) + { + case ComparisonFilter comparison: + return ComparisonClause(comparison.Key, NegateOperator(comparison.Operator), comparison.Value); + case InFilter inFilter: + // NOT(x in [a,b]) => (x != a) AND (x != b) + var notEquals = inFilter.Values.Select(v => ComparisonClause(inFilter.Key, MetadataFilterOperator.Ne, v)).ToList(); + if (notEquals.Count == 1) + return notEquals[0]; + return "{operator: And, operands: [" + string.Join(", ", notEquals) + "]}"; + case ExistsFilter existsFilter: + return "{path: [\"" + MetaPrefix + existsFilter.Key + "\"], operator: IsNull, valueBoolean: true}"; + case NotFilter notFilter: + return BuildClause(notFilter.Operand); + case LogicalFilter logical when logical.Operator == MetadataFilterOperator.And: + return "{operator: Or, operands: [" + string.Join(", ", logical.Operands.Select(NegateClause)) + "]}"; + case LogicalFilter logical: + return "{operator: And, operands: [" + string.Join(", ", logical.Operands.Select(NegateClause)) + "]}"; + default: + throw new NotSupportedException($"Unsupported metadata filter node: {filter.GetType().Name}"); + } + } + + private static string ComparisonClause(string key, MetadataFilterOperator op, object value) + { + var path = "[\"" + MetaPrefix + key + "\"]"; + var opName = WeaviateOperator(op); + var valueArg = ValueArg(value); + return "{path: " + path + ", operator: " + opName + ", " + valueArg + "}"; + } + + private static string ContainsAnyClause(string key, IReadOnlyList values) + { + var path = "[\"" + MetaPrefix + key + "\"]"; + var allNumeric = values.Count > 0 && values.All(MetadataFilter.IsNumeric); + if (allNumeric) + { + var nums = values.Select(v => Convert.ToDouble(v, CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture)); + return "{path: " + path + ", operator: ContainsAny, valueNumber: [" + string.Join(", ", nums) + "]}"; + } + + var texts = values.Select(v => JsonConvert.ToString(v?.ToString() ?? string.Empty)); + return "{path: " + path + ", operator: ContainsAny, valueText: [" + string.Join(", ", texts) + "]}"; + } + + private static string ValueArg(object value) + { + if (value == null) + return "valueText: " + JsonConvert.ToString(string.Empty); + if (value is bool b) + return "valueBoolean: " + (b ? "true" : "false"); + if (MetadataFilter.IsNumeric(value)) + return "valueNumber: " + Convert.ToDouble(value, CultureInfo.InvariantCulture).ToString("R", CultureInfo.InvariantCulture); + return "valueText: " + JsonConvert.ToString(value.ToString() ?? string.Empty); + } + + private static string WeaviateOperator(MetadataFilterOperator op) + { + switch (op) + { + case MetadataFilterOperator.Eq: return "Equal"; + case MetadataFilterOperator.Ne: return "NotEqual"; + case MetadataFilterOperator.Gt: return "GreaterThan"; + case MetadataFilterOperator.Gte: return "GreaterThanEqual"; + case MetadataFilterOperator.Lt: return "LessThan"; + case MetadataFilterOperator.Lte: return "LessThanEqual"; + default: throw new NotSupportedException($"Unsupported comparison operator: {op}"); + } + } + + private static MetadataFilterOperator NegateOperator(MetadataFilterOperator op) + { + switch (op) + { + case MetadataFilterOperator.Eq: return MetadataFilterOperator.Ne; + case MetadataFilterOperator.Ne: return MetadataFilterOperator.Eq; + case MetadataFilterOperator.Gt: return MetadataFilterOperator.Lte; + case MetadataFilterOperator.Gte: return MetadataFilterOperator.Lt; + case MetadataFilterOperator.Lt: return MetadataFilterOperator.Gte; + case MetadataFilterOperator.Lte: return MetadataFilterOperator.Gt; + default: throw new NotSupportedException($"Unsupported comparison operator: {op}"); + } + } + /// protected override Document? GetByIdCore(string documentId) => GetByIdCoreImplAsync(documentId, CancellationToken.None).GetAwaiter().GetResult(); diff --git a/src/RetrievalAugmentedGeneration/Filtering/MetadataFilter.cs b/src/RetrievalAugmentedGeneration/Filtering/MetadataFilter.cs new file mode 100644 index 0000000000..b58a5d5765 --- /dev/null +++ b/src/RetrievalAugmentedGeneration/Filtering/MetadataFilter.cs @@ -0,0 +1,480 @@ +using System; +using System.Collections; +using System.Collections.Generic; +using System.Globalization; +using System.Linq; + +namespace AiDotNet.RetrievalAugmentedGeneration.Filtering; + +/// +/// The kind of node in a expression tree. +/// +/// +/// For Beginners: Every part of a filter is one of these operations. Leaf operations +/// (Eq, Gt, In, ...) test a single metadata field; the combinators +/// (And, Or, Not) glue smaller filters into bigger ones - exactly like the +/// filter languages of Pinecone, Qdrant and Weaviate. +/// +/// +public enum MetadataFilterOperator +{ + /// Field equals a value. + Eq, + + /// Field does not equal a value (also matches when the field is absent). + Ne, + + /// Field is strictly greater than a value. + Gt, + + /// Field is greater than or equal to a value. + Gte, + + /// Field is strictly less than a value. + Lt, + + /// Field is less than or equal to a value. + Lte, + + /// Field equals one of a set of values. + In, + + /// Field is present in the metadata. + Exists, + + /// Logical conjunction of child filters. + And, + + /// Logical disjunction of child filters. + Or, + + /// Logical negation of a child filter. + Not, +} + +/// +/// An immutable, provider-agnostic boolean metadata-filter expression tree. Supports equality, +/// inequality, full ranges, set membership, existence and arbitrary AND/OR/NOT +/// nesting, mirroring the filter capabilities of Pinecone, Qdrant, Weaviate, Milvus, Azure AI Search, +/// pgvector and Elasticsearch. +/// +/// +/// +/// Build filters with the static factory methods (, +/// , , ...) and +/// compose them with the fluent combinators (, +/// , ). The resulting tree can be evaluated +/// in-memory with or translated to a +/// backend's native filter language by a document store. +/// +/// For Beginners: This is a small, database-independent way to describe "which documents +/// do I want". For example: +/// +/// // category == "science" AND year >= 2020 AND (author in ["A","B"]) +/// var filter = MetadataFilter.Eq("category", "science") +/// .And(MetadataFilter.Gte("year", 2020)) +/// .And(MetadataFilter.In("author", new object[] { "A", "B" })); +/// +/// The same filter object works against every store: an in-memory store evaluates it directly, +/// while Pinecone/Qdrant/etc. translate it into their own query syntax. +/// +/// +public abstract class MetadataFilter +{ + /// Gets the operator that identifies this node's kind. + public abstract MetadataFilterOperator Operator { get; } + + /// + /// Evaluates this filter against a metadata dictionary in memory. + /// + /// The document metadata to test; may be null (treated as empty). + /// true when the metadata satisfies the filter; otherwise false. + public abstract bool Matches(IReadOnlyDictionary metadata); + + // ------------------------------------------------------------------ + // Leaf factories + // ------------------------------------------------------------------ + + /// Creates a filter that matches when equals . + public static MetadataFilter Eq(string key, object value) => new ComparisonFilter(MetadataFilterOperator.Eq, key, value); + + /// Creates a filter that matches when is missing or does not equal . + public static MetadataFilter Ne(string key, object value) => new ComparisonFilter(MetadataFilterOperator.Ne, key, value); + + /// Creates a filter that matches when is strictly greater than . + public static MetadataFilter Gt(string key, object value) => new ComparisonFilter(MetadataFilterOperator.Gt, key, value); + + /// Creates a filter that matches when is greater than or equal to . + public static MetadataFilter Gte(string key, object value) => new ComparisonFilter(MetadataFilterOperator.Gte, key, value); + + /// Creates a filter that matches when is strictly less than . + public static MetadataFilter Lt(string key, object value) => new ComparisonFilter(MetadataFilterOperator.Lt, key, value); + + /// Creates a filter that matches when is less than or equal to . + public static MetadataFilter Lte(string key, object value) => new ComparisonFilter(MetadataFilterOperator.Lte, key, value); + + /// + /// Creates an inclusive range filter (lower <= key <= upper), the equivalent of + /// Gte(key, lower).And(Lte(key, upper)). + /// + public static MetadataFilter Range(string key, object lower, object upper) + => And(Gte(key, lower), Lte(key, upper)); + + /// Creates a filter that matches when equals one of . + public static MetadataFilter In(string key, IEnumerable values) => new InFilter(key, values); + + /// Creates a filter that matches when equals one of . + public static MetadataFilter In(string key, params object[] values) => new InFilter(key, values); + + /// Creates a filter that matches when is present in the metadata. + public static MetadataFilter Exists(string key) => new ExistsFilter(key); + + // ------------------------------------------------------------------ + // Combinator factories + // ------------------------------------------------------------------ + + /// Combines two or more filters with logical AND. + public static MetadataFilter And(params MetadataFilter[] filters) => new LogicalFilter(MetadataFilterOperator.And, filters); + + /// Combines two or more filters with logical AND. + public static MetadataFilter And(IEnumerable filters) => new LogicalFilter(MetadataFilterOperator.And, filters); + + /// Combines two or more filters with logical OR. + public static MetadataFilter Or(params MetadataFilter[] filters) => new LogicalFilter(MetadataFilterOperator.Or, filters); + + /// Combines two or more filters with logical OR. + public static MetadataFilter Or(IEnumerable filters) => new LogicalFilter(MetadataFilterOperator.Or, filters); + + // ------------------------------------------------------------------ + // Fluent combinators + // ------------------------------------------------------------------ + + /// Returns a new filter that is the logical AND of this filter and . + public MetadataFilter And(MetadataFilter other) => new LogicalFilter(MetadataFilterOperator.And, new[] { this, other }); + + /// Returns a new filter that is the logical OR of this filter and . + public MetadataFilter Or(MetadataFilter other) => new LogicalFilter(MetadataFilterOperator.Or, new[] { this, other }); + + /// Returns a new filter that is the logical negation of this filter. + public MetadataFilter Not() => new NotFilter(this); + + // ------------------------------------------------------------------ + // Shared evaluation helpers (used by the leaf nodes and reusable by + // in-memory store implementations). + // ------------------------------------------------------------------ + + /// + /// Compares two metadata values for equality using lenient, provider-like semantics: numbers are + /// compared numerically (so 2020 and 2020L and "2020" match), booleans by + /// value, and everything else by ordinal string form. + /// + public static bool ValuesEqual(object? a, object? b) + { + if (a == null || b == null) + return a == null && b == null; + + if (a is bool || b is bool) + return BoolText(a) == BoolText(b); + + if (TryToDouble(a, out var da) && TryToDouble(b, out var db)) + return da.Equals(db); + + if (a.Equals(b)) + return true; + + return string.Equals(ToInvariantString(a), ToInvariantString(b), StringComparison.Ordinal); + } + + /// + /// Orders two metadata values, returning a negative/zero/positive result like + /// , or null when they are not comparable. + /// Numbers are compared numerically; otherwise strings/other comparables are compared naturally. + /// + public static int? CompareValues(object? left, object? right) + { + if (left == null || right == null) + return null; + + if (TryToDouble(left, out var dl) && TryToDouble(right, out var dr)) + return dl.CompareTo(dr); + + // Only order two strings against each other. A string vs a number (that didn't parse numerically + // above) is NOT orderable — fall through to null so e.g. Gt("name", 5) never "matches" "abc". + if (left is string ls && right is string rs) + return string.CompareOrdinal(ls, rs); + + if (left.GetType() == right.GetType() && left is IComparable comparable) + return comparable.CompareTo(right); + + return null; + } + + private static string BoolText(object? value) + { + if (value is bool b) + return b ? "true" : "false"; + return ToInvariantString(value).ToLowerInvariant(); + } + + /// + /// Attempts to interpret a metadata value as a (numeric types and numeric + /// strings). Exposed so provider filter builders can share the same numeric-coercion rules. + /// + public static bool TryToDouble(object? value, out double result) + { + result = 0; + switch (value) + { + case sbyte v: result = v; return true; + case byte v: result = v; return true; + case short v: result = v; return true; + case ushort v: result = v; return true; + case int v: result = v; return true; + case uint v: result = v; return true; + case long v: result = v; return true; + case ulong v: result = v; return true; + case float v: result = v; return true; + case double v: result = v; return true; + case decimal v: result = (double)v; return true; + case string s: + return double.TryParse(s, NumberStyles.Any, CultureInfo.InvariantCulture, out result); + default: + return false; + } + } + + /// + /// Returns true when is a built-in numeric type. Exposed so provider + /// filter builders can share the same numeric detection rules. + /// + public static bool IsNumeric(object? value) + { + return value is sbyte || value is byte || value is short || value is ushort + || value is int || value is uint || value is long || value is ulong + || value is float || value is double || value is decimal; + } + + internal static string ToInvariantString(object? value) + { + if (value == null) + return string.Empty; + if (value is IFormattable formattable) + return formattable.ToString(null, CultureInfo.InvariantCulture); + return value.ToString() ?? string.Empty; + } +} + +/// +/// A leaf node comparing a single metadata field against a scalar value with one of the operators +/// , , +/// , , +/// or . +/// +public sealed class ComparisonFilter : MetadataFilter +{ + /// Initializes a new comparison node. + /// The comparison operator. + /// The metadata field name. + /// The value to compare against. + public ComparisonFilter(MetadataFilterOperator op, string key, object value) + { + if (string.IsNullOrEmpty(key)) + throw new ArgumentException("Filter key cannot be null or empty", nameof(key)); + if (op != MetadataFilterOperator.Eq && op != MetadataFilterOperator.Ne + && op != MetadataFilterOperator.Gt && op != MetadataFilterOperator.Gte + && op != MetadataFilterOperator.Lt && op != MetadataFilterOperator.Lte) + { + throw new ArgumentException($"Operator '{op}' is not a comparison operator", nameof(op)); + } + + Operator = op; + Key = key; + Value = value; + } + + /// + public override MetadataFilterOperator Operator { get; } + + /// Gets the metadata field name being compared. + public string Key { get; } + + /// Gets the value being compared against. + public object Value { get; } + + /// + public override bool Matches(IReadOnlyDictionary metadata) + { + var present = metadata != null && metadata.TryGetValue(Key, out _); + object? actual = present ? metadata![Key] : null; + + switch (Operator) + { + case MetadataFilterOperator.Eq: + return present && ValuesEqual(actual, Value); + case MetadataFilterOperator.Ne: + return !(present && ValuesEqual(actual, Value)); + case MetadataFilterOperator.Gt: + return present && CompareValues(actual, Value) is int gt && gt > 0; + case MetadataFilterOperator.Gte: + return present && CompareValues(actual, Value) is int gte && gte >= 0; + case MetadataFilterOperator.Lt: + return present && CompareValues(actual, Value) is int lt && lt < 0; + case MetadataFilterOperator.Lte: + return present && CompareValues(actual, Value) is int lte && lte <= 0; + default: + return false; + } + } +} + +/// +/// A leaf node matching when a metadata field equals any value in a set (set membership). +/// +public sealed class InFilter : MetadataFilter +{ + /// Initializes a new set-membership node. + /// The metadata field name. + /// The candidate values. + public InFilter(string key, IEnumerable values) + { + if (string.IsNullOrEmpty(key)) + throw new ArgumentException("Filter key cannot be null or empty", nameof(key)); + if (values == null) + throw new ArgumentNullException(nameof(values)); + + Key = key; + Values = values.ToList().AsReadOnly(); + } + + /// + public override MetadataFilterOperator Operator => MetadataFilterOperator.In; + + /// Gets the metadata field name being tested. + public string Key { get; } + + /// Gets the candidate values. + public IReadOnlyList Values { get; } + + /// + public override bool Matches(IReadOnlyDictionary metadata) + { + if (metadata == null || !metadata.TryGetValue(Key, out var actual)) + return false; + + foreach (var candidate in Values) + { + if (ValuesEqual(actual, candidate)) + return true; + } + + return false; + } +} + +/// +/// A leaf node matching when a metadata field is present. +/// +public sealed class ExistsFilter : MetadataFilter +{ + /// Initializes a new existence node. + /// The metadata field name that must be present. + public ExistsFilter(string key) + { + if (string.IsNullOrEmpty(key)) + throw new ArgumentException("Filter key cannot be null or empty", nameof(key)); + + Key = key; + } + + /// + public override MetadataFilterOperator Operator => MetadataFilterOperator.Exists; + + /// Gets the metadata field name that must be present. + public string Key { get; } + + /// + public override bool Matches(IReadOnlyDictionary metadata) + => metadata != null && metadata.ContainsKey(Key); +} + +/// +/// A combinator node joining child filters with logical or +/// . +/// +public sealed class LogicalFilter : MetadataFilter +{ + /// Initializes a new logical combinator node. + /// Either or . + /// The child filters; must contain at least one element. + public LogicalFilter(MetadataFilterOperator op, IEnumerable operands) + { + if (op != MetadataFilterOperator.And && op != MetadataFilterOperator.Or) + throw new ArgumentException($"Operator '{op}' is not a logical combinator", nameof(op)); + if (operands == null) + throw new ArgumentNullException(nameof(operands)); + + var list = operands.ToList(); + if (list.Count == 0) + throw new ArgumentException("A logical filter requires at least one operand", nameof(operands)); + foreach (var operand in list) + { + if (operand == null) + throw new ArgumentException("A logical filter operand cannot be null", nameof(operands)); + } + + Operator = op; + Operands = list.AsReadOnly(); + } + + /// + public override MetadataFilterOperator Operator { get; } + + /// Gets the child filters. + public IReadOnlyList Operands { get; } + + /// + public override bool Matches(IReadOnlyDictionary metadata) + { + if (Operator == MetadataFilterOperator.And) + { + foreach (var operand in Operands) + { + if (!operand.Matches(metadata)) + return false; + } + + return true; + } + + foreach (var operand in Operands) + { + if (operand.Matches(metadata)) + return true; + } + + return false; + } +} + +/// +/// A combinator node negating a single child filter. +/// +public sealed class NotFilter : MetadataFilter +{ + /// Initializes a new negation node. + /// The filter to negate. + public NotFilter(MetadataFilter operand) + { + Operand = operand ?? throw new ArgumentNullException(nameof(operand)); + } + + /// + public override MetadataFilterOperator Operator => MetadataFilterOperator.Not; + + /// Gets the negated child filter. + public MetadataFilter Operand { get; } + + /// + public override bool Matches(IReadOnlyDictionary metadata) + => !Operand.Matches(metadata); +} diff --git a/tests/AiDotNet.Tests/AiDotNetTests.csproj b/tests/AiDotNet.Tests/AiDotNetTests.csproj index 823134d4ae..e79bd40a15 100644 --- a/tests/AiDotNet.Tests/AiDotNetTests.csproj +++ b/tests/AiDotNet.Tests/AiDotNetTests.csproj @@ -63,6 +63,7 @@ + diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/BM25RetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/BM25RetrieverTests.cs index 3293dd337c..8267d1a82b 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/BM25RetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/BM25RetrieverTests.cs @@ -27,6 +27,10 @@ private class MockDocumentStore : IDocumentStore public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + + public IEnumerable> GetSimilarWithFilter(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK) { var __r = new List>(); foreach (var __d in GetSimilar(queryVector, topK * 10)) { if (filter == null || filter.Matches(__d.Metadata)) { __r.Add(__d); if (__r.Count >= topK) break; } } return __r; } + + public System.Threading.Tasks.Task>> GetSimilarWithFilterAsync(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); } public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ColBERTRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ColBERTRetrieverTests.cs index bb99379361..229e2e3b65 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ColBERTRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ColBERTRetrieverTests.cs @@ -27,6 +27,10 @@ private class MockDocumentStore : IDocumentStore public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + + public IEnumerable> GetSimilarWithFilter(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK) { var __r = new List>(); foreach (var __d in GetSimilar(queryVector, topK * 10)) { if (filter == null || filter.Matches(__d.Metadata)) { __r.Add(__d); if (__r.Count >= topK) break; } } return __r; } + + public System.Threading.Tasks.Task>> GetSimilarWithFilterAsync(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); } public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DenseRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DenseRetrieverTests.cs index ec65de9fc3..66ab69062f 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DenseRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DenseRetrieverTests.cs @@ -76,6 +76,10 @@ private class MockDocumentStore : IDocumentStore public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + + public IEnumerable> GetSimilarWithFilter(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK) { var __r = new List>(); foreach (var __d in GetSimilar(queryVector, topK * 10)) { if (filter == null || filter.Matches(__d.Metadata)) { __r.Add(__d); if (__r.Count >= topK) break; } } return __r; } + + public System.Threading.Tasks.Task>> GetSimilarWithFilterAsync(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); } public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/InMemoryDocumentStoreTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/InMemoryDocumentStoreTests.cs index 660946ce15..d1aef62881 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/InMemoryDocumentStoreTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/InMemoryDocumentStoreTests.cs @@ -4,6 +4,7 @@ using System.Threading.Tasks; using AiDotNet.LinearAlgebra; using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; using AiDotNet.RetrievalAugmentedGeneration.Models; using AiDotNet.Tensors.LinearAlgebra; using Xunit; @@ -296,6 +297,93 @@ public async Task ConcurrentAdd_WithMultipleThreads_AllDocumentsAdded() #endregion + #region Rich MetadataFilter (end-to-end) + + private static VectorDocument RichDoc(string id, float[] vector, Dictionary metadata) + => new VectorDocument + { + Document = new Document(id, id + "-content", metadata), + Embedding = new Vector(vector) + }; + + private static InMemoryDocumentStore BuildRichStore() + { + var store = new InMemoryDocumentStore(vectorDimension: 3); + store.AddBatch(new List> + { + RichDoc("d1", new float[] { 1, 0, 0 }, new Dictionary { ["category"] = "science", ["year"] = 2022, ["archived"] = false }), + RichDoc("d2", new float[] { 0.9f, 0.1f, 0 }, new Dictionary { ["category"] = "science", ["year"] = 2015, ["author"] = "A", ["archived"] = false }), + RichDoc("d3", new float[] { 0.8f, 0.2f, 0 }, new Dictionary { ["category"] = "science", ["year"] = 2023, ["archived"] = true }), + RichDoc("d4", new float[] { 0.7f, 0.3f, 0 }, new Dictionary { ["category"] = "food", ["year"] = 2024, ["archived"] = false }), + RichDoc("d5", new float[] { 0.6f, 0.4f, 0 }, new Dictionary { ["category"] = "science", ["year"] = 2010, ["author"] = "Z", ["archived"] = false }), + }); + return store; + } + + [Fact] + public void GetSimilarWithFilter_And_Range_ReturnsOnlyMatching() + { + var store = BuildRichStore(); + // category == science AND year >= 2020 AND NOT archived + var filter = MetadataFilter.Eq("category", "science") + .And(MetadataFilter.Gte("year", 2020)) + .And(MetadataFilter.Eq("archived", true).Not()); + + var results = store.GetSimilarWithFilter(new Vector(new float[] { 1, 0, 0 }), filter, topK: 10).ToList(); + + var ids = results.Select(r => r.Id).OrderBy(x => x).ToList(); + Assert.Equal(new[] { "d1" }, ids); + } + + [Fact] + public void GetSimilarWithFilter_OrAndIn_ReturnsUnion() + { + var store = BuildRichStore(); + // category == science AND (year >= 2020 OR author in [A]) + var filter = MetadataFilter.Eq("category", "science") + .And(MetadataFilter.Gte("year", 2020).Or(MetadataFilter.In("author", new object[] { "A" }))); + + var results = store.GetSimilarWithFilter(new Vector(new float[] { 1, 0, 0 }), filter, topK: 10).ToList(); + + var ids = results.Select(r => r.Id).OrderBy(x => x).ToList(); + // d1 (2022), d2 (author A), d3 (2023) - all science; d4 is food, d5 is 2010/author Z. + Assert.Equal(new[] { "d1", "d2", "d3" }, ids); + } + + [Fact] + public void GetSimilarWithFilter_Exists_And_TopKHonoured() + { + var store = BuildRichStore(); + var filter = MetadataFilter.Exists("author"); + + var results = store.GetSimilarWithFilter(new Vector(new float[] { 1, 0, 0 }), filter, topK: 1).ToList(); + + Assert.Single(results); + Assert.Contains(results[0].Id, new[] { "d2", "d5" }); + Assert.True(results[0].HasRelevanceScore); + } + + [Fact(Timeout = 60000)] + public async Task GetSimilarWithFilterAsync_MatchesSyncResult() + { + var store = BuildRichStore(); + var filter = MetadataFilter.Ne("category", "science"); + + var results = (await store.GetSimilarWithFilterAsync(new Vector(new float[] { 1, 0, 0 }), filter, topK: 10)).ToList(); + + Assert.Equal(new[] { "d4" }, results.Select(r => r.Id).ToList()); + } + + [Fact] + public void GetSimilarWithFilter_NullFilter_BehavesLikeGetSimilar() + { + var store = BuildRichStore(); + var results = store.GetSimilarWithFilter(new Vector(new float[] { 1, 0, 0 }), null!, topK: 3).ToList(); + Assert.Equal(3, results.Count); + } + + #endregion + #region Helper Methods private VectorDocument CreateTestDocument( diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/MetadataFilterTranslationTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/MetadataFilterTranslationTests.cs new file mode 100644 index 0000000000..c47828306a --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/DocumentStores/MetadataFilterTranslationTests.cs @@ -0,0 +1,316 @@ +using System.Collections.Generic; + +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Filtering; + +using Newtonsoft.Json; +using Newtonsoft.Json.Linq; + +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.DocumentStores +{ + /// + /// Unit tests that assert each store translates a expression tree into + /// the correct provider-native filter (JSON / GraphQL / boolean-expr / OData / SQL). These exercise + /// the translation logic directly with no live services. + /// + public class MetadataFilterTranslationTests + { + private static readonly MetadataFilter Sample = + MetadataFilter.Eq("category", "science").And(MetadataFilter.Gte("year", 2020)); + + // ------------------------------------------------------------------ + // Qdrant: must / should / must_not + match / range (nested filters). + // ------------------------------------------------------------------ + + [Fact] + public void Qdrant_Eq_WrapsInMustMatch() + { + var json = JObject.Parse(JsonConvert.SerializeObject( + QdrantDocumentStore.TranslateFilter(MetadataFilter.Eq("category", "science")))); + var cond = (JArray)json["must"]!; + Assert.Equal("metadata.category", (string?)cond[0]!["key"]); + Assert.Equal("science", (string?)cond[0]!["match"]!["value"]); + } + + [Fact] + public void Qdrant_AndOfEqAndRange() + { + var json = JObject.Parse(JsonConvert.SerializeObject(QdrantDocumentStore.TranslateFilter(Sample))); + var must = (JArray)json["must"]!; + Assert.Equal("metadata.category", (string?)must[0]!["key"]); + Assert.Equal(2020.0, (double)must[1]!["range"]!["gte"]!); + } + + [Fact] + public void Qdrant_Not_MapsToMustNot() + { + var json = JObject.Parse(JsonConvert.SerializeObject( + QdrantDocumentStore.TranslateFilter(MetadataFilter.Eq("archived", true).Not()))); + Assert.NotNull(json["must_not"]); + } + + [Fact] + public void Qdrant_Or_MapsToShould_And_In_UsesMatchAny_And_Exists_UsesIsEmpty() + { + var or = JObject.Parse(JsonConvert.SerializeObject(QdrantDocumentStore.TranslateFilter( + MetadataFilter.Eq("a", "x").Or(MetadataFilter.Eq("a", "y"))))); + Assert.NotNull(or["should"]); + + var inJson = JObject.Parse(JsonConvert.SerializeObject(QdrantDocumentStore.TranslateFilter( + MetadataFilter.In("author", new object[] { "A", "B" })))); + var anyArr = (JArray)inJson["must"]![0]!["match"]!["any"]!; + Assert.Equal(2, anyArr.Count); + + var exists = JObject.Parse(JsonConvert.SerializeObject(QdrantDocumentStore.TranslateFilter( + MetadataFilter.Exists("author")))); + Assert.Equal("metadata.author", (string?)((JArray)exists["must_not"]!)[0]!["is_empty"]!["key"]); + } + + // ------------------------------------------------------------------ + // Pinecone: $and/$or/$eq/$ne/$in/$gte/$lte; $not is pushed down. + // ------------------------------------------------------------------ + + [Fact] + public void Pinecone_AndOfEqAndGte() + { + var json = JObject.Parse(JsonConvert.SerializeObject(PineconeDocumentStore.TranslateFilter(Sample))); + var and = (JArray)json["$and"]!; + Assert.Equal("science", (string?)and[0]!["category"]!["$eq"]); + Assert.Equal(2020.0, (double)and[1]!["year"]!["$gte"]!); + } + + [Fact] + public void Pinecone_NotEq_PushesToNe() + { + var json = JObject.Parse(JsonConvert.SerializeObject( + PineconeDocumentStore.TranslateFilter(MetadataFilter.Eq("category", "science").Not()))); + Assert.Equal("science", (string?)json["category"]!["$ne"]); + } + + [Fact] + public void Pinecone_NotIn_PushesToNin_And_NotAnd_UsesOr() + { + var nin = JObject.Parse(JsonConvert.SerializeObject( + PineconeDocumentStore.TranslateFilter(MetadataFilter.In("author", new object[] { "A", "B" }).Not()))); + Assert.Equal(2, ((JArray)nin["author"]!["$nin"]!).Count); + + var demorgan = JObject.Parse(JsonConvert.SerializeObject( + PineconeDocumentStore.TranslateFilter(Sample.Not()))); + Assert.NotNull(demorgan["$or"]); + } + + [Fact] + public void Pinecone_In_And_Exists() + { + var inJson = JObject.Parse(JsonConvert.SerializeObject( + PineconeDocumentStore.TranslateFilter(MetadataFilter.In("author", new object[] { "A", "B" })))); + Assert.Equal(2, ((JArray)inJson["author"]!["$in"]!).Count); + + var exists = JObject.Parse(JsonConvert.SerializeObject( + PineconeDocumentStore.TranslateFilter(MetadataFilter.Exists("author")))); + Assert.True((bool)exists["author"]!["$exists"]!); + } + + // ------------------------------------------------------------------ + // Weaviate GraphQL where (negation pushed down; ContainsAny for In). + // ------------------------------------------------------------------ + + [Fact] + public void Weaviate_And_Equal_And_GreaterThanEqual() + { + var where = WeaviateDocumentStore.TranslateWhere(Sample); + Assert.Contains("operator: And", where); + Assert.Contains("path: [\"m_category\"]", where); + Assert.Contains("operator: Equal", where); + Assert.Contains("valueText: \"science\"", where); + Assert.Contains("operator: GreaterThanEqual", where); + Assert.Contains("valueNumber: 2020", where); + } + + [Fact] + public void Weaviate_NotEq_UsesNotEqual_And_In_UsesContainsAny_And_Exists_UsesIsNull() + { + Assert.Contains("operator: NotEqual", + WeaviateDocumentStore.TranslateWhere(MetadataFilter.Eq("category", "science").Not())); + + var inWhere = WeaviateDocumentStore.TranslateWhere(MetadataFilter.In("author", new object[] { "A", "B" })); + Assert.Contains("operator: ContainsAny", inWhere); + Assert.Contains("valueText: [\"A\", \"B\"]", inWhere); + + var exists = WeaviateDocumentStore.TranslateWhere(MetadataFilter.Exists("author")); + Assert.Contains("operator: IsNull", exists); + Assert.Contains("valueBoolean: false", exists); + } + + // ------------------------------------------------------------------ + // Milvus boolean expression. + // ------------------------------------------------------------------ + + [Fact] + public void Milvus_And_Comparisons() + { + var expr = MilvusDocumentStore.TranslateFilter(Sample); + Assert.Equal("(m_category == \"science\" and m_year >= 2020)", expr); + } + + [Fact] + public void Milvus_Not_In_Exists_Or() + { + Assert.Equal("not (m_archived == true)", + MilvusDocumentStore.TranslateFilter(MetadataFilter.Eq("archived", true).Not())); + Assert.Equal("m_year in [2020, 2021]", + MilvusDocumentStore.TranslateFilter(MetadataFilter.In("year", new object[] { 2020, 2021 }))); + Assert.Equal("exists m_author", + MilvusDocumentStore.TranslateFilter(MetadataFilter.Exists("author"))); + Assert.Equal("(m_a == \"x\" or m_a == \"y\")", + MilvusDocumentStore.TranslateFilter(MetadataFilter.Eq("a", "x").Or(MetadataFilter.Eq("a", "y")))); + } + + // ------------------------------------------------------------------ + // Azure AI Search OData $filter. + // ------------------------------------------------------------------ + + [Fact] + public void Azure_And_EqAndGe() + { + Assert.Equal("(category eq 'science' and year ge 2020)", + AzureSearchDocumentStore.TranslateFilter(Sample)); + } + + [Fact] + public void Azure_Ne_Not_In_Exists() + { + Assert.Equal("category ne 'science'", + AzureSearchDocumentStore.TranslateFilter(MetadataFilter.Ne("category", "science"))); + Assert.Equal("not (category eq 'science')", + AzureSearchDocumentStore.TranslateFilter(MetadataFilter.Eq("category", "science").Not())); + Assert.Equal("search.in(author, 'A|B', '|')", + AzureSearchDocumentStore.TranslateFilter(MetadataFilter.In("author", new object[] { "A", "B" }))); + Assert.Equal("author ne null", + AzureSearchDocumentStore.TranslateFilter(MetadataFilter.Exists("author"))); + } + +#if NET5_0_OR_GREATER + // ------------------------------------------------------------------ + // pgvector: parameterised jsonb WHERE (net10-only Storage.Postgres). + // ------------------------------------------------------------------ + + [Fact] + public void Postgres_And_EqAndGte_Parameterised() + { + var parameters = new Dictionary(); + var sql = PostgresVectorFilterBuilder.Build(Sample, parameters); + + Assert.Equal(" WHERE ((metadata ->> @k0) = @v0 AND (metadata ->> @k1)::numeric >= @v1)", sql); + Assert.Equal("category", parameters["k0"]); + Assert.Equal("science", parameters["v0"]); + Assert.Equal("year", parameters["k1"]); + Assert.Equal(2020.0, parameters["v1"]); + } + + [Fact] + public void Postgres_Not_In_Exists() + { + var p1 = new Dictionary(); + Assert.Equal(" WHERE (NOT ((metadata ->> @k0) = @v0))", + PostgresVectorFilterBuilder.Build(MetadataFilter.Eq("c", "x").Not(), p1)); + + var p2 = new Dictionary(); + Assert.Equal(" WHERE (metadata ->> @k0) = ANY(@v0)", + PostgresVectorFilterBuilder.Build(MetadataFilter.In("a", new object[] { "x", "y" }), p2)); + + var p3 = new Dictionary(); + Assert.Equal(" WHERE jsonb_exists(metadata, @k0)", + PostgresVectorFilterBuilder.Build(MetadataFilter.Exists("a"), p3)); + } + + [Fact] + public void Postgres_NullFilter_ReturnsEmpty() + { + Assert.Equal(string.Empty, PostgresVectorFilterBuilder.Build((MetadataFilter?)null, new Dictionary())); + } + + // ------------------------------------------------------------------ + // RediSearch expression (net10-only Storage.Redis). + // ------------------------------------------------------------------ + + [Fact] + public void Redis_And_Tag_And_NumericRange_FullyPushed() + { + var fields = new Dictionary + { + ["category"] = RedisVectorFieldType.Tag, + ["year"] = RedisVectorFieldType.Numeric, + }; + var expr = RedisVectorQueryBuilder.BuildFilterExpression(Sample, fields, out var fullyPushed); + Assert.True(fullyPushed); + Assert.Equal("(@category:{science} @year:[2020 +inf])", expr); + } + + [Fact] + public void Redis_Or_Not_In() + { + var fields = new Dictionary { ["c"] = RedisVectorFieldType.Tag }; + + var or = RedisVectorQueryBuilder.BuildFilterExpression( + MetadataFilter.Eq("c", "x").Or(MetadataFilter.Eq("c", "y")), fields, out _); + Assert.Equal("(@c:{x} | @c:{y})", or); + + var not = RedisVectorQueryBuilder.BuildFilterExpression( + MetadataFilter.Eq("c", "x").Not(), fields, out _); + Assert.Equal("-(@c:{x})", not); + + var inExpr = RedisVectorQueryBuilder.BuildFilterExpression( + MetadataFilter.In("c", new object[] { "x", "y" }), fields, out _); + Assert.Equal("@c:{x|y}", inExpr); + } + + [Fact] + public void Redis_UndeclaredField_Or_Exists_NotFullyPushed() + { + var fields = new Dictionary { ["c"] = RedisVectorFieldType.Tag }; + + RedisVectorQueryBuilder.BuildFilterExpression(MetadataFilter.Eq("unknown", "x"), fields, out var pushed1); + Assert.False(pushed1); + + RedisVectorQueryBuilder.BuildFilterExpression(MetadataFilter.Exists("c"), fields, out var pushed2); + Assert.False(pushed2); + } + + // ------------------------------------------------------------------ + // Elasticsearch bool query (net10 test reference to Storage.Elasticsearch). + // ------------------------------------------------------------------ + + [Fact] + public void Elasticsearch_And_Term_And_Range() + { + var json = JObject.Parse(JsonConvert.SerializeObject(ElasticsearchVectorFilterBuilder.Build(Sample))); + var must = (JArray)json["bool"]!["must"]!; + Assert.Equal("science", (string?)must[0]!["term"]!["metadata.category"]); + Assert.Equal(2020, (int)must[1]!["range"]!["metadata.year"]!["gte"]!); + } + + [Fact] + public void Elasticsearch_Not_Or_In_Exists() + { + var not = JObject.Parse(JsonConvert.SerializeObject( + ElasticsearchVectorFilterBuilder.Build(MetadataFilter.Eq("c", "x").Not()))); + Assert.NotNull(not["bool"]!["must_not"]); + + var or = JObject.Parse(JsonConvert.SerializeObject( + ElasticsearchVectorFilterBuilder.Build(MetadataFilter.Eq("c", "x").Or(MetadataFilter.Eq("c", "y"))))); + Assert.Equal(1, (int)or["bool"]!["minimum_should_match"]!); + + var inJson = JObject.Parse(JsonConvert.SerializeObject( + ElasticsearchVectorFilterBuilder.Build(MetadataFilter.In("a", new object[] { "x", "y" })))); + Assert.Equal(2, ((JArray)inJson["terms"]!["metadata.a"]!).Count); + + var exists = JObject.Parse(JsonConvert.SerializeObject( + ElasticsearchVectorFilterBuilder.Build(MetadataFilter.Exists("a")))); + Assert.Equal("metadata.a", (string?)exists["exists"]!["field"]); + } +#endif + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Filtering/MetadataFilterTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Filtering/MetadataFilterTests.cs new file mode 100644 index 0000000000..e3751a150f --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Filtering/MetadataFilterTests.cs @@ -0,0 +1,266 @@ +using System; +using System.Collections.Generic; + +using AiDotNet.RetrievalAugmentedGeneration.Filtering; + +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.Filtering +{ + /// + /// Exhaustive unit tests for the in-memory evaluator covering equality, + /// inequality, ranges, set membership, existence and arbitrary AND/OR/NOT nesting. + /// + public class MetadataFilterTests + { + private static Dictionary Meta(params (string Key, object Value)[] pairs) + { + var d = new Dictionary(); + foreach (var (key, value) in pairs) + d[key] = value; + return d; + } + + // ---------------- Eq ---------------- + + [Fact] + public void Eq_MatchesEqualString() + { + var f = MetadataFilter.Eq("category", "science"); + Assert.True(f.Matches(Meta(("category", "science")))); + Assert.False(f.Matches(Meta(("category", "food")))); + } + + [Fact] + public void Eq_MissingKey_DoesNotMatch() + { + var f = MetadataFilter.Eq("category", "science"); + Assert.False(f.Matches(Meta(("other", "x")))); + } + + [Fact] + public void Eq_NumericCoercion_IntLongStringAllEqual() + { + var f = MetadataFilter.Eq("year", 2020); + Assert.True(f.Matches(Meta(("year", 2020)))); + Assert.True(f.Matches(Meta(("year", 2020L)))); + Assert.True(f.Matches(Meta(("year", 2020.0)))); + Assert.True(f.Matches(Meta(("year", "2020")))); + Assert.False(f.Matches(Meta(("year", 2021)))); + } + + [Fact] + public void Eq_Bool_Matches() + { + var f = MetadataFilter.Eq("published", true); + Assert.True(f.Matches(Meta(("published", true)))); + Assert.False(f.Matches(Meta(("published", false)))); + } + + // ---------------- Ne ---------------- + + [Fact] + public void Ne_MatchesWhenDifferent() + { + var f = MetadataFilter.Ne("category", "science"); + Assert.True(f.Matches(Meta(("category", "food")))); + Assert.False(f.Matches(Meta(("category", "science")))); + } + + [Fact] + public void Ne_MissingKey_Matches() + { + // Missing field is "not equal" to the value. + var f = MetadataFilter.Ne("category", "science"); + Assert.True(f.Matches(Meta(("other", "x")))); + } + + [Fact] + public void Ne_IsLogicalNegationOfEq() + { + var meta = Meta(("category", "food")); + var eq = MetadataFilter.Eq("category", "science"); + var ne = MetadataFilter.Ne("category", "science"); + Assert.Equal(!eq.Matches(meta), ne.Matches(meta)); + } + + // ---------------- Ranges ---------------- + + [Theory] + [InlineData(2019, false)] + [InlineData(2020, false)] + [InlineData(2021, true)] + public void Gt_Numeric(int year, bool expected) + { + var f = MetadataFilter.Gt("year", 2020); + Assert.Equal(expected, f.Matches(Meta(("year", year)))); + } + + [Theory] + [InlineData(2019, false)] + [InlineData(2020, true)] + [InlineData(2021, true)] + public void Gte_Numeric(int year, bool expected) + { + var f = MetadataFilter.Gte("year", 2020); + Assert.Equal(expected, f.Matches(Meta(("year", year)))); + } + + [Theory] + [InlineData(2019, true)] + [InlineData(2020, false)] + public void Lt_Numeric(int year, bool expected) + { + var f = MetadataFilter.Lt("year", 2020); + Assert.Equal(expected, f.Matches(Meta(("year", year)))); + } + + [Theory] + [InlineData(2020, true)] + [InlineData(2021, false)] + public void Lte_Numeric(int year, bool expected) + { + var f = MetadataFilter.Lte("year", 2020); + Assert.Equal(expected, f.Matches(Meta(("year", year)))); + } + + [Theory] + [InlineData(2019, false)] + [InlineData(2020, true)] + [InlineData(2022, true)] + [InlineData(2024, true)] + [InlineData(2025, false)] + public void Range_InclusiveBounds(int year, bool expected) + { + var f = MetadataFilter.Range("year", 2020, 2024); + Assert.Equal(expected, f.Matches(Meta(("year", year)))); + } + + [Fact] + public void Comparison_NonComparableValue_DoesNotMatch() + { + // Comparing a string field with a numeric threshold is not orderable numerically. + var f = MetadataFilter.Gt("name", 5); + Assert.False(f.Matches(Meta(("name", "abc")))); + } + + // ---------------- In ---------------- + + [Fact] + public void In_MatchesAnyMember() + { + var f = MetadataFilter.In("author", new object[] { "A", "B", "C" }); + Assert.True(f.Matches(Meta(("author", "B")))); + Assert.False(f.Matches(Meta(("author", "Z")))); + } + + [Fact] + public void In_NumericMembers() + { + var f = MetadataFilter.In("year", new object[] { 2020, 2021 }); + Assert.True(f.Matches(Meta(("year", 2021L)))); + Assert.False(f.Matches(Meta(("year", 2019)))); + } + + [Fact] + public void In_MissingKey_DoesNotMatch() + { + var f = MetadataFilter.In("author", new object[] { "A" }); + Assert.False(f.Matches(Meta(("other", "A")))); + } + + // ---------------- Exists ---------------- + + [Fact] + public void Exists_MatchesWhenPresent() + { + var f = MetadataFilter.Exists("author"); + Assert.True(f.Matches(Meta(("author", "A")))); + Assert.False(f.Matches(Meta(("other", "x")))); + } + + // ---------------- And / Or / Not ---------------- + + [Fact] + public void And_RequiresAllOperands() + { + var f = MetadataFilter.Eq("category", "science").And(MetadataFilter.Gte("year", 2020)); + Assert.True(f.Matches(Meta(("category", "science"), ("year", 2021)))); + Assert.False(f.Matches(Meta(("category", "science"), ("year", 2019)))); + Assert.False(f.Matches(Meta(("category", "food"), ("year", 2021)))); + } + + [Fact] + public void Or_RequiresAnyOperand() + { + var f = MetadataFilter.Eq("category", "science").Or(MetadataFilter.Eq("category", "math")); + Assert.True(f.Matches(Meta(("category", "science")))); + Assert.True(f.Matches(Meta(("category", "math")))); + Assert.False(f.Matches(Meta(("category", "food")))); + } + + [Fact] + public void Not_NegatesOperand() + { + var f = MetadataFilter.Eq("archived", true).Not(); + Assert.True(f.Matches(Meta(("archived", false)))); + Assert.True(f.Matches(Meta(("other", "x")))); + Assert.False(f.Matches(Meta(("archived", true)))); + } + + [Fact] + public void Nested_AndOrNot_EvaluatesCorrectly() + { + // category == "science" AND (year >= 2020 OR author in ["A","B"]) AND NOT archived + var f = MetadataFilter.Eq("category", "science") + .And(MetadataFilter.Gte("year", 2020).Or(MetadataFilter.In("author", new object[] { "A", "B" }))) + .And(MetadataFilter.Eq("archived", true).Not()); + + Assert.True(f.Matches(Meta(("category", "science"), ("year", 2022), ("archived", false)))); + Assert.True(f.Matches(Meta(("category", "science"), ("year", 2000), ("author", "A"), ("archived", false)))); + Assert.False(f.Matches(Meta(("category", "science"), ("year", 2000), ("author", "Z"), ("archived", false)))); + Assert.False(f.Matches(Meta(("category", "science"), ("year", 2022), ("archived", true)))); + Assert.False(f.Matches(Meta(("category", "food"), ("year", 2022), ("archived", false)))); + } + + [Fact] + public void Or_ThreeOperands_ViaParams() + { + var f = MetadataFilter.Or( + MetadataFilter.Eq("c", "a"), + MetadataFilter.Eq("c", "b"), + MetadataFilter.Eq("c", "c")); + Assert.True(f.Matches(Meta(("c", "c")))); + Assert.False(f.Matches(Meta(("c", "d")))); + } + + [Fact] + public void NullMetadata_TreatedAsEmpty() + { + Assert.False(MetadataFilter.Eq("k", "v").Matches(null!)); + Assert.True(MetadataFilter.Ne("k", "v").Matches(null!)); + Assert.False(MetadataFilter.Exists("k").Matches(null!)); + } + + // ---------------- Factory guards ---------------- + + [Fact] + public void Factories_RejectEmptyKey() + { + Assert.Throws(() => MetadataFilter.Eq("", "v")); + Assert.Throws(() => MetadataFilter.Exists("")); + Assert.Throws(() => MetadataFilter.In("", new object[] { 1 })); + } + + [Fact] + public void Operator_PropertyReflectsNodeKind() + { + Assert.Equal(MetadataFilterOperator.Eq, MetadataFilter.Eq("k", 1).Operator); + Assert.Equal(MetadataFilterOperator.In, MetadataFilter.In("k", new object[] { 1 }).Operator); + Assert.Equal(MetadataFilterOperator.Exists, MetadataFilter.Exists("k").Operator); + Assert.Equal(MetadataFilterOperator.And, MetadataFilter.Eq("k", 1).And(MetadataFilter.Eq("j", 2)).Operator); + Assert.Equal(MetadataFilterOperator.Or, MetadataFilter.Eq("k", 1).Or(MetadataFilter.Eq("j", 2)).Operator); + Assert.Equal(MetadataFilterOperator.Not, MetadataFilter.Eq("k", 1).Not().Operator); + } + } +} diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRetrieverTests.cs index 22ba81fb34..357b6ae4ff 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/GraphRetrieverTests.cs @@ -75,6 +75,10 @@ private class MockDocumentStore : IDocumentStore public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + + public IEnumerable> GetSimilarWithFilter(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK) { var __r = new List>(); foreach (var __d in GetSimilar(queryVector, topK * 10)) { if (filter == null || filter.Matches(__d.Metadata)) { __r.Add(__d); if (__r.Count >= topK) break; } } return __r; } + + public System.Threading.Tasks.Task>> GetSimilarWithFilterAsync(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); } public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiVectorRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiVectorRetrieverTests.cs index 824b02548f..f8c8245685 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiVectorRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/MultiVectorRetrieverTests.cs @@ -28,6 +28,10 @@ private class MockDocumentStore : IDocumentStore public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + + public IEnumerable> GetSimilarWithFilter(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK) { var __r = new List>(); foreach (var __d in GetSimilar(queryVector, topK * 10)) { if (filter == null || filter.Matches(__d.Metadata)) { __r.Add(__d); if (__r.Count >= topK) break; } } return __r; } + + public System.Threading.Tasks.Task>> GetSimilarWithFilterAsync(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); } public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ParentDocumentRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ParentDocumentRetrieverTests.cs index 02dd4f391c..eccee3b4a2 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ParentDocumentRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/ParentDocumentRetrieverTests.cs @@ -75,6 +75,10 @@ private class MockDocumentStore : IDocumentStore public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + + public IEnumerable> GetSimilarWithFilter(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK) { var __r = new List>(); foreach (var __d in GetSimilar(queryVector, topK * 10)) { if (filter == null || filter.Matches(__d.Metadata)) { __r.Add(__d); if (__r.Count >= topK) break; } } return __r; } + + public System.Threading.Tasks.Task>> GetSimilarWithFilterAsync(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); } public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/TestHelpers.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/TestHelpers.cs index 6fbfdcd8ac..83b398d8a3 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/TestHelpers.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/Retrievers/TestHelpers.cs @@ -145,6 +145,10 @@ public IEnumerable> GetAll() public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + + public IEnumerable> GetSimilarWithFilter(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK) { var __r = new List>(); foreach (var __d in GetSimilar(queryVector, topK * 10)) { if (filter == null || filter.Matches(__d.Metadata)) { __r.Add(__d); if (__r.Count >= topK) break; } } return __r; } + + public System.Threading.Tasks.Task>> GetSimilarWithFilterAsync(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); } public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/TFIDFRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/TFIDFRetrieverTests.cs index b260d1d0e9..6bf4c6717a 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/TFIDFRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/TFIDFRetrieverTests.cs @@ -27,6 +27,10 @@ private class MockDocumentStore : IDocumentStore public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + + public IEnumerable> GetSimilarWithFilter(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK) { var __r = new List>(); foreach (var __d in GetSimilar(queryVector, topK * 10)) { if (filter == null || filter.Matches(__d.Metadata)) { __r.Add(__d); if (__r.Count >= topK) break; } } return __r; } + + public System.Threading.Tasks.Task>> GetSimilarWithFilterAsync(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); } public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorRetrieverTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorRetrieverTests.cs index adb772c3fd..b26993a9d4 100644 --- a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorRetrieverTests.cs +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorRetrieverTests.cs @@ -95,6 +95,10 @@ private class MockDocumentStore : IDocumentStore public System.Threading.Tasks.Task AddBatchAsync(IEnumerable> vectorDocuments, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); AddBatch(vectorDocuments); return System.Threading.Tasks.Task.CompletedTask; } public System.Threading.Tasks.Task>> GetSimilarAsync(Vector queryVector, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilar(queryVector, topK)); } public System.Threading.Tasks.Task>> GetSimilarWithFiltersAsync(Vector queryVector, int topK, Dictionary metadataFilters, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilters(queryVector, topK, metadataFilters)); } + + public IEnumerable> GetSimilarWithFilter(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK) { var __r = new List>(); foreach (var __d in GetSimilar(queryVector, topK * 10)) { if (filter == null || filter.Matches(__d.Metadata)) { __r.Add(__d); if (__r.Count >= topK) break; } } return __r; } + + public System.Threading.Tasks.Task>> GetSimilarWithFilterAsync(Vector queryVector, AiDotNet.RetrievalAugmentedGeneration.Filtering.MetadataFilter filter, int topK, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetSimilarWithFilter(queryVector, filter, topK)); } public System.Threading.Tasks.Task?> GetByIdAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(GetById(documentId)); } public System.Threading.Tasks.Task RemoveAsync(string documentId, System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); return System.Threading.Tasks.Task.FromResult(Remove(documentId)); } public System.Threading.Tasks.Task ClearAsync(System.Threading.CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Clear(); return System.Threading.Tasks.Task.CompletedTask; } From b9222e9ec1c1decbe80a9d367a8f3ae513b4c9e7 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Sun, 19 Jul 2026 23:37:11 -0400 Subject: [PATCH 24/25] feat(rag): native FAISS vector store (IVF/HNSW/PQ) via managed wrapper Adds a real FAISS-backed IDocumentStore in a new opt-in storage project so the core AiDotNet package takes no native dependency (mirrors the AiDotNet.Storage.Elasticsearch/Sqlite extraction pattern). - New net10.0-only project src/AiDotNet.Storage.Faiss referencing the FaissNet managed wrapper (native Facebook AI Similarity Search). Wired into Directory.Packages.props (central versioning), AiDotNet.sln, excluded from the core csproj compile glob, and referenced net10-conditionally by the test project. - FaissDocumentStore supporting Flat / IVFFlat / HNSW / IVFPQ index types and Cosine (normalized inner product) / InnerProduct / L2 metrics, selected via ctor. - FaissSidecar maps FAISS int64 ids -> {docId, content, metadata, embedding}; drives id round-tripping, metadata filtering, index rebuild on delete for index types without in-place removal (HNSW), and JSON persistence alongside the FAISS index file (Save/Load). - Metadata filtering via over-fetch (topK * oversample) + reuse of the base DocumentStoreBase.MatchesFilters evaluator (FAISS has no native metadata filter). - Kept the in-memory FAISSDocumentStore (brute-force simulation) but documented it clearly and pointed to the new native FaissDocumentStore; no silent duplication. - Tests: pure sidecar/id-mapping + over-fetch planner unit tests (no native); gated [Trait Integration] tests that build a real Flat/HNSW index + search + persist, self-skipping when the native lib can't load. IVF/IVFPQ training tests are opt-in (AIDOTNET_FAISS_IVF=1) since FAISS k-means needs a complete Intel MKL runtime. Co-Authored-By: Claude Opus 4.8 (1M context) --- AiDotNet.sln | 15 + Directory.Packages.props | 3 + .../AiDotNet.Storage.Faiss.csproj | 43 ++ src/AiDotNet.Storage.Faiss/GlobalUsings.cs | 15 + .../DocumentStores/FaissDocumentStore.cs | 541 ++++++++++++++++++ .../DocumentStores/FaissSidecar.cs | 274 +++++++++ src/AiDotNet.csproj | 7 + .../DocumentStores/FAISSDocumentStore.cs | 12 +- tests/AiDotNet.Tests/AiDotNetTests.csproj | 4 + .../Storage/FaissDocumentStoreTests.cs | 253 ++++++++ .../Storage/FaissSidecarTests.cs | 165 ++++++ 11 files changed, 1331 insertions(+), 1 deletion(-) create mode 100644 src/AiDotNet.Storage.Faiss/AiDotNet.Storage.Faiss.csproj create mode 100644 src/AiDotNet.Storage.Faiss/GlobalUsings.cs create mode 100644 src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissDocumentStore.cs create mode 100644 src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissSidecar.cs create mode 100644 tests/AiDotNet.Tests/Storage/FaissDocumentStoreTests.cs create mode 100644 tests/AiDotNet.Tests/Storage/FaissSidecarTests.cs diff --git a/AiDotNet.sln b/AiDotNet.sln index 46453d8ce6..d2f2c365ed 100644 --- a/AiDotNet.sln +++ b/AiDotNet.sln @@ -41,6 +41,8 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AiDotNet.Storage.Sqlite", " EndProject Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AiDotNet.Storage.Neo4j", "src\AiDotNet.Storage.Neo4j\AiDotNet.Storage.Neo4j.csproj", "{AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}" EndProject +Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AiDotNet.Storage.Faiss", "src\AiDotNet.Storage.Faiss\AiDotNet.Storage.Faiss.csproj", "{A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}" +EndProject Global GlobalSection(SolutionConfigurationPlatforms) = preSolution Debug|Any CPU = Debug|Any CPU @@ -183,6 +185,18 @@ Global {BBC4C51F-3BFD-419C-9B59-ADD7AB006095}.Release|x64.Build.0 = Release|Any CPU {BBC4C51F-3BFD-419C-9B59-ADD7AB006095}.Release|x86.ActiveCfg = Release|Any CPU {BBC4C51F-3BFD-419C-9B59-ADD7AB006095}.Release|x86.Build.0 = Release|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|Any CPU.ActiveCfg = Debug|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|Any CPU.Build.0 = Debug|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|x64.ActiveCfg = Debug|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|x64.Build.0 = Debug|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|x86.ActiveCfg = Debug|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|x86.Build.0 = Debug|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|Any CPU.ActiveCfg = Release|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|Any CPU.Build.0 = Release|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|x64.ActiveCfg = Release|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|x64.Build.0 = Release|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|x86.ActiveCfg = Release|Any CPU + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|x86.Build.0 = Release|Any CPU {B1427BFB-8A17-4735-8EC5-E594A612AB81}.Debug|Any CPU.ActiveCfg = Debug|Any CPU {B1427BFB-8A17-4735-8EC5-E594A612AB81}.Debug|Any CPU.Build.0 = Debug|Any CPU {B1427BFB-8A17-4735-8EC5-E594A612AB81}.Debug|x64.ActiveCfg = Debug|Any CPU @@ -241,6 +255,7 @@ Global {035B7687-2236-4D7F-97C5-8BA20567DD37} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} {C1E3CF04-768E-46AE-A6EB-C9F9334365E6} = {66320409-64EC-F7C5-3DEF-65E7510DAAD1} {BBC4C51F-3BFD-419C-9B59-ADD7AB006095} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} + {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} {B1427BFB-8A17-4735-8EC5-E594A612AB81} = {66320409-64EC-F7C5-3DEF-65E7510DAAD1} {D71746F5-5342-474C-B23E-4A07B29FA4F7} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} {251CF282-40BA-452B-A563-FAE6B9826B8B} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} diff --git a/Directory.Packages.props b/Directory.Packages.props index 19ab29f26b..1fb8e6ee9d 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -308,6 +308,9 @@ + + diff --git a/src/AiDotNet.Storage.Faiss/AiDotNet.Storage.Faiss.csproj b/src/AiDotNet.Storage.Faiss/AiDotNet.Storage.Faiss.csproj new file mode 100644 index 0000000000..9c1946831f --- /dev/null +++ b/src/AiDotNet.Storage.Faiss/AiDotNet.Storage.Faiss.csproj @@ -0,0 +1,43 @@ + + + + + net10.0 + enable + enable + True + 0.204.0 + AiDotNet.Storage.Faiss + Native FAISS (Flat / IVF / HNSW / PQ ANN) document-store backend for AiDotNet's retrieval-augmented generation pipeline. Implements IDocumentStore over the FaissNet managed wrapper of Facebook AI Similarity Search. Licensed under BSL 1.1 — see https://aidotnet.dev/license. + Ooples Finance + ooples + Ooples Finance LLC 2023 + https://github.com/ooples/AiDotNet + git + https://github.com/ooples/AiDotNet + ai; aidotnet; faiss; ann; vector-search; rag; document-store; retrieval-augmented-generation + LICENSE + True + latest + AiDotNet.RetrievalAugmentedGeneration.DocumentStores + + + + + + + + + + + + + + + diff --git a/src/AiDotNet.Storage.Faiss/GlobalUsings.cs b/src/AiDotNet.Storage.Faiss/GlobalUsings.cs new file mode 100644 index 0000000000..4069439749 --- /dev/null +++ b/src/AiDotNet.Storage.Faiss/GlobalUsings.cs @@ -0,0 +1,15 @@ +// Mirrors the AiDotNet core + auto-generated AiDotNet.Tensors usings so this +// extracted package compiles without per-file boilerplate. See +// src/GlobalUsings.cs and src/obj/Release//AiDotNet.GlobalUsings.g.cs. +global using AiDotNet.Engines; +global using AiDotNet.Enums; +global using AiDotNet.Interfaces; +global using AiDotNet.LinearAlgebra; +global using AiDotNet.MetaLearning.Models; +global using AiDotNet.NeuralNetworks.Layers; +global using AiDotNet.Tensors.Engines; +global using AiDotNet.Tensors.Helpers; +global using AiDotNet.Tensors.Interfaces; +global using AiDotNet.Tensors.LinearAlgebra; +global using AiDotNet.Tensors.NumericOperations; +global using System.Globalization; diff --git a/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissDocumentStore.cs b/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissDocumentStore.cs new file mode 100644 index 0000000000..1732933d35 --- /dev/null +++ b/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissDocumentStore.cs @@ -0,0 +1,541 @@ +using System; +using System.Collections.Generic; +using System.IO; +using System.Linq; + +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json; + +using FaissIndex = FaissNet.Index; +using FaissMetricType = FaissNet.MetricType; + +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; + +/// +/// The approximate-nearest-neighbor index structure backing a . +/// +/// +/// For Beginners: This picks the trade-off between search speed, memory, and accuracy. +/// +/// Flat — exact brute-force search. Perfectly accurate, simplest, no training. Best for small/medium collections. +/// IVFFlat — clusters vectors into cells and only searches the closest cells. Much faster on large collections; requires a one-time training step and returns approximate results. +/// HNSW — a navigable small-world graph. Very fast queries with high recall, no training, but higher memory and no in-place deletion. +/// IVFPQ — IVF plus product quantization, which compresses vectors so billions fit in memory. Smallest footprint; requires training and is the most approximate. +/// +/// +/// +public enum FaissIndexType +{ + /// Exact brute-force search (IDMap2,Flat). No training required. + Flat = 0, + + /// Inverted-file index with flat storage (IVF{nlist},Flat). Requires training. + IVFFlat = 1, + + /// Hierarchical Navigable Small World graph (IDMap2,HNSW{M}). No training; no in-place deletion. + HNSW = 2, + + /// Inverted-file index with product quantization (IVF{nlist},PQ{m}). Requires training; compresses vectors. + IVFPQ = 3 +} + +/// +/// The distance metric used for similarity search. +/// +/// +/// For Beginners: This is how "closeness" between two vectors is measured. +/// +/// Cosine — angle between vectors, ignoring length. Implemented as inner product over L2-normalized vectors. The usual choice for text embeddings. +/// InnerProduct — raw dot product (no normalization). Higher is more similar. +/// L2 — Euclidean distance. Lower is more similar. +/// +/// +/// +public enum FaissDistanceMetric +{ + /// Cosine similarity via inner product over L2-normalized vectors. + Cosine = 0, + + /// Raw inner (dot) product. + InnerProduct = 1, + + /// Euclidean (L2) distance. + L2 = 2 +} + +/// +/// A real, native FAISS-backed vector store (Facebook AI Similarity Search) exposing IVF / HNSW / PQ +/// approximate-nearest-neighbor indexes through the managed FaissNet wrapper. +/// +/// The numeric data type used for vectors and relevance scores (typically float or double). +/// +/// +/// Unlike the in-memory (a brute-force simulation kept for +/// dependency-free scenarios), this store delegates indexing and search to the native FAISS +/// library and therefore matches FAISS on raw ANN throughput. It lives in the opt-in +/// AiDotNet.Storage.Faiss package so the core AiDotNet package takes no native dependency. +/// +/// +/// FAISS stores only vectors keyed by an int64 id, so this store keeps a +/// mapping each int64 id to the document's string id, content, metadata, and (metric-adjusted) +/// embedding. The sidecar drives id round-tripping, metadata filtering, index rebuilds on deletion +/// for index types that lack in-place removal (HNSW), and persistence. +/// +/// Native requirements: FaissNet ships a win-x64 native runtime. On platforms/TFMs where +/// the native library cannot be loaded, constructing this store throws (the wrapper raises a +/// / type-initialization error); callers that need graceful +/// degradation should catch that and fall back to another IDocumentStore implementation. +/// +[ComponentType(ComponentType.DocumentStore)] +[PipelineStage(PipelineStage.Indexing)] +public sealed class FaissDocumentStore : DocumentStoreBase, IDisposable +{ + private const int PersistenceVersion = 1; + + private readonly FaissIndexType _indexType; + private readonly FaissDistanceMetric _metric; + private readonly FaissMetricType _faissMetric; + private readonly int _vectorDimension; + private readonly int _nlist; + private readonly int _hnswM; + private readonly int _pqM; + private readonly int _searchOversample; + private readonly bool _normalize; + private readonly bool _requiresTraining; + + private FaissIndex _index; + private FaissSidecar _sidecar; + private bool _isTrained; + private bool _disposed; + + /// + public override int DocumentCount => _sidecar.Count; + + /// + public override int VectorDimension => _vectorDimension; + + /// The ANN index structure this store was created with. + public FaissIndexType IndexType => _indexType; + + /// The distance metric this store was created with. + public FaissDistanceMetric Metric => _metric; + + /// + /// Creates a native FAISS-backed document store. + /// + /// Dimensionality of the embeddings (must be positive). + /// Which ANN index structure to build. Defaults to exact . + /// Distance metric. Defaults to . + /// Number of IVF cells (only used by IVF-based index types). + /// Number of neighbor links per node (only used by ). + /// Number of PQ sub-quantizers (only used by ; must divide ). + /// + /// Over-fetch multiplier for metadata-filtered searches. FAISS cannot filter on metadata, so + /// filtered queries request topK * searchOversample candidates and filter them in managed code. + /// + public FaissDocumentStore( + int vectorDimension, + FaissIndexType indexType = FaissIndexType.Flat, + FaissDistanceMetric metric = FaissDistanceMetric.Cosine, + int nlist = 100, + int hnswM = 32, + int pqM = 8, + int searchOversample = 4) + { + if (vectorDimension <= 0) + throw new ArgumentOutOfRangeException(nameof(vectorDimension), "Vector dimension must be positive"); + if (nlist <= 0) + throw new ArgumentOutOfRangeException(nameof(nlist), "nlist must be positive"); + if (hnswM <= 0) + throw new ArgumentOutOfRangeException(nameof(hnswM), "hnswM must be positive"); + if (pqM <= 0) + throw new ArgumentOutOfRangeException(nameof(pqM), "pqM must be positive"); + if (indexType == FaissIndexType.IVFPQ && vectorDimension % pqM != 0) + throw new ArgumentException($"For IVFPQ, vectorDimension ({vectorDimension}) must be divisible by pqM ({pqM}).", nameof(pqM)); + if (searchOversample < 1) + throw new ArgumentOutOfRangeException(nameof(searchOversample), "searchOversample must be at least 1"); + + _vectorDimension = vectorDimension; + _indexType = indexType; + _metric = metric; + _nlist = nlist; + _hnswM = hnswM; + _pqM = pqM; + _searchOversample = searchOversample; + _faissMetric = ToFaissMetric(metric); + _normalize = metric == FaissDistanceMetric.Cosine; + _requiresTraining = indexType is FaissIndexType.IVFFlat or FaissIndexType.IVFPQ; + + _sidecar = new FaissSidecar(); + _index = CreateIndex(); + _isTrained = !_requiresTraining; + } + + // Private ctor used by Load: adopts an already-populated index + sidecar. + private FaissDocumentStore( + FaissIndex index, + FaissSidecar sidecar, + int vectorDimension, + FaissIndexType indexType, + FaissDistanceMetric metric, + int nlist, + int hnswM, + int pqM, + int searchOversample, + bool isTrained) + { + _index = index; + _sidecar = sidecar; + _vectorDimension = vectorDimension; + _indexType = indexType; + _metric = metric; + _nlist = nlist; + _hnswM = hnswM; + _pqM = pqM; + _searchOversample = searchOversample; + _faissMetric = ToFaissMetric(metric); + _normalize = metric == FaissDistanceMetric.Cosine; + _requiresTraining = indexType is FaissIndexType.IVFFlat or FaissIndexType.IVFPQ; + _isTrained = isTrained; + } + + private static FaissMetricType ToFaissMetric(FaissDistanceMetric metric) => metric switch + { + FaissDistanceMetric.Cosine => FaissMetricType.METRIC_INNER_PRODUCT, + FaissDistanceMetric.InnerProduct => FaissMetricType.METRIC_INNER_PRODUCT, + FaissDistanceMetric.L2 => FaissMetricType.METRIC_L2, + _ => throw new ArgumentOutOfRangeException(nameof(metric), metric, "Unsupported FAISS metric") + }; + + /// + /// Builds the FAISS index-factory string for the configured index type. + /// See https://github.com/facebookresearch/faiss/wiki/The-index-factory. + /// + private string BuildFactoryString() => _indexType switch + { + // IDMap2 lets Flat/HNSW carry arbitrary int64 ids and reconstruct vectors. + // IVF-based indexes carry ids natively, so no IDMap wrapper is used there. + FaissIndexType.Flat => "IDMap2,Flat", + FaissIndexType.HNSW => $"IDMap2,HNSW{_hnswM}", + FaissIndexType.IVFFlat => $"IVF{_nlist},Flat", + FaissIndexType.IVFPQ => $"IVF{_nlist},PQ{_pqM}", + _ => throw new ArgumentOutOfRangeException(nameof(_indexType), _indexType, "Unsupported FAISS index type") + }; + + private FaissIndex CreateIndex() => FaissIndex.Create(_vectorDimension, BuildFactoryString(), _faissMetric); + + /// + protected override void AddCore(VectorDocument vectorDocument) + { + AddBatchCore(new[] { vectorDocument }); + } + + /// + protected override void AddBatchCore(IList> vectorDocuments) + { + if (vectorDocuments.Count == 0) + return; + + var n = vectorDocuments.Count; + var vectors = new float[n][]; + var ids = new long[n]; + var replaced = new List(); + + for (int i = 0; i < n; i++) + { + var vd = vectorDocuments[i]; + if (vd.Embedding.Length != _vectorDimension) + throw new ArgumentException( + $"Vector dimension mismatch. Expected {_vectorDimension}, got {vd.Embedding.Length}", nameof(vectorDocuments)); + + var vec = ToFloatVector(vd.Embedding); + var id = _sidecar.Upsert( + vd.Document.Id, + vd.Document.Content, + vd.Document.Metadata, + vec, + out var replacedId); + + vectors[i] = vec; + ids[i] = id; + if (replacedId.HasValue) + replaced.Add(replacedId.Value); + } + + // Evict any documents that were replaced by an id-collision (re-add / update). + if (replaced.Count > 0) + EvictFromIndex(replaced.ToArray()); + + EnsureTrained(vectors); + + var flat = Flatten(vectors); + _index.AddWithIdsFlat(n, flat, ids); + } + + /// + protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) + { + if (_sidecar.Count == 0) + return Array.Empty>(); + + var hasFilters = metadataFilters != null && metadataFilters.Count > 0; + var fetch = FaissRetrievalPlanner.ComputeFetchCount(topK, _searchOversample, _sidecar.Count, hasFilters); + if (fetch <= 0) + return Array.Empty>(); + + var query = ToFloatVector(queryVector); + var (distances, resultIds) = _index.SearchFlat(1, query, fetch); + + // FAISS returns hits best-first, padding with id == -1 when fewer than k exist. + var ranked = new List<(FaissSidecarEntry Entry, double RawScore)>(resultIds.Length); + for (int i = 0; i < resultIds.Length; i++) + { + var id = resultIds[i]; + if (id < 0) + continue; + if (_sidecar.TryGetByFaissId(id, out var entry)) + ranked.Add((entry, distances[i])); + } + + return FaissRetrievalPlanner.SelectTopK( + ranked, + metadataFilters ?? new Dictionary(), + topK, + RawScoreToRelevance, + MetadataMatches); + } + + // Reuses the shared DocumentStoreBase.MatchesFilters evaluator against just the metadata. + private bool MetadataMatches(Dictionary metadata, Dictionary filters) + => MatchesFilters(new Document { Metadata = metadata }, filters); + + /// + /// Converts a raw FAISS distance/score into a relevance score where higher always means more relevant. + /// Inner-product/cosine scores are used directly; L2 (squared) distances map to 1/(1+d). + /// + private T RawScoreToRelevance(double raw) => _metric == FaissDistanceMetric.L2 + ? NumOps.FromDouble(1.0 / (1.0 + raw)) + : NumOps.FromDouble(raw); + + /// + protected override Document? GetByIdCore(string documentId) + => _sidecar.TryGetByDocumentId(documentId, out var entry) + ? new Document(entry.DocumentId, entry.Content, new Dictionary(entry.Metadata)) + : null; + + /// + protected override bool RemoveCore(string documentId) + { + if (!_sidecar.RemoveByDocumentId(documentId, out var removedFaissId)) + return false; + + EvictFromIndex(new[] { removedFaissId }); + return true; + } + + /// + protected override IEnumerable> GetAllCore() + => _sidecar.Entries + .Select(e => new Document(e.DocumentId, e.Content, new Dictionary(e.Metadata))) + .ToList(); + + /// + public override void Clear() + { + _sidecar.Clear(); + var old = _index; + _index = CreateIndex(); + old.Dispose(); + _isTrained = !_requiresTraining; + } + + /// + /// Removes the given int64 ids from the FAISS index. Index types that implement in-place + /// deletion (Flat, IVF*) use RemoveIds; those that do not (HNSW) trigger a full + /// rebuild from the sidecar, which is the source of truth for live documents. + /// + private void EvictFromIndex(long[] faissIds) + { + if (faissIds.Length == 0) + return; + + if (_indexType == FaissIndexType.HNSW) + { + RebuildIndex(); + return; + } + + try + { + _index.RemoveIds(faissIds); + } + catch + { + // Fall back to a rebuild if this index type cannot delete in place. + RebuildIndex(); + } + } + + /// + /// Rebuilds the FAISS index from scratch using the sidecar's live embeddings. Used when a + /// deletion cannot be applied in place, so the index and sidecar stay consistent. + /// + private void RebuildIndex() + { + var fresh = CreateIndex(); + var entries = _sidecar.Entries.ToList(); + _isTrained = !_requiresTraining; + + if (entries.Count > 0) + { + var vectors = entries.Select(e => e.Embedding).ToArray(); + var ids = entries.Select(e => e.FaissId).ToArray(); + + if (_requiresTraining && !_isTrained) + { + fresh.Train(vectors.Length, Flatten(vectors)); + _isTrained = true; + } + + fresh.AddWithIdsFlat(vectors.Length, Flatten(vectors), ids); + } + + var old = _index; + _index = fresh; + old.Dispose(); + } + + /// Trains IVF/PQ indexes on first population; a no-op for Flat/HNSW or once trained. + private void EnsureTrained(float[][] vectors) + { + if (!_requiresTraining || _isTrained || vectors.Length == 0) + return; + + _index.Train(vectors.Length, Flatten(vectors)); + _isTrained = true; + } + + private float[] ToFloatVector(Vector vector) + { + var raw = vector.ToArray(); + var result = new float[raw.Length]; + for (int i = 0; i < raw.Length; i++) + result[i] = (float)Convert.ToDouble(raw[i], CultureInfo.InvariantCulture); + + if (_normalize) + NormalizeInPlace(result); + return result; + } + + private static void NormalizeInPlace(float[] vector) + { + double sumSq = 0.0; + for (int i = 0; i < vector.Length; i++) + sumSq += (double)vector[i] * vector[i]; + + if (sumSq <= 0.0) + return; + + var inv = 1.0 / Math.Sqrt(sumSq); + for (int i = 0; i < vector.Length; i++) + vector[i] = (float)(vector[i] * inv); + } + + private static float[] Flatten(float[][] vectors) + { + if (vectors.Length == 0) + return Array.Empty(); + + var dim = vectors[0].Length; + var flat = new float[vectors.Length * dim]; + for (int i = 0; i < vectors.Length; i++) + Buffer.BlockCopy(vectors[i], 0, flat, i * dim * sizeof(float), dim * sizeof(float)); + return flat; + } + + /// + /// Persists the store to disk: the native FAISS index to {path}.faissindex and the + /// sidecar + configuration to {path}.faissmeta.json. + /// + /// Base path (without extension) for the two output files. + public void Save(string path) + { + if (string.IsNullOrWhiteSpace(path)) + throw new ArgumentException("Path cannot be null or empty", nameof(path)); + + _index.Save(IndexPath(path)); + + var meta = new PersistedMeta + { + Version = PersistenceVersion, + IndexType = _indexType, + Metric = _metric, + VectorDimension = _vectorDimension, + Nlist = _nlist, + HnswM = _hnswM, + PqM = _pqM, + SearchOversample = _searchOversample, + IsTrained = _isTrained, + SidecarJson = _sidecar.ToJson() + }; + File.WriteAllText(MetaPath(path), JsonConvert.SerializeObject(meta)); + } + + /// + /// Loads a store previously written by . + /// + /// The same base path passed to . + public static FaissDocumentStore Load(string path) + { + if (string.IsNullOrWhiteSpace(path)) + throw new ArgumentException("Path cannot be null or empty", nameof(path)); + + var metaJson = File.ReadAllText(MetaPath(path)); + var meta = JsonConvert.DeserializeObject(metaJson) + ?? throw new InvalidOperationException("Failed to deserialize FAISS store metadata."); + + var index = FaissIndex.Load(IndexPath(path)); + var sidecar = FaissSidecar.FromJson(meta.SidecarJson); + + return new FaissDocumentStore( + index, + sidecar, + meta.VectorDimension, + meta.IndexType, + meta.Metric, + meta.Nlist, + meta.HnswM, + meta.PqM, + meta.SearchOversample, + meta.IsTrained); + } + + private static string IndexPath(string basePath) => basePath + ".faissindex"; + private static string MetaPath(string basePath) => basePath + ".faissmeta.json"; + + /// Releases the native FAISS index. + public void Dispose() + { + if (_disposed) + return; + _disposed = true; + _index.Dispose(); + } + + private sealed class PersistedMeta + { + public int Version { get; set; } + public FaissIndexType IndexType { get; set; } + public FaissDistanceMetric Metric { get; set; } + public int VectorDimension { get; set; } + public int Nlist { get; set; } + public int HnswM { get; set; } + public int PqM { get; set; } + public int SearchOversample { get; set; } + public bool IsTrained { get; set; } + public string SidecarJson { get; set; } = string.Empty; + } +} diff --git a/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissSidecar.cs b/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissSidecar.cs new file mode 100644 index 0000000000..c8211fbc4a --- /dev/null +++ b/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissSidecar.cs @@ -0,0 +1,274 @@ +using System; +using System.Collections.Generic; +using System.Linq; + +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Newtonsoft.Json; + +namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; + +/// +/// One record in the FAISS sidecar: everything FAISS itself does NOT store. +/// +/// +/// +/// FAISS stores only raw vectors keyed by an id. It has no notion of +/// a document's string id, its text content, or its metadata, and lossy index types +/// (e.g. product quantization) cannot reconstruct the original vector. The sidecar keeps +/// all of that alongside the FAISS index so the store can map search hits back to real +/// documents, apply metadata filters, rebuild the index on removal (for index types that +/// do not support in-place deletion), and round-trip through persistence. +/// +/// +public sealed class FaissSidecarEntry +{ + /// The int64 id this document was assigned inside the FAISS index. + public long FaissId { get; set; } + + /// The caller-supplied unique document id. + public string DocumentId { get; set; } = string.Empty; + + /// The document's text content. + public string Content { get; set; } = string.Empty; + + /// The document's metadata (used for over-fetch filtering). + public Dictionary Metadata { get; set; } = new(); + + /// + /// The (already metric-adjusted, e.g. L2-normalized for cosine) embedding as it was + /// added to FAISS. Retained so the index can be rebuilt from scratch when a document + /// is removed from an index type that does not implement in-place deletion (HNSW). + /// + public float[] Embedding { get; set; } = Array.Empty(); +} + +/// +/// Managed, native-free sidecar for FaissDocumentStore: owns the +/// string-id <-> int64-id mapping, the per-document payload, and JSON persistence. +/// +/// +/// +/// This type deliberately performs no FAISS/native calls, so its id-assignment, +/// upsert/remove, and serialization behavior are fully unit-testable without the +/// native library being loadable. +/// +/// +public sealed class FaissSidecar +{ + private readonly Dictionary _docIdToFaissId; + private readonly Dictionary _entries; + + /// + /// The next int64 id to hand out. Monotonically increasing and never reused, even + /// after removals, so a stale FAISS hit for a deleted id can never collide with a + /// live document. + /// + public long NextId { get; private set; } + + /// Number of live documents tracked by the sidecar. + public int Count => _entries.Count; + + /// All live entries (unordered). + public IReadOnlyCollection Entries => _entries.Values; + + /// Creates an empty sidecar. + public FaissSidecar() + { + _docIdToFaissId = new Dictionary(); + _entries = new Dictionary(); + NextId = 0; + } + + private FaissSidecar(Dictionary docIdToFaissId, Dictionary entries, long nextId) + { + _docIdToFaissId = docIdToFaissId; + _entries = entries; + NextId = nextId; + } + + /// + /// Inserts a document, or replaces an existing one with the same . + /// A brand-new int64 id is always allocated; when replacing, the previous id is returned via + /// so the caller can evict it from the FAISS index. + /// + /// The newly allocated FAISS id for this document. + public long Upsert(string documentId, string content, Dictionary? metadata, float[] embedding, out long? replacedFaissId) + { + if (string.IsNullOrWhiteSpace(documentId)) + throw new ArgumentException("Document id cannot be null or empty", nameof(documentId)); + if (embedding == null) + throw new ArgumentNullException(nameof(embedding)); + + replacedFaissId = null; + if (_docIdToFaissId.TryGetValue(documentId, out var existing)) + { + replacedFaissId = existing; + _entries.Remove(existing); + } + + var id = NextId++; + _entries[id] = new FaissSidecarEntry + { + FaissId = id, + DocumentId = documentId, + Content = content ?? string.Empty, + Metadata = metadata ?? new Dictionary(), + Embedding = embedding + }; + _docIdToFaissId[documentId] = id; + return id; + } + + /// Attempts to look up an entry by its FAISS int64 id. + public bool TryGetByFaissId(long faissId, out FaissSidecarEntry entry) => _entries.TryGetValue(faissId, out entry!); + + /// Attempts to look up an entry by its caller-supplied document id. + public bool TryGetByDocumentId(string documentId, out FaissSidecarEntry entry) + { + entry = null!; + return documentId != null + && _docIdToFaissId.TryGetValue(documentId, out var faissId) + && _entries.TryGetValue(faissId, out entry!); + } + + /// + /// Removes the entry for . Returns true and the freed + /// FAISS id when the document existed; false otherwise. + /// + public bool RemoveByDocumentId(string documentId, out long removedFaissId) + { + removedFaissId = -1; + if (documentId == null || !_docIdToFaissId.TryGetValue(documentId, out var faissId)) + return false; + + removedFaissId = faissId; + _docIdToFaissId.Remove(documentId); + _entries.Remove(faissId); + return true; + } + + /// Removes all entries and resets the id counter. + public void Clear() + { + _docIdToFaissId.Clear(); + _entries.Clear(); + NextId = 0; + } + + /// Serializes the sidecar to JSON. + public string ToJson() => JsonConvert.SerializeObject(new PersistShape + { + NextId = NextId, + Entries = _entries.Values.ToList() + }); + + /// Rehydrates a sidecar from JSON produced by . + public static FaissSidecar FromJson(string json) + { + if (string.IsNullOrWhiteSpace(json)) + return new FaissSidecar(); + + var shape = JsonConvert.DeserializeObject(json) ?? new PersistShape(); + var byFaissId = new Dictionary(); + var byDocId = new Dictionary(); + long maxId = -1; + foreach (var e in shape.Entries) + { + e.Metadata ??= new Dictionary(); + e.Embedding ??= Array.Empty(); + byFaissId[e.FaissId] = e; + byDocId[e.DocumentId] = e.FaissId; + if (e.FaissId > maxId) maxId = e.FaissId; + } + + // NextId must exceed every persisted id so re-added documents never reuse an id. + var nextId = Math.Max(shape.NextId, maxId + 1); + return new FaissSidecar(byDocId, byFaissId, nextId); + } + + private sealed class PersistShape + { + public long NextId { get; set; } + public List Entries { get; set; } = new(); + } +} + +/// +/// Pure, native-free helpers for the over-fetch-then-filter retrieval plan used by +/// FaissDocumentStore. Extracted so the fetch-count math and the +/// filter/top-k selection can be unit-tested without touching FAISS. +/// +public static class FaissRetrievalPlanner +{ + /// + /// Computes how many neighbors to request from FAISS. FAISS cannot filter on + /// metadata, so when a filter is present we over-fetch topK * oversample + /// candidates and filter them in managed code. The result is always clamped to the + /// number of documents actually in the index (asking FAISS for more than it holds + /// just yields -1 padding ids). + /// + /// Requested number of results. + /// Over-fetch multiplier applied when filtering (>= 1). + /// Number of documents currently indexed. + /// Whether a metadata filter will be applied. + public static int ComputeFetchCount(int topK, int oversample, int totalDocs, bool hasFilters) + { + if (topK <= 0 || totalDocs <= 0) + return 0; + + long want = hasFilters + ? (long)topK * Math.Max(1, oversample) + : topK; + + if (want > totalDocs) want = totalDocs; + if (want < 1) want = 1; + return (int)want; + } + + /// + /// Maps ranked FAISS hits back to documents, applies the metadata filter, and takes + /// the top — preserving FAISS's (best-first) ordering. + /// + /// Numeric type of the document relevance score. + /// + /// FAISS hits in best-first order: the resolved sidecar entry and the raw distance/score FAISS returned. + /// + /// Metadata filters to apply (empty = keep all). + /// Maximum number of documents to return. + /// Converts a raw FAISS distance/score into a T relevance score. + /// + /// Metadata evaluator (document metadata, filters) -> keep. Injected so the store can + /// reuse the shared DocumentStoreBase.MatchesFilters evaluator. + /// + public static List> SelectTopK( + IEnumerable<(FaissSidecarEntry Entry, double RawScore)> ranked, + Dictionary filters, + int topK, + Func scoreConverter, + Func, Dictionary, bool> matches) + { + var results = new List>(Math.Max(0, topK)); + if (topK <= 0) + return results; + + var hasFilters = filters != null && filters.Count > 0; + foreach (var (entry, rawScore) in ranked) + { + if (entry == null) + continue; + if (hasFilters && !matches(entry.Metadata, filters!)) + continue; + + results.Add(new Document(entry.DocumentId, entry.Content, new Dictionary(entry.Metadata)) + { + RelevanceScore = scoreConverter(rawScore), + HasRelevanceScore = true + }); + + if (results.Count >= topK) + break; + } + + return results; + } +} diff --git a/src/AiDotNet.csproj b/src/AiDotNet.csproj index 492528f402..6da862e0e7 100644 --- a/src/AiDotNet.csproj +++ b/src/AiDotNet.csproj @@ -189,6 +189,13 @@ + + + + + + + diff --git a/src/RetrievalAugmentedGeneration/DocumentStores/FAISSDocumentStore.cs b/src/RetrievalAugmentedGeneration/DocumentStores/FAISSDocumentStore.cs index c2b74ba9fb..f409b7ecb5 100644 --- a/src/RetrievalAugmentedGeneration/DocumentStores/FAISSDocumentStore.cs +++ b/src/RetrievalAugmentedGeneration/DocumentStores/FAISSDocumentStore.cs @@ -10,10 +10,20 @@ namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores { /// - /// FAISS-inspired document store with indexed vectors for efficient similarity search. + /// In-memory, brute-force document store that mimics FAISS-style integer indexing. This is NOT + /// the native FAISS library. /// /// /// + /// Looking for real FAISS? This type is a dependency-free, pure-managed simulation: + /// every query is an exact O(n) brute-force cosine scan, so it does not deliver FAISS-level ANN + /// throughput. For a genuine native FAISS backend with IVF / HNSW / PQ indexes, reference the opt-in + /// AiDotNet.Storage.Faiss package and use + /// AiDotNet.RetrievalAugmentedGeneration.DocumentStores.FaissDocumentStore<T> + /// (note the Pascal-cased name). This in-memory type is retained for tests, prototyping, and + /// environments where the native FAISS runtime is unavailable. + /// + /// /// This implementation provides an in-memory simulation of Facebook AI Similarity Search (FAISS), /// using integer-based indexing for fast vector lookup. It maintains both document storage and /// a separate index mapping for optimized retrieval operations. diff --git a/tests/AiDotNet.Tests/AiDotNetTests.csproj b/tests/AiDotNet.Tests/AiDotNetTests.csproj index e79bd40a15..ea9f155efc 100644 --- a/tests/AiDotNet.Tests/AiDotNetTests.csproj +++ b/tests/AiDotNet.Tests/AiDotNetTests.csproj @@ -64,6 +64,10 @@ + + diff --git a/tests/AiDotNet.Tests/Storage/FaissDocumentStoreTests.cs b/tests/AiDotNet.Tests/Storage/FaissDocumentStoreTests.cs new file mode 100644 index 0000000000..6171d61eca --- /dev/null +++ b/tests/AiDotNet.Tests/Storage/FaissDocumentStoreTests.cs @@ -0,0 +1,253 @@ +#if NET10_0_OR_GREATER +using System; +using System.Collections.Generic; +using System.IO; +using System.Linq; + +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Xunit; + +namespace AiDotNetTests.Storage; + +/// +/// Integration tests that build a real native FAISS index and search it through +/// . They self-skip when the native FAISS runtime cannot be +/// loaded (e.g. non-win-x64 CI), so the suite stays green without the native library. +/// +/// +/// IVF / IVFPQ tests additionally require FAISS's k-means training, which links Intel MKL. Some +/// packaged FAISS builds ship an incomplete MKL redistributable (missing mkl_def.*.dll), +/// and a missing MKL kernel raises a PROCESS-FATAL error that a try/catch cannot recover — it would +/// abort the whole test host. Those tests are therefore opt-in via the +/// AIDOTNET_FAISS_IVF=1 environment variable and are only meaningful where a complete MKL +/// is present. Flat and HNSW require no training and run whenever the native lib loads. +/// +[Trait("Category", "Integration")] +public class FaissDocumentStoreTests +{ + private const int Dim = 8; + + /// Probes whether the native FAISS library loads by exercising a tiny Flat index. + private static bool NativeAvailable() + { + try + { + using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.L2); + store.Add(Doc("probe", Unit(0))); + _ = store.GetSimilar(Vec(Unit(0)), 1).ToList(); + return true; + } + catch + { + return false; + } + } + + private static bool IvfEnabled() => + Environment.GetEnvironmentVariable("AIDOTNET_FAISS_IVF") == "1"; + + [SkippableFact] + public void Flat_Cosine_FindsNearestNeighbor() + { + Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); + + using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.Cosine); + store.AddBatch(new[] + { + Doc("a", Unit(0)), + Doc("b", Unit(1)), + Doc("c", Unit(2)), + }); + + Assert.Equal(3, store.DocumentCount); + + // A query aligned with axis 1 (scaled) should return "b" first regardless of magnitude (cosine). + var results = store.GetSimilar(Vec(Scaled(1, 5f)), 2).ToList(); + Assert.Equal("b", results[0].Id); + Assert.True(results[0].HasRelevanceScore); + Assert.True(Convert.ToDouble(results[0].RelevanceScore) >= Convert.ToDouble(results[1].RelevanceScore)); + } + + [SkippableFact] + public void Flat_L2_OrdersByEuclideanDistance() + { + Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); + + using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.L2); + store.AddBatch(new[] + { + Doc("near", Custom(0.0f)), + Doc("far", Custom(9.0f)), + }); + + var results = store.GetSimilar(Vec(Custom(0.1f)), 2).ToList(); + Assert.Equal("near", results[0].Id); + Assert.Equal("far", results[1].Id); + } + + [SkippableFact] + public void GetSimilarWithFilters_OverFetchesThenFiltersMetadata() + { + Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); + + using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.Cosine); + // Closest vector is "wrong-cat"; the filter must push retrieval to the next matching doc. + store.AddBatch(new[] + { + DocMeta("wrong-cat", Unit(1), new Dictionary { ["cat"] = "history" }), + DocMeta("right-cat", Scaled(1, 0.9f), new Dictionary { ["cat"] = "science" }), + DocMeta("other", Unit(2), new Dictionary { ["cat"] = "science" }), + }); + + var filters = new Dictionary { ["cat"] = "science" }; + var results = store.GetSimilarWithFilters(Vec(Unit(1)), 1, filters).ToList(); + + Assert.Single(results); + Assert.Equal("right-cat", results[0].Id); + } + + [SkippableFact] + public void Remove_Flat_DeletesFromIndexAndSidecar() + { + Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); + + using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.L2); + store.AddBatch(new[] { Doc("a", Unit(0)), Doc("b", Unit(1)) }); + + Assert.True(store.Remove("a")); + Assert.Equal(1, store.DocumentCount); + Assert.Null(store.GetById("a")); + Assert.NotNull(store.GetById("b")); + + var results = store.GetSimilar(Vec(Unit(0)), 5).ToList(); + Assert.DoesNotContain(results, r => r.Id == "a"); + } + + [SkippableFact] + public void Hnsw_AddSearch_AndRemoveViaRebuild() + { + Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); + + using var store = new FaissDocumentStore(Dim, FaissIndexType.HNSW, FaissDistanceMetric.L2); + store.AddBatch(Enumerable.Range(0, 10).Select(i => Doc("d" + i, Custom(i))).ToArray()); + + var results = store.GetSimilar(Vec(Custom(3)), 3).ToList(); + Assert.Equal("d3", results[0].Id); + + // HNSW has no in-place delete; the store must rebuild from the sidecar and stay consistent. + Assert.True(store.Remove("d3")); + Assert.Equal(9, store.DocumentCount); + var after = store.GetSimilar(Vec(Custom(3)), 3).ToList(); + Assert.DoesNotContain(after, r => r.Id == "d3"); + } + + [SkippableFact] + public void SaveAndLoad_RoundTripsIndexAndSidecar() + { + Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); + + var basePath = Path.Combine(Path.GetTempPath(), "faiss_test_" + Guid.NewGuid().ToString("N")); + try + { + using (var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.Cosine)) + { + store.AddBatch(new[] + { + DocMeta("a", Unit(0), new Dictionary { ["k"] = "v" }), + DocMeta("b", Unit(1), new Dictionary()), + }); + store.Save(basePath); + } + + using var loaded = FaissDocumentStore.Load(basePath); + Assert.Equal(2, loaded.DocumentCount); + Assert.Equal(FaissIndexType.Flat, loaded.IndexType); + Assert.Equal(FaissDistanceMetric.Cosine, loaded.Metric); + + var doc = loaded.GetById("a"); + Assert.NotNull(doc); + Assert.Equal("v", doc!.Metadata["k"]); + + var results = loaded.GetSimilar(Vec(Unit(0)), 1).ToList(); + Assert.Equal("a", results[0].Id); + } + finally + { + TryDelete(basePath + ".faissindex"); + TryDelete(basePath + ".faissmeta.json"); + } + } + + [SkippableFact] + public void IvfFlat_TrainAddSearch_OptIn() + { + Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); + Skip.IfNot(IvfEnabled(), "IVF training requires a complete Intel MKL runtime; set AIDOTNET_FAISS_IVF=1 to enable"); + + using var store = new FaissDocumentStore(Dim, FaissIndexType.IVFFlat, FaissDistanceMetric.L2, nlist: 4); + var rnd = new Random(0); + var docs = Enumerable.Range(0, 256) + .Select(i => Doc("d" + i, Enumerable.Range(0, Dim).Select(_ => (float)rnd.NextDouble()).ToArray())) + .ToArray(); + store.AddBatch(docs); + + Assert.Equal(256, store.DocumentCount); + var results = store.GetSimilar(Vec(docs[0].Embedding.ToArray()), 3).ToList(); + Assert.NotEmpty(results); + } + + [SkippableFact] + public void IvfPq_TrainAddSearch_OptIn() + { + Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); + Skip.IfNot(IvfEnabled(), "IVFPQ training requires a complete Intel MKL runtime; set AIDOTNET_FAISS_IVF=1 to enable"); + + const int dim = 16; + using var store = new FaissDocumentStore(dim, FaissIndexType.IVFPQ, FaissDistanceMetric.L2, nlist: 8, pqM: 4); + var rnd = new Random(1); + var docs = Enumerable.Range(0, 512) + .Select(i => Doc("d" + i, Enumerable.Range(0, dim).Select(_ => (float)rnd.NextDouble()).ToArray())) + .ToArray(); + store.AddBatch(docs); + + Assert.Equal(512, store.DocumentCount); + var results = store.GetSimilar(Vec(docs[0].Embedding.ToArray()), 3).ToList(); + Assert.NotEmpty(results); + } + + // ---- helpers ------------------------------------------------------------- + + private static float[] Unit(int axis) + { + var v = new float[Dim]; + v[axis] = 1f; + return v; + } + + private static float[] Scaled(int axis, float scale) + { + var v = new float[Dim]; + v[axis] = scale; + return v; + } + + private static float[] Custom(float baseValue) + => Enumerable.Range(0, Dim).Select(j => baseValue + j * 0.01f).ToArray(); + + private static Vector Vec(float[] data) => new(data); + + private static VectorDocument Doc(string id, float[] embedding) + => new(new Document(id, "content-" + id), new Vector(embedding)); + + private static VectorDocument DocMeta(string id, float[] embedding, Dictionary metadata) + => new(new Document(id, "content-" + id, metadata), new Vector(embedding)); + + private static void TryDelete(string path) + { + try { if (File.Exists(path)) File.Delete(path); } + catch { /* best effort */ } + } +} +#endif diff --git a/tests/AiDotNet.Tests/Storage/FaissSidecarTests.cs b/tests/AiDotNet.Tests/Storage/FaissSidecarTests.cs new file mode 100644 index 0000000000..fcea88feb2 --- /dev/null +++ b/tests/AiDotNet.Tests/Storage/FaissSidecarTests.cs @@ -0,0 +1,165 @@ +#if NET10_0_OR_GREATER +using System.Collections.Generic; +using System.Linq; + +using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; +using AiDotNet.RetrievalAugmentedGeneration.Models; +using Xunit; + +namespace AiDotNetTests.Storage; + +/// +/// Pure, native-free unit tests for the FAISS sidecar (id mapping + persistence) and the +/// over-fetch/filter retrieval planner from the opt-in AiDotNet.Storage.Faiss package. +/// These exercise the managed logic that maps FAISS's int64 ids back to real documents and the +/// metadata-filter over-fetch selection — no native FAISS library is loaded, so they run on any +/// host where the net10 build is available. +/// +public class FaissSidecarTests +{ + [Fact] + public void Upsert_AssignsMonotonicIdsAndTracksCount() + { + var sidecar = new FaissSidecar(); + + var id0 = sidecar.Upsert("a", "content-a", null, new float[] { 1, 0 }, out var replaced0); + var id1 = sidecar.Upsert("b", "content-b", null, new float[] { 0, 1 }, out var replaced1); + + Assert.Equal(0, id0); + Assert.Equal(1, id1); + Assert.Null(replaced0); + Assert.Null(replaced1); + Assert.Equal(2, sidecar.Count); + Assert.Equal(2, sidecar.NextId); + } + + [Fact] + public void Upsert_SameDocumentId_ReplacesAndReportsOldId() + { + var sidecar = new FaissSidecar(); + var first = sidecar.Upsert("dup", "v1", null, new float[] { 1 }, out _); + + var second = sidecar.Upsert("dup", "v2", null, new float[] { 2 }, out var replaced); + + Assert.Equal(first, replaced); // old FAISS id reported for eviction + Assert.NotEqual(first, second); // a fresh id is always allocated + Assert.Equal(1, sidecar.Count); // still a single live document + Assert.True(sidecar.TryGetByDocumentId("dup", out var entry)); + Assert.Equal("v2", entry.Content); + Assert.Equal(second, entry.FaissId); + Assert.False(sidecar.TryGetByFaissId(first, out _)); // stale id gone + } + + [Fact] + public void RemoveByDocumentId_FreesEntry_AndNeverReusesId() + { + var sidecar = new FaissSidecar(); + var id = sidecar.Upsert("x", "c", null, new float[] { 1 }, out _); + + Assert.True(sidecar.RemoveByDocumentId("x", out var removedId)); + Assert.Equal(id, removedId); + Assert.Equal(0, sidecar.Count); + Assert.False(sidecar.TryGetByDocumentId("x", out _)); + Assert.False(sidecar.RemoveByDocumentId("x", out _)); // idempotent + + // Re-adding the same document id must not reuse the freed int64 id. + var newId = sidecar.Upsert("x", "c2", null, new float[] { 1 }, out _); + Assert.NotEqual(id, newId); + } + + [Fact] + public void JsonRoundTrip_PreservesEntriesAndIdCounter() + { + var sidecar = new FaissSidecar(); + sidecar.Upsert("a", "content-a", new Dictionary { ["year"] = 2024, ["cat"] = "sci" }, new float[] { 1, 2, 3 }, out _); + sidecar.Upsert("b", "content-b", new Dictionary { ["cat"] = "hist" }, new float[] { 4, 5, 6 }, out _); + sidecar.RemoveByDocumentId("a", out _); + var expectedNextId = sidecar.NextId; + + var restored = FaissSidecar.FromJson(sidecar.ToJson()); + + Assert.Equal(1, restored.Count); + Assert.True(expectedNextId <= restored.NextId); // never hands back a used id + Assert.True(restored.TryGetByDocumentId("b", out var b)); + Assert.Equal("content-b", b.Content); + Assert.Equal("hist", b.Metadata["cat"]); + Assert.Equal(new float[] { 4, 5, 6 }, b.Embedding); + Assert.False(restored.TryGetByDocumentId("a", out _)); + } + + [Fact] + public void FromJson_EmptyOrNull_ReturnsEmptySidecar() + { + Assert.Equal(0, FaissSidecar.FromJson("").Count); + Assert.Equal(0, FaissSidecar.FromJson(" ").Count); + } + + [Theory] + [InlineData(5, 4, 100, false, 5)] // no filter: fetch exactly topK + [InlineData(5, 4, 100, true, 20)] // filter: over-fetch topK * oversample + [InlineData(5, 4, 12, true, 12)] // over-fetch clamped to total docs + [InlineData(5, 4, 3, false, 3)] // topK clamped to total docs + [InlineData(0, 4, 100, true, 0)] // topK <= 0 + [InlineData(5, 4, 0, true, 0)] // empty index + public void ComputeFetchCount_AppliesOversampleAndClamp(int topK, int oversample, int total, bool hasFilters, int expected) + { + Assert.Equal(expected, FaissRetrievalPlanner.ComputeFetchCount(topK, oversample, total, hasFilters)); + } + + [Fact] + public void SelectTopK_OverFetchesThenFiltersAndTakesTopK_PreservingOrder() + { + // Candidates already in FAISS best-first order; only some pass the metadata filter. + var ranked = new List<(FaissSidecarEntry Entry, double RawScore)> + { + (Entry("d1", 2024), 0.99), + (Entry("d2", 2019), 0.95), // filtered out (year != 2024) + (Entry("d3", 2024), 0.90), + (Entry("d4", 2024), 0.80), + }; + var filters = new Dictionary { ["year"] = 2024 }; + + var result = FaissRetrievalPlanner.SelectTopK( + ranked, + filters, + topK: 2, + scoreConverter: d => d, + matches: EqualityMatch); + + Assert.Equal(2, result.Count); + Assert.Equal("d1", result[0].Id); // order preserved, d2 skipped + Assert.Equal("d3", result[1].Id); + Assert.Equal(0.99, result[0].RelevanceScore, 5); + Assert.True(result[0].HasRelevanceScore); + } + + [Fact] + public void SelectTopK_NoFilters_ReturnsAllUpToTopK() + { + var ranked = new List<(FaissSidecarEntry Entry, double RawScore)> + { + (Entry("d1", 2024), 0.9), + (Entry("d2", 2019), 0.8), + }; + + var result = FaissRetrievalPlanner.SelectTopK( + ranked, + new Dictionary(), + topK: 10, + scoreConverter: d => d, + matches: EqualityMatch); + + Assert.Equal(new[] { "d1", "d2" }, result.Select(r => r.Id).ToArray()); + } + + private static FaissSidecarEntry Entry(string id, int year) => new() + { + DocumentId = id, + Content = "content-" + id, + Metadata = new Dictionary { ["year"] = year } + }; + + private static bool EqualityMatch(Dictionary metadata, Dictionary filters) + => filters.All(kv => metadata.TryGetValue(kv.Key, out var v) && Equals(v, kv.Value)); +} +#endif From 2bae3bbf9198705fa71cf71f053055bdb3424932 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Mon, 20 Jul 2026 01:02:16 -0400 Subject: [PATCH 25/25] chore: native ANN vector index on the Tensors stack; drop external FaissNet MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the broken external FaissNet backend (its IVF/PQ path dies with a process-fatal "Cannot load mkl_def.2.dll") with a dependency-free native ANN index built on AiDotNet.Tensors' fused ANN kernels. - AnnVectorIndex : IVectorIndex — Flat/IVF/PQ/IVFPQ over Tensors' AnnIndex. Bridges the incremental Add/Remove/Search index contract onto a train-then- populate ANN index via a live-vector source of truth + lazy rebuild, and AttachGpu() dispatches to the fused GPU ANN kernels across all 7 backends, falling back to the managed CPU reference. Reuses the existing VectorIndexDocumentStore adapter, so it plugs into the RAG facade unchanged. - AiModelBuilder.ConfigureNativeAnnIndex(...) (on both IAiModelBuilder and the concrete builder) surfaces it from the facade with optional GPU acceleration. - Bump AiDotNet.Tensors 0.114.0 -> 0.118.0 (brings Tensors #824, the native ANN kernel stack — the gating dependency). - Remove the AiDotNet.Storage.Faiss project, its tests, the FaissNet package reference, and all solution/csproj wiring: the native path supersedes it and takes no external native dependency. Tests: AnnVectorIndexTests 9/9 green on net471 + net10 (exact flat NN, in-cluster recall for IVF/PQ/IVFPQ on separable clusters, Add/Remove/Clear/ AddBatch/cosine-normalization). Co-Authored-By: Claude Opus 4.8 (1M context) --- AiDotNet.sln | 15 - Directory.Packages.props | 10 +- .../AiDotNet.Storage.Faiss.csproj | 43 -- src/AiDotNet.Storage.Faiss/GlobalUsings.cs | 15 - .../DocumentStores/FaissDocumentStore.cs | 541 ------------------ .../DocumentStores/FaissSidecar.cs | 274 --------- src/AiDotNet.csproj | 7 - src/AiModelBuilder.Workflows.cs | 45 ++ src/Interfaces/IAiModelBuilder.cs | 27 + .../VectorSearch/Indexes/AnnVectorIndex.cs | 283 +++++++++ tests/AiDotNet.Tests/AiDotNetTests.csproj | 4 - .../Storage/FaissDocumentStoreTests.cs | 253 -------- .../Storage/FaissSidecarTests.cs | 165 ------ .../Indexes/AnnVectorIndexTests.cs | 128 +++++ 14 files changed, 490 insertions(+), 1320 deletions(-) delete mode 100644 src/AiDotNet.Storage.Faiss/AiDotNet.Storage.Faiss.csproj delete mode 100644 src/AiDotNet.Storage.Faiss/GlobalUsings.cs delete mode 100644 src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissDocumentStore.cs delete mode 100644 src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissSidecar.cs create mode 100644 src/RetrievalAugmentedGeneration/VectorSearch/Indexes/AnnVectorIndex.cs delete mode 100644 tests/AiDotNet.Tests/Storage/FaissDocumentStoreTests.cs delete mode 100644 tests/AiDotNet.Tests/Storage/FaissSidecarTests.cs create mode 100644 tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/AnnVectorIndexTests.cs diff --git a/AiDotNet.sln b/AiDotNet.sln index d2f2c365ed..46453d8ce6 100644 --- a/AiDotNet.sln +++ b/AiDotNet.sln @@ -41,8 +41,6 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AiDotNet.Storage.Sqlite", " EndProject Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AiDotNet.Storage.Neo4j", "src\AiDotNet.Storage.Neo4j\AiDotNet.Storage.Neo4j.csproj", "{AC1D46B1-1711-4DC4-AD97-DBC6FA2BF396}" EndProject -Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AiDotNet.Storage.Faiss", "src\AiDotNet.Storage.Faiss\AiDotNet.Storage.Faiss.csproj", "{A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}" -EndProject Global GlobalSection(SolutionConfigurationPlatforms) = preSolution Debug|Any CPU = Debug|Any CPU @@ -185,18 +183,6 @@ Global {BBC4C51F-3BFD-419C-9B59-ADD7AB006095}.Release|x64.Build.0 = Release|Any CPU {BBC4C51F-3BFD-419C-9B59-ADD7AB006095}.Release|x86.ActiveCfg = Release|Any CPU {BBC4C51F-3BFD-419C-9B59-ADD7AB006095}.Release|x86.Build.0 = Release|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|Any CPU.ActiveCfg = Debug|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|Any CPU.Build.0 = Debug|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|x64.ActiveCfg = Debug|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|x64.Build.0 = Debug|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|x86.ActiveCfg = Debug|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Debug|x86.Build.0 = Debug|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|Any CPU.ActiveCfg = Release|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|Any CPU.Build.0 = Release|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|x64.ActiveCfg = Release|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|x64.Build.0 = Release|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|x86.ActiveCfg = Release|Any CPU - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53}.Release|x86.Build.0 = Release|Any CPU {B1427BFB-8A17-4735-8EC5-E594A612AB81}.Debug|Any CPU.ActiveCfg = Debug|Any CPU {B1427BFB-8A17-4735-8EC5-E594A612AB81}.Debug|Any CPU.Build.0 = Debug|Any CPU {B1427BFB-8A17-4735-8EC5-E594A612AB81}.Debug|x64.ActiveCfg = Debug|Any CPU @@ -255,7 +241,6 @@ Global {035B7687-2236-4D7F-97C5-8BA20567DD37} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} {C1E3CF04-768E-46AE-A6EB-C9F9334365E6} = {66320409-64EC-F7C5-3DEF-65E7510DAAD1} {BBC4C51F-3BFD-419C-9B59-ADD7AB006095} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} - {A1F5C7D2-9E3B-4C6A-8D2F-1B7E4A9C0D53} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} {B1427BFB-8A17-4735-8EC5-E594A612AB81} = {66320409-64EC-F7C5-3DEF-65E7510DAAD1} {D71746F5-5342-474C-B23E-4A07B29FA4F7} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} {251CF282-40BA-452B-A563-FAE6B9826B8B} = {827E0CD3-B72D-47B6-A68D-7590B98EB39B} diff --git a/Directory.Packages.props b/Directory.Packages.props index 1fb8e6ee9d..3530323db1 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -248,8 +248,13 @@ PersistentInputRegistry}; the src/Training mirror classes stay as the consumer-side primitives (the fused-primitive centralization in #1847/#1848 wires every consumer through them) and now run against the fixed engine. 0.113.0 is PUBLISHED to NuGet and is a superset of 0.112.0 (also - carries #765 compiled-ReduceMax axis fill), so all rationale above is retained. --> - + carries #765 compiled-ReduceMax axis fill), so all rationale above is retained. + + Bumped 0.114.0 -> 0.118.0: brings AiDotNet.Tensors #824 — the native fused ANN kernel stack + (IAnnBackend + AnnPrimitives + AnnIndex: Flat/IVF/PQ/IVFPQ across all 7 backends). This is the + GATING dependency for the dependency-free NativeAnn vector index (AnnVectorIndex) that + replaces the broken external FaissNet IVF/PQ path. Requires publishing Tensors 0.118.0. --> + @@ -310,7 +315,6 @@ - diff --git a/src/AiDotNet.Storage.Faiss/AiDotNet.Storage.Faiss.csproj b/src/AiDotNet.Storage.Faiss/AiDotNet.Storage.Faiss.csproj deleted file mode 100644 index 9c1946831f..0000000000 --- a/src/AiDotNet.Storage.Faiss/AiDotNet.Storage.Faiss.csproj +++ /dev/null @@ -1,43 +0,0 @@ - - - - - net10.0 - enable - enable - True - 0.204.0 - AiDotNet.Storage.Faiss - Native FAISS (Flat / IVF / HNSW / PQ ANN) document-store backend for AiDotNet's retrieval-augmented generation pipeline. Implements IDocumentStore over the FaissNet managed wrapper of Facebook AI Similarity Search. Licensed under BSL 1.1 — see https://aidotnet.dev/license. - Ooples Finance - ooples - Ooples Finance LLC 2023 - https://github.com/ooples/AiDotNet - git - https://github.com/ooples/AiDotNet - ai; aidotnet; faiss; ann; vector-search; rag; document-store; retrieval-augmented-generation - LICENSE - True - latest - AiDotNet.RetrievalAugmentedGeneration.DocumentStores - - - - - - - - - - - - - - - diff --git a/src/AiDotNet.Storage.Faiss/GlobalUsings.cs b/src/AiDotNet.Storage.Faiss/GlobalUsings.cs deleted file mode 100644 index 4069439749..0000000000 --- a/src/AiDotNet.Storage.Faiss/GlobalUsings.cs +++ /dev/null @@ -1,15 +0,0 @@ -// Mirrors the AiDotNet core + auto-generated AiDotNet.Tensors usings so this -// extracted package compiles without per-file boilerplate. See -// src/GlobalUsings.cs and src/obj/Release//AiDotNet.GlobalUsings.g.cs. -global using AiDotNet.Engines; -global using AiDotNet.Enums; -global using AiDotNet.Interfaces; -global using AiDotNet.LinearAlgebra; -global using AiDotNet.MetaLearning.Models; -global using AiDotNet.NeuralNetworks.Layers; -global using AiDotNet.Tensors.Engines; -global using AiDotNet.Tensors.Helpers; -global using AiDotNet.Tensors.Interfaces; -global using AiDotNet.Tensors.LinearAlgebra; -global using AiDotNet.Tensors.NumericOperations; -global using System.Globalization; diff --git a/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissDocumentStore.cs b/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissDocumentStore.cs deleted file mode 100644 index 1732933d35..0000000000 --- a/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissDocumentStore.cs +++ /dev/null @@ -1,541 +0,0 @@ -using System; -using System.Collections.Generic; -using System.IO; -using System.Linq; - -using AiDotNet.Attributes; -using AiDotNet.Enums; -using AiDotNet.LinearAlgebra; -using AiDotNet.RetrievalAugmentedGeneration.Models; -using Newtonsoft.Json; - -using FaissIndex = FaissNet.Index; -using FaissMetricType = FaissNet.MetricType; - -namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; - -/// -/// The approximate-nearest-neighbor index structure backing a . -/// -/// -/// For Beginners: This picks the trade-off between search speed, memory, and accuracy. -/// -/// Flat — exact brute-force search. Perfectly accurate, simplest, no training. Best for small/medium collections. -/// IVFFlat — clusters vectors into cells and only searches the closest cells. Much faster on large collections; requires a one-time training step and returns approximate results. -/// HNSW — a navigable small-world graph. Very fast queries with high recall, no training, but higher memory and no in-place deletion. -/// IVFPQ — IVF plus product quantization, which compresses vectors so billions fit in memory. Smallest footprint; requires training and is the most approximate. -/// -/// -/// -public enum FaissIndexType -{ - /// Exact brute-force search (IDMap2,Flat). No training required. - Flat = 0, - - /// Inverted-file index with flat storage (IVF{nlist},Flat). Requires training. - IVFFlat = 1, - - /// Hierarchical Navigable Small World graph (IDMap2,HNSW{M}). No training; no in-place deletion. - HNSW = 2, - - /// Inverted-file index with product quantization (IVF{nlist},PQ{m}). Requires training; compresses vectors. - IVFPQ = 3 -} - -/// -/// The distance metric used for similarity search. -/// -/// -/// For Beginners: This is how "closeness" between two vectors is measured. -/// -/// Cosine — angle between vectors, ignoring length. Implemented as inner product over L2-normalized vectors. The usual choice for text embeddings. -/// InnerProduct — raw dot product (no normalization). Higher is more similar. -/// L2 — Euclidean distance. Lower is more similar. -/// -/// -/// -public enum FaissDistanceMetric -{ - /// Cosine similarity via inner product over L2-normalized vectors. - Cosine = 0, - - /// Raw inner (dot) product. - InnerProduct = 1, - - /// Euclidean (L2) distance. - L2 = 2 -} - -/// -/// A real, native FAISS-backed vector store (Facebook AI Similarity Search) exposing IVF / HNSW / PQ -/// approximate-nearest-neighbor indexes through the managed FaissNet wrapper. -/// -/// The numeric data type used for vectors and relevance scores (typically float or double). -/// -/// -/// Unlike the in-memory (a brute-force simulation kept for -/// dependency-free scenarios), this store delegates indexing and search to the native FAISS -/// library and therefore matches FAISS on raw ANN throughput. It lives in the opt-in -/// AiDotNet.Storage.Faiss package so the core AiDotNet package takes no native dependency. -/// -/// -/// FAISS stores only vectors keyed by an int64 id, so this store keeps a -/// mapping each int64 id to the document's string id, content, metadata, and (metric-adjusted) -/// embedding. The sidecar drives id round-tripping, metadata filtering, index rebuilds on deletion -/// for index types that lack in-place removal (HNSW), and persistence. -/// -/// Native requirements: FaissNet ships a win-x64 native runtime. On platforms/TFMs where -/// the native library cannot be loaded, constructing this store throws (the wrapper raises a -/// / type-initialization error); callers that need graceful -/// degradation should catch that and fall back to another IDocumentStore implementation. -/// -[ComponentType(ComponentType.DocumentStore)] -[PipelineStage(PipelineStage.Indexing)] -public sealed class FaissDocumentStore : DocumentStoreBase, IDisposable -{ - private const int PersistenceVersion = 1; - - private readonly FaissIndexType _indexType; - private readonly FaissDistanceMetric _metric; - private readonly FaissMetricType _faissMetric; - private readonly int _vectorDimension; - private readonly int _nlist; - private readonly int _hnswM; - private readonly int _pqM; - private readonly int _searchOversample; - private readonly bool _normalize; - private readonly bool _requiresTraining; - - private FaissIndex _index; - private FaissSidecar _sidecar; - private bool _isTrained; - private bool _disposed; - - /// - public override int DocumentCount => _sidecar.Count; - - /// - public override int VectorDimension => _vectorDimension; - - /// The ANN index structure this store was created with. - public FaissIndexType IndexType => _indexType; - - /// The distance metric this store was created with. - public FaissDistanceMetric Metric => _metric; - - /// - /// Creates a native FAISS-backed document store. - /// - /// Dimensionality of the embeddings (must be positive). - /// Which ANN index structure to build. Defaults to exact . - /// Distance metric. Defaults to . - /// Number of IVF cells (only used by IVF-based index types). - /// Number of neighbor links per node (only used by ). - /// Number of PQ sub-quantizers (only used by ; must divide ). - /// - /// Over-fetch multiplier for metadata-filtered searches. FAISS cannot filter on metadata, so - /// filtered queries request topK * searchOversample candidates and filter them in managed code. - /// - public FaissDocumentStore( - int vectorDimension, - FaissIndexType indexType = FaissIndexType.Flat, - FaissDistanceMetric metric = FaissDistanceMetric.Cosine, - int nlist = 100, - int hnswM = 32, - int pqM = 8, - int searchOversample = 4) - { - if (vectorDimension <= 0) - throw new ArgumentOutOfRangeException(nameof(vectorDimension), "Vector dimension must be positive"); - if (nlist <= 0) - throw new ArgumentOutOfRangeException(nameof(nlist), "nlist must be positive"); - if (hnswM <= 0) - throw new ArgumentOutOfRangeException(nameof(hnswM), "hnswM must be positive"); - if (pqM <= 0) - throw new ArgumentOutOfRangeException(nameof(pqM), "pqM must be positive"); - if (indexType == FaissIndexType.IVFPQ && vectorDimension % pqM != 0) - throw new ArgumentException($"For IVFPQ, vectorDimension ({vectorDimension}) must be divisible by pqM ({pqM}).", nameof(pqM)); - if (searchOversample < 1) - throw new ArgumentOutOfRangeException(nameof(searchOversample), "searchOversample must be at least 1"); - - _vectorDimension = vectorDimension; - _indexType = indexType; - _metric = metric; - _nlist = nlist; - _hnswM = hnswM; - _pqM = pqM; - _searchOversample = searchOversample; - _faissMetric = ToFaissMetric(metric); - _normalize = metric == FaissDistanceMetric.Cosine; - _requiresTraining = indexType is FaissIndexType.IVFFlat or FaissIndexType.IVFPQ; - - _sidecar = new FaissSidecar(); - _index = CreateIndex(); - _isTrained = !_requiresTraining; - } - - // Private ctor used by Load: adopts an already-populated index + sidecar. - private FaissDocumentStore( - FaissIndex index, - FaissSidecar sidecar, - int vectorDimension, - FaissIndexType indexType, - FaissDistanceMetric metric, - int nlist, - int hnswM, - int pqM, - int searchOversample, - bool isTrained) - { - _index = index; - _sidecar = sidecar; - _vectorDimension = vectorDimension; - _indexType = indexType; - _metric = metric; - _nlist = nlist; - _hnswM = hnswM; - _pqM = pqM; - _searchOversample = searchOversample; - _faissMetric = ToFaissMetric(metric); - _normalize = metric == FaissDistanceMetric.Cosine; - _requiresTraining = indexType is FaissIndexType.IVFFlat or FaissIndexType.IVFPQ; - _isTrained = isTrained; - } - - private static FaissMetricType ToFaissMetric(FaissDistanceMetric metric) => metric switch - { - FaissDistanceMetric.Cosine => FaissMetricType.METRIC_INNER_PRODUCT, - FaissDistanceMetric.InnerProduct => FaissMetricType.METRIC_INNER_PRODUCT, - FaissDistanceMetric.L2 => FaissMetricType.METRIC_L2, - _ => throw new ArgumentOutOfRangeException(nameof(metric), metric, "Unsupported FAISS metric") - }; - - /// - /// Builds the FAISS index-factory string for the configured index type. - /// See https://github.com/facebookresearch/faiss/wiki/The-index-factory. - /// - private string BuildFactoryString() => _indexType switch - { - // IDMap2 lets Flat/HNSW carry arbitrary int64 ids and reconstruct vectors. - // IVF-based indexes carry ids natively, so no IDMap wrapper is used there. - FaissIndexType.Flat => "IDMap2,Flat", - FaissIndexType.HNSW => $"IDMap2,HNSW{_hnswM}", - FaissIndexType.IVFFlat => $"IVF{_nlist},Flat", - FaissIndexType.IVFPQ => $"IVF{_nlist},PQ{_pqM}", - _ => throw new ArgumentOutOfRangeException(nameof(_indexType), _indexType, "Unsupported FAISS index type") - }; - - private FaissIndex CreateIndex() => FaissIndex.Create(_vectorDimension, BuildFactoryString(), _faissMetric); - - /// - protected override void AddCore(VectorDocument vectorDocument) - { - AddBatchCore(new[] { vectorDocument }); - } - - /// - protected override void AddBatchCore(IList> vectorDocuments) - { - if (vectorDocuments.Count == 0) - return; - - var n = vectorDocuments.Count; - var vectors = new float[n][]; - var ids = new long[n]; - var replaced = new List(); - - for (int i = 0; i < n; i++) - { - var vd = vectorDocuments[i]; - if (vd.Embedding.Length != _vectorDimension) - throw new ArgumentException( - $"Vector dimension mismatch. Expected {_vectorDimension}, got {vd.Embedding.Length}", nameof(vectorDocuments)); - - var vec = ToFloatVector(vd.Embedding); - var id = _sidecar.Upsert( - vd.Document.Id, - vd.Document.Content, - vd.Document.Metadata, - vec, - out var replacedId); - - vectors[i] = vec; - ids[i] = id; - if (replacedId.HasValue) - replaced.Add(replacedId.Value); - } - - // Evict any documents that were replaced by an id-collision (re-add / update). - if (replaced.Count > 0) - EvictFromIndex(replaced.ToArray()); - - EnsureTrained(vectors); - - var flat = Flatten(vectors); - _index.AddWithIdsFlat(n, flat, ids); - } - - /// - protected override IEnumerable> GetSimilarCore(Vector queryVector, int topK, Dictionary metadataFilters) - { - if (_sidecar.Count == 0) - return Array.Empty>(); - - var hasFilters = metadataFilters != null && metadataFilters.Count > 0; - var fetch = FaissRetrievalPlanner.ComputeFetchCount(topK, _searchOversample, _sidecar.Count, hasFilters); - if (fetch <= 0) - return Array.Empty>(); - - var query = ToFloatVector(queryVector); - var (distances, resultIds) = _index.SearchFlat(1, query, fetch); - - // FAISS returns hits best-first, padding with id == -1 when fewer than k exist. - var ranked = new List<(FaissSidecarEntry Entry, double RawScore)>(resultIds.Length); - for (int i = 0; i < resultIds.Length; i++) - { - var id = resultIds[i]; - if (id < 0) - continue; - if (_sidecar.TryGetByFaissId(id, out var entry)) - ranked.Add((entry, distances[i])); - } - - return FaissRetrievalPlanner.SelectTopK( - ranked, - metadataFilters ?? new Dictionary(), - topK, - RawScoreToRelevance, - MetadataMatches); - } - - // Reuses the shared DocumentStoreBase.MatchesFilters evaluator against just the metadata. - private bool MetadataMatches(Dictionary metadata, Dictionary filters) - => MatchesFilters(new Document { Metadata = metadata }, filters); - - /// - /// Converts a raw FAISS distance/score into a relevance score where higher always means more relevant. - /// Inner-product/cosine scores are used directly; L2 (squared) distances map to 1/(1+d). - /// - private T RawScoreToRelevance(double raw) => _metric == FaissDistanceMetric.L2 - ? NumOps.FromDouble(1.0 / (1.0 + raw)) - : NumOps.FromDouble(raw); - - /// - protected override Document? GetByIdCore(string documentId) - => _sidecar.TryGetByDocumentId(documentId, out var entry) - ? new Document(entry.DocumentId, entry.Content, new Dictionary(entry.Metadata)) - : null; - - /// - protected override bool RemoveCore(string documentId) - { - if (!_sidecar.RemoveByDocumentId(documentId, out var removedFaissId)) - return false; - - EvictFromIndex(new[] { removedFaissId }); - return true; - } - - /// - protected override IEnumerable> GetAllCore() - => _sidecar.Entries - .Select(e => new Document(e.DocumentId, e.Content, new Dictionary(e.Metadata))) - .ToList(); - - /// - public override void Clear() - { - _sidecar.Clear(); - var old = _index; - _index = CreateIndex(); - old.Dispose(); - _isTrained = !_requiresTraining; - } - - /// - /// Removes the given int64 ids from the FAISS index. Index types that implement in-place - /// deletion (Flat, IVF*) use RemoveIds; those that do not (HNSW) trigger a full - /// rebuild from the sidecar, which is the source of truth for live documents. - /// - private void EvictFromIndex(long[] faissIds) - { - if (faissIds.Length == 0) - return; - - if (_indexType == FaissIndexType.HNSW) - { - RebuildIndex(); - return; - } - - try - { - _index.RemoveIds(faissIds); - } - catch - { - // Fall back to a rebuild if this index type cannot delete in place. - RebuildIndex(); - } - } - - /// - /// Rebuilds the FAISS index from scratch using the sidecar's live embeddings. Used when a - /// deletion cannot be applied in place, so the index and sidecar stay consistent. - /// - private void RebuildIndex() - { - var fresh = CreateIndex(); - var entries = _sidecar.Entries.ToList(); - _isTrained = !_requiresTraining; - - if (entries.Count > 0) - { - var vectors = entries.Select(e => e.Embedding).ToArray(); - var ids = entries.Select(e => e.FaissId).ToArray(); - - if (_requiresTraining && !_isTrained) - { - fresh.Train(vectors.Length, Flatten(vectors)); - _isTrained = true; - } - - fresh.AddWithIdsFlat(vectors.Length, Flatten(vectors), ids); - } - - var old = _index; - _index = fresh; - old.Dispose(); - } - - /// Trains IVF/PQ indexes on first population; a no-op for Flat/HNSW or once trained. - private void EnsureTrained(float[][] vectors) - { - if (!_requiresTraining || _isTrained || vectors.Length == 0) - return; - - _index.Train(vectors.Length, Flatten(vectors)); - _isTrained = true; - } - - private float[] ToFloatVector(Vector vector) - { - var raw = vector.ToArray(); - var result = new float[raw.Length]; - for (int i = 0; i < raw.Length; i++) - result[i] = (float)Convert.ToDouble(raw[i], CultureInfo.InvariantCulture); - - if (_normalize) - NormalizeInPlace(result); - return result; - } - - private static void NormalizeInPlace(float[] vector) - { - double sumSq = 0.0; - for (int i = 0; i < vector.Length; i++) - sumSq += (double)vector[i] * vector[i]; - - if (sumSq <= 0.0) - return; - - var inv = 1.0 / Math.Sqrt(sumSq); - for (int i = 0; i < vector.Length; i++) - vector[i] = (float)(vector[i] * inv); - } - - private static float[] Flatten(float[][] vectors) - { - if (vectors.Length == 0) - return Array.Empty(); - - var dim = vectors[0].Length; - var flat = new float[vectors.Length * dim]; - for (int i = 0; i < vectors.Length; i++) - Buffer.BlockCopy(vectors[i], 0, flat, i * dim * sizeof(float), dim * sizeof(float)); - return flat; - } - - /// - /// Persists the store to disk: the native FAISS index to {path}.faissindex and the - /// sidecar + configuration to {path}.faissmeta.json. - /// - /// Base path (without extension) for the two output files. - public void Save(string path) - { - if (string.IsNullOrWhiteSpace(path)) - throw new ArgumentException("Path cannot be null or empty", nameof(path)); - - _index.Save(IndexPath(path)); - - var meta = new PersistedMeta - { - Version = PersistenceVersion, - IndexType = _indexType, - Metric = _metric, - VectorDimension = _vectorDimension, - Nlist = _nlist, - HnswM = _hnswM, - PqM = _pqM, - SearchOversample = _searchOversample, - IsTrained = _isTrained, - SidecarJson = _sidecar.ToJson() - }; - File.WriteAllText(MetaPath(path), JsonConvert.SerializeObject(meta)); - } - - /// - /// Loads a store previously written by . - /// - /// The same base path passed to . - public static FaissDocumentStore Load(string path) - { - if (string.IsNullOrWhiteSpace(path)) - throw new ArgumentException("Path cannot be null or empty", nameof(path)); - - var metaJson = File.ReadAllText(MetaPath(path)); - var meta = JsonConvert.DeserializeObject(metaJson) - ?? throw new InvalidOperationException("Failed to deserialize FAISS store metadata."); - - var index = FaissIndex.Load(IndexPath(path)); - var sidecar = FaissSidecar.FromJson(meta.SidecarJson); - - return new FaissDocumentStore( - index, - sidecar, - meta.VectorDimension, - meta.IndexType, - meta.Metric, - meta.Nlist, - meta.HnswM, - meta.PqM, - meta.SearchOversample, - meta.IsTrained); - } - - private static string IndexPath(string basePath) => basePath + ".faissindex"; - private static string MetaPath(string basePath) => basePath + ".faissmeta.json"; - - /// Releases the native FAISS index. - public void Dispose() - { - if (_disposed) - return; - _disposed = true; - _index.Dispose(); - } - - private sealed class PersistedMeta - { - public int Version { get; set; } - public FaissIndexType IndexType { get; set; } - public FaissDistanceMetric Metric { get; set; } - public int VectorDimension { get; set; } - public int Nlist { get; set; } - public int HnswM { get; set; } - public int PqM { get; set; } - public int SearchOversample { get; set; } - public bool IsTrained { get; set; } - public string SidecarJson { get; set; } = string.Empty; - } -} diff --git a/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissSidecar.cs b/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissSidecar.cs deleted file mode 100644 index c8211fbc4a..0000000000 --- a/src/AiDotNet.Storage.Faiss/RetrievalAugmentedGeneration/DocumentStores/FaissSidecar.cs +++ /dev/null @@ -1,274 +0,0 @@ -using System; -using System.Collections.Generic; -using System.Linq; - -using AiDotNet.RetrievalAugmentedGeneration.Models; -using Newtonsoft.Json; - -namespace AiDotNet.RetrievalAugmentedGeneration.DocumentStores; - -/// -/// One record in the FAISS sidecar: everything FAISS itself does NOT store. -/// -/// -/// -/// FAISS stores only raw vectors keyed by an id. It has no notion of -/// a document's string id, its text content, or its metadata, and lossy index types -/// (e.g. product quantization) cannot reconstruct the original vector. The sidecar keeps -/// all of that alongside the FAISS index so the store can map search hits back to real -/// documents, apply metadata filters, rebuild the index on removal (for index types that -/// do not support in-place deletion), and round-trip through persistence. -/// -/// -public sealed class FaissSidecarEntry -{ - /// The int64 id this document was assigned inside the FAISS index. - public long FaissId { get; set; } - - /// The caller-supplied unique document id. - public string DocumentId { get; set; } = string.Empty; - - /// The document's text content. - public string Content { get; set; } = string.Empty; - - /// The document's metadata (used for over-fetch filtering). - public Dictionary Metadata { get; set; } = new(); - - /// - /// The (already metric-adjusted, e.g. L2-normalized for cosine) embedding as it was - /// added to FAISS. Retained so the index can be rebuilt from scratch when a document - /// is removed from an index type that does not implement in-place deletion (HNSW). - /// - public float[] Embedding { get; set; } = Array.Empty(); -} - -/// -/// Managed, native-free sidecar for FaissDocumentStore: owns the -/// string-id <-> int64-id mapping, the per-document payload, and JSON persistence. -/// -/// -/// -/// This type deliberately performs no FAISS/native calls, so its id-assignment, -/// upsert/remove, and serialization behavior are fully unit-testable without the -/// native library being loadable. -/// -/// -public sealed class FaissSidecar -{ - private readonly Dictionary _docIdToFaissId; - private readonly Dictionary _entries; - - /// - /// The next int64 id to hand out. Monotonically increasing and never reused, even - /// after removals, so a stale FAISS hit for a deleted id can never collide with a - /// live document. - /// - public long NextId { get; private set; } - - /// Number of live documents tracked by the sidecar. - public int Count => _entries.Count; - - /// All live entries (unordered). - public IReadOnlyCollection Entries => _entries.Values; - - /// Creates an empty sidecar. - public FaissSidecar() - { - _docIdToFaissId = new Dictionary(); - _entries = new Dictionary(); - NextId = 0; - } - - private FaissSidecar(Dictionary docIdToFaissId, Dictionary entries, long nextId) - { - _docIdToFaissId = docIdToFaissId; - _entries = entries; - NextId = nextId; - } - - /// - /// Inserts a document, or replaces an existing one with the same . - /// A brand-new int64 id is always allocated; when replacing, the previous id is returned via - /// so the caller can evict it from the FAISS index. - /// - /// The newly allocated FAISS id for this document. - public long Upsert(string documentId, string content, Dictionary? metadata, float[] embedding, out long? replacedFaissId) - { - if (string.IsNullOrWhiteSpace(documentId)) - throw new ArgumentException("Document id cannot be null or empty", nameof(documentId)); - if (embedding == null) - throw new ArgumentNullException(nameof(embedding)); - - replacedFaissId = null; - if (_docIdToFaissId.TryGetValue(documentId, out var existing)) - { - replacedFaissId = existing; - _entries.Remove(existing); - } - - var id = NextId++; - _entries[id] = new FaissSidecarEntry - { - FaissId = id, - DocumentId = documentId, - Content = content ?? string.Empty, - Metadata = metadata ?? new Dictionary(), - Embedding = embedding - }; - _docIdToFaissId[documentId] = id; - return id; - } - - /// Attempts to look up an entry by its FAISS int64 id. - public bool TryGetByFaissId(long faissId, out FaissSidecarEntry entry) => _entries.TryGetValue(faissId, out entry!); - - /// Attempts to look up an entry by its caller-supplied document id. - public bool TryGetByDocumentId(string documentId, out FaissSidecarEntry entry) - { - entry = null!; - return documentId != null - && _docIdToFaissId.TryGetValue(documentId, out var faissId) - && _entries.TryGetValue(faissId, out entry!); - } - - /// - /// Removes the entry for . Returns true and the freed - /// FAISS id when the document existed; false otherwise. - /// - public bool RemoveByDocumentId(string documentId, out long removedFaissId) - { - removedFaissId = -1; - if (documentId == null || !_docIdToFaissId.TryGetValue(documentId, out var faissId)) - return false; - - removedFaissId = faissId; - _docIdToFaissId.Remove(documentId); - _entries.Remove(faissId); - return true; - } - - /// Removes all entries and resets the id counter. - public void Clear() - { - _docIdToFaissId.Clear(); - _entries.Clear(); - NextId = 0; - } - - /// Serializes the sidecar to JSON. - public string ToJson() => JsonConvert.SerializeObject(new PersistShape - { - NextId = NextId, - Entries = _entries.Values.ToList() - }); - - /// Rehydrates a sidecar from JSON produced by . - public static FaissSidecar FromJson(string json) - { - if (string.IsNullOrWhiteSpace(json)) - return new FaissSidecar(); - - var shape = JsonConvert.DeserializeObject(json) ?? new PersistShape(); - var byFaissId = new Dictionary(); - var byDocId = new Dictionary(); - long maxId = -1; - foreach (var e in shape.Entries) - { - e.Metadata ??= new Dictionary(); - e.Embedding ??= Array.Empty(); - byFaissId[e.FaissId] = e; - byDocId[e.DocumentId] = e.FaissId; - if (e.FaissId > maxId) maxId = e.FaissId; - } - - // NextId must exceed every persisted id so re-added documents never reuse an id. - var nextId = Math.Max(shape.NextId, maxId + 1); - return new FaissSidecar(byDocId, byFaissId, nextId); - } - - private sealed class PersistShape - { - public long NextId { get; set; } - public List Entries { get; set; } = new(); - } -} - -/// -/// Pure, native-free helpers for the over-fetch-then-filter retrieval plan used by -/// FaissDocumentStore. Extracted so the fetch-count math and the -/// filter/top-k selection can be unit-tested without touching FAISS. -/// -public static class FaissRetrievalPlanner -{ - /// - /// Computes how many neighbors to request from FAISS. FAISS cannot filter on - /// metadata, so when a filter is present we over-fetch topK * oversample - /// candidates and filter them in managed code. The result is always clamped to the - /// number of documents actually in the index (asking FAISS for more than it holds - /// just yields -1 padding ids). - /// - /// Requested number of results. - /// Over-fetch multiplier applied when filtering (>= 1). - /// Number of documents currently indexed. - /// Whether a metadata filter will be applied. - public static int ComputeFetchCount(int topK, int oversample, int totalDocs, bool hasFilters) - { - if (topK <= 0 || totalDocs <= 0) - return 0; - - long want = hasFilters - ? (long)topK * Math.Max(1, oversample) - : topK; - - if (want > totalDocs) want = totalDocs; - if (want < 1) want = 1; - return (int)want; - } - - /// - /// Maps ranked FAISS hits back to documents, applies the metadata filter, and takes - /// the top — preserving FAISS's (best-first) ordering. - /// - /// Numeric type of the document relevance score. - /// - /// FAISS hits in best-first order: the resolved sidecar entry and the raw distance/score FAISS returned. - /// - /// Metadata filters to apply (empty = keep all). - /// Maximum number of documents to return. - /// Converts a raw FAISS distance/score into a T relevance score. - /// - /// Metadata evaluator (document metadata, filters) -> keep. Injected so the store can - /// reuse the shared DocumentStoreBase.MatchesFilters evaluator. - /// - public static List> SelectTopK( - IEnumerable<(FaissSidecarEntry Entry, double RawScore)> ranked, - Dictionary filters, - int topK, - Func scoreConverter, - Func, Dictionary, bool> matches) - { - var results = new List>(Math.Max(0, topK)); - if (topK <= 0) - return results; - - var hasFilters = filters != null && filters.Count > 0; - foreach (var (entry, rawScore) in ranked) - { - if (entry == null) - continue; - if (hasFilters && !matches(entry.Metadata, filters!)) - continue; - - results.Add(new Document(entry.DocumentId, entry.Content, new Dictionary(entry.Metadata)) - { - RelevanceScore = scoreConverter(rawScore), - HasRelevanceScore = true - }); - - if (results.Count >= topK) - break; - } - - return results; - } -} diff --git a/src/AiDotNet.csproj b/src/AiDotNet.csproj index 6da862e0e7..492528f402 100644 --- a/src/AiDotNet.csproj +++ b/src/AiDotNet.csproj @@ -189,13 +189,6 @@ - - - - - - - diff --git a/src/AiModelBuilder.Workflows.cs b/src/AiModelBuilder.Workflows.cs index b0085ba944..0531df177c 100644 --- a/src/AiModelBuilder.Workflows.cs +++ b/src/AiModelBuilder.Workflows.cs @@ -620,6 +620,51 @@ public IAiModelBuilder ConfigureVectorIndex( return this; } + /// + /// Configures a dependency-free native ANN index (Flat / IVF / PQ / IVFPQ) implemented on the AiDotNet + /// Tensors fused-kernel stack, adapts it into a document store, and builds a dense retriever over it. + /// + /// + /// + /// This is the self-contained replacement for the external FaissNet-backed store: it takes no FAISS / MKL + /// native dependency and, when is set, dispatches distance / assignment / PQ scans + /// to the fused ANN GPU kernels (IAnnBackend) across every supported backend, falling back to the + /// managed CPU reference otherwise. The underlying index trains IVF/PQ codebooks and (re)builds automatically. + /// + /// For Beginners: Use this to get FAISS-style IVF/PQ/HNSW-scale search entirely on AiDotNet's + /// own stack — no external native library to install and nothing that can fail to load at runtime. + /// + /// + public IAiModelBuilder ConfigureNativeAnnIndex( + RetrievalAugmentedGeneration.VectorSearch.Indexes.AnnVectorIndexType indexType = RetrievalAugmentedGeneration.VectorSearch.Indexes.AnnVectorIndexType.Flat, + int vectorDimension = 0, + RetrievalAugmentedGeneration.VectorSearch.Indexes.AnnVectorMetric metric = RetrievalAugmentedGeneration.VectorSearch.Indexes.AnnVectorMetric.Cosine, + int nlist = 64, + int nprobe = 8, + int m = 8, + int ksub = 256, + bool useGpu = false, + int defaultTopK = 5) + { + var index = new RetrievalAugmentedGeneration.VectorSearch.Indexes.AnnVectorIndex( + indexType, vectorDimension, metric, nlist, nprobe, m, ksub); + + if (useGpu) + { + // Best-effort: attach the best available GPU backend so ANN ops use the fused kernels. A missing or + // failing backend leaves the index on its managed CPU path (AnnVectorIndex tolerates a null backend). + try { index.AttachGpu(AiDotNet.Tensors.Engines.DirectGpu.DirectGpuBackendFactory.Create()); } + catch { /* CPU path */ } + } + + var store = new VectorIndexDocumentStore(index, vectorDimension); + _ragDocumentStore = store; + + var embeddingModel = _configuredEmbeddingModel ?? BuildDefaultEmbeddingModel(vectorDimension); + _ragRetriever = new RetrievalAugmentedGeneration.Retrievers.VectorRetriever(store, embeddingModel, defaultTopK); + return this; + } + /// /// Constructs the selected in-memory vector index using the supplied similarity metric. /// diff --git a/src/Interfaces/IAiModelBuilder.cs b/src/Interfaces/IAiModelBuilder.cs index e56d547779..fe77f68345 100644 --- a/src/Interfaces/IAiModelBuilder.cs +++ b/src/Interfaces/IAiModelBuilder.cs @@ -2422,6 +2422,33 @@ IAiModelBuilder ConfigureVectorIndex( RetrievalAugmentedGeneration.VectorSearch.ISimilarityMetric? metric = null, int defaultTopK = 5); + /// + /// Configures a dependency-free native ANN index (Flat / IVF / PQ / IVFPQ) implemented on the AiDotNet + /// Tensors fused-kernel stack, adapts it into a document store, and builds a dense retriever over it. This + /// is the self-contained replacement for the external FaissNet backend — no FAISS / MKL native dependency — + /// and dispatches to the GPU ANN kernels across all supported backends when is set. + /// + /// Which native ANN structure to build (default exact Flat). + /// The embedding dimension (0 = inferred on first add). + /// Distance metric (default cosine). + /// IVF coarse lists (IVF/IVFPQ only). + /// IVF lists probed per query (IVF/IVFPQ only). + /// PQ subspaces (PQ/IVFPQ only; must divide the dimension). + /// PQ sub-centroids per subspace (PQ/IVFPQ only). + /// When true, attaches the best available GPU backend so ANN ops use the fused kernels. + /// Default number of documents the retriever returns per query. + /// The builder instance for method chaining. + IAiModelBuilder ConfigureNativeAnnIndex( + RetrievalAugmentedGeneration.VectorSearch.Indexes.AnnVectorIndexType indexType = RetrievalAugmentedGeneration.VectorSearch.Indexes.AnnVectorIndexType.Flat, + int vectorDimension = 0, + RetrievalAugmentedGeneration.VectorSearch.Indexes.AnnVectorMetric metric = RetrievalAugmentedGeneration.VectorSearch.Indexes.AnnVectorMetric.Cosine, + int nlist = 64, + int nprobe = 8, + int m = 8, + int ksub = 256, + bool useGpu = false, + int defaultTopK = 5); + /// /// Materializes a declarative RAG configuration into the builder's RAG components (chunking, embedding, /// document store + retriever, reranking, context compression). diff --git a/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/AnnVectorIndex.cs b/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/AnnVectorIndex.cs new file mode 100644 index 0000000000..aa2b558a5f --- /dev/null +++ b/src/RetrievalAugmentedGeneration/VectorSearch/Indexes/AnnVectorIndex.cs @@ -0,0 +1,283 @@ +using System; +using System.Collections.Generic; +using System.Globalization; +using System.Linq; +using AiDotNet.Attributes; +using AiDotNet.Enums; +using AiDotNet.LinearAlgebra; +using AiDotNet.Tensors.Ann; +using AiDotNet.Tensors.Engines.DirectGpu; +using AiDotNet.Tensors.Helpers; + +namespace AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Indexes +{ + /// + /// The approximate-nearest-neighbour structure backing an . + /// Mirrors the FAISS index families but runs entirely on the AiDotNet Tensors stack. + /// + /// + /// For Beginners: This is the speed/memory/accuracy trade-off. + /// + /// Flat — exact brute-force. Perfectly accurate, no training. Best for small/medium sets. + /// Ivf — clusters vectors and only scans the closest clusters. Much faster on large sets; approximate; trained. + /// Pq — product quantization: compressed codes searched with asymmetric distance. Tiny footprint; trained. + /// IvfPq — IVF over PQ-compressed residuals (FAISS's workhorse). Smallest + fast; most approximate; trained. + /// + /// + public enum AnnVectorIndexType + { + /// Exact brute-force scan. No training. + Flat = 0, + /// Inverted file (coarse k-means) — probe the nearest lists only. + Ivf = 1, + /// Product quantization — compressed codes + ADC search. + Pq = 2, + /// IVF over PQ-compressed residuals. + IvfPq = 3 + } + + /// The distance metric used by an . + public enum AnnVectorMetric + { + /// Cosine similarity via inner product over L2-normalized vectors (the usual choice for text embeddings). + Cosine = 0, + /// Raw inner (dot) product — higher is more similar. + InnerProduct = 1, + /// Euclidean (L2) distance — lower is more similar. + L2 = 2 + } + + /// + /// A dependency-free approximate-nearest-neighbour index (Flat / IVF / PQ / IVFPQ) implemented on the + /// AiDotNet Tensors stack via . This is the native replacement for the external + /// FaissNet-backed index whose IVF/PQ path is blocked by an incomplete MKL redistribution: it takes no + /// external native dependency and, when a GPU backend is attached, dispatches its distance / assignment / + /// PQ scans to the fused ANN kernels () across all supported backends, falling + /// back to the managed CPU reference otherwise. + /// + /// + /// + /// The underlying trains its coarse quantizer / PQ codebooks once on a representative + /// sample and does not support in-place deletion, so this adapter keeps the live vectors as the source of + /// truth and lazily (re)builds the index on the next after any mutation. That keeps the + /// simple incremental contract (Add / Remove / Search) correct over an index + /// that is fundamentally train-then-populate. + /// + /// For Beginners: You just add and remove vectors by id and search — the training and rebuilding + /// that IVF/PQ need happen automatically behind the scenes the first time you query after a change. + /// + /// The numeric data type used for vectors and scores (typically float or double). + [ComponentType(ComponentType.VectorIndex)] + [PipelineStage(PipelineStage.Retrieval)] + public sealed class AnnVectorIndex : IVectorIndex + { + private readonly AnnVectorIndexType _type; + private readonly AnnVectorMetric _metric; + private readonly int _annMetric; // AnnPrimitives.Metric* code + private readonly bool _normalize; // cosine => normalize + inner product + private readonly int _nlist; + private readonly int _nprobe; + private readonly int _m; + private readonly int _ksub; + private readonly int _seed; + private readonly bool _requiresTraining; + private readonly INumericOperations _numOps; + + // Live vectors are the source of truth for lazy (re)builds; insertion order gives deterministic ordinals. + private readonly Dictionary _vectors = new(); + private readonly List _order = new(); + + private int _dim; + private AnnIndex? _index; + private List? _ordinalToId; + private bool _dirty = true; + private IDirectGpuBackend? _gpu; + + /// + public int Count => _vectors.Count; + + /// The ANN index structure this index was created with. + public AnnVectorIndexType IndexType => _type; + + /// The distance metric this index was created with. + public AnnVectorMetric Metric => _metric; + + /// + /// Creates a native ANN vector index. + /// + /// Which ANN structure to build. Defaults to exact . + /// Vector dimensionality (0 = inferred from the first added vector). + /// Distance metric. Defaults to . + /// Number of IVF coarse lists (IVF/IVFPQ only). + /// Number of IVF lists probed per query (IVF/IVFPQ only). + /// Number of PQ subspaces (PQ/IVFPQ only; must divide when known). + /// PQ sub-centroids per subspace (PQ/IVFPQ only; typically 256). + /// Deterministic k-means seed. + public AnnVectorIndex( + AnnVectorIndexType indexType = AnnVectorIndexType.Flat, + int dimension = 0, + AnnVectorMetric metric = AnnVectorMetric.Cosine, + int nlist = 64, + int nprobe = 8, + int m = 8, + int ksub = 256, + int seed = 42) + { + if (dimension < 0) throw new ArgumentOutOfRangeException(nameof(dimension)); + if (nlist <= 0) throw new ArgumentOutOfRangeException(nameof(nlist)); + if (nprobe <= 0) throw new ArgumentOutOfRangeException(nameof(nprobe)); + if (m <= 0) throw new ArgumentOutOfRangeException(nameof(m)); + if (ksub <= 0) throw new ArgumentOutOfRangeException(nameof(ksub)); + + _type = indexType; + _metric = metric; + _dim = dimension; + _normalize = metric == AnnVectorMetric.Cosine; + _annMetric = metric == AnnVectorMetric.L2 ? AnnPrimitives.MetricL2 : AnnPrimitives.MetricInnerProduct; + _nlist = nlist; + _nprobe = nprobe; + _m = m; + _ksub = ksub; + _seed = seed; + _requiresTraining = indexType != AnnVectorIndexType.Flat; + _numOps = MathHelper.GetNumericOperations(); + } + + /// + /// Attaches a GPU backend so search/training dispatch to the fused ANN kernels when available. Pass a + /// backend from DirectGpuBackendFactory.Create(), or null to force the managed CPU path. + /// Takes effect on the next (re)build. + /// + public void AttachGpu(IDirectGpuBackend? backend) + { + _gpu = backend; + _dirty = true; + } + + /// + public void Add(string id, Vector vector) + { + if (string.IsNullOrEmpty(id)) throw new ArgumentException("Id must be non-empty.", nameof(id)); + if (vector == null) throw new ArgumentNullException(nameof(vector)); + if (_dim == 0) _dim = vector.Length; + if (vector.Length != _dim) + throw new ArgumentException($"Vector dimension mismatch. Expected {_dim}, got {vector.Length}.", nameof(vector)); + + if (!_vectors.ContainsKey(id)) _order.Add(id); + _vectors[id] = ToFloat(vector); + _dirty = true; + } + + /// + public void AddBatch(Dictionary> vectors) + { + if (vectors == null) throw new ArgumentNullException(nameof(vectors)); + foreach (var kvp in vectors) + Add(kvp.Key, kvp.Value); + } + + /// + public bool Remove(string id) + { + if (id == null || !_vectors.Remove(id)) return false; + _order.Remove(id); + _dirty = true; + return true; + } + + /// + public void Clear() + { + _vectors.Clear(); + _order.Clear(); + _index = null; + _ordinalToId = null; + _dirty = true; + } + + /// + public List<(string Id, T Score)> Search(Vector query, int k) + { + if (query == null) throw new ArgumentNullException(nameof(query)); + if (k <= 0 || _vectors.Count == 0) return new List<(string Id, T Score)>(); + + EnsureBuilt(); + if (_index == null || _ordinalToId == null) return new List<(string Id, T Score)>(); + + var q = ToFloat(query); + var (ids, distances) = _index.Search(q, k); + + var results = new List<(string Id, T Score)>(ids.Length); + for (int i = 0; i < ids.Length; i++) + { + int ordinal = (int)ids[i]; + if (ordinal < 0 || ordinal >= _ordinalToId.Count) continue; + results.Add((_ordinalToId[ordinal], ToScore(distances[i]))); + } + return results; + } + + /// Lazily (re)builds the underlying from the live vectors after any mutation. + private void EnsureBuilt() + { + if (_index != null && !_dirty) return; + + var live = _order.Where(_vectors.ContainsKey).ToList(); + if (live.Count == 0 || _dim == 0) + { + _index = null; + _ordinalToId = null; + _dirty = false; + return; + } + + var idx = new AnnIndex((AnnIndexType)_type, _dim, _annMetric, _nlist, _nprobe, _m, _ksub, _seed); + idx.AttachGpu(_gpu); + + if (_requiresTraining) + { + var flat = new float[(long)live.Count * _dim]; + for (int i = 0; i < live.Count; i++) + Array.Copy(_vectors[live[i]], 0, flat, (long)i * _dim, _dim); + idx.Train(flat, live.Count); + } + + var ordinalToId = new List(live.Count); + for (int i = 0; i < live.Count; i++) + { + idx.Add(i, _vectors[live[i]]); + ordinalToId.Add(live[i]); + } + + _index = idx; + _ordinalToId = ordinalToId; + _dirty = false; + } + + private float[] ToFloat(Vector vector) + { + var result = new float[vector.Length]; + for (int i = 0; i < vector.Length; i++) + result[i] = (float)Convert.ToDouble(vector[i], CultureInfo.InvariantCulture); + if (_normalize) NormalizeInPlace(result); + return result; + } + + private static void NormalizeInPlace(float[] vector) + { + double sumSq = 0.0; + for (int i = 0; i < vector.Length; i++) + sumSq += (double)vector[i] * vector[i]; + if (sumSq <= 0.0) return; + var inv = 1.0 / Math.Sqrt(sumSq); + for (int i = 0; i < vector.Length; i++) + vector[i] = (float)(vector[i] * inv); + } + + // AnnIndex returns metric-native distances (L2: smaller nearer; inner product: larger nearer) already + // ordered best-first. Map to a relevance score where higher always means more similar. + private T ToScore(float distance) => _annMetric == AnnPrimitives.MetricL2 + ? _numOps.FromDouble(1.0 / (1.0 + distance)) + : _numOps.FromDouble(distance); + } +} diff --git a/tests/AiDotNet.Tests/AiDotNetTests.csproj b/tests/AiDotNet.Tests/AiDotNetTests.csproj index ea9f155efc..e79bd40a15 100644 --- a/tests/AiDotNet.Tests/AiDotNetTests.csproj +++ b/tests/AiDotNet.Tests/AiDotNetTests.csproj @@ -64,10 +64,6 @@ - - diff --git a/tests/AiDotNet.Tests/Storage/FaissDocumentStoreTests.cs b/tests/AiDotNet.Tests/Storage/FaissDocumentStoreTests.cs deleted file mode 100644 index 6171d61eca..0000000000 --- a/tests/AiDotNet.Tests/Storage/FaissDocumentStoreTests.cs +++ /dev/null @@ -1,253 +0,0 @@ -#if NET10_0_OR_GREATER -using System; -using System.Collections.Generic; -using System.IO; -using System.Linq; - -using AiDotNet.LinearAlgebra; -using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; -using AiDotNet.RetrievalAugmentedGeneration.Models; -using Xunit; - -namespace AiDotNetTests.Storage; - -/// -/// Integration tests that build a real native FAISS index and search it through -/// . They self-skip when the native FAISS runtime cannot be -/// loaded (e.g. non-win-x64 CI), so the suite stays green without the native library. -/// -/// -/// IVF / IVFPQ tests additionally require FAISS's k-means training, which links Intel MKL. Some -/// packaged FAISS builds ship an incomplete MKL redistributable (missing mkl_def.*.dll), -/// and a missing MKL kernel raises a PROCESS-FATAL error that a try/catch cannot recover — it would -/// abort the whole test host. Those tests are therefore opt-in via the -/// AIDOTNET_FAISS_IVF=1 environment variable and are only meaningful where a complete MKL -/// is present. Flat and HNSW require no training and run whenever the native lib loads. -/// -[Trait("Category", "Integration")] -public class FaissDocumentStoreTests -{ - private const int Dim = 8; - - /// Probes whether the native FAISS library loads by exercising a tiny Flat index. - private static bool NativeAvailable() - { - try - { - using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.L2); - store.Add(Doc("probe", Unit(0))); - _ = store.GetSimilar(Vec(Unit(0)), 1).ToList(); - return true; - } - catch - { - return false; - } - } - - private static bool IvfEnabled() => - Environment.GetEnvironmentVariable("AIDOTNET_FAISS_IVF") == "1"; - - [SkippableFact] - public void Flat_Cosine_FindsNearestNeighbor() - { - Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); - - using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.Cosine); - store.AddBatch(new[] - { - Doc("a", Unit(0)), - Doc("b", Unit(1)), - Doc("c", Unit(2)), - }); - - Assert.Equal(3, store.DocumentCount); - - // A query aligned with axis 1 (scaled) should return "b" first regardless of magnitude (cosine). - var results = store.GetSimilar(Vec(Scaled(1, 5f)), 2).ToList(); - Assert.Equal("b", results[0].Id); - Assert.True(results[0].HasRelevanceScore); - Assert.True(Convert.ToDouble(results[0].RelevanceScore) >= Convert.ToDouble(results[1].RelevanceScore)); - } - - [SkippableFact] - public void Flat_L2_OrdersByEuclideanDistance() - { - Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); - - using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.L2); - store.AddBatch(new[] - { - Doc("near", Custom(0.0f)), - Doc("far", Custom(9.0f)), - }); - - var results = store.GetSimilar(Vec(Custom(0.1f)), 2).ToList(); - Assert.Equal("near", results[0].Id); - Assert.Equal("far", results[1].Id); - } - - [SkippableFact] - public void GetSimilarWithFilters_OverFetchesThenFiltersMetadata() - { - Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); - - using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.Cosine); - // Closest vector is "wrong-cat"; the filter must push retrieval to the next matching doc. - store.AddBatch(new[] - { - DocMeta("wrong-cat", Unit(1), new Dictionary { ["cat"] = "history" }), - DocMeta("right-cat", Scaled(1, 0.9f), new Dictionary { ["cat"] = "science" }), - DocMeta("other", Unit(2), new Dictionary { ["cat"] = "science" }), - }); - - var filters = new Dictionary { ["cat"] = "science" }; - var results = store.GetSimilarWithFilters(Vec(Unit(1)), 1, filters).ToList(); - - Assert.Single(results); - Assert.Equal("right-cat", results[0].Id); - } - - [SkippableFact] - public void Remove_Flat_DeletesFromIndexAndSidecar() - { - Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); - - using var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.L2); - store.AddBatch(new[] { Doc("a", Unit(0)), Doc("b", Unit(1)) }); - - Assert.True(store.Remove("a")); - Assert.Equal(1, store.DocumentCount); - Assert.Null(store.GetById("a")); - Assert.NotNull(store.GetById("b")); - - var results = store.GetSimilar(Vec(Unit(0)), 5).ToList(); - Assert.DoesNotContain(results, r => r.Id == "a"); - } - - [SkippableFact] - public void Hnsw_AddSearch_AndRemoveViaRebuild() - { - Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); - - using var store = new FaissDocumentStore(Dim, FaissIndexType.HNSW, FaissDistanceMetric.L2); - store.AddBatch(Enumerable.Range(0, 10).Select(i => Doc("d" + i, Custom(i))).ToArray()); - - var results = store.GetSimilar(Vec(Custom(3)), 3).ToList(); - Assert.Equal("d3", results[0].Id); - - // HNSW has no in-place delete; the store must rebuild from the sidecar and stay consistent. - Assert.True(store.Remove("d3")); - Assert.Equal(9, store.DocumentCount); - var after = store.GetSimilar(Vec(Custom(3)), 3).ToList(); - Assert.DoesNotContain(after, r => r.Id == "d3"); - } - - [SkippableFact] - public void SaveAndLoad_RoundTripsIndexAndSidecar() - { - Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); - - var basePath = Path.Combine(Path.GetTempPath(), "faiss_test_" + Guid.NewGuid().ToString("N")); - try - { - using (var store = new FaissDocumentStore(Dim, FaissIndexType.Flat, FaissDistanceMetric.Cosine)) - { - store.AddBatch(new[] - { - DocMeta("a", Unit(0), new Dictionary { ["k"] = "v" }), - DocMeta("b", Unit(1), new Dictionary()), - }); - store.Save(basePath); - } - - using var loaded = FaissDocumentStore.Load(basePath); - Assert.Equal(2, loaded.DocumentCount); - Assert.Equal(FaissIndexType.Flat, loaded.IndexType); - Assert.Equal(FaissDistanceMetric.Cosine, loaded.Metric); - - var doc = loaded.GetById("a"); - Assert.NotNull(doc); - Assert.Equal("v", doc!.Metadata["k"]); - - var results = loaded.GetSimilar(Vec(Unit(0)), 1).ToList(); - Assert.Equal("a", results[0].Id); - } - finally - { - TryDelete(basePath + ".faissindex"); - TryDelete(basePath + ".faissmeta.json"); - } - } - - [SkippableFact] - public void IvfFlat_TrainAddSearch_OptIn() - { - Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); - Skip.IfNot(IvfEnabled(), "IVF training requires a complete Intel MKL runtime; set AIDOTNET_FAISS_IVF=1 to enable"); - - using var store = new FaissDocumentStore(Dim, FaissIndexType.IVFFlat, FaissDistanceMetric.L2, nlist: 4); - var rnd = new Random(0); - var docs = Enumerable.Range(0, 256) - .Select(i => Doc("d" + i, Enumerable.Range(0, Dim).Select(_ => (float)rnd.NextDouble()).ToArray())) - .ToArray(); - store.AddBatch(docs); - - Assert.Equal(256, store.DocumentCount); - var results = store.GetSimilar(Vec(docs[0].Embedding.ToArray()), 3).ToList(); - Assert.NotEmpty(results); - } - - [SkippableFact] - public void IvfPq_TrainAddSearch_OptIn() - { - Skip.IfNot(NativeAvailable(), "Native FAISS library not available on this host"); - Skip.IfNot(IvfEnabled(), "IVFPQ training requires a complete Intel MKL runtime; set AIDOTNET_FAISS_IVF=1 to enable"); - - const int dim = 16; - using var store = new FaissDocumentStore(dim, FaissIndexType.IVFPQ, FaissDistanceMetric.L2, nlist: 8, pqM: 4); - var rnd = new Random(1); - var docs = Enumerable.Range(0, 512) - .Select(i => Doc("d" + i, Enumerable.Range(0, dim).Select(_ => (float)rnd.NextDouble()).ToArray())) - .ToArray(); - store.AddBatch(docs); - - Assert.Equal(512, store.DocumentCount); - var results = store.GetSimilar(Vec(docs[0].Embedding.ToArray()), 3).ToList(); - Assert.NotEmpty(results); - } - - // ---- helpers ------------------------------------------------------------- - - private static float[] Unit(int axis) - { - var v = new float[Dim]; - v[axis] = 1f; - return v; - } - - private static float[] Scaled(int axis, float scale) - { - var v = new float[Dim]; - v[axis] = scale; - return v; - } - - private static float[] Custom(float baseValue) - => Enumerable.Range(0, Dim).Select(j => baseValue + j * 0.01f).ToArray(); - - private static Vector Vec(float[] data) => new(data); - - private static VectorDocument Doc(string id, float[] embedding) - => new(new Document(id, "content-" + id), new Vector(embedding)); - - private static VectorDocument DocMeta(string id, float[] embedding, Dictionary metadata) - => new(new Document(id, "content-" + id, metadata), new Vector(embedding)); - - private static void TryDelete(string path) - { - try { if (File.Exists(path)) File.Delete(path); } - catch { /* best effort */ } - } -} -#endif diff --git a/tests/AiDotNet.Tests/Storage/FaissSidecarTests.cs b/tests/AiDotNet.Tests/Storage/FaissSidecarTests.cs deleted file mode 100644 index fcea88feb2..0000000000 --- a/tests/AiDotNet.Tests/Storage/FaissSidecarTests.cs +++ /dev/null @@ -1,165 +0,0 @@ -#if NET10_0_OR_GREATER -using System.Collections.Generic; -using System.Linq; - -using AiDotNet.RetrievalAugmentedGeneration.DocumentStores; -using AiDotNet.RetrievalAugmentedGeneration.Models; -using Xunit; - -namespace AiDotNetTests.Storage; - -/// -/// Pure, native-free unit tests for the FAISS sidecar (id mapping + persistence) and the -/// over-fetch/filter retrieval planner from the opt-in AiDotNet.Storage.Faiss package. -/// These exercise the managed logic that maps FAISS's int64 ids back to real documents and the -/// metadata-filter over-fetch selection — no native FAISS library is loaded, so they run on any -/// host where the net10 build is available. -/// -public class FaissSidecarTests -{ - [Fact] - public void Upsert_AssignsMonotonicIdsAndTracksCount() - { - var sidecar = new FaissSidecar(); - - var id0 = sidecar.Upsert("a", "content-a", null, new float[] { 1, 0 }, out var replaced0); - var id1 = sidecar.Upsert("b", "content-b", null, new float[] { 0, 1 }, out var replaced1); - - Assert.Equal(0, id0); - Assert.Equal(1, id1); - Assert.Null(replaced0); - Assert.Null(replaced1); - Assert.Equal(2, sidecar.Count); - Assert.Equal(2, sidecar.NextId); - } - - [Fact] - public void Upsert_SameDocumentId_ReplacesAndReportsOldId() - { - var sidecar = new FaissSidecar(); - var first = sidecar.Upsert("dup", "v1", null, new float[] { 1 }, out _); - - var second = sidecar.Upsert("dup", "v2", null, new float[] { 2 }, out var replaced); - - Assert.Equal(first, replaced); // old FAISS id reported for eviction - Assert.NotEqual(first, second); // a fresh id is always allocated - Assert.Equal(1, sidecar.Count); // still a single live document - Assert.True(sidecar.TryGetByDocumentId("dup", out var entry)); - Assert.Equal("v2", entry.Content); - Assert.Equal(second, entry.FaissId); - Assert.False(sidecar.TryGetByFaissId(first, out _)); // stale id gone - } - - [Fact] - public void RemoveByDocumentId_FreesEntry_AndNeverReusesId() - { - var sidecar = new FaissSidecar(); - var id = sidecar.Upsert("x", "c", null, new float[] { 1 }, out _); - - Assert.True(sidecar.RemoveByDocumentId("x", out var removedId)); - Assert.Equal(id, removedId); - Assert.Equal(0, sidecar.Count); - Assert.False(sidecar.TryGetByDocumentId("x", out _)); - Assert.False(sidecar.RemoveByDocumentId("x", out _)); // idempotent - - // Re-adding the same document id must not reuse the freed int64 id. - var newId = sidecar.Upsert("x", "c2", null, new float[] { 1 }, out _); - Assert.NotEqual(id, newId); - } - - [Fact] - public void JsonRoundTrip_PreservesEntriesAndIdCounter() - { - var sidecar = new FaissSidecar(); - sidecar.Upsert("a", "content-a", new Dictionary { ["year"] = 2024, ["cat"] = "sci" }, new float[] { 1, 2, 3 }, out _); - sidecar.Upsert("b", "content-b", new Dictionary { ["cat"] = "hist" }, new float[] { 4, 5, 6 }, out _); - sidecar.RemoveByDocumentId("a", out _); - var expectedNextId = sidecar.NextId; - - var restored = FaissSidecar.FromJson(sidecar.ToJson()); - - Assert.Equal(1, restored.Count); - Assert.True(expectedNextId <= restored.NextId); // never hands back a used id - Assert.True(restored.TryGetByDocumentId("b", out var b)); - Assert.Equal("content-b", b.Content); - Assert.Equal("hist", b.Metadata["cat"]); - Assert.Equal(new float[] { 4, 5, 6 }, b.Embedding); - Assert.False(restored.TryGetByDocumentId("a", out _)); - } - - [Fact] - public void FromJson_EmptyOrNull_ReturnsEmptySidecar() - { - Assert.Equal(0, FaissSidecar.FromJson("").Count); - Assert.Equal(0, FaissSidecar.FromJson(" ").Count); - } - - [Theory] - [InlineData(5, 4, 100, false, 5)] // no filter: fetch exactly topK - [InlineData(5, 4, 100, true, 20)] // filter: over-fetch topK * oversample - [InlineData(5, 4, 12, true, 12)] // over-fetch clamped to total docs - [InlineData(5, 4, 3, false, 3)] // topK clamped to total docs - [InlineData(0, 4, 100, true, 0)] // topK <= 0 - [InlineData(5, 4, 0, true, 0)] // empty index - public void ComputeFetchCount_AppliesOversampleAndClamp(int topK, int oversample, int total, bool hasFilters, int expected) - { - Assert.Equal(expected, FaissRetrievalPlanner.ComputeFetchCount(topK, oversample, total, hasFilters)); - } - - [Fact] - public void SelectTopK_OverFetchesThenFiltersAndTakesTopK_PreservingOrder() - { - // Candidates already in FAISS best-first order; only some pass the metadata filter. - var ranked = new List<(FaissSidecarEntry Entry, double RawScore)> - { - (Entry("d1", 2024), 0.99), - (Entry("d2", 2019), 0.95), // filtered out (year != 2024) - (Entry("d3", 2024), 0.90), - (Entry("d4", 2024), 0.80), - }; - var filters = new Dictionary { ["year"] = 2024 }; - - var result = FaissRetrievalPlanner.SelectTopK( - ranked, - filters, - topK: 2, - scoreConverter: d => d, - matches: EqualityMatch); - - Assert.Equal(2, result.Count); - Assert.Equal("d1", result[0].Id); // order preserved, d2 skipped - Assert.Equal("d3", result[1].Id); - Assert.Equal(0.99, result[0].RelevanceScore, 5); - Assert.True(result[0].HasRelevanceScore); - } - - [Fact] - public void SelectTopK_NoFilters_ReturnsAllUpToTopK() - { - var ranked = new List<(FaissSidecarEntry Entry, double RawScore)> - { - (Entry("d1", 2024), 0.9), - (Entry("d2", 2019), 0.8), - }; - - var result = FaissRetrievalPlanner.SelectTopK( - ranked, - new Dictionary(), - topK: 10, - scoreConverter: d => d, - matches: EqualityMatch); - - Assert.Equal(new[] { "d1", "d2" }, result.Select(r => r.Id).ToArray()); - } - - private static FaissSidecarEntry Entry(string id, int year) => new() - { - DocumentId = id, - Content = "content-" + id, - Metadata = new Dictionary { ["year"] = year } - }; - - private static bool EqualityMatch(Dictionary metadata, Dictionary filters) - => filters.All(kv => metadata.TryGetValue(kv.Key, out var v) && Equals(v, kv.Value)); -} -#endif diff --git a/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/AnnVectorIndexTests.cs b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/AnnVectorIndexTests.cs new file mode 100644 index 0000000000..1877bfd73b --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/RetrievalAugmentedGeneration/VectorSearch/Indexes/AnnVectorIndexTests.cs @@ -0,0 +1,128 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using AiDotNet.LinearAlgebra; +using AiDotNet.RetrievalAugmentedGeneration.VectorSearch.Indexes; +using Xunit; + +namespace AiDotNetTests.UnitTests.RetrievalAugmentedGeneration.VectorSearch.Indexes +{ + /// + /// Tests for — the dependency-free native ANN index (Flat/IVF/PQ/IVFPQ) on the + /// AiDotNet Tensors stack that replaces the external FaissNet backend. Validated on separable synthetic + /// clusters so the expected nearest neighbour is unambiguous, plus the incremental Add/Remove/Clear contract. + /// + public class AnnVectorIndexTests + { + private const int Dim = 16; + private const int Clusters = 8; + private const int PerCluster = 40; + + // Deterministic well-separated clusters: cluster c centred at a far-apart lattice point. + private static (List<(string id, Vector v)> data, Vector[] centers) MakeData(int seed) + { + var rng = new Random(seed); + var data = new List<(string, Vector)>(); + var centers = new Vector[Clusters]; + for (int c = 0; c < Clusters; c++) + { + var center = new double[Dim]; + for (int d = 0; d < Dim; d++) center[d] = c * 50.0 + d; + centers[c] = new Vector(center); + for (int p = 0; p < PerCluster; p++) + { + var v = new double[Dim]; + for (int d = 0; d < Dim; d++) v[d] = center[d] + (rng.NextDouble() - 0.5) * 2.0; + data.Add(($"c{c}_p{p}", new Vector(v))); + } + } + return (data, centers); + } + + private static int ClusterOf(string id) => int.Parse(id.Substring(1, id.IndexOf('_') - 1)); + + [Theory] + [InlineData(AnnVectorIndexType.Flat)] + [InlineData(AnnVectorIndexType.Ivf)] + [InlineData(AnnVectorIndexType.Pq)] + [InlineData(AnnVectorIndexType.IvfPq)] + public void Search_ReturnsNeighboursFromTheQueryCluster(AnnVectorIndexType type) + { + var (data, centers) = MakeData(seed: 7); + var index = new AnnVectorIndex(type, Dim, AnnVectorMetric.L2, nlist: Clusters, nprobe: 3, m: 4, ksub: 32); + foreach (var (id, v) in data) index.Add(id, v); + + Assert.Equal(data.Count, index.Count); + + int target = 5; + var hits = index.Search(centers[target], 5); + Assert.NotEmpty(hits); + Assert.Equal(target, ClusterOf(hits[0].Id)); // exact/approx top-1 in the query cluster + Assert.True(hits.Count(h => ClusterOf(h.Id) == target) >= 3); // majority of top-5 in-cluster + } + + [Fact] + public void Flat_ReturnsStoredVectorAsItsOwnNearestNeighbour() + { + var (data, _) = MakeData(seed: 11); + var index = new AnnVectorIndex(AnnVectorIndexType.Flat, Dim, AnnVectorMetric.L2); + foreach (var (id, v) in data) index.Add(id, v); + + var (probeId, probeVec) = data[123]; + var hits = index.Search(probeVec, 1); + Assert.Single(hits); + Assert.Equal(probeId, hits[0].Id); + } + + [Fact] + public void Remove_ExcludesTheDocumentAndDecrementsCount() + { + var (data, _) = MakeData(seed: 3); + var index = new AnnVectorIndex(AnnVectorIndexType.Flat, Dim, AnnVectorMetric.L2); + foreach (var (id, v) in data) index.Add(id, v); + + var (probeId, probeVec) = data[200]; + Assert.True(index.Remove(probeId)); + Assert.Equal(data.Count - 1, index.Count); + Assert.False(index.Remove(probeId)); // idempotent + + var hits = index.Search(probeVec, 5); + Assert.DoesNotContain(hits, h => h.Id == probeId); + } + + [Fact] + public void Clear_EmptiesTheIndex() + { + var (data, _) = MakeData(seed: 1); + var index = new AnnVectorIndex(AnnVectorIndexType.Ivf, Dim, AnnVectorMetric.L2, nlist: Clusters); + foreach (var (id, v) in data) index.Add(id, v); + Assert.True(index.Count > 0); + + index.Clear(); + Assert.Equal(0, index.Count); + Assert.Empty(index.Search(data[0].v, 5)); + } + + [Fact] + public void Cosine_NormalizesSoDirectionWinsOverMagnitude() + { + var index = new AnnVectorIndex(AnnVectorIndexType.Flat, dimension: 3, metric: AnnVectorMetric.Cosine); + index.Add("a", new Vector(new[] { 1.0, 0.0, 0.0 })); + index.Add("b", new Vector(new[] { 0.0, 1.0, 0.0 })); + + // Query points along +x with a large magnitude; cosine must pick "a" regardless of length. + var hits = index.Search(new Vector(new[] { 100.0, 1.0, 0.0 }), 1); + Assert.Single(hits); + Assert.Equal("a", hits[0].Id); + } + + [Fact] + public void AddBatch_AddsAllVectors() + { + var (data, _) = MakeData(seed: 9); + var index = new AnnVectorIndex(AnnVectorIndexType.Flat, Dim, AnnVectorMetric.L2); + index.AddBatch(data.ToDictionary(d => d.id, d => d.v)); + Assert.Equal(data.Count, index.Count); + } + } +}