fix: Better prompting and context provided for codebase structure awareness (#85)

* fix: Better prompting and context provided for codebase structure awareness

* cr

apps/open-swe/src/nodes/diagnose-error.ts

@@ -40,8 +40,9 @@ Here is the current task you're working on:
 And here are all of the tasks you've completed so far, along with their summaries:
 {PLAN_PROMPT}
 
-Finally, here is a summary of general context about the codebase:
-{CODEBASE_CONTEXT}
+Below is an up to date tree of the codebase (going 3 levels deep). This is up to date, and is updated after every action you take. Always assume this is the most up to date context about the codebase.
+It was generated by using the \`tree\` command, passing in the gitignore file to ignore files and directories you should not have access to (\`git ls-files | tree --fromfile -L 3\`). It is always executed inside the repo directory: {REPO_DIRECTORY}
+{CODEBASE_TREE}
 
 Please carefully go over all of this information, and provide a helpful explanation of exactly what the issue is, and how you can fix it. When you are ready to provide your diagnosis, call the \`diagnose_error\` tool.
 `;
@@ -65,7 +66,7 @@ const diagnoseErrorTool = {
 const formatSystemPrompt = (
   lastFailedActionContent: string,
   plan: PlanItem[],
-  codebaseContext: string,
+  codebaseTree: string,
 ): string => {
   const currentPlanItem = getCurrentPlanItem(plan);
   const completedTasks = getCompletedPlanItems(plan);
@@ -80,7 +81,10 @@ const formatSystemPrompt = (
       `<current-task index="${currentPlanItem.index}">${currentPlanItem.plan}</current-task>`,
     )
     .replace("{PLAN_PROMPT}", formatPlanPromptWithSummaries(completedTasks))
-    .replace("{CODEBASE_CONTEXT}", codebaseContext);
+    .replace(
+      "{CODEBASE_TREE}",
+      `<codebase-tree>\n${codebaseTree || "No codebase tree generated yet."}\n</codebase-tree>`,
+    );
 };
 
 const formatUserPrompt = (messages: BaseMessage[]): string => {
@@ -114,7 +118,7 @@ export async function diagnoseError(
       content: formatSystemPrompt(
         getMessageContentString(lastFailedAction.content),
         getActivePlanItems(state.plan),
-        state.codebaseContext,
+        state.codebaseTree,
       ),
     },
     {

apps/open-swe/src/nodes/generate-message.ts

@@ -12,6 +12,7 @@ import { createLogger, LogLevel } from "../utils/logger.js";
 import { getCurrentPlanItem } from "../utils/current-task.js";
 import { getMessageContentString } from "../utils/message/content.js";
 import { getActivePlanItems } from "../utils/task-plan.js";
+import { SANDBOX_ROOT_DIR } from "../constants.js";
 
 const logger = createLogger(LogLevel.INFO, "GenerateMessageNode");
 
@@ -52,7 +53,7 @@ You MUST adhere to the following criteria when executing the task:
 - Take advantage of the task summaries from completed tasks in the prompt above. Ensure you always read these summaries to avoid duplicate work, and so you always have up to date context on the codebase, and tasks you've completed.
   - Each summary message will 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'.
   - 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.
-- Additionally, you're also provided with a section titled 'Codebase context' which contains an up to date list of files, and descriptions of the files' contents. If you need context on a file, or directory, ensure you first check to see if you can find it in the codebase context before performing an action to read/find it.
+- Additionally, you're also provided with a section titled 'Codebase tree' which contains an up to date list of files three levels deep, ignoring gitignore.
 - 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.
 - You should NOT try to create empty files with \`apply_patch\`. If you need to create a file, use the \`shell\` tool, and pass \`touch <file path>\` to create the file.
@@ -92,8 +93,11 @@ You MUST adhere to the following criteria when executing the task:
 - Only make changes to the existing Git repo ({REPO_DIRECTORY}). Any changes outside this repo will not be detected, so do not attempt to create new files or directories outside of this repo.
 - You do NOT have access to the \`set_task_status\` or \`diagnose_error\` tools. NEVER attempt to call them.
 
-Below, is a collection of useful context about the codebase. It is updated after each completed task, and is provided to you to help you make decisions, and avoid duplicate work:
-{CODEBASE_CONTEXT}
+Below is an up to date tree of the codebase (going 3 levels deep). This is up to date, and is updated after every action you take. Always assume this is the most up to date context about the codebase.
+It was generated by using the \`tree\` command, passing in the gitignore file to ignore files and directories you should not have access to (\`git ls-files | tree --fromfile -L 3\`). It is always executed inside the repo directory: {REPO_DIRECTORY}
+{CODEBASE_TREE}
+
+Your current working directory is: {CURRENT_WORKING_DIRECTORY}
 
 Once again, here are the completed tasks, remaining tasks, and the current task you're working on:
 {PLAN_PROMPT}
@@ -118,9 +122,10 @@ const formatPrompt = (state: GraphState): string => {
       `<plan-generation-summary>\n${state.planContextSummary}\n</plan-generation-summary>`,
     )
     .replaceAll(
-      "{CODEBASE_CONTEXT}",
-      `<codebase-context>\n${state.codebaseContext || "No codebase context generated yet. Please use the conversation below as context."}\n</codebase-context>`,
-    );
+      "{CODEBASE_TREE}",
+      `<codebase-tree>\n${state.codebaseTree || "No codebase tree generated yet."}\n</codebase-tree>`,
+    )
+    .replaceAll("{CURRENT_WORKING_DIRECTORY}", SANDBOX_ROOT_DIR);
 };
 
 export async function generateAction(

apps/open-swe/src/nodes/initialize.ts

@@ -11,6 +11,7 @@ import {
 import { daytonaClient } from "../utils/sandbox.js";
 import { SNAPSHOT_NAME } from "../constants.js";
 import { getGitHubTokensFromConfig } from "../utils/github-tokens.js";
+import { getCodebaseTree } from "../utils/tree.js";
 
 const logger = createLogger(LogLevel.INFO, "Initialize");
 
@@ -36,8 +37,10 @@ export async function initialize(
       // Resume the sandbox if the session ID is in the config.
       const existingSandbox = await daytonaClient().get(sandboxSessionId);
       await pullLatestChanges(absoluteRepoDir, existingSandbox);
+      const codebaseTree = await getCodebaseTree(existingSandbox.id);
       return {
         sandboxSessionId: existingSandbox.id,
+        codebaseTree,
       };
     } catch (e) {
       // Error thrown, log it and continue. Will create a new sandbox session since the resumption failed.
@@ -82,8 +85,11 @@ export async function initialize(
     throw new Error("Failed to checkout branch");
   }
 
+  const codebaseTree = await getCodebaseTree(sandbox.id);
+
   return {
     sandboxSessionId: sandbox.id,
     targetRepository,
+    codebaseTree,
   };
 }

apps/open-swe/src/nodes/summarize-task-steps.ts

@@ -28,7 +28,7 @@ Here are all of your tasks you've completed, remaining, and the current task you
 
 You MUST adhere to the following criteria when summarizing the conversation history:
   - Include insights, and learnings you've discovered about the codebase or specific files while completing the task.
-    - You should NOT document scripts, file structure, or other context which could be categorized as 'general codebase context'. General codebase context (e.g. scripts, file structure, package managers, etc.) will be generated in a separate step. Inspect the codebase context string provided below for this information.
+    - You should NOT document scripts, file structure, or other context which could be categorized as 'general codebase context'. General codebase context is automatically included via the \`tree\` command.
   - If files were created or modified, include short summaries of the changes made.
     - What file(s) were modified/created.
     - What content was added/removed.
@@ -41,9 +41,6 @@ You MUST adhere to the following criteria when summarizing the conversation hist
   - You do not need to include specific codebase context here, as codebase context will be generated in a separate step. Your sole task is to generate a concise summary of this specific task you just completed.
   - Ensure your summary is as concise as possible, but useful for future context.
 
-Here is the current state of the codebase context you've accumulated. Remember YOU SHOULD NOT INCLUDE ANY GENERAL CODEBASE CONTEXT IN YOUR TASK SUMMARY.
-{CODEBASE_CONTEXT}
-
 Ensure you do NOT include codebase context in your task summary, as we want to avoid including duplicate information.
 
 With all of this in mind, please carefully summarize and condense the conversation history of the task you just completed, provided by the user below. Remember that this summary should ONLY include details about the completed task, and should NOT include any general codebase context.
@@ -59,55 +56,16 @@ With this in mind, please use the following conversation history to generate a c
 Conversation history:
 {CONVERSATION_HISTORY}`;
 
-const updateCodebaseContextSysPrompt = `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.
-
-Your current task is to update the codebase context, given the recent actions taken by the agent.
-
-The codebase context should contain:
-  - Up to date information on the codebase file paths, and their contents.
-    - Do not include entire file contents, but rather high level descriptions of what a file contains, and what it does.
-  - Information on the software installed, and used in the codebase, including information such as version numbers, and dependencies.
-  - High level context about the codebase structure, and style.
-  - Any other relevant codebase information which may be useful for future context.
-  - There should be NO task specific context here. ONLY include context about the codebase. This context should be generally applicable and not tied to the specifics of the task.
-
-You have the following codebase context:
-{CODEBASE_CONTEXT}
-
-Please inspect this context, and given the rules above, please respond with a full, complete codebase context I can use for future context.
-When responding, ensure:
-  - You do not duplicate information.
-  - You remove old/stale context from the existing codebase context string if recent messages contradict it.
-  - You do NOT remove any information from the existing codebase context string if recent messages do not contradict it. We want to ensure we always have a complete picture of the codebase.
-  - You modify/combine information from the existing codebase context string if if new information is provided which warrants a change.
-
-Please be concise, clear and helpful. Omit any extraneous information. Respond ONLY with the codebase context. Do not include any additional information, or text before or after the codebase context.
-`;
-
-const updateCodebaseContextUserMessage = `Here is the task you just completed:
-{COMPLETED_TASK}
-
-The first message in the conversation history is the user's request. Messages from previously completed tasks have already been removed, in favor of task summaries.
-With this in mind, please use the following conversation history to update the codebase context to include new relevant information.
-
-Conversation history:
-{CONVERSATION_HISTORY}`;
-
 const logger = createLogger(LogLevel.INFO, "SummarizeTaskSteps");
 
-const formatPrompt = (plan: PlanItem[], codebaseContext: string): string =>
-  taskSummarySysPrompt
-    .replace(
-      "{PLAN_PROMPT}",
-      formatPlanPrompt(plan, {
-        useLastCompletedTask: true,
-        includeSummaries: true,
-      }),
-    )
-    .replace(
-      "{CODEBASE_CONTEXT}",
-      `<codebase-context>\n${codebaseContext || "No codebase context generated yet."}\n</codebase-context>`,
-    );
+const formatPrompt = (plan: PlanItem[]): string =>
+  taskSummarySysPrompt.replace(
+    "{PLAN_PROMPT}",
+    formatPlanPrompt(plan, {
+      useLastCompletedTask: true,
+      includeSummaries: true,
+    }),
+  );
 
 const formatUserMessage = (
   messages: BaseMessage[],
@@ -128,31 +86,6 @@ const formatUserMessage = (
     );
 };
 
-const formatCodebaseContextPrompt = (codebaseContext: string): string =>
-  updateCodebaseContextSysPrompt.replace(
-    "{CODEBASE_CONTEXT}",
-    `<codebase-context>\n${codebaseContext || "No codebase context generated yet."}\n</codebase-context>`,
-  );
-
-const formatUserCodebaseContextMessage = (
-  messages: BaseMessage[],
-  plans: PlanItem[],
-): string => {
-  const completedTask = plans.find((p) => p.completed);
-  if (!completedTask) {
-    throw new Error(
-      "No completed task found when trying to format user message for task summary.",
-    );
-  }
-
-  return updateCodebaseContextUserMessage
-    .replace("{COMPLETED_TASK}", completedTask.plan)
-    .replace(
-      "{CONVERSATION_HISTORY}",
-      messages.map(getMessageString).join("\n"),
-    );
-};
-
 async function generateTaskSummaryFunc(
   state: GraphState,
   model: ConfigurableModel,
@@ -167,7 +100,7 @@ async function generateTaskSummaryFunc(
   const response = await model.withConfig({ tags: ["nostream"] }).invoke([
     {
       role: "system",
-      content: formatPrompt(activePlanItems, state.codebaseContext),
+      content: formatPrompt(activePlanItems),
     },
     {
       role: "user",
@@ -182,33 +115,7 @@ async function generateTaskSummaryFunc(
 }
 
 const generateTaskSummary = traceable(generateTaskSummaryFunc, {
-  name: "generate_task_summary",
-});
-
-async function updateCodebaseContextFunc(
-  state: GraphState,
-  model: ConfigurableModel,
-): Promise<string> {
-  logger.info(`Updating codebase context...`);
-  const response = await model.withConfig({ tags: ["nostream"] }).invoke([
-    {
-      role: "system",
-      content: formatCodebaseContextPrompt(state.codebaseContext),
-    },
-    {
-      role: "user",
-      content: formatUserCodebaseContextMessage(
-        state.messages,
-        getActivePlanItems(state.plan),
-      ),
-    },
-  ]);
-  const contentString = getMessageContentString(response.content);
-  return contentString;
-}
-
-const updateCodebaseContext = traceable(updateCodebaseContextFunc, {
-  name: "update_codebase_context",
+  name: "generate-task-summary",
 });
 
 export async function summarizeTaskSteps(
@@ -222,10 +129,7 @@ export async function summarizeTaskSteps(
   }
 
   const model = await loadModel(config, Task.SUMMARIZER);
-  const [taskSummary, updatedCodebaseContext] = await Promise.all([
-    generateTaskSummary(state, model),
-    updateCodebaseContext(state, model),
-  ]);
+  const taskSummary = await generateTaskSummary(state, model);
   const updatedTaskPlan = completePlanItem(
     state.plan,
     getActiveTask(state.plan).id,
@@ -250,7 +154,6 @@ export async function summarizeTaskSteps(
     const commandUpdate: GraphUpdate = {
       messages: newMessagesStateUpdate,
       plan: updatedTaskPlan,
-      codebaseContext: updatedCodebaseContext,
     };
     return new Command({
       goto: "generate-conclusion",
@@ -261,7 +164,6 @@ export async function summarizeTaskSteps(
   const commandUpdate: GraphUpdate = {
     messages: newMessagesStateUpdate,
     plan: updatedTaskPlan,
-    codebaseContext: updatedCodebaseContext,
   };
   return new Command({
     goto: "generate-action",

apps/open-swe/src/nodes/take-action.ts

@@ -18,6 +18,7 @@ import {
 import { Command } from "@langchain/langgraph";
 import { truncateOutput } from "../utils/truncate-outputs.js";
 import { daytonaClient } from "../utils/sandbox.js";
+import { getCodebaseTree } from "../utils/tree.js";
 
 const logger = createLogger(LogLevel.INFO, "TakeAction");
 
@@ -154,11 +155,15 @@ export async function takeAction(
         isToolMessage(m) && !m.additional_kwargs?.is_diagnosis,
     ),
   );
+
+  const codebaseTree = await getCodebaseTree();
+
   return new Command({
     goto: shouldRouteDiagnoseNode ? "diagnose-error" : "progress-plan-step",
     update: {
       messages: [toolMessage],
       ...(branchName && { branchName }),
+      codebaseTree,
     },
   });
 }

apps/open-swe/src/subgraphs/planner/nodes/generate-message.ts

@@ -8,6 +8,7 @@ import { getUserRequest } from "../../../utils/user-request.js";
 import { isHumanMessage } from "@langchain/core/messages";
 import { formatFollowupMessagePrompt } from "../utils/followup-prompt.js";
 import { getRepoAbsolutePath } from "../../../utils/git.js";
+import { SANDBOX_ROOT_DIR } from "../../../constants.js";
 
 const logger = createLogger(LogLevel.INFO, "GeneratePlanningMessageNode");
 
@@ -25,6 +26,12 @@ You MUST adhere to the following criteria when gathering context for the plan:
 - Not generating a tool call will be interpreted as an indication that you've gathered enough context to generate a plan.
 - The repo is already cloned, and located inside {REPO_DIRECTORY}
 
+Below is an up to date tree of the codebase (going 3 levels deep). This is up to date, and is updated after every action you take. Always assume this is the most up to date context about the codebase.
+It was generated by using the \`tree\` command, passing in the gitignore file to ignore files and directories you should not have access to (\`git ls-files | tree --fromfile -L 3\`). It is always executed inside the repo directory: {REPO_DIRECTORY}
+{CODEBASE_TREE}
+
+Your current working directory is: {CURRENT_WORKING_DIRECTORY}
+
 The user's request is the first user message in the conversation below. Ensure you generate your plan in accordance with the user's request.
 `;
 
@@ -37,7 +44,12 @@ function formatSystemPrompt(state: PlannerGraphState): string {
       "{FOLLOWUP_MESSAGE_PROMPT}",
       isFollowup ? formatFollowupMessagePrompt(state.plan) : "",
     )
-    .replace("{REPO_DIRECTORY}", getRepoAbsolutePath(state.targetRepository));
+    .replaceAll("{REPO_DIRECTORY}", getRepoAbsolutePath(state.targetRepository))
+    .replaceAll(
+      "{CODEBASE_TREE}",
+      `<codebase-tree>\n${state.codebaseTree || "No codebase tree generated yet."}\n</codebase-tree>`,
+    )
+    .replaceAll("{CURRENT_WORKING_DIRECTORY}", SANDBOX_ROOT_DIR);
 }
 
 export async function generateAction(

apps/open-swe/src/tools/shell.ts

@@ -6,12 +6,10 @@ import { getCurrentTaskInput } from "@langchain/langgraph";
 import { getSandboxErrorFields } from "../utils/sandbox-error-fields.js";
 import { createLogger, LogLevel } from "../utils/logger.js";
 import { daytonaClient } from "../utils/sandbox.js";
-import { SANDBOX_ROOT_DIR } from "../constants.js";
+import { SANDBOX_ROOT_DIR, TIMEOUT_SEC } from "../constants.js";
 
 const logger = createLogger(LogLevel.INFO, "ShellTool");
 
-const DEFAULT_COMMAND_TIMEOUT = 60_000; // 1 minute
-
 const DEFAULT_ENV = {
   // Prevents corepack from showing a y/n download prompt which causes the command to hang
   COREPACK_ENABLE_DOWNLOAD_PROMPT: "0",
@@ -28,9 +26,9 @@ const shellToolSchema = z.object({
   timeout: z
     .number()
     .optional()
-    .default(DEFAULT_COMMAND_TIMEOUT)
+    .default(TIMEOUT_SEC)
     .describe(
-      "The maximum time to wait for the command to complete in milliseconds.",
+      "The maximum time to wait for the command to complete in seconds.",
     ),
 });
 
@@ -55,7 +53,7 @@ export const shellTool = tool(
         command.join(" "),
         workdir,
         DEFAULT_ENV,
-        timeout ?? DEFAULT_COMMAND_TIMEOUT,
+        timeout ?? TIMEOUT_SEC,
       );
 
       if (response.exitCode !== 0) {

apps/open-swe/src/types.ts

@@ -129,10 +129,6 @@ export const GraphAnnotation = z.object({
     .string()
     .default(() => "")
     .langgraph.reducer((_state, update) => update),
-  codebaseContext: z
-    .string()
-    .default(() => "")
-    .langgraph.reducer((_state, update) => update),
   /**
    * The session ID of the Sandbox to use.
    */
@@ -153,6 +149,10 @@ export const GraphAnnotation = z.object({
   targetRepository: z
     .custom<TargetRepository>()
     .langgraph.reducer((_state, update) => update),
+  codebaseTree: z
+    .string()
+    .default(() => "")
+    .langgraph.reducer((_state, update) => update),
 });
 
 export type GraphState = z.infer<typeof GraphAnnotation>;

apps/open-swe/src/utils/tree.ts

@@ -0,0 +1,36 @@
+import { getCurrentTaskInput } from "@langchain/langgraph";
+import { GraphState } from "../types.js";
+import { daytonaClient } from "./sandbox.js";
+import { createLogger, LogLevel } from "./logger.js";
+import path from "node:path";
+import { SANDBOX_ROOT_DIR, TIMEOUT_SEC } from "../constants.js";
+
+const logger = createLogger(LogLevel.INFO, "Tree");
+
+export async function getCodebaseTree(sandboxSessionId_?: string) {
+  const command = `git ls-files | tree --fromfile -L 3`;
+  const state = getCurrentTaskInput<GraphState>();
+  const sandboxSessionId = sandboxSessionId_ || state.sandboxSessionId;
+  if (!sandboxSessionId) {
+    logger.error("Failed to generate tree: No sandbox session ID provided");
+    throw new Error("Failed generate tree: No sandbox session ID provided");
+  }
+  const sandbox = await daytonaClient().get(sandboxSessionId);
+  const repoDir = path.join(SANDBOX_ROOT_DIR, state.targetRepository.repo);
+  const response = await sandbox.process.executeCommand(
+    command,
+    repoDir,
+    undefined,
+    TIMEOUT_SEC,
+  );
+
+  if (response.exitCode !== 0) {
+    logger.error("Failed to generate tree", {
+      error: response.result,
+      error_result: response,
+    });
+    throw new Error(`Failed to generate tree: ${response.result}`);
+  }
+
+  return response.result;
+}