lc-embeddings 0.22.4

// lc-embeddings/src/token_level.rs
//! Optional token-level embedding capability (0.21.0 S5.2, native-async pilot).
//!
//! Pooled (one-vector-per-text) embeddings lose cross-chunk references: after
//! chunking, "the company"/"this plan" style pronouns embed without their
//! antecedent. Token-level output — one vector per token plus its byte span in
//! the original text — is the prerequisite for **late chunking** (embed the
//! whole document token by token first, then pool along chunk boundaries).
//!
//! This module provides the capability trait; whether a provider can honor it
//! is detected by implementing the trait (capability-by-type, no runtime
//! probing). Requires a long-context, token-output-capable embedder
//! (Qwen3-Embedding, BGE-M3, Jina v3, ...). Pooled models cannot implement it.
//!
//! 0.21.0 S3.2 pilot — **native `async fn` in trait** instead of `#[async_trait]`:
//! - RPITIT (stable since 1.75) gives static dispatch with zero boxing;
//! - the `Send` bound needed for concurrent use is generated by
//!   [`trait_variant`](https://docs.rs/trait_variant), the current standard
//!   stopgap until native async dyn dispatch lands;
//! - RPITIT is not dyn-compatible, so this trait is designed for *generic*
//!   (static) dispatch — see late_chunk in `lc-rag`, which takes the
//!   embedder as a generic parameter. Existing `#[async_trait]` traits
//!   (`Embeddings`, `BaseChatModel`) are NOT migrated.

/// Byte span of a token within the input text (start inclusive, end exclusive).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TokenSpan {
    /// Byte offset of the first byte of the token.
    pub start: usize,
    /// Byte offset one past the last byte of the token.
    pub end: usize,
}

impl TokenSpan {
    /// Creates a new span.
    pub fn new(start: usize, end: usize) -> Self {
        Self { start, end }
    }

    /// Whether the span intersects `[range_start, range_end)`.
    pub fn intersects(&self, range_start: usize, range_end: usize) -> bool {
        self.start < range_end && range_start < self.end
    }
}

/// Converts HuggingFace tokenizers offsets (char positions of the original
/// text) into byte spans for [`TokenSpan`].
///
/// Pure (0.21.0 S5.2): the token-level ONNX path gets char offsets from
/// `tokenizers::Encoding::get_offsets`, while [`TokenSpan`] is byte-based (so
/// late chunking can slice the original string directly). Out-of-range or
/// reversed spans are clamped/skipped defensively rather than panicking.
pub fn char_spans_to_byte_spans(text: &str, offsets: &[(usize, usize)]) -> Vec<TokenSpan> {
    let char_to_byte: Vec<usize> = text.char_indices().map(|(b, _)| b).collect();
    let text_len_chars = char_to_byte.len();
    let mut out = Vec::with_capacity(offsets.len());
    for &(start, end) in offsets {
        if end <= start {
            continue; // empty/reversed span → skip
        }
        let start = start.min(text_len_chars);
        let end = end.min(text_len_chars);
        if end <= start {
            continue;
        }
        let byte_start = char_to_byte[start];
        let byte_end = if end == text_len_chars {
            text.len()
        } else {
            char_to_byte[end]
        };
        out.push(TokenSpan::new(byte_start, byte_end));
    }
    out
}

/// One token's embedding plus its byte span in the input text.
#[derive(Debug, Clone, PartialEq)]
pub struct TokenEmbedding {
    /// Byte span of the token in the input text.
    pub span: TokenSpan,
    /// L2-normalized embedding vector for the token (same contract as
    /// [`crate::Embeddings`]).
    pub vector: Vec<f32>,
}

