webfind 0.5.0

Self-hosted native search engine for AI agents — pure CLI, zero cost
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use std::collections::HashMap;

use dashmap::DashMap;

/// In-memory dense vector index with brute-force cosine search.
///
/// This is intentionally simple: it is correct for small-to-medium indexes and
/// avoids the complexity and build-time cost of an HNSW dependency. Swap in
/// `hnsw_rs` later if profiling shows it is needed.
pub struct VectorStore {
    vectors: DashMap<String, Vec<f32>>,
}

impl VectorStore {
    pub fn new() -> Self {
        Self {
            vectors: DashMap::new(),
        }
    }

    pub fn insert(&self, id: String, vector: Vec<f32>) {
        self.vectors.insert(id, vector);
    }

    pub fn len(&self) -> usize {
        self.vectors.len()
    }

    pub fn is_empty(&self) -> bool {
        self.vectors.is_empty()
    }

    /// Search for the `top_k` most similar vectors using cosine similarity.
    pub fn search(&self, query: &[f32], top_k: usize) -> Vec<(String, f64)> {
        let qnorm = norm(query).max(1e-8);
        let mut scored: Vec<(String, f64)> = self
            .vectors
            .iter()
            .map(|entry| {
                let id = entry.key().clone();
                let vec = entry.value();
                let score = cosine(query, vec, qnorm);
                (id, score)
            })
            .collect();
        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        scored.truncate(top_k);
        scored
    }

    /// Return a copy of all stored vectors.
    pub fn all(&self) -> HashMap<String, Vec<f32>> {
        self.vectors
            .iter()
            .map(|entry| (entry.key().clone(), entry.value().clone()))
            .collect()
    }
}

impl Default for VectorStore {
    fn default() -> Self {
        Self::new()
    }
}

fn cosine(a: &[f32], b: &[f32], a_norm: f64) -> f64 {
    let min_len = a.len().min(b.len());
    let dot: f64 = (0..min_len).map(|i| (a[i] as f64) * (b[i] as f64)).sum();
    let b_norm = norm(b).max(1e-8);
    dot / (a_norm * b_norm)
}

fn norm(v: &[f32]) -> f64 {
    v.iter()
        .map(|x| (*x as f64) * (*x as f64))
        .sum::<f64>()
        .sqrt()
}

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

    #[test]
    fn test_vector_store_cosine() {
        let store = VectorStore::new();
        store.insert("a".to_string(), vec![1.0, 0.0, 0.0]);
        store.insert("b".to_string(), vec![0.0, 1.0, 0.0]);
        store.insert("c".to_string(), vec![1.0, 1.0, 0.0]);

        let results = store.search(&[1.0, 0.0, 0.0], 2);
        assert_eq!(results[0].0, "a");
        assert!(results[0].1 > 0.99);
        assert_eq!(results.len(), 2);
    }
}