feat: Add explicit step for generating conclusion and ending (#19)

src/index.ts

@@ -8,6 +8,7 @@ import {
   interruptPlan,
   progressPlanStep,
   summarizeTaskSteps,
+  generateConclusion,
 } from "./nodes/index.js";
 import { isAIMessage } from "@langchain/core/messages";
 import { plannerGraph } from "./subgraphs/index.js";
@@ -62,7 +63,10 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
   .addNode("progress-plan-step", progressPlanStep, {
     ends: ["summarize-task-steps", "generate-action"],
   })
-  .addNode("summarize-task-steps", summarizeTaskSteps)
+  .addNode("summarize-task-steps", summarizeTaskSteps, {
+    ends: ["generate-action", "generate-conclusion"],
+  })
+  .addNode("generate-conclusion", generateConclusion)
   .addEdge(START, "initialize")
   .addEdge("initialize", "generate-plan-subgraph")
   // TODO: Update routing to work w/ new interrupt node.
@@ -74,7 +78,7 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
   .addEdge("rewrite-plan", "interrupt-plan")
   .addConditionalEdges("generate-action", takeActionOrEnd, ["take-action", END])
   .addEdge("take-action", "progress-plan-step")
-  .addEdge("summarize-task-steps", "generate-action");
+  .addEdge("generate-conclusion", END);
 
 // Zod types are messed up
 export const graph = workflow.compile() as any;

src/nodes/generate-conclusion.ts

@@ -0,0 +1,62 @@
+import { isHumanMessage } from "@langchain/core/messages";
+import { GraphConfig, GraphState, GraphUpdate, PlanItem } from "../types.js";
+import { loadModel, Task } from "../utils/load-model.js";
+import {
+  getMessageContentString,
+  getMessageString,
+} from "../utils/message/content.js";
+import { createLogger, LogLevel } from "../utils/logger.js";
+
+const logger = createLogger(LogLevel.INFO, "GenerateConclusionNode");
+
+const prompt = `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 have just completed all of the tasks in the plan:
+{COMPLETED_TASKS}
+
+Since you've successfully completed the user's request, you should now generate a short, concise concision. It can be helpful here to outline all of the changes you've made to the codebase, any additional steps you think the user should take, any relevant informatioon from the conversation hostiry below, etc.
+Your concision message should be concise and to the point, you do NOT want to include any details which are not ABSOLUTELY NECESSARY.
+`;
+
+const formatPrompt = (plan: PlanItem[]): string => {
+  return prompt.replace(
+    "{COMPLETED_TASKS}",
+    plan.map((p) => `${p.index}. ${p.plan}`).join("\n"),
+  );
+};
+
+export async function generateConclusion(
+  state: GraphState,
+  config: GraphConfig,
+): Promise<GraphUpdate> {
+  const model = await loadModel(config, Task.SUMMARIZER);
+
+  const firstUserMessage = state.messages.find(isHumanMessage);
+
+  const userMessage = `The user's initial request is as follows:
+${getMessageContentString(firstUserMessage?.content ?? "No user message found")}
+
+The conversation history is as follows:
+${state.messages.map(getMessageString).join("\n")}
+
+Given all of this, please respond with the concise conclusion. Do not include any additional text besides the conclusion.`;
+
+  logger.info("Generating conclusion");
+
+  const response = await model.invoke([
+    {
+      role: "system",
+      content: formatPrompt(state.plan),
+    },
+    {
+      role: "user",
+      content: userMessage,
+    },
+  ]);
+
+  logger.info("✅ Successfully generated conclusion. Ending run. 👋");
+
+  return {
+    messages: [response],
+  };
+}

src/nodes/generate-message.ts

@@ -19,7 +19,12 @@ You work based on a plan which was generated in a previous step. The plan items
 
 {PLAN_PROMPT}
 
-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.
+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:
+First, read through the conversation history to see if you have already searched for the file or information you need. Pay extra close attention to the condensed context tool call messages in the conversation history. These contain summarized/condensed context from previously completed steps. Ensure you always read these messages to avoid duplicate work (e.g.: searching for file paths).
+If you are still not sure, 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.
 

src/nodes/index.ts

@@ -5,3 +5,4 @@ export * from "./rewrite-plan.js";
 export * from "./interrupt-plan.js";
 export * from "./progress-plan-step.js";
 export * from "./summarize-task-steps.js";
+export * from "./generate-conclusion.js";

src/nodes/summarize-task-steps.ts

@@ -13,6 +13,7 @@ import {
   removeFirstHumanMessage,
   removeLastTaskMessages,
 } from "../utils/message/modify-array.js";
+import { Command } from "@langchain/langgraph";
 
 const logger = createLogger(LogLevel.INFO, "SummarizeTaskSteps");
 
@@ -56,7 +57,7 @@ const condenseContextTool = {
 export async function summarizeTaskSteps(
   state: GraphState,
   config: GraphConfig,
-): Promise<GraphUpdate> {
+): Promise<Command> {
   const model = await loadModel(config, Task.SUMMARIZER);
   const modelWithTools = model.bindTools([condenseContextTool], {
     tool_choice: condenseContextTool.name,
@@ -100,17 +101,34 @@ Given this full conversation history please generate a concise, and useful summa
 
   const removedMessages = removeLastTaskMessages(state.messages);
   logger.info(`Removing ${removedMessages.length} message(s) from state.`);
-  return {
-    messages: [
-      ...removedMessages,
-      new AIMessage({
-        ...response,
-        additional_kwargs: {
-          ...response.additional_kwargs,
-          summary_message: true,
-        },
-      }),
-      toolMessage,
-    ],
-  };
+
+  const allTasksCompleted = state.plan.every((p) => p.completed);
+
+  const newMessagesStateUpdate = [
+    ...removedMessages,
+    new AIMessage({
+      ...response,
+      additional_kwargs: {
+        ...response.additional_kwargs,
+        summary_message: true,
+      },
+    }),
+    toolMessage,
+  ];
+
+  if (!allTasksCompleted) {
+    return new Command({
+      goto: "generate-action",
+      update: {
+        messages: newMessagesStateUpdate,
+      },
+    });
+  }
+
+  return new Command({
+    goto: "generate-conclusion",
+    update: {
+      messages: newMessagesStateUpdate,
+    },
+  });
 }