skippy-server 0.76.1

Embedded Skippy staged runtime server
use std::path::Path;

use mesh_native_serving_plugin_api as native_plugin_api;
use model_artifact::gguf::scan_gguf_tokenizer_inventory;
use skippy_protocol::StageConfig;
use skippy_tokenizer::TokenizerIdentity;

use super::{TOKENIZER_VERSION, TokenizerSource};

pub(super) fn source_gguf_path(config: &StageConfig) -> Option<&Path> {
    [
        config.source_model_path.as_deref(),
        config.model_path.as_deref(),
    ]
    .into_iter()
    .flatten()
    .map(Path::new)
    .find(|path| path.is_file())
}

pub(super) fn inventory_from_stage(
    config: &StageConfig,
    identity: &TokenizerIdentity,
    source: &dyn TokenizerSource,
) -> Option<native_plugin_api::TokenizerInventory> {
    let source_path = source_gguf_path(config)?;
    let vocabulary = scan_gguf_tokenizer_inventory(source_path)?;
    let token_ids = (0..vocabulary.tokens.len())
        .map(|id| i32::try_from(id).ok())
        .collect::<Option<Vec<_>>>()?;
    let token_pieces = source.token_pieces(&token_ids).ok()?;
    if token_pieces.len() != vocabulary.tokens.len() {
        return None;
    }
    let mut tokens = Vec::with_capacity(vocabulary.tokens.len());
    for (id, (token, bytes)) in vocabulary.tokens.into_iter().zip(token_pieces).enumerate() {
        let id = u32::try_from(id).ok()?;
        let piece = if token.is_control {
            native_plugin_api::TokenizerInventoryPiece::Control {
                descriptor: token.raw,
            }
        } else {
            native_plugin_api::TokenizerInventoryPiece::Bytes { bytes }
        };
        tokens.push(native_plugin_api::TokenizerInventoryToken { id, piece });
    }
    Some(native_plugin_api::TokenizerInventory {
        schema_version: native_plugin_api::TOKENIZER_INVENTORY_SCHEMA,
        model_id: identity.model_id.clone(),
        source_model_sha256: identity.source_model_sha256.clone(),
        tokenizer_id: identity.tokenizer_id.clone(),
        tokens,
    })
}

pub(super) fn tokenizer_binding_digest(
    identity: &TokenizerIdentity,
    inventory: &native_plugin_api::TokenizerInventory,
) -> [u8; 32] {
    let mut hasher = blake3::Hasher::new();
    hasher.update(b"mesh-native-tokenizer-binding-v1\0");
    update_string(&mut hasher, &identity.model_id);
    update_string(&mut hasher, &identity.source_model_sha256);
    update_string(&mut hasher, &identity.tokenizer_id);
    update_string(&mut hasher, TOKENIZER_VERSION);
    hasher.update(&inventory.schema_version.to_le_bytes());
    hasher.update(&(inventory.tokens.len() as u64).to_le_bytes());
    for token in &inventory.tokens {
        hasher.update(&token.id.to_le_bytes());
        match &token.piece {
            native_plugin_api::TokenizerInventoryPiece::Bytes { bytes } => {
                hasher.update(&[0]);
                hasher.update(&(bytes.len() as u64).to_le_bytes());
                hasher.update(bytes);
            }
            native_plugin_api::TokenizerInventoryPiece::Control { descriptor } => {
                hasher.update(&[1]);
                hasher.update(&(descriptor.len() as u64).to_le_bytes());
                hasher.update(descriptor);
            }
        }
    }
    *hasher.finalize().as_bytes()
}

fn update_string(hasher: &mut blake3::Hasher, value: &str) {
    hasher.update(&(value.len() as u64).to_le_bytes());
    hasher.update(value.as_bytes());
}