Skip to main content

post_cortex_embeddings/rerank/
mod.rs

1// Copyright (c) 2026 Julius ML
2// Licensed under the MIT License. See LICENSE at the workspace root.
3
4//! Rerankers: score each candidate document against the query so the
5//! candidate pool from vector + full-text search can be reordered.
6//!
7//! - [`OnnxCrossEncoder`] (feature `rerank`) — a BERT cross-encoder such as
8//!   `cross-encoder/ms-marco-MiniLM-L4-v2`, int8 on ONNX Runtime, local CPU.
9//! - [`VoyageReranker`] (feature `rerank-voyage`) — Voyage AI's rerank API;
10//!   candidate text leaves the machine.
11//! - [`FallbackReranker`] — one reranker backed by another.
12
13use std::sync::Arc;
14
15use anyhow::Result;
16use async_trait::async_trait;
17use tracing::warn;
18
19#[cfg(feature = "rerank")]
20mod onnx;
21#[cfg(feature = "rerank-voyage")]
22mod voyage;
23
24#[cfg(feature = "rerank")]
25pub use onnx::OnnxCrossEncoder;
26#[cfg(feature = "rerank-voyage")]
27pub use voyage::VoyageReranker;
28
29/// Scores (query, document) pairs.
30#[async_trait]
31pub trait Reranker: Send + Sync {
32    /// One relevance in [0, 1] per document, in input order; higher is more
33    /// relevant.
34    async fn score(&self, query: &str, documents: &[String]) -> Result<Vec<f32>>;
35}
36
37/// `primary`, or `fallback` when `primary` fails (e.g. a remote reranker
38/// that is down or rate limited, backed by a local one).
39pub struct FallbackReranker {
40    primary: Arc<dyn Reranker>,
41    fallback: Arc<dyn Reranker>,
42}
43
44impl FallbackReranker {
45    /// Rerank with `primary`, falling back to `fallback` on error.
46    pub fn new(primary: Arc<dyn Reranker>, fallback: Arc<dyn Reranker>) -> Self {
47        Self { primary, fallback }
48    }
49}
50
51#[async_trait]
52impl Reranker for FallbackReranker {
53    async fn score(&self, query: &str, documents: &[String]) -> Result<Vec<f32>> {
54        match self.primary.score(query, documents).await {
55            Ok(scores) => Ok(scores),
56            Err(e) => {
57                warn!("Reranker failed, using the fallback: {:#}", e);
58                self.fallback.score(query, documents).await
59            }
60        }
61    }
62}