chore: import upstream snapshot with attribution

packages/zai/e2e/check.test.ts

@@ -0,0 +1,216 @@
+import { describe, it, expect, afterAll, beforeEach, afterEach, vi } from 'vitest'
+import { BotpressDocumentation, getClient, getZai, metadata } from './utils'
+import { TableAdapter } from '../src/adapters/botpress-table'
+import { getCachedCognitiveClient, getCognitiveClient } from './client'
+
+describe('zai.check', { timeout: 60_000 }, () => {
+  const zai = getZai()
+
+  it('basic check on a string', async () => {
+    const value = await zai.check('This text is very clearly written in English.', 'is an english sentence')
+    expect(value).toBe(true)
+  })
+
+  it('basic check on a string (full)', async () => {
+    const { output, usage } = await zai
+      .check('This text is very clearly written in English.', 'is an english sentence')
+      .result()
+
+    expect(output.value).toBe(true)
+    expect(output.explanation).toBeTypeOf('string')
+    expect(output.explanation.length).toBeGreaterThan(5)
+    expect(usage.requests.requests).toBeGreaterThanOrEqual(1)
+    expect(usage.requests.responses).toBeGreaterThanOrEqual(1)
+  })
+
+  it('can abort', async () => {
+    // no caching
+    const request = getZai(getCognitiveClient()).check(
+      'This text is very clearly written in English.',
+      'is an english sentence'
+    )
+
+    request.abort('CANCEL')
+    await expect(request).rejects.toThrow('CANCEL')
+  })
+
+  it('can abort via external signal', async () => {
+    // no caching
+    const controller = new AbortController()
+    const request = getZai(getCognitiveClient())
+      .check('This text is very clearly written in English.', 'is an english sentence')
+      .bindSignal(controller.signal)
+
+    controller.abort('CANCEL2')
+    await expect(request).rejects.toThrow('CANCEL2')
+  })
+
+  it('text that is too long gets truncated', async () => {
+    const isBotpressDocumentation = await zai.check(
+      BotpressDocumentation,
+      'is about botpress and looks like documentation'
+    )
+
+    const isAboutBirds = await zai.check(BotpressDocumentation, 'is a book about birds and their species')
+
+    expect(isBotpressDocumentation).toBe(true)
+    expect(isAboutBirds).toBe(false)
+  })
+
+  it('works with any input type', async () => {
+    const sly = { name: 'Sylvain Perron', age: 30, job: 'CEO', company: 'Botpress', location: 'Quebec' }
+    const american = await zai.check(sly, 'person lives in north america')
+    const european = await zai.check(sly, 'person lives in europe')
+
+    expect(american).toBe(true)
+    expect(european).toBe(false)
+  })
+
+  it('retries on generation failure', async () => {
+    const cognitive = getCachedCognitiveClient()
+    const mocked = getCachedCognitiveClient()
+
+    let callCount = 0
+
+    const mock = vi.fn().mockImplementation(async (input) => {
+      const output = await cognitive.generateContent(input)
+
+      if (callCount++ < 1) {
+        output.output.choices[0].content = '...'
+      }
+
+      return output
+    })
+
+    mocked.clone = vi.fn().mockReturnValue(mocked)
+    mocked.generateContent = mock
+
+    const isEnglish = await getZai(mocked).check(
+      'This text is very clearly written in English',
+      'is an english sentence'
+    )
+
+    expect(isEnglish).toBe(true)
+    expect(mock).toHaveBeenCalledTimes(2)
+  })
+
+  it('check with examples', async () => {
+    const examples = [
+      {
+        input: 'Rasa (framework)',
+        check: true,
+        reason: 'Rasa is a chatbot framework, so it competes with us (Botpress).',
+      },
+      {
+        input: 'Rasa (coffee company)',
+        check: false,
+        reason:
+          'Rasa (coffee company) is not in the chatbot or AI agent industry, therefore it does not compete with us (Botpress).',
+      },
+      {
+        input: 'Dialogflow',
+        check: true,
+        reason: 'Dialogflow is a chatbot development product, so it competes with us (Botpress).',
+      },
+    ]
+
+    const moveworks = await zai.check('Moveworks (AI chatbot)', 'competes with us', { examples })
+    const ada = await zai.check('Ada.cx', 'competes with us', { examples })
+    const voiceflow = await zai.check('Voiceflow', 'competes with us', { examples })
+    const nike = await zai.check('Nike', 'competes with us', { examples })
+    const adidas = await zai.check('Adidas', 'competes with us', { examples })
+
+    expect(moveworks).toBe(true)
+    expect(ada).toBe(true)
+    expect(voiceflow).toBe(true)
+
+    expect(nike).toBe(false)
+    expect(adidas).toBe(false)
+  })
+})
+
+describe('zai.learn.check', { timeout: 60_000, sequential: true }, () => {
+  const client = getClient()
+  let tableName = 'ZaiTestCheckInternalTable'
+  let taskId = 'check'
+  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 contradiction from examples', async () => {
+    const adapter = new TableAdapter({
+      client,
+      tableName,
+    })
+
+    const value = await zai.learn(taskId).check(`What's up`, 'is a greeting')
+    expect(value).toBe(true)
+
+    let rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBe(1)
+    expect(rows.rows[0].output.value).toEqual(value)
+
+    await adapter.saveExample({
+      key: 't1',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.check',
+      instructions: 'is a greeting',
+      input: 'what is up',
+      output: false,
+      explanation: `"What's up" in our business scenario is NOT considered an official greeting.`,
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 't2',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.check',
+      instructions: 'is a greeting',
+      input: 'hello! how are you?',
+      output: true,
+      explanation: `"hello!" is a common greeting in English.`,
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 't3',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.check',
+      instructions: 'is a greeting',
+      input: 'wassup',
+      output: false,
+      explanation: `"wassup" is a slang term and not considered a formal greeting. It is therefore NOT considered a greeting.`,
+      metadata,
+      status: 'approved',
+    })
+
+    const second = await zai.learn(taskId).check(`What's up`, 'is a greeting')
+    expect(second).toBe(false)
+    rows = await client.findTableRows({ table: tableName })
+
+    expect(rows.rows.length).toBe(4)
+    expect(rows.rows[0].output.value).toEqual(second)
+  })
+})

packages/zai/e2e/client.ts

@@ -0,0 +1,148 @@
+import { Client } from '@botpress/client'
+import { Cognitive } from '@botpress/cognitive'
+import { diffLines } from 'diff'
+import fs from 'node:fs'
+import path from 'node:path'
+import { expect } from 'vitest'
+
+function stringifyWithSortedKeys(obj: any, space?: number): string {
+  function sortKeys(input: any): any {
+    if (Array.isArray(input)) {
+      return input.map(sortKeys)
+    } else if (input && typeof input === 'object' && input.constructor === Object) {
+      return Object.keys(input)
+        .sort()
+        .reduce(
+          (acc, key) => {
+            acc[key] = sortKeys(input[key])
+            return acc
+          },
+          {} as Record<string, any>
+        )
+    } else {
+      return input
+    }
+  }
+
+  return JSON.stringify(sortKeys(obj), null, space)
+}
+
+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) {
+    try {
+      const obj = JSON.parse(line) as T
+      const value = obj[keyProperty]
+      if (value === undefined || value === null) {
+        continue
+      }
+      const key = String(value)
+      map.set(key, obj)
+    } catch {}
+  }
+
+  return map
+}
+
+type CacheEntry = { key: string; value: any; test?: string; input: string }
+
+const cache: Map<string, CacheEntry> = readJSONL(path.resolve(__dirname, './data/cache.jsonl'), 'key')
+const cacheByTest: Map<string, CacheEntry> = readJSONL(path.resolve(__dirname, './data/cache.jsonl'), 'test')
+
+class CachedClient extends Client {
+  #client: Client
+  #callsByTest: Record<string, number> = {}
+
+  public constructor(options: ConstructorParameters<typeof Client>[0]) {
+    super(options)
+    this.#client = new Client(options)
+  }
+
+  public callAction = async (...args: Parameters<Client['callAction']>) => {
+    const currentTestName = expect.getState().currentTestName ?? 'default'
+    this.#callsByTest[currentTestName] ||= 0
+    this.#callsByTest[currentTestName]++
+
+    const testKey = `${currentTestName}-${this.#callsByTest[currentTestName]}`
+
+    const key = fastHash(stringifyWithSortedKeys(args))
+    const cached = cache.get(key)
+
+    if (cached) {
+      return cached.value
+    }
+
+    if (process.env.CI) {
+      const previous = cacheByTest.get(testKey)
+      if (previous) {
+        console.info(`Cache miss for ${key} in test ${testKey}`)
+        console.info(
+          diffLines(
+            JSON.stringify(JSON.parse(previous.input), null, 2),
+            JSON.stringify(JSON.parse(stringifyWithSortedKeys(args)), null, 2)
+          )
+        )
+      }
+
+      throw new Error(
+        `Missing cached Botpress action response for ${key} in ${testKey}. Run pnpm test:e2e:update to refresh packages/zai/e2e/data/cache.jsonl.`
+      )
+    }
+
+    const response = await this.#client.callAction(...args)
+    cache.set(key, { key, value: response, test: testKey, input: stringifyWithSortedKeys(args) })
+
+    fs.appendFileSync(
+      path.resolve(__dirname, './data/cache.jsonl'),
+      JSON.stringify({
+        test: testKey,
+        key,
+        input: stringifyWithSortedKeys(args),
+        value: response,
+      }) + '\n'
+    )
+
+    return response
+  }
+
+  public clone() {
+    return this
+  }
+}
+
+export const getCognitiveClient = () => {
+  const cognitive = new Cognitive({
+    client: new Client({
+      apiUrl: process.env.CLOUD_API_ENDPOINT ?? 'https://api.botpress.dev',
+      botId: process.env.CLOUD_BOT_ID,
+      token: process.env.CLOUD_PAT,
+    }),
+    __experimental_beta: true,
+  })
+  return cognitive
+}
+
+export const getCachedCognitiveClient = () => {
+  const cognitive = new Cognitive({
+    client: new CachedClient({
+      // Using mock server (MSW) - actual URL doesn't matter as requests are intercepted
+      apiUrl: process.env.CLOUD_API_ENDPOINT ?? 'https://api.botpress.dev',
+      botId: process.env.CLOUD_BOT_ID,
+      token: process.env.CLOUD_PAT,
+    }),
+    __experimental_beta: true,
+  })
+  return cognitive
+}
+
+function fastHash(str: string): string {
+  let hash = 0
+  for (let i = 0; i < str.length; i++) {
+    hash = (hash << 5) - hash + str.charCodeAt(i)
+    hash |= 0 // Convert to 32bit integer
+  }
+  return (hash >>> 0).toString(16) // Convert to unsigned and then to hex
+}

packages/zai/e2e/errors.test.ts

@@ -0,0 +1,287 @@
+import { describe, it, expect } from 'vitest'
+import { z } from '@bpinternal/zui'
+import { JsonParsingError } from '../src/operations/errors'
+
+describe('JsonParsingError', () => {
+  it('formats simple zod validation error in LLM-friendly way', () => {
+    const schema = z.object({
+      name: z.string(),
+      age: z.number(),
+    })
+
+    try {
+      schema.parse({ name: 'John', age: 'not a number' })
+    } catch (error) {
+      const jsonParsingError = new JsonParsingError('{"name":"John","age":"not a number"}', error as Error)
+
+      expect(jsonParsingError.message).toMatchInlineSnapshot(`
+        "Error parsing JSON:
+
+        ---JSON---
+        {"name":"John","age":"not a number"}
+
+        ---Validation Errors---
+
+        1. Field: "age"
+           Problem: Expected number, but received string
+           Message: Expected number, received string
+        "
+      `)
+    }
+  })
+
+  it('formats complex nested zod validation error in LLM-friendly way', () => {
+    const schema = z.object({
+      name: z.string().min(3),
+      email: z.string().email(),
+      age: z.number().positive(),
+      address: z.object({
+        street: z.string(),
+        city: z.string(),
+        zipCode: z.string().length(5),
+      }),
+      tags: z.array(z.string()),
+    })
+
+    try {
+      schema.parse({
+        name: 'Jo',
+        email: 'invalid-email',
+        age: -5,
+        address: {
+          street: 'Main St',
+          city: '',
+          zipCode: '1234',
+        },
+        tags: ['valid', 123, null],
+      })
+    } catch (error) {
+      const jsonParsingError = new JsonParsingError(
+        JSON.stringify({
+          name: 'Jo',
+          email: 'invalid-email',
+          age: -5,
+          address: { street: 'Main St', city: '', zipCode: '1234' },
+          tags: ['valid', 123, null],
+        }),
+        error as Error
+      )
+
+      expect(jsonParsingError.message).toMatchInlineSnapshot(`
+        "Error parsing JSON:
+
+        ---JSON---
+        {"name":"Jo","email":"invalid-email","age":-5,"address":{"street":"Main St","city":"","zipCode":"1234"},"tags":["valid",123,null]}
+
+        ---Validation Errors---
+
+        1. Field: "name"
+           Problem: String must be at least 3 characters
+           Message: String must contain at least 3 character(s)
+
+        2. Field: "email"
+           Problem: Invalid email format
+           Message: Invalid email
+
+        3. Field: "age"
+           Problem: Number must be greater than 0
+           Message: Number must be greater than 0
+
+        4. Field: "address.zipCode"
+           Problem: String must be exactly 5 characters
+           Message: String must contain exactly 5 character(s)
+
+        5. Field: "tags.1"
+           Problem: Expected string, but received number
+           Message: Expected string, received number
+
+        6. Field: "tags.2"
+           Problem: Expected string, but received null
+           Message: Expected string, received null
+        "
+      `)
+    }
+  })
+
+  it('formats array validation error in LLM-friendly way', () => {
+    const schema = z.array(
+      z.object({
+        id: z.number(),
+        name: z.string(),
+      })
+    )
+
+    try {
+      schema.parse([
+        { id: 1, name: 'Valid' },
+        { id: 'invalid', name: 'Invalid ID' },
+        { id: 3, name: 123 },
+      ])
+    } catch (error) {
+      const jsonParsingError = new JsonParsingError(
+        JSON.stringify([
+          { id: 1, name: 'Valid' },
+          { id: 'invalid', name: 'Invalid ID' },
+          { id: 3, name: 123 },
+        ]),
+        error as Error
+      )
+
+      expect(jsonParsingError.message).toMatchInlineSnapshot(`
+        "Error parsing JSON:
+
+        ---JSON---
+        [{"id":1,"name":"Valid"},{"id":"invalid","name":"Invalid ID"},{"id":3,"name":123}]
+
+        ---Validation Errors---
+
+        1. Field: "1.id"
+           Problem: Expected number, but received string
+           Message: Expected number, received string
+
+        2. Field: "2.name"
+           Problem: Expected string, but received number
+           Message: Expected string, received number
+        "
+      `)
+    }
+  })
+
+  it('formats missing required field error in LLM-friendly way', () => {
+    const schema = z.object({
+      requiredField: z.string(),
+      optionalField: z.string().optional(),
+    })
+
+    try {
+      schema.parse({ optionalField: 'present' })
+    } catch (error) {
+      const jsonParsingError = new JsonParsingError('{"optionalField":"present"}', error as Error)
+
+      expect(jsonParsingError.message).toMatchInlineSnapshot(`
+        "Error parsing JSON:
+
+        ---JSON---
+        {"optionalField":"present"}
+
+        ---Validation Errors---
+
+        1. Field: "requiredField"
+           Problem: Expected string, but received undefined
+           Message: Required
+        "
+      `)
+    }
+  })
+
+  it('formats union/enum validation error in LLM-friendly way', () => {
+    const schema = z.object({
+      status: z.enum(['pending', 'approved', 'rejected']),
+      priority: z.union([z.literal('low'), z.literal('medium'), z.literal('high')]),
+    })
+
+    try {
+      schema.parse({ status: 'invalid-status', priority: 'urgent' })
+    } catch (error) {
+      const jsonParsingError = new JsonParsingError('{"status":"invalid-status","priority":"urgent"}', error as Error)
+
+      expect(jsonParsingError.message).toMatchInlineSnapshot(`
+        "Error parsing JSON:
+
+        ---JSON---
+        {"status":"invalid-status","priority":"urgent"}
+
+        ---Validation Errors---
+
+        1. Field: "status"
+           Problem: Invalid value "invalid-status"
+           Allowed values: "pending", "approved", "rejected"
+           Message: Invalid enum value. Expected 'pending' | 'approved' | 'rejected', received 'invalid-status'
+
+        2. Field: "priority"
+           Problem: Value doesn't match any of the expected formats
+           Message: Invalid input
+        "
+      `)
+    }
+  })
+
+  it('formats type mismatch errors in LLM-friendly way', () => {
+    const schema = z.object({
+      string: z.string(),
+      number: z.number(),
+      boolean: z.boolean(),
+      array: z.array(z.string()),
+      object: z.object({ nested: z.string() }),
+    })
+
+    try {
+      schema.parse({
+        string: 123,
+        number: 'not a number',
+        boolean: 'yes',
+        array: 'not an array',
+        object: 'not an object',
+      })
+    } catch (error) {
+      const jsonParsingError = new JsonParsingError(
+        JSON.stringify({
+          string: 123,
+          number: 'not a number',
+          boolean: 'yes',
+          array: 'not an array',
+          object: 'not an object',
+        }),
+        error as Error
+      )
+
+      expect(jsonParsingError.message).toMatchInlineSnapshot(`
+        "Error parsing JSON:
+
+        ---JSON---
+        {"string":123,"number":"not a number","boolean":"yes","array":"not an array","object":"not an object"}
+
+        ---Validation Errors---
+
+        1. Field: "string"
+           Problem: Expected string, but received number
+           Message: Expected string, received number
+
+        2. Field: "number"
+           Problem: Expected number, but received string
+           Message: Expected number, received string
+
+        3. Field: "boolean"
+           Problem: Expected boolean, but received string
+           Message: Expected boolean, received string
+
+        4. Field: "array"
+           Problem: Expected array, but received string
+           Message: Expected array, received string
+
+        5. Field: "object"
+           Problem: Expected object, but received string
+           Message: Expected object, received string
+        "
+      `)
+    }
+  })
+
+  it('handles non-zod errors gracefully', () => {
+    const regularError = new Error('This is a regular parsing error')
+    const jsonParsingError = new JsonParsingError('{"invalid json"}', regularError)
+
+    expect(jsonParsingError.message).toMatchInlineSnapshot(`
+      "Error parsing JSON:
+
+      ---JSON---
+      {"invalid json"}
+
+      ---Error---
+
+      The JSON provided is not valid JSON.
+      Details: This is a regular parsing error
+      "
+    `)
+  })
+})

