mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
Fuse recent and lexical skills in shadow selection (#38204)
## What changed - Add the `lru_plus_lexical_v1` selector, combining the 50 most recent skills with weighted lexical matches through reciprocal rank fusion. - Include the combined selector in the shadow-selection experiment and propagate query and candidate truncation metadata from both inputs. - Generalize the existing reciprocal-rank fusion helper to accept a rank constant while preserving the existing lexical/character selector behavior. ## Testing - Cover lexical-only selection, merging recent and matching skills, promotion from agreement between both rankings, and weak overlaps. - Extend the skills extension test to verify invocation metrics for the new selector alongside `lru_v1`. GitOrigin-RevId: 1b7b82bc97ea919c719cd05a009351b701019306
This commit is contained in:
1 parent
9dd22890f5
commit
3d7f9b4637
6 files changed
+204
-8
No files matched your search
@@ -2,6 +2,7 @@ mod character_ngram;
|
|||||||
mod character_routing_card;
|
mod character_routing_card;
|
||||||
mod fielded_bm25;
|
mod fielded_bm25;
|
||||||
mod lru;
|
mod lru;
|
||||||
|
mod lru_plus_lexical;
|
||||||
mod multi_query_lexical;
|
mod multi_query_lexical;
|
||||||
mod routing_card_lexical;
|
mod routing_card_lexical;
|
||||||
mod rrf_lexical_char;
|
mod rrf_lexical_char;
|
||||||
@@ -11,6 +12,7 @@ pub(crate) use character_routing_card::CharacterRoutingCardSkillSelector;
|
|||||||
use codex_skills::SkillDependencies;
|
use codex_skills::SkillDependencies;
|
||||||
pub(crate) use fielded_bm25::FieldedBm25SkillSelector;
|
pub(crate) use fielded_bm25::FieldedBm25SkillSelector;
|
||||||
pub(crate) use lru::LruSkillSelector;
|
pub(crate) use lru::LruSkillSelector;
|
||||||
|
pub(crate) use lru_plus_lexical::LruPlusLexicalSkillSelector;
|
||||||
pub(crate) use multi_query_lexical::MultiQueryLexicalSkillSelector;
|
pub(crate) use multi_query_lexical::MultiQueryLexicalSkillSelector;
|
||||||
pub(crate) use routing_card_lexical::RoutingCardLexicalSkillSelector;
|
pub(crate) use routing_card_lexical::RoutingCardLexicalSkillSelector;
|
||||||
pub(crate) use rrf_lexical_char::RrfLexicalCharSkillSelector;
|
pub(crate) use rrf_lexical_char::RrfLexicalCharSkillSelector;
|
||||||
|
|||||||
@@ -0,0 +1,60 @@
|
|||||||
|
use super::CheapSkillSelection;
|
||||||
|
use super::CheapSkillSelector;
|
||||||
|
use super::LruSkillSelector;
|
||||||
|
use super::SkillSelectionDocument;
|
||||||
|
use super::WeightedLexicalSkillSelector;
|
||||||
|
use super::rrf_lexical_char::fuse_rankings_with_constant;
|
||||||
|
|
||||||
|
const MAX_RESULTS: usize = 50;
|
||||||
|
const RRF_K: usize = 10;
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Default)]
|
||||||
|
pub(crate) struct LruPlusLexicalSkillSelector {
|
||||||
|
lru_selector: LruSkillSelector,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LruPlusLexicalSkillSelector {
|
||||||
|
pub(crate) fn new(lru_selector: LruSkillSelector) -> Self {
|
||||||
|
Self { lru_selector }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CheapSkillSelector for LruPlusLexicalSkillSelector {
|
||||||
|
fn method(&self) -> &'static str {
|
||||||
|
"lru_plus_lexical_v1"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn select(
|
||||||
|
&self,
|
||||||
|
query: &str,
|
||||||
|
documents: &[SkillSelectionDocument<'_>],
|
||||||
|
limit: usize,
|
||||||
|
) -> CheapSkillSelection {
|
||||||
|
if limit == 0 {
|
||||||
|
return CheapSkillSelection::default();
|
||||||
|
}
|
||||||
|
|
||||||
|
let recent = self
|
||||||
|
.lru_selector
|
||||||
|
.select(query, documents, /*limit*/ MAX_RESULTS);
|
||||||
|
let lexical =
|
||||||
|
WeightedLexicalSkillSelector.select(query, documents, /*limit*/ MAX_RESULTS);
|
||||||
|
let candidate_ids = fuse_rankings_with_constant(
|
||||||
|
[&recent.candidate_ids, &lexical.candidate_ids],
|
||||||
|
limit.min(MAX_RESULTS),
|
||||||
|
/*rank_constant*/ RRF_K,
|
||||||
|
);
|
||||||
|
|
||||||
|
CheapSkillSelection {
|
||||||
|
candidate_ids,
|
||||||
|
query_term_count: recent.query_term_count.max(lexical.query_term_count),
|
||||||
|
query_truncated: recent.query_truncated || lexical.query_truncated,
|
||||||
|
candidate_set_truncated: recent.candidate_set_truncated
|
||||||
|
|| lexical.candidate_set_truncated,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
#[path = "lru_plus_lexical_tests.rs"]
|
||||||
|
mod tests;
|
||||||
@@ -0,0 +1,119 @@
|
|||||||
|
use pretty_assertions::assert_eq;
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn combined_selector_uses_lexical_results_before_invocation_history_exists() {
|
||||||
|
let documents = [
|
||||||
|
SkillSelectionDocument {
|
||||||
|
id: 10,
|
||||||
|
name: "ci",
|
||||||
|
short_description: None,
|
||||||
|
description: "Investigate failing checks.",
|
||||||
|
dependencies: None,
|
||||||
|
},
|
||||||
|
SkillSelectionDocument {
|
||||||
|
id: 20,
|
||||||
|
name: "python-tools",
|
||||||
|
short_description: None,
|
||||||
|
description: "Manage Python environments.",
|
||||||
|
dependencies: None,
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
let selection = LruPlusLexicalSkillSelector::new(LruSkillSelector::default()).select(
|
||||||
|
"manage python environments",
|
||||||
|
&documents,
|
||||||
|
/*limit*/ 50,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(vec![20], selection.candidate_ids);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn combined_selector_merges_recent_and_newly_matching_skills() {
|
||||||
|
let documents = [
|
||||||
|
SkillSelectionDocument {
|
||||||
|
id: 10,
|
||||||
|
name: "ci",
|
||||||
|
short_description: None,
|
||||||
|
description: "Investigate failing checks.",
|
||||||
|
dependencies: None,
|
||||||
|
},
|
||||||
|
SkillSelectionDocument {
|
||||||
|
id: 20,
|
||||||
|
name: "python-tools",
|
||||||
|
short_description: None,
|
||||||
|
description: "Manage Python environments.",
|
||||||
|
dependencies: None,
|
||||||
|
},
|
||||||
|
SkillSelectionDocument {
|
||||||
|
id: 30,
|
||||||
|
name: "monorepo",
|
||||||
|
short_description: None,
|
||||||
|
description: "Follow repository conventions.",
|
||||||
|
dependencies: None,
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
let selection = LruPlusLexicalSkillSelector::new(LruSkillSelector::new(vec![30, 10])).select(
|
||||||
|
"manage python environments",
|
||||||
|
&documents,
|
||||||
|
/*limit*/ 50,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(vec![20, 30, 10], selection.candidate_ids);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn combined_selector_promotes_skills_supported_by_both_signals() {
|
||||||
|
let documents = [
|
||||||
|
SkillSelectionDocument {
|
||||||
|
id: 10,
|
||||||
|
name: "ci",
|
||||||
|
short_description: None,
|
||||||
|
description: "Investigate failing checks.",
|
||||||
|
dependencies: None,
|
||||||
|
},
|
||||||
|
SkillSelectionDocument {
|
||||||
|
id: 30,
|
||||||
|
name: "monorepo",
|
||||||
|
short_description: None,
|
||||||
|
description: "Follow repository conventions.",
|
||||||
|
dependencies: None,
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
let selection = LruPlusLexicalSkillSelector::new(LruSkillSelector::new(vec![30, 10])).select(
|
||||||
|
"investigate failing ci checks",
|
||||||
|
&documents,
|
||||||
|
/*limit*/ 50,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(vec![10, 30], selection.candidate_ids);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn combined_fusion_keeps_a_top_ranked_skill_ahead_of_a_weak_overlap() {
|
||||||
|
let recent = std::iter::once(1)
|
||||||
|
.chain(100..148)
|
||||||
|
.chain(std::iter::once(2))
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let lexical = (200..249).chain(std::iter::once(2)).collect::<Vec<_>>();
|
||||||
|
|
||||||
|
let fused = fuse_rankings_with_constant(
|
||||||
|
[&recent, &lexical],
|
||||||
|
/*limit*/ MAX_RESULTS,
|
||||||
|
/*rank_constant*/ RRF_K,
|
||||||
|
);
|
||||||
|
let top_ranked = fused
|
||||||
|
.iter()
|
||||||
|
.position(|candidate| *candidate == 1)
|
||||||
|
.expect("top-ranked skill should be selected");
|
||||||
|
let weak_overlap = fused
|
||||||
|
.iter()
|
||||||
|
.position(|candidate| *candidate == 2)
|
||||||
|
.expect("weak overlapping skill should be selected");
|
||||||
|
|
||||||
|
assert!(top_ranked < weak_overlap);
|
||||||
|
}
|
||||||
@@ -53,12 +53,20 @@ impl CheapSkillSelector for RrfLexicalCharSkillSelector {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn fuse_rankings<const N: usize>(rankings: [&[usize]; N], limit: usize) -> Vec<usize> {
|
fn fuse_rankings<const N: usize>(rankings: [&[usize]; N], limit: usize) -> Vec<usize> {
|
||||||
|
fuse_rankings_with_constant(rankings, limit, RRF_K)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn fuse_rankings_with_constant<const N: usize>(
|
||||||
|
rankings: [&[usize]; N],
|
||||||
|
limit: usize,
|
||||||
|
rank_constant: usize,
|
||||||
|
) -> Vec<usize> {
|
||||||
let mut scores = HashMap::<usize, f64>::new();
|
let mut scores = HashMap::<usize, f64>::new();
|
||||||
let mut best_ranks = HashMap::<usize, usize>::new();
|
let mut best_ranks = HashMap::<usize, usize>::new();
|
||||||
for ranking in rankings {
|
for ranking in rankings {
|
||||||
for (index, id) in ranking.iter().copied().enumerate() {
|
for (index, id) in ranking.iter().copied().enumerate() {
|
||||||
let rank = index + 1;
|
let rank = index + 1;
|
||||||
*scores.entry(id).or_default() += 1.0 / (RRF_K + rank) as f64;
|
*scores.entry(id).or_default() += 1.0 / (rank_constant + rank) as f64;
|
||||||
best_ranks
|
best_ranks
|
||||||
.entry(id)
|
.entry(id)
|
||||||
.and_modify(|best_rank| *best_rank = (*best_rank).min(rank))
|
.and_modify(|best_rank| *best_rank = (*best_rank).min(rank))
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ use crate::dynamic_skill_selector::CharacterRoutingCardSkillSelector;
|
|||||||
use crate::dynamic_skill_selector::CheapSkillSelection;
|
use crate::dynamic_skill_selector::CheapSkillSelection;
|
||||||
use crate::dynamic_skill_selector::CheapSkillSelector;
|
use crate::dynamic_skill_selector::CheapSkillSelector;
|
||||||
use crate::dynamic_skill_selector::FieldedBm25SkillSelector;
|
use crate::dynamic_skill_selector::FieldedBm25SkillSelector;
|
||||||
|
use crate::dynamic_skill_selector::LruPlusLexicalSkillSelector;
|
||||||
use crate::dynamic_skill_selector::LruSkillSelector;
|
use crate::dynamic_skill_selector::LruSkillSelector;
|
||||||
use crate::dynamic_skill_selector::MultiQueryLexicalSkillSelector;
|
use crate::dynamic_skill_selector::MultiQueryLexicalSkillSelector;
|
||||||
use crate::dynamic_skill_selector::RoutingCardLexicalSkillSelector;
|
use crate::dynamic_skill_selector::RoutingCardLexicalSkillSelector;
|
||||||
@@ -120,7 +121,8 @@ impl ShadowSelectionExperiment {
|
|||||||
.collect();
|
.collect();
|
||||||
let routing_selector = CharacterRoutingCardSkillSelector::new(catalog, host_snapshot);
|
let routing_selector = CharacterRoutingCardSkillSelector::new(catalog, host_snapshot);
|
||||||
let lru_selector = LruSkillSelector::new(recent_skill_ids);
|
let lru_selector = LruSkillSelector::new(recent_skill_ids);
|
||||||
let mut ranked_selections = Vec::with_capacity(self.selectors.len() + 2);
|
let lru_plus_lexical_selector = LruPlusLexicalSkillSelector::new(lru_selector.clone());
|
||||||
|
let mut ranked_selections = Vec::with_capacity(self.selectors.len() + 3);
|
||||||
|
|
||||||
for selector in self
|
for selector in self
|
||||||
.selectors
|
.selectors
|
||||||
@@ -129,6 +131,7 @@ impl ShadowSelectionExperiment {
|
|||||||
.chain([
|
.chain([
|
||||||
&routing_selector as &dyn CheapSkillSelector,
|
&routing_selector as &dyn CheapSkillSelector,
|
||||||
&lru_selector as &dyn CheapSkillSelector,
|
&lru_selector as &dyn CheapSkillSelector,
|
||||||
|
&lru_plus_lexical_selector as &dyn CheapSkillSelector,
|
||||||
])
|
])
|
||||||
{
|
{
|
||||||
let start = Instant::now();
|
let start = Instant::now();
|
||||||
|
|||||||
@@ -932,7 +932,7 @@ async fn shadow_lru_selector_recovers_a_skill_invoked_on_an_earlier_turn() -> Te
|
|||||||
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
|
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
|
||||||
.find(|metric| metric.name() == "codex.skills.shadow_selection.invocation")
|
.find(|metric| metric.name() == "codex.skills.shadow_selection.invocation")
|
||||||
.ok_or("shadow invocation metric should be recorded")?;
|
.ok_or("shadow invocation metric should be recorded")?;
|
||||||
let mut lru_hits = match metric.data() {
|
let mut selector_hits = match metric.data() {
|
||||||
AggregatedMetrics::U64(MetricData::Sum(sum)) => sum
|
AggregatedMetrics::U64(MetricData::Sum(sum)) => sum
|
||||||
.data_points()
|
.data_points()
|
||||||
.filter_map(|point| {
|
.filter_map(|point| {
|
||||||
@@ -941,7 +941,7 @@ async fn shadow_lru_selector_recovers_a_skill_invoked_on_an_earlier_turn() -> Te
|
|||||||
.find(|attribute| attribute.key.as_str() == "method")?
|
.find(|attribute| attribute.key.as_str() == "method")?
|
||||||
.value
|
.value
|
||||||
.as_str();
|
.as_str();
|
||||||
if method != "lru_v1" {
|
if method != "lru_v1" && method != "lru_plus_lexical_v1" {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
let hit = point
|
let hit = point
|
||||||
@@ -950,16 +950,20 @@ async fn shadow_lru_selector_recovers_a_skill_invoked_on_an_earlier_turn() -> Te
|
|||||||
.value
|
.value
|
||||||
.as_str()
|
.as_str()
|
||||||
.to_string();
|
.to_string();
|
||||||
Some((hit, point.value()))
|
Some((method.to_string(), hit, point.value()))
|
||||||
})
|
})
|
||||||
.collect::<Vec<_>>(),
|
.collect::<Vec<_>>(),
|
||||||
data => panic!("unexpected shadow invocation metric data: {data:?}"),
|
data => panic!("unexpected shadow invocation metric data: {data:?}"),
|
||||||
};
|
};
|
||||||
lru_hits.sort();
|
selector_hits.sort();
|
||||||
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
vec![("false".to_string(), 1), ("true".to_string(), 1)],
|
vec![
|
||||||
lru_hits
|
("lru_plus_lexical_v1".to_string(), "true".to_string(), 2),
|
||||||
|
("lru_v1".to_string(), "false".to_string(), 1),
|
||||||
|
("lru_v1".to_string(), "true".to_string(), 1),
|
||||||
|
],
|
||||||
|
selector_hits
|
||||||
);
|
);
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in new issue
Block a user