haem_frames/
embedding.rs

1//! One vocabulary and one set of checks for both sides of an embedding exchange.
2
3use serde::{Deserialize, Serialize};
4
5mod checks;
6mod codec;
7mod vector;
8
9pub use checks::{
10    check_embed, check_model_limits, check_model_limits_windowed, check_ready, check_start,
11    check_vectors, check_vectors_windowed, input_refusal, vectors_message_size,
12    windows_message_size,
13};
14pub use codec::{Message, decode_message, read_message, read_start, write_message};
15pub use vector::{decode_vector, encode_vector};
16
17pub const PROTOCOL: u32 = 1;
18/// EMB-E2: the protocol a Start names when it asks for windows. A worker
19/// answers Ready with the protocol its Start named.
20pub const PROTOCOL_WINDOWS: u32 = 2;
21pub const MINIMUM_FRAME_BYTES: u32 = 8_192;
22pub const FAILED_MESSAGE_CHARS: usize = 1_024;
23pub const MEMORY_LIMIT_EXIT: i32 = 86;
24
25#[derive(Clone, Copy, Debug, Eq, PartialEq)]
26#[repr(u8)]
27pub enum Tag {
28    Start = 16,
29    Ready = 17,
30    Embed = 18,
31    Vectors = 19,
32    Failed = 20,
33}
34
35impl TryFrom<u8> for Tag {
36    type Error = Error;
37    fn try_from(value: u8) -> Result<Self, Error> {
38        match value {
39            16 => Ok(Self::Start),
40            17 => Ok(Self::Ready),
41            18 => Ok(Self::Embed),
42            19 => Ok(Self::Vectors),
43            20 => Ok(Self::Failed),
44            _ => Err(Error::new(
45                FailedCode::EmbedProtocol,
46                "unknown embedding frame tag",
47            )),
48        }
49    }
50}
51
52#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
53#[serde(rename_all = "snake_case")]
54pub enum FailedCode {
55    EmbedModelMismatch,
56    EmbedModelMissing,
57    EmbedLimits,
58    EmbedBatchTooLarge,
59    EmbedOutputInvalid,
60    EmbedProtocol,
61    EmbedInferenceFailed,
62}
63
64impl FailedCode {
65    #[must_use]
66    pub const fn as_str(self) -> &'static str {
67        match self {
68            Self::EmbedModelMismatch => "embed_model_mismatch",
69            Self::EmbedModelMissing => "embed_model_missing",
70            Self::EmbedLimits => "embed_limits",
71            Self::EmbedBatchTooLarge => "embed_batch_too_large",
72            Self::EmbedOutputInvalid => "embed_output_invalid",
73            Self::EmbedProtocol => "embed_protocol",
74            Self::EmbedInferenceFailed => "embed_inference_failed",
75        }
76    }
77}
78
79#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
80#[serde(rename_all = "snake_case")]
81pub enum ItemCode {
82    EmbedInputEmpty,
83    EmbedInputTooLarge,
84    TextLongerThanModel,
85    /// EMB-E2: a text whose whole cover needs more windows than
86    /// `Windows::max_windows`. It gets no vector; it is never embedded as its
87    /// first windows.
88    WindowsOverCap,
89}
90
91#[derive(Debug)]
92pub struct Error {
93    pub code: FailedCode,
94    pub message: String,
95}
96
97impl Error {
98    pub(super) fn new(code: FailedCode, message: impl Into<String>) -> Self {
99        Self {
100            code,
101            message: message.into(),
102        }
103    }
104}
105
106impl std::fmt::Display for Error {
107    fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108        write!(output, "{}: {}", self.code.as_str(), self.message)
109    }
110}
111
112impl std::error::Error for Error {}
113
114#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
115#[serde(deny_unknown_fields)]
116pub struct Limits {
117    pub memory_bytes: u64,
118    pub threads: u64,
119    pub batch_items: u64,
120    pub input_bytes: u64,
121    pub tokens: u64,
122    pub frame_bytes: u32,
123}
124
125#[derive(Debug, Deserialize, Serialize)]
126#[serde(deny_unknown_fields)]
127pub struct Start {
128    pub protocol: u32,
129    pub model_dir: String,
130    #[serde(deserialize_with = "codec::object")]
131    pub limits: Limits,
132    /// EMB-E2: present exactly when `protocol` is [`PROTOCOL_WINDOWS`].
133    #[serde(default, skip_serializing_if = "Option::is_none")]
134    pub windows: Option<Windows>,
135}
136
137/// EMB-E2: how a text longer than the model reads is embedded whole.
138///
139/// Each window holds at most the model's content tokens; consecutive windows
140/// share at least `overlap_tokens`; a text needing more than `max_windows` is
141/// refused `windows_over_cap`. No window is ever truncated.
142#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
143#[serde(deny_unknown_fields)]
144pub struct Windows {
145    pub overlap_tokens: u64,
146    pub max_windows: u64,
147}
148
149#[derive(Debug, Deserialize, Serialize)]
150#[serde(deny_unknown_fields)]
151pub struct Ready {
152    pub protocol: u32,
153    pub worker: String,
154    pub core_limit: u64,
155    pub descriptors_closed: u64,
156    pub environment: u64,
157    #[serde(deserialize_with = "codec::object")]
158    pub model: Model,
159}
160
161#[derive(Debug, Deserialize, Serialize)]
162#[serde(deny_unknown_fields)]
163pub struct Model {
164    pub name: String,
165    pub revision: String,
166    pub manifest_sha256: String,
167    pub dimensions: u64,
168    pub max_tokens: u64,
169    pub special_tokens: u64,
170    pub prefix_tokens: u64,
171    pub normalised: bool,
172    pub distance: Distance,
173}
174
175#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
176#[serde(rename_all = "snake_case")]
177pub enum Distance {
178    Cosine,
179    Dot,
180    Euclidean,
181}
182
183#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
184#[serde(rename_all = "snake_case")]
185pub enum Kind {
186    Document,
187    Query,
188}
189
190#[derive(Debug, Deserialize, Serialize)]
191#[serde(deny_unknown_fields)]
192pub struct Embed {
193    pub kind: Kind,
194    #[serde(deserialize_with = "codec::objects")]
195    pub items: Vec<Input>,
196}
197
198#[derive(Debug, Deserialize, Serialize)]
199#[serde(deny_unknown_fields)]
200pub struct Input {
201    pub id: String,
202    pub text: String,
203}
204
205#[derive(Debug, Deserialize, Serialize)]
206#[serde(deny_unknown_fields)]
207pub struct Vectors {
208    #[serde(deserialize_with = "codec::objects")]
209    pub items: Vec<Outcome>,
210}
211
212#[derive(Debug, Deserialize, Serialize)]
213#[serde(tag = "outcome", rename_all = "snake_case", deny_unknown_fields)]
214pub enum Outcome {
215    Vector {
216        id: String,
217        vector: String,
218        tokens_read: u64,
219        tokens_total: u64,
220    },
221    /// EMB-E2: a text longer than the model reads, embedded whole as
222    /// overlapping windows in text order.
223    Windows {
224        id: String,
225        #[serde(deserialize_with = "codec::objects")]
226        windows: Vec<Window>,
227        /// The content tokens of the whole text.
228        tokens_total: u64,
229    },
230    Refused {
231        id: String,
232        code: ItemCode,
233        /// The content tokens of the whole text, without special or prefix tokens; only with
234        /// `text_longer_than_model`.
235        #[serde(default, skip_serializing_if = "Option::is_none")]
236        tokens_total: Option<u64>,
237        /// The model's limit as declared in Ready; only with `text_longer_than_model`.
238        #[serde(default, skip_serializing_if = "Option::is_none")]
239        tokens_limit: Option<u64>,
240    },
241}
242
243impl Outcome {
244    #[must_use]
245    pub fn id(&self) -> &str {
246        match self {
247            Self::Vector { id, .. } | Self::Windows { id, .. } | Self::Refused { id, .. } => id,
248        }
249    }
250}
251
252/// EMB-E2: one window of a text: its bytes `[start, end)`, on char and token
253/// boundaries, and their vector.
254#[derive(Debug, Deserialize, Serialize)]
255#[serde(deny_unknown_fields)]
256pub struct Window {
257    pub start: u64,
258    pub end: u64,
259    pub vector: String,
260    pub tokens_read: u64,
261}
262
263#[derive(Debug, Deserialize, Serialize)]
264#[serde(deny_unknown_fields)]
265pub struct Failed {
266    pub code: FailedCode,
267    pub message: String,
268}
haem_frames/embedding.rs

haem_frames/
embedding.rs
