zeppelin_embed/manifest/
mod.rs

1//! Atomic manifest commit point for every visible persisted artifact.
2
3pub mod io;
4
5#[cfg(feature = "graph-cypher")]
6mod graph;
7#[cfg(feature = "graph-cypher")]
8pub use graph::{GraphManifest, GraphObject};
9
10use std::collections::BTreeSet;
11use std::path::PathBuf;
12
13use crate::epoch::{
14    ComputeUnits, EmbeddingEpoch, EmbeddingRuntime, EmbeddingTower, EpochId, EpochIdentity,
15    Normalization,
16};
17use crate::format::FormatFamily;
18use crate::format::frame::{FormatError, decode_artifact, encode_artifact_version};
19use crate::fts::tokenizer::TokenizerEpoch;
20use crate::meta::{ColumnDefinition, ColumnId, ColumnType, Schema};
21use crate::segment::{ClusteringKeyRange, SegmentId, SegmentMeta};
22
23const CLUSTERING_RANGE_EXTENSION_MAGIC: [u8; 4] = *b"TSR1";
24const CLUSTERING_RANGE_RECORD_LEN: usize = 24;
25
26#[cfg(feature = "graph-cypher")]
27impl Manifest {
28    // Every manifest-only increment is recorded at the current durable WAL
29    // boundary in the same commit. Watermarks fold older offsets into the base;
30    // offsets at the watermark still precede the next unabsorbed batch.
31    pub(crate) fn record_generation_bump(&mut self, boundary: u64) -> Result<(), ManifestError> {
32        if let Some(graph) = &mut self.graph {
33            let watermark = self.log_seq.max(graph.graph_absorbed_through);
34            graph
35                .generation_bumps
36                .retain(|(sequence, _)| *sequence >= watermark);
37            graph.generation_absorbed_through = Some(boundary);
38            if let Some((sequence, count)) = graph.generation_bumps.last_mut()
39                && *sequence == boundary
40            {
41                *count = count
42                    .checked_add(1)
43                    .ok_or_else(|| ManifestError::Decode("generation bump overflow".to_owned()))?;
44            } else {
45                graph.generation_bumps.push((boundary, 1));
46            }
47        }
48        Ok(())
49    }
50}
51
52/// Interpretation-critical model/tokenizer identity for one coexistence epoch.
53#[derive(Clone, Debug, Eq, PartialEq)]
54pub struct EpochMeta {
55    /// Stable epoch identifier.
56    pub id: EpochId,
57    /// Complete authoritative embedding-model identity.
58    pub embedding: EmbeddingEpoch,
59    /// Authoritative tokenizer identity.
60    pub tokenizer: TokenizerEpoch,
61}
62
63/// The single atomic snapshot commit point.
64#[derive(Clone, Debug, Eq, PartialEq)]
65pub struct Manifest {
66    /// Optional graph checkpoint and live immutable object inventory.
67    #[cfg(feature = "graph-cypher")]
68    pub graph: Option<GraphManifest>,
69    /// Monotonic committed generation.
70    pub generation: u64,
71    /// Durable WAL sequence covered by this snapshot.
72    pub log_seq: u64,
73    /// Complete reachable immutable segment set.
74    pub segments: Vec<SegmentMeta>,
75    /// Complete interpretation epoch set.
76    pub epochs: Vec<EpochMeta>,
77    /// Atomically published embedding/tokenizer identity used by queries and writes.
78    pub epoch_alias: Option<EpochIdentity>,
79    /// Typed metadata schema interpreted by every listed segment.
80    pub schema: Schema,
81}
82
83/// Typed manifest encode/load/open failure.
84#[derive(Debug)]
85pub enum ManifestError {
86    /// Filesystem operation failed for a named path.
87    Io {
88        /// Path involved in the operation.
89        path: PathBuf,
90        /// Underlying platform error.
91        source: std::io::Error,
92    },
93    /// Framing, version, length, or checksum validation failed.
94    Format(FormatError),
95    /// Hand-written payload decoding rejected malformed bytes.
96    Decode(String),
97    /// This build cannot open a v3 manifest containing graph state.
98    GraphUnsupportedBuild,
99    /// Snapshot log sequence exceeds the durable WAL end.
100    AheadOfLog {
101        /// Sequence covered by the snapshot.
102        snapshot: u64,
103        /// Sequence known durable in the WAL.
104        durable: u64,
105    },
106    /// A referenced segment header failed validation.
107    Segment(crate::segment::SegmentError),
108    /// The published alias did not name a registry entry.
109    UnknownEpochAlias {
110        /// Rejected alias.
111        alias: EpochIdentity,
112    },
113    /// A segment epoch tag did not name any registry entry.
114    UnknownSegmentEpoch {
115        /// Segment carrying the bad tag.
116        segment: SegmentId,
117        /// Rejected embedding epoch id.
118        epoch: EpochId,
119    },
120}
121
122impl ManifestError {
123    pub(crate) fn io(path: impl Into<PathBuf>, source: std::io::Error) -> Self {
124        Self::Io {
125            path: path.into(),
126            source,
127        }
128    }
129}
130
131impl std::fmt::Display for ManifestError {
132    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133        match self {
134            Self::Io { path, source } => {
135                write!(formatter, "manifest I/O {}: {source}", path.display())
136            }
137            Self::Format(error) => error.fmt(formatter),
138            Self::GraphUnsupportedBuild => {
139                write!(formatter, "graph manifest requires graph-cypher")
140            }
141            Self::Decode(detail) => write!(formatter, "manifest decode failed: {detail}"),
142            Self::AheadOfLog { snapshot, durable } => write!(
143                formatter,
144                "manifest log sequence {snapshot} is ahead of durable WAL end {durable}"
145            ),
146            Self::Segment(error) => {
147                write!(formatter, "manifest segment validation failed: {error}")
148            }
149            Self::UnknownEpochAlias { alias } => write!(
150                formatter,
151                "manifest epoch alias ({}, {}) does not name a registry entry",
152                alias.embedding, alias.tokenizer
153            ),
154            Self::UnknownSegmentEpoch { segment, epoch } => write!(
155                formatter,
156                "manifest segment {segment} names unknown embedding epoch {epoch}"
157            ),
158        }
159    }
160}
161
162impl std::error::Error for ManifestError {
163    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
164        match self {
165            Self::Io { source, .. } => Some(source),
166            Self::Format(error) => Some(error),
167            Self::Segment(error) => Some(error),
168            Self::GraphUnsupportedBuild
169            | Self::Decode(_)
170            | Self::AheadOfLog { .. }
171            | Self::UnknownEpochAlias { .. }
172            | Self::UnknownSegmentEpoch { .. } => None,
173        }
174    }
175}
176
177impl From<FormatError> for ManifestError {
178    fn from(error: FormatError) -> Self {
179        Self::Format(error)
180    }
181}
182
183impl From<crate::segment::SegmentError> for ManifestError {
184    fn from(error: crate::segment::SegmentError) -> Self {
185        Self::Segment(error)
186    }
187}
188
189/// Encodes a complete manifest through a single checksummed frame.
190pub fn encode_manifest(manifest: &Manifest) -> Result<Vec<u8>, ManifestError> {
191    validate_epoch_references(manifest)?;
192    let mut payload = Vec::new();
193    payload.extend_from_slice(&manifest.generation.to_le_bytes());
194    payload.extend_from_slice(&manifest.log_seq.to_le_bytes());
195    append_u32_len(manifest.segments.len(), "segments", &mut payload)?;
196    append_u32_len(manifest.epochs.len(), "epochs", &mut payload)?;
197    append_u32_len(
198        manifest.schema.user_column_count(),
199        "schema columns",
200        &mut payload,
201    )?;
202    payload.extend_from_slice(&0_u32.to_le_bytes());
203    append_optional_identity(manifest.epoch_alias, &mut payload);
204    for segment in &manifest.segments {
205        payload.extend_from_slice(segment.id.as_bytes());
206        payload.extend_from_slice(&segment.row_count.to_le_bytes());
207        payload.extend_from_slice(&segment.scheme.to_le_bytes());
208        payload.extend_from_slice(&0_u16.to_le_bytes());
209        payload.extend_from_slice(&segment.dims.to_le_bytes());
210        payload.extend_from_slice(&segment.file_size.to_le_bytes());
211        append_optional_epoch_id(segment.epoch_id, &mut payload);
212    }
213    for epoch in &manifest.epochs {
214        payload.extend_from_slice(&epoch.id.value().to_le_bytes());
215        payload.extend_from_slice(&epoch.tokenizer.value().to_le_bytes());
216        append_embedding_epoch(&epoch.embedding, &mut payload)?;
217    }
218    for definition in manifest.schema.columns().iter().skip(1) {
219        payload.extend_from_slice(&definition.id().get().to_le_bytes());
220        payload.extend_from_slice(&column_type_id(definition.column_type()).to_le_bytes());
221        payload.extend_from_slice(&u16::from(definition.is_nullable()).to_le_bytes());
222        append_string(definition.name(), &mut payload)?;
223    }
224    append_clustering_ranges(&manifest.segments, &mut payload)?;
225    #[cfg(feature = "graph-cypher")]
226    let version = if let Some(graph) = &manifest.graph {
227        graph::append(graph, &mut payload)?;
228        3
229    } else {
230        2
231    };
232    #[cfg(not(feature = "graph-cypher"))]
233    let version = 2;
234    Ok(encode_artifact_version(
235        FormatFamily::Manifest,
236        version,
237        0,
238        &payload,
239    ))
240}
241
242fn append_clustering_ranges(
243    segments: &[SegmentMeta],
244    output: &mut Vec<u8>,
245) -> Result<(), ManifestError> {
246    if segments
247        .iter()
248        .all(|segment| segment.clustering_key_range == ClusteringKeyRange::Unstamped)
249    {
250        return Ok(());
251    }
252    output.extend_from_slice(&CLUSTERING_RANGE_EXTENSION_MAGIC);
253    append_u32_len(segments.len(), "clustering ranges", output)?;
254    for segment in segments {
255        let (tag, min_ts, max_ts) = match segment.clustering_key_range {
256            ClusteringKeyRange::Unstamped => (0_u8, 0_i64, 0_i64),
257            ClusteringKeyRange::Empty => (1_u8, 0_i64, 0_i64),
258            ClusteringKeyRange::Bounded { min_ts, max_ts } => {
259                if min_ts > max_ts {
260                    return Err(ManifestError::Decode(format!(
261                        "segment {} clustering range {min_ts}..={max_ts} is inverted",
262                        segment.id
263                    )));
264                }
265                (2_u8, min_ts, max_ts)
266            }
267        };
268        output.push(tag);
269        output.extend_from_slice(&[0_u8; 7]);
270        output.extend_from_slice(&min_ts.to_le_bytes());
271        output.extend_from_slice(&max_ts.to_le_bytes());
272    }
273    Ok(())
274}
275
276/// Decodes a complete manifest only after both xxh3-64 checksums validate.
277pub fn decode_manifest(artifact: &str, bytes: &[u8]) -> Result<Manifest, ManifestError> {
278    #[cfg(any(test, feature = "test-seams"))]
279    let actual_family = bytes
280        .get(8..10)
281        .and_then(|family| <[u8; 2]>::try_from(family).ok())
282        .map(u16::from_le_bytes);
283    let framed = decode_artifact(artifact, FormatFamily::Manifest, bytes).map_err(|error| {
284        #[cfg(any(test, feature = "test-seams"))]
285        crate::lifecycle::record_storage_manifest_format_fault(&error, actual_family);
286        ManifestError::Format(error)
287    })?;
288    if framed.header.flags != 0 {
289        return Err(ManifestError::Decode(
290            "manifest flags are non-zero".to_owned(),
291        ));
292    }
293    #[cfg(not(feature = "graph-cypher"))]
294    if framed.header.version == 3 {
295        return Err(ManifestError::GraphUnsupportedBuild);
296    }
297    let mut cursor = ManifestCursor::new(framed.payload);
298    let generation = cursor.u64()?;
299    let log_seq = cursor.u64()?;
300    let segment_count = cursor.usize_from_u32()?;
301    let epoch_count = cursor.usize_from_u32()?;
302    let schema_count = cursor.usize_from_u32()?;
303    if cursor.u32()? != 0 {
304        return Err(ManifestError::Decode(
305            "manifest reserved field is non-zero".to_owned(),
306        ));
307    }
308    let epoch_alias = cursor.optional_identity()?;
309    let mut segments = Vec::with_capacity(segment_count);
310    for _ in 0..segment_count {
311        let id_bytes: [u8; 16] = cursor
312            .take(16)?
313            .try_into()
314            .map_err(|_| ManifestError::Decode("invalid segment id".to_owned()))?;
315        let row_count = cursor.u32()?;
316        let scheme = cursor.u16()?;
317        crate::format::FormatRegistry::require_scheme(scheme)
318            .map_err(|error| ManifestError::Decode(error.to_string()))?;
319        if cursor.u16()? != 0 {
320            return Err(ManifestError::Decode(
321                "segment reserved field is non-zero".to_owned(),
322            ));
323        }
324        segments.push(SegmentMeta {
325            id: SegmentId::from_bytes(id_bytes),
326            row_count,
327            scheme,
328            dims: cursor.u32()?,
329            file_size: cursor.u64()?,
330            epoch_id: cursor.optional_epoch_id()?,
331            clustering_key_range: crate::segment::ClusteringKeyRange::Unstamped,
332        });
333    }
334    let mut epochs = Vec::with_capacity(epoch_count);
335    for _ in 0..epoch_count {
336        epochs.push(EpochMeta {
337            id: EpochId::from_value(cursor.u64()?),
338            tokenizer: TokenizerEpoch::from_value(cursor.u64()?),
339            embedding: cursor.embedding_epoch()?,
340        });
341    }
342    let mut definitions = Vec::with_capacity(schema_count);
343    for _ in 0..schema_count {
344        let id = ColumnId::new(cursor.u32()?);
345        let column_type = column_type_from_id(cursor.u16()?)?;
346        let nullable = match cursor.u16()? {
347            0 => false,
348            1 => true,
349            value => {
350                return Err(ManifestError::Decode(format!(
351                    "invalid nullable flag {value}"
352                )));
353            }
354        };
355        definitions.push(ColumnDefinition::new(
356            id,
357            cursor.string()?,
358            column_type,
359            nullable,
360        ));
361    }
362    if cursor
363        .bytes
364        .get(cursor.position..)
365        .is_some_and(|tail| tail.starts_with(&CLUSTERING_RANGE_EXTENSION_MAGIC))
366    {
367        decode_clustering_ranges(&mut cursor, &mut segments)?;
368    }
369    #[cfg(feature = "graph-cypher")]
370    let graph = if framed.header.version == 3 {
371        Some(graph::decode(&mut cursor)?)
372    } else {
373        None
374    };
375    cursor.finish()?;
376    let schema =
377        Schema::new(definitions).map_err(|error| ManifestError::Decode(error.to_string()))?;
378    let manifest = Manifest {
379        #[cfg(feature = "graph-cypher")]
380        graph,
381        generation,
382        log_seq,
383        segments,
384        epochs,
385        epoch_alias,
386        schema,
387    };
388    validate_epoch_references(&manifest)?;
389    Ok(manifest)
390}
391
392fn validate_epoch_references(manifest: &Manifest) -> Result<(), ManifestError> {
393    let mut identities = BTreeSet::new();
394    for epoch in &manifest.epochs {
395        let identity = EpochIdentity::from_meta(epoch)
396            .map_err(|error| ManifestError::Decode(error.to_string()))?;
397        if !identities.insert(identity) {
398            return Err(ManifestError::Decode(format!(
399                "duplicate epoch registry identity ({}, {})",
400                identity.embedding, identity.tokenizer
401            )));
402        }
403    }
404
405    match manifest.epoch_alias {
406        Some(alias) if !identities.contains(&alias) => {
407            return Err(ManifestError::UnknownEpochAlias { alias });
408        }
409        None if !identities.is_empty() => {
410            return Err(ManifestError::Decode(
411                "non-empty epoch registry has no published alias".to_owned(),
412            ));
413        }
414        Some(_) | None => {}
415    }
416
417    for segment in &manifest.segments {
418        match segment.epoch_id {
419            Some(epoch)
420                if !identities
421                    .iter()
422                    .any(|identity| identity.embedding == epoch) =>
423            {
424                return Err(ManifestError::UnknownSegmentEpoch {
425                    segment: segment.id,
426                    epoch,
427                });
428            }
429            None if !identities.is_empty() => {
430                return Err(ManifestError::Decode(format!(
431                    "segment {} has no embedding epoch tag",
432                    segment.id
433                )));
434            }
435            Some(_) | None => {}
436        }
437    }
438    Ok(())
439}
440
441fn append_optional_identity(identity: Option<EpochIdentity>, output: &mut Vec<u8>) {
442    let (tag, embedding, tokenizer) = match identity {
443        Some(identity) => (1_u8, identity.embedding.value(), identity.tokenizer.value()),
444        None => (0, 0, 0),
445    };
446    output.push(tag);
447    output.extend_from_slice(&[0_u8; 7]);
448    output.extend_from_slice(&embedding.to_le_bytes());
449    output.extend_from_slice(&tokenizer.to_le_bytes());
450}
451
452fn append_optional_epoch_id(epoch: Option<EpochId>, output: &mut Vec<u8>) {
453    let (tag, value) = epoch.map_or((0_u8, 0_u64), |epoch| (1, epoch.value()));
454    output.push(tag);
455    output.extend_from_slice(&[0_u8; 7]);
456    output.extend_from_slice(&value.to_le_bytes());
457}
458
459fn append_embedding_epoch(
460    epoch: &EmbeddingEpoch,
461    output: &mut Vec<u8>,
462) -> Result<(), ManifestError> {
463    append_embedding_tower(&epoch.document, output)?;
464    append_embedding_tower(&epoch.query, output)?;
465    append_bytes(&epoch.alignment_digest, output)
466}
467
468fn append_embedding_tower(
469    tower: &EmbeddingTower,
470    output: &mut Vec<u8>,
471) -> Result<(), ManifestError> {
472    append_string(&tower.model_id, output)?;
473    append_string(&tower.model_version, output)?;
474    append_bytes(&tower.weights_digest, output)?;
475    output.extend_from_slice(&tower.dims.to_le_bytes());
476    output.extend_from_slice(&(tower.normalization as u16).to_le_bytes());
477    append_string(&tower.prompt_prefix, output)?;
478    output.extend_from_slice(&tower.max_tokens.to_le_bytes());
479    output.extend_from_slice(&(tower.runtime as u16).to_le_bytes());
480    output.extend_from_slice(&(tower.compute_units as u16).to_le_bytes());
481    match &tower.os_build {
482        Some(build) => {
483            output.push(1);
484            output.extend_from_slice(&[0_u8; 3]);
485            append_string(build, output)?;
486        }
487        None => {
488            output.push(0);
489            output.extend_from_slice(&[0_u8; 3]);
490        }
491    }
492    Ok(())
493}
494
495fn decode_clustering_ranges(
496    cursor: &mut ManifestCursor<'_>,
497    segments: &mut [SegmentMeta],
498) -> Result<(), ManifestError> {
499    if cursor.take(CLUSTERING_RANGE_EXTENSION_MAGIC.len())? != CLUSTERING_RANGE_EXTENSION_MAGIC {
500        return Err(ManifestError::Decode(
501            "unknown manifest extension after schema".to_owned(),
502        ));
503    }
504    let count = cursor.usize_from_u32()?;
505    if count != segments.len() {
506        return Err(ManifestError::Decode(format!(
507            "clustering range count {count} does not match segment count {}",
508            segments.len()
509        )));
510    }
511    let expected_bytes = count
512        .checked_mul(CLUSTERING_RANGE_RECORD_LEN)
513        .ok_or_else(|| ManifestError::Decode("clustering range bytes overflow".to_owned()))?;
514    if cursor.remaining() < expected_bytes {
515        return Err(ManifestError::Decode(format!(
516            "clustering range extension has {} bytes, expected {expected_bytes}",
517            cursor.remaining()
518        )));
519    }
520    for segment in segments {
521        let tag = cursor.u8()?;
522        if cursor.take(7)?.iter().any(|byte| *byte != 0) {
523            return Err(ManifestError::Decode(
524                "clustering range reserved bytes are non-zero".to_owned(),
525            ));
526        }
527        let min_ts = cursor.i64()?;
528        let max_ts = cursor.i64()?;
529        segment.clustering_key_range = match tag {
530            0 if min_ts == 0 && max_ts == 0 => ClusteringKeyRange::Unstamped,
531            1 if min_ts == 0 && max_ts == 0 => ClusteringKeyRange::Empty,
532            2 if min_ts <= max_ts => ClusteringKeyRange::Bounded { min_ts, max_ts },
533            0 | 1 => {
534                return Err(ManifestError::Decode(format!(
535                    "clustering range tag {tag} requires zero bounds"
536                )));
537            }
538            2 => {
539                return Err(ManifestError::Decode(format!(
540                    "clustering range {min_ts}..={max_ts} is inverted"
541                )));
542            }
543            unknown => {
544                return Err(ManifestError::Decode(format!(
545                    "unknown clustering range tag {unknown}"
546                )));
547            }
548        };
549    }
550    Ok(())
551}
552
553fn append_u32_len(length: usize, name: &str, output: &mut Vec<u8>) -> Result<(), ManifestError> {
554    let value = u32::try_from(length)
555        .map_err(|_| ManifestError::Decode(format!("{name} length exceeds u32")))?;
556    output.extend_from_slice(&value.to_le_bytes());
557    Ok(())
558}
559
560fn append_string(value: &str, output: &mut Vec<u8>) -> Result<(), ManifestError> {
561    append_bytes(value.as_bytes(), output)
562}
563
564fn append_bytes(value: &[u8], output: &mut Vec<u8>) -> Result<(), ManifestError> {
565    append_u32_len(value.len(), "byte field", output)?;
566    output.extend_from_slice(value);
567    Ok(())
568}
569
570fn column_type_id(column_type: ColumnType) -> u16 {
571    match column_type {
572        ColumnType::Id128 => 7,
573        ColumnType::U64 => 1,
574        ColumnType::I64 => 2,
575        ColumnType::F64 => 3,
576        ColumnType::Bool => 4,
577        ColumnType::DictionaryString => 5,
578        ColumnType::RawString => 6,
579    }
580}
581
582fn column_type_from_id(id: u16) -> Result<ColumnType, ManifestError> {
583    match id {
584        7 => Ok(ColumnType::Id128),
585        1 => Ok(ColumnType::U64),
586        2 => Ok(ColumnType::I64),
587        3 => Ok(ColumnType::F64),
588        4 => Ok(ColumnType::Bool),
589        5 => Ok(ColumnType::DictionaryString),
590        6 => Ok(ColumnType::RawString),
591        _ => Err(ManifestError::Decode(format!("unknown column type {id}"))),
592    }
593}
594
595struct ManifestCursor<'a> {
596    bytes: &'a [u8],
597    position: usize,
598}
599
600impl<'a> ManifestCursor<'a> {
601    const fn new(bytes: &'a [u8]) -> Self {
602        Self { bytes, position: 0 }
603    }
604
605    fn take(&mut self, length: usize) -> Result<&'a [u8], ManifestError> {
606        let end = self
607            .position
608            .checked_add(length)
609            .ok_or_else(|| ManifestError::Decode("payload offset overflow".to_owned()))?;
610        let value = self.bytes.get(self.position..end).ok_or_else(|| {
611            ManifestError::Decode(format!(
612                "payload truncated at {}, need {length}, total {}",
613                self.position,
614                self.bytes.len()
615            ))
616        })?;
617        self.position = end;
618        Ok(value)
619    }
620
621    fn u16(&mut self) -> Result<u16, ManifestError> {
622        let raw: [u8; 2] = self
623            .take(2)?
624            .try_into()
625            .map_err(|_| ManifestError::Decode("invalid u16".to_owned()))?;
626        Ok(u16::from_le_bytes(raw))
627    }
628
629    fn u8(&mut self) -> Result<u8, ManifestError> {
630        self.take(1)?
631            .first()
632            .copied()
633            .ok_or_else(|| ManifestError::Decode("invalid u8".to_owned()))
634    }
635
636    fn u32(&mut self) -> Result<u32, ManifestError> {
637        let raw: [u8; 4] = self
638            .take(4)?
639            .try_into()
640            .map_err(|_| ManifestError::Decode("invalid u32".to_owned()))?;
641        Ok(u32::from_le_bytes(raw))
642    }
643
644    fn u64(&mut self) -> Result<u64, ManifestError> {
645        let raw: [u8; 8] = self
646            .take(8)?
647            .try_into()
648            .map_err(|_| ManifestError::Decode("invalid u64".to_owned()))?;
649        Ok(u64::from_le_bytes(raw))
650    }
651
652    fn optional_identity(&mut self) -> Result<Option<EpochIdentity>, ManifestError> {
653        let tag = self.u8()?;
654        if self.take(7)?.iter().any(|byte| *byte != 0) {
655            return Err(ManifestError::Decode(
656                "epoch alias reserved bytes are non-zero".to_owned(),
657            ));
658        }
659        let embedding = self.u64()?;
660        let tokenizer = self.u64()?;
661        match (tag, embedding, tokenizer) {
662            (0, 0, 0) => Ok(None),
663            (1, embedding, tokenizer) => Ok(Some(EpochIdentity {
664                embedding: EpochId::from_value(embedding),
665                tokenizer: TokenizerEpoch::from_value(tokenizer),
666            })),
667            (0, _, _) => Err(ManifestError::Decode(
668                "absent epoch alias carries non-zero identity bytes".to_owned(),
669            )),
670            (unknown, _, _) => Err(ManifestError::Decode(format!(
671                "unknown epoch alias presence tag {unknown}"
672            ))),
673        }
674    }
675
676    fn optional_epoch_id(&mut self) -> Result<Option<EpochId>, ManifestError> {
677        let tag = self.u8()?;
678        if self.take(7)?.iter().any(|byte| *byte != 0) {
679            return Err(ManifestError::Decode(
680                "segment epoch reserved bytes are non-zero".to_owned(),
681            ));
682        }
683        let epoch = self.u64()?;
684        match (tag, epoch) {
685            (0, 0) => Ok(None),
686            (1, epoch) => Ok(Some(EpochId::from_value(epoch))),
687            (0, _) => Err(ManifestError::Decode(
688                "absent segment epoch tag carries a non-zero id".to_owned(),
689            )),
690            (unknown, _) => Err(ManifestError::Decode(format!(
691                "unknown segment epoch presence tag {unknown}"
692            ))),
693        }
694    }
695
696    fn embedding_epoch(&mut self) -> Result<EmbeddingEpoch, ManifestError> {
697        Ok(EmbeddingEpoch {
698            document: self.embedding_tower()?,
699            query: self.embedding_tower()?,
700            alignment_digest: self.bytes()?,
701        })
702    }
703
704    fn embedding_tower(&mut self) -> Result<EmbeddingTower, ManifestError> {
705        let model_id = self.string()?;
706        let model_version = self.string()?;
707        let weights_digest = self.bytes()?;
708        let dims = self.u32()?;
709        let normalization = match self.u16()? {
710            0 => Normalization::None,
711            1 => Normalization::L2,
712            unknown => {
713                return Err(ManifestError::Decode(format!(
714                    "unknown embedding normalization {unknown}"
715                )));
716            }
717        };
718        let prompt_prefix = self.string()?;
719        let max_tokens = self.u32()?;
720        let runtime = match self.u16()? {
721            1 => EmbeddingRuntime::CoreMl,
722            2 => EmbeddingRuntime::Mlx,
723            3 => EmbeddingRuntime::CpuReference,
724            unknown => {
725                return Err(ManifestError::Decode(format!(
726                    "unknown embedding runtime {unknown}"
727                )));
728            }
729        };
730        let compute_units = match self.u16()? {
731            1 => ComputeUnits::Cpu,
732            2 => ComputeUnits::CpuAndGpu,
733            3 => ComputeUnits::CpuAndNeuralEngine,
734            4 => ComputeUnits::All,
735            unknown => {
736                return Err(ManifestError::Decode(format!(
737                    "unknown embedding compute units {unknown}"
738                )));
739            }
740        };
741        let os_tag = self.u8()?;
742        if self.take(3)?.iter().any(|byte| *byte != 0) {
743            return Err(ManifestError::Decode(
744                "embedding OS-build reserved bytes are non-zero".to_owned(),
745            ));
746        }
747        let os_build = match os_tag {
748            0 => None,
749            1 => Some(self.string()?),
750            unknown => {
751                return Err(ManifestError::Decode(format!(
752                    "unknown embedding OS-build presence tag {unknown}"
753                )));
754            }
755        };
756        Ok(EmbeddingTower {
757            model_id,
758            model_version,
759            weights_digest,
760            dims,
761            normalization,
762            prompt_prefix,
763            max_tokens,
764            runtime,
765            compute_units,
766            os_build,
767        })
768    }
769
770    fn i64(&mut self) -> Result<i64, ManifestError> {
771        let raw: [u8; 8] = self
772            .take(8)?
773            .try_into()
774            .map_err(|_| ManifestError::Decode("invalid i64".to_owned()))?;
775        Ok(i64::from_le_bytes(raw))
776    }
777
778    const fn remaining(&self) -> usize {
779        self.bytes.len().saturating_sub(self.position)
780    }
781
782    fn usize_from_u32(&mut self) -> Result<usize, ManifestError> {
783        usize::try_from(self.u32()?)
784            .map_err(|_| ManifestError::Decode("u32 length exceeds usize".to_owned()))
785    }
786
787    fn string(&mut self) -> Result<String, ManifestError> {
788        let length = self.usize_from_u32()?;
789        std::str::from_utf8(self.take(length)?)
790            .map(str::to_owned)
791            .map_err(|error| ManifestError::Decode(format!("invalid UTF-8: {error}")))
792    }
793
794    fn bytes(&mut self) -> Result<Vec<u8>, ManifestError> {
795        let length = self.usize_from_u32()?;
796        Ok(self.take(length)?.to_vec())
797    }
798
799    fn finish(&self) -> Result<(), ManifestError> {
800        if self.position == self.bytes.len() {
801            Ok(())
802        } else {
803            Err(ManifestError::Decode(format!(
804                "{} trailing payload bytes",
805                self.bytes.len().saturating_sub(self.position)
806            )))
807        }
808    }
809}
zeppelin_embed/manifest/mod.rs

zeppelin_embed/manifest/
mod.rs
