use crate::ast::{Action, Qail};
impl Qail {
pub fn search(collection: &str) -> Self {
Self {
action: Action::Search,
table: collection.to_string(),
..Default::default()
}
}
pub fn upsert(collection: &str) -> Self {
Self {
action: Action::Upsert,
table: collection.to_string(),
..Default::default()
}
}
pub fn scroll(collection: &str) -> Self {
Self {
action: Action::Scroll,
table: collection.to_string(),
..Default::default()
}
}
pub fn vector(mut self, embedding: Vec<f32>) -> Self {
self.vector = Some(embedding);
self
}
pub fn score_threshold(mut self, threshold: f32) -> Self {
self.score_threshold = Some(threshold);
self
}
pub fn vector_name(mut self, name: &str) -> Self {
self.vector_name = Some(name.to_string());
self
}
pub fn with_vectors(mut self) -> Self {
self.with_vector = true;
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_search_builder() {
let cmd = Qail::search("products")
.vector(vec![0.1, 0.2, 0.3])
.score_threshold(0.8)
.limit(10);
assert_eq!(cmd.action, Action::Search);
assert_eq!(cmd.table, "products");
assert_eq!(cmd.vector, Some(vec![0.1, 0.2, 0.3]));
assert_eq!(cmd.score_threshold, Some(0.8));
}
#[test]
fn test_vector_name() {
let cmd = Qail::search("articles")
.vector_name("title")
.vector(vec![0.5, 0.5]);
assert_eq!(cmd.vector_name, Some("title".to_string()));
}
#[test]
fn test_with_vectors() {
let cmd = Qail::search("products").vector(vec![0.1]).with_vectors();
assert!(cmd.with_vector);
}
}