a3s_vec/
embedding.rs

1//! Caller-owned embedding interfaces.
2//!
3//! The database never downloads a model or performs network I/O implicitly.
4//! Applications can opt into these traits when they already own an embedding
5//! provider, which keeps the core deterministic and easy to run on the
6//! supported Linux, Windows, and macOS platforms.
7
8use crate::error::Result;
9
10/// Input accepted by an embedding provider.
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub enum EmbeddingInput {
13    Text(String),
14    Document(String),
15}
16
17impl From<&str> for EmbeddingInput {
18    fn from(value: &str) -> Self {
19        Self::Text(value.to_string())
20    }
21}
22
23impl From<String> for EmbeddingInput {
24    fn from(value: String) -> Self {
25        Self::Text(value)
26    }
27}
28
29/// Dense embedding provider.  Implementations should be pure with respect to
30/// the collection; retries and network policy belong in the adapter.
31pub trait DenseEmbedding: Send + Sync {
32    fn embed(&self, input: &EmbeddingInput) -> Result<Vec<f32>>;
33}
34
35/// Sparse embedding provider.
36pub trait SparseEmbedding: Send + Sync {
37    fn embed_sparse(&self, input: &EmbeddingInput) -> Result<Vec<(u32, f32)>>;
38}
39
40/// Optional convenience for executing a query after embedding text.
41pub trait QueryExecutor: Send + Sync {
42    fn execute_text(
43        &self,
44        input: &EmbeddingInput,
45        field_name: &str,
46        topk: usize,
47    ) -> Result<Vec<crate::Doc>>;
48}
49
50#[cfg(test)]
51#[allow(clippy::cast_precision_loss)]
52mod tests {
53    use super::{DenseEmbedding, EmbeddingInput, QueryExecutor, SparseEmbedding};
54    use crate::error::{Error, Result};
55    use crate::Doc;
56
57    struct FixedDense;
58
59    impl DenseEmbedding for FixedDense {
60        fn embed(&self, input: &EmbeddingInput) -> Result<Vec<f32>> {
61            match input {
62                EmbeddingInput::Text(text) | EmbeddingInput::Document(text) => {
63                    Ok(vec![text.len() as f32, 1.0])
64                }
65            }
66        }
67    }
68
69    struct FixedSparse;
70
71    impl SparseEmbedding for FixedSparse {
72        fn embed_sparse(&self, input: &EmbeddingInput) -> Result<Vec<(u32, f32)>> {
73            match input {
74                EmbeddingInput::Text(text) | EmbeddingInput::Document(text) => {
75                    Ok(vec![(0, text.len() as f32)])
76                }
77            }
78        }
79    }
80
81    struct EchoExecutor;
82
83    impl QueryExecutor for EchoExecutor {
84        fn execute_text(
85            &self,
86            input: &EmbeddingInput,
87            field_name: &str,
88            topk: usize,
89        ) -> Result<Vec<Doc>> {
90            if field_name.is_empty() || topk == 0 {
91                return Err(Error::invalid_argument("field/topk"));
92            }
93            let _ = input;
94            Ok(Vec::new())
95        }
96    }
97
98    #[test]
99    fn embedding_input_and_provider_traits_are_callable() {
100        let from_str: EmbeddingInput = "hello".into();
101        let from_string: EmbeddingInput = String::from("world").into();
102        assert_eq!(from_str, EmbeddingInput::Text("hello".into()));
103        assert_eq!(from_string, EmbeddingInput::Text("world".into()));
104        assert_eq!(
105            FixedDense
106                .embed(&EmbeddingInput::Document("ab".into()))
107                .unwrap(),
108            vec![2.0, 1.0]
109        );
110        assert_eq!(
111            FixedSparse
112                .embed_sparse(&EmbeddingInput::Text("xyz".into()))
113                .unwrap(),
114            vec![(0, 3.0)]
115        );
116        assert!(EchoExecutor
117            .execute_text(&EmbeddingInput::Text("q".into()), "embedding", 3)
118            .unwrap()
119            .is_empty());
120        assert!(EchoExecutor
121            .execute_text(&EmbeddingInput::Text("q".into()), "", 3)
122            .is_err());
123    }
124}
a3s_vec/embedding.rs

a3s_vec/
embedding.rs
