graphrag_core/graphrag/
stats.rs

1#![allow(unused_imports)]
2
3use crate::config::Config;
4use crate::core::{
5    ChunkId, Document, DocumentId, Entity, EntityId, GraphRAGError, KnowledgeGraph, Relationship,
6    Result, TextChunk,
7};
8use crate::{critic, ollama, persistence, query, retrieval};
9
10#[cfg(feature = "parallel-processing")]
11#[allow(unused_imports)]
12use crate::parallel;
13
14use super::GraphRAG;
15
16impl GraphRAG {
17    /// Get a reference to the current configuration
18    pub fn config(&self) -> &Config {
19        &self.config
20    }
21
22    /// Check if system is initialized
23    pub fn is_initialized(&self) -> bool {
24        self.knowledge_graph.is_some() && self.retrieval_system.is_some()
25    }
26
27    /// Check if documents have been added
28    pub fn has_documents(&self) -> bool {
29        if let Some(graph) = &self.knowledge_graph {
30            graph.chunks().count() > 0
31        } else {
32            false
33        }
34    }
35
36    /// Check if graph has been built
37    pub fn has_graph(&self) -> bool {
38        if let Some(graph) = &self.knowledge_graph {
39            graph.entities().count() > 0
40        } else {
41            false
42        }
43    }
44
45    /// Get a reference to the knowledge graph
46    pub fn knowledge_graph(&self) -> Option<&KnowledgeGraph> {
47        self.knowledge_graph.as_ref()
48    }
49
50    /// Get entity details by ID
51    pub fn get_entity(&self, entity_id: &str) -> Option<&Entity> {
52        if let Some(graph) = &self.knowledge_graph {
53            graph.entities().find(|e| e.id.0 == entity_id)
54        } else {
55            None
56        }
57    }
58
59    /// Get all relationships involving an entity
60    pub fn get_entity_relationships(&self, entity_id: &str) -> Vec<&Relationship> {
61        if let Some(graph) = &self.knowledge_graph {
62            let entity_id_obj = EntityId::new(entity_id.to_string());
63            graph
64                .relationships()
65                .filter(|r| r.source == entity_id_obj || r.target == entity_id_obj)
66                .collect()
67        } else {
68            Vec::new()
69        }
70    }
71
72    /// Get chunk by ID
73    pub fn get_chunk(&self, chunk_id: &str) -> Option<&TextChunk> {
74        if let Some(graph) = &self.knowledge_graph {
75            graph.chunks().find(|c| c.id.0 == chunk_id)
76        } else {
77            None
78        }
79    }
80    /// Get a mutable reference to the knowledge graph
81    pub fn knowledge_graph_mut(&mut self) -> Option<&mut KnowledgeGraph> {
82        self.knowledge_graph.as_mut()
83    }
84}
graphrag_core/graphrag/stats.rs

graphrag_core/graphrag/
stats.rs
