trace-utils.test.ts159 lines · main
1import { describe, expect, it } from 'vitest'
2
3import { getThreadPartsFromThread } from './trace-utils'
4
5// Sanitized mock of the thread shape returned by trace.getThread().
6const MOCK_THREAD = [
7 {
8 role: 'system',
9 content: 'System instructions omitted for fixture.',
10 },
11 {
12 role: 'assistant',
13 content: "The user's current project is Acme Analytics.",
14 },
15 {
16 role: 'user',
17 content: 'What did we decide earlier?',
18 },
19 {
20 role: 'assistant',
21 content: [
22 {
23 type: 'text',
24 text: 'We decided to add an orders table with RLS policies before generating sample data.',
25 },
26 ],
27 },
28 {
29 role: 'user',
30 content: 'Can you create that orders table now?',
31 },
32 {
33 role: 'assistant',
34 id: null,
35 content: [
36 {
37 type: 'tool_call',
38 tool_name: 'rename_chat',
39 tool_call_id: 'call_dummy_rename',
40 arguments: {
41 type: 'valid',
42 value: {
43 newName: 'Create Orders Table',
44 },
45 },
46 },
47 ],
48 },
49 {
50 role: 'tool',
51 content: [
52 {
53 type: 'tool_result',
54 tool_name: 'rename_chat',
55 tool_call_id: 'call_dummy_rename',
56 output: {
57 status: 'Chat request sent to client',
58 },
59 },
60 ],
61 },
62 {
63 role: 'assistant',
64 id: null,
65 content: [
66 {
67 type: 'tool_call',
68 tool_name: 'load_knowledge',
69 tool_call_id: 'call_dummy_knowledge',
70 arguments: {
71 type: 'valid',
72 value: {
73 name: 'database',
74 },
75 },
76 },
77 {
78 type: 'tool_call',
79 tool_name: 'execute_sql',
80 tool_call_id: 'call_dummy_sql',
81 arguments: {
82 type: 'valid',
83 value: {
84 sql: 'create table public.orders (id bigint generated by default as identity primary key);',
85 },
86 },
87 },
88 ],
89 },
90 {
91 role: 'tool',
92 content: [
93 {
94 type: 'tool_result',
95 tool_name: 'load_knowledge',
96 tool_call_id: 'call_dummy_knowledge',
97 output: 'Knowledge fixture omitted.',
98 },
99 {
100 type: 'tool_result',
101 tool_name: 'execute_sql',
102 tool_call_id: 'call_dummy_sql',
103 output: {
104 type: 'text',
105 text: 'SQL executed successfully.',
106 },
107 },
108 ],
109 },
110 {
111 role: 'assistant',
112 id: null,
113 content:
114 'I created the public.orders table. You should add RLS policies before exposing it to users.',
115 },
116]
117
118describe('getThreadPartsFromThread', () => {
119 it('parses a sanitized Braintrust trace.getThread payload', () => {
120 expect(getThreadPartsFromThread(MOCK_THREAD)).toEqual({
121 projectContext: "The user's current project is Acme Analytics.",
122 priorConversation:
123 '[user]\nWhat did we decide earlier?\n\n[assistant]\nWe decided to add an orders table with RLS policies before generating sample data.',
124 currentUserInput: 'Can you create that orders table now?',
125 lastAssistantTurn:
126 '[assistant]\n[called rename_chat]\n\n[assistant]\n[called load_knowledge]\n[called execute_sql]\n\n[assistant]\nI created the public.orders table. You should add RLS policies before exposing it to users.',
127 })
128 })
129
130 it('uses the most recent project context message', () => {
131 expect(
132 getThreadPartsFromThread([
133 {
134 role: 'assistant',
135 content: "The user's current project is Old Project.",
136 },
137 ...MOCK_THREAD,
138 ])
139 ).toMatchObject({
140 projectContext: "The user's current project is Acme Analytics.",
141 })
142 })
143
144 it('returns prior conversation without current turn parts when there is no user message', () => {
145 expect(
146 getThreadPartsFromThread([
147 {
148 role: 'assistant',
149 content: 'I can help with your Briven project.',
150 },
151 ])
152 ).toEqual({
153 projectContext: null,
154 priorConversation: '[assistant]\nI can help with your Briven project.',
155 currentUserInput: null,
156 lastAssistantTurn: null,
157 })
158 })
159})