SurrealDB

rig::surrealdb (feature surrealdb) implements Rig’s InsertDocuments and VectorStoreIndex traits on SurrealDB, using its built-in vector functions (such as vector::similarity::cosine). It works with an in-memory instance for development and a remote instance over WebSocket in production.

Setup

[dependencies]
rig = { version = "0.44.0", features = ["surrealdb"] }
surrealdb = "3"
tokio = { version = "1", features = ["full"] }

rig::surrealdb re-exports SurrealDB’s Mem (in-memory), Ws and Wss engines; the surrealdb crate gives you the Surreal client.

Example

use rig::prelude::*;
use serde::{Deserialize, Serialize};
use rig::providers::openai::{self, OpenAI};
use rig::surrealdb::{Mem, SurrealVectorStore};
use rig::vector_store::InsertDocuments;
use surrealdb::Surreal;

#[derive(Embed, Serialize, Deserialize, Clone, Debug, Default)]
struct WordDefinition {
    word: String,
    #[embed]
    definition: String,
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let model = OpenAI::from_env()?.embedding(openai::TEXT_EMBEDDING_3_SMALL, None);

    // In-memory instance; use `Surreal::new::<Ws>("localhost:8000")` for a server.
    let surreal = Surreal::new::<Mem>(()).await?;
    surreal.use_ns("example").use_db("example").await?;

    let words = vec![
        WordDefinition { word: "flurbo".into(), definition: "A fictional currency from Rick and Morty.".into() },
        WordDefinition { word: "glarb-glarb".into(), definition: "A creature from the marshlands of Glibbo.".into() },
    ];
    let documents = EmbeddingsBuilder::new(model.clone())
        .documents(words)?
        .build()
        .await?;

    let vector_store = SurrealVectorStore::with_defaults(model, surreal);
    vector_store.insert_documents(documents).await?;

    let req = VectorSearchRequest::builder()
        .query("weird alien creature")
        .samples(2)
        .build();

    for result in vector_store.top_n::<WordDefinition>(req).await? {
        println!("{:.3} {}", result.score, result.document.word);
    }

    Ok(())
}

How it works

SurrealVectorStore::with_defaults(model, surreal) uses the documents table and cosine similarity.
insert_documents writes one record per embedding: the document serialized as JSON, the vector, and the embedded text.
top_n::<T>(req) embeds the query, scores every record with the store’s function, and deserializes the stored document into T. result.score is the function’s value and result.id the record id.
inner_client() returns the Surreal client for your own queries.

Thresholds

A request threshold keeps results whose score is at least the threshold, and results are ordered highest first. That suits the similarity functions (cosine, Jaccard). For the distance functions (Euclidean, Hamming, KNN), where lower is closer, both the ordering and the threshold are inverted, so prefer a similarity function.

let req = VectorSearchRequest::builder()
    .query("weird alien creature")
    .samples(1)
    .threshold(0.5)
    .build();

let results = vector_store.top_n::<WordDefinition>(req).await?;

Table and distance function

To choose the table and the scoring function, build the store with new:

use rig::providers::openai::{self, OpenAI};
use rig::surrealdb::{Mem, SurrealDistanceFunction, SurrealVectorStore};
use surrealdb::Surreal;

let model = OpenAI::from_env()?.embedding(openai::TEXT_EMBEDDING_3_SMALL, None);
let surreal = Surreal::new::<Mem>(()).await?;

let store = SurrealVectorStore::new(
    model,
    surreal,
    Some("word_definitions".into()),
    SurrealDistanceFunction::Cosine,
);

The functions are Cosine (the default), Jaccard, Euclidean, Hamming and Knn.