From 36825cb9a447ada09ff809ae94c3ddf92bc4e2b4 Mon Sep 17 00:00:00 2001 From: Ali Ibrahim Jr <48456829+IBJunior@users.noreply.github.com> Date: Fri, 12 Sep 2025 19:00:18 +0200 Subject: [PATCH 1/2] feat(langchain): add comprehensive ToolMessage support with metadata preservation --- .gitignore | 4 +- CHANGELOG.md | 22 ++ package.json | 2 +- src/adapters/langchain.ts | 88 +++++++- src/interfaces.ts | 2 + tests/langchain.test.ts | 428 +++++++++++++++++++++++++++++++++++++- 6 files changed, 538 insertions(+), 8 deletions(-) diff --git a/.gitignore b/.gitignore index 4ec5993..179bc01 100644 --- a/.gitignore +++ b/.gitignore @@ -54,4 +54,6 @@ coverage/ locales/ -*.tgz \ No newline at end of file +*.tgz + +.claude \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index c0304e4..6577e5f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,28 @@ All notable changes to this project will be documented in this file. The format is based on Keep a Changelog, and this project adheres to Semantic Versioning. +## [2.1.3] - 2025-09-12 + +### Added + +- Enhanced LangChain adapter with comprehensive ToolMessage support +- Metadata preservation system for maintaining message properties during compression +- Content preservation for complex message types (tool calls, function calls, etc.) +- Extended SlimContextMessage interface with optional metadata field +- Comprehensive test coverage for tool message handling and metadata preservation + +### Changed + +- Improved LangChain message conversion with robust metadata handling +- Enhanced content extraction with fallback preservation for complex content types +- Better roundtrip fidelity for LangChain BaseMessage conversions + +### Fixed + +- Tool message conversion now properly preserves tool_call_id and other tool-specific fields +- Complex message content (arrays, objects) is now correctly preserved during compression +- Message metadata fields are properly maintained throughout compression workflows + ## [2.1.2] - 2025-09-06 ### Added diff --git a/package.json b/package.json index b32e78c..9bbc838 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "slimcontext", - "version": "2.1.2", + "version": "2.1.3", "description": "Lightweight, model-agnostic chat history compression (trim + summarize) for AI assistants.", "main": "dist/index.js", "types": "dist/index.d.ts", diff --git a/src/adapters/langchain.ts b/src/adapters/langchain.ts index ff97503..a363f12 100644 --- a/src/adapters/langchain.ts +++ b/src/adapters/langchain.ts @@ -5,9 +5,13 @@ import { HumanMessage, isAIMessage, SystemMessage, + ToolMessage, isHumanMessage, isSystemMessage, isToolMessage, + AIMessageFields, + SystemMessageFields, + HumanMessageFields, } from '@langchain/core/messages'; import { @@ -79,23 +83,97 @@ export function getLangChainMessageType(msg: BaseMessage): string { /** Convert a LangChain BaseMessage to a SlimContextMessage used by compression. */ export function baseToSlim(msg: BaseMessage): SlimContextMessage { const type = getLangChainMessageType(msg); - return { + + // Collect all metadata fields to preserve during conversion + const metadata: Record = {}; + + // Preserve common BaseMessage fields + if (msg.id !== undefined) metadata.id = msg.id; + if (msg.name !== undefined) metadata.name = msg.name; + if (msg.additional_kwargs && Object.keys(msg.additional_kwargs).length > 0) { + metadata.additional_kwargs = msg.additional_kwargs; + } + if (msg.response_metadata && Object.keys(msg.response_metadata).length > 0) { + metadata.response_metadata = msg.response_metadata; + } + + // Preserve tool-specific fields + if (isToolMessage(msg)) { + metadata.tool_call_id = msg.tool_call_id; + if (msg.status !== undefined) metadata.status = msg.status; + if (msg.artifact !== undefined) metadata.artifact = msg.artifact; + if (msg.metadata !== undefined) metadata.langchain_metadata = msg.metadata; + } + + const extractedContent = extractContent(msg.content as unknown); + + // If extractContent returns empty but original content exists, preserve original content + // Note: msg.content can be undefined, empty string, or have content that extractContent can't process + if (!extractedContent && msg.content !== undefined && msg.content !== null) { + metadata.original_content = msg.content; + } + + const result: SlimContextMessage = { role: roleFromMessageType(type), - content: extractContent(msg.content as unknown), + content: extractedContent, }; + + // Only add metadata if there are fields to preserve + if (Object.keys(metadata).length > 0) { + result.metadata = metadata; + } + + return result; } /** Map SlimContextMessage back to a LangChain BaseMessage class. */ export function slimToLangChain(msg: SlimContextMessage): BaseMessage { + // Use original content if preserved, otherwise use processed content + const content = msg.metadata?.original_content || msg.content; + + // Prepare fields with preserved metadata + const fields: Record = { + content, + }; + + // Restore preserved metadata fields + if (msg.metadata) { + if (msg.metadata.id !== undefined) fields.id = msg.metadata.id; + if (msg.metadata.name !== undefined) fields.name = msg.metadata.name; + if (msg.metadata.additional_kwargs !== undefined) { + fields.additional_kwargs = msg.metadata.additional_kwargs; + } + if (msg.metadata.response_metadata !== undefined) { + fields.response_metadata = msg.metadata.response_metadata; + } + } + switch (msg.role) { case 'assistant': - return new AIMessage(msg.content); + return new AIMessage(fields as AIMessageFields); case 'system': - return new SystemMessage(msg.content); + return new SystemMessage(fields as SystemMessageFields); + case 'tool': { + // For tool messages, we need the tool_call_id and other tool-specific fields + const toolFields = { ...fields }; + if (msg.metadata?.tool_call_id) { + toolFields.tool_call_id = msg.metadata.tool_call_id; + } else { + // Fallback to empty string if no tool_call_id preserved (shouldn't happen with new logic) + toolFields.tool_call_id = ''; + } + if (msg.metadata?.status !== undefined) toolFields.status = msg.metadata.status; + if (msg.metadata?.artifact !== undefined) toolFields.artifact = msg.metadata.artifact; + if (msg.metadata?.langchain_metadata !== undefined) { + toolFields.metadata = msg.metadata.langchain_metadata; + } + // eslint-disable-next-line @typescript-eslint/no-explicit-any + return new ToolMessage(toolFields as any); + } case 'user': case 'human': default: - return new HumanMessage(msg.content); + return new HumanMessage(fields as HumanMessageFields); } } diff --git a/src/interfaces.ts b/src/interfaces.ts index 1299397..e2cd4db 100644 --- a/src/interfaces.ts +++ b/src/interfaces.ts @@ -5,6 +5,8 @@ export interface SlimContextMessage { role: 'system' | 'user' | 'assistant' | 'tool' | 'human'; content: string; + /** Optional metadata to preserve additional message fields during conversion */ + metadata?: Record; } export interface SlimContextModelResponse { diff --git a/tests/langchain.test.ts b/tests/langchain.test.ts index 28eff3a..5bf7f65 100644 --- a/tests/langchain.test.ts +++ b/tests/langchain.test.ts @@ -1,5 +1,11 @@ import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; -import { AIMessage, BaseMessage, HumanMessage, SystemMessage } from '@langchain/core/messages'; +import { + AIMessage, + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, +} from '@langchain/core/messages'; import type { MessageContentComplex } from '@langchain/core/messages'; import { Runnable, type RunnableConfig } from '@langchain/core/runnables'; import { describe, it, expect, vi } from 'vitest'; @@ -73,6 +79,265 @@ describe('LangChain Adapter', () => { expect(baseMessage).toBeInstanceOf(AIMessage); expect(baseMessage.content).toBe('Hi there'); }); + + it('should convert ToolMessage to SlimContextMessage and back', () => { + const toolMessage = new ToolMessage('Function result: 42', 'call_123'); + const slimMessage = baseToSlim(toolMessage); + expect(slimMessage.role).toBe('tool'); + expect(slimMessage.content).toBe('Function result: 42'); + expect(slimMessage.metadata?.tool_call_id).toBe('call_123'); + + const backToBase = slimToLangChain(slimMessage); + expect(backToBase).toBeInstanceOf(ToolMessage); + expect(backToBase.content).toBe('Function result: 42'); + expect((backToBase as ToolMessage).tool_call_id).toBe('call_123'); + }); + + it('should preserve tool message content through roundtrip conversion', () => { + const originalContent = + 'Complex tool response with\nmultiple lines\nand special chars: {"result": "success"}'; + const toolMessage = new ToolMessage(originalContent, 'call_456'); + + const slimMessage = baseToSlim(toolMessage); + const backToLangChain = slimToLangChain(slimMessage); + + expect(slimMessage.role).toBe('tool'); + expect(slimMessage.content).toBe(originalContent); + expect(backToLangChain).toBeInstanceOf(ToolMessage); + expect(backToLangChain.content).toBe(originalContent); + expect((backToLangChain as ToolMessage).tool_call_id).toBe('call_456'); + }); + + it('should preserve all message metadata fields', () => { + const complexMessage = new AIMessage({ + content: 'Complex response', + id: 'msg_123', + name: 'assistant_bot', + additional_kwargs: { model: 'gpt-4', temperature: 0.7 }, + response_metadata: { tokens: 150, latency: 250 }, + }); + + const slimMessage = baseToSlim(complexMessage); + const backToLangChain = slimToLangChain(slimMessage); + + expect(slimMessage.role).toBe('assistant'); + expect(slimMessage.content).toBe('Complex response'); + expect(slimMessage.metadata?.id).toBe('msg_123'); + expect(slimMessage.metadata?.name).toBe('assistant_bot'); + expect(slimMessage.metadata?.additional_kwargs).toEqual({ + model: 'gpt-4', + temperature: 0.7, + }); + expect(slimMessage.metadata?.response_metadata).toEqual({ + tokens: 150, + latency: 250, + }); + + expect(backToLangChain).toBeInstanceOf(AIMessage); + expect(backToLangChain.content).toBe('Complex response'); + expect(backToLangChain.id).toBe('msg_123'); + expect(backToLangChain.name).toBe('assistant_bot'); + expect(backToLangChain.additional_kwargs).toEqual({ + model: 'gpt-4', + temperature: 0.7, + }); + expect(backToLangChain.response_metadata).toEqual({ + tokens: 150, + latency: 250, + }); + }); + + it('should preserve tool message with all metadata', () => { + const toolMessage = new ToolMessage({ + content: '{"result": 42, "status": "success"}', + tool_call_id: 'call_xyz_789', + name: 'calculator', + id: 'tool_msg_001', + status: 'success' as const, + artifact: { raw_output: 42, computation_time: 0.05 }, + additional_kwargs: { model_version: 'v2.1' }, + response_metadata: { execution_time: 100 }, + }); + + const slimMessage = baseToSlim(toolMessage); + const backToLangChain = slimToLangChain(slimMessage); + + // Check SlimContext message preserves metadata + expect(slimMessage.role).toBe('tool'); + expect(slimMessage.content).toBe('{"result": 42, "status": "success"}'); + expect(slimMessage.metadata?.tool_call_id).toBe('call_xyz_789'); + expect(slimMessage.metadata?.name).toBe('calculator'); + expect(slimMessage.metadata?.id).toBe('tool_msg_001'); + expect(slimMessage.metadata?.status).toBe('success'); + expect(slimMessage.metadata?.artifact).toEqual({ + raw_output: 42, + computation_time: 0.05, + }); + + // Check LangChain message is fully restored + expect(backToLangChain).toBeInstanceOf(ToolMessage); + const restoredTool = backToLangChain as ToolMessage; + expect(restoredTool.content).toBe('{"result": 42, "status": "success"}'); + expect(restoredTool.tool_call_id).toBe('call_xyz_789'); + expect(restoredTool.name).toBe('calculator'); + expect(restoredTool.id).toBe('tool_msg_001'); + expect(restoredTool.status).toBe('success'); + expect(restoredTool.artifact).toEqual({ + raw_output: 42, + computation_time: 0.05, + }); + expect(restoredTool.additional_kwargs).toEqual({ model_version: 'v2.1' }); + expect(restoredTool.response_metadata).toEqual({ execution_time: 100 }); + }); + + it('should handle messages without metadata gracefully', () => { + const simpleMessage = new HumanMessage('Simple question'); + const slimMessage = baseToSlim(simpleMessage); + const backToLangChain = slimToLangChain(slimMessage); + + expect(slimMessage.role).toBe('user'); + expect(slimMessage.content).toBe('Simple question'); + expect(slimMessage.metadata).toBeUndefined(); + + expect(backToLangChain).toBeInstanceOf(HumanMessage); + expect(backToLangChain.content).toBe('Simple question'); + }); + + it('should preserve complex content when extractContent returns empty', () => { + // Create AI message with complex content (tool calls, function calls, no plain text) + const complexContent = [ + { + functionCall: { + name: 'duckduckgo__search', + args: { query: 'AI agents trends and future' }, + }, + thoughtSignature: 'CikB0e2Kbyvu3hrfNhjy...', + }, + ]; + + const aiMessage = new AIMessage({ + content: complexContent, + id: 'run-3073efed-c743-4ef2-b305-a52435fba26f', + additional_kwargs: { + tool_calls: [ + { + name: 'duckduckgo__search', + args: { query: 'AI agents trends and future' }, + id: '696ecb8f-d532-4172-abfa-98b72023674a', + type: 'tool_call', + }, + ], + }, + }); + + const slimMessage = baseToSlim(aiMessage); + const backToLangChain = slimToLangChain(slimMessage); + + // Content should be empty string from extractContent + expect(slimMessage.content).toBe(''); + // But original content should be preserved in metadata + expect(slimMessage.metadata?.original_content).toEqual(complexContent); + expect(slimMessage.metadata?.id).toBe('run-3073efed-c743-4ef2-b305-a52435fba26f'); + + // When restored, should get back the original complex content + expect(backToLangChain).toBeInstanceOf(AIMessage); + expect(backToLangChain.content).toEqual(complexContent); + expect(backToLangChain.id).toBe('run-3073efed-c743-4ef2-b305-a52435fba26f'); + }); + + it('should preserve AI message with tool call chunks and no text content', () => { + const complexMessageData = { + content: [ + { + functionCall: { + name: 'duckduckgo__search', + args: { query: 'AI agents trends and future' }, + }, + thoughtSignature: 'CikB0e2Kbyvu3hrfNhjy...', + }, + ], + tool_call_chunks: [ + { + name: 'duckduckgo__search', + args: '{"query":"AI agents trends and future"}', + index: 0, + type: 'tool_call_chunk', + id: '696ecb8f-d532-4172-abfa-98b72023674a', + }, + ], + additional_kwargs: {}, + tool_calls: [ + { + name: 'duckduckgo__search', + args: { query: 'AI agents trends and future' }, + id: '696ecb8f-d532-4172-abfa-98b72023674a', + type: 'tool_call', + }, + ], + invalid_tool_calls: [], + response_metadata: {}, + id: 'run-3073efed-c743-4ef2-b305-a52435fba26f', + }; + + const aiMessage = new AIMessage(complexMessageData); + const slimMessage = baseToSlim(aiMessage); + const backToLangChain = slimToLangChain(slimMessage); + + // Should preserve the original complex content structure + expect(slimMessage.content).toBe(''); + expect(slimMessage.metadata?.original_content).toEqual(complexMessageData.content); + + // When restored, should be identical to original + expect(backToLangChain).toBeInstanceOf(AIMessage); + expect(backToLangChain.content).toEqual(complexMessageData.content); + expect(backToLangChain.id).toBe(complexMessageData.id); + + // Check that additional_kwargs is preserved (empty objects are not preserved to avoid bloat) + expect(slimMessage.metadata?.additional_kwargs).toBeUndefined(); + }); + + it('should handle mixed content with both text and complex structures', () => { + // Test with plain string content which should extract normally + const simpleMessage = new AIMessage('I will search for information about AI agents.'); + const slimMessage = baseToSlim(simpleMessage); + const backToLangChain = slimToLangChain(slimMessage); + + // Should extract the text content + expect(slimMessage.content).toBe('I will search for information about AI agents.'); + // Should NOT preserve original content since extractContent succeeded + expect(slimMessage.metadata?.original_content).toBeUndefined(); + + // Should restore with extracted text content + expect(backToLangChain.content).toBe('I will search for information about AI agents.'); + }); + + it('should handle complex array content by preserving original when needed', () => { + // Test with array containing objects that may not be extractable by LangChain processing + const mixedContent = [ + { text: 'I will search for information about AI agents.' }, + { + functionCall: { + name: 'search_tool', + args: { query: 'AI agents' }, + }, + }, + ]; + + const aiMessage = new AIMessage({ content: mixedContent }); + const slimMessage = baseToSlim(aiMessage); + const backToLangChain = slimToLangChain(slimMessage); + + // Behavior depends on how LangChain processes the content internally + // If extractContent returns empty, original content should be preserved + if (slimMessage.content === '') { + expect(slimMessage.metadata?.original_content).toEqual(mixedContent); + expect(backToLangChain.content).toEqual(mixedContent); + } else { + // If content was extracted successfully, original shouldn't be preserved + expect(slimMessage.metadata?.original_content).toBeUndefined(); + expect(backToLangChain.content).toBe(slimMessage.content); + } + }); }); describe('compressLangChainHistory', () => { @@ -143,5 +408,166 @@ describe('LangChain Adapter', () => { expect(compressed[0].content).toBe('Custom summary'); expect(compressed[1].content).toBe('Message 4'); }); + + it('should preserve tool messages during compression', async () => { + const historyWithTools: BaseMessage[] = [ + new SystemMessage('You are a helpful assistant.'), + new HumanMessage('What is 2 + 2?'), + new AIMessage('Let me calculate that for you.'), + new ToolMessage('{"result": 4}', 'calc_123'), + new AIMessage('The answer is 4.'), + new HumanMessage('What about 5 + 5?'), + ]; + + const compressed = await compressLangChainHistory(historyWithTools, { + strategy: 'trim', + maxModelTokens: 300, + thresholdPercent: 0.5, + minRecentMessages: 3, + estimateTokens: () => 100, // Total: 600 tokens, threshold: 150 tokens + }); + + // Should preserve system + last 3 messages (including the tool message) + expect(compressed).toHaveLength(4); + expect(compressed[0]).toBeInstanceOf(SystemMessage); + expect(compressed[1]).toBeInstanceOf(ToolMessage); + expect(compressed[1].content).toBe('{"result": 4}'); + expect((compressed[1] as ToolMessage).tool_call_id).toBe('calc_123'); + expect(compressed[2]).toBeInstanceOf(AIMessage); + expect(compressed[2].content).toBe('The answer is 4.'); + expect(compressed[3]).toBeInstanceOf(HumanMessage); + expect(compressed[3].content).toBe('What about 5 + 5?'); + }); + + it('should preserve rich metadata during compression workflow', async () => { + const historyWithMetadata: BaseMessage[] = [ + new SystemMessage({ + content: 'You are a helpful assistant.', + id: 'system_001', + name: 'system', + }), + new HumanMessage({ + content: 'Calculate 10 + 15', + id: 'user_001', + name: 'user123', + additional_kwargs: { session_id: 'abc123' }, + }), + new AIMessage({ + content: 'I will calculate that using a tool.', + id: 'ai_001', + response_metadata: { model: 'gpt-4', tokens: 12 }, + }), + new ToolMessage({ + content: '{"result": 25}', + tool_call_id: 'calc_operation_456', + id: 'tool_001', + name: 'calculator', + status: 'success' as const, + artifact: { operation: 'addition', operands: [10, 15] }, + }), + new AIMessage({ + content: 'The result is 25.', + id: 'ai_002', + response_metadata: { tokens: 8 }, + }), + ]; + + const compressed = await compressLangChainHistory(historyWithMetadata, { + strategy: 'trim', + maxModelTokens: 200, + thresholdPercent: 0.6, + minRecentMessages: 2, + estimateTokens: () => 50, // Total: 250, threshold: 120 + }); + + // Should preserve system + last 2 messages + expect(compressed).toHaveLength(3); + + // Check system message metadata preserved + expect(compressed[0]).toBeInstanceOf(SystemMessage); + expect(compressed[0].id).toBe('system_001'); + expect(compressed[0].name).toBe('system'); + + // Check tool message metadata preserved + expect(compressed[1]).toBeInstanceOf(ToolMessage); + const toolMsg = compressed[1] as ToolMessage; + expect(toolMsg.content).toBe('{"result": 25}'); + expect(toolMsg.tool_call_id).toBe('calc_operation_456'); + expect(toolMsg.id).toBe('tool_001'); + expect(toolMsg.name).toBe('calculator'); + expect(toolMsg.status).toBe('success'); + expect(toolMsg.artifact).toEqual({ + operation: 'addition', + operands: [10, 15], + }); + + // Check AI message metadata preserved + expect(compressed[2]).toBeInstanceOf(AIMessage); + expect(compressed[2].id).toBe('ai_002'); + expect(compressed[2].response_metadata).toEqual({ tokens: 8 }); + }); + + it('should preserve complex content during compression', async () => { + const complexContent = [ + { + functionCall: { + name: 'search_tool', + args: { query: 'test search' }, + }, + thoughtSignature: 'abc123...', + }, + ]; + + const historyWithComplexContent: BaseMessage[] = [ + new SystemMessage('You are a helpful assistant.'), + new HumanMessage('Search for something'), + new AIMessage({ + content: complexContent, + id: 'complex_ai_msg', + additional_kwargs: { + tool_calls: [ + { + name: 'search_tool', + args: { query: 'test search' }, + id: 'tool_call_123', + type: 'tool_call', + }, + ], + }, + }), + new HumanMessage('Thanks!'), + ]; + + const compressed = await compressLangChainHistory(historyWithComplexContent, { + strategy: 'trim', + maxModelTokens: 200, + thresholdPercent: 0.5, + minRecentMessages: 2, + estimateTokens: () => 80, // Total: 320, threshold: 100 + }); + + // Should preserve system + last 2 messages + expect(compressed).toHaveLength(3); + expect(compressed[0]).toBeInstanceOf(SystemMessage); + + // Complex AI message should be preserved with original content structure + expect(compressed[1]).toBeInstanceOf(AIMessage); + const aiMsg = compressed[1] as AIMessage; + expect(aiMsg.content).toEqual(complexContent); + expect(aiMsg.id).toBe('complex_ai_msg'); + expect(aiMsg.additional_kwargs).toEqual({ + tool_calls: [ + { + name: 'search_tool', + args: { query: 'test search' }, + id: 'tool_call_123', + type: 'tool_call', + }, + ], + }); + + expect(compressed[2]).toBeInstanceOf(HumanMessage); + expect(compressed[2].content).toBe('Thanks!'); + }); }); }); From 6b222039d68fbc59b445f55a7ab5722d55014796 Mon Sep 17 00:00:00 2001 From: Ali Ibrahim Jr <48456829+IBJunior@users.noreply.github.com> Date: Sat, 13 Sep 2025 13:42:34 +0200 Subject: [PATCH 2/2] feat(compression): add smart compression timing to prevent tool workflow disruption --- CHANGELOG.md | 6 +- src/strategies/common.ts | 10 +++ src/strategies/summarize.ts | 7 +- src/strategies/trim.ts | 7 +- tests/langchain.test.ts | 46 ++++++------- tests/summarize.test.ts | 125 +++++++++++++++++++++++++++++++++++- tests/trim.test.ts | 90 ++++++++++++++++++++++++++ 7 files changed, 262 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6577e5f..98e3ead 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ All notable changes to this project will be documented in this file. The format is based on Keep a Changelog, and this project adheres to Semantic Versioning. -## [2.1.3] - 2025-09-12 +## [2.1.3] - 2025-09-13 ### Added @@ -12,19 +12,23 @@ The format is based on Keep a Changelog, and this project adheres to Semantic Ve - Metadata preservation system for maintaining message properties during compression - Content preservation for complex message types (tool calls, function calls, etc.) - Extended SlimContextMessage interface with optional metadata field +- **Smart compression timing**: Added `shouldAllowCompression` function to prevent compression during active tool use cycles - Comprehensive test coverage for tool message handling and metadata preservation +- Enhanced test coverage for compression timing and tool interaction patterns ### Changed - Improved LangChain message conversion with robust metadata handling - Enhanced content extraction with fallback preservation for complex content types - Better roundtrip fidelity for LangChain BaseMessage conversions +- **TrimCompressor and SummarizeCompressor now only compress when the last message is from a user/human** to avoid disrupting tool use workflows ### Fixed - Tool message conversion now properly preserves tool_call_id and other tool-specific fields - Complex message content (arrays, objects) is now correctly preserved during compression - Message metadata fields are properly maintained throughout compression workflows +- **Compression timing**: Prevents inappropriate compression during tool interactions that could break multi-turn tool workflows ## [2.1.2] - 2025-09-06 diff --git a/src/strategies/common.ts b/src/strategies/common.ts index 854f7e4..972a179 100644 --- a/src/strategies/common.ts +++ b/src/strategies/common.ts @@ -46,3 +46,13 @@ export function estimateTotalTokens( for (const m of messages) total += estimateTokens(m); return total; } + +/** + * Check if compression should be allowed based on the last message type. + * Only compress when the last message is from a user to avoid disrupting tool use cycles. + */ +export function shouldAllowCompression(messages: SlimContextMessage[]): boolean { + if (messages.length === 0) return false; + const lastMessage = messages[messages.length - 1]; + return lastMessage.role === 'user' || lastMessage.role === 'human'; +} diff --git a/src/strategies/summarize.ts b/src/strategies/summarize.ts index 869e878..e3dbd45 100644 --- a/src/strategies/summarize.ts +++ b/src/strategies/summarize.ts @@ -4,7 +4,7 @@ import { SlimContextMessage, type TokenBudgetConfig, } from '../interfaces'; -import { normalizeBudgetConfig, computeThresholdTokens } from './common'; +import { normalizeBudgetConfig, computeThresholdTokens, shouldAllowCompression } from './common'; const DEFAULT_SUMMARY_PROMPT = ` You are a conversation summarizer. @@ -73,6 +73,11 @@ export class SummarizeCompressor implements SlimContextCompressor { * Compress the conversation history by summarizing the middle portion. */ async compress(messages: SlimContextMessage[]): Promise { + // Only compress when the last message is from a user to avoid disrupting tool use cycles + if (!shouldAllowCompression(messages)) { + return messages; + } + const thresholdTokens = computeThresholdTokens( this.cfg.maxModelTokens, this.cfg.thresholdPercent, diff --git a/src/strategies/trim.ts b/src/strategies/trim.ts index 34bc124..7a001e2 100644 --- a/src/strategies/trim.ts +++ b/src/strategies/trim.ts @@ -1,5 +1,5 @@ import { SlimContextCompressor, SlimContextMessage, type TokenBudgetConfig } from '../interfaces'; -import { normalizeBudgetConfig, computeThresholdTokens } from './common'; +import { normalizeBudgetConfig, computeThresholdTokens, shouldAllowCompression } from './common'; /** * Trim configuration options for the TrimCompressor using token thresholding. @@ -19,6 +19,11 @@ export class TrimCompressor implements SlimContextCompressor { } async compress(messages: SlimContextMessage[]): Promise { + // Only compress when the last message is from a user to avoid disrupting tool use cycles + if (!shouldAllowCompression(messages)) { + return messages; + } + const thresholdTokens = computeThresholdTokens( this.cfg.maxModelTokens, this.cfg.thresholdPercent, diff --git a/tests/langchain.test.ts b/tests/langchain.test.ts index 5bf7f65..0aaae81 100644 --- a/tests/langchain.test.ts +++ b/tests/langchain.test.ts @@ -347,6 +347,7 @@ describe('LangChain Adapter', () => { new AIMessage('Message 2'), new HumanMessage('Message 3'), new AIMessage('Message 4'), + new HumanMessage('Message 5'), // End with user message to allow compression ]; it('should compress history with the trim strategy', async () => { @@ -361,10 +362,10 @@ describe('LangChain Adapter', () => { // our TrimCompressor preserves last two regardless; will drop earlier non-systems. expect(compressed.length).toBe(3); expect(compressed[0]).toBeInstanceOf(SystemMessage); - expect(compressed[1]).toBeInstanceOf(HumanMessage); - expect(compressed[1].content).toBe('Message 3'); - expect(compressed[2]).toBeInstanceOf(AIMessage); - expect(compressed[2].content).toBe('Message 4'); + expect(compressed[1]).toBeInstanceOf(AIMessage); + expect(compressed[1].content).toBe('Message 4'); + expect(compressed[2]).toBeInstanceOf(HumanMessage); + expect(compressed[2].content).toBe('Message 5'); }); it('should compress history with the summarize strategy', async () => { @@ -386,8 +387,8 @@ describe('LangChain Adapter', () => { expect(compressed[0]).toBeInstanceOf(SystemMessage); expect(compressed[1]).toBeInstanceOf(SystemMessage); // Summary is an System Message expect(compressed[1].content).toContain('This is a summary of messages 1 and 2.'); - expect(compressed[2].content).toBe('Message 3'); - expect(compressed[3].content).toBe('Message 4'); + expect(compressed[2].content).toBe('Message 4'); + expect(compressed[3].content).toBe('Message 5'); }); it('should work with a pre-created compressor', async () => { @@ -406,7 +407,7 @@ describe('LangChain Adapter', () => { expect(compressed).toHaveLength(2); expect(compressed[0].content).toBe('Custom summary'); - expect(compressed[1].content).toBe('Message 4'); + expect(compressed[1].content).toBe('Message 5'); }); it('should preserve tool messages during compression', async () => { @@ -470,6 +471,11 @@ describe('LangChain Adapter', () => { id: 'ai_002', response_metadata: { tokens: 8 }, }), + new HumanMessage({ + content: 'Thanks!', + id: 'user_002', + name: 'user123', + }), ]; const compressed = await compressLangChainHistory(historyWithMetadata, { @@ -488,23 +494,17 @@ describe('LangChain Adapter', () => { expect(compressed[0].id).toBe('system_001'); expect(compressed[0].name).toBe('system'); - // Check tool message metadata preserved - expect(compressed[1]).toBeInstanceOf(ToolMessage); - const toolMsg = compressed[1] as ToolMessage; - expect(toolMsg.content).toBe('{"result": 25}'); - expect(toolMsg.tool_call_id).toBe('calc_operation_456'); - expect(toolMsg.id).toBe('tool_001'); - expect(toolMsg.name).toBe('calculator'); - expect(toolMsg.status).toBe('success'); - expect(toolMsg.artifact).toEqual({ - operation: 'addition', - operands: [10, 15], - }); + // Check AI message metadata preserved (2nd to last) + expect(compressed[1]).toBeInstanceOf(AIMessage); + const aiMsg = compressed[1] as AIMessage; + expect(aiMsg.content).toBe('The result is 25.'); + expect(aiMsg.id).toBe('ai_002'); - // Check AI message metadata preserved - expect(compressed[2]).toBeInstanceOf(AIMessage); - expect(compressed[2].id).toBe('ai_002'); - expect(compressed[2].response_metadata).toEqual({ tokens: 8 }); + // Check last human message + expect(compressed[2]).toBeInstanceOf(HumanMessage); + const humanMsg = compressed[2] as HumanMessage; + expect(humanMsg.content).toBe('Thanks!'); + expect(humanMsg.id).toBe('user_002'); }); it('should preserve complex content during compression', async () => { diff --git a/tests/summarize.test.ts b/tests/summarize.test.ts index 0af9abb..6cf4659 100644 --- a/tests/summarize.test.ts +++ b/tests/summarize.test.ts @@ -53,7 +53,8 @@ describe('SummarizeCompressor', () => { // Start with user instead of system, then alternate and end with user const history: SlimContextMessage[] = []; - for (let i = 0; i < 10; i++) { + for (let i = 0; i < 11; i++) { + // Changed to 11 to ensure it ends with user const role = i % 2 === 0 ? 'user' : 'assistant'; history.push({ role: role as 'user' | 'assistant', content: `${role[0]}${i}` }); } @@ -61,7 +62,125 @@ describe('SummarizeCompressor', () => { const result = await summarize.compress(history); expect(result[0].role).toBe('system'); // summary system (no original system preserved) expect(result[0].content).toContain('fake summary'); - expect(result.at(-2)?.role).toBe('user'); - expect(result.at(-1)?.role).toBe('assistant'); + expect(result.at(-2)?.role).toBe('assistant'); + expect(result.at(-1)?.role).toBe('user'); + }); + + it('skips compression when last message is not from user', async () => { + const fakeModel: SlimContextChatModel = { + async invoke(_msgs: SlimContextMessage[]): Promise { + return { content: 'fake summary' }; + }, + }; + + const summarize = new SummarizeCompressor({ + model: fakeModel, + maxModelTokens: 400, + thresholdPercent: 0.5, // 200 tokens + estimateTokens: () => 100, // each message 100 tokens + minRecentMessages: 2, + }); + + const history: SlimContextMessage[] = [ + { role: 'user', content: 'u1' }, // 100 + { role: 'assistant', content: 'a1' }, // 100 + { role: 'user', content: 'u2' }, // 100 + { role: 'tool', content: 'tool result' }, // 100 - last message is tool + ]; // total 400 > threshold 200 + + const result = await summarize.compress(history); + // Should return unchanged since last message is tool + expect(result).toEqual(history); + expect(result.length).toBe(4); + }); + + it('skips compression when last message is from assistant', async () => { + const fakeModel: SlimContextChatModel = { + async invoke(_msgs: SlimContextMessage[]): Promise { + return { content: 'fake summary' }; + }, + }; + + const summarize = new SummarizeCompressor({ + model: fakeModel, + maxModelTokens: 400, + thresholdPercent: 0.5, // 200 tokens + estimateTokens: () => 100, // each message 100 tokens + minRecentMessages: 2, + }); + + const history: SlimContextMessage[] = [ + { role: 'user', content: 'u1' }, // 100 + { role: 'assistant', content: 'a1' }, // 100 + { role: 'user', content: 'u2' }, // 100 + { role: 'assistant', content: 'a2' }, // 100 - last message is assistant + ]; // total 400 > threshold 200 + + const result = await summarize.compress(history); + // Should return unchanged since last message is assistant + expect(result).toEqual(history); + expect(result.length).toBe(4); + }); + + it('compresses normally when last message is from user', async () => { + const fakeModel: SlimContextChatModel = { + async invoke(_msgs: SlimContextMessage[]): Promise { + return { content: 'fake summary' }; + }, + }; + + const summarize = new SummarizeCompressor({ + model: fakeModel, + maxModelTokens: 400, + thresholdPercent: 0.5, // 200 tokens + estimateTokens: () => 100, // each message 100 tokens + minRecentMessages: 2, + }); + + const history: SlimContextMessage[] = [ + { role: 'assistant', content: 'a1' }, // 100 + { role: 'user', content: 'u1' }, // 100 + { role: 'assistant', content: 'a2' }, // 100 + { role: 'user', content: 'u2' }, // 100 - last message is user + ]; // total 400 > threshold 200 + + const result = await summarize.compress(history); + // Should compress since last message is user + expect(result.length).toBe(3); // summary + last 2 messages + expect(result[0].role).toBe('system'); // summary + expect(result[0].content).toContain('fake summary'); + expect(result[1]).toEqual({ role: 'assistant', content: 'a2' }); + expect(result[2]).toEqual({ role: 'user', content: 'u2' }); + }); + + it('compresses normally when last message is human role', async () => { + const fakeModel: SlimContextChatModel = { + async invoke(_msgs: SlimContextMessage[]): Promise { + return { content: 'fake summary' }; + }, + }; + + const summarize = new SummarizeCompressor({ + model: fakeModel, + maxModelTokens: 400, + thresholdPercent: 0.5, // 200 tokens + estimateTokens: () => 100, // each message 100 tokens + minRecentMessages: 2, + }); + + const history: SlimContextMessage[] = [ + { role: 'assistant', content: 'a1' }, // 100 + { role: 'user', content: 'u1' }, // 100 + { role: 'assistant', content: 'a2' }, // 100 + { role: 'human', content: 'h1' }, // 100 - last message is human (user equivalent) + ]; // total 400 > threshold 200 + + const result = await summarize.compress(history); + // Should compress since last message is human (equivalent to user) + expect(result.length).toBe(3); // summary + last 2 messages + expect(result[0].role).toBe('system'); // summary + expect(result[0].content).toContain('fake summary'); + expect(result[1]).toEqual({ role: 'assistant', content: 'a2' }); + expect(result[2]).toEqual({ role: 'human', content: 'h1' }); }); }); diff --git a/tests/trim.test.ts b/tests/trim.test.ts index aaa33b1..2dcf13e 100644 --- a/tests/trim.test.ts +++ b/tests/trim.test.ts @@ -28,4 +28,94 @@ describe('TrimCompressor', () => { // Older non-system messages should be dropped expect(trimmed.length).toBe(3); }); + + it('skips compression when last message is not from user', async () => { + const estimate = (_m: SlimContextMessage) => 100; + const trim = new TrimCompressor({ + maxModelTokens: 400, + thresholdPercent: 0.5, // threshold = 200 + estimateTokens: estimate, + minRecentMessages: 2, + }); + + const history: SlimContextMessage[] = [ + { role: 'user', content: 'u1' }, // 100 + { role: 'assistant', content: 'a1' }, // 100 + { role: 'user', content: 'u2' }, // 100 + { role: 'tool', content: 'tool result' }, // 100 - last message is tool + ]; // total 400 > threshold 200 + + const result = await trim.compress(history); + // Should return unchanged since last message is tool + expect(result).toEqual(history); + expect(result.length).toBe(4); + }); + + it('skips compression when last message is from assistant', async () => { + const estimate = (_m: SlimContextMessage) => 100; + const trim = new TrimCompressor({ + maxModelTokens: 400, + thresholdPercent: 0.5, // threshold = 200 + estimateTokens: estimate, + minRecentMessages: 2, + }); + + const history: SlimContextMessage[] = [ + { role: 'user', content: 'u1' }, // 100 + { role: 'assistant', content: 'a1' }, // 100 + { role: 'user', content: 'u2' }, // 100 + { role: 'assistant', content: 'a2' }, // 100 - last message is assistant + ]; // total 400 > threshold 200 + + const result = await trim.compress(history); + // Should return unchanged since last message is assistant + expect(result).toEqual(history); + expect(result.length).toBe(4); + }); + + it('compresses normally when last message is from user', async () => { + const estimate = (_m: SlimContextMessage) => 100; + const trim = new TrimCompressor({ + maxModelTokens: 400, + thresholdPercent: 0.5, // threshold = 200 + estimateTokens: estimate, + minRecentMessages: 2, + }); + + const history: SlimContextMessage[] = [ + { role: 'assistant', content: 'a1' }, // 100 + { role: 'user', content: 'u1' }, // 100 + { role: 'assistant', content: 'a2' }, // 100 + { role: 'user', content: 'u2' }, // 100 - last message is user + ]; // total 400 > threshold 200 + + const result = await trim.compress(history); + // Should compress since last message is user + expect(result.length).toBe(2); // Only last 2 messages preserved + expect(result[0]).toEqual({ role: 'assistant', content: 'a2' }); + expect(result[1]).toEqual({ role: 'user', content: 'u2' }); + }); + + it('compresses normally when last message is human role', async () => { + const estimate = (_m: SlimContextMessage) => 100; + const trim = new TrimCompressor({ + maxModelTokens: 400, + thresholdPercent: 0.5, // threshold = 200 + estimateTokens: estimate, + minRecentMessages: 2, + }); + + const history: SlimContextMessage[] = [ + { role: 'assistant', content: 'a1' }, // 100 + { role: 'user', content: 'u1' }, // 100 + { role: 'assistant', content: 'a2' }, // 100 + { role: 'human', content: 'h1' }, // 100 - last message is human (user equivalent) + ]; // total 400 > threshold 200 + + const result = await trim.compress(history); + // Should compress since last message is human (equivalent to user) + expect(result.length).toBe(2); // Only last 2 messages preserved + expect(result[0]).toEqual({ role: 'assistant', content: 'a2' }); + expect(result[1]).toEqual({ role: 'human', content: 'h1' }); + }); });