a3s_code_core/embedding/
executor.rs

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

a3s_code_core/embedding/
executor.rs
