use crate::models::{Chunk, Document, Node};
use async_trait::async_trait;
use thiserror::Error;
#[derive(Error, Debug)]
pub enum CerebroError {
#[error("Storage error: {0}")]
StorageError(String),
#[error("Embedding error: {0}")]
EmbeddingError(String),
#[error("Ingestion error: {0}")]
IngestionError(String),
}
pub type Result<T> = std::result::Result<T, CerebroError>;
#[async_trait]
pub trait Ingestor: Send + Sync {
async fn ingest(&self, source: &str) -> Result<Vec<Document>>;
}
pub trait Chunker: Send + Sync {
fn chunk(&self, document: &Document) -> Result<Vec<Chunk>>;
}
#[async_trait]
pub trait Embedder: Send + Sync {
async fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
async fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
let res = self.embed(&[text]).await?;
Ok(res.into_iter().next().unwrap_or_default())
}
}
#[async_trait]
pub trait VectorStore: Send + Sync {
async fn upsert(&self, nodes: Vec<Node>) -> Result<()>;
async fn get(&self, node_ids: &[&str]) -> Result<Vec<Node>>;
async fn search(&self, text_query: &str, embedding: &[f32], top_k: usize) -> Result<Vec<(Node, f32)>>;
async fn delete_document(&self, doc_id: &str) -> Result<()>;
async fn get_all_nodes(&self) -> Result<Vec<Node>>;
}
#[async_trait]
pub trait KVStore: Send + Sync {
async fn set(&self, key: &str, value: &str) -> Result<()>;
async fn get(&self, key: &str) -> Result<Option<String>>;
}