supabase/apps/studio/evals/assistant.eval.ts

import { openai } from '@ai-sdk/openai'
import { Eval } from 'braintrust'
import { generateAssistantResponse } from 'lib/ai/generate-assistant-response'
import { getMockTools } from 'lib/ai/tools/mock-tools'
import assert from 'node:assert'
import { dataset } from './dataset'
import {
  completenessScorer,
  concisenessScorer,
  correctnessScorer,
  docsFaithfulnessScorer,
  goalCompletionScorer,
  sqlIdentifierQuotingScorer,
  sqlSyntaxScorer,
  toolUsageScorer,
  urlValidityScorer,
} from './scorer'
import { ToolSet, TypedToolCall, TypedToolResult } from 'ai'

assert(process.env.BRAINTRUST_PROJECT_ID, 'BRAINTRUST_PROJECT_ID is not set')
assert(process.env.OPENAI_API_KEY, 'OPENAI_API_KEY is not set')

Eval('Assistant', {
  projectId: process.env.BRAINTRUST_PROJECT_ID,
  trialCount: process.env.CI ? 3 : 1,
  data: () => dataset,
  task: async (input) => {
    const result = await generateAssistantResponse({
      model: openai('gpt-5-mini'),
      messages: [{ id: '1', role: 'user', parts: [{ type: 'text', text: input.prompt }] }],
      tools: await getMockTools(input.mockTables ? { list_tables: input.mockTables } : undefined),
    })

    const finishReason = await result.finishReason

    // `result.toolCalls` only shows the last step, instead aggregate tools across all steps
    const steps = await result.steps

    const simplifiedSteps = steps.map((step) => ({
      text: step.text,
      toolCalls: step.toolCalls.map((call) => ({
        toolName: call.toolName,
        input: call.input,
      })),
    }))

    const toolNames: string[] = []
    const sqlQueries: string[] = []
    const docs: string[] = []

    for (const step of steps) {
      for (const [i, toolCall] of step.toolCalls.entries()) {
        toolNames.push(toolCall.toolName)

        const toolResult = step.toolResults.at(i)
        if (!toolResult) {
          continue
        }

        const parsed = parseToolCall(toolCall, toolResult)

        if (parsed.sqlQuery) {
          sqlQueries.push(parsed.sqlQuery)
        }
        if (parsed.docs) {
          docs.push(...parsed.docs)
        }
      }
    }

    return {
      finishReason,
      steps: simplifiedSteps,
      toolNames,
      sqlQueries,
      docs,
    }
  },
  scores: [
    toolUsageScorer,
    sqlSyntaxScorer,
    sqlIdentifierQuotingScorer,
    goalCompletionScorer,
    concisenessScorer,
    completenessScorer,
    docsFaithfulnessScorer,
    correctnessScorer,
    urlValidityScorer,
  ],
})

type ParsedToolCall = {
  /** Query generated by `execute_sql` */
  sqlQuery?: string

  /** Docs text pulled in from `search_docs` */
  docs?: string[]
}

/**
 * Validate and extract relevant info from a tool call/result
 */
function parseToolCall(
  toolCall: TypedToolCall<ToolSet>,
  toolResult: TypedToolResult<ToolSet>
): ParsedToolCall {
  switch (toolCall.toolName) {
    case 'execute_sql': {
      const sqlQuery = toolCall.input.sql
      if (typeof sqlQuery !== 'string') {
        return {}
      }

      return { sqlQuery }
    }
    case 'search_docs': {
      const content = toolResult.output.content
      if (!content || !Array.isArray(content)) {
        return {}
      }

      const docs = content.map((item) => item?.text).filter((text) => typeof text === 'string')
      if (docs.length === 0) {
        return {}
      }

      return { docs }
    }
  }

  return {}
}