Implement COMPRESS_FIRST strategy for IntelligentContextManager

When context is <10% over budget, try deeper compression of tool messages
before dropping. Uses ContentRouter integration for intelligent routing to
SmartCrusher, CodeAwareCompressor, SearchCompressor, or LogCompressor.

- Add _get_content_router() with lazy loading and aggressive config
- Add _apply_compress_first() to compress tool messages via ContentRouter
- Add _get_tool_source_hint() to extract hints from tool calls
- Add _compress_content_blocks() for Anthropic-style content blocks
- Falls back to DROP_BY_SCORE if compression isn't enough

Adds 14 comprehensive integration tests (no mocks):
- TestCompressFirstStrategy: core functionality (8 tests)
- TestCompressFirstWithContentBlocks: Anthropic format
- TestCompressFirstIntegrationWithTOIN: TOIN integration
- TestCompressFirstEdgeCases: edge cases (4 tests)
This commit is contained in:
chopratejas 2026-01-18 22:55:32 -08:00
parent 14ecab6bfa
commit 57b2de525c
2 changed files with 837 additions and 1 deletions

View file

@ -8,6 +8,12 @@ All importance signals are derived from:
1. Computed metrics (recency, density, references)
2. TOIN-learned patterns (field_semantics, retrieval_rate)
3. Embedding similarity (optional)
Strategy Selection:
- NONE: Under budget, no action needed
- COMPRESS_FIRST: When <compress_threshold over budget, try deeper compression
of tool outputs using ContentRouter before dropping messages
- DROP_BY_SCORE: When significantly over budget, drop lowest-scored messages
"""
from __future__ import annotations
@ -25,6 +31,7 @@ from .scoring import MessageScore, MessageScorer
if TYPE_CHECKING:
from ..telemetry.toin import ToolIntelligenceNetwork
from .content_router import ContentRouter
logger = logging.getLogger(__name__)
@ -85,6 +92,9 @@ class IntelligentContextManager(Transform):
recency_decay_rate=self.config.recency_decay_rate,
)
# Lazy-loaded content router for COMPRESS_FIRST strategy
self._content_router: ContentRouter | None = None
def should_apply(
self,
messages: list[dict[str, Any]],
@ -145,8 +155,52 @@ class IntelligentContextManager(Transform):
strategy = self._select_strategy(current_tokens, available)
logger.debug(f"IntelligentContextManager: selected strategy {strategy.value}")
# Get protected indices and tool units
# Get protected indices
protected = self._get_protected_indices(result_messages)
# ========== COMPRESS_FIRST STRATEGY ==========
# Try to compress tool messages before dropping anything
if strategy == ContextStrategy.COMPRESS_FIRST:
result_messages, compress_transforms, tokens_saved = self._apply_compress_first(
result_messages, tokenizer, protected
)
transforms_applied.extend(compress_transforms)
# Recheck token count after compression
current_tokens = tokenizer.count_messages(result_messages)
# If now under budget, we're done!
if current_tokens <= available:
logger.info(
"IntelligentContextManager: COMPRESS_FIRST succeeded, "
"saved %d tokens: %d -> %d",
tokens_saved,
tokens_before,
current_tokens,
)
return TransformResult(
messages=result_messages,
tokens_before=tokens_before,
tokens_after=current_tokens,
transforms_applied=transforms_applied,
markers_inserted=markers_inserted,
warnings=warnings,
)
# Still over budget, fall through to DROP_BY_SCORE
logger.debug(
"IntelligentContextManager: COMPRESS_FIRST saved %d tokens but still "
"over budget (%d > %d), proceeding to DROP_BY_SCORE",
tokens_saved,
current_tokens,
available,
)
strategy = ContextStrategy.DROP_BY_SCORE
# Need to recalculate protected indices after compression
protected = self._get_protected_indices(result_messages)
# ========== DROP_BY_SCORE STRATEGY ==========
# Get tool units for atomic dropping
tool_units = find_tool_units(result_messages)
tool_unit_indices = self._get_tool_unit_indices(tool_units)
@ -258,6 +312,223 @@ class IntelligentContextManager(Transform):
return ContextStrategy.DROP_BY_SCORE
def _get_content_router(self) -> ContentRouter | None:
"""Get or create content router for COMPRESS_FIRST strategy (lazy load)."""
if self._content_router is None:
try:
from .content_router import ContentRouter, ContentRouterConfig
# Configure for aggressive compression in COMPRESS_FIRST context
router_config = ContentRouterConfig(
enable_code_aware=True,
enable_llmlingua=True,
enable_smart_crusher=True,
enable_search_compressor=True,
enable_log_compressor=True,
skip_user_messages=True,
protect_recent_code=0, # Don't protect in COMPRESS_FIRST
protect_analysis_context=False, # We're over budget
min_section_tokens=20,
ccr_enabled=True,
)
self._content_router = ContentRouter(config=router_config)
except ImportError:
logger.debug("ContentRouter not available for COMPRESS_FIRST")
return self._content_router
def _apply_compress_first(
self,
messages: list[dict[str, Any]],
tokenizer: Tokenizer,
protected: set[int],
) -> tuple[list[dict[str, Any]], list[str], int]:
"""Apply deeper compression to tool messages using ContentRouter.
This is the COMPRESS_FIRST strategy: try to compress tool outputs
more aggressively before falling back to dropping messages.
Args:
messages: List of messages to compress.
tokenizer: Tokenizer for counting.
protected: Set of protected message indices.
Returns:
Tuple of (compressed_messages, transforms_applied, tokens_saved).
"""
router = self._get_content_router()
if router is None:
return messages, [], 0
compressed_messages = deep_copy_messages(messages)
transforms_applied: list[str] = []
total_tokens_saved = 0
for i, msg in enumerate(compressed_messages):
# Skip protected messages
if i in protected:
continue
role = msg.get("role")
content = msg.get("content")
# Focus on tool messages (highest compression potential)
if role == "tool" and isinstance(content, str) and len(content) > 100:
try:
# Get source hint from tool call context
tool_call_id = msg.get("tool_call_id", "")
source_hint = self._get_tool_source_hint(messages, tool_call_id)
# Compress using ContentRouter
result = router.compress(
content,
source_hint=source_hint,
context="", # No specific context in compress-first mode
)
# Check if compression was effective
if result.compression_ratio < 0.9: # At least 10% savings
tokens_before = tokenizer.count_text(content)
tokens_after = tokenizer.count_text(result.compressed)
tokens_saved = tokens_before - tokens_after
if tokens_saved > 0:
compressed_messages[i] = {
**msg,
"content": result.compressed,
}
transforms_applied.append(
f"compress_first:{result.strategy_used.value}:{i}"
)
total_tokens_saved += tokens_saved
logger.debug(
"COMPRESS_FIRST: message %d compressed %d%d tokens (%s)",
i,
tokens_before,
tokens_after,
result.strategy_used.value,
)
except Exception as e:
logger.debug("COMPRESS_FIRST: compression failed for message %d: %s", i, e)
continue
# Also try to compress assistant messages with tool results in content blocks
elif role == "assistant" and isinstance(content, list):
compressed_blocks, block_transforms, block_saved = self._compress_content_blocks(
content, router, tokenizer
)
if block_saved > 0:
compressed_messages[i] = {**msg, "content": compressed_blocks}
transforms_applied.extend(block_transforms)
total_tokens_saved += block_saved
if total_tokens_saved > 0:
logger.info(
"COMPRESS_FIRST: saved %d tokens across %d compressions",
total_tokens_saved,
len(transforms_applied),
)
return compressed_messages, transforms_applied, total_tokens_saved
def _get_tool_source_hint(self, messages: list[dict[str, Any]], tool_call_id: str) -> str:
"""Extract source hint from the tool call that produced this result.
Args:
messages: List of all messages.
tool_call_id: The ID of the tool call to find.
Returns:
Source hint string (e.g., "tool:grep", "file:main.py").
"""
if not tool_call_id:
return ""
# Find the assistant message with this tool call
for msg in messages:
if msg.get("role") == "assistant" and msg.get("tool_calls"):
for tc in msg.get("tool_calls", []):
if tc.get("id") == tool_call_id:
func = tc.get("function", {})
tool_name = func.get("name", "")
# Import here to avoid circular imports
try:
import json
from .content_router import generate_source_hint
args_str = func.get("arguments", "{}")
try:
args = (
json.loads(args_str) if isinstance(args_str, str) else args_str
)
except json.JSONDecodeError:
args = {}
return generate_source_hint(tool_name, args)
except ImportError:
return f"tool:{tool_name}" if tool_name else ""
return ""
def _compress_content_blocks(
self,
content_blocks: list[Any],
router: ContentRouter,
tokenizer: Tokenizer,
) -> tuple[list[Any], list[str], int]:
"""Compress content blocks (Anthropic format) using ContentRouter.
Args:
content_blocks: List of content blocks.
router: ContentRouter instance.
tokenizer: Tokenizer for counting.
Returns:
Tuple of (compressed_blocks, transforms_applied, tokens_saved).
"""
compressed_blocks: list[Any] = []
transforms_applied: list[str] = []
total_saved = 0
for block in content_blocks:
if not isinstance(block, dict):
compressed_blocks.append(block)
continue
block_type = block.get("type")
# Handle tool_result blocks
if block_type == "tool_result":
tool_content = block.get("content", "")
if isinstance(tool_content, str) and len(tool_content) > 200:
try:
result = router.compress(tool_content, context="")
if result.compression_ratio < 0.9:
tokens_before = tokenizer.count_text(tool_content)
tokens_after = tokenizer.count_text(result.compressed)
saved = tokens_before - tokens_after
if saved > 0:
compressed_blocks.append(
{
**block,
"content": result.compressed,
}
)
transforms_applied.append(
f"compress_first:block:{result.strategy_used.value}"
)
total_saved += saved
continue
except Exception:
pass
compressed_blocks.append(block)
return compressed_blocks, transforms_applied, total_saved
def _get_protected_indices(self, messages: list[dict[str, Any]]) -> set[int]:
"""Get indices that should never be dropped."""
protected: set[int] = set()

View file

@ -714,3 +714,568 @@ class TestCustomWeights:
# Should complete successfully
assert len(result.messages) < len(long_conversation)
# =============================================================================
# Test COMPRESS_FIRST Strategy - Integration Tests
# =============================================================================
class TestCompressFirstStrategy:
"""Integration tests for COMPRESS_FIRST strategy.
These tests verify that:
1. COMPRESS_FIRST is selected when slightly over budget
2. ContentRouter actually compresses tool messages
3. Compression can bring context under budget
4. Fallback to DROP_BY_SCORE works when compression isn't enough
"""
@pytest.fixture
def conversation_with_large_tool_outputs(self) -> list[dict[str, Any]]:
"""Conversation with large JSON tool outputs (compressible)."""
import json
# Generate a large JSON array that SmartCrusher can compress
large_results = [
{
"id": i,
"name": f"Item {i}",
"status": "active" if i % 2 == 0 else "inactive",
"value": i * 100,
"description": f"This is a description for item number {i} with some extra text",
}
for i in range(100)
]
return [
{"role": "system", "content": "You are a helpful assistant with search tools."},
{"role": "user", "content": "Search for items in the database."},
{
"role": "assistant",
"content": "I'll search the database for you.",
"tool_calls": [
{
"id": "call_db_1",
"type": "function",
"function": {
"name": "database_search",
"arguments": '{"query": "items"}',
},
}
],
},
{
"role": "tool",
"tool_call_id": "call_db_1",
"content": json.dumps(large_results),
},
{"role": "assistant", "content": "I found 100 items in the database."},
{"role": "user", "content": "Great, can you show me more details?"},
]
@pytest.fixture
def conversation_with_search_output(self) -> list[dict[str, Any]]:
"""Conversation with grep-style search output (compressible)."""
# Generate search results in grep format
search_lines = [
f"src/module{i}.py:{i * 10}: def function_{i}(self, param):" for i in range(50)
]
return [
{"role": "system", "content": "You are a code assistant."},
{"role": "user", "content": "Search for function definitions."},
{
"role": "assistant",
"content": "Searching...",
"tool_calls": [
{
"id": "call_grep_1",
"type": "function",
"function": {
"name": "Grep",
"arguments": '{"pattern": "def function"}',
},
}
],
},
{
"role": "tool",
"tool_call_id": "call_grep_1",
"content": "\n".join(search_lines),
},
{"role": "assistant", "content": "Found 50 function definitions."},
{"role": "user", "content": "Thanks!"},
]
def test_compress_first_selected_for_small_overage(self, tokenizer: Tokenizer):
"""COMPRESS_FIRST should be selected when <10% over budget."""
config = IntelligentContextConfig(compress_threshold=0.10)
manager = IntelligentContextManager(config=config)
# 5% over budget should select COMPRESS_FIRST
strategy = manager._select_strategy(current_tokens=2100, available=2000)
assert strategy == ContextStrategy.COMPRESS_FIRST
# 9% over budget should still select COMPRESS_FIRST
strategy = manager._select_strategy(current_tokens=2180, available=2000)
assert strategy == ContextStrategy.COMPRESS_FIRST
# 15% over budget should select DROP_BY_SCORE
strategy = manager._select_strategy(current_tokens=2300, available=2000)
assert strategy == ContextStrategy.DROP_BY_SCORE
def test_compress_first_compresses_json_tool_output(
self,
conversation_with_large_tool_outputs: list[dict[str, Any]],
tokenizer: Tokenizer,
):
"""COMPRESS_FIRST should compress JSON tool outputs using ContentRouter."""
config = IntelligentContextConfig(compress_threshold=0.15)
manager = IntelligentContextManager(config=config)
tokens_before = tokenizer.count_messages(conversation_with_large_tool_outputs)
# Set limit to be slightly over (within COMPRESS_FIRST range)
# We want tokens_before to be ~5-10% over the limit
target_limit = int(tokens_before / 1.05) # ~5% over
result = manager.apply(
conversation_with_large_tool_outputs,
tokenizer,
model_limit=target_limit,
output_buffer=50,
)
# Should have compression transforms or be under budget
if result.tokens_after <= target_limit - 50:
# If under budget, compression worked!
assert result.tokens_after < result.tokens_before
else:
# May have needed to drop as well
assert result.tokens_after <= result.tokens_before
def test_compress_first_with_search_output(
self,
conversation_with_search_output: list[dict[str, Any]],
tokenizer: Tokenizer,
):
"""COMPRESS_FIRST should work with search-style output."""
config = IntelligentContextConfig(compress_threshold=0.15)
manager = IntelligentContextManager(config=config)
tokens_before = tokenizer.count_messages(conversation_with_search_output)
target_limit = int(tokens_before / 1.08) # ~8% over
result = manager.apply(
conversation_with_search_output,
tokenizer,
model_limit=target_limit,
output_buffer=50,
)
# Should reduce tokens
assert result.tokens_after <= result.tokens_before
def test_compress_first_fallback_to_drop(
self,
tokenizer: Tokenizer,
):
"""When compression isn't enough, should fall back to dropping."""
import json
# Create a conversation with multiple tool calls where even compression
# won't be enough - use small non-JSON content that can't compress well
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Do multiple searches."},
]
# Add 10 tool calls with results that won't compress much
for i in range(10):
messages.append(
{
"role": "assistant",
"content": f"Searching for item {i}...",
"tool_calls": [
{
"id": f"call_{i}",
"type": "function",
"function": {
"name": "search",
"arguments": json.dumps({"q": f"item{i}"}),
},
}
],
}
)
messages.append(
{
"role": "tool",
"tool_call_id": f"call_{i}",
"content": f"Found result for item {i}: some important data here that cannot be compressed easily",
}
)
messages.append({"role": "assistant", "content": "Here are all the results."})
messages.append({"role": "user", "content": "Thanks!"})
# Use keep_last_turns=1 to allow more messages to be dropped
config = IntelligentContextConfig(
compress_threshold=0.50, # High threshold
keep_last_turns=1, # Only protect last turn
)
manager = IntelligentContextManager(config=config)
tokens_before = tokenizer.count_messages(messages)
# Very small limit that will require dropping
very_small_limit = tokens_before // 4
result = manager.apply(
messages,
tokenizer,
model_limit=very_small_limit,
output_buffer=50,
)
# Should have reduced tokens
assert result.tokens_after < result.tokens_before
# Should have dropped some messages
assert len(result.messages) < len(messages)
def test_compress_first_preserves_message_structure(
self,
conversation_with_large_tool_outputs: list[dict[str, Any]],
tokenizer: Tokenizer,
):
"""COMPRESS_FIRST should preserve message structure integrity."""
config = IntelligentContextConfig(compress_threshold=0.20)
manager = IntelligentContextManager(config=config)
tokens_before = tokenizer.count_messages(conversation_with_large_tool_outputs)
target_limit = int(tokens_before / 1.05)
result = manager.apply(
conversation_with_large_tool_outputs,
tokenizer,
model_limit=target_limit,
output_buffer=50,
)
# Verify structure
for msg in result.messages:
assert "role" in msg
role = msg["role"]
assert role in ("system", "user", "assistant", "tool")
# Tool messages should have tool_call_id
if role == "tool":
assert "tool_call_id" in msg or "content" in msg
# Assistant messages with tool_calls should have that structure
if role == "assistant" and "tool_calls" in msg:
for tc in msg["tool_calls"]:
assert "id" in tc
assert "function" in tc
def test_compress_first_no_compression_when_under_budget(
self, simple_conversation: list[dict[str, Any]], tokenizer: Tokenizer
):
"""COMPRESS_FIRST should not be applied when under budget."""
manager = IntelligentContextManager()
result = manager.apply(
simple_conversation,
tokenizer,
model_limit=128000,
output_buffer=4000,
)
# No compression transforms should be applied
compression_transforms = [
t for t in result.transforms_applied if t.startswith("compress_first:")
]
assert len(compression_transforms) == 0
assert result.tokens_before == result.tokens_after
def test_content_router_lazy_loading(self):
"""ContentRouter should be lazy-loaded only when needed."""
manager = IntelligentContextManager()
# Initially None
assert manager._content_router is None
# Get router
router = manager._get_content_router()
# Should now be set
assert manager._content_router is not None
assert router is manager._content_router
# Second call should return same instance
router2 = manager._get_content_router()
assert router is router2
def test_source_hint_extraction(self):
"""Source hints should be extracted from tool calls."""
manager = IntelligentContextManager()
messages = [
{
"role": "assistant",
"tool_calls": [
{
"id": "call_1",
"function": {
"name": "Read",
"arguments": '{"file_path": "/src/main.py"}',
},
}
],
},
{
"role": "assistant",
"tool_calls": [
{
"id": "call_2",
"function": {
"name": "Grep",
"arguments": '{"pattern": "def"}',
},
}
],
},
]
# Test file read hint
hint1 = manager._get_tool_source_hint(messages, "call_1")
assert "file:" in hint1 or hint1 == "" # May not have content_router import
# Test grep hint
hint2 = manager._get_tool_source_hint(messages, "call_2")
assert "grep" in hint2.lower() or hint2 == ""
# Test unknown tool call
hint3 = manager._get_tool_source_hint(messages, "unknown")
assert hint3 == ""
class TestCompressFirstWithContentBlocks:
"""Tests for COMPRESS_FIRST with Anthropic-style content blocks."""
@pytest.fixture
def conversation_with_content_blocks(self) -> list[dict[str, Any]]:
"""Conversation with Anthropic-style content blocks."""
import json
large_result = json.dumps([{"id": i, "data": f"item_{i}" * 20} for i in range(50)])
return [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Search for data."},
{
"role": "assistant",
"content": [
{"type": "text", "text": "Here are the results:"},
{
"type": "tool_result",
"tool_use_id": "tool_1",
"content": large_result,
},
],
},
{"role": "user", "content": "Thanks!"},
]
def test_compress_first_handles_content_blocks(
self,
conversation_with_content_blocks: list[dict[str, Any]],
tokenizer: Tokenizer,
):
"""COMPRESS_FIRST should handle content blocks format."""
config = IntelligentContextConfig(compress_threshold=0.20)
manager = IntelligentContextManager(config=config)
tokens_before = tokenizer.count_messages(conversation_with_content_blocks)
target_limit = int(tokens_before / 1.08)
result = manager.apply(
conversation_with_content_blocks,
tokenizer,
model_limit=target_limit,
output_buffer=50,
)
# Should complete without error
assert result.messages is not None
assert result.tokens_after <= result.tokens_before
class TestCompressFirstIntegrationWithTOIN:
"""Integration tests for COMPRESS_FIRST with TOIN patterns."""
def test_compress_first_works_without_toin(self, tokenizer: Tokenizer):
"""COMPRESS_FIRST should work without TOIN integration."""
import json
messages = [
{"role": "system", "content": "System"},
{"role": "user", "content": "Search"},
{
"role": "assistant",
"tool_calls": [
{
"id": "c1",
"type": "function",
"function": {"name": "search", "arguments": "{}"},
}
],
"content": "",
},
{
"role": "tool",
"tool_call_id": "c1",
"content": json.dumps([{"x": i} for i in range(50)]),
},
{"role": "assistant", "content": "Done"},
{"role": "user", "content": "Thanks"},
]
# Without TOIN
config = IntelligentContextConfig(
compress_threshold=0.15,
toin_integration=False,
)
manager = IntelligentContextManager(config=config, toin=None)
tokens_before = tokenizer.count_messages(messages)
target_limit = int(tokens_before / 1.08)
result = manager.apply(
messages,
tokenizer,
model_limit=target_limit,
output_buffer=50,
)
# Should work
assert result.messages is not None
assert result.tokens_after <= result.tokens_before
class TestCompressFirstEdgeCases:
"""Edge case tests for COMPRESS_FIRST strategy."""
def test_empty_tool_content(self, tokenizer: Tokenizer):
"""COMPRESS_FIRST should handle empty tool content gracefully."""
messages = [
{"role": "system", "content": "System"},
{"role": "user", "content": "Do something"},
{
"role": "assistant",
"tool_calls": [
{
"id": "c1",
"type": "function",
"function": {"name": "tool", "arguments": "{}"},
}
],
"content": "",
},
{"role": "tool", "tool_call_id": "c1", "content": ""},
{"role": "assistant", "content": "Done"},
]
config = IntelligentContextConfig(compress_threshold=0.50)
manager = IntelligentContextManager(config=config)
# Very small limit to trigger compression
result = manager.apply(
messages,
tokenizer,
model_limit=50,
output_buffer=10,
)
# Should handle gracefully
assert result.messages is not None
def test_non_json_tool_content(self, tokenizer: Tokenizer):
"""COMPRESS_FIRST should handle non-JSON tool content."""
messages = [
{"role": "system", "content": "System"},
{"role": "user", "content": "Read a file"},
{
"role": "assistant",
"tool_calls": [
{
"id": "c1",
"type": "function",
"function": {"name": "Read", "arguments": '{"file_path": "test.py"}'},
}
],
"content": "",
},
{
"role": "tool",
"tool_call_id": "c1",
"content": "def hello():\n print('Hello World')\n" * 20,
},
{"role": "assistant", "content": "Here's the file"},
{"role": "user", "content": "Thanks"},
]
config = IntelligentContextConfig(compress_threshold=0.20)
manager = IntelligentContextManager(config=config)
tokens_before = tokenizer.count_messages(messages)
target_limit = int(tokens_before / 1.08)
result = manager.apply(
messages,
tokenizer,
model_limit=target_limit,
output_buffer=50,
)
# Should handle gracefully
assert result.messages is not None
assert result.tokens_after <= result.tokens_before
def test_protected_tool_messages_not_compressed(self, tokenizer: Tokenizer):
"""Protected tool messages should not be compressed."""
import json
messages = [
{"role": "system", "content": "System"},
{"role": "user", "content": "Search"},
{
"role": "assistant",
"tool_calls": [
{"id": "c1", "type": "function", "function": {"name": "s", "arguments": "{}"}}
],
"content": "",
},
{
"role": "tool",
"tool_call_id": "c1",
"content": json.dumps([{"x": i} for i in range(100)]),
},
{"role": "assistant", "content": "Found results"},
{"role": "user", "content": "More please"},
]
# Protect last 5 turns (should include the tool message)
config = IntelligentContextConfig(
keep_last_turns=5,
compress_threshold=0.50,
)
manager = IntelligentContextManager(config=config)
# Get protected indices
protected = manager._get_protected_indices(messages)
# The recent messages should be protected
# With 6 messages and keep_last_turns=5, most should be protected
assert len(protected) > 0