mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(transforms): adaptive Otsu KEEP/DROP threshold (+ land relevance split on main) (#1726)
## Description Lands the prompt-conditioned relevance split **on `main`** and makes its KEEP/DROP threshold **adaptive**. Context: the Stage B work (#1722) was merged into the feature branch `tejas/proxy-lossless-mode` rather than `main`, so `relevance_split.py` never reached `main`. This PR cherry-picks that work onto `main` and adds the adaptive threshold on top, in three commits: 1. Prompt-conditioned KEEP/DROP tail split (Stage B) — segment LOG/SEARCH output into records, score each against the request's information need (user prompt + triggering tool-call args) via `headroom/relevance/`, keep relevant records verbatim, Kompress the low-relevance tail. Mode-agnostic (marker-free in lossless, retrieval-marker in CCR). 2. On by default with hot-path rails — background embedding-model pre-warm (BM25 until warm, never blocks a request) + optional `relevance_max_records` cap (default 0 = no cap). 3. **Adaptive Otsu threshold** (this PR's new work) — see below. ### Adaptive threshold The keep/drop cut is no longer a fixed constant. For each output we compute the natural relevant/irrelevant break in *its own* score distribution via **Otsu's method** (parameter-free — candidate cuts are the data's own values, no bins or magic numbers), floored by `relevance.relevance_threshold` so absolutely irrelevant records are never kept verbatim. The bar therefore moves with the content + prompt: a highly-relevant output keeps its top cluster and compresses the merely-moderate tail; a mostly-irrelevant output drops almost everything. All-equal scores fall back to the floor. Toggle via `relevance_adaptive_threshold` (default `True`). Closes # ## Type of Change - [x] New feature (non-breaking change that adds functionality) ## Changes Made - `relevance_split.py`: `adaptive_threshold()` + `_otsu_threshold()`; `plan_relevance_split(..., adaptive=True)` uses the adaptive cut, floored by `threshold`. - `content_router.py`: `relevance_adaptive_threshold` config (default `True`), threaded into the split. (Plus the Stage B split + default-on rails from the cherry-picked commits.) - `tests/test_relevance_split.py`: adaptive-threshold cases (bimodal split, floored, all-equal, moves-with-distribution) on top of the Stage B suite. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy`) - [x] New tests added for new functionality ### Test Output ```text $ pytest tests/test_relevance_split.py tests/test_transforms_content_router.py tests/test_lossless_mode.py -q 80 passed, 1 warning in 3.41s $ ruff check headroom/transforms/relevance_split.py headroom/transforms/content_router.py tests/test_relevance_split.py All checks passed! $ ruff format --check <changed files> 3 files already formatted $ mypy headroom/transforms/relevance_split.py headroom/transforms/content_router.py Success: no issues found in 2 source files ``` ## Real Behavior Proof - **Environment:** local, Python 3.12.6. - **Steps:** `adaptive_threshold()` exercised directly on synthetic score distributions; `plan_relevance_split(adaptive=True)` and the real `ContentRouter._apply_strategy_to_content` path driven with a deterministic scorer + Kompress-tail stub (offline). - **Observed:** - Bimodal scores `[0.92, 0.88, 0.12, 0.05]` → cut lands in the valley (`0.12 < t < 0.88`), keeping the high cluster. - Mostly-irrelevant `[0.30, 0.28, 0.05, 0.03]` → cut floored at `0.25`. - All-equal scores → floor. - Higher-scoring distribution yields a higher cut than a lower one (bar adapts). - Router split still fires in both lossless and CCR mode; DIFF stays pure lossless; disabling the flag is byte-identical. - **Not tested:** live embedding model warm/latency at scale; end-to-end `/v1/retrieve` resolution of the CCR tail marker (marker plumbing itself is covered upstream). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Additional Notes - Supersedes the orphaned #1722 merge (which landed on the feature branch, not `main`); this PR is the canonical path onto `main`. - **Follow-ups discussed:** TEXT-strategy extension (relevance split for plain prose, currently whole-block Kompress); batch multiple DROP runs into one Kompress call; eval of savings/fidelity on live traffic. - N/A: CHANGELOG (feature not yet released).
This commit is contained in:
Tejas Chopra
GitHub
parent
c9d717c13c
commit
eea667a720
4 changed files with 666 additions and 12 deletions
|
|
@ -52,6 +52,7 @@ from typing import Any
|
|||
from ..config import (
|
||||
DEFAULT_EXCLUDE_TOOLS,
|
||||
ReadLifecycleConfig,
|
||||
RelevanceScorerConfig,
|
||||
TransformResult,
|
||||
is_tool_excluded,
|
||||
)
|
||||
|
|
@ -61,6 +62,7 @@ from .base import Transform
|
|||
from .content_detector import ContentType, DetectionResult, _try_detect_log, _try_detect_search
|
||||
from .content_detector import detect_content_type as _regex_detect_content_type
|
||||
from .error_detection import content_has_strong_error_indicators
|
||||
from .relevance_split import build_relevance_query, plan_relevance_split
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -74,6 +76,23 @@ def _router_debug_dumps(value: Any) -> str:
|
|||
return json.dumps(value, ensure_ascii=False, default=str, separators=(",", ":"))
|
||||
|
||||
|
||||
def _tool_call_args_text(raw: Any) -> str:
|
||||
"""Compact, query-usable text from a tool call's args.
|
||||
|
||||
Anthropic passes ``input`` as a dict ({"command": "grep …"}); OpenAI passes
|
||||
``arguments`` as a JSON string. Either way we want the scalar values (the
|
||||
grep pattern, the read path) as a short query fragment. Capped so a giant
|
||||
arg blob can't dominate the relevance query.
|
||||
"""
|
||||
if isinstance(raw, str):
|
||||
text = raw
|
||||
elif isinstance(raw, dict):
|
||||
text = " ".join(str(v) for v in raw.values() if isinstance(v, (str, int, float, bool)))
|
||||
else:
|
||||
return ""
|
||||
return " ".join(text.split())[:300]
|
||||
|
||||
|
||||
def _log_router_debug(event: str, **payload: Any) -> None:
|
||||
if not logger.isEnabledFor(logging.DEBUG):
|
||||
return
|
||||
|
|
@ -739,7 +758,6 @@ class ContentRouterConfig:
|
|||
enable_tabular_compressor: Enable CSV/TSV/markdown-table compression.
|
||||
enable_image_optimizer: Enable image token optimization.
|
||||
prefer_code_aware_for_code: Use CodeAware over Kompress for code.
|
||||
mixed_content_threshold: Min distinct types to consider "mixed".
|
||||
min_section_tokens: Minimum tokens for a section to compress.
|
||||
fallback_strategy: Strategy when no compressor matches.
|
||||
skip_user_messages: Never compress user messages (they're the subject).
|
||||
|
|
@ -769,7 +787,6 @@ class ContentRouterConfig:
|
|||
# emits a `<<ccr:…>>` / `Retrieve …` retrieval marker. SmartCrusher is
|
||||
# additionally forced marker-free via smart_crusher_lossless_only.
|
||||
lossless: bool = False
|
||||
mixed_content_threshold: int = 2 # Min types to consider mixed
|
||||
min_section_tokens: int = 20 # Min tokens to compress a section
|
||||
|
||||
# Fallback: Kompress handles unknown/mixed content instead of passing through
|
||||
|
|
@ -816,12 +833,18 @@ class ContentRouterConfig:
|
|||
0.0 # 0.0 = protect ALL excluded-tool outputs (safest for coding agents)
|
||||
)
|
||||
|
||||
# Adaptive compression ratio: scales with context pressure.
|
||||
# At low pressure (<30% full), use the relaxed threshold (reject marginal).
|
||||
# At high pressure (>80% full), use the aggressive threshold (accept anything helpful).
|
||||
# Linearly interpolates between the two.
|
||||
min_ratio_relaxed: float = 0.85 # when context is mostly empty
|
||||
min_ratio_aggressive: float = 0.65 # when context is nearly full
|
||||
# Adaptive acceptance threshold, scaling with context pressure. The gate
|
||||
# accepts a compression when compression_ratio < min_ratio (ratio =
|
||||
# compressed/original, so LOWER ratio = bigger savings). Thus a HIGHER
|
||||
# min_ratio is MORE lenient (accepts marginal wins) and a LOWER one is
|
||||
# stricter (only big wins clear it). min_ratio is interpolated
|
||||
# 0.85 (empty context) -> 0.65 (full), i.e. acceptance gets STRICTER as
|
||||
# context fills, so only large savings justify busting the prefix cache
|
||||
# under pressure. (Direction is deliberate; whether a full context should
|
||||
# instead accept *more* to reclaim space is a design question flagged for
|
||||
# an eval — do not flip without measuring.)
|
||||
min_ratio_relaxed: float = 0.85 # low pressure: lenient, accept marginal wins
|
||||
min_ratio_aggressive: float = 0.65 # high pressure: strict, big wins only
|
||||
|
||||
# CCR (Compress-Cache-Retrieve) settings for SmartCrusher
|
||||
ccr_enabled: bool = True # Enable CCR marker injection for reversible compression
|
||||
|
|
@ -834,6 +857,28 @@ class ContentRouterConfig:
|
|||
# can run marker-free without constructing the crusher by hand.
|
||||
smart_crusher_lossless_only: bool | None = None
|
||||
|
||||
# Prompt-conditioned relevance split for the KEEP/DROP tail. When enabled,
|
||||
# LOG/SEARCH output is segmented into records, each scored against the
|
||||
# request's information need (user prompt + triggering tool-call args) via
|
||||
# `relevance` below; high-relevance records are kept verbatim and the
|
||||
# low-relevance tail is Kompressed. Works in both modes: in lossless mode
|
||||
# the tail is marker-free; in CCR mode it carries a retrieval marker (via
|
||||
# ccr_inject_marker) so dropped detail stays retrievable. On by default; the
|
||||
# embedding model is pre-warmed in the background (BM25 scores until it's
|
||||
# cached) so no request ever blocks on the download.
|
||||
relevance_split: bool = True
|
||||
relevance: RelevanceScorerConfig = field(default_factory=RelevanceScorerConfig)
|
||||
# Optional latency guard: skip the split when an output segments into more
|
||||
# than this many records, capping embedding work on the request thread.
|
||||
# 0 = no cap (default): every record is scored regardless of size. Set a
|
||||
# positive value to bound per-request embedding cost on very large outputs.
|
||||
relevance_max_records: int = 0
|
||||
# Adaptive KEEP/DROP cut: when True (default), the threshold is the natural
|
||||
# relevant/irrelevant break in each output's score distribution (Otsu),
|
||||
# floored by relevance.relevance_threshold — it moves with the content
|
||||
# instead of a fixed constant. False uses the fixed threshold exactly.
|
||||
relevance_adaptive_threshold: bool = True
|
||||
|
||||
# Tag protection: preserve custom/workflow XML tags from text compression.
|
||||
# When False (default), entire <custom-tag>content</custom-tag> blocks are
|
||||
# protected verbatim. When True, only the tag markers are protected and
|
||||
|
|
@ -1140,6 +1185,13 @@ class ContentRouter(Transform):
|
|||
self._html_extractor: Any = None
|
||||
self._tabular_compressor: Any = None
|
||||
self._kompress: Any = None
|
||||
# Stage B relevance split (lazy; None until first use, sentinel-checked
|
||||
# via _relevance_scorer_tried so a failed load isn't retried per call).
|
||||
self._relevance_scorer: Any = None
|
||||
self._relevance_scorer_tried: bool = False
|
||||
self._relevance_prewarm_started: bool = False
|
||||
# tool_call_id → compact args text, populated by _build_tool_name_map.
|
||||
self._tool_call_args: dict[str, str] = {}
|
||||
|
||||
# Phase 0 (#1171): cap the input size handed to kompress (ModernBERT
|
||||
# ONNX). Its inference scales O(tokens) and runs synchronously on the
|
||||
|
|
@ -1630,6 +1682,24 @@ class ContentRouter(Transform):
|
|||
strategy_chain: list[str] = [strategy.value]
|
||||
error: str | None = None
|
||||
|
||||
# Stage B/C: prompt-conditioned relevance split for LOG/SEARCH — keep
|
||||
# relevant records verbatim, compress the low-value tail. Runs in BOTH
|
||||
# modes; the tail's marker behavior follows the mode automatically via
|
||||
# _try_ml_compressor: lossless → marker-free Kompress; CCR → Kompress
|
||||
# with a retrieval marker so the dropped detail stays retrievable (a
|
||||
# safety net if the scorer is wrong). DIFF is excluded — Kompressing
|
||||
# hunks breaks `git apply`. Returns None (falls through to the normal
|
||||
# path) when disabled, unavailable, or no better than plain compression.
|
||||
if self.config.relevance_split and strategy in (
|
||||
CompressionStrategy.LOG,
|
||||
CompressionStrategy.SEARCH,
|
||||
):
|
||||
kind = "log" if strategy is CompressionStrategy.LOG else "search"
|
||||
split = self._relevance_split_compress(content, kind, context)
|
||||
if split is not None:
|
||||
label = f"lossless_{kind}" if self.config.lossless else kind
|
||||
return split, len(split.split()), [label, "relevance_split"]
|
||||
|
||||
# No-CCR lossless mode: LOG/SEARCH/DIFF get format-native lossless
|
||||
# compaction instead of the lossy Rust drop path, so the output stays
|
||||
# marker-free (no `<<ccr:…>>` / `Retrieve …`) and fully recoverable.
|
||||
|
|
@ -2117,6 +2187,112 @@ class ContentRouter(Transform):
|
|||
logger.debug("LogCompressor not available")
|
||||
return self._log_compressor
|
||||
|
||||
def _get_relevance_scorer(self) -> Any:
|
||||
"""Get the relevance scorer for the split (lazy, cached, non-blocking).
|
||||
|
||||
Tier comes from ``config.relevance``. For ``bm25`` this is instant. For
|
||||
``hybrid``/``embedding`` the scorer serves **BM25 immediately** and the
|
||||
embedding model is warmed in a background thread; once it's cached the
|
||||
scorer is swapped in (GIL-atomic ref write), so a request never blocks
|
||||
on the ~30MB download. Returns None (cached) on failure. Never raises.
|
||||
"""
|
||||
if self._relevance_scorer is not None or self._relevance_scorer_tried:
|
||||
return self._relevance_scorer
|
||||
self._relevance_scorer_tried = True
|
||||
tier = (self.config.relevance.tier or "hybrid").lower()
|
||||
try:
|
||||
from ..relevance import BM25Scorer
|
||||
|
||||
if tier == "bm25":
|
||||
self._relevance_scorer = BM25Scorer()
|
||||
else:
|
||||
# Serve BM25 now; swap to the embedding-backed scorer once warm.
|
||||
self._relevance_scorer = BM25Scorer()
|
||||
self._start_relevance_prewarm(tier)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("relevance scorer unavailable: %s", exc)
|
||||
self._relevance_scorer = None
|
||||
return self._relevance_scorer
|
||||
|
||||
def _start_relevance_prewarm(self, tier: str) -> None:
|
||||
"""Warm the embedding model off the request thread, then swap it in.
|
||||
|
||||
Idempotent. On failure (fastembed missing, download error) the router
|
||||
just stays on the BM25 scorer set by ``_get_relevance_scorer``.
|
||||
"""
|
||||
if getattr(self, "_relevance_prewarm_started", False):
|
||||
return
|
||||
self._relevance_prewarm_started = True
|
||||
|
||||
def _warm() -> None:
|
||||
try:
|
||||
from ..relevance import create_scorer
|
||||
|
||||
scorer = create_scorer(tier)
|
||||
# Force the model download+load and a first embed here, in the
|
||||
# background — so the first real request finds it warm.
|
||||
scorer.score_batch(["warmup"], "warmup")
|
||||
self._relevance_scorer = scorer # GIL-atomic ref swap
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("relevance model prewarm failed; staying on BM25: %s", exc)
|
||||
|
||||
threading.Thread(target=_warm, name="relevance-prewarm", daemon=True).start()
|
||||
|
||||
def _relevance_split_compress(self, content: str, kind: str, query: str) -> str | None:
|
||||
"""Prompt-conditioned KEEP/DROP split for the compression tail.
|
||||
|
||||
Keeps high-relevance records byte-verbatim (lossless-compacted) and
|
||||
Kompresses the low-relevance tail (identifiers pinned by Kompress
|
||||
MUST_KEEP). Mode-agnostic: the tail's marker behavior is decided by
|
||||
``_try_ml_compressor`` — marker-free in lossless mode, retrieval-marker
|
||||
in CCR mode. Returns the spliced output, or None to fall back to the
|
||||
normal path when the scorer is unavailable, the query is empty, nothing
|
||||
is dropped, or the split doesn't beat plain compaction. Never raises.
|
||||
|
||||
Embedding cost is bounded two ways: the model is pre-warmed off the
|
||||
request thread (BM25 until it's ready, see _get_relevance_scorer) and
|
||||
outputs segmenting into more than ``relevance_max_records`` records skip
|
||||
the split entirely.
|
||||
"""
|
||||
scorer = self._get_relevance_scorer()
|
||||
if scorer is None or not query.strip():
|
||||
return None
|
||||
from .lossless_compaction import compact_lossless
|
||||
|
||||
try:
|
||||
runs = plan_relevance_split(
|
||||
content,
|
||||
query,
|
||||
scorer,
|
||||
threshold=self.config.relevance.relevance_threshold,
|
||||
adaptive=self.config.relevance_adaptive_threshold,
|
||||
max_records=self.config.relevance_max_records,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("relevance split failed (%s); falling back", exc)
|
||||
return None
|
||||
|
||||
# No low-relevance tail → plain compaction is already optimal here.
|
||||
if not any(not keep for keep, _ in runs):
|
||||
return None
|
||||
|
||||
out_parts: list[str] = []
|
||||
for keep, text in runs:
|
||||
if keep:
|
||||
out_parts.append(compact_lossless(text, kind))
|
||||
continue
|
||||
try:
|
||||
compressed, _ = self._try_ml_compressor(text, query)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("kompress tail failed (%s); keeping verbatim", exc)
|
||||
compressed = compact_lossless(text, kind)
|
||||
out_parts.append(compressed)
|
||||
|
||||
result = "".join(out_parts)
|
||||
# Adopt only when it beats plain whole-block lossless compaction.
|
||||
baseline = compact_lossless(content, kind)
|
||||
return result if len(result) < len(baseline) else None
|
||||
|
||||
def _get_text_crusher(self) -> Any:
|
||||
"""Get TextCrusher (Phase 2, lazy load). Returns None when disabled, or
|
||||
when the native ``headroom._core`` extension is not built (mirrors the
|
||||
|
|
@ -2426,9 +2602,15 @@ class ContentRouter(Transform):
|
|||
"""Build mapping from tool_call_id to tool_name.
|
||||
|
||||
Scans assistant messages to find tool calls and extract their names.
|
||||
Supports both OpenAI and Anthropic message formats.
|
||||
Supports both OpenAI and Anthropic message formats. Also populates
|
||||
``self._tool_call_args`` (id → compact args text) in the same scan, so
|
||||
the relevance split can score a tool output against the *precise* ask
|
||||
that triggered it (grep pattern, read path, …), not just the user
|
||||
prompt. Read-only after build → safe to read from the parallel
|
||||
compression pass.
|
||||
"""
|
||||
mapping: dict[str, str] = {}
|
||||
args_map: dict[str, str] = {}
|
||||
|
||||
for msg in messages:
|
||||
if msg.get("role") != "assistant":
|
||||
|
|
@ -2438,9 +2620,13 @@ class ContentRouter(Transform):
|
|||
for tc in msg.get("tool_calls", []):
|
||||
if isinstance(tc, dict):
|
||||
tc_id = tc.get("id", "")
|
||||
name = tc.get("function", {}).get("name", "")
|
||||
fn = tc.get("function", {})
|
||||
name = fn.get("name", "")
|
||||
if tc_id and name:
|
||||
mapping[tc_id] = name
|
||||
args = _tool_call_args_text(fn.get("arguments"))
|
||||
if args:
|
||||
args_map[tc_id] = args
|
||||
|
||||
# Anthropic format: content blocks with type=tool_use
|
||||
content = msg.get("content", [])
|
||||
|
|
@ -2451,7 +2637,11 @@ class ContentRouter(Transform):
|
|||
name = block.get("name", "")
|
||||
if tc_id and name:
|
||||
mapping[tc_id] = name
|
||||
args = _tool_call_args_text(block.get("input"))
|
||||
if args:
|
||||
args_map[tc_id] = args
|
||||
|
||||
self._tool_call_args = args_map
|
||||
return mapping
|
||||
|
||||
def _net_cost_allows(
|
||||
|
|
@ -3350,6 +3540,15 @@ class ContentRouter(Transform):
|
|||
tool_name = (tool_name_map or {}).get(tool_use_id, "")
|
||||
bias = self._get_tool_bias(tool_name) if tool_name else 1.0
|
||||
|
||||
# Enrich the relevance query with the triggering tool call's
|
||||
# args (grep pattern, read path, …) — the sharpest, per-output
|
||||
# signal. Gated so default behavior is byte-identical.
|
||||
block_context = context
|
||||
if self.config.relevance_split and tool_use_id:
|
||||
call_args = self._tool_call_args.get(tool_use_id, "")
|
||||
if call_args:
|
||||
block_context = build_relevance_query(context, tool_name, call_args)
|
||||
|
||||
tool_content = block.get("content", "")
|
||||
|
||||
# Protection: failed tool calls / error outputs stay verbatim
|
||||
|
|
@ -3394,7 +3593,7 @@ class ContentRouter(Transform):
|
|||
content_key=hash(
|
||||
(tool_content, getattr(self, "_runtime_target_ratio", None))
|
||||
),
|
||||
context=context,
|
||||
context=block_context,
|
||||
bias=bias,
|
||||
min_ratio=min_ratio,
|
||||
compressor_timing=compressor_timing,
|
||||
|
|
@ -3403,6 +3602,7 @@ class ContentRouter(Transform):
|
|||
compressed_details=compressed_details,
|
||||
strategy_label="tool_result",
|
||||
details_prefix="tool",
|
||||
enforce_reversibility=True,
|
||||
)
|
||||
if compressed_content is not None:
|
||||
new_blocks.append({**block, "content": compressed_content})
|
||||
|
|
@ -3479,6 +3679,7 @@ class ContentRouter(Transform):
|
|||
compressed_details: list[str] | None,
|
||||
strategy_label: str,
|
||||
details_prefix: str,
|
||||
enforce_reversibility: bool = False,
|
||||
) -> tuple[str | None, bool]:
|
||||
"""Apply two-tier cache lookup + compression to a single content string.
|
||||
|
||||
|
|
@ -3543,6 +3744,23 @@ class ContentRouter(Transform):
|
|||
key = f"compressor:{result.strategy_used.value}"
|
||||
compressor_timing[key] = compressor_timing.get(key, 0.0) + compress_ms
|
||||
if result.compression_ratio < min_ratio:
|
||||
# Tool ground truth must stay reversible: a lossy summarizer
|
||||
# (kompress/text/code) that emitted no CCR retrieve marker is
|
||||
# unrecoverable, so the agent would act on a fabricated summary
|
||||
# (#1307). The string/`role=="tool"` path guards this; mirror it
|
||||
# here for tool_result blocks (never cached, so the Tier-2 path
|
||||
# above can't serve a poisoned entry).
|
||||
if (
|
||||
enforce_reversibility
|
||||
and result.strategy_used in self.LOSSY_UNMARKED_STRATEGIES
|
||||
and not CCR_RETRIEVAL_MARKER_RE.search(result.compressed)
|
||||
):
|
||||
self._cache.mark_skip(content_key)
|
||||
if route_counts is not None:
|
||||
route_counts["lossy_unrecoverable_skipped"] = (
|
||||
route_counts.get("lossy_unrecoverable_skipped", 0) + 1
|
||||
)
|
||||
return None, False
|
||||
# Compressed — store in result cache
|
||||
self._cache.put(
|
||||
content_key,
|
||||
|
|
|
|||
174
headroom/transforms/relevance_split.py
Normal file
174
headroom/transforms/relevance_split.py
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
"""Prompt-conditioned relevance split for KEEP/DROP compression decisions.
|
||||
|
||||
Segments tool output into coherent records, scores each against the request's
|
||||
*information need* (user prompt + the triggering tool call's args) using the
|
||||
existing :class:`~headroom.relevance.RelevanceScorer` (BM25 / bge-small
|
||||
embeddings / hybrid), and partitions the content into ordered KEEP/DROP runs.
|
||||
|
||||
The split is **mode-agnostic**: this module decides *what* is worth keeping
|
||||
verbatim vs. what is a low-value tail; the caller applies the disposition. In
|
||||
lossless (no-CCR) mode the KEEP runs stay byte-verbatim and the DROP tail is
|
||||
Kompressed marker-free; in CCR mode the same DROP tail can be dropped with a
|
||||
retrieval marker. Nothing here emits markers or calls a compressor.
|
||||
|
||||
Segmentation is boundary-aware, not line-based: blank lines delimit records,
|
||||
indented continuation lines stay attached to their parent (so stack traces and
|
||||
pretty-printed blobs are scored as one unit), and dense blank-free streams
|
||||
(grep, tight logs) are packed into small fixed windows. The partition is
|
||||
lossless -- ``"".join(segment(content)) == content`` -- so KEEP runs
|
||||
reconstruct the original bytes exactly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from headroom.relevance import RelevanceScorer
|
||||
|
||||
__all__ = ["adaptive_threshold", "build_relevance_query", "plan_relevance_split", "segment"]
|
||||
|
||||
|
||||
def build_relevance_query(user_query: str, tool_name: str = "", tool_args: str = "") -> str:
|
||||
"""Compose the information-need query for relevance scoring.
|
||||
|
||||
The user's prompt is the high-level intent; the triggering tool call's args
|
||||
(a grep pattern, a read path, a search query) are the *precise*, per-output
|
||||
ask and usually the sharpest signal. Both are included so the lexical (BM25)
|
||||
half locks onto exact tokens (e.g. the grep pattern) while the semantic half
|
||||
tracks the intent.
|
||||
"""
|
||||
parts: list[str] = []
|
||||
q = (user_query or "").strip()
|
||||
if q:
|
||||
parts.append(q)
|
||||
call = " ".join(p for p in ((tool_name or "").strip(), (tool_args or "").strip()) if p)
|
||||
if call:
|
||||
parts.append(call)
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def segment(content: str, *, window: int = 8, max_chars: int = 1200) -> list[str]:
|
||||
"""Partition ``content`` into coherent records.
|
||||
|
||||
Lossless partition: ``"".join(segment(content)) == content``. Blank lines
|
||||
delimit records; oversized or dense blank-free blocks are packed into
|
||||
windows of at most ``window`` lines / ``max_chars`` chars, with indented
|
||||
continuation lines held to their window so multi-line units aren't cut.
|
||||
"""
|
||||
lines = content.splitlines(keepends=True)
|
||||
if len(lines) <= 1:
|
||||
return [content] if content else []
|
||||
|
||||
# Pass 1: blank-line-delimited blocks (paragraphs / record gaps).
|
||||
blocks: list[list[str]] = []
|
||||
cur: list[str] = []
|
||||
for ln in lines:
|
||||
cur.append(ln)
|
||||
if ln.strip() == "":
|
||||
blocks.append(cur)
|
||||
cur = []
|
||||
if cur:
|
||||
blocks.append(cur)
|
||||
|
||||
# Pass 2: pack/window each block. Dense blank-free streams (grep, tight
|
||||
# logs) become fixed windows; indented continuation lines stay attached to
|
||||
# their window so stack traces / pretty JSON aren't split mid-unit.
|
||||
segments: list[str] = []
|
||||
for block in blocks:
|
||||
if len(block) <= window and sum(len(x) for x in block) <= max_chars:
|
||||
segments.append("".join(block))
|
||||
continue
|
||||
i = 0
|
||||
n = len(block)
|
||||
while i < n:
|
||||
j = min(i + window, n)
|
||||
while j < n and block[j][:1] in (" ", "\t"):
|
||||
j += 1 # don't cut off an indented continuation run
|
||||
segments.append("".join(block[i:j]))
|
||||
i = j
|
||||
return segments
|
||||
|
||||
|
||||
def _otsu_threshold(values: list[float]) -> float:
|
||||
"""Otsu's method: the cut between two classes that maximizes between-class
|
||||
variance. Parameter-free — the candidate cuts are the data's own values, so
|
||||
there's no bin size or magic constant. Returns the midpoint of the winning
|
||||
adjacent pair.
|
||||
"""
|
||||
xs = sorted(values)
|
||||
n = len(xs)
|
||||
total = sum(xs)
|
||||
w0 = 0.0
|
||||
sum0 = 0.0
|
||||
best_t = xs[0]
|
||||
best_var = -1.0
|
||||
for i in range(n - 1):
|
||||
w0 += 1
|
||||
sum0 += xs[i]
|
||||
w1 = n - w0
|
||||
m0 = sum0 / w0
|
||||
m1 = (total - sum0) / w1
|
||||
between = w0 * w1 * (m0 - m1) ** 2
|
||||
if between > best_var:
|
||||
best_var = between
|
||||
best_t = (xs[i] + xs[i + 1]) / 2.0
|
||||
return best_t
|
||||
|
||||
|
||||
def adaptive_threshold(values: list[float], floor: float) -> float:
|
||||
"""Data-driven KEEP/DROP cut for one output's relevance scores.
|
||||
|
||||
The operative cut is the natural relevant/irrelevant break (Otsu) *for this
|
||||
output and query* — so it moves with the score distribution instead of a
|
||||
fixed constant. It's floored by ``floor`` so records below the absolute
|
||||
minimum relevance are never kept verbatim (an all-irrelevant output drops
|
||||
entirely). When every record scores the same there's no break to find, so
|
||||
the floor decides (all-in or all-out).
|
||||
"""
|
||||
if len({round(v, 9) for v in values}) < 2:
|
||||
return floor
|
||||
return max(_otsu_threshold(values), floor)
|
||||
|
||||
|
||||
def plan_relevance_split(
|
||||
content: str,
|
||||
query: str,
|
||||
scorer: RelevanceScorer,
|
||||
*,
|
||||
threshold: float,
|
||||
adaptive: bool = True,
|
||||
window: int = 8,
|
||||
max_chars: int = 1200,
|
||||
max_records: int | None = None,
|
||||
) -> list[tuple[bool, str]]:
|
||||
"""Split ``content`` into ordered ``(keep, text)`` runs by relevance to ``query``.
|
||||
|
||||
A record is KEEP when its relevance score clears the cut. With
|
||||
``adaptive`` (default), the cut is the natural relevant/irrelevant break in
|
||||
*this* output's score distribution (Otsu), floored by ``threshold`` — so it
|
||||
adapts to the content instead of being a fixed constant. With
|
||||
``adaptive=False`` the cut is ``threshold`` exactly. Either way *which*
|
||||
records clear it is entirely prompt-driven, so the KEEP fraction ranges
|
||||
from 0% to 100% with the content, not a fixed quota. Consecutive
|
||||
same-disposition records are merged into runs (order preserved) so the
|
||||
caller applies one disposition per run. Returns a single KEEP run -- i.e.
|
||||
no split -- when the query is empty, the content is a single record, or it
|
||||
segments into more than ``max_records`` records (a latency guard).
|
||||
"""
|
||||
if not query.strip():
|
||||
return [(True, content)]
|
||||
segs = segment(content, window=window, max_chars=max_chars)
|
||||
if len(segs) < 2 or (max_records and len(segs) > max_records):
|
||||
return [(True, content)]
|
||||
|
||||
scores = scorer.score_batch(segs, query)
|
||||
cut = adaptive_threshold([s.score for s in scores], threshold) if adaptive else threshold
|
||||
runs: list[tuple[bool, str]] = []
|
||||
for seg, sc in zip(segs, scores):
|
||||
keep = sc.score >= cut
|
||||
if runs and runs[-1][0] == keep:
|
||||
runs[-1] = (keep, runs[-1][1] + seg)
|
||||
else:
|
||||
runs.append((keep, seg))
|
||||
return runs
|
||||
205
tests/test_relevance_split.py
Normal file
205
tests/test_relevance_split.py
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
"""Unit tests for the prompt-conditioned relevance split (Stage B core).
|
||||
|
||||
Uses a deterministic fake scorer -- no embedding model / network needed -- so
|
||||
these run fast and pin the segmentation + partition logic, not the ML model.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from headroom.relevance.base import RelevanceScore, RelevanceScorer
|
||||
from headroom.transforms.relevance_split import (
|
||||
adaptive_threshold,
|
||||
build_relevance_query,
|
||||
plan_relevance_split,
|
||||
segment,
|
||||
)
|
||||
|
||||
|
||||
class KeywordScorer(RelevanceScorer):
|
||||
"""Score = fraction of query terms present in the item. No model."""
|
||||
|
||||
def score(self, item: str, context: str) -> RelevanceScore:
|
||||
terms = context.lower().split()
|
||||
if not terms:
|
||||
return RelevanceScore(score=0.0)
|
||||
hits = sum(1 for t in terms if t in item.lower())
|
||||
return RelevanceScore(score=hits / len(terms))
|
||||
|
||||
def score_batch(self, items: list[str], context: str) -> list[RelevanceScore]:
|
||||
return [self.score(it, context) for it in items]
|
||||
|
||||
|
||||
def test_segment_partition_is_lossless():
|
||||
text = "a\nb\n\n cont\nc\n"
|
||||
assert "".join(segment(text)) == text
|
||||
|
||||
|
||||
def test_segment_windows_dense_stream_losslessly():
|
||||
text = "".join(f"line{i}\n" for i in range(20))
|
||||
segs = segment(text, window=5)
|
||||
assert "".join(segs) == text
|
||||
assert len(segs) > 1 # dense blank-free stream got windowed
|
||||
|
||||
|
||||
def test_segment_keeps_indented_continuation_attached():
|
||||
# window=1 forces splitting, but indented continuation lines must stay
|
||||
# with their head line (stack-trace / pretty-JSON safety).
|
||||
text = "ERROR boom\n File a.py line 1\n File b.py line 2\nnext record\n"
|
||||
segs = segment(text, window=1)
|
||||
assert "".join(segs) == text
|
||||
for s in segs:
|
||||
assert not s.startswith((" ", "\t")) # every segment starts at a head line
|
||||
|
||||
|
||||
def test_split_keeps_relevant_drops_irrelevant():
|
||||
content = (
|
||||
"the oauth token refresh failed here\n"
|
||||
"\n"
|
||||
"unrelated debug noise about widgets\n"
|
||||
"\n"
|
||||
"another oauth token line\n"
|
||||
)
|
||||
runs = plan_relevance_split(content, "oauth token", KeywordScorer(), threshold=0.5)
|
||||
kept = "".join(t for k, t in runs if k)
|
||||
dropped = "".join(t for k, t in runs if not k)
|
||||
assert "oauth token" in kept
|
||||
assert "widgets" in dropped
|
||||
# partition stays lossless regardless of keep/drop labels
|
||||
assert "".join(t for _, t in runs) == content
|
||||
|
||||
|
||||
def test_empty_query_yields_no_split():
|
||||
assert plan_relevance_split("x\ny\n", "", KeywordScorer(), threshold=0.5) == [(True, "x\ny\n")]
|
||||
|
||||
|
||||
def test_single_record_yields_no_split():
|
||||
assert plan_relevance_split("solo", "anything", KeywordScorer(), threshold=0.5) == [
|
||||
(True, "solo")
|
||||
]
|
||||
|
||||
|
||||
def test_build_query_composes_prompt_and_tool_args():
|
||||
q = build_relevance_query("I need entities", "Bash", "grep -rn 'class .*Entity' src/")
|
||||
assert "entities" in q
|
||||
assert "grep" in q
|
||||
assert "Entity" in q
|
||||
|
||||
|
||||
def test_build_query_handles_missing_pieces():
|
||||
assert build_relevance_query("", "", "") == ""
|
||||
assert build_relevance_query("just a prompt") == "just a prompt"
|
||||
|
||||
|
||||
# --- Adaptive threshold (Otsu) --------------------------------------------------
|
||||
|
||||
|
||||
def test_adaptive_threshold_splits_at_the_natural_gap():
|
||||
# Bimodal: cut lands in the valley between the high and low clusters, so the
|
||||
# high cluster is kept and the low one dropped -- not at a fixed constant.
|
||||
t = adaptive_threshold([0.92, 0.88, 0.12, 0.05], floor=0.25)
|
||||
assert 0.12 < t < 0.88
|
||||
|
||||
|
||||
def test_adaptive_threshold_is_floored():
|
||||
# A mostly-irrelevant output: the natural break is low, but the floor keeps
|
||||
# us from retaining absolute junk verbatim.
|
||||
assert adaptive_threshold([0.30, 0.28, 0.05, 0.03], floor=0.25) == 0.25
|
||||
|
||||
|
||||
def test_adaptive_threshold_all_equal_uses_floor():
|
||||
assert adaptive_threshold([0.4, 0.4, 0.4], floor=0.25) == 0.25
|
||||
|
||||
|
||||
def test_adaptive_threshold_moves_with_distribution():
|
||||
# High-scoring output → higher cut than a low-scoring one: the bar adapts.
|
||||
high = adaptive_threshold([0.95, 0.9, 0.6, 0.55], floor=0.1)
|
||||
low = adaptive_threshold([0.4, 0.35, 0.08, 0.05], floor=0.1)
|
||||
assert high > low
|
||||
|
||||
|
||||
# --- Router integration (real _apply_strategy_to_content path) -----------------
|
||||
# Fake scorer + stubbed Kompress tail → deterministic and offline (no model).
|
||||
|
||||
from headroom.config import RelevanceScorerConfig # noqa: E402
|
||||
from headroom.transforms.content_router import ( # noqa: E402
|
||||
CompressionStrategy,
|
||||
ContentRouter,
|
||||
ContentRouterConfig,
|
||||
)
|
||||
|
||||
_SEARCH = (
|
||||
"src/auth.py:12:oauth token refresh\n"
|
||||
"src/auth.py:13:validate oauth token here\n"
|
||||
"\n"
|
||||
"src/widget.py:5:render the widget layout\n"
|
||||
"src/widget.py:6:widget styling code\n"
|
||||
)
|
||||
|
||||
|
||||
def _router(split_on: bool, *, lossless: bool = True) -> ContentRouter:
|
||||
cfg = ContentRouterConfig(
|
||||
lossless=lossless,
|
||||
relevance_split=split_on,
|
||||
relevance=RelevanceScorerConfig(tier="bm25", relevance_threshold=0.5),
|
||||
)
|
||||
r = ContentRouter(cfg)
|
||||
# Inject deterministic scorer + Kompress-tail stub (no model / network).
|
||||
r._relevance_scorer = KeywordScorer()
|
||||
r._relevance_scorer_tried = True
|
||||
r._try_ml_compressor = lambda text, ctx, question=None: ("[TAIL]", 1) # type: ignore[assignment]
|
||||
return r
|
||||
|
||||
|
||||
def test_router_relevance_split_fires_for_search():
|
||||
r = _router(split_on=True) # lossless mode
|
||||
out, _, chain = r._apply_strategy_to_content(_SEARCH, CompressionStrategy.SEARCH, "oauth token")
|
||||
assert chain == ["lossless_search", "relevance_split"]
|
||||
assert "oauth token" in out # relevant records kept verbatim
|
||||
assert "widget" not in out # irrelevant tail replaced by Kompress stub
|
||||
assert "[TAIL]" in out
|
||||
|
||||
|
||||
def test_router_relevance_split_fires_in_ccr_mode():
|
||||
# lossless=False → CCR mode. Same split, unprefixed label. The DROP tail's
|
||||
# retrieval marker is emitted by Kompress when ccr_inject_marker is on (see
|
||||
# #1721); the _try_ml_compressor stub stands in for it here. Proves the
|
||||
# split is mode-agnostic, not lossless-only.
|
||||
r = _router(split_on=True, lossless=False)
|
||||
out, _, chain = r._apply_strategy_to_content(_SEARCH, CompressionStrategy.SEARCH, "oauth token")
|
||||
assert chain == ["search", "relevance_split"]
|
||||
assert "oauth token" in out
|
||||
assert "[TAIL]" in out
|
||||
|
||||
|
||||
def test_router_diff_stays_pure_lossless():
|
||||
r = _router(split_on=True)
|
||||
diff = "diff --git a/x b/x\nindex 111..222 100644\n@@ -1 +1 @@\n-old widget\n+new oauth token\n"
|
||||
_, _, chain = r._apply_strategy_to_content(diff, CompressionStrategy.DIFF, "oauth token")
|
||||
assert "relevance_split" not in chain # Kompressing hunks would break apply
|
||||
assert chain == ["lossless_diff"]
|
||||
|
||||
|
||||
def test_router_split_can_be_disabled():
|
||||
r = _router(split_on=False)
|
||||
_, _, chain = r._apply_strategy_to_content(_SEARCH, CompressionStrategy.SEARCH, "oauth token")
|
||||
assert "relevance_split" not in chain
|
||||
|
||||
|
||||
def test_relevance_split_on_by_default_and_non_blocking(monkeypatch):
|
||||
from headroom.relevance.bm25 import BM25Scorer
|
||||
|
||||
r = ContentRouter(ContentRouterConfig())
|
||||
assert r.config.relevance_split is True
|
||||
# Stub the background warm-up so this is deterministic: with a warm HF cache
|
||||
# the prewarm thread could otherwise swap in the hybrid scorer before we
|
||||
# read it. We assert the *synchronous* hot path serves BM25 without loading
|
||||
# the embedding model on the request thread (the swap happens later, in the
|
||||
# background thread — proven separately).
|
||||
monkeypatch.setattr(r, "_start_relevance_prewarm", lambda tier: None)
|
||||
assert isinstance(r._get_relevance_scorer(), BM25Scorer)
|
||||
|
||||
|
||||
def test_split_respects_max_records_cap():
|
||||
content = "".join(f"rec {i} widget\n\n" for i in range(10)) # 10 blank-sep records
|
||||
runs = plan_relevance_split(content, "widget", KeywordScorer(), threshold=0.5, max_records=3)
|
||||
assert runs == [(True, content)] # over the cap → no split, caller falls back
|
||||
|
|
@ -158,7 +158,11 @@ def test_force_kompress_routes_anthropic_tool_result_to_targeted_kompress(
|
|||
|
||||
def compress(self, content, **kwargs):
|
||||
captured.update(kwargs)
|
||||
compressed = " ".join(content.split()[:20])
|
||||
# Real Kompress appends a CCR retrieval marker when CCR is enabled,
|
||||
# keeping the lossy result recoverable. Include one so the router's
|
||||
# reversibility gate (tool ground truth must stay recoverable, #1307)
|
||||
# accepts the compression instead of reverting to verbatim.
|
||||
compressed = " ".join(content.split()[:20]) + " Retrieve more: hash=deadbeef"
|
||||
return SimpleNamespace(
|
||||
compressed=compressed,
|
||||
compressed_tokens=len(compressed.split()),
|
||||
|
|
@ -197,6 +201,59 @@ def test_force_kompress_routes_anthropic_tool_result_to_targeted_kompress(
|
|||
assert captured["target_ratio"] == 0.10
|
||||
|
||||
|
||||
def test_anthropic_tool_result_lossy_without_marker_stays_verbatim(router, tokenizer, monkeypatch):
|
||||
"""Reversibility gate (#1307): a lossy Kompress result on a tool_result block
|
||||
with no CCR retrieval marker is unrecoverable, so the router must keep the
|
||||
original verbatim rather than hand the agent a fabricated summary. This is
|
||||
the block-path counterpart to the string/`role=="tool"` guard."""
|
||||
|
||||
class FakeKompress:
|
||||
def is_ready(self) -> bool:
|
||||
return True
|
||||
|
||||
def ensure_background_load(self) -> None:
|
||||
pass
|
||||
|
||||
def compress(self, content, **kwargs):
|
||||
# Lossy summary with NO retrieval marker → unrecoverable.
|
||||
compressed = " ".join(content.split()[:20])
|
||||
return SimpleNamespace(
|
||||
compressed=compressed,
|
||||
compressed_tokens=len(compressed.split()),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(router, "_get_kompress", lambda: FakeKompress())
|
||||
tool_content = " ".join(
|
||||
f'{{"file":"src/module_{i}.py","line":{i},"text":"repeated search payload"}}'
|
||||
for i in range(160)
|
||||
)
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_search_1",
|
||||
"content": tool_content,
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
result = router.apply(
|
||||
messages,
|
||||
tokenizer,
|
||||
force_kompress=True,
|
||||
target_ratio=0.10,
|
||||
compress_user_messages=True,
|
||||
min_tokens_to_compress=10,
|
||||
read_protection_window=0,
|
||||
)
|
||||
|
||||
# Unrecoverable lossy compression is rejected → original kept verbatim.
|
||||
assert result.messages[0]["content"][0]["content"] == tool_content
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestContentRouterConfig
|
||||
# =============================================================================
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue