diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector.rs b/codex-rs/ext/skills/src/dynamic_skill_selector.rs index 6a380a8aabf4..33f2e18d6a36 100644 --- a/codex-rs/ext/skills/src/dynamic_skill_selector.rs +++ b/codex-rs/ext/skills/src/dynamic_skill_selector.rs @@ -1,8 +1,10 @@ mod character_ngram; mod fielded_bm25; +mod multi_query_lexical; 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 weighted_lexical::WeightedLexicalSkillSelector; /// Metadata searched by a cheap skill selector. diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector/multi_query_lexical.rs b/codex-rs/ext/skills/src/dynamic_skill_selector/multi_query_lexical.rs new file mode 100644 index 000000000000..cd697b8a59ec --- /dev/null +++ b/codex-rs/ext/skills/src/dynamic_skill_selector/multi_query_lexical.rs @@ -0,0 +1,166 @@ +use std::collections::HashMap; +use std::collections::HashSet; + +use super::CheapSkillSelection; +use super::CheapSkillSelector; +use super::SkillSelectionDocument; +use super::WeightedLexicalSkillSelector; + +const MAX_QUERY_VIEWS: usize = 8; +const MAX_DECOMPOSITION_BYTES: usize = 4 * 1024; +const MAX_RESULTS: usize = 50; +const CONNECTORS: &[&str] = &[" and then ", " and ", " then ", " also "]; + +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct MultiQueryLexicalSkillSelector; + +impl CheapSkillSelector for MultiQueryLexicalSkillSelector { + fn method(&self) -> &'static str { + "multi_query_lexical_v1" + } + + fn select( + &self, + query: &str, + documents: &[SkillSelectionDocument<'_>], + limit: usize, + ) -> CheapSkillSelection { + let full_selection = + WeightedLexicalSkillSelector.select(query, documents, limit.min(MAX_RESULTS)); + let views = query_views(query); + if views.len() <= 1 || limit == 0 { + return full_selection; + } + + let mut candidates = HashMap::new(); + record_candidates(&mut candidates, &full_selection, /*view_index*/ 0); + let mut query_truncated = full_selection.query_truncated; + let mut candidate_set_truncated = full_selection.candidate_set_truncated; + for (view_index, view) in views.into_iter().enumerate().skip(1) { + let selection = + WeightedLexicalSkillSelector.select(view, documents, limit.min(MAX_RESULTS)); + query_truncated |= selection.query_truncated; + candidate_set_truncated |= selection.candidate_set_truncated; + record_candidates(&mut candidates, &selection, view_index); + } + + let mut candidates = candidates.into_values().collect::>(); + candidates.sort_by(|left, right| { + left.best_rank + .cmp(&right.best_rank) + .then_with(|| { + left.full_query_rank + .unwrap_or(usize::MAX) + .cmp(&right.full_query_rank.unwrap_or(usize::MAX)) + }) + .then_with(|| right.view_count.cmp(&left.view_count)) + .then_with(|| left.first_view.cmp(&right.first_view)) + .then_with(|| left.id.cmp(&right.id)) + }); + + CheapSkillSelection { + candidate_ids: candidates + .into_iter() + .take(limit.min(MAX_RESULTS)) + .map(|candidate| candidate.id) + .collect(), + query_term_count: full_selection.query_term_count, + query_truncated, + candidate_set_truncated, + } + } +} + +struct RankedCandidate { + id: usize, + best_rank: usize, + full_query_rank: Option, + view_count: usize, + first_view: usize, +} + +fn record_candidates( + candidates: &mut HashMap, + selection: &CheapSkillSelection, + view_index: usize, +) { + for (rank, id) in selection.candidate_ids.iter().copied().enumerate() { + let rank = rank + 1; + candidates + .entry(id) + .and_modify(|candidate| { + candidate.best_rank = candidate.best_rank.min(rank); + candidate.view_count = candidate.view_count.saturating_add(1); + if view_index == 0 { + candidate.full_query_rank = Some(rank); + } + }) + .or_insert(RankedCandidate { + id, + best_rank: rank, + full_query_rank: (view_index == 0).then_some(rank), + view_count: 1, + first_view: view_index, + }); + } +} + +fn query_views(query: &str) -> Vec<&str> { + let full_query = bounded(query, MAX_DECOMPOSITION_BYTES).trim(); + if full_query.is_empty() { + return Vec::new(); + } + + let mut views = vec![full_query]; + let mut seen = HashSet::from([full_query]); + for sentence in full_query.split(['\n', '\r', '.', '!', '?', ';']) { + for clause in split_connectors(sentence) { + let clause = clause.trim(); + if clause.chars().count() >= 2 && seen.insert(clause) { + views.push(clause); + if views.len() == MAX_QUERY_VIEWS { + return views; + } + } + } + } + views +} + +fn bounded(value: &str, max_bytes: usize) -> &str { + if value.len() <= max_bytes { + return value; + } + let mut end = max_bytes; + while !value.is_char_boundary(end) { + end = end.saturating_sub(1); + } + &value[..end] +} + +fn split_connectors(value: &str) -> Vec<&str> { + let lowercase = value.to_ascii_lowercase(); + let mut segments = Vec::new(); + let mut start = 0; + while start < value.len() { + let next = CONNECTORS + .iter() + .filter_map(|connector| { + lowercase[start..] + .find(connector) + .map(|offset| (start + offset, connector.len())) + }) + .min_by_key(|(position, _)| *position); + let Some((position, connector_length)) = next else { + break; + }; + segments.push(&value[start..position]); + start = position + connector_length; + } + segments.push(&value[start..]); + segments +} + +#[cfg(test)] +#[path = "multi_query_lexical_tests.rs"] +mod tests; diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector/multi_query_lexical_tests.rs b/codex-rs/ext/skills/src/dynamic_skill_selector/multi_query_lexical_tests.rs new file mode 100644 index 000000000000..faf60d640666 --- /dev/null +++ b/codex-rs/ext/skills/src/dynamic_skill_selector/multi_query_lexical_tests.rs @@ -0,0 +1,90 @@ +use super::*; +use pretty_assertions::assert_eq; + +#[test] +fn multi_query_promotes_each_clause_leader() { + let documents = [ + document(/*id*/ 1, "rust-format", "Format Rust source code."), + document(/*id*/ 2, "rust-lint", "Fix Rust source lint errors."), + document(/*id*/ 3, "rust-review", "Review Rust source code."), + document( + /*id*/ 4, + "ci-fix", + "Diagnose failing GitHub Actions checks.", + ), + ]; + + let selection = MultiQueryLexicalSkillSelector.select( + "format and review Rust source code, and then diagnose failing GitHub Actions checks", + &documents, + /*limit*/ 4, + ); + + assert!(selection.candidate_ids[..3].contains(&1)); + assert!(selection.candidate_ids[..3].contains(&4)); +} + +#[test] +fn single_query_matches_the_underlying_selector() { + let documents = [ + document(/*id*/ 1, "presentations", "Create visual decks."), + document(/*id*/ 2, "spreadsheets", "Analyze tabular data."), + ]; + + let expected = + WeightedLexicalSkillSelector.select("create presentations", &documents, /*limit*/ 20); + let actual = MultiQueryLexicalSkillSelector.select( + "create presentations", + &documents, + /*limit*/ 20, + ); + + assert_eq!(expected, actual); +} + +#[test] +fn query_views_split_sentences_and_connectors() { + assert_eq!( + vec![ + "format code and then fix tests; write a summary", + "format code", + "fix tests", + "write a summary", + ], + query_views("format code and then fix tests; write a summary"), + ); +} + +#[test] +fn multi_query_preserves_bounded_input_signals() { + let long_query = format!("{}\nand inspect logs", "match ".repeat(4 * 1024)); + let names = (0..=1_000) + .map(|index| format!("candidate-{index}")) + .collect::>(); + let documents = names + .iter() + .enumerate() + .map(|(id, name)| SkillSelectionDocument { + id, + name, + short_description: None, + description: "match logs", + }) + .collect::>(); + + let selection = + MultiQueryLexicalSkillSelector.select(&long_query, &documents, /*limit*/ 20); + + assert!(selection.query_truncated); + assert!(selection.candidate_set_truncated); + assert_eq!(20, selection.candidate_ids.len()); +} + +fn document<'a>(id: usize, name: &'a str, description: &'a str) -> SkillSelectionDocument<'a> { + SkillSelectionDocument { + id, + name, + short_description: None, + description, + } +} diff --git a/codex-rs/ext/skills/src/shadow_selection_experiment.rs b/codex-rs/ext/skills/src/shadow_selection_experiment.rs index 46f91edd848f..77d59334729f 100644 --- a/codex-rs/ext/skills/src/shadow_selection_experiment.rs +++ b/codex-rs/ext/skills/src/shadow_selection_experiment.rs @@ -15,6 +15,7 @@ use crate::dynamic_skill_selector::CharacterNgramSkillSelector; 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::SkillSelectionDocument; use crate::dynamic_skill_selector::WeightedLexicalSkillSelector; @@ -41,6 +42,7 @@ impl ShadowSelectionExperiment { Box::new(WeightedLexicalSkillSelector), Box::new(FieldedBm25SkillSelector), Box::new(CharacterNgramSkillSelector), + Box::new(MultiQueryLexicalSkillSelector), ], metrics_client, }