@botpress/zai 1.1.0 → 1.2.0

package/src/operations/extract.ts

@@ -1,4 +1,5 @@
-import { z } from '@bpinternal/zui'
+// eslint-disable consistent-type-definitions
+import { z, ZodObject } from '@bpinternal/zui'
 
 import JSON5 from 'json5'
 import { jsonrepair } from 'jsonrepair'
@@ -18,18 +19,13 @@ const Options = z.object({
     .max(100_000)
     .optional()
     .describe('The maximum number of tokens per chunk')
-    .default(16_000)
+    .default(16_000),
 })
 
 declare module '@botpress/zai' {
   interface Zai {
     /** Extracts one or many elements from an arbitrary input */
-    extract<S extends z.AnyZodObject>(input: unknown, schema: S, options?: Options): Promise<z.infer<S>>
-    extract<S extends z.AnyZodObject>(
-      input: unknown,
-      schema: z.ZodArray<S>,
-      options?: Options
-    ): Promise<Array<z.infer<S>>>
+    extract<S extends z.AnyZodObject | z.ZodArray>(input: unknown, schema: S, options?: Options): Promise<z.TypeOf<S>>
   }
 }
 
@@ -40,21 +36,29 @@ const NO_MORE = '■NO_MORE_ELEMENT■'
 Zai.prototype.extract = async function (this: Zai, input, schema, _options) {
   const options = Options.parse(_options ?? {})
   const tokenizer = await this.getTokenizer()
+  await this.fetchModelDetails()
 
   const taskId = this.taskId
   const taskType = 'zai.extract'
 
-  const PROMPT_COMPONENT = Math.max(this.Model.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
+  const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
 
   let isArrayOfObjects = false
   const originalSchema = schema
 
-  if (schema instanceof z.ZodObject) {
+  const baseType = (schema.naked ? schema.naked() : schema)?.constructor?.name ?? 'unknown'
+
+  if (baseType === 'ZodObject') {
     // Do nothing
-  } else if (schema instanceof z.ZodArray) {
-    if (schema._def.type instanceof z.ZodObject) {
+  } else if (baseType === 'ZodArray') {
+    let elementType = (schema as any).element
+    if (elementType.naked) {
+      elementType = elementType.naked()
+    }
+
+    if (elementType?.constructor?.name === 'ZodObject') {
       isArrayOfObjects = true
-      schema = schema._def.type
+      schema = elementType
     } else {
       throw new Error('Schema must be a ZodObject or a ZodArray<ZodObject>')
     }
@@ -65,9 +69,12 @@ Zai.prototype.extract = async function (this: Zai, input, schema, _options) {
   const schemaTypescript = schema.toTypescript({ declaration: false })
   const schemaLength = tokenizer.count(schemaTypescript)
 
-  options.chunkLength = Math.min(options.chunkLength, this.Model.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength)
+  options.chunkLength = Math.min(
+    options.chunkLength,
+    this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength
+  )
 
-  const keys = Object.keys(schema.shape)
+  const keys = Object.keys((schema as ZodObject).shape)
 
   let inputAsString = stringify(input)
 
@@ -116,7 +123,7 @@ Zai.prototype.extract = async function (this: Zai, input, schema, _options) {
       taskType,
       taskId,
       input: inputAsString,
-      instructions: options.instructions
+      instructions: options.instructions,
     })
   )
 
@@ -124,7 +131,7 @@ Zai.prototype.extract = async function (this: Zai, input, schema, _options) {
     ? await this.adapter.getExamples<string, unknown>({
         input: inputAsString,
         taskType,
-        taskId
+        taskId,
       })
     : []
 
@@ -144,13 +151,13 @@ The end.`,
         extracted: [
           {
             name: 'Alice',
-            age: 30
+            age: 30,
           },
           {
             name: 'Bob',
-            age: 25
-          }
-        ]
+            age: 25,
+          },
+        ],
       }
     : {
         input: `The story goes as follow.
@@ -158,14 +165,14 @@ Once upon a time, there was a person named Alice who was 30 years old.
 The end.`,
         schema: '{ name: string, age: number }',
         instructions: 'Extract the person',
-        extracted: { name: 'Alice', age: 30 }
+        extracted: { name: 'Alice', age: 30 },
       }
 
   const userExamples = examples.map((e) => ({
     input: e.input,
     extracted: e.output,
     schema: schemaTypescript,
-    instructions: options.instructions
+    instructions: options.instructions,
   }))
 
   let exampleId = 1
@@ -211,13 +218,13 @@ ${END}`.trim()
     {
       type: 'text' as const,
       content: formatInput(stringify(example.input ?? null), example.schema, example.instructions),
-      role: 'user' as const
+      role: 'user' as const,
     },
     {
       type: 'text' as const,
       content: formatOutput(example.extracted),
-      role: 'assistant' as const
-    }
+      role: 'assistant' as const,
+    },
   ]
 
   const allExamples = takeUntilTokens(
@@ -228,7 +235,7 @@ ${END}`.trim()
     .map(formatExample)
     .flat()
 
-  const output = await this.callModel({
+  const { output, meta } = await this.callModel({
     systemPrompt: `
 Extract the following information from the input:
 ${schemaTypescript}
@@ -242,9 +249,9 @@ ${instructions.map((x) => `• ${x}`).join('\n')}
       {
         role: 'user',
         type: 'text',
-        content: formatInput(inputAsString, schemaTypescript, options.instructions ?? '')
-      }
-    ]
+        content: formatInput(inputAsString, schemaTypescript, options.instructions ?? ''),
+      },
+    ],
   })
 
   const answer = output.choices[0]?.content as string
