@botpress/vai 0.0.1-beta.1

package/src/assertions/extension.ts

@@ -0,0 +1,51 @@
+import json5 from 'json5'
+import { expect } from 'vitest'
+import { getCurrentTest } from 'vitest/suite'
+
+import { asyncExpect } from '../utils/asyncAssertion'
+import { Output } from '../utils/predictJson'
+
+export type ExtendedPromise<T> = PromiseLike<Output<T>> & {
+  value: PromiseLike<T>
+}
+
+export const toAssertion = <T>(promise: Promise<Output<T>>): ExtendedPromise<T> => {
+  return {
+    then: promise.then.bind(promise),
+    value: promise.then((value) => value.result)
+  }
+}
+
+export const makeToMatchInlineSnapshot =
+  <T>(promise: Promise<Output<T>>) =>
+  async (expected?: string) => {
+    const stack = new Error().stack!.split('\n')[2]
+    const newStack = `
+at __INLINE_SNAPSHOT__ (node:internal/process/task_queues:1:1)
+at randomLine (node:internal/process/task_queues:1:1)
+${stack}
+`.trim()
+
+    const obj = json5.parse(expected ?? '""')
+    const expectation = asyncExpect(promise, (expect) => expect.toMatchObject(obj)).catch(() => {
+      // we swallow the error here, as we're going to throw a new one with the correct stack
+      // this is just to make vitest happy and show a nice error message
+    })
+
+    try {
+      expect((await promise).result).toMatchObject(obj)
+    } catch (err) {
+      const newError = new Error()
+      newError.stack = newStack
+
+      expect.getState().snapshotState.match({
+        isInline: true,
+        received: (await promise).result,
+        testName: getCurrentTest()!.name,
+        error: newError,
+        inlineSnapshot: expected
+      })
+    }
+
+    return expectation
+  }

package/src/assertions/extract.ts

@@ -0,0 +1,39 @@
+import { z } from '@botpress/sdk'
+
+import { Context } from '../context'
+import { asyncExpect } from '../utils/asyncAssertion'
+import { Input, predictJson } from '../utils/predictJson'
+import { makeToMatchInlineSnapshot, toAssertion } from './extension'
+
+export type ExtractOptions<T, S> = {
+  description?: string
+  examples?: { value: T; extracted: S; reason: string }[]
+}
+
+export function extract<T extends Input, S extends z.AnyZodObject>(
+  value: T,
+  shape: S,
+  options?: ExtractOptions<T, z.infer<S>>
+) {
+  const additionalMessage = options?.description
+    ? `\nIn order to extract the right information, follow these instructions:\n${options.description}\n`
+    : ''
+  const promise = predictJson({
+    systemMessage:
+      'From the given input, extract the required information into the requested format.' + additionalMessage.trim(),
+    examples: options?.examples?.map(({ value, reason, extracted }) => ({
+      input: value,
+      output: { reason, result: extracted }
+    })),
+    outputSchema: shape,
+    model: Context.evaluatorModel,
+    input: value
+  })
+
+  return {
+    ...toAssertion(promise),
+    toBe: (expected: z.infer<S>) => asyncExpect(promise, (expect) => expect.toEqual(expected)),
+    toMatchObject: (expected: Partial<z.infer<S>>) => asyncExpect(promise, (expect) => expect.toMatchObject(expected)),
+    toMatchInlineSnapshot: makeToMatchInlineSnapshot(promise)
+  }
+}

package/src/assertions/filter.ts

@@ -0,0 +1,86 @@
+import { literal, z } from '@botpress/sdk'
+
+import { Context } from '../context'
+import { asyncExpect } from '../utils/asyncAssertion'
+import { predictJson } from '../utils/predictJson'
+import { makeToMatchInlineSnapshot, toAssertion } from './extension'
+
+export type FilterOptions<T> = {
+  examples?: { value: T; reason: string; keep: boolean }[]
+}
+
+export function filter<U>(values: U[], condition: string, options?: FilterOptions<U>) {
+  const mappedValues = values.map((_, idx) =>
+    z.object({
+      index: literal(idx),
+      reason: z.string(),
+      keep: z.boolean()
+    })
+  )
+
+  const input = values.map((value, idx) => ({
+    index: idx,
+    value
+  }))
+
+  const schema = z
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    .tuple(mappedValues as any)
+    .describe(
+      'An array of the objects with the index and a boolean value indicating if the object should be kept or not'
+    )
+
+  const promise = predictJson({
+    systemMessage: `
+Based on the following qualification criteria, you need to filter the given list of objects.
+Citeria: ${condition}
+
+---
+You need to return an array of objects with the index and a boolean value indicating if the object should be kept or not.
+`.trim(),
+    examples: options?.examples
+      ? [
+          {
+            input: options?.examples?.map((v, index) => ({
+              index,
+              value: v.value
+            })),
+            output: {
+              reason: 'Here are some examples',
+              result: options?.examples?.map((v, idx) => ({
+                index: idx,
+                reason: v.reason,
+                keep: v.keep
+              }))
+            }
+          }
+        ]
+      : undefined,
+    outputSchema: schema,
+    model: Context.evaluatorModel,
+    input
+  }).then((x) => {
+    const results = schema.parse(x.result) as { index: number; keep: boolean }[]
+    return {
+      result: values.filter((_, idx) => results.find((r) => r.index === idx)?.keep),
+      reason: x.reason
+    }
+  })
+
+  return {
+    ...toAssertion(promise),
+    toBe: (expected: U[]) => asyncExpect(promise, (expect) => expect.toEqual(expected)),
+    toMatchInlineSnapshot: makeToMatchInlineSnapshot(promise),
+    toHaveNoneFiltered: () => asyncExpect(promise, (expect) => expect.toEqual(values)),
+    toHaveSomeFiltered: () => asyncExpect(promise, (expect) => expect.not.toEqual(values)),
+    toBeEmpty: () => asyncExpect(promise, (expect) => expect.toHaveLength(0)),
+    length: {
+      toBe: (expected: number) => asyncExpect(promise, (expect) => expect.toHaveLength(expected)),
+      toBeGreaterThanOrEqual: (expected: number) =>
+        asyncExpect(promise, (expect) => expect.length.greaterThanOrEqual(expected)),
+      toBeLessThanOrEqual: (expected: number) =>
+        asyncExpect(promise, (expect) => expect.length.lessThanOrEqual(expected)),
+      toBeBetween: (min: number, max: number) => asyncExpect(promise, (expect) => expect.length.within(min, max))
+    }
+  }
+}

