diff --git a/src/subgraphs/planner/index.ts b/src/subgraphs/planner/index.ts index 51bc0c66..0c42424e 100644 --- a/src/subgraphs/planner/index.ts +++ b/src/subgraphs/planner/index.ts @@ -8,7 +8,12 @@ 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 { + generateAction, + generatePlan, + summarizer, + takeAction, +} from "./nodes/index.js"; import { isAIMessage, ToolMessage } from "@langchain/core/messages"; function takeActionOrGeneratePlan( @@ -52,6 +57,7 @@ const workflow = new StateGraph(PlannerGraphStateObj, GraphConfiguration) .addNode("generate-plan-context-action", generateAction) .addNode("take-plan-action", takeAction) .addNode("generate-plan", generatePlan) + .addNode("summarizer", summarizer) .addEdge(START, "generate-plan-context-action") .addConditionalEdges( "generate-plan-context-action", @@ -59,7 +65,8 @@ const workflow = new StateGraph(PlannerGraphStateObj, GraphConfiguration) ["take-plan-action", "generate-plan"], ) .addEdge("take-plan-action", "generate-plan-context-action") - .addEdge("generate-plan", END); + .addEdge("generate-plan", "summarizer") + .addEdge("summarizer", END); // TODO: Fix zod types export const plannerGraph = workflow.compile() as any; diff --git a/src/subgraphs/planner/nodes/index.ts b/src/subgraphs/planner/nodes/index.ts index f99d35ec..26a054e9 100644 --- a/src/subgraphs/planner/nodes/index.ts +++ b/src/subgraphs/planner/nodes/index.ts @@ -1,3 +1,4 @@ export * from "./generate-message.js"; export * from "./take-action.js"; export * from "./generate-plan.js"; +export * from "./summarizer.js"; diff --git a/src/subgraphs/planner/nodes/summarizer.ts b/src/subgraphs/planner/nodes/summarizer.ts new file mode 100644 index 00000000..3efa2bde --- /dev/null +++ b/src/subgraphs/planner/nodes/summarizer.ts @@ -0,0 +1,78 @@ +import { z } from "zod"; +import { GraphConfig } from "../../../types.js"; +import { PlannerGraphState, PlannerGraphUpdate } from "../types.js"; +import { loadModel, Task } from "../../../utils/load-model.js"; +import { isHumanMessage, ToolMessage } from "@langchain/core/messages"; +import { getMessageContentString } from "../../../utils/message-content.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. + +You've been given a task to summarize the messages in your conversation history. You just finished gathering context to be used when generating an development plan to address the user's request. +You do not want to keep the entire conversation history, but instead you want to keep the most relevant and important snippets for future context + +You MUST adhere to the following criteria when summarizing the conversation history: +- Retain context such as file paths, versions, and installed software. +- Do not retain any full code snippets. +- Do not retain any full file contents. +- Ensure your summary is concise, but useful for future context. + +Here is the user's initial request +## User request: +{USER_REQUEST} + +With all of this in mind, please carefully summarize and condense the following conversation history. Ensure you pass this condensed context to the \`condense_planning_context\` tool. +`; + +const formatPrompt = (userRequest: string): string => + systemPrompt.replace("{USER_REQUEST}", userRequest); + +const condenseContextToolSchema = z.object({ + context: z + .string() + .describe("The condensed context to be used when generating a plan."), +}); +const condenseContextTool = { + name: "condense_planning_context", + description: + "Condense the conversation history into a concise summary, while still retaining the most relevant and important snippets.", + schema: condenseContextToolSchema, +}; + +export async function summarizer( + state: PlannerGraphState, + config: GraphConfig, +): Promise { + const model = await loadModel(config, Task.PLANNER); + const modelWithTools = model.bindTools([condenseContextTool], { + tool_choice: condenseContextTool.name, + }); + + const firstUserMessage = state.messages.find(isHumanMessage); + + const response = await modelWithTools.invoke([ + { + role: "system", + content: formatPrompt( + getMessageContentString( + firstUserMessage?.content ?? "No user request provided.", + ), + ), + }, + ...state.plannerMessages, + ]); + + const toolCall = response.tool_calls?.[0]; + if (!toolCall) { + throw new Error("Failed to generate plan"); + } + + const toolMessage = new ToolMessage({ + tool_call_id: toolCall.id ?? "", + name: toolCall.name, + content: `Successfully summarized planning context.`, + }); + + return { + messages: [response, toolMessage], + }; +} diff --git a/src/utils/message-content.ts b/src/utils/message-content.ts new file mode 100644 index 00000000..566ab0af --- /dev/null +++ b/src/utils/message-content.ts @@ -0,0 +1,10 @@ +import { MessageContent } from "@langchain/core/messages"; + +export function getMessageContentString(content: MessageContent): string { + if (typeof content === "string") return content; + + return content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map((c) => c.text) + .join(" "); +}