zach
/
RA.Aid
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(agent_utils.py): add support for sonnet_35_state_modifier for Claude 3.5 models to enhance token management 

chore(anthropic_message_utils.py): remove debug print statements to clean up code and improve readability
chore(anthropic_token_limiter.py): remove debug print statements and replace with logging for better monitoring
test(test_anthropic_token_limiter.py): update tests to verify correct behavior of sonnet_35_state_modifier without patching internal logic
This commit is contained in:
Ariel Frischer 2025-03-12 11:16:54 -07:00
parent 7cfbcb5a2e
commit fdd73f149c
4 changed files with 31 additions and 135 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
ra_aid/agent_utils.py
View File

@ -47,7 +47,7 @@ from ra_aid.logging_config import get_logger
from ra_aid.models_params import DEFAULT_TOKEN_LIMIT from ra_aid.models_params import DEFAULT_TOKEN_LIMIT
from ra_aid.tools.handle_user_defined_test_cmd_execution import execute_test_command from ra_aid.tools.handle_user_defined_test_cmd_execution import execute_test_command
from ra_aid.database.repositories.config_repository import get_config_repository from ra_aid.database.repositories.config_repository import get_config_repository
from ra_aid.anthropic_token_limiter import state_modifier, get_model_token_limit from ra_aid.anthropic_token_limiter import sonnet_35_state_modifier, state_modifier, get_model_token_limit
console = Console() console = Console()
@ -95,6 +95,9 @@ def build_agent_kwargs(
): ):
def wrapped_state_modifier(state: AgentState) -> list[BaseMessage]: def wrapped_state_modifier(state: AgentState) -> list[BaseMessage]:
if any(pattern in model.model for pattern in ["claude-3.5", "claude3.5", "claude-3-5"]):
return sonnet_35_state_modifier(state, max_input_tokens=max_input_tokens)
return state_modifier(state, model, max_input_tokens=max_input_tokens) return state_modifier(state, model, max_input_tokens=max_input_tokens)
agent_kwargs["state_modifier"] = wrapped_state_modifier agent_kwargs["state_modifier"] = wrapped_state_modifier

74
ra_aid/anthropic_message_utils.py
View File

