@botpress/zai 1.0.1 → 1.1.0

package/src/operations/filter.ts

@@ -0,0 +1,231 @@
+import { z } from '@bpinternal/zui'
+
+import { clamp } from 'lodash-es'
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+
+type Example = z.input<typeof Example>
+const Example = z.object({
+  input: z.any(),
+  filter: z.boolean(),
+  reason: z.string().optional()
+})
+
+export type Options = z.input<typeof Options>
+const Options = z.object({
+  tokensPerItem: z
+    .number()
+    .min(1)
+    .max(100_000)
+    .optional()
+    .describe('The maximum number of tokens per item')
+    .default(250),
+  examples: z.array(Example).describe('Examples to filter the condition against').default([])
+})
+
+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>>
+  }
+}
+
+const END = '■END■'
+
+Zai.prototype.filter = async function (this: Zai, input, condition, _options) {
+  const options = Options.parse(_options ?? {})
+  const tokenizer = await this.getTokenizer()
+
+  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_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
+
+  condition = tokenizer.truncate(condition, TOKENS_CONDITION_MAX)
+
+  let chunks: Array<typeof input> = []
+  let currentChunk: typeof input = []
+  let currentChunkTokens = 0
+
+  for (const element of input) {
+    const elementAsString = tokenizer.truncate(stringify(element, false), options.tokensPerItem)
+    const elementTokens = tokenizer.count(elementAsString)
+
+    if (currentChunkTokens + elementTokens > TOKENS_INPUT_ARRAY_MAX || currentChunk.length >= MAX_ITEMS_PER_CHUNK) {
+      chunks.push(currentChunk)
+      currentChunk = []
+      currentChunkTokens = 0
+    }
+
+    currentChunk.push(element)
+    currentChunkTokens += elementTokens
+  }
+
+  if (currentChunk.length > 0) {
+    chunks.push(currentChunk)
+  }
+
+  chunks = chunks.filter((x) => x.length > 0)
+
+  // ■1:true■2:true■3:true
+
+  const formatInput = (input: Example[], condition: string) => {
+    return `
+Condition to check:
+${condition}
+
+Items (from ■0 to ■${input.length - 1})
+==============================
+${input.map((x, idx) => `■${idx} = ${stringify(x.input ?? null, false)}`).join('\n')}
+`.trim()
+  }
+
+  const formatExamples = (examples: Example[]) => {
+    return `
+${examples.map((x, idx) => `■${idx}:${!!x.filter ? 'true' : 'false'}`).join('')}
+${END}
+====
+Here's the reasoning behind each example:
+${examples.map((x, idx) => `■${idx}:${!!x.filter ? 'true' : 'false'}:${x.reason ?? 'No reason provided'}`).join('\n')}
+`.trim()
+  }
+
+  const genericExamples: Example[] = [
+    {
+      input: 'apple',
+      filter: true,
+      reason: 'Apples are fruits'
+    },
+    {
+      input: 'Apple Inc.',
+      filter: false,
+      reason: 'Apple Inc. is a company, not a fruit'
+    },
+    {
+      input: 'banana',
+      filter: true,
+      reason: 'Bananas are fruits'
+    },
+    {
+      input: 'potato',
+      filter: false,
+      reason: 'Potatoes are vegetables'
+    }
+  ]
+
+  const genericExamplesMessages = [
+    {
+      type: 'text' as const,
+      content: formatInput(genericExamples, 'is a fruit'),
+      role: 'user' as const
+    },
+    {
+      type: 'text' as const,
+      content: formatExamples(genericExamples),
+      role: 'assistant' as const
+    }
+  ]
+
+  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 allExamples = takeUntilTokens([...examples, ...(options.examples ?? [])], TOKENS_EXAMPLES_MAX, (el) =>
+      tokenizer.count(stringify(el.input))
+    )
+
+    const exampleMessages = [
+      {
+        type: 'text' as const,
+        content: formatInput(allExamples, condition),
+        role: 'user' as const
+      },
+      {
+        type: 'text' as const,
+        content: formatExamples(allExamples),
+        role: 'assistant' as const
+      }
+    ]
+
+    const output = 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:
+■x:true■y:false■z:true (where x, y, z are the indices of the items in the list)
+You need to start with "■0" and go up to the last index "■${chunk.length - 1}".
+If an item meets the condition, you should return ":true", otherwise ":false".
+
+IMPORTANT: Make sure to read the condition and the examples carefully before making your decision.
+The condition is: "${condition}"
+`.trim(),
+      stopSequences: [END],
+      messages: [
+        ...(exampleMessages.length ? exampleMessages : genericExamplesMessages),
+        {
+          type: 'text',
+          content: formatInput(
+            chunk.map((x) => ({ input: x } as Example)),
+            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' }
+      })
+
+    const partial = chunk.filter((_, idx) => {
+      return indices.find((x) => x.idx === idx)?.filter ?? false
+    })
+
+    if (taskId) {
+      const key = fastHash(
+        stringify({
+          taskId,
+          taskType,
+          input: JSON.stringify(chunk),
+          condition
+        })
+      )
+
+      await this.adapter.saveExample({
+        key,
+        taskType,
+        taskId,
+        input: JSON.stringify(chunk),
+        output: partial,
+        instructions: condition,
+        metadata: output.metadata
+      })
+    }
+
+    return partial
+  }
+
+  const filteredChunks = await Promise.all(chunks.map(filterChunk))
+
+  return filteredChunks.flat()
+}

package/src/operations/label.ts

@@ -0,0 +1,332 @@
+import { z } from '@bpinternal/zui'
+
+import { clamp, chunk } from 'lodash-es'
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER } from './constants'
+
+type Label = keyof typeof LABELS
+const LABELS = {
+  ABSOLUTELY_NOT: 'ABSOLUTELY_NOT',
+  PROBABLY_NOT: 'PROBABLY_NOT',
+  AMBIGUOUS: 'AMBIGUOUS',
+  PROBABLY_YES: 'PROBABLY_YES',
+  ABSOLUTELY_YES: 'ABSOLUTELY_YES'
+} as const
+const ALL_LABELS = Object.values(LABELS).join(' | ')
+
+type Example<T extends string> = {
+  input: unknown
+  labels: Partial<Record<T, { label: Label; explanation?: string }>>
+}
+
+export type Options<T extends string> = Omit<z.input<typeof Options>, 'examples'> & {
+  examples?: Array<Partial<Example<T>>>
+}
+
+const Options = z.object({
+  examples: z
+    .array(
+      z.object({
+        input: z.any(),
+        labels: z.record(z.object({ label: z.enum(ALL_LABELS as never), explanation: z.string().optional() }))
+      })
+    )
+    .default([])
+    .describe('Examples to help the user make a decision'),
+  instructions: z.string().optional().describe('Instructions to guide the user on how to extract the data'),
+  chunkLength: z
+    .number()
+    .min(100)
+    .max(100_000)
+    .optional()
+    .describe('The maximum number of tokens per chunk')
+    .default(16_000)
+})
+
+type Labels<T extends string> = Record<T, string>
+
+const Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((labels, ctx) => {
+  const keys = Object.keys(labels)
+
+  for (const key of keys) {
+    if (key.length < 1 || key.length > 250) {
+      ctx.addIssue({ message: `The label key "${key}" must be between 1 and 250 characters long`, code: 'custom' })
+    }
+
+    if (keys.lastIndexOf(key) !== keys.indexOf(key)) {
+      ctx.addIssue({ message: `Duplicate label: ${labels[key]}`, code: 'custom' })
+    }
+
+    if (/[^a-zA-Z0-9_]/.test(key)) {
+      ctx.addIssue({
+        message: `The label key "${key}" must only contain alphanumeric characters and underscores`,
+        code: 'custom'
+      })
+    }
+  }
+
+  return true
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /** Tags the provided input with a list of predefined labels */
+    label<T extends string>(
+      input: unknown,
+      labels: Labels<T>,
+      options?: Options<T>
+    ): Promise<{
+      [K in T]: boolean
+    }>
+  }
+}
+
+const parseLabel = (label: string): Label => {
+  label = label.toUpperCase().replace(/\s+/g, '_').replace(/_{2,}/g, '_').trim()
+  if (label.includes('ABSOLUTELY') && label.includes('NOT')) {
+    return LABELS.ABSOLUTELY_NOT
+  } else if (label.includes('NOT')) {
+    return LABELS.PROBABLY_NOT
+  } else if (label.includes('AMBIGUOUS')) {
+    return LABELS.AMBIGUOUS
+  }
+  if (label.includes('YES')) {
+    return LABELS.PROBABLY_YES
+  } else if (label.includes('ABSOLUTELY') && label.includes('YES')) {
+    return LABELS.ABSOLUTELY_YES
+  }
+  return LABELS.AMBIGUOUS
+}
+
+Zai.prototype.label = async function <T extends string>(this: Zai, input, _labels, _options) {
+  const options = Options.parse(_options ?? {})
+  const labels = Labels.parse(_labels)
+  const tokenizer = await this.getTokenizer()
+
+  const taskId = this.taskId
+  const taskType = 'zai.label'
+
+  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1000, this.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,
+    TOTAL_MAX_TOKENS * 0.5,
+    TOTAL_MAX_TOKENS
+  )
+
+  const inputAsString = stringify(input)
+
+  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)))
+
+    // Merge all the labels together (those who are true will remain true)
+    return allLabels.reduce((acc, x) => {
+      Object.keys(x).forEach((key) => {
+        if (acc[key] === true) {
+          acc[key] = true
+        } else {
+          acc[key] = acc[key] || x[key]
+        }
+      })
+      return acc
+    }, {}) as {
+      [K in T]: boolean
+    }
+  }
+
+  const END = '■END■'
+
+  const Key = fastHash(
+    JSON.stringify({
+      taskType,
+      taskId,
+      input: inputAsString,
+      instructions: options.instructions ?? ''
+    })
+  )
+
+  const convertToAnswer = (mapping: { [K in T]: { explanation: string; label: Label } }) => {
+    return Object.keys(labels).reduce((acc, key) => {
+      acc[key] = mapping[key]?.label === 'ABSOLUTELY_YES' || mapping[key]?.label === 'PROBABLY_YES'
+      return acc
+    }, {}) as { [K in T]: boolean }
+  }
+
+  const examples = taskId
+    ? await this.adapter.getExamples<
+        string,
+        {
+          [K in T]: {
+            explanation: string
+            label: Label
+          }
+        }
+      >({
+        input: inputAsString,
+        taskType,
+        taskId
+      })
+    : []
+
+  options.examples.forEach((example) => {
+    examples.push({
+      key: fastHash(JSON.stringify(example)),
+      input: example.input,
+      similarity: 1,
+      explanation: '',
+      output: example.labels as unknown as {
+        [K in T]: {
+          explanation: string
+          label: Label
+        }
+      }
+    })
+  })
+
+  const exactMatch = examples.find((x) => x.key === Key)
+  if (exactMatch) {
+    return convertToAnswer(exactMatch.output)
+  }
+
+  const allExamples = takeUntilTokens(
+    examples,
+    CHUNK_EXAMPLES_MAX_TOKENS,
+    (el) =>
+      tokenizer.count(stringify(el.input)) +
+      tokenizer.count(stringify(el.output)) +
+      tokenizer.count(el.explanation ?? '') +
+      100
+  )
+    .map((example, idx) => [
+      {
+        type: 'text' as const,
+        role: 'user' as const,
+        content: `
+Expert Example #${idx + 1}
+
+<|start_input|>
+${stringify(example.input)}
+<|end_input|>`.trim()
+      },
+      {
+        type: 'text' as const,
+        role: 'assistant' as const,
+        content: `
+Expert Example #${idx + 1}
+============
+${Object.keys(example.output)
+  .map((key) =>
+    `
+■${key}:【${example.output[key]?.explanation}】:${example.output[key]?.label}■
+`.trim()
+  )
+  .join('\n')}
+${END}
+`.trim()
+      }
+    ])
+    .flat()
+
+  const format = Object.keys(labels)
+    .map((key) => {
+      return `
+■${key}:【explanation (where "explanation" is answering the question "${labels[key]}")】:x■ (where x is ${ALL_LABELS})
+`.trim()
+    })
+    .join('\n\n')
+
+  const output = await this.callModel({
+    stopSequences: [END],
+    systemPrompt: `
+You need to tag the input with the following labels based on the question asked:
+${LABELS.ABSOLUTELY_NOT}: You are absolutely sure that the answer is "NO" to the question.
+${LABELS.PROBABLY_NOT}: You are leaning towards "NO" to the question.
+${LABELS.AMBIGUOUS}: You are unsure about the answer to the question.
+${LABELS.PROBABLY_YES}: You are leaning towards "YES" to the question.
+${LABELS.ABSOLUTELY_YES}: You are absolutely sure that the answer is "YES" to the question.
+
+You need to return a mapping of the labels, an explanation and the answer for each label following the format below:
+\`\`\`
+${format}
+${END}
+\`\`\`
+
+${options.instructions}
+
+===
+You should consider the Expert Examples below to help you make your decision.
+In your "Analysis", please refer to the Expert Examples # to justify your decision.
+`.trim(),
+    messages: [
+      ...allExamples,
+      {
+        type: 'text',
+        role: 'user',
+        content: `
+Input to tag:
+<|start_input|>
+${inputAsString}
+<|end_input|>
+
+Answer with this following format:
+\`\`\`
+${format}
+${END}
+\`\`\`
+
+Format cheatsheet:
+\`\`\`
+■label:【explanation】:x■
+\`\`\`
+
+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()
+      }
+    ]
+  })
+
+  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({
+      key: Key,
+      taskType,
+      taskId,
+      instructions: options.instructions ?? '',
+      metadata: output.metadata,
+      input: inputAsString,
+      output: final
+    })
+  }
+
+  return convertToAnswer(final)
+}

package/src/operations/rewrite.ts

@@ -0,0 +1,148 @@
+import { z } from '@bpinternal/zui'
+
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER } from './constants'
+
+type Example = z.input<typeof Example> & { instructions?: string }
+const Example = z.object({
+  input: 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')
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /** Rewrites a string according to match the prompt */
+    rewrite(original: string, prompt: string, options?: Options): Promise<string>
+  }
+}
+
+const START = '■START■'
+const END = '■END■'
+
+Zai.prototype.rewrite = async function (this: Zai, original, prompt, _options) {
+  const options = Options.parse(_options ?? {})
+  const tokenizer = await this.getTokenizer()
+
+  const taskId = this.taskId
+  const taskType = 'zai.rewrite'
+
+  const INPUT_COMPONENT_SIZE = Math.max(100, (this.Model.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
+  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)`
+    )
+  }
+
+  const instructions: string[] = []
+
+  const originalSize = tokenizer.count(original)
+  if (options.length && originalSize > options.length) {
+    instructions.push(`The original text is ${originalSize} tokens long – it should be less than ${options.length}`)
+    instructions.push(
+      `The text must be standalone and complete in less than ${options.length} tokens, so it has to be shortened to fit the length as well`
+    )
+  }
+
+  const format = (before: string, prompt: string) => {
+    return `
+Prompt: ${prompt}
+
+${START}
+${before}
+${END}
+`.trim()
+  }
+
+  const Key = fastHash(
+    stringify({
+      taskId,
+      taskType,
+      input: original,
+      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}` }
+    ]
+  }
+
+  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' }
+  ]
+
+  const tableExamples = taskId
+    ? await this.adapter.getExamples<string, string>({
+        input: original,
+        taskId,
+        taskType
+      })
+    : []
+
+  const exactMatch = tableExamples.find((x) => x.key === Key)
+  if (exactMatch) {
+    return exactMatch.output
+  }
+
+  const savedExamples: Example[] = [
+    ...tableExamples.map((x) => ({ input: x.input as string, output: x.output as string })),
+    ...options.examples
+  ]
+
+  const REMAINING_TOKENS = this.Model.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER
+  const examples = takeUntilTokens(
+    savedExamples.length ? savedExamples : defaultExamples,
+    REMAINING_TOKENS,
+    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.output))
+  )
+    .map(formatExample)
+    .flat()
+
+  const output = 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]
+  })
+
+  let result = output.choices[0]?.content as string
+
+  if (result.includes(START)) {
+    result = result.slice(result.indexOf(START) + START.length)
+  }
+
+  if (result.includes(END)) {
+    result = result.slice(0, result.indexOf(END))
+  }
+
+  if (taskId) {
+    await this.adapter.saveExample({
+      key: Key,
+      metadata: output.metadata,
+      instructions: prompt,
+      input: original,
+      output: result,
+      taskType,
+      taskId
+    })
+  }
+
+  return result
+}