haematite-cli 0.21.7

//! Vector fields and exact search at immutable captured roots.

use std::cmp::Ordering;
use std::num::NonZeroUsize;

use haematite::db::ReadOnlyDatabase;
use serde::Serialize;

use super::embedded::{self, Field, ModelName, VectorKind};
use super::{Captured, Failure, document_key, identifier, mutations::PreparedWrite, reads};
use crate::serve::access::Access;
use crate::serve::protocol::{Metric, VectorDeclaration, VectorDelete, VectorSearch, VectorWrite};

#[derive(Serialize)]
pub(in crate::serve) struct Vector {
    document_id: String,
    vector: Vec<f32>,
}

#[derive(Serialize)]
pub(in crate::serve) struct Hit {
    pub(in crate::serve) document_id: String,
    pub(in crate::serve) score: f64,
}

#[derive(Serialize)]
pub(in crate::serve) struct VectorField {
    collection: String,
    field: String,
    dimensions: usize,
    count: u64,
}

#[derive(Serialize)]
pub(in crate::serve) struct EmbeddedField {
    collection: String,
    field: String,
    kind: VectorKind,
    source: String,
    model: ModelName,
    dimensions: usize,
    distance: Metric,
    status: &'static str,
    /// EMB1: the records given no vector because the model reads only part
    /// of their text, as a service that embeds the field counts them since
    /// it started; absent where the store alone answers.
    #[serde(skip_serializing_if = "Option::is_none")]
    longer_than_model: Option<u64>,
    /// EMB-E1: the records passed with no vector because their text is
    /// empty or only whitespace, counted as `longer_than_model` is.
    #[serde(skip_serializing_if = "Option::is_none")]
    embed_input_empty: Option<u64>,
    /// EMB-E2: the records given windows, and those given none for needing
    /// more windows than the cap, counted as `longer_than_model` is.
    #[serde(skip_serializing_if = "Option::is_none")]
    windowed: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    windows_over_cap: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    failure: Option<serde_json::Value>,
}

/// What a service that embeds a field says of it.
///
/// Its status (EWR1), the records the model reads only part of (EMB1), those
/// passed for empty text (EMB-E1), those given windows or none past the cap
/// (EMB-E2), and its last failure.
pub(in crate::serve) struct FieldState {
    pub status: &'static str,
    pub longer_than_model: u64,
    pub embed_input_empty: u64,
    pub windows: (u64, u64),
    pub failure: Option<serde_json::Value>,
}

pub(in crate::serve) struct FieldListing {
    pub fields: Vec<VectorField>,
    pub embedded_fields: Vec<EmbeddedField>,
}

impl FieldListing {
    /// Each embedded field's state as `status` names it from the field's
    /// collection and name.
    pub fn statused(&mut self, status: impl Fn(&str, &str) -> FieldState) {
        for embedded in &mut self.embedded_fields {
            let state = status(&embedded.collection, &embedded.field);
            embedded.status = state.status;
            embedded.longer_than_model = Some(state.longer_than_model);
            embedded.embed_input_empty = Some(state.embed_input_empty);
            embedded.windowed = Some(state.windows.0);
            embedded.windows_over_cap = Some(state.windows.1);
            embedded.failure = state.failure;
        }
    }
}

pub(in crate::serve) fn fields(
    access: &Access,
    reader: &ReadOnlyDatabase,
    at: &Captured,
) -> Result<FieldListing, Failure> {
    let mut listing = FieldListing {
        fields: Vec::new(),
        embedded_fields: Vec::new(),
    };
    reads::scan_rows(reader, at, b"w\0", 2, |parts, bytes| {
        match embedded::read_field(bytes)? {
            Field::Supplied { dimensions } => listing.fields.push(VectorField {
                collection: parts[0].to_owned(),
                field: parts[1].to_owned(),
                dimensions,
                count: visible_count(access, reader, at, parts[0], parts[1])?,
            }),
            Field::Embedded(declaration) => listing.embedded_fields.push(EmbeddedField {
                collection: parts[0].to_owned(),
                field: parts[1].to_owned(),
                kind: declaration.kind,
                source: declaration.source,
                model: declaration.model,
                dimensions: declaration.dimensions.get(),
                distance: declaration.distance,
                // EWR1: what the store alone knows; a service that embeds
                // the field puts its driver's status here.
                status: "embed_not_configured",
                longer_than_model: None,
                embed_input_empty: None,
                windowed: None,
                windows_over_cap: None,
                failure: None,
            }),
        }
        Ok(())
    })?;
    Ok(listing)
}