packages/zai/e2e/extract.test.ts

@@ -0,0 +1,368 @@
+import { describe, it, expect, beforeEach, afterEach, afterAll } from 'vitest'
+
+import { BotpressDocumentation, getCachedClient, getClient, getZai, metadata } from './utils'
+
+import { z } from '@bpinternal/zui'
+import { check } from '@botpress/vai'
+
+import { TableAdapter } from '../src/adapters/botpress-table'
+
+describe('zai.extract', () => {
+  let cognitive = getCachedClient()
+  let zai = getZai(cognitive)
+
+  beforeEach(() => {
+    cognitive = getCachedClient()
+    zai = getZai(cognitive)
+  })
+
+  it('extract simple object from paragraph', async () => {
+    const person = await zai.extract(
+      'My name is John Doe, I am 30 years old and I live in Quebec',
+      z.object({
+        name: z.string().describe('The full name of the person'),
+        age: z.number(),
+        location: z.string(),
+      })
+    )
+
+    expect(person).toMatchInlineSnapshot(`
+      {
+        "age": 30,
+        "location": "Quebec",
+        "name": "John Doe",
+      }
+    `)
+  })
+
+  it('rejects non-zui schemas', async () => {
+    const schema = {
+      _output: undefined,
+      safeParse: () => ({ success: true, data: { name: 'John Doe', age: 30 } }),
+    }
+
+    await expect(
+      zai.extract('My name is John Doe, I am 30 years old and I live in Quebec', schema as any).result()
+    ).rejects.toThrow('@bpinternal/zui')
+  })
+
+  it('extract an array of objects from paragraph', async () => {
+    const people = await zai.extract(
+      `
+      My name is John Doe, I am 30 years old and I live in Quebec.
+      My name is Jane Doe, I am 25 years old and I live in Montreal.
+      His name is Jack Doe, he is 35 years old and he lives in Toronto.
+      Her name is Jill Doe, she is 40 years old and she lives in Vancouver.`,
+      z.array(
+        z.object({
+          name: z.string(),
+          age: z.number(),
+          location: z.string(),
+        })
+      )
+    )
+
+    expect(people).toMatchInlineSnapshot(`
+      [
+        {
+          "age": 30,
+          "location": "Quebec",
+          "name": "John Doe",
+        },
+        {
+          "age": 25,
+          "location": "Montreal",
+          "name": "Jane Doe",
+        },
+        {
+          "age": 35,
+          "location": "Toronto",
+          "name": "Jack Doe",
+        },
+        {
+          "age": 40,
+          "location": "Vancouver",
+          "name": "Jill Doe",
+        },
+      ]
+    `)
+  })
+
+  it('extract an object from anything as input', async () => {
+    const person = await zai.extract(
+      {
+        person: { first: 'John', last: 'Doe', age: 30 },
+      },
+      z.object({
+        a: z.string().describe('The full name of the person in the text'),
+        b: z.number().describe('The age of the person in the text'),
+      })
+    )
+
+    expect(person).toMatchInlineSnapshot(`
+      {
+        "a": "John Doe",
+        "b": 30,
+      }
+    `)
+  })
+
+  it('extract age', async () => {
+    const age = await zai.extract(
+      `Countries are Canada, Russia and Pakistan. My favorite colors are red, green, and blue. Dog, cat, fish. I am thirty years old. I was born in 1990.`,
+      z.number().describe('Age of the person')
+    )
+
+    expect(age).toMatchInlineSnapshot(`30`)
+  })
+
+  it('extract an array of string', async () => {
+    const colors = await zai.extract(
+      `Countries are Canada, Russia and Pakistan. My favorite colors are red, green, and blue. Dog, cat, fish.`,
+      z.array(z.string().describe('Color'))
+    )
+
+    expect(colors).toMatchInlineSnapshot(`
+      [
+        "red",
+        "green",
+        "blue",
+      ]
+    `)
+  })
+
+  it('extract a fragmented object from a long text (multi-chunks)', async () => {
+    const TOKEN = 'TOKEN '
+    let text = `Name: John Doe
+\n${TOKEN.repeat(500)}
+Age: 30
+\n${TOKEN.repeat(500)}
+Address: 123 Main St, Anytown, USA
+\n${TOKEN.repeat(500)}
+Phone: (123) 456-7890`
+
+    const { output, usage } = await zai
+      .extract(
+        text,
+        z.object({
+          name: z.string().describe('The name of the person'),
+          age: z.number().describe('The age of the person'),
+          address: z.string().describe('The address of the person'),
+          phone: z.string().describe('The phone number of the person'),
+        }),
+        { chunkLength: 250, strict: true }
+      )
+      .result()
+
+    expect(usage.requests.responses).toBeGreaterThan(5)
+    expect(output).toMatchInlineSnapshot(`
+      {
+        "address": "123 Main St, Anytown, USA",
+        "age": 30,
+        "name": "John Doe",
+        "phone": "(123) 456-7890",
+      }
+    `)
+  })
+
+  it('extract an object of array from a long text (multi-chunks)', async () => {
+    const TOKEN = 'TOKEN '
+    let text = `Feature 1: Tables
+\n${TOKEN.repeat(500)}
+Feature 2: HITL (Human in the Loop)
+\n${TOKEN.repeat(500)}
+Feature 3: Analytics
+\n${TOKEN.repeat(500)}
+Feature 4: Integrations`
+
+    const result = await zai
+      .extract(text, z.object({ features: z.array(z.string()) }), {
+        instructions: 'Extract all features from the text',
+        chunkLength: 250,
+      })
+      .result()
+
+    expect(result.usage.requests.responses).toBeGreaterThan(5)
+    expect(result.output.features.length).toBeGreaterThanOrEqual(4)
+  })
+
+  it('extract an array of discriminated union', async () => {
+    cognitive.on('request', (req) => {
+      console.log(req.input.messages)
+    })
+
+    const schema = z
+      .array(
+        z.discriminatedUnion('type', [
+          z
+            .object({
+              type: z.literal('book'),
+              title: z.string(),
+              author: z.string().describe('The author of the book'),
+            })
+            .describe('A book'),
+          z
+            .object({
+              type: z.literal('animal'),
+              species: z.string(),
+              name: z.string().describe('The name of the animal'),
+            })
+            .describe('An animal'),
+        ])
+      )
+      .describe('An array of books and animals')
+
+    const items = await zai.extract(
+      `I have a book called "The Great Gatsby" by F. Scott Fitzgerald and a pet dog named "Buddy".`,
+      schema
+    )
+
+    expect(items).toMatchInlineSnapshot(`
+      [
+        {
+          "author": "F. Scott Fitzgerald",
+          "title": "The Great Gatsby",
+          "type": "book",
+        },
+        {
+          "name": "Buddy",
+          "species": "dog",
+          "type": "animal",
+        },
+      ]
+    `)
+  })
+
+  it('extract zero elements when none found', async () => {
+    const text = `The quick brown fox jumps over the lazy dog.`
+
+    const tags = await zai.extract(
+      text,
+      z.array(
+        z.object({
+          city: z.string().describe('The name of the city'),
+        })
+      ),
+      {
+        instructions: 'Extract all the cities mentioned in the text. If there is none, return an empty array.',
+      }
+    )
+
+    expect(tags).toMatchInlineSnapshot(`[]`)
+  })
+
+  it('extract an array of objects from a super long text', async () => {
+    const features = await zai.extract(
+      BotpressDocumentation,
+      z.array(
+        z
+          .object({
+            feature: z.string().describe('The name of the feature'),
+            parent: z.string().optional().describe('The parent feature').nullable(),
+            description: z.string().describe('The description of the feature'),
+          })
+          .describe('A feature of Botpress')
+      ),
+      {
+        instructions:
+          'Extract all things that looks like a Botpress feature in the provided input. You must extract a minimum of one element.',
+      }
+    )
+
+    expect(features.length).toBeGreaterThanOrEqual(5)
+    check(features, 'Contains botpress related features, like dashboard or studio or tables').toBe(true)
+  })
+})
+
+describe.sequential('zai.learn.extract', () => {
+  const client = getClient()
+  let tableName = 'ZaiTestExtractInternalTable'
+  let taskId = 'extract'
+  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 extraction format from examples', async () => {
+    const adapter = new TableAdapter({
+      client,
+      tableName,
+    })
+
+    const value = await zai.learn(taskId).extract(
+      `I really liked Casino Royale`,
+      z.object({
+        name: z.string(),
+        movie: z.string(),
+      }),
+      { instructions: 'extract the name of the movie and name of the main character' }
+    )
+
+    check(value, 'extracted james bond and casino royale').toBe(true)
+    check(value, 'the values are NOT IN ALL CAPS').toBe(true)
+
+    let rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBeGreaterThanOrEqual(1)
+
+    await adapter.saveExample({
+      key: 't1',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.extract',
+      instructions: 'extract name of movie and main character',
+      input: `I went to see the Titanic yesterday and I fell asleep`,
+      output: { name: 'JACK DAWSON', movie: 'TITANIC' },
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 't2',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.extract',
+      instructions: 'extract name of movie and main character',
+      input: `Did you know that the gladiator movie has a lot of fighting scenes?`,
+      output: { name: 'MAXIMUS DECIMUS MERIDIUS', movie: 'GLADIATOR' },
+      metadata,
+      status: 'approved',
+    })
+
+    const second = await zai.learn(taskId).extract(
+      `I really liked Casino Royale`,
+      z.object({
+        name: z.string(),
+        movie: z.string(),
+      }),
+      { instructions: 'extract the name of the movie and name of the main character' }
+    )
+
+    expect(second).toMatchInlineSnapshot(`
+      {
+        "movie": "CASINO ROYALE",
+        "name": "JAMES BOND",
+      }
+    `)
+
+    rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBe(3)
+    expect(rows.rows[0].output.value).toMatchObject(second)
+  })
+})

packages/zai/e2e/filter.test.ts

@@ -0,0 +1,155 @@
+import { describe, it, expect, beforeEach, afterEach, afterAll, vi } from 'vitest'
+
+import { getClient, getZai, metadata } from './utils'
+import { TableAdapter } from '../src/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('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 (chatbot)' },
+        { name: 'Ada.cx' },
+        { name: 'Nike' },
+        { name: 'Voiceflow' },
+        { name: 'Adidas' },
+      ],
+      'competes with us',
+      { examples }
+    )
+
+    const names = value.map((v) => v.name)
+    expect(names).toMatchInlineSnapshot(`
+      [
+        "Moveworks (chatbot)",
+        "Ada.cx",
+        "Voiceflow",
+      ]
+    `)
+  })
+})
+
+describe.sequential('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).toBeGreaterThanOrEqual(3)
+    expect(rows.rows.at(-1)!.output.value).toEqual(second)
+  })
+})

packages/zai/e2e/json-utils.test.ts

@@ -0,0 +1,143 @@
+import { describe, it, expect, assert } from 'vitest'
+import { isJsonFile, validateOrRepairJson } from '../src/json-utils'
+
+describe('json-utils', () => {
+  describe('isJsonFile', () => {
+    it('detects .json extension from path', () => {
+      expect(isJsonFile('config.json', 'config.json')).toBe(true)
+      expect(isJsonFile('src/data/config.json', 'config.json')).toBe(true)
+      expect(isJsonFile('package.json', 'package.json')).toBe(true)
+    })
+
+    it('detects .json extension from name when path differs', () => {
+      expect(isJsonFile('/some/path/file', 'file.json')).toBe(true)
+    })
+
+    it('rejects non-json files', () => {
+      expect(isJsonFile('src/hello.ts', 'hello.ts')).toBe(false)
+      expect(isJsonFile('readme.md', 'readme.md')).toBe(false)
+      expect(isJsonFile('data.jsonl', 'data.jsonl')).toBe(false)
+      expect(isJsonFile('config.yaml', 'config.yaml')).toBe(false)
+    })
+
+    it('rejects files with json in the name but not as extension', () => {
+      expect(isJsonFile('json-parser.ts', 'json-parser.ts')).toBe(false)
+      expect(isJsonFile('myjsonfile.txt', 'myjsonfile.txt')).toBe(false)
+    })
+  })
+
+  describe('validateOrRepairJson', () => {
+    // --- valid JSON ---
+
+    it('accepts a valid JSON object', () => {
+      const result = validateOrRepairJson('{"name": "test", "value": 42}')
+      assert(result.valid)
+      expect(result.data).toEqual({ name: 'test', value: 42 })
+      expect(result.repaired).toBe(false)
+    })
+
+    it('accepts a valid JSON array', () => {
+      const result = validateOrRepairJson('[1, 2, 3]')
+      assert(result.valid)
+      expect(result.data).toEqual([1, 2, 3])
+    })
+
+    it('accepts JSON primitives', () => {
+      for (const [input, expected] of [
+        ['"hello"', 'hello'],
+        ['42', 42],
+        ['true', true],
+        ['null', null],
+      ] as const) {
+        const result = validateOrRepairJson(input)
+        assert(result.valid)
+        expect(result.data).toEqual(expected)
+      }
+    })
+
+    it('accepts pretty-printed JSON', () => {
+      const content = JSON.stringify({ name: 'app', version: '1.0.0' }, null, 2)
+      const result = validateOrRepairJson(content)
+      assert(result.valid)
+      expect(result.data).toEqual({ name: 'app', version: '1.0.0' })
+      expect(result.repaired).toBe(false)
+      expect(result.content).toBe(content)
+    })
+
+    it('accepts complex nested JSON', () => {
+      const content = JSON.stringify(
+        {
+          database: { host: 'localhost', port: 5432, pool: { min: 2, max: 10 } },
+          features: ['auth', 'logging'],
+        },
+        null,
+        2
+      )
+      const result = validateOrRepairJson(content)
+      assert(result.valid)
+      expect(result.repaired).toBe(false)
+    })
+
+    // --- repairable JSON ---
+
+    it('repairs missing quotes on keys', () => {
+      const result = validateOrRepairJson('{name: "test"}')
+      assert(result.valid)
+      expect(result.repaired).toBe(true)
+      expect(result.data).toEqual({ name: 'test' })
+    })
+
+    it('repairs trailing commas', () => {
+      const result = validateOrRepairJson('{"a": 1, "b": 2,}')
+      assert(result.valid)
+      expect(result.repaired).toBe(true)
+      expect(result.data).toEqual({ a: 1, b: 2 })
+    })
+
+    it('repairs single quotes', () => {
+      const result = validateOrRepairJson("{'name': 'test'}")
+      assert(result.valid)
+      expect(result.repaired).toBe(true)
+      expect(result.data).toEqual({ name: 'test' })
+    })
+
+    it('repairs missing closing brace', () => {
+      const result = validateOrRepairJson('{"name": "test"')
+      assert(result.valid)
+      expect(result.repaired).toBe(true)
+      expect(result.data).toEqual({ name: 'test' })
+    })
+
+    it('repairs missing closing bracket in array', () => {
+      const result = validateOrRepairJson('[1, 2, 3')
+      assert(result.valid)
+      expect(result.repaired).toBe(true)
+      expect(result.data).toEqual([1, 2, 3])
+    })
+
+    it('repairs trailing comma in nested object', () => {
+      const result = validateOrRepairJson('{"app": {"name": "test",}, "version": "1.0.0"}')
+      assert(result.valid)
+      expect(result.repaired).toBe(true)
+      expect(result.data).toEqual({ app: { name: 'test' }, version: '1.0.0' })
+    })
+
+    // --- invalid JSON (unrepairable) → formatted error ---
+
+    it('returns a formatted error with line numbers when unrepairable', () => {
+      const result = validateOrRepairJson('{}{}{}')
+      assert(!result.valid)
+      expect(result.error).toContain('JSON Syntax Error:')
+      expect(result.error).toMatch(/\d+\|/)
+    })
+
+    it('returns an error marking the exact error line and column', () => {
+      // jsonrepair cannot fix multiple adjacent root-level braced objects without content
+      const content = '{\n  "a": 1\n}{'
+      const result = validateOrRepairJson(content)
+      assert(!result.valid)
+      expect(result.error).toContain('← ERROR')
+      expect(result.error).toContain('^')
+    })
+  })
+})

packages/zai/e2e/label.test.ts

