agcore/src/llm/embedding.rs

//! Embedding 抽象 —— 文本向量化接口。
//!
//! 提供 [`Embedding`] trait 和零依赖的 [`MockEmbedding`] 引用实现。
//! 上层可实现此 trait 以对接真实 Embedding Provider（OpenAI、Cohere 等）。
//!
//! 所有实现使用 [`LlmError`] 作为统一错误类型，与 llm 模块保持一致。

use async_trait::async_trait;

use crate::llm::error::LlmError;

/// 文本向量化抽象接口。
///
/// 将文本字符串转换为固定维度的浮点向量，用于语义相似度计算。
/// 设计为异步以支持网络 IO（如 OpenAI Embedding API）。
///
/// 使用 [`LlmError`] 作为统一错误类型，与 llm 模块保持一致。
///
/// # 实现要求
///
/// - `embed()` 返回的向量外层的 Vec 长度必须等于输入切片长度（一对一映射）
/// - 内层 Vec 长度必须等于 `dim()` 返回值
/// - 调用方应保证输入非空（空切片返回空外层 Vec，不报错）
///
/// # 稳定性
///
/// 实验性 API（v0.3.x），方法签名可能在 v0.4 中调整。
#[async_trait]
pub trait Embedding: Send + Sync {
    /// 批量向量化。
    ///
    /// 返回 `Vec<Vec<f32>>`，第 i 个内层向量对应 `input[i]`。
    async fn embed(&self, input: &[String]) -> Result<Vec<Vec<f32>>, LlmError>;

    /// 返回向量维度。
    fn dim(&self) -> usize;
}

/// 确定性 Mock Embedding —— 零依赖伪随机单位向量。
///
/// 使用 sin 哈希将输入字符串映射到单位球面上的一个点：
/// 1. 对输入字符串计算简单哈希（字符字节和 + 长度）作为种子
/// 2. 用 `f32::sin(seed + i) * 10000` 生成第 i 个维度的值
/// 3. 归一化到单位长度（L2 norm = 1.0）
///
/// 特性：
/// - **确定性**：相同输入 → 相同向量
/// - **有区分度**：不同输入产生不同向量（高概率）
/// - **单位范数**：余弦相似度等价于点积
/// - **开销极低**：不分配额外内存，无 IO
///
/// # 已知限制
///
/// `f32::sin(seed + i) * 10000` 在维度较高时（如 1536，OpenAI Embedding 维度）
/// 可能出现周期性模式——相邻维度取值在 `sin` 周期 2π 约束下呈规律性重复。
/// MockEmbedding 仅用于测试验证，**不应用于生产级相似度排序**；
/// 做严肃验证时建议使用真实 Embedding Provider 或显式随机初始化。
pub struct MockEmbedding {
    dim: usize,
}

impl MockEmbedding {
    /// 创建 Mock Embedding，输出向量维度为 `dim`。
    pub fn new(dim: usize) -> Self {
        Self { dim }
    }
}

#[async_trait]
impl Embedding for MockEmbedding {
    async fn embed(&self, input: &[String]) -> Result<Vec<Vec<f32>>, LlmError> {
        let results: Vec<Vec<f32>> = input
            .iter()
            .map(|text| {
                // 简单哈希：字符字节值和 + 文本长度作为种子
                let seed: f64 = text.bytes().map(|b| b as f64).sum::<f64>() + text.len() as f64;
                let mut vec: Vec<f32> = (0..self.dim)
                    .map(|i| f32::sin(seed as f32 + i as f32) * 10000.0)
                    .collect();
                l2_normalize(&mut vec);
                vec
            })
            .collect();
        Ok(results)
    }

    fn dim(&self) -> usize {
        self.dim
    }
}

/// L2 归一化（in-place）。
///
/// 零向量（norm == 0）保持全零 —— 防除零保护。
fn l2_normalize(vec: &mut [f32]) {
    let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > f32::EPSILON {
        for x in vec.iter_mut() {
            *x /= norm;
        }
    }
}

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

    /// 计算向量的 L2 范数。
    fn l2_norm(v: &[f32]) -> f32 {
        v.iter().map(|x| x * x).sum::<f32>().sqrt()
    }

    #[tokio::test]
    async fn embed_correct_dim() {
        let embedder = MockEmbedding::new(8);
        let inputs = vec!["hello".to_string(), "world".to_string()];
        let result = embedder.embed(&inputs).await.unwrap();
        assert_eq!(result.len(), 2);
        for vec in &result {
            assert_eq!(vec.len(), 8);
        }
    }

    #[tokio::test]
    async fn embed_batch_size_match() {
        let embedder = MockEmbedding::new(4);
        let inputs = vec![
            "a".to_string(),
            "b".to_string(),
            "c".to_string(),
            "d".to_string(),
            "e".to_string(),
        ];
        let result = embedder.embed(&inputs).await.unwrap();
        assert_eq!(result.len(), inputs.len());
    }

    #[tokio::test]
    async fn embed_deterministic() {
        let embedder = MockEmbedding::new(4);
        let inputs = vec!["deterministic test".to_string()];
        let r1 = embedder.embed(&inputs).await.unwrap();
        let r2 = embedder.embed(&inputs).await.unwrap();
        assert_eq!(r1, r2);
    }

    #[tokio::test]
    async fn embed_unit_vector_norm() {
        let embedder = MockEmbedding::new(16);
        let inputs = vec!["any text".to_string(), "another".to_string()];
        let result = embedder.embed(&inputs).await.unwrap();
        for vec in &result {
            let norm = l2_norm(vec);
            assert!(
                (norm - 1.0).abs() < 1e-5,
                "vector norm should be ~1.0, got {}",
                norm
            );
        }
    }

    #[tokio::test]
    async fn embed_different_inputs_different_vectors() {
        let embedder = MockEmbedding::new(16);
        let r1 = embedder.embed(&["hello world".to_string()]).await.unwrap();
        let r2 = embedder
            .embed(&["completely different".to_string()])
            .await
            .unwrap();
        assert_ne!(r1, r2);
    }

    #[tokio::test]
    async fn embed_empty_string() {
        // 空字符串输入应不 panic，且向量范数仍≈1.0（防除零路径）
        let embedder = MockEmbedding::new(4);
        let inputs = vec!["".to_string()];
        let result = embedder.embed(&inputs).await.unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].len(), 4);
        let norm = l2_norm(&result[0]);
        assert!(
            (norm - 1.0).abs() < 1e-5,
            "empty-string vector norm should be ~1.0, got {}",
            norm
        );
    }
}