pub mod ai;
pub mod builder;
pub mod database;
pub mod relationships;
pub mod types;
pub mod utils;
#[cfg(test)]
mod tests;
pub use builder::GraphBuilder;
pub use types::{CodeGraph, CodeNode, CodeRelationship, FunctionInfo};
pub use utils::{
cosine_similarity, detect_project_root, graphrag_nodes_to_markdown, graphrag_nodes_to_text,
render_graphrag_nodes_json, to_relative_path,
};
use crate::config::Config;
use anyhow::Result;
#[derive(Clone)]
pub struct GraphRAG {
config: Config,
}
impl GraphRAG {
pub fn new(config: Config) -> Self {
Self { config }
}
pub async fn search(&self, query: &str) -> Result<String> {
let builder = GraphBuilder::new_with_quiet(self.config.clone(), true).await?;
let nodes = builder.search_nodes(query).await?;
Ok(graphrag_nodes_to_text(&nodes))
}
pub async fn get_node(&self, node_id: &str) -> Result<String> {
let builder = GraphBuilder::new_with_quiet(self.config.clone(), true).await?;
let graph = builder.get_graph().await?;
match graph.nodes.get(node_id) {
Some(node) => Ok(format!(
"Node: {}\nID: {}\nKind: {}\nPath: {}\nDescription: {}\nSymbols: {}\n",
node.name,
node.id,
node.kind,
node.path,
node.description,
node.symbols.join(", ")
)),
None => Err(anyhow::anyhow!("Node not found: {}", node_id)),
}
}
pub async fn get_relationships(&self, node_id: &str) -> Result<String> {
let builder = GraphBuilder::new_with_quiet(self.config.clone(), true).await?;
let graph = builder.get_graph().await?;
if !graph.nodes.contains_key(node_id) {
return Err(anyhow::anyhow!("Node not found: {}", node_id));
}
let relationships: Vec<_> = graph
.relationships
.iter()
.filter(|rel| rel.source == *node_id || rel.target == *node_id)
.collect();
if relationships.is_empty() {
return Ok(format!("No relationships found for node: {}", node_id));
}
let mut output = format!(
"Relationships for {} ({} total):\n\n",
node_id,
relationships.len()
);
let outgoing: Vec<_> = relationships
.iter()
.filter(|rel| rel.source == *node_id)
.collect();
if !outgoing.is_empty() {
output.push_str("Outgoing:\n");
for rel in outgoing {
let target_name = graph
.nodes
.get(&rel.target)
.map(|n| n.name.clone())
.unwrap_or_else(|| rel.target.clone());
output.push_str(&format!(
" {} → {} ({}): {}\n",
rel.relation_type, target_name, rel.target, rel.description
));
}
output.push('\n');
}
let incoming: Vec<_> = relationships
.iter()
.filter(|rel| rel.target == *node_id)
.collect();
if !incoming.is_empty() {
output.push_str("Incoming:\n");
for rel in incoming {
let source_name = graph
.nodes
.get(&rel.source)
.map(|n| n.name.clone())
.unwrap_or_else(|| rel.source.clone());
output.push_str(&format!(
" {} ← {} ({}): {}\n",
rel.relation_type, source_name, rel.source, rel.description
));
}
}
Ok(output)
}
pub async fn find_path(
&self,
source_id: &str,
target_id: &str,
max_depth: usize,
) -> Result<String> {
let builder = GraphBuilder::new_with_quiet(self.config.clone(), true).await?;
let paths = builder.find_paths(source_id, target_id, max_depth).await?;
let graph = builder.get_graph().await?;
if paths.is_empty() {
return Ok(format!(
"No paths found between {} and {} within depth {}",
source_id, target_id, max_depth
));
}
let mut output = format!(
"Paths from {} to {} ({} found):\n\n",
source_id,
target_id,
paths.len()
);
for (i, path) in paths.iter().enumerate() {
output.push_str(&format!("Path {}:\n", i + 1));
for (j, node_id) in path.iter().enumerate() {
let node_name = graph
.nodes
.get(node_id)
.map(|n| n.name.clone())
.unwrap_or_else(|| node_id.clone());
if j > 0 {
let prev_id = &path[j - 1];
let rel = graph
.relationships
.iter()
.find(|r| r.source == *prev_id && r.target == *node_id);
if let Some(rel) = rel {
output.push_str(&format!(" --{}-> ", rel.relation_type));
} else {
output.push_str(" -> ");
}
}
output.push_str(&format!("{} ({})", node_name, node_id));
}
output.push_str("\n\n");
}
Ok(output)
}
pub async fn overview(&self) -> Result<String> {
let builder = GraphBuilder::new_with_quiet(self.config.clone(), true).await?;
let graph = builder.get_graph().await?;
let node_count = graph.nodes.len();
let relationship_count = graph.relationships.len();
let mut node_types = std::collections::HashMap::new();
for node in graph.nodes.values() {
*node_types.entry(node.kind.clone()).or_insert(0) += 1;
}
let mut rel_types = std::collections::HashMap::new();
for rel in &graph.relationships {
*rel_types.entry(rel.relation_type.clone()).or_insert(0) += 1;
}
let mut output = format!(
"GraphRAG Overview: {} nodes, {} relationships\n\n",
node_count, relationship_count
);
output.push_str("Node Types:\n");
for (kind, count) in node_types.iter() {
output.push_str(&format!(" {}: {}\n", kind, count));
}
output.push_str("\nRelationship Types:\n");
for (rel_type, count) in rel_types.iter() {
output.push_str(&format!(" {}: {}\n", rel_type, count));
}
Ok(output)
}
pub fn config(&self) -> &Config {
&self.config
}
}