pub(super) fn declare(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    write: &VectorDeclaration,
) -> Result<Vec<PreparedWrite>, Failure> {
    identifier(&write.collection)?;
    identifier(&write.field)?;
    super::field_index::require_non_vector(reader, at, &write.collection, &write.field)?;
    identifier(&write.declaration.source)?;
    identifier(&write.declaration.model.name)?;
    identifier(&write.declaration.model.revision)?;
    let manifest = &write.declaration.model.manifest_sha256;
    if manifest.len() != 64
        || !manifest
            .bytes()
            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
    {
        return Err(Failure::new(
            "embed_manifest_invalid",
            "model manifest must be 64 lowercase hexadecimal digits",
        ));
    }
    if reads::collection(reader, at, &write.collection)?.is_none() {
        return Err(Failure::new("collection_missing", &write.collection));
    }
    match recorded_field(reader, at, &write.collection, &write.field)? {
        Some(Field::Supplied { .. }) => {
            return Err(Failure::new(
                "vector_field_supplied",
                "a supplied field cannot become embedded",
            ));
        }
        Some(Field::Embedded(existing)) if existing != write.declaration => {
            return Err(Failure::new(
                "embed_model_change_not_available",
                "an embedded declaration cannot be replaced",
            ));
        }
        Some(Field::Embedded(_)) | None => {}
    }
    Ok(vec![(
        dimension_key(&write.collection, &write.field),
        Some(embedded::encode_declaration(&write.declaration)?),
    )])
}

fn visible_count(
    access: &Access,
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
    field: &str,
) -> Result<u64, Failure> {
    let mut count = 0_u64;
    reads::scan_rows(
        reader,
        at,
        &vector_prefix(collection, field),
        1,
        |parts, _| {
            if access.can_read_document(collection, parts[0]) {
                count = count.checked_add(1).ok_or_else(|| {
                    Failure::new("count_overflow", "stored key count exceeds u64")
                })?;
            }
            Ok(())
        },
    )?;
    Ok(count)
}

pub(super) fn put(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    write: &VectorWrite,
) -> Result<Vec<PreparedWrite>, Failure> {
    parts(&write.collection, &write.field, &write.document_id)?;
    super::field_index::require_non_vector(reader, at, &write.collection, &write.field)?;
    finite_nonempty(&write.vector)?;
    if reads::get_bytes(
        reader,
        at,
        &document_key(&write.collection, &write.document_id),
    )?
    .is_none()
    {
        return Err(Failure::new("document_missing", &write.document_id));
    }
    let recorded = match recorded_field(reader, at, &write.collection, &write.field)? {
        Some(Field::Supplied { dimensions }) => Some(dimensions),
        Some(Field::Embedded(_)) => {
            return Err(Failure::new(
                "vector_field_embedded",
                "an embedded field does not accept supplied vectors",
            ));
        }
        None => None,
    };
    if recorded.is_some_and(|dimension| dimension != write.vector.len()) {
        return Err(Failure::new(
            "vector_dimension",
            "vector dimension differs from the field",
        ));
    }
    let mut writes = vec![(
        vector_key(&write.collection, &write.field, &write.document_id),
        Some(
            write
                .vector
                .iter()
                .flat_map(|member| member.to_le_bytes())
                .collect(),
        ),
    )];
    if recorded.is_none() {
        writes.push((
            dimension_key(&write.collection, &write.field),
            Some(write.vector.len().to_string().into_bytes()),
        ));
    }
    Ok(writes)
}

