Skip to main content

xz_embed/index/
builder.rs

1use async_trait::async_trait;
2use std::fmt::Debug;
3use std::sync::Arc;
4use std::time::Duration;
5use tracing::{debug, info, warn};
6
7use crate::error::StoreError;
8use crate::traits::VectorStore;
9use crate::types::StoreStats;
10
11/// 索引构建 throttle 配置
12#[derive(Debug, Clone)]
13pub struct IndexThrottleConfig {
14    /// 重建间隔(冷却期)
15    pub cooldown: Duration,
16    /// 触发重建的新增条目阈值
17    pub count_threshold: usize,
18}
19
20impl Default for IndexThrottleConfig {
21    fn default() -> Self {
22        Self {
23            cooldown: Duration::from_secs(60),
24            count_threshold: 1000,
25        }
26    }
27}
28
29/// 索引构建器
30///
31/// 管理向量存储的索引生命周期。
32#[derive(Debug)]
33pub struct IndexBuilder {
34    store: Arc<dyn VectorStore>,
35    config: IndexThrottleConfig,
36}
37
38impl IndexBuilder {
39    pub fn new(store: Arc<dyn VectorStore>, config: IndexThrottleConfig) -> Self {
40        Self { store, config }
41    }
42
43    /// 重建索引
44    pub async fn rebuild(&self) -> Result<(), StoreError> {
45        info!(target: "xz_embed", "starting index rebuild");
46        self.store.rebuild_index().await?;
47        info!(target: "xz_embed", "index rebuild completed");
48        Ok(())
49    }
50
51    /// 条件重建(仅当条目超过阈值时)
52    pub async fn rebuild_if_needed(&self) -> Result<(), StoreError> {
53        let count = self.store.count().await?;
54        if count >= self.config.count_threshold {
55            self.rebuild().await?;
56        } else {
57            debug!(target: "xz_embed", count, threshold = self.config.count_threshold, "skipping rebuild");
58        }
59        Ok(())
60    }
61
62    /// 后台定期重建任务
63    pub async fn start_background_rebuild(self: Arc<Self>) {
64        tokio::spawn(async move {
65            loop {
66                tokio::time::sleep(self.config.cooldown).await;
67                match self.rebuild_if_needed().await {
68                    Ok(_) => {}
69                    Err(e) => warn!(target: "xz_embed", error = %e, "background rebuild failed"),
70                }
71            }
72        });
73    }
74}