@botpress/zai 2.0.7 → 2.0.10

package/e2e/utils.ts

@@ -4,10 +4,9 @@ import fs from 'node:fs'
 import path from 'node:path'
 import { beforeAll } from 'vitest'
 import { Zai } from '../src'
-import { fastHash } from '../src/utils'
+import { getCachedCognitiveClient } from './client'
 
 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 = () => {
@@ -18,59 +17,8 @@ export const getClient = () => {
   })
 }
 
-function readJSONL<T>(filePath: string, keyProperty: keyof T): Map<string, T> {
-  const lines = fs.readFileSync(filePath, 'utf-8').split(/\r?\n/).filter(Boolean)
-
-  const map = new Map<string, T>()
-
-  for (const line of lines) {
-    const obj = JSON.parse(line) as T
-    const key = String(obj[keyProperty])
-    map.set(key, obj)
-  }
-
-  return map
-}
-
-const cache: Map<string, { key: string; value: any }> = readJSONL(CACHE_PATH, 'key')
-
 export const getCachedClient = () => {
-  const client = getClient()
-
-  const proxy = new Proxy(client, {
-    get(target, prop) {
-      if (prop === 'callAction') {
-        return async (...args: Parameters<Client['callAction']>) => {
-          const key = fastHash(JSON.stringify(args))
-          const cached = cache.get(key)
-
-          if (cached) {
-            return cached.value
-          }
-
-          const response = await target.callAction(...args)
-          cache.set(key, { key, value: response })
-
-          fs.appendFileSync(
-            CACHE_PATH,
-            JSON.stringify({
-              key,
-              value: response,
-            }) + '\n'
-          )
-
-          return response
-        }
-      }
-      return Reflect.get(target, prop)
-    },
-  })
-
-  ;(proxy as any).clone = () => {
-    return getCachedClient()
-  }
-
-  return proxy
+  return getCachedCognitiveClient()
 }
 
 export const getZai = () => {

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": "2.0.7",
+  "version": "2.0.10",
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "exports": {
@@ -23,7 +23,7 @@
   "author": "",
   "license": "ISC",
   "dependencies": {
-    "@botpress/cognitive": "0.1.22",
+    "@botpress/cognitive": "0.1.24",
     "json5": "^2.2.3",
     "jsonrepair": "^3.10.0",
     "lodash-es": "^4.17.21"
@@ -33,6 +33,7 @@
     "@botpress/common": "workspace:*",
     "@botpress/vai": "workspace:*",
     "@types/lodash-es": "^4.17.12",
+    "diff": "^8.0.1",
     "dotenv": "^16.4.4",
     "esbuild": "^0.16.12",
     "glob": "^9.3.4",
@@ -41,7 +42,7 @@
   },
   "peerDependencies": {
     "@bpinternal/thicktoken": "^1.0.0",
-    "@bpinternal/zui": "^0.17.1"
+    "@bpinternal/zui": "^0.22.5"
   },
   "engines": {
     "node": ">=18.0.0"

package/src/adapters/botpress-table.ts

@@ -20,6 +20,11 @@ const OPTIONAL_TAGS = {
 
 const FACTOR = 30
 
+type Props = {
+  client: Client
+  tableName: string
+}
+
 const Props = z.object({
   client: z.custom(() => true),
   tableName: z
@@ -30,7 +35,22 @@ const Props = z.object({
     ),
 })
 
-export type TableSchema = (typeof TableSchema)['_input']
+export type TableSchema = {
+  taskType: string
+  taskId: string
+  key: string
+  instructions: string
+  input: Record<string, unknown>
+  output: Record<string, unknown>
+  explanation: string | null
+  metadata: GenerationMetadata
+  status: 'pending' | 'rejected' | 'approved'
+  feedback: {
+    rating: 'very-bad' | 'bad' | 'good' | 'very-good'
+    comment: string | null
+  } | null
+}
+
 const TableSchema = z.object({
   taskType: z.string().describe('The type of the task (filter, extract, etc.)'),
   taskId: z.string(),
@@ -39,7 +59,7 @@ const TableSchema = z.object({
   input: z.object({}).passthrough().describe('The input to the task'),
   output: z.object({}).passthrough().describe('The expected output'),
   explanation: z.string().nullable(),
-  metadata: GenerationMetadata,
+  metadata: z.object({}).passthrough(),
   status: z.enum(['pending', 'rejected', 'approved']),
   feedback: z
     .object({
@@ -65,9 +85,9 @@ export class TableAdapter extends Adapter {
 
   private _status: 'initialized' | 'ready' | 'error'
 
-  public constructor(props: (typeof Props)['_input']) {
+  public constructor(props: Props) {
     super()
-    props = Props.parse(props)
+    props = Props.parse(props) as Props
     this._client = props.client
     this._tableName = props.tableName
     this._status = 'ready'
@@ -131,6 +151,7 @@ export class TableAdapter extends Adapter {
             explanation: explanation ?? null,
             status,
             metadata,
+            feedback: null, // Feedback is not provided at this point
           } satisfies TableSchema,
         ],
       })

package/src/operations/check.ts

@@ -5,15 +5,27 @@ import { fastHash, stringify, takeUntilTokens } from '../utils'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER } from './constants'
 
-const Example = z.object({
+const _Example = z.object({
   input: z.any(),
   check: z.boolean(),
   reason: z.string().optional(),
+  condition: z.string().optional(),
 })
 
-export type Options = (typeof Options)['_input']
-const Options = z.object({
-  examples: z.array(Example).describe('Examples to check the condition against').default([]),
+type Example = {
+  input: unknown
+  check: boolean
+  reason?: string
+  condition?: string
+}
+
+export type Options = {
+  /** Examples to check the condition against */
+  examples?: Array<Example>
+}
+
+const _Options = z.object({
+  examples: z.array(_Example).describe('Examples to check the condition against').default([]),
 })
 
 declare module '@botpress/zai' {
@@ -36,8 +48,8 @@ const TRUE = '■TRUE■'
 const FALSE = '■FALSE■'
 const END = '■END■'
 
-Zai.prototype.check = async function (this: Zai, input, condition, _options) {
-  const options = Options.parse(_options ?? {})
+Zai.prototype.check = async function (this: Zai, input: unknown, condition: string, _options: Options | undefined) {
+  const options = _Options.parse(_options ?? {}) as Options
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
   const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
@@ -80,14 +92,20 @@ Zai.prototype.check = async function (this: Zai, input, condition, _options) {
   }
 
   const defaultExamples = [
-    { input: '50 Cent', check: true, reason: '50 Cent is widely recognized as a public personality.' },
+    {
+      input: '50 Cent',
+      check: true,
+      reason: '50 Cent is widely recognized as a public personality.',
+      condition: 'Is the input a public personality?',
+    },
     {
       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.',
+      condition: 'Is the input exclusively a list of fruits?',
     },
-  ]
+  ] satisfies Example[]
 
   const userExamples = [
     ...examples.map((e) => ({ input: e.input, check: e.output, reason: e.explanation })),
@@ -115,8 +133,12 @@ ${END}
 `.trim()
   }
 
-  const formatExample = (example: { input?: any; check: boolean; reason?: string }) => [
-    { type: 'text' as const, content: formatInput(stringify(example.input ?? null), condition), role: 'user' as const },
+  const formatExample = (example: Example) => [
+    {
+      type: 'text' as const,
+      content: formatInput(stringify(example.input ?? null), example.condition ?? condition),
+      role: 'user' as const,
+    },
     {
       type: 'text' as const,
       content: formatOutput(example.check, example.reason ?? ''),

package/src/operations/extract.ts

@@ -10,7 +10,13 @@ import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER } from './constants'
 import { JsonParsingError } from './errors'
 
-export type Options = (typeof Options)['_input']
+export type Options = {
+  /** Instructions to guide the user on how to extract the data */
+  instructions?: string
+  /** The maximum number of tokens per chunk */
+  chunkLength?: number
+}
+
 const Options = z.object({
   instructions: z.string().optional().describe('Instructions to guide the user on how to extract the data'),
   chunkLength: z

package/src/operations/filter.ts

@@ -6,15 +6,26 @@ import { fastHash, stringify, takeUntilTokens } from '../utils'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
 
-type Example = (typeof Example)['_input']
-const Example = z.object({
+type Example = {
+  input: unknown
+  filter: boolean
+  reason?: string
+}
+
+const _Example = z.object({
   input: z.any(),
   filter: z.boolean(),
   reason: z.string().optional(),
 })
 
-export type Options = (typeof Options)['_input']
-const Options = z.object({
+export type Options = {
+  /** The maximum number of tokens per item */
+  tokensPerItem?: number
+  /** Examples to filter the condition against */
+  examples?: Array<Example>
+}
+
+const _Options = z.object({
   tokensPerItem: z
     .number()
     .min(1)
@@ -22,7 +33,7 @@ const Options = z.object({
     .optional()
     .describe('The maximum number of tokens per item')
     .default(250),
-  examples: z.array(Example).describe('Examples to filter the condition against').default([]),
+  examples: z.array(_Example).describe('Examples to filter the condition against').default([]),
 })
 
 declare module '@botpress/zai' {
@@ -35,7 +46,7 @@ declare module '@botpress/zai' {
 const END = '■END■'
 
 Zai.prototype.filter = async function (this: Zai, input, condition, _options) {
-  const options = Options.parse(_options ?? {})
+  const options = _Options.parse(_options ?? {}) as Options
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
 

package/src/operations/label.ts

@@ -22,11 +22,16 @@ type Example<T extends string> = {
   labels: Partial<Record<T, { label: Label; explanation?: string }>>
 }
 
-export type Options<T extends string> = Omit<(typeof Options)['_input'], 'examples'> & {
-  examples?: Array<Partial<Example<T>>>
+export type Options<T extends string> = {
+  /** Examples to help the user make a decision */
+  examples?: Array<Example<T>>
+  /** Instructions to guide the user on how to extract the data */
+  instructions?: string
+  /** The maximum number of tokens per chunk */
+  chunkLength?: number
 }
 
-const Options = z.object({
+const _Options = z.object({
   examples: z
     .array(
       z.object({
@@ -48,7 +53,7 @@ const Options = z.object({
 
 type Labels<T extends string> = Record<T, string>
 
-const Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((labels, ctx) => {
+const _Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((labels, ctx) => {
   const keys = Object.keys(labels)
 
   for (const key of keys) {
@@ -119,9 +124,14 @@ const getConfidence = (label: Label) => {
   }
 }
 
-Zai.prototype.label = async function <T extends string>(this: Zai, input, _labels, _options) {
-  const options = Options.parse(_options ?? {})
-  const labels = Labels.parse(_labels)
+Zai.prototype.label = async function <T extends string>(
+  this: Zai,
+  input: unknown,
+  _labels: Labels<T>,
+  _options: Options<T> | undefined
+) {
+  const options = _Options.parse(_options ?? {}) as unknown as Options<T>
+  const labels = _Labels.parse(_labels) as Labels<T>
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
 
@@ -211,7 +221,7 @@ Zai.prototype.label = async function <T extends string>(this: Zai, input, _label
   options.examples.forEach((example) => {
     examples.push({
       key: fastHash(JSON.stringify(example)),
-      input: example.input,
+      input: stringify(example.input),
       similarity: 1,
       explanation: '',
       output: example.labels as unknown as {

package/src/operations/rewrite.ts

@@ -5,15 +5,26 @@ import { fastHash, stringify, takeUntilTokens } from '../utils'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER } from './constants'
 
-type Example = (typeof Example)['_input'] & { instructions?: string }
-const Example = z.object({
+type Example = {
+  input: string
+  output: string
+  instructions?: string
+}
+
+const _Example = z.object({
   input: z.string(),
   output: z.string(),
 })
 
-export type Options = (typeof Options)['_input']
+export type Options = {
+  /** Examples to guide the rewriting */
+  examples?: Array<Example>
+  /** The maximum number of tokens to generate */
+  length?: number
+}
+
 const Options = z.object({
-  examples: z.array(Example).default([]),
+  examples: z.array(_Example).default([]),
   length: z.number().min(10).max(16_000).optional().describe('The maximum number of tokens to generate'),
 })
 
@@ -28,7 +39,7 @@ const START = '■START■'
 const END = '■END■'
 
 Zai.prototype.rewrite = async function (this: Zai, original, prompt, _options) {
-  const options = Options.parse(_options ?? {})
+  const options = Options.parse(_options ?? {}) as Options
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
 
@@ -101,7 +112,7 @@ ${END}
   }
 
   const savedExamples: Example[] = [
-    ...tableExamples.map((x) => ({ input: x.input as string, output: x.output as string })),
+    ...tableExamples.map((x) => ({ input: x.input as string, output: x.output as string }) satisfies Example),
     ...options.examples,
   ]
 

package/src/operations/summarize.ts

@@ -5,7 +5,24 @@ import { chunk } from 'lodash-es'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
 
-export type Options = (typeof Options)['_input']
+export type Options = {
+  /** What should the text be summarized to? */
+  prompt?: string
+  /** How to format the example text */
+  format?: string
+  /** The length of the summary in tokens */
+  length?: number
+  /** How many times longer (than final length) are the intermediate summaries generated */
+  intermediateFactor?: number
+  /** The maximum number of iterations to perform */
+  maxIterations?: number
+  /** Sliding window options */
+  sliding?: {
+    window: number
+    overlap: number
+  }
+}
+
 const Options = z.object({
   prompt: z
     .string()
@@ -45,7 +62,7 @@ const START = '■START■'
 const END = '■END■'
 
 Zai.prototype.summarize = async function (this: Zai, original, _options) {
-  const options = Options.parse(_options ?? {})
+  const options = Options.parse(_options ?? {}) as Options
   const tokenizer = await this.getTokenizer()
   await this.fetchModelDetails()
 

package/src/operations/text.ts

@@ -5,7 +5,11 @@ import { clamp } from 'lodash-es'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
 
-export type Options = (typeof Options)['_input']
+export type Options = {
+  /** The maximum number of tokens to generate */
+  length?: number
+}
+
 const Options = z.object({
   length: z.number().min(1).max(100_000).optional().describe('The maximum number of tokens to generate'),
 })

package/src/utils.ts

@@ -1,5 +1,3 @@
-import { z } from '@bpinternal/zui'
-
 export const stringify = (input: unknown, beautify = true) => {
   return typeof input === 'string' && !!input.length
     ? input
@@ -33,20 +31,15 @@ export const takeUntilTokens = <T>(arr: T[], tokens: number, count: (el: T) => n
   return result
 }
 
-export type GenerationMetadata = (typeof GenerationMetadata)['_input']
-export const GenerationMetadata = z.object({
-  model: z.string(),
-  cost: z
-    .object({
-      input: z.number(),
-      output: z.number(),
-    })
-    .describe('Cost in $USD'),
-  latency: z.number().describe('Latency in milliseconds'),
-  tokens: z
-    .object({
-      input: z.number(),
-      output: z.number(),
-    })
-    .describe('Number of tokens used'),
-})
+export type GenerationMetadata = {
+  model: string
+  cost: {
+    input: number
+    output: number
+  }
+  latency: number
+  tokens: {
+    input: number
+    output: number
+  }
+}

package/src/zai.ts

@@ -1,3 +1,4 @@
+import { Client } from '@botpress/client'
 import { BotpressClientLike, Cognitive, Model } from '@botpress/cognitive'
 
 import { type TextTokenizer, getWasmTokenizer } from '@bpinternal/thicktoken'
@@ -9,8 +10,13 @@ import { MemoryAdapter } from './adapters/memory'
 
 type ModelId = Required<Parameters<Cognitive['generateContent']>[0]['model']>
 
-type ActiveLearning = (typeof ActiveLearning)['_input']
-const ActiveLearning = z.object({
+type ActiveLearning = {
+  enable: boolean
+  tableName: string
+  taskId: string
+}
+
+const _ActiveLearning = z.object({
   enable: z.boolean().describe('Whether to enable active learning').default(false),
   tableName: z
     .string()
@@ -30,8 +36,15 @@ const ActiveLearning = z.object({
     .default('default'),
 })
 
-type ZaiConfig = (typeof ZaiConfig)['_input']
-const ZaiConfig = z.object({
+type ZaiConfig = {
+  client: BotpressClientLike | Cognitive
+  userId?: string
+  modelId?: ModelId | string
+  activeLearning?: ActiveLearning
+  namespace?: string
+}
+
+const _ZaiConfig = z.object({
   client: z.custom<BotpressClientLike | Cognitive>(),
   userId: z.string().describe('The ID of the user consuming the API').optional(),
   modelId: z
@@ -53,7 +66,7 @@ const ZaiConfig = z.object({
     )
     .describe('The ID of the model you want to use')
     .default('best' satisfies ModelId),
-  activeLearning: ActiveLearning.default({ enable: false }),
+  activeLearning: _ActiveLearning.default({ enable: false }),
   namespace: z
     .string()
     .regex(
@@ -79,7 +92,7 @@ export class Zai {
 
   public constructor(config: ZaiConfig) {
     this._originalConfig = config
-    const parsed = ZaiConfig.parse(config)
+    const parsed = _ZaiConfig.parse(config) as ZaiConfig
 
     this.client = Cognitive.isCognitiveClient(parsed.client)
       ? (parsed.client as unknown as Cognitive)
@@ -88,10 +101,13 @@ export class Zai {
     this.namespace = parsed.namespace
     this._userId = parsed.userId
     this.Model = parsed.modelId as ModelId
-    this.activeLearning = parsed.activeLearning
+    this.activeLearning = parsed.activeLearning as ActiveLearning
 
     this.adapter = parsed.activeLearning?.enable
-      ? new TableAdapter({ client: this.client.client, tableName: parsed.activeLearning.tableName })
+      ? new TableAdapter({
+          client: this.client.client as unknown as Client,
+          tableName: parsed.activeLearning.tableName,
+        })
       : new MemoryAdapter([])
   }