feat(indexer): add synonym expansion and cross-reference resolution

- Add enable_synonym_expansion option to IndexOptions for expanding
  keywords with LLM-generated synonyms during indexing to improve
  recall for differently-worded queries
- Implement cross-reference resolution in enrich stage to extract
  and resolve in-document references like "see Section 2.1" or
  "Appendix G" to actual node IDs in the document tree
- Add DocumentTree::set_references method for managing node
  references and children_with_refs for including referenced nodes
- Include resolved reference count in indexing metrics

Author: zTgx

python/src/lib.rs

@@ -107,6 +107,9 @@ fn parse_format(format: &str) -> PyResult<DocumentFormat> {
 ///     generate_description: Whether to generate document description. Default: False.
 ///     include_text: Whether to include node text in the tree. Default: True.
 ///     generate_ids: Whether to generate node IDs. Default: True.
+///     enable_synonym_expansion: Whether to expand keywords with LLM-generated
+///         synonyms during indexing. Improves recall for differently-worded queries.
+///         Default: False.
 #[pyclass(name = "IndexOptions", skip_from_py_object)]
 #[derive(Clone)]
 pub struct PyIndexOptions {
@@ -116,13 +119,14 @@ pub struct PyIndexOptions {
 #[pymethods]
 impl PyIndexOptions {
     #[new]
-    #[pyo3(signature = (mode="default", generate_summaries=true, generate_description=false, include_text=true, generate_ids=true))]
+    #[pyo3(signature = (mode="default", generate_summaries=true, generate_description=false, include_text=true, generate_ids=true, enable_synonym_expansion=false))]
     fn new(
         mode: &str,
         generate_summaries: bool,
         generate_description: bool,
         include_text: bool,
         generate_ids: bool,
+        enable_synonym_expansion: bool,
     ) -> PyResult<Self> {
         let mut opts = IndexOptions::new();
         match mode {
@@ -140,12 +144,13 @@ impl PyIndexOptions {
         opts.generate_description = generate_description;
         opts.include_text = include_text;
         opts.generate_ids = generate_ids;
+        opts.enable_synonym_expansion = enable_synonym_expansion;
         Ok(Self { inner: opts })
     }
 
     fn __repr__(&self) -> String {
         format!(
-            "IndexOptions(mode='{}', generate_summaries={}, generate_description={}, include_text={}, generate_ids={})",
+            "IndexOptions(mode='{}', generate_summaries={}, generate_description={}, include_text={}, generate_ids={}, enable_synonym_expansion={})",
             match self.inner.mode {
                 IndexMode::Default => "default",
                 IndexMode::Force => "force",
@@ -155,6 +160,7 @@ impl PyIndexOptions {
             self.inner.generate_description,
             self.inner.include_text,
             self.inner.generate_ids,
+            self.inner.enable_synonym_expansion,
         )
     }
 }

rust/src/client/indexer.rs

@@ -28,7 +28,7 @@ use uuid::Uuid;
 
 use crate::error::{Error, Result};
 use crate::index::parse::DocumentFormat;
-use crate::index::{IndexInput, IndexMode, PipelineExecutor, PipelineOptions, SummaryStrategy};
+use crate::index::{IndexInput, IndexMode, PipelineExecutor, PipelineOptions, ReasoningIndexConfig, SummaryStrategy};
 use crate::llm::LlmClient;
 use crate::storage::{DocumentMeta, PersistedDocument};
 
@@ -285,6 +285,10 @@ impl IndexerClient {
                 SummaryStrategy::none()
             },
             generate_description: options.generate_description,
+            reasoning_index: ReasoningIndexConfig {
+                enable_synonym_expansion: options.enable_synonym_expansion,
+                ..ReasoningIndexConfig::default()
+            },
             existing_tree,
             ..Default::default()
         }

rust/src/client/types.rs

@@ -206,6 +206,11 @@ pub struct IndexOptions {
 
     /// Whether to generate document description.
     pub generate_description: bool,
+
+    /// Whether to expand keywords with LLM-generated synonyms
+    /// during reasoning index construction. Improves recall for
+    /// queries that use different wording than the document.
+    pub enable_synonym_expansion: bool,
 }
 
 impl Default for IndexOptions {
@@ -216,6 +221,7 @@ impl Default for IndexOptions {
             include_text: true,
             generate_ids: true,
             generate_description: false,
+            enable_synonym_expansion: true,
         }
     }
 }

rust/src/document/tree.rs

@@ -622,6 +622,13 @@ impl DocumentTree {
         }
     }
 
+    /// Set the references for a node.
+    pub fn set_references(&mut self, id: NodeId, references: Vec<super::reference::NodeReference>) {
+        if let Some(node) = self.get_mut(id) {
+            node.references = references;
+        }
+    }
+
     /// Export the tree structure to JSON format.
     pub fn to_structure_json(&self, doc_name: &str) -> DocumentStructure {
         let structure = self.build_structure_nodes(self.root_id);
@@ -795,3 +802,88 @@ impl Default for DocumentTree {
         Self::new("Root", "")
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::document::reference::{NodeReference, RefType};
+
+    #[test]
+    fn test_children_with_refs_no_references() {
+        let mut tree = DocumentTree::new("Root", "root content");
+        let child1 = tree.add_child(tree.root(), "Section 1", "content 1");
+        let child2 = tree.add_child(tree.root(), "Section 2", "content 2");
+
+        let children = tree.children_with_refs(tree.root());
+        assert_eq!(children.len(), 2);
+        assert!(children.contains(&child1));
+        assert!(children.contains(&child2));
+    }
+
+    #[test]
+    fn test_children_with_refs_includes_resolved_references() {
+        let mut tree = DocumentTree::new("Root", "root content");
+        let section1 = tree.add_child(tree.root(), "Section 1", "content 1");
+        let section2 = tree.add_child(tree.root(), "Section 2", "content 2");
+        let appendix = tree.add_child(tree.root(), "Appendix A", "appendix content");
+
+        // Add a resolved reference from Section 1 to Appendix A
+        let refs = vec![NodeReference::resolved(
+            "see Appendix A".to_string(),
+            "A".to_string(),
+            RefType::Appendix,
+            10,
+            appendix,
+            0.9,
+        )];
+        tree.set_references(section1, refs);
+
+        // section1's children_with_refs should include appendix as a reference target
+        let children = tree.children_with_refs(section1);
+        // section1 has no direct children, but has a resolved reference to appendix
+        assert_eq!(children.len(), 1);
+        assert!(children.contains(&appendix));
+    }
+
+    #[test]
+    fn test_children_with_refs_deduplicates() {
+        let mut tree = DocumentTree::new("Root", "root content");
+        let child = tree.add_child(tree.root(), "Section 1", "content 1");
+
+        // Add a reference that points to the same node as an existing child
+        let refs = vec![NodeReference::resolved(
+            "see Section 1".to_string(),
+            "1".to_string(),
+            RefType::Section,
+            5,
+            child,
+            0.8,
+        )];
+        tree.set_references(tree.root(), refs);
+
+        let children = tree.children_with_refs(tree.root());
+        // Should not duplicate
+        assert_eq!(children.len(), 1);
+        assert!(children.contains(&child));
+    }
+
+    #[test]
+    fn test_children_with_refs_unresolved_ignored() {
+        let mut tree = DocumentTree::new("Root", "root content");
+        let child = tree.add_child(tree.root(), "Section 1", "content 1");
+
+        // Add an unresolved reference (target_node = None)
+        let refs = vec![NodeReference::new(
+            "see Section 5".to_string(),
+            "5".to_string(),
+            RefType::Section,
+            5,
+        )];
+        tree.set_references(tree.root(), refs);
+
+        let children = tree.children_with_refs(tree.root());
+        // Unresolved reference should not be included
+        assert_eq!(children.len(), 1);
+        assert!(children.contains(&child));
+    }
+}

rust/src/index/mod.rs

@@ -64,6 +64,7 @@ pub use pipeline::{IndexInput, IndexMetrics, PipelineExecutor, PipelineResult};
 
 // Re-export config types
 pub use config::{IndexMode, PipelineOptions, ThinningConfig};
+pub use crate::document::ReasoningIndexConfig;
 
 // Re-export summary
 pub use summary::SummaryStrategy;

rust/src/index/stages/enrich.rs

@@ -7,7 +7,7 @@ use super::async_trait;
 use std::time::Instant;
 use tracing::info;
 
-use crate::document::{DocumentTree, NodeId, TocView};
+use crate::document::{DocumentTree, NodeId, RefType, ReferenceExtractor, TocView};
 use crate::error::Result;
 
 use super::{AccessPattern, IndexStage, StageResult};
@@ -93,6 +93,46 @@ impl EnrichStage {
             }
         }
     }
+
+    /// Extract and resolve in-document cross-references for all nodes.
+    ///
+    /// Parses content for patterns like "see Section 2.1", "Appendix G", etc.
+    /// and resolves them to actual `NodeId`s in the tree using the retrieval
+    /// index for fast lookup.
+    fn resolve_references(tree: &mut DocumentTree) -> usize {
+        let retrieval_index = tree.build_retrieval_index();
+        let node_ids: Vec<NodeId> = tree.traverse().into_iter().collect();
+        let mut total_resolved = 0;
+
+        for node_id in node_ids {
+            let content = tree.get(node_id).map(|n| n.content.clone()).unwrap_or_default();
+            if content.is_empty() {
+                continue;
+            }
+
+            // Quick check: skip nodes without any reference-like patterns
+            let content_lower = content.to_lowercase();
+            let has_ref_pattern = content_lower.contains("section")
+                || content_lower.contains("appendix")
+                || content_lower.contains("table")
+                || content_lower.contains("figure")
+                || content_lower.contains("page")
+                || content_lower.contains("equation");
+
+            if !has_ref_pattern {
+                continue;
+            }
+
+            let refs = ReferenceExtractor::extract_and_resolve(&content, tree, &retrieval_index);
+            let resolved = refs.iter().filter(|r| r.is_resolved()).count();
+            if resolved > 0 {
+                total_resolved += resolved;
+            }
+            tree.set_references(node_id, refs);
+        }
+
+        total_resolved
+    }
 }
 
 impl Default for EnrichStage {
@@ -142,7 +182,13 @@ impl IndexStage for EnrichStage {
         let (total_tokens, node_count) = Self::calculate_token_stats(tree);
         info!("Total tokens: {}, nodes: {}", total_tokens, node_count);
 
-        // 4. Generate document description
+        // 4. Extract and resolve cross-references
+        let resolved_refs = Self::resolve_references(tree);
+        if resolved_refs > 0 {
+            info!("Resolved {} cross-references", resolved_refs);
+        }
+
+        // 5. Generate document description
         self.generate_description(ctx);
 
         let duration = start.elapsed().as_millis() as u64;
@@ -158,7 +204,42 @@ impl IndexStage for EnrichStage {
         stage_result
             .metadata
             .insert("node_count".to_string(), serde_json::json!(node_count));
+        stage_result
+            .metadata
+            .insert("resolved_references".to_string(), serde_json::json!(resolved_refs));
 
         Ok(stage_result)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_resolve_references_section_ref() {
+        let mut tree = DocumentTree::new("Root", "root content");
+        let s1 = tree.add_child(tree.root(), "Introduction", "Introduction text.");
+        tree.set_structure(s1, "1");
+        let s2 = tree.add_child(tree.root(), "Details", "For details, see Section 1 for more info");
+        tree.set_structure(s2, "2");
+
+        let resolved = EnrichStage::resolve_references(&mut tree);
+        assert_eq!(resolved, 1);
+
+        // Verify the reference was stored on s2 and resolved to s1
+        let refs = tree.get(s2).unwrap().references.clone();
+        assert_eq!(refs.len(), 1);
+        assert_eq!(refs[0].ref_type, RefType::Section);
+        assert_eq!(refs[0].target_node, Some(s1));
+    }
+
+    #[test]
+    fn test_resolve_references_no_refs() {
+        let mut tree = DocumentTree::new("Root", "root content");
+        tree.add_child(tree.root(), "Section 1", "No references here.");
+
+        let resolved = EnrichStage::resolve_references(&mut tree);
+        assert_eq!(resolved, 0);
+    }
+}