Hello.
I am currently writing a post processing Middleware function to automatically summarize and remove unnecessary messages after summarization. I know I can use the summarizationMiddleware, which works great.. but it doesn’t auto-remove summarized messages from the state after summarization, and also I need to store the summarization message that was created in the state, so that it can be re-used for future summarization messages. Otherwise the summarization re-runs after every subsequent createAgent LLM call, which is obviously not ideal. Both of these additional features would be much appreciated in the future, as this would save writing customMiddleware. Unless there is some other way to achieve this?
So I have written a customMiddleware to do this, but it is not working as expected. The problem is that, even though my graph stateSchema includes a custom state field for messagesSummary (see code below).. LangSmith is showing that the middleWare sucessfully outputs the result of messagesSummary as expected, but this change to the state never passes through to the final output of the createAgent call See: LangSmith for full trace.
Here is the stateSchema I am using for the graph.
export const DynamicAnnotation = Annotation.Root({
…MessagesAnnotation.spec, // Spread in the messages state
// Message Summary (used when summarising the conversation for shorter LLM)
messagesSummary: Annotation<string>({
reducer: (x, y) => {
// Handle concurrent updates by taking the last non-empty value
if (Array.isArray(y)) {
// Multiple updates in same step - take the last non-empty one
const lastUpdate = y.filter(val => val != null && val.trim() !== '').pop();
return lastUpdate ?? x;
}
// Single update - replace with new value if provided, otherwise keep existing
return y ?? x;
}
}),
// Main dataset used for the graph (not appended, replaced);
data: Annotation<Record<string, any>>(),
// Temp data store for various purposes (e.g.: Evaluation iterators) - not appended, replaced;
_temp: Annotation<Record<string, any>>({
default: () => ({}),
reducer: (x, y) => ({ ...x, ...y }),
}),
// Next node to route too
next: Annotation<string>()
})
So to attempt to fix this, I updated the middleware.. to define stateSchema.. aka
createMiddleware({
name: "SummarizationMiddleware",
stateSchema: z.object({
messagesSummary: z.string().optional(),
}),
.....
This matches the return statement for createMiddleware…
.....
const summaryResponse = await model.invoke(summaryPrompt);
const updatedSummary = "SUMMARY OF PREVIOUS CONVERSATION HISTORY: " +
(typeof summaryResponse.content === ‘string’ ? summaryResponse.content : JSON.stringify(summaryResponse.content));
// Create RemoveMessage objects for all messages being removed
const messagesToRemove: RemoveMessage[] = [];
try {
for (const messageToSummarize of messagesToSummarize) {
const msgId = (messageToSummarize as any).id;
if (msgId) {
// Verify the message exists in the original messages array
const messageExists = messages.some((m: any) => m.id === msgId);
if (messageExists) {
messagesToRemove.push(
new RemoveMessage({
id: msgId as string,
})
);
}
}
}
} catch (error) {
console.warn("Error creating RemoveMessage objects:", error);
}
trimmerRunning[trimmerId] = false;
// Return updated state with:
// 1. Preserved messages + RemoveMessage objects (to remove old
messages)
// 2. Updated messagesSummary
return {
messages: [...messagesToRemove],
messagesSummary: updatedSummary
};
But then i get the following error when the middleware tries to run..
Error in AgentGraphNode simple_agent: TypeError: keyValidator._parse is not a function
at ZodObject._parse (file:///home/runner/workspace/node_modules/zod/v3/types.js:1963:37)
at ZodObject._parseSync (file:///home/runner/workspace/node_modules/zod/v3/types.js:100:29)
at ZodObject.safeParse (file:///home/runner/workspace/node_modules/zod/v3/types.js:129:29)
at ZodObject.parse (file:///home/runner/workspace/node_modules/zod/v3/types.js:111:29)
at interopParse (file:///home/runner/workspace/node_modules/@langchain/core/dist/utils/types/zod.js:136:43)
at CompiledStateGraph._validateInput (file:///home/runner/workspace/node_modules/@langchain/langgraph/dist/graph/state.js:438:30)
at CompiledStateGraph._streamIterator (file:///home/runner/workspace/node_modules/@langchain/langgraph/dist/pregel/index.js:964:33)
at _streamIterator.next (<anonymous>)
at file:///home/runner/workspace/node_modules/@langchain/core/dist/utils/stream.js:136:41
at AsyncLocalStorage.run (node:internal/async_local_storage/async_hooks:91:14)
❌ Failed to test template undefined: TypeError: keyValidator._parse is not a function
at ZodObject._parse (file:///home/runner/workspace/node_modules/zod/v3/types.js:1963:37)
at ZodObject._parseSync (file:///home/runner/workspace/node_modules/zod/v3/types.js:100:29)
at ZodObject.safeParse (file:///home/runner/workspace/node_modules/zod/v3/types.js:129:29)
at ZodObject.parse (file:///home/runner/workspace/node_modules/zod/v3/types.js:111:29)
at interopParse (file:///home/runner/workspace/node_modules/@langchain/core/dist/utils/types/zod.js:136:43)
at CompiledStateGraph._validateInput (file:///home/runner/workspace/node_modules/@langchain/langgraph/dist/graph/state.js:438:30)
at CompiledStateGraph._streamIterator (file:///home/runner/workspace/node_modules/@langchain/langgraph/dist/pregel/index.js:964:33)
at _streamIterator.next (<anonymous>)
at file:///home/runner/workspace/node_modules/@langchain/core/dist/utils/stream.js:136:41
at AsyncLocalStorage.run (node:internal/async_local_storage/async_hooks:91:14) {
pregelTaskId: 'eb9eae4e-9173-5cf5-823e-5002e5367b1d'
}
I even attempted to run the default code from LangChain Docs (section “Custom state schema), and I get the same error as above.. Custom middleware - Docs by LangChain
Any help to resolve this would be much appreciated.