@@ -0,0 +1,247 @@
+import { describe, it, expect, beforeEach, afterEach, afterAll } from 'vitest'
+
+import { BotpressDocumentation, getClient, getZai, metadata } from './utils'
+import { TableAdapter } from '../src/adapters/botpress-table'
+import { check } from '@botpress/vai'
+
+const getValues = <T extends Record<string, { value: boolean }>>(records: T) =>
+  Object.entries(records).reduce((acc, [key, value]) => {
+    acc[key] = value.value
+    return acc
+  }, {}) as Record<keyof T, boolean>
+
+describe('zai.label', { timeout: 60_000 }, () => {
+  const zai = getZai()
+
+  it('simple labels on small text', async () => {
+    const { output: labels } = await zai
+      .label(
+        {
+          name: 'John',
+          story: ['John donated to charity last month.', 'John is loved by his community.'],
+          criminal_record: 'John has no criminal record.',
+        },
+        {
+          is_human: 'is the person a human?',
+          good_person: 'is the person a good person?',
+          bad_person: 'is the person a bad person?',
+          is_criminal: 'is the person a criminal?',
+        }
+      )
+      .result()
+
+    expect(getValues(labels)).toMatchInlineSnapshot(`
+      {
+        "bad_person": false,
+        "good_person": true,
+        "is_criminal": false,
+        "is_human": true,
+      }
+    `)
+  })
+
+  it('simple labels with example', async () => {
+    const labels = {
+      is_human: 'is the person a human?',
+      good_person: 'is the person a good person?',
+      bad_person: 'is the person a bad person?',
+      canadian: 'is the person canadian?',
+      is_french: 'is the person french?',
+    }
+
+    const initial = await zai.label(`Sylvain Perron has no criminal record.`, labels)
+
+    expect(initial.canadian).toBe(false)
+    expect(initial.is_french).toBe(false)
+    expect(initial.bad_person).toBe(false)
+    expect(initial.is_human).toBe(true)
+
+    const second = await zai.label(`Sylvain Perron has no criminal record.`, labels, {
+      examples: [
+        {
+          input: 'Sylvain Pellerin has no criminal record.',
+          labels: {
+            is_french: {
+              label: 'ABSOLUTELY_YES',
+              explanation: 'Important: Sylvain Pellerin is a common French name.',
+            },
+            canadian: {
+              label: 'ABSOLUTELY_YES',
+              explanation: 'Important: We assume all person named Sylvain are Canadian (business rule).',
+            },
+          },
+        },
+        {
+          input: 'Sylvain Bouchard is a criminal.',
+          labels: {
+            bad_person: {
+              label: 'PROBABLY_YES',
+              explanation: 'Important: Sylvain Bouchard is a criminal, so probably a bad person.',
+            },
+            is_french: {
+              label: 'ABSOLUTELY_YES',
+              explanation: 'Important: Sylvain is a common French name.',
+            },
+            canadian: {
+              label: 'ABSOLUTELY_YES',
+              explanation: 'Important: We assume all person named Sylvain are Canadian (business rule).',
+            },
+          },
+        },
+      ],
+    })
+
+    expect(second.canadian).toBe(true)
+    expect(second.is_french).toBe(true)
+    expect(second.is_human).toBe(true)
+    expect(second.bad_person).toBe(false)
+  })
+
+  it('label a huge text', async () => {
+    const labels = await zai.label(BotpressDocumentation, {
+      is_about_animals: 'is the text about animals?',
+      contains_lua_code: 'does the text contain Lua code?',
+      contains_python_code: 'does the text contain Python code?',
+      contains_js_code: 'does the text contain JavaScript code?',
+      is_botpress: 'is the text about Botpress?',
+      is_rasa: 'is the text about Rasa?',
+      has_flows: 'does the text mention flows?',
+      has_api: 'does the text mention the Botpress API?',
+      has_enterprise: 'does the text mention Botpress Enterprise?',
+      has_workspaces: 'does the text mention workspaces?',
+      has_webchat: 'does the text mention the Webchat?',
+      has_hitl: 'does the text mention HITL (human in the loop)?',
+    })
+
+    expect(labels).toMatchInlineSnapshot(`
+      {
+        "contains_js_code": true,
+        "contains_lua_code": false,
+        "contains_python_code": false,
+        "has_api": true,
+        "has_enterprise": true,
+        "has_flows": true,
+        "has_hitl": true,
+        "has_webchat": true,
+        "has_workspaces": true,
+        "is_about_animals": false,
+        "is_botpress": true,
+        "is_rasa": false,
+      }
+    `)
+  })
+})
+
+describe.sequential('zai.learn.label', { timeout: 60_000 }, () => {
+  const client = getClient()
+  let tableName = 'ZaiTestLabelInternalTable'
+  let taskId = 'label'
+  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 labelling rule from examples', async () => {
+    const adapter = new TableAdapter({
+      client,
+      tableName,
+    })
+
+    const value = await zai.learn(taskId).label(`Sylvain Perron has no criminal record.`, {
+      is_human: 'is the person a human?',
+      good_person: 'is the person a good person?',
+      bad_person: 'is the person a bad person?',
+      canadian: 'is the person canadian?',
+      is_french: 'is the person french?',
+    })
+
+    expect(value.is_human).toBe(true)
+    expect(value.is_french).toBe(false)
+    expect(value.canadian).toBe(false)
+
+    let rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBe(1)
+
+    await adapter.saveExample({
+      key: 't1',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.label',
+      instructions: 'label the sentence',
+      input: 'Sylvain Pellerin has no criminal record.',
+      output: {
+        is_french: {
+          label: 'ABSOLUTELY_YES',
+          explanation: 'Important: Sylvain is a common French name.',
+        },
+        canadian: {
+          label: 'ABSOLUTELY_YES',
+          explanation: 'Since we are doing business only in Canada, we assume all users are Canadians.',
+        },
+      },
+      // The below doesn't make sense on purpose, it's just to test the influence of the explanation on the next prediction
+      explanation: `IMPORTANT: Sylvain is a common French name and since we're doing business only in Canada, we assume ALL users are Canadians as soon as they mention a French name.`,
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 't2',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.label',
+      instructions: 'label the sentence',
+      input: 'Joannie Côté has a dog.',
+      output: {
+        is_french: {
+          label: 'ABSOLUTELY_YES',
+          explanation: 'Important: Joannie is a common French name and Côté is a common French last name.',
+        },
+        canadian: {
+          label: 'ABSOLUTELY_YES',
+          explanation: 'Since we are doing business only in Canada, we assume all users are Canadians.',
+        },
+      },
+      metadata,
+      status: 'approved',
+    })
+
+    rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBe(3)
+
+    const second = await zai.learn(taskId).label(`Sylvain Perron has no criminal record.`, {
+      is_human: 'is the person a human?',
+      good_person: 'is the person a good person?',
+      bad_person: 'is the person a bad person?',
+      canadian: 'is the person canadian?',
+      is_french: 'is the person french?',
+    })
+
+    expect(second.is_human).toBe(true)
+    expect(second.is_french).toBe(true)
+    expect(second.canadian).toBe(true)
+
+    rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBe(3)
+
+    check(rows.rows[0].output.value.canadian, 'label is positive (yes) and there is an explanation of why').toBe(true)
+    check(rows.rows[0].output.value.is_french, 'label is positive (yes) and there is an explanation of why').toBe(true)
+  })
+})

packages/zai/e2e/memoize.test.ts

@@ -0,0 +1,250 @@
+import { describe, it, expect, vi } from 'vitest'
+import { z } from '@bpinternal/zui'
+
+import { Zai, type Memoizer } from '../src'
+
+let _responseText = ''
+
+const createMockCognitive = () => {
+  const mock: any = {
+    $$IS_COGNITIVE: true,
+    clone() {
+      return { ...mock, clone: mock.clone, on: mock.on }
+    },
+    on: vi.fn(() => () => {}),
+    generateContent: vi.fn(async () => ({
+      output: { choices: [{ content: _responseText }] },
+      meta: { cached: false, tokens: { input: 10, output: 10 }, cost: { input: 0.001, output: 0.001 } },
+    })),
+    getModelDetails: vi.fn(async () => ({
+      input: { maxTokens: 100_000 },
+      output: { maxTokens: 4096 },
+    })),
+  }
+  return mock
+}
+
+const createSequenceMock = (responses: string[]) => {
+  let callCount = 0
+  const mock = createMockCognitive()
+  mock.generateContent = vi.fn(async () => {
+    const text = responses[callCount] ?? responses[responses.length - 1]
+    callCount++
+    return {
+      output: { choices: [{ content: text }] },
+      meta: { cached: false, tokens: { input: 10, output: 10 }, cost: { input: 0.001, output: 0.001 } },
+    }
+  })
+  return mock
+}
+
+const createSpyMemoizer = () => {
+  const spy = vi.fn(async (_id: string, fn: () => Promise<any>) => fn())
+  const factory = () => ({ run: spy }) as Memoizer
+  return { factory, spy }
+}
+
+describe('memoizer integration', { timeout: 10_000 }, () => {
+  it('zai.check calls memoizer at least 1 time', async () => {
+    _responseText = 'The text is in English.\n■TRUE■\n■END■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    await zai.check('Hello world', 'is english')
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.check:/)
+  })
+
+  it('zai.extract calls memoizer at least 1 time', async () => {
+    _responseText = '■json_start■\n{"name": "John", "age": 30}\n■json_end■\n■NO_MORE_ELEMENT■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    const schema = z.object({ name: z.string(), age: z.number() })
+    await zai.extract('John is 30 years old', schema)
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.extract:/)
+  })
+
+  it('zai.text calls memoizer at least 1 time', async () => {
+    _responseText = 'Hello, this is generated text.'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    await zai.text('Write a greeting')
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.text:/)
+  })
+
+  it('zai.check calls memoizer at least 1 time (non-factory)', async () => {
+    _responseText = 'The text is in English.\n■TRUE■\n■END■'
+    const spy = vi.fn(async (_id: string, fn: () => Promise<any>) => fn())
+    const memoizer: Memoizer = { run: spy }
+    const zai = new Zai({ client: createMockCognitive(), memoize: memoizer })
+    await zai.check('Hello world', 'is english')
+    expect(spy).toHaveBeenCalledTimes(1)
+  })
+
+  it('zai.filter calls memoizer at least 1 time', async () => {
+    _responseText = '■0:true■1:false■END■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    await zai.filter(['apple', 'rock'], 'is a fruit')
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.filter:/)
+  })
+
+  it('zai.label calls memoizer at least 1 time', async () => {
+    _responseText = '■positive:【The text sounds happy】:ABSOLUTELY_YES■\n■END■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    await zai.label('I love this!', { positive: 'is positive sentiment' })
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.label:/)
+  })
+
+  it('zai.rate calls memoizer at least 1 time', async () => {
+    // rate has 2 phases: criteria generation (expects JSON) + scoring
+    const mock = createSequenceMock([
+      '```json\n{"quality": {"very_bad": "terrible", "bad": "poor", "average": "ok", "good": "nice", "very_good": "excellent"}}\n```',
+      '■0:quality=good■\n■END■',
+    ])
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: mock, memoize: factory })
+    await zai.rate(['item1'], 'quality')
+    expect(spy.mock.calls.length).toBeGreaterThanOrEqual(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.rate:/)
+  })
+
+  it('zai.sort calls memoizer at least 1 time', async () => {
+    // sort has 2 phases: criteria generation + scoring
+    const mock = createSequenceMock([
+      '■relevance■\nlow;medium;high\n■END■',
+      '■0:relevance=high■\n■1:relevance=low■\n■END■',
+    ])
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: mock, memoize: factory })
+    await zai.sort(['banana', 'apple'], 'alphabetical order')
+    expect(spy.mock.calls.length).toBeGreaterThanOrEqual(2)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.sort:/)
+  })
+
+  it('zai.answer calls memoizer at least 1 time', async () => {
+    _responseText = '■answer\nBotpress was founded in 2016 ■001.\n■end■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    await zai.answer(['Botpress was founded in 2016.'], 'When was Botpress founded?')
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.answer:/)
+  })
+
+  it('zai.summarize calls memoizer at least 1 time', async () => {
+    _responseText = '■START■\nThis is a summary of the text.\n■END■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    await zai.summarize('This is a long text that needs to be summarized.', { length: 20 })
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:summarize:/)
+  })
+
+  it('zai.rewrite calls memoizer at least 1 time', async () => {
+    _responseText = '■START■\nBonjour le monde\n■END■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    await zai.rewrite('Hello world', 'Translate to French')
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.rewrite:/)
+  })
+
+  it('zai.patch calls memoizer at least 1 time', async () => {
+    _responseText = '<FILE path="hello.ts">\n◼︎=1|console.log("Hi")\n</FILE>'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    await zai.patch([{ name: 'hello.ts', path: 'hello.ts', content: 'console.log("Hello")' }], 'Change Hello to Hi')
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.patch:/)
+  })
+
+  it('zai.group calls memoizer at least 1 time', async () => {
+    _responseText = '■0:Fruits■\n■1:Vegetables■\n■END■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    await zai.group(['apple', 'carrot'])
+    // group may do multiple passes (initial + merge)
+    expect(spy.mock.calls.length).toBeGreaterThanOrEqual(1)
+    expect(spy.mock.calls[0]![0]).toMatch(/^zai:memo:zai\.group:/)
+  })
+
+  it('memoizer key is deterministic for same inputs', async () => {
+    _responseText = 'The text is in English.\n■TRUE■\n■END■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+
+    await zai.check('Hello world', 'is english')
+    await zai.check('Hello world', 'is english')
+
+    expect(spy).toHaveBeenCalledTimes(2)
+    expect(spy.mock.calls[0]![0]).toBe(spy.mock.calls[1]![0])
+  })
+
+  it('memoizer key differs for different inputs', async () => {
+    _responseText = 'The text is in English.\n■TRUE■\n■END■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+
+    await zai.check('Hello world', 'is english')
+    await zai.check('Bonjour le monde', 'is english')
+
+    expect(spy).toHaveBeenCalledTimes(2)
+    expect(spy.mock.calls[0]![0]).not.toBe(spy.mock.calls[1]![0])
+  })
+
+  it('memoizer key differs for different operations', async () => {
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+
+    _responseText = 'The text is in English.\n■TRUE■\n■END■'
+    await zai.check('Hello', 'is english')
+
+    _responseText = 'Generated text.'
+    await zai.text('Hello')
+
+    expect(spy).toHaveBeenCalledTimes(2)
+    expect(spy.mock.calls[0]![0]).not.toBe(spy.mock.calls[1]![0])
+  })
+
+  it('memoizer can short-circuit and return cached result', async () => {
+    // The check transform returns { finalAnswer, explanation }, so the cached result must match
+    const cachedResult = {
+      meta: { cached: true, tokens: { input: 0, output: 0 }, cost: { input: 0, output: 0 } },
+      output: { choices: [{ content: '' }] },
+      text: '',
+      extracted: { finalAnswer: true, explanation: 'cached' },
+    }
+
+    const spy = vi.fn(async (_id: string, _fn: () => Promise<any>) => cachedResult)
+    const memoizer: Memoizer = { run: spy as any }
+    const mock = createMockCognitive()
+    const zai = new Zai({ client: mock, memoize: memoizer })
+
+    const result = await zai.check('anything', 'anything')
+    expect(result).toBe(true)
+    expect(spy).toHaveBeenCalledTimes(1)
+    expect(mock.generateContent).not.toHaveBeenCalled()
+  })
+
+  it('works with .with() chaining', async () => {
+    _responseText = 'The text is in English.\n■TRUE■\n■END■'
+    const { factory, spy } = createSpyMemoizer()
+    const zai = new Zai({ client: createMockCognitive(), memoize: factory })
+    const fast = zai.with({ modelId: 'fast' })
+
+    await fast.check('Hello', 'is english')
+    expect(spy).toHaveBeenCalledTimes(1)
+  })
+
+  it('no memoizer means no wrapping (noop)', async () => {
+    _responseText = 'The text is in English.\n■TRUE■\n■END■'
+    const mock = createMockCognitive()
+    const zai = new Zai({ client: mock })
+    await zai.check('Hello', 'is english')
+    expect(mock.generateContent).toHaveBeenCalled()
+  })
+})

packages/zai/e2e/micropatch.test.ts

