lc_rag/graph_rag/
mod.rs

1// src/retrieval/graph_rag/mod.rs
2//! GraphRAG (Knowledge Graph RAG) module.
3//!
4//! Builds a knowledge graph from documents via LLM-based entity and relation
5//! extraction, detects communities, and supports Global / Local / Hybrid
6//! query modes.
7//!
8//! # Example
9//! ```ignore
10//! use langchainrust::retrieval::graph_rag::{GraphRAG, GraphRAGConfig, QueryMode};
11//! use langchainrust::OpenAIChat;
12//!
13//! let llm = OpenAIChat::new(config);
14//! let graph_rag = GraphRAG::new(llm).with_config(GraphRAGConfig::default());
15//!
16//! graph_rag.add_documents(&docs).await?;
17//! graph_rag.build_communities().await?;
18//! let result = graph_rag.query("What is Rust?", QueryMode::Local).await?;
19//! println!("{}", result.answer);
20//! ```
21
22pub mod community;
23pub mod extractor;
24pub mod graph_store;
25pub mod matcher;
26pub mod query;
27
28pub use graph_store::{Community, Entity, GraphStore, Relation};
29pub use matcher::{EmbeddingMatcher, EntityMatcher, KeywordMatcher};
30pub use query::{GraphRAGResult, QueryMode};
31
32use lc_core::language_models::BaseChatModel;
33use lc_vector_stores::Document;
34use tokio::sync::RwLock;
35
36/// GraphRAG error type.
37#[derive(Debug, thiserror::Error)]
38pub enum GraphRAGError {
39    #[error("LLM error: {0}")]
40    LLMError(String),
41
42    #[error("Extraction error: {0}")]
43    ExtractionError(String),
44
45    #[error("Query error: {0}")]
46    QueryError(String),
47
48    #[error("Community error: {0}")]
49    CommunityError(String),
50}
51
52/// Configuration for GraphRAG.
53pub struct GraphRAGConfig {
54    pub max_entities_per_doc: usize,
55    pub max_relations_per_doc: usize,
56    /// 社区大小分桶数(P1-3)。
57    ///
58    /// 注意:这是**大小分桶**(size tiering),不是经典层级社区(父-子包含关系)。
59    /// 社区按大小排名均分为这么多档,`Community::level` 表示"第几档大小"。
60    pub community_size_tiers: usize,
61    /// Maximum number of tokens for context in query prompts.
62    /// When set, community summaries or subgraph context is truncated to fit.
63    pub max_context_tokens: Option<usize>,
64    /// Custom entity matcher for local/hybrid queries.
65    /// When None, uses the default KeywordMatcher.
66    pub entity_matcher: Option<Box<dyn EntityMatcher>>,
67}
68
69impl Default for GraphRAGConfig {
70    fn default() -> Self {
71        Self {
72            max_entities_per_doc: 10,
73            max_relations_per_doc: 10,
74            community_size_tiers: 3,
75            max_context_tokens: None,
76            entity_matcher: None,
77        }
78    }
79}
80
81impl GraphRAGConfig {
82    pub fn new() -> Self {
83        Self::default()
84    }
85
86    pub fn with_max_entities_per_doc(mut self, n: usize) -> Self {
87        self.max_entities_per_doc = n;
88        self
89    }
90
91    pub fn with_max_relations_per_doc(mut self, n: usize) -> Self {
92        self.max_relations_per_doc = n;
93        self
94    }
95
96    /// 设置社区大小分桶数(P1-3,见 [`GraphRAGConfig::community_size_tiers`])。
97    pub fn with_community_size_tiers(mut self, n: usize) -> Self {
98        self.community_size_tiers = n;
99        self
100    }
101
102    /// Sets the maximum number of tokens for context in query prompts.
103    ///
104    /// When set, community summaries (Global/Hybrid) or subgraph context
105    /// (Local/Hybrid) are truncated from lowest-priority items to fit
106    /// within this budget.
107    pub fn with_max_context_tokens(mut self, tokens: usize) -> Self {
108        self.max_context_tokens = Some(tokens);
109        self
110    }
111
112    /// Sets a custom entity matcher for local/hybrid queries.
113    ///
114    /// When set, the matcher is used instead of the default keyword-based
115    /// matching to find relevant entities.
116    pub fn with_entity_matcher(mut self, matcher: Box<dyn EntityMatcher>) -> Self {
117        self.entity_matcher = Some(matcher);
118        self
119    }
120}
121
122/// GraphRAG: Knowledge Graph-based Retrieval Augmented Generation.
123///
124/// Wraps an LLM for entity/relation extraction and community summarization,
125/// and an in-memory [`GraphStore`] for graph operations.
126pub struct GraphRAG<M: BaseChatModel> {
127    llm: M,
128    store: RwLock<GraphStore>,
129    config: GraphRAGConfig,
130}
131
132impl<M: BaseChatModel> GraphRAG<M> {
133    /// Creates a new GraphRAG instance with the given LLM.
134    pub fn new(llm: M) -> Self {
135        Self {
136            llm,
137            store: RwLock::new(GraphStore::new()),
138            config: GraphRAGConfig::default(),
139        }
140    }
141
142    /// Sets a custom configuration.
143    pub fn with_config(mut self, config: GraphRAGConfig) -> Self {
144        self.config = config;
145        self
146    }
147
148    /// Adds documents to the knowledge graph by extracting entities and
149    /// relations from each document via the LLM.
150    pub async fn add_documents(&self, docs: &[Document]) -> Result<(), GraphRAGError> {
151        for doc in docs {
152            let extraction = extractor::extract(
153                &self.llm,
154                &doc.content,
155                self.config.max_entities_per_doc,
156                self.config.max_relations_per_doc,
157            )
158            .await?;
159
160            let doc_id = doc.id.clone();
161            let mut store = self.store.write().await;
162
163            // Build a name-to-id map for deduplication.
164            let mut name_to_id: std::collections::HashMap<String, String> = store
165                .all_entities()
166                .values()
167                .map(|e| (e.name.to_lowercase(), e.id.clone()))
168                .collect();
169
170            // Insert extracted entities (deduplicate by name).
171            for ext_ent in &extraction.entities {
172                let key = ext_ent.name.to_lowercase();
173                if let Some(_existing_id) = name_to_id.get(&key) {
174                    // Entity already exists; skip (M57: log instead of silent discard).
175                    log::info!("GraphRAG: skipping duplicate entity '{}'", ext_ent.name);
176                    continue;
177                }
178
179                let id = format!("e_{}", uuid::Uuid::new_v4().as_simple());
180                name_to_id.insert(key, id.clone());
181
182                store.add_entity(Entity {
183                    id,
184                    name: ext_ent.name.clone(),
185                    entity_type: ext_ent.entity_type.clone(),
186                    description: ext_ent.description.clone(),
187                });
188            }
189
190            // Insert extracted relations (resolve names to ids).
191            for ext_rel in &extraction.relations {
192                let source_key = ext_rel.source.to_lowercase();
193                let target_key = ext_rel.target.to_lowercase();
194
195                let source_id = match name_to_id.get(&source_key) {
196                    Some(id) => id.clone(),
197                    None => {
198                        log::info!(
199                            "GraphRAG: skipping relation with unknown source entity '{}'",
200                            ext_rel.source
201                        );
202                        continue;
203                    }
204                };
205                let target_id = match name_to_id.get(&target_key) {
206                    Some(id) => id.clone(),
207                    None => {
208                        log::info!(
209                            "GraphRAG: skipping relation with unknown target entity '{}'",
210                            ext_rel.target
211                        );
212                        continue;
213                    }
214                };
215
216                store.add_relation(Relation {
217                    source: source_id,
218                    target: target_id,
219                    relation_type: ext_rel.relation_type.clone(),
220                    description: ext_rel.description.clone(),
221                    doc_id: doc_id.clone(),
222                });
223            }
224        }
225
226        Ok(())
227    }
228
229    /// Runs community detection and generates community summaries via the LLM.
230    pub async fn build_communities(&self) -> Result<(), GraphRAGError> {
231        let communities = {
232            let store = self.store.read().await;
233            community::detect_communities(&store, self.config.community_size_tiers)
234        };
235
236        // Generate summaries for each community.
237        let mut summaries = Vec::with_capacity(communities.len());
238        for comm in &communities {
239            let store_clone = {
240                let store = self.store.read().await;
241                store.clone()
242            };
243            let summary = community::summarize_community(&self.llm, &store_clone, comm).await?;
244            summaries.push(summary);
245        }
246
247        // Write communities and summaries back.
248        let mut store = self.store.write().await;
249        store.set_communities(communities);
250        store.set_community_summaries(summaries);
251
252        Ok(())
253    }
254
255    /// Queries the knowledge graph using the specified mode.
256    pub async fn query(&self, q: &str, mode: QueryMode) -> Result<GraphRAGResult, GraphRAGError> {
257        let store = {
258            let guard = self.store.read().await;
259            guard.clone()
260        };
261
262        let max_tokens = self.config.max_context_tokens;
263
264        match mode {
265            QueryMode::Global => query::global_query(&self.llm, &store, q, max_tokens).await,
266            QueryMode::Local => {
267                let matcher = self.config.entity_matcher.as_deref();
268                query::local_query(&self.llm, &store, q, max_tokens, matcher).await
269            }
270            QueryMode::Hybrid => {
271                let matcher = self.config.entity_matcher.as_deref();
272                query::hybrid_query(&self.llm, &store, q, max_tokens, matcher).await
273            }
274        }
275    }
276
277    /// Returns the number of entities in the graph.
278    pub async fn entity_count(&self) -> usize {
279        let store = self.store.read().await;
280        store.entity_count()
281    }
282
283    /// Returns the number of relations in the graph.
284    pub async fn relation_count(&self) -> usize {
285        let store = self.store.read().await;
286        store.relation_count()
287    }
288
289    /// Returns the number of communities.
290    pub async fn community_count(&self) -> usize {
291        let store = self.store.read().await;
292        store.communities().len()
293    }
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn test_graph_rag_config_default() {
302        let config = GraphRAGConfig::default();
303        assert_eq!(config.max_entities_per_doc, 10);
304        assert_eq!(config.max_relations_per_doc, 10);
305        assert_eq!(config.community_size_tiers, 3);
306        assert!(config.max_context_tokens.is_none());
307    }
308
309    #[test]
310    fn test_graph_rag_config_builder() {
311        let config = GraphRAGConfig::new()
312            .with_max_entities_per_doc(5)
313            .with_max_relations_per_doc(8)
314            .with_community_size_tiers(2);
315
316        assert_eq!(config.max_entities_per_doc, 5);
317        assert_eq!(config.max_relations_per_doc, 8);
318        assert_eq!(config.community_size_tiers, 2);
319    }
320
321    #[test]
322    fn test_graph_error_display() {
323        let err = GraphRAGError::LLMError("timeout".into());
324        assert!(err.to_string().contains("timeout"));
325
326        let err = GraphRAGError::ExtractionError("bad json".into());
327        assert!(err.to_string().contains("bad json"));
328
329        let err = GraphRAGError::QueryError("no entities".into());
330        assert!(err.to_string().contains("no entities"));
331    }
332}
lc_rag/graph_rag/mod.rs

lc_rag/graph_rag/
mod.rs
