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
parent
a3284c9d7e
commit
376d486db8
|
|
@ -57,7 +57,9 @@ def has_tool_use(message: BaseMessage) -> bool:
|
|||
return True
|
||||
|
||||
# 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 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
|
||||
"""
|
||||
return (
|
||||
isinstance(message1, AIMessage) and
|
||||
isinstance(message2, ToolMessage) and
|
||||
has_tool_use(message1)
|
||||
isinstance(message1, AIMessage)
|
||||
and isinstance(message2, ToolMessage)
|
||||
and has_tool_use(message1)
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -129,6 +131,8 @@ def anthropic_trim_messages(
|
|||
) -> List[BaseMessage]:
|
||||
"""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
|
||||
handling for Anthropic message formats to avoid API errors.
|
||||
|
||||
|
|
@ -220,14 +224,20 @@ def anthropic_trim_messages(
|
|||
for start, end in pairs:
|
||||
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
|
||||
# But we need to be careful about the first message if it has tool_use
|
||||
result = []
|
||||
|
||||
# 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
|
||||
for i, (ai_idx, tool_idx) in enumerate(pairs):
|
||||
if messages[ai_idx] is kept_messages[-1]:
|
||||
|
|
@ -268,7 +278,9 @@ def anthropic_trim_messages(
|
|||
print(f"DEBUG - Added complete pair ({ai_idx}, {tool_idx})")
|
||||
else:
|
||||
# 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
|
||||
|
||||
# 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")
|
||||
else:
|
||||
# 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
|
||||
|
||||
final_result = kept_messages + result
|
||||
|
|
@ -339,16 +353,24 @@ def anthropic_trim_messages(
|
|||
current_msg = final_result[i]
|
||||
|
||||
# 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 (
|
||||
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
|
||||
valid_result.append(current_msg)
|
||||
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
|
||||
else:
|
||||
# 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
|
||||
i += 1
|
||||
else:
|
||||
|
|
@ -358,8 +380,14 @@ def anthropic_trim_messages(
|
|||
i += 1
|
||||
|
||||
# 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]):
|
||||
print("WARNING: Last message is AIMessage with tool_use but no following ToolMessage")
|
||||
if (
|
||||
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
|
||||
|
||||
print("\nDEBUG - Final validated result:")
|
||||
|
|
|
|||
Loading…
Reference in New Issue