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> {
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,
}
}
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,
}
}
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
}
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
}
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,
})
}
}