From 20408af629b610aba3f83bd8a0e3373e52acc796 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 15:27:49 +0000 Subject: [PATCH 01/12] Add Leiden community detection and LightRAG dual-level retrieval Extend the port with two optional subsystems from graphrag-rs: - Leiden community detection (Graph/Leiden.swift): greedy modularity local-moving plus a refinement pass that splits internally disconnected communities. Made deterministic (stable node ordering) and weighted (uses relationship confidence as edge weight); reports Newman modularity. Only result-affecting config is exposed. - LightRAG dual-level retrieval (LightRAG/): extract high-/low-level keywords from a query via the LLM with a deterministic offline fallback, search a low-level chunk store and a high-level per-community summary store, and merge (interleave/highFirst/lowFirst/weighted) with dedup by id. LightRAGEngine ties it together with detectCommunities/retrieve/ask. Wire both into GraphRAG via detectCommunities() and lightRAG(), document them in the README, and add unit tests covering community separation, determinism, keyword extraction/fallback, merge strategies, and an offline end-to-end ask. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- README.md | 54 +++- Sources/GraphRAG/Graph/Leiden.swift | 259 ++++++++++++++++++ Sources/GraphRAG/GraphRAG/Engine.swift | 17 ++ Sources/GraphRAG/LightRAG/DualRetrieval.swift | 151 ++++++++++ Sources/GraphRAG/LightRAG/Keywords.swift | 141 ++++++++++ .../GraphRAG/LightRAG/LightRAGEngine.swift | 83 ++++++ Sources/GraphRAG/LightRAG/Searchers.swift | 77 ++++++ Tests/GraphRAGTests/LightRAGTests.swift | 148 ++++++++++ 8 files changed, 925 insertions(+), 5 deletions(-) create mode 100644 Sources/GraphRAG/Graph/Leiden.swift create mode 100644 Sources/GraphRAG/LightRAG/DualRetrieval.swift create mode 100644 Sources/GraphRAG/LightRAG/Keywords.swift create mode 100644 Sources/GraphRAG/LightRAG/LightRAGEngine.swift create mode 100644 Sources/GraphRAG/LightRAG/Searchers.swift create mode 100644 Tests/GraphRAGTests/LightRAGTests.swift diff --git a/README.md b/README.md index a073aa7..35b1ba5 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,8 @@ summary of the top chunks. | Extraction | `PatternEntityExtractor`, `LLMEntityExtractor`, `Prompts` | | Embeddings | `HashEmbedder` (offline, deterministic), `OllamaEmbedder` | | LLM | `OllamaClient` | +| Communities | `LeidenCommunityDetector` (weighted, deterministic), `Community` | +| LightRAG | `LightRAGEngine`, `DualLevelRetriever`, `KeywordExtractor`, `SemanticSearcher` | | Orchestration | `GraphRAG` (actor), `GraphRAGBuilder`, `Config` | ## Design notes / port fidelity @@ -84,10 +86,51 @@ summary of the top chunks. - **Unicode safety**: the Rust chunker works on UTF-8 byte offsets guarded by `is_char_boundary`. This port operates on `Character` (grapheme) arrays, which are always valid boundaries; sizes and offsets are measured in characters. -- **Scope**: this is the portable core pipeline. The Rust workspace's - server/WASM/CLI crates and heavier optional subsystems (LightRAG, ROGRAG, - Leiden communities, distributed caching, persistence backends) are out of - scope for this port. +- **Scope**: this is the portable core pipeline plus the LightRAG dual-level + retrieval and Leiden community-detection subsystems (see below). The Rust + workspace's server/WASM/CLI crates and other optional subsystems (ROGRAG, + distributed caching, persistence backends) remain out of scope for this port. + +## Community detection (Leiden) + +`LeidenCommunityDetector` partitions the knowledge graph into communities via +greedy modularity local-moving plus a refinement pass that splits internally +disconnected communities. It ports the structure of the Rust crate's +single-level Leiden, but makes it **deterministic** (stable node ordering, so +repeated runs give identical assignments) and **weighted** — it uses each +relationship's `confidence` as an edge weight, which the Rust version ignored. + +```swift +let graph = await rag.knowledgeGraph() +let result = LeidenCommunityDetector().detect(graph) +for community in result.communities { + print("community \(community.id): \(community.members.count) members") +} +print("modularity:", result.modularity) // Newman modularity of the partition +``` + +Only knobs that affect the result are exposed via `LeidenConfig`: `resolution` +(higher → more, smaller communities), `maxIterations`, and `minModularityGain`. + +## Dual-level retrieval (LightRAG) + +`LightRAGEngine` answers queries by searching two levels at once: a **low-level** +store over document chunks (entity/detail-centric) and a **high-level** store +over per-community theme summaries derived from Leiden (global/relationship- +centric). A `KeywordExtractor` splits the query into high- and low-level +keywords using the LLM (with a deterministic offline fallback), each level is +searched independently, and the hits are merged — `interleave` (default), +`highFirst`, `lowFirst`, or `weighted`. + +```swift +let engine = await rag.lightRAG() +let answer = try await engine.ask("Who worked on the Analytical Engine?") +print(answer.text) + +// Or inspect both levels directly: +let results = try await engine.retrieve("...", topK: 10) +print(results.highLevelChunks, results.lowLevelChunks, results.mergedChunks) +``` ## Testing @@ -97,4 +140,5 @@ swift test The suite covers chunking, keyword extraction, BM25 ranking, cosine/vector search, the knowledge graph, PageRank, traversal, analytics, pattern extraction, -and the end-to-end offline build/ask pipeline. +Leiden community detection, LightRAG dual-level retrieval, and the end-to-end +offline build/ask pipeline. diff --git a/Sources/GraphRAG/Graph/Leiden.swift b/Sources/GraphRAG/Graph/Leiden.swift new file mode 100644 index 0000000..9f9f80a --- /dev/null +++ b/Sources/GraphRAG/Graph/Leiden.swift @@ -0,0 +1,259 @@ +// Leiden.swift +// Community detection over the knowledge graph, ported from graphrag-rs +// `graph::leiden`. +// +// The Rust implementation is a simplified single-level Leiden: greedy modularity +// local-moving followed by one refinement pass that splits internally +// disconnected communities. This port keeps that structure but makes it +// deterministic (stable node ordering) and *weighted* — it uses each +// relationship's confidence as an edge weight, since the `KnowledgeGraph` +// carries them (the Rust version ignored weights). Only configuration that +// actually affects the result is exposed. + +import Foundation + +/// Tunables for `LeidenCommunityDetector`. +public struct LeidenConfig: Sendable { + /// Higher resolution → more, smaller communities (default 1.0). + public var resolution: Double + /// Maximum local-moving passes (default 100). + public var maxIterations: Int + /// Stop when a pass improves modularity by less than this (default 1e-6). + public var minModularityGain: Double + + public init(resolution: Double = 1.0, maxIterations: Int = 100, minModularityGain: Double = 1e-6) { + self.resolution = resolution + self.maxIterations = maxIterations + self.minModularityGain = minModularityGain + } +} + +/// A detected community: a contiguous integer id and its member entities. +public struct Community: Sendable, Equatable, Identifiable { + public var id: Int + public var members: [EntityID] + + public init(id: Int, members: [EntityID]) { + self.id = id + self.members = members + } + + public var size: Int { members.count } +} + +/// The output of community detection. +public struct CommunityDetectionResult: Sendable { + /// Communities ordered by id, each with its members (in node order). + public var communities: [Community] + /// Map from entity to its community id. + public var assignment: [EntityID: Int] + /// Modularity of the final partition. + public var modularity: Double + + public init(communities: [Community], assignment: [EntityID: Int], modularity: Double) { + self.communities = communities + self.assignment = assignment + self.modularity = modularity + } + + public var communityCount: Int { communities.count } +} + +/// Weighted, deterministic Leiden community detection. +public struct LeidenCommunityDetector: Sendable { + public var config: LeidenConfig + + public init(config: LeidenConfig = LeidenConfig()) { + self.config = config + } + + public func detect(_ graph: KnowledgeGraph) -> CommunityDetectionResult { + let nodes = graph.entities.map(\.id) + let n = nodes.count + guard n > 0 else { + return CommunityDetectionResult(communities: [], assignment: [:], modularity: 0) + } + + var indexOf: [EntityID: Int] = [:] + for (i, id) in nodes.enumerated() { indexOf[id] = i } + + // Build an undirected, weighted adjacency (parallel edges summed, self + // loops dropped, non-positive weights skipped). + var adjacency: [[(node: Int, weight: Double)]] = Array(repeating: [], count: n) + var mergedWeight: [Int: [Int: Double]] = [:] + for rel in graph.relationships { + guard let s = indexOf[rel.source], let t = indexOf[rel.target], s != t else { continue } + let w = Double(rel.confidence) + guard w > 0 else { continue } + mergedWeight[s, default: [:]][t, default: 0] += w + mergedWeight[t, default: [:]][s, default: 0] += w + } + var degree = [Double](repeating: 0, count: n) + for (i, neighbors) in mergedWeight { + for (j, w) in neighbors { + adjacency[i].append((j, w)) + degree[i] += w + } + } + let twoM = degree.reduce(0, +) + + // No edges → every node is its own community. + guard twoM > 0 else { + return singletons(nodes) + } + + // Phase 1: greedy modularity local moving. + var communityOf = Array(0.. bestScore + 1e-12 { + bestScore = score + bestCommunity = comm + } + } + + communityOf[i] = bestCommunity + sigmaTot[bestCommunity] += ki + if bestCommunity != ci { moved = true } + } + + let currentModularity = modularity(communityOf, adjacency, degree, twoM) + if !moved || currentModularity - previousModularity < config.minModularityGain { + previousModularity = currentModularity + break + } + previousModularity = currentModularity + } + + // Phase 2: refinement — split internally disconnected communities. + refine(&communityOf, adjacency: adjacency, n: n) + + // Renumber to contiguous ids ordered by first member appearance. + return finalize(communityOf, nodes: nodes, adjacency: adjacency, degree: degree, twoM: twoM) + } + + // MARK: - Internals + + private func singletons(_ nodes: [EntityID]) -> CommunityDetectionResult { + var communities: [Community] = [] + var assignment: [EntityID: Int] = [:] + for (i, id) in nodes.enumerated() { + communities.append(Community(id: i, members: [id])) + assignment[id] = i + } + return CommunityDetectionResult(communities: communities, assignment: assignment, modularity: 0) + } + + private func refine( + _ communityOf: inout [Int], adjacency: [[(node: Int, weight: Double)]], n: Int + ) { + var members: [Int: [Int]] = [:] + for i in 0.. = [] + var components: [[Int]] = [] + for start in nodes where !visited.contains(start) { + var component: [Int] = [] + var queue = [start] + visited.insert(start) + var head = 0 + while head < queue.count { + let u = queue[head] + head += 1 + component.append(u) + for edge in adjacency[u] + where nodeSet.contains(edge.node) && !visited.contains(edge.node) { + visited.insert(edge.node) + queue.append(edge.node) + } + } + components.append(component) + } + // Keep the first component under the original id; give the rest new ids. + if components.count > 1 { + for componentIndex in 1.. CommunityDetectionResult { + var canonical: [Int: Int] = [:] + var finalOf = [Int](repeating: 0, count: nodes.count) + var nextId = 0 + for i in 0.. Double { + guard twoM > 0 else { return 0 } + var internalWeight: [Int: Double] = [:] // Σ_in per community (edges counted twice) + var totalDegree: [Int: Double] = [:] + for i in 0.. KnowledgeGraph { graph } + /// Detect entity communities in the current graph via Leiden. + public func detectCommunities(config: LeidenConfig = LeidenConfig()) -> CommunityDetectionResult { + LeidenCommunityDetector(config: config).detect(graph) + } + + /// A LightRAG dual-level retrieval engine over the current graph snapshot, + /// wired to this instance's embedder and (optional) language model. + public func lightRAG( + config: DualRetrievalConfig = DualRetrievalConfig(), + keywordConfig: KeywordExtractorConfig = KeywordExtractorConfig(), + leidenConfig: LeidenConfig = LeidenConfig() + ) -> LightRAGEngine { + LightRAGEngine( + graph: graph, embedder: embedder, languageModel: languageModel, + config: config, keywordConfig: keywordConfig, leidenConfig: leidenConfig) + } + /// Persist the knowledge graph to JSON. public func save(toJSON path: String) throws { try graph.save(toJSON: path) } diff --git a/Sources/GraphRAG/LightRAG/DualRetrieval.swift b/Sources/GraphRAG/LightRAG/DualRetrieval.swift new file mode 100644 index 0000000..47dbaf9 --- /dev/null +++ b/Sources/GraphRAG/LightRAG/DualRetrieval.swift @@ -0,0 +1,151 @@ +// DualRetrieval.swift +// Dual-level retrieval, ported from graphrag-rs `lightrag::dual_retrieval`. + +import Foundation + +/// A single retrieval hit at one abstraction level. +public struct LightRAGResult: Sendable, Equatable { + public var id: String + public var content: String + public var score: Float + public var entities: [String] + public var sourceChunks: [String] + + public init( + id: String, content: String, score: Float, + entities: [String] = [], sourceChunks: [String] = [] + ) { + self.id = id + self.content = content + self.score = score + self.entities = entities + self.sourceChunks = sourceChunks + } +} + +/// A store that can be searched semantically for one retrieval level. +public protocol SemanticSearcher: Sendable { + func search(_ query: String, topK: Int) async throws -> [LightRAGResult] +} + +/// How high- and low-level result sets are combined. +public enum MergeStrategy: Sendable, Equatable { + /// Alternate one high, one low, … (default). + case interleave + /// All high-level first, then low-level. + case highFirst + /// All low-level first, then high-level. + case lowFirst + /// Sort by level-weighted score, descending. + case weighted +} + +/// Configuration for `DualLevelRetriever`. +public struct DualRetrievalConfig: Sendable { + public var highLevelWeight: Float + public var lowLevelWeight: Float + public var mergeStrategy: MergeStrategy + + public init( + highLevelWeight: Float = 0.6, + lowLevelWeight: Float = 0.4, + mergeStrategy: MergeStrategy = .interleave + ) { + self.highLevelWeight = highLevelWeight + self.lowLevelWeight = lowLevelWeight + self.mergeStrategy = mergeStrategy + } +} + +/// The output of a dual-level retrieval. +public struct DualRetrievalResults: Sendable { + public var highLevelChunks: [LightRAGResult] + public var lowLevelChunks: [LightRAGResult] + public var mergedChunks: [LightRAGResult] + public var keywords: DualLevelKeywords + + public init( + highLevelChunks: [LightRAGResult], lowLevelChunks: [LightRAGResult], + mergedChunks: [LightRAGResult], keywords: DualLevelKeywords + ) { + self.highLevelChunks = highLevelChunks + self.lowLevelChunks = lowLevelChunks + self.mergedChunks = mergedChunks + self.keywords = keywords + } +} + +/// Runs LightRAG dual-level retrieval: extract high-/low-level keywords, search +/// each level's store, and merge the results. +public struct DualLevelRetriever: Sendable { + public var keywordExtractor: KeywordExtractor + public var highLevelStore: any SemanticSearcher + public var lowLevelStore: any SemanticSearcher + public var config: DualRetrievalConfig + + public init( + keywordExtractor: KeywordExtractor, + highLevelStore: any SemanticSearcher, + lowLevelStore: any SemanticSearcher, + config: DualRetrievalConfig = DualRetrievalConfig() + ) { + self.keywordExtractor = keywordExtractor + self.highLevelStore = highLevelStore + self.lowLevelStore = lowLevelStore + self.config = config + } + + public func retrieve(_ query: String, topK: Int = 10) async throws -> DualRetrievalResults { + let keywords = await keywordExtractor.extract(query) + + // Each level searches with its joined keywords as a single query. + let highQuery = keywords.highLevel.joined(separator: " ") + let lowQuery = keywords.lowLevel.joined(separator: " ") + let high = highQuery.isEmpty ? [] : try await highLevelStore.search(highQuery, topK: topK) + let low = lowQuery.isEmpty ? [] : try await lowLevelStore.search(lowQuery, topK: topK) + + let merged = merge(high: high, low: low, topK: topK) + return DualRetrievalResults( + highLevelChunks: high, lowLevelChunks: low, mergedChunks: merged, keywords: keywords) + } + + // MARK: - Merge + + private func merge(high: [LightRAGResult], low: [LightRAGResult], topK: Int) -> [LightRAGResult] { + guard topK > 0 else { return [] } + var seen: Set = [] + var merged: [LightRAGResult] = [] + + func take(_ result: LightRAGResult) { + guard merged.count < topK, seen.insert(result.id).inserted else { return } + merged.append(result) + } + + switch config.mergeStrategy { + case .interleave: + var i = 0 + while merged.count < topK && (i < high.count || i < low.count) { + if i < high.count { take(high[i]) } + if i < low.count { take(low[i]) } + i += 1 + } + case .highFirst: + for r in high { take(r) } + for r in low { take(r) } + case .lowFirst: + for r in low { take(r) } + for r in high { take(r) } + case .weighted: + let weighted = + high.map { ($0, $0.score * config.highLevelWeight) } + + low.map { ($0, $0.score * config.lowLevelWeight) } + for (result, _) in weighted.sorted(by: { lhs, rhs in + if lhs.1 == rhs.1 { return lhs.0.id < rhs.0.id } + return lhs.1 > rhs.1 + }) { + take(result) + } + } + return merged + } +} diff --git a/Sources/GraphRAG/LightRAG/Keywords.swift b/Sources/GraphRAG/LightRAG/Keywords.swift new file mode 100644 index 0000000..21cde2f --- /dev/null +++ b/Sources/GraphRAG/LightRAG/Keywords.swift @@ -0,0 +1,141 @@ +// Keywords.swift +// Dual-level keyword extraction, ported from graphrag-rs +// `lightrag::keyword_extraction`. + +import Foundation + +/// Keywords extracted at two levels of abstraction. +public struct DualLevelKeywords: Sendable, Equatable, Decodable { + /// Broad topics / concepts / themes. + public var highLevel: [String] + /// Specific entities / attributes / details. + public var lowLevel: [String] + + public init(highLevel: [String] = [], lowLevel: [String] = []) { + self.highLevel = highLevel + self.lowLevel = lowLevel + } + + enum CodingKeys: String, CodingKey { + case highLevel = "high_level" + case lowLevel = "low_level" + } + + public init(from decoder: Decoder) throws { + let c = try decoder.container(keyedBy: CodingKeys.self) + highLevel = (try? c.decode([String].self, forKey: .highLevel)) ?? [] + lowLevel = (try? c.decode([String].self, forKey: .lowLevel)) ?? [] + } + + public var total: Int { highLevel.count + lowLevel.count } + public var isEmpty: Bool { total == 0 } +} + +/// Configuration for `KeywordExtractor`. +public struct KeywordExtractorConfig: Sendable { + public var maxKeywords: Int + public var language: String + + public init(maxKeywords: Int = 20, language: String = "English") { + self.maxKeywords = maxKeywords + self.language = language + } +} + +/// Extracts high-/low-level keywords from a query using an LLM, with a +/// deterministic non-LLM fallback. +public struct KeywordExtractor: Sendable { + public let model: (any LanguageModel)? + public var config: KeywordExtractorConfig + + public init(model: (any LanguageModel)? = nil, config: KeywordExtractorConfig = KeywordExtractorConfig()) { + self.model = model + self.config = config + } + + /// Extract dual-level keywords. Never throws — any LLM/parse failure falls + /// back to a simple query tokenization. + public func extract(_ query: String) async -> DualLevelKeywords { + guard let model, await model.isAvailable() else { + return KeywordExtractor.fallback(query) + } + let prompt = buildPrompt(query) + do { + let response = try await model.complete(prompt) + if let parsed = KeywordExtractor.parse(response), !parsed.isEmpty { + return capped(parsed) + } + } catch { + // fall through to fallback + } + return KeywordExtractor.fallback(query) + } + + /// Cap the combined keyword count at `maxKeywords`, keeping high-level first. + private func capped(_ keywords: DualLevelKeywords) -> DualLevelKeywords { + guard keywords.total > config.maxKeywords else { return keywords } + var high = keywords.highLevel + var low = keywords.lowLevel + if high.count > config.maxKeywords { high = Array(high.prefix(config.maxKeywords)) } + let remaining = max(0, config.maxKeywords - high.count) + low = Array(low.prefix(remaining)) + return DualLevelKeywords(highLevel: high, lowLevel: low) + } + + func buildPrompt(_ query: String) -> String { + """ + Extract keywords at two levels from this query: "\(query)" + + Return JSON with this exact structure: + { + "high_level": ["theme1", "theme2", ...], + "low_level": ["entity1", "entity2", ...] + } + + Rules: + 1. high_level: Broader topics, concepts, themes (abstract level) + 2. low_level: Specific entities, attributes, details (concrete level) + 3. LIMIT: Maximum \(config.maxKeywords) total keywords combined + 4. NO duplication between levels + 5. Keep keywords concise (1-3 words each) + + Example 1: + Query: "How did Alice and Bob collaborate on the quantum computing project?" + { + "high_level": ["collaboration", "quantum computing", "teamwork"], + "low_level": ["Alice", "Bob", "project"] + } + + Example 2: + Query: "What are the main themes in the dataset?" + { + "high_level": ["themes", "patterns", "overview"], + "low_level": ["dataset"] + } + + Language: \(config.language) + + Now extract keywords: + """ + } + + /// Parse the JSON object between the first `{` and last `}` of the response. + static func parse(_ response: String) -> DualLevelKeywords? { + guard let first = response.firstIndex(of: "{"), + let last = response.lastIndex(of: "}"), first < last + else { return nil } + let slice = String(response[first...last]) + guard let data = slice.data(using: .utf8) else { return nil } + return try? JSONDecoder().decode(DualLevelKeywords.self, from: data) + } + + /// Deterministic fallback: query words of length >= 4, up to 10, as low-level. + static func fallback(_ query: String) -> DualLevelKeywords { + let words = + query + .split(whereSeparator: { $0.isWhitespace }) + .map { String($0.filter { $0.isLetter || $0.isNumber }) } + .filter { $0.count >= 4 } + return DualLevelKeywords(highLevel: [], lowLevel: Array(words.prefix(10))) + } +} diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift new file mode 100644 index 0000000..e227d5f --- /dev/null +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -0,0 +1,83 @@ +// LightRAGEngine.swift +// High-level LightRAG facade over a knowledge graph: detects communities, +// assembles the two dual-level stores, and answers queries. + +import Foundation + +/// A self-contained LightRAG engine over a snapshot of a `KnowledgeGraph`. +/// +/// The low-level store searches document chunks (entity/detail-centric); the +/// high-level store searches per-community theme summaries derived from Leiden +/// community detection (global/relationship-centric). +public struct LightRAGEngine: Sendable { + public let graph: KnowledgeGraph + private let embedder: any EmbeddingModel + private let languageModel: (any LanguageModel)? + public var config: DualRetrievalConfig + public var keywordConfig: KeywordExtractorConfig + public var leidenConfig: LeidenConfig + + public init( + graph: KnowledgeGraph, + embedder: any EmbeddingModel, + languageModel: (any LanguageModel)? = nil, + config: DualRetrievalConfig = DualRetrievalConfig(), + keywordConfig: KeywordExtractorConfig = KeywordExtractorConfig(), + leidenConfig: LeidenConfig = LeidenConfig() + ) { + self.graph = graph + self.embedder = embedder + self.languageModel = languageModel + self.config = config + self.keywordConfig = keywordConfig + self.leidenConfig = leidenConfig + } + + /// Detect entity communities via Leiden. + public func detectCommunities() -> CommunityDetectionResult { + LeidenCommunityDetector(config: leidenConfig).detect(graph) + } + + /// Run dual-level retrieval for `query`. + public func retrieve(_ query: String, topK: Int = 10) async throws -> DualRetrievalResults { + let communities = detectCommunities() + let lowStore = try await LightRAG.chunkSearcher(graph: graph, embedder: embedder) + let highStore = try await LightRAG.communitySearcher( + graph: graph, communities: communities, embedder: embedder) + let extractor = KeywordExtractor(model: languageModel, config: keywordConfig) + let retriever = DualLevelRetriever( + keywordExtractor: extractor, + highLevelStore: highStore, + lowLevelStore: lowStore, + config: config) + return try await retriever.retrieve(query, topK: topK) + } + + /// 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) + guard !results.mergedChunks.isEmpty else { + return Answer( + text: "I don't have enough information to answer this question.", confidence: 0) + } + + let context = results.mergedChunks.map { result in + "[Relevance: \(String(format: "%.3f", result.score))]\n\(result.content)" + }.joined(separator: "\n\n---\n\n") + let sources = results.mergedChunks.map { ChunkID($0.id) } + let confidence = min(1.0, Float(results.mergedChunks.count) / Float(max(1, topK))) + + if let languageModel, await languageModel.isAvailable() { + let prompt = Prompts.fill(Prompts.answerGeneration, ["context": context, "query": query]) + let raw = try await languageModel.complete(prompt) + return Answer( + text: GraphRAG.stripThinkingTags(raw), confidence: confidence, sources: sources) + } + + let extractive = results.mergedChunks.prefix(3).map(\.content).joined(separator: "\n\n") + return Answer( + text: "Based on the retrieved context:\n\n\(extractive)", + confidence: confidence, sources: sources) + } +} diff --git a/Sources/GraphRAG/LightRAG/Searchers.swift b/Sources/GraphRAG/LightRAG/Searchers.swift new file mode 100644 index 0000000..8d3676e --- /dev/null +++ b/Sources/GraphRAG/LightRAG/Searchers.swift @@ -0,0 +1,77 @@ +// Searchers.swift +// A concrete `SemanticSearcher` (BM25 + cosine over an in-memory corpus) and +// builders that turn a knowledge graph into the two LightRAG stores. + +import Foundation + +/// A `SemanticSearcher` backed by a `HybridRetriever` over a fixed corpus. +public struct InMemorySemanticSearcher: SemanticSearcher { + private let retriever: HybridRetriever + private let embedder: any EmbeddingModel + + private init(retriever: HybridRetriever, embedder: any EmbeddingModel) { + self.retriever = retriever + self.embedder = embedder + } + + /// Build a searcher by embedding and indexing each document. + public static func build( + documents: [(id: String, content: String)], embedder: any EmbeddingModel + ) async throws -> InMemorySemanticSearcher { + var retriever = HybridRetriever() + for doc in documents { + let embedding = try await embedder.embed(doc.content) + retriever.index(id: doc.id, content: doc.content, embedding: embedding) + } + return InMemorySemanticSearcher(retriever: retriever, embedder: embedder) + } + + public func search(_ query: String, topK: Int) async throws -> [LightRAGResult] { + let queryEmbedding = try await embedder.embed(query) + let hits = retriever.search(query: query, queryEmbedding: queryEmbedding, limit: topK) + return hits.map { + LightRAGResult(id: $0.id, content: $0.content, score: $0.score, sourceChunks: [$0.id]) + } + } +} + +/// Builders for the LightRAG stores from a `KnowledgeGraph`. +public enum LightRAG { + /// Low-level (entity/detail) store: the document chunks themselves. + public static func chunkSearcher( + graph: KnowledgeGraph, embedder: any EmbeddingModel + ) async throws -> InMemorySemanticSearcher { + let documents = graph.chunks.map { (id: $0.id.raw, content: $0.content) } + return try await InMemorySemanticSearcher.build(documents: documents, embedder: embedder) + } + + /// High-level (theme/global) store: one short summary per detected community. + public static func communitySearcher( + graph: KnowledgeGraph, communities: CommunityDetectionResult, embedder: any EmbeddingModel + ) async throws -> InMemorySemanticSearcher { + let documents = communities.communities.map { community in + (id: "community_\(community.id)", content: communitySummary(community, graph: graph)) + } + return try await InMemorySemanticSearcher.build(documents: documents, embedder: embedder) + } + + /// A short textual theme for a community: member names plus the relationship + /// types connecting them. + public static func communitySummary(_ community: Community, graph: KnowledgeGraph) -> String { + let names = community.members.compactMap { graph.entity($0)?.name } + let memberSet = Set(community.members) + var relationTypes: Set = [] + for member in community.members { + for (neighbor, relationship) in graph.neighbors(of: member) + where memberSet.contains(neighbor) { + relationTypes.insert(relationship.relationType) + } + } + var parts: [String] = [] + if !names.isEmpty { parts.append("Entities: " + names.joined(separator: ", ")) } + if !relationTypes.isEmpty { + parts.append("Relationships: " + relationTypes.sorted().joined(separator: ", ")) + } + return parts.joined(separator: ". ") + } +} diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift new file mode 100644 index 0000000..0663d46 --- /dev/null +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -0,0 +1,148 @@ +import Testing + +@testable import GraphRAG + +// MARK: - Test doubles + +private struct MockLLM: LanguageModel { + let response: String + func complete(_ prompt: String, params: GenerationParams) async throws -> String { response } + func isAvailable() async -> Bool { true } + var modelInfo: ModelInfo { ModelInfo(name: "mock") } +} + +private struct MockSearcher: SemanticSearcher { + let results: [LightRAGResult] + func search(_ query: String, topK: Int) async throws -> [LightRAGResult] { + Array(results.prefix(topK)) + } +} + +private func triangleGraph() -> KnowledgeGraph { + var graph = KnowledgeGraph() + let names = ["a1", "a2", "a3", "b1", "b2", "b3"] + for name in names { + graph.addEntity(Entity(id: EntityID(name), name: name, entityType: "X")) + } + func link(_ x: String, _ y: String) { + graph.addRelationship( + Relationship(source: EntityID(x), target: EntityID(y), relationType: "R", confidence: 1)) + } + // Two disconnected triangles. + link("a1", "a2"); link("a2", "a3"); link("a1", "a3") + link("b1", "b2"); link("b2", "b3"); link("b1", "b3") + return graph +} + +// MARK: - Leiden + +@Test func leidenSeparatesDisconnectedClusters() { + let result = LeidenCommunityDetector().detect(triangleGraph()) + #expect(result.communityCount == 2) + // Each triangle's nodes share a community. + let a = result.assignment[EntityID("a1")] + #expect(result.assignment[EntityID("a2")] == a) + #expect(result.assignment[EntityID("a3")] == a) + let b = result.assignment[EntityID("b1")] + #expect(result.assignment[EntityID("b2")] == b) + #expect(result.assignment[EntityID("b3")] == b) + #expect(a != b) + #expect(result.modularity > 0) +} + +@Test func leidenEmptyGraphIsEmpty() { + let result = LeidenCommunityDetector().detect(KnowledgeGraph()) + #expect(result.communityCount == 0) + #expect(result.modularity == 0) +} + +@Test func leidenNoEdgesYieldsSingletons() { + var graph = KnowledgeGraph() + for name in ["x", "y", "z"] { + graph.addEntity(Entity(id: EntityID(name), name: name, entityType: "T")) + } + let result = LeidenCommunityDetector().detect(graph) + #expect(result.communityCount == 3) +} + +@Test func leidenIsDeterministic() { + let graph = triangleGraph() + let a = LeidenCommunityDetector().detect(graph) + let b = LeidenCommunityDetector().detect(graph) + #expect(a.assignment == b.assignment) +} + +// MARK: - Keyword extraction + +@Test func keywordFallbackTokenizesQuery() async { + let extractor = KeywordExtractor(model: nil) + let keywords = await extractor.extract("Find the quantum computing project") + #expect(keywords.highLevel.isEmpty) + #expect(keywords.lowLevel.contains("quantum")) + #expect(!keywords.lowLevel.contains("the")) // shorter than 4 chars +} + +@Test func keywordLLMParsesDualLevels() async { + let json = #"{"high_level": ["collaboration"], "low_level": ["Alice", "Bob"]}"# + let extractor = KeywordExtractor(model: MockLLM(response: json)) + let keywords = await extractor.extract("How did Alice and Bob collaborate?") + #expect(keywords.highLevel == ["collaboration"]) + #expect(keywords.lowLevel == ["Alice", "Bob"]) +} + +// MARK: - Dual-level retrieval / merge + +@Test func dualRetrievalInterleavesAndDedups() async throws { + let high = MockSearcher(results: [ + LightRAGResult(id: "h1", content: "theme one", score: 0.9), + LightRAGResult(id: "shared", content: "shared", score: 0.5), + ]) + let low = MockSearcher(results: [ + LightRAGResult(id: "l1", content: "detail one", score: 0.8), + LightRAGResult(id: "shared", content: "shared", score: 0.7), + ]) + let extractor = KeywordExtractor( + model: MockLLM(response: #"{"high_level":["t"],"low_level":["d"]}"#)) + let retriever = DualLevelRetriever( + keywordExtractor: extractor, highLevelStore: high, lowLevelStore: low) + + let results = try await retriever.retrieve("query", topK: 5) + let ids = results.mergedChunks.map(\.id) + // Interleave: h1, l1, shared (deduped once), … no duplicate "shared". + #expect(ids.first == "h1") + #expect(ids.filter { $0 == "shared" }.count == 1) + #expect(Set(ids) == ["h1", "l1", "shared"]) +} + +@Test func dualRetrievalWeightedOrdersByWeightedScore() async throws { + let high = MockSearcher(results: [LightRAGResult(id: "h", content: "h", score: 1.0)]) + let low = MockSearcher(results: [LightRAGResult(id: "l", content: "l", score: 1.0)]) + let extractor = KeywordExtractor( + model: MockLLM(response: #"{"high_level":["t"],"low_level":["d"]}"#)) + // highLevelWeight 0.6 > lowLevelWeight 0.4, equal base scores → high first. + let retriever = DualLevelRetriever( + keywordExtractor: extractor, highLevelStore: high, lowLevelStore: low, + config: DualRetrievalConfig(mergeStrategy: .weighted)) + let results = try await retriever.retrieve("query", topK: 5) + #expect(results.mergedChunks.map(\.id) == ["h", "l"]) +} + +// MARK: - End-to-end via GraphRAG + +@Test func lightRAGEndToEndOffline() async throws { + let rag = try GraphRAGBuilder().withChunkSize(400).withChunkOverlap(50).withTopK(3).build() + await rag.addDocument( + text: """ + Ada Lovelace collaborated with Charles Babbage on the Analytical Engine, + an early mechanical general-purpose computer. + """) + try await rag.build() + + let engine = await rag.lightRAG() + let communities = engine.detectCommunities() + #expect(communities.communityCount >= 1) + + let answer = try await engine.ask("Who worked on the Analytical Engine?") + #expect(!answer.text.isEmpty) + #expect(!answer.sources.isEmpty) +} From d35f34d5497ecff5fefec2eea91df087a2c2063f Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 15:32:54 +0000 Subject: [PATCH 02/12] Address Gemini review: reuse embeddings, determinism, concurrency, robustness - Searchers: InMemorySemanticSearcher.build now reuses each chunk's precomputed embedding from the graph instead of re-embedding on every query (community summaries still embed on demand). Avoids the async-`??` autoclosure pitfall with an explicit branch. - Leiden: build adjacency by node index with sorted neighbors, and sum modularity in sorted community-id order, so neighbor ordering and floating-point sums are deterministic across runs. - Keywords: strip thinking tags before locating the JSON braces, so a `{` in a preamble can't derail parsing; deduplicate fallback keywords case-insensitively in first-seen order. - DualRetrieval / LightRAGEngine: run the two independent searches, and the two store builds, concurrently via async let. - LightRAGEngine: exclude virtual "community_" ids from grounding sources so they only contain real chunk ids. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/Graph/Leiden.swift | 13 ++++++--- Sources/GraphRAG/LightRAG/DualRetrieval.swift | 9 ++++--- Sources/GraphRAG/LightRAG/Keywords.swift | 15 +++++++---- .../GraphRAG/LightRAG/LightRAGEngine.swift | 14 +++++++--- Sources/GraphRAG/LightRAG/Searchers.swift | 27 +++++++++++++++---- 5 files changed, 58 insertions(+), 20 deletions(-) diff --git a/Sources/GraphRAG/Graph/Leiden.swift b/Sources/GraphRAG/Graph/Leiden.swift index 9f9f80a..8b2cd75 100644 --- a/Sources/GraphRAG/Graph/Leiden.swift +++ b/Sources/GraphRAG/Graph/Leiden.swift @@ -88,9 +88,13 @@ public struct LeidenCommunityDetector: Sendable { mergedWeight[s, default: [:]][t, default: 0] += w mergedWeight[t, default: [:]][s, default: 0] += w } + // Build adjacency by node index and sorted neighbor id so both the + // neighbor ordering and the floating-point degree summation are + // deterministic across runs (dictionary iteration order is randomized). var degree = [Double](repeating: 0, count: n) - for (i, neighbors) in mergedWeight { - for (j, w) in neighbors { + for i in 0.. DualRetrievalResults { let keywords = await keywordExtractor.extract(query) - // Each level searches with its joined keywords as a single query. + // Each level searches with its joined keywords as a single query. The two + // levels are independent, so run them concurrently. let highQuery = keywords.highLevel.joined(separator: " ") let lowQuery = keywords.lowLevel.joined(separator: " ") - let high = highQuery.isEmpty ? [] : try await highLevelStore.search(highQuery, topK: topK) - let low = lowQuery.isEmpty ? [] : try await lowLevelStore.search(lowQuery, topK: topK) + async let highResult = highQuery.isEmpty ? [] : highLevelStore.search(highQuery, topK: topK) + async let lowResult = lowQuery.isEmpty ? [] : lowLevelStore.search(lowQuery, topK: topK) + let high = try await highResult + let low = try await lowResult let merged = merge(high: high, low: low, topK: topK) return DualRetrievalResults( diff --git a/Sources/GraphRAG/LightRAG/Keywords.swift b/Sources/GraphRAG/LightRAG/Keywords.swift index 21cde2f..563cc6f 100644 --- a/Sources/GraphRAG/LightRAG/Keywords.swift +++ b/Sources/GraphRAG/LightRAG/Keywords.swift @@ -120,22 +120,27 @@ public struct KeywordExtractor: Sendable { } /// Parse the JSON object between the first `{` and last `}` of the response. + /// Thinking-tag blocks are stripped first, so a `{` inside a `…` + /// preamble can't be mistaken for the start of the JSON payload. static func parse(_ response: String) -> DualLevelKeywords? { - guard let first = response.firstIndex(of: "{"), - let last = response.lastIndex(of: "}"), first < last + let cleaned = GraphRAG.stripThinkingTags(response) + guard let first = cleaned.firstIndex(of: "{"), + let last = cleaned.lastIndex(of: "}"), first < last else { return nil } - let slice = String(response[first...last]) + let slice = String(cleaned[first...last]) guard let data = slice.data(using: .utf8) else { return nil } return try? JSONDecoder().decode(DualLevelKeywords.self, from: data) } - /// Deterministic fallback: query words of length >= 4, up to 10, as low-level. + /// Deterministic fallback: query words of length >= 4, deduplicated + /// case-insensitively in first-seen order, up to 10, as low-level. static func fallback(_ query: String) -> DualLevelKeywords { + var seen: Set = [] let words = query .split(whereSeparator: { $0.isWhitespace }) .map { String($0.filter { $0.isLetter || $0.isNumber }) } - .filter { $0.count >= 4 } + .filter { $0.count >= 4 && seen.insert($0.lowercased()).inserted } return DualLevelKeywords(highLevel: [], lowLevel: Array(words.prefix(10))) } } diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift index e227d5f..2c62755 100644 --- a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -41,11 +41,12 @@ public struct LightRAGEngine: Sendable { /// Run dual-level retrieval for `query`. public func retrieve(_ query: String, topK: Int = 10) async throws -> DualRetrievalResults { let communities = detectCommunities() - let lowStore = try await LightRAG.chunkSearcher(graph: graph, embedder: embedder) - let highStore = try await LightRAG.communitySearcher( + // The two stores are independent — build them concurrently. + async let lowStore = LightRAG.chunkSearcher(graph: graph, embedder: embedder) + async let highStore = LightRAG.communitySearcher( graph: graph, communities: communities, embedder: embedder) let extractor = KeywordExtractor(model: languageModel, config: keywordConfig) - let retriever = DualLevelRetriever( + let retriever = try await DualLevelRetriever( keywordExtractor: extractor, highLevelStore: highStore, lowLevelStore: lowStore, @@ -65,7 +66,12 @@ public struct LightRAGEngine: Sendable { let context = results.mergedChunks.map { result in "[Relevance: \(String(format: "%.3f", result.score))]\n\(result.content)" }.joined(separator: "\n\n---\n\n") - let sources = results.mergedChunks.map { ChunkID($0.id) } + // High-level hits carry virtual "community_" ids that don't exist in + // the graph; keep only real chunk ids as grounding sources. + let sources = results.mergedChunks + .map(\.id) + .filter { !$0.hasPrefix("community_") } + .map { ChunkID($0) } let confidence = min(1.0, Float(results.mergedChunks.count) / Float(max(1, topK))) if let languageModel, await languageModel.isAvailable() { diff --git a/Sources/GraphRAG/LightRAG/Searchers.swift b/Sources/GraphRAG/LightRAG/Searchers.swift index 8d3676e..b56c651 100644 --- a/Sources/GraphRAG/LightRAG/Searchers.swift +++ b/Sources/GraphRAG/LightRAG/Searchers.swift @@ -14,13 +14,22 @@ public struct InMemorySemanticSearcher: SemanticSearcher { self.embedder = embedder } - /// Build a searcher by embedding and indexing each document. + /// Build a searcher, reusing each document's precomputed embedding when one + /// is supplied and embedding on demand otherwise. public static func build( - documents: [(id: String, content: String)], embedder: any EmbeddingModel + documents: [(id: String, content: String, embedding: [Float]?)], + embedder: any EmbeddingModel ) async throws -> InMemorySemanticSearcher { var retriever = HybridRetriever() for doc in documents { - let embedding = try await embedder.embed(doc.content) + // `??` can't wrap an async default (its autoclosure isn't async), so + // branch explicitly to embed only when no precomputed vector exists. + let embedding: [Float] + if let precomputed = doc.embedding { + embedding = precomputed + } else { + embedding = try await embedder.embed(doc.content) + } retriever.index(id: doc.id, content: doc.content, embedding: embedding) } return InMemorySemanticSearcher(retriever: retriever, embedder: embedder) @@ -41,7 +50,10 @@ public enum LightRAG { public static func chunkSearcher( graph: KnowledgeGraph, embedder: any EmbeddingModel ) async throws -> InMemorySemanticSearcher { - let documents = graph.chunks.map { (id: $0.id.raw, content: $0.content) } + // 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) + } return try await InMemorySemanticSearcher.build(documents: documents, embedder: embedder) } @@ -49,8 +61,13 @@ public enum LightRAG { public static func communitySearcher( graph: KnowledgeGraph, communities: CommunityDetectionResult, embedder: any EmbeddingModel ) async throws -> InMemorySemanticSearcher { + // Summaries are derived text with no precomputed embedding — embed on demand. let documents = communities.communities.map { community in - (id: "community_\(community.id)", content: communitySummary(community, graph: graph)) + ( + id: "community_\(community.id)", + content: communitySummary(community, graph: graph), + embedding: [Float]?.none + ) } return try await InMemorySemanticSearcher.build(documents: documents, embedder: embedder) } From f98e1ad03128c259ee2217f0d8f79800c41c82e7 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 15:42:49 +0000 Subject: [PATCH 03/12] Address Codex review: cache stores, real source grounding, collision-proof merge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - LightRAGEngine now builds the two dual-level stores once per engine (via a memoizing actor cache) instead of re-embedding the whole corpus on every retrieve/ask — with a remote embedder that was O(corpus) network calls per query. - Searchers thread the real chunk ids each document grounds to: a community summary grounds to the chunks mentioning its member entities, so answers built from high-level hits cite actual evidence chunks rather than synthetic community ids. `ask` grounds `sources` on these real chunk ids. - Merge deduplicates within a retrieval level only (keyed by level), so a chunk id that happens to collide with a "community_" no longer silently evicts the community hit or vice versa. Tests: split the interleave test into within-level dedup + colliding cross-level id cases, and add community→member-chunk grounding coverage. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/LightRAG/DualRetrieval.swift | 34 +++++---- .../GraphRAG/LightRAG/LightRAGEngine.swift | 72 ++++++++++++++----- Sources/GraphRAG/LightRAG/Searchers.swift | 47 +++++++++--- Tests/GraphRAGTests/LightRAGTests.swift | 46 ++++++++++-- 4 files changed, 154 insertions(+), 45 deletions(-) diff --git a/Sources/GraphRAG/LightRAG/DualRetrieval.swift b/Sources/GraphRAG/LightRAG/DualRetrieval.swift index b81321d..5d097ba 100644 --- a/Sources/GraphRAG/LightRAG/DualRetrieval.swift +++ b/Sources/GraphRAG/LightRAG/DualRetrieval.swift @@ -116,11 +116,17 @@ public struct DualLevelRetriever: Sendable { private func merge(high: [LightRAGResult], low: [LightRAGResult], topK: Int) -> [LightRAGResult] { guard topK > 0 else { return [] } + // Dedup within a level only: high- and low-level ids live in different + // spaces (community summaries vs chunks), so a chunk that happens to be + // named like a community id must not evict the community hit (or vice + // versa). Keying by level makes the merge collision-proof. var seen: Set = [] var merged: [LightRAGResult] = [] - func take(_ result: LightRAGResult) { - guard merged.count < topK, seen.insert(result.id).inserted else { return } + func take(_ result: LightRAGResult, _ level: Character) { + guard merged.count < topK, seen.insert("\(level)\u{1F}\(result.id)").inserted else { + return + } merged.append(result) } @@ -128,25 +134,25 @@ public struct DualLevelRetriever: Sendable { case .interleave: var i = 0 while merged.count < topK && (i < high.count || i < low.count) { - if i < high.count { take(high[i]) } - if i < low.count { take(low[i]) } + if i < high.count { take(high[i], "h") } + if i < low.count { take(low[i], "l") } i += 1 } case .highFirst: - for r in high { take(r) } - for r in low { take(r) } + for r in high { take(r, "h") } + for r in low { take(r, "l") } case .lowFirst: - for r in low { take(r) } - for r in high { take(r) } + for r in low { take(r, "l") } + for r in high { take(r, "h") } case .weighted: let weighted = - high.map { ($0, $0.score * config.highLevelWeight) } - + low.map { ($0, $0.score * config.lowLevelWeight) } - for (result, _) in weighted.sorted(by: { lhs, rhs in - if lhs.1 == rhs.1 { return lhs.0.id < rhs.0.id } - return lhs.1 > rhs.1 + high.map { ($0, "h" as Character, $0.score * config.highLevelWeight) } + + low.map { ($0, "l" as Character, $0.score * config.lowLevelWeight) } + for (result, level, _) in weighted.sorted(by: { lhs, rhs in + if lhs.2 == rhs.2 { return lhs.0.id < rhs.0.id } + return lhs.2 > rhs.2 }) { - take(result) + take(result, level) } } return merged diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift index 2c62755..f9ff7ed 100644 --- a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -4,6 +4,43 @@ import Foundation +/// Builds and memoizes the two dual-level stores for a fixed graph snapshot, so +/// repeated queries don't re-embed the whole corpus (which, with a remote +/// embedder, would cost O(corpus) network calls per query). Reference semantics +/// via `actor` let a value-type `LightRAGEngine` share one cache across calls. +private actor LightRAGStoreCache { + private let graph: KnowledgeGraph + private let embedder: any EmbeddingModel + private let leidenConfig: LeidenConfig + private var communitiesCache: CommunityDetectionResult? + private var storesCache: (low: any SemanticSearcher, high: any SemanticSearcher)? + + init(graph: KnowledgeGraph, embedder: any EmbeddingModel, leidenConfig: LeidenConfig) { + self.graph = graph + self.embedder = embedder + self.leidenConfig = leidenConfig + } + + func communities() -> CommunityDetectionResult { + if let communitiesCache { return communitiesCache } + let detected = LeidenCommunityDetector(config: leidenConfig).detect(graph) + communitiesCache = detected + return detected + } + + func stores() async throws -> (low: any SemanticSearcher, high: any SemanticSearcher) { + if let storesCache { return storesCache } + let detected = communities() + async let low = LightRAG.chunkSearcher(graph: graph, embedder: embedder) + async let high = LightRAG.communitySearcher( + graph: graph, communities: detected, embedder: embedder) + let built: (low: any SemanticSearcher, high: any SemanticSearcher) = + (low: try await low, high: try await high) + storesCache = built + return built + } +} + /// A self-contained LightRAG engine over a snapshot of a `KnowledgeGraph`. /// /// The low-level store searches document chunks (entity/detail-centric); the @@ -13,9 +50,10 @@ public struct LightRAGEngine: Sendable { public let graph: KnowledgeGraph private let embedder: any EmbeddingModel private let languageModel: (any LanguageModel)? - public var config: DualRetrievalConfig - public var keywordConfig: KeywordExtractorConfig - public var leidenConfig: LeidenConfig + public let config: DualRetrievalConfig + public let keywordConfig: KeywordExtractorConfig + public let leidenConfig: LeidenConfig + private let cache: LightRAGStoreCache public init( graph: KnowledgeGraph, @@ -31,6 +69,8 @@ public struct LightRAGEngine: Sendable { self.config = config self.keywordConfig = keywordConfig self.leidenConfig = leidenConfig + self.cache = LightRAGStoreCache( + graph: graph, embedder: embedder, leidenConfig: leidenConfig) } /// Detect entity communities via Leiden. @@ -38,18 +78,15 @@ public struct LightRAGEngine: Sendable { LeidenCommunityDetector(config: leidenConfig).detect(graph) } - /// Run dual-level retrieval for `query`. + /// Run dual-level retrieval for `query`. The two stores are built once per + /// engine (cached across calls) rather than rebuilt each query. public func retrieve(_ query: String, topK: Int = 10) async throws -> DualRetrievalResults { - let communities = detectCommunities() - // The two stores are independent — build them concurrently. - async let lowStore = LightRAG.chunkSearcher(graph: graph, embedder: embedder) - async let highStore = LightRAG.communitySearcher( - graph: graph, communities: communities, embedder: embedder) + let stores = try await cache.stores() let extractor = KeywordExtractor(model: languageModel, config: keywordConfig) - let retriever = try await DualLevelRetriever( + let retriever = DualLevelRetriever( keywordExtractor: extractor, - highLevelStore: highStore, - lowLevelStore: lowStore, + highLevelStore: stores.high, + lowLevelStore: stores.low, config: config) return try await retriever.retrieve(query, topK: topK) } @@ -66,11 +103,14 @@ public struct LightRAGEngine: Sendable { let context = results.mergedChunks.map { result in "[Relevance: \(String(format: "%.3f", result.score))]\n\(result.content)" }.joined(separator: "\n\n---\n\n") - // High-level hits carry virtual "community_" ids that don't exist in - // the graph; keep only real chunk ids as grounding sources. + + // Ground on the real chunk ids each hit stands for (a high-level hit + // resolves to its member chunks), deduped in first-seen order, so + // synthetic community ids never surface as unresolvable sources. + var seenSource: Set = [] let sources = results.mergedChunks - .map(\.id) - .filter { !$0.hasPrefix("community_") } + .flatMap(\.sourceChunks) + .filter { seenSource.insert($0).inserted } .map { ChunkID($0) } let confidence = min(1.0, Float(results.mergedChunks.count) / Float(max(1, topK))) diff --git a/Sources/GraphRAG/LightRAG/Searchers.swift b/Sources/GraphRAG/LightRAG/Searchers.swift index b56c651..83e5fcf 100644 --- a/Sources/GraphRAG/LightRAG/Searchers.swift +++ b/Sources/GraphRAG/LightRAG/Searchers.swift @@ -5,22 +5,33 @@ import Foundation /// A `SemanticSearcher` backed by a `HybridRetriever` over a fixed corpus. +/// +/// Each indexed document carries the real chunk ids it stands for +/// (`sourceChunks`), so a high-level community hit still resolves back to the +/// actual graph chunks that ground it rather than to its synthetic id. public struct InMemorySemanticSearcher: SemanticSearcher { private let retriever: HybridRetriever private let embedder: any EmbeddingModel + private let sourceChunksByID: [String: [String]] - private init(retriever: HybridRetriever, embedder: any EmbeddingModel) { + private init( + retriever: HybridRetriever, embedder: any EmbeddingModel, + sourceChunksByID: [String: [String]] + ) { self.retriever = retriever self.embedder = embedder + self.sourceChunksByID = sourceChunksByID } /// Build a searcher, reusing each document's precomputed embedding when one - /// is supplied and embedding on demand otherwise. + /// is supplied and embedding on demand otherwise. `sourceChunks` records the + /// real chunk ids each document grounds to. public static func build( - documents: [(id: String, content: String, embedding: [Float]?)], + documents: [(id: String, content: String, embedding: [Float]?, sourceChunks: [String])], embedder: any EmbeddingModel ) async throws -> InMemorySemanticSearcher { var retriever = HybridRetriever() + var sourceChunksByID: [String: [String]] = [:] for doc in documents { // `??` can't wrap an async default (its autoclosure isn't async), so // branch explicitly to embed only when no precomputed vector exists. @@ -31,33 +42,40 @@ public struct InMemorySemanticSearcher: SemanticSearcher { embedding = try await embedder.embed(doc.content) } retriever.index(id: doc.id, content: doc.content, embedding: embedding) + sourceChunksByID[doc.id] = doc.sourceChunks } - return InMemorySemanticSearcher(retriever: retriever, embedder: embedder) + return InMemorySemanticSearcher( + retriever: retriever, embedder: embedder, sourceChunksByID: sourceChunksByID) } public func search(_ query: String, topK: Int) async throws -> [LightRAGResult] { let queryEmbedding = try await embedder.embed(query) let hits = retriever.search(query: query, queryEmbedding: queryEmbedding, limit: topK) return hits.map { - LightRAGResult(id: $0.id, content: $0.content, score: $0.score, sourceChunks: [$0.id]) + LightRAGResult( + id: $0.id, content: $0.content, score: $0.score, + sourceChunks: sourceChunksByID[$0.id] ?? [$0.id]) } } } /// Builders for the LightRAG stores from a `KnowledgeGraph`. public enum LightRAG { - /// Low-level (entity/detail) store: the document chunks themselves. + /// Low-level (entity/detail) store: the document chunks themselves. Each + /// chunk grounds to itself. public static func chunkSearcher( graph: KnowledgeGraph, embedder: any EmbeddingModel ) 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) + (id: $0.id.raw, content: $0.content, embedding: $0.embedding, sourceChunks: [$0.id.raw]) } return try await InMemorySemanticSearcher.build(documents: documents, embedder: embedder) } /// 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 ) async throws -> InMemorySemanticSearcher { @@ -66,7 +84,8 @@ public enum LightRAG { ( id: "community_\(community.id)", content: communitySummary(community, graph: graph), - embedding: [Float]?.none + embedding: [Float]?.none, + sourceChunks: communitySourceChunks(community, graph: graph) ) } return try await InMemorySemanticSearcher.build(documents: documents, embedder: embedder) @@ -91,4 +110,16 @@ public enum LightRAG { } return parts.joined(separator: ". ") } + + /// The real chunk ids that mention any of a community's member entities, in + /// graph chunk order (deterministic). These are the grounding evidence for a + /// high-level community hit. + public static func communitySourceChunks( + _ community: Community, graph: KnowledgeGraph + ) -> [String] { + let memberSet = Set(community.members) + return graph.chunks.compactMap { chunk in + chunk.entities.contains(where: memberSet.contains) ? chunk.id.raw : nil + } + } } diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift index 0663d46..ac762d2 100644 --- a/Tests/GraphRAGTests/LightRAGTests.swift +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -92,14 +92,14 @@ private func triangleGraph() -> KnowledgeGraph { // MARK: - Dual-level retrieval / merge -@Test func dualRetrievalInterleavesAndDedups() async throws { +@Test func dualRetrievalInterleavesAndDedupsWithinLevel() async throws { + // A store that returns a duplicate id within its own results is deduped. let high = MockSearcher(results: [ LightRAGResult(id: "h1", content: "theme one", score: 0.9), - LightRAGResult(id: "shared", content: "shared", score: 0.5), + LightRAGResult(id: "h1", content: "theme one again", score: 0.4), ]) let low = MockSearcher(results: [ - LightRAGResult(id: "l1", content: "detail one", score: 0.8), - LightRAGResult(id: "shared", content: "shared", score: 0.7), + LightRAGResult(id: "l1", content: "detail one", score: 0.8) ]) let extractor = KeywordExtractor( model: MockLLM(response: #"{"high_level":["t"],"low_level":["d"]}"#)) @@ -108,10 +108,42 @@ private func triangleGraph() -> KnowledgeGraph { let results = try await retriever.retrieve("query", topK: 5) let ids = results.mergedChunks.map(\.id) - // Interleave: h1, l1, shared (deduped once), … no duplicate "shared". #expect(ids.first == "h1") - #expect(ids.filter { $0 == "shared" }.count == 1) - #expect(Set(ids) == ["h1", "l1", "shared"]) + #expect(ids.filter { $0 == "h1" }.count == 1) // within-level dedup + #expect(Set(ids) == ["h1", "l1"]) +} + +@Test func dualRetrievalKeepsCollidingCrossLevelIds() async throws { + // A real chunk id equal to a community id must survive on both levels — the + // high (community) and low (chunk) hits are distinct results and neither + // should silently evict the other. + let high = MockSearcher(results: [ + LightRAGResult(id: "community_0", content: "summary", score: 0.9) + ]) + let low = MockSearcher(results: [ + LightRAGResult(id: "community_0", content: "chunk text", score: 0.8) + ]) + let extractor = KeywordExtractor( + model: MockLLM(response: #"{"high_level":["t"],"low_level":["d"]}"#)) + let retriever = DualLevelRetriever( + keywordExtractor: extractor, highLevelStore: high, lowLevelStore: low) + + let results = try await retriever.retrieve("query", topK: 5) + #expect(results.mergedChunks.count == 2) + #expect(results.mergedChunks.filter { $0.id == "community_0" }.count == 2) +} + +@Test func communityGroundsToMemberChunks() throws { + var graph = triangleGraph() + let chunk = TextChunk( + id: ChunkID("c1"), documentID: DocumentID("d"), content: "about a1", + startOffset: 0, endOffset: 8, entities: [EntityID("a1")]) + graph.addChunk(chunk) + let communities = LeidenCommunityDetector().detect(graph) + let a1Community = try #require( + communities.communities.first { $0.members.contains(EntityID("a1")) }) + let sources = LightRAG.communitySourceChunks(a1Community, graph: graph) + #expect(sources.contains("c1")) } @Test func dualRetrievalWeightedOrdersByWeightedScore() async throws { From 869131784f2bd4cd057093f85f1bec9940032607 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 15:53:21 +0000 Subject: [PATCH 04/12] Address Codex review: gate LightRAG on a successful build, cap fallback keywords - GraphRAG.lightRAG() and detectCommunities() now require isBuilt (throwing GraphRAGError.notInitialized otherwise), mirroring the ask/search guard, so callers can't get answers from a raw, stale, or partially rebuilt graph. - KeywordExtractor.extract now runs the non-LLM fallback through the same `capped` step, so a configured maxKeywords below 10 (or 0) is honored even without an LLM. Update the end-to-end test and README for the now-throwing lightRAG(). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- README.md | 2 +- Sources/GraphRAG/GraphRAG/Engine.swift | 20 ++++++++++++++------ Sources/GraphRAG/LightRAG/Keywords.swift | 5 +++-- Tests/GraphRAGTests/LightRAGTests.swift | 2 +- 4 files changed, 19 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 35b1ba5..0728662 100644 --- a/README.md +++ b/README.md @@ -123,7 +123,7 @@ searched independently, and the hits are merged — `interleave` (default), `highFirst`, `lowFirst`, or `weighted`. ```swift -let engine = await rag.lightRAG() +let engine = try await rag.lightRAG() // requires a successful build() first let answer = try await engine.ask("Who worked on the Analytical Engine?") print(answer.text) diff --git a/Sources/GraphRAG/GraphRAG/Engine.swift b/Sources/GraphRAG/GraphRAG/Engine.swift index 4b6b3d1..111c063 100644 --- a/Sources/GraphRAG/GraphRAG/Engine.swift +++ b/Sources/GraphRAG/GraphRAG/Engine.swift @@ -238,19 +238,27 @@ public actor GraphRAG { /// Direct access to the underlying knowledge graph (a value-type snapshot). public func knowledgeGraph() -> KnowledgeGraph { graph } - /// Detect entity communities in the current graph via Leiden. - public func detectCommunities(config: LeidenConfig = LeidenConfig()) -> CommunityDetectionResult { - LeidenCommunityDetector(config: config).detect(graph) + /// Detect entity communities in the current graph via Leiden. Requires a + /// successful `build()` first, so communities reflect extracted entities and + /// not a raw or stale graph (mirrors the `ask`/`search` guard). + public func detectCommunities(config: LeidenConfig = LeidenConfig()) throws + -> CommunityDetectionResult + { + guard isBuilt else { throw GraphRAGError.notInitialized } + return LeidenCommunityDetector(config: config).detect(graph) } /// A LightRAG dual-level retrieval engine over the current graph snapshot, - /// wired to this instance's embedder and (optional) language model. + /// wired to this instance's embedder and (optional) language model. Requires + /// a successful `build()` first, so the engine can't hand out answers from a + /// raw or partially rebuilt graph (mirrors the `ask`/`search` guard). public func lightRAG( config: DualRetrievalConfig = DualRetrievalConfig(), keywordConfig: KeywordExtractorConfig = KeywordExtractorConfig(), leidenConfig: LeidenConfig = LeidenConfig() - ) -> LightRAGEngine { - LightRAGEngine( + ) throws -> LightRAGEngine { + guard isBuilt else { throw GraphRAGError.notInitialized } + return LightRAGEngine( graph: graph, embedder: embedder, languageModel: languageModel, config: config, keywordConfig: keywordConfig, leidenConfig: leidenConfig) } diff --git a/Sources/GraphRAG/LightRAG/Keywords.swift b/Sources/GraphRAG/LightRAG/Keywords.swift index 563cc6f..1483380 100644 --- a/Sources/GraphRAG/LightRAG/Keywords.swift +++ b/Sources/GraphRAG/LightRAG/Keywords.swift @@ -57,7 +57,7 @@ public struct KeywordExtractor: Sendable { /// back to a simple query tokenization. public func extract(_ query: String) async -> DualLevelKeywords { guard let model, await model.isAvailable() else { - return KeywordExtractor.fallback(query) + return capped(KeywordExtractor.fallback(query)) } let prompt = buildPrompt(query) do { @@ -68,7 +68,8 @@ public struct KeywordExtractor: Sendable { } catch { // fall through to fallback } - return KeywordExtractor.fallback(query) + // The fallback is capped too, so `maxKeywords` holds even without an LLM. + return capped(KeywordExtractor.fallback(query)) } /// Cap the combined keyword count at `maxKeywords`, keeping high-level first. diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift index ac762d2..251100b 100644 --- a/Tests/GraphRAGTests/LightRAGTests.swift +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -170,7 +170,7 @@ private func triangleGraph() -> KnowledgeGraph { """) try await rag.build() - let engine = await rag.lightRAG() + let engine = try await rag.lightRAG() let communities = engine.detectCommunities() #expect(communities.communityCount >= 1) From 0265f520e8c165cab47a703c4275c2f0c9ec53f3 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 16:06:13 +0000 Subject: [PATCH 05/12] Address Codex review: clamp keyword limit, keep short terms, honor 0 iterations - KeywordExtractorConfig clamps maxKeywords to >= 0 (and capped() clamps at the slice site too, since the field is mutable), so a negative limit reads as "no keywords" instead of trapping in prefix(_:). - The no-LLM keyword fallback now filters by a small stopword set instead of a length cutoff, so short but meaningful names/acronyms (Ada, Bob, IBM, AI, ML) survive; if every token is a stopword it falls back to the raw tokens so retrieval still runs rather than searching an empty query. - Leiden iterates 0.. Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/Graph/Leiden.swift | 4 ++- Sources/GraphRAG/LightRAG/Keywords.swift | 45 ++++++++++++++++++------ Tests/GraphRAGTests/LightRAGTests.swift | 31 +++++++++++++++- 3 files changed, 68 insertions(+), 12 deletions(-) diff --git a/Sources/GraphRAG/Graph/Leiden.swift b/Sources/GraphRAG/Graph/Leiden.swift index 8b2cd75..63ef51f 100644 --- a/Sources/GraphRAG/Graph/Leiden.swift +++ b/Sources/GraphRAG/Graph/Leiden.swift @@ -111,7 +111,9 @@ public struct LeidenCommunityDetector: Sendable { var sigmaTot = degree // Σ_tot per community id (ids stay in 0.. DualLevelKeywords { - guard keywords.total > config.maxKeywords else { return keywords } + // Clamp locally too: `maxKeywords` is mutable, so a negative value set + // after init would otherwise trap in `prefix(_:)`. + let limit = max(0, config.maxKeywords) + guard keywords.total > limit else { return keywords } var high = keywords.highLevel var low = keywords.lowLevel - if high.count > config.maxKeywords { high = Array(high.prefix(config.maxKeywords)) } - let remaining = max(0, config.maxKeywords - high.count) + if high.count > limit { high = Array(high.prefix(limit)) } + let remaining = max(0, limit - high.count) low = Array(low.prefix(remaining)) return DualLevelKeywords(highLevel: high, lowLevel: low) } @@ -133,15 +138,35 @@ public struct KeywordExtractor: Sendable { return try? JSONDecoder().decode(DualLevelKeywords.self, from: data) } - /// Deterministic fallback: query words of length >= 4, deduplicated - /// case-insensitively in first-seen order, up to 10, as low-level. + /// A small English stopword set for the no-LLM fallback. Filtering by + /// stopword (rather than a length cutoff) keeps short but meaningful entity + /// names and acronyms like "Ada", "Bob", "IBM", "AI", "ML". + static let stopwords: Set = [ + "a", "an", "and", "are", "as", "at", "be", "been", "being", "but", "by", + "did", "do", "does", "done", "for", "from", "had", "has", "have", "how", + "in", "into", "is", "it", "its", "of", "on", "or", "that", "the", "these", + "this", "those", "to", "was", "were", "what", "when", "where", "which", + "who", "whom", "why", "will", "with", + ] + + /// Deterministic fallback: query tokens minus common stopwords, deduplicated + /// case-insensitively in first-seen order, up to 10, as low-level. If every + /// token is a stopword, fall back to the raw tokens so retrieval still runs + /// rather than searching with an empty query. static func fallback(_ query: String) -> DualLevelKeywords { - var seen: Set = [] - let words = + let tokens = query .split(whereSeparator: { $0.isWhitespace }) .map { String($0.filter { $0.isLetter || $0.isNumber }) } - .filter { $0.count >= 4 && seen.insert($0.lowercased()).inserted } - return DualLevelKeywords(highLevel: [], lowLevel: Array(words.prefix(10))) + .filter { !$0.isEmpty } + + func deduped(_ words: [String]) -> [String] { + var seen: Set = [] + return words.filter { seen.insert($0.lowercased()).inserted } + } + + let meaningful = deduped(tokens.filter { !stopwords.contains($0.lowercased()) }) + let chosen = meaningful.isEmpty ? deduped(tokens) : meaningful + return DualLevelKeywords(highLevel: [], lowLevel: Array(chosen.prefix(10))) } } diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift index 251100b..e93636b 100644 --- a/Tests/GraphRAGTests/LightRAGTests.swift +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -79,7 +79,36 @@ private func triangleGraph() -> KnowledgeGraph { let keywords = await extractor.extract("Find the quantum computing project") #expect(keywords.highLevel.isEmpty) #expect(keywords.lowLevel.contains("quantum")) - #expect(!keywords.lowLevel.contains("the")) // shorter than 4 chars + #expect(!keywords.lowLevel.contains("the")) // common stopword +} + +@Test func keywordFallbackKeepsShortEntityNames() async { + // Short but meaningful names/acronyms must survive (they'd be lost under a + // length cutoff), so retrieval isn't handed an empty query. + let extractor = KeywordExtractor(model: nil) + let keywords = await extractor.extract("Ada Bob IBM") + #expect(keywords.lowLevel == ["Ada", "Bob", "IBM"]) +} + +@Test func keywordFallbackFallsBackToRawWhenAllStopwords() async { + let extractor = KeywordExtractor(model: nil) + let keywords = await extractor.extract("what is the") + #expect(!keywords.lowLevel.isEmpty) // not left empty for an all-stopword query +} + +@Test func keywordExtractorClampsNegativeMaxKeywords() async { + // A negative limit must not trap in prefix(_:); it reads as "no keywords". + let extractor = KeywordExtractor( + model: nil, config: KeywordExtractorConfig(maxKeywords: -1)) + let keywords = await extractor.extract("Ada Bob IBM") + #expect(keywords.isEmpty) +} + +@Test func leidenZeroIterationsYieldsSingletons() { + // maxIterations == 0 disables local moving → each node its own community. + let result = LeidenCommunityDetector(config: LeidenConfig(maxIterations: 0)) + .detect(triangleGraph()) + #expect(result.communityCount == 6) } @Test func keywordLLMParsesDualLevels() async { From cef8821d71b88c83901fd72e133040e5abdb2892 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 16:18:50 +0000 Subject: [PATCH 06/12] Address Codex review: offline high-level search, mention grounding, build race - Keyword fallback now populates both levels with the same terms, so the default no-LLM path actually searches the high-level community store instead of silently degrading global/theme questions to chunk search only. - communitySourceChunks derives grounding from both TextChunk.entities and the members' Entity.mentions, so externally-constructed graphs that carry only mention evidence still resolve high-level hits to real chunks. - LightRAGStoreCache memoizes an in-flight Task (stored before the first await) so concurrent first queries share one build instead of each re-embedding the whole corpus under actor reentrancy; a failed build is not cached. Tests: mention-based community grounding, and updated the fallback test for the now dual-populated offline keywords. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/LightRAG/Keywords.swift | 15 +++++--- .../GraphRAG/LightRAG/LightRAGEngine.swift | 34 +++++++++++++------ Sources/GraphRAG/LightRAG/Searchers.swift | 26 +++++++++++--- Tests/GraphRAGTests/LightRAGTests.swift | 19 ++++++++++- 4 files changed, 74 insertions(+), 20 deletions(-) diff --git a/Sources/GraphRAG/LightRAG/Keywords.swift b/Sources/GraphRAG/LightRAG/Keywords.swift index 99c7aaf..f6e885d 100644 --- a/Sources/GraphRAG/LightRAG/Keywords.swift +++ b/Sources/GraphRAG/LightRAG/Keywords.swift @@ -150,9 +150,14 @@ public struct KeywordExtractor: Sendable { ] /// Deterministic fallback: query tokens minus common stopwords, deduplicated - /// case-insensitively in first-seen order, up to 10, as low-level. If every - /// token is a stopword, fall back to the raw tokens so retrieval still runs - /// rather than searching with an empty query. + /// case-insensitively in first-seen order, up to 10. If every token is a + /// stopword, fall back to the raw tokens so retrieval still runs rather than + /// searching with an empty query. + /// + /// Without an LLM there's no theme/entity split, so the same terms populate + /// both levels — otherwise the high-level (community) store would never be + /// searched in the default offline path, and global/theme questions would + /// silently degrade to chunk search only. static func fallback(_ query: String) -> DualLevelKeywords { let tokens = query @@ -166,7 +171,7 @@ public struct KeywordExtractor: Sendable { } let meaningful = deduped(tokens.filter { !stopwords.contains($0.lowercased()) }) - let chosen = meaningful.isEmpty ? deduped(tokens) : meaningful - return DualLevelKeywords(highLevel: [], lowLevel: Array(chosen.prefix(10))) + let chosen = Array((meaningful.isEmpty ? deduped(tokens) : meaningful).prefix(10)) + return DualLevelKeywords(highLevel: chosen, lowLevel: chosen) } } diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift index f9ff7ed..cefb816 100644 --- a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -9,11 +9,16 @@ import Foundation /// embedder, would cost O(corpus) network calls per query). Reference semantics /// via `actor` let a value-type `LightRAGEngine` share one cache across calls. private actor LightRAGStoreCache { + typealias Stores = (low: any SemanticSearcher, high: any SemanticSearcher) + private let graph: KnowledgeGraph private let embedder: any EmbeddingModel private let leidenConfig: LeidenConfig private var communitiesCache: CommunityDetectionResult? - private var storesCache: (low: any SemanticSearcher, high: any SemanticSearcher)? + /// 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) { self.graph = graph @@ -28,16 +33,25 @@ private actor LightRAGStoreCache { return detected } - func stores() async throws -> (low: any SemanticSearcher, high: any SemanticSearcher) { - if let storesCache { return storesCache } + func stores() async throws -> Stores { + if let storesTask { return try await storesTask.value } let detected = communities() - async let low = LightRAG.chunkSearcher(graph: graph, embedder: embedder) - async let high = LightRAG.communitySearcher( - graph: graph, communities: detected, embedder: embedder) - let built: (low: any SemanticSearcher, high: any SemanticSearcher) = - (low: try await low, high: try await high) - storesCache = built - return built + let embedder = self.embedder + let graph = self.graph + let task = Task { () throws -> Stores in + async let low = LightRAG.chunkSearcher(graph: graph, embedder: embedder) + async let high = LightRAG.communitySearcher( + graph: graph, communities: detected, embedder: embedder) + return (low: try await low, high: try await high) + } + storesTask = task + do { + return try await task.value + } catch { + // Let a later call retry rather than caching the failure forever. + storesTask = nil + throw error + } } } diff --git a/Sources/GraphRAG/LightRAG/Searchers.swift b/Sources/GraphRAG/LightRAG/Searchers.swift index 83e5fcf..2d49d61 100644 --- a/Sources/GraphRAG/LightRAG/Searchers.swift +++ b/Sources/GraphRAG/LightRAG/Searchers.swift @@ -111,15 +111,33 @@ public enum LightRAG { return parts.joined(separator: ". ") } - /// The real chunk ids that mention any of a community's member entities, in - /// graph chunk order (deterministic). These are the grounding evidence for a + /// The real chunk ids that ground a community's member entities, deterministic + /// and deduplicated in first-seen order. These are the evidence for a /// high-level community hit. + /// + /// Uses both representations a `KnowledgeGraph` may carry: chunks annotated + /// with member entities (`TextChunk.entities`) and the members' own mention + /// evidence (`Entity.mentions`). The build pipeline populates the former, but + /// externally-constructed graphs may only carry the latter — either alone + /// still yields grounding. public static func communitySourceChunks( _ community: Community, graph: KnowledgeGraph ) -> [String] { let memberSet = Set(community.members) - return graph.chunks.compactMap { chunk in - chunk.entities.contains(where: memberSet.contains) ? chunk.id.raw : nil + var ids: [String] = [] + var seen: Set = [] + func add(_ id: String) { + if seen.insert(id).inserted { ids.append(id) } } + // Chunks annotated with a member entity, in graph chunk order. + for chunk in graph.chunks where chunk.entities.contains(where: memberSet.contains) { + add(chunk.id.raw) + } + // Member entities' own mention evidence, in member then mention order. + for member in community.members { + guard let entity = graph.entity(member) else { continue } + for mention in entity.mentions { add(mention.chunkID.raw) } + } + return ids } } diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift index e93636b..87f0030 100644 --- a/Tests/GraphRAGTests/LightRAGTests.swift +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -77,8 +77,9 @@ private func triangleGraph() -> KnowledgeGraph { @Test func keywordFallbackTokenizesQuery() async { let extractor = KeywordExtractor(model: nil) let keywords = await extractor.extract("Find the quantum computing project") - #expect(keywords.highLevel.isEmpty) + // Both levels are populated offline so the high-level store is still searched. #expect(keywords.lowLevel.contains("quantum")) + #expect(keywords.highLevel.contains("quantum")) #expect(!keywords.lowLevel.contains("the")) // common stopword } @@ -175,6 +176,22 @@ private func triangleGraph() -> KnowledgeGraph { #expect(sources.contains("c1")) } +@Test func communityGroundsViaEntityMentions() throws { + // A graph carrying only Entity.mentions (no TextChunk.entities annotations) + // must still ground community hits to the mentioned chunks. + var graph = KnowledgeGraph() + let mention = EntityMention(chunkID: ChunkID("c9"), startOffset: 0, endOffset: 3, confidence: 1) + graph.addEntity(Entity(id: EntityID("e1"), name: "e1", entityType: "X", mentions: [mention])) + graph.addEntity(Entity(id: EntityID("e2"), name: "e2", entityType: "X")) + graph.addRelationship( + Relationship(source: EntityID("e1"), target: EntityID("e2"), relationType: "R", confidence: 1)) + let communities = LeidenCommunityDetector().detect(graph) + let community = try #require( + communities.communities.first { $0.members.contains(EntityID("e1")) }) + let sources = LightRAG.communitySourceChunks(community, graph: graph) + #expect(sources.contains("c9")) +} + @Test func dualRetrievalWeightedOrdersByWeightedScore() async throws { let high = MockSearcher(results: [LightRAGResult(id: "h", content: "h", score: 1.0)]) let low = MockSearcher(results: [LightRAGResult(id: "l", content: "l", score: 1.0)]) From 7f1fece1d039dbd389c435aa76985118fb2d2c20 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 16:33:25 +0000 Subject: [PATCH 07/12] Address Codex review: offline fallback tokenization/caps, topK guards, grounding - Keyword fallback splits tokens on non-alphanumerics (matching how BM25 and HashEmbedder tokenize the index), so terms like "graph-based" become ["graph","based"] and match indexed content instead of a collapsed "graphbased". - The fallback caps each level to maxKeywords independently instead of routing the duplicated terms through the combined cap, so a small limit no longer starves the low-level (chunk) store while filling the high-level one. - DualLevelRetriever.retrieve returns early for nonpositive topK, before forwarding it to stores that limit with prefix(topK) (which traps on negatives). - InMemorySemanticSearcher sizes its candidate pool to the corpus, so a topK above the default 100 cap can still return enough matches for large stores. - communitySourceChunks also includes Relationship.context chunks for internal edges, completing evidence coverage for externally-constructed graphs. Tests: punctuation splitting, per-level cap, and nonpositive-topK early return. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/LightRAG/DualRetrieval.swift | 8 +++++ Sources/GraphRAG/LightRAG/Keywords.swift | 33 +++++++++++-------- Sources/GraphRAG/LightRAG/Searchers.swift | 23 +++++++++---- Tests/GraphRAGTests/LightRAGTests.swift | 30 +++++++++++++++++ 4 files changed, 74 insertions(+), 20 deletions(-) diff --git a/Sources/GraphRAG/LightRAG/DualRetrieval.swift b/Sources/GraphRAG/LightRAG/DualRetrieval.swift index 5d097ba..30e59fc 100644 --- a/Sources/GraphRAG/LightRAG/DualRetrieval.swift +++ b/Sources/GraphRAG/LightRAG/DualRetrieval.swift @@ -98,6 +98,14 @@ public struct DualLevelRetriever: Sendable { public func retrieve(_ query: String, topK: Int = 10) async throws -> DualRetrievalResults { let keywords = await keywordExtractor.extract(query) + // A nonpositive topK means "no results" — return before searching, since + // some `SemanticSearcher`s limit with `prefix(topK)`, which traps on a + // negative count (mirrors merge's own `topK > 0` guard). + guard topK > 0 else { + return DualRetrievalResults( + highLevelChunks: [], lowLevelChunks: [], mergedChunks: [], keywords: keywords) + } + // Each level searches with its joined keywords as a single query. The two // levels are independent, so run them concurrently. let highQuery = keywords.highLevel.joined(separator: " ") diff --git a/Sources/GraphRAG/LightRAG/Keywords.swift b/Sources/GraphRAG/LightRAG/Keywords.swift index f6e885d..0a2e8a7 100644 --- a/Sources/GraphRAG/LightRAG/Keywords.swift +++ b/Sources/GraphRAG/LightRAG/Keywords.swift @@ -59,7 +59,7 @@ public struct KeywordExtractor: Sendable { /// back to a simple query tokenization. public func extract(_ query: String) async -> DualLevelKeywords { guard let model, await model.isAvailable() else { - return capped(KeywordExtractor.fallback(query)) + return KeywordExtractor.fallback(query, limit: config.maxKeywords) } let prompt = buildPrompt(query) do { @@ -70,8 +70,9 @@ public struct KeywordExtractor: Sendable { } catch { // fall through to fallback } - // The fallback is capped too, so `maxKeywords` holds even without an LLM. - return capped(KeywordExtractor.fallback(query)) + // The fallback caps each level independently, so `maxKeywords` holds even + // without an LLM. + return KeywordExtractor.fallback(query, limit: config.maxKeywords) } /// Cap the combined keyword count at `maxKeywords`, keeping high-level first. @@ -150,20 +151,24 @@ public struct KeywordExtractor: Sendable { ] /// Deterministic fallback: query tokens minus common stopwords, deduplicated - /// case-insensitively in first-seen order, up to 10. If every token is a - /// stopword, fall back to the raw tokens so retrieval still runs rather than - /// searching with an empty query. + /// case-insensitively in first-seen order, capped to `limit` per level. If + /// every token is a stopword, fall back to the raw tokens so retrieval still + /// runs rather than searching with an empty query. + /// + /// Tokens are split on non-alphanumerics — the same way BM25 and + /// `HashEmbedder` tokenize the index — so a term like "graph-based" becomes + /// `["graph", "based"]` and actually matches indexed content instead of a + /// collapsed "graphbased" that matches nothing. /// /// Without an LLM there's no theme/entity split, so the same terms populate - /// both levels — otherwise the high-level (community) store would never be - /// searched in the default offline path, and global/theme questions would - /// silently degrade to chunk search only. - static func fallback(_ query: String) -> DualLevelKeywords { + /// both levels (each within `limit`) — otherwise the high-level (community) + /// store would never be searched in the default offline path, or one level + /// would be starved when the shared budget is spent on the other. + static func fallback(_ query: String, limit: Int) -> DualLevelKeywords { let tokens = query - .split(whereSeparator: { $0.isWhitespace }) - .map { String($0.filter { $0.isLetter || $0.isNumber }) } - .filter { !$0.isEmpty } + .split(whereSeparator: { !($0.isLetter || $0.isNumber) }) + .map(String.init) func deduped(_ words: [String]) -> [String] { var seen: Set = [] @@ -171,7 +176,7 @@ public struct KeywordExtractor: Sendable { } let meaningful = deduped(tokens.filter { !stopwords.contains($0.lowercased()) }) - let chosen = Array((meaningful.isEmpty ? deduped(tokens) : meaningful).prefix(10)) + let chosen = Array((meaningful.isEmpty ? deduped(tokens) : meaningful).prefix(max(0, limit))) return DualLevelKeywords(highLevel: chosen, lowLevel: chosen) } } diff --git a/Sources/GraphRAG/LightRAG/Searchers.swift b/Sources/GraphRAG/LightRAG/Searchers.swift index 2d49d61..9837359 100644 --- a/Sources/GraphRAG/LightRAG/Searchers.swift +++ b/Sources/GraphRAG/LightRAG/Searchers.swift @@ -30,7 +30,10 @@ public struct InMemorySemanticSearcher: SemanticSearcher { documents: [(id: String, content: String, embedding: [Float]?, sourceChunks: [String])], embedder: any EmbeddingModel ) async throws -> InMemorySemanticSearcher { - var retriever = HybridRetriever() + // Let the candidate pool cover the whole corpus so a large `topK` isn't + // silently capped below the number of available documents. + var retriever = HybridRetriever( + config: HybridConfig(maxCandidates: max(100, documents.count))) var sourceChunksByID: [String: [String]] = [:] for doc in documents { // `??` can't wrap an async default (its autoclosure isn't async), so @@ -115,11 +118,13 @@ public enum LightRAG { /// and deduplicated in first-seen order. These are the evidence for a /// high-level community hit. /// - /// Uses both representations a `KnowledgeGraph` may carry: chunks annotated - /// with member entities (`TextChunk.entities`) and the members' own mention - /// evidence (`Entity.mentions`). The build pipeline populates the former, but - /// externally-constructed graphs may only carry the latter — either alone - /// still yields grounding. + /// Uses every evidence representation a `KnowledgeGraph` may carry: chunks + /// annotated with member entities (`TextChunk.entities`), the members' own + /// mention evidence (`Entity.mentions`), and the context chunks of the + /// relationships connecting members (`Relationship.context`). The build + /// pipeline populates chunk annotations, but externally-constructed graphs + /// may carry only mentions or only relationship context — any one alone still + /// yields grounding. public static func communitySourceChunks( _ community: Community, graph: KnowledgeGraph ) -> [String] { @@ -138,6 +143,12 @@ public enum LightRAG { guard let entity = graph.entity(member) else { continue } for mention in entity.mentions { add(mention.chunkID.raw) } } + // Context chunks of relationships internal to the community — the + // evidence behind the relationship types the summary is built from. + for relationship in graph.relationships + where memberSet.contains(relationship.source) && memberSet.contains(relationship.target) { + for chunkID in relationship.context { add(chunkID.raw) } + } return ids } } diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift index 87f0030..f546f93 100644 --- a/Tests/GraphRAGTests/LightRAGTests.swift +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -97,6 +97,36 @@ private func triangleGraph() -> KnowledgeGraph { #expect(!keywords.lowLevel.isEmpty) // not left empty for an all-stopword query } +@Test func keywordFallbackSplitsOnPunctuation() async { + // Tokens must split like the index (BM25 / HashEmbedder) so they match. + let extractor = KeywordExtractor(model: nil) + let keywords = await extractor.extract("graph-based retrieval") + #expect(keywords.lowLevel.contains("graph")) + #expect(keywords.lowLevel.contains("based")) + #expect(!keywords.lowLevel.contains("graphbased")) +} + +@Test func keywordFallbackCapsEachLevelIndependently() async { + // A small combined-looking limit must not starve one level (both stay used). + let extractor = KeywordExtractor( + model: nil, config: KeywordExtractorConfig(maxKeywords: 2)) + let keywords = await extractor.extract("alpha beta gamma delta epsilon") + #expect(keywords.highLevel.count == 2) + #expect(keywords.lowLevel.count == 2) +} + +@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)]) + let extractor = KeywordExtractor( + model: MockLLM(response: #"{"high_level":["t"],"low_level":["d"]}"#)) + let retriever = DualLevelRetriever( + keywordExtractor: extractor, highLevelStore: high, lowLevelStore: low) + // Must not forward a negative topK to the stores (some use prefix, which traps). + let results = try await retriever.retrieve("query", topK: -1) + #expect(results.mergedChunks.isEmpty) +} + @Test func keywordExtractorClampsNegativeMaxKeywords() async { // A negative limit must not trap in prefix(_:); it reads as "no keywords". let extractor = KeywordExtractor( From 77d0fdd38d857161732adcf5ccebab911d2e26ce Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 16:48:52 +0000 Subject: [PATCH 08/12] Address Codex review: short-circuit topK<=0 in engine, guard embedding dimension - LightRAGEngine.retrieve returns an empty result for topK <= 0 before building the stores, so a no-op request never triggers corpus-wide embedding (or a throw from a remote/unavailable embedder). Mirrors DualLevelRetriever's guard one level up. - InMemorySemanticSearcher.build reuses a precomputed chunk embedding only when its dimension matches the current embedder; on a mismatch (e.g. a graph saved under a different embedder) it re-embeds, since cosine similarity would otherwise silently return 0 and drop the vector signal. Test: retrieve(topK: 0) over an engine with a throwing embedder returns empty without embedding. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/LightRAG/LightRAGEngine.swift | 7 +++++++ Sources/GraphRAG/LightRAG/Searchers.swift | 9 ++++++--- Tests/GraphRAGTests/LightRAGTests.swift | 16 ++++++++++++++++ 3 files changed, 29 insertions(+), 3 deletions(-) diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift index cefb816..57a2ea7 100644 --- a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -95,6 +95,13 @@ 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. public func retrieve(_ query: String, topK: Int = 10) async throws -> DualRetrievalResults { + // A nonpositive topK yields nothing — return before building/embedding the + // stores, so a no-op request never triggers corpus-wide embedding work. + guard topK > 0 else { + return DualRetrievalResults( + highLevelChunks: [], lowLevelChunks: [], mergedChunks: [], + keywords: DualLevelKeywords()) + } let stores = try await cache.stores() let extractor = KeywordExtractor(model: languageModel, config: keywordConfig) let retriever = DualLevelRetriever( diff --git a/Sources/GraphRAG/LightRAG/Searchers.swift b/Sources/GraphRAG/LightRAG/Searchers.swift index 9837359..e831da3 100644 --- a/Sources/GraphRAG/LightRAG/Searchers.swift +++ b/Sources/GraphRAG/LightRAG/Searchers.swift @@ -36,10 +36,13 @@ public struct InMemorySemanticSearcher: SemanticSearcher { config: HybridConfig(maxCandidates: max(100, documents.count))) var sourceChunksByID: [String: [String]] = [:] for doc in documents { - // `??` can't wrap an async default (its autoclosure isn't async), so - // branch explicitly to embed only when no precomputed vector exists. + // Reuse a precomputed vector only when its dimension matches this + // embedder — a mismatch (e.g. a graph saved under a different + // embedder) would make cosine similarity silently return 0. Branch + // explicitly since `??` can't wrap an async default (its autoclosure + // isn't async). let embedding: [Float] - if let precomputed = doc.embedding { + if let precomputed = doc.embedding, precomputed.count == embedder.dimension { embedding = precomputed } else { embedding = try await embedder.embed(doc.content) diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift index f546f93..16ea539 100644 --- a/Tests/GraphRAGTests/LightRAGTests.swift +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -18,6 +18,14 @@ private struct MockSearcher: SemanticSearcher { } } +/// Throws on any embedding call — used to prove a code path never embeds. +private struct FailingEmbedder: EmbeddingModel { + func embed(_ text: String) async throws -> [Float] { + throw GraphRAGError.validation(message: "embedder should not be called") + } + var dimension: Int { 8 } +} + private func triangleGraph() -> KnowledgeGraph { var graph = KnowledgeGraph() let names = ["a1", "a2", "a3", "b1", "b2", "b3"] @@ -115,6 +123,14 @@ private func triangleGraph() -> KnowledgeGraph { #expect(keywords.lowLevel.count == 2) } +@Test func lightRAGRetrieveNonPositiveTopKSkipsStoreBuild() async throws { + // The failing embedder would throw if the stores were built; a topK <= 0 + // request must short-circuit before that. + let engine = LightRAGEngine(graph: KnowledgeGraph(), embedder: FailingEmbedder()) + let results = try await engine.retrieve("q", topK: 0) + #expect(results.mergedChunks.isEmpty) +} + @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 478d62339ec30217d0d6e623243de93aebf0057f Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 16:50:14 +0000 Subject: [PATCH 09/12] Fix CI: add missing isAvailable() to FailingEmbedder test double EmbeddingModel requires isAvailable(); the new test double omitted it, breaking the test build. Source is unaffected. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Tests/GraphRAGTests/LightRAGTests.swift | 1 + 1 file changed, 1 insertion(+) diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift index 16ea539..5c81eb6 100644 --- a/Tests/GraphRAGTests/LightRAGTests.swift +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -23,6 +23,7 @@ private struct FailingEmbedder: EmbeddingModel { func embed(_ text: String) async throws -> [Float] { throw GraphRAGError.validation(message: "embedder should not be called") } + func isAvailable() async -> Bool { true } var dimension: Int { 8 } } From 2bd1644b9602176f30e038c3f9519e37a552ce54 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 17:08:16 +0000 Subject: [PATCH 10/12] Address Codex review: reject keyword-only LightRAG, short-circuit maxKeywords=0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - GraphRAG.lightRAG() rejects approach == "keyword" with a clear error: LightRAG needs chunk embeddings, which a keyword-only build leaves nil, so fail fast instead of forcing corpus-wide re-embedding (which would throw outright on an unavailable embedder). - LightRAGEngine.retrieve also short-circuits when keywordConfig.maxKeywords <= 0 (guaranteed-empty queries) before building the stores — a plain config check, so no embedding or LLM keyword call happens. - Leiden: documented why local moving intentionally omits the empty-singleton candidate (keeps the port aligned with the simplified single-level upstream; splitting connected over-merges is a deliberate non-goal), so the review can skip re-flagging it. Tests: keyword-only rejection, and maxKeywords=0 short-circuit via a throwing embedder. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/Graph/Leiden.swift | 9 ++++++++ Sources/GraphRAG/GraphRAG/Engine.swift | 9 ++++++++ .../GraphRAG/LightRAG/LightRAGEngine.swift | 9 +++++--- Tests/GraphRAGTests/LightRAGTests.swift | 21 +++++++++++++++++++ 4 files changed, 45 insertions(+), 3 deletions(-) diff --git a/Sources/GraphRAG/Graph/Leiden.swift b/Sources/GraphRAG/Graph/Leiden.swift index 63ef51f..044a380 100644 --- a/Sources/GraphRAG/Graph/Leiden.swift +++ b/Sources/GraphRAG/Graph/Leiden.swift @@ -129,6 +129,15 @@ public struct LeidenCommunityDetector: Sendable { // Score staying in ci, then every neighbor community; keep the // best (ties: prefer current, then lowest id — deterministic). + // + // We intentionally do NOT score the empty-singleton candidate + // (Σ_tot = 0) here, so a connected over-merge from an earlier pass + // can't be split back into a singleton by local moving. This keeps + // the port aligned with the upstream simplified single-level Leiden + // (disconnected communities are still split in the refinement pass) + // and avoids introducing an unverifiable change to this + // deterministic core — full Leiden singleton refinement is a + // deliberate non-goal for this port. var bestCommunity = ci var bestScore = (weightToCommunity[ci] ?? 0) - config.resolution * ki * sigmaTot[ci] / twoM for (comm, wIn) in weightToCommunity.sorted(by: { $0.key < $1.key }) where comm != ci { diff --git a/Sources/GraphRAG/GraphRAG/Engine.swift b/Sources/GraphRAG/GraphRAG/Engine.swift index 111c063..98585eb 100644 --- a/Sources/GraphRAG/GraphRAG/Engine.swift +++ b/Sources/GraphRAG/GraphRAG/Engine.swift @@ -258,6 +258,15 @@ public actor GraphRAG { leidenConfig: LeidenConfig = LeidenConfig() ) throws -> LightRAGEngine { guard isBuilt else { throw GraphRAGError.notInitialized } + // LightRAG does dual-level *semantic* retrieval, which needs chunk + // embeddings. A keyword-only build (`approach == "keyword"`) leaves those + // 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 config.approach.lowercased() != "keyword" else { + throw GraphRAGError.validation( + message: "LightRAG requires embeddings and is unavailable with approach == \"keyword\"") + } return LightRAGEngine( graph: graph, embedder: embedder, languageModel: languageModel, config: config, keywordConfig: keywordConfig, leidenConfig: leidenConfig) diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift index 57a2ea7..62f4063 100644 --- a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -95,9 +95,12 @@ 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. public func retrieve(_ query: String, topK: Int = 10) async throws -> DualRetrievalResults { - // A nonpositive topK yields nothing — return before building/embedding the - // stores, so a no-op request never triggers corpus-wide embedding work. - guard topK > 0 else { + // 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 { return DualRetrievalResults( highLevelChunks: [], lowLevelChunks: [], mergedChunks: [], keywords: DualLevelKeywords()) diff --git a/Tests/GraphRAGTests/LightRAGTests.swift b/Tests/GraphRAGTests/LightRAGTests.swift index 5c81eb6..ac36487 100644 --- a/Tests/GraphRAGTests/LightRAGTests.swift +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -132,6 +132,27 @@ private func triangleGraph() -> KnowledgeGraph { #expect(results.mergedChunks.isEmpty) } +@Test func lightRAGRetrieveZeroMaxKeywordsSkipsStoreBuild() async throws { + // maxKeywords == 0 forces empty queries, so retrieve must return empty + // before building/embedding the stores (the failing embedder proves it). + let engine = LightRAGEngine( + graph: KnowledgeGraph(), embedder: FailingEmbedder(), + keywordConfig: KeywordExtractorConfig(maxKeywords: 0)) + let results = try await engine.retrieve("q", topK: 5) + #expect(results.mergedChunks.isEmpty) +} + +@Test func lightRAGRejectsKeywordOnlyApproach() async throws { + // Keyword-only builds leave chunk embeddings nil; LightRAG needs them, so + // exporting the engine must fail fast rather than force re-embedding. + let rag = try GraphRAGBuilder().withApproach("keyword").build() + await rag.addDocument(text: "Ada Lovelace worked with Charles Babbage.") + try await rag.build() + await #expect(throws: GraphRAGError.self) { + _ = try await rag.lightRAG() + } +} + @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 696770c333dd468212354df692107d2653c5b54d Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 17:09:37 +0000 Subject: [PATCH 11/12] Fix CI: reference self.config.approach in lightRAG() The local DualRetrievalConfig parameter named `config` shadowed the actor's `self.config: Config`; qualify with `self.` so the keyword-approach guard reads the GraphRAG Config that actually has `approach`. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/GraphRAG/Engine.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Sources/GraphRAG/GraphRAG/Engine.swift b/Sources/GraphRAG/GraphRAG/Engine.swift index 98585eb..a14b20c 100644 --- a/Sources/GraphRAG/GraphRAG/Engine.swift +++ b/Sources/GraphRAG/GraphRAG/Engine.swift @@ -263,7 +263,7 @@ 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 config.approach.lowercased() != "keyword" else { + guard self.config.approach.lowercased() != "keyword" else { throw GraphRAGError.validation( message: "LightRAG requires embeddings and is unavailable with approach == \"keyword\"") } From 169e5334a03aa1c6164a085eb1ab3999d3fe2e6c Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 17:25:57 +0000 Subject: [PATCH 12/12] Document LightRAG's independence from GraphRAG hybrid-path Config (Codex review) Round of Codex P2s all point at one design question: should LightRAG inherit the owning GraphRAG's hybrid-path Config (approach / similarityThreshold / topKResults) and short-circuit empty queries before building stores? Deliberate decision: LightRAG is a self-contained dual-level retrieval subsystem configured via DualRetrievalConfig / KeywordExtractorConfig and a per-call topK. It intentionally does not inherit the GraphRAG hybrid-path Config, and empty queries are handled by the (cached, built-once) retriever returning no hits. Documented at both flagged sites so the decision is explicit rather than an omission; can be threaded through if inheritance is later preferred. No behavior change. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01VUzx3BzstYuB68txcHvAUh --- Sources/GraphRAG/LightRAG/LightRAGEngine.swift | 10 ++++++++++ Sources/GraphRAG/LightRAG/Searchers.swift | 9 +++++++++ 2 files changed, 19 insertions(+) diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift index 62f4063..4f55312 100644 --- a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -94,6 +94,16 @@ 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 { // Requests that can't produce hits — a nonpositive topK, or a keyword // budget of 0 (which forces empty high/low queries) — return before diff --git a/Sources/GraphRAG/LightRAG/Searchers.swift b/Sources/GraphRAG/LightRAG/Searchers.swift index e831da3..b5eab22 100644 --- a/Sources/GraphRAG/LightRAG/Searchers.swift +++ b/Sources/GraphRAG/LightRAG/Searchers.swift @@ -55,6 +55,15 @@ public struct InMemorySemanticSearcher: SemanticSearcher { } 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). let queryEmbedding = try await embedder.embed(query) let hits = retriever.search(query: query, queryEmbedding: queryEmbedding, limit: topK) return hits.map {