@botpress/zai 1.0.1 → 1.1.0

package/src/operations/__tests/index.ts

@@ -0,0 +1,87 @@
+import { Client } from '@botpress/client'
+import { type TextTokenizer, getWasmTokenizer } from '@botpress/wasm'
+
+import fs from 'node:fs'
+import path from 'node:path'
+import { beforeAll } from 'vitest'
+
+import { Zai } from '../..'
+
+import { fastHash } from '../../utils'
+
+export const getClient = () => {
+  return new Client({
+    apiUrl: process.env.CLOUD_API_ENDPOINT ?? 'https://api.botpress.dev',
+    botId: process.env.CLOUD_BOT_ID,
+    token: process.env.CLOUD_PAT
+  })
+}
+
+function readJSONL<T>(filePath: string, keyProperty: keyof T): Map<string, T> {
+  const lines = fs.readFileSync(filePath, 'utf-8').split(/\r?\n/).filter(Boolean)
+
+  const map = new Map<string, T>()
+
+  for (const line of lines) {
+    const obj = JSON.parse(line) as T
+    const key = String(obj[keyProperty])
+    map.set(key, obj)
+  }
+
+  return map
+}
+
+const cache: Map<string, { key: string; value: any }> = readJSONL(
+  path.resolve(import.meta.dirname, './cache.jsonl'),
+  'key'
+)
+
+export const getCachedClient = () => {
+  const client = getClient()
+
+  const proxy = new Proxy(client, {
+    get(target, prop) {
+      if (prop === 'callAction') {
+        return async (...args: Parameters<Client['callAction']>) => {
+          const key = fastHash(JSON.stringify(args))
+          const cached = cache.get(key)
+
+          if (cached) {
+            return cached.value
+          }
+
+          const response = await target.callAction(...args)
+          cache.set(key, { key, value: response })
+
+          fs.appendFileSync(
+            path.resolve(import.meta.dirname, './cache.jsonl'),
+            JSON.stringify({
+              key,
+              value: response
+            }) + '\n'
+          )
+
+          return response
+        }
+      }
+      return Reflect.get(target, prop)
+    }
+  })
+
+  return proxy
+}
+
+export const getZai = () => {
+  const client = getCachedClient()
+  return new Zai({ client, retry: { maxRetries: 0 } })
+}
+
+export let tokenizer: TextTokenizer = null!
+
+beforeAll(async () => {
+  tokenizer = await getWasmTokenizer()
+})
+
+export const BotpressDocumentation = fs.readFileSync(path.join(__dirname, './botpress_docs.txt'), 'utf-8').trim()
+
+export const metadata = { cost: { input: 1, output: 1 }, latency: 0, model: '', tokens: { input: 1, output: 1 } }

package/src/operations/check.ts

