xz_embed/
batch_manager.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use futures::future;
5use tokio::sync::Semaphore;
6use tracing::{debug, warn};
7
8use crate::config::RetryConfig;
9use crate::error::EmbedError;
10use crate::traits::EmbeddingModel;
11
12/// 并发批量管理器
13///
14/// 将大批次文本拆分为多个小批次，并发发送 Embedding API 请求。
15/// 内置重试、限流、进度回调。
16pub struct ConcurrentBatchManager {
17    embedder: Arc<dyn EmbeddingModel>,
18    /// 每批文本数（不超过 embedder.max_batch_size()）
19    batch_size: usize,
20    /// 最大并发批次
21    max_concurrency: usize,
22    /// 重试策略
23    retry: RetryConfig,
24}
25
26impl ConcurrentBatchManager {
27    pub fn new(
28        embedder: Box<dyn EmbeddingModel>,
29        batch_size: usize,
30        max_concurrency: usize,
31    ) -> Self {
32        Self {
33            embedder: Arc::from(embedder),
34            batch_size,
35            max_concurrency,
36            retry: RetryConfig::default(),
37        }
38    }
39
40    pub fn with_retry(mut self, retry: RetryConfig) -> Self {
41        self.retry = retry;
42        self
43    }
44
45    /// 将文本列表拆分为多个子批次
46    fn chunk_texts(&self, texts: &[impl AsRef<str>]) -> Vec<Vec<String>> {
47        let max_batch = self.batch_size.min(self.embedder.max_batch_size());
48        texts
49            .chunks(max_batch)
50            .map(|chunk| chunk.iter().map(|t| t.as_ref().to_string()).collect())
51            .collect()
52    }
53
54    /// 嵌入全部文本，返回顺序与输入一致
55    pub async fn embed_all(&self, texts: &[impl AsRef<str>]) -> Result<Vec<Vec<f32>>, EmbedError> {
56        self.embed_all_with_progress(texts, |_, _| {}).await
57    }
58
59    /// 带进度回调的嵌入
60    pub async fn embed_all_with_progress(
61        &self,
62        texts: &[impl AsRef<str>],
63        on_batch_done: impl Fn(usize, usize),
64    ) -> Result<Vec<Vec<f32>>, EmbedError> {
65        let batches = self.chunk_texts(texts);
66        let total_batches = batches.len();
67
68        if total_batches == 0 {
69            return Ok(vec![]);
70        }
71
72        let semaphore = Arc::new(Semaphore::new(self.max_concurrency));
73        let mut handles = Vec::with_capacity(total_batches);
74
75        for (i, batch) in batches.into_iter().enumerate() {
76            let permit = semaphore
77                .clone()
78                .acquire_owned()
79                .await
80                .map_err(|e| EmbedError::Config(format!("获取信号量失败: {e}")))?;
81
82            let embedder = self.embedder.clone();
83            let retry = self.retry.clone();
84
85            handles.push(tokio::spawn(async move {
86                let _permit = permit;
87                let texts_refs: Vec<&str> = batch.iter().map(|s| s.as_str()).collect();
88                let result = retry_with_backoff(|| embedder.embed(&texts_refs), &retry).await;
89                (i, result)
90            }));
91        }
92
93        let mut ordered_results: Vec<Option<Vec<Vec<f32>>>> = vec![None; total_batches];
94        let mut errors = Vec::new();
95
96        for handle in handles {
97            match handle.await {
98                Ok((idx, Ok(vectors))) => {
99                    ordered_results[idx] = Some(vectors);
100                    on_batch_done(idx + 1, total_batches);
101                }
102                Ok((idx, Err(e))) => {
103                    warn!(target: "xz_embed", batch = idx, error = %e, "batch embedding failed");
104                    errors.push(e);
105                }
106                Err(e) => {
107                    errors.push(EmbedError::Config(format!("task join error: {e}")));
108                }
109            }
110        }
111
112        if !errors.is_empty() {
113            return Err(errors.remove(0));
114        }
115
116        let all_vectors: Vec<Vec<f32>> =
117            ordered_results.into_iter().filter_map(|r| r).flatten().collect();
118
119        debug!(
120            target: "xz_embed",
121            total_texts = texts.len(),
122            total_batches,
123            total_vectors = all_vectors.len(),
124            "embed_all completed"
125        );
126
127        Ok(all_vectors)
128    }
129}
130
131async fn retry_with_backoff<F, Fut, T>(f: F, config: &RetryConfig) -> Result<T, EmbedError>
132where
133    F: Fn() -> Fut,
134    Fut: std::future::Future<Output = Result<T, EmbedError>>,
135{
136    let mut attempt = 0;
137    let mut backoff_ms = config.initial_backoff_ms;
138
139    loop {
140        match f().await {
141            Ok(result) => return Ok(result),
142            Err(e) if e.is_retryable() && attempt < config.max_retries => {
143                attempt += 1;
144                debug!(
145                    target: "xz_embed",
146                    attempt,
147                    backoff_ms,
148                    error = %e,
149                    "retrying embedding request"
150                );
151                tokio::time::sleep(Duration::from_millis(backoff_ms)).await;
152                backoff_ms = (backoff_ms as f64 * config.backoff_multiplier)
153                    .min(config.max_backoff_ms as f64) as u64;
154            }
155            Err(e) => return Err(e),
156        }
157    }
158}
xz_embed/batch_manager.rs

xz_embed/
batch_manager.rs
