rig_surrealdb/
lib.rs

1//! SurrealDB vector store integration for Rig.
2//!
3//! This crate provides [`SurrealVectorStore`], a Rig vector store backed by
4//! SurrealDB. It supports local in-memory and remote WebSocket connections
5//! through the re-exported SurrealDB engine types.
6//!
7//! The root `rig` facade re-exports this crate as `rig::surrealdb` when the
8//! `surrealdb` feature is enabled.
9
10use std::fmt::Display;
11
12use rig_core::{
13    Embed, OneOrMany,
14    embeddings::{Embedding, EmbeddingModel},
15    vector_store::{
16        InsertDocuments, TopNResults, VectorStoreError, VectorStoreIndex, VectorStoreIndexDyn,
17        request::{Filter, FilterError, SearchFilter, VectorSearchRequest},
18    },
19    wasm_compat::WasmBoxedFuture,
20};
21use serde::{Deserialize, Serialize, de::DeserializeOwned};
22use surrealdb::{
23    Connection, Surreal,
24    types::{RecordId, RecordIdKey, SurrealValue, ToSql, Value},
25};
26
27pub use surrealdb::engine::local::Mem;
28pub use surrealdb::engine::remote::ws::{Ws, Wss};
29
30pub struct SurrealVectorStore<C, Model>
31where
32    C: Connection,
33    Model: EmbeddingModel,
34{
35    model: Model,
36    surreal: Surreal<C>,
37    documents_table: String,
38    distance_function: SurrealDistanceFunction,
39}
40
41/// SurrealDB supported distances
42pub enum SurrealDistanceFunction {
43    Knn,
44    Hamming,
45    Euclidean,
46    Cosine,
47    Jaccard,
48}
49
50impl Display for SurrealDistanceFunction {
51    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
52        match self {
53            SurrealDistanceFunction::Cosine => write!(f, "vector::similarity::cosine"),
54            SurrealDistanceFunction::Knn => write!(f, "vector::distance::knn"),
55            SurrealDistanceFunction::Euclidean => write!(f, "vector::distance::euclidean"),
56            SurrealDistanceFunction::Hamming => write!(f, "vector::distance::hamming"),
57            SurrealDistanceFunction::Jaccard => write!(f, "vector::similarity::jaccard"),
58        }
59    }
60}
61
62#[derive(Debug, Deserialize, SurrealValue)]
63struct SearchResult {
64    id: RecordId,
65    document: String,
66    distance: f64,
67}
68
69#[derive(Debug, Serialize, Deserialize, SurrealValue)]
70pub struct CreateRecord {
71    document: String,
72    embedded_text: String,
73    embedding: Vec<f64>,
74}
75
76#[derive(Debug, Deserialize, SurrealValue)]
77pub struct SearchResultOnlyId {
78    id: RecordId,
79    distance: f64,
80}
81
82impl SearchResult {
83    pub fn into_result<T: DeserializeOwned>(self) -> Result<(f64, String, T), VectorStoreError> {
84        let document: T =
85            serde_json::from_str(&self.document).map_err(VectorStoreError::JsonError)?;
86
87        Ok((self.distance, record_key_to_string(&self.id.key), document))
88    }
89}
90
91fn record_key_to_string(key: &RecordIdKey) -> String {
92    match key {
93        RecordIdKey::Number(value) => value.to_string(),
94        RecordIdKey::String(value) => value.clone(),
95        RecordIdKey::Uuid(value) => value.to_string(),
96        RecordIdKey::Array(_) | RecordIdKey::Object(_) | RecordIdKey::Range(_) => key.to_sql(),
97    }
98}
99
100impl<C, Model> InsertDocuments for SurrealVectorStore<C, Model>
101where
102    C: Connection + Send + Sync,
103    Model: EmbeddingModel + Send + Sync,
104{
105    async fn insert_documents<Doc: Serialize + Embed + Send>(
106        &self,
107        documents: Vec<(Doc, OneOrMany<Embedding>)>,
108    ) -> Result<(), VectorStoreError> {
109        for (document, embeddings) in documents {
110            let json_document: serde_json::Value =
111                serde_json::to_value(&document).map_err(VectorStoreError::JsonError)?;
112            let json_document_as_string =
113                serde_json::to_string(&json_document).map_err(VectorStoreError::JsonError)?;
114
115            for embedding in embeddings {
116                let embedded_text = embedding.document;
117                let embedding: Vec<f64> = embedding.vec;
118
119                let record = CreateRecord {
120                    document: json_document_as_string.clone(),
121                    embedded_text,
122                    embedding,
123                };
124
125                self.surreal
126                    .create::<Option<CreateRecord>>(self.documents_table.clone())
127                    .content(record)
128                    .await
129                    .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?;
130            }
131        }
132
133        Ok(())
134    }
135}
136
137#[derive(Debug, Clone, Serialize, Deserialize)]
138pub struct SurrealSearchFilter(String);
139
140impl SurrealSearchFilter {
141    fn inner(self) -> String {
142        self.0
143    }
144}
145
146impl TryFrom<Filter<serde_json::Value>> for SurrealSearchFilter {
147    type Error = FilterError;
148
149    fn try_from(value: Filter<serde_json::Value>) -> Result<Self, Self::Error> {
150        match value {
151            Filter::Eq(key, value) => Ok(Self::eq(key, Value::from_t(value))),
152            Filter::Gt(key, value) => Ok(Self::gt(key, Value::from_t(value))),
153            Filter::Lt(key, value) => Ok(Self::lt(key, Value::from_t(value))),
154            Filter::And(lhs, rhs) => Ok(Self::try_from(*lhs)?.and(Self::try_from(*rhs)?)),
155            Filter::Or(lhs, rhs) => Ok(Self::try_from(*lhs)?.or(Self::try_from(*rhs)?)),
156        }
157    }
158}
159
160impl std::fmt::Display for SurrealSearchFilter {
161    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162        write!(f, "{}", self.0)
163    }
164}
165
166impl SearchFilter for SurrealSearchFilter {
167    type Value = Value;
168
169    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self {
170        Self(format!("{} = {}", key.as_ref(), value.to_sql()))
171    }
172
173    fn gt(key: impl AsRef<str>, value: Self::Value) -> Self {
174        Self(format!("{} > {}", key.as_ref(), value.to_sql()))
175    }
176
177    fn lt(key: impl AsRef<str>, value: Self::Value) -> Self {
178        Self(format!("{} < {}", key.as_ref(), value.to_sql()))
179    }
180
181    fn and(self, rhs: Self) -> Self {
182        Self(format!("({self}) AND ({rhs})"))
183    }
184
185    fn or(self, rhs: Self) -> Self {
186        Self(format!("({self}) OR ({rhs})"))
187    }
188}
189
190impl SurrealSearchFilter {
191    #[allow(clippy::should_implement_trait)]
192    pub fn not(self) -> Self {
193        Self(format!("NOT ({self})"))
194    }
195
196    /// Test if the value at `key` contains `val`
197    pub fn contains(key: String, val: <Self as SearchFilter>::Value) -> Self {
198        Self(format!("{key} CONTAINS {}", val.to_sql()))
199    }
200
201    /// Test if the value at `key` does *not* contain `val`
202    pub fn does_not_contain(key: String, val: <Self as SearchFilter>::Value) -> Self {
203        Self(format!("{key} CONTAINSNOT {}", val.to_sql()))
204    }
205
206    /// Test if the value at `key` contains every element of `vals`
207    /// `vals` should be a SurrealDB collection
208    pub fn all(key: String, vals: <Self as SearchFilter>::Value) -> Self {
209        Self(format!("{key} CONTAINSALL {}", vals.to_sql()))
210    }
211
212    /// Test if the value at `key` contains any elements of `vals`
213    /// `vals` should be a SurrealDB collection
214    pub fn any(key: String, vals: <Self as SearchFilter>::Value) -> Self {
215        Self(format!("{key} CONTAINSANY {}", vals.to_sql()))
216    }
217
218    /// Test if the value at `key` is a member of `vals`
219    /// `vals` should be a SurrealDB collection
220    pub fn member(key: String, vals: <Self as SearchFilter>::Value) -> Self {
221        Self(format!("{key} IN {}", vals.to_sql()))
222    }
223
224    /// Test if the value at `key` is *not* a member of `vals`
225    /// `vals` should be a SurrealDB collection
226    pub fn not_member(key: String, vals: <Self as SearchFilter>::Value) -> Self {
227        Self(format!("{key} NOTIN {}", vals.to_sql()))
228    }
229
230    // Geospatial filters
231    /// Test if the value at `key` is inside `geometry`
232    pub fn inside(key: String, geometry: <Self as SearchFilter>::Value) -> Self {
233        Self(format!("{key} INSIDE {}", geometry.to_sql()))
234    }
235
236    /// Test if the value at `key` is outside `geometry`
237    pub fn outside(key: String, geometry: <Self as SearchFilter>::Value) -> Self {
238        Self(format!("{key} OUTSIDE {}", geometry.to_sql()))
239    }
240
241    /// Test if the value at `key` intersects `geometry`
242    pub fn intersects(key: String, geometry: <Self as SearchFilter>::Value) -> Self {
243        Self(format!("{key} INTERSECTS {}", geometry.to_sql()))
244    }
245
246    // String ops
247    /// SurrealDB text search
248    pub fn matches<'a, S: AsRef<&'a str>>(key: String, query: S) -> Self {
249        Self(format!("{key} @@ {}", query.as_ref()))
250    }
251
252    /// Check if the value at `key` matches regex `pattern`
253    /// `pattern` should be a valid surrealDB regex
254    pub fn regex<'a, S: AsRef<&'a str>>(key: String, pattern: S) -> Self {
255        Self(format!("{key} = /{}/", pattern.as_ref()))
256    }
257}
258
259impl<C, Model> SurrealVectorStore<C, Model>
260where
261    C: Connection,
262    Model: EmbeddingModel,
263{
264    pub fn new(
265        model: Model,
266        surreal: Surreal<C>,
267        documents_table: Option<String>,
268        distance_function: SurrealDistanceFunction,
269    ) -> Self {
270        Self {
271            model,
272            surreal,
273            documents_table: documents_table.unwrap_or(String::from("documents")),
274            distance_function,
275        }
276    }
277
278    pub fn inner_client(&self) -> &Surreal<C> {
279        &self.surreal
280    }
281
282    pub fn with_defaults(model: Model, surreal: Surreal<C>) -> Self {
283        Self::new(model, surreal, None, SurrealDistanceFunction::Cosine)
284    }
285
286    fn search_query_full(&self) -> String {
287        self.search_query(true)
288    }
289
290    fn search_query_only_ids(&self) -> String {
291        self.search_query(false)
292    }
293
294    fn search_query(&self, with_document: bool) -> String {
295        let document = if with_document { ", document" } else { "" };
296        let embedded_text = if with_document { ", embedded_text" } else { "" };
297
298        let Self {
299            distance_function, ..
300        } = self;
301
302        format!(
303            "
304            SELECT id {document} {embedded_text}, {distance_function}($vec, embedding) as distance \
305              from type::table($tablename) \
306              where {distance_function}($vec, embedding) >= $threshold AND $filter \
307              order by distance desc \
308            LIMIT $limit",
309        )
310    }
311}
312
313impl<C, Model> VectorStoreIndex for SurrealVectorStore<C, Model>
314where
315    C: Connection,
316    Model: EmbeddingModel,
317{
318    type Filter = SurrealSearchFilter;
319
320    /// Get the top n documents based on the distance to the given query.
321    /// The result is a list of tuples of the form (score, id, document)
322    async fn top_n<T: for<'a> Deserialize<'a> + Send>(
323        &self,
324        req: VectorSearchRequest<SurrealSearchFilter>,
325    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
326        let embedded_query: Vec<f64> = self.model.embed_text(req.query()).await?.vec;
327
328        let mut response = self
329            .surreal
330            .query(self.search_query_full().as_str())
331            .bind(("vec", embedded_query))
332            .bind(("tablename", self.documents_table.clone()))
333            .bind(("threshold", req.threshold().unwrap_or(0.)))
334            .bind(("limit", req.samples() as usize))
335            .bind((
336                "filter",
337                req.filter()
338                    .clone()
339                    .map(SurrealSearchFilter::inner)
340                    .unwrap_or("true".into()),
341            ))
342            .await
343            .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?;
344
345        let rows: Vec<SearchResult> = response
346            .take(0)
347            .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?;
348
349        let rows: Vec<(f64, String, T)> = rows
350            .into_iter()
351            .map(SearchResult::into_result)
352            .collect::<Result<Vec<_>, _>>()?;
353
354        Ok(rows)
355    }
356
357    /// Same as `top_n` but returns the document ids only.
358    async fn top_n_ids(
359        &self,
360        req: VectorSearchRequest<SurrealSearchFilter>,
361    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
362        let embedded_query: Vec<f32> = self
363            .model
364            .embed_text(req.query())
365            .await?
366            .vec
367            .iter()
368            .map(|&x| x as f32)
369            .collect();
370
371        let mut response = self
372            .surreal
373            .query(self.search_query_only_ids().as_str())
374            .bind(("vec", embedded_query))
375            .bind(("tablename", self.documents_table.clone()))
376            .bind(("threshold", req.threshold().unwrap_or(0.)))
377            .bind(("limit", req.samples() as usize))
378            .bind((
379                "filter",
380                req.filter()
381                    .clone()
382                    .map(SurrealSearchFilter::inner)
383                    .unwrap_or("true".into()),
384            ))
385            .await
386            .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?;
387
388        let rows: Vec<SearchResultOnlyId> = response
389            .take::<Vec<SearchResultOnlyId>>(0)
390            .map_err(|e| VectorStoreError::DatastoreError(Box::new(e)))?;
391
392        let rows: Vec<(f64, String)> = rows
393            .into_iter()
394            .map(|row| (row.distance, record_key_to_string(&row.id.key)))
395            .collect();
396
397        Ok(rows)
398    }
399}
400
401// SurrealDB keeps a native filter value type, so it cannot use the blanket
402// `VectorStoreIndexDyn` impl that assumes JSON-valued filters.
403impl<C, Model> VectorStoreIndexDyn for SurrealVectorStore<C, Model>
404where
405    C: Connection,
406    Model: EmbeddingModel + Send + Sync,
407{
408    fn top_n<'a>(
409        &'a self,
410        req: VectorSearchRequest<Filter<serde_json::Value>>,
411    ) -> WasmBoxedFuture<'a, TopNResults> {
412        Box::pin(async move {
413            let req = req.try_map_filter(SurrealSearchFilter::try_from)?;
414            let results = <Self as VectorStoreIndex>::top_n::<serde_json::Value>(self, req).await?;
415            Ok(results)
416        })
417    }
418
419    fn top_n_ids<'a>(
420        &'a self,
421        req: VectorSearchRequest<Filter<serde_json::Value>>,
422    ) -> WasmBoxedFuture<'a, Result<Vec<(f64, String)>, VectorStoreError>> {
423        Box::pin(async move {
424            let req = req.try_map_filter(SurrealSearchFilter::try_from)?;
425            <Self as VectorStoreIndex>::top_n_ids(self, req).await
426        })
427    }
428}
429
430#[cfg(test)]
431mod tests {
432    use super::{Mem, SurrealSearchFilter, SurrealVectorStore};
433    use rig_core::{
434        client::Nothing,
435        embeddings::{Embedding, EmbeddingError, EmbeddingModel},
436        vector_store::{VectorStoreIndexDyn, request::Filter},
437    };
438    use serde_json::json;
439    use surrealdb::Surreal;
440
441    #[derive(Clone)]
442    struct MockEmbeddingModel;
443
444    impl EmbeddingModel for MockEmbeddingModel {
445        const MAX_DOCUMENTS: usize = 4;
446
447        type Client = Nothing;
448
449        fn make(_: &Self::Client, _: impl Into<String>, _: Option<usize>) -> Self {
450            Self
451        }
452
453        fn ndims(&self) -> usize {
454            3
455        }
456
457        async fn embed_texts(
458            &self,
459            texts: impl IntoIterator<Item = String> + Send,
460        ) -> Result<Vec<Embedding>, EmbeddingError> {
461            Ok(texts
462                .into_iter()
463                .map(|text| Embedding {
464                    document: text,
465                    vec: vec![0.0, 0.0, 0.0],
466                })
467                .collect())
468        }
469    }
470
471    #[allow(clippy::panic)]
472    #[test]
473    fn filter_from_json_preserves_nested_values() {
474        let filter = match SurrealSearchFilter::try_from(Filter::Eq(
475            "metadata".to_string(),
476            json!({
477                "name": "rig",
478                "flags": { "native": true },
479                "tags": ["surreal", "json"]
480            }),
481        )) {
482            Ok(filter) => filter,
483            Err(err) => panic!("unexpected surreal filter conversion failure: {err}"),
484        };
485
486        let sql = filter.to_string();
487
488        assert!(sql.starts_with("metadata = {"));
489        assert!(sql.contains("name: 'rig'"));
490        assert!(sql.contains("flags: { native: true }"));
491        assert!(sql.contains("tags: ['surreal', 'json']"));
492    }
493
494    #[allow(clippy::panic)]
495    #[tokio::test]
496    async fn surreal_vector_store_supports_type_erased_queries() {
497        fn assert_dyn<T: VectorStoreIndexDyn + Send + Sync + 'static>(_: T) {}
498
499        let surreal = match Surreal::new::<Mem>(()).await {
500            Ok(surreal) => surreal,
501            Err(err) => panic!("failed to create in-memory surreal client: {err}"),
502        };
503        let vector_store = SurrealVectorStore::with_defaults(MockEmbeddingModel, surreal);
504
505        assert_dyn(vector_store);
506    }
507}
rig_surrealdb/lib.rs

rig_surrealdb/
lib.rs
