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    pub community_max_levels: usize,
57    /// Maximum number of tokens for context in query prompts.
58    /// When set, community summaries or subgraph context is truncated to fit.
59    pub max_context_tokens: Option<usize>,
60    /// Custom entity matcher for local/hybrid queries.
61    /// When None, uses the default KeywordMatcher.
62    pub entity_matcher: Option<Box<dyn EntityMatcher>>,
63}
64
65impl Default for GraphRAGConfig {
66    fn default() -> Self {
67        Self {
68            max_entities_per_doc: 10,
69            max_relations_per_doc: 10,
70            community_max_levels: 3,
71            max_context_tokens: None,
72            entity_matcher: None,
73        }
74    }
75}
76
77impl GraphRAGConfig {
78    pub fn new() -> Self {
79        Self::default()
80    }
81
82    pub fn with_max_entities_per_doc(mut self, n: usize) -> Self {
83        self.max_entities_per_doc = n;
84        self
85    }
86
87    pub fn with_max_relations_per_doc(mut self, n: usize) -> Self {
88        self.max_relations_per_doc = n;
89        self
90    }
91
92    pub fn with_community_max_levels(mut self, n: usize) -> Self {
93        self.community_max_levels = n;
94        self
95    }
96
97    /// Sets the maximum number of tokens for context in query prompts.
98    ///
99    /// When set, community summaries (Global/Hybrid) or subgraph context
100    /// (Local/Hybrid) are truncated from lowest-priority items to fit
101    /// within this budget.
102    pub fn with_max_context_tokens(mut self, tokens: usize) -> Self {
103        self.max_context_tokens = Some(tokens);
104        self
105    }
106
107    /// Sets a custom entity matcher for local/hybrid queries.
108    ///
109    /// When set, the matcher is used instead of the default keyword-based
110    /// matching to find relevant entities.
111    pub fn with_entity_matcher(mut self, matcher: Box<dyn EntityMatcher>) -> Self {
112        self.entity_matcher = Some(matcher);
113        self
114    }
115}
116
117/// GraphRAG: Knowledge Graph-based Retrieval Augmented Generation.
118///
119/// Wraps an LLM for entity/relation extraction and community summarization,
120/// and an in-memory [`GraphStore`] for graph operations.
121pub struct GraphRAG<M: BaseChatModel> {
122    llm: M,
123    store: RwLock<GraphStore>,
124    config: GraphRAGConfig,
125}
126
127impl<M: BaseChatModel> GraphRAG<M> {
128    /// Creates a new GraphRAG instance with the given LLM.
129    pub fn new(llm: M) -> Self {
130        Self {
131            llm,
132            store: RwLock::new(GraphStore::new()),
133            config: GraphRAGConfig::default(),
134        }
135    }
136
137    /// Sets a custom configuration.
138    pub fn with_config(mut self, config: GraphRAGConfig) -> Self {
139        self.config = config;
140        self
141    }
142
143    /// Adds documents to the knowledge graph by extracting entities and
144    /// relations from each document via the LLM.
145    pub async fn add_documents(&self, docs: &[Document]) -> Result<(), GraphRAGError> {
146        for doc in docs {
147            let extraction = extractor::extract(
148                &self.llm,
149                &doc.content,
150                self.config.max_entities_per_doc,
151                self.config.max_relations_per_doc,
152            )
153            .await?;
154
155            let doc_id = doc.id.clone();
156            let mut store = self.store.write().await;
157
158            // Build a name-to-id map for deduplication.
159            let mut name_to_id: std::collections::HashMap<String, String> = store
160                .all_entities()
161                .values()
162                .map(|e| (e.name.to_lowercase(), e.id.clone()))
163                .collect();
164
165            // Insert extracted entities (deduplicate by name).
166            for ext_ent in &extraction.entities {
167                let key = ext_ent.name.to_lowercase();
168                if let Some(_existing_id) = name_to_id.get(&key) {
169                    // Entity already exists; skip (M57: log instead of silent discard).
170                    log::info!("GraphRAG: skipping duplicate entity '{}'", ext_ent.name);
171                    continue;
172                }
173
174                let id = format!("e_{}", uuid::Uuid::new_v4().as_simple());
175                name_to_id.insert(key, id.clone());
176
177                store.add_entity(Entity {
178                    id,
179                    name: ext_ent.name.clone(),
180                    entity_type: ext_ent.entity_type.clone(),
181                    description: ext_ent.description.clone(),
182                });
183            }
184
185            // Insert extracted relations (resolve names to ids).
186            for ext_rel in &extraction.relations {
187                let source_key = ext_rel.source.to_lowercase();
188                let target_key = ext_rel.target.to_lowercase();
189
190                let source_id = match name_to_id.get(&source_key) {
191                    Some(id) => id.clone(),
192                    None => {
193                        log::info!(
194                            "GraphRAG: skipping relation with unknown source entity '{}'",
195                            ext_rel.source
196                        );
197                        continue;
198                    }
199                };
200                let target_id = match name_to_id.get(&target_key) {
201                    Some(id) => id.clone(),
202                    None => {
203                        log::info!(
204                            "GraphRAG: skipping relation with unknown target entity '{}'",
205                            ext_rel.target
206                        );
207                        continue;
208                    }
209                };
210
211                store.add_relation(Relation {
212                    source: source_id,
213                    target: target_id,
214                    relation_type: ext_rel.relation_type.clone(),
215                    description: ext_rel.description.clone(),
216                    doc_id: doc_id.clone(),
217                });
218            }
219        }
220
221        Ok(())
222    }
223
224    /// Runs community detection and generates community summaries via the LLM.
225    pub async fn build_communities(&self) -> Result<(), GraphRAGError> {
226        let communities = {
227            let store = self.store.read().await;
228            community::detect_communities(&store, self.config.community_max_levels)
229        };
230
231        // Generate summaries for each community.
232        let mut summaries = Vec::with_capacity(communities.len());
233        for comm in &communities {
234            let store_clone = {
235                let store = self.store.read().await;
236                store.clone()
237            };
238            let summary = community::summarize_community(&self.llm, &store_clone, comm).await?;
239            summaries.push(summary);
240        }
241
242        // Write communities and summaries back.
243        let mut store = self.store.write().await;
244        store.set_communities(communities);
245        store.set_community_summaries(summaries);
246
247        Ok(())
248    }
249
250    /// Queries the knowledge graph using the specified mode.
251    pub async fn query(&self, q: &str, mode: QueryMode) -> Result<GraphRAGResult, GraphRAGError> {
252        let store = {
253            let guard = self.store.read().await;
254            guard.clone()
255        };
256
257        let max_tokens = self.config.max_context_tokens;
258
259        match mode {
260            QueryMode::Global => query::global_query(&self.llm, &store, q, max_tokens).await,
261            QueryMode::Local => {
262                let matcher = self.config.entity_matcher.as_deref();
263                query::local_query(&self.llm, &store, q, max_tokens, matcher).await
264            }
265            QueryMode::Hybrid => {
266                let matcher = self.config.entity_matcher.as_deref();
267                query::hybrid_query(&self.llm, &store, q, max_tokens, matcher).await
268            }
269        }
270    }
271
272    /// Returns the number of entities in the graph.
273    pub async fn entity_count(&self) -> usize {
274        let store = self.store.read().await;
275        store.entity_count()
276    }
277
278    /// Returns the number of relations in the graph.
279    pub async fn relation_count(&self) -> usize {
280        let store = self.store.read().await;
281        store.relation_count()
282    }
283
284    /// Returns the number of communities.
285    pub async fn community_count(&self) -> usize {
286        let store = self.store.read().await;
287        store.communities().len()
288    }
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294
295    #[test]
296    fn test_graph_rag_config_default() {
297        let config = GraphRAGConfig::default();
298        assert_eq!(config.max_entities_per_doc, 10);
299        assert_eq!(config.max_relations_per_doc, 10);
300        assert_eq!(config.community_max_levels, 3);
301        assert!(config.max_context_tokens.is_none());
302    }
303
304    #[test]
305    fn test_graph_rag_config_builder() {
306        let config = GraphRAGConfig::new()
307            .with_max_entities_per_doc(5)
308            .with_max_relations_per_doc(8)
309            .with_community_max_levels(2);
310
311        assert_eq!(config.max_entities_per_doc, 5);
312        assert_eq!(config.max_relations_per_doc, 8);
313        assert_eq!(config.community_max_levels, 2);
314    }
315
316    #[test]
317    fn test_graph_error_display() {
318        let err = GraphRAGError::LLMError("timeout".into());
319        assert!(err.to_string().contains("timeout"));
320
321        let err = GraphRAGError::ExtractionError("bad json".into());
322        assert!(err.to_string().contains("bad json"));
323
324        let err = GraphRAGError::QueryError("no entities".into());
325        assert!(err.to_string().contains("no entities"));
326    }
327}
lc_rag/graph_rag/mod.rs

lc_rag/graph_rag/
mod.rs
