"""Tests for tool crusher transform.""" import json from headroom import OpenAIProvider, Tokenizer, ToolCrusherConfig from headroom.transforms import ToolCrusher # Create a shared provider for tests _provider = OpenAIProvider() def get_tokenizer(model: str = "gpt-4o") -> Tokenizer: """Get a tokenizer for tests using OpenAI provider.""" token_counter = _provider.get_token_counter(model) return Tokenizer(token_counter, model) class TestToolCrusher: """Tests for ToolCrusher transform.""" def test_small_tool_output_unchanged(self): """Small tool outputs should not be modified.""" messages = [ {"role": "system", "content": "You are helpful."}, {"role": "tool", "tool_call_id": "call_1", "content": '{"status": "ok"}'}, ] crusher = ToolCrusher() tokenizer = get_tokenizer() result = crusher.apply(messages, tokenizer) # Should not be modified (too small) assert result.messages[1]["content"] == '{"status": "ok"}' assert len(result.transforms_applied) == 0 def test_large_json_array_truncated(self): """Large arrays should be truncated.""" large_array = [{"id": i, "name": f"Item {i}"} for i in range(50)] large_json = json.dumps({"results": large_array}) messages = [ {"role": "system", "content": "You are helpful."}, {"role": "tool", "tool_call_id": "call_1", "content": large_json}, ] config = ToolCrusherConfig(min_tokens_to_crush=50, max_array_items=5) crusher = ToolCrusher(config) tokenizer = get_tokenizer() result = crusher.apply(messages, tokenizer) # Should be modified tool_content = result.messages[1]["content"] parsed = json.loads(tool_content.split("\n 5: break assert depth <= 4 # Should be limited def test_digest_marker_added(self): """Digest marker should be added to crushed content.""" large_data = {"items": list(range(100))} messages = [ {"role": "system", "content": "You are helpful."}, {"role": "tool", "tool_call_id": "call_1", "content": json.dumps(large_data)}, ] config = ToolCrusherConfig(min_tokens_to_crush=10, max_array_items=5) crusher = ToolCrusher(config) tokenizer = get_tokenizer() result = crusher.apply(messages, tokenizer) tool_content = result.messages[1]["content"] # Should have digest marker assert "