Fix for parallel tool calls summarization #126

src/langmem/short_term/summarization.py

@@ -198,20 +198,23 @@ def _preprocess_messages(
     else:
         messages_to_summarize = messages[total_summarized_messages : idx + 1]
 
-    # If the last message is an AI message with tool calls,
-    # include subsequent corresponding tool messages in the summary as well,
-    # to avoid issues w/ the LLM provider
-    if (
-        messages_to_summarize
-        and isinstance(messages_to_summarize[-1], AIMessage)
-        and (tool_calls := messages_to_summarize[-1].tool_calls)
-    ):
-        # Add any matching tool messages from our dictionary
-        for tool_call in tool_calls:
-            if tool_call["id"] in tool_call_id_to_tool_message:
-                tool_message = tool_call_id_to_tool_message[tool_call["id"]]
-                n_tokens_to_summarize += token_counter([tool_message])
-                messages_to_summarize.append(tool_message)
+    # For all AI message with tool calls, include subsequent corresponding
+    # tool messages in the summary as well, to avoid issues w/ the LLM provider
+    tool_call_ids = []
+    seen_tool_call_ids = set()
+    for msg in messages_to_summarize:
+        if isinstance(msg, AIMessage) and msg.tool_calls:
+            for tool_call in msg.tool_calls:
+                tool_call_ids.append(tool_call["id"])
+        elif isinstance(msg, ToolMessage) and msg.tool_call_id:
+            seen_tool_call_ids.add(msg.tool_call_id)
+
+    for tool_call_id in tool_call_ids:
+        if tool_call_id not in seen_tool_call_ids:
+            # Add matching Tool Messages that are not seen from our dictionary
+            tool_message = tool_call_id_to_tool_message[tool_call_id]
+            n_tokens_to_summarize += token_counter([tool_message])
+            messages_to_summarize.append(tool_message)
 
     return PreprocessedMessages(
         messages_to_summarize=messages_to_summarize,

tests/short_term/test_summarization.py

@@ -3,6 +3,7 @@
     AIMessage,
     HumanMessage,
     SystemMessage,
+    ToolCall,
     ToolMessage,
 )
 from langchain_core.messages.utils import count_tokens_approximately
@@ -574,6 +575,46 @@ def test_last_ai_with_tool_calls():
     )
 
 
+def test_multiple_parallel_tool_calls():
+    model = FakeChatModel(responses=[AIMessage(content="Summary for tool call messages.")])
+    model.invoke_calls = []
+
+    messages = [
+        HumanMessage(content="Generate two long random strings", id="id1"),
+        AIMessage(
+            content="Utilizing `rand_string` tool for this",
+            id="id2",
+            tool_calls=[ToolCall(name="rand_string", args={}, id=str(id)) for id in [1, 2]]
+        ),
+        ToolMessage(content="a", tool_call_id="1", id="id3"),
+        ToolMessage(content="b", tool_call_id="2", id="id4"),
+        AIMessage(content="Generated the two random strings.", id="id5"),
+        HumanMessage(content="Can you explain the generation algorithm?", id="id6"),
+    ]
+
+    summarization_node = SummarizationNode(
+        model=model,
+        token_counter=len,
+        max_tokens=4,
+        max_tokens_before_summary=3,
+        max_summary_tokens=1,
+        output_messages_key="summarized_messages",
+    )
+
+    result = summarization_node.invoke({"messages": messages})
+
+    # ensure the summarization model saw both tool messages tied to the tool calls
+    assert len(model.invoke_calls) == 1
+
+    # Ensure that all tool calls and their corresponding tool messages were summarized
+    running_summary = result["context"]["running_summary"]
+    assert running_summary is not None
+    assert running_summary.last_summarized_message_id == "id4"
+    assert running_summary.summarized_message_ids.issuperset(
+        {"id1", "id2", "id3", "id4"}
+    )
+
+
 def test_missing_message_ids():
     messages = [
         HumanMessage(content="Message 1", id="1"),