scorer-wasm.ts91 lines · main
| 1 | import { EvalScorer, Trace } from 'braintrust' |
| 2 | import { parse } from 'libpg-query' |
| 3 | |
| 4 | import { AssistantEvalInput, AssistantEvalOutput, Expected } from './scorer' |
| 5 | import { getParsedToolSpans } from './trace-utils' |
| 6 | import { executeSqlInputSchema } from '@/lib/ai/tools/studio-tools' |
| 7 | import { extractIdentifiers, isQuotedInSql, needsQuoting } from '@/lib/sql-identifier-quoting' |
| 8 | |
| 9 | /** Extracts SQL strings from all `execute_sql` tool spans in the trace. */ |
| 10 | async 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 | |
| 17 | export 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 | |
| 47 | export 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 | } |