From 6915bac7ba2753277cfb6679b547c03e4fe567ed Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 21 Jul 2026 13:44:31 +0000 Subject: [PATCH] Add reciprocal rank fusion skill selection (#34547) ## What changed - Add a skill selector that uses reciprocal rank fusion to combine the top 50 results from the weighted lexical and character n-gram selectors. - Prefer candidates supported by both rankings, with deterministic rank and identifier tie-breakers, and propagate truncation metadata from both selectors. - Include the fused selector in the shadow-selection experiment under `rrf_lexical_char_v1`. ## Testing - Add unit coverage for cross-ranking support, deterministic tie-breaking, and combined input-bound reporting. GitOrigin-RevId: 3165136b8729ea3a9484f25c7f2ce0f91fb8aaa9 --- .../ext/skills/src/dynamic_skill_selector.rs | 2 + .../rrf_lexical_char.rs | 85 +++++++++++++++++++ .../rrf_lexical_char_tests.rs | 34 ++++++++ .../skills/src/shadow_selection_experiment.rs | 2 + 4 files changed, 123 insertions(+) create mode 100644 codex-rs/ext/skills/src/dynamic_skill_selector/rrf_lexical_char.rs create mode 100644 codex-rs/ext/skills/src/dynamic_skill_selector/rrf_lexical_char_tests.rs diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector.rs b/codex-rs/ext/skills/src/dynamic_skill_selector.rs index 33f2e18d6a..148d99d46b 100644 --- a/codex-rs/ext/skills/src/dynamic_skill_selector.rs +++ b/codex-rs/ext/skills/src/dynamic_skill_selector.rs @@ -1,10 +1,12 @@ mod character_ngram; mod fielded_bm25; mod multi_query_lexical; +mod rrf_lexical_char; mod weighted_lexical; pub(crate) use character_ngram::CharacterNgramSkillSelector; pub(crate) use fielded_bm25::FieldedBm25SkillSelector; pub(crate) use multi_query_lexical::MultiQueryLexicalSkillSelector; +pub(crate) use rrf_lexical_char::RrfLexicalCharSkillSelector; pub(crate) use weighted_lexical::WeightedLexicalSkillSelector; /// Metadata searched by a cheap skill selector. diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector/rrf_lexical_char.rs b/codex-rs/ext/skills/src/dynamic_skill_selector/rrf_lexical_char.rs new file mode 100644 index 0000000000..c09b29c6c5 --- /dev/null +++ b/codex-rs/ext/skills/src/dynamic_skill_selector/rrf_lexical_char.rs @@ -0,0 +1,85 @@ +use std::collections::HashMap; + +use super::CharacterNgramSkillSelector; +use super::CheapSkillSelection; +use super::CheapSkillSelector; +use super::SkillSelectionDocument; +use super::WeightedLexicalSkillSelector; + +const MAX_FUSION_CANDIDATES: usize = 50; +const RRF_K: usize = 60; + +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct RrfLexicalCharSkillSelector; + +impl CheapSkillSelector for RrfLexicalCharSkillSelector { + fn method(&self) -> &'static str { + "rrf_lexical_char_v1" + } + + fn select( + &self, + query: &str, + documents: &[SkillSelectionDocument<'_>], + limit: usize, + ) -> CheapSkillSelection { + if limit == 0 { + return CheapSkillSelection::default(); + } + + let lexical = WeightedLexicalSkillSelector.select( + query, + documents, + /*limit*/ MAX_FUSION_CANDIDATES, + ); + let character = CharacterNgramSkillSelector.select( + query, + documents, + /*limit*/ MAX_FUSION_CANDIDATES, + ); + let candidate_ids = fuse_rankings( + [&lexical.candidate_ids, &character.candidate_ids], + limit.min(MAX_FUSION_CANDIDATES), + ); + + CheapSkillSelection { + candidate_ids, + query_term_count: lexical.query_term_count.max(character.query_term_count), + query_truncated: lexical.query_truncated || character.query_truncated, + candidate_set_truncated: lexical.candidate_set_truncated + || character.candidate_set_truncated, + } + } +} + +fn fuse_rankings(rankings: [&[usize]; N], limit: usize) -> Vec { + let mut scores = HashMap::::new(); + let mut best_ranks = HashMap::::new(); + for ranking in rankings { + for (index, id) in ranking.iter().copied().enumerate() { + let rank = index + 1; + *scores.entry(id).or_default() += 1.0 / (RRF_K + rank) as f64; + best_ranks + .entry(id) + .and_modify(|best_rank| *best_rank = (*best_rank).min(rank)) + .or_insert(rank); + } + } + + let mut candidates = scores.into_iter().collect::>(); + candidates.sort_by(|(left_id, left_score), (right_id, right_score)| { + right_score + .total_cmp(left_score) + .then_with(|| best_ranks[left_id].cmp(&best_ranks[right_id])) + .then_with(|| left_id.cmp(right_id)) + }); + candidates + .into_iter() + .take(limit) + .map(|(id, _)| id) + .collect() +} + +#[cfg(test)] +#[path = "rrf_lexical_char_tests.rs"] +mod tests; diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector/rrf_lexical_char_tests.rs b/codex-rs/ext/skills/src/dynamic_skill_selector/rrf_lexical_char_tests.rs new file mode 100644 index 0000000000..2bc55e3002 --- /dev/null +++ b/codex-rs/ext/skills/src/dynamic_skill_selector/rrf_lexical_char_tests.rs @@ -0,0 +1,34 @@ +use pretty_assertions::assert_eq; + +use super::*; + +#[test] +fn fusion_prefers_candidates_supported_by_both_rankings() { + let fused = fuse_rankings([&[1, 2, 3], &[2, 3, 4]], /*limit*/ 4); + + assert_eq!(vec![2, 3, 1, 4], fused); +} + +#[test] +fn fusion_uses_rank_and_id_to_break_ties() { + let fused = fuse_rankings([&[20, 10], &[30]], /*limit*/ 3); + + assert_eq!(vec![20, 30, 10], fused); +} + +#[test] +fn selector_reports_the_combined_input_bounds() { + let documents = [SkillSelectionDocument { + id: 7, + name: "presentations", + short_description: None, + description: "Create visual decks.", + }]; + let query = "presentation ".repeat(4 * 1024); + + let selection = RrfLexicalCharSkillSelector.select(&query, &documents, /*limit*/ 20); + + assert_eq!(vec![7], selection.candidate_ids); + assert!(selection.query_truncated); + assert!(!selection.candidate_set_truncated); +} diff --git a/codex-rs/ext/skills/src/shadow_selection_experiment.rs b/codex-rs/ext/skills/src/shadow_selection_experiment.rs index 77d5933472..d79c5f9034 100644 --- a/codex-rs/ext/skills/src/shadow_selection_experiment.rs +++ b/codex-rs/ext/skills/src/shadow_selection_experiment.rs @@ -16,6 +16,7 @@ use crate::dynamic_skill_selector::CheapSkillSelection; use crate::dynamic_skill_selector::CheapSkillSelector; use crate::dynamic_skill_selector::FieldedBm25SkillSelector; use crate::dynamic_skill_selector::MultiQueryLexicalSkillSelector; +use crate::dynamic_skill_selector::RrfLexicalCharSkillSelector; use crate::dynamic_skill_selector::SkillSelectionDocument; use crate::dynamic_skill_selector::WeightedLexicalSkillSelector; @@ -43,6 +44,7 @@ impl ShadowSelectionExperiment { Box::new(FieldedBm25SkillSelector), Box::new(CharacterNgramSkillSelector), Box::new(MultiQueryLexicalSkillSelector), + Box::new(RrfLexicalCharSkillSelector), ], metrics_client, }