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 std::sync::Arc;

use anyhow::Result;
use tracing::warn;

use crate::engine::embedder::Embedder;
use crate::storage::vector_store::VectorStore;

/// Dense-vector search engine.
///
/// Wraps an embedder and an in-memory vector store. Callers index documents
/// as they are added to the BM25 index, then search with a query string.
pub struct VectorEngine {
    embedder: Arc<dyn Embedder>,
    store: VectorStore,
}

impl VectorEngine {
    pub fn new(embedder: Arc<dyn Embedder>) -> Self {
        Self {
            embedder,
            store: VectorStore::new(),
        }
    }

    /// Embed and store a single document.
    pub fn index(&self, id: &str, text: &str) -> Result<()> {
        let text = if text.trim().is_empty() { id } else { text };
        let vectors = self.embedder.embed(&[text])?;
        if let Some(vec) = vectors.into_iter().next() {
            self.store.insert(id.to_string(), vec);
        }
        Ok(())
    }

    /// Batch index several documents.
    pub fn index_batch(&self, items: &[(String, String)]) -> Result<()> {
        if items.is_empty() {
            return Ok(());
        }
        let inputs: Vec<&str> = items
            .iter()
            .map(|(id, text)| {
                if text.trim().is_empty() {
                    id.as_str()
                } else {
                    text.as_str()
                }
            })
            .collect();
        let vectors = self.embedder.embed(&inputs)?;
        for ((id, _), vec) in items.iter().zip(vectors) {
            self.store.insert(id.clone(), vec);
        }
        Ok(())
    }

    /// Search for the top-k most similar documents and return a URL -> score map.
    pub fn search(&self, query: &str, top_k: usize) -> Result<HashMap<String, f64>> {
        let qvec = self.embedder.embed(&[query])?;
        let Some(query_vec) = qvec.into_iter().next() else {
            return Ok(HashMap::new());
        };
        if self.store.is_empty() {
            warn!("vector search requested but no vectors are indexed");
            return Ok(HashMap::new());
        }
        Ok(self.store.search(&query_vec, top_k).into_iter().collect())
    }

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

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::engine::embedder::DummyEmbedder;

    #[test]
    fn test_vector_engine_roundtrip() {
        let engine = VectorEngine::new(Arc::new(DummyEmbedder));
        engine.index("a", "the quick brown fox").unwrap();
        engine.index("b", "lazy dog sleeping").unwrap();

        let hits = engine.search("quick fox", 5).unwrap();
        assert!(hits.contains_key("a"));
        assert!(hits.get("a").unwrap() > hits.get("b").unwrap());
    }
}