@@ -0,0 +1,574 @@
+import { test, expect, describe } from 'vitest'
+import { Micropatch } from '../src/micropatch'
+
+describe('Micropatch', () => {
+  describe('EOL detection', () => {
+    test('detects LF', () => {
+      expect(Micropatch.detectEOL('line1\nline2\n')).toBe('lf')
+    })
+
+    test('detects CRLF', () => {
+      expect(Micropatch.detectEOL('line1\r\nline2\r\n')).toBe('crlf')
+    })
+
+    test('defaults to LF when no line breaks', () => {
+      expect(Micropatch.detectEOL('single line')).toBe('lf')
+    })
+  })
+
+  describe('Delete operations', () => {
+    test('deletes single line', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎-2'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+line3
+"
+`)
+    })
+
+    test('deletes range of lines', () => {
+      const source = 'line1\nline2\nline3\nline4\nline5\n'
+      const ops = '◼︎-2-4'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+line5
+"
+`)
+    })
+
+    test('deletes multiple non-contiguous lines', () => {
+      const source = 'line1\nline2\nline3\nline4\n'
+      const ops = '◼︎-2\n◼︎-4'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+line3
+"
+`)
+    })
+
+    test('deletes first line', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎-1'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line2
+line3
+"
+`)
+    })
+
+    test('deletes last line', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎-3'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+line2
+"
+`)
+    })
+  })
+
+  describe('Replace operations', () => {
+    test('replaces single line with single line', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎=2|REPLACED'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+REPLACED
+line3
+"
+`)
+    })
+
+    test('replaces single line with multiple lines', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎=2|REPLACED1\nREPLACED2\nREPLACED3'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+REPLACED1
+REPLACED2
+REPLACED3
+line3
+"
+`)
+    })
+
+    test('replaces range with single line', () => {
+      const source = 'line1\nline2\nline3\nline4\n'
+      const ops = '◼︎=2-3|REPLACED'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+REPLACED
+line4
+"
+`)
+    })
+
+    test('replaces range with multiple lines', () => {
+      const source = 'line1\nline2\nline3\nline4\n'
+      const ops = '◼︎=2-3|NEW1\nNEW2'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+NEW1
+NEW2
+line4
+"
+`)
+    })
+
+    test('replaces empty payload', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎=2|'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+
+line3
+"
+`)
+    })
+  })
+
+  describe('Insert operations', () => {
+    test('inserts before line', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎<2|INSERTED'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+INSERTED
+line2
+line3
+"
+`)
+    })
+
+    test('inserts after line', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎>2|INSERTED'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+line2
+INSERTED
+line3
+"
+`)
+    })
+
+    test('inserts before first line', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎<1|FIRST'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"FIRST
+line1
+line2
+"
+`)
+    })
+
+    test('inserts after last line', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎>2|LAST'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+line2
+LAST
+"
+`)
+    })
+
+    test('multiple inserts at same position', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎>1|INSERT1\n◼︎>1|INSERT2'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+INSERT2
+INSERT1
+line2
+line3
+"
+`)
+    })
+  })
+
+  describe('Complex multi-operation patches', () => {
+    test('combines delete, replace, and insert', () => {
+      const source = 'line1\nline2\nline3\nline4\nline5\n'
+      const ops = `◼︎-3
+◼︎=2|REPLACED
+◼︎>4|INSERTED`
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+REPLACED
+line4
+INSERTED
+line5
+"
+`)
+    })
+
+    test('operations maintain original line references', () => {
+      const source = 'A\nB\nC\nD\nE\n'
+      // Delete B (line 2), then insert after original C (line 3)
+      // After delete: A, C, D, E
+      // Insert after original C should still work
+      const ops = '◼︎-2\n◼︎>3|X'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"A
+C
+X
+D
+E
+"
+`)
+    })
+  })
+
+  describe('Multiline payloads', () => {
+    test('handles multiline replace with embedded marker escape', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎=2|This has \\◼︎ marker\nSecond line\nThird line'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+This has ◼︎ marker
+Second line
+Third line
+line3
+"
+`)
+    })
+
+    test('multiline payload ends at next op marker', () => {
+      const source = 'line1\nline2\nline3\nline4\n'
+      const ops = `◼︎=2|MULTI1
+MULTI2
+◼︎=4|REPLACED`
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+MULTI1
+MULTI2
+line3
+REPLACED
+"
+`)
+    })
+  })
+
+  describe('EOL handling', () => {
+    test('preserves CRLF line endings', () => {
+      const source = 'line1\r\nline2\r\nline3\r\n'
+      const ops = '◼︎=2|REPLACED'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toBe('line1\r\nREPLACED\r\nline3\r\n')
+    })
+
+    test('uses specified EOL style', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎=2|REPLACED'
+      const result = Micropatch.applyText(source, ops, 'crlf')
+      expect(result).toBe('line1\r\nREPLACED\r\nline3\r\n')
+    })
+
+    test('can switch EOL style', () => {
+      const source = 'line1\r\nline2\r\nline3\r\n'
+      const ops = '◼︎=2|REPLACED'
+      const result = Micropatch.applyText(source, ops, 'lf')
+      expect(result).toBe('line1\nREPLACED\nline3\n')
+    })
+  })
+
+  describe('Micropatch instance methods', () => {
+    test('getText returns current text', () => {
+      const mp = new Micropatch('line1\nline2\n')
+      expect(mp.getText()).toBe('line1\nline2\n')
+    })
+
+    test('setText updates text', () => {
+      const mp = new Micropatch('line1\nline2\n')
+      mp.setText('new1\nnew2\n')
+      expect(mp.getText()).toBe('new1\nnew2\n')
+    })
+
+    test('apply updates internal text', () => {
+      const mp = new Micropatch('line1\nline2\nline3\n')
+      mp.apply('◼︎=2|REPLACED')
+      expect(mp.getText()).toBe('line1\nREPLACED\nline3\n')
+    })
+
+    test('can apply multiple patches sequentially', () => {
+      const mp = new Micropatch('line1\nline2\nline3\n')
+      mp.apply('◼︎=2|REPLACED')
+      mp.apply('◼︎-3')
+      expect(mp.getText()).toBe('line1\nREPLACED\n')
+    })
+
+    test('renderNumberedView shows line numbers', () => {
+      const mp = new Micropatch('line1\nline2\nline3\n')
+      const view = mp.renderNumberedView()
+      expect(view).toMatchInlineSnapshot(`
+"001|line1
+002|line2
+003|line3
+004|"
+`)
+    })
+  })
+
+  describe('Validation', () => {
+    test('validate returns ok and count', () => {
+      const ops = '◼︎-1\n◼︎=2|X\n◼︎>3|Y'
+      const result = Micropatch.validate(ops)
+      expect(result.ok).toBe(true)
+      expect(result.count).toBe(3)
+    })
+
+    test('validate ignores non-op lines', () => {
+      const ops = '# comment\n◼︎-1\n\n◼︎=2|X'
+      const result = Micropatch.validate(ops)
+      expect(result.ok).toBe(true)
+      expect(result.count).toBe(2)
+    })
+  })
+
+  describe('Error handling', () => {
+    test('throws on invalid op syntax', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎INVALID'
+      expect(() => Micropatch.applyText(source, ops)).toThrow('Invalid op syntax')
+    })
+
+    test('throws on range for insert', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎<1-2|X'
+      expect(() => Micropatch.applyText(source, ops)).toThrow('Insert cannot target a range')
+    })
+
+    test('throws on delete with payload', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎-1|INVALID'
+      expect(() => Micropatch.applyText(source, ops)).toThrow('Delete must not have a payload')
+    })
+
+    test('throws on invalid line number', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎-0'
+      expect(() => Micropatch.applyText(source, ops)).toThrow('Invalid line/range')
+    })
+
+    test('throws on invalid range', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎-3-2'
+      expect(() => Micropatch.applyText(source, ops)).toThrow('Invalid line/range')
+    })
+  })
+
+  describe('Edge cases', () => {
+    test('handles empty source', () => {
+      const source = ''
+      const ops = '◼︎<1|NEW'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"NEW
+"
+`)
+    })
+
+    test('handles single line source', () => {
+      const source = 'line1'
+      const ops = '◼︎=1|REPLACED'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`"REPLACED"`)
+    })
+
+    test('skips ops targeting non-existent lines', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎-1\n◼︎-1' // Try to delete line 1 twice
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line2
+"
+`)
+    })
+
+    test('skips inserts targeting non-existent lines', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎<5|BEFORE\n◼︎>9|AFTER' // Stale references beyond EOF
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+line2
+"
+`)
+    })
+
+    test('insert before a line deleted at the top of the file lands where it was', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎-1\n◼︎<1|NEW' // Delete + insert as a replace of line 1
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"NEW
+line2
+"
+`)
+    })
+
+    test('insert after a range deleted at the top of the file lands where it was', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '◼︎-1-2\n◼︎>1|NEW'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"NEW
+line3
+"
+`)
+    })
+
+    test('handles empty op text', () => {
+      const source = 'line1\nline2\n'
+      const ops = ''
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+line2
+"
+`)
+    })
+
+    test('handles blank lines in ops', () => {
+      const source = 'line1\nline2\nline3\n'
+      const ops = '\n\n◼︎=2|REPLACED\n\n'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+REPLACED
+
+
+line3
+"
+`)
+    })
+
+    test('two subsequent replace ops', () => {
+      const source = 'line1\nline2\nline3\nline4\nline5\n'
+      const ops = '◼︎=2|REPLACED\n\n◼︎=5|REPLACED2'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+REPLACED
+
+line3
+line4
+REPLACED2
+"
+`)
+    })
+  })
+
+  describe('Operation order', () => {
+    test('applies deletes before replaces', () => {
+      const source = 'A\nB\nC\nD\n'
+      // Delete C (line 3), replace B (line 2)
+      // Should delete first, then replace
+      const ops = '◼︎=2|X\n◼︎-3'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"A
+X
+D
+"
+`)
+    })
+
+    test('applies replaces before inserts', () => {
+      const source = 'A\nB\nC\n'
+      // Insert after B (line 2), replace B (line 2)
+      // Should replace first, then insert
+      const ops = '◼︎>2|Y\n◼︎=2|X'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"A
+X
+Y
+C
+"
+`)
+    })
+
+    test('applies < inserts before > inserts', () => {
+      const source = 'A\nB\n'
+      // <2 is processed before >1 per operation order
+      const ops = '◼︎>1|AFTER\n◼︎<2|BEFORE'
+      const result = Micropatch.applyText(source, ops)
+      // < ops go first (BEFORE at line 2), then > ops (AFTER after line 1)
+      expect(result).toMatchInlineSnapshot(`
+"A
+AFTER
+BEFORE
+B
+"
+`)
+    })
+  })
+
+  describe('Escaping', () => {
+    test('unescapes marker in insert', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎>1|Text with \\◼︎ marker'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+Text with ◼︎ marker
+line2
+"
+`)
+    })
+
+    test('unescapes marker in replace', () => {
+      const source = 'line1\nline2\n'
+      const ops = '◼︎=2|Escaped \\◼︎ here'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"line1
+Escaped ◼︎ here
+"
+`)
+    })
+
+    test('multiple escaped markers', () => {
+      const source = 'line1\n'
+      const ops = '◼︎=1|\\◼︎ start \\◼︎ middle \\◼︎ end'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"◼︎ start ◼︎ middle ◼︎ end
+"
+`)
+    })
+
+    test('no other escape sequences', () => {
+      const source = 'line1\n'
+      const ops = '◼︎=1|\\n \\t \\r stay literal'
+      const result = Micropatch.applyText(source, ops)
+      expect(result).toMatchInlineSnapshot(`
+"\\n \\t \\r stay literal
+"
+`)
+    })
+  })
+})

packages/zai/e2e/mocks/handlers.ts

@@ -0,0 +1,119 @@
+import { http, HttpResponse, bypass } from 'msw'
+import fs from 'node:fs'
+import path from 'node:path'
+
+// Load cached responses
+const loadCache = () => {
+  const cachePath = path.resolve(__dirname, '../data/cache.jsonl')
+  if (!fs.existsSync(cachePath)) {
+    return new Map()
+  }
+
+  const lines = fs.readFileSync(cachePath, 'utf-8').split(/\r?\n/).filter(Boolean)
+  const cache = new Map<string, any>()
+
+  for (const line of lines) {
+    try {
+      const entry = JSON.parse(line)
+      const inputHash = fastHash(entry.input)
+      cache.set(inputHash, entry.value)
+    } catch {}
+  }
+
+  return cache
+}
+
+function fastHash(str: string): string {
+  let hash = 0
+  const input = typeof str === 'string' ? str : JSON.stringify(str)
+  for (let i = 0; i < input.length; i++) {
+    hash = (hash << 5) - hash + input.charCodeAt(i)
+    hash |= 0 // Convert to 32bit integer
+  }
+  return (hash >>> 0).toString(16) // Convert to unsigned and then to hex
+}
+
+const cache = loadCache()
+
+function stringifyWithSortedKeys(obj: any): string {
+  function sortKeys(input: any): any {
+    if (Array.isArray(input)) {
+      return input.map(sortKeys)
+    } else if (input && typeof input === 'object' && input.constructor === Object) {
+      return Object.keys(input)
+        .sort()
+        .reduce(
+          (acc, key) => {
+            acc[key] = sortKeys(input[key])
+            return acc
+          },
+          {} as Record<string, any>
+        )
+    } else {
+      return input
+    }
+  }
+
+  return JSON.stringify(sortKeys(obj))
+}
+
+export const handlers = [
+  // Mock all Botpress Cloud API requests using cache
+  http.all(/^https:\/\/api\.botpress\.(cloud|dev)\/.*/, async ({ request }) => {
+    // Build request key from method, URL, and body
+    const method = request.method
+    const url = request.url
+    const body = request.method !== 'GET' ? await request.clone().text() : null
+    const shouldCache = !url.includes('/tables') && !url.includes('/v1/admin/')
+
+    const requestData = stringifyWithSortedKeys({
+      method,
+      url: url.toString().replace('.dev/', '.cloud/'), // Normalize dev/cloud URLs
+      body: body ? JSON.parse(body) : null,
+    })
+
+    const hash = fastHash(requestData)
+
+    // Check cache
+    const cached = cache.get(hash)
+    if (cached && shouldCache) {
+      return HttpResponse.json(cached)
+    }
+
+    if (process.env.CI && shouldCache) {
+      return HttpResponse.json(
+        {
+          error: `Missing cached Botpress Cloud response for ${hash}. Run pnpm test:e2e:update to refresh packages/zai/e2e/data/cache.jsonl.`,
+        },
+        { status: 500 }
+      )
+    }
+
+    // Not in cache - fetch from real API
+    try {
+      // Use bypass to avoid infinite recursion
+      const response = await fetch(bypass(request))
+
+      const responseData = await response.json()
+
+      // Cache the response
+      if (shouldCache) {
+        cache.set(hash, responseData)
+        const cachePath = path.resolve(__dirname, '../data/cache.jsonl')
+        fs.appendFileSync(
+          cachePath,
+          JSON.stringify({
+            key: hash,
+            input: requestData,
+            value: responseData,
+          }) + '\n'
+        )
+      }
+
+      return HttpResponse.json(responseData, { status: response.status })
+    } catch (error) {
+      console.error('Error fetching from real API:', error)
+      return HttpResponse.json({ error: 'Failed to fetch from real API' }, { status: 500 })
+    }
+  }),
+]

packages/zai/e2e/mocks/server.ts

@@ -0,0 +1,5 @@
+import { setupServer } from 'msw/node'
+import { handlers } from './handlers'
+
+// Setup MSW server for Node.js environment (Vitest)
+export const server = setupServer(...handlers)

packages/zai/e2e/rate.test.ts

