@librechat/agents 2.2.8 → 2.2.9
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/cjs/graphs/Graph.cjs +5 -1
- package/dist/cjs/graphs/Graph.cjs.map +1 -1
- package/dist/cjs/messages/prune.cjs +93 -9
- package/dist/cjs/messages/prune.cjs.map +1 -1
- package/dist/esm/graphs/Graph.mjs +5 -1
- package/dist/esm/graphs/Graph.mjs.map +1 -1
- package/dist/esm/messages/prune.mjs +85 -1
- package/dist/esm/messages/prune.mjs.map +1 -1
- package/dist/types/messages/prune.d.ts +1 -0
- package/dist/types/scripts/code_exec_files.d.ts +1 -0
- package/package.json +2 -1
- package/src/graphs/Graph.ts +7 -1
- package/src/messages/prune.ts +106 -4
- package/src/scripts/code_exec_files.ts +162 -0
- package/src/specs/thinking-prune.test.ts +393 -0
|
@@ -0,0 +1,162 @@
|
|
|
1
|
+
// src/scripts/cli.ts
|
|
2
|
+
import { config } from 'dotenv';
|
|
3
|
+
config();
|
|
4
|
+
import { HumanMessage, BaseMessage } from '@langchain/core/messages';
|
|
5
|
+
import type { RunnableConfig } from '@langchain/core/runnables';
|
|
6
|
+
import type * as t from '@/types';
|
|
7
|
+
import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
|
|
8
|
+
import { ToolEndHandler, ModelEndHandler, createMetadataAggregator } from '@/events';
|
|
9
|
+
import { getLLMConfig } from '@/utils/llmConfig';
|
|
10
|
+
import { getArgs } from '@/scripts/args';
|
|
11
|
+
import { GraphEvents } from '@/common';
|
|
12
|
+
import { Run } from '@/run';
|
|
13
|
+
import { createCodeExecutionTool } from '@/tools/CodeExecutor';
|
|
14
|
+
|
|
15
|
+
const conversationHistory: BaseMessage[] = [];
|
|
16
|
+
|
|
17
|
+
async function testCodeExecution(): Promise<void> {
|
|
18
|
+
const { userName, location, provider, currentDate } = await getArgs();
|
|
19
|
+
const { contentParts, aggregateContent } = createContentAggregator();
|
|
20
|
+
const customHandlers = {
|
|
21
|
+
[GraphEvents.TOOL_END]: new ToolEndHandler(),
|
|
22
|
+
[GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
|
|
23
|
+
[GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
|
|
24
|
+
[GraphEvents.ON_RUN_STEP_COMPLETED]: {
|
|
25
|
+
handle: (event: GraphEvents.ON_RUN_STEP_COMPLETED, data: t.StreamEventData): void => {
|
|
26
|
+
console.log('====== ON_RUN_STEP_COMPLETED ======');
|
|
27
|
+
console.dir(data, { depth: null });
|
|
28
|
+
aggregateContent({ event, data: data as unknown as { result: t.ToolEndEvent } });
|
|
29
|
+
}
|
|
30
|
+
},
|
|
31
|
+
[GraphEvents.ON_RUN_STEP]: {
|
|
32
|
+
handle: (event: GraphEvents.ON_RUN_STEP, data: t.StreamEventData): void => {
|
|
33
|
+
console.log('====== ON_RUN_STEP ======');
|
|
34
|
+
console.dir(data, { depth: null });
|
|
35
|
+
aggregateContent({ event, data: data as t.RunStep });
|
|
36
|
+
}
|
|
37
|
+
},
|
|
38
|
+
[GraphEvents.ON_RUN_STEP_DELTA]: {
|
|
39
|
+
handle: (event: GraphEvents.ON_RUN_STEP_DELTA, data: t.StreamEventData): void => {
|
|
40
|
+
console.log('====== ON_RUN_STEP_DELTA ======');
|
|
41
|
+
console.dir(data, { depth: null });
|
|
42
|
+
aggregateContent({ event, data: data as t.RunStepDeltaEvent });
|
|
43
|
+
}
|
|
44
|
+
},
|
|
45
|
+
[GraphEvents.ON_MESSAGE_DELTA]: {
|
|
46
|
+
handle: (event: GraphEvents.ON_MESSAGE_DELTA, data: t.StreamEventData): void => {
|
|
47
|
+
console.log('====== ON_MESSAGE_DELTA ======');
|
|
48
|
+
console.dir(data, { depth: null });
|
|
49
|
+
aggregateContent({ event, data: data as t.MessageDeltaEvent });
|
|
50
|
+
}
|
|
51
|
+
},
|
|
52
|
+
[GraphEvents.TOOL_START]: {
|
|
53
|
+
handle: (_event: string, data: t.StreamEventData, metadata?: Record<string, unknown>): void => {
|
|
54
|
+
console.log('====== TOOL_START ======');
|
|
55
|
+
console.dir(data, { depth: null });
|
|
56
|
+
}
|
|
57
|
+
},
|
|
58
|
+
};
|
|
59
|
+
|
|
60
|
+
const llmConfig = getLLMConfig(provider);
|
|
61
|
+
|
|
62
|
+
const run = await Run.create<t.IState>({
|
|
63
|
+
runId: 'message-num-1',
|
|
64
|
+
graphConfig: {
|
|
65
|
+
type: 'standard',
|
|
66
|
+
llmConfig,
|
|
67
|
+
tools: [createCodeExecutionTool()],
|
|
68
|
+
instructions: 'You are a friendly AI assistant with coding capabilities. Always address the user by their name.',
|
|
69
|
+
additional_instructions: `The user's name is ${userName} and they are located in ${location}. The current date is ${currentDate}.`,
|
|
70
|
+
},
|
|
71
|
+
returnContent: true,
|
|
72
|
+
customHandlers,
|
|
73
|
+
});
|
|
74
|
+
|
|
75
|
+
const config: Partial<RunnableConfig> & { version: 'v1' | 'v2'; run_id?: string; streamMode: string } = {
|
|
76
|
+
configurable: {
|
|
77
|
+
provider,
|
|
78
|
+
thread_id: 'conversation-num-1',
|
|
79
|
+
},
|
|
80
|
+
streamMode: 'values',
|
|
81
|
+
version: 'v2' as const,
|
|
82
|
+
// recursionLimit: 3,
|
|
83
|
+
};
|
|
84
|
+
|
|
85
|
+
console.log('Test 1: Create Project Plan');
|
|
86
|
+
|
|
87
|
+
const userMessage1 = `
|
|
88
|
+
Hi ${userName} here. We are testing your file capabilities.
|
|
89
|
+
|
|
90
|
+
1. Create a text file named "project_plan.txt" that contains: "This is a project plan for a new software development project."
|
|
91
|
+
|
|
92
|
+
Please generate this file so I can review it.
|
|
93
|
+
`;
|
|
94
|
+
|
|
95
|
+
conversationHistory.push(new HumanMessage(userMessage1));
|
|
96
|
+
|
|
97
|
+
let inputs = {
|
|
98
|
+
messages: conversationHistory,
|
|
99
|
+
};
|
|
100
|
+
const finalContentParts1 = await run.processStream(inputs, config);
|
|
101
|
+
const finalMessages1 = run.getRunMessages();
|
|
102
|
+
if (finalMessages1) {
|
|
103
|
+
conversationHistory.push(...finalMessages1);
|
|
104
|
+
}
|
|
105
|
+
console.log('\n\n====================\n\n');
|
|
106
|
+
console.dir(contentParts, { depth: null });
|
|
107
|
+
|
|
108
|
+
console.log('Test 2: Edit Project Plan');
|
|
109
|
+
|
|
110
|
+
const userMessage2 = `
|
|
111
|
+
Thanks for creating the project plan. Now I'd like you to edit the same plan to:
|
|
112
|
+
|
|
113
|
+
1. Add a new section called "Technology Stack" that contains: "The technology stack for this project includes the following technologies" and nothing more.
|
|
114
|
+
|
|
115
|
+
`;
|
|
116
|
+
|
|
117
|
+
// Make sure to pass the file ID of the previous file you created and explicitly duplicate or rename the file in your code so we can then access it. Also print the contents of the new file to ensure we did what we wanted.`;
|
|
118
|
+
|
|
119
|
+
conversationHistory.push(new HumanMessage(userMessage2));
|
|
120
|
+
|
|
121
|
+
inputs = {
|
|
122
|
+
messages: conversationHistory,
|
|
123
|
+
};
|
|
124
|
+
const finalContentParts2 = await run.processStream(inputs, config);
|
|
125
|
+
const finalMessages2 = run.getRunMessages();
|
|
126
|
+
if (finalMessages2) {
|
|
127
|
+
conversationHistory.push(...finalMessages2);
|
|
128
|
+
}
|
|
129
|
+
console.log('\n\n====================\n\n');
|
|
130
|
+
console.dir(contentParts, { depth: null });
|
|
131
|
+
|
|
132
|
+
const { handleLLMEnd, collected } = createMetadataAggregator();
|
|
133
|
+
const titleResult = await run.generateTitle({
|
|
134
|
+
inputText: userMessage2,
|
|
135
|
+
contentParts,
|
|
136
|
+
chainOptions: {
|
|
137
|
+
callbacks: [{
|
|
138
|
+
handleLLMEnd,
|
|
139
|
+
}],
|
|
140
|
+
},
|
|
141
|
+
});
|
|
142
|
+
console.log('Generated Title:', titleResult);
|
|
143
|
+
console.log('Collected metadata:', collected);
|
|
144
|
+
}
|
|
145
|
+
|
|
146
|
+
process.on('unhandledRejection', (reason, promise) => {
|
|
147
|
+
console.error('Unhandled Rejection at:', promise, 'reason:', reason);
|
|
148
|
+
console.log('Conversation history:');
|
|
149
|
+
console.dir(conversationHistory, { depth: null });
|
|
150
|
+
process.exit(1);
|
|
151
|
+
});
|
|
152
|
+
|
|
153
|
+
process.on('uncaughtException', (err) => {
|
|
154
|
+
console.error('Uncaught Exception:', err);
|
|
155
|
+
});
|
|
156
|
+
|
|
157
|
+
testCodeExecution().catch((err) => {
|
|
158
|
+
console.error(err);
|
|
159
|
+
console.log('Conversation history:');
|
|
160
|
+
console.dir(conversationHistory, { depth: null });
|
|
161
|
+
process.exit(1);
|
|
162
|
+
});
|
|
@@ -0,0 +1,393 @@
|
|
|
1
|
+
// src/specs/thinking-prune.test.ts
|
|
2
|
+
import { HumanMessage, AIMessage, SystemMessage, BaseMessage, ToolMessage } from '@langchain/core/messages';
|
|
3
|
+
import type * as t from '@/types';
|
|
4
|
+
import { createPruneMessages } from '@/messages/prune';
|
|
5
|
+
|
|
6
|
+
// Create a simple token counter for testing
|
|
7
|
+
const createTestTokenCounter = (): t.TokenCounter => {
|
|
8
|
+
return (message: BaseMessage): number => {
|
|
9
|
+
// Use type assertion to help TypeScript understand the type
|
|
10
|
+
const content = message.content as string | Array<any> | undefined;
|
|
11
|
+
|
|
12
|
+
// Handle string content
|
|
13
|
+
if (typeof content === 'string') {
|
|
14
|
+
return content.length;
|
|
15
|
+
}
|
|
16
|
+
|
|
17
|
+
// Handle array content
|
|
18
|
+
if (Array.isArray(content)) {
|
|
19
|
+
let totalLength = 0;
|
|
20
|
+
|
|
21
|
+
for (const item of content) {
|
|
22
|
+
if (typeof item === 'string') {
|
|
23
|
+
totalLength += item.length;
|
|
24
|
+
} else if (item && typeof item === 'object') {
|
|
25
|
+
if (item.type === 'thinking' && typeof item.thinking === 'string') {
|
|
26
|
+
totalLength += item.thinking.length;
|
|
27
|
+
} else if ('text' in item && typeof item.text === 'string') {
|
|
28
|
+
totalLength += item.text.length;
|
|
29
|
+
} else if ('input' in item && typeof item.input === 'string') {
|
|
30
|
+
totalLength += item.input.length;
|
|
31
|
+
}
|
|
32
|
+
}
|
|
33
|
+
}
|
|
34
|
+
|
|
35
|
+
return totalLength;
|
|
36
|
+
}
|
|
37
|
+
|
|
38
|
+
// Default case - if content is null, undefined, or any other type
|
|
39
|
+
return 0;
|
|
40
|
+
};
|
|
41
|
+
};
|
|
42
|
+
|
|
43
|
+
describe('Prune Messages with Thinking Mode Tests', () => {
|
|
44
|
+
jest.setTimeout(30000);
|
|
45
|
+
|
|
46
|
+
it('should preserve thinking blocks when pruning with thinking mode enabled', () => {
|
|
47
|
+
// Create a token counter
|
|
48
|
+
const tokenCounter = createTestTokenCounter();
|
|
49
|
+
|
|
50
|
+
// Create messages based on the example provided
|
|
51
|
+
const userMessage = new HumanMessage({
|
|
52
|
+
content: [
|
|
53
|
+
{
|
|
54
|
+
type: "text",
|
|
55
|
+
text: "/home/danny/LibreChat/gistfile1.txt\n\nread it 200 lines at a time\n\nthere are 5000 lines\n\ndo not stop until done",
|
|
56
|
+
},
|
|
57
|
+
],
|
|
58
|
+
});
|
|
59
|
+
|
|
60
|
+
const assistantMessageWithThinking = new AIMessage({
|
|
61
|
+
content: [
|
|
62
|
+
{
|
|
63
|
+
type: "thinking",
|
|
64
|
+
thinking: "The user is asking me to read a file located at `/home/danny/LibreChat/gistfile1.txt` in chunks of 200 lines at a time, mentioning that the file has 5000 lines total. They want me to continue reading through the entire file without stopping.\n\nI'll need to use the text editor tool to view the file in chunks of 200 lines each. Since the file has 5000 lines, I'll need to view it in 25 chunks (5000 ÷ 200 = 25).\n\nI'll need to make multiple calls to the text editor with the `view` command, specifying different line ranges for each call.\n\nLet me plan out the approach:\n1. Start with lines 1-200\n2. Then 201-400\n3. Then 401-600\n4. And so on until I reach 4801-5000\n\nFor each call, I'll use the `view` command with the specific line range in the `view_range` parameter. I'll continue until I've shown all 5000 lines as requested.",
|
|
65
|
+
},
|
|
66
|
+
{
|
|
67
|
+
type: "text",
|
|
68
|
+
text: "I'll read the file for you in chunks of 200 lines at a time until all 5000 lines are displayed.\n\nLet me start with the first 200 lines:",
|
|
69
|
+
},
|
|
70
|
+
{
|
|
71
|
+
type: "tool_use",
|
|
72
|
+
id: "toolu_01YApWuFsEQCuBFDgYXmcmeZ",
|
|
73
|
+
name: "text_editor_mcp_textEditor",
|
|
74
|
+
input: "{\"command\": \"view\", \"path\": \"/home/danny/LibreChat/gistfile1.txt\", \"description\": \"Viewing lines 1-200 of the file\", \"view_range\": [1,200]}",
|
|
75
|
+
},
|
|
76
|
+
],
|
|
77
|
+
});
|
|
78
|
+
|
|
79
|
+
const toolResponseMessage1 = new ToolMessage({
|
|
80
|
+
content: [
|
|
81
|
+
{
|
|
82
|
+
type: "text",
|
|
83
|
+
text: "{\"success\":true,\"message\":\"File content (truncated):\"",
|
|
84
|
+
},
|
|
85
|
+
],
|
|
86
|
+
tool_call_id: "toolu_01YApWuFsEQCuBFDgYXmcmeZ",
|
|
87
|
+
name: "text_editor_mcp_textEditor",
|
|
88
|
+
});
|
|
89
|
+
|
|
90
|
+
const assistantMessage2 = new AIMessage({
|
|
91
|
+
content: [
|
|
92
|
+
{
|
|
93
|
+
type: "text",
|
|
94
|
+
text: "It looks like the file content is clipped in the response. I'll continue with the next chunks:",
|
|
95
|
+
},
|
|
96
|
+
{
|
|
97
|
+
type: "tool_use",
|
|
98
|
+
id: "toolu_01VnyMQ4CvEd6zLDxxtTd6d4",
|
|
99
|
+
name: "text_editor_mcp_textEditor",
|
|
100
|
+
input: "{\"command\": \"view\", \"path\": \"/home/danny/LibreChat/gistfile1.txt\", \"description\": \"Viewing lines 201-400 of the file\", \"view_range\": [201,400]}",
|
|
101
|
+
},
|
|
102
|
+
],
|
|
103
|
+
});
|
|
104
|
+
|
|
105
|
+
const toolResponseMessage2 = new ToolMessage({
|
|
106
|
+
content: [
|
|
107
|
+
{
|
|
108
|
+
type: "text",
|
|
109
|
+
text: "{\"success\":true,\"message\":\"File content (truncated):\"",
|
|
110
|
+
},
|
|
111
|
+
],
|
|
112
|
+
tool_call_id: "toolu_01VnyMQ4CvEd6zLDxxtTd6d4",
|
|
113
|
+
name: "text_editor_mcp_textEditor",
|
|
114
|
+
});
|
|
115
|
+
|
|
116
|
+
const assistantMessage3 = new AIMessage({
|
|
117
|
+
content: [
|
|
118
|
+
{
|
|
119
|
+
type: "text",
|
|
120
|
+
text: "Continuing with the next chunk:",
|
|
121
|
+
},
|
|
122
|
+
{
|
|
123
|
+
type: "tool_use",
|
|
124
|
+
id: "toolu_01TZKs4nnBc58BYXKz1Mw4fp",
|
|
125
|
+
name: "text_editor_mcp_textEditor",
|
|
126
|
+
input: "{\"command\": \"view\", \"path\": \"/home/danny/LibreChat/gistfile1.txt\", \"description\": \"Viewing lines 401-600 of the file\", \"view_range\": [401,600]}",
|
|
127
|
+
},
|
|
128
|
+
],
|
|
129
|
+
});
|
|
130
|
+
|
|
131
|
+
const toolResponseMessage3 = new ToolMessage({
|
|
132
|
+
content: [
|
|
133
|
+
{
|
|
134
|
+
type: "text",
|
|
135
|
+
text: "{\"success\":true,\"message\":\"File content (truncated):\"",
|
|
136
|
+
},
|
|
137
|
+
],
|
|
138
|
+
tool_call_id: "toolu_01TZKs4nnBc58BYXKz1Mw4fp",
|
|
139
|
+
name: "text_editor_mcp_textEditor",
|
|
140
|
+
});
|
|
141
|
+
|
|
142
|
+
const assistantMessage4 = new AIMessage({
|
|
143
|
+
content: [
|
|
144
|
+
{
|
|
145
|
+
type: "text",
|
|
146
|
+
text: "Let's continue with the next chunk:",
|
|
147
|
+
},
|
|
148
|
+
{
|
|
149
|
+
type: "tool_use",
|
|
150
|
+
id: "toolu_01TZgBacNxjx1QNUpJg9hca5",
|
|
151
|
+
name: "text_editor_mcp_textEditor",
|
|
152
|
+
input: "{\"command\": \"view\", \"path\": \"/home/danny/LibreChat/gistfile1.txt\", \"description\": \"Viewing lines 601-800 of the file\", \"view_range\": [601,800]}",
|
|
153
|
+
},
|
|
154
|
+
],
|
|
155
|
+
});
|
|
156
|
+
|
|
157
|
+
const toolResponseMessage4 = new ToolMessage({
|
|
158
|
+
content: [
|
|
159
|
+
{
|
|
160
|
+
type: "text",
|
|
161
|
+
text: "{\"success\":true,\"message\":\"File content (truncated):\"",
|
|
162
|
+
},
|
|
163
|
+
],
|
|
164
|
+
tool_call_id: "toolu_01TZgBacNxjx1QNUpJg9hca5",
|
|
165
|
+
name: "text_editor_mcp_textEditor",
|
|
166
|
+
});
|
|
167
|
+
|
|
168
|
+
const messages = [
|
|
169
|
+
userMessage,
|
|
170
|
+
assistantMessageWithThinking,
|
|
171
|
+
toolResponseMessage1,
|
|
172
|
+
assistantMessage2,
|
|
173
|
+
toolResponseMessage2,
|
|
174
|
+
assistantMessage3,
|
|
175
|
+
toolResponseMessage3,
|
|
176
|
+
assistantMessage4,
|
|
177
|
+
toolResponseMessage4,
|
|
178
|
+
];
|
|
179
|
+
|
|
180
|
+
// Create indexTokenCountMap based on the example provided
|
|
181
|
+
const indexTokenCountMap = {
|
|
182
|
+
"0": 617, // userMessage
|
|
183
|
+
"1": 52, // assistantMessageWithThinking
|
|
184
|
+
"2": 4995, // toolResponseMessage1
|
|
185
|
+
"3": 307, // assistantMessage2
|
|
186
|
+
"4": 9359, // toolResponseMessage2
|
|
187
|
+
"5": 178, // assistantMessage3
|
|
188
|
+
"6": 5463, // toolResponseMessage3
|
|
189
|
+
"7": 125, // assistantMessage4
|
|
190
|
+
"8": 4264, // toolResponseMessage4
|
|
191
|
+
};
|
|
192
|
+
|
|
193
|
+
// Create pruneMessages function with thinking mode enabled
|
|
194
|
+
const pruneMessages = createPruneMessages({
|
|
195
|
+
maxTokens: 19800,
|
|
196
|
+
startIndex: 0,
|
|
197
|
+
tokenCounter,
|
|
198
|
+
indexTokenCountMap: { ...indexTokenCountMap },
|
|
199
|
+
thinkingEnabled: true,
|
|
200
|
+
});
|
|
201
|
+
|
|
202
|
+
// Prune messages
|
|
203
|
+
const result = pruneMessages({
|
|
204
|
+
messages,
|
|
205
|
+
usageMetadata: {
|
|
206
|
+
input_tokens: 25254,
|
|
207
|
+
output_tokens: 106,
|
|
208
|
+
total_tokens: 25360,
|
|
209
|
+
input_token_details: {
|
|
210
|
+
cache_read: 0,
|
|
211
|
+
cache_creation: 0,
|
|
212
|
+
},
|
|
213
|
+
},
|
|
214
|
+
startOnMessageType: 'human',
|
|
215
|
+
});
|
|
216
|
+
|
|
217
|
+
// Verify that the first assistant message in the pruned context has a thinking block
|
|
218
|
+
expect(result.context.length).toBeGreaterThan(0);
|
|
219
|
+
|
|
220
|
+
// Find the first assistant message in the pruned context
|
|
221
|
+
const firstAssistantIndex = result.context.findIndex(msg => msg.getType() === 'ai');
|
|
222
|
+
expect(firstAssistantIndex).toBe(0);
|
|
223
|
+
|
|
224
|
+
const firstAssistantMsg = result.context[firstAssistantIndex];
|
|
225
|
+
expect(Array.isArray(firstAssistantMsg.content)).toBe(true);
|
|
226
|
+
|
|
227
|
+
// Verify that the first assistant message has a thinking block
|
|
228
|
+
const hasThinkingBlock = (firstAssistantMsg.content as any[]).some(item =>
|
|
229
|
+
item && typeof item === 'object' && item.type === 'thinking');
|
|
230
|
+
expect(hasThinkingBlock).toBe(true);
|
|
231
|
+
|
|
232
|
+
// Verify that the thinking block is from the original assistant message
|
|
233
|
+
const thinkingBlock = (firstAssistantMsg.content as any[]).find(item =>
|
|
234
|
+
item && typeof item === 'object' && item.type === 'thinking');
|
|
235
|
+
expect(thinkingBlock).toBeDefined();
|
|
236
|
+
expect(thinkingBlock.thinking).toContain("The user is asking me to read a file");
|
|
237
|
+
});
|
|
238
|
+
|
|
239
|
+
it('should handle token recalculation when inserting thinking blocks', () => {
|
|
240
|
+
// Create a token counter
|
|
241
|
+
const tokenCounter = createTestTokenCounter();
|
|
242
|
+
|
|
243
|
+
// Create a message with thinking block
|
|
244
|
+
const assistantMessageWithThinking = new AIMessage({
|
|
245
|
+
content: [
|
|
246
|
+
{
|
|
247
|
+
type: "thinking",
|
|
248
|
+
thinking: "This is a thinking block",
|
|
249
|
+
},
|
|
250
|
+
{
|
|
251
|
+
type: "text",
|
|
252
|
+
text: "Response with thinking",
|
|
253
|
+
},
|
|
254
|
+
],
|
|
255
|
+
});
|
|
256
|
+
|
|
257
|
+
// Create a message without thinking block
|
|
258
|
+
const assistantMessageWithoutThinking = new AIMessage({
|
|
259
|
+
content: [
|
|
260
|
+
{
|
|
261
|
+
type: "text",
|
|
262
|
+
text: "Response without thinking",
|
|
263
|
+
},
|
|
264
|
+
],
|
|
265
|
+
});
|
|
266
|
+
|
|
267
|
+
const messages = [
|
|
268
|
+
new SystemMessage("System instruction"),
|
|
269
|
+
new HumanMessage("Hello"),
|
|
270
|
+
assistantMessageWithThinking,
|
|
271
|
+
new HumanMessage("Next message"),
|
|
272
|
+
assistantMessageWithoutThinking,
|
|
273
|
+
];
|
|
274
|
+
|
|
275
|
+
// Calculate token counts for each message
|
|
276
|
+
const indexTokenCountMap: Record<string, number> = {};
|
|
277
|
+
for (let i = 0; i < messages.length; i++) {
|
|
278
|
+
indexTokenCountMap[i] = tokenCounter(messages[i]);
|
|
279
|
+
}
|
|
280
|
+
|
|
281
|
+
// Create pruneMessages function with thinking mode enabled
|
|
282
|
+
const pruneMessages = createPruneMessages({
|
|
283
|
+
maxTokens: 50, // Set a low token limit to force pruning
|
|
284
|
+
startIndex: 0,
|
|
285
|
+
tokenCounter,
|
|
286
|
+
indexTokenCountMap: { ...indexTokenCountMap },
|
|
287
|
+
thinkingEnabled: true,
|
|
288
|
+
});
|
|
289
|
+
|
|
290
|
+
// Prune messages
|
|
291
|
+
const result = pruneMessages({ messages });
|
|
292
|
+
|
|
293
|
+
// Verify that the pruned context has fewer messages than the original
|
|
294
|
+
expect(result.context.length).toBeLessThan(messages.length);
|
|
295
|
+
|
|
296
|
+
// Find the first assistant message in the pruned context
|
|
297
|
+
const firstAssistantIndex = result.context.findIndex(msg => msg.getType() === 'ai');
|
|
298
|
+
expect(firstAssistantIndex).toBeGreaterThan(-1);
|
|
299
|
+
|
|
300
|
+
const firstAssistantMsg = result.context[firstAssistantIndex];
|
|
301
|
+
expect(Array.isArray(firstAssistantMsg.content)).toBe(true);
|
|
302
|
+
|
|
303
|
+
// Verify that the first assistant message has a thinking block
|
|
304
|
+
const hasThinkingBlock = (firstAssistantMsg.content as any[]).some(item =>
|
|
305
|
+
item && typeof item === 'object' && item.type === 'thinking');
|
|
306
|
+
expect(hasThinkingBlock).toBe(true);
|
|
307
|
+
});
|
|
308
|
+
|
|
309
|
+
it('should not modify messages when under token limit', () => {
|
|
310
|
+
// Create a token counter
|
|
311
|
+
const tokenCounter = createTestTokenCounter();
|
|
312
|
+
|
|
313
|
+
// Create a message with thinking block
|
|
314
|
+
const assistantMessageWithThinking = new AIMessage({
|
|
315
|
+
content: [
|
|
316
|
+
{
|
|
317
|
+
type: "thinking",
|
|
318
|
+
thinking: "This is a thinking block",
|
|
319
|
+
},
|
|
320
|
+
{
|
|
321
|
+
type: "text",
|
|
322
|
+
text: "Response with thinking",
|
|
323
|
+
},
|
|
324
|
+
],
|
|
325
|
+
});
|
|
326
|
+
|
|
327
|
+
const messages = [
|
|
328
|
+
new SystemMessage("System instruction"),
|
|
329
|
+
new HumanMessage("Hello"),
|
|
330
|
+
assistantMessageWithThinking,
|
|
331
|
+
];
|
|
332
|
+
|
|
333
|
+
// Calculate token counts for each message
|
|
334
|
+
const indexTokenCountMap: Record<string, number> = {};
|
|
335
|
+
for (let i = 0; i < messages.length; i++) {
|
|
336
|
+
indexTokenCountMap[i] = tokenCounter(messages[i]);
|
|
337
|
+
}
|
|
338
|
+
|
|
339
|
+
// Create pruneMessages function with thinking mode enabled
|
|
340
|
+
const pruneMessages = createPruneMessages({
|
|
341
|
+
maxTokens: 1000, // Set a high token limit to avoid pruning
|
|
342
|
+
startIndex: 0,
|
|
343
|
+
tokenCounter,
|
|
344
|
+
indexTokenCountMap: { ...indexTokenCountMap },
|
|
345
|
+
thinkingEnabled: true,
|
|
346
|
+
});
|
|
347
|
+
|
|
348
|
+
// Prune messages
|
|
349
|
+
const result = pruneMessages({ messages });
|
|
350
|
+
|
|
351
|
+
// Verify that all messages are preserved
|
|
352
|
+
expect(result.context.length).toBe(messages.length);
|
|
353
|
+
expect(result.context).toEqual(messages);
|
|
354
|
+
});
|
|
355
|
+
|
|
356
|
+
it('should handle the case when no thinking blocks are found', () => {
|
|
357
|
+
// Create a token counter
|
|
358
|
+
const tokenCounter = createTestTokenCounter();
|
|
359
|
+
|
|
360
|
+
// Create messages without thinking blocks
|
|
361
|
+
const messages = [
|
|
362
|
+
new SystemMessage("System instruction"),
|
|
363
|
+
new HumanMessage("Hello"),
|
|
364
|
+
new AIMessage("Response without thinking"),
|
|
365
|
+
new HumanMessage("Next message"),
|
|
366
|
+
new AIMessage("Another response without thinking"),
|
|
367
|
+
];
|
|
368
|
+
|
|
369
|
+
// Calculate token counts for each message
|
|
370
|
+
const indexTokenCountMap: Record<string, number> = {};
|
|
371
|
+
for (let i = 0; i < messages.length; i++) {
|
|
372
|
+
indexTokenCountMap[i] = tokenCounter(messages[i]);
|
|
373
|
+
}
|
|
374
|
+
|
|
375
|
+
// Create pruneMessages function with thinking mode enabled
|
|
376
|
+
const pruneMessages = createPruneMessages({
|
|
377
|
+
maxTokens: 50, // Set a low token limit to force pruning
|
|
378
|
+
startIndex: 0,
|
|
379
|
+
tokenCounter,
|
|
380
|
+
indexTokenCountMap: { ...indexTokenCountMap },
|
|
381
|
+
thinkingEnabled: true,
|
|
382
|
+
});
|
|
383
|
+
|
|
384
|
+
// Prune messages
|
|
385
|
+
const result = pruneMessages({ messages });
|
|
386
|
+
|
|
387
|
+
// Verify that the pruned context has fewer messages than the original
|
|
388
|
+
expect(result.context.length).toBeLessThan(messages.length);
|
|
389
|
+
|
|
390
|
+
// The function should not throw an error even though no thinking blocks are found
|
|
391
|
+
expect(() => pruneMessages({ messages })).not.toThrow();
|
|
392
|
+
});
|
|
393
|
+
});
|