llm_brain/

embeddings.rs

use std::sync::OnceLock;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tiktoken_rs::{CoreBPE, cl100k_base};

use crate::error::{LLMBrainError, Result};

/// Cached tokenizer instance for improved performance
static BPE: OnceLock<CoreBPE> = OnceLock::new();

/// Vector embedding provider trait
///
/// Structs implementing this trait can generate vector embeddings for text
#[async_trait]
pub trait EmbeddingProvider: Send + Sync {
    /// Generate embedding vector for a single text
    async fn generate_embedding(&self, text: &str) -> Result<Vec<f32>>;

    /// Generate embedding vectors for multiple texts
    async fn generate_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>>;

    /// Count tokens in text
    fn count_tokens(&self, text: &str) -> Result<usize>;
}

/// Embedding model configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingModelConfig {
    /// Embedding model name
    pub model_name: String,

    /// Embedding vector dimensions
    pub dimensions: usize,

    /// Maximum context window size (in tokens)
    pub max_context_length: usize,
}

impl Default for EmbeddingModelConfig {
    fn default() -> Self {
        Self {
            model_name: "text-embedding-3-small".to_owned(),
            dimensions: 1536,
            max_context_length: 8191,
        }
    }
}

/// Embedding generation middleware
///
/// This struct provides pre-processing and post-processing for embedding
/// requests, such as:
/// - Text normalization
/// - Embedding caching
/// - Batch processing optimization
/// - Vector normalization
pub struct EmbeddingMiddleware<P: EmbeddingProvider> {
    provider: P,
    normalize_vectors: bool,
}

impl<P: EmbeddingProvider> EmbeddingMiddleware<P> {
    /// Create a new embedding middleware instance
    pub fn new(provider: P, normalize_vectors: bool) -> Self {
        Self {
            provider,
            normalize_vectors,
        }
    }

    /// Initialize tokenizer
    pub fn initialize_tokenizer() -> Result<()> {
        BPE.get_or_init(|| cl100k_base().expect("Failed to load cl100k_base tokenizer"));
        Ok(())
    }

    /// Normalize text content
    pub fn normalize_text(&self, text: &str) -> String {
        // Basic text normalization: trim whitespace, merge multiple spaces into one
        text.trim().to_owned()
    }

    /// Normalize vector to make it a unit vector
    pub fn normalize_vector(&self, vector: &mut [f32]) {
        if !self.normalize_vectors {
            return;
        }

        // Calculate Euclidean norm of the vector
        let norm = vector.iter().map(|&x| x * x).sum::<f32>().sqrt();

        // Avoid division by zero
        if norm > 1e-10 {
            for x in vector.iter_mut() {
                *x /= norm;
            }
        }
    }
}

#[async_trait]
impl<P: EmbeddingProvider + Send + Sync> EmbeddingProvider for EmbeddingMiddleware<P> {
    async fn generate_embedding(&self, text: &str) -> Result<Vec<f32>> {
        let normalized_text = self.normalize_text(text);
        let mut embedding = self.provider.generate_embedding(&normalized_text).await?;
        self.normalize_vector(&mut embedding);
        Ok(embedding)
    }

    async fn generate_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>> {
        let normalized_texts = texts
            .iter()
            .map(|text| self.normalize_text(text))
            .collect::<Vec<String>>();

        let mut embeddings = self.provider.generate_embeddings(normalized_texts).await?;

        for embedding in &mut embeddings {
            self.normalize_vector(embedding);
        }

        Ok(embeddings)
    }

    fn count_tokens(&self, text: &str) -> Result<usize> {
        self.provider.count_tokens(text)
    }
}

/// Text chunk-based embedding strategy
///
/// Defines how to handle long text embeddings: splitting, averaging, using
/// specific parts, etc.
#[derive(Default)]
pub enum ChunkingStrategy {
    /// Maximize use of model context window, no chunking needed
    #[default]
    NoChunking,

    /// Split long text into chunks, generate embeddings for each chunk, then
    /// average
    ChunkAndAverage {
        chunk_size: usize,
        chunk_overlap: usize,
    },

    /// Only use the first N tokens of the text
    UsePrefix(usize),

    /// Only use the last N tokens of the text
    UseSuffix(usize),
}

/// Long text embedding handler
pub struct LongTextHandler<P: EmbeddingProvider> {
    provider: P,
    model_config: EmbeddingModelConfig,
    chunking_strategy: ChunkingStrategy,
}

impl<P: EmbeddingProvider> LongTextHandler<P> {
    /// Create a new long text handler
    pub fn new(
        provider: P, model_config: EmbeddingModelConfig, chunking_strategy: ChunkingStrategy,
    ) -> Self {
        Self {
            provider,
            model_config,
            chunking_strategy,
        }
    }

