@botpress/zai 2.0.14 → 2.0.16

package/e2e/utils.ts

@@ -1,4 +1,5 @@
 import { Client } from '@botpress/client'
+import { Cognitive } from '@botpress/cognitive'
 import { type TextTokenizer, getWasmTokenizer } from '@bpinternal/thicktoken'
 import fs from 'node:fs'
 import path from 'node:path'
@@ -21,8 +22,8 @@ export const getCachedClient = () => {
   return getCachedCognitiveClient()
 }
 
-export const getZai = () => {
-  const client = getCachedClient()
+export const getZai = (cognitive?: Cognitive) => {
+  const client = cognitive || getCachedClient()
   return new Zai({ client })
 }
 

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.14",
+  "version": "2.0.16",
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "exports": {

package/src/operations/extract.ts

@@ -15,6 +15,8 @@ export type Options = {
   instructions?: string
   /** The maximum number of tokens per chunk */
   chunkLength?: number
+  /** Whether to strictly follow the schema or not */
+  strict?: boolean
 }
 
 const Options = z.object({
@@ -26,6 +28,7 @@ const Options = z.object({
     .optional()
     .describe('The maximum number of tokens per chunk')
     .default(16_000),
+  strict: z.boolean().optional().default(true).describe('Whether to strictly follow the schema or not'),
 })
 
 type __Z<T extends any = any> = { _output: T }
@@ -35,7 +38,7 @@ type AnyObjectOrArray = Record<string, unknown> | Array<unknown>
 declare module '@botpress/zai' {
   interface Zai {
     /** Extracts one or many elements from an arbitrary input */
-    extract<S extends OfType<AnyObjectOrArray>>(input: unknown, schema: S, options?: Options): Promise<S['_output']>
+    extract<S extends OfType<any>>(input: unknown, schema: S, options?: Options): Promise<S['_output']>
   }
 }
 
@@ -60,26 +63,37 @@ Zai.prototype.extract = async function <S extends OfType<AnyObjectOrArray>>(
   const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
 
   let isArrayOfObjects = false
+  let wrappedValue = false
   const originalSchema = schema
 
   const baseType = (schema.naked ? schema.naked() : schema)?.constructor?.name ?? 'unknown'
 
-  if (baseType === 'ZodObject') {
-    // Do nothing
-  } else if (baseType === 'ZodArray') {
+  if (baseType === 'ZodArray') {
+    isArrayOfObjects = true
     let elementType = (schema as any).element
     if (elementType.naked) {
       elementType = elementType.naked()
     }
 
     if (elementType?.constructor?.name === 'ZodObject') {
-      isArrayOfObjects = true
       schema = elementType
     } else {
-      throw new Error('Schema must be a ZodObject or a ZodArray<ZodObject>')
+      wrappedValue = true
+      schema = z.object({
+        value: elementType,
+      })
     }
-  } else {
-    throw new Error('Schema must be either a ZuiObject or a ZuiArray<ZuiObject>')
+  } else if (baseType !== 'ZodObject') {
+    wrappedValue = true
+    schema = z.object({
+      value: originalSchema,
+    })
+  }
+
+  if (!options.strict) {
+    try {
+      schema = (schema as ZodObject).partial()
+    } catch {}
   }
 
   const schemaTypescript = schema.toTypescriptType({ declaration: false })
@@ -92,20 +106,38 @@ Zai.prototype.extract = async function <S extends OfType<AnyObjectOrArray>>(
 
   const keys = Object.keys((schema as ZodObject).shape)
 
-  let inputAsString = stringify(input)
+  const inputAsString = stringify(input)
 
   if (tokenizer.count(inputAsString) > options.chunkLength) {
-    // If we want to extract an array of objects, we will run this function recursively
-    if (isArrayOfObjects) {
-      const tokens = tokenizer.split(inputAsString)
-      const chunks = chunk(tokens, options.chunkLength).map((x) => x.join(''))
-      const all = await Promise.all(chunks.map((chunk) => this.extract(chunk, originalSchema)))
+    const tokens = tokenizer.split(inputAsString)
+    const chunks = chunk(tokens, options.chunkLength).map((x) => x.join(''))
+    const all = await Promise.allSettled(
+      chunks.map((chunk) =>
+        this.extract(chunk, originalSchema, {
+          ...options,
+          strict: false, // We don't want to fail on strict mode for sub-chunks
+        })
+      )
+    ).then((results) =>
+      results.filter((x) => x.status === 'fulfilled').map((x) => (x as PromiseFulfilledResult<S['_output']>).value)
+    )
 
-      return all.flat() as any as S['_output']
-    } else {
-      // Truncate the input to fit the model's input size
-      inputAsString = tokenizer.truncate(stringify(input), options.chunkLength)
-    }
+    // We run this function recursively until all chunks are merged into a single output
+    const rows = all.map((x, idx) => `<part-${idx + 1}>\n${stringify(x, true)}\n</part-${idx + 1}>`).join('\n')
+    return this.extract(
+      `
+The result has been split into ${all.length} parts. Recursively merge the result into the final result.
+When merging arrays, take unique values.
+When merging conflictual (but defined) information, take the most reasonable and frequent value.
+Non-defined values are OK and normal. Don't delete fields because of null values. Focus on defined values.
+
+Here's the data:
+${rows}
+
+Merge it back into a final result.`.trim(),
+      originalSchema,
+      options
+    )
   }
 
   const instructions: string[] = []
@@ -131,6 +163,10 @@ Zai.prototype.extract = async function <S extends OfType<AnyObjectOrArray>>(
     instructions.push(`The element must be a JSON object with exactly the format: ${START}${shape}${END}`)
   }
 
+  if (!options.strict) {
+    instructions.push('You may ignore any fields that are not present in the input. All keys are optional.')
+  }
+
   // All tokens remaining after the input and condition are accounted can be used for examples
   const EXAMPLES_TOKENS = PROMPT_COMPONENT - tokenizer.count(inputAsString) - tokenizer.count(instructions.join('\n'))
 
@@ -270,18 +306,27 @@ ${instructions.map((x) => `• ${x}`).join('\n')}
     ],
   })
 
-  const answer = output.choices[0]?.content as string
+  const answer = (output.choices[0]?.content ?? '{}') as string
 
   const elements = answer
-    .split(START)
-    .filter((x) => x.trim().length > 0)
+    ?.split(START)
+    .filter((x) => x.trim().length > 0 && x.includes('}'))
     .map((x) => {
       try {
         const json = x.slice(0, x.indexOf(END)).trim()
         const repairedJson = jsonrepair(json)
         const parsedJson = JSON5.parse(repairedJson)
+        const safe = schema.safeParse(parsedJson)
+
+        if (safe.success) {
+          return safe.data
+        }
+
+        if (options.strict) {
+          throw new JsonParsingError(x, safe.error)
+        }
 
-        return schema.parse(parsedJson)
+        return parsedJson
       } catch (error) {
         throw new JsonParsingError(x, error instanceof Error ? error : new Error('Unknown error'))
       }
@@ -293,11 +338,19 @@ ${instructions.map((x) => `• ${x}`).join('\n')}
   if (isArrayOfObjects) {
     final = elements
   } else if (elements.length === 0) {
-    final = schema.parse({})
+    final = options.strict ? schema.parse({}) : {}
   } else {
     final = elements[0]
   }
 
+  if (wrappedValue) {
+    if (Array.isArray(final)) {
+      final = final.map((x) => ('value' in x ? x.value : x))
+    } else {
+      final = 'value' in final ? final.value : final
+    }
+  }
+
   if (taskId) {
     await this.adapter.saveExample({
       key: Key,