use async_trait::async_trait;
use std::collections::HashMap;
use crate::bm25::{config::Bm25Config, error::Result, result::Bm25Result};
use crate::common::namespace::Namespace;
use crate::store::health::HealthReport;
#[derive(Debug, Clone)]
pub struct IndexDocument {
pub id: String,
pub fields: HashMap<String, String>,
}
impl IndexDocument {
pub fn new(id: impl Into<String>, text: impl Into<String>) -> Self {
let mut fields = HashMap::new();
fields.insert("text".to_string(), text.into());
Self { id: id.into(), fields }
}
pub fn with_fields(
id: impl Into<String>,
fields: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
) -> Self {
Self {
id: id.into(),
fields: fields.into_iter().map(|(k, v)| (k.into(), v.into())).collect(),
}
}
}
#[derive(Debug, Clone)]
pub struct SearchOptions {
pub limit: usize,
pub min_score: Option<f32>,
pub fields: Vec<String>,
}
impl Default for SearchOptions {
fn default() -> Self {
Self { limit: 20, min_score: None, fields: Vec::new() }
}
}
impl SearchOptions {
pub fn with_limit(mut self, limit: usize) -> Self {
self.limit = limit;
self
}
pub fn with_min_score(mut self, min: f32) -> Self {
self.min_score = Some(min);
self
}
pub fn with_fields(mut self, fields: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.fields = fields.into_iter().map(Into::into).collect();
self
}
}
#[async_trait]
pub trait Bm25Adapter: Send + Sync {
fn name(&self) -> &'static str;
fn is_connected(&self) -> bool;
fn config(&self) -> &Bm25Config;
async fn index(
&self,
document: IndexDocument,
namespace: Option<&Namespace>,
) -> Result<()>;
async fn remove(
&self,
id: &str,
namespace: Option<&Namespace>,
) -> Result<bool>;
async fn search(
&self,
query: &str,
namespace: Option<&Namespace>,
options: SearchOptions,
) -> Result<Vec<Bm25Result>>;
async fn count(&self, namespace: Option<&Namespace>) -> Result<usize>;
async fn clear(&self, namespace: Option<&Namespace>) -> Result<usize>;
async fn healthcheck(&self) -> HealthReport;
}