post_cortex_embeddings/rerank/
mod.rs1use 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#[async_trait]
31pub trait Reranker: Send + Sync {
32 async fn score(&self, query: &str, documents: &[String]) -> Result<Vec<f32>>;
35}
36
37pub struct FallbackReranker {
40 primary: Arc<dyn Reranker>,
41 fallback: Arc<dyn Reranker>,
42}
43
44impl FallbackReranker {
45 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}