fix(langgraph-checkpoint-aws): clean orphan tool_calls in AgentCoreMemorySaver to prevent Bedrock ValidationException #725
Conversation
…ck ValidationException
murilosimao
changed the title
|
Hey @michaelnchin, could you please take a look at issue #726? I encountered the mentioned error when trying to use the new LangChain version. |
There was a problem hiding this comment.
@murilosimao thanks for submitting both the issue and fix! LGTM in general, some feedback below.
| resolved_tool_call_ids = set() | ||
| for msg in messages: | ||
| if isinstance(msg, ToolMessage) and hasattr(msg, "tool_call_id"): | ||
| resolved_tool_call_ids.add(msg.tool_call_id) |
There was a problem hiding this comment.
Nit: For very large message sets, it would be faster to create resolved_tool_calls directly from set comprehension:
| resolved_tool_call_ids = set() | |
| for msg in messages: | |
| if isinstance(msg, ToolMessage) and hasattr(msg, "tool_call_id"): | |
| resolved_tool_call_ids.add(msg.tool_call_id) | |
| resolved_tool_call_ids = { | |
| msg.tool_call_id | |
| for msg in messages | |
| if isinstance(msg, ToolMessage) and hasattr(msg, "tool_call_id") | |
| } |
| if len(valid_tool_calls) != len(msg.tool_calls): | ||
| cleaned_msg = msg.model_copy(update={"tool_calls": valid_tool_calls}) | ||
| cleaned_messages.append(cleaned_msg) |
There was a problem hiding this comment.
Given that we're dropping possibly useful context from the message history, we should add a logger warning here to ensure that the tool call removals aren't completely silent.
| # Clean up AIMessages with orphaned tool_calls | ||
| cleaned_messages = [] | ||
| for msg in messages: | ||
| if isinstance(msg, AIMessage) and hasattr(msg, "tool_calls") and msg.tool_calls: |
There was a problem hiding this comment.
Nit: AIMessage always has tool_calls attribute:
| if isinstance(msg, AIMessage) and hasattr(msg, "tool_calls") and msg.tool_calls: | |
| if isinstance(msg, AIMessage) and msg.tool_calls: |
| ) | ||
|
|
||
|
|
||
| def clean_orphan_tool_calls(messages: list[Any]) -> list[Any]: |
There was a problem hiding this comment.
Removing the tool calls is OK as an unblocker, but I'm slightly concerned about the conversational history loss prior to the checkpoint.
Maybe not in this PR, but we should consider to something less destructive. One possible solution would be to add dummy result ToolMessages:
...
new_messages = []
for msg in messages:
new_messages.append(msg)
if isinstance(msg, AIMessage) and msg.tool_calls:
# Gather IDs of orphan tool_calls
missing_tool_calls = [
tc for tc in msg.tool_calls
if tc.get("id") and tc["id"] not in resolved_ids
]
# Create filler ToolMessages each orphan ID
for tc in missing_tool_calls:
new_messages.append(
ToolMessage(
tool_call_id=tc["id"],
content="Tool call failed or was cancelled.",
)
)
return new_messages
2 participants
Fixed #726