@@ -0,0 +1,712 @@
+// eslint-disable consistent-type-definitions
+import { afterAll, afterEach, beforeEach, describe, expect, it } from 'vitest'
+import { getClient, getZai, metadata } from './utils'
+import { TableAdapter } from '../src/adapters/botpress-table'
+
+describe('zai.rate', () => {
+  describe('string instructions', () => {
+    it('should rate a single item with string instructions', async () => {
+      const zai = getZai()
+      // Unambiguous: complete student record with all fields filled
+      const students = [
+        {
+          name: 'Alice Johnson',
+          studentId: 'STU-001',
+          email: 'alice@school.edu',
+          phone: '+1-555-0100',
+          address: '123 Main St, City, State 12345',
+          enrollmentDate: '2024-01-15',
+          gpa: 3.8,
+          credits: 120,
+        },
+      ]
+      const ratings = await zai.rate(students, 'how complete is this student record?')
+
+      expect(ratings).toHaveLength(1)
+      // Should be rated high (4-5 per criterion) since all fields are present
+      expect(ratings[0]).toBeGreaterThanOrEqual(12) // 3 criteria * 4 minimum
+      expect(ratings[0]).toBeLessThanOrEqual(25) // max possible with 5 criteria * 5
+    })
+
+    it('should rate multiple items with clear differences', async () => {
+      const zai = getZai()
+      const students = [
+        // Complete record - should rate high
+        {
+          name: 'Alice Johnson',
+          studentId: 'STU-001',
+          email: 'alice@school.edu',
+          phone: '+1-555-0100',
+          address: '123 Main St',
+          enrollmentDate: '2024-01-15',
+        },
+        // Minimal record - should rate low
+        {
+          name: 'Bob',
+          studentId: null,
+          email: null,
+          phone: null,
+          address: null,
+          enrollmentDate: null,
+        },
+        // Partial record - should rate medium
+        {
+          name: 'Carol Smith',
+          studentId: 'STU-003',
+          email: 'carol@school.edu',
+          phone: null,
+          address: null,
+          enrollmentDate: '2024-02-01',
+        },
+      ]
+      const ratings = await zai.rate(students, 'how complete is this student record?')
+
+      expect(ratings).toHaveLength(3)
+      // Alice (complete) should be rated higher than Bob (minimal)
+      expect(ratings[0]).toBeGreaterThan(ratings[1]!)
+      // Carol (partial) should be between Alice and Bob
+      expect(ratings[2]).toBeGreaterThan(ratings[1]!)
+      expect(ratings[0]).toBeGreaterThan(ratings[2]!)
+    })
+
+    it('should provide detailed results with string instructions', async () => {
+      const zai = getZai()
+      const forms = [
+        {
+          firstName: 'John',
+          lastName: 'Doe',
+          email: 'john.doe@example.com',
+          phone: '+1-555-0123',
+          address: '456 Oak Ave',
+          city: 'Springfield',
+          state: 'IL',
+          zipCode: '62701',
+        },
+      ]
+      const response = zai.rate(forms, 'how complete is this form submission?')
+      const { output, usage, elapsed } = await response.result()
+
+      expect(output).toHaveLength(1)
+      expect(output[0]).toHaveProperty('total')
+      expect(output[0]?.total).toBeGreaterThanOrEqual(3)
+      expect(output[0]?.total).toBeLessThanOrEqual(25) // max 5 criteria * 5 rating
+
+      // Should have multiple criteria (3-5)
+      const criteriaKeys = Object.keys(output[0] ?? {}).filter((k) => k !== 'total')
+      expect(criteriaKeys.length).toBeGreaterThanOrEqual(3)
+      expect(criteriaKeys.length).toBeLessThanOrEqual(5)
+
+      // Each criterion should be rated 1-5
+      criteriaKeys.forEach((key) => {
+        const score = output[0]?.[key]
+        expect(score).toBeGreaterThanOrEqual(1)
+        expect(score).toBeLessThanOrEqual(5)
+        expect(Number.isInteger(score)).toBe(true)
+      })
+
+      expect(usage.tokens.total).toBeGreaterThan(0)
+      expect(elapsed).toBeGreaterThan(0)
+    })
+  })
+
+  describe('record instructions', () => {
+    it('should rate items with fixed criteria', async () => {
+      const zai = getZai()
+      const passwords = [
+        {
+          password: 'Tr0ng!P@ssw0rd#2024',
+          length: 20,
+          hasUppercase: true,
+          hasLowercase: true,
+          hasNumbers: true,
+          hasSymbols: true,
+        },
+        { password: 'weak', length: 4, hasUppercase: false, hasLowercase: true, hasNumbers: false, hasSymbols: false },
+      ]
+      const ratings = await zai.rate(passwords, {
+        length: 'password length (12+ chars = very_good, 8-11 = good, 6-7 = average, 4-5 = bad, <4 = very_bad)',
+        complexity:
+          'character variety (uppercase+lowercase+numbers+symbols = very_good, 3 types = good, 2 types = average, 1 type = bad)',
+        strength: 'overall password strength against brute force attacks',
+      })
+
+      expect(ratings).toHaveLength(2)
+
+      // First password is strong
+      expect(ratings[0]).toHaveProperty('length')
+      expect(ratings[0]).toHaveProperty('complexity')
+      expect(ratings[0]).toHaveProperty('strength')
+      expect(ratings[0]).toHaveProperty('total')
+      expect(ratings[0]?.length).toBeGreaterThanOrEqual(4) // 20 chars = very_good or good
+      expect(ratings[0]?.complexity).toBeGreaterThanOrEqual(4) // all types = very_good or good
+      expect(ratings[0]?.strength).toBeGreaterThanOrEqual(4)
+      expect(ratings[0]?.total).toBe(
+        (ratings[0]?.length ?? 0) + (ratings[0]?.complexity ?? 0) + (ratings[0]?.strength ?? 0)
+      )
+
+      // Second password is weak
+      expect(ratings[1]?.length).toBeLessThanOrEqual(2) // 4 chars = bad or very_bad
+      expect(ratings[1]?.complexity).toBeLessThanOrEqual(2) // only lowercase = bad or very_bad
+      expect(ratings[1]?.strength).toBeLessThanOrEqual(2)
+      expect(ratings[1]?.total).toBe(
+        (ratings[1]?.length ?? 0) + (ratings[1]?.complexity ?? 0) + (ratings[1]?.strength ?? 0)
+      )
+
+      // Strong password should have higher total than weak
+      expect(ratings[0]?.total).toBeGreaterThan(ratings[1]?.total ?? 0)
+    })
+
+    it('should rate with single criterion', async () => {
+      const zai = getZai()
+      const emails = [
+        { email: 'valid.email@example.com', hasAt: true, hasDomain: true, hasTLD: true },
+        { email: 'invalid-email', hasAt: false, hasDomain: false, hasTLD: false },
+      ]
+      const ratings = await zai.rate(emails, {
+        validity: 'email format validity (has @ symbol, domain, and TLD)',
+      })
+
+      expect(ratings).toHaveLength(2)
+      expect(ratings[0]).toHaveProperty('validity')
+      expect(ratings[0]).toHaveProperty('total')
+      // Valid email should rate high
+      expect(ratings[0]?.validity).toBeGreaterThanOrEqual(4)
+      expect(ratings[0]?.total).toBe(ratings[0]?.validity)
+
+      // Invalid email should rate low
+      expect(ratings[1]?.validity).toBeLessThanOrEqual(2)
+      expect(ratings[1]?.total).toBe(ratings[1]?.validity)
+    })
+
+    it('should provide detailed results with record instructions', async () => {
+      const zai = getZai()
+      const files = [
+        {
+          filename: 'document.pdf',
+          sizeBytes: 1024000,
+          hasExtension: true,
+          validName: true,
+          reasonableSize: true,
+        },
+      ]
+      const response = zai.rate(files, {
+        naming: 'filename follows conventions (has extension, no special chars)',
+        size: 'file size is reasonable (not too large or too small)',
+      })
+      const { output, usage, elapsed } = await response.result()
+
+      expect(output).toHaveLength(1)
+      expect(output[0]).toHaveProperty('naming')
+      expect(output[0]).toHaveProperty('size')
+      expect(output[0]).toHaveProperty('total')
+      expect(output[0]?.total).toBe((output[0]?.naming ?? 0) + (output[0]?.size ?? 0))
+
+      expect(usage.tokens.total).toBeGreaterThan(0)
+      expect(elapsed).toBeGreaterThan(0)
+    })
+  })
+
+  describe('large arrays', () => {
+    it('should rate many items efficiently with parallelization', async () => {
+      const zai = getZai()
+      // Create 500 items to force chunking and parallel processing
+      const items = Array.from({ length: 500 }, (_, i) => ({
+        id: i,
+        // Every 10th item is complete, others are incomplete
+        requiredField1: i % 10 === 0 ? 'present' : null,
+        requiredField2: i % 10 === 0 ? 'present' : null,
+        requiredField3: i % 10 === 0 ? 'present' : null,
+        optionalField: i % 10 === 0 ? 'present' : null,
+      }))
+
+      const startTime = Date.now()
+      const ratings = await zai.rate(items, 'how complete are the required fields?')
+      const elapsed = Date.now() - startTime
+
+      expect(ratings).toHaveLength(500)
+
+      // Verify ratings are in correct order (matching input order)
+      for (let i = 0; i < 500; i++) {
+        expect(ratings[i]).toBeDefined()
+        expect(typeof ratings[i]).toBe('number')
+
+        if (i % 10 === 0) {
+          // Complete items should rate high
+          expect(ratings[i]).toBeGreaterThanOrEqual(12) // 3 criteria * 4
+        } else {
+          // Incomplete items should rate low
+          expect(ratings[i]).toBeLessThanOrEqual(10) // Should be lower
+        }
+      }
+
+      // With 500 items and max 50 per chunk, we should have ~10 chunks
+      // Parallel processing should complete significantly faster than sequential
+      // (this is a sanity check, not a strict performance test)
+      console.log(`Rated 500 items in ${elapsed}ms`)
+      expect(elapsed).toBeLessThan(300000) // Should complete within 5 minutes
+    })
+
+    it('should maintain order with large parallel chunks', async () => {
+      const zai = getZai()
+      // Create 300 items with predictable patterns
+      const items = Array.from({ length: 300 }, (_, i) => ({
+        id: i,
+        value: i * 2,
+        isEven: i % 2 === 0,
+      }))
+
+      const ratings = await zai.rate(items, {
+        id_presence: 'does it have an id field?',
+        value_presence: 'does it have a value field?',
+      })
+
+      expect(ratings).toHaveLength(300)
+
+      // Verify order is maintained
+      for (let i = 0; i < 300; i++) {
+        expect(ratings[i]).toBeDefined()
+        // All items have id and value, so all should rate high
+        expect(ratings[i]?.id_presence).toBeGreaterThanOrEqual(4)
+        expect(ratings[i]?.value_presence).toBeGreaterThanOrEqual(4)
+        expect(ratings[i]?.total).toBe((ratings[i]?.id_presence ?? 0) + (ratings[i]?.value_presence ?? 0))
+      }
+    })
+
+    it('should rate many items efficiently', async () => {
+      const zai = getZai()
+      // Create 100 items with objective completion scores
+      const items = Array.from({ length: 100 }, (_, i) => ({
+        id: i,
+        // Every 10th item is complete, others are incomplete
+        requiredField1: i % 10 === 0 ? 'present' : null,
+        requiredField2: i % 10 === 0 ? 'present' : null,
+        requiredField3: i % 10 === 0 ? 'present' : null,
+        optionalField: i % 10 === 0 ? 'present' : null,
+      }))
+
+      const ratings = await zai.rate(items, 'how complete are the required fields?')
+
+      expect(ratings).toHaveLength(100)
+
+      // Items at index 0, 10, 20, etc. should be rated higher
+      for (let i = 0; i < 100; i++) {
+        if (i % 10 === 0) {
+          // Complete items should rate high
+          expect(ratings[i]).toBeGreaterThanOrEqual(12) // 3 criteria * 4
+        } else {
+          // Incomplete items should rate low
+          expect(ratings[i]).toBeLessThanOrEqual(10) // Should be lower
+        }
+      }
+    })
+
+    it('should handle chunking with record instructions', async () => {
+      const zai = getZai()
+      const items = Array.from({ length: 75 }, (_, i) => ({
+        id: i,
+        isEven: i % 2 === 0,
+        isDivisibleBy5: i % 5 === 0,
+      }))
+
+      const ratings = await zai.rate(items, {
+        evenness: 'is the ID an even number?',
+        divisibility: 'is the ID divisible by 5?',
+      })
+
+      expect(ratings).toHaveLength(75)
+      ratings.forEach((rating, idx) => {
+        expect(rating).toHaveProperty('evenness')
+        expect(rating).toHaveProperty('divisibility')
+        expect(rating).toHaveProperty('total')
+        expect(rating.total).toBe(rating.evenness + rating.divisibility)
+
+        // Check logical ratings
+        if (idx % 2 === 0) {
+          // Even numbers should rate high on evenness
+          expect(rating.evenness).toBeGreaterThanOrEqual(4)
+        } else {
+          // Odd numbers should rate low on evenness
+          expect(rating.evenness).toBeLessThanOrEqual(2)
+        }
+
+        if (idx % 5 === 0) {
+          // Divisible by 5 should rate high
+          expect(rating.divisibility).toBeGreaterThanOrEqual(4)
+        }
+      })
+    })
+  })
+
+  describe('edge cases', () => {
+    it('should handle empty array', async () => {
+      const zai = getZai()
+      const ratings = await zai.rate([], 'rate this')
+
+      expect(ratings).toHaveLength(0)
+      expect(ratings).toEqual([])
+    })
+
+    it('should handle single item edge case', async () => {
+      const zai = getZai()
+      const items = [{ status: 'complete', fields: 10, validated: true }]
+      const ratings = await zai.rate(items, 'is this item complete?')
+
+      expect(ratings).toHaveLength(1)
+      expect(ratings[0]).toBeGreaterThanOrEqual(3)
+    })
+
+    it('should handle items with clear binary outcomes', async () => {
+      const zai = getZai()
+      const checks = [
+        { testsPassing: true, buildSuccessful: true, lintClean: true, coverageAbove80: true },
+        { testsPassing: false, buildSuccessful: false, lintClean: false, coverageAbove80: false },
+      ]
+      const ratings = await zai.rate(checks, 'how healthy is this codebase status?')
+
+      expect(ratings).toHaveLength(2)
+      // First item (all passing) should rate significantly higher than second (all failing)
+      expect(ratings[0]).toBeGreaterThan(ratings[1]! * 2)
+    })
+  })
+
+  describe('options', () => {
+    it('should respect tokensPerItem option', async () => {
+      const zai = getZai()
+      const largeItem = {
+        description: 'A'.repeat(10000),
+        hasDescription: true,
+        isVeryLong: true,
+      }
+      const ratings = await zai.rate([largeItem], 'does this have a description?', { tokensPerItem: 50 })
+
+      expect(ratings).toHaveLength(1)
+      expect(ratings[0]).toBeGreaterThanOrEqual(3)
+    })
+
+    it('should respect maxItemsPerChunk option', async () => {
+      const zai = getZai()
+      const items = Array.from({ length: 10 }, (_, i) => ({ id: i, isPresent: true }))
+      const ratings = await zai.rate(items, 'is the id field present?', { maxItemsPerChunk: 3 })
+
+      expect(ratings).toHaveLength(10)
+      // All should rate high since id is always present
+      ratings.forEach((rating) => {
+        expect(rating).toBeGreaterThanOrEqual(3)
+      })
+    })
+  })
+
+  describe('progress tracking', () => {
+    it('should emit progress events', async () => {
+      const zai = getZai()
+      const items = Array.from({ length: 20 }, (_, i) => ({ id: i, value: i }))
+      const response = zai.rate(items, 'rate these')
+
+      const progressEvents: number[] = []
+      response.on('progress', (usage) => {
+        progressEvents.push(usage.requests.percentage)
+      })
+
+      await response
+
+      expect(progressEvents.length).toBeGreaterThan(0)
+      // Last event should be 100% or close to it
+      const lastProgress = progressEvents[progressEvents.length - 1]
+      expect(lastProgress).toBeGreaterThan(0)
+      expect(lastProgress).toBeLessThanOrEqual(1)
+    })
+  })
+
+  describe('abort signal', () => {
+    it('should support abort signal', async () => {
+      const zai = getZai()
+      // Use more items to ensure the operation takes longer
+      const items = Array.from({ length: 200 }, (_, i) => ({ id: i, data: 'x'.repeat(100) }))
+      const controller = new AbortController()
+      const response = zai.rate(items, 'rate these items')
+      response.bindSignal(controller.signal)
+
+      // Abort immediately
+      controller.abort()
+
+      await expect(response).rejects.toThrow()
+    })
+  })
+
+  describe('objective numerical criteria', () => {
+    it('should rate based on clear numerical thresholds', async () => {
+      const zai = getZai()
+      const servers = [
+        { cpu: 95, memory: 90, disk: 85, uptime: 30 }, // Critical - should rate low
+        { cpu: 50, memory: 55, disk: 60, uptime: 99.9 }, // Healthy - should rate high
+        { cpu: 75, memory: 70, disk: 65, uptime: 95 }, // Warning - should rate medium
+      ]
+
+      const ratings = await zai.rate(servers, {
+        cpu_health:
+          'CPU usage health (0-50% = very_good, 51-70% = good, 71-80% = average, 81-90% = bad, 90%+ = very_bad)',
+        memory_health:
+          'Memory usage health (0-50% = very_good, 51-70% = good, 71-80% = average, 81-90% = bad, 90%+ = very_bad)',
+        disk_health:
+          'Disk usage health (0-50% = very_good, 51-70% = good, 71-80% = average, 81-90% = bad, 90%+ = very_bad)',
+      })
+
+      expect(ratings).toHaveLength(3)
+
+      // Server 0 (critical): all metrics >85% should rate badly
+      expect(ratings[0]?.cpu_health).toBeLessThanOrEqual(2)
+      expect(ratings[0]?.memory_health).toBeLessThanOrEqual(2)
+      expect(ratings[0]?.disk_health).toBeLessThanOrEqual(2)
+
+      // Server 1 (healthy): all metrics 50-60% should rate well
+      expect(ratings[1]?.cpu_health).toBeGreaterThanOrEqual(4)
+      expect(ratings[1]?.memory_health).toBeGreaterThanOrEqual(4)
+      expect(ratings[1]?.disk_health).toBeGreaterThanOrEqual(3)
+
+      // Healthy server should have higher total than critical server
+      expect(ratings[1]?.total).toBeGreaterThan(ratings[0]?.total ?? 0)
+    })
+
+    it('should rate based on count criteria', async () => {
+      const zai = getZai()
+      const repos = [
+        { name: 'repo-a', openIssues: 2, contributors: 50, stars: 1000 },
+        { name: 'repo-b', openIssues: 200, contributors: 1, stars: 5 },
+      ]
+
+      const ratings = await zai.rate(repos, {
+        issue_management:
+          'open issues count (0-10 = very_good, 11-50 = good, 51-100 = average, 101-200 = bad, 200+ = very_bad)',
+        community: 'contributor count (50+ = very_good, 20-49 = good, 10-19 = average, 5-9 = bad, <5 = very_bad)',
+        popularity: 'star count (1000+ = very_good, 500-999 = good, 100-499 = average, 50-99 = bad, <50 = very_bad)',
+      })
+
+      expect(ratings).toHaveLength(2)
+
+      // repo-a has good metrics
+      expect(ratings[0]?.issue_management).toBeGreaterThanOrEqual(4) // 2 issues = very_good
+      expect(ratings[0]?.community).toBeGreaterThanOrEqual(4) // 50 contributors = very_good
+      expect(ratings[0]?.popularity).toBeGreaterThanOrEqual(4) // 1000 stars = very_good
+
+      // repo-b has poor metrics
+      expect(ratings[1]?.issue_management).toBeLessThanOrEqual(2) // 200 issues = bad
+      expect(ratings[1]?.community).toBeLessThanOrEqual(2) // 1 contributor = very_bad
+      expect(ratings[1]?.popularity).toBeLessThanOrEqual(2) // 5 stars = very_bad
+
+      expect(ratings[0]?.total).toBeGreaterThan(ratings[1]?.total ?? 0)
+    })
+  })
+})
+
+describe.sequential('zai.learn.rate', () => {
+  const client = getClient()
+  const tableName = 'ZaiTestRateInternalTable'
+  const taskId = 'rate'
+  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) {}
+  })
+
+  // TODO: fix and re-enable
+  it.skip('learns counterintuitive rating pattern that LLM cannot guess', async () => {
+    const adapter = new TableAdapter({
+      client,
+      tableName,
+    })
+
+    // Test a COUNTERINTUITIVE pattern: Custom email domain ratings
+    // Custom company-specific domain importance rules that LLM cannot know without learning:
+    // @sequoia.vc = our investor (highest importance - rate 5/5/5)
+    // @a16z.com = potential investor (high importance - rate 4/4/4)
+    // @google.com = competitor (medium importance - rate 3/3/3)
+    // analyst@* = spam regardless of domain (lowest importance - rate 1/1/1)
+
+    // Add approved examples showing the domain pattern
+    await adapter.saveExample({
+      key: 'email_domain1',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rate',
+      instructions: 'rate email sender importance',
+      input: JSON.stringify([{ from: 'partner@sequoia.vc', subject: 'Q4 Review' }]),
+      output: [{ trustworthiness: 5, priority: 5, relevance: 5, total: 15 }],
+      explanation: 'RULE: @sequoia.vc is our investor - highest importance rating',
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 'email_domain2',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rate',
+      instructions: 'rate email sender importance',
+      input: JSON.stringify([{ from: 'analyst@bankofamerica.com', subject: 'Market Report' }]),
+      output: [{ trustworthiness: 1, priority: 1, relevance: 1, total: 3 }],
+      explanation: 'RULE: analyst@* prefix is spam - lowest importance rating',
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 'email_domain3',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rate',
+      instructions: 'rate email sender importance',
+      input: JSON.stringify([{ from: 'ben@a16z.com', subject: 'Investment Discussion' }]),
+      output: [{ trustworthiness: 4, priority: 4, relevance: 4, total: 12 }],
+      explanation: 'RULE: @a16z.com is potential investor - high importance rating',
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 'email_domain4',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rate',
+      instructions: 'rate email sender importance',
+      input: JSON.stringify([{ from: 'team@google.com', subject: 'Partnership Proposal' }]),
+      output: [{ trustworthiness: 3, priority: 3, relevance: 3, total: 9 }],
+      explanation: 'RULE: @google.com is competitor - medium importance rating',
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 'email_domain5',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rate',
+      instructions: 'rate email sender importance',
+      input: JSON.stringify([{ from: 'roelof@sequoia.vc', subject: 'Portfolio Update' }]),
+      output: [{ trustworthiness: 5, priority: 5, relevance: 5, total: 15 }],
+      explanation: 'RULE: @sequoia.vc is our investor - highest importance rating',
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 'email_domain6',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rate',
+      instructions: 'rate email sender importance',
+      input: JSON.stringify([{ from: 'analyst@jpmorgan.com', subject: 'Stock Analysis' }]),
+      output: [{ trustworthiness: 1, priority: 1, relevance: 1, total: 3 }],
+      explanation: 'RULE: analyst@* prefix is spam - lowest importance rating',
+      metadata,
+      status: 'approved',
+    })
+
+    // Now test with learned examples
+    const ratings = await zai.learn(taskId).rate(
+      [
+        { from: 'sarah@sequoia.vc', subject: 'Board Meeting' }, // investor - should rate HIGH
+        { from: 'analyst@goldmansachs.com', subject: 'Earnings Report' }, // analyst spam - should rate LOW
+        { from: 'marc@a16z.com', subject: 'Funding Round' }, // potential investor - should rate HIGH
+        { from: 'recruiter@google.com', subject: 'Hiring' }, // competitor - should rate MEDIUM
+      ],
+      'rate email sender importance'
+    )
+
+    expect(ratings).toHaveLength(4)
+
+    // With the pattern learned, domain ratings should follow the rules
+    const sequoiaScore = ratings[0]! // @sequoia.vc (investor - HIGH)
+    const analystScore = ratings[1]! // analyst@* (spam - LOW)
+    const a16zScore = ratings[2]! // @a16z.com (potential investor - HIGH)
+    const googleScore = ratings[3]! // @google.com (competitor - MEDIUM)
+
+    // Most reliable pattern: analyst@ should ALWAYS score lowest
+    expect(analystScore).toBeLessThan(sequoiaScore)
+    expect(analystScore).toBeLessThan(a16zScore)
+    expect(analystScore).toBeLessThanOrEqual(googleScore)
+
+    // Investors should score higher than competitor
+    const investorAvg = (sequoiaScore + a16zScore) / 2
+    expect(investorAvg).toBeGreaterThan(googleScore)
+
+    // Sequoia (our investor) should score very high
+    expect(sequoiaScore).toBeGreaterThanOrEqual(10) // At least 10/15
+
+    const rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBeGreaterThanOrEqual(6) // 6 examples + new ratings
+  })
+
+  it('learns rating patterns from examples with record instructions', async () => {
+    const adapter = new TableAdapter({
+      client,
+      tableName,
+    })
+
+    // Add approved examples with specific criteria
+    await adapter.saveExample({
+      key: 'password1',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rate',
+      instructions: 'rate password strength',
+      input: JSON.stringify([{ password: 'Str0ng!P@ss#2024', length: 16, hasAll: true }]),
+      output: [{ length: 5, complexity: 5, strength: 5, total: 15 }],
+      explanation: 'Strong password: 16 chars, all character types',
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 'password2',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rate',
+      instructions: 'rate password strength',
+      input: JSON.stringify([{ password: 'weak', length: 4, hasAll: false }]),
+      output: [{ length: 1, complexity: 1, strength: 1, total: 3 }],
+      explanation: 'Weak password: only 4 chars, missing character types',
+      metadata,
+      status: 'approved',
+    })
+
+    // Rate passwords with learned patterns
+    const ratings = await zai.learn(taskId).rate(
+      [
+        { password: 'MyStr0ng!Pass', length: 13, hasAll: true }, // Strong
+        { password: 'bad', length: 3, hasAll: false }, // Weak
+      ],
+      {
+        length: 'password length (12+ = very_good, 8-11 = good, 6-7 = average, 4-5 = bad, <4 = very_bad)',
+        complexity: 'character variety (all types = very_good)',
+        strength: 'overall password strength',
+      }
+    )
+
+    expect(ratings).toHaveLength(2)
+
+    // Strong password should rate high (learned pattern)
+    expect(ratings[0]?.length).toBeGreaterThanOrEqual(4) // 13 chars
+    expect(ratings[0]?.complexity).toBeGreaterThanOrEqual(4) // has all types
+    expect(ratings[0]?.strength).toBeGreaterThanOrEqual(4)
+    expect(ratings[0]?.total).toBeGreaterThanOrEqual(12)
+
+    // Weak password should rate low (learned pattern)
+    expect(ratings[1]?.length).toBeLessThanOrEqual(2) // 3 chars
+    expect(ratings[1]?.complexity).toBeLessThanOrEqual(2) // missing types
+    expect(ratings[1]?.strength).toBeLessThanOrEqual(2)
+    expect(ratings[1]?.total).toBeLessThanOrEqual(6)
+
+    const rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBeGreaterThanOrEqual(2)
+  })
+})

