refactor: enable using mem as db

0xPlaygrounds · Feb 7, 2025 · 38bc1bf · 38bc1bf
1 parent 50ba885
commit 38bc1bf
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 22 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/rig-surrealdb/Cargo.toml b/rig-surrealdb/Cargo.toml
@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-surrealdb = { version = "2.1.4", features = ["protocol-ws"] }
+surrealdb = { version = "2.1.4", features = ["kv-mem", "protocol-ws"] }
 rig-core = { path = "../rig-core", version = "0.7.0",  features = ["derive"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"

diff --git a/rig-surrealdb/README.md b/rig-surrealdb/README.md
@@ -1,13 +1,12 @@
 # Rig SurrealDB integration
 This crate integrates SurrealDB into Rig, allowing you to easily use RAG with this database.
 
-## How to run the example
-To run the example, spin up a docker  container:
-```bash
-docker run --rm --pull always -p 9999:8000 surrealdb/surrealdb:latest start -u root --password root
-```
+## How to use
+To use this integration, you need a SurrealDB instance that has a table with the following fields:
+- `
 
-Once done, open another tab then add your OpenAI API key as an environment variable:
+## How to run the example
+To run the example, add your OpenAI API key as an environment variable:
 ```bash
 export OPENAI_API_KEY=my_key
 ```

diff --git a/rig-surrealdb/examples/vector_search_surreal.rs b/rig-surrealdb/examples/vector_search_surreal.rs
@@ -1,5 +1,5 @@
 use rig::{embeddings::EmbeddingsBuilder, vector_store::VectorStoreIndex, Embed};
-use rig_surrealdb::{SurrealVectorStore, Ws};
+use rig_surrealdb::{Mem, SurrealVectorStore};
 use serde::{Deserialize, Serialize};
 use surrealdb::{opt::auth::Root, Surreal};
 
@@ -29,22 +29,10 @@ async fn main() -> Result<(), anyhow::Error> {
     let openai_client = rig::providers::openai::Client::from_env();
     let model = openai_client.embedding_model(rig::providers::openai::TEXT_EMBEDDING_3_SMALL);
 
-    let surreal = Surreal::new::<Ws>("localhost:9999").await?;
-
-    surreal
-        .signin(Root {
-            username: "root",
-            password: "root",
-        })
-        .await?;
+    let surreal = Surreal::new::<Mem>(()).await?;
 
     surreal.use_ns("example").use_db("example").await?;
 
-    surreal
-        .query(include_str!("./migrations.surql"))
-        .await
-        .unwrap();
-
     // create test documents with mocked embeddings
     let words = vec![
         WordDefinition {

diff --git a/rig-surrealdb/src/lib.rs b/rig-surrealdb/src/lib.rs
@@ -8,6 +8,7 @@ use rig::{
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 use surrealdb::{sql::Thing, Connection, Surreal};
 
+pub use surrealdb::engine::local::{Db, Mem};
 pub use surrealdb::engine::remote::ws::{Ws, Wss};
 
 pub struct SurrealVectorStore<Model: EmbeddingModel, C: Connection> {