lattice-embed 0.11.1

//! Async embedding-service contract and native implementations.
//!
//! The trait defines generic, query, and passage embedding; native builds additionally expose
//! a lazy local-inference service, an LRU caching wrapper, and dormant ADR-088 attestation/report
//! and resource-budget contracts. The dormant contracts do not yet construct or serve a prepared
//! model. See `docs/service.md` for lifecycle, validation, cache, and contract boundaries.

#[cfg(feature = "native")]
mod cached;
#[cfg(feature = "native")]
mod native;
#[cfg(feature = "native")]
mod prepared;

#[cfg(test)]
mod tests;

use crate::error::{EmbedError, Result};
use crate::model::{EmbeddingModel, ModelConfig};
use async_trait::async_trait;

#[cfg(test)]
std::thread_local! {
    static VALIDATE_TEXTS_CALLS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}

// Re-exports
#[cfg(feature = "native")]
pub use cached::CachedEmbeddingService;
#[cfg(feature = "native")]
pub use native::NativeEmbeddingService;
#[cfg(feature = "native")]
pub use prepared::{
    AttestationAlgorithm, CheckpointAttestor, MAX_SUPPLEMENTARY_ATTESTATION_EVIDENCE_BYTES,
    MIN_SUPPLEMENTARY_ATTESTATION_EVIDENCE_BYTES, NativeResourceBudget,
    SupplementaryAttestationEvidence,
};

/// **Stable**: default maximum batch size to prevent OOM.
///
/// This limit prevents accidentally passing huge batches that could exhaust memory.
/// Can be overridden by using chunked calls if larger batches are needed.
pub const DEFAULT_MAX_BATCH_SIZE: usize = 1000;

/// **Stable**: maximum allowed text length in UTF-8 bytes.
///
/// This limit prevents OOM attacks via extremely large input texts.
/// 32KB is sufficient for most embedding use cases while preventing abuse.
///
/// Multibyte text can exceed this limit well before its `chars()` count does;
/// the guard checks `str::len()` (bytes), not character count.
pub const MAX_TEXT_BYTES: usize = 32768;

/// Deprecated alias for [`MAX_TEXT_BYTES`] — the old name implied character
/// count, but the guard has always counted UTF-8 bytes.
#[deprecated(
    since = "0.7.0",
    note = "use MAX_TEXT_BYTES; this limit counts UTF-8 bytes, not chars"
)]
pub const MAX_TEXT_CHARS: usize = MAX_TEXT_BYTES;

/// **Stable**: role of text in asymmetric retrieval.
///
/// Selects query, passage, or generic preparation and cache-key namespace.
/// See [`docs/service.md`](../../docs/service.md#trait-api-details) for retrieval-role semantics.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum EmbeddingRole {
    /// Query / question text — may receive a query-side prompt prefix.
    Query,
    /// Document / passage text — may receive a passage-side prompt prefix.
    Passage,
    /// Generic text with no role-specific prefix (backwards-compatible default).
    Generic,
}

impl EmbeddingRole {
    /// Returns this role's short cache-key tag.
    #[inline]
    pub(crate) const fn cache_tag(self) -> &'static str {
        match self {
            EmbeddingRole::Query => "role:query",
            EmbeddingRole::Passage => "role:passage",
            EmbeddingRole::Generic => "role:generic",
        }
    }

    /// Returns the instruction this role prepends for `model`, if any.
    #[inline]
    pub(crate) const fn instruction(self, model: EmbeddingModel) -> Option<&'static str> {
        match self {
            EmbeddingRole::Query => model.query_instruction(),
            EmbeddingRole::Passage => model.document_instruction(),
            EmbeddingRole::Generic => None,
        }
    }
}

/// Enforces the published request bounds against caller-supplied text.
///
/// This is the contract check, so it runs on what the caller actually passed
/// and before any instruction is prepended. Guards that run downstream of
/// preparation are memory backstops and size themselves with
/// [`EmbeddingModel::max_instruction_bytes`]; they are not this check.
pub(crate) fn validate_texts(texts: &[String]) -> Result<()> {
    #[cfg(test)]
    VALIDATE_TEXTS_CALLS.set(VALIDATE_TEXTS_CALLS.get() + 1);
    validate_texts_bounded(texts, MAX_TEXT_BYTES)
}

