konnosDashboard
test-fixtures.ts114 lines · main
1import type { ToolUIPart, UIMessage } from 'ai'
2
3export function createUserMessage(content: string, id = 'user-msg-1'): UIMessage {
4 return {
5 id,
6 role: 'user',
7 parts: [
8 {
9 type: 'text',
10 text: content,
11 },
12 ],
13 }
14}
15
16export function createAssistantTextMessage(content: string, id = 'assistant-msg-1'): UIMessage {
17 return {
18 id,
19 role: 'assistant',
20 parts: [
21 {
22 type: 'text',
23 text: content,
24 },
25 ],
26 }
27}
28
29export function createAssistantMessageWithExecuteSqlTool(
30 query: string,
31 results: Array<Record<string, any>> = [{ id: 1, name: 'test' }],
32 id = 'assistant-tool-msg-1'
33): UIMessage {
34 return {
35 id,
36 role: 'assistant',
37 parts: [
38 {
39 type: 'text',
40 text: "I'll run that SQL query for you.",
41 },
42 {
43 type: 'tool-execute_sql',
44 state: 'output-available',
45 toolCallId: 'call-123',
46 input: { sql: query },
47 output: results,
48 } satisfies ToolUIPart,
49 ],
50 }
51}
52
53export function createAssistantMessageWithMultipleTools(
54 id = 'assistant-multi-tool-msg-1'
55): UIMessage {
56 return {
57 id,
58 role: 'assistant',
59 parts: [
60 {
61 type: 'text',
62 text: 'Let me check the database structure and run some queries.',
63 },
64 {
65 type: 'tool-execute_sql',
66 state: 'output-available',
67 toolCallId: 'call-456',
68 input: { sql: 'SELECT * FROM users LIMIT 5' },
69 output: [
70 { id: 1, email: 'user1@example.com' },
71 { id: 2, email: 'user2@example.com' },
72 ],
73 } satisfies ToolUIPart,
74 {
75 type: 'tool-execute_sql',
76 state: 'output-available',
77 toolCallId: 'call-789',
78 toolName: 'execute_sql',
79 input: { sql: 'DESCRIBE users' },
80 output: [
81 { column: 'id', type: 'integer', nullable: false },
82 { column: 'email', type: 'varchar', nullable: false },
83 ],
84 } as ToolUIPart,
85 ],
86 }
87}
88
89export function createLongConversation(): Array<UIMessage> {
90 return [
91 createUserMessage('Show me all users', 'msg-1'),
92 createAssistantMessageWithExecuteSqlTool('SELECT * FROM users', [{ id: 1 }], 'msg-2'),
93 createUserMessage('How many users are there?', 'msg-3'),
94 createAssistantMessageWithExecuteSqlTool(
95 'SELECT COUNT(*) FROM users',
96 [{ count: 100 }],
97 'msg-4'
98 ),
99 createUserMessage('Show me the schema', 'msg-5'),
100 createAssistantTextMessage("Here's the database schema...", 'msg-6'),
101 createUserMessage('Create a new table', 'msg-7'),
102 createAssistantMessageWithExecuteSqlTool(
103 'CREATE TABLE posts (id SERIAL PRIMARY KEY)',
104 [],
105 'msg-8'
106 ),
107 createUserMessage('Add some data', 'msg-9'),
108 createAssistantMessageWithExecuteSqlTool(
109 "INSERT INTO posts (title) VALUES ('Test')",
110 [],
111 'msg-10'
112 ),
113 ]
114}