akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(rust): port CJK-aware relevance-query matching to CodeCompressor (#2634)

Browse source 
## Description

The Rust port of `CodeCompressor` (#1154, parity-only) did not carry
over the CJK-aware relevance-query matching from
`headroom/transforms/code_compressor.py` (`_CONTEXT_DELIMS` /
`_CJK_CHARS` / `_query_context_tokens()` / `_symbol_in_context()`, lines
2353-2387, called from lines 987/1009):

- Rust tokenized the context with an ASCII-only delimiter class
`[\s,;:.()\[\]{}"']+`, so a CJK query (no spaces, CJK punctuation)
collapses into a single blob and never isolates an ASCII symbol name.
- The substring-fallback guard `chars().count() > 3` had no CJK
relaxation, so a short ASCII name glued to CJK text (e.g. `run` in
`修复run函数的报错`, `db` in `请保留db相关的逻辑`) could never receive the +3.0 context
boost — while Python does boost it. Same `(code, context)` input,
different `symbol_scores`.

This PR ports the two Python helpers with identical semantics:

- `query_context_tokens()` — delimiter class extended with the
CJK/full-width punctuation and ideographic space from Python's
`_CONTEXT_DELIMS`; returns `(words, lowered, has_cjk)` with CJK
detection over U+3000-U+9FFF, U+AC00-U+D7AF, U+FF00-U+FFEF (Python's
`_CJK_CHARS`).
- `symbol_in_context()` — exact token match, plus the substring fallback
gated by `> 3` **characters** (Python `len()`, not bytes), relaxed when
the query contains CJK.

The call site in `analyze_symbols` now uses these helpers; no other
behavior changed. Pure-ASCII query behavior is identical to before
(exact token match, `>3`-gated substring fallback), which the tests pin
down.

Closes #2630

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `crates/headroom-core/src/transforms/code_compressor.rs`: extract
`query_context_tokens()` / `symbol_in_context()` free functions
mirroring the Python helpers (CJK/full-width delimiter class, CJK
detection, CJK-relaxed `>3`-character guard); replace the inline
ASCII-only tokenization + guard in `analyze_symbols` with calls to them.
- Unit tests mirroring
`tests/test_transforms/test_code_compressor_cjk.py` case-for-case, plus
a character-vs-byte guard test and an end-to-end `compress_with` test
asserting `symbol_scores`.

## Testing

- [x] Unit tests pass (`cargo test -p headroom-core` — Rust-only change;
Python untouched)
- [x] Linting passes (`cargo fmt --check`, `cargo clippy -p
headroom-core --all-targets` — no new warnings)
- [ ] Type checking passes (`mypy headroom`) — N/A, no Python changes
- [x] New tests added for new functionality
- [x] Manual testing performed

New tests (all in `code_compressor.rs` `mod tests`):

- `cjk_query_isolates_wrapped_ascii_symbol` — full-width parens isolate
`parse_config`
- `cjk_query_matches_short_ascii_name_glued_to_cjk` — `db` (len 2) glued
to CJK matches via the relaxed guard
- `english_short_name_substring_still_gated` — `db` vs "keep the
database helper" must NOT match (ASCII guard unchanged)
- `english_exact_token_match_unchanged`,
`english_long_name_substring_fallback_unchanged`,
`empty_context_matches_nothing`
- `guard_counts_chars_not_bytes` — the guard is a character count,
matching Python `len()`
- `cjk_context_boosts_named_symbol_end_to_end` — full `compress_with`
run asserting `symbol_scores` (red on main, green here — see proof)

### Test Output

```text
$ cargo test -p headroom-core --lib -- code_compressor::tests
test transforms::code_compressor::tests::empty_and_short_passthrough ... ok
test transforms::code_compressor::tests::empty_context_matches_nothing ... ok
test transforms::code_compressor::tests::estimate_tokens_uses_chars_div_4_min_1 ... ok
test transforms::code_compressor::tests::py_round3_matches_cpython ... ok
test transforms::code_compressor::tests::py_round_int_is_half_to_even ... ok
test transforms::code_compressor::tests::cjk_query_isolates_wrapped_ascii_symbol ... ok
test transforms::code_compressor::tests::english_short_name_substring_still_gated ... ok
test transforms::code_compressor::tests::cjk_query_matches_short_ascii_name_glued_to_cjk ... ok
test transforms::code_compressor::tests::english_exact_token_match_unchanged ... ok
test transforms::code_compressor::tests::english_long_name_substring_fallback_unchanged ... ok
test transforms::code_compressor::tests::guard_counts_chars_not_bytes ... ok
test transforms::code_compressor::tests::cjk_context_boosts_named_symbol_end_to_end ... ok
test transforms::code_compressor::tests::detect_language_basic ... ok
test result: ok. 13 passed; 0 failed; 0 ignored; 0 measured; 899 filtered out; finished in 0.05s

$ cargo test -p headroom-core   # per-binary summaries
lib .......................... ok. 911 passed; 0 failed; 1 ignored
auth_mode .................... ok.  16 passed; 0 failed
cache_control ................ ok.  14 passed; 0 failed
ccr_backends ................. ok.   7 passed; 0 failed
ccr_roundtrip ................ ok.  15 passed; 0 failed
code_compressor_parity ....... ok.   1 passed; 0 failed   (recorded byte-parity fixtures)
live_zone_ccr ................ ok.   3 passed; 0 failed
live_zone_dispatch ........... ok.   6 passed; 0 failed
live_zone_thresholds ......... ok.   2 passed; 0 failed
live_zone_token_validation ... ok.   3 passed; 0 failed
recommendations_loader ....... ok.   4 passed; 0 failed
tokenizer_proptest ........... ok.   5 passed; 0 failed
doc-tests .................... ok.   1 passed; 0 failed; 2 ignored

$ cargo fmt --check    # clean
$ cargo clippy -p headroom-core --all-targets   # no new warnings
```

## Real Behavior Proof

- Environment: macOS (Darwin 25.5.0), repo-pinned Rust toolchain
(`rust-toolchain.toml`), branch based on current `main`.
- Exact command / steps: the end-to-end test was written first and run
against unmodified `main` (red), then after the fix (green). Input:
Python source with two signal-symmetric functions `run` and `keep`;
context `修复run函数的报错`. The Python reference gives `run` the boost
(`_symbol_in_context('run', ...) == True`, `_symbol_in_context('keep',
...) == False`, verified against the live Python implementation), so
expected normalized scores are `run = 1.0`, `keep = 0.0`.
- Observed result: on unmodified main the end-to-end test fails (`left:
0.5, right: 1.0` — the CJK query `修复run函数的报错` gives `run` no boost, both
symbols collapse to 0.5, while Python scores `run=1.0, keep=0.0`); on
this branch all 8 new tests pass and the same query boosts `run` to 1.0,
matching Python. Full output:

Before (unmodified `main` + new test only — Rust gives no boost, both
symbols collapse to 0.5):

  ```text
----
transforms::code_compressor::tests::cjk_context_boosts_named_symbol_end_to_end
stdout ----

thread '...cjk_context_boosts_named_symbol_end_to_end' panicked at
crates/headroom-core/src/transforms/code_compressor.rs:1890:9:
  assertion `left == right` failed: run must get the context boost
    left: 0.5
   right: 1.0

test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 904
filtered out
  ```

After (this branch): the same test passes, including its ASCII control
case (`fix the runner` must NOT boost `run` — scores stay 0.5/0.5),
proving pure-ASCII behavior is unchanged. The recorded byte-parity
fixture suite (`code_compressor_parity`) also still passes.

- Not tested: real proxy traffic end-to-end (change is confined to the
symbol-scoring context boost inside the Rust compressor; the Python
implementation is the behavioral reference and is untouched).
`kompress_parity` was not run locally — it is model-gated and my sandbox
blocks the model fetch; it is unrelated to this change and CI covers its
skip path.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation — N/A
(internal behavior fix)
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md` — it is generated by
release-please from my Conventional Commit PR title (a CI guard enforces
this)

## Additional Notes

For background, the Python-side CJK handling comes from the merged CJK
sweep work (#2220 and follow-ups); #1154 predates part of it, which is
likely how the port missed it. Longer names wrapped in full-width
punctuation happened to still match in Rust via the substring fallback,
but the token set itself was wrong; this PR restores exact-token
semantics for those too.
This commit is contained in:
Zhenjia ZHOU 2026-07-30 00:14:29 +08:00 • committed by GitHub
parent 22b707fd31
commit e86c6390ce
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 164 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

183
crates/headroom-core/src/transforms/code_compressor.rs
View file

@ -413,6 +413,59 @@ fn get_definition_name(node: Node, code: &str) -> Option<String> {
None
}
/// Tokenize a relevance query for symbol-name matching (CJK-aware).
/// Mirrors `_query_context_tokens`: returns (word set, lowercased query,
/// has_cjk). Symbol names are ASCII identifiers; CJK relevance queries have
/// no spaces and use CJK/full-width punctuation, so an ASCII-only delimiter
/// class would collapse the whole query into one blob and never isolate an
/// ASCII name the user asked to keep. CJK/full-width punctuation and the
/// ideographic space are therefore delimiters too.
fn query_context_tokens(context: &str) -> (BTreeSet<String>, String, bool) {
if context.is_empty() {
return (BTreeSet::new(), String::new(), false);
}
static DELIMS: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
let delims = DELIMS.get_or_init(|| {
// Same class as Python `_CONTEXT_DELIMS`.
regex::Regex::new(r#"[\s,;:.()\[\]{}"'，、；：。．！？（）【】「」『』《》〈〉·…—　]+"#)
.unwrap()
});
static CJK: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
let cjk = CJK.get_or_init(|| {
// Same class as Python `_CJK_CHARS`:
// U+3000-U+9FFF, U+AC00-U+D7AF (Hangul), U+FF00-U+FFEF (full-width).
regex::Regex::new(r"[\u{3000}-\u{9FFF}\u{AC00}-\u{D7AF}\u{FF00}-\u{FFEF}]").unwrap()
});
let lowered = context.to_lowercase();
let words: BTreeSet<String> = delims
.split(&lowered)
.filter(|s| !s.is_empty())
.map(|s| s.to_string())
.collect();
let has_cjk = cjk.is_match(&lowered);
(words, lowered, has_cjk)
}
/// Whether the relevance query names this symbol. Mirrors `_symbol_in_context`:
/// exact token match, or a substring fallback gated by len>3 (in characters,
/// like Python's `len`) for ASCII queries but relaxed for CJK queries — a
/// short ASCII name glued to CJK has no delimiter to isolate it, so the
/// exact match can't fire and the guard would wrongly drop it.
fn symbol_in_context(
name_lower: &str,
words: &BTreeSet<String>,
context_lower: &str,
has_cjk: bool,
) -> bool {
if words.is_empty() || name_lower.is_empty() {
return false;
}
if words.contains(name_lower) {
return true;
}
context_lower.contains(name_lower) && (name_lower.chars().count() > 3 || has_cjk)
}
fn is_public_symbol(name: &str, language: CodeLanguage) -> bool {
if name.is_empty() {
return false;
@ -1066,18 +1119,8 @@ impl CodeAwareCompressor {
ref_counts.insert(qname.clone(), (count - def_count).max(0));
}
// Context words (empty when context is "").
let context_lower = context.to_lowercase();
let context_words: BTreeSet<String> = if context.is_empty() {
BTreeSet::new()
} else {
static SPLIT: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
let re = SPLIT.get_or_init(|| regex::Regex::new(r#"[\s,;:.()\[\]{}"']+"#).unwrap());
re.split(&context_lower)
.filter(|s| !s.is_empty())
.map(|s| s.to_string())
.collect()
};
// Context words (empty when context is ""). Mirrors `_query_context_tokens`.
let (context_words, context_lower, context_has_cjk) = query_context_tokens(context);
// Raw importance signals per symbol.
let mut raw_signals: Vec<(String, f64)> = Vec::new();
@ -1106,13 +1149,14 @@ impl CodeAwareCompressor {
raw += 1.0;
}
if !context_words.is_empty() {
let name_lower = short.to_lowercase();
if context_words.contains(&name_lower)
|| (name_lower.chars().count() > 3 && context_lower.contains(&name_lower))
{
raw += 3.0;
}
// Context boost: the relevance query named this symbol.
if symbol_in_context(
&short.to_lowercase(),
&context_words,
&context_lower,
context_has_cjk,
) {
raw += 3.0;
}
raw_signals.push((qname.clone(), raw));
}
@ -1866,6 +1910,107 @@ mod tests {
assert_eq!(lang, CodeLanguage::Unknown);
}
// CJK-aware relevance-query matching. Mirrors
// tests/test_transforms/test_code_compressor_cjk.py (Python reference).
#[test]
fn cjk_query_isolates_wrapped_ascii_symbol() {
// Full-width parens around the name must still tokenize parse_config out.
let (words, lowered, has_cjk) =
query_context_tokens("请重点保留（parse_config）的解析配置");
assert!(has_cjk);
assert!(words.contains("parse_config"));
assert!(symbol_in_context("parse_config", &words, &lowered, has_cjk));
}
#[test]
fn cjk_query_matches_short_ascii_name_glued_to_cjk() {
// 'db' (len 2) glued to CJK has no delimiter to isolate it; the len>3
// guard is relaxed for CJK so the substring fallback still matches.
let (words, lowered, has_cjk) = query_context_tokens("请保留db相关的逻辑");
assert!(has_cjk);
assert!(symbol_in_context("db", &words, &lowered, has_cjk));
}
#[test]
fn english_short_name_substring_still_gated() {
// ASCII query unchanged: a short name that is only a substring (not a
// token) of an English query must NOT match (avoids spurious boosts).
let (words, lowered, has_cjk) = query_context_tokens("keep the database helper");
assert!(!has_cjk);
assert!(!symbol_in_context("db", &words, &lowered, has_cjk));
}
#[test]
fn english_exact_token_match_unchanged() {
let (words, lowered, has_cjk) = query_context_tokens("keep parse_config and helper");
assert!(!has_cjk);
assert!(symbol_in_context("parse_config", &words, &lowered, has_cjk));
assert!(symbol_in_context("helper", &words, &lowered, has_cjk));
}
#[test]
fn english_long_name_substring_fallback_unchanged() {
// ASCII path, len>3 substring fallback: 'parse_config' is not a
// standalone token but is a substring of 'parse_configs' -> must match.
let (words, lowered, has_cjk) = query_context_tokens("parse_configs and related helpers");
assert!(!has_cjk);
assert!(!words.contains("parse_config"));
assert!(symbol_in_context("parse_config", &words, &lowered, has_cjk));
}
#[test]
fn empty_context_matches_nothing() {
let (words, lowered, has_cjk) = query_context_tokens("");
assert!(words.is_empty());
assert_eq!(lowered, "");
assert!(!has_cjk);
assert!(!symbol_in_context("foo", &words, &lowered, has_cjk));
}
#[test]
fn guard_counts_chars_not_bytes() {
// Python's len() counts characters. A 4-char name that is >3 in chars
// must take the substring fallback on an ASCII query even though a
// byte-length comparison would agree here; conversely a 3-char name
// must not, even when it is many bytes away from any CJK.
let (words, lowered, has_cjk) = query_context_tokens("prefer the runs_fast variant");
assert!(!has_cjk);
assert!(symbol_in_context("runs", &words, &lowered, has_cjk));
assert!(!symbol_in_context("run", &words, &lowered, has_cjk));
}
/// Symmetric pair of Python functions: identical raw importance signals,
/// so any score difference comes only from the context boost.
const CJK_BOOST_CODE: &str = "import os\n\n\
def run(config):\n value = config.get(\"alpha\")\n result = value + 1\n total = result * 2\n scaled = total - value\n merged = scaled + result\n print(merged)\n print(scaled)\n print(total)\n return merged\n\n\
def keep(config):\n value = config.get(\"beta\")\n result = value + 2\n total = result * 3\n scaled = total - value\n merged = scaled + result\n print(merged)\n print(scaled)\n print(total)\n return merged\n";
fn score_of(result: &CodeCompressionResult, name: &str) -> f64 {
result
.symbol_scores
.iter()
.find(|(k, _)| k == name)
.map(|(_, v)| *v)
.unwrap_or_else(|| panic!("no score for {name}: {:?}", result.symbol_scores))
}
#[test]
fn cjk_context_boosts_named_symbol_end_to_end() {
// Python reference: a CJK query with no spaces still boosts the ASCII
// symbol it names ("run" glued to CJK, len 3 <= guard, has_cjk relaxes it).
let c = CodeAwareCompressor::new(CodeCompressorConfig::default());
let r = c.compress_with(CJK_BOOST_CODE, Some("python"), "修复run函数的报错");
assert_eq!(score_of(&r, "run"), 1.0, "run must get the context boost");
assert_eq!(score_of(&r, "keep"), 0.0);
// ASCII query unchanged: "run" is only a substring of "runner" and the
// len>3 guard is NOT relaxed without CJK -> no boost, symmetric scores.
let r = c.compress_with(CJK_BOOST_CODE, Some("python"), "fix the runner");
assert_eq!(score_of(&r, "run"), 0.5);
assert_eq!(score_of(&r, "keep"), 0.5);
}
#[test]
fn empty_and_short_passthrough() {
let c = CodeAwareCompressor::new(CodeCompressorConfig::default());