/// Optional capability trait: token-level (late-chunking-ready) embeddings.
///
/// Implementations must return one [`TokenEmbedding`] per token, in ascending
/// span order, covering the whole input text. Vectors follow the same
/// L2-normalization contract as [`crate::Embeddings`].
#[trait_variant::make(TokenLevelEmbeddings: Send)]
pub trait LocalTokenLevelEmbeddings {
    /// Embeds `text` token by token, returning each token's span and vector.
    ///
    /// Tokenization is provider-specific; `span` is a byte offset into `text`
    /// (so UTF-8 multibyte characters are handled by byte offsets, never char
    /// indices).
    async fn embed_tokens(&self, text: &str) -> Result<Vec<TokenEmbedding>, crate::EmbeddingError>;
}

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

    /// Deterministic mock: tokenizes on whitespace and derives the vector from
    /// the token bytes — enough to exercise trait ergonomics and the
    /// late-chunking math without a real model.
    ///
    /// Implements the generated `Send` variant (`TokenLevelEmbeddings`); the
    /// `!Send` base (`LocalTokenLevelEmbeddings`) is then auto-implemented by
    /// trait_variant's blanket impl.
    struct MockTokenEmbeddings;

    impl TokenLevelEmbeddings for MockTokenEmbeddings {
        async fn embed_tokens(
            &self,
            text: &str,
        ) -> Result<Vec<TokenEmbedding>, crate::EmbeddingError> {
            if text.trim().is_empty() {
                return Err(crate::EmbeddingError::EmptyInput);
            }
            let mut out = Vec::new();
            let mut cursor = 0usize;
            for word in text.split_whitespace() {
                let start = text[cursor..]
                    .find(word)
                    .map(|p| cursor + p)
                    .unwrap_or(cursor);
                let end = start + word.len();
                cursor = end;
                let vector = vec![word.bytes().map(|b| b as f32).sum::<f32>(), 1.0];
                out.push(TokenEmbedding {
                    span: TokenSpan::new(start, end),
                    vector,
                });
            }
            Ok(out)
        }
    }

    /// Generic (static) dispatch helper — method-call syntax is ambiguous for
    /// types implementing both the variant and the blanket-implemented base,
    /// so tests route through generics (which is also the intended usage).
    async fn embed<E: TokenLevelEmbeddings>(
        e: &E,
        text: &str,
    ) -> Result<Vec<TokenEmbedding>, crate::EmbeddingError> {
        e.embed_tokens(text).await
    }

    /// Native async fn in trait (RPITIT) compiles and dispatches statically —
    /// the S3.2 pilot's core assertion.
    #[tokio::test]
    async fn embed_tokens_returns_spans_and_vectors() {
        let embedder = MockTokenEmbeddings;
        let tokens = embed(&embedder, "hello late world").await.unwrap();
        assert_eq!(tokens.len(), 3);
        assert_eq!(tokens[0].span, TokenSpan::new(0, 5));
        assert_eq!(tokens[1].span, TokenSpan::new(6, 10));
        assert_eq!(tokens[2].span, TokenSpan::new(11, 16));
        assert_eq!(
            tokens[0].vector,
            vec![b"hello".iter().map(|b| *b as f32).sum::<f32>(), 1.0]
        );
    }

    /// Generic (static) dispatch over the generated `Send` variant.
    async fn generic_len<E: TokenLevelEmbeddings>(e: &E, text: &str) -> usize {
        e.embed_tokens(text).await.unwrap().len()
    }

    #[tokio::test]
    async fn trait_variant_send_is_generically_usable() {
        let embedder = MockTokenEmbeddings;
        assert_eq!(generic_len(&embedder, "a b c").await, 3);
    }

    /// Spans are byte offsets that survive UTF-8 multibyte content.
    #[tokio::test]
    async fn spans_are_byte_offsets() {
        let embedder = MockTokenEmbeddings;
        // "你好" is 6 bytes; the second token starts after the space at byte 6.
        let tokens = embed(&embedder, "你好 world").await.unwrap();
        assert_eq!(tokens.len(), 2);
        assert_eq!(tokens[0].span, TokenSpan::new(0, 6));
        assert_eq!(tokens[1].span, TokenSpan::new(7, 12));
        // The span slices back into the original text at byte boundaries.
        assert_eq!(
            &"你好 world"[tokens[1].span.start..tokens[1].span.end],
            "world"
        );
    }

    /// Empty input is rejected by contract (same as `Embeddings`).
    #[tokio::test]
    async fn empty_input_rejected() {
        let embedder = MockTokenEmbeddings;
        let err = embed(&embedder, "   ").await.unwrap_err();
        assert!(matches!(err, crate::EmbeddingError::EmptyInput));
    }

    /// `TokenSpan::intersects` drives the chunk-boundary pooling in
    /// [`crate::late_chunk`] — spot-check the boundary semantics.
    #[test]
    fn span_intersects_boundaries() {
        let span = TokenSpan::new(5, 10);
        assert!(span.intersects(4, 6));
        assert!(span.intersects(9, 11));
        assert!(span.intersects(0, 100));
        assert!(!span.intersects(0, 5), "end-exclusive");
        assert!(!span.intersects(10, 20), "start-inclusive");
    }

    /// 0.21.0 S5.2: char offsets → byte spans, surviving multibyte content.
    #[test]
    fn char_spans_to_byte_spans_ascii() {
        let spans = char_spans_to_byte_spans("hello world", &[(0, 5), (6, 11)]);
        assert_eq!(spans, vec![TokenSpan::new(0, 5), TokenSpan::new(6, 11)]);
    }

    /// Multibyte: "你好 world" — 你 is 3 bytes, 好 is 3 bytes; char offsets
    /// (2, 3) (the space.. no — chars: 你=0 好=1 space=2 w=3..). Span (3, 8)
    /// covers "world" starting at char 3 (w) ending at char 8 (past 'd').
    #[test]
    fn char_spans_to_byte_spans_multibyte() {
        let text = "你好 world";
        // char index: 0=你 1=好 2=' ' 3=w 4=o 5=r 6=l 7=d → byte 7..12.
        let spans = char_spans_to_byte_spans(text, &[(3, 8)]);
        assert_eq!(spans, vec![TokenSpan::new(7, 12)]);
        // The byte span slices the original text at a char boundary.
        assert_eq!(&text[spans[0].start..spans[0].end], "world");
    }

    /// Out-of-range and reversed offsets are clamped/skipped, never panicked.
    #[test]
    fn char_spans_to_byte_spans_out_of_range() {
        let spans = char_spans_to_byte_spans("abc", &[(0, 10), (5, 2), (1, 1)]);
        // (0, 10) clamps to the whole text; (5, 2) and (1, 1) skipped.
        assert_eq!(spans.len(), 1);
        assert_eq!(spans[0], TokenSpan::new(0, 3));
    }

    /// Empty input → empty output.
    #[test]
    fn char_spans_to_byte_spans_empty() {
        assert!(char_spans_to_byte_spans("", &[]).is_empty());
        assert!(char_spans_to_byte_spans("abc", &[]).is_empty());
    }
}