From 1736c18a199d022a2ce4adc7d3765406e4405ac8 Mon Sep 17 00:00:00 2001 From: Brace Sproul Date: Mon, 9 Jun 2025 15:06:20 -0700 Subject: [PATCH] feat: Add update plan tool (#98) * feat: Add update plan tool * cr --- apps/open-swe/src/index.ts | 21 ++- apps/open-swe/src/nodes/generate-message.ts | 11 +- apps/open-swe/src/nodes/index.ts | 1 + apps/open-swe/src/nodes/update-plan.ts | 181 ++++++++++++++++++++ apps/open-swe/src/tools/index.ts | 1 + apps/open-swe/src/tools/update-plan.ts | 19 ++ 6 files changed, 230 insertions(+), 4 deletions(-) create mode 100644 apps/open-swe/src/nodes/update-plan.ts create mode 100644 apps/open-swe/src/tools/update-plan.ts diff --git a/apps/open-swe/src/index.ts b/apps/open-swe/src/index.ts index 6b1d6557..0d398f0a 100644 --- a/apps/open-swe/src/index.ts +++ b/apps/open-swe/src/index.ts @@ -1,4 +1,4 @@ -import { END, START, StateGraph } from "@langchain/langgraph"; +import { END, Send, START, StateGraph } from "@langchain/langgraph"; import { GraphAnnotation, GraphConfiguration, @@ -16,6 +16,7 @@ import { openPullRequest, diagnoseError, requestHelp, + updatePlan, } from "./nodes/index.js"; import { isAIMessage } from "@langchain/core/messages"; import { plannerGraph } from "./subgraphs/index.js"; @@ -26,11 +27,11 @@ import { plannerGraph } from "./subgraphs/index.js"; * Otherwise, it ends the process. * * @param {GraphState} state - The current graph state. - * @returns {"open-pr" | "take-action" | "request-help"} The next node to execute, or END if the process should stop. + * @returns {"open-pr" | "take-action" | "request-help" | Send} The next node to execute, or END if the process should stop. */ async function routeGeneratedAction( state: GraphState, -): Promise<"open-pr" | "take-action" | "request-help"> { +): Promise<"open-pr" | "take-action" | "request-help" | Send> { const { messages } = state; const lastMessage = messages[messages.length - 1]; @@ -41,6 +42,17 @@ async function routeGeneratedAction( return "request-help"; } + if ( + toolCall.name === "update_plan" && + "update_plan_reasoning" in toolCall.args && + typeof toolCall.args?.update_plan_reasoning === "string" + ) { + // Need to return a `Send` here so that we can update the state to include the plan change request. + return new Send("update-plan", { + planChangeRequest: toolCall.args?.update_plan_reasoning, + }); + } + return "take-action"; } @@ -59,6 +71,7 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration) .addNode("take-action", takeAction, { ends: ["progress-plan-step", "diagnose-error"], }) + .addNode("update-plan", updatePlan) .addNode("progress-plan-step", progressPlanStep, { ends: ["summarize-task-steps", "generate-action", "generate-conclusion"], }) @@ -80,7 +93,9 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration) "take-action", "request-help", "open-pr", + "update-plan", ]) + .addEdge("update-plan", "generate-action") .addEdge("generate-conclusion", "open-pr") .addEdge("diagnose-error", "generate-action") .addEdge("open-pr", END); diff --git a/apps/open-swe/src/nodes/generate-message.ts b/apps/open-swe/src/nodes/generate-message.ts index 1562af45..19d983fd 100644 --- a/apps/open-swe/src/nodes/generate-message.ts +++ b/apps/open-swe/src/nodes/generate-message.ts @@ -8,6 +8,7 @@ import { shellTool, applyPatchTool, requestHumanHelpTool, + updatePlanTool, } from "../tools/index.js"; import { getRepoAbsolutePath } from "../utils/git.js"; import { formatPlanPrompt } from "../utils/plan-prompt.js"; @@ -66,6 +67,9 @@ You MUST adhere to the following criteria when executing the task: - If the package manager fails to install, or you have issues installing dependencies, do not try to use a different package manager. Instead, skip installing dependencies. - If you are lacking enough context to complete the user's task, you may call the \`request_human_help\` tool to request help from the human. - This tool should only be used if you have already tried to gather all the context you need, and are still unable to complete the user's task. +- If you determine your current plan is not appropriate, or you need to update/remove/add steps to your plan, you may call the \`update_plan\` tool. + - This tool should only be called to make major changes, such as removing a task, or adding new tasks. For small changes, you do not necessarily need to call this tool, and can instead just act on those small updates. + - The \`update_plan\` tool can only update/remove/add plans to the list of tasks which are not yet completed (this includes the current task). - If completing the user's task requires writing or modifying files: - Your code and final answer should follow these *CODING GUIDELINES*: - Avoid writing to files which you have not already read. @@ -137,7 +141,12 @@ export async function generateAction( config: GraphConfig, ): Promise { const model = await loadModel(config, Task.ACTION_GENERATOR); - const tools = [shellTool, applyPatchTool, requestHumanHelpTool]; + const tools = [ + shellTool, + applyPatchTool, + requestHumanHelpTool, + updatePlanTool, + ]; const modelWithTools = model.bindTools(tools, { tool_choice: "auto" }); const response = await modelWithTools.invoke([ diff --git a/apps/open-swe/src/nodes/index.ts b/apps/open-swe/src/nodes/index.ts index c9b9d0a7..5391722c 100644 --- a/apps/open-swe/src/nodes/index.ts +++ b/apps/open-swe/src/nodes/index.ts @@ -9,3 +9,4 @@ export * from "./generate-conclusion.js"; export * from "./open-pr.js"; export * from "./diagnose-error.js"; export * from "./request-help.js"; +export * from "./update-plan.js"; diff --git a/apps/open-swe/src/nodes/update-plan.ts b/apps/open-swe/src/nodes/update-plan.ts new file mode 100644 index 00000000..42f25c85 --- /dev/null +++ b/apps/open-swe/src/nodes/update-plan.ts @@ -0,0 +1,181 @@ +import { + GraphState, + GraphConfig, + PlanItem, + GraphUpdate, +} from "@open-swe/shared/open-swe/types"; +import { loadModel, Task } from "../utils/load-model.js"; +import { z } from "zod"; +import { getActiveTask, updateTaskPlanItems } from "../utils/task-plan.js"; +import { + BaseMessage, + isAIMessage, + ToolMessage, +} from "@langchain/core/messages"; +import { getMessageString } from "../utils/message/content.js"; +import { formatPlanPrompt } from "../utils/plan-prompt.js"; +import { createLogger, LogLevel } from "../utils/logger.js"; + +const logger = createLogger(LogLevel.INFO, "UpdatePlanNode"); + +const systemPrompt = `You are operating as an agentic coding assistant built by LangChain. You've decided that the current plan you're working through needs to be updated. +To aid in this process, you've generated some reasoning and additional context into which plan steps you should update, remove, or whether to add new step(s). + +Here is the user's initial request which you used to generate the initial plan: +{USER_REQUEST} + +Here is the full plan you generated, which should have changes made to it: +{PLAN} + +Here is the reasoning and context you generated for which plan steps to update, remove, or add: +{REASONING} + +Given this context, update, remove or add plan steps as needed. + +You MUST adhere to the following criteria when generating the plan: +- Make as few changes as possible to the tasks, while still following the users request. +- You are only allowed to update plan items which are remaining, including the current task. Plan items which have already been completed are not allowed to be modified. +- The user will provide the full conversation history which led up to your deciding you need to update the plan. Use this conversation as context when making changes. +- The plan items listed above will include: + - The index of the plan item. This is the order in which the plan items should be executed in. + - The actual plan of the individual task. + - If it's been completed, it will include a summary of the completed task. +- To update the plan, you MUST pass every updated/added/untouched plan item to the \`update_plan\` tool. + - These will replace all of the existing plan items. + - This means you still need to include all of the unmodified plan items in the \`update_plan\` tool call. +- You should call the \`update_plan\` tool, passing in each plan item in the order they should be executed in. +- To remove an item from the plan, you should not include it in the \`update_plan\` tool call. + +With all of this in mind, please call the \`update_plan\` tool with the updated plan. +`; + +const updatePlanToolSchema = z.object({ + plan: z + .array(z.string()) + .describe( + "The updated, or new plan, including any changes to the plan items, as well as any new plan items you've added.", + ), +}); + +const updatePlanTool = { + name: "update_plan", + description: + "The updated plan, including any changes to the plan items, as well as any new plan items you've added, and the unchanged plan items. This should NOT include any of the completed plan items.", + schema: updatePlanToolSchema, +}; + +const formatSystemPrompt = ( + userRequest: string, + reasoning: string, + planItems: PlanItem[], +) => { + return systemPrompt + .replace("{USER_REQUEST}", userRequest) + .replace("{PLAN}", formatPlanPrompt(planItems, { includeSummaries: true })) + .replace("{REASONING}", reasoning); +}; + +const formatUserMessage = (messages: BaseMessage[]): string => { + return `Here is the full conversation history you should use as context when making changes to the plan: + +${messages.map(getMessageString).join("\n")}`; +}; + +export async function updatePlan( + state: GraphState, + config: GraphConfig, +): Promise { + if (!state.planChangeRequest) { + throw new Error("No plan change request found."); + } + const lastMessage = state.messages[state.messages.length - 1]; + if ( + !lastMessage || + !isAIMessage(lastMessage) || + !lastMessage.tool_calls?.length || + lastMessage.tool_calls[0].name !== updatePlanTool.name || + !lastMessage.tool_calls[0].id + ) { + throw new Error("Last message was not an update plan tool call."); + } + const updatePlanToolCallId = lastMessage.tool_calls[0].id; + + logger.info("Updating plan", { + updatePlanToolCallId, + planChangeRequest: state.planChangeRequest, + }); + + const model = await loadModel(config, Task.PLANNER); + const modelWithTools = model.bindTools([updatePlanTool], { + tool_choice: updatePlanTool.name, + }); + + const activeTask = getActiveTask(state.plan); + const request = activeTask.request; + const activePlanItems = activeTask.planRevisions.find( + (pr) => pr.revisionIndex === activeTask.activeRevisionIndex, + )?.plans; + if (!activePlanItems?.length) { + throw new Error("No active plan items found."); + } + + const systemPrompt = formatSystemPrompt( + request, + state.planChangeRequest, + activePlanItems, + ); + const userMessage = formatUserMessage(state.messages); + + const response = await modelWithTools.invoke([ + { + role: "system", + content: systemPrompt, + }, + { + role: "user", + content: userMessage, + }, + ]); + const toolCall = response.tool_calls?.[0]; + if (!toolCall) { + throw new Error("No tool call found."); + } + + const { plan } = toolCall.args as z.infer; + const completedPlanItems = activePlanItems.filter((item) => item.completed); + const totalCompletedPlanItems = completedPlanItems.length; + const newPlanItems: PlanItem[] = [ + ...completedPlanItems, + ...plan.map((p, index) => ({ + index: totalCompletedPlanItems + index, + plan: p, + completed: false, + summary: undefined, + })), + ]; + + const newTaskPlan = updateTaskPlanItems( + state.plan, + activeTask.id, + newPlanItems, + "agent", + ); + + return { + messages: [ + new ToolMessage({ + tool_call_id: updatePlanToolCallId, + content: + "Successfully updated the plan. The complete updated plan items are as follow:\n\n" + + newPlanItems + .map( + (p) => + `${p.plan}`, + ) + .join("\n"), + }), + ], + plan: newTaskPlan, + planChangeRequest: null, + }; +} diff --git a/apps/open-swe/src/tools/index.ts b/apps/open-swe/src/tools/index.ts index 2473dbec..e50358f8 100644 --- a/apps/open-swe/src/tools/index.ts +++ b/apps/open-swe/src/tools/index.ts @@ -2,3 +2,4 @@ export * from "./apply-patch.js"; export * from "./shell.js"; export * from "./session-plan.js"; export * from "./request-human-help.js"; +export * from "./update-plan.js"; diff --git a/apps/open-swe/src/tools/update-plan.ts b/apps/open-swe/src/tools/update-plan.ts new file mode 100644 index 00000000..c27e63ee --- /dev/null +++ b/apps/open-swe/src/tools/update-plan.ts @@ -0,0 +1,19 @@ +import { z } from "zod"; + +const updatePlanSchema = z.object({ + update_plan_reasoning: z + .string() + .describe( + "The reasoning for why you are updating the plan. This should include context which will be useful when actually updating the plan, such as what plan items to update, edit, or remove, along with any other context that would be useful when updating the plan.", + ), +}); + +export const updatePlanTool = { + name: "update_plan", + schema: updatePlanSchema, + description: + "Call this tool to update the current plan. This should ONLY be called if you want to remove, edit, or add plan items to the current plan." + + "\nDo NOT call this tool to mark a plan item as completed, or add a summary." + + "\nYou can not edit/remove completed plan items. This tool can only be used to update/add/remove plan items from the remaining and current plan items." + + "\nThe reasoning you pass to this tool will be used in the step that actually updates the plan, so ensure it is useful and concise.", +};