packages/zai/e2e/rewrite.test.ts

@@ -0,0 +1,114 @@
+import { describe, it, expect, beforeEach, afterEach, afterAll } from 'vitest'
+import { check } from '@botpress/vai'
+
+import { getClient, getZai, metadata, tokenizer } from './utils'
+import { TableAdapter } from '../src/adapters/botpress-table'
+
+const Zoe = `
+Part 1. Zoe walks to the park.
+Part 2. She meets her friend.
+Part 3. They play together.
+Part 4. They have a picnic.
+Part 5. They go home.
+`.trim()
+
+describe('zai.rewrite', { timeout: 60_000 }, () => {
+  const zai = getZai()
+
+  it('transforms text to all caps', async () => {
+    const result = await zai.rewrite(`Hello, what is the time today?`, 'write in all caps')
+    expect(result).toBe(`HELLO, WHAT IS THE TIME TODAY?`)
+  })
+
+  it('transforms text to all caps and respects tokens restrictions', async () => {
+    const result = await zai.rewrite(Zoe, 'write in all caps', { length: 20 })
+    expect(tokenizer.count(result)).toBeLessThanOrEqual(30)
+    expect(result).toContain(`ZOE`)
+    expect(result).toContain(`PARK`)
+    expect(result).toBe(result.toUpperCase())
+    expect(result).not.toContain(`PART 3`)
+  })
+
+  it('french translation of the story', async () => {
+    const result = await zai.rewrite(Zoe, 'translate to french')
+    check(result, 'is a french story about Zeo and with 5 parts').toBe(true)
+  })
+
+  it('Throws if input is bigger than the model max tokens', async () => {
+    await expect(zai.rewrite(Zoe.repeat(100_000), 'translate to french')).rejects.toThrow(/tokens/i)
+  })
+})
+
+describe.sequential('zai.learn.rewrite', { timeout: 60_000 }, () => {
+  const client = getClient()
+  let tableName = 'ZaiTestRewriteInternalTable'
+  let taskId = 'rewrite'
+  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 rewrite rules from examples', async () => {
+    const adapter = new TableAdapter({
+      client,
+      tableName,
+    })
+
+    const value = await zai.learn(taskId).rewrite(`Botpress is awesome`, 'write it like we want it')
+
+    let rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBeGreaterThanOrEqual(1)
+    expect(rows.rows[0].output.value).toBe(value)
+
+    await adapter.saveExample({
+      key: 't1',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rewrite',
+      instructions: 'write it like we want it',
+      input: 'Microsoft is a big company',
+      output: `# MICROSOFT IS A BIG COMPANY`,
+      metadata,
+      status: 'approved',
+    })
+
+    await adapter.saveExample({
+      key: 't2',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.rewrite',
+      instructions: 'write it like we want it',
+      input: 'Google is an evil company',
+      output: `# GOOGLE IS AN EVIL COMPANY`,
+      metadata,
+      status: 'approved',
+    })
+
+    const second = await zai.learn(taskId).rewrite(`Botpress is awesome`, 'write it like we want it')
+
+    rows = await client.findTableRows({ table: tableName })
+    expect(rows.rows.length).toBe(3)
+
+    check(second, `The text is "BOTPRESS IS AWESOME" and starts with a hashtag`).toBe(true)
+
+    expect(rows.rows.length).toBe(3)
+    expect(rows.rows[0].output.value).toBe(second)
+  })
+})

packages/zai/e2e/sort.test.ts

