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

rig_surrealdb/
lib.rs