    /// Truncate text to maximum context window size
    pub fn truncate_text(&self, text: &str) -> Result<String> {
        let token_count = self.provider.count_tokens(text)?;

        if token_count <= self.model_config.max_context_length {
            return Ok(text.to_owned());
        }

        // Simple truncation implementation - in practice, can be more precisely
        // truncated at token boundaries
        let bpe = BPE.get().expect("BPE Tokenizer not initialized");
        let tokens = bpe.encode_with_special_tokens(text);
        let truncated_tokens = tokens[0..self.model_config.max_context_length].to_vec();

        // Decode back to text
        bpe.decode(truncated_tokens)
            .map_err(|e| LLMBrainError::InputError(format!("Failed to decode tokens: {e}")))
    }

    /// Chunk text into pieces
    pub fn chunk_text(&self, text: &str, chunk_size: usize, overlap: usize) -> Result<Vec<String>> {
        let bpe = BPE.get().expect("BPE Tokenizer not initialized");
        let tokens = bpe.encode_with_special_tokens(text);

        if tokens.len() <= chunk_size {
            return Ok(vec![text.to_owned()]);
        }

        let mut chunks = Vec::new();
        let mut start = 0;

        while start < tokens.len() {
            let end = (start + chunk_size).min(tokens.len());
            let chunk_tokens = tokens[start..end].to_vec();

            // Decode the current chunk
            let chunk = bpe
                .decode(chunk_tokens)
                .map_err(|e| LLMBrainError::InputError(format!("Failed to decode tokens: {e}")))?;

            chunks.push(chunk);

            if end >= tokens.len() {
                break;
            }

            // Calculate the starting position of the next chunk, considering overlap
            start = end - overlap;
        }

        Ok(chunks)
    }

    /// Process embedding vectors
    pub async fn process_embeddings(&self, text: &str) -> Result<Vec<f32>> {
        match &self.chunking_strategy {
            ChunkingStrategy::NoChunking => {
                let truncated = self.truncate_text(text)?;
                self.provider.generate_embedding(&truncated).await
            }

            ChunkingStrategy::ChunkAndAverage {
                chunk_size,
                chunk_overlap,
            } => {
                let chunks = self.chunk_text(text, *chunk_size, *chunk_overlap)?;
                if chunks.is_empty() {
                    return Err(LLMBrainError::InputError(
                        "No chunks generated from text".to_owned(),
                    ));
                }

                let embeddings = self.provider.generate_embeddings(chunks).await?;

                // Average all embedding vectors
                if embeddings.is_empty() {
                    return Err(LLMBrainError::ApiError(
                        "No embeddings generated".to_owned(),
                    ));
                }

                let dimensions = embeddings[0].len();
                let mut average = vec![0.0; dimensions];

                for embedding in &embeddings {
                    for (i, &value) in embedding.iter().enumerate() {
                        average[i] += value / embeddings.len() as f32;
                    }
                }

                Ok(average)
            }

            ChunkingStrategy::UsePrefix(size) => {
                let bpe = BPE.get().expect("BPE Tokenizer not initialized");
                let tokens = bpe.encode_with_special_tokens(text);

                let prefix_tokens = tokens.iter().take(*size).cloned().collect::<Vec<_>>();

                let prefix = bpe.decode(prefix_tokens).map_err(|e| {
                    LLMBrainError::InputError(format!("Failed to decode tokens: {e}"))
                })?;

                self.provider.generate_embedding(&prefix).await
            }

            ChunkingStrategy::UseSuffix(size) => {
                let bpe = BPE.get().expect("BPE Tokenizer not initialized");
                let tokens = bpe.encode_with_special_tokens(text);

                let suffix_tokens = tokens.iter().rev().take(*size).cloned().collect::<Vec<_>>();

                let suffix = bpe.decode(suffix_tokens).map_err(|e| {
                    LLMBrainError::InputError(format!("Failed to decode tokens: {e}"))
                })?;

                self.provider.generate_embedding(&suffix).await
            }
        }
    }
}

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

    // Mock embedding provider implementation for testing
    struct MockEmbeddingProvider;

    #[async_trait]
    impl EmbeddingProvider for MockEmbeddingProvider {
        async fn generate_embedding(&self, _text: &str) -> Result<Vec<f32>> {
            // Return a fixed test embedding vector
            Ok(vec![0.1, 0.2, 0.3, 0.4])
        }

        async fn generate_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>> {
            // Generate test embeddings for each input text
            let mut result = Vec::new();
            for _ in 0..texts.len() {
                result.push(vec![0.1, 0.2, 0.3, 0.4]);
            }
            Ok(result)
        }

        fn count_tokens(&self, text: &str) -> Result<usize> {
            // Simple estimation: assume each word is one token
            Ok(text.split_whitespace().count())
        }
    }

    #[tokio::test]
    async fn test_embedding_middleware() {
        let provider = MockEmbeddingProvider;
        let middleware = EmbeddingMiddleware::new(provider, true);

        // Test vector normalization
        let embedding = middleware.generate_embedding("test text").await.unwrap();

        // Calculate the length of the normalized vector (should be close to 1.0)
        let norm = embedding.iter().map(|&x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 1e-6);
    }
}