@botpress/zai 1.0.1 → 1.2.0

package/src/operations/extract.ts

@@ -0,0 +1,309 @@
+// eslint-disable consistent-type-definitions
+import { z, ZodObject } 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 | z.ZodArray>(input: unknown, schema: S, options?: Options): Promise<z.TypeOf<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()
+  await this.fetchModelDetails()
+
+  const taskId = this.taskId
+  const taskType = 'zai.extract'
+
+  const PROMPT_COMPONENT = Math.max(this.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
+
+  let isArrayOfObjects = false
+  const originalSchema = schema
+
+  const baseType = (schema.naked ? schema.naked() : schema)?.constructor?.name ?? 'unknown'
+
+  if (baseType === 'ZodObject') {
+    // Do nothing
+  } else if (baseType === 'ZodArray') {
+    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>')
+    }
+  } 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.ModelDetails.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength
+  )
+
+  const keys = Object.keys((schema as ZodObject).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, meta } = 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: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output,
+        },
+        latency: meta.latency,
+        model: this.Model,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output,
+        },
+      },
+    })
+  }
+
+  return final
+}

package/src/operations/filter.ts

@@ -0,0 +1,244 @@
+// eslint-disable consistent-type-definitions
+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()
+  await this.fetchModelDetails()
+
+  const taskId = this.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_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, 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:
+■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: {
+          cost: {
+            input: meta.cost.input,
+            output: meta.cost.output,
+          },
+          latency: meta.latency,
+          model: this.Model,
+          tokens: {
+            input: meta.tokens.input,
+            output: meta.tokens.output,
+          },
+        },
+      })
+    }
+
+    return partial
+  }
+
+  const filteredChunks = await Promise.all(chunks.map(filterChunk))
+
+  return filteredChunks.flat()
+}