@botpress/zai 2.1.19 → 2.1.20

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.1.19",
+  "version": "2.1.20",
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "exports": {
@@ -35,7 +35,8 @@
     "@botpress/cognitive": "0.1.50",
     "json5": "^2.2.3",
     "jsonrepair": "^3.10.0",
-    "lodash-es": "^4.17.21"
+    "lodash-es": "^4.17.21",
+    "p-limit": "^7.2.0"
   },
   "devDependencies": {
     "@botpress/client": "workspace:^",
@@ -53,7 +54,7 @@
   },
   "peerDependencies": {
     "@bpinternal/thicktoken": "^1.0.0",
-    "@bpinternal/zui": "1.2.1"
+    "@bpinternal/zui": "^1.2.2"
   },
   "engines": {
     "node": ">=18.0.0"

package/src/operations/errors.ts

@@ -1,9 +1,104 @@
+import { ZodError } from '@bpinternal/zui'
+
 export class JsonParsingError extends Error {
   public constructor(
     public json: unknown,
     public error: Error
   ) {
-    const message = `Error parsing JSON:\n\n---JSON---\n${json}\n\n---Error---\n\n ${error}`
+    const message = JsonParsingError._formatError(json, error)
     super(message)
   }
+
+  private static _formatError(json: unknown, error: Error): string {
+    let errorMessage = 'Error parsing JSON:\n\n'
+    errorMessage += `---JSON---\n${json}\n\n`
+
+    if (error instanceof ZodError) {
+      errorMessage += '---Validation Errors---\n\n'
+      errorMessage += JsonParsingError._formatZodError(error)
+    } else {
+      errorMessage += '---Error---\n\n'
+      errorMessage += 'The JSON provided is not valid JSON.\n'
+      errorMessage += `Details: ${error.message}\n`
+    }
+
+    return errorMessage
+  }
+
+  private static _formatZodError(zodError: ZodError): string {
+    const issues = zodError.issues
+    if (issues.length === 0) {
+      return 'Unknown validation error\n'
+    }
+
+    let message = ''
+    for (let i = 0; i < issues.length; i++) {
+      const issue = issues[i]
+      const path = issue.path.length > 0 ? issue.path.join('.') : 'root'
+
+      message += `${i + 1}. Field: "${path}"\n`
+
+      switch (issue.code) {
+        case 'invalid_type':
+          message += `   Problem: Expected ${issue.expected}, but received ${issue.received}\n`
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'invalid_string':
+          if ('validation' in issue) {
+            message += `   Problem: Invalid ${issue.validation} format\n`
+          }
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'too_small':
+          if (issue.type === 'string') {
+            if (issue.exact) {
+              message += `   Problem: String must be exactly ${issue.minimum} characters\n`
+            } else {
+              message += `   Problem: String must be at least ${issue.minimum} characters\n`
+            }
+          } else if (issue.type === 'number') {
+            message += `   Problem: Number must be ${issue.inclusive ? 'at least' : 'greater than'} ${issue.minimum}\n`
+          } else if (issue.type === 'array') {
+            message += `   Problem: Array must contain ${issue.inclusive ? 'at least' : 'more than'} ${issue.minimum} items\n`
+          }
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'too_big':
+          if (issue.type === 'string') {
+            if (issue.exact) {
+              message += `   Problem: String must be exactly ${issue.maximum} characters\n`
+            } else {
+              message += `   Problem: String must be at most ${issue.maximum} characters\n`
+            }
+          } else if (issue.type === 'number') {
+            message += `   Problem: Number must be ${issue.inclusive ? 'at most' : 'less than'} ${issue.maximum}\n`
+          } else if (issue.type === 'array') {
+            message += `   Problem: Array must contain ${issue.inclusive ? 'at most' : 'fewer than'} ${issue.maximum} items\n`
+          }
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'invalid_enum_value':
+          message += `   Problem: Invalid value "${issue.received}"\n`
+          message += `   Allowed values: ${issue.options.map((o: any) => `"${o}"`).join(', ')}\n`
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'invalid_literal':
+          message += `   Problem: Expected the literal value "${issue.expected}", but received "${issue.received}"\n`
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'invalid_union':
+          message += "   Problem: Value doesn't match any of the expected formats\n"
+          message += `   Message: ${issue.message}\n`
+          break
+        default:
+          message += `   Problem: ${issue.message}\n`
+      }
+
+      if (i < issues.length - 1) {
+        message += '\n'
+      }
+    }
+
+    return message
+  }
 }

package/src/operations/extract.ts

@@ -1,10 +1,11 @@
 // eslint-disable consistent-type-definitions
-import { z, ZodObject } from '@bpinternal/zui'
+import { z, ZodObject, transforms } from '@bpinternal/zui'
 
 import JSON5 from 'json5'
 import { jsonrepair } from 'jsonrepair'
 
 import { chunk, isArray } from 'lodash-es'
+import pLimit from 'p-limit'
 import { ZaiContext } from '../context'
 import { Response } from '../response'
 import { getTokenizer } from '../tokenizer'
@@ -48,6 +49,7 @@ declare module '@botpress/zai' {
 const START = '■json_start■'
 const END = '■json_end■'
 const NO_MORE = '■NO_MORE_ELEMENT■'
+const ZERO_ELEMENTS = '■ZERO_ELEMENTS■'
 
 const extract = async <S extends OfType<AnyObjectOrArray>>(
   input: unknown,
@@ -56,7 +58,9 @@ const extract = async <S extends OfType<AnyObjectOrArray>>(
   ctx: ZaiContext
 ): Promise<S['_output']> => {
   ctx.controller.signal.throwIfAborted()
-  let schema = _schema as any as z.ZodType
+
+  let schema = transforms.fromJSONSchema(transforms.toJSONSchema(_schema as any as z.ZodType))
+
   const options = Options.parse(_options ?? {})
   const tokenizer = await getTokenizer()
   const model = await ctx.getModel()
@@ -110,18 +114,21 @@ const extract = async <S extends OfType<AnyObjectOrArray>>(
   const inputAsString = stringify(input)
 
   if (tokenizer.count(inputAsString) > options.chunkLength) {
+    const limit = pLimit(10) // Limit to 10 concurrent extraction operations
     const tokens = tokenizer.split(inputAsString)
     const chunks = chunk(tokens, options.chunkLength).map((x) => x.join(''))
     const all = await Promise.allSettled(
       chunks.map((chunk) =>
-        extract(
-          chunk,
-          originalSchema,
-          {
-            ...options,
-            strict: false, // We don't want to fail on strict mode for sub-chunks
-          },
-          ctx
+        limit(() =>
+          extract(
+            chunk,
+            originalSchema,
+            {
+              ...options,
+              strict: false, // We don't want to fail on strict mode for sub-chunks
+            },
+            ctx
+          )
         )
       )
     ).then((results) =>
@@ -162,8 +169,11 @@ Merge it back into a final result.`.trim(),
     instructions.push('You may have multiple elements, or zero elements in the input.')
     instructions.push('You must extract each element separately.')
     instructions.push(`Each element must be a JSON object with exactly the format: ${START}${shape}${END}`)
+    instructions.push(`If there are no elements to extract, respond with ${ZERO_ELEMENTS}.`)
     instructions.push(`When you are done extracting all elements, type "${NO_MORE}" to finish.`)
-    instructions.push(`For example, if you have zero elements, the output should look like this: ${NO_MORE}`)
+    instructions.push(
+      `For example, if you have zero elements, the output should look like this: ${ZERO_ELEMENTS}${NO_MORE}`
+    )
     instructions.push(
       `For example, if you have two elements, the output should look like this: ${START}${abbv}${END}${START}${abbv}${END}${NO_MORE}`
     )

package/src/operations/filter.ts

@@ -2,6 +2,7 @@
 import { z } from '@bpinternal/zui'
 
 import { clamp } from 'lodash-es'
+import pLimit from 'p-limit'
 import { ZaiContext } from '../context'
 import { Response } from '../response'
 import { getTokenizer } from '../tokenizer'
@@ -259,7 +260,8 @@ The condition is: "${condition}"
     return partial
   }
 
-  const filteredChunks = await Promise.all(chunks.map(filterChunk))
+  const limit = pLimit(10) // Limit to 10 concurrent filtering operations
+  const filteredChunks = await Promise.all(chunks.map((chunk) => limit(() => filterChunk(chunk))))
 
   return filteredChunks.flat()
 }

package/src/operations/label.ts

@@ -2,6 +2,7 @@
 import { z } from '@bpinternal/zui'
 
 import { chunk, clamp } from 'lodash-es'
+import pLimit from 'p-limit'
 import { ZaiContext } from '../context'
 import { Response } from '../response'
 import { getTokenizer } from '../tokenizer'
@@ -162,9 +163,10 @@ const label = async <T extends string>(
   const inputAsString = stringify(input)
 
   if (tokenizer.count(inputAsString) > CHUNK_INPUT_MAX_TOKENS) {
+    const limit = pLimit(10) // Limit to 10 concurrent labeling operations
     const tokens = tokenizer.split(inputAsString)
     const chunks = chunk(tokens, CHUNK_INPUT_MAX_TOKENS).map((x) => x.join(''))
-    const allLabels = await Promise.all(chunks.map((chunk) => label(chunk, _labels, _options, ctx)))
+    const allLabels = await Promise.all(chunks.map((chunk) => limit(() => label(chunk, _labels, _options, ctx))))
 
     // Merge all the labels together (those who are true will remain true)
     return allLabels.reduce((acc, x) => {

package/src/operations/summarize.ts

@@ -2,6 +2,7 @@
 import { z } from '@bpinternal/zui'
 
 import { chunk } from 'lodash-es'
+import pLimit from 'p-limit'
 import { ZaiContext } from '../context'
 import { Response } from '../response'
 
@@ -115,9 +116,9 @@ ${newText}
   const chunkSize = Math.ceil(tokens.length / (parts * N))
 
   if (useMergeSort) {
-    // TODO: use pLimit here to not have too many chunks
+    const limit = pLimit(10) // Limit to 10 concurrent summarization operations
     const chunks = chunk(tokens, chunkSize).map((x) => x.join(''))
-    const allSummaries = (await Promise.allSettled(chunks.map((chunk) => summarize(chunk, options, ctx))))
+    const allSummaries = (await Promise.allSettled(chunks.map((chunk) => limit(() => summarize(chunk, options, ctx)))))
       .filter((x) => x.status === 'fulfilled')
       .map((x) => x.value)
     return summarize(allSummaries.join('\n\n============\n\n'), options, ctx)

package/src/zai.ts

@@ -1,5 +1,5 @@
 import { Client } from '@botpress/client'
-import { BotpressClientLike, Cognitive, Model } from '@botpress/cognitive'
+import { BotpressClientLike, Cognitive, Model, Models } from '@botpress/cognitive'
 
 import { type TextTokenizer, getWasmTokenizer } from '@bpinternal/thicktoken'
 import { z } from '@bpinternal/zui'
@@ -8,8 +8,6 @@ import { Adapter } from './adapters/adapter'
 import { TableAdapter } from './adapters/botpress-table'
 import { MemoryAdapter } from './adapters/memory'
 
-type ModelId = Required<Parameters<Cognitive['generateContent']>[0]['model']>
-
 type ActiveLearning = {
   enable: boolean
   tableName: string
@@ -39,7 +37,7 @@ const _ActiveLearning = z.object({
 type ZaiConfig = {
   client: BotpressClientLike | Cognitive
   userId?: string
-  modelId?: ModelId | string
+  modelId?: Models
   activeLearning?: ActiveLearning
   namespace?: string
 }
@@ -48,7 +46,7 @@ const _ZaiConfig = z.object({
   client: z.custom<BotpressClientLike | Cognitive>(),
   userId: z.string().describe('The ID of the user consuming the API').optional(),
   modelId: z
-    .custom<ModelId | string>(
+    .custom<Models>(
       (value) => {
         if (typeof value !== 'string') {
           return false
@@ -65,7 +63,7 @@ const _ZaiConfig = z.object({
       }
     )
     .describe('The ID of the model you want to use')
-    .default('best' satisfies ModelId),
+    .default('best' satisfies Models),
   activeLearning: _ActiveLearning.default({ enable: false }),
   namespace: z
     .string()
@@ -84,7 +82,7 @@ export class Zai {
 
   private _userId: string | undefined
 
-  protected Model: ModelId
+  protected Model: Models
   protected ModelDetails: Model
   protected namespace: string
   protected adapter: Adapter
@@ -100,7 +98,7 @@ export class Zai {
 
     this.namespace = parsed.namespace
     this._userId = parsed.userId
-    this.Model = parsed.modelId as ModelId
+    this.Model = parsed.modelId as Models
     this.activeLearning = parsed.activeLearning as ActiveLearning
 
     this.adapter = parsed.activeLearning?.enable
@@ -117,7 +115,7 @@ export class Zai {
   ): ReturnType<Cognitive['generateContent']> {
     return this.client.generateContent({
       ...props,
-      model: this.Model,
+      model: this.Model as Required<Parameters<Cognitive['generateContent']>[0]>['model'],
       userId: this._userId,
     })
   }