diff --git a/examples/src/createAgent/middleware/contextEditing.ts b/examples/src/createAgent/middleware/contextEditing.ts new file mode 100644 index 000000000000..eb75e8d3f4d7 --- /dev/null +++ b/examples/src/createAgent/middleware/contextEditing.ts @@ -0,0 +1,208 @@ +/** + * Example demonstrating the Context Editing Middleware. + * + * This middleware automatically clears older tool results when the conversation + * grows beyond a configurable token threshold, helping manage context size. + */ + +import { ChatOpenAI } from "@langchain/openai"; +import { tool } from "@langchain/core/tools"; +import { z } from "zod"; +import { HumanMessage } from "@langchain/core/messages"; +import { + createAgent, + contextEditingMiddleware, + ToolMessage, + ClearToolUsesEdit, +} from "langchain"; + +/** + * Define a simple search tool that returns large results + */ +const searchTool = tool( + async ({ query }) => { + /** + * Simulate a large search result (each repeat is ~30 characters = ~7-8 tokens) + * 100 repeats = ~750 tokens + */ + return `Search results for "${query}":\n${"Lorem ipsum dolor sit amet consectetur. ".repeat( + 100 + )}`; + }, + { + name: "search", + description: "Search for information", + schema: z.object({ + query: z.string().describe("The search query"), + }), + } +); + +/** + * Define a calculator tool + */ +const calculatorTool = tool( + async ({ a, b }) => { + return `The result of ${a} + ${b} is ${a + b}`; + }, + { + name: "calculator", + description: "Perform simple calculations", + schema: z.object({ + a: z.number().describe("First number"), + b: z.number().describe("Second number"), + }), + } +); + +/** + * Create an agent with context editing middleware + */ +const agent = createAgent({ + model: new ChatOpenAI({ model: "gpt-4o-mini" }), + tools: [searchTool, calculatorTool], + middleware: [ + contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + /** + * Trigger clearing when context exceeds 500 tokens (low threshold for demo) + */ + triggerTokens: 500, + /** + * Clear at least 100 tokens when triggered + */ + clearAtLeast: 100, + /** + * Keep only the 1 most recent tool result + */ + keep: 1, + /** + * Don't clear calculator tool results + */ + excludeTools: ["calculator"], + /** + * Clear tool inputs as well to save more space + */ + clearToolInputs: true, + }), + ], + /** + * Use approximate token counting (faster for demo) + */ + tokenCountMethod: "approx", + }), + ], + systemPrompt: "You are a helpful assistant. Use tools when needed.", +}); + +console.log("=== Context Editing Middleware Example ===\n"); + +/** + * First turn: Search for React agents + */ +console.log("Turn 1: Searching for React agents..."); +let result = await agent.invoke({ + messages: [ + { + role: "user", + content: "Search for 'React agents'", + }, + ], +}); + +console.log(`- Accumulated ${result.messages.length} messages`); +console.log( + `- Last message: ${result.messages[result.messages.length - 1].content + .toString() + .slice(0, 100)}...` +); + +/** + * Second turn: Search for LangChain + * This should trigger context editing since we'll exceed 500 tokens + */ +console.log("\nTurn 2: Searching for LangChain..."); +result = await agent.invoke({ + messages: result.messages.concat([ + new HumanMessage("Now search for 'LangChain'"), + ]), +}); + +console.log(`- Accumulated ${result.messages.length} messages`); + +/** + * Third turn: One more search + * This should definitely trigger clearing + */ +console.log("\nTurn 3: Searching for TypeScript..."); +result = await agent.invoke({ + messages: result.messages.concat([ + new HumanMessage("Finally, search for 'TypeScript' and calculate 42 + 58"), + ]), +}); + +console.log(`- Accumulated ${result.messages.length} messages`); +console.log( + `\nFinal response: ${result.messages[result.messages.length - 1].content + .toString() + .slice(0, 150)}...` +); + +/** + * Check how many tool results were cleared + */ +const clearedMessages = result.messages.filter( + (msg) => + ToolMessage.isInstance(msg) && + (msg.response_metadata as any)?.context_editing?.cleared +); + +console.log("\n=== Context Editing Results ==="); +console.log(`Total messages: ${result.messages.length}`); +console.log(`Tool results cleared: ${clearedMessages.length}`); + +if (clearedMessages.length > 0) { + console.log("\nCleared tool messages:"); + clearedMessages.forEach((msg, idx) => { + const toolMsg = msg as typeof ToolMessage.prototype; + console.log( + ` ${idx + 1}. Tool call ID: ${toolMsg.tool_call_id}, Content: "${ + toolMsg.content + }"` + ); + }); +} else { + console.log( + "\nNote: No tool results were cleared. Try lowering triggerTokens or adding more turns." + ); +} + +/** + * Final response: + * ``` + * === Context Editing Middleware Example === + * + * Turn 1: Searching for React agents... + * - Accumulated 4 messages + * - Last message: It seems that the search results for "React agents" returned generic placeholder text rather than sp... + * + * Turn 2: Searching for LangChain... + * - Accumulated 8 messages + * + * Turn 3: Searching for TypeScript... + * - Accumulated 13 messages + * + * Final response: The search for "TypeScript" again returned generic placeholder text, which is not informative. + * + * However, the calculation result for \( 42 + 58 \) is \... + * + * === Context Editing Results === + * Total messages: 13 + * Tool results cleared: 2 + * + * Cleared tool messages: + * 1. Tool call ID: call_M9KmodGVBjPktgjM8tkaNjWG, Content: "[cleared]" + * 2. Tool call ID: call_jIMx9NyquDdVqjs0QjDT1EAI, Content: "[cleared]" + * ``` + */ diff --git a/libs/langchain/src/agents/middlewareAgent/middleware/contextEditing.ts b/libs/langchain/src/agents/middlewareAgent/middleware/contextEditing.ts new file mode 100644 index 000000000000..6d86f34e22c8 --- /dev/null +++ b/libs/langchain/src/agents/middlewareAgent/middleware/contextEditing.ts @@ -0,0 +1,509 @@ +/** + * Context editing middleware. + * + * This middleware mirrors Anthropic's context editing capabilities by clearing + * older tool results once the conversation grows beyond a configurable token + * threshold. The implementation is intentionally model-agnostic so it can be used + * with any LangChain chat model. + */ + +import type { BaseMessage } from "@langchain/core/messages"; +import type { LanguageModelLike } from "@langchain/core/language_models/base"; +import { + AIMessage, + ToolMessage, + SystemMessage, +} from "@langchain/core/messages"; + +import { countTokensApproximately } from "./utils.js"; +import { createMiddleware } from "../middleware.js"; +import type { ModelRequest } from "../types.js"; + +const DEFAULT_TOOL_PLACEHOLDER = "[cleared]"; + +/** + * Function type for counting tokens in a sequence of messages. + */ +export type TokenCounter = ( + messages: BaseMessage[] +) => number | Promise; + +/** + * Protocol describing a context editing strategy. + * + * Implement this interface to create custom strategies for managing + * conversation context size. The `apply` method should modify the + * messages array in-place and return the updated token count. + * + * @example + * ```ts + * import { SystemMessage } from "langchain"; + * + * class RemoveOldSystemMessages implements ContextEdit { + * async apply({ tokens, messages, countTokens }) { + * // Remove old system messages if over limit + * if (tokens > 50000) { + * messages = messages.filter(SystemMessage.isInstance); + * return await countTokens(messages); + * } + * return tokens; + * } + * } + * ``` + */ +export interface ContextEdit { + /** + * Apply an edit to the message list, returning the new token count. + * + * This method should: + * 1. Check if editing is needed based on `tokens` parameter + * 2. Modify the `messages` array in-place (if needed) + * 3. Return the new token count after modifications + * + * @param params - Parameters for the editing operation + * @returns The updated token count after applying edits + */ + apply(params: { + /** + * Current token count of all messages + */ + tokens: number; + /** + * Array of messages to potentially edit (modify in-place) + */ + messages: BaseMessage[]; + /** + * Function to count tokens in a message array + */ + countTokens: TokenCounter; + }): number | Promise; +} + +/** + * Configuration for clearing tool outputs when token limits are exceeded. + */ +export interface ClearToolUsesEditConfig { + /** + * Token count that triggers the edit. + * @default 100000 + */ + triggerTokens?: number; + + /** + * Minimum number of tokens to reclaim when the edit runs. + * @default 0 + */ + clearAtLeast?: number; + + /** + * Number of most recent tool results that must be preserved. + * @default 3 + */ + keep?: number; + + /** + * Whether to clear the originating tool call parameters on the AI message. + * @default false + */ + clearToolInputs?: boolean; + + /** + * List of tool names to exclude from clearing. + * @default [] + */ + excludeTools?: string[]; + + /** + * Placeholder text inserted for cleared tool outputs. + * @default "[cleared]" + */ + placeholder?: string; +} + +/** + * Strategy for clearing tool outputs when token limits are exceeded. + * + * This strategy mirrors Anthropic's `clear_tool_uses_20250919` behavior by + * replacing older tool results with a placeholder text when the conversation + * grows too large. It preserves the most recent tool results and can exclude + * specific tools from being cleared. + * + * @example + * ```ts + * import { ClearToolUsesEdit } from "langchain"; + * + * const edit = new ClearToolUsesEdit({ + * triggerTokens: 100000, // Start clearing at 100K tokens + * clearAtLeast: 0, // Clear as much as needed + * keep: 3, // Always keep 3 most recent results + * excludeTools: ["important"], // Never clear "important" tool + * clearToolInputs: false, // Keep tool call arguments + * placeholder: "[cleared]", // Replacement text + * }); + * ``` + */ +export class ClearToolUsesEdit implements ContextEdit { + triggerTokens: number; + clearAtLeast: number; + keep: number; + clearToolInputs: boolean; + excludeTools: Set; + placeholder: string; + + constructor(config: ClearToolUsesEditConfig = {}) { + this.triggerTokens = config.triggerTokens ?? 100000; + this.clearAtLeast = config.clearAtLeast ?? 0; + this.keep = config.keep ?? 3; + this.clearToolInputs = config.clearToolInputs ?? false; + this.excludeTools = new Set(config.excludeTools ?? []); + this.placeholder = config.placeholder ?? DEFAULT_TOOL_PLACEHOLDER; + } + + async apply(params: { + tokens: number; + messages: BaseMessage[]; + countTokens: TokenCounter; + }): Promise { + const { tokens, messages, countTokens } = params; + + if (tokens <= this.triggerTokens) { + return tokens; + } + + /** + * Find all tool message candidates with their actual indices in the messages array + */ + const candidates: Array<{ idx: number; msg: ToolMessage }> = []; + for (let i = 0; i < messages.length; i++) { + const msg = messages[i]; + if (ToolMessage.isInstance(msg)) { + candidates.push({ idx: i, msg }); + } + } + + /** + * Keep the most recent tool messages + */ + const candidatesToClear = + this.keep >= candidates.length + ? [] + : this.keep > 0 + ? candidates.slice(0, -this.keep) + : candidates; + + let clearedTokens = 0; + for (const { idx, msg: toolMessage } of candidatesToClear) { + /** + * Stop if we've cleared enough tokens + */ + if (this.clearAtLeast > 0 && clearedTokens >= this.clearAtLeast) { + break; + } + + /** + * Skip if already cleared + */ + const contextEditing = toolMessage.response_metadata?.context_editing as + | { cleared?: boolean } + | undefined; + if (contextEditing?.cleared) { + continue; + } + + /** + * Find the corresponding AI message + */ + const aiMessage = this.#findAIMessageForToolCall( + messages.slice(0, idx), + toolMessage.tool_call_id + ); + + if (!aiMessage) { + continue; + } + + /** + * Find the corresponding tool call + */ + const toolCall = aiMessage.tool_calls?.find( + (call) => call.id === toolMessage.tool_call_id + ); + + if (!toolCall) { + continue; + } + + /** + * Skip if tool is excluded + */ + const toolName = toolMessage.name || toolCall.name; + if (this.excludeTools.has(toolName)) { + continue; + } + + /** + * Clear the tool message + */ + messages[idx] = new ToolMessage({ + tool_call_id: toolMessage.tool_call_id, + content: this.placeholder, + name: toolMessage.name, + artifact: undefined, + response_metadata: { + ...toolMessage.response_metadata, + context_editing: { + cleared: true, + strategy: "clear_tool_uses", + }, + }, + }); + + /** + * Optionally clear the tool inputs + */ + if (this.clearToolInputs) { + const aiMsgIdx = messages.indexOf(aiMessage); + if (aiMsgIdx >= 0) { + messages[aiMsgIdx] = this.#buildClearedToolInputMessage( + aiMessage, + toolMessage.tool_call_id + ); + } + } + + /** + * Recalculate tokens + */ + const newTokenCount = await countTokens(messages); + clearedTokens = Math.max(0, tokens - newTokenCount); + } + + return tokens - clearedTokens; + } + + #findAIMessageForToolCall( + previousMessages: BaseMessage[], + toolCallId: string + ): AIMessage | null { + // Search backwards through previous messages + for (let i = previousMessages.length - 1; i >= 0; i--) { + const msg = previousMessages[i]; + if (AIMessage.isInstance(msg)) { + const hasToolCall = msg.tool_calls?.some( + (call) => call.id === toolCallId + ); + if (hasToolCall) { + return msg; + } + } + } + return null; + } + + #buildClearedToolInputMessage( + message: AIMessage, + toolCallId: string + ): AIMessage { + const updatedToolCalls = message.tool_calls?.map((toolCall) => { + if (toolCall.id === toolCallId) { + return { ...toolCall, args: {} }; + } + return toolCall; + }); + + const metadata = { ...message.response_metadata }; + const contextEntry = { + ...(metadata.context_editing as Record), + }; + + const clearedIds = new Set( + contextEntry.cleared_tool_inputs as string[] | undefined + ); + clearedIds.add(toolCallId); + contextEntry.cleared_tool_inputs = Array.from(clearedIds).sort(); + metadata.context_editing = contextEntry; + + return new AIMessage({ + content: message.content, + tool_calls: updatedToolCalls, + response_metadata: metadata, + id: message.id, + name: message.name, + additional_kwargs: message.additional_kwargs, + }); + } +} + +/** + * Configuration for the Context Editing Middleware. + */ +export interface ContextEditingMiddlewareConfig { + /** + * Sequence of edit strategies to apply. Defaults to a single + * ClearToolUsesEdit mirroring Anthropic defaults. + */ + edits?: ContextEdit[]; + + /** + * Whether to use approximate token counting (faster, less accurate) + * or exact counting implemented by the chat model (potentially slower, more accurate). + * Currently only OpenAI models support exact counting. + * @default "approx" + */ + tokenCountMethod?: "approx" | "model"; +} + +/** + * Middleware that automatically prunes tool results to manage context size. + * + * This middleware applies a sequence of edits when the total input token count + * exceeds configured thresholds. By default, it uses the `ClearToolUsesEdit` strategy + * which mirrors Anthropic's `clear_tool_uses_20250919` behaviour by clearing older + * tool results once the conversation exceeds 100,000 tokens. + * + * ## Basic Usage + * + * Use the middleware with default settings to automatically manage context: + * + * @example Basic usage with defaults + * ```ts + * import { contextEditingMiddleware } from "langchain"; + * import { createAgent } from "langchain"; + * + * const agent = createAgent({ + * model: "anthropic:claude-3-5-sonnet", + * tools: [searchTool, calculatorTool], + * middleware: [ + * contextEditingMiddleware(), + * ], + * }); + * ``` + * + * The default configuration: + * - Triggers when context exceeds **100,000 tokens** + * - Keeps the **3 most recent** tool results + * - Uses **approximate token counting** (fast) + * - Does not clear tool call arguments + * + * ## Custom Configuration + * + * Customize the clearing behavior with `ClearToolUsesEdit`: + * + * @example Custom ClearToolUsesEdit configuration + * ```ts + * import { contextEditingMiddleware, ClearToolUsesEdit } from "langchain"; + * + * const agent = createAgent({ + * model: "anthropic:claude-3-5-sonnet", + * tools: [searchTool, calculatorTool], + * middleware: [ + * contextEditingMiddleware({ + * edits: [ + * new ClearToolUsesEdit({ + * triggerTokens: 50000, // Clear when exceeding 50K tokens + * clearAtLeast: 1000, // Reclaim at least 1K tokens + * keep: 5, // Keep 5 most recent tool results + * excludeTools: ["search"], // Never clear search results + * clearToolInputs: true, // Also clear tool call arguments + * }), + * ], + * tokenCountMethod: "approx", // Use approximate counting (or "model") + * }), + * ], + * }); + * ``` + * + * ## Custom Editing Strategies + * + * Implement your own context editing strategy by creating a class that + * implements the `ContextEdit` interface: + * + * @example Custom editing strategy + * ```ts + * import { contextEditingMiddleware, type ContextEdit, type TokenCounter } from "langchain"; + * import type { BaseMessage } from "@langchain/core/messages"; + * + * class CustomEdit implements ContextEdit { + * async apply(params: { + * tokens: number; + * messages: BaseMessage[]; + * countTokens: TokenCounter; + * }): Promise { + * // Implement your custom editing logic here + * // and apply it to the messages array, then + * // return the new token count after edits + * return countTokens(messages); + * } + * } + * ``` + * + * @param config - Configuration options for the middleware + * @returns A middleware instance that can be used with `createAgent` + */ +export function contextEditingMiddleware( + config: ContextEditingMiddlewareConfig = {} +) { + const edits = config.edits ?? [new ClearToolUsesEdit()]; + const tokenCountMethod = config.tokenCountMethod ?? "approx"; + + return createMiddleware({ + name: "ContextEditingMiddleware", + modifyModelRequest: async (request: ModelRequest) => { + if (!request.messages || request.messages.length === 0) { + return request; + } + + /** + * Use model's token counting method + */ + const systemMsg = request.systemPrompt + ? [new SystemMessage(request.systemPrompt)] + : []; + + const countTokens: TokenCounter = + tokenCountMethod === "approx" + ? countTokensApproximately + : async (messages: BaseMessage[]): Promise => { + const allMessages = [...systemMsg, ...messages]; + + /** + * Check if model has getNumTokensFromMessages method + * currently only OpenAI models have this method + */ + if ("getNumTokensFromMessages" in request.model) { + return ( + request.model as LanguageModelLike & { + getNumTokensFromMessages: ( + messages: BaseMessage[] + ) => Promise<{ + totalCount: number; + countPerMessage: number[]; + }>; + } + ) + .getNumTokensFromMessages(allMessages) + .then(({ totalCount }) => totalCount); + } + + throw new Error( + `Model "${request.model.getName()}" does not support token counting` + ); + }; + + let tokens = await countTokens(request.messages); + + /** + * Apply each edit in sequence + */ + for (const edit of edits) { + tokens = await edit.apply({ + tokens, + messages: request.messages, + countTokens, + }); + } + + return request; + }, + }); +} diff --git a/libs/langchain/src/agents/middlewareAgent/middleware/index.ts b/libs/langchain/src/agents/middlewareAgent/middleware/index.ts index 9adf708220b9..1d1184904c0a 100644 --- a/libs/langchain/src/agents/middlewareAgent/middleware/index.ts +++ b/libs/langchain/src/agents/middlewareAgent/middleware/index.ts @@ -1,6 +1,5 @@ export { summarizationMiddleware, - countTokensApproximately, type SummarizationMiddlewareConfig, } from "./summarization.js"; export { @@ -26,6 +25,14 @@ export { piiRedactionMiddleware, type PIIRedactionMiddlewareConfig, } from "./piiRedaction.js"; +export { + contextEditingMiddleware, + ClearToolUsesEdit, + type ContextEditingMiddlewareConfig, + type ContextEdit, + type ClearToolUsesEditConfig, + type TokenCounter, +} from "./contextEditing.js"; export { toolCallLimitMiddleware, ToolCallLimitExceededError, @@ -37,3 +44,4 @@ export { } from "./callLimit.js"; export { modelFallbackMiddleware } from "./modelFallback.js"; export { type AgentMiddleware } from "../types.js"; +export { countTokensApproximately } from "./utils.js"; diff --git a/libs/langchain/src/agents/middlewareAgent/middleware/summarization.ts b/libs/langchain/src/agents/middlewareAgent/middleware/summarization.ts index 4b92ad22b229..18c0ccb1c373 100644 --- a/libs/langchain/src/agents/middlewareAgent/middleware/summarization.ts +++ b/libs/langchain/src/agents/middlewareAgent/middleware/summarization.ts @@ -16,6 +16,7 @@ import { } from "@langchain/core/utils/types"; import { REMOVE_ALL_MESSAGES } from "@langchain/langgraph"; import { createMiddleware } from "../middleware.js"; +import { countTokensApproximately } from "./utils.js"; const DEFAULT_SUMMARY_PROMPT = ` Context Extraction Assistant @@ -71,34 +72,6 @@ export type SummarizationMiddlewareConfig = InferInteropZodInput< typeof contextSchema >; -/** - * Default token counter that approximates based on character count - * @param messages Messages to count tokens for - * @returns Approximate token count - */ -export function countTokensApproximately(messages: BaseMessage[]): number { - let totalChars = 0; - for (const msg of messages) { - let textContent: string; - if (typeof msg.content === "string") { - textContent = msg.content; - } else if (Array.isArray(msg.content)) { - textContent = msg.content - .map((item) => { - if (typeof item === "string") return item; - if (item.type === "text" && "text" in item) return item.text; - return ""; - }) - .join(""); - } else { - textContent = ""; - } - totalChars += textContent.length; - } - // Approximate 1 token = 4 characters - return Math.ceil(totalChars / 4); -} - /** * Summarization middleware that automatically summarizes conversation history when token limits are approached. * diff --git a/libs/langchain/src/agents/middlewareAgent/middleware/tests/contextEditing.test.ts b/libs/langchain/src/agents/middlewareAgent/middleware/tests/contextEditing.test.ts new file mode 100644 index 000000000000..1f2ccbd978cf --- /dev/null +++ b/libs/langchain/src/agents/middlewareAgent/middleware/tests/contextEditing.test.ts @@ -0,0 +1,658 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +import { describe, it, expect } from "vitest"; +import { HumanMessage, AIMessage, ToolMessage } from "@langchain/core/messages"; +import type { BaseMessage } from "@langchain/core/messages"; +import { + contextEditingMiddleware, + ClearToolUsesEdit, + type ContextEdit, + type TokenCounter, +} from "../contextEditing.js"; +import { createAgent } from "../../index.js"; +import { FakeToolCallingChatModel } from "../../../tests/utils.js"; + +describe("contextEditingMiddleware", () => { + /** + * Helper to create a conversation with tool calls + */ + function createToolCallConversation() { + const messages: BaseMessage[] = [ + new HumanMessage("Search for 'React'"), + new AIMessage({ + content: "I'll search for that.", + tool_calls: [ + { + id: "call_1", + name: "search", + args: { query: "React" }, + }, + ], + }), + new ToolMessage({ + content: "x".repeat(1000), // Large result + tool_call_id: "call_1", + }), + new AIMessage("Found React information."), + new HumanMessage("Now search for 'TypeScript'"), + new AIMessage({ + content: "I'll search for TypeScript.", + tool_calls: [ + { + id: "call_2", + name: "search", + args: { query: "TypeScript" }, + }, + ], + }), + new ToolMessage({ + content: "y".repeat(1000), // Large result + tool_call_id: "call_2", + }), + new AIMessage("Found TypeScript information."), + new HumanMessage("Search for 'JavaScript'"), + ]; + + return messages; + } + + describe("default behavior", () => { + it("should use default ClearToolUsesEdit with Anthropic defaults when no params provided", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Final response")], + }); + + const middleware = contextEditingMiddleware(); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + // Create a conversation that doesn't exceed default 100K token threshold + const messages = [ + new HumanMessage("Hello"), + new AIMessage({ + content: "Let me search.", + tool_calls: [ + { id: "call_1", name: "search", args: { query: "test" } }, + ], + }), + new ToolMessage({ + content: "Results", + tool_call_id: "call_1", + }), + new AIMessage("Done."), + new HumanMessage("Thanks"), + ]; + + const result = await agent.invoke({ messages }); + + // With default 100K token threshold, nothing should be cleared + const toolMessages = result.messages.filter(ToolMessage.isInstance); + const clearedMessages = toolMessages.filter( + (msg) => (msg.response_metadata as any)?.context_editing?.cleared + ); + + expect(clearedMessages.length).toBe(0); + }); + + it("should clear tool results when exceeding default trigger threshold", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Final response")], + }); + + // Create middleware with low threshold to test clearing + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 100, // Very low threshold + keep: 1, // Keep only 1 most recent + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages = createToolCallConversation(); + const result = await agent.invoke({ messages }); + + // Should clear the older tool message + const toolMessages = result.messages.filter(ToolMessage.isInstance); + const clearedMessages = toolMessages.filter( + (msg) => (msg.response_metadata as any)?.context_editing?.cleared + ); + + expect(clearedMessages.length).toBeGreaterThan(0); + + // Verify cleared message has placeholder + const clearedMsg = clearedMessages[0]; + expect(clearedMsg.content).toBe("[cleared]"); + expect( + (clearedMsg.response_metadata as any).context_editing.cleared + ).toBe(true); + expect( + (clearedMsg.response_metadata as any).context_editing.strategy + ).toBe("clear_tool_uses"); + }); + }); + + describe("custom ClearToolUsesEdit configuration", () => { + it("should respect custom triggerTokens threshold", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 50, // Very low threshold + keep: 0, // Clear all + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages = createToolCallConversation(); + const result = await agent.invoke({ messages }); + + const clearedMessages = result.messages.filter( + (msg) => + ToolMessage.isInstance(msg) && + (msg.response_metadata as any)?.context_editing?.cleared + ); + + // Should clear tool messages + expect(clearedMessages.length).toBeGreaterThan(0); + }); + + it("should keep the specified number of most recent tool results", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const keepCount = 1; + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 50, + keep: keepCount, + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages = createToolCallConversation(); + const result = await agent.invoke({ messages }); + + const toolMessages = result.messages.filter(ToolMessage.isInstance); + const unclearedMessages = toolMessages.filter( + (msg) => !(msg.response_metadata as any)?.context_editing?.cleared + ); + + // Should keep at least 'keepCount' uncleared tool messages + expect(unclearedMessages.length).toBeGreaterThanOrEqual(keepCount); + }); + + it("should exclude specified tools from clearing", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 50, + keep: 0, + excludeTools: ["important_search"], + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages: BaseMessage[] = [ + new HumanMessage("Search"), + new AIMessage({ + content: "Searching...", + tool_calls: [ + { + id: "call_1", + name: "important_search", + args: { query: "test" }, + }, + ], + }), + new ToolMessage({ + content: "x".repeat(1000), + tool_call_id: "call_1", + name: "important_search", + }), + new AIMessage("Found it."), + new HumanMessage("Search again"), + new AIMessage({ + content: "Searching...", + tool_calls: [ + { id: "call_2", name: "regular_search", args: { query: "test" } }, + ], + }), + new ToolMessage({ + content: "y".repeat(1000), + tool_call_id: "call_2", + name: "regular_search", + }), + new AIMessage("Done."), + new HumanMessage("One more"), + ]; + + const result = await agent.invoke({ messages }); + + // Find the excluded tool message + const excludedToolMsg = result.messages.find( + (msg) => ToolMessage.isInstance(msg) && msg.tool_call_id === "call_1" + ) as ToolMessage; + + // Should NOT be cleared (excluded tools are never cleared) + expect( + (excludedToolMsg.response_metadata as any)?.context_editing?.cleared + ).toBeFalsy(); + expect(excludedToolMsg.content).not.toBe("[cleared]"); + }); + + it("should clear tool inputs when clearToolInputs is true", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 50, + keep: 0, + clearToolInputs: true, + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages = createToolCallConversation(); + const result = await agent.invoke({ messages }); + + // Find an AI message whose tool output was cleared + const clearedToolMsg = result.messages.find( + (msg) => + ToolMessage.isInstance(msg) && + (msg.response_metadata as any)?.context_editing?.cleared + ) as ToolMessage; + + if (clearedToolMsg) { + // Find the corresponding AI message + const aiMsg = result.messages.find( + (msg) => + AIMessage.isInstance(msg) && + msg.tool_calls?.some((tc) => tc.id === clearedToolMsg.tool_call_id) + ) as AIMessage; + + if (aiMsg) { + const toolCall = aiMsg.tool_calls?.find( + (tc) => tc.id === clearedToolMsg.tool_call_id + ); + // Tool call args should be cleared + expect(toolCall?.args).toEqual({}); + expect( + (aiMsg.response_metadata as any)?.context_editing + ?.cleared_tool_inputs + ).toContain(clearedToolMsg.tool_call_id); + } + } + }); + + it("should use custom placeholder text", async () => { + const customPlaceholder = "[REDACTED]"; + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 50, + keep: 0, + placeholder: customPlaceholder, + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages = createToolCallConversation(); + const result = await agent.invoke({ messages }); + + const clearedMessages = result.messages.filter( + (msg) => + ToolMessage.isInstance(msg) && + (msg.response_metadata as any)?.context_editing?.cleared + ); + + if (clearedMessages.length > 0) { + expect(clearedMessages[0].content).toBe(customPlaceholder); + } + }); + + it("should clear at least the specified number of tokens", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 100, + clearAtLeast: 500, + keep: 0, + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages = createToolCallConversation(); + const result = await agent.invoke({ messages }); + + const clearedMessages = result.messages.filter( + (msg) => + ToolMessage.isInstance(msg) && + (msg.response_metadata as any)?.context_editing?.cleared + ); + + // Should clear messages until clearAtLeast tokens are reclaimed + expect(clearedMessages.length).toBeGreaterThan(0); + }); + }); + + describe("custom editing strategies", () => { + it("should support custom ContextEdit implementation", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + let customEditCalled = false; + + /** + * Custom strategy that removes all human messages over threshold + */ + class RemoveHumanMessages implements ContextEdit { + async apply(params: { + tokens: number; + messages: BaseMessage[]; + countTokens: TokenCounter; + }): Promise { + customEditCalled = true; + + if (params.tokens > 100) { + // Remove all human messages except the last one + const humanIndices: number[] = []; + params.messages.forEach((msg, idx) => { + if (HumanMessage.isInstance(msg)) { + humanIndices.push(idx); + } + }); + + // Keep only the last human message + if (humanIndices.length > 1) { + // Remove from the end backwards to maintain indices + for (let i = humanIndices.length - 2; i >= 0; i--) { + params.messages.splice(humanIndices[i], 1); + } + } + + return await params.countTokens(params.messages); + } + + return params.tokens; + } + } + + const middleware = contextEditingMiddleware({ + edits: [new RemoveHumanMessages()], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages = createToolCallConversation(); + const humanCountBefore = messages.filter(HumanMessage.isInstance).length; + + const result = await agent.invoke({ messages }); + + expect(customEditCalled).toBe(true); + + const humanCountAfter = result.messages.filter( + HumanMessage.isInstance + ).length; + + // Should have fewer human messages + expect(humanCountAfter).toBeLessThan(humanCountBefore); + }); + + it("should chain multiple editing strategies", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + let strategy1Called = false; + let strategy2Called = false; + + class Strategy1 implements ContextEdit { + async apply(params: { + tokens: number; + messages: BaseMessage[]; + countTokens: TokenCounter; + }): Promise { + strategy1Called = true; + return params.tokens; + } + } + + class Strategy2 implements ContextEdit { + async apply(params: { + tokens: number; + messages: BaseMessage[]; + countTokens: TokenCounter; + }): Promise { + strategy2Called = true; + return params.tokens; + } + } + + const middleware = contextEditingMiddleware({ + edits: [new Strategy1(), new Strategy2()], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages = [new HumanMessage("Hello"), new AIMessage("Hi there!")]; + + await agent.invoke({ messages }); + + // Both strategies should be called in sequence + expect(strategy1Called).toBe(true); + expect(strategy2Called).toBe(true); + }); + }); + + describe("token counting methods", () => { + it("should use approximate token counting by default", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const middleware = contextEditingMiddleware({ + tokenCountMethod: "approx", + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages = [ + new HumanMessage("Test message"), + new AIMessage("Response"), + ]; + + // Should not throw even without model token counting support + await expect(agent.invoke({ messages })).resolves.toBeDefined(); + }); + + it("should handle messages with empty content", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 50, + keep: 0, + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + const messages: BaseMessage[] = [ + new HumanMessage(""), + new AIMessage({ + content: "", + tool_calls: [{ id: "call_1", name: "tool", args: {} }], + }), + new ToolMessage({ + content: "", + tool_call_id: "call_1", + }), + ]; + + await expect(agent.invoke({ messages })).resolves.toBeDefined(); + }); + }); + + describe("edge cases", () => { + it("should not clear already cleared messages", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 50, + keep: 0, + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + // Create a message that's already marked as cleared + const messages: BaseMessage[] = [ + new HumanMessage("Test"), + new AIMessage({ + content: "Testing", + tool_calls: [{ id: "call_1", name: "test", args: {} }], + }), + new ToolMessage({ + content: "[cleared]", + tool_call_id: "call_1", + response_metadata: { + context_editing: { + cleared: true, + strategy: "clear_tool_uses", + }, + }, + }), + new AIMessage("Done"), + new HumanMessage("More"), + ]; + + const result = await agent.invoke({ messages }); + + // Should not try to clear an already cleared message + const clearedMsg = result.messages.find( + (msg) => ToolMessage.isInstance(msg) && msg.tool_call_id === "call_1" + ) as ToolMessage; + + expect(clearedMsg.content).toBe("[cleared]"); + expect( + (clearedMsg.response_metadata as any).context_editing.cleared + ).toBe(true); + }); + + it("should handle messages with no corresponding AI message", async () => { + const model = new FakeToolCallingChatModel({ + responses: [new AIMessage("Response")], + }); + + const middleware = contextEditingMiddleware({ + edits: [ + new ClearToolUsesEdit({ + triggerTokens: 50, + keep: 0, + }), + ], + }); + + const agent = createAgent({ + model, + middleware: [middleware] as const, + }); + + // Tool message without a corresponding AI message (malformed conversation) + const messages: BaseMessage[] = [ + new HumanMessage("Test"), + new ToolMessage({ + content: "Result", + tool_call_id: "orphan_call", + }), + new AIMessage("Done"), + ]; + + // Should handle gracefully without throwing + await expect(agent.invoke({ messages })).resolves.toBeDefined(); + }); + }); +}); diff --git a/libs/langchain/src/agents/middlewareAgent/middleware/tests/summarization.test.ts b/libs/langchain/src/agents/middlewareAgent/middleware/tests/summarization.test.ts index 64e143ed6dad..9d4000010d1a 100644 --- a/libs/langchain/src/agents/middlewareAgent/middleware/tests/summarization.test.ts +++ b/libs/langchain/src/agents/middlewareAgent/middleware/tests/summarization.test.ts @@ -6,10 +6,8 @@ import { SystemMessage, ToolMessage, } from "@langchain/core/messages"; -import { - summarizationMiddleware, - countTokensApproximately, -} from "../summarization.js"; +import { summarizationMiddleware } from "../summarization.js"; +import { countTokensApproximately } from "../utils.js"; import { createAgent } from "../../index.js"; import { FakeToolCallingChatModel } from "../../../tests/utils.js"; diff --git a/libs/langchain/src/agents/middlewareAgent/middleware/utils.ts b/libs/langchain/src/agents/middlewareAgent/middleware/utils.ts new file mode 100644 index 000000000000..d8e992a552cf --- /dev/null +++ b/libs/langchain/src/agents/middlewareAgent/middleware/utils.ts @@ -0,0 +1,29 @@ +import type { BaseMessage } from "@langchain/core/messages"; + +/** + * Default token counter that approximates based on character count + * @param messages Messages to count tokens for + * @returns Approximate token count + */ +export function countTokensApproximately(messages: BaseMessage[]): number { + let totalChars = 0; + for (const msg of messages) { + let textContent: string; + if (typeof msg.content === "string") { + textContent = msg.content; + } else if (Array.isArray(msg.content)) { + textContent = msg.content + .map((item) => { + if (typeof item === "string") return item; + if (item.type === "text" && "text" in item) return item.text; + return ""; + }) + .join(""); + } else { + textContent = ""; + } + totalChars += textContent.length; + } + // Approximate 1 token = 4 characters + return Math.ceil(totalChars / 4); +}