From c983a53f2035bb9b6e999b1355b0056946671e9e Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 16 Jul 2026 14:07:09 +0000 Subject: [PATCH] Add character n-gram skill selection (#33613) ## What changed - Add a character n-gram selector that scores matches across skill names, short descriptions, and descriptions using field weights and inverse document frequency. - Include the selector in the shadow selection experiment alongside the existing lexical selectors. - Bound query, document, candidate, and result processing, and report query or candidate truncation through the existing selection metadata. ## Testing Add unit coverage for related word forms, typos, CJK text without word boundaries, and bounded inputs. GitOrigin-RevId: 4df60cc661338050eeb7741c716bc1c9f730fedd --- .../ext/skills/src/dynamic_skill_selector.rs | 2 + .../dynamic_skill_selector/character_ngram.rs | 243 ++++++++++++++++++ .../character_ngram_tests.rs | 77 ++++++ .../skills/src/shadow_selection_experiment.rs | 2 + 4 files changed, 324 insertions(+) create mode 100644 codex-rs/ext/skills/src/dynamic_skill_selector/character_ngram.rs create mode 100644 codex-rs/ext/skills/src/dynamic_skill_selector/character_ngram_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 fdb9531006..6a380a8aab 100644 --- a/codex-rs/ext/skills/src/dynamic_skill_selector.rs +++ b/codex-rs/ext/skills/src/dynamic_skill_selector.rs @@ -1,5 +1,7 @@ +mod character_ngram; mod fielded_bm25; mod weighted_lexical; +pub(crate) use character_ngram::CharacterNgramSkillSelector; pub(crate) use fielded_bm25::FieldedBm25SkillSelector; pub(crate) use weighted_lexical::WeightedLexicalSkillSelector; diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector/character_ngram.rs b/codex-rs/ext/skills/src/dynamic_skill_selector/character_ngram.rs new file mode 100644 index 0000000000..8cd628f249 --- /dev/null +++ b/codex-rs/ext/skills/src/dynamic_skill_selector/character_ngram.rs @@ -0,0 +1,243 @@ +use std::collections::HashMap; +use std::collections::HashSet; + +use super::CheapSkillSelection; +use super::CheapSkillSelector; +use super::SkillSelectionDocument; + +const MAX_QUERY_BYTES: usize = 4 * 1024; +const MAX_QUERY_TERMS: usize = 64; +const MAX_QUERY_GRAMS: usize = 512; +const MAX_DOCUMENT_BYTES: usize = 4 * 1024; +const MAX_DOCUMENT_TERMS: usize = 256; +const MAX_DOCUMENT_GRAMS: usize = 512; +const MAX_CANDIDATES: usize = 1_000; +const MAX_RESULTS: usize = 50; +const MIN_GRAM_CHARS: usize = 2; +const MAX_GRAM_CHARS: usize = 5; +const FIELD_WEIGHTS: [f64; 3] = [8.0, 4.0, 1.0]; + +const STOP_WORDS: &[&str] = &[ + "a", "an", "and", "are", "as", "at", "be", "by", "do", "for", "from", "how", "i", "in", "is", + "it", "me", "my", "of", "on", "or", "please", "that", "the", "this", "to", "use", "we", "what", + "when", "where", "which", "with", "you", "your", +]; + +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct CharacterNgramSkillSelector; + +impl CheapSkillSelector for CharacterNgramSkillSelector { + fn method(&self) -> &'static str { + "character_ngram_v1" + } + + fn select( + &self, + query: &str, + documents: &[SkillSelectionDocument<'_>], + limit: usize, + ) -> CheapSkillSelection { + let (query, query_bytes_truncated) = bounded(query, MAX_QUERY_BYTES); + let (query_terms, query_terms_truncated) = query_terms(query); + let (query_grams, query_grams_truncated) = grams(&query_terms, MAX_QUERY_GRAMS); + let query_truncated = + query_bytes_truncated || query_terms_truncated || query_grams_truncated; + let candidate_set_truncated = documents.len() > MAX_CANDIDATES; + if query_grams.is_empty() || limit == 0 { + return CheapSkillSelection { + query_term_count: query_terms.len(), + query_truncated, + candidate_set_truncated, + ..Default::default() + }; + } + + let prepared = documents + .iter() + .take(MAX_CANDIDATES) + .map(PreparedDocument::new) + .collect::>(); + let document_frequencies = document_frequencies(&prepared); + let document_count = prepared.len() as f64; + let minimum_matches = query_grams.len().min(3); + let mut scored = prepared + .iter() + .filter_map(|document| { + let (score, matched_grams) = score_document( + document, + &query_grams, + &document_frequencies, + document_count, + ); + (matched_grams >= minimum_matches).then_some((score, document.id, document.name)) + }) + .collect::>(); + scored.sort_by(|left, right| { + right + .0 + .total_cmp(&left.0) + .then_with(|| left.2.cmp(right.2)) + .then_with(|| left.1.cmp(&right.1)) + }); + + CheapSkillSelection { + candidate_ids: scored + .into_iter() + .take(limit.min(MAX_RESULTS)) + .map(|(_, id, _)| id) + .collect(), + query_term_count: query_terms.len(), + query_truncated, + candidate_set_truncated, + } + } +} + +struct PreparedDocument<'a> { + id: usize, + name: &'a str, + fields: [HashSet; 3], +} + +impl<'a> PreparedDocument<'a> { + fn new(document: &'a SkillSelectionDocument<'a>) -> Self { + Self { + id: document.id, + name: document.name, + fields: [ + document_grams(document.name), + document_grams(document.short_description.unwrap_or_default()), + document_grams(document.description), + ], + } + } +} + +fn score_document( + document: &PreparedDocument<'_>, + query_grams: &[String], + document_frequencies: &HashMap, + document_count: f64, +) -> (f64, usize) { + query_grams.iter().fold((0.0, 0), |(score, matched), gram| { + let frequency = document_frequencies.get(gram).copied().unwrap_or_default() as f64; + if frequency == 0.0 { + return (score, matched); + } + let inverse_document_frequency = + (1.0 + (document_count - frequency + 0.5) / (frequency + 0.5)).ln(); + let field_weight = document + .fields + .iter() + .enumerate() + .filter(|(_, field)| field.contains(gram)) + .map(|(index, _)| FIELD_WEIGHTS[index]) + .sum::(); + if field_weight == 0.0 { + (score, matched) + } else { + ( + score + inverse_document_frequency * field_weight, + matched + 1, + ) + } + }) +} + +fn document_frequencies(documents: &[PreparedDocument<'_>]) -> HashMap { + let mut frequencies = HashMap::new(); + for document in documents { + let grams = document + .fields + .iter() + .flatten() + .map(String::as_str) + .collect::>(); + for gram in grams { + *frequencies.entry(gram.to_string()).or_default() += 1; + } + } + frequencies +} + +fn query_terms(query: &str) -> (Vec, bool) { + let mut seen = HashSet::new(); + let mut terms = Vec::new(); + for term in normalized_terms(query) + .into_iter() + .filter(|term| term.chars().count() >= 2 && !STOP_WORDS.contains(&term.as_str())) + { + if !seen.insert(term.clone()) { + continue; + } + if terms.len() == MAX_QUERY_TERMS { + return (terms, true); + } + terms.push(term); + } + (terms, false) +} + +fn document_grams(value: &str) -> HashSet { + let (value, _) = bounded(value, MAX_DOCUMENT_BYTES); + let terms = normalized_terms(value) + .into_iter() + .take(MAX_DOCUMENT_TERMS) + .collect::>(); + grams(&terms, MAX_DOCUMENT_GRAMS).0.into_iter().collect() +} + +fn grams(terms: &[String], limit: usize) -> (Vec, bool) { + let mut seen = HashSet::new(); + let mut grams = Vec::new(); + for term in terms { + let characters = term.chars().collect::>(); + let minimum_gram_chars = if term.is_ascii() && characters.len() > MIN_GRAM_CHARS { + MIN_GRAM_CHARS + 1 + } else { + MIN_GRAM_CHARS + }; + for gram_size in minimum_gram_chars..=MAX_GRAM_CHARS.min(characters.len()) { + for start in 0..=characters.len() - gram_size { + let gram = characters[start..start + gram_size] + .iter() + .collect::(); + if !seen.insert(gram.clone()) { + continue; + } + if grams.len() == limit { + return (grams, true); + } + grams.push(gram); + } + } + } + (grams, false) +} + +fn normalized_terms(value: &str) -> Vec { + let mut normalized = String::with_capacity(value.len()); + for character in value.chars() { + if character.is_alphanumeric() { + normalized.extend(character.to_lowercase()); + } else { + normalized.push(' '); + } + } + normalized.split_whitespace().map(str::to_string).collect() +} + +fn bounded(value: &str, max_bytes: usize) -> (&str, bool) { + if value.len() <= max_bytes { + return (value, false); + } + let mut end = max_bytes; + while !value.is_char_boundary(end) { + end = end.saturating_sub(1); + } + (&value[..end], true) +} + +#[cfg(test)] +#[path = "character_ngram_tests.rs"] +mod tests; diff --git a/codex-rs/ext/skills/src/dynamic_skill_selector/character_ngram_tests.rs b/codex-rs/ext/skills/src/dynamic_skill_selector/character_ngram_tests.rs new file mode 100644 index 0000000000..4621746f53 --- /dev/null +++ b/codex-rs/ext/skills/src/dynamic_skill_selector/character_ngram_tests.rs @@ -0,0 +1,77 @@ +use super::*; +use pretty_assertions::assert_eq; + +#[test] +fn ngrams_match_related_word_forms() { + let documents = [ + document(/*id*/ 1, "presentations", "Create visual decks."), + document(/*id*/ 2, "spreadsheets", "Analyze tabular data."), + ]; + + let selection = + CharacterNgramSkillSelector.select("create a presentation", &documents, /*limit*/ 20); + + assert_eq!(vec![1], selection.candidate_ids); +} + +#[test] +fn ngrams_tolerate_a_typo() { + let documents = [ + document(/*id*/ 1, "postgresql", "Manage a relational database."), + document(/*id*/ 2, "postscript", "Render printable documents."), + ]; + + let selection = CharacterNgramSkillSelector.select( + "repair my postgrez database", + &documents, + /*limit*/ 20, + ); + + assert_eq!(vec![1, 2], selection.candidate_ids); +} + +#[test] +fn ngrams_match_cjk_without_word_boundaries() { + let documents = [ + document(/*id*/ 1, "演示文稿", "创建幻灯片。"), + document(/*id*/ 2, "电子表格", "分析表格数据。"), + ]; + + let selection = + CharacterNgramSkillSelector.select("帮我制作演示文稿", &documents, /*limit*/ 20); + + assert_eq!(vec![1], selection.candidate_ids); +} + +#[test] +fn ngrams_report_bounded_inputs() { + let long_query = "match ".repeat(MAX_QUERY_BYTES); + let names = (0..=MAX_CANDIDATES) + .map(|index| format!("candidate-{index}")) + .collect::>(); + let documents = names + .iter() + .enumerate() + .map(|(id, name)| SkillSelectionDocument { + id, + name, + short_description: None, + description: "match", + }) + .collect::>(); + + let selection = CharacterNgramSkillSelector.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 b2fba0d1e3..46f91edd84 100644 --- a/codex-rs/ext/skills/src/shadow_selection_experiment.rs +++ b/codex-rs/ext/skills/src/shadow_selection_experiment.rs @@ -11,6 +11,7 @@ use codex_protocol::user_input::UserInput; use crate::catalog::SkillCatalog; use crate::catalog::SkillSourceKind; +use crate::dynamic_skill_selector::CharacterNgramSkillSelector; use crate::dynamic_skill_selector::CheapSkillSelection; use crate::dynamic_skill_selector::CheapSkillSelector; use crate::dynamic_skill_selector::FieldedBm25SkillSelector; @@ -39,6 +40,7 @@ impl ShadowSelectionExperiment { selectors: vec![ Box::new(WeightedLexicalSkillSelector), Box::new(FieldedBm25SkillSelector), + Box::new(CharacterNgramSkillSelector), ], metrics_client, }