mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
14ecab6bfa
commit
57b2de525c
2 changed files with 837 additions and 1 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue