diff --git a/CLAUDE.md b/CLAUDE.md index d6871d1..b2bec0a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1025,6 +1025,32 @@ The items below remain open. `POST /outcome` is the answer for streamed traffic: it arrives afterwards, so it works identically either way. +5. **`local_encoder` noise isolation — built for the confirmed shapes; residuals below.** + The raw task's fenced code blocks, `Tool result:`-shaped lines, and closed + `` spans are now stripped by `_isolate_task_text` (pure, + stdlib-only, deterministic) before the fit + embed pass — `classify_zero_shot` + runs isolate → fit → prefix, so cleaning happens first and a noisy task often + fits the token window outright. Measured on the real `BAAI/bge-large-en-v1.5` + (2026-09-19): 8 clean one-sentence tasks scored 8/8, but the same instructions + wrapped in that noise scored 2/8 with the truncation fix already in place — + and tail-biased windowing ALONE also scored 2/8 on long noisy pairs, so the + fit does not subsume isolation. Post-isolation: 8/8 on short and long noisy + pairs, clean-vs-noisy pair consistency 8/8 + 8/8; `'[code]'`/`'[elided]'` + placeholder tokens measured worse than pure removal (7/8, 6/8 on long pairs) + and were rejected; a 20%-ratio floor guard measured harmful (3/8 — it reverts + exactly the short noisy inputs isolation exists to fix) in favor of an + absolute 24-char floor that only catches near-all-code inputs. Offline + regression: a noise tripwire (same instruction bare vs wrapped must classify + identically) fails against pre-isolation code and passes after. Still not + built: attention-masking de-weighting as an alternative to stripping (it + would preserve the noise tokens' presence without letting them dominate); + unclosed `` tags and un-fenced diff hunks are left in place + (only closed-tag spans, fenced blocks, and marker-prefixed lines are + stripped); and a code-grounded instruction whose pasted snippet is the + subject can still land on a near-category — measured 3/4 on a 4-task + grounded set, the residual miss being description similarity + ("refactor this helper" + code → `debugging`), not noise dominance. + ## Gaming mode, and the backoff that used to do nothing **The classifier circuit breaker did not break the circuit.** diff --git a/src/local_encoder.py b/src/local_encoder.py index ce327b1..ae189ae 100644 --- a/src/local_encoder.py +++ b/src/local_encoder.py @@ -34,7 +34,10 @@ so the router itself has no UI dependency." A deployment that never selects from __future__ import annotations -from typing import Any, Optional +import re +from typing import Any, Final + +import logs _MISSING_DEPENDENCY_MESSAGE = ( "classifier.mode is 'local_encoder' but the 'transformers'/'torch' " @@ -87,6 +90,228 @@ _BGE_QUERY_INSTRUCTION = "Represent this sentence for searching relevant passage # confidence being a genuine probability, not a monotonic score. _SOFTMAX_TEMPERATURE = 0.10 +# ── Task-text isolation (structural agent-session noise) ────────────── +# +# Production task text is raw user_content (dispatcher passes the RAW +# task, not the clamped/framed user_content): fenced code blocks, tool +# call/result stanzas, and system-reminder-style wrappers surround the +# prose instruction. Mean pooling weights every token equally, so when +# noise tokens outnumber the instruction words the pooled embedding — +# and therefore the verdict — describes the noise, not the intent. +# Observed on the real BAAI/bge-large-en-v1.5 (2026-09-19): 8 clean +# one-sentence tasks scored 8/8 correct, but the same instructions +# wrapped in a fenced code block + a "Tool result:" line + a +# tag classified 2/8 — every non-summarization +# instruction flipped to a file/code-flavored category. Tail-biased +# windowing (the fit below) does NOT fix this: measured 2/8 on long +# noisy inputs too, because the window keeps code and instruction +# mixed. The helper below strips that structural noise BEFORE the fit + +# embed pass; cleaning first also means a noisy task often fits the +# window outright, so the fit sees prose, not a code wall. +# +# Deliberately conservative, pure, stdlib-only (``re``), deterministic, +# and free of false strip targets for ordinary prose: inline code spans +# (``like this``) are KEPT — instructions reference real file names and +# commands, which are signal, not noise. + +# Fenced code blocks: strip the content outright. Empirically decided +# placeholder vs pure removal on the real model (2026-09-19, 8 noisy +# pairs x short+long): removal scored 8/8 on both, while '[code]' and +# '[elided]' placeholders scored 7/8 and 6/8 on long pairs — the +# placeholder token itself pulls the pooled embedding toward code- +# flavored categories and costs accuracy, so the winner is removal. +_FENCE_PAIR_RE: Final = re.compile(r"```.*?```", re.DOTALL) +_FENCE_OPEN_TAIL_RE: Final = re.compile(r"```.*\Z", re.DOTALL) + +# System-reminder-style wrapper tags: agent harnesses inject session +# metadata inside custom tags; the content is context, not the task. +# Closed-tag spans only — an unclosed tag is left alone rather than +# guessing where it ends (the floor guard still protects the verdict). +_ISOLATE_WRAPPER_TAGS: Final[tuple[str, ...]] = ("system-reminder",) +_WRAPPER_TAG_RE: Final = re.compile( + r"<({})\b[^>]*>.*?".format("|".join(_ISOLATE_WRAPPER_TAGS)), + re.DOTALL | re.IGNORECASE, +) + +# Tool-result lines: "Tool result: file read successfully, 12 lines" and +# the call/output variants seen in agent-session transcripts. Line-level +# and conservative: the line must START with the marker (case- +# insensitive), so prose merely mentioning a tool result survives. +_TOOL_LINE_RE: Final = re.compile( + r"^[ \t]*tool[\s_-]*(?:result|output|call)[\s_-]*:[^\n]*\n?", + re.IGNORECASE | re.MULTILINE, +) + +# Floor guard: if stripping leaves (almost) nothing, the input was +# itself one big code block — e.g. a bare file dump with no instruction. +# Embedding a bare placeholder would score a meaningless stub, so fall +# back to the original text, whose code content is then the only signal +# available. Measured (2026-09-19, real model): code-GROUNDED tasks +# (instruction + pasted code) classify BETTER stripped even though +# stripping keeps well under 20% of the characters, so the guard keys on +# an absolute character floor only — a ratio guard would revert exactly +# the noise-wrapped short inputs the isolation exists to fix. +_ISOLATE_MIN_CHARS: Final[int] = 24 + + +def _isolate_task_text(task: str) -> str: + """Strip structural agent-session noise from *task* before embedding. + + Removes fenced code block bodies, system-reminder-style wrapper tag + spans, and tool call/result lines — the shapes that dominate a + mean-pooled embedding when they outnumber the prose instruction. + Inline code spans and ordinary prose are kept. Pure string in, + string out; no I/O, no state. + + If stripping leaves fewer than _ISOLATE_MIN_CHARS — the input was + essentially only noise (a bare code dump, no instruction) — the + original text is returned: the code content is then the only signal + there is, and a stripped stub would be a worse embedding target. + + Idempotent: already-isolated text contains none of the stripped + shapes, so a second pass returns it unchanged. + """ + if not task: + return task + text = _FENCE_PAIR_RE.sub("", task) + text = _FENCE_OPEN_TAIL_RE.sub("", text) + text = _WRAPPER_TAG_RE.sub("", text) + text = _TOOL_LINE_RE.sub("", text) + if len(text) < _ISOLATE_MIN_CHARS: + return task + return text + + +# ── Token-budget fitting (pre-tokenization window selection) ────────── +# +# The final ``tokenizer(...)`` call in ``classify_zero_shot`` passes +# ``truncation=True`` with no ``max_length``, so HuggingFace truncates to +# the tokenizer's ``model_max_length`` — 512 tokens for the BERT-family +# BGE models — and, with the default ``truncation_side='right'``, KEEPS +# the head. The head of an agent-session prompt is thousands of chars of +# shared boilerplate while the task-specific content lives in the tail, +# so every long request scored the same truncated prefix and returned the +# same verdict (observed live on 2026-09-18: six distinct 6.4k-6.6k-char +# requests, all classified ``diff_checking`` at confidence 0.192363). +# The helpers below fit the task to the real token window BEFORE that +# final tokenize, choosing ourselves which tokens survive. + +# Tail-biased head+tail split: when the task overflows the budget, the +# head slice keeps this share of the content budget and the tail gets the +# rest. A head slice survives at all because the raw task text can carry +# its instruction at either end ("Translate this: " vs " — +# translate this") — the same rationale as dispatcher.clamp_for_classifier, +# and it applies here directly because the local_encoder path receives the +# RAW task (dispatcher._classify_via_configured_mode), not the clamped/ +# framed user_content the chat modes see. The head share stays small +# because the live failure mode was the opposite collapse: shared head +# boilerplate crowding out the tail. +_FIT_HEAD_FRACTION: Final[float] = 0.25 + +# Marker spliced between the kept head and tail, mirroring +# dispatcher.clamp_for_classifier's elision convention. Its own token +# cost is measured per tokenizer and reserved from the budget. +_FIT_ELISION_MARKER: Final[str] = "\n\n[... content elided ...]\n\n" + +# decode() -> re-tokenize is not guaranteed token-identical for every +# input (unicode normalization, odd graphemes), so the fitted text is +# built a few tokens UNDER budget. If drift ever exceeds the margin, the +# final ``truncation=True`` clips a token or two off the tail — a +# bounded, graceful degradation, never a crash. +_FIT_SAFETY_MARGIN_TOKENS: Final[int] = 8 + +# model_max_length values some tokenizers report when the true window is +# unknown (HF's VERY_LARGE_INTEGER is int(1e30)). Unusable for budget +# math; falls back to the documented constant below. +_MODEL_MAX_LENGTH_SENTINEL: Final[int] = int(1e30) + +# Conservative window used when model_max_length is missing or a +# sentinel: 512 is the BERT-family window the BGE models actually have. +_FALLBACK_MODEL_MAX_LENGTH: Final[int] = 512 + +# Set once the fallback above has been logged, so a misconfigured +# tokenizer logs one warning instead of one per request. +_budget_fallback_logged = False + + +def _special_tokens_allowance(tokenizer: Any) -> int: + """Tokens the final tokenize adds around the text ([CLS]/[SEP] for + BERT-family tokenizers, including the BGE models).""" + add_fn = getattr(tokenizer, "num_special_tokens_to_add", None) + if add_fn is None: + return 2 + return max(0, int(add_fn())) + + +def _usable_token_budget(tokenizer: Any) -> int: + """Tokens available for task content under the model's real window. + + usable = model_max_length - len(BGE instruction prefix tokens) + - special tokens added by the final tokenize. + + A missing or sentinel ``model_max_length`` falls back to + _FALLBACK_MODEL_MAX_LENGTH (logged once) rather than raising — an + unusual tokenizer must degrade conservatively, not kill the request. + """ + global _budget_fallback_logged + max_len = getattr(tokenizer, "model_max_length", None) + if ( + not isinstance(max_len, int) + or max_len <= 0 + or max_len >= _MODEL_MAX_LENGTH_SENTINEL + ): + if not _budget_fallback_logged: + _budget_fallback_logged = True + logs.warning( + "encoder_token_budget_fallback", + reported_model_max_length=str(max_len), + fallback=_FALLBACK_MODEL_MAX_LENGTH, + ) + max_len = _FALLBACK_MODEL_MAX_LENGTH + prefix_ids = tokenizer( + _BGE_QUERY_INSTRUCTION, add_special_tokens=False, + )["input_ids"] + return max_len - len(list(prefix_ids)) - _special_tokens_allowance(tokenizer) + + +def _fit_task_to_token_budget(task: str, tokenizer: Any) -> str: + """Fit *task* into the model's token window, biased toward the tail. + + Tokenizes the bare task (no special tokens, no truncation) and, when + it overflows the usable budget, keeps a small head slice plus a + tail-majority window (see _FIT_HEAD_FRACTION for why a head slice + survives at all), spliced with the elision marker. The window is + selected on real token ids and decoded back to text — no + chars-per-token arithmetic anywhere. + + Returns *task* unchanged when it already fits (including the empty + string). Never mutates tokenizer state — the cached tokenizer is + shared across requests — and the final ``truncation=True`` in + classify_zero_shot stays as a last-resort backstop. + """ + budget = _usable_token_budget(tokenizer) + if budget <= 0: + # Pathological window (smaller than prefix + specials): leave the + # task alone; the final tokenize does whatever it can. + return task + ids = list(tokenizer(task, add_special_tokens=False)["input_ids"]) + if len(ids) <= budget: + return task + marker_tokens = len(list( + tokenizer(_FIT_ELISION_MARKER, add_special_tokens=False)["input_ids"] + )) + content_budget = budget - marker_tokens - _FIT_SAFETY_MARGIN_TOKENS + if content_budget <= 0: + # Degenerate window: no room for the marker or the margin — keep + # a pure tail, which is where task-specific content lives. + return tokenizer.decode(ids[-budget:], skip_special_tokens=True) + head_n = int(content_budget * _FIT_HEAD_FRACTION) + tail_n = content_budget - head_n + head_text = tokenizer.decode(ids[:head_n], skip_special_tokens=True) + tail_text = tokenizer.decode(ids[-tail_n:], skip_special_tokens=True) + return head_text + _FIT_ELISION_MARKER + tail_text + + # Cached loaded model + tokenizer, keyed by (model_id, device) so a config # change to either picks up a fresh instance rather than silently reusing # one built for a different model. @@ -246,6 +471,15 @@ def classify_zero_shot( Categories are matched under their natural-language description (``_CATEGORY_DESCRIPTIONS``), never the raw config identifier — see that map's docstring for why. + + Long inputs are first fitted to the model's real token window (see + ``_fit_task_to_token_budget``): the tokenizer's own silent truncation + would otherwise keep only the shared head of a long prompt and drop + the task-specific tail. Before that, structural agent-session noise + (fenced code blocks, tool call/result lines, system-reminder-style + tags) is stripped (see ``_isolate_task_text``): mean pooling weights + every token equally, so unwrapped noise outnumbers the instruction + and the verdict scores the noise, not the intent. """ import torch @@ -257,8 +491,16 @@ def classify_zero_shot( model, tokenizer, tensor_device, model_id, device, categories, ) - # Encode the input task with BGE's query instruction prefix - prefixed_task = _BGE_QUERY_INSTRUCTION + task + # Isolate then fit, in that order: stripping the noise first means + # the window sees prose and a noisy task often fits outright; the + # fit then protects whatever long prose survives. Without isolation + # the tokenizer below silently right-truncates to model_max_length + # (512 for the BGE models), which for long agent-session prompts + # scores only the shared head boilerplate and discards the + # task-specific tail: every long request returns the same verdict. + isolated_task = _isolate_task_text(task) + fitted_task = _fit_task_to_token_budget(isolated_task, tokenizer) + prefixed_task = _BGE_QUERY_INSTRUCTION + fitted_task encoded = tokenizer( prefixed_task, padding=True, diff --git a/tests/test_local_encoder.py b/tests/test_local_encoder.py index 91334a9..09b66ac 100644 --- a/tests/test_local_encoder.py +++ b/tests/test_local_encoder.py @@ -267,29 +267,64 @@ def _install_fake_modules(): class MockTokenizer: """Acts as transformers.AutoTokenizer. - Always returns fixed ``input_ids`` and unit ``attention_mask``, - both as ``MockTensor``. + Emulates the parts of the HF contract the encoder path relies + on: a known ``model_max_length`` (512, BERT/BGE family), + right-side truncation when called with ``truncation=True`` and + no ``max_length`` (HF's default ``truncation_side='right'``), + and plain Python id lists when called without ``return_tensors`` + (the real BatchEncoding behaves the same). Ids are word lengths, + so token counts stay content-dependent. """ + model_max_length = 512 + truncation_side = "right" + @classmethod def from_pretrained(cls, model_id, **kwargs): return cls() - def __call__(self, text, **_kw): - # Return a single token with attention_mask = 1 + def __call__( + self, text, add_special_tokens=True, truncation=False, + padding=True, return_tensors=None, max_length=None, **_kw, + ): + ids = [len(word) for word in str(text).split()] + if truncation: + keep = self.model_max_length - (2 if add_special_tokens else 0) + ids = ids[:keep] + if add_special_tokens: + ids = [0] + ids + [1] # [CLS] ... [SEP] + input_ids = MockTensor([ids]) if return_tensors == "pt" else ids return { - "input_ids": MockTensor([[1]]), - "attention_mask": MockTensor([[1]]), + "input_ids": input_ids, + "attention_mask": MockTensor([[1] * len(ids)]), } + def num_special_tokens_to_add(self) -> int: + return 2 + + def decode(self, ids, skip_special_tokens=True): + kept = [i for i in ids if not (skip_special_tokens and i in (0, 1))] + return " ".join(str(i) for i in kept) + # ── Mock torch ───────────────────────────────────────────────── + def _flat_values(x): + """Read a 1-D value list from a MockTensor of either orientation. + + torch.tensor(flat_list) makes a (1, n) ROW, while softmax/argmax + results elsewhere are (n, 1) COLUMNS; the old readers only handled + the column shape, so a multi-category row silently collapsed to + its first value and every classify returned (category_0, 1.0). + """ + if isinstance(x, MockTensor): + if len(x._data) == 1: + return list(x._data[0]) + return [row[0] for row in x._data] + return list(x) + def _softmax(x, dim=0): """Compute softmax over dim 0 using MockTensor values.""" - if isinstance(x, MockTensor): - vals = [row[0] for row in x._data] if x._data else [] - else: - vals = list(x) + vals = _flat_values(x) max_v = max(vals) exps = [math.exp(v - max_v) for v in vals] total = sum(exps) @@ -297,12 +332,8 @@ def _install_fake_modules(): return MockTensor([[p] for p in probs]) def _argmax(x, dim=None): - if isinstance(x, MockTensor): - vals = [row[0] for row in x._data] if x._data else [] - else: - vals = list(x) - idx = vals.index(max(vals)) - return MockTensor([idx]) + vals = _flat_values(x) + return MockTensor([vals.index(max(vals))]) def _make_tensor(data): if isinstance(data, MockTensor): @@ -344,9 +375,11 @@ def _clear_caches(): """All caches in local_encoder must be empty between tests.""" local_encoder._model_cache.clear() local_encoder._desc_cache.clear() + local_encoder._budget_fallback_logged = False yield local_encoder._model_cache.clear() local_encoder._desc_cache.clear() + local_encoder._budget_fallback_logged = False @pytest.fixture @@ -648,4 +681,529 @@ def test_desc_cache_extended_for_new_categories(): cached_after = local_encoder._desc_cache.get(("m", "cpu")) assert cached_after is not None assert "coding_general" in cached_after - assert "debugging" in cached_after \ No newline at end of file + assert "debugging" in cached_after + + +# ── Token-budget fitting (the 2026-09-18 truncation regression) ──────── +# +# classify_zero_shot's final tokenizer call used truncation=True without +# max_length, so HuggingFace clipped the prefixed task to the tokenizer's +# model_max_length from the RIGHT: for BGE (BERT family, 512) a 6000+ +# char agent-session prompt became ~510 tokens of its shared head +# boilerplate, the task-specific tail never reached the model, and every +# long request returned the same verdict. Production evidence: six +# distinct requests (6425-6572 chars) all classified diff_checking at +# confidence 0.192363 for 20+ minutes. +# +# The fakes below emulate the HF contract faithfully enough that the +# regression test FAILS against the pre-fix code and PASSES after. + + +class SessionStyleTokenizer: + """Fake HF tokenizer with the contract that broke production. + + - known ``model_max_length`` (512, the BERT/BGE window); + - right-side truncation when called with ``truncation=True`` and no + ``max_length`` (HF's default ``truncation_side='right'`` keeps the + head and silently drops the tail); + - plain Python id lists when called without ``return_tensors``; + - content-dependent ids: one id per distinct word, allocated in + first-seen order, with a ``decode`` that reconstructs text. + """ + + def __init__(self, model_max_length: int = 512) -> None: + self.model_max_length = model_max_length + self.truncation_side = "right" + self._vocab: dict = {} + self._words: dict = {} + + @classmethod + def from_pretrained(cls, model_id, **kwargs): + return cls() + + def _ids_for(self, text: str) -> list: + ids = [] + for word in text.split(): + if word not in self._vocab: + next_id = len(self._vocab) + 2 # 0=[CLS], 1=[SEP] + self._vocab[word] = next_id + self._words[next_id] = word + ids.append(self._vocab[word]) + return ids + + def __call__( + self, text, add_special_tokens=True, truncation=False, + padding=True, return_tensors=None, max_length=None, **_kw, + ): + ids = self._ids_for(str(text)) + if truncation: + keep = self.model_max_length - (2 if add_special_tokens else 0) + ids = ids[:keep] + if add_special_tokens: + ids = [0] + ids + [1] + input_ids = MockTensor([ids]) if return_tensors == "pt" else ids + return { + "input_ids": input_ids, + "attention_mask": MockTensor([[1] * len(ids)]), + } + + def num_special_tokens_to_add(self) -> int: + return 2 + + def decode(self, ids, skip_special_tokens=True): + words = [ + self._words[i] + for i in ids + if not (skip_special_tokens and i in (0, 1)) + ] + return " ".join(words) + + +class ContentSignModel: + """Embeds a token sequence as the SIGN of its last content token id. + + Provably identical verdicts for inputs that truncate to the same + prefix (same last surviving token), and opposite sides of the + centroid pair for the regression test's two tails: with 100 alpha + words vs 101 beta words after the same boilerplate, the last content + token takes ids of opposite parity by construction. The seeded + [[1.0]] (coding_general) / [[-1.0]] (debugging) descriptions turn + that sign into a category flip. + """ + + @classmethod + def from_pretrained(cls, model_id, **kwargs): + return cls() + + def to(self, device_str): + return self + + def __call__(self, **kwargs): + ids = kwargs["input_ids"] + if isinstance(ids, MockTensor): + ids = ids._data[0] + content = [i for i in ids if i not in (0, 1)] + last = float(content[-1]) if content else 0.0 + sign = 1.0 if last % 2 == 0 else -1.0 + return types.SimpleNamespace(last_hidden_state=MockTensor([[sign]])) + + +def _session_style_prompt(tail_words: list) -> str: + """An agent-session-shaped prompt: huge shared head, varying tail.""" + boilerplate = " ".join( + f"session guideline {i} assume tools available" for i in range(400) + ) + return f"{boilerplate} {' '.join(tail_words)}" + + +def test_long_shared_prefix_tasks_do_not_collapse_to_identical_verdict(): + """THE regression tripwire: two long prompts sharing thousands of + chars of identical head boilerplate but differing only in tail + content must NOT produce the same (category, confidence). + + Pre-fix, both prefixed tasks right-truncated to the same 510 shared + boilerplate tokens -> same embedding -> same verdict, exactly the + production signature. This fails against pre-fix code by + construction: the tail words sit beyond token 510 and are never seen + by the model. Post-fix the tail survives the fit and the verdicts + separate. + """ + tail_a = [f"alpha{i}" for i in range(100)] + tail_b = [f"beta{i}" for i in range(101)] + task_a = _session_style_prompt(tail_a) + task_b = _session_style_prompt(tail_b) + assert len(task_a) > 2000 and len(task_b) > 2000 + assert task_a.startswith("session guideline 0") == task_b.startswith( + "session guideline 0" + ) + assert task_a[: -len(" ".join(tail_a)) - 1] == ( + task_b[: -len(" ".join(tail_b)) - 1] + ) + + sys.modules["transformers"].AutoTokenizer = SessionStyleTokenizer + sys.modules["transformers"].AutoModel = ContentSignModel + local_encoder._desc_cache[("m", "cpu")] = { + "coding_general": MockTensor([[1.0]]), + "debugging": MockTensor([[-1.0]]), + } + + cat_a, conf_a = local_encoder.classify_zero_shot( + task_a, ["coding_general", "debugging"], model_id="m", device="cpu", + ) + cat_b, conf_b = local_encoder.classify_zero_shot( + task_b, ["coding_general", "debugging"], model_id="m", device="cpu", + ) + assert (cat_a, conf_a) != (cat_b, conf_b) + assert cat_a != cat_b + + +def test_fit_task_noop_when_under_budget(): + """A task within the usable budget passes through unchanged.""" + tok = SessionStyleTokenizer(model_max_length=512) + short = "summarize this paragraph carefully" + assert local_encoder._fit_task_to_token_budget(short, tok) == short + + +def test_fit_task_over_budget_keeps_head_and_tail_drops_middle(): + """Over budget: the fitted text keeps a head slice and the tail, + drops the middle, splices the elision marker, and the re-tokenized + prefixed result fits the model window.""" + tok = SessionStyleTokenizer(model_max_length=512) + task = ( + " ".join(f"head{i}" for i in range(200)) + + " " + " ".join(f"mid{i}" for i in range(2000)) + + " " + " ".join(f"tail{i}" for i in range(200)) + ) + fitted = local_encoder._fit_task_to_token_budget(task, tok) + assert fitted != task + assert "head0" in fitted and "head121" in fitted + assert "tail199" in fitted and "mid1999" in fitted + assert "mid1000" not in fitted + assert "[... content elided ...]" in fitted + total = tok(local_encoder._BGE_QUERY_INSTRUCTION + fitted) + assert len(total["input_ids"]) <= tok.model_max_length + + +def test_fit_budget_subtracts_prefix_and_special_tokens(): + """The usable budget is model_max_length - prefix tokens - specials: + a task of exactly that many tokens is a no-op, one token more is + fitted, and the fitted result fits the real window.""" + tok = SessionStyleTokenizer(model_max_length=512) + prefix_words = len(local_encoder._BGE_QUERY_INSTRUCTION.split()) + budget = ( + tok.model_max_length + - prefix_words + - tok.num_special_tokens_to_add() + ) + exact = " ".join(f"w{i}" for i in range(budget)) + assert local_encoder._fit_task_to_token_budget(exact, tok) == exact + + over = exact + " overflow" + fitted = local_encoder._fit_task_to_token_budget(over, tok) + assert fitted != over + total = tok(local_encoder._BGE_QUERY_INSTRUCTION + fitted) + assert len(total["input_ids"]) <= tok.model_max_length + + +@pytest.mark.parametrize("bad_max_length", [None, int(1e30), 0, -5]) +def test_fit_falls_back_when_model_max_length_unusable(bad_max_length): + """A missing/sentinel/non-positive model_max_length falls back to the + conservative 512 window (logged once via the module flag) instead of + blowing up, and the fitted result still fits 512.""" + tok = SessionStyleTokenizer(model_max_length=512) + tok.model_max_length = bad_max_length + task = " ".join(f"w{i}" for i in range(4000)) + fitted = local_encoder._fit_task_to_token_budget(task, tok) + assert fitted != task + assert local_encoder._budget_fallback_logged + total = tok(local_encoder._BGE_QUERY_INSTRUCTION + fitted) + assert len(total["input_ids"]) <= local_encoder._FALLBACK_MODEL_MAX_LENGTH + + +def test_fit_degenerate_window_returns_task_unchanged(): + """A window smaller than prefix + specials leaves the task alone + rather than producing a nonsensical slice.""" + tok = SessionStyleTokenizer(model_max_length=4) # < prefix(7) + 2 + task = "a b c" + assert local_encoder._fit_task_to_token_budget(task, tok) == task + + +def test_classify_zero_shot_leaves_cached_tokenizer_state_untouched(): + """The fit must not mutate the shared cached tokenizer (no + truncation_side flips) — the cache is process-wide.""" + sys.modules["transformers"].AutoTokenizer = SessionStyleTokenizer + sys.modules["transformers"].AutoModel = ContentSignModel + local_encoder.classify_zero_shot( + _session_style_prompt([f"alpha{i}" for i in range(100)]), + ["coding_general", "debugging"], + model_id="m", device="cpu", + ) + tok = local_encoder._model_cache[("m", "cpu")][1] + assert tok.truncation_side == "right" + assert tok.model_max_length == 512 + + +# ── Task-text isolation (the 2026-09-19 noise-dominance regression) ──── +# +# classify_zero_shot embedded the RAW task: a prose instruction wrapped +# in agent-session noise (fenced code block, "Tool result:" line, +# tag) mean-pooled the noise tokens with equal weight, +# and the verdict scored the noise — measured on the real +# BAAI/bge-large-en-v1.5, 8 clean tasks scored 8/8 but the same +# instructions wrapped scored 2/8, on SHORT inputs far under any +# truncation limit (a second, distinct bug; the fit fix does not touch +# it — windowing alone measured 2/8 on long noisy pairs too). +# +# The fakes below reproduce that mechanism in miniature: tokens drawn +# from the code/structural vocabulary push the pooled embedding toward +# the code centroid, prose tokens toward the prose centroid, so +# noise-count dominance flips the verdict exactly like the real model. + +_INSTRUCTION = ( + "Please write a clear docstring and a short readme section " + "explaining the storage module." +) + +_NOISE_REMINDER = ( + "\nsession context: agent harness active\n" + "" +) +_NOISE_CODE = ( + "```python\ndef parse_line(line):\n" + " tokens = line.strip().split(',')\n" + " if not tokens:\n" + " return []\n" + " return [int(t) for t in tokens if t.isdigit()]\n" + "FILES = {'a': 1, 'b': 2}\n```" +) +_NOISE_TOOL = "Tool result: file read successfully, 12 lines" + + +def _norm(word: str) -> str: + """Normalize a whitespace token to its word characters (lowercased).""" + return "".join(ch for ch in word.lower() if ch.isalnum() or ch == "_") + + +# Every multi-char word of the noise fixture is a "code/structural" +# vocabulary entry — the fake tokenizer's stand-in for "this token +# embeds near code", mirroring how real code/tool tokens pull the +# pooled embedding toward code-flavored categories. +_CODE_VOCAB = frozenset( + w + for w in ( + _norm(tok) + for tok in (_NOISE_REMINDER + " " + _NOISE_CODE + " " + _NOISE_TOOL).split() + ) + if len(w) > 1 +) + + +class NoiseDominantTokenizer: + """Fake HF tokenizer whose ids encode word class. + + Words in _CODE_VOCAB get ids >= 1000; prose words get ids from 2 up + (first-seen, like SessionStyleTokenizer). Otherwise the same HF + contract: known model_max_length, right-side truncation under + ``truncation=True`` without ``max_length``, plain id lists without + ``return_tensors``, and a decode that reconstructs text. + """ + + def __init__(self, model_max_length: int = 512) -> None: + self.model_max_length = model_max_length + self.truncation_side = "right" + self._vocab: dict = {} + self._words: dict = {} + + @classmethod + def from_pretrained(cls, model_id, **kwargs): + return cls() + + def _ids_for(self, text: str) -> list: + ids = [] + for word in str(text).split(): + norm = _norm(word) + if not norm: + continue + if norm not in self._vocab: + if norm in _CODE_VOCAB: + next_id = 1000 + len(self._vocab) + else: + next_id = 2 + len(self._vocab) + self._vocab[norm] = next_id + self._words[next_id] = word + ids.append(self._vocab[norm]) + return ids + + def __call__( + self, text, add_special_tokens=True, truncation=False, + padding=True, return_tensors=None, max_length=None, **_kw, + ): + ids = self._ids_for(text) + if truncation: + keep = self.model_max_length - (2 if add_special_tokens else 0) + ids = ids[:keep] + if add_special_tokens: + ids = [0] + ids + [1] + input_ids = MockTensor([ids]) if return_tensors == "pt" else ids + return { + "input_ids": input_ids, + "attention_mask": MockTensor([[1] * len(ids)]), + } + + def num_special_tokens_to_add(self) -> int: + return 2 + + def decode(self, ids, skip_special_tokens=True): + words = [ + self._words[i] + for i in ids + if not (skip_special_tokens and i in (0, 1)) + ] + return " ".join(words) + + +class NoiseDominantModel: + """Embeds a token sequence as the MEAN of per-token class values. + + +1.0 for code-vocabulary ids (>= 1000), -1.0 for prose ids: mean + pooling weights every token equally, so whichever class supplies + more tokens wins the pooled sign — the miniature of the real + failure, where code/tool noise outnumbers the instruction words and + the verdict follows the noise. + """ + + @classmethod + def from_pretrained(cls, model_id, **kwargs): + return cls() + + def to(self, device_str): + return self + + def __call__(self, **kwargs): + ids = kwargs["input_ids"] + if isinstance(ids, MockTensor): + ids = ids._data[0] + content = [i for i in ids if i not in (0, 1)] + values = [1.0 if i >= 1000 else -1.0 for i in content] + mean = sum(values) / len(values) if values else 0.0 + return types.SimpleNamespace(last_hidden_state=MockTensor([[mean]])) + + +def _noisy_instruction() -> str: + """The clean instruction wrapped in the operator's repro noise: + system-reminder tag + fenced code block + tool-result line, the + instruction still clearly present at the END.""" + return "\n\n".join([_NOISE_REMINDER, _NOISE_CODE, _NOISE_TOOL, _INSTRUCTION]) + + +def test_noise_wrapped_task_classifies_same_as_clean_pair(): + """THE noise tripwire: the same prose instruction, bare and wrapped + in agent-session noise, must classify to the SAME category. + + Fails against pre-isolation code by construction: the wrapped task's + code/structural tokens outnumber the prose tokens, so the pooled + mean goes positive (code centroid) while the bare instruction pools + negative (prose centroid) — a category flip, the miniature of the + real-model observation. Post-isolation the noise shapes are stripped + before embedding and both tasks pool to the prose centroid. + """ + clean = _INSTRUCTION + noisy = _noisy_instruction() + # Precondition: the fixture is short — windowing cannot apply, so + # this pair exercises isolation, not the fit. + assert len(noisy) < 1000 + assert not (_CODE_VOCAB & {_norm(w) for w in clean.split()}) + + sys.modules["transformers"].AutoTokenizer = NoiseDominantTokenizer + sys.modules["transformers"].AutoModel = NoiseDominantModel + local_encoder._desc_cache[("m", "cpu")] = { + "docs_writing": MockTensor([[-1.0]]), + "coding_general": MockTensor([[1.0]]), + } + + cat_clean, _ = local_encoder.classify_zero_shot( + clean, ["docs_writing", "coding_general"], model_id="m", device="cpu", + ) + cat_noisy, _ = local_encoder.classify_zero_shot( + noisy, ["docs_writing", "coding_general"], model_id="m", device="cpu", + ) + assert cat_clean == "docs_writing" + assert cat_noisy == cat_clean, ( + f"noise flipped the verdict: clean={cat_clean}, noisy={cat_noisy}" + ) + + +def test_isolate_strips_fenced_code_blocks(): + """A fenced code block is removed outright — body AND fences.""" + task = "Summarize the setup steps.\n```python\nimport os\nos.getcwd()\n```" + out = local_encoder._isolate_task_text(task) + assert "import os" not in out + assert "os.getcwd()" not in out + assert "```" not in out + + +def test_isolate_removal_decision_no_placeholder(): + """The eval-decided placeholder-vs-removal winner is REMOVAL: the + stripped text must not carry any '[code]'-style placeholder — the + placeholder token itself pulled long noisy pairs to 6/8 vs removal's + 8/8 on the real model.""" + task = ( + "Explain what this algorithm computes and when you would use it.\n" + "```\nfor i in range(3):\n print(i)\n```" + ) + out = local_encoder._isolate_task_text(task) + assert "[code]" not in out + assert "[elided]" not in out + assert "for i in range(3)" not in out + + +def test_isolate_strips_tool_result_lines(): + """'Tool result:'-shaped lines are stripped line-level; result and + output/call variants share the marker shape.""" + task = ( + "Tool result: file read successfully, 12 lines\n" + "Tool output: exit code 0\n" + "Now summarize the findings above." + ) + out = local_encoder._isolate_task_text(task) + assert "Tool result:" not in out + assert "Tool output:" not in out + assert "exit code 0" not in out + assert "summarize the findings" in out + + +def test_isolate_strips_system_reminder_content(): + """ spans are removed with their content; the tag + itself does not survive either.""" + task = ( + "session id 42, tools enabled\n" + "Translate the attached paragraph to French." + ) + out = local_encoder._isolate_task_text(task) + assert "session id 42" not in out + assert "system-reminder" not in out + assert "Translate the attached paragraph" in out + + +def test_isolate_preserves_trailing_instruction(): + """The prose instruction survives verbatim when it trails the + noise — the operator's repro shape.""" + assert _noisy_instruction().endswith(_INSTRUCTION) + out = local_encoder._isolate_task_text(_noisy_instruction()) + assert out.strip().endswith(_INSTRUCTION) + + +def test_isolate_keeps_inline_code_spans(): + """Inline code spans are KEPT by default: instructions reference + real file names and commands, which are signal, not noise.""" + task = "Rename `config_loader.py` to `settings.py` and fix the imports." + assert local_encoder._isolate_task_text(task) == task + + +def test_isolate_floor_guard_falls_back_on_all_code_task(): + """A task that is essentially only a code dump (nothing left above + the character floor) falls back to the original text rather than + embedding a stripped stub.""" + bare_dump = "```python\n" + "\n".join(f"x{i} = {i}" for i in range(40)) + "\n```" + assert local_encoder._isolate_task_text(bare_dump) == bare_dump + + +def test_isolate_is_idempotent(): + """Applying isolation twice equals applying it once.""" + once = local_encoder._isolate_task_text(_noisy_instruction()) + assert local_encoder._isolate_task_text(once) == once + + +def test_isolate_unterminated_fence_keeps_leading_instruction(): + """An unterminated fence (truncated tool output) eats to the end of + the string, so a leading instruction survives and a trailing one + cannot — the floor guard then reverts rather than embed a stub.""" + leading = ( + "Summarize the parser module.\n\n```python\nimport csv\nrows = []\n" + ) + out = local_encoder._isolate_task_text(leading) + assert "Summarize the parser module." in out + assert "import csv" not in out + + trailing = "```python\nimport csv\n" + "x = 1\n" * 30 + "summarize it" + assert local_encoder._isolate_task_text(trailing) == trailing \ No newline at end of file