aurora_semantic/embeddings/

mod.rs

//! Embeddings module for generating semantic vectors.
//!
//! This module provides functionality to generate embedding vectors
//! from source code chunks using various embedding providers:
//!
//! - **API-based**: Jina AI, OpenAI, Cohere, Voyage AI
//! - **Local**: HuggingFace models via Candle
//! - **Custom**: Implement the `Embedder` trait

pub mod pooling;
mod providers;

pub use pooling::cosine_similarity;
pub use providers::*;

use crate::error::Result;
use crate::types::Chunk;

/// Trait for embedding generators.
///
/// Implement this trait to add support for custom embedding providers.
///
/// # Example
///
/// ```rust,ignore
/// use aurora_semantic::{Embedder, Result};
///
/// struct MyEmbedder {
///     dimension: usize,
/// }
///
/// impl Embedder for MyEmbedder {
///     fn embed(&self, text: &str) -> Result<Vec<f32>> {
///         // Your embedding logic here
///         Ok(vec![0.0; self.dimension])
///     }
///
///     fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
///         texts.iter().map(|t| self.embed(t)).collect()
///     }
///
///     fn dimension(&self) -> usize {
///         self.dimension
///     }
///
///     fn name(&self) -> &'static str {
///         "my-embedder"
///     }
/// }
/// ```
pub trait Embedder: Send + Sync {
    /// Generate an embedding for a single text.
    fn embed(&self, text: &str) -> Result<Vec<f32>>;

    /// Generate embeddings for multiple texts in batch.
    ///
    /// Default implementation calls `embed` for each text.
    /// Override for more efficient batch processing.
    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        texts.iter().map(|t| self.embed(t)).collect()
    }

    /// Generate an embedding for a search query.
    ///
    /// For asymmetric retrieval models (like Jina Code 1.5B), this uses
    /// query-specific instruction prefixes. For symmetric models, this
    /// just calls `embed()`.
    ///
    /// Use this when embedding user search queries.
    fn embed_for_query(&self, text: &str) -> Result<Vec<f32>> {
        // Default: same as embed()
        // JinaCodeEmbedder overrides this to use query prefix
        self.embed(text)
    }

    /// Get the embedding dimension.
    fn dimension(&self) -> usize;

    /// Get the name of this embedder.
    fn name(&self) -> &'static str;

    /// Get the maximum sequence length supported.
    fn max_sequence_length(&self) -> usize {
        512 // Default, override for models with different limits
    }
}

/// Embed chunks and return (chunk_id, embedding) pairs.
#[allow(dead_code)]
pub fn embed_chunks<E: Embedder>(
    embedder: &E,
    chunks: &[Chunk],
    batch_size: usize,
) -> Result<Vec<(String, Vec<f32>)>> {
    let mut results = Vec::with_capacity(chunks.len());

    for batch in chunks.chunks(batch_size) {
        let texts: Vec<&str> = batch.iter().map(|c| c.content.as_str()).collect();
        let embeddings = embedder.embed_batch(&texts)?;

        for (chunk, embedding) in batch.iter().zip(embeddings.into_iter()) {
            results.push((chunk.id.0.to_string(), embedding));
        }
    }

    Ok(results)
}

/// Preprocess code for embedding.
///
/// This normalizes code to improve embedding quality.
pub fn preprocess_code(content: &str) -> String {
    let mut result = String::with_capacity(content.len());

    // Remove excessive whitespace while preserving structure
    let mut prev_was_space = false;
    let mut prev_was_newline = false;

    for c in content.chars() {
        if c == '\n' {
            if !prev_was_newline {
                result.push('\n');
                prev_was_newline = true;
            }
            prev_was_space = false;
        } else if c.is_whitespace() {
            if !prev_was_space && !prev_was_newline {
                result.push(' ');
                prev_was_space = true;
            }
        } else {
            result.push(c);
            prev_was_space = false;
            prev_was_newline = false;
        }
    }

    result.trim().to_string()
}

/// Create a context-enriched text for embedding.
///
/// This adds metadata to help the embedding model understand the code context.
pub fn create_embedding_text(chunk: &Chunk) -> String {
    let mut text = String::new();

    // Add chunk type context
    let type_name = match chunk.chunk_type {
        crate::types::ChunkType::Function => "function",
        crate::types::ChunkType::Class => "class",
        crate::types::ChunkType::Struct => "struct",
        crate::types::ChunkType::Enum => "enum",
        crate::types::ChunkType::Interface => "interface",
        crate::types::ChunkType::Implementation => "implementation",
        crate::types::ChunkType::Module => "module",
        crate::types::ChunkType::Imports => "imports",
        crate::types::ChunkType::Constant => "constant",
        crate::types::ChunkType::TypeDef => "type definition",
        crate::types::ChunkType::Block => "code block",
        crate::types::ChunkType::Comment => "documentation",
    };

    // Add symbol name if available
    if let Some(ref name) = chunk.symbol_name {
        text.push_str(&format!("{} {} ", type_name, name));
    } else {
        text.push_str(&format!("{} ", type_name));
    }

    // Add parent context if available
    if let Some(ref parent) = chunk.parent_symbol {
        text.push_str(&format!("in {} ", parent));
    }

    // Add the actual code
    text.push_str(&preprocess_code(&chunk.content));

    text
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{ChunkId, ChunkType, DocumentId};

    #[test]
    fn test_preprocess_code() {
        let input = "fn   foo()  {\n\n\n    bar()\n}";
        let output = preprocess_code(input);
        assert!(!output.contains("  ")); // No double spaces
        assert!(!output.contains("\n\n")); // No double newlines
    }

    #[test]
    fn test_create_embedding_text() {
        let chunk = Chunk {
            id: ChunkId::new(),
            document_id: DocumentId::new(),
            content: "fn add(a: i32, b: i32) -> i32 { a + b }".to_string(),
            chunk_type: ChunkType::Function,
            start_line: 1,
            end_line: 1,
            start_byte: 0,
            end_byte: 38,
            symbol_name: Some("add".to_string()),
            parent_symbol: None,
        };

        let text = create_embedding_text(&chunk);
        assert!(text.starts_with("function add"));
        assert!(text.contains("fn add"));
    }
}