@botpress/zai 1.1.0 → 2.0.0

package/{src/operations/__tests/index.ts → e2e/utils.ts}

@@ -1,19 +1,20 @@
 import { Client } from '@botpress/client'
-import { type TextTokenizer, getWasmTokenizer } from '@botpress/wasm'
-
+import { type TextTokenizer, getWasmTokenizer } from '@bpinternal/thicktoken'
 import fs from 'node:fs'
 import path from 'node:path'
 import { beforeAll } from 'vitest'
+import { Zai } from '../src'
+import { fastHash } from '../src/utils'
 
-import { Zai } from '../..'
-
-import { fastHash } from '../../utils'
+const DATA_PATH = path.join(__dirname, 'data')
+const CACHE_PATH = path.join(DATA_PATH, 'cache.jsonl')
+const DOC_PATH = path.join(DATA_PATH, 'botpress_docs.txt')
 
 export const getClient = () => {
   return new Client({
     apiUrl: process.env.CLOUD_API_ENDPOINT ?? 'https://api.botpress.dev',
     botId: process.env.CLOUD_BOT_ID,
-    token: process.env.CLOUD_PAT
+    token: process.env.CLOUD_PAT,
   })
 }
 
@@ -31,10 +32,7 @@ function readJSONL<T>(filePath: string, keyProperty: keyof T): Map<string, T> {
   return map
 }
 
-const cache: Map<string, { key: string; value: any }> = readJSONL(
-  path.resolve(import.meta.dirname, './cache.jsonl'),
-  'key'
-)
+const cache: Map<string, { key: string; value: any }> = readJSONL(CACHE_PATH, 'key')
 
 export const getCachedClient = () => {
   const client = getClient()
@@ -54,10 +52,10 @@ export const getCachedClient = () => {
           cache.set(key, { key, value: response })
 
           fs.appendFileSync(
-            path.resolve(import.meta.dirname, './cache.jsonl'),
+            CACHE_PATH,
             JSON.stringify({
               key,
-              value: response
+              value: response,
             }) + '\n'
           )
 
@@ -65,23 +63,27 @@ export const getCachedClient = () => {
         }
       }
       return Reflect.get(target, prop)
-    }
+    },
   })
 
+  ;(proxy as any).clone = () => {
+    return getCachedClient()
+  }
+
   return proxy
 }
 
 export const getZai = () => {
   const client = getCachedClient()
-  return new Zai({ client, retry: { maxRetries: 0 } })
+  return new Zai({ client })
 }
 
 export let tokenizer: TextTokenizer = null!
 
 beforeAll(async () => {
-  tokenizer = await getWasmTokenizer()
+  tokenizer = (await getWasmTokenizer()) as TextTokenizer
 })
 
-export const BotpressDocumentation = fs.readFileSync(path.join(__dirname, './botpress_docs.txt'), 'utf-8').trim()
+export const BotpressDocumentation = fs.readFileSync(DOC_PATH, 'utf-8').trim()
 
 export const metadata = { cost: { input: 1, output: 1 }, latency: 0, model: '', tokens: { input: 1, output: 1 } }

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@botpress/zai",
   "description": "Zui AI (zai) – An LLM utility library written on top of Zui and the Botpress API",
