#![allow(unused_imports)]
use crate::config::Config;
use crate::core::{
ChunkId, Document, DocumentId, Entity, EntityId, GraphRAGError, KnowledgeGraph, Relationship,
Result, TextChunk,
};
use crate::{critic, ollama, persistence, query, retrieval};
#[cfg(feature = "parallel-processing")]
#[allow(unused_imports)]
use crate::parallel;
use super::GraphRAG;
impl GraphRAG {
pub fn config(&self) -> &Config {
&self.config
}
pub fn is_initialized(&self) -> bool {
self.knowledge_graph.is_some() && self.retrieval_system.is_some()
}
pub fn has_documents(&self) -> bool {
if let Some(graph) = &self.knowledge_graph {
graph.chunks().count() > 0
} else {
false
}
}
pub fn has_graph(&self) -> bool {
if let Some(graph) = &self.knowledge_graph {
graph.entities().count() > 0
} else {
false
}
}
pub fn knowledge_graph(&self) -> Option<&KnowledgeGraph> {
self.knowledge_graph.as_ref()
}
pub fn get_entity(&self, entity_id: &str) -> Option<&Entity> {
if let Some(graph) = &self.knowledge_graph {
graph.entities().find(|e| e.id.0 == entity_id)
} else {
None
}
}
pub fn get_entity_relationships(&self, entity_id: &str) -> Vec<&Relationship> {
if let Some(graph) = &self.knowledge_graph {
let entity_id_obj = EntityId::new(entity_id.to_string());
graph
.relationships()
.filter(|r| r.source == entity_id_obj || r.target == entity_id_obj)
.collect()
} else {
Vec::new()
}
}
pub fn get_chunk(&self, chunk_id: &str) -> Option<&TextChunk> {
if let Some(graph) = &self.knowledge_graph {
graph.chunks().find(|c| c.id.0 == chunk_id)
} else {
None
}
}
pub fn knowledge_graph_mut(&mut self) -> Option<&mut KnowledgeGraph> {
self.knowledge_graph.as_mut()
}
}