mem0_rust/rerankers/
mod.rs

1mod cohere;
2
3pub use cohere::CohereReranker;
4
5use async_trait::async_trait;
6use crate::models::ScoredMemory;
7use crate::errors::MemoryError;
8use crate::config::RerankerConfig;
9use std::sync::Arc;
10
11#[async_trait]
12pub trait Reranker: Send + Sync {
13    async fn rerank(&self, query: &str, results: Vec<ScoredMemory>) -> Result<Vec<ScoredMemory>, MemoryError>;
14    fn model_name(&self) -> &str;
15}
16
17pub fn create_reranker(config: &RerankerConfig) -> Result<Arc<dyn Reranker>, MemoryError> {
18    match config {
19        RerankerConfig::Cohere(cfg) => Ok(Arc::new(CohereReranker::new(cfg.clone())?)),
20    }
21}