diff --git a/README.md b/README.md index a073aa7..0728662 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 = try await rag.lightRAG() // requires a successful build() first +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..044a380 --- /dev/null +++ b/Sources/GraphRAG/Graph/Leiden.swift @@ -0,0 +1,277 @@ +// 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 + } + // 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 in 0.. 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. 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. 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() + ) 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 self.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) + } + /// 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..30e59fc --- /dev/null +++ b/Sources/GraphRAG/LightRAG/DualRetrieval.swift @@ -0,0 +1,168 @@ +// 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) + + // 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: " ") + let lowQuery = keywords.lowLevel.joined(separator: " ") + 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( + 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 [] } + // 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, _ level: Character) { + guard merged.count < topK, seen.insert("\(level)\u{1F}\(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], "h") } + if i < low.count { take(low[i], "l") } + i += 1 + } + case .highFirst: + for r in high { take(r, "h") } + for r in low { take(r, "l") } + case .lowFirst: + for r in low { take(r, "l") } + for r in high { take(r, "h") } + case .weighted: + let weighted = + 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, level) + } + } + return merged + } +} diff --git a/Sources/GraphRAG/LightRAG/Keywords.swift b/Sources/GraphRAG/LightRAG/Keywords.swift new file mode 100644 index 0000000..0a2e8a7 --- /dev/null +++ b/Sources/GraphRAG/LightRAG/Keywords.swift @@ -0,0 +1,182 @@ +// 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") { + // Clamp so a negative limit reads as "no keywords" instead of trapping + // later in `prefix(_:)`. + self.maxKeywords = max(0, 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, limit: config.maxKeywords) + } + 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 + } + // 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. + private func capped(_ keywords: DualLevelKeywords) -> DualLevelKeywords { + // 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 > limit { high = Array(high.prefix(limit)) } + let remaining = max(0, limit - 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. + /// 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? { + let cleaned = GraphRAG.stripThinkingTags(response) + guard let first = cleaned.firstIndex(of: "{"), + let last = cleaned.lastIndex(of: "}"), first < last + else { return nil } + let slice = String(cleaned[first...last]) + guard let data = slice.data(using: .utf8) else { return nil } + return try? JSONDecoder().decode(DualLevelKeywords.self, from: data) + } + + /// 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, 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 (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.isLetter || $0.isNumber) }) + .map(String.init) + + 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 = Array((meaningful.isEmpty ? deduped(tokens) : meaningful).prefix(max(0, limit))) + return DualLevelKeywords(highLevel: chosen, lowLevel: chosen) + } +} diff --git a/Sources/GraphRAG/LightRAG/LightRAGEngine.swift b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift new file mode 100644 index 0000000..4f55312 --- /dev/null +++ b/Sources/GraphRAG/LightRAG/LightRAGEngine.swift @@ -0,0 +1,163 @@ +// LightRAGEngine.swift +// High-level LightRAG facade over a knowledge graph: detects communities, +// assembles the two dual-level stores, and answers queries. + +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 { + 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? + /// 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 + 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 -> Stores { + if let storesTask { return try await storesTask.value } + let detected = communities() + 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 + } + } +} + +/// 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 let config: DualRetrievalConfig + public let keywordConfig: KeywordExtractorConfig + public let leidenConfig: LeidenConfig + private let cache: LightRAGStoreCache + + 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 + self.cache = LightRAGStoreCache( + graph: graph, embedder: embedder, leidenConfig: leidenConfig) + } + + /// Detect entity communities via Leiden. + public func detectCommunities() -> CommunityDetectionResult { + LeidenCommunityDetector(config: leidenConfig).detect(graph) + } + + /// 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 + // 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()) + } + let stores = try await cache.stores() + let extractor = KeywordExtractor(model: languageModel, config: keywordConfig) + let retriever = DualLevelRetriever( + keywordExtractor: extractor, + highLevelStore: stores.high, + lowLevelStore: stores.low, + 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") + + // 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 + .flatMap(\.sourceChunks) + .filter { seenSource.insert($0).inserted } + .map { ChunkID($0) } + 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..b5eab22 --- /dev/null +++ b/Sources/GraphRAG/LightRAG/Searchers.swift @@ -0,0 +1,166 @@ +// 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. +/// +/// 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, + 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. `sourceChunks` records the + /// real chunk ids each document grounds to. + public static func build( + documents: [(id: String, content: String, embedding: [Float]?, sourceChunks: [String])], + embedder: any EmbeddingModel + ) async throws -> InMemorySemanticSearcher { + // Let the candidate pool cover the whole corpus so a large `topK` isn't + // silently capped below the number of available documents. + var retriever = HybridRetriever( + config: HybridConfig(maxCandidates: max(100, documents.count))) + var sourceChunksByID: [String: [String]] = [:] + for doc in documents { + // 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, precomputed.count == embedder.dimension { + embedding = precomputed + } else { + 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, sourceChunksByID: sourceChunksByID) + } + + 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 { + 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. 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, 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 { + // 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), + embedding: [Float]?.none, + sourceChunks: communitySourceChunks(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: ". ") + } + + /// 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 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] { + let memberSet = Set(community.members) + 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) } + } + // 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 new file mode 100644 index 0000000..ac36487 --- /dev/null +++ b/Tests/GraphRAGTests/LightRAGTests.swift @@ -0,0 +1,294 @@ +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)) + } +} + +/// 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") + } + func isAvailable() async -> Bool { true } + var dimension: Int { 8 } +} + +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") + // 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 +} + +@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 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 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 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)]) + 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( + 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 { + 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 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: "h1", content: "theme one again", score: 0.4), + ]) + let low = MockSearcher(results: [ + LightRAGResult(id: "l1", content: "detail one", 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) + let ids = results.mergedChunks.map(\.id) + #expect(ids.first == "h1") + #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 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)]) + 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 = try 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) +}