feat: sql_similar assertion

apps/studio/evals/assistant.eval.ts

@@ -1,5 +1,6 @@
 import { openai } from '@ai-sdk/openai'
 import { Eval } from 'braintrust'
+import { stripIndent } from 'common-tags'
 import { Assertion, createAssertionScorer } from 'lib/ai/evals/scorer'
 import { generateAssistantResponse } from 'lib/ai/generate-assistant-response'
 import { getMockTools } from 'lib/ai/tools/mock-tools'
@@ -33,6 +34,15 @@ Eval('Assistant', {
           { type: 'llm_criteria_met', criteria: 'Response reflects there are RLS issues to fix.' },
         ],
       },
+      {
+        input: 'Create a new table "foods" with columns for "name" and "color"',
+        expected: [
+          {
+            type: 'sql_similar',
+            sql: stripIndent`CREATE TABLE IF NOT EXISTS public.foods ( id bigserial PRIMARY KEY, name text NOT NULL, color text );`,
+          },
+        ],
+      },
     ] satisfies Array<{ input: string; expected: EvalAssertion[] }>
   },
   task: async (input) => {
@@ -50,9 +60,20 @@ Eval('Assistant', {
       step.toolCalls.map((call) => call.toolName)
     )
 
+    const sqlQueries: string[] = []
+
+    for (const step of steps) {
+      for (const call of step.toolCalls) {
+        if (call.toolName === 'execute_sql') {
+          sqlQueries.push(call.input.sql)
+        }
+      }
+    }
+
     return {
       text: textLastStep,
       tools: toolNamesAllStepsFlattened,
+      sqlQueries,
     }
   },
   scores: [AssertionScorer],

apps/studio/lib/ai/evals/scorer.ts

@@ -1,11 +1,12 @@
 import { EvalScorer } from 'braintrust'
-import { ClosedQA } from 'autoevals'
+import { ClosedQA, Sql } from 'autoevals'
 
 type Input = string
 
 type Output = {
   text: string
   tools: string[]
+  sqlQueries?: string[]
 }
 
 export type Assertion<ToolName extends string = string> =
@@ -21,20 +22,26 @@ export type Assertion<ToolName extends string = string> =
       type: 'llm_criteria_met'
       criteria: string
     }
+  | {
+      type: 'sql_similar'
+      sql: string
+    }
 
 export function createAssertionScorer<ToolName extends string>(): EvalScorer<
   Input,
   Output,
   Assertion<ToolName>[]
 > {
-  return async ({ output, expected: assertions }) => {
+  return async ({ input, output, expected: assertions }) => {
     const assertionResults: {
       status: string
+      statusDetail?: string
       assertion: Assertion<ToolName>
     }[] = []
 
     for (const assertion of assertions) {
       let passedTest = false
+      let statusDetail: string | undefined
 
       try {
         switch (assertion.type) {
@@ -55,6 +62,29 @@ export function createAssertionScorer<ToolName extends string>(): EvalScorer<
             passedTest = closedQA.score !== null && closedQA.score > 0.5
             break
           }
+          case 'sql_similar': {
+            const sqlQueries = output.sqlQueries || []
+            if (sqlQueries.length === 0) {
+              passedTest = false
+              break
+            }
+            // Check if any of the generated SQL queries are similar to the expected SQL
+            const similarityScores = await Promise.all(
+              sqlQueries.map(async (generatedSql) => {
+                const sqlScore = await Sql({
+                  input,
+                  output: generatedSql,
+                  expected: assertion.sql,
+                })
+                return sqlScore.score ?? 0
+              })
+            )
+            passedTest = similarityScores.some((score) => score > 0.5)
+            statusDetail = similarityScores
+              .map((score, index) => `SQL ${index + 1}: ${score}`)
+              .join('\n')
+            break
+          }
           default:
             throw new Error(`Unknown assertion type`)
         }
@@ -64,6 +94,7 @@ export function createAssertionScorer<ToolName extends string>(): EvalScorer<
 
       assertionResults.push({
         status: passedTest ? 'passed' : 'failed',
+        statusDetail,
         assertion,
       })
     }