-  "version": "1.1.0",
+  "version": "2.0.0",
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "exports": {
@@ -10,39 +10,41 @@
     "import": "./dist/index.js"
   },
   "scripts": {
-    "build": "npm run build:types && npm run build:neutral",
-    "build:neutral": "esbuild src/**/*.ts src/*.ts --platform=neutral --outdir=dist",
+    "check:type": "tsc --noEmit",
+    "build": "bp add -y && pnpm run build:types && pnpm run build:neutral",
+    "build:neutral": "ts-node -T ./build.ts",
     "build:types": "tsup",
     "watch": "tsup --watch",
-    "test": "vitest run --config vitest.config.ts",
-    "test:update": "vitest -u run --config vitest.config.ts",
-    "test:watch": "vitest --config vitest.config.ts",
-    "build-with-latest-models": "pnpm run update-types && pnpm run update-models && pnpm run build",
-    "update-models": "ts-node ./scripts/update-models.mts",
-    "update-types": "ts-node ./scripts/update-types.mts"
+    "test:e2e": "vitest run --config vitest.config.ts",
+    "test:e2e:update": "vitest -u run --config vitest.config.ts",
+    "test:e2e:watch": "vitest --config vitest.config.ts"
   },
   "keywords": [],
   "author": "",
   "license": "ISC",
   "dependencies": {
+    "@botpress/cognitive": "0.1.16",
     "json5": "^2.2.3",
-    "jsonrepair": "^3.11.2",
+    "jsonrepair": "^3.10.0",
     "lodash-es": "^4.17.21"
   },
   "devDependencies": {
-    "@botpress/vai": "0.0.1-beta.7",
+    "@botpress/client": "workspace:^",
+    "@botpress/common": "workspace:*",
+    "@botpress/vai": "workspace:*",
     "@types/lodash-es": "^4.17.12",
-    "dotenv": "^16.4.7",
-    "esbuild": "^0.24.2",
+    "dotenv": "^16.4.4",
+    "esbuild": "^0.16.12",
+    "glob": "^9.3.4",
     "lodash": "^4.17.21",
-    "ts-node": "^10.9.2",
-    "tsup": "^8.3.5",
-    "typescript": "^5.7.2",
-    "vitest": "^2.1.8"
+    "tsup": "^8.0.2"
   },
   "peerDependencies": {
-    "@botpress/client": "^0.40.0",
-    "@botpress/wasm": "^1.0.0",
-    "@bpinternal/zui": "^0.13.4"
-  }
+    "@bpinternal/thicktoken": "^1.0.0",
+    "@bpinternal/zui": "^0.17.1"
+  },
+  "engines": {
+    "node": ">=18.0.0"
+  },
+  "packageManager": "pnpm@8.6.2"
 }

package/src/adapters/adapter.ts

