From e6d113e65ace5829ad49720aa82a948a4d3fba57 Mon Sep 17 00:00:00 2001 From: Brace Sproul Date: Mon, 26 May 2025 16:33:52 -0700 Subject: [PATCH] fix: Improve rewrite plan and planner subgraph (#25) --- apps/open-swe/package.json | 1 + apps/open-swe/scripts/run-e2e.ts | 12 +- apps/open-swe/src/nodes/generate-message.ts | 5 +- apps/open-swe/src/nodes/rewrite-plan.ts | 265 +++++++++++++++--- .../src/nodes/summarize-task-steps.ts | 7 +- .../planner/nodes/generate-message.ts | 2 +- .../subgraphs/planner/nodes/generate-plan.ts | 8 +- .../src/subgraphs/planner/nodes/summarizer.ts | 3 + apps/open-swe/src/tools/shell.ts | 6 +- yarn.lock | 1 + 10 files changed, 260 insertions(+), 50 deletions(-) diff --git a/apps/open-swe/package.json b/apps/open-swe/package.json index 494c540b..e6ce601b 100644 --- a/apps/open-swe/package.json +++ b/apps/open-swe/package.json @@ -31,6 +31,7 @@ "@langchain/openai": "^0.5.10", "diff": "^8.0.1", "langchain": "^0.3.26", + "langsmith": "^0.3.29", "uuid": "^11.0.5", "zod": "^3.23.8" }, diff --git a/apps/open-swe/scripts/run-e2e.ts b/apps/open-swe/scripts/run-e2e.ts index ad54f978..cd386341 100644 --- a/apps/open-swe/scripts/run-e2e.ts +++ b/apps/open-swe/scripts/run-e2e.ts @@ -15,8 +15,16 @@ async function runE2E() { const threadId = uuidv4(); - const userRequest = - "The system message contains a REPO_DIRECTORY variable. But nothing is formatted in that variable. Please make sure to format the REPO_DIRECTORY variable in the system message. You will probably want to move the repo directory info from config to the State."; + const userRequest = `The project is structured as a monorepo, with some apps located inside the /apps directory. In this directory, there is an /auth directory. This directory only contains the scaffolding for a new app in the monorepo, but is not yet implemented. Please take the following plan/task description and implement it in the /auth directory: +This monorepo is for an AI coding agent. The app runs and edits the code in the cloud in a sandboxed environment. Right now, we require users to generate a GitHub PAT, which we store in a .env file and can use to authenticate with GitHub. This is not idea, and instead we want to have a github oauth app which users can authenticate with. +Please implement a new auth server inside the /auth directory which can do this. +You will not have any access to secrets, so you will not be able to run the server to test it. +I want the server to be able to authenticate users with GitHub, such that we will be able to take the following actions: +1. clone repositories they give us access to +2. checkout existing and create new branches on the repositories they give us access to +3. make pull requests and push changes to the repositories they give us access to +Once you're done, ensure you've documented the development process in the readme of this new app.`; + const configurable: Omit< GraphConfig["configurable"], "thread_id" | "assistant_id" diff --git a/apps/open-swe/src/nodes/generate-message.ts b/apps/open-swe/src/nodes/generate-message.ts index a41225ac..0bae0746 100644 --- a/apps/open-swe/src/nodes/generate-message.ts +++ b/apps/open-swe/src/nodes/generate-message.ts @@ -38,7 +38,10 @@ You MUST adhere to the following criteria when executing the task: - Analyzing code for vulnerabilities is allowed. - Showing user code and tool call details is allowed. - Remember to always properly format and quote your shell commands. -- Take advantage of the condensed context tool call messages in the conversation history. These contain summarized/condensed context from previously completed steps. Ensure you always read these messages to avoid duplicate work (e.g.: searching for file paths). +- Take advantage of the condensed context tool call messages in the conversation history (under the names \`condense_task_context\` and \`condense_planning_context\`). These contain summarized/condensed context from previously completed steps. Ensure you always read these messages to avoid duplicate work (e.g.: searching for file paths). + - These summary messages may include a section called 'Codebase files and descriptions' which contains a list of files, and descriptions of the files' contents. If you need context on a file, or directory, ensure you first check this section of the summary messages to avoid duplicate work. + - The summary messages may also include a section called 'Key repository insights and learnings'. This contains key insights, learnings, and facts the model discovered while completing a task. + - Each summary message will also include a short description of the task it completed, how it did so, and every change it made to the codebase during this task. This section will be titled 'Repository modifications summary'. - All changes are automatically committed, so you should not worry about creating backups, or committing changes. - Use \`apply_patch\` to edit files. This tool accepts diffs and file paths. It will then apply the given diff to the file. - When using the \`shell\` tool, always take advantage of the \`workdir\` parameter to run commands inside the repo directory. You should not try to generate a command with \`cd \` as passing that path to \`workdir\` is much more efficient. diff --git a/apps/open-swe/src/nodes/rewrite-plan.ts b/apps/open-swe/src/nodes/rewrite-plan.ts index af2ce714..2983d11c 100644 --- a/apps/open-swe/src/nodes/rewrite-plan.ts +++ b/apps/open-swe/src/nodes/rewrite-plan.ts @@ -1,33 +1,235 @@ -import { GraphState, GraphConfig, GraphUpdate } from "../types.js"; +import { GraphState, GraphConfig, GraphUpdate, PlanItem } from "../types.js"; import { loadModel, Task } from "../utils/load-model.js"; -import { sessionPlanTool } from "../tools/index.js"; +import { isHumanMessage } from "@langchain/core/messages"; +import { getMessageContentString } from "../utils/message/content.js"; +import { z } from "zod"; +import { tool } from "@langchain/core/tools"; +import { ConfigurableModel } from "langchain/chat_models/universal"; +import { traceable } from "langsmith/traceable"; -const systemPrompt = `You are operating as a terminal-based agentic coding assistant built by LangChain. It wraps LLM models to enable natural language interaction with a local codebase. You are expected to be precise, safe, and helpful. +const systemPromptIdentifyChanges = `You are operating as an agentic coding assistant built by LangChain. You've previously been given a task to generate a plan of action for, to address the user's initial request. -In this step, you are expected to rewrite a high-level plan to address the user's initial request. In a previous step you generated a plan, however the user has requested some changes: -## User Request +Here is the user's initial request: +{USER_INITIAL_REQUEST} + +After generating that plan, the user has submitted some feedback/change requests. You should now identify exactly which tasks in the plan should be modified based on their request. + +Here is their request: {USER_REQUEST} -Here is the previous plan: -## Previous Plan -{PREVIOUS_PLAN} +The plan you generated originally, which they submitted the above request for is as follows: +{PLAN} -The plan must be a list of actions to take, in order, to address the user's request. You should not include any code in the plan, only a list of actions to take. +Please read over the generated plan, and the user's request, and identify exactly which tasks in the plan should be modified/removed. Call the 'identify_plan_changes' tool and use the indices of the tasks listed above when calling the tool.`; + +const systemPrompt = `You are operating as an agentic coding assistant built by LangChain. You've previously been given a task to generate a plan of action for, to address the user's initial request. + +In this step, the user has requested you rewrite/modify parts of a high-level plan. You have already identified the specific tasks in the plan that should be modified/removed based on the user's request. + +Here is the user's initial request which you used to generate the initial plan: +{USER_INITIAL_REQUEST} + +Here is the full plan you generated: +{PLAN} + +Here is the request the user has just made which you should use to rewrite/modify the plan: +{USER_REQUEST} + +And here are the specific tasks in the plan which were identified as tasks the user wants to modify/remove: +{TASKS_TO_MODIFY} + +Given this context, please address the user's request to modify/remove/add the tasks in the plan. You MUST adhere to the following criteria when generating the plan: -- You do not have access to the codebase yet, so you cannot inspect it or make assumptions about it. -- Your plan should be high-level in nature, but should still be specific enough to be actionable. -- Make as few changes as possible to the previous plan, while still addressing the user's request. -- When you are ready to generate the plan, ensure you call the 'session_plan' tool. -- Ensure you generate the full plan in this tool call, not just the changes. +- Make as few changes as possible to the tasks, while still following the users request. +- You should NOT make ANY changes to the tasks in the plan that are NOT listed as tasks to modify/remove. +- Do NOT modify tasks in the plan not listed as tasks to modify/remove. +- When responding, ensure you include the unmodified tasks in the plan, as well as the modified/new tasks. +- To remove a specific task, simply do NOT include it in the response. +- To add a new task, simply include it in the response. `; -const formatSysPrompt = (userRequest: string, previousPlan: string) => { - return systemPrompt +const formatSysPromptIdentifyTasks = ( + userInitialRequest: string, + userRequest: string, + previousPlan: string[], +) => { + return systemPromptIdentifyChanges + .replace("{USER_INITIAL_REQUEST}", userInitialRequest) .replace("{USER_REQUEST}", userRequest) - .replace("{PREVIOUS_PLAN}", previousPlan); + .replace( + "{PLAN}", + previousPlan.map((plan, index) => `${index}: ${plan}`).join("\n"), + ); }; +const formatSysPromptRewritePlan = ( + userInitialRequest: string, + userRequest: string, + previousPlan: string[], + tasksToModify: PlanItem[], +) => { + return systemPrompt + .replace("{USER_INITIAL_REQUEST}", userInitialRequest) + .replace("{USER_REQUEST}", userRequest) + .replace( + "{PLAN}", + previousPlan.map((plan, index) => `${index}: ${plan}`).join("\n"), + ) + .replace( + "{TASKS_TO_MODIFY}", + tasksToModify.map((p) => `${p.index}: ${p.plan}`).join("\n"), + ); +}; + +async function identifyTasksToModifyFunc( + state: GraphState, + model: ConfigurableModel, +): Promise { + if (!state.planChangeRequest) { + throw new Error("No plan change request found."); + } + + const identifyPlanChangesSchema = z.object({ + task_change_indices: z + .array(z.number()) + .describe( + "The indices of the tasks in the plan that should be modified/removed.", + ), + }); + + const identifyPlanChangesTool = tool( + (input): PlanItem[] => { + const { task_change_indices } = input; + const tasksToModify = state.proposedPlan.flatMap((plan, planIndex) => { + const planItem = task_change_indices.some( + (changeIndex) => changeIndex === planIndex, + ); + if (!planItem) { + return []; + } + return { + index: planIndex, + plan: plan, + completed: false, + }; + }); + + return tasksToModify; + }, + { + name: "identify_plan_changes", + schema: identifyPlanChangesSchema, + description: + "Identify which tasks in the plan should be modified/removed based on the user's request.", + }, + ); + + const modelWithIdentifyChangesTool = model.bindTools( + [identifyPlanChangesTool], + { + // The model should always call the tool when identifying plan changes. + tool_choice: identifyPlanChangesTool.name, + }, + ); + + const firstUserMessage = state.messages.find(isHumanMessage); + + const response = await modelWithIdentifyChangesTool.invoke([ + { + role: "user", + content: formatSysPromptIdentifyTasks( + getMessageContentString( + firstUserMessage?.content ?? "No user message found", + ), + state.planChangeRequest, + state.proposedPlan, + ), + }, + ]); + + const toolCall = response.tool_calls?.[0]; + if (!toolCall) { + throw new Error( + "Tool call not returned when attempting to identify plan changes.", + ); + } + + const tasksToModify = await identifyPlanChangesTool.invoke( + toolCall.args as z.infer, + ); + return tasksToModify; +} + +const identifyTasksToModify = traceable(identifyTasksToModifyFunc, { + name: "identify_tasks_to_modify", +}); + +async function updatePlanTasksFunc( + state: GraphState, + tasksToModify: PlanItem[], + model: ConfigurableModel, +): Promise { + if (!state.planChangeRequest) { + throw new Error("No plan change request found."); + } + + const updatePlanTasksSchema = z.object({ + updated_plan_tasks: z + .array( + z + .string() + .describe( + "The updated or unmodified plan for the task. Do NOT include the task index.", + ), + ) + .describe( + "The updated plan tasks. Must be in the order of which they should be executed in.", + ), + }); + const updatePlanTasksTool = { + name: "update_plan_tasks", + description: "Call this tool to respond with the updated plan.", + schema: updatePlanTasksSchema, + }; + + const modelWithUpdatePlanTasksTool = model.bindTools([updatePlanTasksTool], { + // The model should always call the tool when identifying plan changes. + tool_choice: updatePlanTasksTool.name, + }); + + const firstUserMessage = state.messages.find(isHumanMessage); + + const response = await modelWithUpdatePlanTasksTool.invoke([ + { + role: "user", + content: formatSysPromptRewritePlan( + getMessageContentString( + firstUserMessage?.content ?? "No user message found", + ), + state.planChangeRequest, + state.proposedPlan, + tasksToModify, + ), + }, + ]); + + const toolCall = response.tool_calls?.[0]; + if (!toolCall) { + throw new Error( + "Tool call not returned when attempting to update plan tasks.", + ); + } + + return ( + toolCall.args as z.infer + ).updated_plan_tasks.map((p) => p); +} + +const updatePlanTasks = traceable(updatePlanTasksFunc, { + name: "update_plan_tasks", +}); + export async function rewritePlan( state: GraphState, config: GraphConfig, @@ -37,28 +239,11 @@ export async function rewritePlan( } const model = await loadModel(config, Task.PLANNER); - const modelWithTools = model.bindTools([sessionPlanTool], { - // The model should always call the tool when rewriting the plan. - tool_choice: sessionPlanTool.name, - }); + const tasksToModify = await identifyTasksToModify(state, model); + const updatedPlanTasks = await updatePlanTasks(state, tasksToModify, model); - const response = await modelWithTools.invoke([ - { - role: "system", - content: formatSysPrompt( - state.planChangeRequest, - " - " + state.plan.join("\n - "), - ), - }, - ...state.messages, - ]); - - if (response.tool_calls?.length) { - return { - proposedPlan: response.tool_calls[0].args.plan, - plan: [], - }; - } - - throw new Error("Failed to rewrite plan."); + return { + plan: [], + proposedPlan: updatedPlanTasks, + }; } diff --git a/apps/open-swe/src/nodes/summarize-task-steps.ts b/apps/open-swe/src/nodes/summarize-task-steps.ts index 737c8df3..d6bf6274 100644 --- a/apps/open-swe/src/nodes/summarize-task-steps.ts +++ b/apps/open-swe/src/nodes/summarize-task-steps.ts @@ -25,7 +25,12 @@ You do not want to keep the entire conversation history, but instead you want to {PLAN_PROMPT} You MUST adhere to the following criteria when summarizing the conversation history: -- Retain context such as file paths, versions, and installed software. +- Retain context such as file paths, versions, and installed software which future iterations will find useful. + - It is very important to include the file paths of files you've already searched for, along with a description of the file's contents, inside a 'Codebase files and descriptions' section, so that future steps can reuse this information, and will not need to search through the codebase for files again. + - Consider including a section titled 'Key repository insights and learnings' which may include information, insights and learnings you've discovered while completing the task. + - This section should be concise, but still including enough information so following steps will not repeat any mistakes or go down rabbit holes which you already know about. + - If changes were made to the repository during this task, ensure you include a section titled 'Repository modifications summary' which should include a short description of the task it completed, how it did so, and every change it made to the codebase during this task. + - Do not include the actual changes you made, but rather high level bullet points containing context and descriptions on the modifications made. - Do not retain any full code snippets. - Do not retain any full file contents. - Ensure your summary is concise, but useful for future context. diff --git a/apps/open-swe/src/subgraphs/planner/nodes/generate-message.ts b/apps/open-swe/src/subgraphs/planner/nodes/generate-message.ts index c7801349..1475e729 100644 --- a/apps/open-swe/src/subgraphs/planner/nodes/generate-message.ts +++ b/apps/open-swe/src/subgraphs/planner/nodes/generate-message.ts @@ -29,7 +29,7 @@ export async function generateAction( const firstUserMessage = state.messages.find(isHumanMessage); const response = await modelWithTools - .bind({ tags: ["langsmith:nostream"] }) + .withConfig({ tags: ["langsmith:nostream"] }) .invoke([ { role: "system", diff --git a/apps/open-swe/src/subgraphs/planner/nodes/generate-plan.ts b/apps/open-swe/src/subgraphs/planner/nodes/generate-plan.ts index b84edf08..f5c984bf 100644 --- a/apps/open-swe/src/subgraphs/planner/nodes/generate-plan.ts +++ b/apps/open-swe/src/subgraphs/planner/nodes/generate-plan.ts @@ -14,9 +14,11 @@ const systemPrompt = `You are operating as a terminal-based agentic coding assis In this step, you are expected to generate a high-level plan to address the user's request. The plan should be a list of actions to take, in order, to address the user's request. You should not include any code in the plan, only a list of actions to take. You MUST adhere to the following criteria when generating the plan: -- You have already gathered context from the repository the user has requested you take actions on. This context is provided in the conversation history below. +- You have already gathered context from the repository the user has requested you take actions on, and are now ready to generate a plan based on it. + - This context is provided in the conversation history below. - Your plan should be high-level in nature, but should still be specific enough to be actionable. -- Ensure your plan is as concise as possible. Omit any unnecessary details or steps. Your goal is to complete the task in the least number of steps possible. +- Ensure your plan is as concise as possible. Omit any unnecessary details or steps the user did not request, or are not required to complete the task. + - Your goal is to complete the task outlined by the user in the least number of steps possible. - Do not pack multiple complex tasks into a single plan item. Each high level task you'll need to complete should have its own plan item. - When you are ready to generate the plan, ensure you call the 'session_plan' tool. You are REQUIRED to call this tool. - The first user message in this conversation contains the user's request. @@ -45,7 +47,7 @@ export async function generatePlan( } const response = await modelWithTools - .bind({ tags: ["langsmith:nostream"] }) + .withConfig({ tags: ["langsmith:nostream"] }) .invoke([ { role: "system", diff --git a/apps/open-swe/src/subgraphs/planner/nodes/summarizer.ts b/apps/open-swe/src/subgraphs/planner/nodes/summarizer.ts index 014c33ff..f2b2f96f 100644 --- a/apps/open-swe/src/subgraphs/planner/nodes/summarizer.ts +++ b/apps/open-swe/src/subgraphs/planner/nodes/summarizer.ts @@ -19,6 +19,9 @@ You do not want to keep the entire conversation history, but instead you want to You MUST adhere to the following criteria when summarizing the conversation history: - Retain context such as file paths, versions, and installed software. + - It is very important to include the file paths of files you've already searched for, along with a description of the file's contents, inside a 'Codebase files and descriptions' section, so that future steps can reuse this information, and will not need to search through the codebase for files again. +- Consider including a section titled 'Key repository insights and learnings' which may include information, insights and learnings you've discovered while gathering context for the user's request. + - This section should be concise, but still including enough information so following steps will not repeat any mistakes or go down rabbit holes which you already know about. - Do not retain any full code snippets. - Do not retain any full file contents. - Ensure your summary is concise, but useful for future context. diff --git a/apps/open-swe/src/tools/shell.ts b/apps/open-swe/src/tools/shell.ts index ba87a0bc..1cf41d10 100644 --- a/apps/open-swe/src/tools/shell.ts +++ b/apps/open-swe/src/tools/shell.ts @@ -12,11 +12,13 @@ const logger = createLogger(LogLevel.INFO, "ShellTool"); const DEFAULT_COMMAND_TIMEOUT = 60_000; // 1 minute const shellToolSchema = z.object({ - command: z.array(z.string()).describe("The command to run"), workdir: z .string() .optional() - .describe("The working directory for the command."), + .describe( + "The working directory for the command. Ensure this path is NOT included in any command arguments, as it will be added automatically.", + ), + command: z.array(z.string()).describe("The command to run"), timeout: z .number() .optional() diff --git a/yarn.lock b/yarn.lock index 11a744dd..512b2dcd 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2172,6 +2172,7 @@ __metadata: eslint-plugin-prettier: ^4.2.1 jest: ^29.7.0 langchain: ^0.3.26 + langsmith: ^0.3.29 prettier: ^3.5.2 ts-jest: ^29.1.0 typescript: ~5.7.2