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}
175
176impl<S: ChunkedDocumentStoreTrait> ChunkedBM25Index<S> {
177    /// Creates an index with default settings
178    pub fn new(store: Arc<S>) -> Self {
179        Self::with_config(store, AutoMergingConfig::default())
180    }
181
182    /// Creates an index with the given settings
183    pub fn with_config(store: Arc<S>, config: AutoMergingConfig) -> Self {
184        Self {
185            store,
186            chunk_id_list: Vec::new(),
187            chunk_term_freqs: Vec::new(),
188            term_index: HashMap::new(),
189            parent_to_leaves: HashMap::new(),
190            doc_lengths: Vec::new(),
191            avgdl: 0.0,
192            n_docs: 0,
193            idf_cache: HashMap::new(),
194            params: BM25Params::default(),
195            tokenizer: Tokenizer::new(),
196            config,
197        }
198    }
199
200    /// Creates an index with the given BM25 parameters
201    pub fn with_params(store: Arc<S>, params: BM25Params) -> Self {
202        let mut index = Self::new(store);
203        index.params = params;
204        index
205    }
206
207    /// Adds a chunk index (the content is already in the store)
208    pub fn add_chunk_index(
209        &mut self,
210        chunk_id: impl Into<String>,
211        parent_id: impl Into<String>,
212        content: &str,
213    ) {
214        let chunk_idx = self.n_docs;
215        let chunk_id = chunk_id.into();
216        let parent_id = parent_id.into();
217
218        let terms = self.tokenizer.tokenize(content);
219        let term_freq = self.compute_term_freq(&terms);
220
221        // Update the inverted index
222        for (term, freq) in &term_freq {
223            self.term_index
224                .entry(term.clone())
225                .or_default()
226                .push((chunk_idx, *freq));
227        }
228
229        // Update the parent-to-chunk mapping
230        self.parent_to_leaves
231            .entry(parent_id)
232            .or_default()
233            .push(chunk_idx);
234
235        // Store the chunk_id and term frequencies (needed for BM25 scoring)
236        self.chunk_id_list.push(chunk_id);
237        self.chunk_term_freqs.push(term_freq.clone());
238
239        let doc_length: usize = term_freq.values().sum();
240        self.doc_lengths.push(doc_length);
241        self.n_docs += 1;
242        self.update_avgdl();
243        self.idf_cache.clear();
244    }
245
246    /// Adds chunk indexes in batch
247    pub fn add_chunk_indexes(&mut self, chunks: Vec<(String, String, String)>) {
248        for (chunk_id, parent_id, content) in chunks {
249            self.add_chunk_index(chunk_id, parent_id, &content);
250        }
251    }
252
253    fn compute_term_freq(&self, terms: &[String]) -> HashMap<String, usize> {
254        let mut freq = HashMap::new();
255        for term in terms {
256            *freq.entry(term.clone()).or_insert(0) += 1;
257        }
258        freq
259    }
260
261    fn update_avgdl(&mut self) {
262        if self.n_docs == 0 {
263            self.avgdl = 0.0;
264        } else {
265            let total: usize = self.doc_lengths.iter().sum();
266            self.avgdl = total as f64 / self.n_docs as f64;
267        }
268    }
269
270    fn compute_idf_for_term(&mut self, term: &str) -> f64 {
271        if let Some(idf) = self.idf_cache.get(term) {
272            return *idf;
273        }
274
275        let n = self.term_index.get(term).map(|v| v.len()).unwrap_or(0);
276        let idf = compute_idf(n, self.n_docs);
277        self.idf_cache.insert(term.to_string(), idf);
278        idf
279    }
280
281    /// Gets the chunk id by chunk index
282    pub fn get_chunk_id(&self, chunk_idx: usize) -> Option<&String> {
283        self.chunk_id_list.get(chunk_idx)
284    }
285
286    /// Gets all chunk ids under the given Parent
287    pub fn get_chunk_ids_for_parent(&self, parent_id: &str) -> Vec<&String> {
288        self.parent_to_leaves
289            .get(parent_id)
290            .map(|indices| {
291                indices
292                    .iter()
293                    .filter_map(|idx| self.chunk_id_list.get(*idx))
294                    .collect()
295            })
296            .unwrap_or_default()
297    }
298
299    /// Returns the AutoMerging configuration
300    pub fn config(&self) -> &AutoMergingConfig {
301        &self.config
302    }
303
304    /// Returns the number of indexed documents
305    pub fn n_docs(&self) -> usize {
306        self.n_docs
307    }
308
309    /// Returns the underlying document store
310    pub fn store(&self) -> &Arc<S> {
311        &self.store
312    }
313
314    /// Clears the index data
315    pub fn clear(&mut self) {
316        self.chunk_id_list.clear();
317        self.chunk_term_freqs.clear();
318        self.term_index.clear();
319        self.parent_to_leaves.clear();
320        self.doc_lengths.clear();
321        self.avgdl = 0.0;
322        self.n_docs = 0;
323        self.idf_cache.clear();
324    }
325}
326
327impl Default for ChunkedBM25Index<lc_vector_stores::ChunkedDocumentStore> {
328    fn default() -> Self {
329        Self::new(Arc::new(lc_vector_stores::ChunkedDocumentStore::new()))
330    }
331}
332
333// ============================================================================
334// ChunkedBM25Retriever
335// ============================================================================
336
337/// A BM25 retriever based on AutoMerging
338pub struct ChunkedBM25Retriever<
339    S: ChunkedDocumentStoreTrait = lc_vector_stores::ChunkedDocumentStore,
340> {
341    index: ChunkedBM25Index<S>,
342}
343
344impl<S: ChunkedDocumentStoreTrait> ChunkedBM25Retriever<S> {
345    /// Creates a retriever with default settings
346    pub fn new(store: Arc<S>) -> Self {
347        Self {
348            index: ChunkedBM25Index::new(store),
349        }
350    }
351
352    /// Creates a retriever with the given settings
353    pub fn with_config(store: Arc<S>, config: AutoMergingConfig) -> Self {
354        Self {
355            index: ChunkedBM25Index::with_config(store, config),
356        }
357    }
358
359    /// Creates a retriever with the given k1 and b parameters
360    pub fn with_params(store: Arc<S>, k1: f64, b: f64) -> Self {
361        Self {
362            index: ChunkedBM25Index::with_params(store, BM25Params::with_values(k1, b)),
363        }
364    }
365
366    /// Returns the underlying document store
367    pub fn store(&self) -> &Arc<S> {
368        self.index.store()
369    }
370
371    /// Adds a single chunk index (the content is already stored in the store)
372    pub fn add_chunk_index(
373        &mut self,
374        chunk_id: impl Into<String>,
375        parent_id: impl Into<String>,
376        content: &str,
377    ) {
378        self.index.add_chunk_index(chunk_id, parent_id, content);
379    }
380
381    /// Adds chunk indexes in batch
382    pub fn add_chunk_indexes(&mut self, chunks: Vec<(String, String, String)>) {
383        self.index.add_chunk_indexes(chunks);
384    }
385
386    /// Adds a document synchronously: automatically splits Parent/Leaf and builds the index
387    pub fn add_document(&mut self, document: Document) -> Result<(), VectorStoreError> {
388        let parent_id = document
389            .id
390            .clone()
391            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
392
393        // P0-1: a document without an id has the pre-allocated parent_id attached before
394        // storing; otherwise the store generates a fresh uuid, and get_chunks_for_parent
395        // would look up with the wrong key and find nothing.
396        self.index.store.add_parent_document_blocking(
397            document.clone().with_id(parent_id.clone()),
398            self.index.config.leaf_chunk_size,
399        )?;
400
401        let chunks = self
402            .index
403            .store
404            .blocking_get_chunks_for_parent(&parent_id)?;
405
406        for chunk in chunks {
407            self.add_chunk_index(
408                chunk.chunk_id.clone(),
409                chunk.parent_id.clone(),
410                &chunk.content,
411            );
412        }
413
414        Ok(())
415    }
416
417    /// Adds a document asynchronously: automatically splits Parent/Leaf and builds the index
418    pub async fn add_document_async(&mut self, document: Document) -> Result<(), VectorStoreError> {
419        let parent_id = document
420            .id
421            .clone()
422            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
423
424        self.index
425            .store
426            .add_parent_document(
427                document.clone().with_id(parent_id.clone()),
428                self.index.config.leaf_chunk_size,
429            )
430            .await?;
431
432        let chunks = self.index.store.get_chunks_for_parent(&parent_id).await?;
433
434        for chunk in chunks {
435            self.add_chunk_index(
436                chunk.chunk_id.clone(),
437                chunk.parent_id.clone(),
438                &chunk.content,
439            );
440        }
441
442        Ok(())
443    }
444
445    /// Adds documents in batch synchronously
446    pub fn add_documents(&mut self, documents: Vec<Document>) -> Result<(), VectorStoreError> {
447        for doc in documents {
448            self.add_document(doc)?;
449        }
450        Ok(())
451    }
452
453    /// Adds documents in batch asynchronously
454    pub async fn add_documents_async(
455        &mut self,
456        documents: Vec<Document>,
457    ) -> Result<(), VectorStoreError> {
458        for doc in documents {
459            self.add_document_async(doc).await?;
460        }
461        Ok(())
462    }
463
464    /// Runs BM25 retrieval synchronously, returning the top k AutoMerging results
465    pub fn search(&mut self, query: &str, k: usize) -> Vec<ChunkedSearchResult> {
466        if self.index.n_docs == 0 {
467            return Vec::new();
468        }
469
470        let query_terms = self.index.tokenizer.tokenize(query);
471        if query_terms.is_empty() {
472            return Vec::new();
473        }
474
475        let idf_values: HashMap<String, f64> = query_terms
476            .iter()
477            .map(|t| (t.clone(), self.index.compute_idf_for_term(t)))
478            .collect();
479
480        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
481
482        if scored_chunks.is_empty() {
483            return Vec::new();
484        }
485
486        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
487
488        self.auto_merge_sync(top_chunks, k)
489    }
490
491    /// Runs BM25 retrieval asynchronously, returning the top k AutoMerging results
492    pub async fn search_async(&mut self, query: &str, k: usize) -> Vec<ChunkedSearchResult> {
493        if self.index.n_docs == 0 {
494            return Vec::new();
495        }
496
497        let query_terms = self.index.tokenizer.tokenize(query);
498        if query_terms.is_empty() {
499            return Vec::new();
500        }
501
502        let idf_values: HashMap<String, f64> = query_terms
503            .iter()
504            .map(|t| (t.clone(), self.index.compute_idf_for_term(t)))
505            .collect();
506
507        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
508
509        if scored_chunks.is_empty() {
510            return Vec::new();
511        }
512
513        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
514
515        self.auto_merge_async(top_chunks, k).await
516    }
517
518    /// Read-only BM25 retrieval: returns the list of matched parent ids (deduplicated),
519    /// sorted by the best chunk score.
520    ///
521    /// Unlike [`search`](Self::search)/[`search_async`](Self::search_async):
522    /// no AutoMerging ratio gating is applied here — any chunk hit lets its parent through,
523    /// matching the "hit child chunk -> return the whole parent document" semantics that
524    /// [`ParentDocumentRetriever`](crate::parent_document::ParentDocumentRetriever) needs.
525    /// Fully `&self` read-only (idf is not cached), safe to call concurrently.
526    pub fn search_matched_parents(&self, query: &str, k: usize) -> Vec<(String, f32)> {
527        if self.index.n_docs == 0 {
528            return Vec::new();
529        }
530
531        let query_terms = self.index.tokenizer.tokenize(query);
532        if query_terms.is_empty() {
533            return Vec::new();
534        }
535
536        // Read-only idf: does not write idf_cache, avoiding `&mut self`.
537        let idf_values: HashMap<String, f64> = query_terms
538            .iter()
539            .map(|t| {
540                let n = self.index.term_index.get(t).map(|v| v.len()).unwrap_or(0);
541                (t.clone(), compute_idf(n, self.index.n_docs))
542            })
543            .collect();
544
545        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
546        if scored_chunks.is_empty() {
547            return Vec::new();
548        }
549
550        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
551        let parent_stats = self.collect_parent_stats(&top_chunks);
552
553        let mut ranked: Vec<(String, f32)> = parent_stats
554            .into_iter()
555            .map(|(parent_id, leaves)| {
556                let best = leaves.iter().map(|(_, s)| *s as f32).fold(0.0f32, f32::max);
557                (parent_id, best)
558            })
559            .collect();
560        ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
561        ranked.into_iter().take(k).collect()
562    }
563
564    fn auto_merge_sync(
565        &self,
566        scored_chunks: Vec<(usize, f64)>,
567        k: usize,
568    ) -> Vec<ChunkedSearchResult> {
569        let threshold = self.index.config.merge_threshold;
570        let leaves_per_parent = self.index.config.leaves_per_parent;
571
572        let parent_stats = self.collect_parent_stats(&scored_chunks);
573
574        let mut results: Vec<ChunkedSearchResult> = Vec::new();
575
576        for (parent_id, matched_leaves) in parent_stats {
577            let ratio = matched_leaves.len() as f32 / leaves_per_parent as f32;
578
579            let avg_score =
580                matched_leaves.iter().map(|(_, s)| s).sum::<f64>() / matched_leaves.len() as f64;
581
582            let matched_terms = matched_leaves
583                .iter()
584                .filter_map(|(idx, _)| self.index.chunk_term_freqs.get(*idx))
585                .flat_map(|tf| tf.keys().cloned())
586                .collect::<Vec<_>>();
587
588            if ratio >= threshold {
589                let parent_doc = self
590                    .index
591                    .store()
592                    .get_parent_document_blocking(&parent_id)
593                    .ok()
594                    .flatten();
595
596                results.push(ChunkedSearchResult {
597                    merged_parent: parent_doc,
598                    leaf_chunks: Vec::new(),
599                    score: avg_score as f32,
600                    matched_terms,
601                    parent_id,
602                });
603            } else {
604                let leaf_chunks: Vec<ChunkDocument> = matched_leaves
605                    .iter()
606                    .filter_map(|(idx, _)| {
607                        let chunk_id = self.index.get_chunk_id(*idx)?;
608                        let chunk = self
609                            .index
610                            .store()
611                            .get_chunk_blocking(chunk_id)
612                            .ok()
613                            .flatten()?;
614                        Some(chunk)
615                    })
616                    .collect();
617
618                results.push(ChunkedSearchResult {
619                    merged_parent: None,
620                    leaf_chunks,
621                    score: avg_score as f32,
622                    matched_terms,
623                    parent_id,
624                });
625            }
626        }
627
628        results.sort_by(|a, b| {
629            b.score
630                .partial_cmp(&a.score)
631                .unwrap_or(std::cmp::Ordering::Equal)
632        });
633        results.into_iter().take(k).collect()
634    }
635
636    async fn auto_merge_async(
637        &self,
638        scored_chunks: Vec<(usize, f64)>,
639        k: usize,
640    ) -> Vec<ChunkedSearchResult> {
641        let threshold = self.index.config.merge_threshold;
642        let leaves_per_parent = self.index.config.leaves_per_parent;
643
644        let parent_stats = self.collect_parent_stats(&scored_chunks);
645
646        let mut results: Vec<ChunkedSearchResult> = Vec::new();
647
648        for (parent_id, matched_leaves) in parent_stats {
649            let ratio = matched_leaves.len() as f32 / leaves_per_parent as f32;
650
651            let avg_score =
652                matched_leaves.iter().map(|(_, s)| s).sum::<f64>() / matched_leaves.len() as f64;
653
654            let matched_terms = matched_leaves
655                .iter()
656                .filter_map(|(idx, _)| self.index.chunk_term_freqs.get(*idx))
657                .flat_map(|tf| tf.keys().cloned())
658                .collect::<Vec<_>>();
659
660            if ratio >= threshold {
661                let parent_doc = self
662                    .index
663                    .store()
664                    .get_parent_document(&parent_id)
665                    .await
666                    .ok()
667                    .flatten();
668
669                results.push(ChunkedSearchResult {
670                    merged_parent: parent_doc,
671                    leaf_chunks: Vec::new(),
672                    score: avg_score as f32,
673                    matched_terms,
674                    parent_id,
675                });
676            } else {
677                let mut leaf_chunks = Vec::new();
678                for (idx, _) in matched_leaves {
679                    if let Some(chunk_id) = self.index.get_chunk_id(idx) {
680                        match self.index.store().get_chunk(chunk_id).await {
681                            Ok(Some(chunk)) => leaf_chunks.push(chunk),
682                            Ok(None) => {}
683                            Err(e) => {
684                                // No longer swallow errors silently: a failed read is logged,
685                                // and the chunk is missing from the results
686                                log::error!(
687                                    "failed to read chunk `{}` during retrieval (chunk missing from results): {}",
688                                    chunk_id,
689                                    e
690                                );
691                            }
692                        }
693                    }
694                }
695
696                results.push(ChunkedSearchResult {
697                    merged_parent: None,
698                    leaf_chunks,
699                    score: avg_score as f32,
700                    matched_terms,
701                    parent_id,
702                });
703            }
704        }
705
706        results.sort_by(|a, b| {
707            b.score
708                .partial_cmp(&a.score)
709                .unwrap_or(std::cmp::Ordering::Equal)
710        });
711        results.into_iter().take(k).collect()
712    }
713
714    fn score_chunks(
715        &self,
716        query_terms: &[String],
717        idf_values: &HashMap<String, f64>,
718    ) -> Vec<(usize, f64)> {
719        let mut scored = Vec::new();
720
721        for chunk_idx in 0..self.index.n_docs {
722            if let Some(term_freqs) = self.index.chunk_term_freqs.get(chunk_idx) {
723                let doc_length = *self.index.doc_lengths.get(chunk_idx).unwrap_or(&0);
724
725                let score = bm25_score(
726                    query_terms,
727                    term_freqs,
728                    doc_length,
729                    self.index.avgdl,
730                    idf_values,
731                    &self.index.params,
732                );
733
734                if score > 0.0 {
735                    scored.push((chunk_idx, score));
736                }
737            }
738        }
739
740        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
741        scored
742    }
743
744    fn collect_parent_stats(
745        &self,
746        scored_chunks: &[(usize, f64)],
747    ) -> HashMap<String, Vec<(usize, f64)>> {
748        let mut stats: HashMap<String, Vec<(usize, f64)>> = HashMap::new();
749
750        for (chunk_idx, score) in scored_chunks {
751            if let Some(chunk_id) = self.index.chunk_id_list.get(*chunk_idx) {
752                let parent_id = chunk_id.split("::").next().unwrap_or_default().to_string();
753                stats
754                    .entry(parent_id)
755                    .or_default()
756                    .push((*chunk_idx, *score));
757            }
758        }
759
760        stats
761    }
762
763    /// Gets the parent document by Parent id
764    pub fn get_parent_document(&self, parent_id: &str) -> Option<Document> {
765        self.index
766            .store()
767            .get_parent_document_blocking(parent_id)
768            .ok()
769            .flatten()
770    }
771
772    /// Returns the number of documents in the index
773    pub fn len(&self) -> usize {
774        self.index.n_docs()
775    }
776
777    /// Whether the index is empty
778    pub fn is_empty(&self) -> bool {
779        self.index.n_docs() == 0
780    }
781
782    /// Clears the index
783    pub fn clear(&mut self) {
784        self.index.clear();
785    }
786
787    /// Returns the AutoMerging configuration
788    pub fn config(&self) -> &AutoMergingConfig {
789        self.index.config()
790    }
791
792    // Persistence methods
793    /// Serializes the index data to Bincode and saves it to the given path
794    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), Box<dyn std::error::Error>> {
795        let data = ChunkedIndexData {
796            chunk_id_list: self.index.chunk_id_list.clone(),
797            chunk_term_freqs: self.index.chunk_term_freqs.clone(),
798            term_index: self.index.term_index.clone(),
799            parent_to_leaves: self.index.parent_to_leaves.clone(),
800            doc_lengths: self.index.doc_lengths.clone(),
801            avgdl: self.index.avgdl,
802            n_docs: self.index.n_docs,
803            params: BM25ParamsData::from(self.index.params.clone()),
804            config: self.index.config.clone(),
805        };
806        let encoded = bincode::serialize(&data)?;
807        std::fs::write(path.as_ref(), encoded)?;
808        Ok(())
809    }
810}
811
812impl ChunkedBM25Retriever<lc_vector_stores::ChunkedDocumentStore> {
813    /// Loads Bincode-serialized index data from the given path
814    pub fn load(
815        store: Arc<lc_vector_stores::ChunkedDocumentStore>,
816        path: impl AsRef<Path>,
817    ) -> Result<Self, Box<dyn std::error::Error>> {
818        let bytes = std::fs::read(path.as_ref())?;
819        let data: ChunkedIndexData = bincode::deserialize(&bytes)?;
820        let params: BM25Params = data.params.into();
821
822        Ok(Self {
823            index: ChunkedBM25Index {
824                store,
825                chunk_id_list: data.chunk_id_list,
826                chunk_term_freqs: data.chunk_term_freqs,
827                term_index: data.term_index,
828                parent_to_leaves: data.parent_to_leaves,
829                doc_lengths: data.doc_lengths,
830                avgdl: data.avgdl,
831                n_docs: data.n_docs,
832                idf_cache: HashMap::new(),
833                params,
834                tokenizer: Tokenizer::new(),
835                config: data.config,
836            },
837        })
838    }
839}
840
841impl Default for ChunkedBM25Retriever<lc_vector_stores::ChunkedDocumentStore> {
842    fn default() -> Self {
843        Self::new(Arc::new(lc_vector_stores::ChunkedDocumentStore::new()))
844    }
845}
lc_rag/bm25/chunked.rs

lc_rag/bm25/
chunked.rs
