semantic_operators/
rerank.rs

1//! Independent relevance scoring and stable ranking without synthetic scores.
2use crate::each::{BatchOptions, ItemResult, ask_each, validate_ids};
3use crate::{Answer, AnswerValue, Error, ErrorKind, Provider, Result, Score, State};
4use indexmap::IndexMap;
5use serde_json::{Map, Value};
6use std::collections::BTreeSet;
7
8/// A caller-retrieved document and its local identifier.
9#[derive(Clone)]
10pub struct Candidate {
11    id: String,
12    document: Value,
13}
14impl Candidate {
15    /// Validate a local identifier; documents may be any JSON value.
16    ///
17    /// # Errors
18    /// Returns `InvalidArgument` if `id` is empty. [`rerank`] separately checks
19    /// identifiers for uniqueness.
20    pub fn new(id: impl Into<String>, document: Value) -> Result<Self> {
21        let id = id.into();
22        validate_ids([id.as_str()])?;
23        Ok(Self { id, document })
24    }
25    /// Local identifier, never added to the provider's state.
26    pub fn id(&self) -> &str {
27        &self.id
28    }
29    /// Caller-provided document value.
30    pub fn document(&self) -> &Value {
31        &self.document
32    }
33}
34impl std::fmt::Debug for Candidate {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        f.debug_struct("Candidate").finish_non_exhaustive()
37    }
38}
39/// One decided candidate and its ranking key.
40#[derive(Clone, Debug)]
41pub struct Ranked {
42    id: String,
43    answer: Answer<AnswerValue>,
44    score: f64,
45    position: usize,
46}
47impl Ranked {
48    /// Local candidate identifier.
49    pub fn id(&self) -> &str {
50        &self.id
51    }
52    /// Original normalized score evidence and call provenance.
53    pub fn answer(&self) -> &Answer<AnswerValue> {
54        &self.answer
55    }
56    /// Expected rubric level, or expected custom weight after reweighting.
57    pub fn score(&self) -> f64 {
58        self.score
59    }
60    /// Retrieval position used to preserve exact ties.
61    pub fn position(&self) -> usize {
62        self.position
63    }
64}
65/// Sorted decided candidates and every unscored candidate's evidence.
66#[derive(Clone, Debug)]
67pub struct Ranking {
68    question: Score,
69    ranked: Vec<Ranked>,
70    unscored: IndexMap<String, ItemResult>,
71}
72impl Ranking {
73    /// Original scoring rubric.
74    pub fn question(&self) -> &Score {
75        &self.question
76    }
77    /// Candidates sorted by descending score, with exact ties in retrieval order.
78    pub fn ranked(&self) -> &[Ranked] {
79        &self.ranked
80    }
81    /// Undecided answers and provider failures, in retrieval order.
82    pub fn unscored(&self) -> &IndexMap<String, ItemResult> {
83        &self.unscored
84    }
85    /// Whether every candidate has a decided score.
86    pub fn complete(&self) -> bool {
87        self.unscored.is_empty()
88    }
89    /// Nonempty model identities actually reported, including undecided answers.
90    pub fn models(&self) -> BTreeSet<&str> {
91        self.ranked
92            .iter()
93            .map(|r| &r.answer)
94            .chain(self.unscored.values().filter_map(ItemResult::answer))
95            .filter_map(|a| a.call())
96            .filter_map(|c| c.model.as_deref())
97            .filter(|m| !m.is_empty())
98            .collect()
99    }
100    /// Return at most `k` candidates, checking completeness and reported-model consistency.
101    ///
102    /// `allow_partial` permits unscored candidates; it never permits mixed models.
103    /// Unknown model metadata stays unknown rather than claiming one common model.
104    ///
105    /// # Errors
106    /// Returns `InvalidArgument` if candidates are unscored and `allow_partial`
107    /// is false, or if answers report more than one nonempty model identity.
108    /// The mixed-model check applies even when partial results are allowed.
109    pub fn top(&self, k: usize, allow_partial: bool) -> Result<&[Ranked]> {
110        if !self.complete() && !allow_partial {
111            return Err(invalid(
112                "ranking has unscored candidates; explicitly allow partial results",
113            ));
114        }
115        if self.models().len() > 1 {
116            return Err(invalid(
117                "ranking contains answers from different reported models",
118            ));
119        }
120        Ok(&self.ranked[..k.min(self.ranked.len())])
121    }
122    /// Re-sort decided answers using finite strictly increasing rubric weights.
123    ///
124    /// Uses normalized `sum(weight * probability) / sum(probability)` and retains
125    /// the original answers and unscored entries. No provider call occurs.
126    ///
127    /// # Errors
128    /// Returns `InvalidArgument` if weights do not match the rubric length, are
129    /// nonfinite or not strictly increasing, or weighted arithmetic is nonfinite.
130    pub fn reweighted(&self, weights: &[f64]) -> Result<Self> {
131        if weights.len() != self.question.levels().len()
132            || weights.iter().any(|w| !w.is_finite())
133            || weights.windows(2).any(|w| w[0] >= w[1])
134        {
135            return Err(invalid(
136                "need one finite strictly increasing weight per rubric level",
137            ));
138        }
139        let mut ranked = Vec::with_capacity(self.ranked.len());
140        for old in &self.ranked {
141            let probabilities: Vec<_> = self
142                .question
143                .levels()
144                .iter()
145                .map(|label| old.answer.probabilities()[label])
146                .collect();
147            let total = accurate_sum(probabilities.iter().copied())?;
148            let weighted = accurate_sum(weights.iter().zip(&probabilities).map(|(w, p)| w * p))?;
149            let score = weighted / total;
150            if !score.is_finite() {
151                return Err(invalid("reweighted score is not finite"));
152            }
153            ranked.push(Ranked {
154                id: old.id.clone(),
155                answer: old.answer.clone(),
156                score,
157                position: old.position,
158            });
159        }
160        sort(&mut ranked);
161        Ok(Self {
162            question: self.question.clone(),
163            ranked,
164            unscored: self.unscored.clone(),
165        })
166    }
167}
168/// Score each candidate independently against a nonblank query.
169///
170/// Each request contains `query`, `document`, and optional non-null `context`.
171/// Local IDs and retrieval positions are never added. Failures and abstentions
172/// remain unscored, while exact score ties preserve retrieval order.
173///
174/// # Errors
175/// Returns `InvalidArgument` for a blank query or duplicate candidate IDs,
176/// and propagates local validation or unsupported-capability errors. Provider,
177/// timeout, and malformed-response failures are retained in `unscored` rather
178/// than converted into scores.
179pub async fn rerank(
180    provider: &dyn Provider,
181    question: &Score,
182    query: &str,
183    candidates: &[Candidate],
184    options: &BatchOptions,
185) -> Result<Ranking> {
186    if crate::question::is_blank(query) {
187        return Err(invalid("query must be a non-empty string"));
188    }
189    validate_ids(candidates.iter().map(Candidate::id))?;
190    let mut states = Vec::with_capacity(candidates.len());
191    for candidate in candidates {
192        let mut value = Map::new();
193        value.insert("query".into(), Value::String(query.into()));
194        value.insert("document".into(), candidate.document.clone());
195        if let Some(context) = options.context.as_ref().filter(|v| !v.is_null()) {
196            value.insert("context".into(), context.clone());
197        }
198        states.push(State::try_from(Value::Object(value))?);
199    }
200    let results = ask_each(provider, &question.clone().into(), &states, options).await?;
201    let mut ranked = Vec::new();
202    let mut unscored = IndexMap::new();
203    for (position, (candidate, result)) in candidates.iter().zip(results).enumerate() {
204        match result {
205            ItemResult::Answer(answer) if answer.decided() => {
206                let Some(AnswerValue::Score(score)) = answer.value() else {
207                    return Err(Error::new(
208                        ErrorKind::InvalidResponse,
209                        "rerank answer is not a score",
210                    ));
211                };
212                ranked.push(Ranked {
213                    id: candidate.id.clone(),
214                    score: *score,
215                    answer,
216                    position,
217                });
218            }
219            other => {
220                unscored.insert(candidate.id.clone(), other);
221            }
222        }
223    }
224    sort(&mut ranked);
225    Ok(Ranking {
226        question: question.clone(),
227        ranked,
228        unscored,
229    })
230}
231/// Re-rank existing answers with no additional calls.
232///
233/// # Errors
234/// Propagates the weight-validation and arithmetic errors from
235/// [`Ranking::reweighted`].
236pub fn reweighted(ranking: &Ranking, weights: &[f64]) -> Result<Ranking> {
237    ranking.reweighted(weights)
238}
239fn sort(items: &mut [Ranked]) {
240    items.sort_by(|a, b| {
241        b.score
242            .partial_cmp(&a.score)
243            .unwrap_or(std::cmp::Ordering::Equal)
244            .then(a.position.cmp(&b.position))
245    });
246}
247fn invalid(message: &str) -> Error {
248    Error::new(ErrorKind::InvalidArgument, message)
249}
250// Error-free partial summation plus half-even correction, matching math.fsum.
251// Adapted from CPython math.fsum; see NOTICE.md and LICENSE-PYTHON.
252fn accurate_sum(values: impl IntoIterator<Item = f64>) -> Result<f64> {
253    let mut partials: Vec<f64> = Vec::new();
254    for mut x in values {
255        if !x.is_finite() {
256            return Err(invalid("reweighted sum overflow"));
257        }
258        let mut i = 0;
259        for j in 0..partials.len() {
260            let mut y = partials[j];
261            if x.abs() < y.abs() {
262                std::mem::swap(&mut x, &mut y);
263            }
264            let hi = x + y;
265            if !hi.is_finite() {
266                return Err(invalid("reweighted sum overflow"));
267            }
268            let lo = y - (hi - x);
269            if lo != 0.0 {
270                partials[i] = lo;
271                i += 1;
272            }
273            x = hi;
274        }
275        partials.truncate(i);
276        partials.push(x);
277    }
278    let mut hi = partials.pop().unwrap_or(0.0);
279    let mut lo = 0.0;
280    while let Some(y) = partials.pop() {
281        let x = hi;
282        hi = x + y;
283        lo = y - (hi - x);
284        if lo != 0.0 {
285            break;
286        }
287    }
288    if let Some(last) = partials.last() {
289        if (lo < 0.0 && *last < 0.0) || (lo > 0.0 && *last > 0.0) {
290            let y = lo * 2.0;
291            let x = hi + y;
292            if y == x - hi {
293                hi = x;
294            }
295        }
296    }
297    Ok(hi)
298}
semantic_operators/rerank.rs

semantic_operators/
rerank.rs
