haem-frames 0.1.1

//! One vocabulary and one set of checks for both sides of an embedding exchange.

use serde::{Deserialize, Serialize};

mod checks;
mod codec;
mod vector;

pub use checks::{
    check_embed, check_model_limits, check_model_limits_windowed, check_ready, check_start,
    check_vectors, check_vectors_windowed, input_refusal, vectors_message_size,
    windows_message_size,
};
pub use codec::{Message, decode_message, read_message, read_start, write_message};
pub use vector::{decode_vector, encode_vector};

pub const PROTOCOL: u32 = 1;
/// EMB-E2: the protocol a Start names when it asks for windows. A worker
/// answers Ready with the protocol its Start named.
pub const PROTOCOL_WINDOWS: u32 = 2;
pub const MINIMUM_FRAME_BYTES: u32 = 8_192;
pub const FAILED_MESSAGE_CHARS: usize = 1_024;
pub const MEMORY_LIMIT_EXIT: i32 = 86;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u8)]
pub enum Tag {
    Start = 16,
    Ready = 17,
    Embed = 18,
    Vectors = 19,
    Failed = 20,
}

impl TryFrom<u8> for Tag {
    type Error = Error;
    fn try_from(value: u8) -> Result<Self, Error> {
        match value {
            16 => Ok(Self::Start),
            17 => Ok(Self::Ready),
            18 => Ok(Self::Embed),
            19 => Ok(Self::Vectors),
            20 => Ok(Self::Failed),
            _ => Err(Error::new(
                FailedCode::EmbedProtocol,
                "unknown embedding frame tag",
            )),
        }
    }
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum FailedCode {
    EmbedModelMismatch,
    EmbedModelMissing,
    EmbedLimits,
    EmbedBatchTooLarge,
    EmbedOutputInvalid,
    EmbedProtocol,
    EmbedInferenceFailed,
}

impl FailedCode {
    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::EmbedModelMismatch => "embed_model_mismatch",
            Self::EmbedModelMissing => "embed_model_missing",
            Self::EmbedLimits => "embed_limits",
            Self::EmbedBatchTooLarge => "embed_batch_too_large",
            Self::EmbedOutputInvalid => "embed_output_invalid",
            Self::EmbedProtocol => "embed_protocol",
            Self::EmbedInferenceFailed => "embed_inference_failed",
        }
    }
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ItemCode {
    EmbedInputEmpty,
    EmbedInputTooLarge,
    TextLongerThanModel,
    /// EMB-E2: a text whose whole cover needs more windows than
    /// `Windows::max_windows`. It gets no vector; it is never embedded as its
    /// first windows.
    WindowsOverCap,
}

#[derive(Debug)]
pub struct Error {
    pub code: FailedCode,
    pub message: String,
}

impl Error {
    pub(super) fn new(code: FailedCode, message: impl Into<String>) -> Self {
        Self {
            code,
            message: message.into(),
        }
    }
}

impl std::fmt::Display for Error {
    fn fmt(&self, output: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(output, "{}: {}", self.code.as_str(), self.message)
    }
}

impl std::error::Error for Error {}

#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Limits {
    pub memory_bytes: u64,
    pub threads: u64,
    pub batch_items: u64,
    pub input_bytes: u64,
    pub tokens: u64,
    pub frame_bytes: u32,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Start {
    pub protocol: u32,
    pub model_dir: String,
    #[serde(deserialize_with = "codec::object")]
    pub limits: Limits,
    /// EMB-E2: present exactly when `protocol` is [`PROTOCOL_WINDOWS`].
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub windows: Option<Windows>,
}

/// EMB-E2: how a text longer than the model reads is embedded whole.
///
/// Each window holds at most the model's content tokens; consecutive windows
/// share at least `overlap_tokens`; a text needing more than `max_windows` is
/// refused `windows_over_cap`. No window is ever truncated.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Windows {
    pub overlap_tokens: u64,
    pub max_windows: u64,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Ready {
    pub protocol: u32,
    pub worker: String,
    pub core_limit: u64,
    pub descriptors_closed: u64,
    pub environment: u64,
    #[serde(deserialize_with = "codec::object")]
    pub model: Model,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Model {
    pub name: String,
    pub revision: String,
    pub manifest_sha256: String,
    pub dimensions: u64,
    pub max_tokens: u64,
    pub special_tokens: u64,
    pub prefix_tokens: u64,
    pub normalised: bool,
    pub distance: Distance,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum Distance {
    Cosine,
    Dot,
    Euclidean,
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum Kind {
    Document,
    Query,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Embed {
    pub kind: Kind,
    #[serde(deserialize_with = "codec::objects")]
    pub items: Vec<Input>,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Input {
    pub id: String,
    pub text: String,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Vectors {
    #[serde(deserialize_with = "codec::objects")]
    pub items: Vec<Outcome>,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(tag = "outcome", rename_all = "snake_case", deny_unknown_fields)]
pub enum Outcome {
    Vector {
        id: String,
        vector: String,
        tokens_read: u64,
        tokens_total: u64,
    },
    /// EMB-E2: a text longer than the model reads, embedded whole as
    /// overlapping windows in text order.
    Windows {
        id: String,
        #[serde(deserialize_with = "codec::objects")]
        windows: Vec<Window>,
        /// The content tokens of the whole text.
        tokens_total: u64,
    },
    Refused {
        id: String,
        code: ItemCode,
        /// The content tokens of the whole text, without special or prefix tokens; only with
        /// `text_longer_than_model`.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        tokens_total: Option<u64>,
        /// The model's limit as declared in Ready; only with `text_longer_than_model`.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        tokens_limit: Option<u64>,
    },
}

impl Outcome {
    #[must_use]
    pub fn id(&self) -> &str {
        match self {
            Self::Vector { id, .. } | Self::Windows { id, .. } | Self::Refused { id, .. } => id,
        }
    }
}

/// EMB-E2: one window of a text: its bytes `[start, end)`, on char and token
/// boundaries, and their vector.
#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Window {
    pub start: u64,
    pub end: u64,
    pub vector: String,
    pub tokens_read: u64,
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Failed {
    pub code: FailedCode,
    pub message: String,
}