seyk_embed/
embedder.rs

1use serde::{Deserialize, Serialize};
2
3use crate::{Dimension, Embedding, EmbeddingTask, Modality, Result};
4
5/// A borrowed input to embed. The request selects its task independently.
6#[derive(Debug, Clone)]
7#[non_exhaustive]
8pub enum ResourceInput<'a> {
9    /// A text segment. `title` is a heading or file name, when one is known.
10    Text {
11        title: Option<&'a str>,
12        text: &'a str,
13    },
14    /// An encoded image file (PNG, JPEG, ...).
15    Image { bytes: &'a [u8] },
16    /// Mono PCM samples at the given sample rate.
17    Audio {
18        samples: &'a [f32],
19        sample_rate: u32,
20    },
21    /// Decoded frames with their timestamps. Reserved: no shipped embedder accepts video yet.
22    Video { frames: &'a [VideoFrame<'a>] },
23}
24
25impl ResourceInput<'_> {
26    pub fn modality(&self) -> Modality {
27        match self {
28            ResourceInput::Text { .. } => Modality::Text,
29            ResourceInput::Image { .. } => Modality::Image,
30            ResourceInput::Audio { .. } => Modality::Audio,
31            ResourceInput::Video { .. } => Modality::Video,
32        }
33    }
34}
35
36/// One decoded video frame (RGB8, row-major).
37#[derive(Debug, Clone)]
38pub struct VideoFrame<'a> {
39    pub timestamp_ms: u64,
40    pub width: u32,
41    pub height: u32,
42    pub rgb: &'a [u8],
43}
44
45/// Shared embedding-space configuration used by every modality.
46/// Modality-specific weights/preprocessing additionally use [`Embedder::input_identity`].
47///
48/// Persisted alongside cached vectors. If any field changes, cached vectors
49/// from the old profile are never used.
50#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
51pub struct ModelIdentity {
52    /// For example `google/embeddinggemma-2`.
53    pub model_id: String,
54    /// Exact weights revision (a commit hash when known).
55    pub revision: String,
56    /// Runtime and precision, for example `onnx-q8`.
57    pub variant: String,
58    pub dimension: Dimension,
59    /// How vectors are normalised, for example `l2-after-truncation`.
60    pub normalization: String,
61    pub query_prompt: String,
62    pub document_prompt: String,
63    /// Inputs longer than this many tokens are truncated.
64    pub max_tokens: usize,
65}
66
67/// One borrowed input and its embedding purpose.
68#[derive(Debug, Clone)]
69pub struct EmbeddingRequest<'a> {
70    pub input: ResourceInput<'a>,
71    pub task: EmbeddingTask,
72}
73
74/// Synchronous, task-aware local embedding interface.
75/// Successful batches return one result per request, in input order.
76/// Async callers can use their executor's blocking pool.
77pub trait Embedder: Send + Sync {
78    fn identity(&self) -> &ModelIdentity;
79    fn dimension(&self) -> Dimension {
80        self.identity().dimension
81    }
82    /// Extra modality-specific weights/preparation identity, included in that
83    /// input's cache key without invalidating other modalities' vectors.
84    /// Resolve this before cache admission, even if inference is unnecessary.
85    /// Backends may lazily read model assets here; failed preparation is retryable.
86    /// Once successful, the identity must remain stable until the model is reloaded.
87    fn input_identity(&self, _modality: Modality) -> Result<&str> {
88        Ok("")
89    }
90    /// Why a configured modality is unavailable, including recovery instructions.
91    fn unavailable_reason(&self, _modality: Modality) -> Option<String> {
92        None
93    }
94    fn supports(&self, modality: Modality) -> bool;
95    fn embed_batch(&self, requests: &[EmbeddingRequest<'_>]) -> Result<Vec<Embedding>>;
96
97    fn embed(&self, input: ResourceInput<'_>, task: EmbeddingTask) -> Result<Embedding> {
98        let mut out = self.embed_batch(&[EmbeddingRequest { input, task }])?;
99        if out.len() != 1 {
100            return Err(crate::Error::Embedding(
101                "embedder returned the wrong number of vectors".into(),
102            ));
103        }
104        Ok(out.remove(0))
105    }
106    fn embed_text(&self, text: &str, task: EmbeddingTask) -> Result<Embedding> {
107        self.embed(ResourceInput::Text { title: None, text }, task)
108    }
109    /// Retrieval convenience adapter.
110    fn embed_query(&self, query: &str) -> Result<Embedding> {
111        self.embed_text(query, EmbeddingTask::SearchQuery)
112    }
113    /// Retrieval convenience adapter. Text titles are preserved.
114    fn embed_resources(&self, inputs: &[ResourceInput<'_>]) -> Result<Vec<Embedding>> {
115        let requests: Vec<_> = inputs
116            .iter()
117            .cloned()
118            .map(|input| EmbeddingRequest {
119                input,
120                task: EmbeddingTask::SearchDocument,
121            })
122            .collect();
123        self.embed_batch(&requests)
124    }
125    fn embed_text_document(&self, title: Option<&str>, text: &str) -> Result<Embedding> {
126        self.embed(
127            ResourceInput::Text { title, text },
128            EmbeddingTask::SearchDocument,
129        )
130    }
131}
132
133impl<E: Embedder + ?Sized> Embedder for Box<E> {
134    fn identity(&self) -> &ModelIdentity {
135        (**self).identity()
136    }
137    fn input_identity(&self, modality: Modality) -> Result<&str> {
138        (**self).input_identity(modality)
139    }
140    fn unavailable_reason(&self, modality: Modality) -> Option<String> {
141        (**self).unavailable_reason(modality)
142    }
143    fn supports(&self, modality: Modality) -> bool {
144        (**self).supports(modality)
145    }
146    fn embed_batch(&self, requests: &[EmbeddingRequest<'_>]) -> Result<Vec<Embedding>> {
147        (**self).embed_batch(requests)
148    }
149    fn dimension(&self) -> Dimension {
150        (**self).dimension()
151    }
152    fn embed(&self, input: ResourceInput<'_>, task: EmbeddingTask) -> Result<Embedding> {
153        (**self).embed(input, task)
154    }
155    fn embed_text(&self, text: &str, task: EmbeddingTask) -> Result<Embedding> {
156        (**self).embed_text(text, task)
157    }
158    fn embed_query(&self, query: &str) -> Result<Embedding> {
159        (**self).embed_query(query)
160    }
161    fn embed_resources(&self, inputs: &[ResourceInput<'_>]) -> Result<Vec<Embedding>> {
162        (**self).embed_resources(inputs)
163    }
164    fn embed_text_document(&self, title: Option<&str>, text: &str) -> Result<Embedding> {
165        (**self).embed_text_document(title, text)
166    }
167}
168impl<E: Embedder + ?Sized> Embedder for std::sync::Arc<E> {
169    fn identity(&self) -> &ModelIdentity {
170        (**self).identity()
171    }
172    fn input_identity(&self, modality: Modality) -> Result<&str> {
173        (**self).input_identity(modality)
174    }
175    fn unavailable_reason(&self, modality: Modality) -> Option<String> {
176        (**self).unavailable_reason(modality)
177    }
178    fn supports(&self, modality: Modality) -> bool {
179        (**self).supports(modality)
180    }
181    fn embed_batch(&self, requests: &[EmbeddingRequest<'_>]) -> Result<Vec<Embedding>> {
182        (**self).embed_batch(requests)
183    }
184    fn dimension(&self) -> Dimension {
185        (**self).dimension()
186    }
187    fn embed(&self, input: ResourceInput<'_>, task: EmbeddingTask) -> Result<Embedding> {
188        (**self).embed(input, task)
189    }
190    fn embed_text(&self, text: &str, task: EmbeddingTask) -> Result<Embedding> {
191        (**self).embed_text(text, task)
192    }
193    fn embed_query(&self, query: &str) -> Result<Embedding> {
194        (**self).embed_query(query)
195    }
196    fn embed_resources(&self, inputs: &[ResourceInput<'_>]) -> Result<Vec<Embedding>> {
197        (**self).embed_resources(inputs)
198    }
199    fn embed_text_document(&self, title: Option<&str>, text: &str) -> Result<Embedding> {
200        (**self).embed_text_document(title, text)
201    }
202}
seyk_embed/embedder.rs

seyk_embed/
embedder.rs
