mirror of https://github.com/aliasrobotics/cai.git
275 lines
10 KiB
Python
275 lines
10 KiB
Python
"""
|
|
Tests for dynamic tool output weighting and context optimization.
|
|
|
|
These tests verify that:
|
|
1. Tool outputs are truncated when context usage exceeds thresholds
|
|
2. Older tool outputs are truncated more aggressively than newer ones
|
|
3. Small outputs pass through unchanged
|
|
4. Binary and large outputs are handled correctly
|
|
"""
|
|
|
|
import pytest
|
|
|
|
|
|
# Test the pure logic functions directly without heavy module imports
|
|
class TestTruncateToolOutputPureLogic:
|
|
"""Tests for truncation logic."""
|
|
|
|
def _truncate_tool_output(self, content: str, max_chars: int) -> str:
|
|
"""Local implementation for testing."""
|
|
if len(content) <= max_chars:
|
|
return content
|
|
head_size = max_chars // 2 - 100
|
|
tail_size = max_chars // 2 - 100
|
|
head = content[:head_size]
|
|
tail = content[-tail_size:]
|
|
omitted = len(content) - head_size - tail_size
|
|
return f"{head}\n\n[... TRUNCATED {omitted:,} chars for context optimization ...]\n\n{tail}"
|
|
|
|
def test_small_content_unchanged(self):
|
|
"""Small content should pass through unchanged."""
|
|
content = "Small output"
|
|
result = self._truncate_tool_output(content, 1000)
|
|
assert result == content
|
|
|
|
def test_large_content_truncated(self):
|
|
"""Large content should be truncated to max_chars."""
|
|
content = "X" * 10000
|
|
max_chars = 1000
|
|
result = self._truncate_tool_output(content, max_chars)
|
|
assert len(result) <= max_chars + 200
|
|
assert "TRUNCATED" in result
|
|
|
|
def test_preserves_head_and_tail(self):
|
|
"""Truncation should preserve both head and tail of content."""
|
|
content = "HEAD" + ("X" * 10000) + "TAIL"
|
|
result = self._truncate_tool_output(content, 1000)
|
|
assert result.startswith("HEAD")
|
|
assert result.endswith("TAIL")
|
|
|
|
|
|
class TestDynamicToolOutputWeightingPureLogic:
|
|
"""Tests for the weighting function logic."""
|
|
|
|
TOOL_OUTPUT_WEIGHT_TOKEN_THRESHOLD = 60000
|
|
TOOL_OUTPUT_WEIGHT_CONTEXT_THRESHOLD = 0.50
|
|
TOOL_OUTPUT_MAX_CHARS = 50000
|
|
TOOL_OUTPUT_MIN_CHARS = 2000
|
|
|
|
def _count_tokens(self, messages):
|
|
"""Simple char-based token estimate."""
|
|
total = 0
|
|
for msg in messages:
|
|
content = msg.get("content", "")
|
|
if isinstance(content, str):
|
|
total += len(content) // 4
|
|
return total
|
|
|
|
def _truncate_tool_output(self, content: str, max_chars: int) -> str:
|
|
if len(content) <= max_chars:
|
|
return content
|
|
head_size = max_chars // 2 - 100
|
|
tail_size = max_chars // 2 - 100
|
|
head = content[:head_size]
|
|
tail = content[-tail_size:]
|
|
omitted = len(content) - head_size - tail_size
|
|
return f"{head}\n\n[... TRUNCATED {omitted:,} chars ...]\n\n{tail}"
|
|
|
|
def _apply_weighting(self, messages, max_context_tokens):
|
|
"""Local implementation for testing."""
|
|
if not messages:
|
|
return messages
|
|
|
|
current_tokens = self._count_tokens(messages)
|
|
context_usage = current_tokens / max_context_tokens if max_context_tokens > 0 else 0
|
|
|
|
needs_weighting = (
|
|
current_tokens > self.TOOL_OUTPUT_WEIGHT_TOKEN_THRESHOLD
|
|
or context_usage > self.TOOL_OUTPUT_WEIGHT_CONTEXT_THRESHOLD
|
|
)
|
|
|
|
if not needs_weighting:
|
|
return messages
|
|
|
|
tool_indices = []
|
|
for i, msg in enumerate(messages):
|
|
if msg.get("role") == "tool":
|
|
content = msg.get("content", "")
|
|
if isinstance(content, str) and len(content) > self.TOOL_OUTPUT_MIN_CHARS:
|
|
tool_indices.append((i, len(content)))
|
|
|
|
if not tool_indices:
|
|
return messages
|
|
|
|
if context_usage > 0.85:
|
|
pressure = 1.0
|
|
elif context_usage > 0.70:
|
|
pressure = 0.75
|
|
elif context_usage > self.TOOL_OUTPUT_WEIGHT_CONTEXT_THRESHOLD:
|
|
pressure = 0.5
|
|
elif current_tokens > 100000:
|
|
pressure = 0.6
|
|
elif current_tokens > self.TOOL_OUTPUT_WEIGHT_TOKEN_THRESHOLD:
|
|
pressure = 0.3
|
|
else:
|
|
pressure = 0.2
|
|
|
|
num_tools = len(tool_indices)
|
|
modified_messages = [m.copy() for m in messages]
|
|
|
|
for position, (msg_idx, content_len) in enumerate(tool_indices):
|
|
position_weight = position / num_tools if num_tools > 1 else 1.0
|
|
allocation = (1 - pressure) + (pressure * position_weight)
|
|
max_chars = int(self.TOOL_OUTPUT_MAX_CHARS * allocation)
|
|
max_chars = max(max_chars, self.TOOL_OUTPUT_MIN_CHARS)
|
|
|
|
msg = modified_messages[msg_idx]
|
|
content = msg.get("content", "")
|
|
if isinstance(content, str) and len(content) > max_chars:
|
|
truncated = self._truncate_tool_output(content, max_chars)
|
|
modified_messages[msg_idx] = {**msg, "content": truncated}
|
|
|
|
return modified_messages
|
|
|
|
def test_empty_messages_unchanged(self):
|
|
"""Empty message list should return unchanged."""
|
|
result = self._apply_weighting([], 200000)
|
|
assert result == []
|
|
|
|
def test_small_messages_unchanged(self):
|
|
"""Messages below threshold should pass through unchanged."""
|
|
messages = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there"},
|
|
]
|
|
result = self._apply_weighting(messages, 200000)
|
|
assert result == messages
|
|
|
|
def test_large_tool_outputs_truncated(self):
|
|
"""Tool outputs should be truncated when over threshold."""
|
|
large_output = "X" * 300000 # ~75k tokens
|
|
messages = [
|
|
{"role": "user", "content": "Test"},
|
|
{"role": "tool", "tool_call_id": "1", "content": large_output},
|
|
]
|
|
|
|
result = self._apply_weighting(messages, 200000)
|
|
assert len(result[1]["content"]) < len(large_output)
|
|
|
|
def test_older_outputs_truncated_more(self):
|
|
"""Older tool outputs should be truncated more aggressively."""
|
|
large_output = "X" * 300000
|
|
|
|
messages = [
|
|
{"role": "tool", "tool_call_id": "1", "content": large_output}, # Oldest
|
|
{"role": "tool", "tool_call_id": "2", "content": large_output}, # Middle
|
|
{"role": "tool", "tool_call_id": "3", "content": large_output}, # Newest
|
|
]
|
|
|
|
result = self._apply_weighting(messages, 200000)
|
|
|
|
oldest_len = len(result[0]["content"])
|
|
middle_len = len(result[1]["content"])
|
|
newest_len = len(result[2]["content"])
|
|
|
|
assert oldest_len <= middle_len <= newest_len
|
|
|
|
def test_preserves_original(self):
|
|
"""Function should not modify the original messages."""
|
|
large_output = "X" * 300000
|
|
messages = [{"role": "tool", "tool_call_id": "1", "content": large_output}]
|
|
original_len = len(messages[0]["content"])
|
|
|
|
result = self._apply_weighting(messages, 200000)
|
|
|
|
assert len(messages[0]["content"]) == original_len
|
|
assert result is not messages
|
|
|
|
def test_token_reduction(self):
|
|
"""Weighting should reduce overall token count."""
|
|
large_output = "X" * 300000
|
|
messages = [
|
|
{"role": "user", "content": "Test"},
|
|
{"role": "tool", "tool_call_id": "1", "content": large_output},
|
|
]
|
|
|
|
original_tokens = self._count_tokens(messages)
|
|
result = self._apply_weighting(messages, 200000)
|
|
new_tokens = self._count_tokens(result)
|
|
|
|
assert new_tokens < original_tokens
|
|
|
|
|
|
class TestGenericLinuxCommandCompression:
|
|
"""Tests for the output compression in generic_linux_command."""
|
|
|
|
MAX_OUTPUT_CHARS = 50000
|
|
|
|
def _is_binary_content(self, data: str) -> bool:
|
|
"""Local implementation for testing."""
|
|
if not data:
|
|
return False
|
|
sample = data[:8192]
|
|
non_printable = sum(1 for c in sample if ord(c) < 32 and c not in "\n\r\t")
|
|
if "\x00" in sample:
|
|
return True
|
|
ratio = non_printable / len(sample) if sample else 0
|
|
return ratio > 0.10
|
|
|
|
def test_binary_detection_text(self):
|
|
"""Text content should not be detected as binary."""
|
|
assert self._is_binary_content("Hello World\nThis is text\n") is False
|
|
|
|
def test_binary_detection_null_bytes(self):
|
|
"""Content with null bytes should be detected as binary."""
|
|
assert self._is_binary_content("data\x00\x01\x02more\x00") is True
|
|
|
|
def test_binary_detection_high_non_printable(self):
|
|
"""High ratio of non-printable chars should be detected as binary."""
|
|
binary_like = "".join(chr(i) for i in range(32)) * 100
|
|
assert self._is_binary_content(binary_like) is True
|
|
|
|
|
|
class TestModuleImport:
|
|
"""Test that the actual module functions can be imported."""
|
|
|
|
def test_import_truncate_function(self):
|
|
"""Should be able to import _truncate_tool_output."""
|
|
from cai.sdk.agents.models.openai_chatcompletions import _truncate_tool_output
|
|
|
|
result = _truncate_tool_output("small", 1000)
|
|
assert result == "small"
|
|
|
|
def test_import_weighting_function(self):
|
|
"""Should be able to import _apply_dynamic_tool_output_weighting."""
|
|
from cai.sdk.agents.models.openai_chatcompletions import (
|
|
_apply_dynamic_tool_output_weighting,
|
|
)
|
|
|
|
result = _apply_dynamic_tool_output_weighting([], "gpt-4", 200000)
|
|
assert result == []
|
|
|
|
def test_import_constants(self):
|
|
"""Should be able to import threshold constants."""
|
|
from cai.sdk.agents.models.openai_chatcompletions import (
|
|
TOOL_OUTPUT_WEIGHT_TOKEN_THRESHOLD,
|
|
TOOL_OUTPUT_WEIGHT_CONTEXT_THRESHOLD,
|
|
TOOL_OUTPUT_MAX_CHARS,
|
|
TOOL_OUTPUT_MIN_CHARS,
|
|
)
|
|
|
|
assert TOOL_OUTPUT_WEIGHT_TOKEN_THRESHOLD > 0
|
|
assert 0 < TOOL_OUTPUT_WEIGHT_CONTEXT_THRESHOLD <= 1.0
|
|
assert TOOL_OUTPUT_MAX_CHARS > TOOL_OUTPUT_MIN_CHARS
|
|
|
|
def test_import_generic_linux_helpers(self):
|
|
"""Should be able to import generic_linux_command helpers."""
|
|
from cai.tools.reconnaissance.generic_linux_command import (
|
|
_is_binary_content,
|
|
_compress_output_for_model,
|
|
MAX_OUTPUT_CHARS,
|
|
)
|
|
|
|
assert MAX_OUTPUT_CHARS > 0
|
|
assert _is_binary_content("text") is False
|