feat: Implement write file tool (#3)

* feat: Implement write file tool

* cr

* feat: Implement generate message and take action nodes

e2b.js.Dockerfile

@@ -2,7 +2,7 @@
 FROM e2bdev/code-interpreter:latest
 
 # Ensure git is installed
-RUN apt-get update && apt-get install -y git curl
+RUN apt-get update && apt-get install -y git curl ripgrep
 
 # Set the working directory
 WORKDIR /app

src/nodes/generate-message.ts

@@ -1,12 +1,90 @@
+import { spawnSync } from "node:child_process";
 import { GraphState, GraphConfig, GraphUpdate } from "../types.js";
+import { loadModel } from "../utils/load-model.js";
+import { shellTool, applyPatchTool } from "../tools/index.js";
+
+const dynamicLines: Array<string> = [];
+if (spawnSync("rg", ["--version"], { stdio: "ignore" }).status === 0) {
+  dynamicLines.push(
+    "- Always use rg instead of grep/ls -R because it is much faster and respects gitignore",
+  );
+}
+const dynamicPrefix = dynamicLines.join("\n");
+const prefix = `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.
+
+You can:
+- Receive user prompts, project context, and files.
+- Stream responses and emit function calls (e.g., shell commands, code edits).
+- Apply patches, run commands, and manage user approvals based on policy.
+- Work inside a sandboxed, git-backed workspace with rollback support.
+
+You work based on a plan which was generated in a previous step.
+The following tasks have already been completed:
+{COMPLETED_TASKS}
+
+The following tasks remain to be completed:
+{PENDING_TASKS}
+
+You are currently working on the following task:
+{CURRENT_TASK}
+
+You are an agent - please keep going until the user's query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved. If you are not sure about file content or codebase structure pertaining to the user's request, use your tools to read files and gather the relevant information: do NOT guess or make up an answer.
+
+Please resolve the user's task by editing and testing the code files in your current code execution session. You are a deployed coding agent. Your session allows for you to modify and run code. The repo(s) are already cloned in your working directory, and you must fully solve the problem for your answer to be considered correct. You are permitted to take as long as you need to complete the current task.
+
+You MUST adhere to the following criteria when executing the task:
+- Working on the repo(s) in the current environment is allowed, even if they are proprietary.
+- Analyzing code for vulnerabilities is allowed.
+- Showing user code and tool call details is allowed.
+- User instructions may overwrite the *CODING GUIDELINES* section in this developer message.
+- Use \`apply_patch\` to edit files. This tool accepts diffs and file paths. It will then apply the given diff to the file.
+- If completing the user's task requires writing or modifying files:
+    - Your code and final answer should follow these *CODING GUIDELINES*:
+        - Fix the problem at the root cause rather than applying surface-level patches, when possible.
+        - Avoid unneeded complexity in your solution.
+            - Ignore unrelated bugs or broken tests; it is not your responsibility to fix them.
+        - Update documentation as necessary.
+        - Keep changes consistent with the style of the existing codebase. Changes should be minimal and focused on the task.
+            - Use \`git log\` and \`git blame\` to search the history of the codebase if additional context is required; internet access is disabled.
+        - NEVER add copyright or license headers unless specifically requested.
+        - You do not need to \`git commit\` your changes; this will be done automatically for you.
+        - If there is a .pre-commit-config.yaml, use \`pre-commit run --files ...\` to check that your changes pass the pre-commit checks. However, do not fix pre-existing errors on lines you didn't touch.
+            - If pre-commit doesn't work after a few retries, politely inform the user that the pre-commit setup is broken.
+        - Once you finish coding, you must
+            - Remove all inline comments you added as much as possible, even if they look normal. Check using \`git diff\`. Inline comments must be generally avoided, unless active maintainers of the repo, after long careful study of the code and the issue, will still misinterpret the code without the comments.
+            - Check if you accidentally add copyright or license headers. If so, remove them.
+            - Try to run pre-commit if it is available.
+            - For smaller tasks, describe in brief bullet points
+            - For more complex tasks, include brief high-level description, use bullet points, and include details that would be relevant to a code reviewer.
+- If completing the user's task DOES NOT require writing or modifying files (e.g., the user asks a question about the code base):
+    - Respond in a friendly tone as a remote teammate, who is knowledgeable, capable and eager to help with coding.
+- When your task involves writing or modifying files:
+    - Do NOT tell the user to "save the file" or "copy the code into a file" if you already created or modified the file using \`apply_patch\`. Instead, reference the file as already saved.
+    - Do NOT show the full contents of large files you have already written, unless the user explicitly asks for them.
+
+${dynamicPrefix}`;
 
 /**
  * Initializes the session. This ensures there's an active VM session, and that
  * the proper credentials are provided for taking actions on GitHub.
  */
-export function generateAction(
+export async function generateAction(
   state: GraphState,
   config: GraphConfig,
 ): Promise<GraphUpdate> {
-  throw new Error("Not implemented");
+  const model = await loadModel(config);
+  const tools = [shellTool, applyPatchTool];
+  const modelWithTools = model.bindTools(tools, { tool_choice: "auto" });
+
+  const response = await modelWithTools.invoke([
+    {
+      role: "system",
+      content: prefix,
+    },
+    ...state.messages,
+  ]);
+
+  return {
+    messages: response,
+  };
 }

src/nodes/take-action.ts

@@ -1,12 +1,42 @@
+import { isAIMessage, ToolMessage } from "@langchain/core/messages";
+import { applyPatchTool, shellTool } from "../tools/index.js";
 import { GraphState, GraphConfig, GraphUpdate } from "../types.js";
 
 /**
  * Initializes the session. This ensures there's an active VM session, and that
  * the proper credentials are provided for taking actions on GitHub.
  */
-export function takeAction(
+export async function takeAction(
   state: GraphState,
-  config: GraphConfig,
+  _config: GraphConfig,
 ): Promise<GraphUpdate> {
-  throw new Error("Not implemented");
+  const lastMessage = state.messages[state.messages.length - 1];
+
+  if (!isAIMessage(lastMessage) || !lastMessage.tool_calls?.length) {
+    throw new Error("Last message is not an AI message with tool calls.");
+  }
+
+  const toolsMap = {
+    [applyPatchTool.name]: applyPatchTool,
+    [shellTool.name]: shellTool,
+  };
+
+  const toolCall = lastMessage.tool_calls[0];
+
+  if (!toolCall) {
+    throw new Error("No tool call found.");
+  }
+
+  const tool = toolsMap[toolCall.name];
+
+  if (!tool) {
+    throw new Error(`Unknown tool: ${toolCall.name}`);
+  }
+
+  // @ts-expect-error tool.invoke types are weird here...
+  const result: ToolMessage = await tool.invoke(toolCall.args);
+
+  return {
+    messages: [...state.messages, result],
+  };
 }

src/tools/apply-patch.ts

@@ -5,12 +5,12 @@ import { GraphConfig } from "../types.js";
 import { Sandbox } from "@e2b/code-interpreter";
 import { readFile, writeFile } from "../utils/read-write.js";
 
-const writeFileToolSchema = z.object({
-  diff: z.string().describe("The diff to apply"),
-  file_path: z.string().describe("The file path to apply the diff to"),
+const applyPatchToolSchema = z.object({
+  diff: z.string().describe("The diff to apply. Use a standard diff format."),
+  file_path: z.string().describe("The file path to apply the diff to."),
 });
 
-export const writeFileTool = tool(
+export const applyPatchTool = tool(
   async (input, config: GraphConfig) => {
     const { diff, file_path } = input;
     const sessionId = config.configurable?.sandbox_session_id;
@@ -44,9 +44,8 @@ export const writeFileTool = tool(
     return `Successfully applied diff to \`${file_path}\` and saved changes.`;
   },
   {
-    name: "write_file",
-    description:
-      "Writes a file given a file path and diff content. Can be used to create or update files.",
-    schema: writeFileToolSchema,
+    name: "apply_patch",
+    description: "Applies a diff to a file given a file path and diff content.",
+    schema: applyPatchToolSchema,
   },
 );

src/tools/index.ts

@@ -0,0 +1,2 @@
+export * from "./apply-patch.js";
+export * from "./shell.js";

src/types.ts

@@ -66,6 +66,82 @@ export const GraphConfiguration = z.object({
    * The language of the sandbox to use.
    */
   sandbox_language: z.enum(["js", "python"]).optional().langgraph.metadata({}),
+  /**
+   * The model ID to use for the reflection generation.
+   * Should be in the format `provider:model_name`.
+   * Defaults to `anthropic:claude-3-7-sonnet-latest`.
+   */
+  modelName: z
+    .string()
+    .optional()
+    .langgraph.metadata({
+      x_lg_ui_config: {
+        type: "select",
+        default: "anthropic:claude-3-7-sonnet-latest",
+        description: "The model to use in all generations",
+        options: [
+          {
+            label: "Claude 3.7 Sonnet",
+            value: "anthropic:claude-3-7-sonnet-latest",
+          },
+          {
+            label: "Claude 3.5 Sonnet",
+            value: "anthropic:claude-3-5-sonnet-latest",
+          },
+          {
+            label: "GPT 4o",
+            value: "openai:gpt-4o",
+          },
+          {
+            label: "GPT 4.1",
+            value: "openai:gpt-4.1",
+          },
+          {
+            label: "o3",
+            value: "openai:o3",
+          },
+          {
+            label: "o3 mini",
+            value: "openai:o3-mini",
+          },
+          {
+            label: "o4",
+            value: "openai:o4",
+          },
+        ],
+      },
+    }),
+  /**
+   * The temperature to use for the reflection generation.
+   * Defaults to `0.7`.
+   */
+  temperature: z
+    .number()
+    .optional()
+    .langgraph.metadata({
+      x_lg_ui_config: {
+        type: "slider",
+        default: 0,
+        min: 0,
+        max: 2,
+        step: 0.1,
+        description: "Controls randomness (0 = deterministic, 2 = creative)",
+      },
+    }),
+  /**
+   * The maximum number of tokens to generate.
+   * Defaults to `1000`.
+   */
+  maxTokens: z
+    .number()
+    .optional()
+    .langgraph.metadata({
+      x_lg_ui_config: {
+        type: "number",
+        min: 1,
+        description: "The maximum number of tokens to generate",
+      },
+    }),
 });
 
 export type GraphConfig = LangGraphRunnableConfig<

src/utils/load-model.ts

@@ -0,0 +1,16 @@
+import { initChatModel } from "langchain/chat_models/universal";
+import { GraphConfig } from "../types.js";
+
+export async function loadModel(config: GraphConfig) {
+  const modelStr =
+    config.configurable?.modelName ?? "anthropic:claude-3-7-sonnet-latest";
+  const [modelProvider, ...modelNameParts] = modelStr.split(":");
+  const modelName = modelNameParts.join(":");
+  const model = await initChatModel(modelName, {
+    modelProvider,
+    temperature: config.configurable?.temperature ?? 0,
+    maxTokens: config.configurable?.maxTokens ?? undefined,
+  });
+
+  return model;
+}