pub(super) fn delete(write: &VectorDelete) -> Result<Vec<PreparedWrite>, Failure> {
    parts(&write.collection, &write.field, &write.document_id)?;
    Ok(vec![(
        vector_key(&write.collection, &write.field, &write.document_id),
        None,
    )])
}

pub(in crate::serve) fn get(
    access: &Access,
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
    field: &str,
    document_id: &str,
) -> Result<Option<Vec<f32>>, Failure> {
    parts(collection, field, document_id)?;
    if !access.can_read_document(collection, document_id) {
        return Ok(None);
    }
    let recorded = dimension(reader, at, collection, field)?;
    let bytes = reads::get_bytes(reader, at, &vector_key(collection, field, document_id))?;
    let Some(bytes) = bytes else { return Ok(None) };
    let recorded = recorded.ok_or_else(|| {
        Failure::new("vector_dimension_missing", "stored vector has no dimension")
    })?;
    decode(&bytes, recorded).map(Some)
}

pub(in crate::serve) fn list(
    access: &Access,
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
    field: &str,
    after: Option<&str>,
    limits: (NonZeroUsize, NonZeroUsize),
) -> Result<(Vec<Vector>, Option<String>), Failure> {
    identifier(collection)?;
    identifier(field)?;
    let recorded = dimension(reader, at, collection, field)?;
    let prefix = vector_prefix(collection, field);
    let from = boundary(&prefix, after)?;
    reads::visible_key_page(
        reader,
        at,
        reads::VisibleKeyPage {
            access,
            collection,
            prefix: &prefix,
            from: &from,
            limits: reads::PageLimits {
                rows: limits.0,
                bytes: limits.1,
            },
        },
        |document_id, bytes| stored_vector(document_id, &bytes, recorded),
        |vector| {
            if limits.1 == NonZeroUsize::MAX {
                return Ok(0);
            }
            serde_json::to_vec(vector)
                .map(|bytes| bytes.len())
                .map_err(Failure::storage)
        },
    )
}

fn stored_vector(
    document_id: String,
    bytes: &[u8],
    recorded: Option<usize>,
) -> Result<Vector, Failure> {
    let recorded = recorded.ok_or_else(|| {
        Failure::new("vector_dimension_missing", "stored vector has no dimension")
    })?;
    Ok(Vector {
        document_id,
        vector: decode(bytes, recorded)?,
    })
}

fn search_page(
    access: &Access,
    reader: &ReadOnlyDatabase,
    at: &Captured,
    request: &VectorSearch,
    after: Option<&str>,
    recorded: Option<usize>,
    rows: NonZeroUsize,
) -> Result<(Vec<Option<Vector>>, Option<String>), Failure> {
    let prefix = vector_prefix(&request.collection, &request.field);
    let from = boundary(&prefix, after)?;
    reads::bounded_key_page(
        reader,
        at,
        &prefix,
        &from,
        reads::PageLimits {
            rows,
            bytes: NonZeroUsize::MAX,
        },
        |document_id, bytes| {
            identifier(&document_id)?;
            if !access.can_read_document(&request.collection, &document_id) {
                return Ok(None);
            }
            stored_vector(document_id, &bytes, recorded).map(Some)
        },
        |_: &Option<Vector>| Ok(0),
    )
}

pub(in crate::serve) fn search(
    access: &Access,
    reader: &ReadOnlyDatabase,
    at: &Captured,
    request: &VectorSearch,
    k: NonZeroUsize,
) -> Result<(Vec<Hit>, u64), Failure> {
    search_within(access, reader, at, request, k, None)
}