#[cfg(test)]
pub(crate) fn reset_validate_texts_calls() {
    VALIDATE_TEXTS_CALLS.set(0);
}

#[cfg(test)]
pub(crate) fn validate_texts_calls() -> usize {
    VALIDATE_TEXTS_CALLS.get()
}

/// Shared body of the caller-text contract check and the prepared-text backstop.
///
/// `max_bytes` is [`MAX_TEXT_BYTES`] on caller text, and that value plus the
/// model's longest instruction on text that has already been prepared. The
/// reported `max` is `max_bytes` so an error names the bound that actually
/// rejected the input rather than a constant the caller cannot relate to.
pub(crate) fn validate_texts_bounded<T: AsRef<str>>(texts: &[T], max_bytes: usize) -> Result<()> {
    if texts.is_empty() {
        return Err(EmbedError::InvalidInput("no texts provided".into()));
    }
    if texts.len() > DEFAULT_MAX_BATCH_SIZE {
        return Err(EmbedError::InvalidInput(format!(
            "batch size {} exceeds maximum {}",
            texts.len(),
            DEFAULT_MAX_BATCH_SIZE
        )));
    }
    for text in texts {
        let text = text.as_ref();
        if text.len() > max_bytes {
            return Err(EmbedError::TextTooLong {
                length: text.len(),
                max: max_bytes,
            });
        }
    }
    Ok(())
}

/// **Stable**: token counts for one input text, measured on the sequence the model sees.
///
/// Special tokens and any role instruction are counted, because they occupy the model's
/// window. `before_truncation` is the length the sequence would have without a limit and
/// `embedded` is the length the model actually consumed, so `embedded` never exceeds it.
/// The limit is the one the loaded model applies, which can differ from
/// [`EmbeddingModel::max_input_tokens`].
///
/// See [`docs/service.md`](../../docs/service.md#token-counts-and-truncation-reports) for how
/// each native model family produces the counts.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct TokenCount {
    /// Sequence length before any length limit was applied.
    pub before_truncation: usize,
    /// Sequence length the model embedded. At most `before_truncation`.
    pub embedded: usize,
}

impl TokenCount {
    /// **Stable**: builds a count for implementations of
    /// [`EmbeddingService::count_tokens`] and [`EmbeddingService::embed_with_report`].
    ///
    /// `embedded` must not exceed `before_truncation`.
    #[inline]
    pub const fn new(before_truncation: usize, embedded: usize) -> Self {
        Self {
            before_truncation,
            embedded,
        }
    }

    /// **Stable**: whether the model embedded only a prefix of the input.
    #[inline]
    pub const fn truncated(&self) -> bool {
        self.embedded < self.before_truncation
    }
}

/// **Stable**: embeddings returned together with a per-input token report.
///
/// See [`EmbeddingService::embed_with_report`].
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct EmbeddingsWithReport {
    /// One embedding per input, in input order. Identical to what
    /// [`EmbeddingService::embed_with_role`] returns for the same request.
    pub embeddings: Vec<Vec<f32>>,
    /// One token count per input, in input order, same length as `embeddings`.
    pub token_counts: Vec<TokenCount>,
}

impl EmbeddingsWithReport {
    /// **Stable**: pairs embeddings with their token counts, one count per embedding.
    pub fn new(embeddings: Vec<Vec<f32>>, token_counts: Vec<TokenCount>) -> Self {
        Self {
            embeddings,
            token_counts,
        }
    }
}

enum ValidatedTextBatchInner<'a> {
    Contiguous(&'a [String]),
    Borrowed(&'a [&'a str]),
}

/// Internal capability proving the caller-text request bounds were checked.
///
/// The constructors are crate-private, so external callers cannot manufacture
/// this capability to bypass [`EmbeddingService`] input validation.
#[doc(hidden)]
pub struct ValidatedTextBatch<'a> {
    inner: ValidatedTextBatchInner<'a>,
}

