lc_rag/bm25/
chunked.rs

1// src/retrieval/bm25/chunked.rs
2//! BM25 Chunked Retriever - a BM25 retriever supporting the Parent-Child document structure
3//!
4//! Implements the LlamaIndex AutoMerging pattern:
5//! - Documents are split into Parent + Leaf layers
6//! - BM25 searches at the Leaf layer
7//! - AutoMerging merges multiple Leaves under the same Parent
8//! - Supports Bincode persistence
9
10use super::algorithm::{bm25_score, compute_idf, BM25Params};
11use super::tokenizer::Tokenizer;
12use lc_vector_stores::document_store::{ChunkDocument, ChunkedDocumentStoreTrait};
13use lc_vector_stores::{Document, VectorStoreError};
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16use std::path::Path;
17use std::sync::Arc;
18
19// ============================================================================
20// Data structure definitions
21// ============================================================================
22
23// ChunkDocument is now defined in document_store.rs and used directly by BM25
24
25/// AutoMerging configuration
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct AutoMergingConfig {
28    /// Merge threshold: when the ratio of hit Leaves under the same Parent reaches this
29    /// value, they are merged into a Parent document
30    pub merge_threshold: f32,
31    /// The size of a Leaf chunk (in characters)
32    pub leaf_chunk_size: usize,
33    /// The size of a Parent chunk (in characters)
34    pub parent_chunk_size: usize,
35    /// The expected number of Leaves per Parent
36    pub leaves_per_parent: usize,
37}
38
39impl Default for AutoMergingConfig {
40    fn default() -> Self {
41        Self {
42            merge_threshold: 0.5,
43            leaf_chunk_size: 400,
44            parent_chunk_size: 2000,
45            leaves_per_parent: 5,
46        }
47    }
48}
49
50impl AutoMergingConfig {
51    /// Creates an `AutoMergingConfig` with default settings
52    pub fn new() -> Self {
53        Self::default()
54    }
55
56    /// Sets the merge threshold
57    pub fn with_threshold(mut self, threshold: f32) -> Self {
58        self.merge_threshold = threshold;
59        self
60    }
61
62    /// Sets the Leaf chunk size
63    pub fn with_leaf_size(mut self, size: usize) -> Self {
64        self.leaf_chunk_size = size;
65        self
66    }
67
68    /// Sets the Parent chunk size
69    pub fn with_parent_size(mut self, size: usize) -> Self {
70        self.parent_chunk_size = size;
71        self
72    }
73}
74
75/// AutoMerging search result
76#[derive(Debug, Clone)]
77pub struct ChunkedSearchResult {
78    /// The merged Parent document (`None` when merging was not triggered)
79    pub merged_parent: Option<Document>,
80    /// The hit Leaf chunks
81    pub leaf_chunks: Vec<ChunkDocument>,
82    /// The BM25 score of this result
83    pub score: f32,
84    /// The matched query terms
85    pub matched_terms: Vec<String>,
86    /// The id of the owning Parent
87    pub parent_id: String,
88}
89
90impl ChunkedSearchResult {
91    /// Returns the merged result's content: prefers the Parent content, otherwise
92    /// concatenates all Leaf content
93    pub fn content(&self) -> String {
94        if let Some(parent) = &self.merged_parent {
95            parent.content.clone()
96        } else {
97            self.leaf_chunks
98                .iter()
99                .map(|c| c.content.as_str())
100                .collect::<Vec<_>>()
101                .join("\n")
102        }
103    }
104
105    /// Whether AutoMerging merging was triggered
106    pub fn is_merged(&self) -> bool {
107        self.merged_parent.is_some()
108    }
109}
110
111/// Serializable version of the BM25 parameters
112#[derive(Debug, Clone, Serialize, Deserialize)]
113pub struct BM25ParamsData {
114    pub k1: f64,
115    pub b: f64,
116}
117
118impl From<BM25Params> for BM25ParamsData {
119    fn from(params: BM25Params) -> Self {
120        Self {
121            k1: params.k1,
122            b: params.b,
123        }
124    }
125}
126
127impl From<BM25ParamsData> for BM25Params {
128    fn from(data: BM25ParamsData) -> Self {
129        BM25Params::with_values(data.k1, data.b)
130    }
131}
132
133/// Serializable index data (no content; content lives in ChunkedDocumentStore)
134#[derive(Debug, Clone, Serialize, Deserialize)]
135pub struct ChunkedIndexData {
136    /// The list of chunk ids
137    pub chunk_id_list: Vec<String>,
138    /// Each chunk's term-frequency table
139    pub chunk_term_freqs: Vec<HashMap<String, usize>>,
140    /// Inverted index: term -> list of (chunk index, term frequency)
141    pub term_index: HashMap<String, Vec<(usize, usize)>>,
142    /// Parent id -> the list of Leaf chunk indices under that Parent
143    pub parent_to_leaves: HashMap<String, Vec<usize>>,
144    /// Each chunk's document length
145    pub doc_lengths: Vec<usize>,
146    /// The average document length
147    pub avgdl: f64,
148    /// The number of documents
149    pub n_docs: usize,
150    /// BM25 parameters
151    pub params: BM25ParamsData,
152    /// AutoMerging configuration
153    pub config: AutoMergingConfig,
154}
155
156// ============================================================================
157// ChunkedBM25Index index structure
158// ============================================================================
159
160/// A BM25 inverted index supporting the Parent-Child structure
161pub struct ChunkedBM25Index<S: ChunkedDocumentStoreTrait = lc_vector_stores::ChunkedDocumentStore> {
162    store: Arc<S>,
163    chunk_id_list: Vec<String>,
164    chunk_term_freqs: Vec<HashMap<String, usize>>,
165    term_index: HashMap<String, Vec<(usize, usize)>>,
166    parent_to_leaves: HashMap<String, Vec<usize>>,
167    doc_lengths: Vec<usize>,
168    avgdl: f64,
169    n_docs: usize,
170    idf_cache: HashMap<String, f64>,
171    params: BM25Params,
172    tokenizer: Tokenizer,
173    config: AutoMergingConfig,
174    /// 0.22.0 C5 fix: chunk_id → slot. Re-adding an existing chunk id
175    /// **overwrites** its slot (postings / term freqs / doc length) instead
176    /// of appending a duplicate that inflates `n_docs`, skews avgdl/IDF and
177    /// makes the same parent surface multiple times in results.
178    chunk_id_slots: HashMap<String, usize>,
179}
180
181impl<S: ChunkedDocumentStoreTrait> ChunkedBM25Index<S> {
182    /// Creates an index with default settings
183    pub fn new(store: Arc<S>) -> Self {
184        Self::with_config(store, AutoMergingConfig::default())
185    }
186
187    /// Creates an index with the given settings
188    pub fn with_config(store: Arc<S>, config: AutoMergingConfig) -> Self {
189        Self {
190            store,
191            chunk_id_list: Vec::new(),
192            chunk_term_freqs: Vec::new(),
193            term_index: HashMap::new(),
194            parent_to_leaves: HashMap::new(),
195            doc_lengths: Vec::new(),
196            avgdl: 0.0,
197            n_docs: 0,
198            idf_cache: HashMap::new(),
199            params: BM25Params::default(),
200            tokenizer: Tokenizer::new(),
201            config,
202            chunk_id_slots: HashMap::new(),
203        }
204    }
205
206    /// Creates an index with the given BM25 parameters
207    pub fn with_params(store: Arc<S>, params: BM25Params) -> Self {
208        let mut index = Self::new(store);
209        index.params = params;
210        index
211    }
212
213    /// Adds a chunk index (the content is already in the store)
214    ///
215    /// 0.22.0 C5 fix: re-adding an existing `chunk_id` overwrites its slot in
216    /// place (idempotent re-ingest) instead of appending a duplicate that
217    /// inflated `n_docs` and skewed `avgdl` / IDF.
218    pub fn add_chunk_index(
219        &mut self,
220        chunk_id: impl Into<String>,
221        parent_id: impl Into<String>,
222        content: &str,
223    ) {
224        let chunk_id = chunk_id.into();
225        let parent_id = parent_id.into();
226
227        let terms = self.tokenizer.tokenize(content);
228        let term_freq = self.compute_term_freq(&terms);
229        let doc_length: usize = term_freq.values().sum();
230
231        // C5: idempotent overwrite for a known chunk id.
232        if let Some(&slot) = self.chunk_id_slots.get(&chunk_id) {
233            for (term, _) in &self.chunk_term_freqs[slot] {
234                if let Some(postings) = self.term_index.get_mut(term) {
235                    postings.retain(|(idx, _)| *idx != slot);
236                    if postings.is_empty() {
237                        self.term_index.remove(term);
238                    }
239                }
240            }
241            self.chunk_term_freqs[slot] = term_freq;
242            self.doc_lengths[slot] = doc_length;
243            self.update_avgdl();
244            self.idf_cache.clear();
245            return;
246        }
247
248        let chunk_idx = self.n_docs;
249
250        // Update the inverted index
251        for (term, freq) in &term_freq {
252            self.term_index
253                .entry(term.clone())
254                .or_default()
255                .push((chunk_idx, *freq));
256        }
257
258        // Update the parent-to-chunk mapping
259        self.parent_to_leaves
260            .entry(parent_id)
261            .or_default()
262            .push(chunk_idx);
263
264        // Store the chunk_id and term frequencies (needed for BM25 scoring)
265        self.chunk_id_slots.insert(chunk_id.clone(), chunk_idx);
266        self.chunk_id_list.push(chunk_id);
267        self.chunk_term_freqs.push(term_freq.clone());
268
269        self.doc_lengths.push(doc_length);
270        self.n_docs += 1;
271        self.update_avgdl();
272        self.idf_cache.clear();
273    }
274
275    /// Adds chunk indexes in batch
276    pub fn add_chunk_indexes(&mut self, chunks: Vec<(String, String, String)>) {
277        for (chunk_id, parent_id, content) in chunks {
278            self.add_chunk_index(chunk_id, parent_id, &content);
279        }
280    }
281
282    fn compute_term_freq(&self, terms: &[String]) -> HashMap<String, usize> {
283        let mut freq = HashMap::new();
284        for term in terms {
285            *freq.entry(term.clone()).or_insert(0) += 1;
286        }
287        freq
288    }
289
290    fn update_avgdl(&mut self) {
291        if self.n_docs == 0 {
292            self.avgdl = 0.0;
293        } else {
294            let total: usize = self.doc_lengths.iter().sum();
295            self.avgdl = total as f64 / self.n_docs as f64;
296        }
297    }
298
299    fn compute_idf_for_term(&mut self, term: &str) -> f64 {
300        if let Some(idf) = self.idf_cache.get(term) {
301            return *idf;
302        }
303
304        let n = self.term_index.get(term).map(|v| v.len()).unwrap_or(0);
305        let idf = compute_idf(n, self.n_docs);
306        self.idf_cache.insert(term.to_string(), idf);
307        idf
308    }
309
310    /// Gets the chunk id by chunk index
311    pub fn get_chunk_id(&self, chunk_idx: usize) -> Option<&String> {
312        self.chunk_id_list.get(chunk_idx)
313    }
314
315    /// Gets all chunk ids under the given Parent
316    pub fn get_chunk_ids_for_parent(&self, parent_id: &str) -> Vec<&String> {
317        self.parent_to_leaves
318            .get(parent_id)
319            .map(|indices| {
320                indices
321                    .iter()
322                    .filter_map(|idx| self.chunk_id_list.get(*idx))
323                    .collect()
324            })
325            .unwrap_or_default()
326    }
327
328    /// Returns the AutoMerging configuration
329    pub fn config(&self) -> &AutoMergingConfig {
330        &self.config
331    }
332
333    /// Returns the number of indexed documents
334    pub fn n_docs(&self) -> usize {
335        self.n_docs
336    }
337
338    /// Returns the underlying document store
339    pub fn store(&self) -> &Arc<S> {
340        &self.store
341    }
342
343    /// Clears the index data
344    pub fn clear(&mut self) {
345        self.chunk_id_list.clear();
346        self.chunk_term_freqs.clear();
347        self.term_index.clear();
348        self.parent_to_leaves.clear();
349        self.doc_lengths.clear();
350        self.avgdl = 0.0;
351        self.n_docs = 0;
352        self.idf_cache.clear();
353        self.chunk_id_slots.clear();
354    }
355}
356
357impl Default for ChunkedBM25Index<lc_vector_stores::ChunkedDocumentStore> {
358    fn default() -> Self {
359        Self::new(Arc::new(lc_vector_stores::ChunkedDocumentStore::new()))
360    }
361}
362
363// ============================================================================
364// ChunkedBM25Retriever
365// ============================================================================
366
367/// A BM25 retriever based on AutoMerging
368pub struct ChunkedBM25Retriever<
369    S: ChunkedDocumentStoreTrait = lc_vector_stores::ChunkedDocumentStore,
370> {
371    index: ChunkedBM25Index<S>,
372}
373
374impl<S: ChunkedDocumentStoreTrait> ChunkedBM25Retriever<S> {
375    /// Creates a retriever with default settings
376    pub fn new(store: Arc<S>) -> Self {
377        Self {
378            index: ChunkedBM25Index::new(store),
379        }
380    }
381
382    /// Creates a retriever with the given settings
383    pub fn with_config(store: Arc<S>, config: AutoMergingConfig) -> Self {
384        Self {
385            index: ChunkedBM25Index::with_config(store, config),
386        }
387    }
388
389    /// Creates a retriever with the given k1 and b parameters
390    pub fn with_params(store: Arc<S>, k1: f64, b: f64) -> Self {
391        Self {
392            index: ChunkedBM25Index::with_params(store, BM25Params::with_values(k1, b)),
393        }
394    }
395
396    /// Returns the underlying document store
397    pub fn store(&self) -> &Arc<S> {
398        self.index.store()
399    }
400
401    /// Adds a single chunk index (the content is already stored in the store)
402    pub fn add_chunk_index(
403        &mut self,
404        chunk_id: impl Into<String>,
405        parent_id: impl Into<String>,
406        content: &str,
407    ) {
408        self.index.add_chunk_index(chunk_id, parent_id, content);
409    }
410
411    /// Adds chunk indexes in batch
412    pub fn add_chunk_indexes(&mut self, chunks: Vec<(String, String, String)>) {
413        self.index.add_chunk_indexes(chunks);
414    }
415
416    /// Adds a document synchronously: automatically splits Parent/Leaf and builds the index
417    pub fn add_document(&mut self, document: Document) -> Result<(), VectorStoreError> {
418        let parent_id = document
419            .id
420            .clone()
421            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
422
423        // P0-1: a document without an id has the pre-allocated parent_id attached before
424        // storing; otherwise the store generates a fresh uuid, and get_chunks_for_parent
425        // would look up with the wrong key and find nothing.
426        self.index.store.add_parent_document_blocking(
427            document.clone().with_id(parent_id.clone()),
428            self.index.config.leaf_chunk_size,
429        )?;
430
431        let chunks = self
432            .index
433            .store
434            .blocking_get_chunks_for_parent(&parent_id)?;
435
436        for chunk in chunks {
437            self.add_chunk_index(
438                chunk.chunk_id.clone(),
439                chunk.parent_id.clone(),
440                &chunk.content,
441            );
442        }
443
444        Ok(())
445    }
446
447    /// Adds a document asynchronously: automatically splits Parent/Leaf and builds the index
448    pub async fn add_document_async(&mut self, document: Document) -> Result<(), VectorStoreError> {
449        let parent_id = document
450            .id
451            .clone()
452            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
453
454        self.index
455            .store
456            .add_parent_document(
457                document.clone().with_id(parent_id.clone()),
458                self.index.config.leaf_chunk_size,
459            )
460            .await?;
461
462        let chunks = self.index.store.get_chunks_for_parent(&parent_id).await?;
463
464        for chunk in chunks {
465            self.add_chunk_index(
466                chunk.chunk_id.clone(),
467                chunk.parent_id.clone(),
468                &chunk.content,
469            );
470        }
471
472        Ok(())
473    }
474
475    /// Adds documents in batch synchronously
476    pub fn add_documents(&mut self, documents: Vec<Document>) -> Result<(), VectorStoreError> {
477        for doc in documents {
478            self.add_document(doc)?;
479        }
480        Ok(())
481    }
482
483    /// Adds documents in batch asynchronously
484    pub async fn add_documents_async(
485        &mut self,
486        documents: Vec<Document>,
487    ) -> Result<(), VectorStoreError> {
488        for doc in documents {
489            self.add_document_async(doc).await?;
490        }
491        Ok(())
492    }
493
494    /// Runs BM25 retrieval synchronously, returning the top k AutoMerging results
495    pub fn search(&mut self, query: &str, k: usize) -> Vec<ChunkedSearchResult> {
496        if self.index.n_docs == 0 {
497            return Vec::new();
498        }
499
500        let query_terms = self.index.tokenizer.tokenize(query);
501        if query_terms.is_empty() {
502            return Vec::new();
503        }
504
505        let idf_values: HashMap<String, f64> = query_terms
506            .iter()
507            .map(|t| (t.clone(), self.index.compute_idf_for_term(t)))
508            .collect();
509
510        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
511
512        if scored_chunks.is_empty() {
513            return Vec::new();
514        }
515
516        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
517
518        self.auto_merge_sync(top_chunks, k)
519    }
520
521    /// Runs BM25 retrieval asynchronously, returning the top k AutoMerging results
522    pub async fn search_async(&mut self, query: &str, k: usize) -> Vec<ChunkedSearchResult> {
523        if self.index.n_docs == 0 {
524            return Vec::new();
525        }
526
527        let query_terms = self.index.tokenizer.tokenize(query);
528        if query_terms.is_empty() {
529            return Vec::new();
530        }
531
532        let idf_values: HashMap<String, f64> = query_terms
533            .iter()
534            .map(|t| (t.clone(), self.index.compute_idf_for_term(t)))
535            .collect();
536
537        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
538
539        if scored_chunks.is_empty() {
540            return Vec::new();
541        }
542
543        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
544
545        self.auto_merge_async(top_chunks, k).await
546    }
547
548    /// Read-only BM25 retrieval: returns the list of matched parent ids (deduplicated),
549    /// sorted by the best chunk score.
550    ///
551    /// Unlike [`search`](Self::search)/[`search_async`](Self::search_async):
552    /// no AutoMerging ratio gating is applied here — any chunk hit lets its parent through,
553    /// matching the "hit child chunk -> return the whole parent document" semantics that
554    /// [`ParentDocumentRetriever`](crate::parent_document::ParentDocumentRetriever) needs.
555    /// Fully `&self` read-only (idf is not cached), safe to call concurrently.
556    pub fn search_matched_parents(&self, query: &str, k: usize) -> Vec<(String, f32)> {
557        if self.index.n_docs == 0 {
558            return Vec::new();
559        }
560
561        let query_terms = self.index.tokenizer.tokenize(query);
562        if query_terms.is_empty() {
563            return Vec::new();
564        }
565
566        // Read-only idf: does not write idf_cache, avoiding `&mut self`.
567        let idf_values: HashMap<String, f64> = query_terms
568            .iter()
569            .map(|t| {
570                let n = self.index.term_index.get(t).map(|v| v.len()).unwrap_or(0);
571                (t.clone(), compute_idf(n, self.index.n_docs))
572            })
573            .collect();
574
575        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
576        if scored_chunks.is_empty() {
577            return Vec::new();
578        }
579
580        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
581        let parent_stats = self.collect_parent_stats(&top_chunks);
582
583        let mut ranked: Vec<(String, f32)> = parent_stats
584            .into_iter()
585            .map(|(parent_id, leaves)| {
586                let best = leaves.iter().map(|(_, s)| *s as f32).fold(0.0f32, f32::max);
587                (parent_id, best)
588            })
589            .collect();
590        ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
591        ranked.into_iter().take(k).collect()
592    }
593
594    fn auto_merge_sync(
595        &self,
596        scored_chunks: Vec<(usize, f64)>,
597        k: usize,
598    ) -> Vec<ChunkedSearchResult> {
599        let threshold = self.index.config.merge_threshold;
600        let leaves_per_parent = self.index.config.leaves_per_parent;
601
602        let parent_stats = self.collect_parent_stats(&scored_chunks);
603
604        let mut results: Vec<ChunkedSearchResult> = Vec::new();
605
606        for (parent_id, matched_leaves) in parent_stats {
607            // 0.22.0 H-R1: AutoMerging is a "hit-coverage" semantic — the ratio
608            // is matched leaves over the parent's ACTUAL leaf count, not the
609            // fixed `leaves_per_parent` config value. The old divisor made a
610            // 2-leaf parent's 2/2 full hit 0.4 (< threshold → never merged)
611            // while a 10-leaf parent merged on just 5/10 hits. Use the real
612            // `parent_to_leaves` count, falling back to the config only when
613            // the parent is unknown (avoids a divide-by-zero).
614            let total_leaves = self
615                .index
616                .parent_to_leaves
617                .get(&parent_id)
618                .map_or(leaves_per_parent, |leaves| leaves.len())
619                .max(1);
620            let ratio = matched_leaves.len() as f32 / total_leaves as f32;
621
622            let avg_score =
623                matched_leaves.iter().map(|(_, s)| s).sum::<f64>() / matched_leaves.len() as f64;
624
625            let matched_terms = matched_leaves
626                .iter()
627                .filter_map(|(idx, _)| self.index.chunk_term_freqs.get(*idx))
628                .flat_map(|tf| tf.keys().cloned())
629                .collect::<Vec<_>>();
630
631            if ratio >= threshold {
632                let parent_doc = self
633                    .index
634                    .store()
635                    .get_parent_document_blocking(&parent_id)
636                    .ok()
637                    .flatten();
638
639                results.push(ChunkedSearchResult {
640                    merged_parent: parent_doc,
641                    leaf_chunks: Vec::new(),
642                    score: avg_score as f32,
643                    matched_terms,
644                    parent_id,
645                });
646            } else {
647                let leaf_chunks: Vec<ChunkDocument> = matched_leaves
648                    .iter()
649                    .filter_map(|(idx, _)| {
650                        let chunk_id = self.index.get_chunk_id(*idx)?;
651                        let chunk = self
652                            .index
653                            .store()
654                            .get_chunk_blocking(chunk_id)
655                            .ok()
656                            .flatten()?;
657                        Some(chunk)
658                    })
659                    .collect();
660
661                results.push(ChunkedSearchResult {
662                    merged_parent: None,
663                    leaf_chunks,
664                    score: avg_score as f32,
665                    matched_terms,
666                    parent_id,
667                });
668            }
669        }
670
671        results.sort_by(|a, b| {
672            b.score
673                .partial_cmp(&a.score)
674                .unwrap_or(std::cmp::Ordering::Equal)
675        });
676        results.into_iter().take(k).collect()
677    }
678
679    async fn auto_merge_async(
680        &self,
681        scored_chunks: Vec<(usize, f64)>,
682        k: usize,
683    ) -> Vec<ChunkedSearchResult> {
684        let threshold = self.index.config.merge_threshold;
685        let leaves_per_parent = self.index.config.leaves_per_parent;
686
687        let parent_stats = self.collect_parent_stats(&scored_chunks);
688
689        let mut results: Vec<ChunkedSearchResult> = Vec::new();
690
691        for (parent_id, matched_leaves) in parent_stats {
692            // 0.22.0 H-R1: see auto_merge_sync — ratio is hit-coverage over the
693            // parent's actual leaf count, not the fixed config value.
694            let total_leaves = self
695                .index
696                .parent_to_leaves
697                .get(&parent_id)
698                .map_or(leaves_per_parent, |leaves| leaves.len())
699                .max(1);
700            let ratio = matched_leaves.len() as f32 / total_leaves as f32;
701
702            let avg_score =
703                matched_leaves.iter().map(|(_, s)| s).sum::<f64>() / matched_leaves.len() as f64;
704
705            let matched_terms = matched_leaves
706                .iter()
707                .filter_map(|(idx, _)| self.index.chunk_term_freqs.get(*idx))
708                .flat_map(|tf| tf.keys().cloned())
709                .collect::<Vec<_>>();
710
711            if ratio >= threshold {
712                let parent_doc = self
713                    .index
714                    .store()
715                    .get_parent_document(&parent_id)
716                    .await
717                    .ok()
718                    .flatten();
719
720                results.push(ChunkedSearchResult {
721                    merged_parent: parent_doc,
722                    leaf_chunks: Vec::new(),
723                    score: avg_score as f32,
724                    matched_terms,
725                    parent_id,
726                });
727            } else {
728                let mut leaf_chunks = Vec::new();
729                for (idx, _) in matched_leaves {
730                    if let Some(chunk_id) = self.index.get_chunk_id(idx) {
731                        match self.index.store().get_chunk(chunk_id).await {
732                            Ok(Some(chunk)) => leaf_chunks.push(chunk),
733                            Ok(None) => {}
734                            Err(e) => {
735                                // No longer swallow errors silently: a failed read is logged,
736                                // and the chunk is missing from the results
737                                log::error!(
738                                    "failed to read chunk `{}` during retrieval (chunk missing from results): {}",
739                                    chunk_id,
740                                    e
741                                );
742                            }
743                        }
744                    }
745                }
746
747                results.push(ChunkedSearchResult {
748                    merged_parent: None,
749                    leaf_chunks,
750                    score: avg_score as f32,
751                    matched_terms,
752                    parent_id,
753                });
754            }
755        }
756
757        results.sort_by(|a, b| {
758            b.score
759                .partial_cmp(&a.score)
760                .unwrap_or(std::cmp::Ordering::Equal)
761        });
762        results.into_iter().take(k).collect()
763    }
764
765    fn score_chunks(
766        &self,
767        query_terms: &[String],
768        idf_values: &HashMap<String, f64>,
769    ) -> Vec<(usize, f64)> {
770        let mut scored = Vec::new();
771
772        for chunk_idx in 0..self.index.n_docs {
773            if let Some(term_freqs) = self.index.chunk_term_freqs.get(chunk_idx) {
774                let doc_length = *self.index.doc_lengths.get(chunk_idx).unwrap_or(&0);
775
776                let score = bm25_score(
777                    query_terms,
778                    term_freqs,
779                    doc_length,
780                    self.index.avgdl,
781                    idf_values,
782                    &self.index.params,
783                );
784
785                if score > 0.0 {
786                    scored.push((chunk_idx, score));
787                }
788            }
789        }
790
791        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
792        scored
793    }
794
795    fn collect_parent_stats(
796        &self,
797        scored_chunks: &[(usize, f64)],
798    ) -> HashMap<String, Vec<(usize, f64)>> {
799        let mut stats: HashMap<String, Vec<(usize, f64)>> = HashMap::new();
800
801        for (chunk_idx, score) in scored_chunks {
802            if let Some(chunk_id) = self.index.chunk_id_list.get(*chunk_idx) {
803                let parent_id = chunk_id.split("::").next().unwrap_or_default().to_string();
804                stats
805                    .entry(parent_id)
806                    .or_default()
807                    .push((*chunk_idx, *score));
808            }
809        }
810
811        stats
812    }
813
814    /// Gets the parent document by Parent id
815    pub fn get_parent_document(&self, parent_id: &str) -> Option<Document> {
816        self.index
817            .store()
818            .get_parent_document_blocking(parent_id)
819            .ok()
820            .flatten()
821    }
822
823    /// Returns the number of documents in the index
824    pub fn len(&self) -> usize {
825        self.index.n_docs()
826    }
827
828    /// Whether the index is empty
829    pub fn is_empty(&self) -> bool {
830        self.index.n_docs() == 0
831    }
832
833    /// Clears the index
834    pub fn clear(&mut self) {
835        self.index.clear();
836    }
837
838    /// Returns the AutoMerging configuration
839    pub fn config(&self) -> &AutoMergingConfig {
840        self.index.config()
841    }
842
843    // Persistence methods
844    /// Serializes the index data to Bincode and saves it to the given path
845    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), Box<dyn std::error::Error>> {
846        let data = ChunkedIndexData {
847            chunk_id_list: self.index.chunk_id_list.clone(),
848            chunk_term_freqs: self.index.chunk_term_freqs.clone(),
849            term_index: self.index.term_index.clone(),
850            parent_to_leaves: self.index.parent_to_leaves.clone(),
851            doc_lengths: self.index.doc_lengths.clone(),
852            avgdl: self.index.avgdl,
853            n_docs: self.index.n_docs,
854            params: BM25ParamsData::from(self.index.params.clone()),
855            config: self.index.config.clone(),
856        };
857        let encoded = bincode::serialize(&data)?;
858        std::fs::write(path.as_ref(), encoded)?;
859        Ok(())
860    }
861}
862
863impl ChunkedBM25Retriever<lc_vector_stores::ChunkedDocumentStore> {
864    /// Loads Bincode-serialized index data from the given path
865    pub fn load(
866        store: Arc<lc_vector_stores::ChunkedDocumentStore>,
867        path: impl AsRef<Path>,
868    ) -> Result<Self, Box<dyn std::error::Error>> {
869        let bytes = std::fs::read(path.as_ref())?;
870        let data: ChunkedIndexData = bincode::deserialize(&bytes)?;
871        let params: BM25Params = data.params.into();
872
873        // Rebuild the chunk_id → slot map from the persisted id list (C5).
874        let chunk_id_slots: HashMap<String, usize> = data
875            .chunk_id_list
876            .iter()
877            .enumerate()
878            .map(|(idx, id)| (id.clone(), idx))
879            .collect();
880
881        Ok(Self {
882            index: ChunkedBM25Index {
883                store,
884                chunk_id_list: data.chunk_id_list,
885                chunk_term_freqs: data.chunk_term_freqs,
886                term_index: data.term_index,
887                parent_to_leaves: data.parent_to_leaves,
888                doc_lengths: data.doc_lengths,
889                avgdl: data.avgdl,
890                n_docs: data.n_docs,
891                idf_cache: HashMap::new(),
892                params,
893                tokenizer: Tokenizer::new(),
894                config: data.config,
895                chunk_id_slots,
896            },
897        })
898    }
899}
900
901impl Default for ChunkedBM25Retriever<lc_vector_stores::ChunkedDocumentStore> {
902    fn default() -> Self {
903        Self::new(Arc::new(lc_vector_stores::ChunkedDocumentStore::new()))
904    }
905}
lc_rag/bm25/chunked.rs

lc_rag/bm25/
chunked.rs