@@ -0,0 +1,966 @@
+import { afterAll, describe, expect, it } from 'vitest'
+import { getClient, getZai, metadata } from './utils'
+import { TableAdapter } from '../src/adapters/botpress-table'
+
+describe('zai.sort', () => {
+  const zai = getZai()
+
+  describe('animal sorting', () => {
+    it('should sort animals by speed (ascending - slowest to fastest)', async () => {
+      const animals = [
+        { name: 'Cheetah', type: 'land' },
+        { name: 'Sloth', type: 'land' },
+        { name: 'Horse', type: 'land' },
+        { name: 'Snail', type: 'land' },
+        { name: 'Human', type: 'land' },
+        { name: 'Turtle', type: 'land' },
+        { name: 'Lion', type: 'land' },
+        { name: 'Elephant', type: 'land' },
+      ]
+
+      const sorted = await zai.sort(animals, 'from slowest to fastest')
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(animals.length)
+
+      // Verify all elements are present
+      expect(sorted.map((a) => a.name).sort()).toEqual(animals.map((a) => a.name).sort())
+
+      // Verify order: slowest animals should be first
+      const slowestIndex = sorted.findIndex((a) => a.name === 'Sloth' || a.name === 'Snail')
+      const fastestIndex = sorted.findIndex((a) => a.name === 'Cheetah')
+      expect(slowestIndex).toBeLessThan(fastestIndex)
+
+      // Cheetah should be last (fastest)
+      const cheetahIndex = sorted.findIndex((a) => a.name === 'Cheetah')
+      expect(cheetahIndex).toBeGreaterThan(sorted.length / 2)
+    })
+
+    it('should sort animals by speed (descending - fastest to slowest)', async () => {
+      const animals = [
+        { name: 'Snail', type: 'land' },
+        { name: 'Cheetah', type: 'land' },
+        { name: 'Turtle', type: 'land' },
+        { name: 'Horse', type: 'land' },
+        { name: 'Human', type: 'land' },
+        { name: 'Sloth', type: 'land' },
+        { name: 'Lion', type: 'land' },
+      ]
+
+      const sorted = await zai.sort(animals, 'from fastest to slowest')
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(animals.length)
+
+      // Cheetah should be first (fastest)
+      const cheetahIndex = sorted.findIndex((a) => a.name === 'Cheetah')
+      expect(cheetahIndex).toBeLessThan(sorted.length / 2)
+
+      // Sloth or Snail should be last
+      const lastAnimal = sorted[sorted.length - 1]
+      expect(['Sloth', 'Snail', 'Turtle']).toContain(lastAnimal.name)
+    })
+
+    it('should sort animals by danger level (ascending - least to most dangerous)', async () => {
+      const animals = [
+        { name: 'Grizzly Bear', habitat: 'forest' },
+        { name: 'Rabbit', habitat: 'grassland' },
+        { name: 'Hippopotamus', habitat: 'river' },
+        { name: 'Deer', habitat: 'forest' },
+        { name: 'Great White Shark', habitat: 'ocean' },
+        { name: 'Butterfly', habitat: 'garden' },
+        { name: 'Crocodile', habitat: 'swamp' },
+        { name: 'Lion', habitat: 'savanna' },
+        { name: 'Lamb', habitat: 'farm' },
+      ]
+
+      const sorted = await zai.sort(animals, 'from least dangerous to most dangerous')
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(animals.length)
+
+      // Harmless animals should be first
+      const harmlessIndex = Math.max(
+        sorted.findIndex((a) => a.name === 'Rabbit'),
+        sorted.findIndex((a) => a.name === 'Butterfly'),
+        sorted.findIndex((a) => a.name === 'Lamb')
+      )
+
+      // Dangerous animals should be last
+      const dangerousIndex = Math.min(
+        sorted.findIndex((a) => a.name === 'Great White Shark'),
+        sorted.findIndex((a) => a.name === 'Grizzly Bear'),
+        sorted.findIndex((a) => a.name === 'Crocodile')
+      )
+
+      expect(harmlessIndex).toBeLessThan(dangerousIndex)
+
+      // Most dangerous should be in second half
+      const sharkIndex = sorted.findIndex((a) => a.name === 'Great White Shark')
+      expect(sharkIndex).toBeGreaterThan(sorted.length / 2)
+    })
+
+    it('should sort animals by danger level (descending - most to least dangerous)', async () => {
+      const animals = [
+        { name: 'Rabbit', type: 'mammal' },
+        { name: 'Lion', type: 'mammal' },
+        { name: 'Butterfly', type: 'insect' },
+        { name: 'Crocodile', type: 'reptile' },
+        { name: 'Deer', type: 'mammal' },
+        { name: 'Cobra', type: 'reptile' },
+        { name: 'Squirrel', type: 'mammal' },
+      ]
+
+      const sorted = await zai.sort(animals, 'from most dangerous to least dangerous')
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(animals.length)
+
+      // Dangerous animals should be first
+      const firstAnimal = sorted[0]
+      expect(['Lion', 'Crocodile', 'Cobra']).toContain(firstAnimal.name)
+
+      // Harmless animals should be last
+      const lastAnimal = sorted[sorted.length - 1]
+      expect(['Rabbit', 'Butterfly', 'Squirrel']).toContain(lastAnimal.name)
+    })
+  })
+
+  describe('email priority sorting', () => {
+    it('should sort emails by priority (spam to urgent bills)', async () => {
+      const emails = [
+        { subject: 'Electricity bill due tomorrow', from: 'utility@electric.com', date: '2024-01-15' },
+        { subject: 'Hot singles in your area!', from: 'spam@xyz.com', date: '2024-01-14' },
+        { subject: 'Team meeting notes', from: 'colleague@work.com', date: '2024-01-15' },
+        { subject: 'Your credit card payment is overdue', from: 'bank@chase.com', date: '2024-01-15' },
+        { subject: 'Get rich quick scheme', from: 'scam@fake.com', date: '2024-01-13' },
+        { subject: 'Newsletter: Weekly updates', from: 'newsletter@blog.com', date: '2024-01-14' },
+        { subject: 'Rent payment due in 2 days', from: 'landlord@property.com', date: '2024-01-15' },
+        { subject: 'Congratulations! You won the lottery!', from: 'notreal@scam.com', date: '2024-01-12' },
+        { subject: 'Project deadline reminder', from: 'manager@work.com', date: '2024-01-15' },
+      ]
+
+      const sorted = await zai.sort(emails, 'from least urgent (spam) to most urgent (bills to pay)')
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(emails.length)
+
+      // Spam should be first
+      const firstEmail = sorted[0]
+      expect(
+        ['Hot singles', 'Get rich quick', 'Congratulations', 'lottery'].some((spam) =>
+          firstEmail.subject.includes(spam)
+        )
+      ).toBe(true)
+
+      // Urgent bills should be last
+      const lastEmail = sorted[sorted.length - 1]
+      expect(['bill', 'payment', 'overdue', 'Rent'].some((urgent) => lastEmail.subject.includes(urgent))).toBe(true)
+
+      // Verify bills are in the last third
+      const billIndices = sorted
+        .map((e, i) => (['bill', 'payment', 'overdue', 'Rent'].some((urgent) => e.subject.includes(urgent)) ? i : -1))
+        .filter((i) => i >= 0)
+
+      billIndices.forEach((idx) => {
+        expect(idx).toBeGreaterThan(sorted.length * 0.6)
+      })
+    })
+
+    it('should sort emails by priority (urgent bills to spam)', async () => {
+      const emails = [
+        { subject: 'Newsletter subscription', from: 'news@site.com', date: '2024-01-14' },
+        { subject: 'Insurance payment required', from: 'insurance@company.com', date: '2024-01-15' },
+        { subject: 'Win a free iPhone!', from: 'spam@ads.com', date: '2024-01-13' },
+        { subject: 'Mortgage payment due', from: 'bank@mortgage.com', date: '2024-01-16' },
+        { subject: 'Social media notification', from: 'noreply@social.com', date: '2024-01-15' },
+      ]
+
+      const sorted = await zai.sort(emails, 'from most urgent (bills to pay) to least urgent (spam)')
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(emails.length)
+
+      // Bills should be first
+      const firstEmail = sorted[0]
+      expect(['Insurance', 'Mortgage', 'payment'].some((bill) => firstEmail.subject.includes(bill))).toBe(true)
+
+      // Spam should be last
+      const lastEmail = sorted[sorted.length - 1]
+      expect(
+        ['Win', 'free', 'iPhone', 'Newsletter', 'notification'].some((spam) => lastEmail.subject.includes(spam))
+      ).toBe(true)
+    })
+  })
+
+  describe('help desk ticket sorting', () => {
+    it('should sort tickets by urgency (based on ARR and customer tenure)', async () => {
+      const tickets = [
+        {
+          id: 'T-001',
+          customer: 'StartupCo',
+          issue: 'Login not working',
+          arr: 5000, // $5K ARR
+          customerSince: '2023-01-15', // 1 year
+          status: 'open',
+          createdAt: '2024-01-15T10:00:00Z',
+        },
+        {
+          id: 'T-002',
+          customer: 'EnterpriseCorp',
+          issue: 'System down - production outage',
+          arr: 500000, // $500K ARR
+          customerSince: '2020-03-10', // 4 years
+          status: 'open',
+          createdAt: '2024-01-15T09:00:00Z',
+        },
+        {
+          id: 'T-003',
+          customer: 'MediumBiz',
+          issue: 'Feature request for dashboard',
+          arr: 50000, // $50K ARR
+          customerSince: '2022-06-20', // 1.5 years
+          status: 'open',
+          createdAt: '2024-01-15T11:00:00Z',
+        },
+        {
+          id: 'T-004',
+          customer: 'NewCustomer',
+          issue: 'Cannot access account',
+          arr: 2000, // $2K ARR
+          customerSince: '2024-01-01', // Brand new
+          status: 'open',
+          createdAt: '2024-01-15T08:00:00Z',
+        },
+        {
+          id: 'T-005',
+          customer: 'LoyalClient',
+          issue: 'Critical bug affecting workflow',
+          arr: 120000, // $120K ARR
+          customerSince: '2019-01-01', // 5 years
+          status: 'open',
+          createdAt: '2024-01-15T10:30:00Z',
+        },
+      ]
+
+      const sorted = await zai.sort(
+        tickets,
+        'from least urgent to most urgent (consider both ARR and how long they have been customers - high ARR + long tenure = most urgent)'
+      )
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(tickets.length)
+
+      // EnterpriseCorp should be near the end (highest ARR + long tenure + production outage)
+      const enterpriseIndex = sorted.findIndex((t) => t.id === 'T-002')
+      expect(enterpriseIndex).toBeGreaterThan(sorted.length / 2)
+
+      // NewCustomer should be near the beginning (lowest priority)
+      const newCustomerIndex = sorted.findIndex((t) => t.id === 'T-004')
+      expect(newCustomerIndex).toBeLessThan(sorted.length / 2)
+
+      // LoyalClient should also be high priority (good ARR + 5 years + critical bug)
+      const loyalIndex = sorted.findIndex((t) => t.id === 'T-005')
+      expect(loyalIndex).toBeGreaterThan(sorted.length / 2)
+    })
+
+    it('should sort tickets by creation date (oldest to newest)', async () => {
+      const tickets = [
+        { id: 'T-101', title: 'Bug in checkout', createdAt: '2024-01-15T14:00:00Z', status: 'open' },
+        { id: 'T-102', title: 'Feature request', createdAt: '2024-01-10T09:00:00Z', status: 'open' },
+        { id: 'T-103', title: 'Login issue', createdAt: '2024-01-18T11:00:00Z', status: 'open' },
+        { id: 'T-104', title: 'Payment failed', createdAt: '2024-01-12T16:30:00Z', status: 'open' },
+        { id: 'T-105', title: 'UI glitch', createdAt: '2024-01-08T08:00:00Z', status: 'open' },
+      ]
+
+      const sorted = await zai.sort(tickets, 'from oldest to newest by creation date')
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(tickets.length)
+
+      // T-105 (Jan 8) should be first
+      expect(sorted[0].id).toBe('T-105')
+
+      // T-103 (Jan 18) should be last
+      expect(sorted[sorted.length - 1].id).toBe('T-103')
+
+      // Verify chronological order
+      for (let i = 1; i < sorted.length; i++) {
+        const prev = new Date(sorted[i - 1].createdAt).getTime()
+        const curr = new Date(sorted[i].createdAt).getTime()
+        expect(curr).toBeGreaterThanOrEqual(prev)
+      }
+    })
+
+    it('should sort tickets by creation date (newest to oldest)', async () => {
+      const tickets = [
+        { id: 'T-201', title: 'Issue A', createdAt: '2024-01-10T10:00:00Z', status: 'open' },
+        { id: 'T-202', title: 'Issue B', createdAt: '2024-01-15T12:00:00Z', status: 'open' },
+        { id: 'T-203', title: 'Issue C', createdAt: '2024-01-08T09:00:00Z', status: 'open' },
+        { id: 'T-204', title: 'Issue D', createdAt: '2024-01-20T14:00:00Z', status: 'open' },
+      ]
+
+      const sorted = await zai.sort(tickets, 'from newest to oldest by creation date')
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(tickets.length)
+
+      // T-204 (Jan 20) should be first
+      expect(sorted[0].id).toBe('T-204')
+
+      // T-203 (Jan 8) should be last
+      expect(sorted[sorted.length - 1].id).toBe('T-203')
+
+      // Verify reverse chronological order
+      for (let i = 1; i < sorted.length; i++) {
+        const prev = new Date(sorted[i - 1].createdAt).getTime()
+        const curr = new Date(sorted[i].createdAt).getTime()
+        expect(curr).toBeLessThanOrEqual(prev)
+      }
+    })
+
+    it('should sort tickets by status and time open (priority: open+old, then open+new, then closed)', async () => {
+      const tickets = [
+        {
+          id: 'T-301',
+          title: 'Old open ticket',
+          status: 'open',
+          createdAt: '2024-01-05T10:00:00Z',
+          resolvedAt: null,
+        },
+        {
+          id: 'T-302',
+          title: 'Recently closed',
+          status: 'closed',
+          createdAt: '2024-01-10T10:00:00Z',
+          resolvedAt: '2024-01-18T15:00:00Z',
+        },
+        {
+          id: 'T-303',
+          title: 'Very old open ticket',
+          status: 'open',
+          createdAt: '2024-01-01T08:00:00Z',
+          resolvedAt: null,
+        },
+        {
+          id: 'T-304',
+          title: 'New open ticket',
+          status: 'open',
+          createdAt: '2024-01-18T14:00:00Z',
+          resolvedAt: null,
+        },
+        {
+          id: 'T-305',
+          title: 'Old closed ticket',
+          status: 'closed',
+          createdAt: '2024-01-03T09:00:00Z',
+          resolvedAt: '2024-01-15T10:00:00Z',
+        },
+        {
+          id: 'T-306',
+          title: 'Medium age open ticket',
+          status: 'open',
+          createdAt: '2024-01-12T11:00:00Z',
+          resolvedAt: null,
+        },
+      ]
+
+      const sorted = await zai.sort(
+        tickets,
+        'prioritize by status and age: highest priority = open tickets that have been open longest; lowest priority = closed tickets'
+      )
+
+      // Verify array has same length
+      expect(sorted).toHaveLength(tickets.length)
+
+      // Find indices of open vs closed tickets
+      const openIndices = sorted.map((t, i) => (t.status === 'open' ? i : -1)).filter((i) => i >= 0)
+      const closedIndices = sorted.map((t, i) => (t.status === 'closed' ? i : -1)).filter((i) => i >= 0)
+
+      // All open tickets should come before closed tickets
+      const maxOpenIndex = Math.max(...openIndices)
+      const minClosedIndex = Math.min(...closedIndices)
+      expect(maxOpenIndex).toBeLessThan(minClosedIndex)
+
+      // T-303 (oldest open, Jan 1) should be first (highest priority, like top of todo list)
+      const oldestOpenIndex = sorted.findIndex((t) => t.id === 'T-303')
+      expect(oldestOpenIndex).toBeLessThan(sorted.length / 2)
+
+      // Closed tickets should be last (lowest priority, like bottom of todo list)
+      expect(closedIndices.every((i) => i >= sorted.length / 2)).toBe(true)
+    })
+  })
+
+  describe('edge cases', () => {
+    it('should handle empty array', async () => {
+      const sorted = await zai.sort([], 'from smallest to largest')
+      expect(sorted).toEqual([])
+    })
+
+    it('should handle single element', async () => {
+      const items = [{ name: 'Only one' }]
+      const sorted = await zai.sort(items, 'alphabetically')
+      expect(sorted).toEqual(items)
+    })
+
+    it('should handle large array (500 items)', async () => {
+      // Generate 500 deterministic values (using index-based formula for consistent results)
+      const items = Array.from({ length: 500 }, (_, i) => ({
+        id: i,
+        value: (i * 17 + 23) % 1000, // Deterministic pseudo-random pattern
+      }))
+
+      const sorted = await zai.sort(items, 'from smallest value to largest value')
+
+      // Verify all items present
+      expect(sorted).toHaveLength(500)
+      expect(sorted.map((i) => i.id).sort((a, b) => a - b)).toEqual(items.map((i) => i.id).sort((a, b) => a - b))
+
+      // Verify general trend (first quartile average < last quartile average)
+      const firstQuartile = sorted.slice(0, 125)
+      const lastQuartile = sorted.slice(375)
+
+      const firstAvg = firstQuartile.reduce((sum, item) => sum + item.value, 0) / firstQuartile.length
+      const lastAvg = lastQuartile.reduce((sum, item) => sum + item.value, 0) / lastQuartile.length
+
+      expect(firstAvg).toBeLessThan(lastAvg)
+    })
+
+    it('should handle identical items', async () => {
+      const items = [
+        { name: 'Apple', color: 'red' },
+        { name: 'Apple', color: 'red' },
+        { name: 'Apple', color: 'red' },
+      ]
+
+      const sorted = await zai.sort(items, 'alphabetically by name')
+      expect(sorted).toHaveLength(3)
+    })
+  })
+
+  describe('token truncation', () => {
+    it('should handle items with very large content by truncating', async () => {
+      // Create items with massive text content (way beyond token limits)
+      const items = [
+        {
+          id: 'doc-1',
+          priority: 'low',
+          content:
+            'This is a low priority document. '.repeat(500) + // ~3500 tokens
+            'Priority: LOW. This should be sorted first.',
+        },
+        {
+          id: 'doc-2',
+          priority: 'high',
+          content:
+            'This is a high priority document. '.repeat(500) + // ~3500 tokens
+            'Priority: HIGH. This should be sorted last.',
+        },
+        {
+          id: 'doc-3',
+          priority: 'medium',
+          content:
+            'This is a medium priority document. '.repeat(500) + // ~3500 tokens
+            'Priority: MEDIUM. This should be in the middle.',
+        },
+      ]
+
+      const sorted = await zai.sort(items, 'from lowest to highest priority', {
+        tokensPerItem: 100, // Force aggressive truncation
+      })
+
+      // Verify all items present
+      expect(sorted).toHaveLength(3)
+      expect(sorted.map((i) => i.id).sort()).toEqual(['doc-1', 'doc-2', 'doc-3'])
+
+      // Verify sorting still works despite truncation
+      // Low priority should come before high priority
+      const lowIndex = sorted.findIndex((i) => i.id === 'doc-1')
+      const highIndex = sorted.findIndex((i) => i.id === 'doc-2')
+      expect(lowIndex).toBeLessThan(highIndex)
+    })
+
+    it('should sort articles by length with token truncation', async () => {
+      const items = [
+        {
+          title: 'Short Article',
+          content: 'This is a short article with minimal content.',
+          wordCount: 8,
+        },
+        {
+          title: 'Very Long Article',
+          content:
+            'This is a very long article. '.repeat(1000) + // Massive content
+            'It has tons of text and goes on forever.',
+          wordCount: 10000,
+        },
+        {
+          title: 'Medium Article',
+          content: 'This is a medium-sized article. '.repeat(50) + 'It has a moderate amount of text.',
+          wordCount: 200,
+        },
+        {
+          title: 'Tiny Article',
+          content: 'Tiny.',
+          wordCount: 1,
+        },
+      ]
+
+      const sorted = await zai.sort(items, 'from shortest to longest based on word count', {
+        tokensPerItem: 150, // Truncate to reasonable size
+      })
+
+      // Verify all items present
+      expect(sorted).toHaveLength(4)
+
+      // Tiny should be first
+      expect(sorted[0].title).toBe('Tiny Article')
+
+      // Very Long should be last
+      expect(sorted[sorted.length - 1].title).toBe('Very Long Article')
+
+      // Verify order by word count
+      for (let i = 1; i < sorted.length; i++) {
+        expect(sorted[i].wordCount).toBeGreaterThanOrEqual(sorted[i - 1].wordCount)
+      }
+    })
+
+    it('should handle mixed content sizes with consistent truncation', async () => {
+      const items = [
+        { id: 1, description: 'Small', size: 'small' },
+        { id: 2, description: 'X'.repeat(10000), size: 'huge' }, // 10K characters
+        { id: 3, description: 'Medium length description here', size: 'medium' },
+        { id: 4, description: 'Y'.repeat(5000), size: 'large' }, // 5K characters
+        { id: 5, description: 'Tiny', size: 'tiny' },
+      ]
+
+      const sorted = await zai.sort(items, 'from smallest to largest size', {
+        tokensPerItem: 50, // Very aggressive truncation
+      })
+
+      // Verify all items present
+      expect(sorted).toHaveLength(5)
+      expect(sorted.map((i) => i.id).sort()).toEqual([1, 2, 3, 4, 5])
+
+      // Verify small/tiny come before huge
+      const smallIndices = sorted.filter((i) => i.size === 'small' || i.size === 'tiny').map((i) => sorted.indexOf(i))
+      const hugeIndex = sorted.findIndex((i) => i.size === 'huge')
+
+      smallIndices.forEach((idx) => {
+        expect(idx).toBeLessThan(hugeIndex)
+      })
+    })
+
+    it('should sort code snippets by complexity despite truncation', async () => {
+      const items = [
+        {
+          name: 'Simple Function',
+          code: `function add(a, b) { return a + b; }`,
+          complexity: 1,
+        },
+        {
+          name: 'Complex Algorithm',
+          code: `
+            function complexSort(arr) {
+              ${'// This is very complex code\n'.repeat(500)}
+              return arr.sort();
+            }
+          `,
+          complexity: 10,
+        },
+        {
+          name: 'Moderate Function',
+          code: `
+            function processData(data) {
+              ${'// Some processing logic\n'.repeat(50)}
+              return data.map(x => x * 2);
+            }
+          `,
+          complexity: 5,
+        },
+      ]
+
+      const sorted = await zai.sort(items, 'from least complex to most complex', {
+        tokensPerItem: 100, // Force truncation of large code blocks
+      })
+
+      // Verify all items present
+      expect(sorted).toHaveLength(3)
+
+      // Verify order matches complexity
+      expect(sorted[0].name).toBe('Simple Function')
+      expect(sorted[sorted.length - 1].name).toBe('Complex Algorithm')
+    })
+  })
+
+  describe('detailed results', () => {
+    it('should return detailed scoring information with .result()', async () => {
+      const animals = [{ name: 'Cheetah' }, { name: 'Sloth' }, { name: 'Horse' }]
+
+      const { output, usage, elapsed } = await zai.sort(animals, 'from slowest to fastest').result()
+
+      // Verify output is sorted array
+      expect(output).toHaveLength(3)
+      expect(Array.isArray(output)).toBe(true)
+
+      // Verify usage metadata
+      expect(usage.requests.responses).toBeGreaterThan(0)
+
+      // Verify elapsed time
+      expect(elapsed).toBeGreaterThan(0)
+    })
+  })
+})
+
+describe.sequential('zai.learn.sort', () => {
+  const taskId = 'sort-counterintuitive-test'
+  const tableName = 'TestSortLearningTable'
+  const client = getClient()
+
+  afterAll(async () => {
+    // Cleanup: delete the table after all tests
+    try {
+      await client.deleteTable({ table: tableName })
+    } catch {
+      // Table might not exist
+    }
+  })
+
+  it('learns custom vocabulary that maps nonsense phrases to urgency levels', async () => {
+    const adapter = new TableAdapter({
+      client,
+      tableName,
+    })
+
+    // Custom vocabulary that LLM cannot know without examples:
+    // "blazing sun" = critical urgency (appears first)
+    // "frosting teeth" = spam (appears last)
+    // "crimson moon" = high priority
+    // "velvet rain" = medium priority
+    // "copper wind" = low priority
+
+    // Example 1: Critical, spam, high priority
+    await adapter.saveExample({
+      key: 'vocab_example1',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort emails by priority',
+      input: JSON.stringify([
+        { subject: 'Meeting reminder - blazing sun', from: 'boss@company.com' },
+        { subject: 'Get rich quick - frosting teeth', from: 'spam@fake.com' },
+        { subject: 'Project deadline - crimson moon', from: 'manager@company.com' },
+      ]),
+      output: [
+        { subject: 'Meeting reminder - blazing sun', from: 'boss@company.com' },
+        { subject: 'Project deadline - crimson moon', from: 'manager@company.com' },
+        { subject: 'Get rich quick - frosting teeth', from: 'spam@fake.com' },
+      ],
+      explanation: 'blazing sun = critical (first), crimson moon = high, frosting teeth = spam (last)',
+      metadata,
+      status: 'approved',
+    })
+
+    // Example 2: Medium, low, critical
+    await adapter.saveExample({
+      key: 'vocab_example2',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort emails by priority',
+      input: JSON.stringify([
+        { subject: 'Newsletter - velvet rain', from: 'news@blog.com' },
+        { subject: 'Invoice - copper wind', from: 'billing@service.com' },
+        { subject: 'URGENT: Server down - blazing sun', from: 'alerts@monitoring.com' },
+      ]),
+      output: [
+        { subject: 'URGENT: Server down - blazing sun', from: 'alerts@monitoring.com' },
+        { subject: 'Newsletter - velvet rain', from: 'news@blog.com' },
+        { subject: 'Invoice - copper wind', from: 'billing@service.com' },
+      ],
+      explanation: 'blazing sun = critical (first), velvet rain = medium, copper wind = low',
+      metadata,
+      status: 'approved',
+    })
+
+    // Example 3: Spam, low, high
+    await adapter.saveExample({
+      key: 'vocab_example3',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort emails by priority',
+      input: JSON.stringify([
+        { subject: 'Win a prize - frosting teeth', from: 'scam@xyz.com' },
+        { subject: 'Subscription renewal - copper wind', from: 'auto@service.com' },
+        { subject: 'Action required - crimson moon', from: 'hr@company.com' },
+      ]),
+      output: [
+        { subject: 'Action required - crimson moon', from: 'hr@company.com' },
+        { subject: 'Subscription renewal - copper wind', from: 'auto@service.com' },
+        { subject: 'Win a prize - frosting teeth', from: 'scam@xyz.com' },
+      ],
+      explanation: 'crimson moon = high (first), copper wind = low, frosting teeth = spam (last)',
+      metadata,
+      status: 'approved',
+    })
+
+    // Example 4: All levels
+    await adapter.saveExample({
+      key: 'vocab_example4',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort emails by priority',
+      input: JSON.stringify([
+        { subject: 'Report - velvet rain', from: 'team@company.com' },
+        { subject: 'System alert - blazing sun', from: 'ops@company.com' },
+        { subject: 'Promotional - frosting teeth', from: 'ads@marketing.com' },
+        { subject: 'Review needed - crimson moon', from: 'lead@company.com' },
+        { subject: 'FYI - copper wind', from: 'info@company.com' },
+      ]),
+      output: [
+        { subject: 'System alert - blazing sun', from: 'ops@company.com' },
+        { subject: 'Review needed - crimson moon', from: 'lead@company.com' },
+        { subject: 'Report - velvet rain', from: 'team@company.com' },
+        { subject: 'FYI - copper wind', from: 'info@company.com' },
+        { subject: 'Promotional - frosting teeth', from: 'ads@marketing.com' },
+      ],
+      explanation:
+        'blazing sun (critical) → crimson moon (high) → velvet rain (medium) → copper wind (low) → frosting teeth (spam)',
+      metadata,
+      status: 'approved',
+    })
+
+    // Now test with new emails using the custom vocabulary
+    const testEmails = [
+      { subject: 'Offer - frosting teeth', from: 'deals@shop.com' }, // spam - should be last
+      { subject: 'Update - velvet rain', from: 'team@work.com' }, // medium - should be middle
+      { subject: 'CRITICAL - blazing sun', from: 'security@company.com' }, // critical - should be first
+    ]
+
+    const sorted = await getZai().learn(taskId).sort(testEmails, 'sort emails by priority')
+
+    // Verify the learned vocabulary was applied
+    expect(sorted).toHaveLength(3)
+    expect(sorted[0].subject).toContain('blazing sun') // Critical first
+    expect(sorted[1].subject).toContain('velvet rain') // Medium middle
+    expect(sorted[2].subject).toContain('frosting teeth') // Spam last
+
+    // Verify exact order
+    expect(sorted[0].from).toBe('security@company.com')
+    expect(sorted[1].from).toBe('team@work.com')
+    expect(sorted[2].from).toBe('deals@shop.com')
+  })
+
+  it.skip('learns custom domain rules and sender patterns for email prioritization', async () => {
+    const adapter = new TableAdapter({
+      client,
+      tableName,
+    })
+
+    // Custom domain rules that LLM cannot know:
+    // @sequoia.vc = our investor (critical - highest priority)
+    // @a16z.com = potential investor (high priority)
+    // @google.com = competitor (medium priority)
+    // @random-startup.io = other startups (low priority)
+    // analyst@* = spam regardless of domain (lowest priority)
+
+    // Example 1: Investor emails vs competitor
+    await adapter.saveExample({
+      key: 'domain_example1',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort by sender importance',
+      input: JSON.stringify([
+        { subject: 'Q4 metrics review', from: 'partner@sequoia.vc' },
+        { subject: 'Partnership opportunity', from: 'bd@google.com' },
+        { subject: 'Coffee chat?', from: 'analyst@somebank.com' },
+      ]),
+      output: [
+        { subject: 'Q4 metrics review', from: 'partner@sequoia.vc' },
+        { subject: 'Partnership opportunity', from: 'bd@google.com' },
+        { subject: 'Coffee chat?', from: 'analyst@somebank.com' },
+      ],
+      explanation: 'sequoia.vc (our investor) > google.com (competitor) > analyst@ (spam)',
+      metadata,
+      status: 'approved',
+    })
+
+    // Example 2: Potential investor vs spam
+    await adapter.saveExample({
+      key: 'domain_example2',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort by sender importance',
+      input: JSON.stringify([
+        { subject: 'Market research', from: 'analyst@research-firm.com' },
+        { subject: 'Investment opportunity', from: 'marc@a16z.com' },
+        { subject: 'Quick question', from: 'founder@random-startup.io' },
+      ]),
+      output: [
+        { subject: 'Investment opportunity', from: 'marc@a16z.com' },
+        { subject: 'Quick question', from: 'founder@random-startup.io' },
+        { subject: 'Market research', from: 'analyst@research-firm.com' },
+      ],
+      explanation: 'a16z.com (potential investor) > random-startup.io (other) > analyst@ (spam)',
+      metadata,
+      status: 'approved',
+    })
+
+    // Example 3: Mixed priorities
+    await adapter.saveExample({
+      key: 'domain_example3',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort by sender importance',
+      input: JSON.stringify([
+        { subject: 'Industry report', from: 'analyst@jpmorgan.com' },
+        { subject: 'Board meeting notes', from: 'roelof@sequoia.vc' },
+        { subject: 'Cloud platform update', from: 'product@google.com' },
+        { subject: 'Seed round', from: 'investor@a16z.com' },
+        { subject: 'Demo request', from: 'ceo@small-company.io' },
+      ]),
+      output: [
+        { subject: 'Board meeting notes', from: 'roelof@sequoia.vc' },
+        { subject: 'Seed round', from: 'investor@a16z.com' },
+        { subject: 'Cloud platform update', from: 'product@google.com' },
+        { subject: 'Demo request', from: 'ceo@small-company.io' },
+        { subject: 'Industry report', from: 'analyst@jpmorgan.com' },
+      ],
+      explanation:
+        'sequoia.vc (critical) > a16z.com (high) > google.com (medium) > small-company.io (low) > analyst@ (spam)',
+      metadata,
+      status: 'approved',
+    })
+
+    // Example 4: Multiple analysts (all spam)
+    await adapter.saveExample({
+      key: 'domain_example4',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort by sender importance',
+      input: JSON.stringify([
+        { subject: 'Funding news', from: 'sarah@sequoia.vc' },
+        { subject: 'Survey request', from: 'analyst@bigbank.com' },
+        { subject: 'Data report', from: 'analyst@consulting.com' },
+      ]),
+      output: [
+        { subject: 'Funding news', from: 'sarah@sequoia.vc' },
+        { subject: 'Survey request', from: 'analyst@bigbank.com' },
+        { subject: 'Data report', from: 'analyst@consulting.com' },
+      ],
+      explanation: 'sequoia.vc (critical) always first, all analyst@ emails are spam (last)',
+      metadata,
+      status: 'approved',
+    })
+
+    // Example 5: Competitor vs potential investor
+    await adapter.saveExample({
+      key: 'domain_example5',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort by sender importance',
+      input: JSON.stringify([
+        { subject: 'API integration', from: 'eng@google.com' },
+        { subject: 'Series A discussion', from: 'partner@a16z.com' },
+      ]),
+      output: [
+        { subject: 'Series A discussion', from: 'partner@a16z.com' },
+        { subject: 'API integration', from: 'eng@google.com' },
+      ],
+      explanation: 'RULE: a16z.com (potential investor) always before google.com (competitor)',
+      metadata,
+      status: 'approved',
+    })
+
+    // Example 6: Our investor always first
+    await adapter.saveExample({
+      key: 'domain_example6',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort by sender importance',
+      input: JSON.stringify([
+        { subject: 'Update', from: 'john@google.com' },
+        { subject: 'Check-in', from: 'jane@sequoia.vc' },
+        { subject: 'Meeting', from: 'alice@a16z.com' },
+      ]),
+      output: [
+        { subject: 'Check-in', from: 'jane@sequoia.vc' },
+        { subject: 'Meeting', from: 'alice@a16z.com' },
+        { subject: 'Update', from: 'john@google.com' },
+      ],
+      explanation: 'RULE: sequoia.vc (our investor) ALWAYS first, then a16z.com, then google.com',
+      metadata,
+      status: 'approved',
+    })
+
+    // Example 7: analyst@ pattern
+    await adapter.saveExample({
+      key: 'domain_example7',
+      taskId: `zai/${taskId}`,
+      taskType: 'zai.sort',
+      instructions: 'sort by sender importance',
+      input: JSON.stringify([
+        { subject: 'Stats', from: 'analyst@anywhere.com' },
+        { subject: 'Data', from: 'analyst@bigcorp.com' },
+        { subject: 'Note', from: 'person@sequoia.vc' },
+      ]),
+      output: [
+        { subject: 'Note', from: 'person@sequoia.vc' },
+        { subject: 'Stats', from: 'analyst@anywhere.com' },
+        { subject: 'Data', from: 'analyst@bigcorp.com' },
+      ],
+      explanation: 'RULE: analyst@ prefix ALWAYS means spam (last), regardless of domain',
+      metadata,
+      status: 'approved',
+    })
+
+    // Now test with new emails - should apply learned domain rules
+    const testEmails = [
+      { subject: 'Growth metrics', from: 'analyst@goldmansachs.com' }, // spam - should be last
+      { subject: 'Partnership', from: 'team@google.com' }, // competitor - medium
+      { subject: 'Portfolio update', from: 'team@sequoia.vc' }, // our investor - should be first
+      { subject: 'Funding round', from: 'ben@a16z.com' }, // potential investor - high
+      { subject: 'Collab opportunity', from: 'cto@new-startup.io' }, // other startup - low
+    ]
+
+    const sorted = await getZai()
+      .with({
+        activeLearning: {
+          enable: true,
+          tableName,
+          taskId,
+        },
+      })
+      .sort(testEmails, 'sort by sender importance')
+
+    // Verify the learned domain rules were applied
+    expect(sorted).toHaveLength(5)
+
+    // Find positions of each email type
+    const sequoiaIndex = sorted.findIndex((e) => e.from.includes('sequoia.vc'))
+    const a16zIndex = sorted.findIndex((e) => e.from.includes('a16z.com'))
+    const googleIndex = sorted.findIndex((e) => e.from.includes('google.com'))
+    const startupIndex = sorted.findIndex((e) => e.from.includes('new-startup.io'))
+    const analystIndex = sorted.findIndex((e) => e.from.includes('analyst@'))
+
+    // MOST IMPORTANT LEARNED PATTERN: analyst@ should ALWAYS be last
+    // This is the clearest pattern and should be learned consistently
+    console.log(sorted)
+    expect(analystIndex).toBe(4)
+
+    // All non-spam should come before spam (analyst@)
+    expect(sequoiaIndex).toBeLessThan(analystIndex)
+    expect(a16zIndex).toBeLessThan(analystIndex)
+    expect(googleIndex).toBeLessThan(analystIndex)
+    expect(startupIndex).toBeLessThan(analystIndex)
+
+    // Known companies/investors should come before random startups
+    // (sequoia, a16z, google are all "known" vs random-startup.io)
+    expect(sequoiaIndex).toBeLessThan(startupIndex)
+    expect(a16zIndex).toBeLessThan(startupIndex)
+    expect(googleIndex).toBeLessThan(startupIndex)
+
+    // The relative order of sequoia vs a16z vs google may vary,
+    // but at least one investor should be in top 3
+    const investorsInTop3 = [sequoiaIndex, a16zIndex].filter((idx) => idx <= 2)
+    expect(investorsInTop3.length).toBeGreaterThanOrEqual(1)
+  })
+})