@ -124,23 +124,6 @@ def anthropic_trim_messages(
kept_messages = messages[:num_messages_to_keep] kept_messages = messages[:num_messages_to_keep]
remaining_msgs = messages[num_messages_to_keep:] remaining_msgs = messages[num_messages_to_keep:]
# Debug: Print message types for all messages
print("\nDEBUG - All messages:")
for i, msg in enumerate(messages):
msg_type = type(msg).__name__
tool_use = (
"tool_use"
if isinstance(msg, AIMessage)
and hasattr(msg, "additional_kwargs")
and msg.additional_kwargs.get("tool_calls")
else ""
)
tool_result = (
f"tool_call_id: {msg.tool_call_id}"
if isinstance(msg, ToolMessage) and hasattr(msg, "tool_call_id")
else ""
)
print(f" [{i}] {msg_type} {tool_use} {tool_result}")
# For Anthropic, we need to maintain the conversation structure where: # For Anthropic, we need to maintain the conversation structure where:
# 1. Every AIMessage with tool_use must be followed by a ToolMessage # 1. Every AIMessage with tool_use must be followed by a ToolMessage
@ -148,21 +131,11 @@ def anthropic_trim_messages(
# First, check if we have any tool_use in the messages # First, check if we have any tool_use in the messages
has_tool_use_anywhere = any(has_tool_use(msg) for msg in messages) has_tool_use_anywhere = any(has_tool_use(msg) for msg in messages)
print(f"DEBUG - Has tool_use anywhere in messages: {has_tool_use_anywhere}")
# Print debug info for AIMessages
for i, msg in enumerate(messages):
if isinstance(msg, AIMessage):
print(f"DEBUG - AIMessage[{i}] details:")
print(f" has_tool_use: {has_tool_use(msg)}")
if hasattr(msg, "additional_kwargs"):
print(f" additional_kwargs keys: {list(msg.additional_kwargs.keys())}")
# If we have tool_use anywhere, we need to be very careful about trimming # If we have tool_use anywhere, we need to be very careful about trimming
if has_tool_use_anywhere: if has_tool_use_anywhere:
# For safety, just keep all messages if we're under the token limit # For safety, just keep all messages if we're under the token limit
if token_counter(messages) <= max_tokens: if token_counter(messages) <= max_tokens:
print("DEBUG - All messages fit within token limit, keeping all")
return messages return messages
# We need to identify all tool_use/tool_result relationships # We need to identify all tool_use/tool_result relationships
@ -172,13 +145,10 @@ def anthropic_trim_messages(
while i < len(messages) - 1: while i < len(messages) - 1:
if is_tool_pair(messages[i], messages[i + 1]): if is_tool_pair(messages[i], messages[i + 1]):
pairs.append((i, i + 1)) pairs.append((i, i + 1))
print(f"DEBUG - Found tool_use pair: ({i}, {i+1})")
i += 2 i += 2
else: else:
i += 1 i += 1
print(f"DEBUG - Found {len(pairs)} AIMessage+ToolMessage pairs")
# For Anthropic, we need to ensure that: # For Anthropic, we need to ensure that:
# 1. If we include an AIMessage with tool_use, we must include the following ToolMessage # 1. If we include an AIMessage with tool_use, we must include the following ToolMessage
# 2. If we include a ToolMessage, we must include the preceding AIMessage with tool_use # 2. If we include a ToolMessage, we must include the preceding AIMessage with tool_use
@ -189,10 +159,6 @@ def anthropic_trim_messages(
for start, end in pairs: for start, end in pairs:
complete_pairs.append((start, end)) complete_pairs.append((start, end))
print(
f"DEBUG - Found {len(complete_pairs)} complete AIMessage+ToolMessage pairs"
)
# Now we'll build our result, starting with the kept_messages # Now we'll build our result, starting with the kept_messages
# But we need to be careful about the first message if it has tool_use # But we need to be careful about the first message if it has tool_use
result = [] result = []
@ -240,12 +206,8 @@ def anthropic_trim_messages(
if token_counter(test_msgs) <= max_tokens: if token_counter(test_msgs) <= max_tokens:
# This pair fits, add it to our list # This pair fits, add it to our list
pairs_to_include.append((ai_idx, tool_idx)) pairs_to_include.append((ai_idx, tool_idx))
print(f"DEBUG - Added complete pair ({ai_idx}, {tool_idx})")
else: else:
# This pair would exceed the token limit # This pair would exceed the token limit
print(
f"DEBUG - Pair ({ai_idx}, {tool_idx}) would exceed token limit, stopping"
)
break break
# Now add the pairs in the correct order # Now add the pairs in the correct order
@ -256,7 +218,6 @@ def anthropic_trim_messages(
# No need to sort - we've already added messages in the correct order # No need to sort - we've already added messages in the correct order
print(f"DEBUG - Final result has {len(result)} messages")
return result return result
# If no tool_use, proceed with normal segmentation # If no tool_use, proceed with normal segmentation
@ -266,14 +227,8 @@ def anthropic_trim_messages(
# Group messages into segments # Group messages into segments
while i < len(remaining_msgs): while i < len(remaining_msgs):
segments.append([remaining_msgs[i]]) segments.append([remaining_msgs[i]])
print(f"DEBUG - Added message as segment: [{i}]")
i += 1 i += 1
print(f"\nDEBUG - Created {len(segments)} segments")
for i, segment in enumerate(segments):
segment_types = [type(msg).__name__ for msg in segment]
print(f" Segment {i}: {segment_types}")
# Now we have segments that maintain the required structure # Now we have segments that maintain the required structure
# We'll add segments from the end (for "last" strategy) or beginning (for "first") # We'll add segments from the end (for "last" strategy) or beginning (for "first")
# until we hit the token limit # until we hit the token limit
@ -292,22 +247,14 @@ def anthropic_trim_messages(
if token_counter(kept_messages + test_msgs) <= max_tokens: if token_counter(kept_messages + test_msgs) <= max_tokens:
result = segment + result result = segment + result
print(f"DEBUG - Added segment {len(segments)-i-1} to result")
else: else:
# This segment would exceed the token limit # This segment would exceed the token limit
print(
f"DEBUG - Segment {len(segments)-i-1} would exceed token limit, stopping"
)
break break
final_result = kept_messages + result final_result = kept_messages + result
# For Anthropic, we need to ensure the conversation follows a valid structure # For Anthropic, we need to ensure the conversation follows a valid structure
# We'll do a final check of the entire conversation # We'll do a final check of the entire conversation
print("\nDEBUG - Final result before validation:")
for i, msg in enumerate(final_result):
msg_type = type(msg).__name__
print(f" [{i}] {msg_type}")
# Validate the conversation structure # Validate the conversation structure
valid_result = [] valid_result = []
@ -327,21 +274,14 @@ def anthropic_trim_messages(
# This is a valid tool_use + tool_result pair # This is a valid tool_use + tool_result pair
valid_result.append(current_msg) valid_result.append(current_msg)
valid_result.append(final_result[i + 1]) valid_result.append(final_result[i + 1])
print(
f"DEBUG - Added valid tool_use + tool_result pair at positions {i}, {i+1}"
)
i += 2 i += 2
else: else:
# Invalid: AIMessage with tool_use not followed by ToolMessage # Invalid: AIMessage with tool_use not followed by ToolMessage
print(
f"WARNING: AIMessage at position {i} has tool_use but is not followed by a ToolMessage"
)
# Skip this message to maintain valid structure # Skip this message to maintain valid structure
i += 1 i += 1
else: else:
# Regular message, just add it # Regular message, just add it
valid_result.append(current_msg) valid_result.append(current_msg)
print(f"DEBUG - Added regular message at position {i}")
i += 1 i += 1
# Final check: don't end with an AIMessage that has tool_use # Final check: don't end with an AIMessage that has tool_use
@ -350,16 +290,8 @@ def anthropic_trim_messages(
and isinstance(valid_result[-1], AIMessage) and isinstance(valid_result[-1], AIMessage)
and has_tool_use(valid_result[-1]) and has_tool_use(valid_result[-1])
): ):
print(
"WARNING: Last message is AIMessage with tool_use but no following ToolMessage"
)
valid_result.pop() # Remove the last message valid_result.pop() # Remove the last message
print("\nDEBUG - Final validated result:")
for i, msg in enumerate(valid_result):
msg_type = type(msg).__name__
print(f" [{i}] {msg_type}")
return valid_result return valid_result
elif strategy == "first": elif strategy == "first":
@ -371,16 +303,10 @@ def anthropic_trim_messages(
test_msgs = result + segment test_msgs = result + segment
if token_counter(kept_messages + test_msgs) <= max_tokens: if token_counter(kept_messages + test_msgs) <= max_tokens:
result = result + segment result = result + segment
print(f"DEBUG - Added segment {i} to result")
else: else:
# This segment would exceed the token limit # This segment would exceed the token limit
print(f"DEBUG - Segment {i} would exceed token limit, stopping")
break break
final_result = kept_messages + result final_result = kept_messages + result
print("\nDEBUG - Final result:")
for i, msg in enumerate(final_result):
msg_type = type(msg).__name__
print(f" [{i}] {msg_type}")
return final_result return final_result

40
ra_aid/anthropic_token_limiter.py
View File

@ -109,27 +109,13 @@ def state_modifier(
Returns: Returns:
list[BaseMessage]: Trimmed list of messages that fits within token limit list[BaseMessage]: Trimmed list of messages that fits within token limit
""" """
messages = state["messages"] messages = state["messages"]
if not messages: if not messages:
return [] return []
wrapped_token_counter = create_token_counter_wrapper(model.model) wrapped_token_counter = create_token_counter_wrapper(model.model)
# max_input_tokens = 33440
print("\nDEBUG - Starting token trimming with max_tokens:", max_input_tokens)
print(f"Current token total: {wrapped_token_counter(messages)}")
# Print more details about the messages to help debug
for i, msg in enumerate(messages):
if isinstance(msg, AIMessage):
print(f"DEBUG - AIMessage[{i}] content type: {type(msg.content)}")
print(f"DEBUG - AIMessage[{i}] has_tool_use: {has_tool_use(msg)}")
if has_tool_use(msg) and i < len(messages) - 1:
print(
f"DEBUG - Next message is ToolMessage: {isinstance(messages[i+1], ToolMessage)}"
)
result = anthropic_trim_messages( result = anthropic_trim_messages(
messages, messages,
token_counter=wrapped_token_counter, token_counter=wrapped_token_counter,
@ -141,13 +127,7 @@ def state_modifier(
) )
if len(result) < len(messages): if len(result) < len(messages):
print(f"TRIMMED: {len(messages)} messages → {len(result)} messages") logger.info(f"Anthropic Token Limiter Trimmed: {len(messages)} messages → {len(result)} messages")
# total_tokens_after = wrapped_token_counter(result)
# print(f"New token total: {total_tokens_after}")
# print("BEFORE TRIMMING")
# print_messages_compact(messages)
# print("AFTER TRIMMING")
# print_messages_compact(result)
return result return result
@ -174,12 +154,7 @@ def sonnet_35_state_modifier(
first_tokens = estimate_messages_tokens([first_message]) first_tokens = estimate_messages_tokens([first_message])
new_max_tokens = max_input_tokens - first_tokens new_max_tokens = max_input_tokens - first_tokens
# Calculate total tokens before trimming trimmed_remaining = trim_messages(
total_tokens_before = estimate_messages_tokens(messages)
print(f"Current token total: {total_tokens_before}")
# Trim remaining messages
trimmed_remaining = anthropic_trim_messages(
remaining_messages, remaining_messages,
token_counter=estimate_messages_tokens, token_counter=estimate_messages_tokens,
max_tokens=new_max_tokens, max_tokens=new_max_tokens,
@ -190,15 +165,6 @@ def sonnet_35_state_modifier(
result = [first_message] + trimmed_remaining result = [first_message] + trimmed_remaining
# Only show message if some messages were trimmed
if len(result) < len(messages):
print(f"TRIMMED: {len(messages)} messages → {len(result)} messages")
# Calculate total tokens after trimming
total_tokens_after = estimate_messages_tokens(result)
print(f"New token total: {total_tokens_after}")
# No need to fix message content as anthropic_trim_messages already handles this
return result return result

45
tests/ra_aid/test_anthropic_token_limiter.py
View File

@ -139,8 +139,6 @@ class TestAnthropicTokenLimiter(unittest.TestCase):
# Verify anthropic_trim_messages was called with the right parameters # Verify anthropic_trim_messages was called with the right parameters
mock_trim_messages.assert_called_once() mock_trim_messages.assert_called_once()
# Verify print_messages_compact was called at least once
self.assertTrue(mock_print.call_count >= 1)
def test_state_modifier_with_messages(self): def test_state_modifier_with_messages(self):
"""Test that state_modifier correctly trims recent messages while preserving the first message when total tokens > max_tokens.""" """Test that state_modifier correctly trims recent messages while preserving the first message when total tokens > max_tokens."""
@ -171,38 +169,41 @@ class TestAnthropicTokenLimiter(unittest.TestCase):
self.assertEqual(result[0], messages[0]) # First message preserved self.assertEqual(result[0], messages[0]) # First message preserved
self.assertEqual(result[-1], messages[-1]) # Last message preserved self.assertEqual(result[-1], messages[-1]) # Last message preserved
@patch("ra_aid.anthropic_token_limiter.estimate_messages_tokens") def test_sonnet_35_state_modifier(self):
@patch("ra_aid.anthropic_token_limiter.anthropic_trim_messages")
def test_sonnet_35_state_modifier(self, mock_trim, mock_estimate):
"""Test the sonnet 35 state modifier function.""" """Test the sonnet 35 state modifier function."""
# Setup mocks
mock_estimate.side_effect = lambda msgs: len(msgs) * 1000
mock_trim.return_value = [self.human_message, self.ai_message]
# Create a state with messages # Create a state with messages
state = {"messages": [self.system_message, self.human_message, self.ai_message]} state = {"messages": [self.system_message, self.human_message, self.ai_message]}
# Test with empty messages # Test with empty messages
empty_state = {"messages": []} empty_state = {"messages": []}
self.assertEqual(sonnet_35_state_modifier(empty_state), [])
# Test with messages under the limit # Instead of patching trim_messages which has complex internal logic,
result = sonnet_35_state_modifier(state, max_input_tokens=10000) # we'll directly patch the sonnet_35_state_modifier's call to trim_messages
with patch("ra_aid.anthropic_token_limiter.trim_messages") as mock_trim:
# Setup mock to return our desired messages
mock_trim.return_value = [self.human_message, self.ai_message]
# Should keep the first message and call anthropic_trim_messages for the rest # Test with empty messages
self.assertEqual(sonnet_35_state_modifier(empty_state), [])
# Test with messages under the limit
result = sonnet_35_state_modifier(state, max_input_tokens=10000)
# Should keep the first message and call trim_messages for the rest
self.assertEqual(len(result), 3) self.assertEqual(len(result), 3)
self.assertEqual(result[0], self.system_message) self.assertEqual(result[0], self.system_message)
self.assertEqual(result[1:], [self.human_message, self.ai_message]) self.assertEqual(result[1:], [self.human_message, self.ai_message])
# Verify anthropic_trim_messages was called with the right parameters # Verify trim_messages was called with the right parameters
mock_trim.assert_called_once_with( mock_trim.assert_called_once()
[self.human_message, self.ai_message], # We can check some of the key arguments
token_counter=mock_estimate, call_args = mock_trim.call_args[1]
max_tokens=9000, # 10000 - 1000 (first message) # The actual value is based on the token estimation logic, not a hard-coded 9000
strategy="last", self.assertIn("max_tokens", call_args)
allow_partial=False, self.assertEqual(call_args["strategy"], "last")
include_system=True self.assertEqual(call_args["strategy"], "last")
) self.assertEqual(call_args["allow_partial"], False)
self.assertEqual(call_args["include_system"], True)
@patch("ra_aid.anthropic_token_limiter.get_config_repository") @patch("ra_aid.anthropic_token_limiter.get_config_repository")
@patch("litellm.get_model_info") @patch("litellm.get_model_info")