@@ -283,7 +290,18 @@ ${instructions.map((x) => `• ${x}`).join('\n')}
       instructions: options.instructions ?? 'No specific instructions',
       input: inputAsString,
       output: final,
-      metadata: output.metadata
+      metadata: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output,
+        },
+        latency: meta.latency,
+        model: this.Model,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output,
+        },
+      },
     })
   }
 

package/src/operations/filter.ts

@@ -1,3 +1,4 @@
+// eslint-disable consistent-type-definitions
 import { z } from '@bpinternal/zui'
 
 import { clamp } from 'lodash-es'
@@ -9,7 +10,7 @@ type Example = z.input<typeof Example>
 const Example = z.object({
   input: z.any(),
   filter: z.boolean(),
-  reason: z.string().optional()
+  reason: z.string().optional(),
 })
 
 export type Options = z.input<typeof Options>
@@ -21,7 +22,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' {
@@ -36,12 +37,13 @@ const END = '■END■'
 Zai.prototype.filter = async function (this: Zai, input, condition, _options) {
   const options = Options.parse(_options ?? {})
   const tokenizer = await this.getTokenizer()
+  await this.fetchModelDetails()
 
   const taskId = this.taskId
   const taskType = 'zai.filter'
 
   const MAX_ITEMS_PER_CHUNK = 50
-  const TOKENS_TOTAL_MAX = this.Model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
+  const TOKENS_TOTAL_MAX = this.ModelDetails.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
@@ -99,36 +101,36 @@ ${examples.map((x, idx) => `■${idx}:${!!x.filter ? 'true' : 'false'}:${x.reaso
     {
       input: 'apple',
       filter: true,
-      reason: 'Apples are fruits'
+      reason: 'Apples are fruits',
     },
     {
       input: 'Apple Inc.',
       filter: false,
-      reason: 'Apple Inc. is a company, not a fruit'
+      reason: 'Apple Inc. is a company, not a fruit',
     },
     {
       input: 'banana',
       filter: true,
-      reason: 'Bananas are fruits'
+      reason: 'Bananas are fruits',
     },
     {
       input: 'potato',
       filter: false,
-      reason: 'Potatoes are vegetables'
-    }
+      reason: 'Potatoes are vegetables',
+    },
   ]
 
   const genericExamplesMessages = [
     {
       type: 'text' as const,
       content: formatInput(genericExamples, 'is a fruit'),
-      role: 'user' as const
+      role: 'user' as const,
     },
     {
       type: 'text' as const,
       content: formatExamples(genericExamples),
-      role: 'assistant' as const
-    }
+      role: 'assistant' as const,
+    },
   ]
 
   const filterChunk = async (chunk: typeof input) => {
@@ -138,10 +140,10 @@ ${examples.map((x, idx) => `■${idx}:${!!x.filter ? 'true' : 'false'}:${x.reaso
             // The Table API can't search for a huge input string
             input: JSON.stringify(chunk).slice(0, 1000),
             taskType,
-            taskId
+            taskId,
           })
           .then((x) =>
-            x.map((y) => ({ filter: y.output as boolean, input: y.input, reason: y.explanation } satisfies Example))
+            x.map((y) => ({ filter: y.output as boolean, input: y.input, reason: y.explanation }) satisfies Example)
           )
       : []
 
@@ -153,16 +155,16 @@ ${examples.map((x, idx) => `■${idx}:${!!x.filter ? 'true' : 'false'}:${x.reaso
       {
         type: 'text' as const,
         content: formatInput(allExamples, condition),
-        role: 'user' as const
+        role: 'user' as const,
       },
       {
         type: 'text' as const,
         content: formatExamples(allExamples),
-        role: 'assistant' as const
-      }
+        role: 'assistant' as const,
+      },
     ]
 
-    const output = await this.callModel({
+    const { output, meta } = await this.callModel({
       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:
@@ -179,12 +181,12 @@ The condition is: "${condition}"
         {
           type: 'text',
           content: formatInput(
-            chunk.map((x) => ({ input: x } as Example)),
+            chunk.map((x) => ({ input: x }) as Example),
             condition
           ),
-          role: 'user'
-        }
-      ]
+          role: 'user',
+        },
+      ],
     })
 
     const answer = output.choices[0]?.content as string
@@ -207,7 +209,7 @@ The condition is: "${condition}"
           taskId,
           taskType,
           input: JSON.stringify(chunk),
-          condition
+          condition,
         })
       )
 
@@ -218,7 +220,18 @@ The condition is: "${condition}"
         input: JSON.stringify(chunk),
         output: partial,
         instructions: condition,
-        metadata: output.metadata
+        metadata: {
+          cost: {
+            input: meta.cost.input,
+            output: meta.cost.output,
+          },
+          latency: meta.latency,
+          model: this.Model,
+          tokens: {
+            input: meta.tokens.input,
+            output: meta.tokens.output,
+          },
+        },
       })
     }
 

package/src/operations/label.ts

@@ -1,3 +1,4 @@
+// eslint-disable consistent-type-definitions
 import { z } from '@bpinternal/zui'
 
 import { clamp, chunk } from 'lodash-es'
@@ -11,7 +12,7 @@ const LABELS = {
   PROBABLY_NOT: 'PROBABLY_NOT',
   AMBIGUOUS: 'AMBIGUOUS',
   PROBABLY_YES: 'PROBABLY_YES',
-  ABSOLUTELY_YES: 'ABSOLUTELY_YES'
+  ABSOLUTELY_YES: 'ABSOLUTELY_YES',
 } as const
 const ALL_LABELS = Object.values(LABELS).join(' | ')
 
@@ -29,7 +30,7 @@ const Options = z.object({
     .array(
       z.object({
         input: z.any(),
-        labels: z.record(z.object({ label: z.enum(ALL_LABELS as never), explanation: z.string().optional() }))
+        labels: z.record(z.object({ label: z.enum(ALL_LABELS as never), explanation: z.string().optional() })),
       })
     )
     .default([])
@@ -41,7 +42,7 @@ const Options = z.object({
     .max(100_000)
     .optional()
     .describe('The maximum number of tokens per chunk')
-    .default(16_000)
+    .default(16_000),
 })
 
 type Labels<T extends string> = Record<T, string>
@@ -61,7 +62,7 @@ const Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((lab
     if (/[^a-zA-Z0-9_]/.test(key)) {
       ctx.addIssue({
         message: `The label key "${key}" must only contain alphanumeric characters and underscores`,
-        code: 'custom'
+        code: 'custom',
       })
     }
   }
@@ -103,11 +104,12 @@ Zai.prototype.label = async function <T extends string>(this: Zai, input, _label
   const options = Options.parse(_options ?? {})
   const labels = Labels.parse(_labels)
   const tokenizer = await this.getTokenizer()
+  await this.fetchModelDetails()
 
   const taskId = this.taskId
   const taskType = 'zai.label'
 
-  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1000, this.Model.input.maxTokens - PROMPT_INPUT_BUFFER)
+  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1000, this.ModelDetails.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,
@@ -144,7 +146,7 @@ Zai.prototype.label = async function <T extends string>(this: Zai, input, _label
       taskType,
       taskId,
       input: inputAsString,
-      instructions: options.instructions ?? ''
+      instructions: options.instructions ?? '',
     })
   )
 
@@ -167,7 +169,7 @@ Zai.prototype.label = async function <T extends string>(this: Zai, input, _label
       >({
         input: inputAsString,
         taskType,
-        taskId
+        taskId,
       })
     : []
 
@@ -182,7 +184,7 @@ Zai.prototype.label = async function <T extends string>(this: Zai, input, _label
           explanation: string
           label: Label
         }
-      }
+      },
     })
   })
 
@@ -209,7 +211,7 @@ Expert Example #${idx + 1}
 
 <|start_input|>
 ${stringify(example.input)}
-<|end_input|>`.trim()
+<|end_input|>`.trim(),
       },
       {
         type: 'text' as const,
@@ -225,8 +227,8 @@ ${Object.keys(example.output)
   )
   .join('\n')}
 ${END}
-`.trim()
-      }
+`.trim(),
+      },
     ])
     .flat()
 
@@ -238,7 +240,7 @@ ${END}
     })
     .join('\n\n')
 
-  const output = await this.callModel({
+  const { output, meta } = await this.callModel({
     stopSequences: [END],
     systemPrompt: `
 You need to tag the input with the following labels based on the question asked:
@@ -286,9 +288,9 @@ Where \`x\` is one of the following: ${ALL_LABELS}
 
 Remember: In your \`explanation\`, please refer to the Expert Examples # (and quote them) that are relevant to ground your decision-making process.
 The Expert Examples are there to help you make your decision. They have been provided by experts in the field and their answers (and reasoning) are considered the ground truth and should be used as a reference to make your decision when applicable.
-For example, you can say: "According to Expert Example #1, ..."`.trim()
-      }
-    ]
+For example, you can say: "According to Expert Example #1, ..."`.trim(),
+      },
+    ],
   })
 
   const answer = output.choices[0].content as string
@@ -300,12 +302,12 @@ For example, you can say: "According to Expert Example #1, ..."`.trim()
       const label = parseLabel(match[2])
       acc[key] = {
         explanation,
-        label
+        label,
       }
     } else {
       acc[key] = {
         explanation: '',
-        label: LABELS.AMBIGUOUS
+        label: LABELS.AMBIGUOUS,
       }
     }
     return acc
