haematite-cli 0.21.8

use std::num::{NonZeroU64, NonZeroUsize};

use serde_json::value::RawValue;
use serde_json::{Value, json};

use super::{Service, view};
use crate::serve::access::Access;
use crate::serve::embed::question;
use crate::serve::protocol::{Metric, QueryRequest, VectorSearchText};
use crate::serve::query::probe::Prepared;
use crate::serve::query::run::{self, Context};
use crate::serve::storage::embed_search::{self, Embedding, Found, Searched};
use crate::serve::storage::{Captured, Failure};

/// One question searched under a field's declaration at a captured moment.
pub(in crate::serve) struct Words<'a> {
    pub collection: &'a str,
    pub field: &'a str,
    pub text: &'a str,
    pub measure: Option<Metric>,
    pub k: NonZeroUsize,
    pub bound: u64,
}

impl Service {
    /// Reads records at one capture; text search briefly captures its field.
    pub(super) fn query_run(
        &self,
        access: &Access,
        request: &QueryRequest,
        response_byte_limit: Option<NonZeroUsize>,
    ) -> Result<Box<RawValue>, Failure> {
        let context = Context {
            access,
            reader: &self.reader,
            catalogue: &self.catalogue,
            slots: &self.query_slots,
            settings: self.query.as_ref().map(Prepared::settings),
            max_page: self.max_page,
            seal: &self.query_seal,
        };
        run::served(&context, request, response_byte_limit, Some(self))
    }

    /// Embed outside writer admission and search exactly the declared field.
    pub(in crate::serve) fn search_words(
        &self,
        access: &Access,
        at: &Captured,
        words: &Words<'_>,
    ) -> Result<(Found, Embedding), Failure> {
        let driver = self.embedding.as_ref().ok_or_else(not_configured)?;
        let asker = driver.asker()?;
        let (declaration, companion, coverage) = {
            let writer = self.admission()?;
            let (declaration, companion) =
                writer.embedded_field(at, words.collection, words.field)?;
            if let Some(measure) = words.measure
                && measure != declaration.distance
            {
                return Err(Failure::new(
                    "query_measure_mismatch",
                    format!(
                        "similar names {} for .{}; the field declares {}",
                        metric(measure),
                        words.field,
                        metric(declaration.distance),
                    ),
                ));
            }
            let coverage = writer.embedding_state(at, words.collection, words.field)?;
            drop(writer);
            (declaration, companion, coverage)
        };
        let vector = question::question_for(asker, words.text, &declaration)?;
        let found = embed_search::search(
            access,
            &self.reader,
            &Searched {
                collection: words.collection,
                field: words.field,
                declaration: &declaration,
                at,
                companion: companion.as_ref(),
            },
            &vector,
            words.k,
            words.bound,
        )?;
        Ok((found, coverage))
    }

    /// `text` embedded without a lock, searched with the field's coverage.
    pub(super) fn search_text(
        &self,
        access: &Access,
        search: &VectorSearchText,
    ) -> Result<Value, Failure> {
        let k = self.limit(search.k)?;
        let driver = self.embedding.as_ref().ok_or_else(not_configured)?;
        let bound = self
            .embed_search_bound
            .map(NonZeroU64::get)
            .ok_or_else(|| {
                Failure::new(
                    "embed_search_unbounded",
                    "text search reads at most query_scan_records vectors; set it",
                )
            })?;
        // Readiness is named before capturing or reading the field.
        driver.asker()?;
        let at = self.capture(&search.selection, false)?;
        let (found, coverage) = self.search_words(
            access,
            &at,
            &Words {
                collection: &search.collection,
                field: &search.field,
                text: &search.text,
                measure: None,
                k,
                bound,
            },
        )?;
        let walk_open = driver.walk_open(&at.branch, &search.collection, &search.field)
            || coverage.toward.is_some();
        let hits: Vec<Value> = found
            .hits
            .iter()
            .map(|hit| {
                let mut answer = json!({"document_id": hit.id, "score": hit.score});
                if let (Some(span), Some(fields)) = (hit.span, answer.as_object_mut()) {
                    fields.insert(
                        "span".to_owned(),
                        json!({"start": span.start, "end": span.end}),
                    );
                }
                answer
            })
            .collect();
        Ok(json!({
            "selection": view(&at), "hits": hits, "compared": found.compared,
            "stale": found.stale,
            "coverage": {
                "state": coverage.state, "branch": coverage.branch,
                "through": coverage.through, "toward": coverage.toward,
                "walk": if walk_open { "open" } else { "done" },
            },
        }))
    }
}

pub(in crate::serve) fn not_configured() -> Failure {
    Failure::new(
        "embed_not_configured",
        "text search needs the embed settings in serve.json",
    )
}

const fn metric(measure: Metric) -> &'static str {
    match measure {
        Metric::Cosine => "cosine",
        Metric::L2 => "l2",
        Metric::Dot => "dot",
    }
}