mongreldb_core/
embedding.rs

1//! Bounded, cancellable embedding generation over materialized vectors.
2//!
3//! Schema stores portable provider/model identity. Node-local configuration
4//! resolves that identity to model files or remote endpoints. Secrets and
5//! filesystem paths never enter replicated schema.
6
7use std::collections::BTreeMap;
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::{Arc, RwLock};
11
12use serde::{Deserialize, Serialize};
13
14pub const DEFAULT_MAX_EMBEDDING_TEXTS: usize = 256;
15pub const DEFAULT_MAX_EMBEDDING_TEXT_BYTES: usize = 64 * 1024;
16pub const DEFAULT_MAX_EMBEDDING_INPUT_BYTES: usize = 1024 * 1024;
17pub const DEFAULT_MAX_EMBEDDING_OUTPUT_BYTES: usize = 16 * 1024 * 1024;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
20#[serde(rename_all = "snake_case")]
21pub enum EmbeddingNormalization {
22    #[default]
23    None,
24    L2,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
28#[serde(rename_all = "snake_case")]
29pub enum EmbeddingFailurePolicy {
30    /// Provider failure aborts the source write. Source and vector commit
31    /// atomically or not at all.
32    #[default]
33    AbortWrite,
34}
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
37#[serde(rename_all = "snake_case")]
38pub enum EmbeddingGenerationStatus {
39    Ready,
40    Pending,
41    Failed,
42}
43
44/// Durable provenance for a generated embedding value.
45#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
46pub struct GeneratedEmbeddingMetadata {
47    pub provider_id: String,
48    pub model_id: String,
49    pub model_version: String,
50    pub preprocessing_version: String,
51    pub source_fingerprint: [u8; 32],
52    pub status: EmbeddingGenerationStatus,
53    pub last_error_category: Option<mongreldb_types::errors::ErrorCategory>,
54    pub attempt_count: u32,
55}
56
57/// Generated vector and its durable provenance.
58#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
59pub struct GeneratedEmbeddingValue {
60    pub vector: Vec<f32>,
61    pub metadata: GeneratedEmbeddingMetadata,
62}
63
64/// Portable generated-column contract stored in schema.
65#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
66pub struct GeneratedEmbeddingSpec {
67    pub provider_id: String,
68    pub model_id: String,
69    pub model_version: String,
70    pub source_columns: Vec<u16>,
71    /// Template placeholders use `{column_name}`. Empty means source values
72    /// joined with one newline in `source_columns` order.
73    pub input_template: String,
74    pub dimension: u32,
75    pub normalization: EmbeddingNormalization,
76    pub failure_policy: EmbeddingFailurePolicy,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
80#[serde(tag = "kind", rename_all = "snake_case")]
81pub enum EmbeddingSource {
82    #[default]
83    SuppliedByApplication,
84    /// Compatibility shape for 0.60.3 Kit clients. The node-local path is
85    /// accepted in memory but never serialized into replicated schema.
86    LocalModel {
87        model_id: String,
88        #[serde(skip)]
89        model_path: std::path::PathBuf,
90    },
91    /// Compatibility shape for explicit Kit embedding calls. Ordinary writes
92    /// do not materialize this legacy source.
93    GeneratedColumn {
94        provider: String,
95    },
96    /// A local provider resolved from node configuration. No path is stored.
97    ConfiguredModel {
98        provider_id: String,
99        model_id: String,
100        model_version: String,
101    },
102    GeneratedColumnSpec {
103        spec: GeneratedEmbeddingSpec,
104    },
105}
106
107impl EmbeddingSource {
108    pub fn label(&self) -> &str {
109        match self {
110            Self::SuppliedByApplication => "supplied_by_application",
111            Self::LocalModel { .. } | Self::ConfiguredModel { .. } => "local_model",
112            Self::GeneratedColumn { .. } | Self::GeneratedColumnSpec { .. } => "generated_column",
113        }
114    }
115
116    pub fn provider_id(&self) -> Option<&str> {
117        match self {
118            Self::SuppliedByApplication => None,
119            Self::LocalModel { model_id, .. } => Some(model_id),
120            Self::GeneratedColumn { provider } => Some(provider),
121            Self::ConfiguredModel { provider_id, .. } => Some(provider_id),
122            Self::GeneratedColumnSpec { spec } => Some(&spec.provider_id),
123        }
124    }
125
126    pub fn model_id(&self) -> Option<&str> {
127        match self {
128            Self::SuppliedByApplication => None,
129            Self::LocalModel { model_id, .. } => Some(model_id),
130            Self::GeneratedColumn { .. } => None,
131            Self::ConfiguredModel { model_id, .. } => Some(model_id),
132            Self::GeneratedColumnSpec { spec } => Some(&spec.model_id),
133        }
134    }
135
136    pub fn model_version(&self) -> Option<&str> {
137        match self {
138            Self::SuppliedByApplication => None,
139            Self::LocalModel { .. } | Self::GeneratedColumn { .. } => None,
140            Self::ConfiguredModel { model_version, .. } => Some(model_version),
141            Self::GeneratedColumnSpec { spec } => Some(&spec.model_version),
142        }
143    }
144
145    pub fn requires_provider(&self) -> bool {
146        !matches!(self, Self::SuppliedByApplication)
147    }
148}
149
150#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
151pub enum EmbeddingError {
152    #[error("embedding provider {0:?} is not registered")]
153    ProviderNotFound(String),
154    #[error("embedding provider {0:?} is already registered")]
155    ProviderAlreadyRegistered(String),
156    #[error(
157        "embedding provider {provider:?} generation mismatch: expected {expected}, got {actual}"
158    )]
159    ProviderGenerationMismatch {
160        provider: String,
161        expected: u64,
162        actual: u64,
163    },
164    #[error("embedding provider {provider:?} is referenced by {references} schema columns")]
165    ProviderInUse { provider: String, references: usize },
166    #[error("embedding provider {provider:?}: {message}")]
167    ProviderFailed { provider: String, message: String },
168    #[error("embedding source is SuppliedByApplication")]
169    SuppliedByApplication,
170    #[error("embedding provider identity mismatch: expected {expected:?}, got {got:?}")]
171    ProviderIdentityMismatch { expected: String, got: String },
172    #[error("embedding model identity mismatch: expected {expected:?}, got {got:?}")]
173    ModelIdentityMismatch { expected: String, got: String },
174    #[error("embedding model version mismatch: expected {expected:?}, got {got:?}")]
175    ModelVersionMismatch { expected: String, got: String },
176    #[error("embedding output count mismatch: expected {expected}, got {got}")]
177    OutputCountMismatch { expected: usize, got: usize },
178    #[error("embedding dimension mismatch: expected {expected}, got {got}")]
179    DimensionMismatch { expected: u32, got: u32 },
180    #[error("embedding request exceeds {resource}: requested {requested}, limit {limit}")]
181    LimitExceeded {
182        resource: &'static str,
183        requested: usize,
184        limit: usize,
185    },
186    #[error("embedding output is not finite")]
187    NonFiniteOutput,
188    #[error("embedding output is not L2 normalized")]
189    NormalizationMismatch,
190    #[error("embedding execution cancelled or timed out: {0}")]
191    Execution(String),
192}
193
194#[derive(Debug, Clone, Copy, PartialEq, Eq)]
195pub enum ProviderHealth {
196    Ready,
197    Degraded,
198    Unavailable,
199}
200
201/// Where provider work executes.
202#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
203pub enum ProviderExecutionMode {
204    /// Short work that cooperatively checks [`crate::ExecutionControl`].
205    #[default]
206    Cooperative,
207    /// CPU-heavy or blocking local work. The registry runs it on its bounded
208    /// blocking pool.
209    Blocking,
210    /// A genuinely asynchronous remote transport implemented by
211    /// [`EmbeddingProvider::embed_async`].
212    Remote,
213}
214
215#[derive(Debug, Clone, Copy, PartialEq, Eq)]
216pub struct EmbeddingLimits {
217    pub max_texts: usize,
218    pub max_text_bytes: usize,
219    pub max_total_input_bytes: usize,
220    pub max_dimension: u32,
221    pub max_output_bytes: usize,
222}
223
224impl Default for EmbeddingLimits {
225    fn default() -> Self {
226        Self {
227            max_texts: DEFAULT_MAX_EMBEDDING_TEXTS,
228            max_text_bytes: DEFAULT_MAX_EMBEDDING_TEXT_BYTES,
229            max_total_input_bytes: DEFAULT_MAX_EMBEDDING_INPUT_BYTES,
230            max_dimension: crate::schema::Schema::MAX_EMBEDDING_DIM,
231            max_output_bytes: DEFAULT_MAX_EMBEDDING_OUTPUT_BYTES,
232        }
233    }
234}
235
236pub struct EmbeddingRequest<'a> {
237    pub texts: &'a [&'a str],
238    pub control: &'a crate::ExecutionControl,
239    pub trace_id: &'a str,
240}
241
242#[derive(Debug, Clone, PartialEq)]
243pub struct EmbeddingResponse {
244    pub vectors: Vec<Vec<f32>>,
245}
246
247pub type EmbeddingFuture<'a> =
248    Pin<Box<dyn Future<Output = Result<EmbeddingResponse, EmbeddingError>> + Send + 'a>>;
249
250pub trait EmbeddingProvider: Send + Sync {
251    fn provider_id(&self) -> &str;
252    fn model_id(&self) -> &str;
253    fn model_version(&self) -> &str;
254    fn dimension(&self) -> u32;
255    fn normalization(&self) -> EmbeddingNormalization;
256    fn preprocessing_version(&self) -> &str;
257
258    fn execution_mode(&self) -> ProviderExecutionMode {
259        ProviderExecutionMode::Cooperative
260    }
261
262    fn health(&self) -> ProviderHealth {
263        ProviderHealth::Ready
264    }
265
266    /// Controlled synchronous inference used by transactional generated
267    /// columns. Implementations must checkpoint during long work.
268    fn embed(&self, request: EmbeddingRequest<'_>) -> Result<EmbeddingResponse, EmbeddingError>;
269
270    /// Async transport hook for remote providers. The default is suitable only
271    /// for short cooperative local inference; heavy local implementations must
272    /// override this and use a bounded blocking executor.
273    fn embed_async<'a>(&'a self, request: EmbeddingRequest<'a>) -> EmbeddingFuture<'a> {
274        Box::pin(async move { self.embed(request) })
275    }
276}
277
278#[derive(Clone)]
279struct ProviderEntry {
280    generation: u64,
281    provider: Arc<dyn EmbeddingProvider>,
282}
283
284#[derive(Debug, Clone, PartialEq, Eq)]
285pub struct ProviderStatus {
286    pub provider_id: String,
287    pub model_id: String,
288    pub model_version: String,
289    pub generation: u64,
290    pub health: ProviderHealth,
291}
292
293#[derive(Clone)]
294pub struct EmbeddingProviderRegistry {
295    inner: Arc<RwLock<BTreeMap<String, ProviderEntry>>>,
296    blocking_slots: Arc<tokio::sync::Semaphore>,
297}
298
299impl Default for EmbeddingProviderRegistry {
300    fn default() -> Self {
301        let concurrency = std::thread::available_parallelism()
302            .map(usize::from)
303            .unwrap_or(1)
304            .clamp(1, 32);
305        Self {
306            inner: Arc::new(RwLock::new(BTreeMap::new())),
307            blocking_slots: Arc::new(tokio::sync::Semaphore::new(concurrency)),
308        }
309    }
310}
311
312impl std::fmt::Debug for EmbeddingProviderRegistry {
313    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
314        f.debug_struct("EmbeddingProviderRegistry")
315            .field("providers", &self.list_ids())
316            .finish()
317    }
318}
319
320impl EmbeddingProviderRegistry {
321    pub fn new() -> Self {
322        Self::default()
323    }
324
325    /// Registry with an explicit blocking-provider concurrency ceiling.
326    pub fn with_max_blocking_concurrency(max_concurrency: usize) -> Self {
327        Self {
328            inner: Arc::new(RwLock::new(BTreeMap::new())),
329            blocking_slots: Arc::new(tokio::sync::Semaphore::new(max_concurrency.max(1))),
330        }
331    }
332
333    pub fn register_new(
334        &self,
335        provider: Arc<dyn EmbeddingProvider>,
336    ) -> Result<u64, EmbeddingError> {
337        let id = provider.provider_id().to_owned();
338        let mut providers = self
339            .inner
340            .write()
341            .unwrap_or_else(|error| error.into_inner());
342        if providers.contains_key(&id) {
343            return Err(EmbeddingError::ProviderAlreadyRegistered(id));
344        }
345        providers.insert(
346            id,
347            ProviderEntry {
348                generation: 1,
349                provider,
350            },
351        );
352        Ok(1)
353    }
354
355    /// Compatibility entry point for 0.60.3 Kit clients. Duplicate IDs keep
356    /// the existing provider instead of silently replacing it.
357    #[deprecated(note = "use register_new and handle duplicate provider IDs")]
358    pub fn register(&self, provider: Arc<dyn EmbeddingProvider>) {
359        let _ = self.register_new(provider);
360    }
361
362    pub fn replace(
363        &self,
364        expected_generation: u64,
365        provider: Arc<dyn EmbeddingProvider>,
366    ) -> Result<u64, EmbeddingError> {
367        let id = provider.provider_id().to_owned();
368        let mut providers = self
369            .inner
370            .write()
371            .unwrap_or_else(|error| error.into_inner());
372        let entry = providers
373            .get_mut(&id)
374            .ok_or_else(|| EmbeddingError::ProviderNotFound(id.clone()))?;
375        if entry.generation != expected_generation {
376            return Err(EmbeddingError::ProviderGenerationMismatch {
377                provider: id,
378                expected: expected_generation,
379                actual: entry.generation,
380            });
381        }
382        entry.generation = entry.generation.saturating_add(1);
383        entry.provider = provider;
384        Ok(entry.generation)
385    }
386
387    pub(crate) fn unregister_unreferenced(
388        &self,
389        provider_id: &str,
390        references: usize,
391    ) -> Result<bool, EmbeddingError> {
392        if references != 0 {
393            return Err(EmbeddingError::ProviderInUse {
394                provider: provider_id.to_owned(),
395                references,
396            });
397        }
398        Ok(self
399            .inner
400            .write()
401            .unwrap_or_else(|error| error.into_inner())
402            .remove(provider_id)
403            .is_some())
404    }
405
406    pub fn get(&self, provider_id: &str) -> Option<Arc<dyn EmbeddingProvider>> {
407        self.inner
408            .read()
409            .unwrap_or_else(|error| error.into_inner())
410            .get(provider_id)
411            .map(|entry| Arc::clone(&entry.provider))
412    }
413
414    pub fn list_ids(&self) -> Vec<String> {
415        self.inner
416            .read()
417            .unwrap_or_else(|error| error.into_inner())
418            .keys()
419            .cloned()
420            .collect()
421    }
422
423    pub fn status(&self, provider_id: &str) -> Option<ProviderStatus> {
424        let providers = self.inner.read().unwrap_or_else(|error| error.into_inner());
425        let entry = providers.get(provider_id)?;
426        Some(ProviderStatus {
427            provider_id: provider_id.to_owned(),
428            model_id: entry.provider.model_id().to_owned(),
429            model_version: entry.provider.model_version().to_owned(),
430            generation: entry.generation,
431            health: entry.provider.health(),
432        })
433    }
434
435    pub fn resolve(
436        &self,
437        source: &EmbeddingSource,
438    ) -> Result<Arc<dyn EmbeddingProvider>, EmbeddingError> {
439        let provider_id = source
440            .provider_id()
441            .ok_or(EmbeddingError::SuppliedByApplication)?;
442        self.get(provider_id)
443            .ok_or_else(|| EmbeddingError::ProviderNotFound(provider_id.to_owned()))
444    }
445
446    pub fn embed(
447        &self,
448        source: &EmbeddingSource,
449        texts: &[&str],
450        expected_dim: u32,
451    ) -> Result<Vec<Vec<f32>>, EmbeddingError> {
452        self.embed_controlled(
453            source,
454            texts,
455            expected_dim,
456            &crate::ExecutionControl::new(None),
457            "embedding",
458            EmbeddingLimits::default(),
459        )
460    }
461
462    pub fn embed_controlled(
463        &self,
464        source: &EmbeddingSource,
465        texts: &[&str],
466        expected_dim: u32,
467        control: &crate::ExecutionControl,
468        trace_id: &str,
469        limits: EmbeddingLimits,
470    ) -> Result<Vec<Vec<f32>>, EmbeddingError> {
471        control
472            .checkpoint()
473            .map_err(|error| EmbeddingError::Execution(error.to_string()))?;
474        validate_request(texts, expected_dim, limits)?;
475        let provider = self.resolve(source)?;
476        validate_identity(source, provider.as_ref())?;
477        if provider.health() == ProviderHealth::Unavailable {
478            return Err(EmbeddingError::ProviderFailed {
479                provider: provider.provider_id().to_owned(),
480                message: "provider unavailable".into(),
481            });
482        }
483        if provider.dimension() != expected_dim {
484            return Err(EmbeddingError::DimensionMismatch {
485                expected: expected_dim,
486                got: provider.dimension(),
487            });
488        }
489        let response = provider.embed(EmbeddingRequest {
490            texts,
491            control,
492            trace_id,
493        })?;
494        control
495            .checkpoint()
496            .map_err(|error| EmbeddingError::Execution(error.to_string()))?;
497        validate_response(
498            provider.as_ref(),
499            texts.len(),
500            expected_dim,
501            limits,
502            &response.vectors,
503        )?;
504        Ok(response.vectors)
505    }
506
507    /// Async provider execution for server/runtime callers.
508    ///
509    /// Blocking local providers run under a bounded semaphore on Tokio's
510    /// blocking pool. Remote providers must implement `embed_async`; both
511    /// paths stop waiting when the execution control is cancelled or expires.
512    pub async fn embed_async_controlled(
513        &self,
514        source: &EmbeddingSource,
515        texts: &[&str],
516        expected_dim: u32,
517        control: &crate::ExecutionControl,
518        trace_id: &str,
519        limits: EmbeddingLimits,
520    ) -> Result<Vec<Vec<f32>>, EmbeddingError> {
521        control
522            .checkpoint()
523            .map_err(|error| EmbeddingError::Execution(error.to_string()))?;
524        validate_request(texts, expected_dim, limits)?;
525        let provider = self.resolve(source)?;
526        validate_identity(source, provider.as_ref())?;
527        if provider.health() == ProviderHealth::Unavailable {
528            return Err(EmbeddingError::ProviderFailed {
529                provider: provider.provider_id().to_owned(),
530                message: "provider unavailable".into(),
531            });
532        }
533        if provider.dimension() != expected_dim {
534            return Err(EmbeddingError::DimensionMismatch {
535                expected: expected_dim,
536                got: provider.dimension(),
537            });
538        }
539
540        let response = match provider.execution_mode() {
541            ProviderExecutionMode::Cooperative | ProviderExecutionMode::Remote => {
542                let request = EmbeddingRequest {
543                    texts,
544                    control,
545                    trace_id,
546                };
547                tokio::select! {
548                    result = provider.embed_async(request) => result?,
549                    _ = control.cancelled() => {
550                        return Err(execution_stopped(control));
551                    }
552                }
553            }
554            ProviderExecutionMode::Blocking => {
555                let permit = tokio::select! {
556                    result = Arc::clone(&self.blocking_slots).acquire_owned() => {
557                        result.map_err(|_| EmbeddingError::Execution(
558                            "embedding blocking executor closed".into()
559                        ))?
560                    }
561                    _ = control.cancelled() => {
562                        return Err(execution_stopped(control));
563                    }
564                };
565                let provider = Arc::clone(&provider);
566                let owned_texts = texts
567                    .iter()
568                    .map(|text| (*text).to_owned())
569                    .collect::<Vec<_>>();
570                let owned_control = control.clone();
571                let owned_trace_id = trace_id.to_owned();
572                let task = tokio::task::spawn_blocking(move || {
573                    let _permit = permit;
574                    let text_refs = owned_texts.iter().map(String::as_str).collect::<Vec<_>>();
575                    provider.embed(EmbeddingRequest {
576                        texts: &text_refs,
577                        control: &owned_control,
578                        trace_id: &owned_trace_id,
579                    })
580                });
581                tokio::select! {
582                    result = task => {
583                        result
584                            .map_err(|error| EmbeddingError::Execution(
585                                format!("embedding blocking task failed: {error}")
586                            ))??
587                    }
588                    _ = control.cancelled() => {
589                        return Err(execution_stopped(control));
590                    }
591                }
592            }
593        };
594
595        control
596            .checkpoint()
597            .map_err(|error| EmbeddingError::Execution(error.to_string()))?;
598        validate_response(
599            provider.as_ref(),
600            texts.len(),
601            expected_dim,
602            limits,
603            &response.vectors,
604        )?;
605        Ok(response.vectors)
606    }
607}
608
609fn execution_stopped(control: &crate::ExecutionControl) -> EmbeddingError {
610    let message = control
611        .checkpoint()
612        .expect_err("cancelled future completed without an execution error")
613        .to_string();
614    EmbeddingError::Execution(message)
615}
616
617fn validate_identity(
618    source: &EmbeddingSource,
619    provider: &dyn EmbeddingProvider,
620) -> Result<(), EmbeddingError> {
621    if let EmbeddingSource::GeneratedColumnSpec { spec } = source {
622        if spec.normalization != provider.normalization() {
623            return Err(EmbeddingError::NormalizationMismatch);
624        }
625        if spec.dimension != provider.dimension() {
626            return Err(EmbeddingError::DimensionMismatch {
627                expected: spec.dimension,
628                got: provider.dimension(),
629            });
630        }
631    }
632    if source
633        .provider_id()
634        .is_some_and(|id| id != provider.provider_id())
635    {
636        return Err(EmbeddingError::ProviderIdentityMismatch {
637            expected: source.provider_id().unwrap_or_default().to_owned(),
638            got: provider.provider_id().to_owned(),
639        });
640    }
641    if source
642        .model_id()
643        .is_some_and(|id| id != provider.model_id())
644    {
645        return Err(EmbeddingError::ModelIdentityMismatch {
646            expected: source.model_id().unwrap_or_default().to_owned(),
647            got: provider.model_id().to_owned(),
648        });
649    }
650    if source
651        .model_version()
652        .is_some_and(|version| version != provider.model_version())
653    {
654        return Err(EmbeddingError::ModelVersionMismatch {
655            expected: source.model_version().unwrap_or_default().to_owned(),
656            got: provider.model_version().to_owned(),
657        });
658    }
659    Ok(())
660}
661
662fn validate_request(
663    texts: &[&str],
664    dimension: u32,
665    limits: EmbeddingLimits,
666) -> Result<(), EmbeddingError> {
667    if texts.len() > limits.max_texts {
668        return Err(EmbeddingError::LimitExceeded {
669            resource: "text count",
670            requested: texts.len(),
671            limit: limits.max_texts,
672        });
673    }
674    if dimension > limits.max_dimension {
675        return Err(EmbeddingError::LimitExceeded {
676            resource: "dimension",
677            requested: dimension as usize,
678            limit: limits.max_dimension as usize,
679        });
680    }
681    let mut total = 0usize;
682    for text in texts {
683        if text.len() > limits.max_text_bytes {
684            return Err(EmbeddingError::LimitExceeded {
685                resource: "text bytes",
686                requested: text.len(),
687                limit: limits.max_text_bytes,
688            });
689        }
690        total = total.saturating_add(text.len());
691    }
692    if total > limits.max_total_input_bytes {
693        return Err(EmbeddingError::LimitExceeded {
694            resource: "total input bytes",
695            requested: total,
696            limit: limits.max_total_input_bytes,
697        });
698    }
699    Ok(())
700}
701
702fn validate_response(
703    provider: &dyn EmbeddingProvider,
704    expected_count: usize,
705    expected_dim: u32,
706    limits: EmbeddingLimits,
707    vectors: &[Vec<f32>],
708) -> Result<(), EmbeddingError> {
709    if vectors.len() != expected_count {
710        return Err(EmbeddingError::OutputCountMismatch {
711            expected: expected_count,
712            got: vectors.len(),
713        });
714    }
715    let output_bytes = vectors
716        .len()
717        .saturating_mul(expected_dim as usize)
718        .saturating_mul(std::mem::size_of::<f32>());
719    if output_bytes > limits.max_output_bytes {
720        return Err(EmbeddingError::LimitExceeded {
721            resource: "output bytes",
722            requested: output_bytes,
723            limit: limits.max_output_bytes,
724        });
725    }
726    for vector in vectors {
727        if vector.len() as u32 != expected_dim {
728            return Err(EmbeddingError::DimensionMismatch {
729                expected: expected_dim,
730                got: vector.len() as u32,
731            });
732        }
733        if vector.iter().any(|value| !value.is_finite()) {
734            return Err(EmbeddingError::NonFiniteOutput);
735        }
736        if provider.normalization() == EmbeddingNormalization::L2 {
737            let norm = vector
738                .iter()
739                .map(|value| (*value as f64) * (*value as f64))
740                .sum::<f64>()
741                .sqrt();
742            if (norm - 1.0).abs() > 1e-4 {
743                return Err(EmbeddingError::NormalizationMismatch);
744            }
745        }
746    }
747    Ok(())
748}
749
750#[derive(Debug, Clone)]
751pub struct FixedVectorProvider {
752    pub id: String,
753    pub model_id: String,
754    pub model_version: String,
755    pub normalization: EmbeddingNormalization,
756    pub vector: Vec<f32>,
757}
758
759impl EmbeddingProvider for FixedVectorProvider {
760    fn provider_id(&self) -> &str {
761        &self.id
762    }
763    fn model_id(&self) -> &str {
764        &self.model_id
765    }
766    fn model_version(&self) -> &str {
767        &self.model_version
768    }
769    fn dimension(&self) -> u32 {
770        self.vector.len() as u32
771    }
772    fn normalization(&self) -> EmbeddingNormalization {
773        self.normalization
774    }
775    fn preprocessing_version(&self) -> &str {
776        "fixed-test-v1"
777    }
778    fn embed(&self, request: EmbeddingRequest<'_>) -> Result<EmbeddingResponse, EmbeddingError> {
779        request
780            .control
781            .checkpoint()
782            .map_err(|error| EmbeddingError::Execution(error.to_string()))?;
783        Ok(EmbeddingResponse {
784            vectors: request.texts.iter().map(|_| self.vector.clone()).collect(),
785        })
786    }
787}
788
789#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
790pub struct EmbeddingModelMeta {
791    pub provider_id: Option<String>,
792    pub model_id: Option<String>,
793    pub model_version: Option<String>,
794    pub preprocessing_version: Option<String>,
795    pub normalization: EmbeddingNormalization,
796    pub dimension: u32,
797    pub source_kind: Option<String>,
798}
799
800impl EmbeddingModelMeta {
801    pub fn from_provider(source: &EmbeddingSource, provider: &dyn EmbeddingProvider) -> Self {
802        Self {
803            provider_id: Some(provider.provider_id().to_owned()),
804            model_id: Some(provider.model_id().to_owned()),
805            model_version: Some(provider.model_version().to_owned()),
806            preprocessing_version: Some(provider.preprocessing_version().to_owned()),
807            normalization: provider.normalization(),
808            dimension: provider.dimension(),
809            source_kind: Some(source.label().to_owned()),
810        }
811    }
812}
813
814#[cfg(test)]
815mod tests {
816    use super::*;
817
818    struct SlowBlockingProvider;
819
820    impl EmbeddingProvider for SlowBlockingProvider {
821        fn provider_id(&self) -> &str {
822            "slow-blocking"
823        }
824        fn model_id(&self) -> &str {
825            "slow"
826        }
827        fn model_version(&self) -> &str {
828            "1"
829        }
830        fn dimension(&self) -> u32 {
831            2
832        }
833        fn normalization(&self) -> EmbeddingNormalization {
834            EmbeddingNormalization::None
835        }
836        fn preprocessing_version(&self) -> &str {
837            "slow-v1"
838        }
839        fn execution_mode(&self) -> ProviderExecutionMode {
840            ProviderExecutionMode::Blocking
841        }
842        fn embed(
843            &self,
844            request: EmbeddingRequest<'_>,
845        ) -> Result<EmbeddingResponse, EmbeddingError> {
846            std::thread::sleep(std::time::Duration::from_millis(50));
847            Ok(EmbeddingResponse {
848                vectors: request.texts.iter().map(|_| vec![1.0, 2.0]).collect(),
849            })
850        }
851    }
852
853    fn provider() -> Arc<FixedVectorProvider> {
854        Arc::new(FixedVectorProvider {
855            id: "local-test".into(),
856            model_id: "fixed".into(),
857            model_version: "1".into(),
858            normalization: EmbeddingNormalization::L2,
859            vector: vec![0.0, 1.0],
860        })
861    }
862
863    fn source() -> EmbeddingSource {
864        EmbeddingSource::GeneratedColumnSpec {
865            spec: GeneratedEmbeddingSpec {
866                provider_id: "local-test".into(),
867                model_id: "fixed".into(),
868                model_version: "1".into(),
869                source_columns: vec![1],
870                input_template: "{text}".into(),
871                dimension: 2,
872                normalization: EmbeddingNormalization::L2,
873                failure_policy: EmbeddingFailurePolicy::AbortWrite,
874            },
875        }
876    }
877
878    #[test]
879    fn registration_requires_expected_generation() {
880        let registry = EmbeddingProviderRegistry::new();
881        assert_eq!(registry.register_new(provider()).unwrap(), 1);
882        assert!(matches!(
883            registry.register_new(provider()),
884            Err(EmbeddingError::ProviderAlreadyRegistered(_))
885        ));
886        assert!(matches!(
887            registry.replace(9, provider()),
888            Err(EmbeddingError::ProviderGenerationMismatch { .. })
889        ));
890        assert_eq!(registry.replace(1, provider()).unwrap(), 2);
891        assert_eq!(registry.status("local-test").unwrap().generation, 2);
892    }
893
894    #[test]
895    fn validates_count_dimension_finiteness_limits_and_cancellation() {
896        let registry = EmbeddingProviderRegistry::new();
897        registry.register_new(provider()).unwrap();
898        assert_eq!(registry.embed(&source(), &["a", "b"], 2).unwrap().len(), 2);
899        let control = crate::ExecutionControl::new(None);
900        control.cancel(crate::CancellationReason::ClientRequest);
901        assert!(matches!(
902            registry.embed_controlled(
903                &source(),
904                &["a"],
905                2,
906                &control,
907                "cancelled",
908                EmbeddingLimits::default(),
909            ),
910            Err(EmbeddingError::Execution(_))
911        ));
912        assert!(matches!(
913            registry.embed_controlled(
914                &source(),
915                &["too long"],
916                2,
917                &crate::ExecutionControl::new(None),
918                "bounded",
919                EmbeddingLimits {
920                    max_text_bytes: 2,
921                    ..EmbeddingLimits::default()
922                },
923            ),
924            Err(EmbeddingError::LimitExceeded { .. })
925        ));
926    }
927
928    #[test]
929    fn legacy_local_model_path_is_process_local_only() {
930        let source = EmbeddingSource::LocalModel {
931            model_id: "legacy".into(),
932            model_path: "/node-only/model.onnx".into(),
933        };
934        let encoded = serde_json::to_string(&source).unwrap();
935        assert!(!encoded.contains("node-only"));
936        assert_eq!(
937            serde_json::from_str::<EmbeddingSource>(&encoded).unwrap(),
938            EmbeddingSource::LocalModel {
939                model_id: "legacy".into(),
940                model_path: std::path::PathBuf::new(),
941            }
942        );
943    }
944
945    struct WrongCount;
946    impl EmbeddingProvider for WrongCount {
947        fn provider_id(&self) -> &str {
948            "wrong"
949        }
950        fn model_id(&self) -> &str {
951            "wrong-model"
952        }
953        fn model_version(&self) -> &str {
954            "1"
955        }
956        fn dimension(&self) -> u32 {
957            2
958        }
959        fn normalization(&self) -> EmbeddingNormalization {
960            EmbeddingNormalization::None
961        }
962        fn preprocessing_version(&self) -> &str {
963            "1"
964        }
965        fn embed(
966            &self,
967            _request: EmbeddingRequest<'_>,
968        ) -> Result<EmbeddingResponse, EmbeddingError> {
969            Ok(EmbeddingResponse {
970                vectors: Vec::new(),
971            })
972        }
973    }
974
975    #[test]
976    fn rejects_wrong_output_count() {
977        let registry = EmbeddingProviderRegistry::new();
978        registry.register_new(Arc::new(WrongCount)).unwrap();
979        let source = EmbeddingSource::ConfiguredModel {
980            provider_id: "wrong".into(),
981            model_id: "wrong-model".into(),
982            model_version: "1".into(),
983        };
984        assert!(matches!(
985            registry.embed(&source, &["a"], 2),
986            Err(EmbeddingError::OutputCountMismatch {
987                expected: 1,
988                got: 0
989            })
990        ));
991    }
992
993    #[tokio::test(flavor = "current_thread")]
994    async fn async_path_bounds_blocking_provider_and_honors_deadline() {
995        let registry = EmbeddingProviderRegistry::with_max_blocking_concurrency(1);
996        registry
997            .register_new(Arc::new(SlowBlockingProvider))
998            .unwrap();
999        let source = EmbeddingSource::ConfiguredModel {
1000            provider_id: "slow-blocking".into(),
1001            model_id: "slow".into(),
1002            model_version: "1".into(),
1003        };
1004        let control = crate::ExecutionControl::with_timeout(std::time::Duration::from_millis(5));
1005        assert!(matches!(
1006            registry
1007                .embed_async_controlled(
1008                    &source,
1009                    &["text"],
1010                    2,
1011                    &control,
1012                    "deadline-test",
1013                    EmbeddingLimits::default(),
1014                )
1015                .await,
1016            Err(EmbeddingError::Execution(_))
1017        ));
1018    }
1019}
mongreldb_core/embedding.rs

mongreldb_core/
embedding.rs
