mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(search-compressor): CJK-aware relevance + harden Rust/Python parity (#1749)
## Description
The search compressor's relevance scorer (`score_matches`, present in
both the Rust runtime path and the Python legacy mirror) split the query
on whitespace. A spaceless CJK query therefore matched a result line
only when the WHOLE query was a literal substring of that line — partial
overlaps never boosted relevant lines, so correct matches got dropped
when the result set was over budget.
This adds CJK character bigrams to the query match set, so a longer CJK
query boosts lines that share a substring. It also fixes two latent
Rust/Python parity divergences the ASCII-only fixtures had masked:
- **Length filter**: Rust counted word length in BYTES (`w.len()`),
Python in codepoints (`len(w)`), so a CJK word crossed the `> 2`
threshold differently. Rust now uses `chars().count()`.
- **Dedup**: Rust collected words into a `Vec` (no dedup), Python into a
`set`, so a repeated query word double-counted in Rust. Rust now uses a
`BTreeSet`.
Both scorers are byte-exact now; non-CJK output is unchanged (the 53
existing tests and the parity fixtures stay green).
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- `crates/headroom-core/src/transforms/search_compressor.rs` +
`headroom/transforms/search_compressor.py`: add
`is_cjk_char`/`_is_cjk_char` and `cjk_bigrams`/`_cjk_bigrams` (identical
ranges + logic), union CJK bigrams into the query match set, and align
the Rust word set to Python (`chars().count()` length, `BTreeSet`
dedup).
- `tests/test_search_compressor_cjk.py` + a Rust unit test: CJK bigram
extraction (same input/expected in both languages) and a CJK query
boosting a partially-overlapping line.
- Corrected a stale `_score_matches` docstring that referenced a
non-existent parity assertion; it now states honestly how the two sides
are pinned (test-equal for word-overlap + CJK bigrams; a few error-boost
keywords still diverge, fixed only Rust-side).
## Testing
- [x] Unit tests pass (`cargo test` + `pytest`)
- [x] Linting passes (`cargo clippy` / `cargo fmt` / `ruff` / `mypy`)
- [x] New tests added for new functionality
- [x] Manual testing performed (see Real Behavior Proof)
### Test Output
```text
$ cargo test -p headroom-core --lib search_compressor
test result: ok. 16 passed; 0 failed
$ .venv/bin/python -m pytest tests/test_search_compressor_cjk.py \
tests/test_transforms_search_compressor.py tests/test_search_compressor.py
55 passed # 2 new CJK tests + 53 existing (no regression)
```
## Real Behavior Proof
- Environment: macOS (Darwin), Rust via cargo, Python in a uv venv
(`_core` rebuilt on this branch), branch `feat/search-compressor-cjk`
off `main`.
- Exact command / steps: scored a CJK content line against a longer CJK
query whose whole form is not a substring of the line.
- Observed result: for content `src/a.py:10:认证令牌已过期需要重新登录` and query
`认证令牌缓存淘汰策略` (the whole query is NOT a substring of the line, but its
bigrams are), the line now scores `> 0` (bigrams 认证 / 证令 / 令牌 match);
before, it scored `0`. An ASCII-only line still scores `0`. All 53
existing search-compressor tests are unchanged. `cjk_bigrams("认证令牌")`
returns `{认证, 证令, 令牌}` in **both** Rust and Python.
## 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 relevance scoring)
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md — N/A: internal relevance-scoring
fix, no user-facing surface change
## Additional Notes
- The two parity divergences (byte-vs-codepoint length, `Vec`-vs-`set`
dedup) were pre-existing and only reachable with non-ASCII or
repeated-word queries — the all-ASCII fixtures never exercised them.
This PR brings both sides back to byte-exact for the word-overlap +
CJK-bigram scoring. The remaining error-boost keyword divergence is
pre-existing (fixed only Rust-side in the 3e.1 port) and is now
documented in the code rather than glossed over.
This commit is contained in:
parent
c85731dc23
commit
985621d60e
3 changed files with 122 additions and 6 deletions
|
|
@ -70,6 +70,36 @@ use md5::{Digest, Md5};
|
|||
|
||||
use crate::ccr::CcrStore;
|
||||
use crate::signals::{ImportanceContext, LineImportanceDetector};
|
||||
|
||||
/// True for CJK ideographs, kana, and Hangul. Code-point ranges kept
|
||||
/// byte-identical with the Python `_is_cjk_char` for search-compressor parity.
|
||||
fn is_cjk_char(c: char) -> bool {
|
||||
matches!(
|
||||
c as u32,
|
||||
0x3040..=0x30FF | 0x3400..=0x4DBF | 0x4E00..=0x9FFF | 0xAC00..=0xD7AF | 0xF900..=0xFAFF
|
||||
)
|
||||
}
|
||||
|
||||
/// CJK character bigrams from the CJK runs of a (lowercased) query, so a
|
||||
/// spaceless CJK query can match content. Mirrors the Python `_cjk_bigrams`.
|
||||
fn cjk_bigrams(text: &str) -> BTreeSet<String> {
|
||||
let mut out = BTreeSet::new();
|
||||
let mut run: Vec<char> = Vec::new();
|
||||
for c in text.chars() {
|
||||
if is_cjk_char(c) {
|
||||
run.push(c);
|
||||
} else {
|
||||
for w in run.windows(2) {
|
||||
out.insert(w.iter().collect::<String>());
|
||||
}
|
||||
run.clear();
|
||||
}
|
||||
}
|
||||
for w in run.windows(2) {
|
||||
out.insert(w.iter().collect::<String>());
|
||||
}
|
||||
out
|
||||
}
|
||||
use crate::transforms::adaptive_sizer::compute_optimal_k;
|
||||
|
||||
// ─── Types ──────────────────────────────────────────────────────────────
|
||||
|
|
@ -363,10 +393,15 @@ impl SearchCompressor {
|
|||
|
||||
pub fn score_matches(&self, files: &mut BTreeMap<String, FileMatches>, context: &str) {
|
||||
let context_lower = context.to_ascii_lowercase();
|
||||
let context_words: Vec<&str> = context_lower
|
||||
// Dedup like Python's `set`; count length in CHARS (not bytes) to match
|
||||
// Python codepoints; and add CJK char bigrams so a spaceless CJK query
|
||||
// (no whitespace words to split on) can still match content.
|
||||
let mut context_words: BTreeSet<String> = context_lower
|
||||
.split_whitespace()
|
||||
.filter(|w| w.len() > 2)
|
||||
.filter(|w| w.chars().count() > 2)
|
||||
.map(|w| w.to_string())
|
||||
.collect();
|
||||
context_words.extend(cjk_bigrams(&context_lower));
|
||||
|
||||
for fm in files.values_mut() {
|
||||
for m in &mut fm.matches {
|
||||
|
|
@ -374,7 +409,7 @@ impl SearchCompressor {
|
|||
let content_lower = m.content.to_ascii_lowercase();
|
||||
|
||||
for w in &context_words {
|
||||
if content_lower.contains(w) {
|
||||
if content_lower.contains(w.as_str()) {
|
||||
score += 0.3;
|
||||
}
|
||||
}
|
||||
|
|
@ -682,6 +717,14 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cjk_bigrams_from_runs() {
|
||||
let b = cjk_bigrams("认证令牌");
|
||||
assert!(b.contains("认证") && b.contains("证令") && b.contains("令牌") && b.len() == 3);
|
||||
assert!(cjk_bigrams("hello").is_empty());
|
||||
assert!(cjk_bigrams("a认b证").is_empty()); // isolated CJK chars -> no pair
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_ripgrep_context_line() {
|
||||
assert_eq!(
|
||||
|
|
|
|||
|
|
@ -50,6 +50,36 @@ from typing import Any, cast
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _is_cjk_char(c: str) -> bool:
|
||||
"""True for CJK ideographs, kana, and Hangul. Code-point ranges kept
|
||||
byte-identical with the Rust `is_cjk_char` for search-compressor parity."""
|
||||
o = ord(c)
|
||||
return (
|
||||
0x3040 <= o <= 0x30FF
|
||||
or 0x3400 <= o <= 0x4DBF
|
||||
or 0x4E00 <= o <= 0x9FFF
|
||||
or 0xAC00 <= o <= 0xD7AF
|
||||
or 0xF900 <= o <= 0xFAFF
|
||||
)
|
||||
|
||||
|
||||
def _cjk_bigrams(text: str) -> set[str]:
|
||||
"""CJK character bigrams from the CJK runs of a (lowercased) query, so a
|
||||
spaceless CJK query can match content. Mirrors the Rust `cjk_bigrams`."""
|
||||
out: set[str] = set()
|
||||
run: list[str] = []
|
||||
for c in text:
|
||||
if _is_cjk_char(c):
|
||||
run.append(c)
|
||||
else:
|
||||
for i in range(len(run) - 1):
|
||||
out.add(run[i] + run[i + 1])
|
||||
run = []
|
||||
for i in range(len(run) - 1):
|
||||
out.add(run[i] + run[i + 1])
|
||||
return out
|
||||
|
||||
|
||||
# ─── Public dataclasses (preserve existing import surface) ──────────────────
|
||||
|
||||
|
||||
|
|
@ -223,14 +253,21 @@ class SearchCompressor:
|
|||
|
||||
Stays Python so the legacy direct-call test surface keeps
|
||||
working without rebuilding through Rust on every test. The
|
||||
scoring constants must mirror Rust `SearchCompressor::score_matches`
|
||||
— Rust unit tests pin Rust's behavior; the parity assertion at
|
||||
the bottom of this module pins both sides agree.
|
||||
scoring constants mirror Rust `SearchCompressor::score_matches`,
|
||||
pinned by Rust unit tests and Python tests over the same inputs:
|
||||
word-overlap and CJK-bigram scoring are byte-equal. (The error-
|
||||
boost keyword set still diverges for a few terms fixed only on
|
||||
the Rust side -- see keyword_detector; there is no cross-impl
|
||||
assertion, so this equality is test-pinned, not mechanically
|
||||
enforced.)
|
||||
"""
|
||||
from headroom.transforms.error_detection import PRIORITY_PATTERNS_SEARCH
|
||||
|
||||
context_lower = context.lower()
|
||||
# Dedup whitespace words (len>2 by codepoints), and add CJK char bigrams
|
||||
# so a spaceless CJK query can match content.
|
||||
context_words = {w for w in context_lower.split() if len(w) > 2}
|
||||
context_words |= _cjk_bigrams(context_lower)
|
||||
|
||||
for fm in file_matches.values():
|
||||
for match in fm.matches:
|
||||
|
|
|
|||
36
tests/test_search_compressor_cjk.py
Normal file
36
tests/test_search_compressor_cjk.py
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
"""CJK-aware relevance scoring in the search compressor.
|
||||
|
||||
The relevance scorer tokenized the query on whitespace, so a spaceless CJK query
|
||||
matched content only when the WHOLE query was a literal substring of a line. CJK
|
||||
char bigrams now let a longer CJK query boost lines that share a substring. The
|
||||
Rust<->Python parity was also hardened (dedup like Python's set; char-length
|
||||
filter instead of bytes). These exercise the Python legacy scorer that mirrors
|
||||
Rust.
|
||||
"""
|
||||
|
||||
from headroom.transforms.search_compressor import (
|
||||
SearchCompressor,
|
||||
SearchCompressorConfig,
|
||||
_cjk_bigrams,
|
||||
)
|
||||
|
||||
|
||||
def test_cjk_bigrams_from_runs():
|
||||
assert _cjk_bigrams("认证令牌") == {"认证", "证令", "令牌"}
|
||||
assert _cjk_bigrams("hello world") == set() # ASCII -> no CJK bigrams
|
||||
assert _cjk_bigrams("a认b证") == set() # isolated CJK chars -> no adjacent pair
|
||||
|
||||
|
||||
def test_score_matches_cjk_query_bigrams_boost():
|
||||
compressor = SearchCompressor(SearchCompressorConfig(boost_errors=False, context_keywords=[]))
|
||||
content = "\n".join(
|
||||
[
|
||||
"src/a.py:10:认证令牌已过期需要重新登录",
|
||||
"src/b.py:2:plain ascii content here",
|
||||
]
|
||||
)
|
||||
parsed = compressor._parse_search_results(content)
|
||||
# the whole query is NOT a substring of the content line, but its bigrams are
|
||||
compressor._score_matches(parsed, "认证令牌缓存淘汰策略")
|
||||
assert parsed["src/a.py"].matches[0].score > 0 # 认证/证令/令牌 bigrams match
|
||||
assert parsed["src/b.py"].matches[0].score == 0
|
||||
Loading…
Add table
Add a link
Reference in a new issue