@botpress/zai 2.0.16 → 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'
@@ -38,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<any>>(input: unknown, schema: S, options?: Options): Promise<S['_output']>
+    extract<S extends OfType<any>>(input: unknown, schema: S, options?: Options): Response<S['_output']>
   }
 }
 
@@ -46,21 +49,22 @@ 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
@@ -99,10 +103,7 @@ Zai.prototype.extract = async function <S extends OfType<AnyObjectOrArray>>(
   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)
 
@@ -113,18 +114,25 @@ Zai.prototype.extract = async function <S extends OfType<AnyObjectOrArray>>(
     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
-        })
+        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
     const rows = all.map((x, idx) => `<part-${idx + 1}>\n${stringify(x, true)}\n</part-${idx + 1}>`).join('\n')
-    return this.extract(
+    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.
@@ -136,7 +144,8 @@ ${rows}
 
 Merge it back into a final result.`.trim(),
       originalSchema,
-      options
+      options,
+      ctx
     )
   }
 
@@ -179,13 +188,14 @@ Merge it back into a final result.`.trim(),
     })
   )
 
-  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) {
@@ -287,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}
@@ -304,43 +314,41 @@ ${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 && 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)
-
   let final: any
 
   if (isArrayOfObjects) {
-    final = elements
-  } else if (elements.length === 0) {
+    final = extracted
+  } else if (extracted.length === 0) {
     final = options.strict ? schema.parse({}) : {}
   } else {
-    final = elements[0]
+    final = extracted[0]
   }
 
   if (wrappedValue) {
@@ -351,8 +359,8 @@ ${instructions.map((x) => `• ${x}`).join('\n')}
     }
   }
 
-  if (taskId) {
-    await this.adapter.saveExample({
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
       key: Key,
       taskId: `zai/${taskId}`,
       taskType,
@@ -365,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,
@@ -376,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)
+}

package/src/operations/label.ts

@@ -1,7 +1,10 @@
 // eslint-disable consistent-type-definitions
 import { z } from '@bpinternal/zui'
 
-import { clamp, chunk } from 'lodash-es'
+import { chunk, 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 } from './constants'
@@ -83,13 +86,16 @@ declare module '@botpress/zai' {
       input: unknown,
       labels: Labels<T>,
       options?: Options<T>
-    ): Promise<{
-      [K in T]: {
-        explanation: string
-        value: boolean
-        confidence: number
-      }
-    }>
+    ): Response<
+      {
+        [K in T]: {
+          explanation: string
+          value: boolean
+          confidence: number
+        }
+      },
+      { [K in T]: boolean }
+    >
   }
 }
 
@@ -124,21 +130,28 @@ const getConfidence = (label: Label) => {
   }
 }
 
-Zai.prototype.label = async function <T extends string>(
-  this: Zai,
+const label = async <T extends string>(
   input: unknown,
   _labels: Labels<T>,
-  _options: Options<T> | undefined
-) {
+  _options: Options<T> | undefined,
+  ctx: ZaiContext
+): Promise<{
+  [K in T]: {
+    explanation: string
+    value: boolean
+    confidence: number
+  }
+}> => {
+  ctx.controller.signal.throwIfAborted()
   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()
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
 
-  const taskId = this.taskId
+  const taskId = ctx.taskId
   const taskType = 'zai.label'
 
-  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1000, this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER)
+  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1000, model.input.maxTokens - PROMPT_INPUT_BUFFER)
   const CHUNK_EXAMPLES_MAX_TOKENS = clamp(Math.floor(TOTAL_MAX_TOKENS * 0.5), 250, 10_000)
   const CHUNK_INPUT_MAX_TOKENS = clamp(
     TOTAL_MAX_TOKENS - CHUNK_EXAMPLES_MAX_TOKENS,
@@ -151,7 +164,7 @@ Zai.prototype.label = async function <T extends string>(
   if (tokenizer.count(inputAsString) > CHUNK_INPUT_MAX_TOKENS) {
     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) => this.label(chunk, _labels)))
