mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
refactor(cache): isolate compression strategy outcomes (#1938)
## Description Extracts local compression strategy accounting out of `CompressionFeedback` into a pure cache-domain object. This keeps strategy counters, retrieval-rate math, pruning, and best-strategy selection independently testable while preserving the existing `LocalToolPattern` public API. Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [x] Code refactoring (no functional changes) ## Changes Made - Added `CompressionStrategyOutcomes` as the strategy-outcome domain for compression/retrieval counters, pruning, retrieval rates, and recommendation selection. - Updated `LocalToolPattern` and `CompressionFeedback` to delegate strategy accounting to that domain while keeping existing fields and methods intact. - Added direct unit coverage for strategy outcome math and bounded pruning behavior. - Updated the LiteLLM callback hook signature to remain compatible with current LiteLLM typing and the existing three-argument call shape. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text python -m ruff check . All checks passed! python -m ruff format --check . 1095 files already formatted python -m mypy headroom --ignore-missing-imports headroom\proxy\server.py:1457: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs [annotation-unchecked] headroom\proxy\server.py:1468: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs [annotation-unchecked] Success: no issues found in 409 source files python -m pytest tests/test_compression_strategy_outcomes.py tests/test_ccr_feedback.py tests/test_toin_fixes.py tests/test_litellm_callback.py tests/test_compress_api.py::TestLiteLLMCallback -q collected 54 items 46 passed, 8 skipped in 6.58s ``` ## Real Behavior Proof - Environment: Windows, Python 3.13.13, clean worktree `C:\git\headroom-pr-slice5` - Exact command / steps: ran the lint, format, type-check, and focused pytest commands listed above. - Observed result: strategy outcome tests and existing feedback/TOIN/LiteLLM compatibility tests pass; repo-wide lint/type validation passes. - Not tested: full pytest suite and Docker/native CI jobs are left to GitHub Actions. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A. ## Additional Notes - Documentation and changelog are N/A for this internal refactor. - Manual UI testing is N/A; this is cache feedback and integration callback logic. - Comment checklist is unchecked because the extracted object is intentionally straightforward and covered by tests.
This commit is contained in:
JD Davis
GitHub
parent
41af39d769
commit
b5aa8a358e
3 changed files with 174 additions and 67 deletions
97
headroom/cache/compression_feedback.py
vendored
97
headroom/cache/compression_feedback.py
vendored
|
|
@ -33,6 +33,8 @@ import time
|
|||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .compression_strategy_outcomes import CompressionStrategyOutcomes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .compression_store import CompressionStore, RetrievalEvent
|
||||
|
||||
|
|
@ -94,28 +96,33 @@ class LocalToolPattern:
|
|||
|
||||
def strategy_retrieval_rate(self, strategy: str) -> float:
|
||||
"""Get retrieval rate for a specific compression strategy."""
|
||||
compressions = self.strategy_compressions.get(strategy, 0)
|
||||
if compressions == 0:
|
||||
return 0.0
|
||||
retrievals = self.strategy_retrievals.get(strategy, 0)
|
||||
return retrievals / compressions
|
||||
return self.strategy_outcomes.retrieval_rate(strategy)
|
||||
|
||||
def best_strategy(self) -> str | None:
|
||||
"""Find the strategy with lowest retrieval rate (most successful)."""
|
||||
if not self.strategy_compressions:
|
||||
return None
|
||||
return self.strategy_outcomes.best_strategy()
|
||||
|
||||
best = None
|
||||
best_rate = 1.0
|
||||
@property
|
||||
def strategy_outcomes(self) -> CompressionStrategyOutcomes:
|
||||
"""Strategy outcome view backed by this pattern's public counters."""
|
||||
return CompressionStrategyOutcomes(
|
||||
compressions=self.strategy_compressions,
|
||||
retrievals=self.strategy_retrievals,
|
||||
)
|
||||
|
||||
for strategy in self.strategy_compressions:
|
||||
rate = self.strategy_retrieval_rate(strategy)
|
||||
# Only consider strategies with enough samples
|
||||
if self.strategy_compressions[strategy] >= 3 and rate < best_rate:
|
||||
best_rate = rate
|
||||
best = strategy
|
||||
def record_strategy_compression(self, strategy: str) -> None:
|
||||
"""Record strategy compression outcome."""
|
||||
outcomes = self.strategy_outcomes
|
||||
outcomes.record_compression(strategy)
|
||||
self.strategy_compressions = outcomes.compressions
|
||||
self.strategy_retrievals = outcomes.retrievals
|
||||
|
||||
return best
|
||||
def record_strategy_retrieval(self, strategy: str) -> None:
|
||||
"""Record strategy retrieval outcome."""
|
||||
outcomes = self.strategy_outcomes
|
||||
outcomes.record_retrieval(strategy)
|
||||
self.strategy_compressions = outcomes.compressions
|
||||
self.strategy_retrievals = outcomes.retrievals
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -235,15 +242,7 @@ class CompressionFeedback:
|
|||
|
||||
# Track strategy usage
|
||||
if strategy:
|
||||
pattern.strategy_compressions[strategy] = (
|
||||
pattern.strategy_compressions.get(strategy, 0) + 1
|
||||
)
|
||||
|
||||
# CRITICAL FIX: When truncating strategy dicts, keep them in sync
|
||||
# to prevent desync between compressions and retrievals.
|
||||
# Both dicts must have the same keys for accurate retrieval rate calculation.
|
||||
if len(pattern.strategy_compressions) > 50:
|
||||
self._truncate_strategy_dicts(pattern)
|
||||
pattern.record_strategy_compression(strategy)
|
||||
|
||||
# Track signature hash for TOIN correlation
|
||||
if tool_signature_hash:
|
||||
|
|
@ -291,14 +290,7 @@ class CompressionFeedback:
|
|||
|
||||
# Track strategy retrievals (for success rate calculation)
|
||||
if strategy:
|
||||
pattern.strategy_retrievals[strategy] = (
|
||||
pattern.strategy_retrievals.get(strategy, 0) + 1
|
||||
)
|
||||
|
||||
# CRITICAL FIX: When truncating strategy dicts, keep them in sync
|
||||
# to prevent desync between compressions and retrievals.
|
||||
if len(pattern.strategy_retrievals) > 50:
|
||||
self._truncate_strategy_dicts(pattern)
|
||||
pattern.record_strategy_retrieval(strategy)
|
||||
|
||||
# Track query patterns
|
||||
if event.query:
|
||||
|
|
@ -318,40 +310,11 @@ class CompressionFeedback:
|
|||
self._extract_field_hints(pattern, event.query)
|
||||
|
||||
def _truncate_strategy_dicts(self, pattern: LocalToolPattern) -> None:
|
||||
"""Truncate strategy_compressions and strategy_retrievals in sync.
|
||||
|
||||
CRITICAL FIX: Both dicts must have the same keys for accurate retrieval
|
||||
rate calculation. When truncating, we keep the union of top strategies
|
||||
from both dicts, then truncate both to the same key set.
|
||||
"""
|
||||
# Get top 40 strategies from each dict (using 40 to allow union to stay under 50)
|
||||
top_compressions = {
|
||||
k
|
||||
for k, _ in sorted(
|
||||
pattern.strategy_compressions.items(),
|
||||
key=lambda x: x[1],
|
||||
reverse=True,
|
||||
)[:40]
|
||||
}
|
||||
top_retrievals = {
|
||||
k
|
||||
for k, _ in sorted(
|
||||
pattern.strategy_retrievals.items(),
|
||||
key=lambda x: x[1],
|
||||
reverse=True,
|
||||
)[:40]
|
||||
}
|
||||
|
||||
# Keep union of top strategies from both
|
||||
keys_to_keep = top_compressions | top_retrievals
|
||||
|
||||
# Truncate both dicts to same keys
|
||||
pattern.strategy_compressions = {
|
||||
k: v for k, v in pattern.strategy_compressions.items() if k in keys_to_keep
|
||||
}
|
||||
pattern.strategy_retrievals = {
|
||||
k: v for k, v in pattern.strategy_retrievals.items() if k in keys_to_keep
|
||||
}
|
||||
"""Truncate strategy counters using the shared strategy outcome domain."""
|
||||
outcomes = pattern.strategy_outcomes
|
||||
outcomes.prune()
|
||||
pattern.strategy_compressions = outcomes.compressions
|
||||
pattern.strategy_retrievals = outcomes.retrievals
|
||||
|
||||
def _extract_field_hints(self, pattern: LocalToolPattern, query: str) -> None:
|
||||
"""Extract potential field names from search queries.
|
||||
|
|
|
|||
99
headroom/cache/compression_strategy_outcomes.py
vendored
Normal file
99
headroom/cache/compression_strategy_outcomes.py
vendored
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
"""Strategy outcome accounting for local compression feedback."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompressionStrategyOutcomes:
|
||||
"""Track compression and retrieval outcomes by compression strategy."""
|
||||
|
||||
compressions: dict[str, int] = field(default_factory=dict)
|
||||
retrievals: dict[str, int] = field(default_factory=dict)
|
||||
max_strategies: int = 50
|
||||
top_strategies_per_counter: int = 40
|
||||
minimum_samples_for_recommendation: int = 3
|
||||
|
||||
def record_compression(self, strategy: str) -> None:
|
||||
"""Record one compression for a strategy."""
|
||||
self.compressions[strategy] = self.compressions.get(strategy, 0) + 1
|
||||
self.prune()
|
||||
|
||||
def record_retrieval(self, strategy: str) -> None:
|
||||
"""Record one retrieval for a strategy."""
|
||||
self.retrievals[strategy] = self.retrievals.get(strategy, 0) + 1
|
||||
self.prune()
|
||||
|
||||
def retrieval_rate(self, strategy: str) -> float:
|
||||
"""Return the retrievals-per-compression rate for one strategy."""
|
||||
compressions = self.compressions.get(strategy, 0)
|
||||
if compressions == 0:
|
||||
return 0.0
|
||||
return self.retrievals.get(strategy, 0) / compressions
|
||||
|
||||
def best_strategy(self) -> str | None:
|
||||
"""Return the sampled strategy with the lowest retrieval rate."""
|
||||
best = None
|
||||
best_rate = 1.0
|
||||
|
||||
for strategy, compression_count in self.compressions.items():
|
||||
if compression_count < self.minimum_samples_for_recommendation:
|
||||
continue
|
||||
|
||||
rate = self.retrieval_rate(strategy)
|
||||
if rate < best_rate:
|
||||
best = strategy
|
||||
best_rate = rate
|
||||
|
||||
return best
|
||||
|
||||
def prune(self) -> None:
|
||||
"""Bound counters while preserving the highest-signal strategies."""
|
||||
if (
|
||||
len(self.compressions) <= self.max_strategies
|
||||
and len(self.retrievals) <= self.max_strategies
|
||||
):
|
||||
return
|
||||
|
||||
keys_to_keep = self._keys_to_keep()
|
||||
self.compressions = {
|
||||
strategy: count
|
||||
for strategy, count in self.compressions.items()
|
||||
if strategy in keys_to_keep
|
||||
}
|
||||
self.retrievals = {
|
||||
strategy: count
|
||||
for strategy, count in self.retrievals.items()
|
||||
if strategy in keys_to_keep
|
||||
}
|
||||
|
||||
def _keys_to_keep(self) -> set[str]:
|
||||
top_compressions = self._top_keys(self.compressions)
|
||||
top_retrievals = self._top_keys(self.retrievals)
|
||||
candidate_keys = top_compressions | top_retrievals
|
||||
|
||||
if len(candidate_keys) <= self.max_strategies:
|
||||
return candidate_keys
|
||||
|
||||
ranked_keys = sorted(
|
||||
candidate_keys,
|
||||
key=lambda strategy: (
|
||||
self.compressions.get(strategy, 0) + self.retrievals.get(strategy, 0),
|
||||
self.compressions.get(strategy, 0),
|
||||
self.retrievals.get(strategy, 0),
|
||||
strategy,
|
||||
),
|
||||
reverse=True,
|
||||
)
|
||||
return set(ranked_keys[: self.max_strategies])
|
||||
|
||||
def _top_keys(self, counts: dict[str, int]) -> set[str]:
|
||||
return {
|
||||
strategy
|
||||
for strategy, _ in sorted(
|
||||
counts.items(),
|
||||
key=lambda item: (item[1], item[0]),
|
||||
reverse=True,
|
||||
)[: self.top_strategies_per_counter]
|
||||
}
|
||||
45
tests/test_compression_strategy_outcomes.py
Normal file
45
tests/test_compression_strategy_outcomes.py
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
from headroom.cache.compression_strategy_outcomes import CompressionStrategyOutcomes
|
||||
|
||||
|
||||
def test_retrieval_rate_is_zero_without_strategy_compressions():
|
||||
outcomes = CompressionStrategyOutcomes(retrievals={"sample": 2})
|
||||
|
||||
assert outcomes.retrieval_rate("sample") == 0.0
|
||||
|
||||
|
||||
def test_best_strategy_requires_minimum_samples():
|
||||
outcomes = CompressionStrategyOutcomes(
|
||||
compressions={"under_sampled": 2, "sampled": 3},
|
||||
retrievals={"under_sampled": 0, "sampled": 1},
|
||||
)
|
||||
|
||||
assert outcomes.best_strategy() == "sampled"
|
||||
|
||||
|
||||
def test_best_strategy_uses_lowest_retrieval_rate():
|
||||
outcomes = CompressionStrategyOutcomes(
|
||||
compressions={"top_n": 10, "smart_sample": 10},
|
||||
retrievals={"top_n": 7, "smart_sample": 2},
|
||||
)
|
||||
|
||||
assert outcomes.retrieval_rate("smart_sample") == 0.2
|
||||
assert outcomes.best_strategy() == "smart_sample"
|
||||
|
||||
|
||||
def test_recording_prunes_strategy_counters_to_bounded_high_signal_set():
|
||||
outcomes = CompressionStrategyOutcomes(max_strategies=10, top_strategies_per_counter=8)
|
||||
|
||||
for index in range(30):
|
||||
strategy = f"strategy_{index:02d}"
|
||||
for _ in range(index + 1):
|
||||
outcomes.record_compression(strategy)
|
||||
|
||||
for index in range(30):
|
||||
strategy = f"strategy_{index:02d}"
|
||||
for _ in range(30 - index):
|
||||
outcomes.record_retrieval(strategy)
|
||||
|
||||
assert len(outcomes.compressions) <= 10
|
||||
assert len(outcomes.retrievals) <= 10
|
||||
assert "strategy_29" in outcomes.compressions
|
||||
assert "strategy_00" in outcomes.retrievals
|
||||
Loading…
Add table
Add a link
Reference in a new issue