From 3eb8f9b44d3410e91da70becce0bc73bf73ac955 Mon Sep 17 00:00:00 2001 From: pragnyanramtha Date: Sun, 17 May 2026 04:38:11 +0000 Subject: [PATCH 1/2] fix(python): ignore empty tool messages Signed-off-by: pragnyanramtha --- python/beeai_framework/memory/utils.py | 4 ++- python/tests/memory/test_utils.py | 38 ++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) create mode 100644 python/tests/memory/test_utils.py diff --git a/python/beeai_framework/memory/utils.py b/python/beeai_framework/memory/utils.py index 0dbe3988b..9d75b3126 100644 --- a/python/beeai_framework/memory/utils.py +++ b/python/beeai_framework/memory/utils.py @@ -18,11 +18,13 @@ def extract_last_tool_call_pair(memory: BaseMemory) -> tuple[AssistantMessage, T return None tool_call: AssistantMessage = memory.messages[tool_call_index] # type: ignore + tool_call_id = tool_call.get_tool_calls()[0].id tool_response_index = find_index( memory.messages, lambda msg: bool( - isinstance(msg, ToolMessage) and msg.get_tool_results()[0].tool_call_id == tool_call.get_tool_calls()[0].id + isinstance(msg, ToolMessage) + and any(result.tool_call_id == tool_call_id for result in msg.get_tool_results()) ), reverse_traversal=True, fallback=-1, diff --git a/python/tests/memory/test_utils.py b/python/tests/memory/test_utils.py new file mode 100644 index 000000000..76cefd7f9 --- /dev/null +++ b/python/tests/memory/test_utils.py @@ -0,0 +1,38 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +from beeai_framework.backend.message import ( + AssistantMessage, + MessageToolCallContent, + MessageToolResultContent, + ToolMessage, +) +from beeai_framework.memory.unconstrained_memory import UnconstrainedMemory +from beeai_framework.memory.utils import extract_last_tool_call_pair + + +def test_extract_last_tool_call_pair_returns_matching_tool_response() -> None: + memory = UnconstrainedMemory() + assistant_message = AssistantMessage( + MessageToolCallContent(id="call_1", tool_name="weather", args='{"city":"Paris"}') + ) + tool_message = ToolMessage(MessageToolResultContent(tool_name="weather", tool_call_id="call_1", result="sunny")) + + asyncio.run(memory.add_many([assistant_message, tool_message])) + + assert extract_last_tool_call_pair(memory) == (assistant_message, tool_message) + + +def test_extract_last_tool_call_pair_ignores_empty_tool_messages() -> None: + memory = UnconstrainedMemory() + assistant_message = AssistantMessage( + MessageToolCallContent(id="call_1", tool_name="weather", args='{"city":"Paris"}') + ) + empty_tool_message = ToolMessage([]) + tool_message = ToolMessage(MessageToolResultContent(tool_name="weather", tool_call_id="call_1", result="sunny")) + + asyncio.run(memory.add_many([assistant_message, empty_tool_message, tool_message])) + + assert extract_last_tool_call_pair(memory) == (assistant_message, tool_message) From 48d44adabd0c12f68905c5d05918c15e7cde3156 Mon Sep 17 00:00:00 2001 From: pragnyanramtha Date: Sun, 17 May 2026 05:05:56 +0000 Subject: [PATCH 2/2] fix(python): use last tool call pair Signed-off-by: pragnyanramtha --- python/beeai_framework/memory/utils.py | 5 ++--- python/tests/memory/test_utils.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/python/beeai_framework/memory/utils.py b/python/beeai_framework/memory/utils.py index 9d75b3126..0c893798d 100644 --- a/python/beeai_framework/memory/utils.py +++ b/python/beeai_framework/memory/utils.py @@ -18,13 +18,12 @@ def extract_last_tool_call_pair(memory: BaseMemory) -> tuple[AssistantMessage, T return None tool_call: AssistantMessage = memory.messages[tool_call_index] # type: ignore - tool_call_id = tool_call.get_tool_calls()[0].id + tool_call_id = tool_call.get_tool_calls()[-1].id tool_response_index = find_index( memory.messages, lambda msg: bool( - isinstance(msg, ToolMessage) - and any(result.tool_call_id == tool_call_id for result in msg.get_tool_results()) + isinstance(msg, ToolMessage) and any(result.tool_call_id == tool_call_id for result in msg.content) ), reverse_traversal=True, fallback=-1, diff --git a/python/tests/memory/test_utils.py b/python/tests/memory/test_utils.py index 76cefd7f9..482872808 100644 --- a/python/tests/memory/test_utils.py +++ b/python/tests/memory/test_utils.py @@ -36,3 +36,23 @@ def test_extract_last_tool_call_pair_ignores_empty_tool_messages() -> None: asyncio.run(memory.add_many([assistant_message, empty_tool_message, tool_message])) assert extract_last_tool_call_pair(memory) == (assistant_message, tool_message) + + +def test_extract_last_tool_call_pair_uses_last_tool_call_from_message() -> None: + memory = UnconstrainedMemory() + assistant_message = AssistantMessage( + [ + MessageToolCallContent(id="call_1", tool_name="weather", args='{"city":"Paris"}'), + MessageToolCallContent(id="call_2", tool_name="weather", args='{"city":"Berlin"}'), + ] + ) + first_tool_message = ToolMessage( + MessageToolResultContent(tool_name="weather", tool_call_id="call_1", result="sunny") + ) + last_tool_message = ToolMessage( + MessageToolResultContent(tool_name="weather", tool_call_id="call_2", result="rainy") + ) + + asyncio.run(memory.add_many([assistant_message, first_tool_message, last_tool_message])) + + assert extract_last_tool_call_pair(memory) == (assistant_message, last_tool_message)