khive_runtime/
embedding_reranker.rs

1//! Token-scoped cosine reranking of entity candidates in one named engine.
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use khive_retrieval::error::RetrievalError;
8use khive_retrieval::hybrid::Reranker;
9use khive_score::{cmp_desc_then_id, DeterministicScore};
10use uuid::Uuid;
11
12use crate::{KhiveRuntime, NamespaceToken, RuntimeError, RuntimeResult};
13
14/// One input entry after cosine scoring or missing-vector fallback.
15#[derive(Clone, Copy, Debug, Eq, PartialEq)]
16pub struct EmbeddingRerankHit {
17    pub id: Uuid,
18    pub score: DeterministicScore,
19    /// No entity vector was found in the selected engine and token namespace.
20    /// This says nothing about vectors in other namespaces, engines or kinds.
21    pub missing_vector: bool,
22}
23
24/// Cosine reranker bound to one runtime, caller token and named engine.
25///
26/// Query embedding uses the runtime's token-aware query-role path. Candidate
27/// scoring uses [`KhiveRuntime::rerank_in`], without a namespace-wide KNN cut.
28/// The signed canonical cosine score replaces each stored candidate's score;
29/// missing vectors retain the exact incoming score instead.
30///
31/// [`Self::rerank_detailed`] exposes the fallback flag and typed runtime errors.
32/// The [`Reranker`] implementation projects hits to tuples and maps failures to
33/// [`RetrievalError::Rerank`], losing their typed causes and retry classification.
34#[derive(Clone)]
35pub struct EmbeddingCosineReranker {
36    runtime: Arc<KhiveRuntime>,
37    token: NamespaceToken,
38    engine: String,
39}
40
41impl EmbeddingCosineReranker {
42    /// Bind a registered engine, canonicalizing built-in aliases without loading
43    /// its embedding service. Unknown or request-excluded engines are refused.
44    pub fn new(
45        runtime: Arc<KhiveRuntime>,
46        token: NamespaceToken,
47        engine: impl Into<String>,
48    ) -> RuntimeResult<Self> {
49        let engine = engine.into();
50        let (engine, _) = runtime.vector_model_metadata(&engine)?;
51        Ok(Self {
52            runtime,
53            token,
54            engine,
55        })
56    }
57
58    /// Score every supplied entry, retaining missing vectors and duplicate IDs.
59    ///
60    /// Vector reads deduplicate IDs; output does not. Missing duplicates each
61    /// retain their own incoming score. Results sort by descending final score,
62    /// ascending UUID, then original position for identical ID/score ties. The
63    /// limit applies after scored and missing entries are merged.
64    ///
65    /// Engine and generated-vector validation still run for empty candidates or
66    /// a zero limit. Embedding and storage failures fail the whole call, rather
67    /// than marking candidates missing. More than `u32::MAX` distinct IDs refuse
68    /// because the underlying scorer's result bound cannot represent that set.
69    pub async fn rerank_detailed(
70        &self,
71        query: &str,
72        results: Vec<(Uuid, DeterministicScore)>,
73        top_k: usize,
74    ) -> RuntimeResult<Vec<EmbeddingRerankHit>> {
75        let ids: Vec<_> = results
76            .iter()
77            .map(|(id, _)| *id)
78            .collect::<BTreeSet<_>>()
79            .into_iter()
80            .collect();
81        let count = u32::try_from(ids.len()).map_err(|_| {
82            RuntimeError::InvalidInput("rerank candidate count exceeds u32::MAX".into())
83        })?;
84        let query_vector = self
85            .runtime
86            .embed_query_with_model_for_token(&self.token, &self.engine, query)
87            .await?;
88        let scores: BTreeMap<_, _> = self
89            .runtime
90            .rerank_in(&self.token, &self.engine, &query_vector, &ids, count)
91            .await?
92            .into_iter()
93            .map(|hit| (hit.subject_id, hit.score))
94            .collect();
95        let mut hits: Vec<_> = results
96            .into_iter()
97            .enumerate()
98            .map(|(position, (id, incoming_score))| {
99                let score = scores.get(&id).copied();
100                (
101                    position,
102                    EmbeddingRerankHit {
103                        id,
104                        score: score.unwrap_or(incoming_score),
105                        missing_vector: score.is_none(),
106                    },
107                )
108            })
109            .collect();
110        hits.sort_by(|(a_position, a), (b_position, b)| {
111            cmp_desc_then_id(a.score, &a.id, b.score, &b.id)
112                .then_with(|| a_position.cmp(b_position))
113        });
114        Ok(hits.into_iter().take(top_k).map(|(_, hit)| hit).collect())
115    }
116}
117
118#[async_trait]
119impl Reranker<Uuid> for EmbeddingCosineReranker {
120    async fn rerank(
121        &self,
122        query: &str,
123        results: Vec<(Uuid, DeterministicScore)>,
124        top_k: usize,
125    ) -> khive_retrieval::error::Result<Vec<(Uuid, DeterministicScore)>> {
126        self.rerank_detailed(query, results, top_k)
127            .await
128            .map(|hits| hits.into_iter().map(|hit| (hit.id, hit.score)).collect())
129            .map_err(|error| RetrievalError::Rerank(error.to_string()))
130    }
131}
khive_runtime/embedding_reranker.rs

khive_runtime/
embedding_reranker.rs
