@botpress/zai 1.0.0-beta.2 → 1.0.0-beta.3

package/src/operations/extract.ts

@@ -1,291 +0,0 @@
-import { ZodArray, ZodObject, z } from '@bpinternal/zui'
-
-import JSON5 from 'json5'
-import { jsonrepair } from 'jsonrepair'
-
-import _ from 'lodash'
-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 '../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 ZodObject) {
-    // Do nothing
-  } else if (schema instanceof ZodArray) {
-    if (schema._def.type instanceof 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)))
-
-      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) =>
-          `
-${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)
-      }
-    })
-    .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
-}

package/src/operations/filter.test.ts

@@ -1,182 +0,0 @@
-import { describe, it, expect, beforeEach, afterEach, afterAll, vi } from 'vitest'
-
-import { getClient, getZai, metadata } from './__tests'
-import { TableAdapter } from '../adapters/botpress-table'
-import { Client } from '@botpress/client'
-
-describe('zai.filter', { timeout: 60_000 }, () => {
-  let zai = getZai()
-
-  beforeEach(async () => {
-    zai = getZai()
-  })
-
-  it('basic filter with small items', async () => {
-    const value = await zai.filter(
-      [
-        { name: 'John', description: 'is a bad person' },
-        { name: 'Alice', description: 'is a good person' },
-        { name: 'Bob', description: 'is a good person' },
-        { name: 'Eve', description: 'is a bad person' },
-        { name: 'Alex', description: 'is a good person' },
-        { name: 'Sara', description: 'donates to charity every month' },
-        { name: 'Tom', description: 'commits crimes and is in jail' }
-      ],
-      'generally good people'
-    )
-
-    const names = value.map((v) => v.name)
-    expect(names).toMatchInlineSnapshot(`
-      [
-        "Alice",
-        "Bob",
-        "Alex",
-        "Sara",
-      ]
-    `)
-  })
-
-  it('filtering huge array chunks it up', async () => {
-    const callAction = vi.fn()
-    const client = { ...getClient(), callAction } as unknown as Client
-
-    zai = getZai().with({
-      client
-    })
-
-    const hugeArray = Array.from({ length: 100 }, (_, i) => ({
-      name: `Person #${i}#`,
-      description: 'blah blah '.repeat(50_000)
-    }))
-
-    try {
-      await zai.filter(hugeArray, 'generally good people', { tokensPerItem: 100_000 })
-    } catch (err) {}
-
-    expect(callAction.mock.calls.length).toBeGreaterThan(20)
-    expect(JSON.stringify(callAction.mock.calls.at(0))).toContain('Person #0#')
-    expect(JSON.stringify(callAction.mock.calls.at(0))).not.toContain('Person #99#')
-
-    expect(JSON.stringify(callAction.mock.calls.at(-1))).not.toContain('Person #0#')
-    expect(JSON.stringify(callAction.mock.calls.at(-1))).toContain('Person #99#')
-
-    callAction.mockReset()
-
-    try {
-      await zai.filter(hugeArray, 'generally good people', { tokensPerItem: 100 })
-    } catch (err) {}
-
-    expect(callAction.mock.calls.length).toBe(2)
-  })
-
-  it('filter with examples', async () => {
-    const examples = [
-      {
-        input: 'Rasa (framework)',
-        filter: true,
-        reason: 'Rasa is a chatbot framework, so it competes with us (Botpress).'
-      },
-      {
-        input: 'Rasa (coffee company)',
-        filter: false,
-        reason:
-          'Rasa (coffee company) is not in the chatbot or AI agent industry, therefore it does not compete with us (Botpress).'
-      },
-      {
-        input: 'Dialogflow',
-        filter: true,
-        reason: 'Dialogflow is a chatbot development product, so it competes with us (Botpress).'
-      }
-    ]
-
-    const value = await zai.filter(
-      [{ name: 'Moveworks' }, { name: 'Ada.cx' }, { name: 'Nike' }, { name: 'Voiceflow' }, { name: 'Adidas' }],
-      'competes with us',
-      { examples }
-    )
-
-    const names = value.map((v) => v.name)
-    expect(names).toMatchInlineSnapshot(`
-      [
-        "Moveworks",
-        "Ada.cx",
-        "Voiceflow",
-      ]
-    `)
-  })
-})
-
-describe('zai.learn.filter', { timeout: 60_000 }, () => {
-  const client = getClient()
-  let tableName = 'ZaiTestFilterInternalTable'
-  let taskId = 'filter'
-  let zai = getZai()
-
-  beforeEach(async () => {
-    zai = getZai().with({
-      activeLearning: {
-        enable: true,
-        taskId,
-        tableName
-      }
-    })
-  })
-
-  afterEach(async () => {
-    try {
-      await client.deleteTableRows({ table: tableName, deleteAllRows: true })
-    } catch (err) {}
-  })
-
-  afterAll(async () => {
-    try {
-      await client.deleteTable({ table: tableName })
-    } catch (err) {}
-  })
-
-  it('learns a filtering rule from examples', async () => {
-    const adapter = new TableAdapter({
-      client,
-      tableName
-    })
-
-    await adapter.saveExample({
-      key: 't1',
-      taskId: `zai/${taskId}`,
-      taskType: 'zai.filter',
-      instructions: 'competes with us?',
-      input: ['Rasa (framework)', 'Rasa (coffee company)'],
-      output: ['Rasa (framework)'],
-      explanation: `Rasa is a chatbot framework, so it competes with us (Botpress). We should keep it. Rasa (coffee company) is not in the chatbot or AI agent industry, therefore it does not compete with us (Botpress). We should filter it out.`,
-      metadata,
-      status: 'approved'
-    })
-
-    await adapter.saveExample({
-      key: 't2',
-      taskId: `zai/${taskId}`,
-      taskType: 'zai.filter',
-      instructions: 'competes with us?',
-      input: ['Voiceflow', 'Dialogflow'],
-      output: ['Voiceflow', 'Dialogflow'],
-      explanation: `Voiceflow is a chatbot development product, so it competes with us (Botpress). We should keep it. Dialogflow is a chatbot development product, so it competes with us (Botpress). We should keep it.`,
-      metadata,
-      status: 'approved'
-    })
-
-    const second = await zai
-      .learn(taskId)
-      .filter(['Nike', 'Ada.cx', 'Adidas', 'Moveworks', 'Lululemon'], 'competes with us? (botpress)')
-
-    expect(second).toMatchInlineSnapshot(`
-      [
-        "Ada.cx",
-        "Moveworks",
-      ]
-    `)
-
-    const rows = await client.findTableRows({ table: tableName })
-    expect(rows.rows.length).toBe(3)
-    expect(rows.rows.at(-1)!.output.value).toEqual(second)
-  })
-})

package/src/operations/filter.ts

@@ -1,231 +0,0 @@
-import { z } from '@bpinternal/zui'
-
-import _ from 'lodash'
-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 '../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()
-}