package/src/assertions/rate.ts

@@ -0,0 +1,40 @@
+import { z } from '@botpress/sdk'
+
+import { Context } from '../context'
+import { asyncExpect } from '../utils/asyncAssertion'
+import { Input, predictJson } from '../utils/predictJson'
+import { makeToMatchInlineSnapshot, toAssertion } from './extension'
+
+export type RatingScore = 1 | 2 | 3 | 4 | 5
+export type RateOptions<T> = {
+  examples?: { value: T; rating: number; reason: string }[]
+}
+
+export function rate<T extends Input>(value: T, condition: string, options?: RateOptions<T>) {
+  const schema = z.number().min(1).max(5).describe('Rating score, higher is better (1 is the worst, 5 is the best)')
+  const promise = predictJson({
+    systemMessage: `Based on the following qualification criteria, you need to rate the given situation from a score of 1 to 5.\nScoring: 1 is the worst score, 5 is the best score possible.\nCriteria: ${condition}`,
+    examples: options?.examples?.map(({ value, reason, rating }) => ({
+      input: value,
+      output: { reason, result: rating }
+    })),
+    outputSchema: schema,
+    model: Context.evaluatorModel,
+    input: value
+  }).then((x) => {
+    return {
+      result: typeof x.result === 'number' ? x.result : parseInt(x.result, 10),
+      reason: x.reason
+    }
+  })
+
+  return {
+    ...toAssertion(promise),
+    toBe: (expected: number) => asyncExpect(promise, (expect) => expect.toEqual(expected)),
+    toMatchInlineSnapshot: makeToMatchInlineSnapshot(promise),
+    toBeGreaterThanOrEqual: (expected: RatingScore) =>
+      asyncExpect(promise, (expect) => expect.toBeGreaterThanOrEqual(expected)),
+    toBeLessThanOrEqual: (expected: RatingScore) =>
+      asyncExpect(promise, (expect) => expect.toBeLessThanOrEqual(expected))
+  }
+}

package/src/context.ts

@@ -0,0 +1,65 @@
+import type { Client } from '@botpress/client'
+import { onTestFinished } from 'vitest'
+import { getCurrentTest } from 'vitest/suite'
+import { Models } from './models'
+
+export type EvaluatorModel = (typeof Models)[number]['id']
+
+export type TestMetadata = {
+  isVaiTest: boolean
+  scenario?: string
+  evaluatorModel?: EvaluatorModel
+}
+
+const getTestMetadata = (): TestMetadata => {
+  const test = getCurrentTest()
+  return (test?.meta ?? {
+    isVaiTest: false
+  }) as TestMetadata
+}
+
+class VaiContext {
+  #client: Client | null = null
+  #wrapError = false
+
+  get wrapError() {
+    return this.#wrapError
+  }
+
+  get client() {
+    if (!this.#client) {
+      throw new Error('Botpress client is not set')
+    }
+
+    return this.#client
+  }
+
+  get evaluatorModel(): EvaluatorModel {
+    return getTestMetadata().evaluatorModel ?? 'openai__gpt-4o-mini-2024-07-18'
+  }
+
+  get scenario() {
+    return getTestMetadata().scenario
+  }
+
+  get isVaiTest() {
+    return getTestMetadata().isVaiTest
+  }
+
+  setClient(cognitive: Client) {
+    this.#client = cognitive
+  }
+
+  swallowErrors() {
+    if (!getCurrentTest()) {
+      throw new Error('cancelBail is a Vitest hook and must be called within a test')
+    }
+
+    this.#wrapError = true
+    onTestFinished(() => {
+      this.#wrapError = false
+    })
+  }
+}
+
+export const Context = new VaiContext()

package/src/hooks/setEvaluator.ts

@@ -0,0 +1,13 @@
+import { getCurrentTest } from 'vitest/suite'
+import { EvaluatorModel, TestMetadata } from '../context'
+
+export const setEvaluator = (model: EvaluatorModel) => {
+  const test = getCurrentTest()
+
+  if (!test) {
+    throw new Error('setEvaluator is a Vitest hook and must be called within a test')
+  }
+
+  const meta = test.meta as TestMetadata
+  meta.evaluatorModel = model
+}

package/src/hooks/setupClient.ts

@@ -0,0 +1,6 @@
+import { Client } from '@botpress/client'
+import { Context } from '../context'
+
+export const setupClient = (client: Client) => {
+  Context.setClient(client)
+}

package/src/index.ts

@@ -0,0 +1,9 @@
+export { compare } from './task/compare'
+
+export { check } from './assertions/check'
+export { extract } from './assertions/extract'
+export { filter } from './assertions/filter'
+export { rate } from './assertions/rate'
+
+export { setEvaluator } from './hooks/setEvaluator'
+export { setupClient } from './hooks/setupClient'