xz_embed/index/
builder.rs1use 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#[derive(Debug, Clone)]
13pub struct IndexThrottleConfig {
14 pub cooldown: Duration,
16 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#[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 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 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 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}