/// [`search`], comparing no more than `bound` vectors when one is given: it
/// stops at the first vector past the bound and refuses `query_scan_limit`,
/// never comparing every vector first. The vectors compared are those the
/// request may read.
pub(in crate::serve) fn search_within(
    access: &Access,
    reader: &ReadOnlyDatabase,
    at: &Captured,
    request: &VectorSearch,
    k: NonZeroUsize,
    bound: Option<u64>,
) -> Result<(Vec<Hit>, u64), Failure> {
    identifier(&request.collection)?;
    identifier(&request.field)?;
    finite_nonempty(&request.query)?;
    let recorded = dimension(reader, at, &request.collection, &request.field)?;
    if recorded.is_some_and(|dimension| dimension != request.query.len()) {
        return Err(Failure::new(
            "vector_dimension",
            "query dimension differs from the field",
        ));
    }
    let mut hits = Vec::new();
    let mut scanned = 0_u64;
    let mut after = None;
    let rows = scan_page_size()?;
    loop {
        let (vectors, next_after) = search_page(
            access,
            reader,
            at,
            request,
            after.as_deref(),
            recorded,
            rows,
        )?;
        for vector in vectors.into_iter().flatten() {
            scanned = scanned
                .checked_add(1)
                .ok_or_else(|| Failure::new("vector_count_overflow", "vector count exceeds u64"))?;
            if let Some(bound) = bound.filter(|bound| scanned > *bound) {
                return Err(Failure::new(
                    "query_scan_limit",
                    format!("the search compared more than {bound} vectors"),
                ));
            }
            let hit = Hit {
                score: score(&request.query, &vector.vector, request.metric),
                document_id: vector.document_id,
            };
            let position =
                match hits.binary_search_by(|candidate| order(candidate, &hit, request.metric)) {
                    Ok(position) | Err(position) => position,
                };
            if position < k.get() {
                hits.insert(position, hit);
                hits.truncate(k.get());
            }
        }
        let Some(next_after) = next_after else {
            return Ok((hits, scanned));
        };
        after = Some(next_after);
    }
}

pub(super) fn removals_for_document(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
    document_id: &str,
) -> Result<Vec<PreparedWrite>, Failure> {
    identifier(collection)?;
    identifier(document_id)?;
    let prefix = format!("w\0{collection}\0").into_bytes();
    let mut after = None;
    let mut removals = Vec::new();
    let rows = scan_page_size()?;
    loop {
        let from = boundary(&prefix, after.as_deref())?;
        let (fields, next_after) = reads::bounded_key_page(
            reader,
            at,
            &prefix,
            &from,
            reads::PageLimits {
                rows,
                bytes: NonZeroUsize::MAX,
            },
            |field, bytes| {
                identifier(&field)?;
                match embedded::read_field(&bytes)? {
                    Field::Supplied { .. } => Ok(Some(field)),
                    Field::Embedded(_) => Ok(None),
                }
            },
            |field| {
                serde_json::to_vec(field)
                    .map(|bytes| bytes.len())
                    .map_err(Failure::storage)
            },
        )?;
        for field in fields.into_iter().flatten() {
            removals.push((vector_key(collection, &field, document_id), None));
        }
        let Some(next_after) = next_after else {
            return Ok(removals);
        };
        after = Some(next_after);
    }
}

fn scan_page_size() -> Result<NonZeroUsize, Failure> {
    NonZeroUsize::new(64)
        .ok_or_else(|| Failure::new("page_limit", "vector scan page size must be nonzero"))
}

fn parts(collection: &str, field: &str, document_id: &str) -> Result<(), Failure> {
    identifier(collection)?;
    identifier(field)?;
    identifier(document_id)
}

fn vector_prefix(collection: &str, field: &str) -> Vec<u8> {
    format!("v\0{collection}\0{field}\0").into_bytes()
}

pub(super) fn vector_key(collection: &str, field: &str, document_id: &str) -> Vec<u8> {
    format!("v\0{collection}\0{field}\0{document_id}").into_bytes()
}

fn dimension_key(collection: &str, field: &str) -> Vec<u8> {
    format!("w\0{collection}\0{field}").into_bytes()
}

