From 179e648a26ea2d5c031f6ca04ea974226e8fe68f Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 23:25:06 +0000 Subject: [PATCH 1/2] LightRAG: inherit GraphRAG retrieval Config when created via lightRAG() An engine created through GraphRAG.lightRAG() now mirrors the owning instance's Config so the LightRAG path stays consistent with GraphRAG.ask/search on the same graph: - topKResults becomes the engine's defaultTopK, used when retrieve/ask are called without an explicit topK (signature is now topK: Int? = nil). - similarityThreshold is applied to semantic hits, and approach == "semantic" drops BM25 (includeKeyword: false), via a new LightRAGSearchOptions threaded into the stores. - approach == "keyword" stays rejected (LightRAG needs embeddings). Constructing LightRAGEngine directly still uses permissive defaults (full hybrid, no threshold, topK 10), so the subsystem remains usable standalone. Tests: topK inheritance (defaultTopK + capped results) and semantic-approach smoke (includeKeyword false, runs without throwing). README documents the inheritance. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- README.md | 8 +++ Sources/GraphRAG/GraphRAG/Engine.swift | 13 +++- .../GraphRAG/LightRAG/LightRAGEngine.swift | 57 +++++++++++------- Sources/GraphRAG/LightRAG/Searchers.swift | 59 +++++++++++++------ Tests/GraphRAGTests/LightRAGTests.swift | 25 ++++++++ 5 files changed, 121 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 0728662..0711771 100644 --- a/README.md +++ b/README.md @@ -132,6 +132,14 @@ let results = try await engine.retrieve("...", topK: 10) print(results.highLevelChunks, results.lowLevelChunks, results.mergedChunks) ``` +An engine created via `rag.lightRAG()` inherits that instance's retrieval +settings from `Config`, so it stays consistent with `rag.ask`/`rag.search`: it +honors `topKResults` (the default when `retrieve`/`ask` are called without an +explicit `topK`), applies `similarityThreshold` to semantic hits, and drops BM25 +when `approach == "semantic"`. `approach == "keyword"` is rejected — LightRAG's +dual-level design requires embeddings. Constructing `LightRAGEngine` directly +uses permissive defaults (full hybrid, no threshold, `topK` 10). + ## Testing ```bash diff --git a/Sources/GraphRAG/GraphRAG/Engine.swift b/Sources/GraphRAG/GraphRAG/Engine.swift index a14b20c..a1070c2 100644 --- a/Sources/GraphRAG/GraphRAG/Engine.swift +++ b/Sources/GraphRAG/GraphRAG/Engine.swift @@ -263,13 +263,22 @@ public actor GraphRAG { // nil on purpose, so reject the mismatch explicitly rather than forcing // corpus-wide re-embedding — which would also throw outright if this // instance's embedder is remote/unavailable. - guard self.config.approach.lowercased() != "keyword" else { + let approach = self.config.approach.lowercased() + guard approach != "keyword" else { throw GraphRAGError.validation( message: "LightRAG requires embeddings and is unavailable with approach == \"keyword\"") } + // Inherit this instance's retrieval settings so the LightRAG path stays + // consistent with `ask`/`search` on the same configured graph: honor the + // similarity threshold, suppress BM25 for `approach == "semantic"`, and + // default `topK` to the configured result cap. + let searchOptions = LightRAGSearchOptions( + semanticThreshold: self.config.similarityThreshold, + includeKeyword: approach != "semantic") return LightRAGEngine( graph: graph, embedder: embedder, languageModel: languageModel, - config: config, keywordConfig: keywordConfig, leidenConfig: leidenConfig) + config: config, keywordConfig: keywordConfig, leidenConfig: leidenConfig, + searchOptions: searchOptions, defaultTopK: self.config.topKResults) } /// Persist the knowledge graph to JSON. diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift index 4f55312..3e48934 100644 --- a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -14,16 +14,21 @@ private actor LightRAGStoreCache { private let graph: KnowledgeGraph private let embedder: any EmbeddingModel private let leidenConfig: LeidenConfig + private let searchOptions: LightRAGSearchOptions private var communitiesCache: CommunityDetectionResult? /// The single in-flight (or completed) build. Stored *before* the first /// await so concurrent first callers await the same task instead of each /// starting a duplicate O(corpus) embedding pass (actor reentrancy). private var storesTask: Task? - init(graph: KnowledgeGraph, embedder: any EmbeddingModel, leidenConfig: LeidenConfig) { + init( + graph: KnowledgeGraph, embedder: any EmbeddingModel, leidenConfig: LeidenConfig, + searchOptions: LightRAGSearchOptions + ) { self.graph = graph self.embedder = embedder self.leidenConfig = leidenConfig + self.searchOptions = searchOptions } func communities() -> CommunityDetectionResult { @@ -38,10 +43,12 @@ private actor LightRAGStoreCache { let detected = communities() let embedder = self.embedder let graph = self.graph + let options = self.searchOptions let task = Task { () throws -> Stores in - async let low = LightRAG.chunkSearcher(graph: graph, embedder: embedder) + async let low = LightRAG.chunkSearcher( + graph: graph, embedder: embedder, options: options) async let high = LightRAG.communitySearcher( - graph: graph, communities: detected, embedder: embedder) + graph: graph, communities: detected, embedder: embedder, options: options) return (low: try await low, high: try await high) } storesTask = task @@ -67,6 +74,12 @@ public struct LightRAGEngine: Sendable { public let config: DualRetrievalConfig public let keywordConfig: KeywordExtractorConfig public let leidenConfig: LeidenConfig + /// Store-level retrieval settings (semantic threshold, keyword toggle), + /// populated from the owning GraphRAG's `Config` when created via `lightRAG()`. + public let searchOptions: LightRAGSearchOptions + /// Default result cap used when `retrieve`/`ask` are called without an explicit + /// `topK`; set from `Config.topKResults` via `lightRAG()`. + public let defaultTopK: Int private let cache: LightRAGStoreCache public init( @@ -75,7 +88,9 @@ public struct LightRAGEngine: Sendable { languageModel: (any LanguageModel)? = nil, config: DualRetrievalConfig = DualRetrievalConfig(), keywordConfig: KeywordExtractorConfig = KeywordExtractorConfig(), - leidenConfig: LeidenConfig = LeidenConfig() + leidenConfig: LeidenConfig = LeidenConfig(), + searchOptions: LightRAGSearchOptions = LightRAGSearchOptions(), + defaultTopK: Int = 10 ) { self.graph = graph self.embedder = embedder @@ -83,8 +98,11 @@ public struct LightRAGEngine: Sendable { self.config = config self.keywordConfig = keywordConfig self.leidenConfig = leidenConfig + self.searchOptions = searchOptions + self.defaultTopK = defaultTopK self.cache = LightRAGStoreCache( - graph: graph, embedder: embedder, leidenConfig: leidenConfig) + graph: graph, embedder: embedder, leidenConfig: leidenConfig, + searchOptions: searchOptions) } /// Detect entity communities via Leiden. @@ -95,22 +113,17 @@ public struct LightRAGEngine: Sendable { /// Run dual-level retrieval for `query`. The two stores are built once per /// engine (cached across calls) rather than rebuilt each query. /// - /// Design note: `topK` is LightRAG's own per-call parameter and is - /// intentionally *not* defaulted from the owning GraphRAG's - /// `Config.topKResults` — LightRAG is configured independently of the - /// hybrid-path Config; pass `topK` explicitly to cap results. Likewise, an - /// empty/degenerate query is handled by the retriever returning no hits: the - /// stores are built once and cached, so an empty *first* query at most - /// triggers that one build early rather than per-query waste. (Only the - /// free config checks below — nonpositive topK, zero keyword budget — are - /// short-circuited ahead of the cache, since they need no keyword extraction.) - public func retrieve(_ query: String, topK: Int = 10) async throws -> DualRetrievalResults { + /// `topK` defaults to `defaultTopK` when omitted — which `GraphRAG.lightRAG()` + /// sets from `Config.topKResults`, so the LightRAG path honors the same result + /// cap as `GraphRAG.ask`/`search` on the configured instance. + public func retrieve(_ query: String, topK: Int? = nil) async throws -> DualRetrievalResults { + let effectiveTopK = topK ?? defaultTopK // Requests that can't produce hits — a nonpositive topK, or a keyword // budget of 0 (which forces empty high/low queries) — return before // building/embedding the stores, so a no-op never triggers corpus-wide // embedding work. (maxKeywords <= 0 is checked here directly, so no LLM // keyword call is made either.) - guard topK > 0, keywordConfig.maxKeywords > 0 else { + guard effectiveTopK > 0, keywordConfig.maxKeywords > 0 else { return DualRetrievalResults( highLevelChunks: [], lowLevelChunks: [], mergedChunks: [], keywords: DualLevelKeywords()) @@ -122,13 +135,15 @@ public struct LightRAGEngine: Sendable { highLevelStore: stores.high, lowLevelStore: stores.low, config: config) - return try await retriever.retrieve(query, topK: topK) + return try await retriever.retrieve(query, topK: effectiveTopK) } /// Answer `query` using dual-level retrieval, synthesizing with the LLM when - /// available and falling back to an extractive summary otherwise. - public func ask(_ query: String, topK: Int = 10) async throws -> Answer { - let results = try await retrieve(query, topK: topK) + /// available and falling back to an extractive summary otherwise. `topK` + /// defaults to `defaultTopK` (see `retrieve`). + public func ask(_ query: String, topK: Int? = nil) async throws -> Answer { + let effectiveTopK = topK ?? defaultTopK + let results = try await retrieve(query, topK: effectiveTopK) guard !results.mergedChunks.isEmpty else { return Answer( text: "I don't have enough information to answer this question.", confidence: 0) @@ -146,7 +161,7 @@ public struct LightRAGEngine: Sendable { .flatMap(\.sourceChunks) .filter { seenSource.insert($0).inserted } .map { ChunkID($0) } - let confidence = min(1.0, Float(results.mergedChunks.count) / Float(max(1, topK))) + let confidence = min(1.0, Float(results.mergedChunks.count) / Float(max(1, effectiveTopK))) if let languageModel, await languageModel.isAvailable() { let prompt = Prompts.fill(Prompts.answerGeneration, ["context": context, "query": query]) diff --git a/Sources/GraphRAG/LightRAG/Searchers.swift b/Sources/GraphRAG/LightRAG/Searchers.swift index b5eab22..aed7797 100644 --- a/Sources/GraphRAG/LightRAG/Searchers.swift +++ b/Sources/GraphRAG/LightRAG/Searchers.swift @@ -4,6 +4,22 @@ import Foundation +/// Retrieval knobs a LightRAG store honors. When the engine is created via +/// `GraphRAG.lightRAG()`, these mirror the owning instance's `Config` so LightRAG +/// stays consistent with `GraphRAG.ask`/`search` on the same graph. Constructed +/// directly, LightRAG uses the permissive defaults (full hybrid, no threshold). +public struct LightRAGSearchOptions: Sendable { + /// Minimum cosine score for a semantic hit (0 keeps all positive matches). + public var semanticThreshold: Float + /// Whether BM25 keyword scoring contributes (`false` = semantic-only). + public var includeKeyword: Bool + + public init(semanticThreshold: Float = 0, includeKeyword: Bool = true) { + self.semanticThreshold = semanticThreshold + self.includeKeyword = includeKeyword + } +} + /// A `SemanticSearcher` backed by a `HybridRetriever` over a fixed corpus. /// /// Each indexed document carries the real chunk ids it stands for @@ -13,22 +29,26 @@ public struct InMemorySemanticSearcher: SemanticSearcher { private let retriever: HybridRetriever private let embedder: any EmbeddingModel private let sourceChunksByID: [String: [String]] + private let options: LightRAGSearchOptions private init( retriever: HybridRetriever, embedder: any EmbeddingModel, - sourceChunksByID: [String: [String]] + sourceChunksByID: [String: [String]], options: LightRAGSearchOptions ) { self.retriever = retriever self.embedder = embedder self.sourceChunksByID = sourceChunksByID + self.options = options } /// Build a searcher, reusing each document's precomputed embedding when one /// is supplied and embedding on demand otherwise. `sourceChunks` records the - /// real chunk ids each document grounds to. + /// real chunk ids each document grounds to; `options` carries the semantic + /// threshold / keyword toggle applied at search time. public static func build( documents: [(id: String, content: String, embedding: [Float]?, sourceChunks: [String])], - embedder: any EmbeddingModel + embedder: any EmbeddingModel, + options: LightRAGSearchOptions = LightRAGSearchOptions() ) async throws -> InMemorySemanticSearcher { // Let the candidate pool cover the whole corpus so a large `topK` isn't // silently capped below the number of available documents. @@ -51,21 +71,20 @@ public struct InMemorySemanticSearcher: SemanticSearcher { sourceChunksByID[doc.id] = doc.sourceChunks } return InMemorySemanticSearcher( - retriever: retriever, embedder: embedder, sourceChunksByID: sourceChunksByID) + retriever: retriever, embedder: embedder, + sourceChunksByID: sourceChunksByID, options: options) } public func search(_ query: String, topK: Int) async throws -> [LightRAGResult] { - // Design note: LightRAG is a self-contained retrieval subsystem with its - // own configuration (DualRetrievalConfig / KeywordExtractorConfig). It - // uses full hybrid (BM25 + cosine) search with the default similarity - // threshold on purpose and does NOT inherit the owning GraphRAG's - // hybrid-path Config (`approach`, `similarityThreshold`) — those govern - // `GraphRAG.ask`/`search`, a different retrieval path. This keeps the two - // strategies independent; a caller who wants semantic-only or a stricter - // threshold uses the GraphRAG path (or we can thread those settings in if - // that inheritance is desired). + // Apply the inherited retrieval settings: the semantic threshold filters + // weak cosine hits, and `includeKeyword == false` (from `approach == + // "semantic"`) suppresses BM25 so results match the owning GraphRAG's + // configured behavior. Defaults keep full hybrid with no threshold. let queryEmbedding = try await embedder.embed(query) - let hits = retriever.search(query: query, queryEmbedding: queryEmbedding, limit: topK) + let hits = retriever.search( + query: query, queryEmbedding: queryEmbedding, limit: topK, + semanticThreshold: options.semanticThreshold, + includeKeyword: options.includeKeyword) return hits.map { LightRAGResult( id: $0.id, content: $0.content, score: $0.score, @@ -79,20 +98,23 @@ public enum LightRAG { /// Low-level (entity/detail) store: the document chunks themselves. Each /// chunk grounds to itself. public static func chunkSearcher( - graph: KnowledgeGraph, embedder: any EmbeddingModel + graph: KnowledgeGraph, embedder: any EmbeddingModel, + options: LightRAGSearchOptions = LightRAGSearchOptions() ) async throws -> InMemorySemanticSearcher { // Reuse the embeddings computed during the graph build, when present. let documents = graph.chunks.map { (id: $0.id.raw, content: $0.content, embedding: $0.embedding, sourceChunks: [$0.id.raw]) } - return try await InMemorySemanticSearcher.build(documents: documents, embedder: embedder) + return try await InMemorySemanticSearcher.build( + documents: documents, embedder: embedder, options: options) } /// High-level (theme/global) store: one short summary per detected community. /// A community grounds to the real chunks that mention its member entities, /// so answers built from high-level hits still cite actual evidence chunks. public static func communitySearcher( - graph: KnowledgeGraph, communities: CommunityDetectionResult, embedder: any EmbeddingModel + graph: KnowledgeGraph, communities: CommunityDetectionResult, embedder: any EmbeddingModel, + options: LightRAGSearchOptions = LightRAGSearchOptions() ) async throws -> InMemorySemanticSearcher { // Summaries are derived text with no precomputed embedding — embed on demand. let documents = communities.communities.map { community in @@ -103,7 +125,8 @@ public enum LightRAG { sourceChunks: communitySourceChunks(community, graph: graph) ) } - return try await InMemorySemanticSearcher.build(documents: documents, embedder: embedder) + return try await InMemorySemanticSearcher.build( + documents: documents, embedder: embedder, options: options) } /// A short textual theme for a community: member names plus the relationship diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift index ac36487..38c6b04 100644 --- a/Tests/GraphRAGTests/LightRAGTests.swift +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -153,6 +153,31 @@ private func triangleGraph() -> KnowledgeGraph { } } +@Test func lightRAGInheritsConfiguredTopK() async throws { + // An engine created via lightRAG() defaults topK to the instance's + // Config.topKResults, so ask/retrieve without an explicit topK honor the cap. + let rag = try GraphRAGBuilder().withChunkSize(24).withChunkOverlap(0).withTopK(1).build() + await rag.addDocument(text: "Alpha region. Bravo region. Charlie region. Delta region.") + try await rag.build() + let engine = try await rag.lightRAG() + #expect(engine.defaultTopK == 1) + let results = try await engine.retrieve("Alpha Bravo Charlie Delta") + #expect(results.mergedChunks.count <= 1) +} + +@Test func lightRAGSemanticApproachRuns() async throws { + // approach == "semantic" suppresses BM25 in the LightRAG stores; the engine + // must still build and answer without throwing (smoke over the inherited path). + let rag = try GraphRAGBuilder().withApproach("semantic").withTopK(3).build() + await rag.addDocument( + text: "Ada Lovelace collaborated with Charles Babbage on the Analytical Engine.") + try await rag.build() + let engine = try await rag.lightRAG() + #expect(!engine.searchOptions.includeKeyword) + let answer = try await engine.ask("Who worked on the Analytical Engine?") + #expect(answer.confidence >= 0) +} + @Test func dualRetrievalNonPositiveTopKReturnsEmpty() async throws { let high = MockSearcher(results: [LightRAGResult(id: "h", content: "h", score: 1)]) let low = MockSearcher(results: [LightRAGResult(id: "l", content: "l", score: 1)]) From 41bc8492fe746c4e3ad83d12666f259b0764d018 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 23:28:35 +0000 Subject: [PATCH 2/2] Build LightRAG stores in a detached task (Gemini review) Run the store build (community summary generation + on-demand embedding) on the global concurrent pool via Task.detached instead of a task inheriting the cache actor's executor, so the actor stays responsive to concurrent callers. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/LightRAG/LightRAGEngine.swift | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift index 3e48934..1183ace 100644 --- a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -44,7 +44,10 @@ private actor LightRAGStoreCache { let embedder = self.embedder let graph = self.graph let options = self.searchOptions - let task = Task { () throws -> Stores in + // Detached so the build (summary generation + on-demand embedding) runs + // on the global pool rather than the cache actor's executor, keeping the + // actor responsive to concurrent callers. + let task = Task.detached { () throws -> Stores in async let low = LightRAG.chunkSearcher( graph: graph, embedder: embedder, options: options) async let high = LightRAG.communitySearcher(