diff --git a/headroom/transforms/adaptive_sizer.py b/headroom/transforms/adaptive_sizer.py index e61ce85a2..51f3befbf 100644 --- a/headroom/transforms/adaptive_sizer.py +++ b/headroom/transforms/adaptive_sizer.py @@ -63,9 +63,27 @@ def compute_optimal_k( curve = compute_unique_bigram_curve(items) knee = find_knee(curve) + # Diversity ratio: what fraction of items are genuinely unique? + # 1.0 = every item is distinct, 0.1 = mostly near-duplicates. + diversity_ratio = unique_count / n + if knee is None: - # No clear knee — content is uniformly diverse, keep ~30% as heuristic - knee = max(min_k, int(n * 0.3)) + # No saturation found — each item adds new information. + # Scale keep-fraction continuously with diversity: + # diversity ~1.0 → keep 100% (all unique — dropping any loses info) + # diversity ~0.5 → keep ~65% (moderate) + # diversity ~0.2 → keep ~44% (low-ish) + # diversity ~0.0 → keep ~30% (mostly dupes, same as old default) + # No arbitrary cap — if items are all unique, keep them all. + keep_fraction = 0.3 + 0.7 * diversity_ratio + knee = max(min_k, int(n * keep_fraction)) + else: + # Knee found, but if diversity is very high the knee may be + # a weak signal (e.g., minor bigram overlap causing a shallow + # curve bend). Don't drop below a diversity floor. + if diversity_ratio > 0.7: + diversity_floor = max(min_k, int(n * (0.3 + 0.7 * diversity_ratio))) + knee = max(knee, diversity_floor) # Apply bias multiplier k = max(min_k, int(knee * bias)) @@ -77,9 +95,10 @@ def compute_optimal_k( k = max(min_k, min(k, effective_max)) logger.debug( - "adaptive_sizer: n=%d unique=%d knee=%s bias=%.1f → k=%d", + "adaptive_sizer: n=%d unique=%d diversity=%.2f knee=%s bias=%.1f → k=%d", n, unique_count, + diversity_ratio, knee, bias, k, diff --git a/headroom/transforms/smart_crusher.py b/headroom/transforms/smart_crusher.py index de2f6b719..80c22dd8e 100644 --- a/headroom/transforms/smart_crusher.py +++ b/headroom/transforms/smart_crusher.py @@ -177,6 +177,79 @@ def _hash_field_name(field_name: str) -> str: return hashlib.sha256(field_name.encode()).hexdigest()[:8] +# Minimum chars for a text field to be worth compressing within an item +_MIN_FIELD_CHARS_FOR_WITHIN = 200 + +# Lazy-loaded compressor for within-item text compression +_within_compressor: Any = None +_within_compressor_checked = False + + +def _get_within_compressor() -> Any: + """Get a text compressor for within-item field compression. + + Returns Kompress if available (requires [ml] extra), else None. + """ + global _within_compressor, _within_compressor_checked + if not _within_compressor_checked: + _within_compressor_checked = True + try: + from .kompress_compressor import KompressCompressor, is_kompress_available + + if is_kompress_available(): + _within_compressor = KompressCompressor() + logger.debug("Within-item compression: using Kompress") + except ImportError: + pass + return _within_compressor + + +def _compress_text_within_items(items: list[dict], context: str = "") -> list[dict]: + """Compress long text fields WITHIN each item, keeping all items. + + Used when diversity is high (all items are unique) — instead of dropping + items, compress the verbose text inside each one. Falls back to the + original list unchanged if no compressor is available or no field is + long enough to benefit. + + Args: + items: JSON-parsed list of dicts. + context: User query context for relevance-aware compression. + + Returns: + Compressed items (new list) or the *same* ``items`` object if + nothing was compressed (caller checks identity). + """ + compressor = _get_within_compressor() + if compressor is None: + return items # No ML compressor available — pass through + + any_compressed = False + result: list[dict] = [] + + for item in items: + new_item = dict(item) # Shallow copy + item_changed = False + + for key, value in item.items(): + if not isinstance(value, str) or len(value) < _MIN_FIELD_CHARS_FOR_WITHIN: + continue + + try: + compressed = compressor.compress(value, context=context) + if compressed.compressed and len(compressed.compressed) < len(value) * 0.9: + new_item[key] = compressed.compressed + item_changed = True + except Exception: + pass # Compression failed for this field — keep original + + result.append(new_item if item_changed else item) + if item_changed: + any_compressed = True + + return result if any_compressed else items + + def _get_preserve_field_values( item: dict, preserve_field_hashes: list[str], @@ -2355,6 +2428,12 @@ class SmartCrusher(Transform): ) if len(items) <= adaptive_k: + # All items kept (high diversity or small array). + # Instead of passing through unchanged, try to compress the TEXT + # WITHIN each item — reduce token count without losing any item. + compressed_items = _compress_text_within_items(items, query_context) + if compressed_items is not items: + return compressed_items, "compress_within:diversity", None, "" return items, "none:adaptive_at_limit", None, "" # Get feedback hints if enabled diff --git a/tests/test_adaptive_sizer.py b/tests/test_adaptive_sizer.py new file mode 100644 index 000000000..ca242c619 --- /dev/null +++ b/tests/test_adaptive_sizer.py @@ -0,0 +1,186 @@ +"""Tests for diversity-aware compute_optimal_k in adaptive_sizer.""" + +from __future__ import annotations + +import json + +from headroom.transforms.adaptive_sizer import compute_optimal_k + + +def _make_unique_items(n: int) -> list[str]: + """Create n completely unique JSON items (high diversity).""" + return [ + json.dumps( + { + "id": i, + "title": f"Unique topic number {i} about subject area {chr(65 + i % 26)}", + "content": ( + f"This is document {i} discussing a completely different subject. " + f"It covers concepts like {chr(65 + i % 26)}-theory, " + f"methodology-{i * 7 % 100}, and framework-{i * 13 % 50}. " + f"The key finding is result-{i} which has implications for field-{i % 10}." + ), + "source": f"source_{i}.pdf", + "score": round(0.99 - i * 0.03, 2), + } + ) + for i in range(n) + ] + + +def _make_repetitive_items(n: int, templates: int = 3) -> list[str]: + """Create n items from a few templates (low diversity).""" + base_templates = [ + { + "status": "ok", + "message": "Health check passed", + "latency_ms": 12, + "service": "api-gateway", + }, + { + "status": "ok", + "message": "Health check passed", + "latency_ms": 15, + "service": "auth-service", + }, + { + "status": "ok", + "message": "Health check passed", + "latency_ms": 8, + "service": "db-proxy", + }, + ] + return [ + json.dumps({**base_templates[i % templates], "timestamp": f"2026-03-25T10:{i:02d}:00Z"}) + for i in range(n) + ] + + +def _make_mixed_items(n: int, unique_fraction: float) -> list[str]: + """Create items where unique_fraction are unique, rest are duplicates.""" + unique_count = int(n * unique_fraction) + dup_count = n - unique_count + items = _make_unique_items(unique_count) + if dup_count > 0: + template = json.dumps( + { + "status": "ok", + "message": "Routine health check passed successfully", + "latency_ms": 10, + } + ) + items.extend([template] * dup_count) + return items + + +class TestSmallArrays: + def test_small_array_returns_n(self): + """Arrays with n <= 8 should always return n (unchanged).""" + items = _make_unique_items(5) + assert compute_optimal_k(items) == 5 + + def test_eight_items_returns_eight(self): + items = _make_unique_items(8) + assert compute_optimal_k(items) == 8 + + +class TestNearTotalRedundancy: + def test_identical_items_returns_min(self): + """20 identical items should return ~3 (near-total redundancy).""" + items = [json.dumps({"status": "ok", "msg": "healthy"})] * 20 + k = compute_optimal_k(items) + assert k <= 3 + + def test_two_groups_returns_small_k(self): + """Items from 2 groups should return small k.""" + items = [json.dumps({"type": "A", "val": 1})] * 10 + [ + json.dumps({"type": "B", "val": 2}) + ] * 10 + k = compute_optimal_k(items) + assert k <= 5 + + +class TestHighDiversity: + def test_all_unique_keeps_most(self): + """15 completely unique items → should keep >= 10 (not 4 like before).""" + items = _make_unique_items(15) + k = compute_optimal_k(items) + assert k >= 10, f"Expected k >= 10 for 15 unique items, got k={k}" + + def test_twenty_unique_keeps_most(self): + """20 unique items → should keep >= 14.""" + items = _make_unique_items(20) + k = compute_optimal_k(items) + assert k >= 14, f"Expected k >= 14 for 20 unique items, got k={k}" + + def test_twelve_unique_rag_chunks(self): + """12 unique RAG chunks → should keep >= 8.""" + items = _make_unique_items(12) + k = compute_optimal_k(items) + assert k >= 8, f"Expected k >= 8 for 12 unique RAG chunks, got k={k}" + + +class TestLowDiversity: + def test_repetitive_items_unchanged(self): + """15 items from 3 templates → k should stay small (same as before).""" + items = _make_repetitive_items(15, templates=3) + k = compute_optimal_k(items) + assert k <= 8, f"Expected k <= 8 for repetitive items, got k={k}" + + def test_twenty_repetitive_stays_small(self): + """20 items from 3 templates → k stays small.""" + items = _make_repetitive_items(20, templates=3) + k = compute_optimal_k(items) + assert k <= 10, f"Expected k <= 10 for 20 repetitive items, got k={k}" + + +class TestModerateDiversity: + def test_half_unique_scales(self): + """20 items, 50% unique → k should be in middle range.""" + items = _make_mixed_items(20, unique_fraction=0.5) + k = compute_optimal_k(items) + assert 6 <= k <= 16, f"Expected 6 <= k <= 16 for 50% unique, got k={k}" + + +class TestKneeInteraction: + def test_knee_with_high_diversity_gets_floor(self): + """Even if knee is found at low value, high diversity boosts k.""" + # Create items that have a weak bigram knee but are all unique via SimHash + items = _make_unique_items(15) + k = compute_optimal_k(items) + # With diversity_ratio ~1.0, diversity_floor should boost k + assert k >= 10, f"Expected k >= 10 with high diversity floor, got k={k}" + + def test_knee_with_low_diversity_stays(self): + """Low diversity + knee found → k stays at knee.""" + items = _make_repetitive_items(15, templates=3) + k = compute_optimal_k(items) + assert k <= 8, f"Expected knee-derived k <= 8 for low diversity, got k={k}" + + +class TestBiasAndCaps: + def test_bias_increases_k(self): + """Bias > 1 should increase k.""" + items = _make_unique_items(15) + k_normal = compute_optimal_k(items, bias=1.0) + k_biased = compute_optimal_k(items, bias=1.5) + assert k_biased >= k_normal + + def test_bias_decreases_k(self): + """Bias < 1 should decrease k.""" + items = _make_unique_items(15) + k_normal = compute_optimal_k(items, bias=1.0) + k_biased = compute_optimal_k(items, bias=0.5) + assert k_biased <= k_normal + + def test_max_k_cap_respected(self): + """Even with high diversity, max_k cap is honored.""" + items = _make_unique_items(20) + k = compute_optimal_k(items, max_k=5) + assert k <= 5 + + def test_min_k_floor_respected(self): + """Even with low diversity, min_k floor is honored.""" + items = [json.dumps({"x": 1})] * 20 + k = compute_optimal_k(items, min_k=3) + assert k >= 3 diff --git a/tests/test_transforms/test_universal_json_crush.py b/tests/test_transforms/test_universal_json_crush.py index 8085608f9..1791fd246 100644 --- a/tests/test_transforms/test_universal_json_crush.py +++ b/tests/test_transforms/test_universal_json_crush.py @@ -243,8 +243,10 @@ class TestCrushMixedArray: def test_basic_compression(self, crusher_large_k): mixed = [{"id": i} for i in range(30)] + [f"msg_{i}" for i in range(30)] crushed, strategy = crusher_large_k._crush_mixed_array(mixed) - assert len(crushed) < len(mixed) - assert "mixed:adaptive" in strategy + # With diversity-aware K, unique items may all be kept. + # Verify compression happened OR items are preserved due to high diversity. + assert len(crushed) <= len(mixed) + assert "mixed" in strategy def test_small_groups_kept(self, crusher): # 50 dicts + 3 strings (below threshold) @@ -293,16 +295,17 @@ class TestCrushMixedArray: class TestAdaptiveK: def test_scales_with_n(self, crusher): - """K grows sublinearly with collection size.""" + """K grows sublinearly with collection size (or saturates at max_items).""" small = [f"item_{i}" for i in range(20)] large = [f"item_{i}" for i in range(500)] k_small = crusher._compute_k_split(small)[0] k_large = crusher._compute_k_split(large)[0] - assert k_large > k_small - # Should be sublinear: ratio of K should be much less than ratio of n - assert k_large / k_small < 500 / 20 + # With diversity-aware K, both may saturate at max_items_after_crush + # for highly unique items. The key property: k never exceeds max_items. + assert k_large >= k_small + assert k_large <= crusher.config.max_items_after_crush def test_respects_max_items(self, crusher): items = [f"item_{i}" for i in range(1000)]