konnosDashboard
trace-utils.ts187 lines · main
1import type { SpanData, Trace } from 'braintrust'
2import { z } from 'zod'
3
4const projectContextPrefix = "The user's current project is "
5
6/**
7 * Matches AI SDK tool spans as Braintrust records them: tool args first,
8 * execution context second.
9 */
10const aiSdkToolSpanInputSchema = z.tuple([
11 z.unknown(),
12 z
13 .object({
14 messages: z.unknown().optional(),
15 toolCallId: z.string().optional(),
16 })
17 .passthrough(),
18])
19
20const threadTextBlockSchema = z.object({ type: z.literal('text'), text: z.string() })
21const threadToolCallBlockSchema = z.object({ type: z.literal('tool_call'), tool_name: z.string() })
22const threadContentBlockSchema = z.union([threadTextBlockSchema, threadToolCallBlockSchema])
23const threadContentSchema = z.union([
24 z.string(),
25 z.array(z.unknown()).transform((blocks) =>
26 blocks.flatMap((block) => {
27 const result = threadContentBlockSchema.safeParse(block)
28 return result.success ? [result.data] : []
29 })
30 ),
31])
32const threadMessageSchema = z.object({
33 role: z.enum(['system', 'user', 'assistant', 'tool']),
34 content: threadContentSchema,
35})
36
37type ThreadMessage = z.infer<typeof threadMessageSchema>
38
39/** Normalized Braintrust tool span with unwrapped tool input and raw output. */
40export type ToolSpan = {
41 span: SpanData
42 input: unknown
43 output: unknown
44}
45
46export type ThreadParts = {
47 projectContext: string | null
48 priorConversation: string | null
49 currentUserInput: string | null
50 lastAssistantTurn: string | null
51}
52
53/** Optional schemas used to validate and type a tool span's input and output. */
54type ToolSpanSchemas<
55 TInputSchema extends z.ZodType | undefined,
56 TOutputSchema extends z.ZodType | undefined,
57> = {
58 inputSchema?: TInputSchema
59 outputSchema?: TOutputSchema
60}
61
62/** Tool span whose input/output types are inferred from provided schemas. */
63type ParsedToolSpan<
64 TInputSchema extends z.ZodType | undefined,
65 TOutputSchema extends z.ZodType | undefined,
66> = {
67 span: SpanData
68 input: TInputSchema extends z.ZodType ? z.infer<TInputSchema> : unknown
69 output: TOutputSchema extends z.ZodType ? z.infer<TOutputSchema> : unknown
70}
71
72/** Extracts the actual tool args from Braintrust's traced function input shape. */
73function getToolSpanInput(span: SpanData): unknown {
74 const result = aiSdkToolSpanInputSchema.safeParse(span.input)
75 return result.success ? result.data[0] : span.input
76}
77
78function serializeMessageContent(message: ThreadMessage | undefined): string | null {
79 if (!message) return null
80 if (typeof message.content === 'string') return message.content || null
81
82 const content = message.content
83 .map((block) => (block.type === 'text' ? block.text : `[called ${block.tool_name}]`))
84 .join('\n')
85
86 return content || null
87}
88
89function serializeMessages(messages: ThreadMessage[]): string | null {
90 const parts = messages.flatMap((message) => {
91 const content = serializeMessageContent(message)
92 return content ? [`[${message.role}]\n${content}`] : []
93 })
94
95 return parts.length > 0 ? parts.join('\n\n') : null
96}
97
98function isProjectContextMessage(message: ThreadMessage): boolean {
99 return (
100 message.role === 'assistant' &&
101 Boolean(serializeMessageContent(message)?.startsWith(projectContextPrefix))
102 )
103}
104
105function findLastUserIndex(messages: ThreadMessage[]): number {
106 for (let i = messages.length - 1; i >= 0; i--) {
107 if (messages[i].role === 'user') return i
108 }
109 return -1
110}
111
112export function getThreadPartsFromThread(thread: unknown[]): ThreadParts {
113 const messages = thread.flatMap((message) => {
114 const result = threadMessageSchema.safeParse(message)
115 if (!result.success || result.data.role === 'system' || result.data.role === 'tool') return []
116 return [result.data]
117 })
118
119 const projectContextMessages = messages.filter(isProjectContextMessage)
120 const chatMessages = messages.filter((message) => !isProjectContextMessage(message))
121 const lastUserIdx = findLastUserIndex(chatMessages)
122 const projectContext = serializeMessageContent(
123 projectContextMessages[projectContextMessages.length - 1]
124 )
125
126 if (lastUserIdx === -1) {
127 return {
128 projectContext,
129 priorConversation: serializeMessages(chatMessages),
130 currentUserInput: null,
131 lastAssistantTurn: null,
132 }
133 }
134
135 return {
136 projectContext,
137 priorConversation: serializeMessages(chatMessages.slice(0, lastUserIdx)),
138 currentUserInput: serializeMessageContent(chatMessages[lastUserIdx]),
139 lastAssistantTurn: serializeMessages(
140 chatMessages.slice(lastUserIdx + 1).filter((message) => message.role === 'assistant')
141 ),
142 }
143}
144
145export async function getThreadParts(trace: Trace): Promise<ThreadParts> {
146 return getThreadPartsFromThread(await trace.getThread())
147}
148
149/** Returns normalized tool spans from the trace, optionally filtered to a specific tool name. */
150export async function getToolSpans(trace: Trace, toolName?: string): Promise<ToolSpan[]> {
151 const spans = await trace.getSpans({ spanType: ['tool'] })
152 const toolSpans = spans.map((span) => ({
153 span,
154 input: getToolSpanInput(span),
155 output: span.output,
156 }))
157 if (!toolName) return toolSpans
158 return toolSpans.filter((s) => s.span.span_attributes?.name === toolName)
159}
160
161/** Returns only tool spans whose normalized input/output match the provided schemas. */
162export async function getParsedToolSpans<
163 TInputSchema extends z.ZodType | undefined = undefined,
164 TOutputSchema extends z.ZodType | undefined = undefined,
165>(
166 trace: Trace,
167 toolName: string,
168 schemas: ToolSpanSchemas<TInputSchema, TOutputSchema> = {}
169): Promise<Array<ParsedToolSpan<TInputSchema, TOutputSchema>>> {
170 const spans = await getToolSpans(trace, toolName)
171
172 return spans.flatMap(({ span, input, output }) => {
173 const parsedInput = schemas.inputSchema?.safeParse(input)
174 if (parsedInput && !parsedInput.success) return []
175
176 const parsedOutput = schemas.outputSchema?.safeParse(output)
177 if (parsedOutput && !parsedOutput.success) return []
178
179 return [
180 {
181 span,
182 input: parsedInput ? parsedInput.data : input,
183 output: parsedOutput ? parsedOutput.data : output,
184 } as ParsedToolSpan<TInputSchema, TOutputSchema>,
185 ]
186 })
187}