impl<'a> ValidatedTextBatch<'a> {
    pub(in crate::service) fn new(texts: &'a [String]) -> Result<Self> {
        validate_texts(texts)?;
        Ok(Self {
            inner: ValidatedTextBatchInner::Contiguous(texts),
        })
    }

    pub(in crate::service) fn borrowed_subset<'b>(
        &'b self,
        texts: &'b [&'b str],
    ) -> ValidatedTextBatch<'b> {
        ValidatedTextBatch {
            inner: ValidatedTextBatchInner::Borrowed(texts),
        }
    }

    pub(in crate::service) fn len(&self) -> usize {
        match self.inner {
            ValidatedTextBatchInner::Contiguous(texts) => texts.len(),
            ValidatedTextBatchInner::Borrowed(texts) => texts.len(),
        }
    }

    pub(in crate::service) fn get(&self, index: usize) -> &str {
        match self.inner {
            ValidatedTextBatchInner::Contiguous(texts) => texts[index].as_str(),
            ValidatedTextBatchInner::Borrowed(texts) => texts[index],
        }
    }

    fn contiguous(&self) -> Option<&'a [String]> {
        match self.inner {
            ValidatedTextBatchInner::Contiguous(texts) => Some(texts),
            ValidatedTextBatchInner::Borrowed(_) => None,
        }
    }

    pub(in crate::service) fn borrowed(&self) -> Option<&'a [&'a str]> {
        match self.inner {
            ValidatedTextBatchInner::Contiguous(_) => None,
            ValidatedTextBatchInner::Borrowed(texts) => Some(texts),
        }
    }

    pub(in crate::service) fn to_owned_with_prefix(&self, prefix: Option<&str>) -> Vec<String> {
        (0..self.len())
            .map(|index| {
                let text = self.get(index);
                match prefix {
                    None => text.to_owned(),
                    Some(prefix) => {
                        let mut prepared = String::with_capacity(prefix.len() + text.len());
                        prepared.push_str(prefix);
                        prepared.push_str(text);
                        prepared
                    }
                }
            })
            .collect()
    }
}

/// **Stable**: external consumers may depend on this; breaking changes require a SemVer bump.
///
/// Async interface for producing one embedding per input text.
///
/// See [`docs/service.md`](../../docs/service.md#trait-api-details) for role handling and implementation requirements.
#[async_trait]
pub trait EmbeddingService: Send + Sync {
    /// **Stable**: generate embeddings for multiple texts.
    ///
    /// Returns a vector of embeddings, one for each input text, in the same order.
    /// Applies no role-specific prompt prefix (equivalent to `Generic` role).
    /// Use [`EmbeddingService::embed_query`] / [`EmbeddingService::embed_passage`]
    /// for asymmetric retrieval models.
    async fn embed(&self, texts: &[String], model: EmbeddingModel) -> Result<Vec<Vec<f32>>>;

    /// **Stable**: generate an embedding for a single text.
    ///
    /// This is a convenience method that calls `embed` with a single-element slice.
    async fn embed_one(&self, text: &str, model: EmbeddingModel) -> Result<Vec<f32>> {
        let texts = vec![text.to_string()];
        let mut embeddings = self.embed(&texts, model).await?;
        embeddings
            .pop()
            .ok_or_else(|| EmbedError::Internal("no embedding generated".into()))
    }

