@lobehub/lobehub 2.0.0-next.40 → 2.0.0-next.41

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.
@@ -1,2632 +0,0 @@
1
- import { DBMessageItem } from '@lobechat/types';
2
- import dayjs from 'dayjs';
3
- import { eq } from 'drizzle-orm';
4
- import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
5
-
6
- import { uuid } from '@/utils/uuid';
7
-
8
- import { getTestDB } from '../../models/__tests__/_util';
9
- import {
10
- agents,
11
- chatGroups,
12
- chunks,
13
- embeddings,
14
- fileChunks,
15
- files,
16
- messagePlugins,
17
- messageQueries,
18
- messageQueryChunks,
19
- messageTTS,
20
- messageTranslates,
21
- messages,
22
- messagesFiles,
23
- sessions,
24
- topics,
25
- users,
26
- } from '../../schemas';
27
- import { LobeChatDatabase } from '../../type';
28
- import { MessageModel } from '../message';
29
- import { codeEmbedding } from './fixtures/embedding';
30
-
31
- const serverDB: LobeChatDatabase = await getTestDB();
32
-
33
- const userId = 'message-db';
34
- const messageModel = new MessageModel(serverDB, userId);
35
- const embeddingsId = uuid();
36
-
37
- beforeEach(async () => {
38
- // Clear tables before each test case
39
- await serverDB.transaction(async (trx) => {
40
- await trx.delete(users);
41
- await trx.insert(users).values([{ id: userId }, { id: '456' }]);
42
-
43
- await trx.insert(sessions).values([
44
- // { id: 'session1', userId },
45
- // { id: 'session2', userId },
46
- { id: '1', userId },
47
- ]);
48
- await trx.insert(files).values({
49
- id: 'f1',
50
- userId: userId,
51
- url: 'abc',
52
- name: 'file-1',
53
- fileType: 'image/png',
54
- size: 1000,
55
- });
56
-
57
- await trx.insert(embeddings).values({
58
- id: embeddingsId,
59
- embeddings: codeEmbedding,
60
- userId,
61
- });
62
- });
63
- });
64
-
65
- afterEach(async () => {
66
- // Clear tables after each test case
67
- await serverDB.delete(users);
68
- });
69
-
70
- describe('MessageModel', () => {
71
- describe('query', () => {
72
- it('should query messages by user ID', async () => {
73
- // Create test data
74
- await serverDB.insert(messages).values([
75
- { id: '1', userId, role: 'user', content: 'message 1', createdAt: new Date('2023-01-01') },
76
- { id: '2', userId, role: 'user', content: 'message 2', createdAt: new Date('2023-02-01') },
77
- {
78
- id: '3',
79
- userId: '456',
80
- role: 'user',
81
- content: 'message 3',
82
- createdAt: new Date('2023-03-01'),
83
- },
84
- ]);
85
-
86
- // Call query method
87
- const result = await messageModel.query();
88
-
89
- // Assert result
90
- expect(result).toHaveLength(2);
91
- expect(result[0].id).toBe('1');
92
- expect(result[1].id).toBe('2');
93
- });
94
-
95
- it('should return empty messages if not match the user ID', async () => {
96
- // Create test data
97
- await serverDB.insert(messages).values([
98
- { id: '1', userId: '456', role: 'user', content: '1', createdAt: new Date('2023-01-01') },
99
- { id: '2', userId: '456', role: 'user', content: '2', createdAt: new Date('2023-02-01') },
100
- { id: '3', userId: '456', role: 'user', content: '3', createdAt: new Date('2023-03-01') },
101
- ]);
102
-
103
- // Call query method
104
- const result = await messageModel.query();
105
-
106
- // Assert result
107
- expect(result).toHaveLength(0);
108
- });
109
-
110
- it('should query messages with pagination', async () => {
111
- // Create test data
112
- await serverDB.insert(messages).values([
113
- { id: '1', userId, role: 'user', content: 'message 1', createdAt: new Date('2023-01-01') },
114
- { id: '2', userId, role: 'user', content: 'message 2', createdAt: new Date('2023-02-01') },
115
- { id: '3', userId, role: 'user', content: 'message 3', createdAt: new Date('2023-03-01') },
116
- ]);
117
-
118
- // Test pagination
119
- const result1 = await messageModel.query({ current: 0, pageSize: 2 });
120
- expect(result1).toHaveLength(2);
121
-
122
- const result2 = await messageModel.query({ current: 1, pageSize: 1 });
123
- expect(result2).toHaveLength(1);
124
- expect(result2[0].id).toBe('2');
125
- });
126
-
127
- it('should filter messages by sessionId', async () => {
128
- // Create test data
129
- await serverDB.insert(sessions).values([
130
- { id: 'session1', userId },
131
- { id: 'session2', userId },
132
- ]);
133
- await serverDB.insert(messages).values([
134
- {
135
- id: '1',
136
- userId,
137
- role: 'user',
138
- sessionId: 'session1',
139
- content: 'message 1',
140
- createdAt: new Date('2022-02-01'),
141
- },
142
- {
143
- id: '2',
144
- userId,
145
- role: 'user',
146
- sessionId: 'session1',
147
- content: 'message 2',
148
- createdAt: new Date('2023-02-02'),
149
- },
150
- { id: '3', userId, role: 'user', sessionId: 'session2', content: 'message 3' },
151
- ]);
152
-
153
- // Test filtering by sessionId
154
- const result = await messageModel.query({ sessionId: 'session1' });
155
- expect(result).toHaveLength(2);
156
- expect(result[0].id).toBe('1');
157
- expect(result[1].id).toBe('2');
158
- });
159
-
160
- it('should filter messages by topicId', async () => {
161
- // Create test data
162
- const sessionId = 'session1';
163
- await serverDB.insert(sessions).values([{ id: sessionId, userId }]);
164
- const topicId = 'topic1';
165
- await serverDB.insert(topics).values([
166
- { id: topicId, sessionId, userId },
167
- { id: 'topic2', sessionId, userId },
168
- ]);
169
-
170
- await serverDB.insert(messages).values([
171
- { id: '1', userId, role: 'user', topicId, content: '1', createdAt: new Date('2022-04-01') },
172
- { id: '2', userId, role: 'user', topicId, content: '2', createdAt: new Date('2023-02-01') },
173
- { id: '3', userId, role: 'user', topicId: 'topic2', content: 'message 3' },
174
- ]);
175
-
176
- // Test filtering by topicId
177
- const result = await messageModel.query({ topicId });
178
- expect(result).toHaveLength(2);
179
- expect(result[0].id).toBe('1');
180
- expect(result[1].id).toBe('2');
181
- });
182
-
183
- it('should filter messages by groupId and expose group metadata', async () => {
184
- await serverDB.transaction(async (trx) => {
185
- await trx.insert(chatGroups).values([
186
- { id: 'group-1', userId, title: 'Group 1' },
187
- { id: 'group-2', userId, title: 'Group 2' },
188
- ]);
189
-
190
- await trx.insert(agents).values([
191
- { id: 'agent-group', userId, title: 'Agent Group' },
192
- { id: 'agent-other', userId, title: 'Agent Other' },
193
- ]);
194
-
195
- await trx.insert(messages).values([
196
- {
197
- id: 'group-message',
198
- userId,
199
- role: 'assistant',
200
- content: 'group message',
201
- groupId: 'group-1',
202
- agentId: 'agent-group',
203
- targetId: 'user',
204
- createdAt: new Date('2024-01-01'),
205
- },
206
- {
207
- id: 'other-message',
208
- userId,
209
- role: 'assistant',
210
- content: 'other group message',
211
- groupId: 'group-2',
212
- agentId: 'agent-other',
213
- targetId: 'user',
214
- createdAt: new Date('2024-01-02'),
215
- },
216
- ]);
217
- });
218
-
219
- const result = await messageModel.query({ groupId: 'group-1' });
220
-
221
- expect(result).toHaveLength(1);
222
- expect(result[0].id).toBe('group-message');
223
- expect(result[0].groupId).toBe('group-1');
224
- expect(result[0].agentId).toBe('agent-group');
225
- expect(result[0].targetId).toBe('user');
226
- });
227
-
228
- it('should query messages with join', async () => {
229
- // Create test data
230
- await serverDB.transaction(async (trx) => {
231
- await trx.insert(messages).values([
232
- {
233
- id: '1',
234
- userId,
235
- role: 'user',
236
- content: 'message 1',
237
- createdAt: new Date('2023-01-01'),
238
- },
239
- {
240
- id: '2',
241
- userId,
242
- role: 'user',
243
- content: 'message 2',
244
- createdAt: new Date('2023-02-01'),
245
- },
246
- {
247
- id: '3',
248
- userId: '456',
249
- role: 'user',
250
- content: 'message 3',
251
- createdAt: new Date('2023-03-01'),
252
- },
253
- ]);
254
- await trx.insert(files).values([
255
- { id: 'f-0', url: 'abc', name: 'file-1', userId, fileType: 'image/png', size: 1000 },
256
- { id: 'f-1', url: 'abc', name: 'file-1', userId, fileType: 'image/png', size: 100 },
257
- { id: 'f-3', url: 'abc', name: 'file-3', userId, fileType: 'image/png', size: 400 },
258
- ]);
259
- await trx.insert(messageTTS).values([
260
- { id: '1', userId },
261
- { id: '2', voice: 'a', fileId: 'f-1', contentMd5: 'abc', userId },
262
- ]);
263
-
264
- await trx.insert(messagesFiles).values([
265
- { fileId: 'f-0', messageId: '1', userId },
266
- { fileId: 'f-3', messageId: '1', userId },
267
- ]);
268
- });
269
-
270
- const domain = 'http://abc.com';
271
- // Call query method
272
- const result = await messageModel.query(
273
- {},
274
- { postProcessUrl: async (path) => `${domain}/${path}` },
275
- );
276
-
277
- // Assert result
278
- expect(result).toHaveLength(2);
279
- expect(result[0].id).toBe('1');
280
- expect(result[0].imageList).toEqual([
281
- { alt: 'file-1', id: 'f-0', url: `${domain}/abc` },
282
- { alt: 'file-3', id: 'f-3', url: `${domain}/abc` },
283
- ]);
284
-
285
- expect(result[1].id).toBe('2');
286
- expect(result[1].imageList).toEqual([]);
287
- });
288
-
289
- it('should include translate, tts and other extra fields in query result', async () => {
290
- // Create test data
291
- await serverDB.transaction(async (trx) => {
292
- await trx.insert(messages).values([
293
- {
294
- id: '1',
295
- userId,
296
- role: 'user',
297
- content: 'message 1',
298
- createdAt: new Date('2023-01-01'),
299
- },
300
- ]);
301
- await trx
302
- .insert(messageTranslates)
303
- .values([{ id: '1', content: 'translated', from: 'en', to: 'zh', userId }]);
304
- await trx
305
- .insert(messageTTS)
306
- .values([{ id: '1', voice: 'voice1', fileId: 'f1', contentMd5: 'md5', userId }]);
307
- });
308
-
309
- // Call query method
310
- const result = await messageModel.query();
311
-
312
- // Assert result
313
- expect(result[0].extra!.translate).toEqual({ content: 'translated', from: 'en', to: 'zh' });
314
- expect(result[0].extra!.tts).toEqual({
315
- contentMd5: 'md5',
316
- file: 'f1',
317
- voice: 'voice1',
318
- });
319
- });
320
-
321
- it('should handle edge cases of pagination parameters', async () => {
322
- // Create test data
323
- await serverDB.insert(messages).values([
324
- { id: '1', userId, role: 'user', content: 'message 1' },
325
- { id: '2', userId, role: 'user', content: 'message 2' },
326
- { id: '3', userId, role: 'user', content: 'message 3' },
327
- ]);
328
-
329
- // 测试 current 和 pageSize 的边界情况
330
- const result1 = await messageModel.query({ current: 0, pageSize: 2 });
331
- expect(result1).toHaveLength(2);
332
-
333
- const result2 = await messageModel.query({ current: 1, pageSize: 2 });
334
- expect(result2).toHaveLength(1);
335
-
336
- const result3 = await messageModel.query({ current: 2, pageSize: 2 });
337
- expect(result3).toHaveLength(0);
338
- });
339
-
340
- describe('query with messageQueries', () => {
341
- it('should include ragQuery, ragQueryId and ragRawQuery in query results', async () => {
342
- // Create test data
343
- const messageId = 'msg-with-query';
344
- const queryId = uuid();
345
-
346
- await serverDB.insert(messages).values({
347
- id: messageId,
348
- userId,
349
- role: 'user',
350
- content: 'test message',
351
- });
352
-
353
- await serverDB.insert(messageQueries).values({
354
- id: queryId,
355
- messageId,
356
- userQuery: 'original query',
357
- rewriteQuery: 'rewritten query',
358
- userId,
359
- });
360
-
361
- // Call query method
362
- const result = await messageModel.query();
363
-
364
- // Assert result
365
- expect(result).toHaveLength(1);
366
- expect(result[0].id).toBe(messageId);
367
- expect(result[0].ragQueryId).toBe(queryId);
368
- expect(result[0].ragQuery).toBe('rewritten query');
369
- expect(result[0].ragRawQuery).toBe('original query');
370
- });
371
-
372
- it.skip('should handle multiple message queries for the same message', async () => {
373
- // Create test data
374
- const messageId = 'msg-multi-query';
375
- const queryId1 = uuid();
376
- const queryId2 = uuid();
377
-
378
- await serverDB.insert(messages).values({
379
- id: messageId,
380
- userId,
381
- role: 'user',
382
- content: 'test message',
383
- });
384
-
385
- // 创建两个查询,但查询结果应该只包含一个(最新的)
386
- await serverDB.insert(messageQueries).values([
387
- {
388
- id: queryId1,
389
- messageId,
390
- userQuery: 'original query 1',
391
- rewriteQuery: 'rewritten query 1',
392
- userId,
393
- },
394
- {
395
- id: queryId2,
396
- messageId,
397
- userQuery: 'original query 2',
398
- rewriteQuery: 'rewritten query 2',
399
- userId,
400
- },
401
- ]);
402
-
403
- // Call query method
404
- const result = await messageModel.query();
405
-
406
- // Assert result - 应该只包含最新的查询
407
- expect(result).toHaveLength(1);
408
- expect(result[0].id).toBe(messageId);
409
- expect(result[0].ragQueryId).toBe(queryId2);
410
- expect(result[0].ragQuery).toBe('rewritten query 2');
411
- expect(result[0].ragRawQuery).toBe('original query 2');
412
- });
413
- });
414
-
415
- it('should handle complex query with multiple joins and file chunks', async () => {
416
- await serverDB.transaction(async (trx) => {
417
- const chunk1Id = uuid();
418
- const query1Id = uuid();
419
- // 创建基础消息
420
- await trx.insert(messages).values({
421
- id: 'msg1',
422
- userId,
423
- role: 'user',
424
- content: 'test message',
425
- createdAt: new Date('2023-01-01'),
426
- });
427
-
428
- // 创建文件
429
- await trx.insert(files).values([
430
- {
431
- id: 'file1',
432
- userId,
433
- name: 'test.txt',
434
- url: 'test-url',
435
- fileType: 'text/plain',
436
- size: 100,
437
- },
438
- ]);
439
-
440
- // 创建文件块
441
- await trx.insert(chunks).values({
442
- id: chunk1Id,
443
- text: 'chunk content',
444
- });
445
-
446
- // 关联消息和文件
447
- await trx.insert(messagesFiles).values({
448
- messageId: 'msg1',
449
- userId,
450
- fileId: 'file1',
451
- });
452
-
453
- // 创建文件块关联
454
- await trx.insert(fileChunks).values({
455
- fileId: 'file1',
456
- userId,
457
- chunkId: chunk1Id,
458
- });
459
-
460
- // 创建消息查询
461
- await trx.insert(messageQueries).values({
462
- id: query1Id,
463
- messageId: 'msg1',
464
- userId,
465
- userQuery: 'original query',
466
- rewriteQuery: 'rewritten query',
467
- });
468
-
469
- // 创建消息查询块关联
470
- await trx.insert(messageQueryChunks).values({
471
- messageId: 'msg1',
472
- queryId: query1Id,
473
- chunkId: chunk1Id,
474
- similarity: '0.95',
475
- userId,
476
- });
477
- });
478
-
479
- const result = await messageModel.query();
480
-
481
- expect(result).toHaveLength(1);
482
- expect(result[0].chunksList).toHaveLength(1);
483
- expect(result[0].chunksList![0]).toMatchObject({
484
- text: 'chunk content',
485
- similarity: 0.95,
486
- });
487
- });
488
-
489
- it('should return empty arrays for files and chunks if none exist', async () => {
490
- await serverDB.insert(messages).values({
491
- id: 'msg1',
492
- userId,
493
- role: 'user',
494
- content: 'test message',
495
- });
496
-
497
- const result = await messageModel.query();
498
-
499
- expect(result).toHaveLength(1);
500
- expect(result[0].fileList).toEqual([]);
501
- expect(result[0].imageList).toEqual([]);
502
- expect(result[0].chunksList).toEqual([]);
503
- });
504
- });
505
-
506
- describe('queryAll', () => {
507
- it('should return all messages belonging to the user in ascending order', async () => {
508
- // Create test data
509
- await serverDB.insert(messages).values([
510
- {
511
- id: '1',
512
- userId,
513
- role: 'user',
514
- content: 'message 1',
515
- createdAt: new Date('2023-01-01'),
516
- },
517
- {
518
- id: '2',
519
- userId,
520
- role: 'user',
521
- content: 'message 2',
522
- createdAt: new Date('2023-02-01'),
523
- },
524
- {
525
- id: '3',
526
- userId: '456',
527
- role: 'user',
528
- content: 'message 3',
529
- createdAt: new Date('2023-03-01'),
530
- },
531
- ]);
532
-
533
- // Call queryAll method
534
- const result = await messageModel.queryAll();
535
-
536
- // Assert result
537
- expect(result).toHaveLength(2);
538
- expect(result[0].id).toBe('1');
539
- expect(result[1].id).toBe('2');
540
- });
541
- });
542
-
543
- describe('findById', () => {
544
- it('should find message by ID', async () => {
545
- // Create test data
546
- await serverDB.insert(messages).values([
547
- { id: '1', userId, role: 'user', content: 'message 1' },
548
- { id: '2', userId: '456', role: 'user', content: 'message 2' },
549
- ]);
550
-
551
- // Call findById method
552
- const result = await messageModel.findById('1');
553
-
554
- // Assert result
555
- expect(result?.id).toBe('1');
556
- expect(result?.content).toBe('message 1');
557
- });
558
-
559
- it('should return undefined if message does not belong to user', async () => {
560
- // Create test data
561
- await serverDB
562
- .insert(messages)
563
- .values([{ id: '1', userId: '456', role: 'user', content: 'message 1' }]);
564
-
565
- // Call findById method
566
- const result = await messageModel.findById('1');
567
-
568
- // Assert result
569
- expect(result).toBeUndefined();
570
- });
571
- });
572
-
573
- describe('queryBySessionId', () => {
574
- it('should query messages by sessionId', async () => {
575
- // Create test data
576
- const sessionId = 'session1';
577
- await serverDB.insert(sessions).values([
578
- { id: 'session1', userId },
579
- { id: 'session2', userId },
580
- ]);
581
- await serverDB.insert(messages).values([
582
- {
583
- id: '1',
584
- userId,
585
- role: 'user',
586
- sessionId,
587
- content: 'message 1',
588
- createdAt: new Date('2022-01-01'),
589
- },
590
- {
591
- id: '2',
592
- userId,
593
- role: 'user',
594
- sessionId,
595
- content: 'message 2',
596
- createdAt: new Date('2023-02-01'),
597
- },
598
- { id: '3', userId, role: 'user', sessionId: 'session2', content: 'message 3' },
599
- ]);
600
-
601
- // Call queryBySessionId method
602
- const result = await messageModel.queryBySessionId(sessionId);
603
-
604
- // Assert result
605
- expect(result).toHaveLength(2);
606
- expect(result[0].id).toBe('1');
607
- expect(result[1].id).toBe('2');
608
- });
609
- });
610
-
611
- describe('queryByKeyWord', () => {
612
- it('should query messages by keyword', async () => {
613
- // Create test data
614
- await serverDB.insert(messages).values([
615
- { id: '1', userId, role: 'user', content: 'apple', createdAt: new Date('2022-02-01') },
616
- { id: '2', userId, role: 'user', content: 'banana' },
617
- { id: '3', userId, role: 'user', content: 'pear' },
618
- { id: '4', userId, role: 'user', content: 'apple pie', createdAt: new Date('2024-02-01') },
619
- ]);
620
-
621
- // Test querying messages with specific keyword
622
- const result = await messageModel.queryByKeyword('apple');
623
-
624
- // Assert result
625
- expect(result).toHaveLength(2);
626
- expect(result[0].id).toBe('4');
627
- expect(result[1].id).toBe('1');
628
- });
629
-
630
- it('should return empty array when keyword is empty', async () => {
631
- // Create test data
632
- await serverDB.insert(messages).values([
633
- { id: '1', userId, role: 'user', content: 'apple' },
634
- { id: '2', userId, role: 'user', content: 'banana' },
635
- { id: '3', userId, role: 'user', content: 'pear' },
636
- { id: '4', userId, role: 'user', content: 'apple pie' },
637
- ]);
638
-
639
- // Test returning empty array when keyword is empty
640
- const result = await messageModel.queryByKeyword('');
641
-
642
- // Assert result
643
- expect(result).toHaveLength(0);
644
- });
645
- });
646
-
647
- describe('createMessage', () => {
648
- it('should create a new message', async () => {
649
- // Call createMessage method
650
- await messageModel.create({ role: 'user', content: 'new message', sessionId: '1' });
651
-
652
- // Assert result
653
- const result = await serverDB.select().from(messages).where(eq(messages.userId, userId));
654
- expect(result).toHaveLength(1);
655
- expect(result[0].content).toBe('new message');
656
- });
657
-
658
- it('should create a message', async () => {
659
- const sessionId = 'session1';
660
- await serverDB.insert(sessions).values([{ id: sessionId, userId }]);
661
-
662
- const result = await messageModel.create({
663
- content: 'message 1',
664
- role: 'user',
665
- sessionId: 'session1',
666
- });
667
-
668
- expect(result.id).toBeDefined();
669
- expect(result.content).toBe('message 1');
670
- expect(result.role).toBe('user');
671
- expect(result.sessionId).toBe('session1');
672
- expect(result.userId).toBe(userId);
673
- });
674
-
675
- it('should generate message ID automatically', async () => {
676
- // Call createMessage method
677
- await messageModel.create({
678
- role: 'user',
679
- content: 'new message',
680
- sessionId: '1',
681
- });
682
-
683
- // Assert result
684
- const result = await serverDB.select().from(messages).where(eq(messages.userId, userId));
685
- expect(result[0].id).toBeDefined();
686
- expect(result[0].id).toHaveLength(18);
687
- });
688
-
689
- it('should create a tool message and insert into messagePlugins table', async () => {
690
- // Call create method
691
- const result = await messageModel.create({
692
- content: 'message 1',
693
- role: 'tool',
694
- sessionId: '1',
695
- tool_call_id: 'tool1',
696
- plugin: {
697
- apiName: 'api1',
698
- arguments: 'arg1',
699
- identifier: 'plugin1',
700
- type: 'default',
701
- },
702
- });
703
-
704
- // Assert result
705
- expect(result.id).toBeDefined();
706
- expect(result.content).toBe('message 1');
707
- expect(result.role).toBe('tool');
708
- expect(result.sessionId).toBe('1');
709
-
710
- const pluginResult = await serverDB
711
- .select()
712
- .from(messagePlugins)
713
- .where(eq(messagePlugins.id, result.id));
714
- expect(pluginResult).toHaveLength(1);
715
- expect(pluginResult[0].identifier).toBe('plugin1');
716
- });
717
-
718
- it('should create tool message ', async () => {
719
- // Call create method
720
- const state = {
721
- query: 'Composio',
722
- answers: [],
723
- results: [
724
- {
725
- url: 'https://www.composio.dev/',
726
- score: 16,
727
- title: 'Composio - Connect 90+ tools to your AI agents',
728
- engine: 'bing',
729
- content:
730
- 'Faster DevelopmentHigher ReliabilityBetter Integrations. Get Started Now. Our platform lets you ditch the specs and seamlessly integrate any tool you need in less than 5 mins.',
731
- engines: ['bing', 'qwant', 'brave', 'duckduckgo'],
732
- category: 'general',
733
- template: 'default.html',
734
- positions: [1, 1, 1, 1],
735
- thumbnail: '',
736
- parsed_url: ['https', 'www.composio.dev', '/', '', '', ''],
737
- publishedDate: null,
738
- },
739
- {
740
- url: 'https://www.composio.co/',
741
- score: 10.75,
742
- title: 'Composio',
743
- engine: 'bing',
744
- content:
745
- 'Composio was created to help streamline the entire book creation process! Writing. Take time out to write / Make a schedule to write consistently. We have writing software that optimizes your books for printing or ebook format. Figure out what you want to write. Collaborate and write with others. Professional editing is a necessity.',
746
- engines: ['qwant', 'duckduckgo', 'google', 'bing', 'brave'],
747
- category: 'general',
748
- template: 'default.html',
749
- positions: [5, 2, 1, 5, 4],
750
- thumbnail: null,
751
- parsed_url: ['https', 'www.composio.co', '/', '', '', ''],
752
- publishedDate: null,
753
- },
754
- ],
755
- unresponsive_engines: [],
756
- };
757
- const result = await messageModel.create({
758
- content: '[{}]',
759
- plugin: {
760
- apiName: 'searchWithSearXNG',
761
- arguments: '{\n "query": "Composio"\n}',
762
- identifier: 'lobe-web-browsing',
763
- type: 'builtin',
764
- },
765
- pluginState: state,
766
- role: 'tool',
767
- tool_call_id: 'tool_call_ymxXC2J0',
768
- sessionId: '1',
769
- });
770
-
771
- // Assert result
772
- expect(result.id).toBeDefined();
773
- expect(result.content).toBe('[{}]');
774
- expect(result.role).toBe('tool');
775
- expect(result.sessionId).toBe('1');
776
-
777
- const pluginResult = await serverDB
778
- .select()
779
- .from(messagePlugins)
780
- .where(eq(messagePlugins.id, result.id));
781
- expect(pluginResult).toHaveLength(1);
782
- expect(pluginResult[0].identifier).toBe('lobe-web-browsing');
783
- expect(pluginResult[0].state!).toMatchObject(state);
784
- });
785
-
786
- describe('create with advanced parameters', () => {
787
- it('should create a message with custom ID', async () => {
788
- const customId = 'custom-msg-id';
789
-
790
- const result = await messageModel.create(
791
- {
792
- role: 'user',
793
- content: 'message with custom ID',
794
- sessionId: '1',
795
- },
796
- customId,
797
- );
798
-
799
- expect(result.id).toBe(customId);
800
-
801
- // Verify database records
802
- const dbResult = await serverDB.select().from(messages).where(eq(messages.id, customId));
803
- expect(dbResult).toHaveLength(1);
804
- expect(dbResult[0].id).toBe(customId);
805
- });
806
-
807
- it.skip('should create a message with file chunks and RAG query ID', async () => {
808
- // Create test data
809
- const chunkId1 = uuid();
810
- const chunkId2 = uuid();
811
- const ragQueryId = uuid();
812
-
813
- await serverDB.insert(chunks).values([
814
- { id: chunkId1, text: 'chunk text 1' },
815
- { id: chunkId2, text: 'chunk text 2' },
816
- ]);
817
-
818
- // Call create method
819
- const result = await messageModel.create({
820
- role: 'assistant',
821
- content: 'message with file chunks',
822
- fileChunks: [
823
- { id: chunkId1, similarity: 0.95 },
824
- { id: chunkId2, similarity: 0.85 },
825
- ],
826
- ragQueryId,
827
- sessionId: '1',
828
- });
829
-
830
- // Verify message created successfully
831
- expect(result.id).toBeDefined();
832
-
833
- // Verify message query chunk associations created successfully
834
- const queryChunks = await serverDB
835
- .select()
836
- .from(messageQueryChunks)
837
- .where(eq(messageQueryChunks.messageId, result.id));
838
-
839
- expect(queryChunks).toHaveLength(2);
840
- expect(queryChunks[0].chunkId).toBe(chunkId1);
841
- expect(queryChunks[0].queryId).toBe(ragQueryId);
842
- expect(queryChunks[0].similarity).toBe('0.95');
843
- expect(queryChunks[1].chunkId).toBe(chunkId2);
844
- expect(queryChunks[1].similarity).toBe('0.85');
845
- });
846
-
847
- it('should create a message with files', async () => {
848
- // Create test data
849
- await serverDB.insert(files).values([
850
- {
851
- id: 'file1',
852
- name: 'file1.txt',
853
- fileType: 'text/plain',
854
- size: 100,
855
- url: 'url1',
856
- userId,
857
- },
858
- {
859
- id: 'file2',
860
- name: 'file2.jpg',
861
- fileType: 'image/jpeg',
862
- size: 200,
863
- url: 'url2',
864
- userId,
865
- },
866
- ]);
867
-
868
- // Call create method
869
- const result = await messageModel.create({
870
- role: 'user',
871
- content: 'message with files',
872
- files: ['file1', 'file2'],
873
- sessionId: '1',
874
- });
875
-
876
- // Verify message created successfully
877
- expect(result.id).toBeDefined();
878
-
879
- // Verify message file associations created successfully
880
- const messageFiles = await serverDB
881
- .select()
882
- .from(messagesFiles)
883
- .where(eq(messagesFiles.messageId, result.id));
884
-
885
- expect(messageFiles).toHaveLength(2);
886
- expect(messageFiles[0].fileId).toBe('file1');
887
- expect(messageFiles[1].fileId).toBe('file2');
888
- });
889
-
890
- it('should create a message with custom timestamps', async () => {
891
- const customCreatedAt = '2022-05-15T10:30:00Z';
892
- const customUpdatedAt = '2022-05-16T11:45:00Z';
893
-
894
- const result = await messageModel.create({
895
- role: 'user',
896
- content: 'message with custom timestamps',
897
- createdAt: customCreatedAt as any,
898
- updatedAt: customUpdatedAt as any,
899
- sessionId: '1',
900
- });
901
-
902
- // Verify database records
903
- const dbResult = await serverDB.select().from(messages).where(eq(messages.id, result.id));
904
-
905
- // Date comparison needs to consider timezone and formatting, so use toISOString for comparison
906
- expect(new Date(dbResult[0].createdAt!).toISOString()).toBe(
907
- new Date(customCreatedAt).toISOString(),
908
- );
909
- expect(new Date(dbResult[0].updatedAt!).toISOString()).toBe(
910
- new Date(customUpdatedAt).toISOString(),
911
- );
912
- });
913
- });
914
- });
915
-
916
- describe('batchCreateMessages', () => {
917
- it('should batch create messages', async () => {
918
- // Prepare test data
919
- const newMessages = [
920
- { id: '1', role: 'user', content: 'message 1' },
921
- { id: '2', role: 'assistant', content: 'message 2' },
922
- ] as DBMessageItem[];
923
-
924
- // Call batchCreateMessages method
925
- await messageModel.batchCreate(newMessages);
926
-
927
- // Assert result
928
- const result = await serverDB.select().from(messages).where(eq(messages.userId, userId));
929
- expect(result).toHaveLength(2);
930
- expect(result[0].content).toBe('message 1');
931
- expect(result[1].content).toBe('message 2');
932
- });
933
- });
934
-
935
- describe('updateMessage', () => {
936
- it('should update message content', async () => {
937
- // Create test data
938
- await serverDB
939
- .insert(messages)
940
- .values([{ id: '1', userId, role: 'user', content: 'message 1' }]);
941
-
942
- // Call updateMessage method
943
- await messageModel.update('1', { content: 'updated message' });
944
-
945
- // Assert result
946
- const result = await serverDB.select().from(messages).where(eq(messages.id, '1'));
947
- expect(result[0].content).toBe('updated message');
948
- });
949
-
950
- it('should only update messages belonging to the user', async () => {
951
- // Create test data
952
- await serverDB
953
- .insert(messages)
954
- .values([{ id: '1', userId: '456', role: 'user', content: 'message 1' }]);
955
-
956
- // Call updateMessage method
957
- await messageModel.update('1', { content: 'updated message' });
958
-
959
- // Assert result
960
- const result = await serverDB.select().from(messages).where(eq(messages.id, '1'));
961
- expect(result[0].content).toBe('message 1');
962
- });
963
-
964
- it('should update message tools', async () => {
965
- // Create test data
966
- await serverDB.insert(messages).values([
967
- {
968
- id: '1',
969
- userId,
970
- role: 'user',
971
- content: 'message 1',
972
- tools: [
973
- {
974
- id: 'call_Z8UU8LedZcoJHFGkfqYecjmT',
975
- type: 'builtin',
976
- apiName: 'searchWithSearXNG',
977
- arguments:
978
- '{"query":"杭州洪水 2023","searchEngines":["google","bing","baidu","duckduckgo","brave"]}',
979
- identifier: 'lobe-web-browsing',
980
- },
981
- ],
982
- },
983
- ]);
984
-
985
- // Call updateMessage method
986
- await messageModel.update('1', {
987
- tools: [
988
- {
989
- id: 'call_Z8UU8LedZcoJHFGkfqYecjmT',
990
- type: 'builtin',
991
- apiName: 'searchWithSearXNG',
992
- arguments: '{"query":"2024 杭州暴雨","searchEngines":["duckduckgo","google","brave"]}',
993
- identifier: 'lobe-web-browsing',
994
- },
995
- ],
996
- });
997
-
998
- // Assert result
999
- const result = await serverDB.select().from(messages).where(eq(messages.id, '1'));
1000
- expect((result[0].tools as any)[0].arguments).toBe(
1001
- '{"query":"2024 杭州暴雨","searchEngines":["duckduckgo","google","brave"]}',
1002
- );
1003
- });
1004
-
1005
- describe('update with imageList', () => {
1006
- it('should update a message and add image files', async () => {
1007
- // Create test data
1008
- await serverDB.insert(messages).values({
1009
- id: 'msg-to-update',
1010
- userId,
1011
- role: 'user',
1012
- content: 'original content',
1013
- });
1014
-
1015
- await serverDB.insert(files).values([
1016
- {
1017
- id: 'img1',
1018
- name: 'image1.jpg',
1019
- fileType: 'image/jpeg',
1020
- size: 100,
1021
- url: 'url1',
1022
- userId,
1023
- },
1024
- { id: 'img2', name: 'image2.png', fileType: 'image/png', size: 200, url: 'url2', userId },
1025
- ]);
1026
-
1027
- // Call update method
1028
- await messageModel.update('msg-to-update', {
1029
- content: 'updated content',
1030
- imageList: [
1031
- { id: 'img1', alt: 'image 1', url: 'url1' },
1032
- { id: 'img2', alt: 'image 2', url: 'url2' },
1033
- ],
1034
- });
1035
-
1036
- // Verify message updated successfully
1037
- const updatedMessage = await serverDB
1038
- .select()
1039
- .from(messages)
1040
- .where(eq(messages.id, 'msg-to-update'));
1041
-
1042
- expect(updatedMessage[0].content).toBe('updated content');
1043
-
1044
- // Verify message file associations created successfully
1045
- const messageFiles = await serverDB
1046
- .select()
1047
- .from(messagesFiles)
1048
- .where(eq(messagesFiles.messageId, 'msg-to-update'));
1049
-
1050
- expect(messageFiles).toHaveLength(2);
1051
- expect(messageFiles[0].fileId).toBe('img1');
1052
- expect(messageFiles[1].fileId).toBe('img2');
1053
- });
1054
-
1055
- it('should handle empty imageList', async () => {
1056
- // Create test data
1057
- await serverDB.insert(messages).values({
1058
- id: 'msg-no-images',
1059
- userId,
1060
- role: 'user',
1061
- content: 'original content',
1062
- });
1063
-
1064
- // Call update method without providing imageList
1065
- await messageModel.update('msg-no-images', {
1066
- content: 'updated content',
1067
- });
1068
-
1069
- // Verify message updated successfully
1070
- const updatedMessage = await serverDB
1071
- .select()
1072
- .from(messages)
1073
- .where(eq(messages.id, 'msg-no-images'));
1074
-
1075
- expect(updatedMessage[0].content).toBe('updated content');
1076
-
1077
- // Verify no message file associations created
1078
- const messageFiles = await serverDB
1079
- .select()
1080
- .from(messagesFiles)
1081
- .where(eq(messagesFiles.messageId, 'msg-no-images'));
1082
-
1083
- expect(messageFiles).toHaveLength(0);
1084
- });
1085
-
1086
- it('should update multiple fields at once', async () => {
1087
- // Create test data
1088
- await serverDB.insert(messages).values({
1089
- id: 'msg-multi-update',
1090
- userId,
1091
- role: 'user',
1092
- content: 'original content',
1093
- model: 'gpt-3.5',
1094
- });
1095
-
1096
- // Call update method to update multiple fields
1097
- await messageModel.update('msg-multi-update', {
1098
- content: 'updated content',
1099
- role: 'assistant',
1100
- model: 'gpt-4',
1101
- metadata: { tps: 1 },
1102
- });
1103
-
1104
- // Verify message updated successfully
1105
- const updatedMessage = await serverDB
1106
- .select()
1107
- .from(messages)
1108
- .where(eq(messages.id, 'msg-multi-update'));
1109
-
1110
- expect(updatedMessage[0].content).toBe('updated content');
1111
- expect(updatedMessage[0].role).toBe('assistant');
1112
- expect(updatedMessage[0].model).toBe('gpt-4');
1113
- expect(updatedMessage[0].metadata).toEqual({ tps: 1 });
1114
- });
1115
- });
1116
-
1117
- describe('update with returnQuery option', () => {
1118
- it('should return updated message list when sessionId is provided', async () => {
1119
- // Create test data
1120
- const sessionId = '1';
1121
- await serverDB.insert(messages).values([
1122
- {
1123
- id: 'msg1',
1124
- userId,
1125
- sessionId,
1126
- role: 'user',
1127
- content: 'message 1',
1128
- },
1129
- {
1130
- id: 'msg2',
1131
- userId,
1132
- sessionId,
1133
- role: 'assistant',
1134
- content: 'message 2',
1135
- },
1136
- ]);
1137
-
1138
- // Call update method with sessionId option
1139
- const result = await messageModel.update(
1140
- 'msg1',
1141
- { content: 'updated message 1' },
1142
- { sessionId },
1143
- );
1144
-
1145
- // Verify return result contains message list
1146
- expect(result.success).toBe(true);
1147
- expect(result.messages).toBeDefined();
1148
- expect(result.messages).toHaveLength(2);
1149
- expect(result.messages![0].content).toBe('updated message 1');
1150
- expect(result.messages![1].content).toBe('message 2');
1151
- });
1152
-
1153
- it('should return updated message list when topicId is provided', async () => {
1154
- // Create test data
1155
- const sessionId = '1';
1156
- const topicId = 'topic-1';
1157
- await serverDB.insert(topics).values({ id: topicId, sessionId, userId });
1158
- await serverDB.insert(messages).values([
1159
- {
1160
- id: 'msg-topic1',
1161
- userId,
1162
- sessionId,
1163
- topicId,
1164
- role: 'user',
1165
- content: 'topic message 1',
1166
- },
1167
- {
1168
- id: 'msg-topic2',
1169
- userId,
1170
- sessionId,
1171
- topicId,
1172
- role: 'assistant',
1173
- content: 'topic message 2',
1174
- },
1175
- ]);
1176
-
1177
- // Call update method with topicId option
1178
- const result = await messageModel.update(
1179
- 'msg-topic1',
1180
- { content: 'updated topic message 1' },
1181
- { topicId, sessionId },
1182
- );
1183
-
1184
- // Verify return result contains message list
1185
- expect(result.success).toBe(true);
1186
- expect(result.messages).toBeDefined();
1187
- expect(result.messages).toHaveLength(2);
1188
- expect(result.messages![0].content).toBe('updated topic message 1');
1189
- expect(result.messages![1].content).toBe('topic message 2');
1190
- });
1191
-
1192
- it('should return success without messages when options not provided', async () => {
1193
- // Create test data
1194
- await serverDB.insert(messages).values({
1195
- id: 'msg-no-options',
1196
- userId,
1197
- role: 'user',
1198
- content: 'original content',
1199
- });
1200
-
1201
- // Call update method,不提供选项
1202
- const result = await messageModel.update('msg-no-options', {
1203
- content: 'updated content',
1204
- });
1205
-
1206
- // 验证返回结果不包含消息列表
1207
- expect(result.success).toBe(true);
1208
- expect(result.messages).toBeUndefined();
1209
- });
1210
- });
1211
- });
1212
-
1213
- describe('deleteMessage', () => {
1214
- it('should delete a message', async () => {
1215
- // Create test data
1216
- await serverDB
1217
- .insert(messages)
1218
- .values([{ id: '1', userId, role: 'user', content: 'message 1' }]);
1219
-
1220
- // 调用 deleteMessage 方法
1221
- await messageModel.deleteMessage('1');
1222
-
1223
- // Assert result
1224
- const result = await serverDB.select().from(messages).where(eq(messages.id, '1'));
1225
- expect(result).toHaveLength(0);
1226
- });
1227
-
1228
- it('should delete a message with tool calls', async () => {
1229
- // Create test data
1230
- await serverDB.transaction(async (trx) => {
1231
- await trx.insert(messages).values([
1232
- { id: '1', userId, role: 'user', content: 'message 1', tools: [{ id: 'tool1' }] },
1233
- { id: '2', userId, role: 'tool', content: 'message 1' },
1234
- ]);
1235
- await trx
1236
- .insert(messagePlugins)
1237
- .values([{ id: '2', toolCallId: 'tool1', identifier: 'plugin-1', userId }]);
1238
- });
1239
-
1240
- // 调用 deleteMessage 方法
1241
- await messageModel.deleteMessage('1');
1242
-
1243
- // Assert result
1244
- const result = await serverDB.select().from(messages).where(eq(messages.id, '1'));
1245
- expect(result).toHaveLength(0);
1246
-
1247
- const result2 = await serverDB
1248
- .select()
1249
- .from(messagePlugins)
1250
- .where(eq(messagePlugins.id, '2'));
1251
-
1252
- expect(result2).toHaveLength(0);
1253
- });
1254
-
1255
- it('should only delete messages belonging to the user', async () => {
1256
- // Create test data
1257
- await serverDB
1258
- .insert(messages)
1259
- .values([{ id: '1', userId: '456', role: 'user', content: 'message 1' }]);
1260
-
1261
- // 调用 deleteMessage 方法
1262
- await messageModel.deleteMessage('1');
1263
-
1264
- // Assert result
1265
- const result = await serverDB.select().from(messages).where(eq(messages.id, '1'));
1266
- expect(result).toHaveLength(1);
1267
- });
1268
- });
1269
-
1270
- describe('deleteMessages', () => {
1271
- it('should delete 2 messages', async () => {
1272
- // Create test data
1273
- await serverDB.insert(messages).values([
1274
- { id: '1', userId, role: 'user', content: 'message 1' },
1275
- { id: '2', userId, role: 'user', content: 'message 2' },
1276
- ]);
1277
-
1278
- // 调用 deleteMessage 方法
1279
- await messageModel.deleteMessages(['1', '2']);
1280
-
1281
- // Assert result
1282
- const result = await serverDB.select().from(messages).where(eq(messages.id, '1'));
1283
- expect(result).toHaveLength(0);
1284
- const result2 = await serverDB.select().from(messages).where(eq(messages.id, '2'));
1285
- expect(result2).toHaveLength(0);
1286
- });
1287
-
1288
- it('should only delete messages belonging to the user', async () => {
1289
- // Create test data
1290
- await serverDB.insert(messages).values([
1291
- { id: '1', userId: '456', role: 'user', content: 'message 1' },
1292
- { id: '2', userId: '456', role: 'user', content: 'message 1' },
1293
- ]);
1294
-
1295
- // 调用 deleteMessage 方法
1296
- await messageModel.deleteMessages(['1', '2']);
1297
-
1298
- // Assert result
1299
- const result = await serverDB.select().from(messages).where(eq(messages.id, '1'));
1300
- expect(result).toHaveLength(1);
1301
- });
1302
- });
1303
-
1304
- describe('deleteAllMessages', () => {
1305
- it('should delete all messages belonging to the user', async () => {
1306
- // Create test data
1307
- await serverDB.insert(messages).values([
1308
- { id: '1', userId, role: 'user', content: 'message 1' },
1309
- { id: '2', userId, role: 'user', content: 'message 2' },
1310
- { id: '3', userId: '456', role: 'user', content: 'message 3' },
1311
- ]);
1312
-
1313
- // 调用 deleteAllMessages 方法
1314
- await messageModel.deleteAllMessages();
1315
-
1316
- // Assert result
1317
- const result = await serverDB.select().from(messages).where(eq(messages.userId, userId));
1318
-
1319
- expect(result).toHaveLength(0);
1320
-
1321
- const otherResult = await serverDB.select().from(messages).where(eq(messages.userId, '456'));
1322
-
1323
- expect(otherResult).toHaveLength(1);
1324
- });
1325
- });
1326
-
1327
- describe('updatePluginState', () => {
1328
- it('should update the state field in messagePlugins table', async () => {
1329
- // Create test data
1330
- await serverDB.insert(messages).values({ id: '1', content: 'abc', role: 'user', userId });
1331
- await serverDB.insert(messagePlugins).values([
1332
- {
1333
- id: '1',
1334
- toolCallId: 'tool1',
1335
- identifier: 'plugin1',
1336
- state: { key1: 'value1' },
1337
- userId,
1338
- },
1339
- ]);
1340
-
1341
- // 调用 updatePluginState 方法
1342
- await messageModel.updatePluginState('1', { key2: 'value2' });
1343
-
1344
- // Assert result
1345
- const result = await serverDB.select().from(messagePlugins).where(eq(messagePlugins.id, '1'));
1346
-
1347
- expect(result[0].state).toEqual({ key1: 'value1', key2: 'value2' });
1348
- });
1349
-
1350
- it('should throw an error if plugin does not exist', async () => {
1351
- // 调用 updatePluginState 方法
1352
- await expect(messageModel.updatePluginState('1', { key: 'value' })).rejects.toThrowError(
1353
- 'Plugin not found',
1354
- );
1355
- });
1356
- });
1357
- describe('updateMessagePlugin', () => {
1358
- it('should update the state field in messagePlugins table', async () => {
1359
- // Create test data
1360
- await serverDB.insert(messages).values({ id: '1', content: 'abc', role: 'user', userId });
1361
- await serverDB.insert(messagePlugins).values([
1362
- {
1363
- id: '1',
1364
- toolCallId: 'tool1',
1365
- identifier: 'plugin1',
1366
- state: { key1: 'value1' },
1367
- userId,
1368
- },
1369
- ]);
1370
-
1371
- // 调用 updatePluginState 方法
1372
- await messageModel.updateMessagePlugin('1', { identifier: 'plugin2' });
1373
-
1374
- // Assert result
1375
- const result = await serverDB.select().from(messagePlugins).where(eq(messagePlugins.id, '1'));
1376
-
1377
- expect(result[0].identifier).toEqual('plugin2');
1378
- });
1379
-
1380
- it('should throw an error if plugin does not exist', async () => {
1381
- // 调用 updatePluginState 方法
1382
- await expect(messageModel.updatePluginState('1', { key: 'value' })).rejects.toThrowError(
1383
- 'Plugin not found',
1384
- );
1385
- });
1386
- });
1387
-
1388
- describe('updateMetadata', () => {
1389
- it('should update metadata for an existing message', async () => {
1390
- // Create test data
1391
- await serverDB.insert(messages).values({
1392
- id: 'msg-with-metadata',
1393
- userId,
1394
- role: 'user',
1395
- content: 'test message',
1396
- metadata: { existingKey: 'existingValue' },
1397
- });
1398
-
1399
- // 调用 updateMetadata 方法
1400
- await messageModel.updateMetadata('msg-with-metadata', { newKey: 'newValue' });
1401
-
1402
- // Assert result
1403
- const result = await serverDB
1404
- .select()
1405
- .from(messages)
1406
- .where(eq(messages.id, 'msg-with-metadata'));
1407
-
1408
- expect(result[0].metadata).toEqual({
1409
- existingKey: 'existingValue',
1410
- newKey: 'newValue',
1411
- });
1412
- });
1413
-
1414
- it('should merge new metadata with existing metadata using lodash merge behavior', async () => {
1415
- // Create test data
1416
- await serverDB.insert(messages).values({
1417
- id: 'msg-merge-metadata',
1418
- userId,
1419
- role: 'assistant',
1420
- content: 'test message',
1421
- metadata: {
1422
- level1: {
1423
- level2a: 'original',
1424
- level2b: { level3: 'deep' },
1425
- },
1426
- array: [1, 2, 3],
1427
- },
1428
- });
1429
-
1430
- // 调用 updateMetadata 方法
1431
- await messageModel.updateMetadata('msg-merge-metadata', {
1432
- level1: {
1433
- level2a: 'updated',
1434
- level2c: 'new',
1435
- },
1436
- newTopLevel: 'value',
1437
- });
1438
-
1439
- // Assert result - 应该使用 lodash merge 行为
1440
- const result = await serverDB
1441
- .select()
1442
- .from(messages)
1443
- .where(eq(messages.id, 'msg-merge-metadata'));
1444
-
1445
- expect(result[0].metadata).toEqual({
1446
- level1: {
1447
- level2a: 'updated',
1448
- level2b: { level3: 'deep' },
1449
- level2c: 'new',
1450
- },
1451
- array: [1, 2, 3],
1452
- newTopLevel: 'value',
1453
- });
1454
- });
1455
-
1456
- it('should handle non-existent message IDs', async () => {
1457
- // 调用 updateMetadata 方法,尝试更新不存在的消息
1458
- const result = await messageModel.updateMetadata('non-existent-id', { key: 'value' });
1459
-
1460
- // Assert result - 应该返回 undefined
1461
- expect(result).toBeUndefined();
1462
- });
1463
-
1464
- it('should handle empty metadata updates', async () => {
1465
- // Create test data
1466
- await serverDB.insert(messages).values({
1467
- id: 'msg-empty-metadata',
1468
- userId,
1469
- role: 'user',
1470
- content: 'test message',
1471
- metadata: { originalKey: 'originalValue' },
1472
- });
1473
-
1474
- // 调用 updateMetadata 方法,传递空对象
1475
- await messageModel.updateMetadata('msg-empty-metadata', {});
1476
-
1477
- // Assert result - 原始 metadata 应该保持不变
1478
- const result = await serverDB
1479
- .select()
1480
- .from(messages)
1481
- .where(eq(messages.id, 'msg-empty-metadata'));
1482
-
1483
- expect(result[0].metadata).toEqual({ originalKey: 'originalValue' });
1484
- });
1485
-
1486
- it('should handle message with null metadata', async () => {
1487
- // Create test data
1488
- await serverDB.insert(messages).values({
1489
- id: 'msg-null-metadata',
1490
- userId,
1491
- role: 'user',
1492
- content: 'test message',
1493
- metadata: null,
1494
- });
1495
-
1496
- // 调用 updateMetadata 方法
1497
- await messageModel.updateMetadata('msg-null-metadata', { key: 'value' });
1498
-
1499
- // Assert result - 应该创建新的 metadata
1500
- const result = await serverDB
1501
- .select()
1502
- .from(messages)
1503
- .where(eq(messages.id, 'msg-null-metadata'));
1504
-
1505
- expect(result[0].metadata).toEqual({ key: 'value' });
1506
- });
1507
-
1508
- it('should only update messages belonging to the current user', async () => {
1509
- // Create test data - 其他用户的消息
1510
- await serverDB.insert(messages).values({
1511
- id: 'msg-other-user',
1512
- userId: '456',
1513
- role: 'user',
1514
- content: 'test message',
1515
- metadata: { originalKey: 'originalValue' },
1516
- });
1517
-
1518
- // 调用 updateMetadata 方法
1519
- const result = await messageModel.updateMetadata('msg-other-user', {
1520
- hackedKey: 'hackedValue',
1521
- });
1522
-
1523
- // Assert result - 应该返回 undefined
1524
- expect(result).toBeUndefined();
1525
-
1526
- // 验证原始 metadata 未被修改
1527
- const dbResult = await serverDB
1528
- .select()
1529
- .from(messages)
1530
- .where(eq(messages.id, 'msg-other-user'));
1531
-
1532
- expect(dbResult[0].metadata).toEqual({ originalKey: 'originalValue' });
1533
- });
1534
-
1535
- it('should handle complex nested metadata updates', async () => {
1536
- // Create test data
1537
- await serverDB.insert(messages).values({
1538
- id: 'msg-complex-metadata',
1539
- userId,
1540
- role: 'assistant',
1541
- content: 'test message',
1542
- metadata: {
1543
- config: {
1544
- settings: {
1545
- enabled: true,
1546
- options: ['a', 'b'],
1547
- },
1548
- version: 1,
1549
- },
1550
- },
1551
- });
1552
-
1553
- // 调用 updateMetadata 方法
1554
- await messageModel.updateMetadata('msg-complex-metadata', {
1555
- config: {
1556
- settings: {
1557
- enabled: false,
1558
- timeout: 5000,
1559
- },
1560
- newField: 'value',
1561
- },
1562
- stats: { count: 10 },
1563
- });
1564
-
1565
- // Assert result
1566
- const result = await serverDB
1567
- .select()
1568
- .from(messages)
1569
- .where(eq(messages.id, 'msg-complex-metadata'));
1570
-
1571
- expect(result[0].metadata).toEqual({
1572
- config: {
1573
- settings: {
1574
- enabled: false,
1575
- options: ['a', 'b'],
1576
- timeout: 5000,
1577
- },
1578
- version: 1,
1579
- newField: 'value',
1580
- },
1581
- stats: { count: 10 },
1582
- });
1583
- });
1584
- });
1585
-
1586
- describe('updateTranslate', () => {
1587
- it('should insert a new record if message does not exist in messageTranslates table', async () => {
1588
- // Create test data
1589
- await serverDB
1590
- .insert(messages)
1591
- .values([{ id: '1', userId, role: 'user', content: 'message 1' }]);
1592
-
1593
- // 调用 updateTranslate 方法
1594
- await messageModel.updateTranslate('1', {
1595
- content: 'translated message 1',
1596
- from: 'en',
1597
- to: 'zh',
1598
- });
1599
-
1600
- // Assert result
1601
- const result = await serverDB
1602
- .select()
1603
- .from(messageTranslates)
1604
- .where(eq(messageTranslates.id, '1'));
1605
-
1606
- expect(result).toHaveLength(1);
1607
- expect(result[0].content).toBe('translated message 1');
1608
- });
1609
-
1610
- it('should update the corresponding fields if message exists in messageTranslates table', async () => {
1611
- // Create test data
1612
- await serverDB.transaction(async (trx) => {
1613
- await trx
1614
- .insert(messages)
1615
- .values([{ id: '1', userId, role: 'user', content: 'message 1' }]);
1616
- await trx
1617
- .insert(messageTranslates)
1618
- .values([{ id: '1', content: 'translated message 1', from: 'en', to: 'zh', userId }]);
1619
- });
1620
-
1621
- // 调用 updateTranslate 方法
1622
- await messageModel.updateTranslate('1', { content: 'updated translated message 1' });
1623
-
1624
- // Assert result
1625
- const result = await serverDB
1626
- .select()
1627
- .from(messageTranslates)
1628
- .where(eq(messageTranslates.id, '1'));
1629
-
1630
- expect(result[0].content).toBe('updated translated message 1');
1631
- });
1632
- });
1633
-
1634
- describe('updateTTS', () => {
1635
- it('should insert a new record if message does not exist in messageTTS table', async () => {
1636
- // Create test data
1637
- await serverDB
1638
- .insert(messages)
1639
- .values([{ id: '1', userId, role: 'user', content: 'message 1' }]);
1640
-
1641
- // 调用 updateTTS 方法
1642
- await messageModel.updateTTS('1', { contentMd5: 'md5', file: 'f1', voice: 'voice1' });
1643
-
1644
- // Assert result
1645
- const result = await serverDB.select().from(messageTTS).where(eq(messageTTS.id, '1'));
1646
-
1647
- expect(result).toHaveLength(1);
1648
- expect(result[0].voice).toBe('voice1');
1649
- });
1650
-
1651
- it('should update the corresponding fields if message exists in messageTTS table', async () => {
1652
- // Create test data
1653
- await serverDB.transaction(async (trx) => {
1654
- await trx
1655
- .insert(messages)
1656
- .values([{ id: '1', userId, role: 'user', content: 'message 1' }]);
1657
- await trx
1658
- .insert(messageTTS)
1659
- .values([{ id: '1', contentMd5: 'md5', fileId: 'f1', voice: 'voice1', userId }]);
1660
- });
1661
-
1662
- // 调用 updateTTS 方法
1663
- await messageModel.updateTTS('1', { voice: 'updated voice1' });
1664
-
1665
- // Assert result
1666
- const result = await serverDB.select().from(messageTTS).where(eq(messageTTS.id, '1'));
1667
-
1668
- expect(result[0].voice).toBe('updated voice1');
1669
- });
1670
- });
1671
-
1672
- describe('deleteMessageTranslate', () => {
1673
- it('should delete the message translate record', async () => {
1674
- // Create test data
1675
- await serverDB.insert(messages).values([{ id: '1', role: 'abc', userId }]);
1676
- await serverDB.insert(messageTranslates).values([{ id: '1', userId }]);
1677
-
1678
- // 调用 deleteMessageTranslate 方法
1679
- await messageModel.deleteMessageTranslate('1');
1680
-
1681
- // Assert result
1682
- const result = await serverDB
1683
- .select()
1684
- .from(messageTranslates)
1685
- .where(eq(messageTranslates.id, '1'));
1686
-
1687
- expect(result).toHaveLength(0);
1688
- });
1689
- });
1690
-
1691
- describe('deleteMessageTTS', () => {
1692
- it('should delete the message TTS record', async () => {
1693
- // Create test data
1694
- await serverDB.insert(messages).values([{ id: '1', role: 'abc', userId }]);
1695
- await serverDB.insert(messageTTS).values([{ userId, id: '1' }]);
1696
-
1697
- // 调用 deleteMessageTTS 方法
1698
- await messageModel.deleteMessageTTS('1');
1699
-
1700
- // Assert result
1701
- const result = await serverDB.select().from(messageTTS).where(eq(messageTTS.id, '1'));
1702
- expect(result).toHaveLength(0);
1703
- });
1704
- });
1705
-
1706
- describe('count', () => {
1707
- it('should return the count of messages belonging to the user', async () => {
1708
- // Create test data
1709
- await serverDB.insert(messages).values([
1710
- { id: '1', userId, role: 'user', content: 'message 1' },
1711
- { id: '2', userId, role: 'user', content: 'message 2' },
1712
- { id: '3', userId: '456', role: 'user', content: 'message 3' },
1713
- ]);
1714
-
1715
- // 调用 count 方法
1716
- const result = await messageModel.count();
1717
-
1718
- // Assert result
1719
- expect(result).toBe(2);
1720
- });
1721
-
1722
- describe('count with date filters', () => {
1723
- beforeEach(async () => {
1724
- // Create test data,包含不同日期的消息
1725
- await serverDB.insert(messages).values([
1726
- {
1727
- id: 'date1',
1728
- userId,
1729
- role: 'user',
1730
- content: 'message 1',
1731
- createdAt: new Date('2023-01-15'),
1732
- },
1733
- {
1734
- id: 'date2',
1735
- userId,
1736
- role: 'user',
1737
- content: 'message 2',
1738
- createdAt: new Date('2023-02-15'),
1739
- },
1740
- {
1741
- id: 'date3',
1742
- userId,
1743
- role: 'user',
1744
- content: 'message 3',
1745
- createdAt: new Date('2023-03-15'),
1746
- },
1747
- {
1748
- id: 'date4',
1749
- userId,
1750
- role: 'user',
1751
- content: 'message 4',
1752
- createdAt: new Date('2023-04-15'),
1753
- },
1754
- ]);
1755
- });
1756
-
1757
- it('should count messages with startDate filter', async () => {
1758
- const result = await messageModel.count({ startDate: '2023-02-01' });
1759
- expect(result).toBe(3); // 2月15日, 3月15日, 4月15日的消息
1760
- });
1761
-
1762
- it('should count messages with endDate filter', async () => {
1763
- const result = await messageModel.count({ endDate: '2023-03-01' });
1764
- expect(result).toBe(2); // 1月15日, 2月15日的消息
1765
- });
1766
-
1767
- it('should count messages with both startDate and endDate filters', async () => {
1768
- const result = await messageModel.count({
1769
- startDate: '2023-02-01',
1770
- endDate: '2023-03-31',
1771
- });
1772
- expect(result).toBe(2); // 2月15日, 3月15日的消息
1773
- });
1774
-
1775
- it('should count messages with range filter', async () => {
1776
- const result = await messageModel.count({
1777
- range: ['2023-02-01', '2023-04-01'],
1778
- });
1779
- expect(result).toBe(2); // 2月15日, 3月15日的消息
1780
- });
1781
-
1782
- it('should handle edge cases in date filters', async () => {
1783
- // 边界日期
1784
- const result1 = await messageModel.count({
1785
- startDate: '2023-01-15',
1786
- endDate: '2023-04-15',
1787
- });
1788
- expect(result1).toBe(4); // 包含所有消息
1789
-
1790
- // 没有消息的日期范围
1791
- const result2 = await messageModel.count({
1792
- startDate: '2023-05-01',
1793
- endDate: '2023-06-01',
1794
- });
1795
- expect(result2).toBe(0);
1796
-
1797
- // 精确到一天
1798
- const result3 = await messageModel.count({
1799
- startDate: '2023-01-15',
1800
- endDate: '2023-01-15',
1801
- });
1802
- expect(result3).toBe(1);
1803
- });
1804
- });
1805
- });
1806
-
1807
- describe('findMessageQueriesById', () => {
1808
- it('should return undefined for non-existent message query', async () => {
1809
- const result = await messageModel.findMessageQueriesById('non-existent-id');
1810
- expect(result).toBeUndefined();
1811
- });
1812
-
1813
- it('should return message query with embeddings', async () => {
1814
- const query1Id = uuid();
1815
- const embeddings1Id = uuid();
1816
-
1817
- await serverDB.transaction(async (trx) => {
1818
- await trx.insert(messages).values({ id: 'msg1', userId, role: 'user', content: 'abc' });
1819
-
1820
- await trx.insert(embeddings).values({
1821
- id: embeddings1Id,
1822
- embeddings: codeEmbedding,
1823
- });
1824
-
1825
- await trx.insert(messageQueries).values({
1826
- id: query1Id,
1827
- messageId: 'msg1',
1828
- userQuery: 'test query',
1829
- rewriteQuery: 'rewritten query',
1830
- embeddingsId: embeddings1Id,
1831
- userId,
1832
- });
1833
- });
1834
-
1835
- const result = await messageModel.findMessageQueriesById('msg1');
1836
-
1837
- expect(result).toBeDefined();
1838
- expect(result).toMatchObject({
1839
- id: query1Id,
1840
- userQuery: 'test query',
1841
- rewriteQuery: 'rewritten query',
1842
- embeddings: codeEmbedding,
1843
- });
1844
- });
1845
- });
1846
-
1847
- describe('deleteMessagesBySession', () => {
1848
- it('should delete messages by session ID', async () => {
1849
- await serverDB.insert(sessions).values([
1850
- { id: 'session1', userId },
1851
- { id: 'session2', userId },
1852
- ]);
1853
-
1854
- await serverDB.insert(messages).values([
1855
- {
1856
- id: '1',
1857
- userId,
1858
- sessionId: 'session1',
1859
- role: 'user',
1860
- content: 'message 1',
1861
- },
1862
- {
1863
- id: '2',
1864
- userId,
1865
- sessionId: 'session1',
1866
- role: 'assistant',
1867
- content: 'message 2',
1868
- },
1869
- {
1870
- id: '3',
1871
- userId,
1872
- sessionId: 'session2',
1873
- role: 'user',
1874
- content: 'message 3',
1875
- },
1876
- ]);
1877
-
1878
- await messageModel.deleteMessagesBySession('session1');
1879
-
1880
- const remainingMessages = await serverDB
1881
- .select()
1882
- .from(messages)
1883
- .where(eq(messages.userId, userId));
1884
-
1885
- expect(remainingMessages).toHaveLength(1);
1886
- expect(remainingMessages[0].id).toBe('3');
1887
- });
1888
-
1889
- it('should delete messages by session ID and topic ID', async () => {
1890
- await serverDB.insert(sessions).values([{ id: 'session1', userId }]);
1891
- await serverDB.insert(topics).values([
1892
- { id: 'topic1', sessionId: 'session1', userId },
1893
- { id: 'topic2', sessionId: 'session1', userId },
1894
- ]);
1895
-
1896
- await serverDB.insert(messages).values([
1897
- {
1898
- id: '1',
1899
- userId,
1900
- sessionId: 'session1',
1901
- topicId: 'topic1',
1902
- role: 'user',
1903
- content: 'message 1',
1904
- },
1905
- {
1906
- id: '2',
1907
- userId,
1908
- sessionId: 'session1',
1909
- topicId: 'topic2',
1910
- role: 'assistant',
1911
- content: 'message 2',
1912
- },
1913
- ]);
1914
-
1915
- await messageModel.deleteMessagesBySession('session1', 'topic1');
1916
-
1917
- const remainingMessages = await serverDB
1918
- .select()
1919
- .from(messages)
1920
- .where(eq(messages.userId, userId));
1921
-
1922
- expect(remainingMessages).toHaveLength(1);
1923
- expect(remainingMessages[0].id).toBe('2');
1924
- });
1925
-
1926
- it('should delete only non-topic messages when topicId is null', async () => {
1927
- await serverDB.insert(sessions).values([{ id: 'session1', userId }]);
1928
- await serverDB.insert(topics).values([
1929
- { id: 'topic1', sessionId: 'session1', userId },
1930
- { id: 'topic2', sessionId: 'session1', userId },
1931
- ]);
1932
-
1933
- await serverDB.insert(messages).values([
1934
- {
1935
- id: '1',
1936
- userId,
1937
- sessionId: 'session1',
1938
- topicId: null,
1939
- role: 'user',
1940
- content: 'message without topic 1',
1941
- },
1942
- {
1943
- id: '2',
1944
- userId,
1945
- sessionId: 'session1',
1946
- topicId: null,
1947
- role: 'assistant',
1948
- content: 'message without topic 2',
1949
- },
1950
- {
1951
- id: '3',
1952
- userId,
1953
- sessionId: 'session1',
1954
- topicId: 'topic1',
1955
- role: 'user',
1956
- content: 'message in topic1',
1957
- },
1958
- {
1959
- id: '4',
1960
- userId,
1961
- sessionId: 'session1',
1962
- topicId: 'topic2',
1963
- role: 'assistant',
1964
- content: 'message in topic2',
1965
- },
1966
- ]);
1967
-
1968
- // Delete messages in session1 with null topicId
1969
- await messageModel.deleteMessagesBySession('session1', null);
1970
-
1971
- const remainingMessages = await serverDB
1972
- .select()
1973
- .from(messages)
1974
- .where(eq(messages.userId, userId))
1975
- .orderBy(messages.id);
1976
-
1977
- // Should only keep messages with topics
1978
- expect(remainingMessages).toHaveLength(2);
1979
- expect(remainingMessages[0].id).toBe('3');
1980
- expect(remainingMessages[1].id).toBe('4');
1981
- });
1982
- });
1983
-
1984
- describe('genId', () => {
1985
- it('should generate unique message IDs', () => {
1986
- const model = new MessageModel(serverDB, userId);
1987
- // @ts-ignore - accessing private method for testing
1988
- const id1 = model.genId();
1989
- // @ts-ignore - accessing private method for testing
1990
- const id2 = model.genId();
1991
-
1992
- expect(id1).toHaveLength(18);
1993
- expect(id2).toHaveLength(18);
1994
- expect(id1).not.toBe(id2);
1995
- expect(id1).toMatch(/^msg_/);
1996
- expect(id2).toMatch(/^msg_/);
1997
- });
1998
- });
1999
-
2000
- describe('countWords', () => {
2001
- it('should count total words of messages belonging to the user', async () => {
2002
- // Create test data
2003
- await serverDB.insert(messages).values([
2004
- { id: '1', userId, role: 'user', content: 'hello world' },
2005
- { id: '2', userId, role: 'user', content: 'test message' },
2006
- { id: '3', userId: '456', role: 'user', content: 'other user message' },
2007
- ]);
2008
-
2009
- // 调用 countWords 方法
2010
- const result = await messageModel.countWords();
2011
-
2012
- // Assert result - 'hello world' + 'test message' = 23 characters
2013
- expect(result).toEqual(23);
2014
- });
2015
-
2016
- it('should count words within date range', async () => {
2017
- // Create test data
2018
- await serverDB.insert(messages).values([
2019
- {
2020
- id: '1',
2021
- userId,
2022
- role: 'user',
2023
- content: 'old message',
2024
- createdAt: new Date('2023-01-01'),
2025
- },
2026
- {
2027
- id: '2',
2028
- userId,
2029
- role: 'user',
2030
- content: 'new message',
2031
- createdAt: new Date('2023-06-01'),
2032
- },
2033
- ]);
2034
-
2035
- // 调用 countWords 方法,设置日期范围
2036
- const result = await messageModel.countWords({
2037
- range: ['2023-05-01', '2023-07-01'],
2038
- });
2039
-
2040
- // Assert result - 只计算 'new message' = 11 characters
2041
- expect(result).toEqual(11);
2042
- });
2043
-
2044
- it('should handle empty content', async () => {
2045
- // Create test data
2046
- await serverDB.insert(messages).values([
2047
- { id: '1', userId, role: 'user', content: '' },
2048
- { id: '2', userId, role: 'user', content: null },
2049
- ]);
2050
-
2051
- // 调用 countWords 方法
2052
- const result = await messageModel.countWords();
2053
-
2054
- // Assert result
2055
- expect(result).toEqual(0);
2056
- });
2057
- });
2058
-
2059
- describe('getHeatmaps', () => {
2060
- it('should return heatmap data for the last year', async () => {
2061
- // 使用固定日期进行测试
2062
- vi.useFakeTimers();
2063
- const fixedDate = new Date('2023-04-07T13:00:00Z');
2064
- vi.setSystemTime(fixedDate);
2065
-
2066
- const today = dayjs(fixedDate);
2067
- const twoDaysAgoDate = today.subtract(2, 'day').format('YYYY-MM-DD');
2068
- const oneDayAgoDate = today.subtract(1, 'day').format('YYYY-MM-DD');
2069
- const todayDate = today.format('YYYY-MM-DD');
2070
-
2071
- // Create test data
2072
- await serverDB.insert(messages).values([
2073
- {
2074
- id: '1',
2075
- userId,
2076
- role: 'user',
2077
- content: 'message 1',
2078
- createdAt: today.subtract(2, 'day').toDate(),
2079
- },
2080
- {
2081
- id: '2',
2082
- userId,
2083
- role: 'user',
2084
- content: 'message 2',
2085
- createdAt: today.subtract(2, 'day').toDate(),
2086
- },
2087
- {
2088
- id: '3',
2089
- userId,
2090
- role: 'user',
2091
- content: 'message 3',
2092
- createdAt: today.subtract(1, 'day').toDate(),
2093
- },
2094
- ]);
2095
-
2096
- // 调用 getHeatmaps 方法
2097
- const result = await messageModel.getHeatmaps();
2098
-
2099
- // Assert result
2100
- expect(result.length).toBeGreaterThanOrEqual(366);
2101
- expect(result.length).toBeLessThan(368);
2102
-
2103
- // 检查两天前的数据
2104
- const twoDaysAgo = result.find((item) => item.date === twoDaysAgoDate);
2105
- expect(twoDaysAgo?.count).toBe(2);
2106
- expect(twoDaysAgo?.level).toBe(1);
2107
-
2108
- // 检查一天前的数据
2109
- const oneDayAgo = result.find((item) => item.date === oneDayAgoDate);
2110
- expect(oneDayAgo?.count).toBe(1);
2111
- expect(oneDayAgo?.level).toBe(1);
2112
-
2113
- // 检查今天的数据
2114
- const todayData = result.find((item) => item.date === todayDate);
2115
- expect(todayData?.count).toBe(0);
2116
- expect(todayData?.level).toBe(0);
2117
-
2118
- vi.useRealTimers();
2119
- });
2120
-
2121
- it('should calculate correct levels based on message count', async () => {
2122
- // 使用固定日期进行测试
2123
- vi.useFakeTimers();
2124
- const fixedDate = new Date('2023-05-15T12:00:00Z');
2125
- vi.setSystemTime(fixedDate);
2126
-
2127
- const today = dayjs(fixedDate);
2128
- const fourDaysAgoDate = today.subtract(4, 'day').format('YYYY-MM-DD');
2129
- const threeDaysAgoDate = today.subtract(3, 'day').format('YYYY-MM-DD');
2130
- const twoDaysAgoDate = today.subtract(2, 'day').format('YYYY-MM-DD');
2131
- const oneDayAgoDate = today.subtract(1, 'day').format('YYYY-MM-DD');
2132
- const todayDate = today.format('YYYY-MM-DD');
2133
-
2134
- // Create test data - 不同数量的消息以测试不同的等级
2135
- await serverDB.insert(messages).values([
2136
- // 1 message - level 1
2137
- {
2138
- id: '1',
2139
- userId,
2140
- role: 'user',
2141
- content: 'message 1',
2142
- createdAt: today.subtract(4, 'day').toDate(),
2143
- },
2144
- // 6 messages - level 2
2145
- ...Array(6)
2146
- .fill(0)
2147
- .map((_, i) => ({
2148
- id: `2-${i}`,
2149
- userId,
2150
- role: 'user',
2151
- content: `message 2-${i}`,
2152
- createdAt: today.subtract(3, 'day').toDate(),
2153
- })),
2154
- // 11 messages - level 3
2155
- ...Array(11)
2156
- .fill(0)
2157
- .map((_, i) => ({
2158
- id: `3-${i}`,
2159
- userId,
2160
- role: 'user',
2161
- content: `message 3-${i}`,
2162
- createdAt: today.subtract(2, 'day').toDate(),
2163
- })),
2164
- // 16 messages - level 4
2165
- ...Array(16)
2166
- .fill(0)
2167
- .map((_, i) => ({
2168
- id: `4-${i}`,
2169
- userId,
2170
- role: 'user',
2171
- content: `message 4-${i}`,
2172
- createdAt: today.subtract(1, 'day').toDate(),
2173
- })),
2174
- // 21 messages - level 4
2175
- ...Array(21)
2176
- .fill(0)
2177
- .map((_, i) => ({
2178
- id: `5-${i}`,
2179
- userId,
2180
- role: 'user',
2181
- content: `message 5-${i}`,
2182
- createdAt: today.toDate(),
2183
- })),
2184
- ]);
2185
-
2186
- // 调用 getHeatmaps 方法
2187
- const result = await messageModel.getHeatmaps();
2188
-
2189
- // 检查不同天数的等级
2190
- const fourDaysAgo = result.find((item) => item.date === fourDaysAgoDate);
2191
- expect(fourDaysAgo?.count).toBe(1);
2192
- expect(fourDaysAgo?.level).toBe(1);
2193
-
2194
- const threeDaysAgo = result.find((item) => item.date === threeDaysAgoDate);
2195
- expect(threeDaysAgo?.count).toBe(6);
2196
- expect(threeDaysAgo?.level).toBe(2);
2197
-
2198
- const twoDaysAgo = result.find((item) => item.date === twoDaysAgoDate);
2199
- expect(twoDaysAgo?.count).toBe(11);
2200
- expect(twoDaysAgo?.level).toBe(3);
2201
-
2202
- const oneDayAgo = result.find((item) => item.date === oneDayAgoDate);
2203
- expect(oneDayAgo?.count).toBe(16);
2204
- expect(oneDayAgo?.level).toBe(4);
2205
-
2206
- const todayData = result.find((item) => item.date === todayDate);
2207
- expect(todayData?.count).toBe(21);
2208
- expect(todayData?.level).toBe(4);
2209
-
2210
- vi.useRealTimers();
2211
- });
2212
-
2213
- it.skip('should return time count correctly when 19:00 time', async () => {
2214
- // 使用固定日期进行测试
2215
- vi.useFakeTimers();
2216
- const fixedDate = new Date('2025-04-02T19:00:00Z');
2217
- vi.setSystemTime(fixedDate);
2218
-
2219
- const today = dayjs(fixedDate);
2220
- const twoDaysAgoDate = today.subtract(2, 'day').format('YYYY-MM-DD');
2221
- const oneDayAgoDate = today.subtract(1, 'day').format('YYYY-MM-DD');
2222
- const todayDate = today.format('YYYY-MM-DD');
2223
-
2224
- // Create test data
2225
- await serverDB.insert(messages).values([
2226
- {
2227
- id: '1',
2228
- userId,
2229
- role: 'user',
2230
- content: 'message 1',
2231
- createdAt: today.subtract(2, 'day').toDate(),
2232
- },
2233
- {
2234
- id: '2',
2235
- userId,
2236
- role: 'user',
2237
- content: 'message 2',
2238
- createdAt: today.subtract(2, 'day').toDate(),
2239
- },
2240
- {
2241
- id: '3',
2242
- userId,
2243
- role: 'user',
2244
- content: 'message 3',
2245
- createdAt: today.subtract(1, 'day').toDate(),
2246
- },
2247
- ]);
2248
-
2249
- // 调用 getHeatmaps 方法
2250
- const result = await messageModel.getHeatmaps();
2251
-
2252
- // Assert result
2253
- expect(result.length).toBeGreaterThanOrEqual(366);
2254
- expect(result.length).toBeLessThan(368);
2255
-
2256
- // 检查两天前的数据
2257
- const twoDaysAgo = result.find((item) => item.date === twoDaysAgoDate);
2258
- expect(twoDaysAgo?.count).toBe(2);
2259
- expect(twoDaysAgo?.level).toBe(1);
2260
-
2261
- // 检查一天前的数据
2262
- const oneDayAgo = result.find((item) => item.date === oneDayAgoDate);
2263
- expect(oneDayAgo?.count).toBe(1);
2264
- expect(oneDayAgo?.level).toBe(1);
2265
-
2266
- // 检查今天的数据
2267
- const todayData = result.find((item) => item.date === todayDate);
2268
- expect(todayData?.count).toBe(0);
2269
- expect(todayData?.level).toBe(0);
2270
-
2271
- vi.useRealTimers();
2272
- });
2273
-
2274
- it('should handle empty data', async () => {
2275
- // 不创建任何消息数据
2276
-
2277
- // 调用 getHeatmaps 方法
2278
- const result = await messageModel.getHeatmaps();
2279
-
2280
- // Assert result
2281
- expect(result.length).toBeGreaterThanOrEqual(366);
2282
- expect(result.length).toBeLessThan(368);
2283
-
2284
- // 检查所有数据的 count 和 level 是否为 0
2285
- result.forEach((item) => {
2286
- expect(item.count).toBe(0);
2287
- expect(item.level).toBe(0);
2288
- });
2289
- });
2290
- });
2291
-
2292
- describe('rankModels', () => {
2293
- it('should rank models by usage count', async () => {
2294
- // Create test data
2295
- await serverDB.insert(messages).values([
2296
- { id: '1', userId, role: 'assistant', content: 'message 1', model: 'gpt-3.5' },
2297
- { id: '2', userId, role: 'assistant', content: 'message 2', model: 'gpt-3.5' },
2298
- { id: '3', userId, role: 'assistant', content: 'message 3', model: 'gpt-4' },
2299
- { id: '4', userId: '456', role: 'assistant', content: 'message 4', model: 'gpt-3.5' }, // 其他用户的消息
2300
- ]);
2301
-
2302
- // 调用 rankModels 方法
2303
- const result = await messageModel.rankModels();
2304
-
2305
- // Assert result
2306
- expect(result).toHaveLength(2);
2307
- expect(result[0]).toEqual({ id: 'gpt-3.5', count: 2 }); // 当前用户使用 gpt-3.5 两次
2308
- expect(result[1]).toEqual({ id: 'gpt-4', count: 1 }); // 当前用户使用 gpt-4 一次
2309
- });
2310
-
2311
- it('should only count messages with model field', async () => {
2312
- // Create test data,包括没有 model 字段的消息
2313
- await serverDB.insert(messages).values([
2314
- { id: '1', userId, role: 'assistant', content: 'message 1', model: 'gpt-3.5' },
2315
- { id: '2', userId, role: 'assistant', content: 'message 2', model: null },
2316
- { id: '3', userId, role: 'user', content: 'message 3' }, // 用户消息通常没有 model
2317
- ]);
2318
-
2319
- // 调用 rankModels 方法
2320
- const result = await messageModel.rankModels();
2321
-
2322
- // Assert result
2323
- expect(result).toHaveLength(1);
2324
- expect(result[0]).toEqual({ id: 'gpt-3.5', count: 1 });
2325
- });
2326
-
2327
- it('should return empty array when no models are used', async () => {
2328
- // Create test data,所有消息都没有 model
2329
- await serverDB.insert(messages).values([
2330
- { id: '1', userId, role: 'user', content: 'message 1' },
2331
- { id: '2', userId, role: 'assistant', content: 'message 2' },
2332
- ]);
2333
-
2334
- // 调用 rankModels 方法
2335
- const result = await messageModel.rankModels();
2336
-
2337
- // Assert result
2338
- expect(result).toHaveLength(0);
2339
- });
2340
-
2341
- it('should order models by count in descending order', async () => {
2342
- // Create test data,使用不同次数的模型
2343
- await serverDB.insert(messages).values([
2344
- { id: '1', userId, role: 'assistant', content: 'message 1', model: 'gpt-4' },
2345
- { id: '2', userId, role: 'assistant', content: 'message 2', model: 'gpt-3.5' },
2346
- { id: '3', userId, role: 'assistant', content: 'message 3', model: 'gpt-3.5' },
2347
- { id: '4', userId, role: 'assistant', content: 'message 4', model: 'claude' },
2348
- { id: '5', userId, role: 'assistant', content: 'message 5', model: 'gpt-3.5' },
2349
- ]);
2350
-
2351
- // 调用 rankModels 方法
2352
- const result = await messageModel.rankModels();
2353
-
2354
- // Assert result
2355
- expect(result).toHaveLength(3);
2356
- expect(result[0]).toEqual({ id: 'gpt-3.5', count: 3 }); // 最多使用
2357
- expect(result[1]).toEqual({ id: 'claude', count: 1 });
2358
- expect(result[2]).toEqual({ id: 'gpt-4', count: 1 });
2359
- });
2360
- });
2361
-
2362
- describe('hasMoreThanN', () => {
2363
- it('should return true when message count is greater than N', async () => {
2364
- // Create test data
2365
- await serverDB.insert(messages).values([
2366
- { id: '1', userId, role: 'user', content: 'message 1' },
2367
- { id: '2', userId, role: 'user', content: 'message 2' },
2368
- { id: '3', userId, role: 'user', content: 'message 3' },
2369
- ]);
2370
-
2371
- // 测试不同的 N 值
2372
- const result1 = await messageModel.hasMoreThanN(2); // 3 > 2
2373
- const result2 = await messageModel.hasMoreThanN(3); // 3 ≯ 3
2374
- const result3 = await messageModel.hasMoreThanN(4); // 3 ≯ 4
2375
-
2376
- expect(result1).toBe(true);
2377
- expect(result2).toBe(false);
2378
- expect(result3).toBe(false);
2379
- });
2380
-
2381
- it('should only count messages belonging to the user', async () => {
2382
- // Create test data,包括其他用户的消息
2383
- await serverDB.insert(messages).values([
2384
- { id: '1', userId, role: 'user', content: 'message 1' },
2385
- { id: '2', userId, role: 'user', content: 'message 2' },
2386
- { id: '3', userId: '456', role: 'user', content: 'message 3' }, // 其他用户的消息
2387
- ]);
2388
-
2389
- const result = await messageModel.hasMoreThanN(2);
2390
-
2391
- expect(result).toBe(false); // 当前用户只有 2 条消息,不大于 2
2392
- });
2393
-
2394
- it('should return false when no messages exist', async () => {
2395
- const result = await messageModel.hasMoreThanN(0);
2396
- expect(result).toBe(false);
2397
- });
2398
-
2399
- it('should handle edge cases', async () => {
2400
- // 创建一条消息
2401
- await serverDB
2402
- .insert(messages)
2403
- .values([{ id: '1', userId, role: 'user', content: 'message 1' }]);
2404
-
2405
- // 测试边界情况
2406
- const result1 = await messageModel.hasMoreThanN(0); // 1 > 0
2407
- const result2 = await messageModel.hasMoreThanN(1); // 1 ≯ 1
2408
- const result3 = await messageModel.hasMoreThanN(-1); // 1 > -1
2409
-
2410
- expect(result1).toBe(true);
2411
- expect(result2).toBe(false);
2412
- expect(result3).toBe(true);
2413
- });
2414
- });
2415
-
2416
- describe('createMessageQuery', () => {
2417
- it('should create a new message query', async () => {
2418
- // Create test data
2419
- await serverDB.insert(messages).values({
2420
- id: 'msg1',
2421
- userId,
2422
- role: 'user',
2423
- content: 'test message',
2424
- });
2425
-
2426
- // 调用 createMessageQuery 方法
2427
- const result = await messageModel.createMessageQuery({
2428
- messageId: 'msg1',
2429
- userQuery: 'original query',
2430
- rewriteQuery: 'rewritten query',
2431
- embeddingsId,
2432
- });
2433
-
2434
- // Assert result
2435
- expect(result).toBeDefined();
2436
- expect(result.id).toBeDefined();
2437
- expect(result.messageId).toBe('msg1');
2438
- expect(result.userQuery).toBe('original query');
2439
- expect(result.rewriteQuery).toBe('rewritten query');
2440
- expect(result.userId).toBe(userId);
2441
-
2442
- // 验证数据库中的记录
2443
- const dbResult = await serverDB
2444
- .select()
2445
- .from(messageQueries)
2446
- .where(eq(messageQueries.id, result.id));
2447
-
2448
- expect(dbResult).toHaveLength(1);
2449
- expect(dbResult[0].messageId).toBe('msg1');
2450
- expect(dbResult[0].userQuery).toBe('original query');
2451
- expect(dbResult[0].rewriteQuery).toBe('rewritten query');
2452
- });
2453
-
2454
- it('should create a message query with embeddings ID', async () => {
2455
- // Create test data
2456
- await serverDB.insert(messages).values({
2457
- id: 'msg2',
2458
- userId,
2459
- role: 'user',
2460
- content: 'test message',
2461
- });
2462
-
2463
- // 调用 createMessageQuery 方法
2464
- const result = await messageModel.createMessageQuery({
2465
- messageId: 'msg2',
2466
- userQuery: 'test query',
2467
- rewriteQuery: 'test rewritten query',
2468
- embeddingsId,
2469
- });
2470
-
2471
- // Assert result
2472
- expect(result).toBeDefined();
2473
- expect(result.embeddingsId).toBe(embeddingsId);
2474
-
2475
- // 验证数据库中的记录
2476
- const dbResult = await serverDB
2477
- .select()
2478
- .from(messageQueries)
2479
- .where(eq(messageQueries.id, result.id));
2480
-
2481
- expect(dbResult[0].embeddingsId).toBe(embeddingsId);
2482
- });
2483
-
2484
- it('should generate a unique ID for each message query', async () => {
2485
- // Create test data
2486
- await serverDB.insert(messages).values({
2487
- id: 'msg3',
2488
- userId,
2489
- role: 'user',
2490
- content: 'test message',
2491
- });
2492
-
2493
- // 连续创建两个消息查询
2494
- const result1 = await messageModel.createMessageQuery({
2495
- messageId: 'msg3',
2496
- userQuery: 'query 1',
2497
- rewriteQuery: 'rewritten query 1',
2498
- embeddingsId,
2499
- });
2500
-
2501
- const result2 = await messageModel.createMessageQuery({
2502
- messageId: 'msg3',
2503
- userQuery: 'query 2',
2504
- rewriteQuery: 'rewritten query 2',
2505
- embeddingsId,
2506
- });
2507
-
2508
- // Assert result
2509
- expect(result1.id).not.toBe(result2.id);
2510
- });
2511
- });
2512
-
2513
- describe('updateMessageRAG', () => {
2514
- it('should insert message query chunks for RAG', async () => {
2515
- // prepare message and query
2516
- const messageId = 'rag-msg-1';
2517
- const queryId = uuid();
2518
- const chunk1 = uuid();
2519
- const chunk2 = uuid();
2520
-
2521
- await serverDB.transaction(async (trx) => {
2522
- await trx.insert(messages).values({ id: messageId, role: 'user', userId, content: 'c' });
2523
- await trx.insert(chunks).values([
2524
- { id: chunk1, text: 'a' },
2525
- { id: chunk2, text: 'b' },
2526
- ]);
2527
- await trx
2528
- .insert(messageQueries)
2529
- .values({ id: queryId, messageId, userId, userQuery: 'q', rewriteQuery: 'rq' });
2530
- });
2531
-
2532
- await messageModel.updateMessageRAG(messageId, {
2533
- ragQueryId: queryId,
2534
- fileChunks: [
2535
- { id: chunk1, similarity: 0.9 },
2536
- { id: chunk2, similarity: 0.8 },
2537
- ],
2538
- });
2539
-
2540
- const rows = await serverDB
2541
- .select()
2542
- .from(messageQueryChunks)
2543
- .where(eq(messageQueryChunks.messageId, messageId));
2544
-
2545
- expect(rows).toHaveLength(2);
2546
- const s1 = rows.find((r) => r.chunkId === chunk1)!;
2547
- const s2 = rows.find((r) => r.chunkId === chunk2)!;
2548
- expect(s1.queryId).toBe(queryId);
2549
- expect(s1.similarity).toBe('0.90000');
2550
- expect(s2.similarity).toBe('0.80000');
2551
- });
2552
- });
2553
-
2554
- describe('deleteMessageQuery', () => {
2555
- it('should delete a message query by ID', async () => {
2556
- // Create test data
2557
- const queryId = uuid();
2558
- await serverDB.insert(messages).values({
2559
- id: 'msg4',
2560
- userId,
2561
- role: 'user',
2562
- content: 'test message',
2563
- });
2564
-
2565
- await serverDB.insert(messageQueries).values({
2566
- id: queryId,
2567
- messageId: 'msg4',
2568
- userQuery: 'test query',
2569
- rewriteQuery: 'rewritten query',
2570
- userId,
2571
- });
2572
-
2573
- // 验证查询已创建
2574
- const beforeDelete = await serverDB
2575
- .select()
2576
- .from(messageQueries)
2577
- .where(eq(messageQueries.id, queryId));
2578
-
2579
- expect(beforeDelete).toHaveLength(1);
2580
-
2581
- // 调用 deleteMessageQuery 方法
2582
- await messageModel.deleteMessageQuery(queryId);
2583
-
2584
- // 验证查询已删除
2585
- const afterDelete = await serverDB
2586
- .select()
2587
- .from(messageQueries)
2588
- .where(eq(messageQueries.id, queryId));
2589
-
2590
- expect(afterDelete).toHaveLength(0);
2591
- });
2592
-
2593
- it('should only delete message queries belonging to the user', async () => {
2594
- // Create test data - 其他用户的查询
2595
- const queryId = uuid();
2596
- await serverDB.insert(messages).values({
2597
- id: 'msg5',
2598
- userId: '456',
2599
- role: 'user',
2600
- content: 'test message',
2601
- });
2602
-
2603
- await serverDB.insert(messageQueries).values({
2604
- id: queryId,
2605
- messageId: 'msg5',
2606
- userQuery: 'test query',
2607
- rewriteQuery: 'rewritten query',
2608
- userId: '456', // 其他用户
2609
- });
2610
-
2611
- // 调用 deleteMessageQuery 方法
2612
- await messageModel.deleteMessageQuery(queryId);
2613
-
2614
- // 验证查询未被删除
2615
- const afterDelete = await serverDB
2616
- .select()
2617
- .from(messageQueries)
2618
- .where(eq(messageQueries.id, queryId));
2619
-
2620
- expect(afterDelete).toHaveLength(1);
2621
- });
2622
-
2623
- it('should throw error when deleting non-existent message query', async () => {
2624
- // 调用 deleteMessageQuery 方法删除不存在的查询
2625
- try {
2626
- await messageModel.deleteMessageQuery('non-existent-id');
2627
- } catch (e) {
2628
- expect(e).toBeInstanceOf(Error);
2629
- }
2630
- });
2631
- });
2632
- });