+    const allLabels = await Promise.all(chunks.map((chunk) => label(chunk, _labels, _options, ctx)))
 
     // Merge all the labels together (those who are true will remain true)
     return allLabels.reduce((acc, x) => {
@@ -202,21 +215,22 @@ Zai.prototype.label = async function <T extends string>(
     }
   }
 
-  const examples = taskId
-    ? await this.adapter.getExamples<
-        string,
-        {
-          [K in T]: {
-            explanation: string
-            label: Label
+  const examples =
+    taskId && ctx.adapter
+      ? await ctx.adapter.getExamples<
+          string,
+          {
+            [K in T]: {
+              explanation: string
+              label: Label
+            }
           }
-        }
-      >({
-        input: inputAsString,
-        taskType,
-        taskId,
-      })
-    : []
+        >({
+          input: inputAsString,
+          taskType,
+          taskId,
+        })
+      : []
 
   options.examples.forEach((example) => {
     examples.push({
@@ -285,7 +299,7 @@ ${END}
     })
     .join('\n\n')
 
-  const { output, meta } = await this.callModel({
+  const { extracted, meta } = await ctx.generateContent({
     stopSequences: [END],
     systemPrompt: `
 You need to tag the input with the following labels based on the question asked:
@@ -336,35 +350,33 @@ The Expert Examples are there to help you make your decision. They have been pro
 For example, you can say: "According to Expert Example #1, ..."`.trim(),
       },
     ],
+    transform: (text) =>
+      Object.keys(labels).reduce((acc, key) => {
+        const match = text.match(new RegExp(`■${key}:【(.+)】:(\\w{2,})■`, 'i'))
+        if (match) {
+          const explanation = match[1].trim()
+          const label = parseLabel(match[2])
+          acc[key] = {
+            explanation,
+            label,
+          }
+        } else {
+          acc[key] = {
+            explanation: '',
+            label: LABELS.AMBIGUOUS,
+          }
+        }
+        return acc
+      }, {}) as {
+        [K in T]: {
+          explanation: string
+          label: Label
+        }
+      },
   })
 
-  const answer = output.choices[0].content as string
-
-  const final = Object.keys(labels).reduce((acc, key) => {
-    const match = answer.match(new RegExp(`■${key}:【(.+)】:(\\w{2,})■`, 'i'))
-    if (match) {
-      const explanation = match[1].trim()
-      const label = parseLabel(match[2])
-      acc[key] = {
-        explanation,
-        label,
-      }
-    } else {
-      acc[key] = {
-        explanation: '',
-        label: LABELS.AMBIGUOUS,
-      }
-    }
-    return acc
-  }, {}) as {
-    [K in T]: {
-      explanation: string
-      label: Label
-    }
-  }
-
-  if (taskId) {
-    await this.adapter.saveExample({
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
       key: Key,
       taskType,
       taskId,
@@ -375,16 +387,59 @@ For example, you can say: "According to Expert Example #1, ..."`.trim(),
           output: meta.cost.output,
         },
         latency: meta.latency,
-        model: this.Model,
+        model: ctx.modelId,
         tokens: {
           input: meta.tokens.input,
           output: meta.tokens.output,
         },
       },
       input: inputAsString,
-      output: final,
+      output: extracted,
     })
   }
 
-  return convertToAnswer(final)
+  return convertToAnswer(extracted)
+}
+
+Zai.prototype.label = function <T extends string>(
+  this: Zai,
+  input: unknown,
+  labels: Labels<T>,
+  _options?: Options<T>
+): Response<
+  {
+    [K in T]: {
+      explanation: string
+      value: boolean
+      confidence: number
+    }
+  },
+  { [K in T]: boolean }
+> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.label',
+    adapter: this.adapter,
+  })
+
+  return new Response<
+    {
+      [K in T]: {
+        explanation: string
+        value: boolean
+        confidence: number
+      }
+    },
+    { [K in T]: boolean }
+  >(context, label(input, labels, _options, context), (result) =>
+    Object.keys(result).reduce(
+      (acc, key) => {
+        acc[key] = result[key].value
+        return acc
+      },
+      {} as { [K in T]: boolean }
+    )
+  )
 }