@@ -322,9 +324,20 @@ For example, you can say: "According to Expert Example #1, ..."`.trim()
       taskType,
       taskId,
       instructions: options.instructions ?? '',
-      metadata: output.metadata,
+      metadata: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output,
+        },
+        latency: meta.latency,
+        model: this.Model,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output,
+        },
+      },
       input: inputAsString,
-      output: final
+      output: final,
     })
   }
 

package/src/operations/rewrite.ts

@@ -1,3 +1,4 @@
+// eslint-disable consistent-type-definitions
 import { z } from '@bpinternal/zui'
 
 import { fastHash, stringify, takeUntilTokens } from '../utils'
@@ -7,13 +8,13 @@ import { PROMPT_INPUT_BUFFER } from './constants'
 type Example = z.input<typeof Example> & { instructions?: string }
 const Example = z.object({
   input: z.string(),
-  output: z.string()
+  output: z.string(),
 })
 
 export type Options = z.input<typeof Options>
 const Options = z.object({
   examples: z.array(Example).default([]),
-  length: z.number().min(10).max(16_000).optional().describe('The maximum number of tokens to generate')
+  length: z.number().min(10).max(16_000).optional().describe('The maximum number of tokens to generate'),
 })
 
 declare module '@botpress/zai' {
@@ -29,18 +30,19 @@ const END = '■END■'
 Zai.prototype.rewrite = async function (this: Zai, original, prompt, _options) {
   const options = Options.parse(_options ?? {})
   const tokenizer = await this.getTokenizer()
+  await this.fetchModelDetails()
 
   const taskId = this.taskId
   const taskType = 'zai.rewrite'
 
-  const INPUT_COMPONENT_SIZE = Math.max(100, (this.Model.input.maxTokens - PROMPT_INPUT_BUFFER) / 2)
+  const INPUT_COMPONENT_SIZE = Math.max(100, (this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER) / 2)
   prompt = tokenizer.truncate(prompt, INPUT_COMPONENT_SIZE)
 
   const inputSize = tokenizer.count(original) + tokenizer.count(prompt)
-  const maxInputSize = this.Model.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER
+  const maxInputSize = this.ModelDetails.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER
   if (inputSize > maxInputSize) {
     throw new Error(
-      `The input size is ${inputSize} tokens long, which is more than the maximum of ${maxInputSize} tokens for this model (${this.Model.name} = ${this.Model.input.maxTokens} tokens)`
+      `The input size is ${inputSize} tokens long, which is more than the maximum of ${maxInputSize} tokens for this model (${this.ModelDetails.name} = ${this.ModelDetails.input.maxTokens} tokens)`
     )
   }
 
@@ -69,27 +71,27 @@ ${END}
       taskId,
       taskType,
       input: original,
-      prompt
+      prompt,
     })
   )
 
   const formatExample = ({ input, output, instructions }: Example) => {
     return [
       { type: 'text' as const, role: 'user' as const, content: format(input, instructions || prompt) },
-      { type: 'text' as const, role: 'assistant' as const, content: `${START}${output}${END}` }
+      { type: 'text' as const, role: 'assistant' as const, content: `${START}${output}${END}` },
     ]
   }
 
   const defaultExamples: Example[] = [
     { input: 'Hello, how are you?', output: 'Bonjour, comment ça va?', instructions: 'translate to French' },
-    { input: '1\n2\n3', output: '3\n2\n1', instructions: 'reverse the order' }
+    { input: '1\n2\n3', output: '3\n2\n1', instructions: 'reverse the order' },
   ]
 
   const tableExamples = taskId
     ? await this.adapter.getExamples<string, string>({
         input: original,
         taskId,
-        taskType
+        taskType,
       })
     : []
 
@@ -100,10 +102,10 @@ ${END}
 
   const savedExamples: Example[] = [
     ...tableExamples.map((x) => ({ input: x.input as string, output: x.output as string })),
-    ...options.examples
+    ...options.examples,
   ]
 
-  const REMAINING_TOKENS = this.Model.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER
+  const REMAINING_TOKENS = this.ModelDetails.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER
   const examples = takeUntilTokens(
     savedExamples.length ? savedExamples : defaultExamples,
     REMAINING_TOKENS,
@@ -112,14 +114,14 @@ ${END}
     .map(formatExample)
     .flat()
 
-  const output = await this.callModel({
+  const { output, meta } = await this.callModel({
     systemPrompt: `
 Rewrite the text between the ${START} and ${END} tags to match the user prompt.
 ${instructions.map((x) => `• ${x}`).join('\n')}
 `.trim(),
     messages: [...examples, { type: 'text', content: format(original, prompt), role: 'user' }],
     maxTokens: options.length,
-    stopSequences: [END]
+    stopSequences: [END],
   })
 
   let result = output.choices[0]?.content as string
@@ -135,12 +137,23 @@ ${instructions.map((x) => `• ${x}`).join('\n')}
   if (taskId) {
     await this.adapter.saveExample({
       key: Key,
-      metadata: output.metadata,
+      metadata: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output,
+        },
+        latency: meta.latency,
+        model: this.Model,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output,
+        },
+      },
       instructions: prompt,
       input: original,
       output: result,
       taskType,
-      taskId
+      taskId,
     })
   }