scorer-wasm.ts91 lines · main
1import { EvalScorer, Trace } from 'braintrust'
2import { parse } from 'libpg-query'
3
4import { AssistantEvalInput, AssistantEvalOutput, Expected } from './scorer'
5import { getParsedToolSpans } from './trace-utils'
6import { executeSqlInputSchema } from '@/lib/ai/tools/studio-tools'
7import { extractIdentifiers, isQuotedInSql, needsQuoting } from '@/lib/sql-identifier-quoting'
8
9/** Extracts SQL strings from all `execute_sql` tool spans in the trace. */
10async function getSqlQueries(trace: Trace): Promise<string[]> {
11 const spans = await getParsedToolSpans(trace, 'execute_sql', {
12 inputSchema: executeSqlInputSchema,
13 })
14 return spans.map((s) => s.input.sql)
15}
16
17export const sqlSyntaxScorer: EvalScorer<
18 AssistantEvalInput,
19 AssistantEvalOutput,
20 Expected
21> = async ({ trace }) => {
22 if (!trace) return null
23
24 const sqlQueries = await getSqlQueries(trace)
25 if (sqlQueries.length === 0) return null
26
27 const errors: string[] = []
28 let validQueries = 0
29
30 for (const sql of sqlQueries) {
31 try {
32 await parse(sql)
33 validQueries++
34 } catch (error) {
35 const errorMessage = error instanceof Error ? error.message : String(error)
36 errors.push(`SQL syntax error: ${errorMessage}`)
37 }
38 }
39
40 return {
41 name: 'SQL Validity',
42 score: validQueries / sqlQueries.length,
43 metadata: errors.length > 0 ? { errors } : undefined,
44 }
45}
46
47export const sqlIdentifierQuotingScorer: EvalScorer<
48 AssistantEvalInput,
49 AssistantEvalOutput,
50 Expected
51> = async ({ trace }) => {
52 if (!trace) return null
53
54 const sqlQueries = await getSqlQueries(trace)
55 if (sqlQueries.length === 0) return null
56
57 const errors: string[] = []
58 let totalNeedingQuotes = 0
59 let properlyQuoted = 0
60
61 for (const sql of sqlQueries) {
62 try {
63 const ast = await parse(sql)
64 const identifiers = extractIdentifiers(ast)
65
66 for (const identifier of identifiers) {
67 if (needsQuoting(identifier)) {
68 totalNeedingQuotes++
69 if (isQuotedInSql(sql, identifier)) {
70 properlyQuoted++
71 } else {
72 const sqlPreview = sql.length > 100 ? `${sql.substring(0, 100)}...` : sql
73 errors.push(
74 `Identifier "${identifier}" needs quoting but is not quoted in: ${sqlPreview}`
75 )
76 }
77 }
78 }
79 } catch {
80 // Skip invalid SQL - already handled by sqlSyntaxScorer
81 }
82 }
83
84 const score = totalNeedingQuotes === 0 ? 1 : properlyQuoted / totalNeedingQuotes
85
86 return {
87 name: 'SQL Identifier Quoting',
88 score,
89 metadata: errors.length > 0 ? { errors } : undefined,
90 }
91}