Add multi-query lexical skill selection

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.

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::<Vec<_>>();
+        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<usize>,
+    view_count: usize,
+    first_view: usize,
+}
+
+fn record_candidates(
+    candidates: &mut HashMap<usize, RankedCandidate>,
+    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;

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::<Vec<_>>();
+    let documents = names
+        .iter()
+        .enumerate()
+        .map(|(id, name)| SkillSelectionDocument {
+            id,
+            name,
+            short_description: None,
+            description: "match logs",
+        })
+        .collect::<Vec<_>>();
+
+    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,
+    }
+}

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,
         }