winx_code_agent/utils/
path_prob.rs

1//! File-path relevance ranking, ported from wcgw's `FastPathAnalyzer`.
2//!
3//! wcgw ships a tiny unigram language model trained over repo paths: a
4//! Hugging Face tokenizer (`paths_tokens.model`) plus a vocab file mapping each
5//! token to its log-probability (`paths_model.vocab`). A path's score is the sum
6//! of the log-probabilities of its tokens — higher (less negative) means the
7//! path looks more like a "real source file worth showing" and less like noise.
8//!
9//! Both assets are embedded so ranking works offline with zero setup, matching
10//! the wcgw package that bundles them alongside `repo_context.py`.
11
12use std::collections::HashMap;
13use std::sync::OnceLock;
14
15use serde::Deserialize;
16use tokie::Tokenizer;
17
18static PATHS_MODEL: &[u8] =
19    include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/assets/paths_tokens.model"));
20static PATHS_VOCAB: &[u8] =
21    include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/assets/paths_model.vocab"));
22
23struct PathAnalyzer {
24    tokenizer: Tokenizer,
25    id_tokens: Vec<String>,
26    vocab_probs: HashMap<String, f64>,
27}
28
29#[derive(Deserialize)]
30struct PathTokenizerFile {
31    model: PathTokenizerModel,
32}
33
34#[derive(Deserialize)]
35struct PathTokenizerModel {
36    vocab: Vec<(String, f64)>,
37}
38
39impl PathAnalyzer {
40    fn load() -> Option<Self> {
41        let model_json = std::str::from_utf8(PATHS_MODEL).ok()?;
42        let tokenizer = match tokie::hf::from_json_str(model_json) {
43            Ok(tokenizer) => tokenizer,
44            Err(error) => {
45                tracing::warn!("Failed to load embedded path-ranking model: {error}");
46                return None;
47            }
48        };
49        let id_tokens = serde_json::from_str::<PathTokenizerFile>(model_json)
50            .ok()?
51            .model
52            .vocab
53            .into_iter()
54            .map(|(token, _)| token)
55            .collect::<Vec<_>>();
56
57        // Vocab lines are `<token>\t<log_prob>`; mirror wcgw's `split()` + len==2 check.
58        let text = std::str::from_utf8(PATHS_VOCAB).ok()?;
59        let mut vocab_probs = HashMap::new();
60        for line in text.lines() {
61            let parts: Vec<&str> = line.split_whitespace().collect();
62            if parts.len() == 2 {
63                if let Ok(prob) = parts[1].parse::<f64>() {
64                    vocab_probs.insert(parts[0].to_string(), prob);
65                }
66            }
67        }
68
69        Some(Self { tokenizer, id_tokens, vocab_probs })
70    }
71
72    fn sum_log_prob(&self, token_ids: &[u32]) -> f64 {
73        token_ids
74            .iter()
75            .filter_map(|token_id| self.id_tokens.get(*token_id as usize))
76            .filter_map(|token| self.vocab_probs.get(token))
77            .sum()
78    }
79}
80
81fn analyzer() -> Option<&'static PathAnalyzer> {
82    static ANALYZER: OnceLock<Option<PathAnalyzer>> = OnceLock::new();
83    ANALYZER.get_or_init(PathAnalyzer::load).as_ref()
84}
85
86/// Score each path by summed token log-probability (higher = more relevant).
87///
88/// Returns `None` if the model failed to load, so callers can fall back to a
89/// heuristic ordering instead of silently mis-ranking everything.
90pub fn score_paths<S: AsRef<str>>(paths: &[S]) -> Option<Vec<f64>> {
91    let analyzer = analyzer()?;
92    let scores = paths
93        .iter()
94        .map(|path| analyzer.sum_log_prob(&analyzer.tokenizer.encode_ids(path.as_ref(), false)))
95        .collect();
96    Some(scores)
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn embedded_path_tokenizer_matches_golden_ids() -> anyhow::Result<()> {
105        let analyzer =
106            analyzer().ok_or_else(|| anyhow::anyhow!("embedded path analyzer failed"))?;
107        let cases: &[(&str, &[u32])] = &[
108            ("src/main.rs", &[13, 3, 21, 4, 506]),
109            ("node_modules/react/index.js", &[14, 5, 7, 3, 75, 3, 63, 4, 9]),
110            (
111                "a/b/c/d/e/f/zzz_tmp_garbage_9f8a.bin",
112                &[
113                    3473, 3, 65, 3, 20, 3, 31, 3, 48, 3, 32, 3, 356, 1193, 5, 514, 5, 2101, 102,
114                    2662, 5, 147, 32, 105, 42, 4, 268,
115                ],
116            ),
117            (
118                "tests/integração/日本語.rs",
119                &[146, 3, 180, 5318, 18526, 18560, 137, 3, 18490, 18436, 18581, 4, 506],
120            ),
121        ];
122        for (path, expected) in cases {
123            let actual = analyzer.tokenizer.encode_ids(path, false);
124            assert_eq!(&actual, expected, "path token ids diverged for {path:?}");
125        }
126        Ok(())
127    }
128
129    #[test]
130    fn ranks_source_above_noise_when_model_present() {
131        let paths =
132            vec!["src/main.rs".to_string(), "a/b/c/d/e/f/zzz_tmp_garbage_9f8a.bin".to_string()];
133        if let Some(scores) = score_paths(&paths) {
134            assert_eq!(scores.len(), 2);
135            // A normal source path should not score worse than deep random noise.
136            assert!(scores[0] >= scores[1], "expected src/main.rs >= noise path");
137        }
138    }
139}
winx_code_agent/utils/path_prob.rs

winx_code_agent/utils/
path_prob.rs
