Spaces:
Paused
Paused
| # -*- coding: utf-8 -*- | |
| """The unittests for the tool result compression.""" | |
| # pylint: disable=protected-access, unused-argument | |
| from unittest.async_case import IsolatedAsyncioTestCase | |
| from utils import MockModel | |
| from agentscope.agent import Agent, ContextConfig | |
| from agentscope.message import ( | |
| ToolResultBlock, | |
| TextBlock, | |
| DataBlock, | |
| Base64Source, | |
| ) | |
| from agentscope.state import AgentState | |
| from agentscope.tool import Toolkit | |
| class ToolResultCompressionTest(IsolatedAsyncioTestCase): | |
| """Test cases for tool result compression.""" | |
| async def asyncSetUp(self) -> None: | |
| """The async setup method.""" | |
| self.mock_model = MockModel() | |
| self.agent = Agent( | |
| name="TestAgent", | |
| system_prompt="Test system prompt", | |
| model=self.mock_model, | |
| toolkit=Toolkit(), | |
| context_config=ContextConfig( | |
| tool_result_limit=100, | |
| ), | |
| state=AgentState(session_id="test_session"), | |
| ) | |
| async def test_below_limit(self) -> None: | |
| """Test when tool result is below the token limit.""" | |
| tool_result = ToolResultBlock( | |
| id="test_1", | |
| name="test_tool", | |
| output=[ | |
| TextBlock(text="Short text 1"), | |
| TextBlock(text="Short text 2"), | |
| ], | |
| ) | |
| async def mock_count_tokens( | |
| messages: list, | |
| tools: list | None = None, | |
| ) -> int: | |
| """Mock token counting function that returns a fixed count.""" | |
| return 50 | |
| self.mock_model.count_tokens = mock_count_tokens | |
| ( | |
| reserved, | |
| offload, | |
| ) = await self.agent._split_tool_result_for_compression( | |
| tool_result, | |
| ) | |
| self.assertEqual(reserved, tool_result) | |
| self.assertIsNone(offload) | |
| async def test_equal_to_limit(self) -> None: | |
| """Test when tool result is exactly at the token limit.""" | |
| tool_result = ToolResultBlock( | |
| id="test_2", | |
| name="test_tool", | |
| output=[TextBlock(text="Text at limit")], | |
| ) | |
| async def mock_count_tokens( | |
| messages: list, | |
| tools: list | None = None, | |
| ) -> int: | |
| """Mock token counting function that returns a fixed count.""" | |
| return 100 | |
| self.mock_model.count_tokens = mock_count_tokens | |
| ( | |
| reserved, | |
| offload, | |
| ) = await self.agent._split_tool_result_for_compression( | |
| tool_result, | |
| ) | |
| self.assertEqual(reserved, tool_result) | |
| self.assertIsNone(offload) | |
| async def test_boundary_last_block_text(self) -> None: | |
| """Test when boundary is the last block and it is a TextBlock.""" | |
| block1 = TextBlock(text="A" * 20, id="block1") | |
| block2 = TextBlock(text="B" * 20, id="block2") | |
| block3 = TextBlock(text="C" * 100, id="block3") | |
| tool_result = ToolResultBlock( | |
| id="test_3", | |
| name="test_tool", | |
| output=[block1, block2, block3], | |
| ) | |
| async def mock_count_tokens( | |
| messages: list, | |
| tools: list | None = None, | |
| ) -> int: | |
| """Mock token counting function that counts text length in | |
| blocks.""" | |
| content = messages[0].content | |
| if isinstance(content, list): | |
| total = sum(len(b.text) for b in content if hasattr(b, "text")) | |
| return total | |
| return 0 | |
| self.mock_model.count_tokens = mock_count_tokens | |
| ( | |
| reserved, | |
| offload, | |
| ) = await self.agent._split_tool_result_for_compression( | |
| tool_result, | |
| ) | |
| # Verify results | |
| self.assertIsNotNone(reserved) | |
| self.assertIsNotNone(offload) | |
| # Verify ToolResultBlock metadata | |
| self.assertEqual(reserved.id, tool_result.id) | |
| self.assertEqual(reserved.name, tool_result.name) | |
| self.assertEqual(reserved.state, tool_result.state) | |
| self.assertEqual(offload.id, tool_result.id) | |
| self.assertEqual(offload.name, tool_result.name) | |
| self.assertEqual(offload.state, tool_result.state) | |
| # Verify results using assertListEqual | |
| expected_reserved = [ | |
| {"type": "text", "text": "A" * 20, "id": "block1"}, | |
| {"type": "text", "text": "B" * 20 + "C" * 60, "id": "block2"}, | |
| ] | |
| expected_offload = [ | |
| {"type": "text", "text": "C" * 40, "id": "block3"}, | |
| ] | |
| self.assertListEqual( | |
| [b.model_dump() for b in reserved.output], | |
| expected_reserved, | |
| ) | |
| self.assertListEqual( | |
| [b.model_dump() for b in offload.output], | |
| expected_offload, | |
| ) | |
| async def test_boundary_last_block_data(self) -> None: | |
| """Test when boundary is the last block and it is a DataBlock.""" | |
| block1 = TextBlock(text="A" * 20, id="block1") | |
| block2 = TextBlock(text="B" * 20, id="block2") | |
| block3 = DataBlock( | |
| source=Base64Source(data="base64data", media_type="image/png"), | |
| id="block3", | |
| ) | |
| tool_result = ToolResultBlock( | |
| id="test_4", | |
| name="test_tool", | |
| output=[block1, block2, block3], | |
| ) | |
| async def mock_count_tokens( | |
| messages: list, | |
| tools: list | None = None, | |
| ) -> int: | |
| """Mock token counting function based on content length.""" | |
| content = messages[0].content | |
| if isinstance(content, list): | |
| if len(content) == 3: | |
| return 150 | |
| elif len(content) == 2: | |
| return 80 | |
| return 50 | |
| self.mock_model.count_tokens = mock_count_tokens | |
| ( | |
| reserved, | |
| offload, | |
| ) = await self.agent._split_tool_result_for_compression( | |
| tool_result, | |
| ) | |
| # Verify results | |
| self.assertIsNotNone(reserved) | |
| self.assertIsNotNone(offload) | |
| # Verify ToolResultBlock metadata | |
| self.assertEqual(reserved.id, tool_result.id) | |
| self.assertEqual(offload.id, tool_result.id) | |
| # Verify results using assertListEqual | |
| expected_reserved = [ | |
| {"type": "text", "text": "A" * 20, "id": "block1"}, | |
| {"type": "text", "text": "B" * 20, "id": "block2"}, | |
| ] | |
| expected_offload = [ | |
| { | |
| "type": "data", | |
| "id": "block3", | |
| "source": { | |
| "type": "base64", | |
| "data": "base64data", | |
| "media_type": "image/png", | |
| }, | |
| "name": None, | |
| }, | |
| ] | |
| self.assertListEqual( | |
| [b.model_dump() for b in reserved.output], | |
| expected_reserved, | |
| ) | |
| self.assertListEqual( | |
| [b.model_dump() for b in offload.output], | |
| expected_offload, | |
| ) | |
| async def test_boundary_first_block_text(self) -> None: | |
| """Test when boundary is the first block and it is a TextBlock.""" | |
| block1 = TextBlock(text="A" * 100, id="block1") | |
| block2 = TextBlock(text="B" * 20, id="block2") | |
| block3 = TextBlock(text="C" * 20, id="block3") | |
| tool_result = ToolResultBlock( | |
| id="test_5", | |
| name="test_tool", | |
| output=[block1, block2, block3], | |
| ) | |
| async def mock_count_tokens( | |
| messages: list, | |
| tools: list | None = None, | |
| ) -> int: | |
| """Mock token counting function that counts text length in | |
| blocks.""" | |
| content = messages[0].content | |
| if isinstance(content, list): | |
| total = sum(len(b.text) for b in content if hasattr(b, "text")) | |
| return total | |
| return 0 | |
| self.mock_model.count_tokens = mock_count_tokens | |
| ( | |
| reserved, | |
| offload, | |
| ) = await self.agent._split_tool_result_for_compression( | |
| tool_result, | |
| ) | |
| # Verify results using assertListEqual | |
| expected_reserved = [ | |
| {"type": "text", "text": "A" * 100, "id": "block1"}, | |
| ] | |
| expected_offload = [ | |
| {"type": "text", "text": "B" * 20, "id": "block2"}, | |
| {"type": "text", "text": "C" * 20, "id": "block3"}, | |
| ] | |
| self.assertListEqual( | |
| [b.model_dump() for b in reserved.output], | |
| expected_reserved, | |
| ) | |
| print(offload.output) | |
| self.assertListEqual( | |
| [b.model_dump() for b in offload.output], | |
| expected_offload, | |
| ) | |
| async def test_boundary_first_block_data(self) -> None: | |
| """Test when boundary is the first block and it is a DataBlock.""" | |
| block1 = DataBlock( | |
| source=Base64Source(data="base64data", media_type="image/png"), | |
| id="block1", | |
| ) | |
| block2 = TextBlock(text="B" * 20, id="block2") | |
| block3 = TextBlock(text="C" * 20, id="block3") | |
| tool_result = ToolResultBlock( | |
| id="test_6", | |
| name="test_tool", | |
| output=[block1, block2, block3], | |
| ) | |
| async def mock_count_tokens( | |
| messages: list, | |
| tools: list | None = None, | |
| ) -> int: | |
| """Mock token counting function based on content length.""" | |
| content = messages[0].content | |
| if isinstance(content, list): | |
| if len(content) == 3: | |
| return 150 | |
| elif len(content) == 2: | |
| return 80 | |
| elif len(content) == 1: | |
| return 60 | |
| return 50 | |
| self.mock_model.count_tokens = mock_count_tokens | |
| ( | |
| reserved, | |
| offload, | |
| ) = await self.agent._split_tool_result_for_compression( | |
| tool_result, | |
| ) | |
| # Verify results | |
| self.assertIsNotNone(reserved) | |
| self.assertIsNotNone(offload) | |
| # Verify ToolResultBlock metadata | |
| self.assertEqual(reserved.id, tool_result.id) | |
| self.assertEqual(offload.id, tool_result.id) | |
| # Verify results using assertListEqual | |
| expected_reserved = [ | |
| { | |
| "type": "data", | |
| "id": "block1", | |
| "source": { | |
| "type": "base64", | |
| "data": "base64data", | |
| "media_type": "image/png", | |
| }, | |
| "name": None, | |
| }, | |
| {"type": "text", "text": "B" * 20 + "C" * 5, "id": "block2"}, | |
| ] | |
| expected_offload = [ | |
| {"type": "text", "text": "C" * 15, "id": "block3"}, | |
| ] | |
| self.assertListEqual( | |
| [b.model_dump() for b in reserved.output], | |
| expected_reserved, | |
| ) | |
| self.assertListEqual( | |
| [b.model_dump() for b in offload.output], | |
| expected_offload, | |
| ) | |
| async def test_boundary_middle_block_text(self) -> None: | |
| """Test when boundary is a middle block and it is a TextBlock.""" | |
| block1 = TextBlock(text="A" * 20, id="block1") | |
| block2 = TextBlock(text="B" * 100, id="block2") | |
| block3 = TextBlock(text="C" * 20, id="block3") | |
| tool_result = ToolResultBlock( | |
| id="test_7", | |
| name="test_tool", | |
| output=[block1, block2, block3], | |
| ) | |
| async def mock_count_tokens( | |
| messages: list, | |
| tools: list | None = None, | |
| ) -> int: | |
| """Mock token counting function that counts text length in | |
| blocks.""" | |
| content = messages[0].content | |
| if isinstance(content, list): | |
| total = sum(len(b.text) for b in content if hasattr(b, "text")) | |
| return total | |
| return 0 | |
| self.mock_model.count_tokens = mock_count_tokens | |
| ( | |
| reserved, | |
| offload, | |
| ) = await self.agent._split_tool_result_for_compression( | |
| tool_result, | |
| ) | |
| # Verify results | |
| self.assertIsNotNone(reserved) | |
| self.assertIsNotNone(offload) | |
| # Verify ToolResultBlock metadata | |
| self.assertEqual(reserved.id, tool_result.id) | |
| self.assertEqual(offload.id, tool_result.id) | |
| # Verify results using assertListEqual | |
| expected_reserved = [ | |
| {"type": "text", "text": "A" * 20 + "B" * 80, "id": "block1"}, | |
| ] | |
| expected_offload = [ | |
| {"type": "text", "text": "B" * 20 + "C" * 20, "id": "block3"}, | |
| ] | |
| self.assertListEqual( | |
| [b.model_dump() for b in reserved.output], | |
| expected_reserved, | |
| ) | |
| self.assertListEqual( | |
| [b.model_dump() for b in offload.output], | |
| expected_offload, | |
| ) | |
| async def test_boundary_middle_block_data(self) -> None: | |
| """Test when boundary is a middle block and it is a DataBlock.""" | |
| block1 = TextBlock(text="A" * 20, id="block1") | |
| block2 = DataBlock( | |
| source=Base64Source(data="base64data", media_type="image/png"), | |
| id="block2", | |
| ) | |
| block3 = TextBlock(text="C" * 20, id="block3") | |
| tool_result = ToolResultBlock( | |
| id="test_8", | |
| name="test_tool", | |
| output=[block1, block2, block3], | |
| ) | |
| async def mock_count_tokens( | |
| messages: list, | |
| tools: list | None = None, | |
| ) -> int: | |
| """Mock token counting function based on content length.""" | |
| content = messages[0].content | |
| if isinstance(content, list): | |
| if len(content) == 3: | |
| return 150 | |
| elif len(content) == 2: | |
| return 80 | |
| elif len(content) == 1: | |
| return 40 | |
| return 50 | |
| self.mock_model.count_tokens = mock_count_tokens | |
| ( | |
| reserved, | |
| offload, | |
| ) = await self.agent._split_tool_result_for_compression( | |
| tool_result, | |
| ) | |
| # Verify results | |
| self.assertIsNotNone(reserved) | |
| self.assertIsNotNone(offload) | |
| # Verify ToolResultBlock metadata | |
| self.assertEqual(reserved.id, tool_result.id) | |
| self.assertEqual(offload.id, tool_result.id) | |
| # Verify results using assertListEqual | |
| expected_reserved = [ | |
| {"type": "text", "text": "A" * 20, "id": "block1"}, | |
| { | |
| "type": "data", | |
| "id": "block2", | |
| "source": { | |
| "type": "base64", | |
| "data": "base64data", | |
| "media_type": "image/png", | |
| }, | |
| "name": None, | |
| }, | |
| {"type": "text", "text": "C" * 5, "id": "block3"}, | |
| ] | |
| expected_offload = [ | |
| {"type": "text", "text": "C" * 15, "id": "block3"}, | |
| ] | |
| self.assertListEqual( | |
| [b.model_dump() for b in reserved.output], | |
| expected_reserved, | |
| ) | |
| self.assertListEqual( | |
| [b.model_dump() for b in offload.output], | |
| expected_offload, | |
| ) | |
| async def asyncTearDown(self) -> None: | |
| """The async teardown method.""" | |