use std::collections::VecDeque;
use anyhow::Result;
use crate::memory::chunks::get_chunks_batch;
use crate::memory::config::MemoryConfig;
use crate::memory::tree::io::{TreeReadHit, TreeReadRequest, TreeReadResult};
use crate::memory::tree::store::{self, SummaryNode};
const DEFAULT_READ_LIMIT: usize = 50;
pub fn read_tree(config: &MemoryConfig, req: &TreeReadRequest) -> Result<TreeReadResult> {
let Some(tree) = store::get_tree(config, &req.tree_id)? else {
return Ok(TreeReadResult {
hits: Vec::new(),
total: 0,
tree_id: req.tree_id.clone(),
});
};
if req.max_depth == 0 {
return Ok(TreeReadResult::empty(&tree));
}
let start_id = req.start_node_id.clone().or_else(|| tree.root_id.clone());
let Some(start_id) = start_id else {
return Ok(TreeReadResult::empty(&tree));
};
let Some(start) = store::get_summary(config, &start_id)? else {
return Ok(TreeReadResult::empty(&tree));
};
let mut hits: Vec<TreeReadHit> = Vec::new();
let mut queue: VecDeque<(SummaryNode, u32)> = VecDeque::new();
queue.push_back((start, 0));
while let Some((node, depth)) = queue.pop_front() {
hits.push(summary_hit(&node));
if depth + 1 >= req.max_depth {
continue;
}
if node.level >= 2 {
let kids = store::get_summaries_batch(config, &node.child_ids)?;
for cid in &node.child_ids {
if let Some(child) = kids.get(cid) {
if !child.deleted {
queue.push_back((child.clone(), depth + 1));
}
}
}
} else {
let chunks = get_chunks_batch(config, &node.child_ids)?;
for cid in &node.child_ids {
if let Some(c) = chunks.get(cid) {
hits.push(TreeReadHit {
node_id: c.id.clone(),
node_kind: "chunk".to_string(),
level: 0,
content: c.content.clone(),
score: 0.0,
});
}
}
}
}
if let Some(q) = req.query.as_deref() {
let terms: Vec<String> = q
.to_lowercase()
.split_whitespace()
.map(|s| s.to_string())
.collect();
for h in &mut hits {
let lc = h.content.to_lowercase();
let matches = terms.iter().filter(|t| lc.contains(t.as_str())).count();
h.score = matches as f32;
}
hits.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
}
let total = hits.len();
let limit = req.limit.unwrap_or(DEFAULT_READ_LIMIT);
hits.truncate(limit);
Ok(TreeReadResult {
hits,
total,
tree_id: tree.id,
})
}
fn summary_hit(node: &SummaryNode) -> TreeReadHit {
TreeReadHit {
node_id: node.id.clone(),
node_kind: "summary".to_string(),
level: node.level,
content: node.content.clone(),
score: 0.0,
}
}
#[cfg(test)]
#[path = "read_tests.rs"]
mod tests;