@botpress/zai 2.0.15 → 2.1.0

package/src/operations/extract.ts

@@ -5,6 +5,9 @@ import JSON5 from 'json5'
 import { jsonrepair } from 'jsonrepair'
 
 import { chunk, isArray } from 'lodash-es'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
 import { fastHash, stringify, takeUntilTokens } from '../utils'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER } from './constants'
@@ -15,6 +18,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 +31,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 +41,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): Response<S['_output']>
   }
 }
 
@@ -43,52 +49,61 @@ const START = '■json_start■'
 const END = '■json_end■'
 const NO_MORE = '■NO_MORE_ELEMENT■'
 
-Zai.prototype.extract = async function <S extends OfType<AnyObjectOrArray>>(
-  this: Zai,
+const extract = async <S extends OfType<AnyObjectOrArray>>(
   input: unknown,
   _schema: S,
-  _options?: Options
-): Promise<S['_output']> {
+  _options: Options | undefined,
+  ctx: ZaiContext
+): Promise<S['_output']> => {
+  ctx.controller.signal.throwIfAborted()
   let schema = _schema as any as z.ZodType
   const options = Options.parse(_options ?? {})
-  const tokenizer = await this.getTokenizer()
-  await this.fetchModelDetails()
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
 
-  const taskId = this.taskId
+  const taskId = ctx.taskId
   const taskType = 'zai.extract'
 
-  const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
+  const PROMPT_COMPONENT = Math.max(model.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 })
   const schemaLength = tokenizer.count(schemaTypescript)
 
-  options.chunkLength = Math.min(
-    options.chunkLength,
-    this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength
-  )
+  options.chunkLength = Math.min(options.chunkLength, model.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength)
 
   const keys = Object.keys((schema as ZodObject).shape)
 
@@ -97,10 +112,41 @@ Zai.prototype.extract = async function <S extends OfType<AnyObjectOrArray>>(
   if (tokenizer.count(inputAsString) > options.chunkLength) {
     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 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
+        )
+      )
+    ).then((results) =>
+      results.filter((x) => x.status === 'fulfilled').map((x) => (x as PromiseFulfilledResult<S['_output']>).value)
+    )
+
+    ctx.controller.signal.throwIfAborted()
 
     // We run this function recursively until all chunks are merged into a single output
-    return this.extract(all, originalSchema, options)
+    const rows = all.map((x, idx) => `<part-${idx + 1}>\n${stringify(x, true)}\n</part-${idx + 1}>`).join('\n')
+    return 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,
+      ctx
+    )
   }
 
   const instructions: string[] = []
@@ -126,6 +172,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'))
 
@@ -138,13 +188,14 @@ Zai.prototype.extract = async function <S extends OfType<AnyObjectOrArray>>(
     })
   )
 
-  const examples = taskId
-    ? await this.adapter.getExamples<string, unknown>({
-        input: inputAsString,
-        taskType,
-        taskId,
-      })
-    : []
+  const examples =
+    taskId && ctx.adapter
+      ? await ctx.adapter.getExamples<string, unknown>({
+          input: inputAsString,
+          taskType,
+          taskId,
+        })
+      : []
 
   const exactMatch = examples.find((x) => x.key === Key)
   if (exactMatch) {
@@ -246,7 +297,7 @@ ${END}`.trim()
     .map(formatExample)
     .flat()
 
-  const { output, meta } = await this.callModel({
+  const { meta, extracted } = await ctx.generateContent({
     systemPrompt: `
 Extract the following information from the input:
 ${schemaTypescript}
@@ -263,38 +314,53 @@ ${instructions.map((x) => `• ${x}`).join('\n')}
         content: formatInput(inputAsString, schemaTypescript, options.instructions ?? ''),
       },
     ],
+    transform: (text) =>
+      (text || '{}')
+        ?.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 parsedJson
+          } catch (error) {
+            throw new JsonParsingError(x, error instanceof Error ? error : new Error('Unknown error'))
+          }
+        })
+        .filter((x) => x !== null),
   })
 
-  const answer = output.choices[0]?.content as string
-
-  const elements = answer
-    .split(START)
-    .filter((x) => x.trim().length > 0)
-    .map((x) => {
-      try {
-        const json = x.slice(0, x.indexOf(END)).trim()
-        const repairedJson = jsonrepair(json)
-        const parsedJson = JSON5.parse(repairedJson)
-
-        return schema.parse(parsedJson)
-      } catch (error) {
-        throw new JsonParsingError(x, error instanceof Error ? error : new Error('Unknown error'))
-      }
-    })
-    .filter((x) => x !== null)
-
   let final: any
 
   if (isArrayOfObjects) {
-    final = elements
-  } else if (elements.length === 0) {
-    final = schema.parse({})
+    final = extracted
+  } else if (extracted.length === 0) {
+    final = options.strict ? schema.parse({}) : {}
   } else {
-    final = elements[0]
+    final = extracted[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({
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
       key: Key,
       taskId: `zai/${taskId}`,
       taskType,
@@ -307,7 +373,7 @@ ${instructions.map((x) => `• ${x}`).join('\n')}
           output: meta.cost.output,
         },
         latency: meta.latency,
-        model: this.Model,
+        model: ctx.modelId,
         tokens: {
           input: meta.tokens.input,
           output: meta.tokens.output,
@@ -318,3 +384,20 @@ ${instructions.map((x) => `• ${x}`).join('\n')}
 
   return final
 }
+
+Zai.prototype.extract = function <S extends OfType<AnyObjectOrArray>>(
+  this: Zai,
+  input: unknown,
+  schema: S,
+  _options?: Options
+): Response<S['_output']> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.extract',
+    adapter: this.adapter,
+  })
+
+  return new Response<S['_output']>(context, extract(input, schema, _options, context), (result) => result)
+}