@@ -0,0 +1,187 @@
+import { z } from '@bpinternal/zui'
+
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER } from './constants'
+
+const Example = z.object({
+  input: z.any(),
+  check: z.boolean(),
+  reason: z.string().optional()
+})
+
+export type Options = z.input<typeof Options>
+const Options = z.object({
+  examples: z.array(Example).describe('Examples to check the condition against').default([])
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /** Checks wether a condition is true or not */
+    check(input: unknown, condition: string, options?: Options): Promise<boolean>
+  }
+}
+
+const TRUE = '■TRUE■'
+const FALSE = '■FALSE■'
+const END = '■END■'
+
+Zai.prototype.check = async function (this: Zai, input, condition, _options) {
+  const options = Options.parse(_options ?? {})
+  const tokenizer = await this.getTokenizer()
+  const PROMPT_COMPONENT = Math.max(this.Model.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
+
+  const taskId = this.taskId
+  const taskType = 'zai.check'
+
+  const PROMPT_TOKENS = {
+    INPUT: Math.floor(0.5 * PROMPT_COMPONENT),
+    CONDITION: Math.floor(0.2 * PROMPT_COMPONENT)
+  }
+
+  // Truncate the input to fit the model's input size
+  const inputAsString = tokenizer.truncate(stringify(input), PROMPT_TOKENS.INPUT)
+  condition = tokenizer.truncate(condition, PROMPT_TOKENS.CONDITION)
+
+  // 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(condition)
+
+  const Key = fastHash(
+    JSON.stringify({
+      taskType,
+      taskId,
+      input: inputAsString,
+      condition
+    })
+  )
+
+  const examples = taskId
+    ? await this.adapter.getExamples<string, boolean>({
+        input: inputAsString,
+        taskType,
+        taskId
+      })
+    : []
+
+  const exactMatch = examples.find((x) => x.key === Key)
+  if (exactMatch) {
+    return exactMatch.output
+  }
+
+  const defaultExamples = [
+    { input: '50 Cent', check: true, reason: '50 Cent is widely recognized as a public personality.' },
+    {
+      input: ['apple', 'banana', 'carrot', 'house'],
+      check: false,
+      reason: 'The list contains a house, which is not a fruit. Also, the list contains a carrot, which is a vegetable.'
+    }
+  ]
+
+  const userExamples = [
+    ...examples.map((e) => ({ input: e.input, check: e.output, reason: e.explanation })),
+    ...options.examples
+  ]
+
+  let exampleId = 1
+
+  const formatInput = (input: string, condition: string) => {
+    const header = userExamples.length ? `Expert Example #${exampleId++}` : `Example of condition: "${condition}"`
+
+    return `
+${header}
+<|start_input|>
+${input.trim()}
+<|end_input|>
+`.trim()
+  }
+
+  const formatOutput = (answer: boolean, justification: string) => {
+    return `
+Analysis: ${justification}
+Final Answer: ${answer ? TRUE : FALSE}
+${END}
+`.trim()
+  }
+
+  const formatExample = (example: { input?: any; check: boolean; reason?: string }) => [
+    { type: 'text' as const, content: formatInput(stringify(example.input ?? null), condition), role: 'user' as const },
+    {
+      type: 'text' as const,
+      content: formatOutput(example.check, example.reason ?? ''),
+      role: 'assistant' as const
+    }
+  ]
+
+  const allExamples = takeUntilTokens(
+    userExamples.length ? userExamples : defaultExamples,
+    EXAMPLES_TOKENS,
+    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(el.reason ?? '')
+  )
+    .map(formatExample)
+    .flat()
+
+  const specialInstructions = userExamples.length
+    ? `
+- You have been provided with examples from previous experts. Make sure to read them carefully before making your decision.
+- Make sure to refer to the examples provided by the experts to justify your decision (when applicable).
+- When in doubt, ground your decision on the examples provided by the experts instead of your own intuition.
+- When no example is similar to the input, make sure to provide a clear justification for your decision while inferring the decision-making process from the examples provided by the experts.
+`.trim()
+    : ''
+
+  const output = await this.callModel({
+    systemPrompt: `
+Check if the following condition is true or false for the given input. Before answering, make sure to read the input and the condition carefully.
+Justify your answer, then answer with either ${TRUE} or ${FALSE} at the very end, then add ${END} to finish the response.
+IMPORTANT: Make sure to answer with either ${TRUE} or ${FALSE} at the end of your response, but NOT both.
+---
+Expert Examples (#1 to #${exampleId - 1}):
+${specialInstructions}
+`.trim(),
+    stopSequences: [END],
+    messages: [
+      ...allExamples,
+      {
+        type: 'text',
+        content: `
+Considering the below input and above examples, is the following condition true or false?
+${formatInput(inputAsString, condition)}
+In your "Analysis", please refer to the Expert Examples # to justify your decision.`.trim(),
+        role: 'user'
+      }
+    ]
+  })
+
+  const answer = output.choices[0]?.content as string
+
+  const hasTrue = answer.includes(TRUE)
+  const hasFalse = answer.includes(FALSE)
+
+  if (!hasTrue && !hasFalse) {
+    throw new Error(`The model did not return a valid answer. The response was: ${answer}`)
+  }
+
+  let finalAnswer: boolean
+
+  if (hasTrue && hasFalse) {
+    // If both TRUE and FALSE are present, we need to check which one was answered last
+    finalAnswer = answer.lastIndexOf(TRUE) > answer.lastIndexOf(FALSE)
+  } else {
+    finalAnswer = hasTrue
+  }
+
+  if (taskId) {
+    await this.adapter.saveExample({
+      key: Key,
+      taskType,
+      taskId,
+      input: inputAsString,
+      instructions: condition,
+      metadata: output.metadata,
+      output: finalAnswer,
+      explanation: answer.replace(TRUE, '').replace(FALSE, '').replace(END, '').replace('Final Answer:', '').trim()
+    })
+  }
+
+  return finalAnswer
+}

package/src/operations/constants.ts

@@ -0,0 +1,2 @@
+export const PROMPT_INPUT_BUFFER = 1048
+export const PROMPT_OUTPUT_BUFFER = 512

package/src/operations/errors.ts

@@ -0,0 +1,9 @@
+export class JsonParsingError extends Error {
+  constructor(
+    public json: unknown,
+    public error: Error
+  ) {
+    const message = `Error parsing JSON:\n\n---JSON---\n${json}\n\n---Error---\n\n ${error}`
+    super(message)
+  }
+}

package/src/operations/extract.ts

@@ -0,0 +1,291 @@
+import { z } from '@bpinternal/zui'
+
+import JSON5 from 'json5'
+import { jsonrepair } from 'jsonrepair'
+
+import { chunk, isArray } from 'lodash-es'
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER } from './constants'
+import { JsonParsingError } from './errors'
+
+export type Options = z.input<typeof Options>
+const Options = z.object({
+  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)
+})
+
+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>>>
+  }
+}
+
+const START = '■json_start■'
+const END = '■json_end■'
+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()
+
+  const taskId = this.taskId
+  const taskType = 'zai.extract'
+
+  const PROMPT_COMPONENT = Math.max(this.Model.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
+
+  let isArrayOfObjects = false
+  const originalSchema = schema
+
+  if (schema instanceof z.ZodObject) {
+    // Do nothing
+  } else if (schema instanceof z.ZodArray) {
+    if (schema._def.type instanceof z.ZodObject) {
+      isArrayOfObjects = true
+      schema = schema._def.type
+    } else {
+      throw new Error('Schema must be a ZodObject or a ZodArray<ZodObject>')
+    }
+  } else {
+    throw new Error('Schema must be either a ZuiObject or a ZuiArray<ZuiObject>')
+  }
+
+  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)
+
+  const keys = Object.keys(schema.shape)
+
+  let 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 as z.AnyZodObject)))
+
+      return all.flat()
+    } else {
+      // Truncate the input to fit the model's input size
+      inputAsString = tokenizer.truncate(stringify(input), options.chunkLength)
+    }
+  }
+
+  const instructions: string[] = []
+
+  if (options.instructions) {
+    instructions.push(options.instructions)
+  }
+
+  const shape = `{ ${keys.map((key) => `"${key}": ...`).join(', ')} }`
+  const abbv = '{ ... }'
+
+  if (isArrayOfObjects) {
+    instructions.push('You may have multiple elements, or zero elements in the input.')
+    instructions.push('You must extract each element separately.')
+    instructions.push(`Each element must be a JSON object with exactly the format: ${START}${shape}${END}`)
+    instructions.push(`When you are done extracting all elements, type "${NO_MORE}" to finish.`)
+    instructions.push(`For example, if you have zero elements, the output should look like this: ${NO_MORE}`)
+    instructions.push(
+      `For example, if you have two elements, the output should look like this: ${START}${abbv}${END}${START}${abbv}${END}${NO_MORE}`
+    )
+  } else {
+    instructions.push('You may have exactly one element in the input.')
+    instructions.push(`The element must be a JSON object with exactly the format: ${START}${shape}${END}`)
+  }
+
+  // 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'))
+
+  const Key = fastHash(
+    JSON.stringify({
+      taskType,
+      taskId,
+      input: inputAsString,
+      instructions: options.instructions
+    })
+  )
+
+  const examples = taskId
+    ? await this.adapter.getExamples<string, unknown>({
+        input: inputAsString,
+        taskType,
+        taskId
+      })
+    : []
+
+  const exactMatch = examples.find((x) => x.key === Key)
+  if (exactMatch) {
+    return exactMatch.output
+  }
+
+  const defaultExample = isArrayOfObjects
+    ? {
+        input: `The story goes as follow.
+Once upon a time, there was a person named Alice who was 30 years old.
+Then, there was a person named Bob who was 25 years old.
+The end.`,
+        schema: 'Array<{ name: string, age: number }>',
+        instructions: 'Extract all people',
+        extracted: [
+          {
+            name: 'Alice',
+            age: 30
+          },
+          {
+            name: 'Bob',
+            age: 25
+          }
+        ]
+      }
+    : {
+        input: `The story goes as follow.
+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 }
+      }
+
+  const userExamples = examples.map((e) => ({
+    input: e.input,
+    extracted: e.output,
+    schema: schemaTypescript,
+    instructions: options.instructions
+  }))
+
+  let exampleId = 1
+
+  const formatInput = (input: string, schema: string, instructions?: string) => {
+    const header = userExamples.length
+      ? `Expert Example #${exampleId++}`
+      : "Here's an example to help you understand the format:"
+
+    return `
+${header}
+
+<|start_schema|>
+${schema}
+<|end_schema|>
+
+<|start_instructions|>
+${instructions ?? 'No specific instructions, just follow the schema above.'}
+<|end_instructions|>
+
+<|start_input|>
+${input.trim()}
+<|end_input|>
+  `.trim()
+  }
+
+  const formatOutput = (extracted: any) => {
+    extracted = isArray(extracted) ? extracted : [extracted]
+
+    return (
+      extracted
+        .map((x: string) =>
+          `
+${START}
+${JSON.stringify(x, null, 2)}
+${END}`.trim()
+        )
+        .join('\n') + NO_MORE
+    )
+  }
+
+  const formatExample = (example: { input?: any; schema: string; instructions?: string; extracted: any }) => [
+    {
+      type: 'text' as const,
+      content: formatInput(stringify(example.input ?? null), example.schema, example.instructions),
+      role: 'user' as const
+    },
+    {
+      type: 'text' as const,
+      content: formatOutput(example.extracted),
+      role: 'assistant' as const
+    }
+  ]
+
+  const allExamples = takeUntilTokens(
+    userExamples.length ? userExamples : [defaultExample],
+    EXAMPLES_TOKENS,
+    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.extracted))
+  )
+    .map(formatExample)
+    .flat()
+
+  const output = await this.callModel({
+    systemPrompt: `
+Extract the following information from the input:
+${schemaTypescript}
+====
+
+${instructions.map((x) => `• ${x}`).join('\n')}
+`.trim(),
+    stopSequences: [isArrayOfObjects ? NO_MORE : END],
+    messages: [
+      ...allExamples,
+      {
+        role: 'user',
+        type: 'text',
+        content: formatInput(inputAsString, schemaTypescript, options.instructions ?? '')
+      }
+    ]
+  })
+
+  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({})
+  } else {
+    final = elements[0]
+  }
+
+  if (taskId) {
+    await this.adapter.saveExample({
+      key: Key,
+      taskId: `zai/${taskId}`,
+      taskType,
+      instructions: options.instructions ?? 'No specific instructions',
+      input: inputAsString,
+      output: final,
+      metadata: output.metadata
+    })
+  }
+
+  return final
+}