Skip to content

fix(planner): handle concatenated plan responses #580

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
March-77 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(planner): handle concatenated plan responses #580

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
24 changes: 23 additions & 1 deletion data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/PlanProcessUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,11 @@
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.dataagent.dto.planner.ExecutionStep;
import com.alibaba.cloud.ai.dataagent.dto.planner.Plan;
import com.fasterxml.jackson.databind.MappingIterator;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.core.ParameterizedTypeReference;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -101,13 +103,33 @@ public static ExecutionStep getCurrentExecutionStep(Plan plan, Integer currentSt
public static Plan getPlan(OverAllState state) {
String plannerNodeOutput = (String) state.value(PLANNER_NODE_OUTPUT)
.orElseThrow(() -> new IllegalStateException("计划节点输出为空"));
Plan plan = converter.convert(plannerNodeOutput);
Plan plan = convertPlan(plannerNodeOutput);
if (plan == null) {
throw new IllegalStateException("计划解析失败");
}
return plan;
}

private static Plan convertPlan(String plannerNodeOutput) {
try (MappingIterator<Plan> plans = JsonUtil.getObjectMapper()
.readerFor(Plan.class)
.readValues(plannerNodeOutput)) {
Plan lastPlan = null;
int planCount = 0;
while (plans.hasNextValue()) {
lastPlan = plans.nextValue();
planCount++;
}
if (planCount > 1) {
return lastPlan;
}
}
catch (IOException ignored) {
// Let the existing converter handle its supported wrappers and report errors.
}
return converter.convert(plannerNodeOutput);
}

/**
* Get the current step number from state
* @param state the overall state
Expand Down
12 changes: 12 additions & 0 deletions ...ent-management/src/test/java/com/alibaba/cloud/ai/dataagent/util/PlanProcessUtilTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,18 @@ void getPlan_withValidJson_returnsParsedPlan() {
assertEquals(SQL_GENERATE_NODE, plan.getExecutionPlan().get(0).getToolToUse());
}

@Test
void getPlan_withConcatenatedPlans_returnsLastCompletePlan() {
String plannerOutput = TestFixtures.createSingleSqlPlanJson() + TestFixtures.createMultiStepPlanJson();
OverAllState state = TestFixtures.createStateWith(Map.of(PLANNER_NODE_OUTPUT, plannerOutput));

Plan plan = PlanProcessUtil.getPlan(state);

assertEquals("Multi-step analysis", plan.getThoughtProcess());
assertEquals(3, plan.getExecutionPlan().size());
assertEquals(REPORT_GENERATOR_NODE, plan.getExecutionPlan().get(2).getToolToUse());
}

@Test
void getPlan_withEmptyState_throwsException() {
OverAllState state = TestFixtures.createStateWith(PLANNER_NODE_OUTPUT);
Expand Down
Loading