lodestone-proto 0.0.16

//! Native dense-vector mapping projection used by compatibility leaves.

use std::error::Error;
use std::fmt::{self, Display, Formatter};

use serde::{Deserialize, Serialize};

use lodestone_types::{
    DenseVectorDefinition, DenseVectorElementType, MAX_VECTOR_DIMENSIONS, VectorMetric, VectorQuantizationScale,
};

pub const NATIVE_VECTOR_MAPPING_FORMAT_VERSION: u16 = 2;

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

/// Native mapping projection. Physical strategy remains server-owned.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct NativeVectorMappingV2 {
    pub format_version: u16,
    #[serde(rename = "type")]
    pub mapping_type: NativeVectorMappingTypeV2,
    pub dimensions: u16,
    pub metric: NativeVectorMetricV2,
    pub element_type: DenseVectorElementType,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub quantization_scale: Option<VectorQuantizationScale>,
}

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

impl NativeVectorMappingV2 {
    pub fn validate(&self) -> Result<(), NativeVectorMappingError> {
        if self.format_version != NATIVE_VECTOR_MAPPING_FORMAT_VERSION {
            return Err(NativeVectorMappingError::UnsupportedVersion(self.format_version));
        }
        if self.dimensions == 0 || self.dimensions > MAX_VECTOR_DIMENSIONS {
            return Err(NativeVectorMappingError::InvalidDimensions(self.dimensions));
        }
        DenseVectorDefinition {
            dimensions: self.dimensions,
            element_type: self.element_type,
            metric: match self.metric {
                NativeVectorMetricV2::Cosine => VectorMetric::Cosine,
                NativeVectorMetricV2::DotProduct => VectorMetric::DotProduct,
                NativeVectorMetricV2::SquaredL2 => VectorMetric::SquaredL2,
            },
            quantization_scale: self.quantization_scale,
        }
        .validate()
        .map_err(|_| NativeVectorMappingError::InvalidElementContract)?;
        Ok(())
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum NativeVectorMappingError {
    UnsupportedVersion(u16),
    InvalidDimensions(u16),
    InvalidElementContract,
}

impl Display for NativeVectorMappingError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::UnsupportedVersion(version) => {
                write!(formatter, "unsupported native vector mapping version {version}")
            }
            Self::InvalidDimensions(dimensions) => write!(
                formatter,
                "native vector mapping dimensions {dimensions} are outside 1..={MAX_VECTOR_DIMENSIONS}"
            ),
            Self::InvalidElementContract => formatter.write_str("native vector mapping element contract is invalid"),
        }
    }
}

impl Error for NativeVectorMappingError {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn mapping_refuses_unknown_version_and_zero_dimensions() {
        let mut mapping = NativeVectorMappingV2 {
            format_version: NATIVE_VECTOR_MAPPING_FORMAT_VERSION,
            mapping_type: NativeVectorMappingTypeV2::Vector,
            dimensions: 2,
            metric: NativeVectorMetricV2::Cosine,
            element_type: DenseVectorElementType::F32,
            quantization_scale: None,
        };
        assert_eq!(mapping.validate(), Ok(()));
        mapping.format_version += 1;
        assert!(matches!(
            mapping.validate(),
            Err(NativeVectorMappingError::UnsupportedVersion(_))
        ));
        mapping.format_version = NATIVE_VECTOR_MAPPING_FORMAT_VERSION;
        mapping.dimensions = 0;
        assert_eq!(mapping.validate(), Err(NativeVectorMappingError::InvalidDimensions(0)));
    }
}