diff --git a/camel/agents/chat_agent.py b/camel/agents/chat_agent.py index c6337408d7..4cfca38096 100644 --- a/camel/agents/chat_agent.py +++ b/camel/agents/chat_agent.py @@ -1480,7 +1480,26 @@ def _clean_snapshot_in_memory( if '- ' in result_text and '[ref=' in result_text: cleaned_result = self._clean_snapshot_content(result_text) - # Update the message in memory storage + chat_history_block = getattr( + self.memory, "_chat_history_block", None + ) + storage = getattr(chat_history_block, "storage", None) + if storage is None: + return + + existing_records = storage.load() + + # Remove records by UUID + updated_records = [ + record + for record in existing_records + if record["uuid"] not in entry.record_uuids + ] + + # Recreate only the function result message with cleaned content. + # The assistant message with tool calls is already recorded + # separately by _record_assistant_tool_calls_from_requests and + # should not be modified here. timestamp = ( entry.record_timestamps[0] if entry.record_timestamps @@ -1495,20 +1514,6 @@ def _clean_snapshot_in_memory( result=cleaned_result, tool_call_id=entry.tool_call_id, ) - - chat_history_block = getattr( - self.memory, "_chat_history_block", None - ) - storage = getattr(chat_history_block, "storage", None) - if storage is None: - return - - existing_records = storage.load() - updated_records = [ - record - for record in existing_records - if record["uuid"] not in entry.record_uuids - ] new_record = MemoryRecord( message=cleaned_message, role_at_backend=OpenAIBackendRole.FUNCTION, diff --git a/test/agents/test_chat_agent.py b/test/agents/test_chat_agent.py index bd654452c8..9699714693 100644 --- a/test/agents/test_chat_agent.py +++ b/test/agents/test_chat_agent.py @@ -129,20 +129,20 @@ def test_chat_agent(model, step_call_count=3): for i in range(step_call_count): for user_msg in [user_msg_bm, user_msg_str]: response = assistant.step(user_msg) - assert isinstance( - response.msgs, list - ), f"Error in round {i + 1}" + assert isinstance(response.msgs, list), ( + f"Error in round {i + 1}" + ) assert len(response.msgs) > 0, f"Error in round {i + 1}" - assert isinstance( - response.terminated, bool - ), f"Error in round {i + 1}" + assert isinstance(response.terminated, bool), ( + f"Error in round {i + 1}" + ) assert response.terminated is False, f"Error in round {i + 1}" - assert isinstance( - response.info, dict - ), f"Error in round {i + 1}" - assert ( - response.info['id'] is not None - ), f"Error in round {i + 1}" + assert isinstance(response.info, dict), ( + f"Error in round {i + 1}" + ) + assert response.info['id'] is not None, ( + f"Error in round {i + 1}" + ) @pytest.mark.model_backend @@ -370,9 +370,9 @@ def test_chat_agent_step_with_external_tools(step_call_count=3): external_tool_call_requests = response.info[ "external_tool_call_requests" ] - assert ( - external_tool_call_requests[0].tool_name == "math_subtract" - ), f"Error in calling round {i + 1}" + assert external_tool_call_requests[0].tool_name == "math_subtract", ( + f"Error in calling round {i + 1}" + ) @pytest.mark.model_backend @@ -514,9 +514,9 @@ async def mock_arun(*args, **kwargs): external_tool_call_requests = response.info[ "external_tool_call_requests" ] - assert ( - external_tool_call_requests[0].tool_name == "math_subtract" - ), f"Error in calling round {i + 1}" + assert external_tool_call_requests[0].tool_name == "math_subtract", ( + f"Error in calling round {i + 1}" + ) @pytest.mark.model_backend @@ -657,18 +657,18 @@ def test_chat_agent_multiple_return_messages(n, step_call_count=3): ) for i in range(step_call_count): - assert ( - assistant_with_sys_msg_response.msgs is not None - ), f"Error in calling round {i + 1}" - assert ( - len(assistant_with_sys_msg_response.msgs) == n - ), f"Error in calling round {i + 1}" - assert ( - assistant_without_sys_msg_response.msgs is not None - ), f"Error in calling round {i + 1}" - assert ( - len(assistant_without_sys_msg_response.msgs) == n - ), f"Error in calling round {i + 1}" + assert assistant_with_sys_msg_response.msgs is not None, ( + f"Error in calling round {i + 1}" + ) + assert len(assistant_with_sys_msg_response.msgs) == n, ( + f"Error in calling round {i + 1}" + ) + assert assistant_without_sys_msg_response.msgs is not None, ( + f"Error in calling round {i + 1}" + ) + assert len(assistant_without_sys_msg_response.msgs) == n, ( + f"Error in calling round {i + 1}" + ) @pytest.mark.model_backend @@ -753,12 +753,12 @@ def test_chat_agent_stream_output(step_call_count=3): assert len(msg.content) > 0, f"Error in calling round {i + 1}" stream_usage = stream_assistant_response.info["usage"] - assert ( - stream_usage["completion_tokens"] > 0 - ), f"Error in calling round {i + 1}" - assert ( - stream_usage["prompt_tokens"] > 0 - ), f"Error in calling round {i + 1}" + assert stream_usage["completion_tokens"] > 0, ( + f"Error in calling round {i + 1}" + ) + assert stream_usage["prompt_tokens"] > 0, ( + f"Error in calling round {i + 1}" + ) assert ( stream_usage["total_tokens"] == stream_usage["completion_tokens"] @@ -1039,12 +1039,12 @@ def test_tool_calling_sync(step_call_count=3): ] assert len(tool_calls) > 0, f"Error in calling round {i + 1}" - assert str(tool_calls[0]).startswith( - "Tool Execution" - ), f"Error in calling round {i + 1}" - assert ( - tool_calls[0].tool_name == "math_multiply" - ), f"Error in calling round {i + 1}" + assert str(tool_calls[0]).startswith("Tool Execution"), ( + f"Error in calling round {i + 1}" + ) + assert tool_calls[0].tool_name == "math_multiply", ( + f"Error in calling round {i + 1}" + ) assert tool_calls[0].args == { "a": 2, "b": 8, @@ -1165,9 +1165,9 @@ async def test_tool_calling_math_async(step_call_count=3): tool_calls = agent_response.info['tool_calls'] - assert ( - tool_calls[0].tool_name == "math_multiply" - ), f"Error in calling round {i + 1}" + assert tool_calls[0].tool_name == "math_multiply", ( + f"Error in calling round {i + 1}" + ) assert tool_calls[0].args == { "a": 2, "b": 8, @@ -1254,16 +1254,16 @@ def mock_run_tool_calling_async(*args, **kwargs): tool_calls = agent_response.info['tool_calls'] assert tool_calls, f"Error in calling round {i + 1}" - assert str(tool_calls[0]).startswith( - "Tool Execution" - ), f"Error in calling round {i + 1}" + assert str(tool_calls[0]).startswith("Tool Execution"), ( + f"Error in calling round {i + 1}" + ) - assert ( - tool_calls[0].tool_name == "async_sleep" - ), f"Error in calling round {i + 1}" - assert tool_calls[0].args == { - 'second': 1 - }, f"Error in calling round {i + 1}" + assert tool_calls[0].tool_name == "async_sleep", ( + f"Error in calling round {i + 1}" + ) + assert tool_calls[0].args == {'second': 1}, ( + f"Error in calling round {i + 1}" + ) assert tool_calls[0].result == 1, f"Error in calling round {i + 1}" @@ -1294,9 +1294,9 @@ def test_response_words_termination(step_call_count=3): assert agent.terminated, f"Error in calling round {i + 1}" assert agent_response.terminated, f"Error in calling round {i + 1}" - assert ( - "goodbye" in agent_response.info['termination_reasons'][0] - ), f"Error in calling round {i + 1}" + assert "goodbye" in agent_response.info['termination_reasons'][0], ( + f"Error in calling round {i + 1}" + ) def test_chat_agent_vision(step_call_count=3): @@ -1362,9 +1362,9 @@ def test_chat_agent_vision(step_call_count=3): for i in range(step_call_count): agent_response = agent.step(user_msg) - assert ( - agent_response.msgs[0].content == "Yes." - ), f"Error in calling round {i + 1}" + assert agent_response.msgs[0].content == "Yes.", ( + f"Error in calling round {i + 1}" + ) @pytest.mark.model_backend @@ -1534,9 +1534,9 @@ async def test_chat_agent_async_stream_with_async_generator(): # Create an async generator that wraps the chunks # This simulates what GeminiModel does with _wrap_async_stream_with_ # thought_preservation - async def mock_async_generator() -> ( - AsyncGenerator[ChatCompletionChunk, None] - ): + async def mock_async_generator() -> AsyncGenerator[ + ChatCompletionChunk, None + ]: for chunk in chunks: yield chunk @@ -1563,12 +1563,12 @@ async def mock_async_generator() -> ( # Verify final response contains the accumulated content final_response = responses[-1] - assert ( - final_response.msg is not None - ), "Final response should have a message" - assert ( - "Hello" in final_response.msg.content - ), "Final content should contain 'Hello'" + assert final_response.msg is not None, ( + "Final response should have a message" + ) + assert "Hello" in final_response.msg.content, ( + "Final content should contain 'Hello'" + ) @pytest.mark.model_backend @@ -1718,9 +1718,9 @@ def test_add(a: int, b: int) -> int: call_count = 0 - async def mock_async_generator() -> ( - AsyncGenerator[ChatCompletionChunk, None] - ): + async def mock_async_generator() -> AsyncGenerator[ + ChatCompletionChunk, None + ]: nonlocal call_count if call_count == 0: call_count += 1 @@ -1837,3 +1837,84 @@ class MathResult(BaseModel): assert len(responses) > 1, "Should receive multiple streaming chunks" assert responses[-1].msg.parsed.answer == 6 assert responses[-1].msg.parsed.explanation + + +def test_clean_snapshot_in_memory(): + """Test that snapshot content is properly cleaned in memory. + + This tests the _clean_snapshot_in_memory functionality which removes + stale snapshot markers and references from tool output messages stored + in memory. The cleaning preserves the assistant message (tool call request) + and only updates the function result message. + """ + from unittest.mock import MagicMock, patch + + from camel.agents.chat_agent import _ToolOutputHistoryEntry + from camel.messages import FunctionCallingMessage + + # Create a mock model to avoid API calls + mock_model = MagicMock() + mock_model.model_type = ModelType.DEFAULT + mock_model.model_config_dict = {} + mock_model.token_counter = None + mock_model.model_platform_name = "openai" + + with patch.object(ChatAgent, '_init_model', return_value=mock_model): + agent = ChatAgent( + system_message="Test agent", + model=mock_model, + ) + + # Manually enable snapshot cleaning + agent._enable_snapshot_clean = True + agent._tool_output_history = [] + + # Create a mock memory storage + mock_storage = MagicMock() + mock_chat_history_block = MagicMock() + mock_chat_history_block.storage = mock_storage + + agent.memory._chat_history_block = mock_chat_history_block + + # Create a test entry with snapshot markers + test_uuid = "test-uuid-123" + test_timestamp = 1234567890.0 + entry = _ToolOutputHistoryEntry( + tool_name="test_tool", + tool_call_id="call_123", + result_text="- Item 1 [ref=abc]\n- Item 2 [ref=def]\n", + record_uuids=[test_uuid], + record_timestamps=[test_timestamp], + cached=False, + ) + agent._tool_output_history.append(entry) + + # Mock the storage to return the existing record + mock_storage.load.return_value = [ + { + "uuid": test_uuid, + "timestamp": test_timestamp, + "message": {"content": "- Item 1 [ref=abc]\n- Item 2 [ref=def]\n"}, + } + ] + + # Call the clean function + agent._clean_snapshot_in_memory(entry) + + # Verify storage was updated + assert mock_storage.clear.called + assert mock_storage.save.called + + # Get the saved records + saved_records = mock_storage.save.call_args[0][0] + + # Should have one record (the cleaned function result) + assert len(saved_records) == 1 + + # The record should be a function message with cleaned content + saved_record = saved_records[0] + assert saved_record["role_at_backend"] == OpenAIBackendRole.FUNCTION.value + + # Verify entry was marked as cached + assert entry.cached is True + assert len(entry.record_uuids) == 1 # Single function result record