@@ -19,7 +19,7 @@ export type GetExamplesProps<TInput> = {
 }
 
 export abstract class Adapter {
-  abstract getExamples<TInput, TOutput>(
+  public abstract getExamples<TInput, TOutput>(
     props: GetExamplesProps<TInput>
   ): Promise<
     Array<{
@@ -31,5 +31,5 @@ export abstract class Adapter {
     }>
   >
 
-  abstract saveExample<TInput, TOutput>(props: SaveExampleProps<TInput, TOutput>): Promise<void>
+  public abstract saveExample<TInput, TOutput>(props: SaveExampleProps<TInput, TOutput>): Promise<void>
 }

package/src/adapters/botpress-table.ts

@@ -1,13 +1,13 @@
 import { type Client } from '@botpress/client'
 import { z } from '@bpinternal/zui'
 
-import { BotpressClient, GenerationMetadata } from '../utils'
+import { GenerationMetadata } from '../utils'
 import { Adapter, GetExamplesProps, SaveExampleProps } from './adapter'
 
 const CRITICAL_TAGS = {
   system: 'true',
   'schema-purpose': 'active-learning',
-  'schema-version': 'Oct-2024'
+  'schema-version': 'Oct-2024',
 } as const
 
 const OPTIONAL_TAGS = {
@@ -15,19 +15,19 @@ const OPTIONAL_TAGS = {
   'x-studio-description': 'Table for storing active learning tasks and examples',
   'x-studio-readonly': 'true',
   'x-studio-icon': 'lucide://atom',
-  'x-studio-color': 'green'
+  'x-studio-color': 'green',
 } as const
 
 const FACTOR = 30
 
 const Props = z.object({
-  client: BotpressClient,
+  client: z.custom(() => true),
   tableName: z
     .string()
     .regex(
       /^[a-zA-Z0-9_]{1,45}Table$/,
       'Table name must be lowercase and contain only letters, numbers and underscores'
-    )
+    ),
 })
 
 export type TableSchema = z.input<typeof TableSchema>
@@ -44,10 +44,10 @@ const TableSchema = z.object({
   feedback: z
     .object({
       rating: z.enum(['very-bad', 'bad', 'good', 'very-good']),
-      comment: z.string().nullable()
+      comment: z.string().nullable(),
     })
     .nullable()
-    .default(null)
+    .default(null),
 })
 
 const searchableColumns = ['input'] as const satisfies Array<keyof typeof TableSchema.shape> as string[]
@@ -60,33 +60,33 @@ const TableJsonSchema = Object.entries(TableSchema.shape).reduce((acc, [key, val
 }, {})
 
 export class TableAdapter extends Adapter {
-  private client: Client
-  private tableName: string
+  private _client: Client
+  private _tableName: string
 
-  private status: 'initialized' | 'ready' | 'error'
+  private _status: 'initialized' | 'ready' | 'error'
 
-  constructor(props: z.input<typeof Props>) {
+  public constructor(props: z.input<typeof Props>) {
     super()
     props = Props.parse(props)
-    this.client = props.client
-    this.tableName = props.tableName
-    this.status = 'ready'
+    this._client = props.client
+    this._tableName = props.tableName
+    this._status = 'ready'
   }
 
   public async getExamples<TInput, TOutput>({ taskType, taskId, input }: GetExamplesProps<TInput>) {
-    await this.assertTableExists()
+    await this._assertTableExists()
 
-    const { rows } = await this.client
+    const { rows } = await this._client
       .findTableRows({
-        table: this.tableName,
+        table: this._tableName,
         search: JSON.stringify({ value: input }).substring(0, 1023), // Search is limited to 1024 characters
         limit: 10, // TODO
         filter: {
           // Proximity match of approved examples
           taskType,
           taskId,
-          status: 'approved'
-        } satisfies Partial<TableSchema>
+          status: 'approved',
+        } satisfies Partial<TableSchema>,
       })
       .catch((err) => {
         // TODO: handle error
@@ -99,7 +99,7 @@ export class TableAdapter extends Adapter {
       input: row.input.value as TInput,
       output: row.output.value as TOutput,
       explanation: row.explanation,
-      similarity: row.similarity ?? 0
+      similarity: row.similarity ?? 0,
     }))
   }
 
@@ -112,13 +112,13 @@ export class TableAdapter extends Adapter {
     output,
     explanation,
     metadata,
-    status = 'pending'
+    status = 'pending',
   }: SaveExampleProps<TInput, TOutput>) {
-    await this.assertTableExists()
+    await this._assertTableExists()
 
-    await this.client
+    await this._client
       .upsertTableRows({
-        table: this.tableName,
+        table: this._tableName,
         keyColumn: 'key',
         rows: [
           {
@@ -130,34 +130,34 @@ export class TableAdapter extends Adapter {
             output: { value: output },
             explanation: explanation ?? null,
             status,
-            metadata
-          } satisfies TableSchema
-        ]
+            metadata,
+          } satisfies TableSchema,
+        ],
       })
       .catch(() => {
         // TODO: handle error
       })
   }
 
-  private async assertTableExists() {
-    if (this.status !== 'ready') {
+  private async _assertTableExists() {
+    if (this._status !== 'ready') {
       return
     }
 
-    const { table, created } = await this.client
+    const { table, created } = await this._client
       .getOrCreateTable({
-        table: this.tableName,
+        table: this._tableName,
         factor: FACTOR,
         frozen: true,
         isComputeEnabled: false,
         tags: {
           ...CRITICAL_TAGS,
-          ...OPTIONAL_TAGS
+          ...OPTIONAL_TAGS,
         },
-        schema: TableJsonSchema
+        schema: TableJsonSchema,
       })
       .catch(() => {
-        this.status = 'error'
+        this._status = 'error'
         return { table: null, created: false }
       })
 
@@ -201,10 +201,10 @@ export class TableAdapter extends Adapter {
       }
 
       if (issues.length) {
-        this.status = 'error'
+        this._status = 'error'
       }
     }
 
-    this.status = 'initialized'
+    this._status = 'initialized'
   }
 }

package/src/adapters/memory.ts

@@ -1,13 +1,13 @@
 import { Adapter } from './adapter'
 
 export class MemoryAdapter extends Adapter {
-  constructor(public examples: any[]) {
+  public constructor(public examples: any[]) {
     super()
   }
 
-  async getExamples() {
+  public async getExamples() {
     return this.examples
   }
 
-  async saveExample() {}
+  public async saveExample() {}
 }

package/src/operations/check.ts

@@ -1,3 +1,4 @@
+// eslint-disable consistent-type-definitions
 import { z } from '@bpinternal/zui'
 
 import { fastHash, stringify, takeUntilTokens } from '../utils'
@@ -7,18 +8,27 @@ import { PROMPT_INPUT_BUFFER } from './constants'
 const Example = z.object({
   input: z.any(),
   check: z.boolean(),
-  reason: z.string().optional()
+  reason: z.string().optional(),
 })
 
 export type Options = z.input<typeof Options>
 const Options = z.object({
-  examples: z.array(Example).describe('Examples to check the condition against').default([])
+  examples: z.array(Example).describe('Examples to check the condition against').default([]),
 })
 
 declare module '@botpress/zai' {
   interface Zai {
     /** Checks wether a condition is true or not */
-    check(input: unknown, condition: string, options?: Options): Promise<boolean>
+    check(
+      input: unknown,
+      condition: string,
+      options?: Options
+    ): Promise<{
+      /** Whether the condition is true or not */
+      value: boolean
+      /** The explanation of the decision */
+      explanation: string
+    }>
   }
 }
 
@@ -29,14 +39,15 @@ const END = '■END■'
 Zai.prototype.check = async function (this: Zai, input, condition, _options) {
   const options = Options.parse(_options ?? {})
   const tokenizer = await this.getTokenizer()
-  const PROMPT_COMPONENT = Math.max(this.Model.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
+  await this.fetchModelDetails()
+  const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
 
   const taskId = this.taskId
   const taskType = 'zai.check'
 
   const PROMPT_TOKENS = {
     INPUT: Math.floor(0.5 * PROMPT_COMPONENT),
-    CONDITION: Math.floor(0.2 * PROMPT_COMPONENT)
+    CONDITION: Math.floor(0.2 * PROMPT_COMPONENT),
   }
 
   // Truncate the input to fit the model's input size
@@ -51,7 +62,7 @@ Zai.prototype.check = async function (this: Zai, input, condition, _options) {
       taskType,
       taskId,
       input: inputAsString,
-      condition
+      condition,
     })
   )
 
@@ -59,13 +70,13 @@ Zai.prototype.check = async function (this: Zai, input, condition, _options) {
     ? await this.adapter.getExamples<string, boolean>({
         input: inputAsString,
         taskType,
-        taskId
+        taskId,
       })
     : []
 
   const exactMatch = examples.find((x) => x.key === Key)
   if (exactMatch) {
-    return exactMatch.output
+    return { explanation: exactMatch.explanation ?? '', value: exactMatch.output }
   }
 
   const defaultExamples = [
@@ -73,13 +84,14 @@ Zai.prototype.check = async function (this: Zai, input, condition, _options) {
     {
       input: ['apple', 'banana', 'carrot', 'house'],
       check: false,
-      reason: 'The list contains a house, which is not a fruit. Also, the list contains a carrot, which is a vegetable.'
-    }
+      reason:
+        'The list contains a house, which is not a fruit. Also, the list contains a carrot, which is a vegetable.',
+    },
   ]
 
   const userExamples = [
     ...examples.map((e) => ({ input: e.input, check: e.output, reason: e.explanation })),
-    ...options.examples
+    ...options.examples,
   ]
 
   let exampleId = 1
@@ -108,8 +120,8 @@ ${END}
     {
       type: 'text' as const,
       content: formatOutput(example.check, example.reason ?? ''),
-      role: 'assistant' as const
-    }
+      role: 'assistant' as const,
+    },
   ]
 
   const allExamples = takeUntilTokens(
@@ -129,7 +141,7 @@ ${END}
 `.trim()
     : ''
 
-  const output = await this.callModel({
+  const { output, meta } = await this.callModel({
     systemPrompt: `
 Check if the following condition is true or false for the given input. Before answering, make sure to read the input and the condition carefully.
 Justify your answer, then answer with either ${TRUE} or ${FALSE} at the very end, then add ${END} to finish the response.
@@ -147,9 +159,9 @@ ${specialInstructions}
 Considering the below input and above examples, is the following condition true or false?
 ${formatInput(inputAsString, condition)}
 In your "Analysis", please refer to the Expert Examples # to justify your decision.`.trim(),
-        role: 'user'
-      }
-    ]
+        role: 'user',
+      },
+    ],
   })
 
   const answer = output.choices[0]?.content as string
