a3s_code_core/embedding/
executor.rs

1use super::{
2    EmbeddingBatchRequest, EmbeddingBatchResponse, EmbeddingError, EmbeddingExecution,
3    EmbeddingFailureKind, EmbeddingInput, EmbeddingNormalization, EmbeddingProvider,
4    EmbeddingProviderDescriptor, EmbeddingProviderError, EmbeddingResult, EmbeddingVector,
5};
6use futures::FutureExt;
7use std::collections::{HashMap, HashSet};
8use std::ops::Range;
9use std::panic::AssertUnwindSafe;
10use std::sync::atomic::{AtomicUsize, Ordering};
11use std::sync::Arc;
12use std::time::Duration;
13use tokio_util::sync::CancellationToken;
14
15const MAX_DESCRIPTOR_TEXT_BYTES: usize = 256;
16const MAX_INPUT_ID_BYTES: usize = 512;
17const MAX_EMBEDDING_DIMENSION: usize = 65_536;
18const MAX_RETRIES: u32 = 8;
19const MAX_OPERATION_DURATION: Duration = Duration::from_secs(5 * 60);
20const UNIT_NORM_TOLERANCE: f64 = 0.01;
21
22/// Hard limits and retry policy for one embedding executor generation.
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
24pub struct EmbeddingExecutorConfig {
25    pub max_batch_inputs: usize,
26    pub max_batch_text_bytes: usize,
27    pub max_input_text_bytes: usize,
28    pub max_request_inputs: usize,
29    pub max_request_text_bytes: usize,
30    pub max_batch_vector_bytes: usize,
31    pub max_request_vector_bytes: usize,
32    pub max_retries: u32,
33    pub base_retry_delay: Duration,
34    pub max_retry_delay: Duration,
35    pub request_timeout: Duration,
36}
37
38impl Default for EmbeddingExecutorConfig {
39    fn default() -> Self {
40        Self {
41            max_batch_inputs: 64,
42            max_batch_text_bytes: 256 * 1024,
43            max_input_text_bytes: 64 * 1024,
44            max_request_inputs: 4_096,
45            max_request_text_bytes: 16 * 1024 * 1024,
46            max_batch_vector_bytes: 32 * 1024 * 1024,
47            max_request_vector_bytes: 64 * 1024 * 1024,
48            max_retries: 2,
49            base_retry_delay: Duration::from_millis(100),
50            max_retry_delay: Duration::from_secs(2),
51            request_timeout: Duration::from_secs(30),
52        }
53    }
54}
55
56impl EmbeddingExecutorConfig {
57    fn validate(self) -> EmbeddingResult<Self> {
58        for (field, value) in [
59            ("max_batch_inputs", self.max_batch_inputs),
60            ("max_batch_text_bytes", self.max_batch_text_bytes),
61            ("max_input_text_bytes", self.max_input_text_bytes),
62            ("max_request_inputs", self.max_request_inputs),
63            ("max_request_text_bytes", self.max_request_text_bytes),
64            ("max_batch_vector_bytes", self.max_batch_vector_bytes),
65            ("max_request_vector_bytes", self.max_request_vector_bytes),
66        ] {
67            if value == 0 {
68                return Err(EmbeddingError::InvalidConfiguration {
69                    field,
70                    reason: "must be greater than zero",
71                });
72            }
73        }
74        if self.max_batch_inputs > self.max_request_inputs {
75            return Err(EmbeddingError::InvalidConfiguration {
76                field: "max_batch_inputs",
77                reason: "must not exceed max_request_inputs",
78            });
79        }
80        if self.max_input_text_bytes > self.max_batch_text_bytes
81            || self.max_batch_text_bytes > self.max_request_text_bytes
82        {
83            return Err(EmbeddingError::InvalidConfiguration {
84                field: "text byte limits",
85                reason: "must be monotonic from input to batch to request",
86            });
87        }
88        if self.max_batch_vector_bytes > self.max_request_vector_bytes {
89            return Err(EmbeddingError::InvalidConfiguration {
90                field: "max_batch_vector_bytes",
91                reason: "must not exceed max_request_vector_bytes",
92            });
93        }
94        if self.request_timeout.is_zero() {
95            return Err(EmbeddingError::InvalidConfiguration {
96                field: "request_timeout",
97                reason: "must be greater than zero",
98            });
99        }
100        if self.request_timeout > MAX_OPERATION_DURATION {
101            return Err(EmbeddingError::InvalidConfiguration {
102                field: "request_timeout",
103                reason: "must not exceed five minutes",
104            });
105        }
106        if self.max_retries > MAX_RETRIES {
107            return Err(EmbeddingError::InvalidConfiguration {
108                field: "max_retries",
109                reason: "must not exceed eight",
110            });
111        }
112        if self.max_retry_delay > MAX_OPERATION_DURATION {
113            return Err(EmbeddingError::InvalidConfiguration {
114                field: "max_retry_delay",
115                reason: "must not exceed five minutes",
116            });
117        }
118        if self.base_retry_delay > self.max_retry_delay {
119            return Err(EmbeddingError::InvalidConfiguration {
120                field: "base_retry_delay",
121                reason: "must not exceed max_retry_delay",
122            });
123        }
124        Ok(self)
125    }
126}
127
128/// Validating, batching wrapper around one host-injected provider generation.
129#[derive(Clone)]
130pub struct EmbeddingExecutor {
131    provider: Arc<dyn EmbeddingProvider>,
132    descriptor: EmbeddingProviderDescriptor,
133    config: EmbeddingExecutorConfig,
134}
135
136impl std::fmt::Debug for EmbeddingExecutor {
137    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
138        formatter
139            .debug_struct("EmbeddingExecutor")
140            .field("descriptor", &self.descriptor)
141            .field("config", &self.config)
142            .finish_non_exhaustive()
143    }
144}
145
146impl EmbeddingExecutor {
147    pub fn new(
148        provider: Arc<dyn EmbeddingProvider>,
149        config: EmbeddingExecutorConfig,
150    ) -> EmbeddingResult<Self> {
151        let descriptor = std::panic::catch_unwind(AssertUnwindSafe(|| provider.descriptor()))
152            .map_err(|_| EmbeddingError::ProviderPanicked {
153                operation: "descriptor",
154            })?;
155        validate_descriptor(&descriptor)?;
156        Ok(Self {
157            provider,
158            descriptor,
159            config: config.validate()?,
160        })
161    }
162
163    pub fn descriptor(&self) -> &EmbeddingProviderDescriptor {
164        &self.descriptor
165    }
166
167    pub fn config(&self) -> EmbeddingExecutorConfig {
168        self.config
169    }
170
171    /// Embed all inputs atomically from the caller's perspective.
172    ///
173    /// No partial vector list is returned if a later batch fails validation.
174    pub async fn embed(
175        &self,
176        inputs: Vec<EmbeddingInput>,
177        cancellation: CancellationToken,
178    ) -> EmbeddingResult<EmbeddingExecution> {
179        self.embed_inner(inputs, cancellation, None).await
180    }
181
182    pub(crate) async fn embed_counted(
183        &self,
184        inputs: Vec<EmbeddingInput>,
185        cancellation: CancellationToken,
186        provider_requests: &AtomicUsize,
187    ) -> EmbeddingResult<EmbeddingExecution> {
188        self.embed_inner(inputs, cancellation, Some(provider_requests))
189            .await
190    }
191
192    async fn embed_inner(
193        &self,
194        inputs: Vec<EmbeddingInput>,
195        cancellation: CancellationToken,
196        provider_requests: Option<&AtomicUsize>,
197    ) -> EmbeddingResult<EmbeddingExecution> {
198        if cancellation.is_cancelled() {
199            return Err(EmbeddingError::Cancelled);
200        }
201        let batches = plan_batches(&inputs, self.descriptor.dimension, self.config)?;
202        let mut vectors = Vec::with_capacity(inputs.len());
203        let mut provider_attempts = 0usize;
204        for range in &batches {
205            let request = EmbeddingBatchRequest::new(inputs[range.clone()].to_vec());
206            let (response, attempts) = self
207                .call_batch(request.clone(), &cancellation, provider_requests)
208                .await?;
209            provider_attempts = provider_attempts.saturating_add(attempts);
210            vectors.extend(validate_response(
211                &self.descriptor,
212                request.inputs(),
213                response,
214            )?);
215        }
216        Ok(EmbeddingExecution {
217            descriptor: self.descriptor.clone(),
218            vectors,
219            batch_count: batches.len(),
220            provider_attempts,
221        })
222    }
223
224    async fn call_batch(
225        &self,
226        request: EmbeddingBatchRequest,
227        cancellation: &CancellationToken,
228        provider_requests: Option<&AtomicUsize>,
229    ) -> EmbeddingResult<(EmbeddingBatchResponse, usize)> {
230        for attempt in 0..=self.config.max_retries {
231            if cancellation.is_cancelled() {
232                return Err(EmbeddingError::Cancelled);
233            }
234            if let Some(provider_requests) = provider_requests {
235                provider_requests.fetch_add(1, Ordering::Relaxed);
236            }
237            let attempt_token = cancellation.child_token();
238            let provider_call =
239                AssertUnwindSafe(self.provider.embed(request.clone(), attempt_token.clone()))
240                    .catch_unwind();
241            let result = tokio::select! {
242                biased;
243                _ = cancellation.cancelled() => {
244                    attempt_token.cancel();
245                    return Err(EmbeddingError::Cancelled);
246                }
247                result = tokio::time::timeout(
248                    self.config.request_timeout,
249                    provider_call,
250                ) => match result {
251                    Ok(Ok(result)) => result,
252                    Ok(Err(_)) => {
253                        return Err(EmbeddingError::ProviderPanicked { operation: "embed" });
254                    }
255                    Err(_) => {
256                        attempt_token.cancel();
257                        Err(EmbeddingProviderError::Timeout)
258                    }
259                }
260            };
261            let attempts = attempt as usize + 1;
262            match result {
263                Ok(response) => return Ok((response, attempts)),
264                Err(EmbeddingProviderError::Cancelled) => return Err(EmbeddingError::Cancelled),
265                Err(error) if error.is_retryable() && attempt < self.config.max_retries => {
266                    let delay = retry_delay(&error, attempt, self.config);
267                    tokio::select! {
268                        biased;
269                        _ = cancellation.cancelled() => return Err(EmbeddingError::Cancelled),
270                        _ = tokio::time::sleep(delay) => {}
271                    }
272                }
273                Err(error) if error.is_retryable() => {
274                    return Err(EmbeddingError::RetriesExhausted {
275                        kind: error.kind(),
276                        attempts,
277                    })
278                }
279                Err(error) => {
280                    return Err(EmbeddingError::ProviderFailure {
281                        kind: error.kind(),
282                        attempts,
283                    })
284                }
285            }
286        }
287        Err(EmbeddingError::RetriesExhausted {
288            kind: EmbeddingFailureKind::Other,
289            attempts: self.config.max_retries as usize + 1,
290        })
291    }
292}
293
294fn validate_descriptor(descriptor: &EmbeddingProviderDescriptor) -> EmbeddingResult<()> {
295    for (field, value) in [
296        ("provider", descriptor.provider.as_str()),
297        ("model", descriptor.model.as_str()),
298    ] {
299        if value.trim().is_empty()
300            || value.len() > MAX_DESCRIPTOR_TEXT_BYTES
301            || value.chars().any(char::is_control)
302        {
303            return Err(EmbeddingError::InvalidDescriptor { field });
304        }
305    }
306    if descriptor.revision.as_ref().is_some_and(|revision| {
307        revision.trim().is_empty()
308            || revision.len() > MAX_DESCRIPTOR_TEXT_BYTES
309            || revision.chars().any(char::is_control)
310    }) {
311        return Err(EmbeddingError::InvalidDescriptor { field: "revision" });
312    }
313    if descriptor.dimension == 0 || descriptor.dimension > MAX_EMBEDDING_DIMENSION {
314        return Err(EmbeddingError::InvalidDescriptor { field: "dimension" });
315    }
316    Ok(())
317}
318
319fn plan_batches(
320    inputs: &[EmbeddingInput],
321    dimension: usize,
322    config: EmbeddingExecutorConfig,
323) -> EmbeddingResult<Vec<Range<usize>>> {
324    if inputs.is_empty() {
325        return Err(EmbeddingError::EmptyRequest);
326    }
327    if inputs.len() > config.max_request_inputs {
328        return Err(EmbeddingError::BudgetExceeded {
329            resource: "request input count",
330            requested: inputs.len(),
331            limit: config.max_request_inputs,
332        });
333    }
334    let mut seen = HashSet::with_capacity(inputs.len());
335    let mut total_text_bytes = 0usize;
336    for (index, input) in inputs.iter().enumerate() {
337        if input.id().is_empty()
338            || input.id().len() > MAX_INPUT_ID_BYTES
339            || input.id().chars().any(char::is_control)
340        {
341            return Err(EmbeddingError::InvalidInput {
342                index,
343                reason: "identifier is empty, oversized, or contains a control character",
344            });
345        }
346        if !seen.insert(input.id()) {
347            return Err(EmbeddingError::InvalidInput {
348                index,
349                reason: "identifier is duplicated",
350            });
351        }
352        if input.text().is_empty() {
353            return Err(EmbeddingError::InvalidInput {
354                index,
355                reason: "text must not be empty",
356            });
357        }
358        if input.text_bytes() > config.max_input_text_bytes {
359            return Err(EmbeddingError::BudgetExceeded {
360                resource: "input text byte",
361                requested: input.text_bytes(),
362                limit: config.max_input_text_bytes,
363            });
364        }
365        total_text_bytes = total_text_bytes.saturating_add(input.text_bytes());
366    }
367    if total_text_bytes > config.max_request_text_bytes {
368        return Err(EmbeddingError::BudgetExceeded {
369            resource: "request text byte",
370            requested: total_text_bytes,
371            limit: config.max_request_text_bytes,
372        });
373    }
374    let vector_bytes_per_input = dimension.saturating_mul(std::mem::size_of::<f32>());
375    let request_vector_bytes = vector_bytes_per_input.saturating_mul(inputs.len());
376    if request_vector_bytes > config.max_request_vector_bytes {
377        return Err(EmbeddingError::BudgetExceeded {
378            resource: "request vector byte",
379            requested: request_vector_bytes,
380            limit: config.max_request_vector_bytes,
381        });
382    }
383
384    let mut batches = Vec::new();
385    let mut start = 0usize;
386    let mut batch_bytes = 0usize;
387    for (index, input) in inputs.iter().enumerate() {
388        let would_exceed_items = index.saturating_sub(start) >= config.max_batch_inputs;
389        let would_exceed_bytes =
390            batch_bytes.saturating_add(input.text_bytes()) > config.max_batch_text_bytes;
391        let batch_items = index.saturating_sub(start).saturating_add(1);
392        let would_exceed_vector_bytes =
393            vector_bytes_per_input.saturating_mul(batch_items) > config.max_batch_vector_bytes;
394        if index > start && (would_exceed_items || would_exceed_bytes || would_exceed_vector_bytes)
395        {
396            batches.push(start..index);
397            start = index;
398            batch_bytes = 0;
399        }
400        if vector_bytes_per_input > config.max_batch_vector_bytes {
401            return Err(EmbeddingError::BudgetExceeded {
402                resource: "batch vector byte",
403                requested: vector_bytes_per_input,
404                limit: config.max_batch_vector_bytes,
405            });
406        }
407        batch_bytes = batch_bytes.saturating_add(input.text_bytes());
408    }
409    if start < inputs.len() {
410        batches.push(start..inputs.len());
411    }
412    Ok(batches)
413}
414
415fn retry_delay(
416    error: &EmbeddingProviderError,
417    attempt: u32,
418    config: EmbeddingExecutorConfig,
419) -> Duration {
420    error
421        .retry_after()
422        .unwrap_or_else(|| {
423            config
424                .base_retry_delay
425                .saturating_mul(1u32 << attempt.min(16))
426        })
427        .min(config.max_retry_delay)
428}
429
430fn validate_response(
431    expected_descriptor: &EmbeddingProviderDescriptor,
432    inputs: &[EmbeddingInput],
433    response: EmbeddingBatchResponse,
434) -> EmbeddingResult<Vec<EmbeddingVector>> {
435    validate_descriptor(&response.descriptor)?;
436    if response.descriptor != *expected_descriptor {
437        return Err(EmbeddingError::DescriptorChanged);
438    }
439    if response.vectors.len() != inputs.len() {
440        return Err(EmbeddingError::OutputCountMismatch {
441            expected: inputs.len(),
442            actual: response.vectors.len(),
443        });
444    }
445    let indices = inputs
446        .iter()
447        .enumerate()
448        .map(|(index, input)| (input.id(), index))
449        .collect::<HashMap<_, _>>();
450    let mut ordered = vec![None; inputs.len()];
451    for vector in response.vectors {
452        let Some(&input_index) = indices.get(vector.id.as_ref()) else {
453            return Err(EmbeddingError::UnexpectedOutput);
454        };
455        if ordered[input_index].is_some() {
456            return Err(EmbeddingError::DuplicateOutput { input_index });
457        }
458        if vector.values.len() != expected_descriptor.dimension {
459            return Err(EmbeddingError::DimensionMismatch {
460                input_index,
461                expected: expected_descriptor.dimension,
462                actual: vector.values.len(),
463            });
464        }
465        if let Some(position) = vector.values.iter().position(|value| !value.is_finite()) {
466            return Err(EmbeddingError::NonFiniteValue {
467                input_index,
468                position,
469            });
470        }
471        if expected_descriptor.normalization == EmbeddingNormalization::Unit {
472            let norm = vector
473                .values
474                .iter()
475                .map(|value| f64::from(*value).powi(2))
476                .sum::<f64>()
477                .sqrt();
478            if (norm - 1.0).abs() > UNIT_NORM_TOLERANCE {
479                return Err(EmbeddingError::NormalizationMismatch { input_index });
480            }
481        }
482        ordered[input_index] = Some(EmbeddingVector::new(
483            Arc::<str>::from(inputs[input_index].id()),
484            vector.values,
485        ));
486    }
487    Ok(ordered.into_iter().flatten().collect())
488}
a3s_code_core/embedding/executor.rs

a3s_code_core/embedding/
executor.rs
