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

refactor(anthropic_message_utils.py): clean up whitespace and improve code readability by removing unnecessary blank lines and aligning code formatting 

fix(anthropic_message_utils.py): add warning in docstring for anthropic_trim_messages function to indicate incomplete implementation and clarify behavior
fix(anthropic_message_utils.py): ensure consistent formatting in conditional statements and improve readability of logical checks
This commit is contained in:
Ariel Frischer 2025-03-11 23:38:31 -07:00
parent a3284c9d7e
commit 376d486db8
1 changed files with 85 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
ra_aid/anthropic_message_utils.py
View File

@ -57,7 +57,9 @@ def has_tool_use(message: BaseMessage) -> bool:
return True return True
# Check additional_kwargs for tool_calls # Check additional_kwargs for tool_calls
if hasattr(message, "additional_kwargs") and message.additional_kwargs.get("tool_calls"): if hasattr(message, "additional_kwargs") and message.additional_kwargs.get(
"tool_calls"
):
return True return True
return False return False
@ -74,9 +76,9 @@ def is_tool_pair(message1: BaseMessage, message2: BaseMessage) -> bool:
bool: True if the messages form a tool use/result pair bool: True if the messages form a tool use/result pair
""" """
return ( return (
isinstance(message1, AIMessage) and isinstance(message1, AIMessage)
isinstance(message2, ToolMessage) and and isinstance(message2, ToolMessage)
has_tool_use(message1) and has_tool_use(message1)
) )
@ -129,6 +131,8 @@ def anthropic_trim_messages(
) -> List[BaseMessage]: ) -> List[BaseMessage]:
"""Trim messages to fit within a token limit, with Anthropic-specific handling. """Trim messages to fit within a token limit, with Anthropic-specific handling.
Warning - not fully implemented - last strategy is supported and test, not
allow partial, not 'first' strategy either.
This function is similar to langchain_core's trim_messages but with special This function is similar to langchain_core's trim_messages but with special
handling for Anthropic message formats to avoid API errors. handling for Anthropic message formats to avoid API errors.
@ -201,8 +205,8 @@ def anthropic_trim_messages(
pairs = [] pairs = []
i = 0 i = 0
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})") print(f"DEBUG - Found tool_use pair: ({i}, {i+1})")
i += 2 i += 2
else: else:
@ -220,14 +224,20 @@ 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") 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 = []
# Check if the last message in kept_messages has tool_use # Check if the last message in kept_messages has tool_use
if kept_messages and isinstance(kept_messages[-1], AIMessage) and has_tool_use(kept_messages[-1]): if (
kept_messages
and isinstance(kept_messages[-1], AIMessage)
and has_tool_use(kept_messages[-1])
):
# We need to find the corresponding ToolMessage # We need to find the corresponding ToolMessage
for i, (ai_idx, tool_idx) in enumerate(pairs): for i, (ai_idx, tool_idx) in enumerate(pairs):
if messages[ai_idx] is kept_messages[-1]: if messages[ai_idx] is kept_messages[-1]:
@ -236,7 +246,7 @@ def anthropic_trim_messages(
# Add the AIMessage and ToolMessage as a pair # Add the AIMessage and ToolMessage as a pair
result.extend([messages[ai_idx], messages[tool_idx]]) result.extend([messages[ai_idx], messages[tool_idx]])
# Remove this pair from the list of pairs to process later # Remove this pair from the list of pairs to process later
pairs = pairs[:i] + pairs[i+1:] pairs = pairs[:i] + pairs[i + 1 :]
break break
else: else:
# If we didn't find a matching pair, just add all kept_messages # If we didn't find a matching pair, just add all kept_messages
@ -268,7 +278,9 @@ def anthropic_trim_messages(
print(f"DEBUG - Added complete pair ({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") 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
@ -318,7 +330,9 @@ def anthropic_trim_messages(
print(f"DEBUG - Added segment {len(segments)-i-1} to 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") 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
@ -339,16 +353,24 @@ def anthropic_trim_messages(
current_msg = final_result[i] current_msg = final_result[i]
# If this is an AIMessage with tool_use, it must be followed by a ToolMessage # If this is an AIMessage with tool_use, it must be followed by a ToolMessage
if i < len(final_result) - 1 and isinstance(current_msg, AIMessage) and has_tool_use(current_msg): if (
if isinstance(final_result[i+1], ToolMessage): i < len(final_result) - 1
and isinstance(current_msg, AIMessage)
and has_tool_use(current_msg)
):
if isinstance(final_result[i + 1], ToolMessage):
# 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}") 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") 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:
@ -358,8 +380,14 @@ def anthropic_trim_messages(
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
if valid_result and isinstance(valid_result[-1], AIMessage) and has_tool_use(valid_result[-1]): if (
print("WARNING: Last message is AIMessage with tool_use but no following ToolMessage") valid_result
and isinstance(valid_result[-1], AIMessage)
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:") print("\nDEBUG - Final validated result:")