Skip to content

fix(agent): use unique block ids for text/thinking stream events #2002

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
liulangjietou wants to merge 1 commit into
base: main
Choose a base branch
from
+128 −8
Open

fix(agent): use unique block ids for text/thinking stream events #2002

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 35 additions & 8 deletions agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2188,14 +2188,19 @@ private void emitBlockEvents(
blockLifecycle.startText(events);
if (tb.getText() != null && !tb.getText().isEmpty()) {
events.add(
new TextBlockDeltaEvent(blockLifecycle.replyId, "text", tb.getText()));
new TextBlockDeltaEvent(
blockLifecycle.replyId,
blockLifecycle.currentTextBlockId(),
tb.getText()));
}
} else if (block instanceof ThinkingBlock tb) {
blockLifecycle.startThinking(events);
if (tb.getThinking() != null && !tb.getThinking().isEmpty()) {
events.add(
new ThinkingBlockDeltaEvent(
blockLifecycle.replyId, "thinking", tb.getThinking()));
blockLifecycle.replyId,
blockLifecycle.currentThinkingBlockId(),
tb.getThinking()));
}
} else if (withToolEvents && block instanceof ToolUseBlock tub) {
String toolId = resolveToolCallId(tub, context);
Expand Down Expand Up @@ -2224,6 +2229,10 @@ private final class ModelCallBlockLifecycle {
private final String replyId;
private final AtomicBoolean textStarted = new AtomicBoolean(false);
private final AtomicBoolean thinkingStarted = new AtomicBoolean(false);
// Unique block ids allocated lazily on start and shared across the block's
// Start/Delta/End events, following the tool-call block precedent (toolId).
private final AtomicReference<String> textBlockId = new AtomicReference<>();
private final AtomicReference<String> thinkingBlockId = new AtomicReference<>();
private final Map<String, String> startedToolCalls = new ConcurrentHashMap<>();

private ModelCallBlockLifecycle(String replyId) {
Expand All @@ -2233,16 +2242,28 @@ private ModelCallBlockLifecycle(String replyId) {
private void startText(List<AgentEvent> events) {
flushThinking(events);
if (textStarted.compareAndSet(false, true)) {
events.add(new TextBlockStartEvent(replyId, "text"));
textBlockId.set(newBlockId());
events.add(new TextBlockStartEvent(replyId, textBlockId.get()));
}
}

private void startThinking(List<AgentEvent> events) {
if (thinkingStarted.compareAndSet(false, true)) {
events.add(new ThinkingBlockStartEvent(replyId, "thinking"));
thinkingBlockId.set(newBlockId());
events.add(new ThinkingBlockStartEvent(replyId, thinkingBlockId.get()));
}
}

/** Current text block id; non-null only while a text block is open. */
private String currentTextBlockId() {
return textBlockId.get();
}

/** Current thinking block id; non-null only while a thinking block is open. */
private String currentThinkingBlockId() {
return thinkingBlockId.get();
}

private void startToolCall(String toolId, String toolName, List<AgentEvent> events) {
if (toolId == null || startedToolCalls.containsKey(toolId)) {
return;
Expand All @@ -2258,13 +2279,13 @@ private void startToolCall(String toolId, String toolName, List<AgentEvent> even

private void flushText(List<AgentEvent> events) {
if (textStarted.compareAndSet(true, false)) {
events.add(new TextBlockEndEvent(replyId, "text"));
events.add(new TextBlockEndEvent(replyId, textBlockId.getAndSet(null)));
}
}

private void flushThinking(List<AgentEvent> events) {
if (thinkingStarted.compareAndSet(true, false)) {
events.add(new ThinkingBlockEndEvent(replyId, "thinking"));
events.add(new ThinkingBlockEndEvent(replyId, thinkingBlockId.getAndSet(null)));
}
}

Expand All @@ -2280,6 +2301,10 @@ private void flushAll(List<AgentEvent> events) {
flushThinking(events);
flushAllToolCalls(events);
}

private String newBlockId() {
return UUID.randomUUID().toString().replace("-", "");
}
}

private String resolveToolCallId(ToolUseBlock tub, ReasoningContext context) {
Expand Down Expand Up @@ -3152,7 +3177,8 @@ private Flux<AgentEvent> summaryModelCallStream(
new TextBlockDeltaEvent(
blockLifecycle
.replyId,
"text",
blockLifecycle
.currentTextBlockId(),
tb.getText()));
}
} else if (block
Expand All @@ -3166,7 +3192,8 @@ private Flux<AgentEvent> summaryModelCallStream(
new ThinkingBlockDeltaEvent(
blockLifecycle
.replyId,
"thinking",
blockLifecycle
.currentThinkingBlockId(),
tb
.getThinking()));
}
Expand Down
93 changes: 93 additions & 0 deletions agentscope-core/src/test/java/io/agentscope/core/agent/ReActAgentNewLoopReplyTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
package io.agentscope.core.agent;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;

Expand All @@ -26,8 +27,10 @@
import io.agentscope.core.event.ExceedMaxItersEvent;
import io.agentscope.core.event.ModelCallEndEvent;
import io.agentscope.core.event.ModelCallStartEvent;
import io.agentscope.core.event.TextBlockDeltaEvent;
import io.agentscope.core.event.TextBlockEndEvent;
import io.agentscope.core.event.TextBlockStartEvent;
import io.agentscope.core.event.ThinkingBlockDeltaEvent;
import io.agentscope.core.event.ThinkingBlockEndEvent;
import io.agentscope.core.event.ThinkingBlockStartEvent;
import io.agentscope.core.event.ToolCallEndEvent;
Expand Down Expand Up @@ -413,6 +416,96 @@ void summaryModelCallClosesThinkingBeforeTextAndFlushesTextBeforeModelEnd() {
assertTrue(summaryModelEnd > summaryTextEnd);
}

@Test
void textAndThinkingBlocksUseUniqueConsistentBlockIds() {
ChatModelBase model =
new ScriptedModel(
List.of(
() ->
Flux.just(
chatResponse(
ThinkingBlock.builder()
.thinking("think")
.build(),
TextBlock.builder()
.text("answer")
.build()))));
ReActAgent agent =
ReActAgent.builder().name("asst").model(model).toolkit(new Toolkit()).build();

List<AgentEvent> events = agent.streamEvents(List.of()).collectList().block();
assertNotNull(events);

String textStartId = firstOf(events, TextBlockStartEvent.class).getBlockId();
String textDeltaId = firstOf(events, TextBlockDeltaEvent.class).getBlockId();
String textEndId = firstOf(events, TextBlockEndEvent.class).getBlockId();
String thinkingStartId = firstOf(events, ThinkingBlockStartEvent.class).getBlockId();
String thinkingDeltaId = firstOf(events, ThinkingBlockDeltaEvent.class).getBlockId();
String thinkingEndId = firstOf(events, ThinkingBlockEndEvent.class).getBlockId();

// Block ids must be non-empty and no longer hardcoded to the block type name.
assertTrue(textStartId != null && !textStartId.isEmpty());
assertNotEquals("text", textStartId);
assertNotEquals("thinking", thinkingStartId);

// Start/Delta/End of the same block must share one id.
assertEquals(textStartId, textDeltaId);
assertEquals(textStartId, textEndId);
assertEquals(thinkingStartId, thinkingDeltaId);
assertEquals(thinkingStartId, thinkingEndId);

// Different block types must have different ids.
assertNotEquals(textStartId, thinkingStartId);
}

@Test
void separateModelCallsAllocateDifferentTextBlockIds() {
ChatModelBase model =
new ScriptedModel(
List.of(
() -> Flux.just(toolUseResponse("tc", "echo", "x")),
() ->
Flux.just(
chatResponse(
TextBlock.builder()
.text("summary")
.build()))));
ReActAgent agent =
ReActAgent.builder()
.name("asst")
.model(model)
.toolkit(toolkitWith(new EchoTool()))
.maxIters(1)
.build();

List<AgentEvent> events = agent.streamEvents(List.of()).collectList().block();
assertNotNull(events);

// First model call emits a tool use with no text block; the summary model call
// emits a text block. Collect all text block start ids to ensure the loop text
// block (if any) and the summary text block never reuse an id.
List<String> textBlockIds =
events.stream()
.filter(TextBlockStartEvent.class::isInstance)
.map(e -> ((TextBlockStartEvent) e).getBlockId())
.distinct()
.toList();
assertEquals(
events.stream().filter(TextBlockStartEvent.class::isInstance).count(),
(long) textBlockIds.size(),
"each text block start must carry a distinct block id");
}

@SuppressWarnings("unchecked")
private static <T> T firstOf(List<AgentEvent> events, Class<T> type) {
for (AgentEvent e : events) {
if (type.isInstance(e)) {
return (T) e;
}
}
return null;
}

private static int indexOf(List<AgentEvent> events, Class<?> type) {
return indexOfFrom(events, type, 0);
}
Expand Down
Loading