fn boundary(prefix: &[u8], after: Option<&str>) -> Result<Vec<u8>, Failure> {
    let mut from = prefix.to_vec();
    if let Some(after) = after {
        identifier(after)?;
        from.extend_from_slice(after.as_bytes());
        from.push(0);
    }
    Ok(from)
}

pub(super) fn finite_nonempty(vector: &[f32]) -> Result<(), Failure> {
    if vector.is_empty() {
        return Err(Failure::new(
            "vector_empty",
            "a vector must have at least one member",
        ));
    }
    if vector.iter().any(|member| !member.is_finite()) {
        return Err(Failure::new(
            "vector_non_finite",
            "vector members must be finite",
        ));
    }
    Ok(())
}

pub(super) fn dimension(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
    field: &str,
) -> Result<Option<usize>, Failure> {
    match recorded_field(reader, at, collection, field)? {
        Some(Field::Supplied { dimensions }) => Ok(Some(dimensions)),
        Some(Field::Embedded(_)) => Err(Failure::new(
            "embed_companion_not_served",
            "embedded vectors are not served by the companion yet",
        )),
        None => Ok(None),
    }
}

pub(super) fn recorded_field(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
    field: &str,
) -> Result<Option<Field>, Failure> {
    reads::get_bytes(reader, at, &dimension_key(collection, field))?
        .map(|bytes| embedded::read_field(&bytes))
        .transpose()
}

pub(super) fn parse_dimension(bytes: &[u8]) -> Result<usize, Failure> {
    let corrupt = || {
        Failure::new(
            "vector_dimension_corrupt",
            "dimension must be positive decimal",
        )
    };
    let text = std::str::from_utf8(bytes).map_err(|_| corrupt())?;
    let dimension = text.parse::<usize>().map_err(|_| corrupt())?;
    if dimension == 0 || text != dimension.to_string() {
        return Err(corrupt());
    }
    Ok(dimension)
}

pub(super) fn decode(bytes: &[u8], dimension: usize) -> Result<Vec<f32>, Failure> {
    let (chunks, rest) = bytes.as_chunks::<4>();
    if !rest.is_empty() || chunks.len() != dimension {
        return Err(Failure::new(
            "vector_corrupt",
            "stored vector has the wrong byte length",
        ));
    }
    let vector = chunks
        .iter()
        .map(|chunk| {
            let member = f32::from_le_bytes(*chunk);
            if !member.is_finite() {
                return Err(Failure::new(
                    "vector_corrupt",
                    "stored vector has a non-finite member",
                ));
            }
            Ok(member)
        })
        .collect::<Result<Vec<_>, Failure>>()?;
    Ok(vector)
}

fn score(query: &[f32], vector: &[f32], metric: Metric) -> f64 {
    let mut dot = 0.0_f64;
    let mut query_norm = 0.0_f64;
    let mut vector_norm = 0.0_f64;
    let mut squared_distance = 0.0_f64;
    for (left, right) in query.iter().zip(vector) {
        let left = f64::from(*left);
        let right = f64::from(*right);
        dot = left.mul_add(right, dot);
        query_norm = left.mul_add(left, query_norm);
        vector_norm = right.mul_add(right, vector_norm);
        let difference = left - right;
        squared_distance = difference.mul_add(difference, squared_distance);
    }
    match metric {
        Metric::Dot => dot,
        Metric::L2 => squared_distance.sqrt(),
        Metric::Cosine if query_norm <= 0.0 || vector_norm <= 0.0 => 0.0,
        Metric::Cosine => dot / (query_norm.sqrt() * vector_norm.sqrt()),
    }
}

fn order(left: &Hit, right: &Hit, metric: Metric) -> Ordering {
    let score = match metric {
        Metric::L2 => left.score.total_cmp(&right.score),
        Metric::Dot | Metric::Cosine => right.score.total_cmp(&left.score),
    };
    score.then_with(|| left.document_id.cmp(&right.document_id))
}