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    pub const fn as_str(self) -> &'static str {
66        match self {
67            Self::EmbedModelMismatch => "embed_model_mismatch",
68            Self::EmbedModelMissing => "embed_model_missing",
69            Self::EmbedLimits => "embed_limits",
70            Self::EmbedBatchTooLarge => "embed_batch_too_large",
71            Self::EmbedOutputInvalid => "embed_output_invalid",
72            Self::EmbedProtocol => "embed_protocol",
73            Self::EmbedInferenceFailed => "embed_inference_failed",
74        }
75    }
76}
77
78#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
79#[serde(rename_all = "snake_case")]
80pub enum ItemCode {
81    EmbedInputEmpty,
82    EmbedInputTooLarge,
83    TextLongerThanModel,
84    /// EMB-E2: a text whose whole cover needs more windows than
85    /// `Windows::max_windows`. It gets no vector; it is never embedded as its
86    /// first windows.
87    WindowsOverCap,
88}
89
90#[derive(Debug)]
91pub struct Error {
92    pub code: FailedCode,
93    pub message: String,
94}
95
96impl Error {
97    pub(super) fn new(code: FailedCode, message: impl Into<String>) -> Self {
98        Self {
99            code,
100            message: message.into(),
101        }
102    }
103}
104
105impl std::fmt::Display for Error {
106    fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107        write!(output, "{}: {}", self.code.as_str(), self.message)
108    }
109}
110
111impl std::error::Error for Error {}
112
113#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
114#[serde(deny_unknown_fields)]
115pub struct Limits {
116    pub memory_bytes: u64,
117    pub threads: u64,
118    pub batch_items: u64,
119    pub input_bytes: u64,
120    pub tokens: u64,
121    pub frame_bytes: u32,
122}
123
124#[derive(Debug, Deserialize, Serialize)]
125#[serde(deny_unknown_fields)]
126pub struct Start {
127    pub protocol: u32,
128    pub model_dir: String,
129    #[serde(deserialize_with = "codec::object")]
130    pub limits: Limits,
131    /// EMB-E2: present exactly when `protocol` is [`PROTOCOL_WINDOWS`].
132    #[serde(default, skip_serializing_if = "Option::is_none")]
133    pub windows: Option<Windows>,
134}
135
136/// EMB-E2: how a text longer than the model reads is embedded whole.
137///
138/// Each window holds at most the model's content tokens; consecutive windows
139/// share at least `overlap_tokens`; a text needing more than `max_windows` is
140/// refused `windows_over_cap`. No window is ever truncated.
141#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
142#[serde(deny_unknown_fields)]
143pub struct Windows {
144    pub overlap_tokens: u64,
145    pub max_windows: u64,
146}
147
148#[derive(Debug, Deserialize, Serialize)]
149#[serde(deny_unknown_fields)]
150pub struct Ready {
151    pub protocol: u32,
152    pub worker: String,
153    pub core_limit: u64,
154    pub descriptors_closed: u64,
155    pub environment: u64,
156    #[serde(deserialize_with = "codec::object")]
157    pub model: Model,
158}
159
160#[derive(Debug, Deserialize, Serialize)]
161#[serde(deny_unknown_fields)]
162pub struct Model {
163    pub name: String,
164    pub revision: String,
165    pub manifest_sha256: String,
166    pub dimensions: u64,
167    pub max_tokens: u64,
168    pub special_tokens: u64,
169    pub prefix_tokens: u64,
170    pub normalised: bool,
171    pub distance: Distance,
172}
173
174#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
175#[serde(rename_all = "snake_case")]
176pub enum Distance {
177    Cosine,
178    Dot,
179    Euclidean,
180}
181
182#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
183#[serde(rename_all = "snake_case")]
184pub enum Kind {
185    Document,
186    Query,
187}
188
189#[derive(Debug, Deserialize, Serialize)]
190#[serde(deny_unknown_fields)]
191pub struct Embed {
192    pub kind: Kind,
193    #[serde(deserialize_with = "codec::objects")]
194    pub items: Vec<Input>,
195}
196
197#[derive(Debug, Deserialize, Serialize)]
198#[serde(deny_unknown_fields)]
199pub struct Input {
200    pub id: String,
201    pub text: String,
202}
203
204#[derive(Debug, Deserialize, Serialize)]
205#[serde(deny_unknown_fields)]
206pub struct Vectors {
207    #[serde(deserialize_with = "codec::objects")]
208    pub items: Vec<Outcome>,
209}
210
211#[derive(Debug, Deserialize, Serialize)]
212#[serde(tag = "outcome", rename_all = "snake_case", deny_unknown_fields)]
213pub enum Outcome {
214    Vector {
215        id: String,
216        vector: String,
217        tokens_read: u64,
218        tokens_total: u64,
219    },
220    /// EMB-E2: a text longer than the model reads, embedded whole as
221    /// overlapping windows in text order.
222    Windows {
223        id: String,
224        #[serde(deserialize_with = "codec::objects")]
225        windows: Vec<Window>,
226        /// The content tokens of the whole text.
227        tokens_total: u64,
228    },
229    Refused {
230        id: String,
231        code: ItemCode,
232        /// The content tokens of the whole text, without special or prefix tokens; only with
233        /// `text_longer_than_model`.
234        #[serde(default, skip_serializing_if = "Option::is_none")]
235        tokens_total: Option<u64>,
236        /// The model's limit as declared in Ready; only with `text_longer_than_model`.
237        #[serde(default, skip_serializing_if = "Option::is_none")]
238        tokens_limit: Option<u64>,
239    },
240}
241
242impl Outcome {
243    pub fn id(&self) -> &str {
244        match self {
245            Self::Vector { id, .. } | Self::Windows { id, .. } | Self::Refused { id, .. } => id,
246        }
247    }
248}
249
250/// EMB-E2: one window of a text: its bytes `[start, end)`, on char and token
251/// boundaries, and their vector.
252#[derive(Debug, Deserialize, Serialize)]
253#[serde(deny_unknown_fields)]
254pub struct Window {
255    pub start: u64,
256    pub end: u64,
257    pub vector: String,
258    pub tokens_read: u64,
259}
260
261#[derive(Debug, Deserialize, Serialize)]
262#[serde(deny_unknown_fields)]
263pub struct Failed {
264    pub code: FailedCode,
265    pub message: String,
266}
haem_frames/embedding.rs

haem_frames/
embedding.rs