    /// **Stable**: embed texts under a retrieval role, applying that role's instruction.
    ///
    /// The published length cap applies to the text the caller supplies, so this
    /// validates first and prepends the instruction second. Doing it the other
    /// way round charges the caller for bytes the service itself added, which
    /// rejects text that is within the documented limit.
    ///
    /// Implementors that enforce a text-length cap inside [`EmbeddingService::embed`]
    /// receive prepared text here, which is longer than the caller's by up to
    /// [`EmbeddingModel::max_instruction_bytes`]. Size that guard accordingly, or
    /// override this method to reach the backend without passing through it.
    ///
    /// See [`docs/service.md`](../../docs/service.md#trait-api-details) for role and cache behavior.
    async fn embed_with_role(
        &self,
        texts: &[String],
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<Vec<Vec<f32>>> {
        validate_texts(texts)?;
        let prepared = apply_prefix(texts, role.instruction(model));
        self.embed(&prepared, model).await
    }

    /// Internal delegation hook for caller text already checked against the public bounds.
    ///
    /// External callers cannot construct [`ValidatedTextBatch`]. The default
    /// preserves external implementors' role overrides by delegating through
    /// [`EmbeddingService::embed_with_role`].
    #[doc(hidden)]
    async fn embed_with_role_prevalidated(
        &self,
        texts: ValidatedTextBatch<'_>,
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<Vec<Vec<f32>>> {
        if let Some(contiguous) = texts.contiguous() {
            return self.embed_with_role(contiguous, model, role).await;
        }
        let owned = texts.to_owned_with_prefix(None);
        self.embed_with_role(&owned, model, role).await
    }

    /// **Stable**: count the tokens each text occupies in `model`, without embedding it.
    ///
    /// Validation and the role instruction match [`EmbeddingService::embed_with_role`], so
    /// the counts describe the sequence that call would embed. [`EmbeddingRole::Generic`]
    /// adds no instruction and counts the same text [`EmbeddingService::embed`] sees. Use it
    /// to chunk text before submitting it. The default returns [`EmbedError::Unsupported`].
    ///
    /// See [`docs/service.md`](../../docs/service.md#token-counts-and-truncation-reports) for the count definition.
    async fn count_tokens(
        &self,
        texts: &[String],
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<Vec<TokenCount>> {
        let _ = (texts, model, role);
        Err(EmbedError::Unsupported(format!(
            "{} does not implement count_tokens",
            self.name()
        )))
    }

    /// **Stable**: embed texts under a retrieval role and report, per input, how many tokens
    /// the model saw.
    ///
    /// The embeddings are identical to those of [`EmbeddingService::embed_with_role`] for the
    /// same request. A caller that bounds input by bytes can read
    /// [`TokenCount::truncated`] to learn that only a prefix was embedded. The default
    /// returns [`EmbedError::Unsupported`].
    ///
    /// See [`docs/service.md`](../../docs/service.md#token-counts-and-truncation-reports) for the count definition.
    async fn embed_with_report(
        &self,
        texts: &[String],
        model: EmbeddingModel,
        role: EmbeddingRole,
    ) -> Result<EmbeddingsWithReport> {
        let _ = (texts, model, role);
        Err(EmbedError::Unsupported(format!(
            "{} does not implement embed_with_report",
            self.name()
        )))
    }

    /// **Stable**: embed query texts after applying the model's query instruction.
    ///
    /// See [`docs/service.md`](../../docs/service.md#trait-api-details) for role and cache behavior.
    async fn embed_query(&self, texts: &[String], model: EmbeddingModel) -> Result<Vec<Vec<f32>>> {
        self.embed_with_role(texts, model, EmbeddingRole::Query)
            .await
    }

    /// **Stable**: embed passages after applying the model's document instruction.
    ///
    /// See [`docs/service.md`](../../docs/service.md#trait-api-details) for role and cache behavior.
    async fn embed_passage(
        &self,
        texts: &[String],
        model: EmbeddingModel,
    ) -> Result<Vec<Vec<f32>>> {
        self.embed_with_role(texts, model, EmbeddingRole::Passage)
            .await
    }

    /// **Unstable**: returns the effective configuration used for this model's cache keys.
    ///
    /// See [`docs/service.md`](../../docs/service.md#trait-api-details) for output-dimension behavior.
    fn model_config(&self, model: EmbeddingModel) -> ModelConfig {
        ModelConfig::new(model)
    }

    /// **Stable**: check if the service supports a given model.
    fn supports_model(&self, model: EmbeddingModel) -> bool;

    /// **Stable**: get the name/identifier of this service.
    fn name(&self) -> &'static str;
}

/// Prepends an optional prompt prefix to each text, cloning unchanged inputs when absent.
pub(crate) fn apply_prefix(texts: &[String], prefix: Option<&str>) -> Vec<String> {
    match prefix {
        None => texts.to_vec(),
        Some(p) => texts.iter().map(|t| format!("{p}{t}")).collect(),
    }
}