mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Diversity-aware SmartCrusher: keep unique items, compress text within
Root fix: compute_optimal_k() now scales k with content diversity using the SimHash uniqueness ratio already computed in the function. diversity ~1.0 → keep 100% of items (all unique, dropping any loses info) diversity ~0.5 → keep ~65% diversity ~0.0 → keep ~30% (same as before for repetitive data) No hardcoded RAG detection. No field name heuristics. Pure statistics — works for any JSON array regardless of source (Pinecone, Chroma, Weaviate, LangChain, custom APIs). When all items are kept (high diversity), SmartCrusher tries to compress text WITHIN each item's long string fields using Kompress (if available). Falls back gracefully when Kompress is not installed. Before: 12 unique RAG chunks → kept 2, dropped 10 (0/6 key concepts) After: 12 unique RAG chunks → kept 12, compressed within (6/6 concepts) Also adds tests/test_adaptive_sizer.py (16 tests covering high/low/moderate diversity, knee interactions, bias, caps). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
chopratejas
parent
bd350f92ba
commit
f582c1932a
4 changed files with 296 additions and 9 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
186
tests/test_adaptive_sizer.py
Normal file
186
tests/test_adaptive_sizer.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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)]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue