From 4a04156dffc7d64203d463c7dd3db44ceebc4f20 Mon Sep 17 00:00:00 2001
From: Matt Rossman <22670878+mattrossman@users.noreply.github.com>
Date: Thu, 11 Dec 2025 12:12:08 -0500
Subject: [PATCH] refactor: generic tool name type in assertions

---
 apps/studio/evals/assistant.eval.ts |  9 ++-
 apps/studio/lib/ai/evals/scorer.ts  | 88 ++++++++++++++---------------
 2 files changed, 50 insertions(+), 47 deletions(-)

diff --git a/apps/studio/evals/assistant.eval.ts b/apps/studio/evals/assistant.eval.ts
index c48fe4a8ae8..bc4a5bae8fd 100644
--- a/apps/studio/evals/assistant.eval.ts
+++ b/apps/studio/evals/assistant.eval.ts
@@ -1,6 +1,6 @@
 import { openai } from '@ai-sdk/openai'
 import { Eval } from 'braintrust'
-import { Assertion, AssertionScorer } from 'lib/ai/evals/scorer'
+import { Assertion, createAssertionScorer } from 'lib/ai/evals/scorer'
 import { generateAssistantResponse } from 'lib/ai/generate-assistant-response'
 import { getMockTools } from 'lib/ai/tools/mock-tools'
 import assert from 'node:assert'
@@ -8,6 +8,11 @@ import assert from 'node:assert'
 assert(process.env.BRAINTRUST_PROJECT_ID, 'BRAINTRUST_PROJECT_ID is not set')
 assert(process.env.OPENAI_API_KEY, 'OPENAI_API_KEY is not set')
 
+type MockToolName = keyof Awaited<ReturnType<typeof getMockTools>>
+type EvalAssertion = Assertion<MockToolName>
+
+const AssertionScorer = createAssertionScorer<MockToolName>()
+
 Eval('Assistant', {
   projectId: process.env.BRAINTRUST_PROJECT_ID,
   data: () => {
@@ -20,7 +25,7 @@ Eval('Assistant', {
         input: 'How do I implement IP address rate limiting? Use search_docs.',
         expected: [{ type: 'tools_include', tool: 'search_docs' }],
       },
-    ] satisfies Array<{ input: string; expected: Assertion[] }>
+    ] satisfies Array<{ input: string; expected: EvalAssertion[] }>
   },
   task: async (input) => {
     const result = await generateAssistantResponse({
diff --git a/apps/studio/lib/ai/evals/scorer.ts b/apps/studio/lib/ai/evals/scorer.ts
index 2f197baa3ab..59ea7d423f8 100644
--- a/apps/studio/lib/ai/evals/scorer.ts
+++ b/apps/studio/lib/ai/evals/scorer.ts
@@ -1,5 +1,4 @@
 import { EvalScorer } from 'braintrust'
-import type { getMockTools } from 'lib/ai/tools/mock-tools'
 
 type Input = string
 
@@ -8,64 +7,63 @@ type Output = {
   tools: string[]
 }
 
-type MockToolName = keyof Awaited<ReturnType<typeof getMockTools>>
-
-export type Assertion =
+export type Assertion<ToolName extends string = string> =
   | {
       type: 'tools_include'
-      tool: MockToolName
+      tool: ToolName
     }
   | {
       type: 'text_includes'
       substring: string
     }
 
-type Expected = Assertion[]
+export function createAssertionScorer<ToolName extends string>(): EvalScorer<
+  Input,
+  Output,
+  Assertion<ToolName>[]
+> {
+  return ({ output, expected: assertions }) => {
+    const assertionResults: {
+      status: string
+      assertion: Assertion<ToolName>
+    }[] = []
 
-export const AssertionScorer: EvalScorer<Input, Output, Expected> = ({
-  output,
-  expected: assertions,
-}) => {
-  const assertionResults: {
-    status: string
-    assertion: Assertion
-  }[] = []
+    let passedTest = false
 
-  let passedTest = false
-
-  for (const assertion of assertions) {
-    try {
-      switch (assertion.type) {
-        case 'tools_include': {
-          passedTest = output.tools.includes(assertion.tool)
-          break
+    for (const assertion of assertions) {
+      try {
+        switch (assertion.type) {
+          case 'tools_include': {
+            passedTest = output.tools.includes(assertion.tool)
+            break
+          }
+          case 'text_includes': {
+            passedTest = output.text.includes(assertion.substring)
+            break
+          }
+          default:
+            throw new Error(`Unknown assertion type`)
         }
-        case 'text_includes': {
-          passedTest = output.text.includes(assertion.substring)
-          break
-        }
-        default:
-          throw new Error(`Unknown assertion type`)
+      } catch (e) {
+        passedTest = false
       }
-    } catch (e) {
-      passedTest = false
+
+      assertionResults.push({
+        status: passedTest ? 'passed' : 'failed',
+        assertion,
+      })
     }
 
-    assertionResults.push({
-      status: passedTest ? 'passed' : 'failed',
-      assertion,
-    })
-  }
+    const passedCount = assertionResults.filter((r) => r.status === 'passed').length
+    const totalCount = assertionResults.length
+    const ratioPassed = totalCount === 0 ? 1 : passedCount / totalCount
 
-  const passedCount = assertionResults.filter((r) => r.status === 'passed').length
-  const totalCount = assertionResults.length
-  const ratioPassed = totalCount === 0 ? 1 : passedCount / totalCount
-
-  return {
-    name: 'Assertions Score',
-    score: ratioPassed,
-    metadata: {
-      assertionResults,
-    },
+    return {
+      name: 'Assertions Score',
+      score: ratioPassed,
+      metadata: {
+        assertionResults,
+      },
+    }
   }
 }