@librechat/agents 2.3.1 → 2.3.3
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 +6 -6
- package/dist/cjs/graphs/Graph.cjs.map +1 -1
- package/dist/cjs/llm/anthropic/llm.cjs +7 -7
- package/dist/cjs/llm/anthropic/llm.cjs.map +1 -1
- package/dist/cjs/llm/anthropic/utils/message_inputs.cjs +6 -6
- package/dist/cjs/llm/anthropic/utils/message_inputs.cjs.map +1 -1
- package/dist/cjs/llm/anthropic/utils/message_outputs.cjs +24 -24
- package/dist/cjs/llm/anthropic/utils/message_outputs.cjs.map +1 -1
- package/dist/cjs/llm/fake.cjs.map +1 -1
- package/dist/cjs/llm/text.cjs.map +1 -1
- package/dist/cjs/main.cjs +3 -0
- package/dist/cjs/main.cjs.map +1 -1
- package/dist/cjs/messages/core.cjs +6 -6
- package/dist/cjs/messages/core.cjs.map +1 -1
- package/dist/cjs/messages/format.cjs +11 -9
- package/dist/cjs/messages/format.cjs.map +1 -1
- package/dist/cjs/messages/prune.cjs +155 -205
- package/dist/cjs/messages/prune.cjs.map +1 -1
- package/dist/cjs/run.cjs.map +1 -1
- package/dist/cjs/stream.cjs +7 -7
- package/dist/cjs/stream.cjs.map +1 -1
- package/dist/cjs/tools/ToolNode.cjs +1 -1
- package/dist/cjs/tools/ToolNode.cjs.map +1 -1
- package/dist/cjs/utils/tokens.cjs +3 -3
- package/dist/cjs/utils/tokens.cjs.map +1 -1
- package/dist/esm/graphs/Graph.mjs +6 -6
- package/dist/esm/graphs/Graph.mjs.map +1 -1
- package/dist/esm/llm/anthropic/llm.mjs +7 -7
- package/dist/esm/llm/anthropic/llm.mjs.map +1 -1
- package/dist/esm/llm/anthropic/utils/message_inputs.mjs +6 -6
- package/dist/esm/llm/anthropic/utils/message_inputs.mjs.map +1 -1
- package/dist/esm/llm/anthropic/utils/message_outputs.mjs +24 -24
- package/dist/esm/llm/anthropic/utils/message_outputs.mjs.map +1 -1
- package/dist/esm/llm/fake.mjs.map +1 -1
- package/dist/esm/llm/text.mjs.map +1 -1
- package/dist/esm/main.mjs +1 -1
- package/dist/esm/messages/core.mjs +6 -6
- package/dist/esm/messages/core.mjs.map +1 -1
- package/dist/esm/messages/format.mjs +11 -9
- package/dist/esm/messages/format.mjs.map +1 -1
- package/dist/esm/messages/prune.mjs +153 -206
- package/dist/esm/messages/prune.mjs.map +1 -1
- package/dist/esm/run.mjs.map +1 -1
- package/dist/esm/stream.mjs +7 -7
- package/dist/esm/stream.mjs.map +1 -1
- package/dist/esm/tools/ToolNode.mjs +1 -1
- package/dist/esm/tools/ToolNode.mjs.map +1 -1
- package/dist/esm/utils/tokens.mjs +3 -3
- package/dist/esm/utils/tokens.mjs.map +1 -1
- package/dist/types/messages/format.d.ts +1 -2
- package/dist/types/messages/prune.d.ts +31 -2
- package/dist/types/types/stream.d.ts +2 -2
- package/dist/types/utils/tokens.d.ts +1 -1
- package/package.json +4 -3
- package/src/graphs/Graph.ts +8 -8
- package/src/llm/anthropic/llm.ts +7 -8
- package/src/llm/anthropic/types.ts +4 -4
- package/src/llm/anthropic/utils/message_inputs.ts +6 -6
- package/src/llm/anthropic/utils/message_outputs.ts +39 -39
- package/src/llm/fake.ts +2 -2
- package/src/llm/text.ts +1 -1
- package/src/messages/core.ts +9 -9
- package/src/messages/format.ts +43 -42
- package/src/messages/formatAgentMessages.test.ts +35 -35
- package/src/messages/formatAgentMessages.tools.test.ts +30 -30
- package/src/messages/prune.ts +182 -255
- package/src/messages/shiftIndexTokenCountMap.test.ts +18 -18
- package/src/mockStream.ts +1 -1
- package/src/run.ts +2 -2
- package/src/specs/prune.test.ts +89 -89
- package/src/specs/reasoning.test.ts +1 -1
- package/src/specs/thinking-prune.test.ts +265 -261
- package/src/specs/tool-error.test.ts +16 -17
- package/src/stream.ts +21 -18
- package/src/tools/ToolNode.ts +1 -1
- package/src/types/stream.ts +4 -3
- package/src/utils/tokens.ts +12 -12
package/src/specs/prune.test.ts
CHANGED
|
@@ -5,40 +5,40 @@ import { HumanMessage, AIMessage, SystemMessage, BaseMessage } from '@langchain/
|
|
|
5
5
|
import type { RunnableConfig } from '@langchain/core/runnables';
|
|
6
6
|
import type { UsageMetadata } from '@langchain/core/messages';
|
|
7
7
|
import type * as t from '@/types';
|
|
8
|
-
import {
|
|
8
|
+
import { createPruneMessages } from '@/messages/prune';
|
|
9
9
|
import { getLLMConfig } from '@/utils/llmConfig';
|
|
10
|
+
import { Providers } from '@/common';
|
|
10
11
|
import { Run } from '@/run';
|
|
11
|
-
import { createPruneMessages } from '@/messages/prune';
|
|
12
12
|
|
|
13
13
|
// Create a simple token counter for testing
|
|
14
14
|
const createTestTokenCounter = (): t.TokenCounter => {
|
|
15
15
|
// This simple token counter just counts characters as tokens for predictable testing
|
|
16
16
|
return (message: BaseMessage): number => {
|
|
17
17
|
// Use type assertion to help TypeScript understand the type
|
|
18
|
-
const content = message.content as string | Array<
|
|
19
|
-
|
|
18
|
+
const content = message.content as string | Array<t.MessageContentComplex | string> | undefined;
|
|
19
|
+
|
|
20
20
|
// Handle string content
|
|
21
21
|
if (typeof content === 'string') {
|
|
22
22
|
return content.length;
|
|
23
23
|
}
|
|
24
|
-
|
|
24
|
+
|
|
25
25
|
// Handle array content
|
|
26
26
|
if (Array.isArray(content)) {
|
|
27
27
|
let totalLength = 0;
|
|
28
|
-
|
|
28
|
+
|
|
29
29
|
for (const item of content) {
|
|
30
30
|
if (typeof item === 'string') {
|
|
31
31
|
totalLength += item.length;
|
|
32
|
-
} else if (
|
|
32
|
+
} else if (typeof item === 'object') {
|
|
33
33
|
if ('text' in item && typeof item.text === 'string') {
|
|
34
34
|
totalLength += item.text.length;
|
|
35
35
|
}
|
|
36
36
|
}
|
|
37
37
|
}
|
|
38
|
-
|
|
38
|
+
|
|
39
39
|
return totalLength;
|
|
40
40
|
}
|
|
41
|
-
|
|
41
|
+
|
|
42
42
|
// Default case - if content is null, undefined, or any other type
|
|
43
43
|
return 0;
|
|
44
44
|
};
|
|
@@ -50,7 +50,7 @@ function calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {
|
|
|
50
50
|
const baseInputTokens = Number(usage.input_tokens) || 0;
|
|
51
51
|
const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;
|
|
52
52
|
const cacheRead = Number(usage.input_token_details?.cache_read) || 0;
|
|
53
|
-
|
|
53
|
+
|
|
54
54
|
const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;
|
|
55
55
|
const totalOutputTokens = Number(usage.output_tokens) || 0;
|
|
56
56
|
|
|
@@ -65,12 +65,12 @@ function getMessagesWithinTokenLimit({
|
|
|
65
65
|
messages: _messages,
|
|
66
66
|
maxContextTokens,
|
|
67
67
|
indexTokenCountMap,
|
|
68
|
-
|
|
68
|
+
startType,
|
|
69
69
|
}: {
|
|
70
70
|
messages: BaseMessage[];
|
|
71
71
|
maxContextTokens: number;
|
|
72
72
|
indexTokenCountMap: Record<string, number>;
|
|
73
|
-
|
|
73
|
+
startType?: string;
|
|
74
74
|
}): {
|
|
75
75
|
context: BaseMessage[];
|
|
76
76
|
remainingContextTokens: number;
|
|
@@ -81,7 +81,7 @@ function getMessagesWithinTokenLimit({
|
|
|
81
81
|
// start with 3 tokens for the label after all messages have been counted.
|
|
82
82
|
let summaryIndex = -1;
|
|
83
83
|
let currentTokenCount = 3;
|
|
84
|
-
const instructions = _messages
|
|
84
|
+
const instructions = _messages[0]?.getType() === 'system' ? _messages[0] : undefined;
|
|
85
85
|
const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;
|
|
86
86
|
let remainingContextTokens = maxContextTokens - instructionsTokenCount;
|
|
87
87
|
const messages = [..._messages];
|
|
@@ -96,7 +96,7 @@ function getMessagesWithinTokenLimit({
|
|
|
96
96
|
}
|
|
97
97
|
const poppedMessage = messages.pop();
|
|
98
98
|
if (!poppedMessage) continue;
|
|
99
|
-
|
|
99
|
+
|
|
100
100
|
const tokenCount = indexTokenCountMap[currentIndex] || 0;
|
|
101
101
|
|
|
102
102
|
if ((currentTokenCount + tokenCount) <= remainingContextTokens) {
|
|
@@ -107,11 +107,11 @@ function getMessagesWithinTokenLimit({
|
|
|
107
107
|
break;
|
|
108
108
|
}
|
|
109
109
|
}
|
|
110
|
-
|
|
111
|
-
// If
|
|
112
|
-
if (
|
|
113
|
-
const requiredTypeIndex = context.findIndex(msg => msg.getType() ===
|
|
114
|
-
|
|
110
|
+
|
|
111
|
+
// If startType is specified, discard messages until we find one of the required type
|
|
112
|
+
if (startType && context.length > 0) {
|
|
113
|
+
const requiredTypeIndex = context.findIndex(msg => msg.getType() === startType);
|
|
114
|
+
|
|
115
115
|
if (requiredTypeIndex > 0) {
|
|
116
116
|
// If we found a message of the required type, discard all messages before it
|
|
117
117
|
const remainingMessages = context.slice(requiredTypeIndex);
|
|
@@ -144,7 +144,7 @@ function checkValidNumber(value: unknown): value is number {
|
|
|
144
144
|
|
|
145
145
|
describe('Prune Messages Tests', () => {
|
|
146
146
|
jest.setTimeout(30000);
|
|
147
|
-
|
|
147
|
+
|
|
148
148
|
describe('calculateTotalTokens', () => {
|
|
149
149
|
it('should calculate total tokens correctly with all fields present', () => {
|
|
150
150
|
const usage: Partial<UsageMetadata> = {
|
|
@@ -155,38 +155,38 @@ describe('Prune Messages Tests', () => {
|
|
|
155
155
|
cache_read: 5
|
|
156
156
|
}
|
|
157
157
|
};
|
|
158
|
-
|
|
158
|
+
|
|
159
159
|
const result = calculateTotalTokens(usage);
|
|
160
|
-
|
|
160
|
+
|
|
161
161
|
expect(result.input_tokens).toBe(115); // 100 + 10 + 5
|
|
162
162
|
expect(result.output_tokens).toBe(50);
|
|
163
163
|
expect(result.total_tokens).toBe(165); // 115 + 50
|
|
164
164
|
});
|
|
165
|
-
|
|
165
|
+
|
|
166
166
|
it('should handle missing fields gracefully', () => {
|
|
167
167
|
const usage: Partial<UsageMetadata> = {
|
|
168
168
|
input_tokens: 100,
|
|
169
169
|
output_tokens: 50
|
|
170
170
|
};
|
|
171
|
-
|
|
171
|
+
|
|
172
172
|
const result = calculateTotalTokens(usage);
|
|
173
|
-
|
|
173
|
+
|
|
174
174
|
expect(result.input_tokens).toBe(100);
|
|
175
175
|
expect(result.output_tokens).toBe(50);
|
|
176
176
|
expect(result.total_tokens).toBe(150);
|
|
177
177
|
});
|
|
178
|
-
|
|
178
|
+
|
|
179
179
|
it('should handle empty usage object', () => {
|
|
180
180
|
const usage: Partial<UsageMetadata> = {};
|
|
181
|
-
|
|
181
|
+
|
|
182
182
|
const result = calculateTotalTokens(usage);
|
|
183
|
-
|
|
183
|
+
|
|
184
184
|
expect(result.input_tokens).toBe(0);
|
|
185
185
|
expect(result.output_tokens).toBe(0);
|
|
186
186
|
expect(result.total_tokens).toBe(0);
|
|
187
187
|
});
|
|
188
188
|
});
|
|
189
|
-
|
|
189
|
+
|
|
190
190
|
describe('getMessagesWithinTokenLimit', () => {
|
|
191
191
|
it('should include all messages when under token limit', () => {
|
|
192
192
|
const messages = [
|
|
@@ -194,26 +194,26 @@ describe('Prune Messages Tests', () => {
|
|
|
194
194
|
new HumanMessage('Hello'),
|
|
195
195
|
new AIMessage('Hi there')
|
|
196
196
|
];
|
|
197
|
-
|
|
197
|
+
|
|
198
198
|
const indexTokenCountMap = {
|
|
199
199
|
0: 17, // "System instruction"
|
|
200
200
|
1: 5, // "Hello"
|
|
201
201
|
2: 8 // "Hi there"
|
|
202
202
|
};
|
|
203
|
-
|
|
203
|
+
|
|
204
204
|
const result = getMessagesWithinTokenLimit({
|
|
205
205
|
messages,
|
|
206
206
|
maxContextTokens: 100,
|
|
207
207
|
indexTokenCountMap
|
|
208
208
|
});
|
|
209
|
-
|
|
209
|
+
|
|
210
210
|
expect(result.context.length).toBe(3);
|
|
211
211
|
expect(result.context[0]).toBe(messages[0]); // System message
|
|
212
212
|
expect(result.context[0].getType()).toBe('system'); // System message
|
|
213
213
|
expect(result.remainingContextTokens).toBe(100 - 17 - 5 - 8 - 3); // -3 for the assistant label tokens
|
|
214
214
|
expect(result.messagesToRefine.length).toBe(0);
|
|
215
215
|
});
|
|
216
|
-
|
|
216
|
+
|
|
217
217
|
it('should prune oldest messages when over token limit', () => {
|
|
218
218
|
const messages = [
|
|
219
219
|
new SystemMessage('System instruction'),
|
|
@@ -222,7 +222,7 @@ describe('Prune Messages Tests', () => {
|
|
|
222
222
|
new HumanMessage('Message 2'),
|
|
223
223
|
new AIMessage('Response 2')
|
|
224
224
|
];
|
|
225
|
-
|
|
225
|
+
|
|
226
226
|
const indexTokenCountMap = {
|
|
227
227
|
0: 17, // "System instruction"
|
|
228
228
|
1: 9, // "Message 1"
|
|
@@ -230,54 +230,54 @@ describe('Prune Messages Tests', () => {
|
|
|
230
230
|
3: 9, // "Message 2"
|
|
231
231
|
4: 10 // "Response 2"
|
|
232
232
|
};
|
|
233
|
-
|
|
233
|
+
|
|
234
234
|
// Set a limit that can only fit the system message and the last two messages
|
|
235
235
|
const result = getMessagesWithinTokenLimit({
|
|
236
236
|
messages,
|
|
237
237
|
maxContextTokens: 40,
|
|
238
238
|
indexTokenCountMap
|
|
239
239
|
});
|
|
240
|
-
|
|
240
|
+
|
|
241
241
|
// Should include system message and the last two messages
|
|
242
242
|
expect(result.context.length).toBe(3);
|
|
243
243
|
expect(result.context[0]).toBe(messages[0]); // System message
|
|
244
244
|
expect(result.context[0].getType()).toBe('system'); // System message
|
|
245
245
|
expect(result.context[1]).toBe(messages[3]); // Message 2
|
|
246
246
|
expect(result.context[2]).toBe(messages[4]); // Response 2
|
|
247
|
-
|
|
247
|
+
|
|
248
248
|
// Should have the first two messages in messagesToRefine
|
|
249
249
|
expect(result.messagesToRefine.length).toBe(2);
|
|
250
250
|
expect(result.messagesToRefine[0]).toBe(messages[1]); // Message 1
|
|
251
251
|
expect(result.messagesToRefine[1]).toBe(messages[2]); // Response 1
|
|
252
252
|
});
|
|
253
|
-
|
|
253
|
+
|
|
254
254
|
it('should always include system message even when at token limit', () => {
|
|
255
255
|
const messages = [
|
|
256
256
|
new SystemMessage('System instruction'),
|
|
257
257
|
new HumanMessage('Hello'),
|
|
258
258
|
new AIMessage('Hi there')
|
|
259
259
|
];
|
|
260
|
-
|
|
260
|
+
|
|
261
261
|
const indexTokenCountMap = {
|
|
262
262
|
0: 17, // "System instruction"
|
|
263
263
|
1: 5, // "Hello"
|
|
264
264
|
2: 8 // "Hi there"
|
|
265
265
|
};
|
|
266
|
-
|
|
266
|
+
|
|
267
267
|
// Set a limit that can only fit the system message
|
|
268
268
|
const result = getMessagesWithinTokenLimit({
|
|
269
269
|
messages,
|
|
270
270
|
maxContextTokens: 20,
|
|
271
271
|
indexTokenCountMap
|
|
272
272
|
});
|
|
273
|
-
|
|
273
|
+
|
|
274
274
|
expect(result.context.length).toBe(1);
|
|
275
275
|
expect(result.context[0]).toBe(messages[0]); // System message
|
|
276
|
-
|
|
276
|
+
|
|
277
277
|
expect(result.messagesToRefine.length).toBe(2);
|
|
278
278
|
});
|
|
279
279
|
|
|
280
|
-
it('should start context with a specific message type when
|
|
280
|
+
it('should start context with a specific message type when startType is specified', () => {
|
|
281
281
|
const messages = [
|
|
282
282
|
new SystemMessage('System instruction'),
|
|
283
283
|
new AIMessage('AI message 1'),
|
|
@@ -285,7 +285,7 @@ describe('Prune Messages Tests', () => {
|
|
|
285
285
|
new AIMessage('AI message 2'),
|
|
286
286
|
new HumanMessage('Human message 2')
|
|
287
287
|
];
|
|
288
|
-
|
|
288
|
+
|
|
289
289
|
const indexTokenCountMap = {
|
|
290
290
|
0: 17, // "System instruction"
|
|
291
291
|
1: 12, // "AI message 1"
|
|
@@ -293,15 +293,15 @@ describe('Prune Messages Tests', () => {
|
|
|
293
293
|
3: 12, // "AI message 2"
|
|
294
294
|
4: 15 // "Human message 2"
|
|
295
295
|
};
|
|
296
|
-
|
|
296
|
+
|
|
297
297
|
// Set a limit that can fit all messages
|
|
298
298
|
const result = getMessagesWithinTokenLimit({
|
|
299
299
|
messages,
|
|
300
300
|
maxContextTokens: 100,
|
|
301
301
|
indexTokenCountMap,
|
|
302
|
-
|
|
302
|
+
startType: 'human'
|
|
303
303
|
});
|
|
304
|
-
|
|
304
|
+
|
|
305
305
|
// All messages should be included since we're under the token limit
|
|
306
306
|
expect(result.context.length).toBe(5);
|
|
307
307
|
expect(result.context[0]).toBe(messages[0]); // System message
|
|
@@ -309,7 +309,7 @@ describe('Prune Messages Tests', () => {
|
|
|
309
309
|
expect(result.context[2]).toBe(messages[2]); // Human message 1
|
|
310
310
|
expect(result.context[3]).toBe(messages[3]); // AI message 2
|
|
311
311
|
expect(result.context[4]).toBe(messages[4]); // Human message 2
|
|
312
|
-
|
|
312
|
+
|
|
313
313
|
// All messages should be included since we're under the token limit
|
|
314
314
|
expect(result.messagesToRefine.length).toBe(0);
|
|
315
315
|
});
|
|
@@ -320,44 +320,44 @@ describe('Prune Messages Tests', () => {
|
|
|
320
320
|
new AIMessage('AI message 1'),
|
|
321
321
|
new AIMessage('AI message 2')
|
|
322
322
|
];
|
|
323
|
-
|
|
323
|
+
|
|
324
324
|
const indexTokenCountMap = {
|
|
325
325
|
0: 17, // "System instruction"
|
|
326
326
|
1: 12, // "AI message 1"
|
|
327
327
|
2: 12 // "AI message 2"
|
|
328
328
|
};
|
|
329
|
-
|
|
329
|
+
|
|
330
330
|
// Set a limit that can fit all messages
|
|
331
331
|
const result = getMessagesWithinTokenLimit({
|
|
332
332
|
messages,
|
|
333
333
|
maxContextTokens: 100,
|
|
334
334
|
indexTokenCountMap,
|
|
335
|
-
|
|
335
|
+
startType: 'human'
|
|
336
336
|
});
|
|
337
|
-
|
|
337
|
+
|
|
338
338
|
// Should include all messages since no human messages exist to start from
|
|
339
339
|
expect(result.context.length).toBe(3);
|
|
340
340
|
expect(result.context[0]).toBe(messages[0]); // System message
|
|
341
341
|
expect(result.context[1]).toBe(messages[1]); // AI message 1
|
|
342
342
|
expect(result.context[2]).toBe(messages[2]); // AI message 2
|
|
343
|
-
|
|
343
|
+
|
|
344
344
|
expect(result.messagesToRefine.length).toBe(0);
|
|
345
345
|
});
|
|
346
346
|
});
|
|
347
|
-
|
|
347
|
+
|
|
348
348
|
describe('checkValidNumber', () => {
|
|
349
349
|
it('should return true for valid positive numbers', () => {
|
|
350
350
|
expect(checkValidNumber(5)).toBe(true);
|
|
351
351
|
expect(checkValidNumber(1.5)).toBe(true);
|
|
352
352
|
expect(checkValidNumber(Number.MAX_SAFE_INTEGER)).toBe(true);
|
|
353
353
|
});
|
|
354
|
-
|
|
354
|
+
|
|
355
355
|
it('should return false for zero, negative numbers, and NaN', () => {
|
|
356
356
|
expect(checkValidNumber(0)).toBe(false);
|
|
357
357
|
expect(checkValidNumber(-5)).toBe(false);
|
|
358
358
|
expect(checkValidNumber(NaN)).toBe(false);
|
|
359
359
|
});
|
|
360
|
-
|
|
360
|
+
|
|
361
361
|
it('should return false for non-number types', () => {
|
|
362
362
|
expect(checkValidNumber('5')).toBe(false);
|
|
363
363
|
expect(checkValidNumber(null)).toBe(false);
|
|
@@ -366,7 +366,7 @@ describe('Prune Messages Tests', () => {
|
|
|
366
366
|
expect(checkValidNumber([])).toBe(false);
|
|
367
367
|
});
|
|
368
368
|
});
|
|
369
|
-
|
|
369
|
+
|
|
370
370
|
describe('createPruneMessages', () => {
|
|
371
371
|
it('should return all messages when under token limit', () => {
|
|
372
372
|
const tokenCounter = createTestTokenCounter();
|
|
@@ -375,26 +375,26 @@ describe('Prune Messages Tests', () => {
|
|
|
375
375
|
new HumanMessage('Hello'),
|
|
376
376
|
new AIMessage('Hi there')
|
|
377
377
|
];
|
|
378
|
-
|
|
378
|
+
|
|
379
379
|
const indexTokenCountMap = {
|
|
380
380
|
0: tokenCounter(messages[0]),
|
|
381
381
|
1: tokenCounter(messages[1]),
|
|
382
382
|
2: tokenCounter(messages[2])
|
|
383
383
|
};
|
|
384
|
-
|
|
384
|
+
|
|
385
385
|
const pruneMessages = createPruneMessages({
|
|
386
386
|
maxTokens: 100,
|
|
387
387
|
startIndex: 0,
|
|
388
388
|
tokenCounter,
|
|
389
389
|
indexTokenCountMap
|
|
390
390
|
});
|
|
391
|
-
|
|
391
|
+
|
|
392
392
|
const result = pruneMessages({ messages });
|
|
393
|
-
|
|
393
|
+
|
|
394
394
|
expect(result.context.length).toBe(3);
|
|
395
395
|
expect(result.context).toEqual(messages);
|
|
396
396
|
});
|
|
397
|
-
|
|
397
|
+
|
|
398
398
|
it('should prune messages when over token limit', () => {
|
|
399
399
|
const tokenCounter = createTestTokenCounter();
|
|
400
400
|
const messages = [
|
|
@@ -404,7 +404,7 @@ describe('Prune Messages Tests', () => {
|
|
|
404
404
|
new HumanMessage('Message 2'),
|
|
405
405
|
new AIMessage('Response 2')
|
|
406
406
|
];
|
|
407
|
-
|
|
407
|
+
|
|
408
408
|
const indexTokenCountMap = {
|
|
409
409
|
0: tokenCounter(messages[0]),
|
|
410
410
|
1: tokenCounter(messages[1]),
|
|
@@ -412,7 +412,7 @@ describe('Prune Messages Tests', () => {
|
|
|
412
412
|
3: tokenCounter(messages[3]),
|
|
413
413
|
4: tokenCounter(messages[4])
|
|
414
414
|
};
|
|
415
|
-
|
|
415
|
+
|
|
416
416
|
// Set a limit that can only fit the system message and the last two messages
|
|
417
417
|
const pruneMessages = createPruneMessages({
|
|
418
418
|
maxTokens: 40,
|
|
@@ -420,9 +420,9 @@ describe('Prune Messages Tests', () => {
|
|
|
420
420
|
tokenCounter,
|
|
421
421
|
indexTokenCountMap
|
|
422
422
|
});
|
|
423
|
-
|
|
423
|
+
|
|
424
424
|
const result = pruneMessages({ messages });
|
|
425
|
-
|
|
425
|
+
|
|
426
426
|
// Should include system message and the last two messages
|
|
427
427
|
expect(result.context.length).toBe(3);
|
|
428
428
|
expect(result.context[0]).toBe(messages[0]); // System message
|
|
@@ -430,7 +430,7 @@ describe('Prune Messages Tests', () => {
|
|
|
430
430
|
expect(result.context[2]).toBe(messages[4]); // Response 2
|
|
431
431
|
});
|
|
432
432
|
|
|
433
|
-
it('should respect
|
|
433
|
+
it('should respect startType parameter', () => {
|
|
434
434
|
const tokenCounter = createTestTokenCounter();
|
|
435
435
|
const messages = [
|
|
436
436
|
new SystemMessage('System instruction'),
|
|
@@ -439,7 +439,7 @@ describe('Prune Messages Tests', () => {
|
|
|
439
439
|
new AIMessage('AI message 2'),
|
|
440
440
|
new HumanMessage('Human message 2')
|
|
441
441
|
];
|
|
442
|
-
|
|
442
|
+
|
|
443
443
|
const indexTokenCountMap = {
|
|
444
444
|
0: tokenCounter(messages[0]),
|
|
445
445
|
1: tokenCounter(messages[1]),
|
|
@@ -447,7 +447,7 @@ describe('Prune Messages Tests', () => {
|
|
|
447
447
|
3: tokenCounter(messages[3]),
|
|
448
448
|
4: tokenCounter(messages[4])
|
|
449
449
|
};
|
|
450
|
-
|
|
450
|
+
|
|
451
451
|
// Set a limit that can fit all messages
|
|
452
452
|
const pruneMessages = createPruneMessages({
|
|
453
453
|
maxTokens: 100,
|
|
@@ -455,12 +455,12 @@ describe('Prune Messages Tests', () => {
|
|
|
455
455
|
tokenCounter,
|
|
456
456
|
indexTokenCountMap: { ...indexTokenCountMap }
|
|
457
457
|
});
|
|
458
|
-
|
|
459
|
-
const result = pruneMessages({
|
|
458
|
+
|
|
459
|
+
const result = pruneMessages({
|
|
460
460
|
messages,
|
|
461
|
-
|
|
461
|
+
startType: 'human'
|
|
462
462
|
});
|
|
463
|
-
|
|
463
|
+
|
|
464
464
|
// All messages should be included since we're under the token limit
|
|
465
465
|
expect(result.context.length).toBe(5);
|
|
466
466
|
expect(result.context[0]).toBe(messages[0]); // System message
|
|
@@ -469,7 +469,7 @@ describe('Prune Messages Tests', () => {
|
|
|
469
469
|
expect(result.context[3]).toBe(messages[3]); // AI message 2
|
|
470
470
|
expect(result.context[4]).toBe(messages[4]); // Human message 2
|
|
471
471
|
});
|
|
472
|
-
|
|
472
|
+
|
|
473
473
|
it('should update token counts when usage metadata is provided', () => {
|
|
474
474
|
const tokenCounter = createTestTokenCounter();
|
|
475
475
|
const messages = [
|
|
@@ -477,47 +477,47 @@ describe('Prune Messages Tests', () => {
|
|
|
477
477
|
new HumanMessage('Hello'),
|
|
478
478
|
new AIMessage('Hi there')
|
|
479
479
|
];
|
|
480
|
-
|
|
480
|
+
|
|
481
481
|
const indexTokenCountMap = {
|
|
482
482
|
0: tokenCounter(messages[0]),
|
|
483
483
|
1: tokenCounter(messages[1]),
|
|
484
484
|
2: tokenCounter(messages[2])
|
|
485
485
|
};
|
|
486
|
-
|
|
486
|
+
|
|
487
487
|
const pruneMessages = createPruneMessages({
|
|
488
488
|
maxTokens: 100,
|
|
489
489
|
startIndex: 0,
|
|
490
490
|
tokenCounter,
|
|
491
491
|
indexTokenCountMap: { ...indexTokenCountMap }
|
|
492
492
|
});
|
|
493
|
-
|
|
493
|
+
|
|
494
494
|
// Provide usage metadata that indicates different token counts
|
|
495
495
|
const usageMetadata: Partial<UsageMetadata> = {
|
|
496
496
|
input_tokens: 50,
|
|
497
497
|
output_tokens: 25,
|
|
498
498
|
total_tokens: 75
|
|
499
499
|
};
|
|
500
|
-
|
|
501
|
-
const result = pruneMessages({
|
|
500
|
+
|
|
501
|
+
const result = pruneMessages({
|
|
502
502
|
messages,
|
|
503
503
|
usageMetadata
|
|
504
504
|
});
|
|
505
|
-
|
|
505
|
+
|
|
506
506
|
// The function should have updated the indexTokenCountMap based on the usage metadata
|
|
507
507
|
expect(result.indexTokenCountMap).not.toEqual(indexTokenCountMap);
|
|
508
|
-
|
|
508
|
+
|
|
509
509
|
// The total of all values in indexTokenCountMap should equal the total_tokens from usageMetadata
|
|
510
510
|
const totalTokens = Object.values(result.indexTokenCountMap).reduce((a, b) => a + b, 0);
|
|
511
511
|
expect(totalTokens).toBe(75);
|
|
512
512
|
});
|
|
513
513
|
});
|
|
514
|
-
|
|
514
|
+
|
|
515
515
|
describe('Integration with Run', () => {
|
|
516
516
|
it('should initialize Run with custom token counter and process messages', async () => {
|
|
517
517
|
const provider = Providers.OPENAI;
|
|
518
518
|
const llmConfig = getLLMConfig(provider);
|
|
519
519
|
const tokenCounter = createTestTokenCounter();
|
|
520
|
-
|
|
520
|
+
|
|
521
521
|
const run = await Run.create<t.IState>({
|
|
522
522
|
runId: 'test-prune-run',
|
|
523
523
|
graphConfig: {
|
|
@@ -527,18 +527,18 @@ describe('Prune Messages Tests', () => {
|
|
|
527
527
|
},
|
|
528
528
|
returnContent: true,
|
|
529
529
|
});
|
|
530
|
-
|
|
530
|
+
|
|
531
531
|
// Override the model to use a fake LLM
|
|
532
532
|
run.Graph?.overrideTestModel(['This is a test response'], 1);
|
|
533
|
-
|
|
533
|
+
|
|
534
534
|
const messages = [
|
|
535
535
|
new HumanMessage('Hello, how are you?')
|
|
536
536
|
];
|
|
537
|
-
|
|
537
|
+
|
|
538
538
|
const indexTokenCountMap = {
|
|
539
539
|
0: tokenCounter(messages[0])
|
|
540
540
|
};
|
|
541
|
-
|
|
541
|
+
|
|
542
542
|
const config: Partial<RunnableConfig> & { version: 'v1' | 'v2'; streamMode: string } = {
|
|
543
543
|
configurable: {
|
|
544
544
|
thread_id: 'test-thread',
|
|
@@ -546,7 +546,7 @@ describe('Prune Messages Tests', () => {
|
|
|
546
546
|
streamMode: 'values',
|
|
547
547
|
version: 'v2' as const,
|
|
548
548
|
};
|
|
549
|
-
|
|
549
|
+
|
|
550
550
|
await run.processStream(
|
|
551
551
|
{ messages },
|
|
552
552
|
config,
|
|
@@ -556,7 +556,7 @@ describe('Prune Messages Tests', () => {
|
|
|
556
556
|
tokenCounter,
|
|
557
557
|
}
|
|
558
558
|
);
|
|
559
|
-
|
|
559
|
+
|
|
560
560
|
const finalMessages = run.getRunMessages();
|
|
561
561
|
expect(finalMessages).toBeDefined();
|
|
562
562
|
expect(finalMessages?.length).toBeGreaterThan(0);
|