konnosDashboard
message-utils.test.ts153 lines · main
1import type { UIMessage } from 'ai'
2import { describe, expect, it } from 'vitest'
3
4import { prepareMessagesForAPI } from './message-utils'
5
6describe('prepareMessagesForAPI', () => {
7 it('should limit messages to last 7 entries', () => {
8 const messages: UIMessage[] = Array.from({ length: 10 }, (_, i) => ({
9 id: `msg-${i}`,
10 role: 'user',
11 parts: [{ type: 'text', text: `Message ${i}` }],
12 }))
13
14 const result = prepareMessagesForAPI(messages)
15
16 expect(result).toHaveLength(7)
17 expect(result[0].parts[0]).toEqual({ type: 'text', text: 'Message 3' })
18 expect(result[6].parts[0]).toEqual({ type: 'text', text: 'Message 9' })
19 })
20
21 it('should remove results property from assistant messages', () => {
22 const messages = [
23 {
24 id: 'msg-1',
25 role: 'assistant',
26 parts: [{ type: 'text', text: 'Response' }],
27 results: { data: 'some data' },
28 },
29 ] as Array<UIMessage & { results?: unknown }>
30
31 const result = prepareMessagesForAPI(messages)
32
33 expect(result).toHaveLength(1)
34 expect(result[0]).not.toHaveProperty('results')
35 expect(result[0].parts[0]).toEqual({ type: 'text', text: 'Response' })
36 })
37
38 it('should preserve messages without results', () => {
39 const messages: UIMessage[] = [
40 {
41 id: 'msg-1',
42 role: 'user',
43 parts: [{ type: 'text', text: 'Question' }],
44 },
45 {
46 id: 'msg-2',
47 role: 'assistant',
48 parts: [{ type: 'text', text: 'Answer' }],
49 },
50 ]
51
52 const result = prepareMessagesForAPI(messages)
53
54 expect(result).toHaveLength(2)
55 expect(result[0]).toEqual(messages[0])
56 expect(result[1]).toEqual(messages[1])
57 })
58
59 it('should handle empty array', () => {
60 const messages: UIMessage[] = []
61
62 const result = prepareMessagesForAPI(messages)
63
64 expect(result).toHaveLength(0)
65 expect(result).toEqual([])
66 })
67
68 it('should handle array with fewer than 7 messages', () => {
69 const messages: UIMessage[] = [
70 { id: 'msg-1', role: 'user', parts: [{ type: 'text', text: 'Message 1' }] },
71 { id: 'msg-2', role: 'assistant', parts: [{ type: 'text', text: 'Message 2' }] },
72 { id: 'msg-3', role: 'user', parts: [{ type: 'text', text: 'Message 3' }] },
73 ]
74
75 const result = prepareMessagesForAPI(messages)
76
77 expect(result).toHaveLength(3)
78 expect(result).toEqual(messages)
79 })
80
81 it('should handle array with exactly 7 messages', () => {
82 const messages: UIMessage[] = Array.from({ length: 7 }, (_, i) => ({
83 id: `msg-${i}`,
84 role: i % 2 === 0 ? 'user' : 'assistant',
85 parts: [{ type: 'text', text: `Message ${i}` }],
86 }))
87
88 const result = prepareMessagesForAPI(messages)
89
90 expect(result).toHaveLength(7)
91 expect(result).toEqual(messages)
92 })
93
94 it('should only remove results from assistant messages, not user messages', () => {
95 const messages = [
96 {
97 id: 'msg-1',
98 role: 'user',
99 parts: [{ type: 'text', text: 'Question' }],
100 results: { data: 'user data' },
101 },
102 {
103 id: 'msg-2',
104 role: 'assistant',
105 parts: [{ type: 'text', text: 'Answer' }],
106 results: { data: 'assistant data' },
107 },
108 ] as Array<UIMessage & { results?: unknown }>
109
110 const result = prepareMessagesForAPI(messages)
111
112 expect(result).toHaveLength(2)
113 // User message keeps results (not removed by the function)
114 expect((result[0] as any).results).toEqual({ data: 'user data' })
115 // Assistant message has results removed
116 expect(result[1]).not.toHaveProperty('results')
117 })
118
119 it('should handle mixed messages with and without results', () => {
120 const messages = [
121 {
122 id: 'msg-1',
123 role: 'assistant',
124 parts: [{ type: 'text', text: 'First' }],
125 results: { data: 'data1' },
126 },
127 {
128 id: 'msg-2',
129 role: 'user',
130 parts: [{ type: 'text', text: 'Second' }],
131 },
132 {
133 id: 'msg-3',
134 role: 'assistant',
135 parts: [{ type: 'text', text: 'Third' }],
136 },
137 {
138 id: 'msg-4',
139 role: 'assistant',
140 parts: [{ type: 'text', text: 'Fourth' }],
141 results: { data: 'data2' },
142 },
143 ] as Array<UIMessage & { results?: unknown }>
144
145 const result = prepareMessagesForAPI(messages)
146
147 expect(result).toHaveLength(4)
148 expect(result[0]).not.toHaveProperty('results')
149 expect(result[1]).toEqual(messages[1])
150 expect(result[2]).toEqual(messages[2])
151 expect(result[3]).not.toHaveProperty('results')
152 })
153})