beeai-framework/python/tests/memory/test_utils.py at 48d44adabd0c12f68905c5d05918c15e7cde3156 · i-am-bee/beeai-framework · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
# 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)


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)