diff --git a/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts b/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts index eaabc42e..57864435 100644 --- a/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts +++ b/apps/open-swe/src/graphs/planner/nodes/generate-message/index.ts @@ -18,6 +18,7 @@ import { filterHiddenMessages } from "../../../../utils/message/filter-hidden.js import { getTaskPlanFromIssue } from "../../../../utils/github/issue-task.js"; import { createRgTool } from "../../../../tools/rg.js"; import { formatCustomRulesPrompt } from "../../../../utils/custom-rules.js"; +import { createPlannerNotesTool } from "../../../../tools/planner-notes.js"; const logger = createLogger(LogLevel.INFO, "GeneratePlanningMessageNode"); @@ -46,7 +47,11 @@ export async function generateAction( config: GraphConfig, ): Promise { const model = await loadModel(config, Task.ACTION_GENERATOR); - const tools = [createRgTool(state), createShellTool(state)]; + const tools = [ + createRgTool(state), + createShellTool(state), + createPlannerNotesTool(), + ]; const modelWithTools = model.bindTools(tools, { tool_choice: "auto", parallel_tool_calls: true, diff --git a/apps/open-swe/src/graphs/planner/nodes/generate-plan.ts b/apps/open-swe/src/graphs/planner/nodes/generate-plan.ts index 6fab293e..cb96a84f 100644 --- a/apps/open-swe/src/graphs/planner/nodes/generate-plan.ts +++ b/apps/open-swe/src/graphs/planner/nodes/generate-plan.ts @@ -15,6 +15,13 @@ import { stopSandbox } from "../../../utils/sandbox.js"; import { filterHiddenMessages } from "../../../utils/message/filter-hidden.js"; import { z } from "zod"; import { formatCustomRulesPrompt } from "../../../utils/custom-rules.js"; +import { getPlannerNotes } from "../utils/get-notes.js"; + +const PLANNER_NOTES_PROMPT = `Here is a collection of technical notes you took while gathering context for the plan. Ensure you take these into account when writing your plan. + + +{PLANNER_NOTES} +`; const systemPrompt = `You are a terminal-based agentic coding assistant built by LangChain, designed to enable natural language interaction with local codebases through wrapped LLM models. @@ -65,13 +72,17 @@ Structure your plan items as clear directives, for example: {CUSTOM_RULES} +{PLANNER_NOTES} + Remember: Your goal is to create a focused, executable plan that efficiently accomplishes the user's request using the context you've already gathered.`; function formatSystemPrompt(state: PlannerGraphState): string { // It's a followup if there's more than one human message. const isFollowup = isFollowupRequest(state.taskPlan, state.proposedPlan); const userRequest = getUserRequest(state.messages); - + const plannerNotes = getPlannerNotes(state.messages) + .map((n) => `- ${n}`) + .join("\n"); return systemPrompt .replace( "{FOLLOWUP_MESSAGE_PROMPT}", @@ -82,7 +93,13 @@ function formatSystemPrompt(state: PlannerGraphState): string { : "", ) .replace("{USER_REQUEST}", userRequest) - .replaceAll("{CUSTOM_RULES}", formatCustomRulesPrompt(state.customRules)); + .replaceAll("{CUSTOM_RULES}", formatCustomRulesPrompt(state.customRules)) + .replaceAll( + "{PLANNER_NOTES}", + plannerNotes.length + ? PLANNER_NOTES_PROMPT.replace("{PLANNER_NOTES}", plannerNotes) + : "", + ); } export async function generatePlan( diff --git a/apps/open-swe/src/graphs/planner/nodes/notetaker.ts b/apps/open-swe/src/graphs/planner/nodes/notetaker.ts index cee40cdd..4d99094f 100644 --- a/apps/open-swe/src/graphs/planner/nodes/notetaker.ts +++ b/apps/open-swe/src/graphs/planner/nodes/notetaker.ts @@ -1,5 +1,5 @@ import { z } from "zod"; -import { CustomRules, GraphConfig } from "@open-swe/shared/open-swe/types"; +import { GraphConfig } from "@open-swe/shared/open-swe/types"; import { PlannerGraphState, PlannerGraphUpdate, @@ -7,8 +7,16 @@ import { import { loadModel, Task } from "../../../utils/load-model.js"; import { getMessageString } from "../../../utils/message/content.js"; import { getUserRequest } from "../../../utils/user-request.js"; -import { BaseMessage } from "@langchain/core/messages"; import { formatCustomRulesPrompt } from "../../../utils/custom-rules.js"; +import { getPlannerNotes } from "../utils/get-notes.js"; + +const PLANNER_NOTES_PROMPT = `You've also taken technical notes throughout the context gathering process. Ensure you include/incorporate these notes, or the highest quality parts of these notes in your conclusion notes. + + +{PLANNER_NOTES} +`; +const CUSTOM_RULES_EXTRA_CONTEXT = + "- Carefully read over the user's custom rules to ensure you don't duplicate or repeat information found in that section, as you will always have access to it (even after the planning step!)."; 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. @@ -19,12 +27,15 @@ The notes you extract should be thoughtful, and should include technical details These notes should not be overly verbose, as you'll be able to gather additional context when executing. Your goal is to generate notes on all of the low-hanging fruit from the conversation history, to speed up the execution so that you don't need to duplicate work to gather context. +{CUSTOM_RULES} + +{PLANNER_NOTES} + You MUST adhere to the following criteria when generating your notes: - Do not retain any full code snippets. - Do not retain any full file contents. - Only take notes on the context provided below, and do not make up, or attempt to infer any information/context which is not explicitly provided. - -{CUSTOM_RULES} +{EXTRA_RULES} Here is the user's request ## User request: @@ -41,20 +52,41 @@ And here is the plan you just generated: With all of this in mind, please carefully inspect the conversation history, and the plan you generated. Then, determine which actions and context from the conversation history will be most useful to you when you execute the plan. After you're done analyzing, call the \`write_technical_notes\` tool. `; -const formatPrompt = ( - userRequest: string, - conversationHistory: BaseMessage[], - proposedPlan: string[], - customRules?: CustomRules, -): string => - systemPrompt +const formatPrompt = (state: PlannerGraphState): string => { + const userRequest = + getUserRequest(state.messages) || "No user request provided."; + const plannerNotes = getPlannerNotes(state.messages) + .map((n) => ` - ${n}`) + .join("\n"); + + return systemPrompt .replace("{USER_REQUEST}", userRequest) .replace( "{CONVERSATION_HISTORY}", - conversationHistory.map(getMessageString).join("\n"), + state.messages.map(getMessageString).join("\n"), ) - .replace("{PROPOSED_PLAN}", ` - ${proposedPlan.join("\n - ")}`) - .replaceAll("{CUSTOM_RULES}", formatCustomRulesPrompt(customRules)); + .replace( + "{PROPOSED_PLAN}", + state.proposedPlan.map((p) => ` - ${p}`).join("\n"), + ) + .replaceAll( + "{CUSTOM_RULES}", + formatCustomRulesPrompt( + state.customRules, + "Keep in mind these user provided rules will always be available to you, so any context present here should NOT be included in your notes as to not duplicate information.", + ), + ) + .replaceAll( + "{PLANNER_NOTES}", + plannerNotes.length + ? PLANNER_NOTES_PROMPT.replace("{PLANNER_NOTES}", plannerNotes) + : "", + ) + .replaceAll( + "{EXTRA_RULES}", + state.customRules ? CUSTOM_RULES_EXTRA_CONTEXT : "", + ); +}; const condenseContextToolSchema = z.object({ notes: z @@ -78,7 +110,6 @@ export async function notetaker( parallel_tool_calls: false, }); - const userRequest = getUserRequest(state.messages); const conversationHistoryStr = `Here is the full conversation history: ${state.messages.map(getMessageString).join("\n")}`; @@ -86,12 +117,7 @@ ${state.messages.map(getMessageString).join("\n")}`; const response = await modelWithTools.invoke([ { role: "system", - content: formatPrompt( - userRequest || "No user request provided.", - state.messages, - state.proposedPlan, - state.customRules, - ), + content: formatPrompt(state), }, { role: "user", diff --git a/apps/open-swe/src/graphs/planner/nodes/take-action.ts b/apps/open-swe/src/graphs/planner/nodes/take-action.ts index 587b35c0..3606c409 100644 --- a/apps/open-swe/src/graphs/planner/nodes/take-action.ts +++ b/apps/open-swe/src/graphs/planner/nodes/take-action.ts @@ -10,6 +10,13 @@ import { zodSchemaToString } from "../../../utils/zod-to-string.js"; import { formatBadArgsError } from "../../../utils/zod-to-string.js"; import { truncateOutput } from "../../../utils/truncate-outputs.js"; import { createRgTool } from "../../../tools/rg.js"; +import { + getChangedFilesStatus, + stashAndClearChanges, +} from "../../../utils/github/git.js"; +import { getRepoAbsolutePath } from "@open-swe/shared/git"; +import { daytonaClient } from "../../../utils/sandbox.js"; +import { createPlannerNotesTool } from "../../../tools/planner-notes.js"; const logger = createLogger(LogLevel.INFO, "TakeAction"); @@ -26,9 +33,11 @@ export async function takeActions( const shellTool = createShellTool(state); const rgTool = createRgTool(state); + const plannerNotesTool = createPlannerNotesTool(); const toolsMap = { [shellTool.name]: shellTool, [rgTool.name]: rgTool, + [plannerNotesTool.name]: plannerNotesTool, }; const toolCalls = lastMessage.tool_calls; @@ -96,7 +105,34 @@ export async function takeActions( return toolMessage; }); - const toolCallResults = await Promise.all(toolCallResultsPromise); + let toolCallResults = await Promise.all(toolCallResultsPromise); + const sandbox = await daytonaClient().get(state.sandboxSessionId); + const repoPath = getRepoAbsolutePath(state.targetRepository); + const changedFiles = await getChangedFilesStatus(repoPath, sandbox); + if (changedFiles?.length > 0) { + logger.warn( + "Changes found in the codebase after taking action. Reverting.", + { + changedFiles, + }, + ); + await stashAndClearChanges(repoPath, sandbox); + + // Rewrite the tool call contents to include a changed files warning. + toolCallResults = toolCallResults.map( + (tc) => + new ToolMessage({ + ...tc, + content: `**WARNING**: THIS TOOL, OR A PREVIOUS TOOL HAS CHANGED FILES IN THE REPO. +Remember that you are only permitted to take **READ** actions during the planning step. The changes have been reverted. + +Please ensure you only take read actions during the planning step to gather context. You may also call the \`take_notes\` tool at any time to record important information for the programmer step. + +Command Output:\n +${tc.content}`, + }), + ); + } logger.info("Completed planner tool action", { ...toolCallResults.map((tc) => ({ diff --git a/apps/open-swe/src/graphs/planner/utils/get-notes.ts b/apps/open-swe/src/graphs/planner/utils/get-notes.ts new file mode 100644 index 00000000..b206a8c0 --- /dev/null +++ b/apps/open-swe/src/graphs/planner/utils/get-notes.ts @@ -0,0 +1,22 @@ +import { BaseMessage, isAIMessage } from "@langchain/core/messages"; +import { createTakePlannerNotesFields } from "@open-swe/shared/open-swe/tools"; +import z from "zod"; + +export function getPlannerNotes(messages: BaseMessage[]): string[] { + const plannerNotesFields = createTakePlannerNotesFields(); + const plannerNotes = messages.flatMap((m) => { + if (!isAIMessage(m)) { + return []; + } + const notesToolCalls = m.tool_calls?.filter( + (tc) => tc.name === plannerNotesFields.name, + ); + if (!notesToolCalls?.length) { + return []; + } + return notesToolCalls.map( + (tc) => (tc.args as z.infer).notes, + ); + }); + return plannerNotes.flat(); +} diff --git a/apps/open-swe/src/tools/planner-notes.ts b/apps/open-swe/src/tools/planner-notes.ts new file mode 100644 index 00000000..1d41b19b --- /dev/null +++ b/apps/open-swe/src/tools/planner-notes.ts @@ -0,0 +1,19 @@ +import { tool } from "@langchain/core/tools"; +import { createTakePlannerNotesFields } from "@open-swe/shared/open-swe/tools"; + +export function createPlannerNotesTool() { + const plannerNotesTool = tool( + async ( + _input, + ): Promise<{ result: string; status: "success" | "error" }> => { + // TODO: This should write to saved state once that feature is released in LangGraph. + return { + result: "Successfully saved notes. Thank you!", + status: "success", + }; + }, + createTakePlannerNotesFields(), + ); + + return plannerNotesTool; +} diff --git a/apps/open-swe/src/utils/custom-rules.ts b/apps/open-swe/src/utils/custom-rules.ts index e3b8ba52..46bff3b6 100644 --- a/apps/open-swe/src/utils/custom-rules.ts +++ b/apps/open-swe/src/utils/custom-rules.ts @@ -26,7 +26,9 @@ const ALL_TAGS = [ TESTING_INSTRUCTIONS_CLOSE_TAG, ]; -export function parseCustomRulesFromString(contents: string): CustomRules { +export function parseCustomRulesFromString( + contents: string, +): CustomRules | undefined { if (ALL_TAGS.every((tag) => !contents.includes(tag))) { // Text file has no custom rules. Return all as general rules return { @@ -78,6 +80,15 @@ export function parseCustomRulesFromString(contents: string): CustomRules { ); } + if ( + !generalRules && + !repositoryStructure && + !dependenciesAndInstallation && + !testingInstructions + ) { + return undefined; + } + return { generalRules, repositoryStructure, @@ -129,20 +140,28 @@ export async function getCustomRules( export const CUSTOM_RULES_PROMPT = ` The following are custom rules provided by the user. +{EXTRA_CONTEXT} {GENERAL_RULES} {REPOSITORY_STRUCTURE} {TESTING_INSTRUCTIONS} {DEPENDENCIES_AND_INSTALLATION} `; -export function formatCustomRulesPrompt(customRules?: CustomRules): string { +export function formatCustomRulesPrompt( + customRules?: CustomRules, + extraContextStr?: string, +): string { if (!customRules) return ""; return CUSTOM_RULES_PROMPT.replace( - "{GENERAL_RULES}", - customRules.generalRules - ? `\n${customRules.generalRules}\n` - : "", + "{EXTRA_CONTEXT}", + extraContextStr ? extraContextStr : "", ) + .replace( + "{GENERAL_RULES}", + customRules.generalRules + ? `\n${customRules.generalRules}\n` + : "", + ) .replace( "{REPOSITORY_STRUCTURE}", customRules.repositoryStructure diff --git a/apps/open-swe/src/utils/github/git.ts b/apps/open-swe/src/utils/github/git.ts index 6187a66f..6b3fc32b 100644 --- a/apps/open-swe/src/utils/github/git.ts +++ b/apps/open-swe/src/utils/github/git.ts @@ -380,6 +380,38 @@ export async function getChangedFilesStatus( .filter((line) => line !== ""); } +export async function stashAndClearChanges( + absoluteRepoDir: string, + sandbox: Sandbox, +): Promise { + try { + const gitStashOutput = await sandbox.process.executeCommand( + "git stash && git reset --hard", + absoluteRepoDir, + undefined, + TIMEOUT_SEC, + ); + + if (gitStashOutput.exitCode !== 0) { + logger.error(`Failed to stash and clear changes`, { + gitStashOutput, + }); + } + return gitStashOutput; + } catch (e) { + const errorFields = getSandboxErrorFields(e); + logger.error(`Failed to stash and clear changes`, { + ...(errorFields && { errorFields }), + ...(e instanceof Error && { + name: e.name, + message: e.message, + stack: e.stack, + }), + }); + return errorFields ?? false; + } +} + export async function checkoutBranchAndCommit( config: GraphConfig, targetRepository: TargetRepository, diff --git a/apps/web/src/components/gen-ui/action-step.tsx b/apps/web/src/components/gen-ui/action-step.tsx index e1357953..fc10c994 100644 --- a/apps/web/src/components/gen-ui/action-step.tsx +++ b/apps/web/src/components/gen-ui/action-step.tsx @@ -18,6 +18,7 @@ import { createApplyPatchToolFields, createShellToolFields, createInstallDependenciesToolFields, + createTakePlannerNotesFields, formatRgCommand, RipgrepCommand, } from "@open-swe/shared/open-swe/tools"; @@ -33,6 +34,8 @@ const installDependenciesTool = createInstallDependenciesToolFields(dummyRepo); type InstallDependenciesToolArgs = z.infer< typeof installDependenciesTool.schema >; +const plannerNotesTool = createTakePlannerNotesFields(); +type PlannerNotesToolArgs = z.infer; // Common props for all action types type BaseActionProps = { @@ -73,12 +76,18 @@ type InstallDependenciesActionProps = BaseActionProps & errorCode?: number; }; +type PlannerNotesActionProps = BaseActionProps & + Partial & { + actionType: "planner_notes"; + }; + export type ActionItemProps = | (BaseActionProps & { status: "loading" }) | ShellActionProps | PatchActionProps | RgActionProps - | InstallDependenciesActionProps; + | InstallDependenciesActionProps + | PlannerNotesActionProps; export type ActionStepProps = { actions: ActionItemProps[]; @@ -126,6 +135,8 @@ function ActionItem(props: ActionItemProps) { return props.success ? "Search completed" : "Search failed"; } else if (props.actionType === "install_dependencies") { return props.success ? "Dependencies installed" : "Installation failed"; + } else if (props.actionType === "planner_notes") { + return props.success ? "Notes saved" : "Failed to save notes"; } } @@ -144,6 +155,8 @@ function ActionItem(props: ActionItemProps) { return !!props.output; } else if (props.actionType === "apply-patch") { return !!props.diff; + } else if (props.actionType === "planner_notes") { + return !!(props.notes && props.notes.length > 0); } return false; @@ -156,7 +169,9 @@ function ActionItem(props: ActionItemProps) { return ; } - if (props.actionType === "install_dependencies") { + if (props.actionType === "planner_notes") { + return ; + } else if (props.actionType === "install_dependencies") { return ; } else if (props.actionType === "apply-patch") { return ; @@ -177,6 +192,16 @@ function ActionItem(props: ActionItemProps) { ); } + if (props.actionType === "planner_notes") { + return ( +
+ + Planner Notes + +
+ ); + } + if ( props.actionType === "shell" || props.actionType === "install_dependencies" @@ -289,6 +314,25 @@ function ActionItem(props: ActionItemProps) { )} ); + } else if ( + props.actionType === "planner_notes" && + props.notes && + props.notes.length > 0 + ) { + return ( +
+
    + {props.notes.map((note, i) => ( +
  • + {note} +
    • + ))} +
+
+ ); } return null; diff --git a/apps/web/src/components/thread/messages/ai.tsx b/apps/web/src/components/thread/messages/ai.tsx index 9ee3db61..70805d88 100644 --- a/apps/web/src/components/thread/messages/ai.tsx +++ b/apps/web/src/components/thread/messages/ai.tsx @@ -29,6 +29,7 @@ import { createRgToolFields, createOpenPrToolFields, createInstallDependenciesToolFields, + createTakePlannerNotesFields, } from "@open-swe/shared/open-swe/tools"; import { z } from "zod"; import { isAIMessageSDK, isToolMessageSDK } from "@/lib/langchain-messages"; @@ -50,6 +51,8 @@ const installDependenciesTool = createInstallDependenciesToolFields(dummyRepo); type InstallDependenciesToolArgs = z.infer< typeof installDependenciesTool.schema >; +const plannerNotesTool = createTakePlannerNotesFields(); +type PlannerNotesToolArgs = z.infer; function CustomComponent({ message, @@ -168,6 +171,15 @@ export function mapToolMessageToActionStepProps( output: getContentString(message.content), reasoningText, }; + } else if (toolCall?.name === plannerNotesTool.name) { + const args = toolCall.args as PlannerNotesToolArgs; + return { + actionType: "planner_notes", + status, + success, + notes: args.notes || [], + reasoningText, + }; } return { status: "loading", @@ -231,7 +243,8 @@ export function AssistantMessage({ tc.name === shellTool.name || tc.name === applyPatchTool.name || tc.name === rgTool.name || - tc.name === installDependenciesTool.name, + tc.name === installDependenciesTool.name || + tc.name === plannerNotesTool.name, ) : []; @@ -350,6 +363,13 @@ export function AssistantMessage({ workdir: args?.workdir || "", output: "", } as ActionItemProps; + } else if (toolCall.name === plannerNotesTool.name) { + const args = toolCall.args as PlannerNotesToolArgs; + return { + actionType: "planner_notes", + status: "generating", + notes: args?.notes || [], + } as ActionItemProps; } else { if (isShellTool) { const args = toolCall.args as ShellToolArgs; diff --git a/packages/shared/src/open-swe/tools.ts b/packages/shared/src/open-swe/tools.ts index 6463ef83..352d2c94 100644 --- a/packages/shared/src/open-swe/tools.ts +++ b/packages/shared/src/open-swe/tools.ts @@ -231,3 +231,22 @@ export function createOpenPrToolFields() { description: "Use this tool to open a pull request.", }; } + +export function createTakePlannerNotesFields() { + const plannerContextNotesSchema = z.object({ + notes: z + .array(z.string()) + .describe( + "The high quality, concise and technical notes you deem important to save for the programmer to use when implementing the plan.", + ), + }); + + return { + name: "take_notes", + schema: plannerContextNotesSchema, + description: + "Use this tool to write & save technical notes on the planner context.\n" + + "This should be called if you come across context you think will be highly useful to the programmer when they are actually implementing the plan, and you want to ensure it's not lost.\n" + + "Do not duplicate any information present in the user provided 'custom rules', as we want to avoid duplicating context.", + }; +}