1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
use std::{error::Error, sync::Arc};
use surrealdb::{Connection, Surreal};
use crate::embedding::embedder_trait::Embedder;
use super::Store;
pub struct StoreBuilder<C: Connection> {
db: Option<Surreal<C>>,
collection_name: String,
collection_table_name: Option<String>,
collection_metadata_key_name: Option<String>,
vector_dimensions: i32,
embedder: Option<Arc<dyn Embedder>>,
schemafull: bool,
}
impl<C: Connection> StoreBuilder<C> {
/// Create a new StoreBuilder optimized for SurrealDB. Refer to `new_with_compatiblity()` if
/// you are looking to connect to store created by python version of langchain.
/// * table is singular - "document" instead of "documents"
/// * uses single table instead of multiple tables
/// * creates a schemafull table required for faster indexing. <https://github.com/surrealdb/surrealdb/issues/2013>
pub fn new() -> Self {
StoreBuilder {
db: None,
collection_name: "document".to_string(),
collection_table_name: Some("document".to_string()),
collection_metadata_key_name: Some("collection".to_string()),
vector_dimensions: 0,
embedder: None,
schemafull: true,
}
}
/// Create a new StoreBuilder with compatibility with python version of langchain
pub fn new_with_compatiblity() -> Self {
StoreBuilder {
db: None,
collection_name: "documents".to_string(),
collection_table_name: None,
collection_metadata_key_name: None,
vector_dimensions: 0,
embedder: None,
schemafull: false,
}
}
/// Use surrealdb
/// ```no_run
/// use oris_runtime::vectorstore::surrealdb::StoreBuilder;
///
/// #[tokio::main]
/// async fn main() {
/// let surrealdb_config = surrealdb::opt::Config::new()
/// .set_strict(true)
/// .capabilities(surrealdb::opt::capabilities::Capabilities::all())
/// .user(surrealdb::opt::auth::Root {
/// username: "username".into(),
/// password: "password".into()
/// });
/// let db = surrealdb::engine::any::connect(("ws://127.0.0.1:8000", surrealdb_config)).await.unwrap();
/// let store = StoreBuilder::new().db(db).vector_dimensions(1000).build().await.unwrap();
/// store.initialize().await.unwrap();
/// }
/// ```
pub fn db(mut self, db: Surreal<C>) -> Self {
self.db = Some(db);
self
}
pub fn collection_name(mut self, collection_name: &str) -> Self {
self.collection_name = collection_name.into();
self
}
/// Setting collection_table_name to None, creates table per collection. Set to some value if
/// you would like to reuse table. Resuing table is not compatible with python version of
/// langchain.
pub fn collection_table_name(mut self, collection_table_name: Option<String>) -> Self {
self.collection_table_name = collection_table_name;
self
}
pub fn collection_metadata_key_name(
mut self,
collection_metadata_key_name: Option<String>,
) -> Self {
self.collection_metadata_key_name = collection_metadata_key_name;
self
}
pub fn vector_dimensions(mut self, vector_dimensions: i32) -> Self {
self.vector_dimensions = vector_dimensions;
self
}
pub fn schemafull(mut self, schemafull: bool) -> Self {
self.schemafull = schemafull;
self
}
pub fn embedder<E: Embedder + 'static>(mut self, embedder: E) -> Self {
self.embedder = Some(Arc::new(embedder));
self
}
// Finalize the builder and construct the Store object
pub async fn build(self) -> Result<Store<C>, Box<dyn Error>> {
if self.embedder.is_none() {
return Err("Embedder is required".into());
}
if self.db.is_none() {
return Err("Db is required".into());
}
Ok(Store {
db: self.db.unwrap(),
collection_name: self.collection_name,
collection_table_name: self.collection_table_name,
collection_metadata_key_name: self.collection_metadata_key_name,
vector_dimensions: self.vector_dimensions,
embedder: self.embedder.unwrap(),
schemafull: self.schemafull,
})
}
}