diff --git a/agentscope-extensions/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java b/agentscope-extensions/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java index 32fb30c5e..6ff3e7335 100644 --- a/agentscope-extensions/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java +++ b/agentscope-extensions/agentscope-extensions-autocontext-memory/src/main/java/io/agentscope/core/memory/autocontext/AutoContextMemory.java @@ -17,6 +17,7 @@ import io.agentscope.core.agent.accumulator.ReasoningContext; import io.agentscope.core.memory.Memory; +import io.agentscope.core.message.ContentBlock; import io.agentscope.core.message.MessageMetadataKeys; import io.agentscope.core.message.Msg; import io.agentscope.core.message.MsgRole; @@ -418,22 +419,114 @@ private boolean summaryCurrentRoundMessages(List rawMessages) { metadata.put("time", compressedMsg.getChatUsage().getTime()); } - // Record compression event (before replacing messages to preserve indices) + // Step 5: Preserve ReAct Structure and Replace + List replacementMsgs = new ArrayList<>(); + int lastToolResultIdx = -1; + ToolResultBlock lastOrigBlock = null; + + for (Msg msg : messagesToCompress) { + if (MsgUtils.isToolUseMessage(msg)) { + replacementMsgs.add(msg); + } else if (MsgUtils.isToolResultMessage(msg)) { + ToolResultBlock origBlock = null; + + for (ContentBlock block : msg.getContent()) { + if (block instanceof ToolResultBlock) { + origBlock = (ToolResultBlock) block; + break; + } + } + + if (origBlock != null) { + Msg placeholder = + Msg.builder() + .role(msg.getRole()) + .name(msg.getName()) + .metadata(msg.getMetadata()) + .content( + List.of( + ToolResultBlock.builder() + .name(origBlock.getName()) + .id(origBlock.getId()) + .metadata(origBlock.getMetadata()) + .output( + List.of( + TextBlock.builder() + .text( + "[Content" + + " compressed]") + .build())) + .build())) + .build(); + replacementMsgs.add(placeholder); + lastToolResultIdx = replacementMsgs.size() - 1; + lastOrigBlock = origBlock; + } else { + replacementMsgs.add(msg); + } + } + } + + Msg actualInsertedSummaryMsg; + + // If there is a tool call, inject the summarized summary into the last ToolResult + if (lastToolResultIdx != -1 && lastOrigBlock != null) { + Msg lastToolMsg = replacementMsgs.get(lastToolResultIdx); + + Map mergedMeta = new HashMap<>(); + if (lastToolMsg.getMetadata() != null) { + mergedMeta.putAll(lastToolMsg.getMetadata()); + } + if (compressedMsg.getMetadata() != null) { + mergedMeta.putAll(compressedMsg.getMetadata()); + } + + Msg finalSummaryMsg = + Msg.builder() + .role(lastToolMsg.getRole()) + .name(lastToolMsg.getName()) + .metadata(mergedMeta) + .content( + List.of( + ToolResultBlock.builder() + .name(lastOrigBlock.getName()) + .id(lastOrigBlock.getId()) + .metadata(lastOrigBlock.getMetadata()) + .output( + List.of( + TextBlock.builder() + .text( + compressedMsg + .getTextContent()) + .build())) + .build())) + .build(); + + replacementMsgs.set(lastToolResultIdx, finalSummaryMsg); + actualInsertedSummaryMsg = finalSummaryMsg; + } else { + replacementMsgs.add(compressedMsg); + actualInsertedSummaryMsg = compressedMsg; + } + + // Record compression event recordCompressionEvent( CompressionEvent.CURRENT_ROUND_MESSAGE_COMPRESS, startIndex, endIndex, rawMessages, - compressedMsg, + actualInsertedSummaryMsg, metadata); - // Step 5: Replace original messages with compressed one + // Clean up old data first, then insert new data rawMessages.subList(startIndex, endIndex + 1).clear(); - rawMessages.add(startIndex, compressedMsg); + rawMessages.addAll(startIndex, replacementMsgs); log.info( - "Replaced {} messages with 1 compressed message at index {}", + "Replaced {} messages with {} structured messages (including summary) starting at" + + " index {}", messagesToCompress.size(), + replacementMsgs.size(), startIndex); return true; } diff --git a/agentscope-extensions/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java b/agentscope-extensions/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java index 4b74e9f66..9430d577a 100644 --- a/agentscope-extensions/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java +++ b/agentscope-extensions/agentscope-extensions-autocontext-memory/src/test/java/io/agentscope/core/memory/autocontext/AutoContextMemoryTest.java @@ -656,41 +656,25 @@ void testMergeAndCompressCurrentRoundMessages() { // Verify that messages were compressed // Original: 8 initial + 1 user + 4 tool messages = 13 messages - // After compression: 8 initial + 1 user + 1 compressed = 10 messages (or less) - assertTrue( - messages.size() <= 10, - "Messages should be compressed. Expected 10 or less, got " + messages.size()); + assertEquals( + 13, + messages.size(), + "Message count should be preserved to maintain ReAct structure"); - // Verify that the compressed message contains the expected format - boolean hasCompressedMessage = false; - for (Msg msg : messages) { - String content = msg.getTextContent(); - if (content != null - && (content.contains("compressed_current_round") - || content.contains("Compressed current round summary"))) { - hasCompressedMessage = true; - break; - } - } - assertTrue(hasCompressedMessage, "Should contain compressed current round message"); + ToolResultBlock lastToolResult = (ToolResultBlock) messages.get(12).getContent().get(0); + String lastToolResultText = lastToolResult.getOutput().get(0).toString(); - // Verify that tool messages were offloaded (can be reloaded) - boolean hasOffloadHint = false; - for (Msg msg : messages) { - String content = msg.getTextContent(); - if (content != null - && (content.contains("uuid:") - || content.contains("uuid=") - || content.contains("CONTEXT_OFFLOAD") - || content.contains("reload") - || content.contains("context_reload") - || content.contains("offloaded"))) { - hasOffloadHint = true; - break; - } - } assertTrue( - hasOffloadHint, + lastToolResultText.contains("Compressed current round summary"), + "Should contain compressed current round message in the last ToolResult"); + + assertTrue( + lastToolResultText.contains("uuid:") + || lastToolResultText.contains("uuid=") + || lastToolResultText.contains("CONTEXT_OFFLOAD") + || lastToolResultText.contains("reload") + || lastToolResultText.contains("context_reload") + || lastToolResultText.contains("offloaded"), "Compressed message should contain offload hint for reloading original tool" + " messages"); @@ -1677,4 +1661,65 @@ void testGetPlanStateContextWithDifferentPlanStates() throws Exception { resultDone.contains("Goal: Test Description"), "Should contain goal for DONE state"); } + + @Test + @DisplayName( + "Should preserve ReAct tool_use/tool_result structure when compressing current round" + + " messages") + void testCurrentRoundSummaryPreservesReActStructure() { + TestModel testModel = new TestModel("Compressed current round summary"); + AutoContextConfig config = + AutoContextConfig.builder() + .msgThreshold(5) + .minConsecutiveToolMessages(10) + .largePayloadThreshold(10000) + .minCompressionTokenThreshold(0) + .build(); + AutoContextMemory testMemory = new AutoContextMemory(config, testModel); + + testMemory.addMessage(createTextMessage("Previous user message", MsgRole.USER)); + testMemory.addMessage(createTextMessage("Previous assistant response", MsgRole.ASSISTANT)); + + // Simulate the current round + testMemory.addMessage(createTextMessage("Current user query", MsgRole.USER)); + + // The simulated large model has initiated two consecutive tool calls in the current round + testMemory.addMessage(createToolUseMessage("search", "call_1")); + testMemory.addMessage(createToolResultMessage("search", "call_1", "Result 1")); + + testMemory.addMessage(createToolUseMessage("read", "call_2")); + testMemory.addMessage(createToolResultMessage("read", "call_2", "Result 2")); + + // Trigger Strategy 6: current round summary + boolean compressed = testMemory.compressIfNeeded(); + assertTrue(compressed, "Compression should be triggered"); + + List messages = testMemory.getMessages(); + assertEquals( + 7, + messages.size(), + "Should have exactly 7 messages, preserving ToolUse/ToolResult pairs"); + + assertTrue(MsgUtils.isToolUseMessage(messages.get(3))); + assertTrue(MsgUtils.isToolResultMessage(messages.get(4))); + assertTrue(MsgUtils.isToolUseMessage(messages.get(5))); + assertTrue(MsgUtils.isToolResultMessage(messages.get(6))); + + ToolResultBlock block1 = (ToolResultBlock) messages.get(4).getContent().get(0); + + String firstToolResultText = block1.getOutput().get(0).toString(); + assertTrue( + firstToolResultText.contains("[Content compressed]"), + "The first tool result should be compressed"); + + ToolResultBlock block2 = (ToolResultBlock) messages.get(6).getContent().get(0); + + String lastToolResultText = block2.getOutput().get(0).toString(); + assertTrue( + lastToolResultText.contains("Compressed current round summary"), + "The last tool result should contain the LLM summary"); + + assertNotNull(messages.get(6).getMetadata()); + assertTrue(messages.get(6).getMetadata().containsKey("_compress_meta")); + } }