From 66f32b4afc9d8da48abbaac64ff6359fc56491d1 Mon Sep 17 00:00:00 2001 From: Brace Sproul Date: Thu, 22 May 2025 19:05:29 -0700 Subject: [PATCH] feat: Implement a planning subgraph (#12) --- langgraph.json | 2 +- src/constants.ts | 3 +- src/index.ts | 33 ++---- src/nodes/generate-message.ts | 14 ++- src/nodes/generate-plan.ts | 48 -------- src/nodes/index.ts | 1 - src/nodes/initialize.ts | 33 ++++-- src/nodes/interrupt-plan.ts | 33 ++++-- src/nodes/take-action.ts | 14 ++- src/subgraphs/index.ts | 1 + src/subgraphs/planner/index.ts | 66 +++++++++++ .../planner/nodes/generate-message.ts | 45 +++++++ src/subgraphs/planner/nodes/generate-plan.ts | 58 +++++++++ src/subgraphs/planner/nodes/index.ts | 3 + src/subgraphs/planner/nodes/take-action.ts | 45 +++++++ src/subgraphs/planner/types.ts | 15 +++ src/tools/apply-patch.ts | 20 ++-- src/tools/shell.ts | 71 ++++++++--- src/types.ts | 111 ++++++++++++++---- src/utils/load-model.ts | 17 ++- src/utils/plan-prompt.ts | 10 +- src/utils/read-write.ts | 6 + src/utils/sandbox.ts | 11 ++ 23 files changed, 516 insertions(+), 144 deletions(-) delete mode 100644 src/nodes/generate-plan.ts create mode 100644 src/subgraphs/index.ts create mode 100644 src/subgraphs/planner/index.ts create mode 100644 src/subgraphs/planner/nodes/generate-message.ts create mode 100644 src/subgraphs/planner/nodes/generate-plan.ts create mode 100644 src/subgraphs/planner/nodes/index.ts create mode 100644 src/subgraphs/planner/nodes/take-action.ts create mode 100644 src/subgraphs/planner/types.ts diff --git a/langgraph.json b/langgraph.json index eab355aa..c3067c6f 100644 --- a/langgraph.json +++ b/langgraph.json @@ -1,7 +1,7 @@ { "node_version": "20", "graphs": { - "agent": "./src/index.ts:graph" + "open-codex": "./src/index.ts:graph" }, "env": ".env", "dependencies": ["."] diff --git a/src/constants.ts b/src/constants.ts index f2ea2700..63af2e5a 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -1,4 +1,5 @@ +export const TIMEOUT_MS = 300_000; // 5 minutes // Always extend timeouts by 5 min after each command execution. export const TIMEOUT_EXTENSION_OPT = { - timeoutMs: 300_000, // 5 minutes + timeoutMs: TIMEOUT_MS, }; diff --git a/src/index.ts b/src/index.ts index 44d22ffd..d270350c 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,12 +1,6 @@ import { END, START, StateGraph } from "@langchain/langgraph"; +import { GraphAnnotation, GraphConfiguration, GraphState } from "./types.js"; import { - GraphAnnotation, - GraphConfig, - GraphConfiguration, - GraphState, -} from "./types.js"; -import { - generatePlan, initialize, generateAction, takeAction, @@ -15,7 +9,7 @@ import { progressPlanStep, } from "./nodes/index.js"; import { isAIMessage } from "@langchain/core/messages"; -import { pauseSandbox } from "./utils/sandbox.js"; +import { plannerGraph } from "./subgraphs/index.js"; /** * @param {GraphState} state - The current graph state. @@ -42,7 +36,6 @@ function routeAfterPlan(state: GraphState): "interrupt-plan" | typeof END { */ async function takeActionOrEnd( state: GraphState, - config: GraphConfig, ): Promise { const { messages } = state; const lastMessage = messages[messages.length - 1]; @@ -52,35 +45,33 @@ async function takeActionOrEnd( return "take-action"; } - // First, pause the sandbox before ending the graph. - if (config.configurable?.sandbox_session_id) { - await pauseSandbox(config.configurable.sandbox_session_id); - } - return END; } const workflow = new StateGraph(GraphAnnotation, GraphConfiguration) - .addNode("generate-plan", generatePlan) + .addNode("initialize", initialize) + .addNode("generate-plan-subgraph", plannerGraph) .addNode("rewrite-plan", rewritePlan) .addNode("interrupt-plan", interruptPlan, { // TODO: Hookup `Command` in interruptPlan node so this actually works. - ends: [END, "rewrite-plan", "initialize"], + ends: [END, "rewrite-plan", "generate-action"], }) - .addNode("initialize", initialize) .addNode("generate-action", generateAction) .addNode("take-action", takeAction) .addNode("progress-plan-step", progressPlanStep) - .addEdge(START, "generate-plan") + .addEdge(START, "initialize") + .addEdge("initialize", "generate-plan-subgraph") // TODO: Update routing to work w/ new interrupt node. - .addConditionalEdges("generate-plan", routeAfterPlan, ["interrupt-plan", END]) + .addConditionalEdges("generate-plan-subgraph", routeAfterPlan, [ + "interrupt-plan", + END, + ]) // Always interrupt after rewriting the plan. .addEdge("rewrite-plan", "interrupt-plan") - .addEdge("initialize", "generate-action") .addConditionalEdges("generate-action", takeActionOrEnd, ["take-action", END]) .addEdge("take-action", "progress-plan-step") .addEdge("progress-plan-step", "generate-action"); // Zod types are messed up export const graph = workflow.compile() as any; -graph.name = "LangGraph ReAct MCP"; +graph.name = "Open Codex"; diff --git a/src/nodes/generate-message.ts b/src/nodes/generate-message.ts index 54b3a1e9..9dd030d2 100644 --- a/src/nodes/generate-message.ts +++ b/src/nodes/generate-message.ts @@ -2,6 +2,7 @@ import { GraphState, GraphConfig, GraphUpdate, PlanItem } from "../types.js"; import { loadModel, Task } from "../utils/load-model.js"; import { shellTool, applyPatchTool } from "../tools/index.js"; import { formatPlanPrompt } from "../utils/plan-prompt.js"; +import { pauseSandbox } from "../utils/sandbox.js"; 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. @@ -48,7 +49,9 @@ You MUST adhere to the following criteria when executing the task: - 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. -- Always use \`rg\` instead of \`grep/ls -R\` because it is much faster and respects gitignore.`; +- Always use \`rg\` instead of \`grep/ls -R\` because it is much faster and respects gitignore. + - Always use glob patterns when searching with \`rg\` for specific file types. For example, to search for all TSX files, use \`rg -i star -g **/*.tsx project-directory/\`. This is because \`rg\` does not have built in file types for every language. +`; const formatPrompt = (plan: PlanItem[]): string => { return systemPrompt.replace("{PLAN_PROMPT}", formatPlanPrompt(plan)); @@ -70,7 +73,16 @@ export async function generateAction( ...state.messages, ]); + const hasToolCalls = !!response.tool_calls?.length; + // No tool calls means the graph is going to end. Pause the sandbox. + let newSandboxSessionId: string | undefined; + if (!hasToolCalls && state.sandboxSessionId) { + console.log("No tool calls found. Pausing sandbox..."); + newSandboxSessionId = await pauseSandbox(state.sandboxSessionId); + } + return { messages: [response], + ...(newSandboxSessionId && { sandboxSessionId: newSandboxSessionId }), }; } diff --git a/src/nodes/generate-plan.ts b/src/nodes/generate-plan.ts deleted file mode 100644 index bad51422..00000000 --- a/src/nodes/generate-plan.ts +++ /dev/null @@ -1,48 +0,0 @@ -import { sessionPlanTool } from "../tools/index.js"; -import { GraphState, GraphConfig, GraphUpdate } from "../types.js"; -import { loadModel, Task } from "../utils/load-model.js"; - -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. - -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 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. -- If you can not generate a plan due to a lack of context, you are permitted to ask the user followup questions. - - If asking followup questions, ensure every question is asked in a single message to avoid back and forth. - - Your questions should be concise and to the point. Remember that you are not including code or technical details in your plan, so your questions should be focused on high-level issues. -- When you are ready to generate the plan, ensure you call the 'session_plan' tool. -`; - -export async function generatePlan( - state: GraphState, - config: GraphConfig, -): Promise { - const model = await loadModel(config, Task.PLANNER); - const modelWithTools = model.bindTools([sessionPlanTool], { - tool_choice: "auto", - }); - - const response = await modelWithTools.invoke([ - { - role: "system", - content: systemPrompt, - }, - ...state.messages, - ]); - - if (response.tool_calls?.length) { - return { - proposedPlan: response.tool_calls[0].args.plan, - plan: [], - }; - } - - // No tool calls generated, instead we should just return the messages. - return { - messages: [response], - proposedPlan: [], - plan: [], - }; -} diff --git a/src/nodes/index.ts b/src/nodes/index.ts index 5aec7113..efd07921 100644 --- a/src/nodes/index.ts +++ b/src/nodes/index.ts @@ -1,4 +1,3 @@ -export * from "./generate-plan.js"; export * from "./initialize.js"; export * from "./generate-message.js"; export * from "./take-action.js"; diff --git a/src/nodes/initialize.ts b/src/nodes/initialize.ts index 010c74b8..5743ad24 100644 --- a/src/nodes/initialize.ts +++ b/src/nodes/initialize.ts @@ -24,6 +24,9 @@ async function cloneRepo(sandbox: Sandbox, targetRepository: TargetRepository) { gitCloneCommand.push(repoUrlWithToken); } + console.log("Cloning repository...", { + command: gitCloneCommand.join(" "), + }); return await sandbox.commands.run( gitCloneCommand.join(" "), TIMEOUT_EXTENSION_OPT, @@ -37,25 +40,35 @@ async function cloneRepo(sandbox: Sandbox, targetRepository: TargetRepository) { * branch. */ export async function initialize( - _state: GraphState, + state: GraphState, config: GraphConfig, ): Promise { if (!config.configurable) { throw new Error("Configuration object not found."); } - const { sandbox_session_id, target_repository, sandbox_language } = - config.configurable; - if (sandbox_session_id) { + const { sandboxSessionId } = state; + + if (sandboxSessionId) { try { + console.log("Sandbox session ID exists. Resuming...", { + sandboxSessionId, + }); // Resume the sandbox if the session ID is in the config. - await Sandbox.resume(sandbox_session_id, TIMEOUT_EXTENSION_OPT); - return {}; + const newSandbox = await Sandbox.resume( + sandboxSessionId, + TIMEOUT_EXTENSION_OPT, + ); + return { + sandboxSessionId: newSandbox.sandboxId, + }; } catch (e) { // Error thrown, log it and continue. Will create a new sandbox session since the resumption failed. console.error("Failed to get sandbox session.", e); } } + const { target_repository, sandbox_language } = config.configurable; + if (!sandbox_language || !target_repository) { throw new Error( "Missing required configuration. Please provide a sandbox language and git repository URL.", @@ -63,18 +76,22 @@ export async function initialize( } if (sandbox_language === "js") { + console.log("Creating JS sandbox..."); const sandbox = await Sandbox.create( JS_SANDBOX_TEMPLATE_ID, TIMEOUT_EXTENSION_OPT, ); - config.configurable.sandbox_session_id = sandbox.sandboxId; const res = await cloneRepo(sandbox, target_repository); if (res.error) { // TODO: This should probably be an interrupt. + console.error("Failed to clone repository.", res.error); throw new Error(`Failed to clone repository.\n${res.error}`); } - return {}; + console.log("Repository cloned successfully."); + return { + sandboxSessionId: sandbox.sandboxId, + }; } if (sandbox_language === "python") { diff --git a/src/nodes/interrupt-plan.ts b/src/nodes/interrupt-plan.ts index 00aa8c94..9f4fa57b 100644 --- a/src/nodes/interrupt-plan.ts +++ b/src/nodes/interrupt-plan.ts @@ -5,9 +5,9 @@ import { HumanInterrupt, HumanResponse, } from "@langchain/langgraph/prebuilt"; -import { v4 as uuidv4 } from "uuid"; +import { resumeSandbox } from "../utils/sandbox.js"; -export function interruptPlan(state: GraphState): Command { +export async function interruptPlan(state: GraphState): Promise { const { proposedPlan } = state; if (!proposedPlan.length) { throw new Error("No proposed plan found."); @@ -30,33 +30,46 @@ export function interruptPlan(state: GraphState): Command { If editing the plan, ensure each step in the plan is separated by ":::".`, })[0]; + if (!state.sandboxSessionId) { + // TODO: This should prob just create a sandbox? + throw new Error("No sandbox session ID found."); + } + if (interruptRes.type === "accept") { - // Plan was accepted, route to the initialize node. + const newSandboxSessionId = (await resumeSandbox(state.sandboxSessionId)) + .sandboxId; + + // Plan was accepted, route to the generate-action node to start taking actions. return new Command({ - goto: "initialize", + goto: "generate-action", update: { - plan: proposedPlan.map((p) => ({ - id: uuidv4(), + plan: proposedPlan.map((p, index) => ({ + index, plan: p, completed: false, })), + sandboxSessionId: newSandboxSessionId, }, }); } if (interruptRes.type === "edit") { - // Plan was edited, route to the initialize node. + const newSandboxSessionId = (await resumeSandbox(state.sandboxSessionId)) + .sandboxId; + + // Plan was edited, route to the generate-action node to start taking actions. const editedPlan = (interruptRes.args as ActionRequest).args.plan .split(":::") .map((step: string) => step.trim()); return new Command({ - goto: "initialize", + goto: "generate-action", update: { - plan: editedPlan.map((p: string) => ({ - id: uuidv4(), + plan: editedPlan.map((p: string, index: number) => ({ + index, plan: p, completed: false, })), + sandboxSessionId: newSandboxSessionId, }, }); } diff --git a/src/nodes/take-action.ts b/src/nodes/take-action.ts index 26576000..c0cc45ab 100644 --- a/src/nodes/take-action.ts +++ b/src/nodes/take-action.ts @@ -28,11 +28,21 @@ export async function takeAction( if (!tool) { throw new Error(`Unknown tool: ${toolCall.name}`); } + if (!state.sandboxSessionId) { + throw new Error( + "Failed to take action: No sandbox session ID found in state.", + ); + } // @ts-expect-error tool.invoke types are weird here... - const result: ToolMessage = await tool.invoke(toolCall.args); + const result: string = await tool.invoke(toolCall.args); + const toolMessage = new ToolMessage({ + tool_call_id: toolCall.id ?? "", + content: result, + name: toolCall.name, + }); return { - messages: [...state.messages, result], + messages: [toolMessage], }; } diff --git a/src/subgraphs/index.ts b/src/subgraphs/index.ts new file mode 100644 index 00000000..a5b79daa --- /dev/null +++ b/src/subgraphs/index.ts @@ -0,0 +1 @@ +export * from "./planner/index.js"; diff --git a/src/subgraphs/planner/index.ts b/src/subgraphs/planner/index.ts new file mode 100644 index 00000000..51bc0c66 --- /dev/null +++ b/src/subgraphs/planner/index.ts @@ -0,0 +1,66 @@ +/** + * Subgraph for gathering context & generating a plan. + * pre-requisites: + * VM is already booted & repo is cloned (init node) + * Steps: + */ + +import { END, Send, START, StateGraph } from "@langchain/langgraph"; +import { PlannerGraphState, PlannerGraphStateObj } from "./types.js"; +import { GraphConfiguration } from "../../types.js"; +import { generateAction, generatePlan, takeAction } from "./nodes/index.js"; +import { isAIMessage, ToolMessage } from "@langchain/core/messages"; + +function takeActionOrGeneratePlan( + state: PlannerGraphState, +): "take-plan-action" | "generate-plan" | Send { + const { plannerMessages } = state; + const lastMessage = plannerMessages[plannerMessages.length - 1]; + // If the last message is a tool call, and we have executed less than 6 actions, take action. + // Max actions is 13, because that's 6 actions (2 messages per action, ai & tool) plus the input message. + const maxActionsCount = 13; + if ( + isAIMessage(lastMessage) && + lastMessage.tool_calls?.length && + plannerMessages.length < maxActionsCount + ) { + return "take-plan-action"; + } + + if (isAIMessage(lastMessage) && lastMessage.tool_calls?.length) { + // If this is true, we need to return a `Command` adding a ToolMessage to the state + // so that the last AI message has a tool message pair + const lastMessageToolCall = lastMessage.tool_calls[0]; + return new Send("generate-plan", { + ...state, + plannerMessages: [ + ...state.plannerMessages, + new ToolMessage({ + tool_call_id: lastMessageToolCall.id ?? "", + name: lastMessageToolCall.name, + content: "Tool call not executed. Max actions reached.", + }), + ], + }); + } + + // If the last message does not have tool calls, continue to generate plan without modifications. + return "generate-plan"; +} + +const workflow = new StateGraph(PlannerGraphStateObj, GraphConfiguration) + .addNode("generate-plan-context-action", generateAction) + .addNode("take-plan-action", takeAction) + .addNode("generate-plan", generatePlan) + .addEdge(START, "generate-plan-context-action") + .addConditionalEdges( + "generate-plan-context-action", + takeActionOrGeneratePlan, + ["take-plan-action", "generate-plan"], + ) + .addEdge("take-plan-action", "generate-plan-context-action") + .addEdge("generate-plan", END); + +// TODO: Fix zod types +export const plannerGraph = workflow.compile() as any; +plannerGraph.name = "Planner"; diff --git a/src/subgraphs/planner/nodes/generate-message.ts b/src/subgraphs/planner/nodes/generate-message.ts new file mode 100644 index 00000000..c7801349 --- /dev/null +++ b/src/subgraphs/planner/nodes/generate-message.ts @@ -0,0 +1,45 @@ +import { loadModel, Task } from "../../../utils/load-model.js"; +import { shellTool, applyPatchTool } from "../../../tools/index.js"; +import { PlannerGraphState, PlannerGraphUpdate } from "../types.js"; +import { GraphConfig } from "../../../types.js"; +import { isHumanMessage } from "@langchain/core/messages"; + +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. + +Your sole task is to gather context from the repository the user has provided which will be helpful when generating a plan to address the user's request. + +You MUST adhere to the following criteria when gathering context for the plan: +- You must ONLY take read actions to gather context. Write actions are NOT allowed. +- Keep in mind you are only permitted to make a maximum of 6 tool calls to gather all your context. Ensure each action is of high quality, and targeted to aid in generating a plan. +- Always use \`rg\` instead of \`grep/ls -R\` because it is much faster and respects gitignore. + - Always use glob patterns when searching with \`rg\` for specific file types. For example, to search for all TSX files, use \`rg -i star -g **/*.tsx project-directory/\`. This is because \`rg\` does not have built in file types for every language. +- If you determine you've gathered enough context to generate a plan, simply reply with 'done' and do NOT call any tools. +- Not generating a tool call will be interpreted as an indication that you've gathered enough context to generate a plan. +- The first user message in this conversation contains the user's request. +`; + +export async function generateAction( + state: PlannerGraphState, + config: GraphConfig, +): Promise { + const model = await loadModel(config, Task.ACTION_GENERATOR); + const tools = [shellTool, applyPatchTool]; + const modelWithTools = model.bindTools(tools, { tool_choice: "auto" }); + + const firstUserMessage = state.messages.find(isHumanMessage); + + const response = await modelWithTools + .bind({ tags: ["langsmith:nostream"] }) + .invoke([ + { + role: "system", + content: systemPrompt, + }, + ...(firstUserMessage ? [firstUserMessage] : []), + ...state.plannerMessages, + ]); + + return { + plannerMessages: [response], + }; +} diff --git a/src/subgraphs/planner/nodes/generate-plan.ts b/src/subgraphs/planner/nodes/generate-plan.ts new file mode 100644 index 00000000..8e17f639 --- /dev/null +++ b/src/subgraphs/planner/nodes/generate-plan.ts @@ -0,0 +1,58 @@ +import { isHumanMessage } from "@langchain/core/messages"; +import { sessionPlanTool } from "../../../tools/index.js"; +import { GraphConfig } from "../../../types.js"; +import { loadModel, Task } from "../../../utils/load-model.js"; +import { PlannerGraphState, PlannerGraphUpdate } from "../types.js"; +import { pauseSandbox } from "../../../utils/sandbox.js"; + +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. + +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. +- 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. +- 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. +`; + +export async function generatePlan( + state: PlannerGraphState, + config: GraphConfig, +): Promise { + const model = await loadModel(config, Task.PLANNER); + const modelWithTools = model.bindTools([sessionPlanTool], { + tool_choice: sessionPlanTool.name, + }); + + const firstUserMessage = state.messages.find(isHumanMessage); + + const response = await modelWithTools + .bind({ tags: ["langsmith:nostream"] }) + .invoke([ + { + role: "system", + content: systemPrompt, + }, + ...(firstUserMessage ? [firstUserMessage] : []), + ...state.plannerMessages, + ]); + + if (!response.tool_calls?.length) { + throw new Error("Failed to generate plan"); + } + + let newSessionId: string | undefined; + if (state.sandboxSessionId) { + // Pause before returning, as the next step will be to interrupt the graph. + newSessionId = await pauseSandbox(state.sandboxSessionId); + } + + return { + proposedPlan: response.tool_calls[0].args.plan, + plan: [], + ...(newSessionId && { sandboxSessionId: newSessionId }), + }; +} diff --git a/src/subgraphs/planner/nodes/index.ts b/src/subgraphs/planner/nodes/index.ts new file mode 100644 index 00000000..f99d35ec --- /dev/null +++ b/src/subgraphs/planner/nodes/index.ts @@ -0,0 +1,3 @@ +export * from "./generate-message.js"; +export * from "./take-action.js"; +export * from "./generate-plan.js"; diff --git a/src/subgraphs/planner/nodes/take-action.ts b/src/subgraphs/planner/nodes/take-action.ts new file mode 100644 index 00000000..d18d2eb9 --- /dev/null +++ b/src/subgraphs/planner/nodes/take-action.ts @@ -0,0 +1,45 @@ +import { isAIMessage, ToolMessage } from "@langchain/core/messages"; +import { applyPatchTool, shellTool } from "../../../tools/index.js"; +import { GraphConfig } from "../../../types.js"; +import { PlannerGraphState, PlannerGraphUpdate } from "../types.js"; + +export async function takeAction( + state: PlannerGraphState, + _config: GraphConfig, +): Promise { + const { plannerMessages: messages } = state; + const lastMessage = messages[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: string = await tool.invoke(toolCall.args); + const toolMessage = new ToolMessage({ + tool_call_id: toolCall.id ?? "", + content: result, + name: toolCall.name, + }); + + return { + plannerMessages: [toolMessage], + }; +} diff --git a/src/subgraphs/planner/types.ts b/src/subgraphs/planner/types.ts new file mode 100644 index 00000000..6dbcc184 --- /dev/null +++ b/src/subgraphs/planner/types.ts @@ -0,0 +1,15 @@ +import "@langchain/langgraph/zod"; +import { z } from "zod"; +import { addMessages, Messages } from "@langchain/langgraph"; +import { BaseMessage } from "@langchain/core/messages"; +import { GraphAnnotation } from "../../types.js"; + +export const PlannerGraphStateObj = GraphAnnotation.extend({ + plannerMessages: z + .custom() + .default(() => []) + .langgraph.reducer((state, update) => addMessages(state, update)), +}); + +export type PlannerGraphState = z.infer; +export type PlannerGraphUpdate = Partial; diff --git a/src/tools/apply-patch.ts b/src/tools/apply-patch.ts index ca33bd0a..7475d1a3 100644 --- a/src/tools/apply-patch.ts +++ b/src/tools/apply-patch.ts @@ -1,9 +1,10 @@ import { tool } from "@langchain/core/tools"; import { z } from "zod"; import { applyPatch } from "diff"; -import { GraphConfig } from "../types.js"; +import { GraphState } from "../types.js"; import { Sandbox } from "@e2b/code-interpreter"; import { readFile, writeFile } from "../utils/read-write.js"; +import { getCurrentTaskInput } from "@langchain/langgraph"; const applyPatchToolSchema = z.object({ diff: z.string().describe("The diff to apply. Use a standard diff format."), @@ -11,14 +12,19 @@ const applyPatchToolSchema = z.object({ }); export const applyPatchTool = tool( - async (input, config: GraphConfig) => { - const { diff, file_path } = input; - const sessionId = config.configurable?.sandbox_session_id; - if (!sessionId) { - return "FAILED TO RUN COMMAND: No sandbox session ID provided"; + async (input) => { + const state = getCurrentTaskInput(); + const { sandboxSessionId } = state; + if (!sandboxSessionId) { + console.error("FAILED TO RUN COMMAND: No sandbox session ID provided", { + input, + }); + throw new Error("FAILED TO RUN COMMAND: No sandbox session ID provided"); } - const sandbox = await Sandbox.connect(sessionId); + const { diff, file_path } = input; + + const sandbox = await Sandbox.connect(sandboxSessionId); const { success: readFileSuccess, output: readFileOutput } = await readFile( sandbox, diff --git a/src/tools/shell.ts b/src/tools/shell.ts index 2c1ead0c..9925f829 100644 --- a/src/tools/shell.ts +++ b/src/tools/shell.ts @@ -1,7 +1,9 @@ import { tool } from "@langchain/core/tools"; import { z } from "zod"; import { Sandbox } from "@e2b/code-interpreter"; -import { GraphConfig } from "../types.js"; +import { GraphState } from "../types.js"; +import { getCurrentTaskInput } from "@langchain/langgraph"; +import { TIMEOUT_MS } from "../constants.js"; const shellToolSchema = z.object({ command: z.array(z.string()).describe("The command to run"), @@ -18,22 +20,59 @@ const shellToolSchema = z.object({ }); export const shellTool = tool( - async (input, config: GraphConfig) => { - const sessionId = config.configurable?.sandbox_session_id; - if (!sessionId) { - return "FAILED TO RUN COMMAND: No sandbox session ID provided"; + async (input) => { + try { + const state = getCurrentTaskInput(); + const { sandboxSessionId } = state; + if (!sandboxSessionId) { + console.error("FAILED TO RUN COMMAND: No sandbox session ID provided", { + input, + }); + throw new Error( + "FAILED TO RUN COMMAND: No sandbox session ID provided", + ); + } + + const sandbox = await Sandbox.connect(sandboxSessionId); + const { command, workdir, timeout } = input; + const result = await sandbox.commands.run(command.join(" "), { + timeoutMs: timeout, + cwd: workdir, + }); + // Add an extra 5 min timeout to the sandbox. + await sandbox.setTimeout(TIMEOUT_MS); + + if (result.error) { + console.error("Failed to run command", { + error: result.error, + error_result: result, + input, + }); + return `Command failed. Exit code: ${result.exitCode}\nError: ${result.error}\nStderr:\n${result.stderr}`; + } + + return result.stdout; + } catch (e: any) { + if (typeof e === "object" && "result" in e && e.result) { + console.error("Failed to run command", { + error: e.message, + error_result: e.result, + input, + }); + return ( + "FAILED TO RUN COMMAND: " + + e.message + + "\n" + + JSON.stringify(e.result, null, 2) + ); + } + + console.error("Failed to run command: " + e.message, { + error: e, + input, + }); + throw new Error("FAILED TO RUN COMMAND: " + e.message); } - - const sandbox = await Sandbox.connect(sessionId); - const { command, workdir, timeout } = input; - const result = await sandbox.commands.run(command.join(" "), { - timeoutMs: timeout, - cwd: workdir, - }); - - return `Command exited with code ${result.exitCode}:\n -${result.error ? `Error: ${result.error}\n` : ""} -Stdout:\n${result.stdout}\n\nStderr:\n${result.stderr}`; }, { name: "shell", diff --git a/src/types.ts b/src/types.ts index 136a4008..a1716d23 100644 --- a/src/types.ts +++ b/src/types.ts @@ -47,6 +47,13 @@ export const GraphAnnotation = z.object({ .nullable() .default(() => null) .langgraph.reducer((_state, update) => update), + /** + * The session ID of the Sandbox to use. + */ + sandboxSessionId: z + .string() + .optional() + .langgraph.reducer((_state, update) => update), }); export type GraphState = z.infer; @@ -111,18 +118,12 @@ const MODEL_OPTIONS = [ }, ]; +const MODEL_OPTIONS_NO_THINKING = MODEL_OPTIONS.filter( + ({ value }) => + !value.includes("extended-thinking") || !value.startsWith("openai:o"), +); + export const GraphConfiguration = z.object({ - /** - * The session ID of the Sandbox to use. - */ - sandbox_session_id: z - .string() - .optional() - .langgraph.metadata({ - x_lg_ui_config: { - type: "hidden", - }, - }), /** * The URL of the repository to clone. */ @@ -132,25 +133,54 @@ export const GraphConfiguration = z.object({ repo: z.string(), branch: z.string().optional(), }) - .langgraph.metadata({}), + .langgraph.metadata({ + x_oap_ui_config: { + type: "json", + default: `{ + "owner": "", + "repo": "", + "branch": "" +}`, + }, + }), /** * The language of the sandbox to use. */ - sandbox_language: z.enum(["js", "python"]).optional().langgraph.metadata({}), + sandbox_language: z + .enum(["js", "python"]) + .optional() + .langgraph.metadata({ + x_oap_ui_config: { + type: "select", + default: "js", + description: "The primary language of the sandbox to use.", + options: [ + { + label: "JavaScript/TypeScript", + value: "js", + }, + { + label: "Python", + value: "python", + }, + ], + }, + }), /** * The model ID to use for the planning step. * This includes initial planning, and rewriting. - * @default "anthropic:extended-thinking:claude-sonnet-4-0" + * @default "anthropic:claude-sonnet-4-0" */ plannerModelName: z .string() .optional() .langgraph.metadata({ - x_lg_ui_config: { + x_oap_ui_config: { type: "select", - default: "anthropic:extended-thinking:claude-sonnet-4-0", + default: "anthropic:claude-sonnet-4-0", description: "The model to use for planning", - options: MODEL_OPTIONS, + // Do not show extended thinking models + options: MODEL_OPTIONS_NO_THINKING, }, }), /** @@ -163,7 +193,42 @@ export const GraphConfiguration = z.object({ .number() .optional() .langgraph.metadata({ - x_lg_ui_config: { + x_oap_ui_config: { + type: "slider", + default: 0, + min: 0, + max: 2, + step: 0.1, + description: "Controls randomness (0 = deterministic, 2 = creative)", + }, + }), + /** + * The model ID to use for the planning step. + * This includes initial planning, and rewriting. + * @default "anthropic:claude-sonnet-4-0" + */ + plannerContextModelName: z + .string() + .optional() + .langgraph.metadata({ + x_oap_ui_config: { + type: "select", + default: "anthropic:claude-sonnet-4-0", + description: "The model to use for planning", + options: MODEL_OPTIONS, + }, + }), + /** + * The temperature to use for the planning step. + * This includes initial planning, and rewriting. + * If selecting a reasoning model, this will be ignored. + * @default 0 + */ + plannerContextTemperature: z + .number() + .optional() + .langgraph.metadata({ + x_oap_ui_config: { type: "slider", default: 0, min: 0, @@ -181,7 +246,7 @@ export const GraphConfiguration = z.object({ .string() .optional() .langgraph.metadata({ - x_lg_ui_config: { + x_oap_ui_config: { type: "select", default: "anthropic:claude-sonnet-4-0", description: "The model to use for action generation", @@ -197,7 +262,7 @@ export const GraphConfiguration = z.object({ .number() .optional() .langgraph.metadata({ - x_lg_ui_config: { + x_oap_ui_config: { type: "slider", default: 0, min: 0, @@ -215,11 +280,11 @@ export const GraphConfiguration = z.object({ .string() .optional() .langgraph.metadata({ - x_lg_ui_config: { + x_oap_ui_config: { type: "select", default: "anthropic:claude-sonnet-4-0", description: "The model to use for progress plan checking", - options: MODEL_OPTIONS, + options: MODEL_OPTIONS_NO_THINKING, }, }), /** @@ -231,7 +296,7 @@ export const GraphConfiguration = z.object({ .number() .optional() .langgraph.metadata({ - x_lg_ui_config: { + x_oap_ui_config: { type: "slider", default: 0, min: 0, diff --git a/src/utils/load-model.ts b/src/utils/load-model.ts index dbab9926..4d8e7641 100644 --- a/src/utils/load-model.ts +++ b/src/utils/load-model.ts @@ -3,13 +3,18 @@ import { GraphConfig } from "../types.js"; export enum Task { PLANNER = "planner", + PLANNER_CONTEXT = "plannerContext", ACTION_GENERATOR = "actionGenerator", PROGRESS_PLAN_CHECKER = "progressPlanChecker", } const TASK_TO_CONFIG_DEFAULTS_MAP = { [Task.PLANNER]: { - modelName: "anthropic:extended-thinking:claude-sonnet-4-0", + modelName: "anthropic:claude-sonnet-4-0", + temperature: 0, + }, + [Task.PLANNER_CONTEXT]: { + modelName: "anthropic:claude-sonnet-4-0", temperature: 0, }, [Task.ACTION_GENERATOR]: { @@ -34,9 +39,11 @@ export async function loadModel(config: GraphConfig, task: Task) { let thinkingModel = false; if (modelNameParts[0] === "extended-thinking") { + console.log("THINKING ENABLED", modelNameParts); // Using a thinking model. Remove it from the model name. modelNameParts.shift(); thinkingModel = true; + console.log("THINKING ENABLED", modelNameParts); } const modelName = modelNameParts.join(":"); @@ -44,11 +51,17 @@ export async function loadModel(config: GraphConfig, task: Task) { thinkingModel = true; } + const thinkingBudgetTokens = 5000; + const maxTokens = thinkingBudgetTokens * 4; + const model = await initChatModel(modelName, { modelProvider, temperature: thinkingModel ? undefined : temperature, ...(thinkingModel && modelProvider === "anthropic" - ? { thinking: { budgetTokens: 5000, type: "enabled" } } + ? { + thinking: { budget_tokens: thinkingBudgetTokens, type: "enabled" }, + maxTokens, + } : {}), }); diff --git a/src/utils/plan-prompt.ts b/src/utils/plan-prompt.ts index 907169fb..2724694f 100644 --- a/src/utils/plan-prompt.ts +++ b/src/utils/plan-prompt.ts @@ -16,11 +16,15 @@ export function formatPlanPrompt(plan: PlanItem[]): string { return PLAN_PROMPT.replace( "{COMPLETED_TASKS}", - completedTasks.map((task) => `${task.index}. ${task.plan}`).join("\n"), + completedTasks?.length + ? completedTasks.map((task) => `${task.index}. ${task.plan}`).join("\n") + : "No completed tasks.", ) .replace( "{REMAINING_TASKS}", - remainingTasks.map((task) => `${task.index}. ${task.plan}`).join("\n"), + remainingTasks?.length + ? remainingTasks.map((task) => `${task.index}. ${task.plan}`).join("\n") + : "No remaining tasks.", ) - .replace("{CURRENT_TASK}", currentTask.plan); + .replace("{CURRENT_TASK}", currentTask?.plan || "No current task."); } diff --git a/src/utils/read-write.ts b/src/utils/read-write.ts index 8ef3922f..945fce9f 100644 --- a/src/utils/read-write.ts +++ b/src/utils/read-write.ts @@ -1,4 +1,5 @@ import { Sandbox } from "@e2b/code-interpreter"; +import { TIMEOUT_MS } from "../constants.js"; export async function readFile( sandbox: Sandbox, @@ -9,6 +10,9 @@ export async function readFile( }> { try { const readOutput = await sandbox.commands.run(`cat "${filePath}"`); + // Add an extra 5 min timeout to the sandbox. + await sandbox.setTimeout(TIMEOUT_MS); + if (readOutput.exitCode !== 0) { console.error( `Error reading file '${filePath}' from sandbox via cat:`, @@ -51,6 +55,8 @@ export async function writeFile( try { const writeCommand = `printf '%s' '${content}' > "${filePath}"`; const writeOutput = await sandbox.commands.run(writeCommand); + // Add an extra 5 min timeout to the sandbox. + await sandbox.setTimeout(TIMEOUT_MS); if (writeOutput.exitCode !== 0) { console.error( diff --git a/src/utils/sandbox.ts b/src/utils/sandbox.ts index 85629bdd..4a5a2607 100644 --- a/src/utils/sandbox.ts +++ b/src/utils/sandbox.ts @@ -15,3 +15,14 @@ export async function pauseSandbox( const sandboxClient = sandbox ?? (await Sandbox.connect(sandboxSessionId)); return await sandboxClient.pause(); } + +/** + * Resumes the sandbox. + * @param sandboxSessionId The ID of the sandbox to resume. + * @returns The sandbox client. + */ +export async function resumeSandbox( + sandboxSessionId: string, +): Promise { + return await Sandbox.resume(sandboxSessionId); +}