Skip to content

Commit 7a45059

Browse files
committed
feat(search): implement cross-document strategy with graph-based boosting
- Add CrossDocumentStrategy and CrossDocumentConfig to support cross-document retrieval with graph-based boosting capabilities - Replace tree.children() with tree.children_with_refs() in ToCNavigator to include reference nodes in tree traversal - Implement ForceCrossDocument strategy preference handling in SearchStage with proper graph attachment when available - Export new cross-document types in strategy module
1 parent 9b123c1 commit 7a45059

3 files changed

Lines changed: 29 additions & 6 deletions

File tree

rust/src/retrieval/search/toc_navigator.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -434,7 +434,7 @@ fn collect_tree_entries(
434434
entries.push((node.title.clone(), summary));
435435
node_ids.push(node_id);
436436

437-
for child_id in tree.children(node_id) {
437+
for child_id in tree.children_with_refs(node_id) {
438438
collect_tree_entries(tree, child_id, entries, node_ids, depth + 1, max_depth);
439439
}
440440
}

rust/src/retrieval/stages/search.rs

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@ use crate::retrieval::search::{
2727
SearchTree, ToCNavigator,
2828
};
2929
use crate::retrieval::strategy::{
30-
HybridConfig, HybridStrategy, KeywordStrategy, LlmStrategy, RetrievalStrategy,
30+
CrossDocumentConfig, CrossDocumentStrategy, DocumentEntry, HybridConfig, HybridStrategy,
31+
KeywordStrategy, LlmStrategy, RetrievalStrategy,
3132
};
3233
use crate::retrieval::types::{
3334
NavigationDecision, ReasoningCandidate, ReasoningStep, StageName, StrategyPreference,
@@ -152,14 +153,35 @@ impl SearchStage {
152153
Arc::new(self.keyword_strategy.clone())
153154
}
154155
}
155-
StrategyPreference::ForceCrossDocument | StrategyPreference::ForcePageRange => {
156+
StrategyPreference::ForceCrossDocument => {
157+
// Build a CrossDocumentStrategy with graph-based boosting
158+
let inner: Box<dyn RetrievalStrategy> =
159+
Box::new(self.keyword_strategy.clone());
160+
161+
let cross_doc =
162+
CrossDocumentStrategy::new(inner).with_config(CrossDocumentConfig::default());
163+
164+
// Attach graph for GraphBoosted merge if available.
165+
// Multi-document trees are collected at the orchestrator level.
166+
let cross_doc = if let Some(ref graph) = ctx.document_graph {
167+
cross_doc.with_graph(graph.clone())
168+
} else {
169+
cross_doc
170+
};
171+
172+
info!(
173+
"Using CrossDocument strategy (graph={})",
174+
ctx.document_graph.is_some()
175+
);
176+
Arc::new(cross_doc)
177+
}
178+
StrategyPreference::ForcePageRange => {
156179
if let Some(ref strategy) = self.hybrid_strategy {
157-
info!("Using Hybrid strategy as fallback for {:?})", preference);
180+
info!("Using Hybrid strategy as fallback for ForcePageRange");
158181
strategy.clone()
159182
} else {
160183
warn!(
161-
"{:?} requires special configuration, falling back to Keyword",
162-
preference
184+
"ForcePageRange requires special configuration, falling back to Keyword"
163185
);
164186
Arc::new(self.keyword_strategy.clone())
165187
}

rust/src/retrieval/strategy/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ mod llm;
1818
mod page_range;
1919
mod r#trait;
2020

21+
pub use cross_document::{CrossDocumentConfig, CrossDocumentStrategy, DocumentEntry};
2122
pub use hybrid::{HybridConfig, HybridStrategy};
2223
pub use keyword::KeywordStrategy;
2324
pub use llm::LlmStrategy;

0 commit comments

Comments
 (0)