package/src/operations/filter.ts

@@ -2,6 +2,9 @@
 import { z } from '@bpinternal/zui'
 
 import { clamp } from 'lodash-es'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
 import { fastHash, stringify, takeUntilTokens } from '../utils'
 import { Zai } from '../zai'
 import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
@@ -39,22 +42,28 @@ const _Options = z.object({
 declare module '@botpress/zai' {
   interface Zai {
     /** Filters elements of an array against a condition */
-    filter<T>(input: Array<T>, condition: string, options?: Options): Promise<Array<T>>
+    filter<T>(input: Array<T>, condition: string, options?: Options): Response<Array<T>>
   }
 }
 
 const END = '■END■'
 
-Zai.prototype.filter = async function (this: Zai, input, condition, _options) {
+const filter = async <T>(
+  input: Array<T>,
+  condition: string,
+  _options: Options | undefined,
+  ctx: ZaiContext
+): Promise<Array<T>> => {
+  ctx.controller.signal.throwIfAborted()
   const options = _Options.parse(_options ?? {}) as Options
-  const tokenizer = await this.getTokenizer()
-  await this.fetchModelDetails()
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
 
-  const taskId = this.taskId
+  const taskId = ctx.taskId
   const taskType = 'zai.filter'
 
   const MAX_ITEMS_PER_CHUNK = 50
-  const TOKENS_TOTAL_MAX = this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
   const TOKENS_EXAMPLES_MAX = Math.floor(Math.max(250, TOKENS_TOTAL_MAX * 0.5))
   const TOKENS_CONDITION_MAX = clamp(TOKENS_TOTAL_MAX * 0.25, 250, tokenizer.count(condition))
   const TOKENS_INPUT_ARRAY_MAX = TOKENS_TOTAL_MAX - TOKENS_EXAMPLES_MAX - TOKENS_CONDITION_MAX
@@ -145,18 +154,19 @@ ${examples.map((x, idx) => `■${idx}:${!!x.filter ? 'true' : 'false'}:${x.reaso
   ]
 
   const filterChunk = async (chunk: typeof input) => {
-    const examples = taskId
-      ? await this.adapter
-          .getExamples<string, unknown>({
-            // The Table API can't search for a huge input string
-            input: JSON.stringify(chunk).slice(0, 1000),
-            taskType,
-            taskId,
-          })
-          .then((x) =>
-            x.map((y) => ({ filter: y.output as boolean, input: y.input, reason: y.explanation }) satisfies Example)
-          )
-      : []
+    const examples =
+      taskId && ctx.adapter
+        ? await ctx.adapter
+            .getExamples<string, unknown>({
+              // The Table API can't search for a huge input string
+              input: JSON.stringify(chunk).slice(0, 1000),
+              taskType,
+              taskId,
+            })
+            .then((x) =>
+              x.map((y) => ({ filter: y.output as boolean, input: y.input, reason: y.explanation }) satisfies Example)
+            )
+        : []
 
     const allExamples = takeUntilTokens([...examples, ...(options.examples ?? [])], TOKENS_EXAMPLES_MAX, (el) =>
       tokenizer.count(stringify(el.input))
@@ -175,7 +185,7 @@ ${examples.map((x, idx) => `■${idx}:${!!x.filter ? 'true' : 'false'}:${x.reaso
       },
     ]
 
-    const { output, meta } = await this.callModel({
+    const { extracted: partial, meta } = await ctx.generateContent({
       systemPrompt: `
 You are given a list of items. Your task is to filter out the items that meet the condition below.
 You need to return the full list of items with the format:
@@ -198,23 +208,23 @@ The condition is: "${condition}"
           role: 'user',
         },
       ],
-    })
-
-    const answer = output.choices[0]?.content as string
-    const indices = answer
-      .trim()
-      .split('■')
-      .filter((x) => x.length > 0)
-      .map((x) => {
-        const [idx, filter] = x.split(':')
-        return { idx: parseInt(idx?.trim() ?? ''), filter: filter?.toLowerCase().trim() === 'true' }
-      })
+      transform: (text) => {
+        const indices = text
+          .trim()
+          .split('■')
+          .filter((x) => x.length > 0)
+          .map((x) => {
+            const [idx, filter] = x.split(':')
+            return { idx: parseInt(idx?.trim() ?? ''), filter: filter?.toLowerCase().trim() === 'true' }
+          })
 
-    const partial = chunk.filter((_, idx) => {
-      return indices.find((x) => x.idx === idx)?.filter ?? false
+        return chunk.filter((_, idx) => {
+          return indices.find((x) => x.idx === idx && x.filter) ?? false
+        })
+      },
     })
 
-    if (taskId) {
+    if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
       const key = fastHash(
         stringify({
           taskId,
@@ -224,7 +234,7 @@ The condition is: "${condition}"
         })
       )
 
-      await this.adapter.saveExample({
+      await ctx.adapter.saveExample({
         key,
         taskType,
         taskId,
@@ -237,7 +247,7 @@ The condition is: "${condition}"
             output: meta.cost.output,
           },
           latency: meta.latency,
-          model: this.Model,
+          model: ctx.modelId,
           tokens: {
             input: meta.tokens.input,
             output: meta.tokens.output,
@@ -253,3 +263,20 @@ The condition is: "${condition}"
 
   return filteredChunks.flat()
 }
+
+Zai.prototype.filter = function <T>(
+  this: Zai,
+  input: Array<T>,
+  condition: string,
+  _options?: Options
+): Response<Array<T>> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.filter',
+    adapter: this.adapter,
+  })
+
+  return new Response<Array<T>>(context, filter(input, condition, _options, context), (result) => result)
+}