packages/zai/e2e/summarize.test.ts

@@ -0,0 +1,30 @@
+import { check } from '@botpress/vai'
+import { describe, it, expect } from 'vitest'
+
+import { BotpressDocumentation, getZai } from './utils'
+
+describe('zai.summarize', () => {
+  const zai = getZai()
+
+  it('summarize long document to a concise 2000 token summary', async () => {
+    let progress: number = 0
+
+    const result = await zai
+      .summarize(BotpressDocumentation, {
+        length: 2000,
+        prompt: `Extract the Table of Contents for the Botpress Documentation. Start with a short sentence explaining what Botpress is. Pay special attention to all the different features. Focus on horizontal coverage of features rather than going in depth into one feature. The goal is to have a complete overview of what the documentation covers.`,
+      })
+      .on('progress', () => (progress = progress + 1))
+
+    expect(progress).toBeGreaterThanOrEqual(10)
+    check(result, 'The text is a summary of the Botpress documentation').toBe(true)
+    check(result, 'The text explains shortly what botpress is').toBe(true)
+    check(result, 'The text uses markdown format').toBe(true)
+    check(result, 'The text has some information about integrations').toBe(true)
+    check(result, 'The text has a section about Flows (or Workflows)').toBe(true)
+    check(result, 'The text has a section about the Botpress API').toBe(true)
+    check(result, 'The text mentions the notion of workspaces').toBe(true)
+    check(result, 'The text has some information about the Webchat').toBe(true)
+    check(result, 'The text has some information about HITL (human in the loop)').toBe(true)
+  })
+})

packages/zai/e2e/text.test.ts

@@ -0,0 +1,60 @@
+import { describe, it, expect } from 'vitest'
+import { check } from '@botpress/vai'
+
+import { getZai, tokenizer } from './utils'
+
+describe('zai.text', { timeout: 60_000 }, () => {
+  const zai = getZai()
+
+  it('generate a horror novel with no params', async () => {
+    const story = await zai.text('write a short horror novel')
+    check(story, 'is a short horror story').toBe(true)
+  })
+
+  it('No fluffy text at the beginning', async () => {
+    const story = await zai.text('write a short horror novel')
+
+    check(story, 'There is no LLM fluff at the beginning', {
+      examples: [
+        {
+          value: 'Title: A horror story\nChapter 1: The woods\nOnce upen a time, ...',
+          expected: true,
+          reason: 'It begins straight with a story, no fluff at the beginning',
+        },
+        { value: 'Once upon a time, a ...', expected: true, reason: 'The story starts directly' },
+        {
+          value: 'Sure, I will generate a story.\nOnce upen a time, a...',
+          expected: false,
+          reason: 'There is some fluff at the beginning',
+        },
+      ],
+    }).toBe(true)
+  })
+
+  it('No fluffy text at the end', async () => {
+    const story = await zai.text('write a short horror novel')
+
+    check(story, 'There is no LLM fluff at the end', {
+      examples: [
+        {
+          value: 'Title: A horror story\nChapter 1: The woods\nOnce upen a time, ... The End.',
+          expected: true,
+          reason: 'The end is clear and direct, no fluff at the end',
+        },
+        {
+          value:
+            'Sure, I will generate a story.\nOnce upen a time, a... The End.\nLet me know if you want more or if you are happy with this.',
+          expected: false,
+          reason: 'There is some fluff from the assistant at the end.',
+        },
+      ],
+    }).toBe(true)
+  })
+
+  it('length/max tokens param', async () => {
+    const story = await zai.text('write a short but complete horror story (with conclusion)', { length: 100 })
+    expect(tokenizer.count(story)).toBeLessThanOrEqual(110)
+
+    check(story, 'could be the beginning of a horror story').toBe(true)
+  })
+})

packages/zai/e2e/utils.ts

@@ -0,0 +1,39 @@
+import { Client } from '@botpress/client'
+import { Cognitive } from '@botpress/cognitive'
+import { getWasmTokenizer } from '@bpinternal/thicktoken/micro'
+import fs from 'node:fs'
+import path from 'node:path'
+import { beforeAll } from 'vitest'
+import { Zai } from '../src'
+import { type TextTokenizer } from '../src/tokenizer'
+import { getCachedCognitiveClient } from './client'
+
+const DATA_PATH = path.join(__dirname, 'data')
+const DOC_PATH = path.join(DATA_PATH, 'botpress_docs.txt')
+
+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,
+  })
+}
+
+export const getCachedClient = () => {
+  return getCachedCognitiveClient()
+}
+
+export const getZai = (cognitive?: Cognitive) => {
+  const client = cognitive || getCachedClient()
+  return new Zai({ client, modelId: 'fast' })
+}
+
+export let tokenizer: TextTokenizer = null!
+
+beforeAll(async () => {
+  tokenizer = await getWasmTokenizer()
+})
+
+export const BotpressDocumentation = fs.readFileSync(DOC_PATH, 'utf-8').trim()
+
+export const metadata = { cost: { input: 1, output: 1 }, latency: 0, model: '', tokens: { input: 1, output: 1 } }

packages/zai/e2e/zai-learn.test.ts

@@ -0,0 +1,84 @@
+import { describe, it, expect, afterAll, beforeEach, afterEach, vi } from 'vitest'
+
+import { getClient, getZai } from './utils'
+
+import { check } from '@botpress/vai'
+
+describe.sequential('zai.learn / generic', { timeout: 60_000 }, () => {
+  const client = getClient()
+  let tableName = 'ZaiTestInternalTable'
+  let taskId = 'test'
+  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('saves examples to tables', async () => {
+    const value = await zai
+      .learn(taskId)
+      .check('This text is very clearly written in English.', 'is an english sentence')
+
+    const { rows } = await client.findTableRows({ table: tableName })
+
+    expect(value).toBe(true)
+    expect(rows.length).toBe(1)
+
+    check(rows[0].explanation, 'is an explanation sentence')
+    expect(rows[0].explanation).not.toContain('Final Answer:')
+    expect(rows[0].output).toMatchObject({ value: true })
+    expect(rows[0].input).toMatchInlineSnapshot(`
+      {
+        "value": "This text is very clearly written in English.",
+      }
+    `)
+    expect(rows[0].taskId).toEqual('zai/test')
+    expect(rows[0].taskType).toBe('zai.check')
+  })
+
+  it.skip('works even if tables are down', async () => {
+    const upsertTableRows = vi.fn(async () => {
+      throw new Error('Table is down')
+    })
+
+    const findTableRows = vi.fn(async () => {
+      throw new Error('Table is down')
+    })
+
+    const client = getClient()
+    Object.assign(client, {
+      upsertTableRows,
+      findTableRows,
+    })
+
+    const value = await zai
+      .with({ client })
+      .learn(taskId)
+      .check('This text is very clearly written in English.', 'Text is in English')
+
+    const { rows } = await getClient().findTableRows({ table: tableName })
+
+    expect(value).toBe(true)
+    expect(rows.length).toBe(0)
+    expect(upsertTableRows).toHaveBeenCalledTimes(1)
+    expect(findTableRows).toHaveBeenCalledTimes(1)
+  })
+})