diff --git a/eslint.config.js b/eslint.config.js index 1dbc09bd..52aa9d97 100644 --- a/eslint.config.js +++ b/eslint.config.js @@ -1,12 +1,12 @@ -import js from "@eslint/js"; +import eslint from "@eslint/js"; import globals from "globals"; import tseslint from "typescript-eslint"; import eslintImport from "eslint-plugin-import"; import eslintNoInstanceof from "eslint-plugin-no-instanceof"; export default tseslint.config( + eslint.configs.recommended, tseslint.configs.recommendedTypeChecked, - js.configs.recommended, { ignores: ["dist", "eslint.config.js"] }, { files: ["**/*.ts"], diff --git a/src/index.ts b/src/index.ts index a57f8bcd..031fef13 100644 --- a/src/index.ts +++ b/src/index.ts @@ -6,30 +6,32 @@ import { generateAction, takeAction, rewritePlan, + interruptPlan, } from "./nodes/index.js"; -import { isAIMessage, isToolMessage } from "@langchain/core/messages"; +import { isAIMessage } from "@langchain/core/messages"; /** - * After generating a plan, ensure there is an approved tool message. - * If there is, route to the generate-action node. Otherwise, route to the rewrite-plan node. + * @param {GraphState} state - The current graph state. + * @returns {"interrupt-plan" | typeof END} The next node to execute, or END if the process should stop. */ -function routeAfterPlan(state: GraphState): "generate-action" | "rewrite-plan" { +function routeAfterPlan(state: GraphState): "interrupt-plan" | typeof END { const { messages } = state; - // Search for a tool message responding to the "session_plan" tool call where the content is "approved" - const planApprovedMessage = messages.find( - (m) => - isToolMessage(m) && m.name === "session_plan" && m.content === "approved", - ); - if (planApprovedMessage) { - return "generate-action"; + const lastMessage = messages[messages.length - 1]; + if (isAIMessage(lastMessage) && !lastMessage.tool_calls) { + // The last message is an AI message without tool calls. This indicates the LLM generated followup questions. + return END; } - // If this does not exist, we should rewrite the plan. - return "rewrite-plan"; + + return "interrupt-plan"; } /** - * After taking action, ensure there is an AI message with tool calls. - * If there is, route to the take-action node. Otherwise, end the graph. + * Routes to the next appropriate node after taking action. + * If the last message is an AI message with tool calls, it routes to "take-action". + * Otherwise, it ends the process. + * + * @param {GraphState} state - The current graph state. + * @returns {typeof END | "take-action"} The next node to execute, or END if the process should stop. */ function takeActionOrEnd(state: GraphState): typeof END | "take-action" { const { messages } = state; @@ -46,18 +48,18 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration) .addNode("initialize", initialize) .addNode("generate-plan", generatePlan) .addNode("rewrite-plan", rewritePlan) + .addNode("interrupt-plan", interruptPlan, { + // TODO: Hookup `Command` in interruptPlan node so this actually works. + ends: [END, "rewrite-plan", "generate-action"], + }) .addNode("generate-action", generateAction) .addNode("take-action", takeAction) .addEdge(START, "initialize") .addEdge("initialize", "generate-plan") - .addConditionalEdges("generate-plan", routeAfterPlan, [ - "generate-action", - "rewrite-plan", - ]) - .addConditionalEdges("rewrite-plan", routeAfterPlan, [ - "generate-action", - "rewrite-plan", - ]) + // TODO: Update routing to work w/ new interrupt node. + .addConditionalEdges("generate-plan", routeAfterPlan, ["interrupt-plan", END]) + // Always interrupt after rewriting the plan. + .addEdge("rewrite-plan", "interrupt-plan") .addEdge("generate-plan", "generate-action") .addConditionalEdges("generate-action", takeActionOrEnd, ["take-action", END]) .addEdge("take-action", "generate-action"); diff --git a/src/nodes/generate-message.ts b/src/nodes/generate-message.ts index 3f7f5780..7c2dcf78 100644 --- a/src/nodes/generate-message.ts +++ b/src/nodes/generate-message.ts @@ -64,10 +64,6 @@ You MUST adhere to the following criteria when executing the task: ${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 async function generateAction( state: GraphState, config: GraphConfig, diff --git a/src/nodes/generate-plan.ts b/src/nodes/generate-plan.ts index d56e0027..5fb0821a 100644 --- a/src/nodes/generate-plan.ts +++ b/src/nodes/generate-plan.ts @@ -1,8 +1,48 @@ +import { sessionPlanTool } from "../tools/index.js"; import { GraphState, GraphConfig, GraphUpdate } from "../types.js"; +import { loadModel } from "../utils/load-model.js"; -export function generatePlan( +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 { - throw new Error("Not implemented"); + const model = await loadModel(config); + 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 e24cfae3..ffafeb47 100644 --- a/src/nodes/index.ts +++ b/src/nodes/index.ts @@ -3,3 +3,4 @@ export * from "./initialize.js"; export * from "./generate-message.js"; export * from "./take-action.js"; export * from "./rewrite-plan.js"; +export * from "./interrupt-plan.js"; diff --git a/src/nodes/interrupt-plan.ts b/src/nodes/interrupt-plan.ts new file mode 100644 index 00000000..6094da73 --- /dev/null +++ b/src/nodes/interrupt-plan.ts @@ -0,0 +1,82 @@ +import { Command, END, interrupt } from "@langchain/langgraph"; +import { GraphState } from "../types.js"; +import { + ActionRequest, + HumanInterrupt, + HumanResponse, +} from "@langchain/langgraph/prebuilt"; +import { v4 as uuidv4 } from "uuid"; + +export function interruptPlan(state: GraphState): Command { + const { proposedPlan } = state; + if (!proposedPlan.length) { + throw new Error("No proposed plan found."); + } + + const interruptRes = interrupt({ + action_request: { + action: "Approve/Edit Plan", + args: { + plan: proposedPlan.join("\n:::\n"), + }, + }, + config: { + allow_accept: true, + allow_edit: true, + allow_respond: true, + allow_ignore: true, + }, + description: `A new plan has been generated for your request. Please review it and either approve it, edit it, respond to it, or ignore it. Responses will be passed to an LLM where it will rewrite then plan. + If editing the plan, ensure each step in the plan is separated by ":::".`, + })[0]; + + if (interruptRes.type === "accept") { + // Plan was accepted, route to the generate action node. + return new Command({ + goto: "generate-action", + update: { + plan: proposedPlan.map((p) => ({ + id: uuidv4(), + plan: p, + completed: false, + })), + }, + }); + } + + if (interruptRes.type === "edit") { + // Plan was edited, route to the generate action node. + const editedPlan = (interruptRes.args as ActionRequest).args.plan + .split(":::") + .map((step: string) => step.trim()); + return new Command({ + goto: "generate-action", + update: { + plan: editedPlan.map((p: string) => ({ + id: uuidv4(), + plan: p, + completed: false, + })), + }, + }); + } + + if (interruptRes.type === "response") { + // Plan was responded to, route to the rewrite plan node. + return new Command({ + goto: "rewrite-plan", + update: { + planChangeRequest: interruptRes.args as string, + }, + }); + } + + if (interruptRes.type === "ignore") { + // Plan was ignored, end the process. + return new Command({ + goto: END, + }); + } + + throw new Error("Unknown interrupt type." + interruptRes.type); +} diff --git a/src/nodes/rewrite-plan.ts b/src/nodes/rewrite-plan.ts index de3384f3..e4cc3ed7 100644 --- a/src/nodes/rewrite-plan.ts +++ b/src/nodes/rewrite-plan.ts @@ -1,12 +1,64 @@ import { GraphState, GraphConfig, GraphUpdate } from "../types.js"; +import { loadModel } from "../utils/load-model.js"; +import { sessionPlanTool } from "../tools/index.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 rewritePlan( +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 rewrite a high-level plan to address the user's initial request. In a previous step you generated a plan, however the user has requested some changes: +## User Request +{USER_REQUEST} + +Here is the previous plan: +## Previous Plan +{PREVIOUS_PLAN} + +The plan must be a list of actions to take, in order, to address the user's request. You should not include any code in the plan, only a list of actions to take. + +You MUST adhere to the following criteria when generating the plan: +- You do not have access to the codebase yet, so you cannot inspect it or make assumptions about it. +- Your plan should be high-level in nature, but should still be specific enough to be actionable. +- Make as few changes as possible to the previous plan, while still addressing the user's request. +- When you are ready to generate the plan, ensure you call the 'session_plan' tool. +- Ensure you generate the full plan in this tool call, not just the changes. +`; + +const formatSysPrompt = (userRequest: string, previousPlan: string) => { + return systemPrompt + .replace("{USER_REQUEST}", userRequest) + .replace("{PREVIOUS_PLAN}", previousPlan); +}; + +export async function rewritePlan( state: GraphState, config: GraphConfig, ): Promise { - throw new Error("Not implemented"); + if (!state.planChangeRequest) { + throw new Error("No plan change request found."); + } + + const model = await loadModel(config); + const modelWithTools = model.bindTools([sessionPlanTool], { + // The model should always call the tool when rewriting the plan. + tool_choice: sessionPlanTool.name, + }); + + const response = await modelWithTools.invoke([ + { + role: "system", + content: formatSysPrompt( + state.planChangeRequest, + " - " + state.plan.join("\n - "), + ), + }, + ...state.messages, + ]); + + if (response.tool_calls?.length) { + return { + proposedPlan: response.tool_calls[0].args.plan, + plan: [], + }; + } + + throw new Error("Failed to rewrite plan."); } diff --git a/src/nodes/take-action.ts b/src/nodes/take-action.ts index f67f61ce..26576000 100644 --- a/src/nodes/take-action.ts +++ b/src/nodes/take-action.ts @@ -2,10 +2,6 @@ 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 async function takeAction( state: GraphState, _config: GraphConfig, diff --git a/src/tools/index.ts b/src/tools/index.ts index 79a06495..2c9f6230 100644 --- a/src/tools/index.ts +++ b/src/tools/index.ts @@ -1,2 +1,3 @@ export * from "./apply-patch.js"; export * from "./shell.js"; +export * from "./session-plan.js"; diff --git a/src/tools/session-plan.ts b/src/tools/session-plan.ts new file mode 100644 index 00000000..8b9d4960 --- /dev/null +++ b/src/tools/session-plan.ts @@ -0,0 +1,11 @@ +import { z } from "zod"; + +const sessionPlanSchema = z.object({ + plan: z.array(z.string()).describe("The plan to address the user's request."), +}); + +export const sessionPlanTool = { + name: "session_plan", + description: "Call this tool when you are ready to generate a plan.", + schema: sessionPlanSchema, +}; diff --git a/src/types.ts b/src/types.ts index 950d05fd..242b518a 100644 --- a/src/types.ts +++ b/src/types.ts @@ -20,10 +20,18 @@ export type TargetRepository = { export const GraphAnnotation = Annotation.Root({ messages: MessagesAnnotation.spec.messages, + proposedPlan: Annotation({ + reducer: (_state, update) => update, + default: () => [], + }), plan: Annotation({ reducer: (_state, update) => update, default: () => [], }), + planChangeRequest: Annotation({ + reducer: (_state, update) => update, + default: () => undefined, + }), }); export type GraphState = typeof GraphAnnotation.State;