@@ -162,6 +174,13 @@ In your "Analysis", please refer to the Expert Examples # to justify your decisi
   }
 
   let finalAnswer: boolean
+  const explanation = answer
+    .replace(TRUE, '')
+    .replace(FALSE, '')
+    .replace(END, '')
+    .replace('Final Answer:', '')
+    .replace('Analysis:', '')
+    .trim()
 
   if (hasTrue && hasFalse) {
     // If both TRUE and FALSE are present, we need to check which one was answered last
@@ -177,11 +196,25 @@ In your "Analysis", please refer to the Expert Examples # to justify your decisi
       taskId,
       input: inputAsString,
       instructions: condition,
-      metadata: output.metadata,
+      metadata: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output,
+        },
+        latency: meta.latency,
+        model: this.Model,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output,
+        },
+      },
       output: finalAnswer,
-      explanation: answer.replace(TRUE, '').replace(FALSE, '').replace(END, '').replace('Final Answer:', '').trim()
+      explanation,
     })
   }
 
-  return finalAnswer
+  return {
+    value: finalAnswer,
+    explanation: explanation.trim(),
+  }
 }

package/src/operations/errors.ts

@@ -1,5 +1,5 @@
 export class JsonParsingError extends Error {
-  constructor(
+  public constructor(
     public json: unknown,
     public error: Error
   ) {