diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/graph/Context/MultiTurnContextManager.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/graph/Context/MultiTurnContextManager.java index 0137a90e6..5e925a10c 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/graph/Context/MultiTurnContextManager.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/graph/Context/MultiTurnContextManager.java @@ -72,7 +72,25 @@ public void appendPlannerChunk(String threadId, String chunk) { } /** - * Finalize current turn and add to history if planner output is available. + * Append assistant output chunk for the current turn. This is used as a + * fallback history entry when a turn ends before PlannerNode, for example when + * the feasibility assessment asks the user to clarify the query. + * @param threadId conversation thread id + * @param chunk assistant streaming chunk + */ + public void appendAssistantChunk(String threadId, String chunk) { + if (StringUtils.isAnyBlank(threadId, chunk)) { + return; + } + PendingTurn pending = pendingTurns.get(threadId); + if (pending != null) { + pending.assistantBuilder.append(chunk); + } + } + + /** + * Finalize current turn and add to history if planner or assistant output is + * available. * @param threadId conversation thread id */ public void finishTurn(String threadId) { @@ -81,18 +99,21 @@ public void finishTurn(String threadId) { return; } String plan = StringUtils.trimToEmpty(pending.planBuilder.toString()); - if (StringUtils.isBlank(plan)) { - log.debug("No planner output recorded for thread {}, skipping history update", threadId); + String assistantOutput = StringUtils.trimToEmpty(pending.assistantBuilder.toString()); + String response = StringUtils.isNotBlank(plan) ? plan : assistantOutput; + if (StringUtils.isBlank(response)) { + log.debug("No planner or assistant output recorded for thread {}, skipping history update", threadId); return; } - String trimmedPlan = StringUtils.abbreviate(plan, properties.getMaxplanlength()); + String responseLabel = StringUtils.isNotBlank(plan) ? "AI计划" : "AI回复"; + String trimmedResponse = StringUtils.abbreviate(response, properties.getMaxplanlength()); Deque deque = history.computeIfAbsent(threadId, k -> new ArrayDeque<>()); synchronized (deque) { while (deque.size() >= properties.getMaxturnhistory()) { deque.pollFirst(); } - deque.addLast(new ConversationTurn(pending.userQuestion, trimmedPlan)); + deque.addLast(new ConversationTurn(pending.userQuestion, responseLabel, trimmedResponse)); } } @@ -135,11 +156,11 @@ public String buildContext(String threadId) { return "(无)"; } return deque.stream() - .map(turn -> "用户: " + turn.userQuestion() + "\nAI计划: " + turn.plan()) + .map(turn -> "用户: " + turn.userQuestion() + "\n" + turn.responseLabel() + ": " + turn.response()) .collect(Collectors.joining("\n")); } - private record ConversationTurn(String userQuestion, String plan) { + private record ConversationTurn(String userQuestion, String responseLabel, String response) { } private static class PendingTurn { @@ -148,6 +169,8 @@ private static class PendingTurn { private final StringBuilder planBuilder = new StringBuilder(); + private final StringBuilder assistantBuilder = new StringBuilder(); + private PendingTurn(String userQuestion) { this.userQuestion = userQuestion; } diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/graph/GraphServiceImpl.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/graph/GraphServiceImpl.java index 6325b9bbf..859d64a18 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/graph/GraphServiceImpl.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/graph/GraphServiceImpl.java @@ -324,6 +324,9 @@ private void handleStreamNodeOutput(GraphRequest request, StreamingOutput output if (PlannerNode.class.getSimpleName().equals(node)) { multiTurnContextManager.appendPlannerChunk(threadId, chunk); } + else { + multiTurnContextManager.appendAssistantChunk(threadId, chunk); + } GraphNodeResponse response = GraphNodeResponse.builder() .agentId(request.getAgentId()) .threadId(threadId) diff --git a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/service/graph/MultiTurnContextManagerTest.java b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/service/graph/MultiTurnContextManagerTest.java index e5e277f35..589925937 100644 --- a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/service/graph/MultiTurnContextManagerTest.java +++ b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/service/graph/MultiTurnContextManagerTest.java @@ -78,7 +78,18 @@ void finishTurn_completesAndStoresInHistory() { } @Test - void finishTurn_noPlannerOutput_skipsHistory() { + void finishTurn_assistantOutputWithoutPlanner_storesHistory() { + contextManager.beginTurn("thread-1", "Query 1"); + contextManager.appendAssistantChunk("thread-1", "Please clarify the metric."); + contextManager.finishTurn("thread-1"); + + String context = contextManager.buildContext("thread-1"); + assertTrue(context.contains("Query 1")); + assertTrue(context.contains("AI回复: Please clarify the metric.")); + } + + @Test + void finishTurn_noOutput_skipsHistory() { contextManager.beginTurn("thread-1", "Query 1"); contextManager.finishTurn("thread-1");