naptrace_embed/
lib.rs

1pub mod local;
2pub mod ollama;
3pub mod voyage;
4
5use anyhow::Result;
6
7/// Boxed future type alias to avoid clippy::type_complexity.
8pub type EmbedFuture<'a> =
9    std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<Vec<f32>>>> + Send + 'a>>;
10
11/// Trait for embedding backends.
12pub trait Embedder: Send + Sync {
13    fn embed(&self, texts: &[String]) -> EmbedFuture<'_>;
14
15    fn dimension(&self) -> usize;
16}
17
18pub fn has_voyage_key() -> bool {
19    std::env::var("VOYAGE_API_KEY").is_ok()
20}
21
22/// Create the best available embedder.
23/// Prefers Voyage if VOYAGE_API_KEY is set, otherwise uses Ollama.
24pub fn create_embedder() -> Box<dyn Embedder> {
25    if has_voyage_key() {
26        Box::new(voyage::VoyageEmbedder::from_env())
27    } else {
28        Box::new(ollama::OllamaEmbedder::from_env())
29    }
30}
31
32/// Cosine similarity between two vectors.
33pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
34    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
35    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
36    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
37    if norm_a == 0.0 || norm_b == 0.0 {
38        return 0.0;
39    }
40    dot / (norm_a * norm_b)
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46
47    #[test]
48    fn cosine_identical() {
49        let v = vec![1.0, 2.0, 3.0];
50        let sim = cosine_similarity(&v, &v);
51        assert!((sim - 1.0).abs() < 1e-6);
52    }
53
54    #[test]
55    fn cosine_orthogonal() {
56        let a = vec![1.0, 0.0];
57        let b = vec![0.0, 1.0];
58        let sim = cosine_similarity(&a, &b);
59        assert!(sim.abs() < 1e-6);
60    }
61
62    #[test]
63    fn cosine_zero_vector() {
64        let a = vec![1.0, 2.0];
65        let b = vec![0.0, 0.0];
66        assert_eq!(cosine_similarity(&a, &b), 0.0);
67    }
68}
naptrace_embed/lib.rs

naptrace_embed/
lib.rs
