lc_rag/bm25/
chunked.rs

1// src/retrieval/bm25/chunked.rs
2//! BM25 Chunked Retriever - 支持 Parent-Child 文档结构的 BM25 检索器
3//!
4//! 基于 LlamaIndex AutoMerging 模式实现：
5//! - 文档拆分为 Parent + Leaf 两层
6//! - BM25 在 Leaf 层搜索
7//! - AutoMerging 合并同一 Parent 的多个 Leaf
8//! - 支持 Bincode 持久化
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// 数据结构定义
21// ============================================================================
22
23// ChunkDocument 现在在 document_store.rs 中定义，BM25 直接使用
24
25/// AutoMerging 配置
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct AutoMergingConfig {
28    /// 合并阈值：同一 Parent 下命中 Leaf 占比达到该比例时合并为 Parent 文档
29    pub merge_threshold: f32,
30    /// Leaf chunk 的大小（字符数）
31    pub leaf_chunk_size: usize,
32    /// Parent chunk 的大小（字符数）
33    pub parent_chunk_size: usize,
34    /// 每个 Parent 下期望的 Leaf 数量
35    pub leaves_per_parent: usize,
36}
37
38impl Default for AutoMergingConfig {
39    fn default() -> Self {
40        Self {
41            merge_threshold: 0.5,
42            leaf_chunk_size: 400,
43            parent_chunk_size: 2000,
44            leaves_per_parent: 5,
45        }
46    }
47}
48
49impl AutoMergingConfig {
50    /// 创建使用默认配置的 `AutoMergingConfig`
51    pub fn new() -> Self {
52        Self::default()
53    }
54
55    /// 设置合并阈值
56    pub fn with_threshold(mut self, threshold: f32) -> Self {
57        self.merge_threshold = threshold;
58        self
59    }
60
61    /// 设置 Leaf chunk 大小
62    pub fn with_leaf_size(mut self, size: usize) -> Self {
63        self.leaf_chunk_size = size;
64        self
65    }
66
67    /// 设置 Parent chunk 大小
68    pub fn with_parent_size(mut self, size: usize) -> Self {
69        self.parent_chunk_size = size;
70        self
71    }
72}
73
74/// AutoMerging 搜索结果
75#[derive(Debug, Clone)]
76pub struct ChunkedSearchResult {
77    /// 合并得到的 Parent 文档（若未触发合并则为 `None`）
78    pub merged_parent: Option<Document>,
79    /// 命中的 Leaf chunks
80    pub leaf_chunks: Vec<ChunkDocument>,
81    /// 该结果的 BM25 评分
82    pub score: f32,
83    /// 命中的查询词项
84    pub matched_terms: Vec<String>,
85    /// 所属 Parent 的 id
86    pub parent_id: String,
87}
88
89impl ChunkedSearchResult {
90    /// 返回合并结果的内容：优先返回 Parent 内容，否则拼接所有 Leaf 内容
91    pub fn content(&self) -> String {
92        if let Some(parent) = &self.merged_parent {
93            parent.content.clone()
94        } else {
95            self.leaf_chunks
96                .iter()
97                .map(|c| c.content.as_str())
98                .collect::<Vec<_>>()
99                .join("\n")
100        }
101    }
102
103    /// 是否触发了 AutoMerging 合并
104    pub fn is_merged(&self) -> bool {
105        self.merged_parent.is_some()
106    }
107}
108
109/// BM25 参数的可序列化版本
110#[derive(Debug, Clone, Serialize, Deserialize)]
111pub struct BM25ParamsData {
112    pub k1: f64,
113    pub b: f64,
114}
115
116impl From<BM25Params> for BM25ParamsData {
117    fn from(params: BM25Params) -> Self {
118        Self {
119            k1: params.k1,
120            b: params.b,
121        }
122    }
123}
124
125impl From<BM25ParamsData> for BM25Params {
126    fn from(data: BM25ParamsData) -> Self {
127        BM25Params::with_values(data.k1, data.b)
128    }
129}
130
131/// 可序列化的索引数据（不含内容，内容在ChunkedDocumentStore中）
132#[derive(Debug, Clone, Serialize, Deserialize)]
133pub struct ChunkedIndexData {
134    /// chunk 的 id 列表
135    pub chunk_id_list: Vec<String>,
136    /// 每个 chunk 的词频表
137    pub chunk_term_freqs: Vec<HashMap<String, usize>>,
138    /// 倒排索引：词项 -> (chunk 下标, 词频) 列表
139    pub term_index: HashMap<String, Vec<(usize, usize)>>,
140    /// Parent id -> 该 Parent 下 Leaf chunk 下标列表
141    pub parent_to_leaves: HashMap<String, Vec<usize>>,
142    /// 每个 chunk 的文档长度
143    pub doc_lengths: Vec<usize>,
144    /// 平均文档长度
145    pub avgdl: f64,
146    /// 文档数量
147    pub n_docs: usize,
148    /// BM25 参数
149    pub params: BM25ParamsData,
150    /// AutoMerging 配置
151    pub config: AutoMergingConfig,
152}
153
154// ============================================================================
155// ChunkedBM25Index 索引结构
156// ============================================================================
157
158/// 支持 Parent-Child 结构的 BM25 倒排索引
159pub struct ChunkedBM25Index<S: ChunkedDocumentStoreTrait = lc_vector_stores::ChunkedDocumentStore> {
160    store: Arc<S>,
161    chunk_id_list: Vec<String>,
162    chunk_term_freqs: Vec<HashMap<String, usize>>,
163    term_index: HashMap<String, Vec<(usize, usize)>>,
164    parent_to_leaves: HashMap<String, Vec<usize>>,
165    doc_lengths: Vec<usize>,
166    avgdl: f64,
167    n_docs: usize,
168    idf_cache: HashMap<String, f64>,
169    params: BM25Params,
170    tokenizer: Tokenizer,
171    config: AutoMergingConfig,
172}
173
174impl<S: ChunkedDocumentStoreTrait> ChunkedBM25Index<S> {
175    /// 使用默认配置创建索引
176    pub fn new(store: Arc<S>) -> Self {
177        Self::with_config(store, AutoMergingConfig::default())
178    }
179
180    /// 使用指定配置创建索引
181    pub fn with_config(store: Arc<S>, config: AutoMergingConfig) -> Self {
182        Self {
183            store,
184            chunk_id_list: Vec::new(),
185            chunk_term_freqs: Vec::new(),
186            term_index: HashMap::new(),
187            parent_to_leaves: HashMap::new(),
188            doc_lengths: Vec::new(),
189            avgdl: 0.0,
190            n_docs: 0,
191            idf_cache: HashMap::new(),
192            params: BM25Params::default(),
193            tokenizer: Tokenizer::new(),
194            config,
195        }
196    }
197
198    /// 使用指定 BM25 参数创建索引
199    pub fn with_params(store: Arc<S>, params: BM25Params) -> Self {
200        let mut index = Self::new(store);
201        index.params = params;
202        index
203    }
204
205    /// 添加chunk索引（内容已在store中）
206    pub fn add_chunk_index(
207        &mut self,
208        chunk_id: impl Into<String>,
209        parent_id: impl Into<String>,
210        content: &str,
211    ) {
212        let chunk_idx = self.n_docs;
213        let chunk_id = chunk_id.into();
214        let parent_id = parent_id.into();
215
216        let terms = self.tokenizer.tokenize(content);
217        let term_freq = self.compute_term_freq(&terms);
218
219        // 更新倒排索引
220        for (term, freq) in &term_freq {
221            self.term_index
222                .entry(term.clone())
223                .or_default()
224                .push((chunk_idx, *freq));
225        }
226
227        // 更新parent到chunk的映射
228        self.parent_to_leaves
229            .entry(parent_id)
230            .or_default()
231            .push(chunk_idx);
232
233        // 存储chunk_id和词频（BM25计算需要）
234        self.chunk_id_list.push(chunk_id);
235        self.chunk_term_freqs.push(term_freq.clone());
236
237        let doc_length: usize = term_freq.values().sum();
238        self.doc_lengths.push(doc_length);
239        self.n_docs += 1;
240        self.update_avgdl();
241        self.idf_cache.clear();
242    }
243
244    /// 批量添加chunk索引
245    pub fn add_chunk_indexes(&mut self, chunks: Vec<(String, String, String)>) {
246        for (chunk_id, parent_id, content) in chunks {
247            self.add_chunk_index(chunk_id, parent_id, &content);
248        }
249    }
250
251    fn compute_term_freq(&self, terms: &[String]) -> HashMap<String, usize> {
252        let mut freq = HashMap::new();
253        for term in terms {
254            *freq.entry(term.clone()).or_insert(0) += 1;
255        }
256        freq
257    }
258
259    fn update_avgdl(&mut self) {
260        if self.n_docs == 0 {
261            self.avgdl = 0.0;
262        } else {
263            let total: usize = self.doc_lengths.iter().sum();
264            self.avgdl = total as f64 / self.n_docs as f64;
265        }
266    }
267
268    fn compute_idf_for_term(&mut self, term: &str) -> f64 {
269        if let Some(idf) = self.idf_cache.get(term) {
270            return *idf;
271        }
272
273        let n = self.term_index.get(term).map(|v| v.len()).unwrap_or(0);
274        let idf = compute_idf(n, self.n_docs);
275        self.idf_cache.insert(term.to_string(), idf);
276        idf
277    }
278
279    /// 按 chunk 下标获取 chunk id
280    pub fn get_chunk_id(&self, chunk_idx: usize) -> Option<&String> {
281        self.chunk_id_list.get(chunk_idx)
282    }
283
284    /// 获取指定 Parent 下的所有 chunk id
285    pub fn get_chunk_ids_for_parent(&self, parent_id: &str) -> Vec<&String> {
286        self.parent_to_leaves
287            .get(parent_id)
288            .map(|indices| {
289                indices
290                    .iter()
291                    .filter_map(|idx| self.chunk_id_list.get(*idx))
292                    .collect()
293            })
294            .unwrap_or_default()
295    }
296
297    /// 返回 AutoMerging 配置
298    pub fn config(&self) -> &AutoMergingConfig {
299        &self.config
300    }
301
302    /// 返回已索引的文档数量
303    pub fn n_docs(&self) -> usize {
304        self.n_docs
305    }
306
307    /// 返回底层文档存储
308    pub fn store(&self) -> &Arc<S> {
309        &self.store
310    }
311
312    /// 清空索引数据
313    pub fn clear(&mut self) {
314        self.chunk_id_list.clear();
315        self.chunk_term_freqs.clear();
316        self.term_index.clear();
317        self.parent_to_leaves.clear();
318        self.doc_lengths.clear();
319        self.avgdl = 0.0;
320        self.n_docs = 0;
321        self.idf_cache.clear();
322    }
323}
324
325impl Default for ChunkedBM25Index<lc_vector_stores::ChunkedDocumentStore> {
326    fn default() -> Self {
327        Self::new(Arc::new(lc_vector_stores::ChunkedDocumentStore::new()))
328    }
329}
330
331// ============================================================================
332// ChunkedBM25Retriever 检索器
333// ============================================================================
334
335/// 基于 AutoMerging 的 BM25 检索器
336pub struct ChunkedBM25Retriever<
337    S: ChunkedDocumentStoreTrait = lc_vector_stores::ChunkedDocumentStore,
338> {
339    index: ChunkedBM25Index<S>,
340}
341
342impl<S: ChunkedDocumentStoreTrait> ChunkedBM25Retriever<S> {
343    /// 使用默认配置创建检索器
344    pub fn new(store: Arc<S>) -> Self {
345        Self {
346            index: ChunkedBM25Index::new(store),
347        }
348    }
349
350    /// 使用指定配置创建检索器
351    pub fn with_config(store: Arc<S>, config: AutoMergingConfig) -> Self {
352        Self {
353            index: ChunkedBM25Index::with_config(store, config),
354        }
355    }
356
357    /// 使用指定的 k1、b 参数创建检索器
358    pub fn with_params(store: Arc<S>, k1: f64, b: f64) -> Self {
359        Self {
360            index: ChunkedBM25Index::with_params(store, BM25Params::with_values(k1, b)),
361        }
362    }
363
364    /// 返回底层文档存储
365    pub fn store(&self) -> &Arc<S> {
366        self.index.store()
367    }
368
369    /// 添加单个 chunk 索引（内容已存储在 store 中）
370    pub fn add_chunk_index(
371        &mut self,
372        chunk_id: impl Into<String>,
373        parent_id: impl Into<String>,
374        content: &str,
375    ) {
376        self.index.add_chunk_index(chunk_id, parent_id, content);
377    }
378
379    /// 批量添加 chunk 索引
380    pub fn add_chunk_indexes(&mut self, chunks: Vec<(String, String, String)>) {
381        self.index.add_chunk_indexes(chunks);
382    }
383
384    /// 以同步方式添加文档：自动拆分 Parent/Leaf 并建立索引
385    pub fn add_document(&mut self, document: Document) -> Result<(), VectorStoreError> {
386        let parent_id = document
387            .id
388            .clone()
389            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
390
391        // P0-1: 无 id 的文档先把预分配的 parent_id 挂到文档上再入库,
392        // 否则 store 内部会再生成一个新 uuid,导致 get_chunks_for_parent 用错 key 查空。
393        self.index.store.add_parent_document_blocking(
394            document.clone().with_id(parent_id.clone()),
395            self.index.config.leaf_chunk_size,
396        )?;
397
398        let chunks = self
399            .index
400            .store
401            .blocking_get_chunks_for_parent(&parent_id)?;
402
403        for chunk in chunks {
404            self.add_chunk_index(
405                chunk.chunk_id.clone(),
406                chunk.parent_id.clone(),
407                &chunk.content,
408            );
409        }
410
411        Ok(())
412    }
413
414    /// 以异步方式添加文档：自动拆分 Parent/Leaf 并建立索引
415    pub async fn add_document_async(&mut self, document: Document) -> Result<(), VectorStoreError> {
416        let parent_id = document
417            .id
418            .clone()
419            .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
420
421        self.index
422            .store
423            .add_parent_document(
424                document.clone().with_id(parent_id.clone()),
425                self.index.config.leaf_chunk_size,
426            )
427            .await?;
428
429        let chunks = self.index.store.get_chunks_for_parent(&parent_id).await?;
430
431        for chunk in chunks {
432            self.add_chunk_index(
433                chunk.chunk_id.clone(),
434                chunk.parent_id.clone(),
435                &chunk.content,
436            );
437        }
438
439        Ok(())
440    }
441
442    /// 批量以同步方式添加文档
443    pub fn add_documents(&mut self, documents: Vec<Document>) -> Result<(), VectorStoreError> {
444        for doc in documents {
445            self.add_document(doc)?;
446        }
447        Ok(())
448    }
449
450    /// 批量以异步方式添加文档
451    pub async fn add_documents_async(
452        &mut self,
453        documents: Vec<Document>,
454    ) -> Result<(), VectorStoreError> {
455        for doc in documents {
456            self.add_document_async(doc).await?;
457        }
458        Ok(())
459    }
460
461    /// 同步执行 BM25 检索，返回前 k 个 AutoMerging 结果
462    pub fn search(&mut self, query: &str, k: usize) -> Vec<ChunkedSearchResult> {
463        if self.index.n_docs == 0 {
464            return Vec::new();
465        }
466
467        let query_terms = self.index.tokenizer.tokenize(query);
468        if query_terms.is_empty() {
469            return Vec::new();
470        }
471
472        let idf_values: HashMap<String, f64> = query_terms
473            .iter()
474            .map(|t| (t.clone(), self.index.compute_idf_for_term(t)))
475            .collect();
476
477        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
478
479        if scored_chunks.is_empty() {
480            return Vec::new();
481        }
482
483        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
484
485        self.auto_merge_sync(top_chunks, k)
486    }
487
488    /// 异步执行 BM25 检索，返回前 k 个 AutoMerging 结果
489    pub async fn search_async(&mut self, query: &str, k: usize) -> Vec<ChunkedSearchResult> {
490        if self.index.n_docs == 0 {
491            return Vec::new();
492        }
493
494        let query_terms = self.index.tokenizer.tokenize(query);
495        if query_terms.is_empty() {
496            return Vec::new();
497        }
498
499        let idf_values: HashMap<String, f64> = query_terms
500            .iter()
501            .map(|t| (t.clone(), self.index.compute_idf_for_term(t)))
502            .collect();
503
504        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
505
506        if scored_chunks.is_empty() {
507            return Vec::new();
508        }
509
510        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
511
512        self.auto_merge_async(top_chunks, k).await
513    }
514
515    /// 只读 BM25 检索:返回命中的 parent id 列表(去重),按最佳 chunk 分排序。
516    ///
517    /// 与 [`search`](Self::search)/[`search_async`](Self::search_async) 不同:
518    /// 这里**不做** AutoMerging 比例门控 —— 任何 chunk 命中即让该 parent 入围,
519    /// 语义即 [`ParentDocumentRetriever`](crate::parent_document::ParentDocumentRetriever)
520    /// 需要的"命中子块 → 返回整篇父文档"。全 `&self` 只读(idf 不缓存),可安全并发。
521    pub fn search_matched_parents(&self, query: &str, k: usize) -> Vec<(String, f32)> {
522        if self.index.n_docs == 0 {
523            return Vec::new();
524        }
525
526        let query_terms = self.index.tokenizer.tokenize(query);
527        if query_terms.is_empty() {
528            return Vec::new();
529        }
530
531        // 只读 idf:不写 idf_cache,避免 `&mut self`。
532        let idf_values: HashMap<String, f64> = query_terms
533            .iter()
534            .map(|t| {
535                let n = self.index.term_index.get(t).map(|v| v.len()).unwrap_or(0);
536                (t.clone(), compute_idf(n, self.index.n_docs))
537            })
538            .collect();
539
540        let scored_chunks = self.score_chunks(&query_terms, &idf_values);
541        if scored_chunks.is_empty() {
542            return Vec::new();
543        }
544
545        let top_chunks: Vec<(usize, f64)> = scored_chunks.into_iter().take(k * 2).collect();
546        let parent_stats = self.collect_parent_stats(&top_chunks);
547
548        let mut ranked: Vec<(String, f32)> = parent_stats
549            .into_iter()
550            .map(|(parent_id, leaves)| {
551                let best = leaves.iter().map(|(_, s)| *s as f32).fold(0.0f32, f32::max);
552                (parent_id, best)
553            })
554            .collect();
555        ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
556        ranked.into_iter().take(k).collect()
557    }
558
559    fn auto_merge_sync(
560        &self,
561        scored_chunks: Vec<(usize, f64)>,
562        k: usize,
563    ) -> Vec<ChunkedSearchResult> {
564        let threshold = self.index.config.merge_threshold;
565        let leaves_per_parent = self.index.config.leaves_per_parent;
566
567        let parent_stats = self.collect_parent_stats(&scored_chunks);
568
569        let mut results: Vec<ChunkedSearchResult> = Vec::new();
570
571        for (parent_id, matched_leaves) in parent_stats {
572            let ratio = matched_leaves.len() as f32 / leaves_per_parent as f32;
573
574            let avg_score =
575                matched_leaves.iter().map(|(_, s)| s).sum::<f64>() / matched_leaves.len() as f64;
576
577            let matched_terms = matched_leaves
578                .iter()
579                .filter_map(|(idx, _)| self.index.chunk_term_freqs.get(*idx))
580                .flat_map(|tf| tf.keys().cloned())
581                .collect::<Vec<_>>();
582
583            if ratio >= threshold {
584                let parent_doc = self
585                    .index
586                    .store()
587                    .get_parent_document_blocking(&parent_id)
588                    .ok()
589                    .flatten();
590
591                results.push(ChunkedSearchResult {
592                    merged_parent: parent_doc,
593                    leaf_chunks: Vec::new(),
594                    score: avg_score as f32,
595                    matched_terms,
596                    parent_id,
597                });
598            } else {
599                let leaf_chunks: Vec<ChunkDocument> = matched_leaves
600                    .iter()
601                    .filter_map(|(idx, _)| {
602                        let chunk_id = self.index.get_chunk_id(*idx)?;
603                        let chunk = self
604                            .index
605                            .store()
606                            .get_chunk_blocking(chunk_id)
607                            .ok()
608                            .flatten()?;
609                        Some(chunk)
610                    })
611                    .collect();
612
613                results.push(ChunkedSearchResult {
614                    merged_parent: None,
615                    leaf_chunks,
616                    score: avg_score as f32,
617                    matched_terms,
618                    parent_id,
619                });
620            }
621        }
622
623        results.sort_by(|a, b| {
624            b.score
625                .partial_cmp(&a.score)
626                .unwrap_or(std::cmp::Ordering::Equal)
627        });
628        results.into_iter().take(k).collect()
629    }
630
631    async fn auto_merge_async(
632        &self,
633        scored_chunks: Vec<(usize, f64)>,
634        k: usize,
635    ) -> Vec<ChunkedSearchResult> {
636        let threshold = self.index.config.merge_threshold;
637        let leaves_per_parent = self.index.config.leaves_per_parent;
638
639        let parent_stats = self.collect_parent_stats(&scored_chunks);
640
641        let mut results: Vec<ChunkedSearchResult> = Vec::new();
642
643        for (parent_id, matched_leaves) in parent_stats {
644            let ratio = matched_leaves.len() as f32 / leaves_per_parent as f32;
645
646            let avg_score =
647                matched_leaves.iter().map(|(_, s)| s).sum::<f64>() / matched_leaves.len() as f64;
648
649            let matched_terms = matched_leaves
650                .iter()
651                .filter_map(|(idx, _)| self.index.chunk_term_freqs.get(*idx))
652                .flat_map(|tf| tf.keys().cloned())
653                .collect::<Vec<_>>();
654
655            if ratio >= threshold {
656                let parent_doc = self
657                    .index
658                    .store()
659                    .get_parent_document(&parent_id)
660                    .await
661                    .ok()
662                    .flatten();
663
664                results.push(ChunkedSearchResult {
665                    merged_parent: parent_doc,
666                    leaf_chunks: Vec::new(),
667                    score: avg_score as f32,
668                    matched_terms,
669                    parent_id,
670                });
671            } else {
672                let mut leaf_chunks = Vec::new();
673                for (idx, _) in matched_leaves {
674                    if let Some(chunk_id) = self.index.get_chunk_id(idx) {
675                        match self.index.store().get_chunk(chunk_id).await {
676                            Ok(Some(chunk)) => leaf_chunks.push(chunk),
677                            Ok(None) => {}
678                            Err(e) => {
679                                // 不再静默吞错:读失败记日志,该 chunk 从结果中缺失
680                                log::error!(
681                                    "failed to read chunk `{}` during retrieval (chunk missing from results): {}",
682                                    chunk_id,
683                                    e
684                                );
685                            }
686                        }
687                    }
688                }
689
690                results.push(ChunkedSearchResult {
691                    merged_parent: None,
692                    leaf_chunks,
693                    score: avg_score as f32,
694                    matched_terms,
695                    parent_id,
696                });
697            }
698        }
699
700        results.sort_by(|a, b| {
701            b.score
702                .partial_cmp(&a.score)
703                .unwrap_or(std::cmp::Ordering::Equal)
704        });
705        results.into_iter().take(k).collect()
706    }
707
708    fn score_chunks(
709        &self,
710        query_terms: &[String],
711        idf_values: &HashMap<String, f64>,
712    ) -> Vec<(usize, f64)> {
713        let mut scored = Vec::new();
714
715        for chunk_idx in 0..self.index.n_docs {
716            if let Some(term_freqs) = self.index.chunk_term_freqs.get(chunk_idx) {
717                let doc_length = *self.index.doc_lengths.get(chunk_idx).unwrap_or(&0);
718
719                let score = bm25_score(
720                    query_terms,
721                    term_freqs,
722                    doc_length,
723                    self.index.avgdl,
724                    idf_values,
725                    &self.index.params,
726                );
727
728                if score > 0.0 {
729                    scored.push((chunk_idx, score));
730                }
731            }
732        }
733
734        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
735        scored
736    }
737
738    fn collect_parent_stats(
739        &self,
740        scored_chunks: &[(usize, f64)],
741    ) -> HashMap<String, Vec<(usize, f64)>> {
742        let mut stats: HashMap<String, Vec<(usize, f64)>> = HashMap::new();
743
744        for (chunk_idx, score) in scored_chunks {
745            if let Some(chunk_id) = self.index.chunk_id_list.get(*chunk_idx) {
746                let parent_id = chunk_id.split("::").next().unwrap_or_default().to_string();
747                stats
748                    .entry(parent_id)
749                    .or_default()
750                    .push((*chunk_idx, *score));
751            }
752        }
753
754        stats
755    }
756
757    /// 按 Parent id 获取父文档
758    pub fn get_parent_document(&self, parent_id: &str) -> Option<Document> {
759        self.index
760            .store()
761            .get_parent_document_blocking(parent_id)
762            .ok()
763            .flatten()
764    }
765
766    /// 返回索引中的文档数量
767    pub fn len(&self) -> usize {
768        self.index.n_docs()
769    }
770
771    /// 索引是否为空
772    pub fn is_empty(&self) -> bool {
773        self.index.n_docs() == 0
774    }
775
776    /// 清空索引
777    pub fn clear(&mut self) {
778        self.index.clear();
779    }
780
781    /// 返回 AutoMerging 配置
782    pub fn config(&self) -> &AutoMergingConfig {
783        self.index.config()
784    }
785
786    // 持久化方法
787    /// 将索引数据序列化为 Bincode 并保存到指定路径
788    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), Box<dyn std::error::Error>> {
789        let data = ChunkedIndexData {
790            chunk_id_list: self.index.chunk_id_list.clone(),
791            chunk_term_freqs: self.index.chunk_term_freqs.clone(),
792            term_index: self.index.term_index.clone(),
793            parent_to_leaves: self.index.parent_to_leaves.clone(),
794            doc_lengths: self.index.doc_lengths.clone(),
795            avgdl: self.index.avgdl,
796            n_docs: self.index.n_docs,
797            params: BM25ParamsData::from(self.index.params.clone()),
798            config: self.index.config.clone(),
799        };
800        let encoded = bincode::serialize(&data)?;
801        std::fs::write(path.as_ref(), encoded)?;
802        Ok(())
803    }
804}
805
806impl ChunkedBM25Retriever<lc_vector_stores::ChunkedDocumentStore> {
807    /// 从指定路径加载 Bincode 序列化的索引数据
808    pub fn load(
809        store: Arc<lc_vector_stores::ChunkedDocumentStore>,
810        path: impl AsRef<Path>,
811    ) -> Result<Self, Box<dyn std::error::Error>> {
812        let bytes = std::fs::read(path.as_ref())?;
813        let data: ChunkedIndexData = bincode::deserialize(&bytes)?;
814        let params: BM25Params = data.params.into();
815
816        Ok(Self {
817            index: ChunkedBM25Index {
818                store,
819                chunk_id_list: data.chunk_id_list,
820                chunk_term_freqs: data.chunk_term_freqs,
821                term_index: data.term_index,
822                parent_to_leaves: data.parent_to_leaves,
823                doc_lengths: data.doc_lengths,
824                avgdl: data.avgdl,
825                n_docs: data.n_docs,
826                idf_cache: HashMap::new(),
827                params,
828                tokenizer: Tokenizer::new(),
829                config: data.config,
830            },
831        })
832    }
833}
834
835impl Default for ChunkedBM25Retriever<lc_vector_stores::ChunkedDocumentStore> {
836    fn default() -> Self {
837        Self::new(Arc::new(lc_vector_stores::ChunkedDocumentStore::new()))
838    }
839}
lc_rag/bm25/chunked.rs

lc_rag/bm25/
chunked.rs
