edgequake-llm 0.8.0

Multi-provider LLM abstraction library with caching, rate limiting, and cost tracking
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

//! Bi-encoder reranker using embedding cosine similarity.
//!
//! Fallback when HTTP cross-encoders (Jina/Cohere/Aliyun) are unavailable but an
//! [`EmbeddingProvider`] is configured.

use std::sync::Arc;

use async_trait::async_trait;

use crate::error::{LlmError, Result};
use crate::traits::EmbeddingProvider;

use super::result::RerankResult;
use super::traits::Reranker;

/// Rerank documents by cosine similarity between query and document embeddings.
pub struct BiEncoderReranker {
    embedding: Arc<dyn EmbeddingProvider>,
}

impl BiEncoderReranker {
    pub fn new(embedding: Arc<dyn EmbeddingProvider>) -> Self {
        Self { embedding }
    }
}

#[async_trait]
impl Reranker for BiEncoderReranker {
    fn name(&self) -> &str {
        "bi-encoder"
    }

    fn model(&self) -> &str {
        "embedding-cosine"
    }

    async fn rerank(
        &self,
        query: &str,
        documents: &[String],
        top_n: Option<usize>,
    ) -> Result<Vec<RerankResult>> {
        if documents.is_empty() {
            return Ok(Vec::new());
        }

        let mut texts = Vec::with_capacity(documents.len() + 1);
        texts.push(query.to_string());
        texts.extend(documents.iter().cloned());

        let embeddings = self.embedding.embed(&texts).await?;
        let query_emb = embeddings.first().ok_or_else(|| {
            LlmError::ProviderError("empty embedding response from bi-encoder reranker".into())
        })?;

        let mut scored: Vec<(usize, f64)> = documents
            .iter()
            .enumerate()
            .filter_map(|(idx, _)| {
                embeddings.get(idx + 1).map(|doc_emb| {
                    let score = cosine_similarity(query_emb, doc_emb) as f64;
                    (idx, score)
                })
            })
            .collect();

        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        let limit = top_n.unwrap_or(documents.len());
        Ok(scored
            .into_iter()
            .take(limit)
            .map(|(index, relevance_score)| RerankResult {
                index,
                relevance_score,
            })
            .collect())
    }
}

pub(crate) fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    if a.len() != b.len() || a.is_empty() {
        return 0.0;
    }
    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm_a <= 0.0 || norm_b <= 0.0 {
        0.0
    } else {
        dot / (norm_a * norm_b)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cosine_identical_vectors_score_one() {
        let v = vec![1.0, 0.0, 0.0];
        assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn cosine_orthogonal_vectors_score_zero() {
        let a = vec![1.0, 0.0];
        let b = vec![0.0, 1.0];
        assert!(cosine_similarity(&a, &b).abs() < 1e-6);
    }
}