use std::collections::HashSet;
use crate::graph::walk::{ancestors_inclusion, descendants_inclusion, outbound_reference};
use crate::graph::{Graph, GraphContext};
use crate::model::node::{NodeIter, NodePointer};
use crate::model::{Key, NodeId};
use crate::query::{self, Filter};
use crate::tokens::{
apply_budget, count_tokens, truncate_to_tokens, truncation_marker, Budget, Truncation,
};
use itertools::Itertools;
use serde::Serialize;
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct EdgeRef {
pub key: String,
pub title: String,
pub section_path: Vec<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DocumentOutput {
pub key: String,
pub title: String,
pub content: String,
pub references: Vec<EdgeRef>,
pub includes: Vec<EdgeRef>,
pub referenced_by: Vec<EdgeRef>,
pub included_by: Vec<EdgeRef>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RetrieveOutput {
pub documents: Vec<DocumentOutput>,
#[serde(skip)]
pub truncation: Truncation,
}
#[derive(Debug, Clone, Default)]
pub struct RetrieveOptions {
pub depth: u8,
pub context: u8,
pub links: bool,
pub backlinks: bool,
pub exclude: HashSet<Key>,
pub children: bool,
pub filter: Option<Filter>,
pub limit: Option<usize>,
pub max_tokens: Option<usize>,
pub max_document_tokens: Option<usize>,
}
pub struct DocumentReader<'a> {
graph: &'a Graph,
}
impl<'a> DocumentReader<'a> {
pub fn new(graph: &'a Graph) -> Self {
Self { graph }
}
pub fn retrieve(&self, key: &Key, options: &RetrieveOptions) -> RetrieveOutput {
let mut documents = Vec::new();
let mut seen_keys = HashSet::new();
let keys_to_process = self.collect_document_keys(key, options);
for doc_key in keys_to_process {
if seen_keys.contains(&doc_key) || options.exclude.contains(&doc_key) {
continue;
}
seen_keys.insert(doc_key.clone());
let doc_output = self.build_document_output(&doc_key, options);
documents.push(doc_output);
}
let truncation = apply_document_budget(&mut documents, options);
RetrieveOutput {
documents,
truncation,
}
}
pub fn retrieve_many(&self, keys: &[Key], options: &RetrieveOptions) -> RetrieveOutput {
let effective_keys: Vec<Key> = match (&options.filter, keys.is_empty()) {
(Some(f), true) => query::evaluate(f, self.graph),
(Some(f), false) => {
let set: HashSet<Key> = query::evaluate(f, self.graph).into_iter().collect();
keys.iter().filter(|k| set.contains(k)).cloned().collect()
}
(None, _) => keys.to_vec(),
};
let mut documents = Vec::new();
let mut seen_keys = HashSet::new();
for key in &effective_keys {
let keys_to_process = self.collect_document_keys(key, options);
for doc_key in keys_to_process {
if seen_keys.contains(&doc_key) || options.exclude.contains(&doc_key) {
continue;
}
seen_keys.insert(doc_key.clone());
let doc_output = self.build_document_output(&doc_key, options);
documents.push(doc_output);
}
}
let truncation = apply_document_budget(&mut documents, options);
RetrieveOutput {
documents,
truncation,
}
}
fn collect_document_keys(&self, key: &Key, options: &RetrieveOptions) -> Vec<Key> {
let mut result: Vec<Key> = vec![key.clone()];
let mut seen: HashSet<Key> = HashSet::from([key.clone()]);
let push = |k: Key, result: &mut Vec<Key>, seen: &mut HashSet<Key>| {
if seen.insert(k.clone()) {
result.push(k);
}
};
if options.depth > 0 {
let mut desc: Vec<(Key, u32)> =
descendants_inclusion(self.graph, key, options.depth as u32)
.into_iter()
.collect();
desc.sort_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
for (k, _) in desc {
push(k, &mut result, &mut seen);
}
}
if options.context > 0 {
let mut anc: Vec<(Key, u32)> =
ancestors_inclusion(self.graph, key, options.context as u32)
.into_iter()
.collect();
anc.sort_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
for (k, _) in anc {
push(k, &mut result, &mut seen);
}
if options.depth > 0 {
let mut sub_doc_keys: Vec<Key> = descendants_inclusion(self.graph, key, 1)
.into_keys()
.collect();
sub_doc_keys.sort();
for sub_key in sub_doc_keys {
let mut sub_anc: Vec<(Key, u32)> =
ancestors_inclusion(self.graph, &sub_key, options.context as u32)
.into_iter()
.collect();
sub_anc.sort_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
for (k, _) in sub_anc {
push(k, &mut result, &mut seen);
}
}
}
}
if options.links {
let mut links: Vec<(Key, u32)> =
outbound_reference(self.graph, key, 1).into_iter().collect();
links.sort_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
for (k, _) in links {
push(k, &mut result, &mut seen);
}
}
result
}
fn build_document_output(&self, key: &Key, options: &RetrieveOptions) -> DocumentOutput {
let title = self
.graph
.get_key_title(key)
.unwrap_or_else(|| key.to_string());
let content = self.get_document_content(key);
let included_by = self.get_parent_documents(key);
let includes = if options.children {
self.get_child_documents(key)
} else {
Vec::new()
};
let referenced_by = if options.backlinks {
self.get_backlinks(key)
} else {
Vec::new()
};
let references = if options.links {
crate::query::edges::references(self.graph, key)
} else {
Vec::new()
};
DocumentOutput {
key: key.to_string(),
title,
content,
references,
includes,
referenced_by,
included_by,
}
}
fn get_document_content(&self, key: &Key) -> String {
self.graph.to_markdown_skip_frontmatter(key)
}
fn get_parent_documents(&self, key: &Key) -> Vec<EdgeRef> {
let refs = self.graph.get_inclusion_edges_to(key);
let mut parents = Vec::new();
for ref_id in refs {
let node = self.graph.node(ref_id);
if let Some(doc_node) = node.to_document() {
if let Some(doc_key) = doc_node.document_key() {
let title = self
.graph
.get_key_title(&doc_key)
.unwrap_or_else(|| doc_key.to_string());
let section_path = self.get_section_path(ref_id);
parents.push(EdgeRef {
key: doc_key.to_string(),
title,
section_path,
});
}
}
}
let mut parents: Vec<EdgeRef> = parents.into_iter().unique_by(|p| p.key.clone()).collect();
parents.sort_by(|a, b| a.key.cmp(&b.key));
parents
}
fn get_child_documents(&self, key: &Key) -> Vec<EdgeRef> {
let refs = self.graph.get_inclusion_edges_in(key);
let mut children = Vec::new();
for ref_id in refs {
if let Some(ref_key) = self.graph.graph_node(ref_id).ref_key() {
let title = self
.graph
.get_key_title(&ref_key)
.unwrap_or_else(|| ref_key.to_string());
let section_path = self.get_section_path(ref_id);
children.push(EdgeRef {
key: ref_key.to_string(),
title,
section_path,
});
}
}
let mut children: Vec<EdgeRef> =
children.into_iter().unique_by(|c| c.key.clone()).collect();
children.sort_by(|a, b| a.key.cmp(&b.key));
children
}
fn get_backlinks(&self, key: &Key) -> Vec<EdgeRef> {
let inline_refs = self.graph.get_reference_edges_to(key);
let mut backlinks = Vec::new();
let mut seen_keys = HashSet::new();
for ref_id in inline_refs {
let node = self.graph.node(ref_id);
if let Some(doc_node) = node.to_document() {
if let Some(doc_key) = doc_node.document_key() {
if seen_keys.contains(&doc_key) {
continue;
}
seen_keys.insert(doc_key.clone());
let title = self
.graph
.get_key_title(&doc_key)
.unwrap_or_else(|| doc_key.to_string());
let section_path = self.get_section_path(ref_id);
backlinks.push(EdgeRef {
key: doc_key.to_string(),
title,
section_path,
});
}
}
}
backlinks.sort_by(|a, b| a.key.cmp(&b.key));
backlinks
}
fn get_section_path(&self, node_id: NodeId) -> Vec<String> {
let mut path = Vec::new();
let mut current = self.graph.node(node_id);
while let Some(parent) = current.to_parent() {
if parent.is_section() && parent.is_header() {
if let Some(grandparent) = parent.to_parent() {
if !grandparent.is_document() {
let text = parent.plain_text().trim().to_string();
path.push(text);
}
}
}
if parent.is_document() {
break;
}
current = parent;
}
path.reverse();
path
}
}
fn apply_document_budget(
documents: &mut Vec<DocumentOutput>,
options: &RetrieveOptions,
) -> Truncation {
let matched = documents.len();
let budget = Budget {
limit: options.limit,
max_tokens: options.max_tokens,
max_document_tokens: options.max_document_tokens,
};
apply_budget(
documents,
&budget,
matched,
|doc| doc.key.clone(),
document_payload_tokens,
|doc, max| {
let (head, omitted) = truncate_to_tokens(&doc.content, max);
if omitted > 0 {
doc.content = format!("{}{}", head, truncation_marker(omitted));
Some(omitted)
} else {
None
}
},
)
}
fn document_payload_tokens(doc: &DocumentOutput) -> usize {
count_tokens(&doc.content) + edge_tokens(doc)
}
fn edge_tokens(doc: &DocumentOutput) -> usize {
[
&doc.references,
&doc.includes,
&doc.referenced_by,
&doc.included_by,
]
.into_iter()
.filter(|edges| !edges.is_empty())
.filter_map(|edges| serde_yaml::to_string(edges).ok())
.map(|s| count_tokens(&s))
.sum()
}