chore: import upstream snapshot with attribution

packages/zai/.npmignore

@@ -0,0 +1,14 @@
+vitest.config.ts
+vitest.setup.ts
+tsup.config.ts
+
+node_modules/
+scripts/
+.turbo
+
+*.test.ts
+*.test.js
+*.mts
+
+.env
+.env.local

packages/zai/README.md

@@ -0,0 +1,330 @@
+# Zai - AI Operations Made Simple
+
+Zai is a powerful LLM utility library that provides a clean, type-safe API for common AI operations. Built on Zod schemas and the Botpress API, it makes AI operations simple, intuitive, and production-ready.
+
+## ✨ Key Features
+
+- **🎯 Simple API** - One-liner operations for common AI tasks
+- **🔒 Type Safety** - Full TypeScript support with Zod schema validation
+- **🧠 Active Learning** - Learn from examples and improve over time
+- **⚡ Performance** - Built-in retries, caching, and error handling
+- **♾️ Infinite Documents** - Handle any document size with automatic chunking
+- **📊 Usage Tracking** - Monitor tokens, costs, and performance
+
+## 📦 Installation
+
+```bash
+npm install @botpress/zai @botpress/client @bpinternal/zui
+```
+
+## 🚀 Quick Start
+
+```typescript
+import { Client } from '@botpress/client'
+import { Zai } from '@botpress/zai'
+import { z } from '@bpinternal/zui'
+
+// Initialize
+const client = new Client({ botId: 'YOUR_BOT_ID', token: 'YOUR_TOKEN' })
+const zai = new Zai({ client })
+
+// Extract structured data
+const person = await zai.extract(
+  'John Doe is 30 years old and lives in New York',
+  z.object({
+    name: z.string(),
+    age: z.number(),
+    location: z.string(),
+  })
+)
+// Result: { name: 'John Doe', age: 30, location: 'New York' }
+
+// Check content
+const isPositive = await zai.check('This product is amazing!', 'expresses positive sentiment')
+// Result: true
+
+// Generate text
+const story = await zai.text('Write a short story about AI', { length: 200 })
+
+// Summarize documents
+const summary = await zai.summarize(longDocument, { length: 500 })
+```
+
+## 📚 Core Operations
+
+### 1. Extract - Get structured data from text
+
+```typescript
+// Extract single object
+const product = await zai.extract(
+  text,
+  z.object({
+    name: z.string(),
+    price: z.number(),
+    inStock: z.boolean(),
+  })
+)
+
+// Extract array of items
+const products = await zai.extract(text, z.array(productSchema))
+```
+
+### 2. Check - Verify boolean conditions
+
+```typescript
+const result = await zai.check(email, 'is spam')
+const { value, explanation } = await result.full()
+```
+
+### 3. Label - Apply multiple labels
+
+```typescript
+const labels = await zai.label(review, {
+  positive: 'expresses positive sentiment',
+  technical: 'mentions technical details',
+  verified: 'from verified purchaser',
+})
+// Result: { positive: true, technical: false, verified: true }
+```
+
+### 4. Rewrite - Transform text
+
+```typescript
+// Translate
+const french = await zai.rewrite(text, 'translate to French')
+
+// Change tone
+const formal = await zai.rewrite('Hey! What's up?', 'make it professional')
+```
+
+### 5. Filter - Filter arrays with natural language
+
+```typescript
+const techCompanies = await zai.filter(companies, 'are technology companies')
+const recentPosts = await zai.filter(posts, 'were published this week')
+```
+
+### 6. Group - Organize items into categories
+
+```typescript
+// Group items automatically
+const grouped = await zai.group(tasks, {
+  instructions: 'Group by priority level',
+})
+// Result: { 'High Priority': [...], 'Medium Priority': [...], 'Low Priority': [...] }
+
+// Group with initial categories
+const categorized = await zai.group(emails, {
+  instructions: 'Group by topic',
+  initialGroups: [
+    { id: 'work', label: 'Work' },
+    { id: 'personal', label: 'Personal' },
+  ],
+})
+
+// Group large datasets efficiently
+const organized = await zai.group(largeArray, {
+  instructions: 'Group by date',
+  chunkLength: 8000, // Process in chunks for better performance
+})
+```
+
+### 7. Rate - Score items on a 1-5 scale
+
+```typescript
+// Auto-generate criteria (returns total score)
+const scores = await zai.rate(products, 'is it a good value product?')
+// Result: [12, 8, 15] (total scores for each item)
+
+// Get detailed ratings
+const { output } = await zai.rate(products, 'is it a good value product?').result()
+// Result: [
+//   { affordability: 4, quality: 5, features: 3, total: 12 },
+//   { affordability: 3, quality: 2, features: 3, total: 8 },
+//   ...
+// ]
+
+// Use fixed criteria
+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 (all types = very_good, 3 types = good, 2 types = average, 1 type = bad)',
+  strength: 'overall password strength',
+})
+// Result: [
+//   { length: 5, complexity: 5, strength: 5, total: 15 },
+//   { length: 1, complexity: 1, strength: 1, total: 3 },
+// ]
+
+// Rate large datasets efficiently (parallelized)
+const allRatings = await zai.rate(Array(500).fill(item), 'how complete is this?')
+// Processes ~500 items in ~120ms with automatic chunking
+```
+
+### 8. Sort - Order items with natural language
+
+```typescript
+// Sort by natural criteria
+const sorted = await zai.sort(emails, 'sort by urgency')
+// LLM determines criteria and orders items accordingly
+
+// Sort with detailed results
+const { output } = await zai.sort(tasks, 'sort by priority').result()
+// output includes scoring breakdown for each item
+
+// Complex multi-criteria sorting
+const prioritized = await zai.sort(tickets, 'sort by customer importance and issue severity')
+
+// Sort large datasets efficiently (parallelized with chunking)
+const orderedItems = await zai.sort(Array(500).fill(item), 'sort by relevance')
+```
+
+### 9. Text - Generate content
+
+```typescript
+const blogPost = await zai.text('Write about the future of AI', {
+  length: 1000,
+  temperature: 0.7,
+})
+```
+
+### 10. Summarize - Create summaries
+
+```typescript
+// Simple summary
+const summary = await zai.summarize(article)
+
+// With custom prompt
+const technicalSummary = await zai.summarize(paper, {
+  length: 500,
+  prompt: 'Focus on technical implementation details',
+})
+```
+
+## 🧠 Active Learning
+
+Enable active learning to improve accuracy over time:
+
+```typescript
+const zai = new Zai({
+  client,
+  activeLearning: {
+    enable: true,
+    tableName: 'ai_learning_data',
+    taskId: 'sentiment-analysis',
+  },
+})
+
+// Use with task ID for learning
+const result = await zai.learn('sentiment-analysis').check(text, 'is positive')
+```
+
+## ⚙️ Configuration
+
+### Model Selection
+
+```typescript
+// Use the best model (default)
+const zai = new Zai({ client, model: 'best' })
+
+// Use fast model for speed
+const fastZai = new Zai({ client, model: 'fast' })
+
+// Use specific model
+const customZai = new Zai({ client, model: 'gpt-4-turbo' })
+```
+
+### Progress Tracking
+
+```typescript
+const response = zai.summarize(veryLongDocument)
+
+// Track progress
+response.on('progress', (progress) => {
+  console.log(`${progress.percent}% complete`)
+})
+
+const summary = await response
+```
+
+### Usage Monitoring
+
+```typescript
+const result = await zai.extract(text, schema)
+const usage = await result.usage()
+
+console.log({
+  tokens: usage.totalTokens,
+  cost: usage.totalCost,
+  latency: usage.totalLatency,
+})
+```
+
+## 🎯 Benefits
+
+1. **Production Ready** - Built-in error handling, retries, and rate limiting
+2. **Type Safe** - Full TypeScript support with runtime validation
+3. **Scalable** - Handle documents of any size with automatic chunking
+4. **Cost Effective** - Track usage and optimize with active learning
+5. **Developer Friendly** - Clean API with method chaining and events
+
+## 📖 Advanced Usage
+
+### Chaining Operations
+
+```typescript
+const processedData = await zai.with({ temperature: 0.3 }).learn('data-extraction').extract(document, complexSchema)
+```
+
+### Handling Large Documents
+
+```typescript
+// Automatically chunks and processes in parallel
+const extractedData = await zai.extract(
+  hugeDocument, // 100k+ tokens
+  z.array(recordSchema),
+  { chunkSize: 4000 }
+)
+```
+
+### Custom Abort Signals
+
+```typescript
+const controller = new AbortController()
+const response = zai.summarize(document, { signal: controller.signal })
+
+// Cancel if needed
+setTimeout(() => controller.abort(), 5000)
+```
+
+## 🛠️ API Reference
+
+### Zai Class
+
+- `new Zai(options)` - Create instance with client and configuration
+- `.with(config)` - Create new instance with merged configuration
+- `.learn(taskId)` - Enable active learning for specific task
+
+### Operations
+
+- `.extract(content, schema, options?)` - Extract structured data
+- `.check(content, condition, options?)` - Verify boolean condition
+- `.label(content, criteria, options?)` - Apply multiple labels
+- `.rewrite(content, instruction, options?)` - Transform text
+- `.filter(items, condition, options?)` - Filter array items
+- `.group(items, options?)` - Organize items into categories
+- `.rate(items, instructions, options?)` - Rate items on 1-5 scale
+- `.sort(items, instructions, options?)` - Order items with natural language
+- `.text(prompt, options?)` - Generate text
+- `.summarize(content, options?)` - Create summary
+
+### Response Methods
+
+- `await response` - Get simple result
+- `await response.full()` - Get detailed result with metadata
+- `await response.usage()` - Get usage statistics
+- `response.on('progress', handler)` - Track progress
+- `response.abort()` - Cancel operation
+
+## 📝 License
+
+MIT - See LICENSE file for details

packages/zai/build.ts

@@ -0,0 +1,9 @@
+import esbuild from 'esbuild'
+import glob from 'glob'
+
+const entryPoints = glob.sync('./src/**/*.ts')
+void esbuild.build({
+  entryPoints,
+  platform: 'neutral',
+  outdir: './dist',
+})

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)
+  })
+})

packages/zai/ensure-env.cjs

@@ -0,0 +1,9 @@
+;(function () {
+  if (!process.env.CLOUD_BOT_ID) {
+    throw new Error('Env: CLOUD_BOT_ID is required')
+  }
+
+  if (!process.env.CLOUD_PAT) {
+    throw new Error('Env: CLOUD_PAT is required')
+  }
+})()

packages/zai/eslint.config.mjs

@@ -0,0 +1,13 @@
+import rootConfig from '../../eslint.config.mjs'
+
+export default [
+  ...rootConfig,
+  {
+    languageOptions: {
+      parserOptions: {
+        project: ['./tsconfig.json'],
+        tsconfigRootDir: import.meta.dirname,
+      },
+    },
+  },
+]

packages/zai/package.json

@@ -0,0 +1,66 @@
+{
+  "name": "@botpress/zai",
+  "description": "Zui AI (zai) – An LLM utility library written on top of Zui and the Botpress API",
+  "version": "2.8.1",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    "types": "./dist/index.d.ts",
+    "require": "./dist/index.js",
+    "import": "./dist/index.js"
+  },
+  "repository": {
+    "url": "https://github.com/botpress/botpress"
+  },
+  "scripts": {
+    "check:type": "tsc --noEmit",
+    "build": "bp add -y && pnpm run build:types && pnpm run build:neutral && size-limit",
+    "build:neutral": "ts-node -T ./build.ts",
+    "build:types": "tsdown",
+    "watch": "tsdown --watch",
+    "test:e2e": "vitest run --config vitest.config.ts",
+    "test:e2e:update": "vitest -u run --config vitest.config.ts",
+    "test:e2e:watch": "vitest --config vitest.config.ts"
+  },
+  "size-limit": [
+    {
+      "limit": "50 kB",
+      "path": "dist/**/*.js"
+    }
+  ],
+  "keywords": [],
+  "author": "",
+  "license": "ISC",
+  "dependencies": {
+    "@botpress/cognitive": "0.6.1",
+    "json5": "^2.2.3",
+    "jsonrepair": "^3.10.0",
+    "lodash-es": "^4.17.21",
+    "p-limit": "^7.2.0"
+  },
+  "devDependencies": {
+    "@botpress/cli": "workspace:*",
+    "@botpress/client": "workspace:*",
+    "@botpress/common": "workspace:*",
+    "@botpress/vai": "workspace:*",
+    "@size-limit/file": "^11.1.6",
+    "@types/lodash-es": "^4.17.12",
+    "diff": "^8.0.1",
+    "dotenv": "^16.4.4",
+    "esbuild": "^0.25.10",
+    "glob": "^9.3.4",
+    "lodash": "^4.17.21",
+    "msw": "^2.12.0",
+    "size-limit": "^11.1.6",
+    "tsdown": "^0.22.1",
+    "unrun": "^0.3.0"
+  },
+  "peerDependencies": {
+    "@bpinternal/thicktoken": "^3.0.0",
+    "@bpinternal/zui": "^2.3.0"
+  },
+  "engines": {
+    "node": ">=18.0.0"
+  },
+  "packageManager": "pnpm@10.29.3"
+}

packages/zai/src/adapters/adapter.ts

@@ -0,0 +1,35 @@
+import { GenerationMetadata } from '../utils'
+
+export type SaveExampleProps<TInput, TOutput> = {
+  key: string
+  taskType: string
+  taskId: string
+  instructions: string
+  input: TInput
+  output: TOutput
+  explanation?: string
+  metadata: GenerationMetadata
+  status?: 'pending' | 'approved'
+}
+
+export type GetExamplesProps<TInput> = {
+  taskType: string
+  taskId: string
+  input: TInput
+}
+
+export abstract class Adapter {
+  public abstract getExamples<TInput, TOutput>(
+    props: GetExamplesProps<TInput>
+  ): Promise<
+    Array<{
+      key: string
+      input: TInput
+      output: TOutput
+      explanation?: string
+      similarity: number
+    }>
+  >
+
+  public abstract saveExample<TInput, TOutput>(props: SaveExampleProps<TInput, TOutput>): Promise<void>
+}

packages/zai/src/adapters/botpress-table.ts

@@ -0,0 +1,231 @@
+import { type Client } from '@botpress/client'
+import { transforms, z } from '@bpinternal/zui'
+
+import { GenerationMetadata } from '../utils'
+import { Adapter, GetExamplesProps, SaveExampleProps } from './adapter'
+
+const CRITICAL_TAGS = {
+  system: 'true',
+  'schema-purpose': 'active-learning',
+  'schema-version': 'Oct-2024',
+} as const
+
+const OPTIONAL_TAGS = {
+  'x-studio-title': 'Active Learning',
+  'x-studio-description': 'Table for storing active learning tasks and examples',
+  'x-studio-readonly': 'true',
+  'x-studio-icon': 'lucide://atom',
+  'x-studio-color': 'green',
+} as const
+
+const FACTOR = 30
+
+type Props = {
+  client: Client
+  tableName: string
+}
+
+const Props = z.object({
+  client: z.custom(() => true),
+  tableName: z
+    .string()
+    .regex(
+      /^[a-zA-Z0-9_]{1,45}Table$/i,
+      'Table name must be lowercase and contain only letters, numbers and underscores. It must also end with "Table". Example: "ActiveLearningTable"'
+    ),
+})
+
+export type TableSchema = {
+  taskType: string
+  taskId: string
+  key: string
+  instructions: string
+  input: Record<string, unknown>
+  output: Record<string, unknown>
+  explanation: string | null
+  metadata: GenerationMetadata
+  status: 'pending' | 'rejected' | 'approved'
+  feedback: {
+    rating: 'very-bad' | 'bad' | 'good' | 'very-good'
+    comment: string | null
+  } | null
+}
+
+const TableSchema = z.object({
+  taskType: z.string().describe('The type of the task (filter, extract, etc.)'),
+  taskId: z.string(),
+  key: z.string().describe('A unique key for the task (e.g. a hash of the input, taskId, taskType and instructions)'),
+  instructions: z.string(),
+  input: z.object({}).passthrough().describe('The input to the task'),
+  output: z.object({}).passthrough().describe('The expected output'),
+  explanation: z.string().nullable(),
+  metadata: z.object({}).passthrough(),
+  status: z.enum(['pending', 'rejected', 'approved']),
+  feedback: z
+    .object({
+      rating: z.enum(['very-bad', 'bad', 'good', 'very-good']),
+      comment: z.string().nullable(),
+    })
+    .nullable()
+    .default(null),
+})
+
+const searchableColumns = ['input'] as const satisfies Array<keyof typeof TableSchema.shape> as string[]
+
+const TableJsonSchema = Object.entries(TableSchema.shape).reduce((acc, [key, value]) => {
+  acc[key] = transforms.toJSONSchemaLegacy(value)
+  acc[key]['x-zui'] ??= {}
+  acc[key]['x-zui'].searchable = searchableColumns.includes(key)
+  return acc
+}, {})
+
+export class TableAdapter extends Adapter {
+  private _client: Client
+  private _tableName: string
+
+  private _status: 'initialized' | 'ready' | 'error'
+
+  public constructor(props: Props) {
+    super()
+    props = Props.parse(props) as Props
+    this._client = props.client
+    this._tableName = props.tableName
+    this._status = 'ready'
+  }
+
+  public async getExamples<TInput, TOutput>({ taskType, taskId, input }: GetExamplesProps<TInput>) {
+    await this._assertTableExists()
+
+    const { rows } = await this._client
+      .findTableRows({
+        table: this._tableName,
+        search: JSON.stringify({ value: input }).substring(0, 1023), // Search is limited to 1024 characters
+        limit: 10, // TODO
+        filter: {
+          // Proximity match of approved examples
+          taskType,
+          taskId,
+          status: 'approved',
+        } satisfies Partial<TableSchema>,
+      })
+      .catch((err) => {
+        // TODO: handle error
+        console.error(`Error fetching examples: ${err.message}`)
+        return { rows: [] }
+      })
+
+    return rows.map((row) => ({
+      key: row.key,
+      input: row.input.value as TInput,
+      output: row.output.value as TOutput,
+      explanation: row.explanation,
+      similarity: row.similarity ?? 0,
+    }))
+  }
+
+  public async saveExample<TInput, TOutput>({
+    key,
+    taskType,
+    taskId,
+    instructions,
+    input,
+    output,
+    explanation,
+    metadata,
+    status = 'pending',
+  }: SaveExampleProps<TInput, TOutput>) {
+    await this._assertTableExists()
+
+    await this._client
+      .upsertTableRows({
+        table: this._tableName,
+        keyColumn: 'key',
+        rows: [
+          {
+            key,
+            taskType,
+            taskId,
+            instructions,
+            input: { value: input },
+            output: { value: output },
+            explanation: explanation ?? null,
+            status,
+            metadata,
+            feedback: null, // Feedback is not provided at this point
+          } satisfies TableSchema,
+        ],
+      })
+      .catch(() => {
+        // TODO: handle error
+      })
+  }
+
+  private async _assertTableExists() {
+    if (this._status !== 'ready') {
+      return
+    }
+
+    const { table, created } = await this._client
+      .getOrCreateTable({
+        table: this._tableName,
+        factor: FACTOR,
+        frozen: true,
+        isComputeEnabled: false,
+        tags: {
+          ...CRITICAL_TAGS,
+          ...OPTIONAL_TAGS,
+        },
+        schema: TableJsonSchema,
+      })
+      .catch(() => {
+        this._status = 'error'
+        return { table: null, created: false }
+      })
+
+    if (!table) {
+      return
+    }
+
+    if (!created) {
+      const issues: string[] = []
+
+      if (table.factor !== FACTOR) {
+        issues.push(`Factor is ${table.factor} instead of ${FACTOR}`)
+      }
+
+      if (table.frozen !== true) {
+        issues.push('Table is not frozen')
+      }
+
+      for (const [key, value] of Object.entries(CRITICAL_TAGS)) {
+        if (table.tags?.[key] !== value) {
+          issues.push(`Tag ${key} is ${table.tags?.[key]} instead of ${value}`)
+        }
+      }
+
+      for (const key of Object.keys(TableJsonSchema)) {
+        const column = table.schema?.properties[key]
+        const expected = TableJsonSchema[key] as { type: string }
+
+        if (!column) {
+          issues.push(`Column ${key} is missing`)
+          continue
+        }
+
+        if (column.type !== expected.type) {
+          issues.push(`Column ${key} has type ${column.type} instead of ${expected.type}`)
+        }
+
+        if (expected['x-zui'].searchable && !column['x-zui'].searchable) {
+          issues.push(`Column ${key} is not searchable but should be`)
+        }
+      }
+
+      if (issues.length) {
+        this._status = 'error'
+      }
+    }
+
+    this._status = 'initialized'
+  }
+}

packages/zai/src/adapters/memory.ts

@@ -0,0 +1,13 @@
+import { Adapter } from './adapter'
+
+export class MemoryAdapter extends Adapter {
+  public constructor(public examples: any[]) {
+    super()
+  }
+
+  public async getExamples() {
+    return this.examples
+  }
+
+  public async saveExample() {}
+}

packages/zai/src/context.ts

@@ -0,0 +1,262 @@
+import { Cognitive, Model, GenerateContentInput, GenerateContentOutput } from '@botpress/cognitive'
+import { Adapter } from './adapters/adapter'
+import { EventEmitter } from './emitter'
+import { fastHash } from './utils'
+import type { Memoizer } from './zai'
+
+type Meta = Awaited<ReturnType<Cognitive['generateContent']>>['meta']
+
+type GenerateContentProps<T> = Omit<GenerateContentInput, 'model' | 'signal'> & {
+  maxRetries?: number
+  transform?: (text: string | undefined, output: GenerateContentOutput) => T
+}
+
+export type ZaiContextProps = {
+  client: Cognitive
+  taskType: string
+  taskId: string
+  modelId: string | string[]
+  adapter?: Adapter
+  source?: GenerateContentInput['meta']
+  memoizer?: Memoizer
+}
+
+/**
+ * Usage statistics tracking tokens, cost, and request metrics for an operation.
+ *
+ * This type is returned via Response events and the `.result()` method, providing
+ * real-time visibility into:
+ * - Token consumption (input/output/total)
+ * - Cost in USD (input/output/total)
+ * - Request statistics (count, errors, cache hits, progress percentage)
+ *
+ * @example
+ * ```typescript
+ * const { usage } = await zai.extract(text, schema).result()
+ *
+ * console.log(usage.tokens.total)    // 1250
+ * console.log(usage.cost.total)      // 0.0075 (USD)
+ * console.log(usage.requests.cached) // 0
+ * ```
+ */
+export type Usage = {
+  /** Request statistics */
+  requests: {
+    /** Total number of requests initiated */
+    requests: number
+    /** Number of requests that failed with errors */
+    errors: number
+    /** Number of successful responses received */
+    responses: number
+    /** Number of responses served from cache (no tokens used) */
+    cached: number
+    /** Operation progress as a decimal (0.0 to 1.0) */
+    percentage: number
+  }
+  /** Cost statistics in USD */
+  cost: {
+    /** Cost for input tokens */
+    input: number
+    /** Cost for output tokens */
+    output: number
+    /** Total cost (input + output) */
+    total: number
+  }
+  /** Token usage statistics */
+  tokens: {
+    /** Input tokens consumed */
+    input: number
+    /** Output tokens generated */
+    output: number
+    /** Total tokens (input + output) */
+    total: number
+  }
+}
+
+type ContextEvents = {
+  update: Usage
+}
+
+export class ZaiContext {
+  private _startedAt = Date.now()
+
+  private _inputCost = 0
+  private _outputCost = 0
+  private _inputTokens = 0
+  private _outputTokens = 0
+  private _totalCachedResponses = 0
+
+  private _totalRequests = 0
+  private _totalErrors = 0
+  private _totalResponses = 0
+
+  public taskId: string
+  public taskType: string
+  public modelId: GenerateContentInput['model']
+  public adapter?: Adapter
+  public source?: GenerateContentInput['meta']
+
+  private _eventEmitter: EventEmitter<ContextEvents>
+  private _memoizer: Memoizer
+
+  public controller: AbortController = new AbortController()
+  private _client: Cognitive
+
+  private static _noopMemoizer: Memoizer = { run: (_id, fn) => fn() }
+
+  public constructor(props: ZaiContextProps) {
+    this._client = props.client.clone()
+    this.taskId = props.taskId
+    this.modelId = props.modelId
+    this.adapter = props.adapter
+    this.source = props.source
+    this.taskType = props.taskType
+    this._memoizer = props.memoizer ?? ZaiContext._noopMemoizer
+    this._eventEmitter = new EventEmitter<ContextEvents>()
+
+    this._client.on('request', () => {
+      this._totalRequests++
+      this._eventEmitter.emit('update', this.usage)
+    })
+
+    this._client.on('response', (_req, res) => {
+      this._totalResponses++
+
+      if (res.meta.cached) {
+        this._totalCachedResponses++
+      } else {
+        this._inputTokens += res.meta.tokens.input || 0
+        this._outputTokens += res.meta.tokens.output || 0
+        this._inputCost += res.meta.cost.input || 0
+        this._outputCost += res.meta.cost.output || 0
+      }
+
+      this._eventEmitter.emit('update', this.usage)
+    })
+
+    this._client.on('error', () => {
+      this._totalErrors++
+      this._eventEmitter.emit('update', this.usage)
+    })
+  }
+
+  public async getModel(): Promise<Model> {
+    // getModelDetails resolves a single model; for a fallback array we report
+    // the primary entry's details. Fallbacks are honored at the request layer.
+    const primary = Array.isArray(this.modelId) ? this.modelId[0] : this.modelId
+    return this._client.getModelDetails(primary)
+  }
+
+  public on<K extends keyof ContextEvents>(type: K, listener: (event: ContextEvents[K]) => void) {
+    this._eventEmitter.on(type, listener)
+    return this
+  }
+
+  public clear() {
+    this._eventEmitter.clear()
+  }
+
+  public async generateContent<Out = string>(
+    props: GenerateContentProps<Out>
+  ): Promise<{ meta: Meta; output: GenerateContentOutput; text: string | undefined; extracted: Out }> {
+    const memoKey = `zai:memo:${this.taskType}:${this.taskId || 'default'}:${fastHash(
+      JSON.stringify({
+        s: props.systemPrompt,
+        m: props.messages?.map((m) => ('content' in m ? m.content : '')),
+        st: props.stopSequences,
+        mt: props.maxTokens,
+      })
+    )}`
+
+    return this._memoizer.run(memoKey, () => this._generateContentInner(props))
+  }
+
+  private async _generateContentInner<Out = string>(
+    props: GenerateContentProps<Out>
+  ): Promise<{ meta: Meta; output: GenerateContentOutput; text: string | undefined; extracted: Out }> {
+    const maxRetries = Math.max(props.maxRetries ?? 3, 0)
+    const transform = props.transform
+    let lastError: Error | null = null
+    const messages = [...(props.messages || [])]
+
+    for (let attempt = 0; attempt <= maxRetries && !this.controller.signal.aborted; attempt++) {
+      try {
+        const response = await this._client.generateContent({
+          ...props,
+          messages,
+          signal: this.controller.signal,
+          model: this.modelId,
+          meta: {
+            integrationName: props.meta?.integrationName || 'zai',
+            promptCategory: props.meta?.promptCategory || `zai:${this.taskType}`,
+            promptSource: props.meta?.promptSource || `zai:${this.taskType}:${this.taskId ?? 'default'}`,
+          },
+        })
+
+        if (this.controller.signal.aborted) {
+          throw this.controller.signal.reason
+        }
+
+        const content = response.output.choices[0]?.content
+        const str = typeof content === 'string' ? content : content?.[0]?.text || ''
+        let output: Out
+
+        messages.push({
+          role: 'assistant',
+          content: str || '<Invalid output, no content provided>',
+        })
+
+        if (!transform) {
+          output = str as Out
+        } else {
+          output = transform(str, response.output)
+        }
+
+        return { meta: response.meta, output: response.output, text: str, extracted: output }
+      } catch (error) {
+        if (this.controller.signal.aborted) {
+          throw this.controller.signal.reason
+        }
+
+        lastError = error as Error
+
+        if (attempt === maxRetries) {
+          throw lastError
+        }
+
+        messages.push({
+          role: 'user',
+          content: `ERROR PARSING OUTPUT\n\n${lastError.message}.\n\nPlease return a valid response addressing the error above.`,
+        })
+      }
+    }
+
+    throw lastError
+  }
+
+  public get elapsedTime(): number {
+    return Date.now() - this._startedAt
+  }
+
+  public get usage(): Usage {
+    return {
+      requests: {
+        errors: this._totalErrors,
+        requests: this._totalRequests,
+        responses: this._totalResponses,
+        cached: this._totalCachedResponses,
+        percentage: this._totalRequests > 0 ? (this._totalResponses + this._totalErrors) / this._totalRequests : 0,
+      },
+      tokens: {
+        input: this._inputTokens,
+        output: this._outputTokens,
+        total: this._inputTokens + this._outputTokens,
+      },
+      cost: {
+        input: this._inputCost,
+        output: this._outputCost,
+        total: this._inputCost + this._outputCost,
+      },
+    }
+  }
+}

packages/zai/src/emitter.ts

@@ -0,0 +1,49 @@
+export class EventEmitter<E extends object> {
+  private _listeners: {
+    [K in keyof E]?: ((event: E[K]) => void)[]
+  } = {}
+
+  public emit<K extends keyof E>(type: K, event: E[K]) {
+    const listeners = this._listeners[type]
+    if (!listeners) {
+      return
+    }
+    for (const listener of listeners) {
+      listener(event)
+    }
+  }
+
+  public once<K extends keyof E>(type: K, listener: (event: E[K]) => void) {
+    const wrapped = (event: E[K]) => {
+      this.off(type, wrapped)
+      listener(event)
+    }
+    this.on(type, wrapped)
+  }
+
+  public on<K extends keyof E>(type: K, listener: (event: E[K]) => void) {
+    if (!this._listeners[type]) {
+      this._listeners[type] = []
+    }
+    this._listeners[type]!.push(listener)
+  }
+
+  public off<K extends keyof E>(type: K, listener: (event: E[K]) => void) {
+    const listeners = this._listeners[type]
+    if (!listeners) {
+      return
+    }
+    const index = listeners.indexOf(listener)
+    if (index !== -1) {
+      listeners.splice(index, 1)
+    }
+  }
+
+  public clear<K extends keyof E>(type?: K) {
+    if (type) {
+      delete this._listeners[type]
+    } else {
+      this._listeners = {}
+    }
+  }
+}

packages/zai/src/index.ts

@@ -0,0 +1,17 @@
+import { Zai, type ZaiConfig, type Memoizer } from './zai'
+
+import './operations/text'
+import './operations/rewrite'
+import './operations/summarize'
+import './operations/check'
+import './operations/filter'
+import './operations/extract'
+import './operations/label'
+import './operations/group'
+import './operations/rate'
+import './operations/sort'
+import './operations/answer'
+import './operations/patch'
+
+export { Zai }
+export type { Memoizer, ZaiConfig }

packages/zai/src/json-utils.ts

@@ -0,0 +1,105 @@
+import { jsonrepair } from 'jsonrepair'
+
+/**
+ * Check if a file is a JSON file based on its path or name extension.
+ */
+export function isJsonFile(path: string, name: string): boolean {
+  return path.endsWith('.json') || name.endsWith('.json')
+}
+
+export type JsonValidResult = { valid: true; data: unknown; repaired: boolean; content: string; error: never }
+export type JsonInvalidResult = { valid: false; error: string }
+export type JsonParseResult = JsonValidResult | JsonInvalidResult
+
+/**
+ * Try to parse a string as JSON.
+ */
+function tryParseJson(content: string): { valid: true; data: unknown } | { valid: false; error: string } {
+  try {
+    const data = JSON.parse(content)
+    return { valid: true, data }
+  } catch (err) {
+    return { valid: false, error: err instanceof Error ? err.message : String(err) }
+  }
+}
+
+/**
+ * Attempt to repair invalid JSON using jsonrepair, then validate.
+ * Returns the repaired string if successful, or null if repair also fails.
+ */
+function repairJson(content: string): string | null {
+  try {
+    const repaired = jsonrepair(content)
+    JSON.parse(repaired)
+    return repaired
+  } catch {
+    return null
+  }
+}
+
+/**
+ * Build an LLM-friendly error message with the JSON content,
+ * line numbers, and the error location highlighted.
+ */
+function formatError(content: string, rawError: string): string {
+  const lines = content.split(/\r?\n/)
+
+  // Extract character position from V8/Bun error messages ("at position N")
+  let errorLine: number | null = null
+  let errorCol: number | null = null
+  const posMatch = rawError.match(/at position (\d+)/)
+  if (posMatch) {
+    const position = parseInt(posMatch[1], 10)
+    let offset = 0
+    for (let i = 0; i < lines.length; i++) {
+      const lineLen = lines[i]!.length + 1
+      if (offset + lineLen > position) {
+        errorLine = i
+        errorCol = position - offset
+        break
+      }
+      offset += lineLen
+    }
+  }
+
+  let out = `JSON Syntax Error: ${rawError}\n\n`
+
+  const contextRadius = 3
+  const startLine = errorLine !== null ? Math.max(0, errorLine - contextRadius) : 0
+  const endLine = errorLine !== null ? Math.min(lines.length - 1, errorLine + contextRadius) : lines.length - 1
+  const padWidth = String(endLine + 1).length
+
+  for (let i = startLine; i <= endLine; i++) {
+    const lineNum = String(i + 1).padStart(padWidth, ' ')
+    const marker = i === errorLine ? ' ← ERROR' : ''
+    out += `${lineNum}|${lines[i]}${marker}\n`
+
+    if (i === errorLine && errorCol !== null) {
+      out += `${' '.repeat(padWidth + 1 + errorCol)}^\n`
+    }
+  }
+
+  return out
+}
+
+/**
+ * Validate and optionally repair JSON content.
+ *
+ * 1. If content is valid JSON, return `{ valid: true, data, repaired: false, content }`.
+ * 2. If invalid, attempt repair with jsonrepair.
+ * 3. If repair succeeds, return `{ valid: true, data, repaired: true, content }`.
+ * 4. If repair fails, return `{ valid: false, error }` where `error` is an
+ *    LLM-friendly string with numbered lines and the syntax error highlighted.
+ */
+export function validateOrRepairJson(content: string): JsonParseResult {
+  const parsed = tryParseJson(content)
+  if (parsed.valid === false) {
+    const repaired = repairJson(content)
+    if (repaired !== null) {
+      return { valid: true, data: JSON.parse(repaired), repaired: true, content: repaired } as JsonValidResult
+    }
+    return { valid: false, error: formatError(content, parsed.error) }
+  }
+
+  return { valid: true, data: parsed.data, repaired: false, content } as JsonValidResult
+}

packages/zai/src/micropatch.ts

@@ -0,0 +1,365 @@
+/**
+ * Micropatch v0.3
+ *
+ * A tiny engine to parse and apply ultra-compact line-based patches designed for LLMs.
+ *
+ * Protocol recap:
+ * - Source lines are referenced by ORIGINAL 1-based numbers (pre-edit).
+ * - Ops start at the beginning of a line with the marker `◼︎`.
+ * - Allowed ops:
+ *    ◼︎<NNN|text        → insert single line BEFORE original NNN
+ *    ◼︎>NNN|text        → insert single line AFTER  original NNN
+ *    ◼︎=NNN|line        → replace line NNN with one or more lines (multiline payload allowed)
+ *    ◼︎=NNN-MMM|lines   → replace inclusive range NNN..MMM with one or more lines
+ *    ◼︎-NNN             → delete line NNN
+ *    ◼︎-NNN-MMM         → delete inclusive range NNN..MMM
+ * - Multiline `=` payload: starts after `|` and continues **until the next line that begins with `◼︎`** or EOF.
+ * - Only escaping rule: `\◼︎` inside payload/text becomes a literal `◼︎`. No other escaping is recognized.
+ * - Ranges are allowed **only** for `-` (delete) and `=` (replace).
+ * - Deterministic apply order on ORIGINAL addresses:
+ *    1) Deletes `-` (desc by start)
+ *    2) Single-line replaces `=NNN` (asc)
+ *    3) Range replaces `=NNN-MMM` (asc by start)
+ *    4) Inserts `<` (asc)
+ *    5) Inserts `>` (asc)
+ *
+ * Notes:
+ * - The engine maintains a live index map so ORIGINAL references remain valid despite shifting.
+ * - Idempotency-friendly: if a target original line no longer maps into current text (e.g., already deleted), the op is skipped.
+ * - Whitespace is preserved; EOL style can be chosen or auto-detected.
+ */
+
+export type EOL = 'lf' | 'crlf'
+
+/** Parsed operations (internal normalized form). */
+type Op =
+  | { k: '<'; n: number; s: string } // insert before (single-line)
+  | { k: '>'; n: number; s: string } // insert after (single-line)
+  | { k: '='; n: number; s: string[] } // replace single line with N lines
+  | { k: '=-'; a: number; b: number; s: string[] } // replace range with N lines
+  | { k: '-'; n: number } // delete single
+  | { k: '--'; a: number; b: number } // delete range
+
+/**
+ * Micropatch: parse & apply patches to text.
+ */
+export class Micropatch {
+  private _text: string
+  private _eol: EOL
+
+  /**
+   * Create a Micropatch instance.
+   * @param source The file contents.
+   * @param eol Line ending style. If omitted, it is auto-detected from `source` (CRLF if any CRLF is present; otherwise LF).
+   */
+  public constructor(source: string, eol?: EOL) {
+    this._text = source
+    this._eol = eol ?? Micropatch.detectEOL(source)
+  }
+
+  /** Get current text. */
+  public getText(): string {
+    return this._text
+  }
+
+  /**
+   * Replace current text.
+   * Useful if you want to "load" a new snapshot without reconstructing the class.
+   */
+  public setText(source: string, eol?: EOL): void {
+    this._text = source
+    this._eol = eol ?? Micropatch.detectEOL(source)
+  }
+
+  /**
+   * Apply ops text to current buffer.
+   * @param opsText One or more operations in the v0.3 syntax.
+   * @returns The updated text (also stored internally).
+   * @throws If the patch contains invalid syntax (e.g., range on insert).
+   */
+  public apply(opsText: string): string {
+    const ops = Micropatch.parseOps(opsText)
+    this._text = Micropatch._applyOps(this._text, ops, this._eol)
+    return this._text
+  }
+
+  /**
+   * Render a numbered view of the current buffer (token-cheap preview for models).
+   * Format: `NNN|<line>`, starting at 001.
+   */
+  public renderNumberedView(): string {
+    const NL = this._eol === 'lf' ? '\n' : '\r\n'
+    const lines = Micropatch._splitEOL(this._text)
+    return lines.map((l, i) => `${String(i + 1).padStart(3, '0')}|${l}`).join(NL)
+  }
+
+  // ---------------------- Static helpers ----------------------
+
+  /** Detect EOL style from content. */
+  public static detectEOL(source: string): EOL {
+    return /\r\n/.test(source) ? 'crlf' : 'lf'
+  }
+
+  /** Split text into lines, preserving empty last line if present. */
+  private static _splitEOL(text: string): string[] {
+    // Use universal split, but keep trailing empty line if final NL exists.
+    const parts = text.split(/\r?\n/)
+    // If text ends with NL, split leaves an extra empty string we want to keep.
+    return parts
+  }
+
+  /** Join lines with the chosen EOL. */
+  private static _joinEOL(lines: string[], eol: EOL): string {
+    const NL = eol === 'lf' ? '\n' : '\r\n'
+    return lines.join(NL)
+  }
+
+  /** Unescape payload text: `\◼︎` → `◼︎`. */
+  private static _unescapeMarker(s: string): string {
+    return s.replace(/\\◼︎/g, '◼︎')
+  }
+
+  /**
+   * Parse ops text (v0.3).
+   * - Ignores blank lines and lines not starting with `◼︎` (you can keep comments elsewhere).
+   * - Validates ranges for allowed ops.
+   */
+  public static parseOps(opsText: string): Op[] {
+    const lines = opsText.split(/\r?\n/)
+    const ops: Op[] = []
+
+    // Regex for op header line:
+    // ◼︎([<>=-])(\d+)(?:-(\d+))?(?:\|(.*))?$
+    const headerRe = /^◼︎([<>=-])(\d+)(?:-(\d+))?(?:\|(.*))?$/
+
+    let i = 0
+    while (i < lines.length) {
+      const line = lines[i]
+      if (!line) {
+        i++
+        continue
+      }
+
+      if (!line.startsWith('◼︎')) {
+        i++
+        continue // ignore non-op lines
+      }
+
+      const m = headerRe.exec(line)
+      if (!m || !m[1] || !m[2]) {
+        throw new Error(`Invalid op syntax at line ${i + 1}: ${line}`)
+      }
+
+      const op = m[1] as '<' | '>' | '=' | '-'
+      const aNum = parseInt(m[2], 10)
+      const bNum = m[3] ? parseInt(m[3], 10) : undefined
+      // LLMs sometimes "double" the leading `|` of a payload that itself begins with `|`
+      // (e.g. markdown table rows), as if escaping it. The protocol has no such escape,
+      // so a payload starting with `||` is treated as a single literal leading `|`.
+      const firstPayload = (m[4] ?? '').replace(/^\|\|/, '|')
+
+      if (aNum < 1 || (bNum !== undefined && bNum < aNum)) {
+        throw new Error(`Invalid line/range at line ${i + 1}: ${line}`)
+      }
+
+      // Inserts: single-line payload only; ranges are illegal.
+      if (op === '<' || op === '>') {
+        if (bNum !== undefined) {
+          throw new Error(`Insert cannot target a range (line ${i + 1})`)
+        }
+        const text = Micropatch._unescapeMarker(firstPayload)
+        ops.push({ k: op, n: aNum, s: text })
+        i++
+        continue
+      }
+
+      // Deletes: may be single or range; no payload allowed.
+      if (op === '-') {
+        if (firstPayload !== '') {
+          throw new Error(`Delete must not have a payload (line ${i + 1})`)
+        }
+        if (bNum === undefined) {
+          ops.push({ k: '-', n: aNum })
+        } else {
+          ops.push({ k: '--', a: aNum, b: bNum })
+        }
+        i++
+        continue
+      }
+
+      // Replaces: '=' supports single or range, and MULTILINE payload.
+      if (op === '=') {
+        // Collect multiline payload until next line starting with '◼︎' or EOF.
+        const payload: string[] = [Micropatch._unescapeMarker(firstPayload)]
+        let j = i + 1
+        while (j < lines.length) {
+          const nextLine = lines[j]
+          if (nextLine.startsWith('◼︎')) break
+          payload.push(Micropatch._unescapeMarker(nextLine))
+          j++
+        }
+        // Normalize potential trailing empty due to splitting; keep as-is (exact payload).
+        if (bNum === undefined) {
+          ops.push({ k: '=', n: aNum, s: payload })
+        } else {
+          ops.push({ k: '=-', a: aNum, b: bNum, s: payload })
+        }
+        i = j
+        continue
+      }
+
+      // Should be unreachable.
+      i++
+    }
+
+    return Micropatch._canonicalizeOrder(ops)
+  }
+
+  /** Order ops deterministically according to the spec. */
+  private static _canonicalizeOrder(ops: Op[]): Op[] {
+    const delS: Array<Extract<Op, { k: '-' }>> = []
+    const delR: Array<Extract<Op, { k: '--' }>> = []
+    const eqS: Array<Extract<Op, { k: '=' }>> = []
+    const eqR: Array<Extract<Op, { k: '=-' }>> = []
+    const insB: Array<Extract<Op, { k: '<' }>> = []
+    const insA: Array<Extract<Op, { k: '>' }>> = []
+
+    for (const o of ops) {
+      switch (o.k) {
+        case '-':
+          delS.push(o)
+          break
+        case '--':
+          delR.push(o)
+          break
+        case '=':
+          eqS.push(o)
+          break
+        case '=-':
+          eqR.push(o)
+          break
+        case '<':
+          insB.push(o)
+          break
+        case '>':
+          insA.push(o)
+          break
+        default:
+          break
+      }
+    }
+
+    delS.sort((a, b) => b.n - a.n)
+    delR.sort((a, b) => b.a - a.a)
+    eqS.sort((a, b) => a.n - b.n)
+    eqR.sort((a, b) => a.a - b.a)
+    insB.sort((a, b) => a.n - b.n)
+    insA.sort((a, b) => a.n - b.n)
+
+    return ([] as Op[]).concat(delS, delR, eqS, eqR, insB, insA)
+  }
+
+  /**
+   * Apply normalized ops to given source.
+   * - Uses a live index map from ORIGINAL 1-based addresses → current positions.
+   * - Skips ops whose targets can no longer be mapped (idempotency-friendly).
+   */
+  private static _applyOps(source: string, ops: Op[], eol: EOL): string {
+    const lines = Micropatch._splitEOL(source)
+
+    // idx[i] = current position (0-based) of original line (i+1).
+    const idx: number[] = Array.from({ length: lines.length }, (_, i) => i)
+
+    const map = (n: number): number => idx[n - 1] ?? -1
+    const bump = (from: number, delta: number) => {
+      // Shift all tracked indices at or after `from` by delta.
+      for (let i = 0; i < idx.length; i++) {
+        const current = idx[i]
+        if (current !== undefined && current >= from) {
+          idx[i] = current + delta
+        }
+      }
+    }
+
+    for (const o of ops) {
+      switch (o.k) {
+        case '-': {
+          const i = map(o.n)
+          if (i >= 0 && i < lines.length) {
+            lines.splice(i, 1)
+            bump(i, -1)
+          }
+          break
+        }
+        case '--': {
+          const a = map(o.a)
+          const b = map(o.b)
+          if (a >= 0 && b >= a && b < lines.length) {
+            lines.splice(a, b - a + 1)
+            bump(a, -(b - a + 1))
+          }
+          break
+        }
+        case '=': {
+          const i = map(o.n)
+          if (i >= 0 && i < lines.length) {
+            const rep = o.s
+            lines.splice(i, 1, ...rep)
+            bump(i + 1, rep.length - 1)
+          }
+          break
+        }
+        case '=-': {
+          const a = map(o.a)
+          const b = map(o.b)
+          if (a >= 0 && b >= a && b < lines.length) {
+            const rep = o.s
+            lines.splice(a, b - a + 1, ...rep)
+            bump(a + 1, rep.length - (b - a + 1))
+          }
+          break
+        }
+        case '<': {
+          if (o.n > idx.length) break // line never existed in the original → stale ref, skip
+          const i = Math.max(0, Math.min(map(o.n), lines.length))
+          lines.splice(i, 0, o.s)
+          bump(i, +1)
+          break
+        }
+        case '>': {
+          if (o.n > idx.length) break
+          const i = Math.max(0, Math.min(map(o.n) + 1, lines.length))
+          lines.splice(i, 0, o.s)
+          bump(i, +1)
+          break
+        }
+        default:
+          break
+      }
+    }
+
+    return Micropatch._joinEOL(lines, eol)
+  }
+
+  // ---------------------- Convenience APIs ----------------------
+
+  /**
+   * Convenience: one-shot apply.
+   * @param source Text to patch.
+   * @param opsText Operations text.
+   * @param eol EOL style (auto-detected if omitted).
+   */
+  public static applyText(source: string, opsText: string, eol?: EOL): string {
+    const inst = new Micropatch(source, eol)
+    return inst.apply(opsText)
+  }
+
+  /**
+   * Convenience: parse only.
+   * Useful for validation without applying.
+   */
+  public static validate(opsText: string): { ok: true; count: number } {
+    const ops = Micropatch.parseOps(opsText)
+    return { ok: true, count: ops.length }
+  }
+}

packages/zai/src/operations/answer.ts

@@ -0,0 +1,827 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+import pLimit from 'p-limit'
+
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { stringify } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER } from './constants'
+
+/**
+ * Citation referencing a specific line or range of lines in support documents
+ */
+export type Citation<T> = {
+  /** The character offset where this citation appears in the answer */
+  offset: number
+  /** The support document item(s) used at this offset */
+  item: T
+  /** The line contents joined together as a snippet */
+  snippet: string
+}
+
+/**
+ * A single answer with citations
+ */
+export type AnswerWithCitations<T> = {
+  /** The answer text */
+  answer: string
+  /** Citations mapping answer text to support documents */
+  citations: Citation<T>[]
+}
+
+/**
+ * Response type when a clear answer can be provided
+ */
+export type AnswerResponse<T> = {
+  type: 'answer'
+} & AnswerWithCitations<T>
+
+/**
+ * Response type when the question is ambiguous and multiple interpretations exist
+ */
+export type AmbiguousResponse<T> = {
+  type: 'ambiguous'
+  /** What is the ambiguity? What concepts clash or are unclear? */
+  ambiguity: string
+  /** A follow-up question to clear out the ambiguity */
+  follow_up: string
+  /** Possible answers for different interpretations (2-3 answers) */
+  answers: AnswerWithCitations<T>[]
+}
+
+/**
+ * Response type when the question is out of topic
+ */
+export type OutOfTopicResponse = {
+  type: 'out_of_topic'
+  /** Why is this question considered out of topic? */
+  reason: string
+}
+
+/**
+ * Response type when the question is invalid or malformed
+ */
+export type InvalidQuestionResponse = {
+  type: 'invalid_question'
+  /** What makes this an invalid question? */
+  reason: string
+}
+
+/**
+ * Response type when there is insufficient knowledge to answer
+ */
+export type MissingKnowledgeResponse = {
+  type: 'missing_knowledge'
+  /** What knowledge is missing to generate a high-quality answer? */
+  reason: string
+}
+
+/**
+ * All possible response types from zai.answer
+ */
+export type AnswerResult<T> =
+  | AnswerResponse<T>
+  | AmbiguousResponse<T>
+  | OutOfTopicResponse
+  | InvalidQuestionResponse
+  | MissingKnowledgeResponse
+
+/**
+ * Example for few-shot learning
+ */
+export type AnswerExample<T> = {
+  /** Support documents for this example */
+  documents: T[]
+  /** The question asked */
+  question: string
+  /** The expected answer result */
+  result: AnswerResult<T>
+}
+
+export type Options<T> = {
+  /** Examples to help guide answer generation */
+  examples?: AnswerExample<T>[]
+  /** Additional instructions for answer generation */
+  instructions?: string
+  /** Maximum number of tokens per document chunk */
+  chunkLength?: number
+  /**
+   * Maximum number of refinement iterations when merging chunked results
+   * @default 3
+   */
+  maxRefinementPasses?: number
+}
+
+const _Options = z.object({
+  examples: z.array(z.any()).default([]).describe('Examples to help guide answer generation'),
+  instructions: z.string().optional().describe('Additional instructions for answer generation'),
+  chunkLength: z
+    .number()
+    .min(250)
+    .max(100_000)
+    .optional()
+    .describe('Maximum number of tokens per document chunk')
+    .default(16_000),
+  maxRefinementPasses: z
+    .number()
+    .min(1)
+    .max(10)
+    .optional()
+    .describe('Maximum number of refinement iterations when merging chunked results')
+    .default(3),
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Answers questions from documents with citations and intelligent handling of edge cases.
+     *
+     * This operation provides a production-ready question-answering system that:
+     * - Cites sources with precise line references
+     * - Handles ambiguous questions with multiple interpretations
+     * - Detects out-of-topic or invalid questions
+     * - Identifies missing knowledge
+     * - Automatically chunks and processes large document sets
+     *
+     * @param documents - Array of documents to search (strings, objects, or any type)
+     * @param question - The question to answer
+     * @param options - Configuration for chunking, examples, and instructions
+     * @returns Response with answer + citations, or error states (ambiguous, out_of_topic, invalid, missing_knowledge)
+     *
+     * @example Basic usage with string documents
+     * ```typescript
+     * const documents = [
+     *   'Botpress was founded in 2016.',
+     *   'The company is based in Quebec, Canada.',
+     *   'Botpress provides an AI agent platform.'
+     * ]
+     *
+     * const result = await zai.answer(documents, 'When was Botpress founded?')
+     * if (result.type === 'answer') {
+     *   console.log(result.answer) // "Botpress was founded in 2016."
+     *   console.log(result.citations) // [{ offset: 30, item: documents[0], snippet: '...' }]
+     * }
+     * ```
+     *
+     * @example With object documents
+     * ```typescript
+     * const products = [
+     *   { id: 1, name: 'Pro Plan', price: 99, features: ['AI', 'Analytics'] },
+     *   { id: 2, name: 'Enterprise', price: 499, features: ['AI', 'Support', 'SLA'] }
+     * ]
+     *
+     * const result = await zai.answer(products, 'What features does the Pro Plan include?')
+     * // Returns answer with citations pointing to the product objects
+     * ```
+     *
+     * @example Handling different response types
+     * ```typescript
+     * const result = await zai.answer(documents, question)
+     *
+     * switch (result.type) {
+     *   case 'answer':
+     *     console.log('Answer:', result.answer)
+     *     console.log('Sources:', result.citations)
+     *     break
+     *
+     *   case 'ambiguous':
+     *     console.log('Question is ambiguous:', result.ambiguity)
+     *     console.log('Clarifying question:', result.follow_up)
+     *     console.log('Possible answers:', result.answers)
+     *     break
+     *
+     *   case 'out_of_topic':
+     *     console.log('Question unrelated:', result.reason)
+     *     break
+     *
+     *   case 'invalid_question':
+     *     console.log('Invalid question:', result.reason)
+     *     break
+     *
+     *   case 'missing_knowledge':
+     *     console.log('Insufficient info:', result.reason)
+     *     break
+     * }
+     * ```
+     *
+     * @example With custom instructions
+     * ```typescript
+     * const result = await zai.answer(documents, 'What is the pricing?', {
+     *   instructions: 'Provide detailed pricing breakdown including all tiers',
+     *   chunkLength: 8000 // Process in smaller chunks for accuracy
+     * })
+     * ```
+     *
+     * @example Large document sets (auto-chunking)
+     * ```typescript
+     * // Handles thousands of documents automatically
+     * const manyDocs = await loadDocuments() // 1000+ documents
+     * const result = await zai.answer(manyDocs, 'What is the refund policy?')
+     * // Automatically chunks, processes in parallel, and merges results
+     * ```
+     *
+     * @example Tracking citations
+     * ```typescript
+     * const result = await zai.answer(documents, question)
+     * if (result.type === 'answer') {
+     *   result.citations.forEach(citation => {
+     *     console.log(`At position ${citation.offset}:`)
+     *     console.log(`  Cited: "${citation.snippet}"`)
+     *     console.log(`  From document:`, citation.item)
+     *   })
+     * }
+     * ```
+     */
+    answer<T>(documents: T[], question: string, options?: Options<T>): Response<AnswerResult<T>, AnswerResult<T>>
+  }
+}
+
+// Markers for parsing LLM output
+const ANSWER_START = '■answer'
+const AMBIGUOUS_START = '■ambiguous'
+const OUT_OF_TOPIC_START = '■out_of_topic'
+const INVALID_QUESTION_START = '■invalid_question'
+const MISSING_KNOWLEDGE_START = '■missing_knowledge'
+const END = '■end■'
+
+/**
+ * Maps line numbers to document index and line within that document
+ */
+type LineMapping<T> = {
+  lineNumber: number
+  documentIndex: number
+  lineInDocument: number
+  text: string
+  document: T
+}
+
+/**
+ * Format documents with line numbers and create mappings
+ */
+const formatDocumentsWithLineNumbers = <T>(documents: T[]): { formatted: string; mappings: LineMapping<T>[] } => {
+  const mappings: LineMapping<T>[] = []
+  const allLines: string[] = []
+  let globalLineNumber = 1
+
+  // First pass: count total lines to determine padding
+  let totalLines = 0
+  documents.forEach((doc) => {
+    const docString = stringify(doc)
+    const lines = docString.split('\n')
+    totalLines += lines.length
+  })
+
+  const padding = Math.max(3, totalLines.toString().length)
+
+  // Second pass: format with padding
+  documents.forEach((doc, docIndex) => {
+    const docString = stringify(doc)
+    const lines = docString.split('\n')
+
+    lines.forEach((line, lineInDoc) => {
+      const paddedNumber = globalLineNumber.toString().padStart(padding, '0')
+      const formattedLine = `■${paddedNumber} | ${line}`
+      allLines.push(formattedLine)
+
+      mappings.push({
+        lineNumber: globalLineNumber,
+        documentIndex: docIndex,
+        lineInDocument: lineInDoc,
+        text: line,
+        document: doc,
+      })
+
+      globalLineNumber++
+    })
+  })
+
+  return {
+    formatted: allLines.join('\n'),
+    mappings,
+  }
+}
+
+/**
+ * Parse citations from answer text
+ * Format: ■001, ■001-005, ■001■003■005
+ */
+const parseCitations = <T>(
+  answerText: string,
+  mappings: LineMapping<T>[]
+): { cleanAnswer: string; citations: Citation<T>[] } => {
+  const citations: Citation<T>[] = []
+  const citationPattern = /■(\d+)(?:-(\d+))?/g
+  let match: RegExpExecArray | null
+  const processedRanges = new Set<string>()
+
+  // Find all citation markers and replace them
+  let cleanAnswer = answerText
+  let offsetAdjustment = 0
+
+  while ((match = citationPattern.exec(answerText)) !== null) {
+    const fullMatch = match[0]
+    const offset = match.index - offsetAdjustment
+    const startLine = parseInt(match[1], 10)
+    const endLine = match[2] ? parseInt(match[2], 10) : startLine
+
+    // Generate range key to avoid duplicates
+    const rangeKey = `${offset}:${startLine}-${endLine}`
+    if (processedRanges.has(rangeKey)) {
+      continue
+    }
+    processedRanges.add(rangeKey)
+
+    // Collect all line numbers in this citation
+    const lineNumbers: number[] = []
+    for (let i = startLine; i <= endLine; i++) {
+      lineNumbers.push(i)
+    }
+
+    // Find the corresponding documents
+    const relevantMappings = mappings.filter((m) => lineNumbers.includes(m.lineNumber))
+    if (relevantMappings.length > 0) {
+      // Group by document to create citations
+      const documentMap = new Map<T, string[]>()
+      relevantMappings.forEach((mapping) => {
+        const existing = documentMap.get(mapping.document) || []
+        existing.push(mapping.text)
+        documentMap.set(mapping.document, existing)
+      })
+
+      documentMap.forEach((lines, document) => {
+        citations.push({
+          offset,
+          item: document,
+          snippet: lines.join('\n'),
+        })
+      })
+    }
+
+    // Remove citation marker from answer text
+    cleanAnswer = cleanAnswer.slice(0, offset) + cleanAnswer.slice(offset + fullMatch.length)
+    offsetAdjustment += fullMatch.length
+  }
+
+  return {
+    cleanAnswer: cleanAnswer.trim(),
+    citations: citations.sort((a, b) => a.offset - b.offset),
+  }
+}
+
+/**
+ * Process a single chunk of documents (all fit in one LLM call)
+ */
+const processSingleChunk = async <T>(
+  formattedDocs: string,
+  mappings: LineMapping<T>[],
+  question: string,
+  options: Options<T>,
+  ctx: ZaiContext
+): Promise<AnswerResult<T>> => {
+  ctx.controller.signal.throwIfAborted()
+  const result = await callLLM(formattedDocs, question, options, mappings, ctx)
+  ctx.controller.signal.throwIfAborted()
+  return result
+}
+
+/**
+ * Process multiple chunks and merge results
+ */
+const processMultipleChunks = async <T>(
+  documents: T[],
+  question: string,
+  chunkTokenLimit: number,
+  options: Options<T>,
+  ctx: ZaiContext
+): Promise<AnswerResult<T>> => {
+  const tokenizer = await getTokenizer()
+
+  // Split documents into chunks
+  const chunks: T[][] = []
+  let currentChunk: T[] = []
+  let currentTokens = 0
+
+  for (const doc of documents) {
+    const docString = stringify(doc)
+    const docTokens = tokenizer.count(docString)
+
+    if (currentTokens + docTokens > chunkTokenLimit && currentChunk.length > 0) {
+      chunks.push([...currentChunk])
+      currentChunk = [doc]
+      currentTokens = docTokens
+    } else {
+      currentChunk.push(doc)
+      currentTokens += docTokens
+    }
+  }
+
+  if (currentChunk.length > 0) {
+    chunks.push(currentChunk)
+  }
+
+  // Process a single chunk
+  const processChunk = async (chunk: T[]): Promise<AnswerResult<T>> => {
+    ctx.controller.signal.throwIfAborted()
+    const { formatted, mappings } = formatDocumentsWithLineNumbers(chunk)
+    const result = await callLLM(formatted, question, options, mappings, ctx)
+    ctx.controller.signal.throwIfAborted()
+    return result
+  }
+
+  // Process all chunks in parallel
+  const limit = pLimit(10) // Limit to 10 concurrent operations
+  const chunkResults = await Promise.all(chunks.map((chunk) => limit(() => processChunk(chunk))))
+
+  ctx.controller.signal.throwIfAborted()
+
+  // Merge results
+  return mergeChunkResults(chunkResults, question, chunkTokenLimit, options, ctx)
+}
+
+/**
+ * Merge results from multiple chunks
+ */
+const mergeChunkResults = async <T>(
+  results: AnswerResult<T>[],
+  question: string,
+  chunkTokenLimit: number,
+  options: Options<T>,
+  ctx: ZaiContext
+): Promise<AnswerResult<T>> => {
+  ctx.controller.signal.throwIfAborted()
+  // Filter out non-answer results
+  const answers = results.filter((r): r is AnswerResponse<T> => r.type === 'answer')
+
+  if (answers.length === 0) {
+    // No answers found, return first non-answer result
+    const nonAnswer = results.find((r) => r.type !== 'answer')
+    return nonAnswer || { type: 'missing_knowledge', reason: 'No relevant information found in documents.' }
+  }
+
+  if (answers.length === 1) {
+    return answers[0]
+  }
+
+  // Collect all cited documents
+  const citedDocs = new Set<T>()
+  answers.forEach((answer) => {
+    answer.citations.forEach((citation) => {
+      citedDocs.add(citation.item)
+    })
+  })
+
+  const citedDocsArray = Array.from(citedDocs)
+
+  // Try to merge by re-running with only cited documents
+  const tokenizer = await getTokenizer()
+  const { formatted, mappings } = formatDocumentsWithLineNumbers(citedDocsArray)
+
+  if (tokenizer.count(formatted) <= chunkTokenLimit) {
+    // Cited docs fit in one chunk, get unified answer
+    return await callLLM(formatted, question, options, mappings, ctx)
+  }
+
+  // Still too large, return best answer from chunks
+  // Prefer the answer with most citations
+  return answers.reduce((best, current) => (current.citations.length > best.citations.length ? current : best))
+}
+
+/**
+ * Call LLM to generate answer with automatic retry on citation errors
+ */
+const callLLM = async <T>(
+  formattedDocs: string,
+  question: string,
+  options: Options<T>,
+  mappings: LineMapping<T>[],
+  ctx: ZaiContext
+): Promise<AnswerResult<T>> => {
+  const systemPrompt = `You are an expert research assistant specialized in answering questions using only the information provided in documents.
+
+# Task
+Answer the user's question based ONLY on the information in the provided documents. You MUST cite your sources using line numbers.
+
+# Document Format
+Documents are provided with line numbers:
+■001 | First line of text
+■002 | Second line of text
+■003 | Third line of text
+
+# Citation Format
+You MUST include citations immediately after statements. Use these formats:
+- Single line: ■035
+- Range: ■005-010
+- Multiple: ■035■046■094
+
+# Response Format
+
+Choose ONE of these response types:
+
+**TYPE 1 - ANSWER** (Use this when you can answer the question)
+■answer
+[Your answer with inline citations■001-003. Make sure each part is cited correctly■013. More text. ■015]
+■end■
+
+**TYPE 2 - AMBIGUOUS** (Use when the question has multiple valid interpretations)
+■ambiguous
+[Explain the ambiguity]
+■follow_up
+[Ask a clarifying question]
+■answer
+[First interpretation with citations ■001 and part 2 as well.■002]
+■answer
+[Second interpretation with citations ■005 and part 2 of the answer.■006]
+■end■
+
+**TYPE 3 - OUT OF TOPIC** (Use when question is completely unrelated to documents)
+■out_of_topic
+[Explain why it's unrelated]
+■end■
+
+**TYPE 4 - INVALID QUESTION** (Use when input is not a proper question, e.g., gibberish, malformed or nonsensical)
+■invalid_question
+[Explain why it's invalid, e.g., "The question is incomplete" or "The question contains nonsensical terms", or "Received gibberish"]
+■end■
+
+**TYPE 5 - MISSING KNOWLEDGE** (Use ONLY when documents lack specific details needed)
+■missing_knowledge
+[Explain what specific information is missing]
+■end■
+
+# Important Rules
+- PREFER answering when possible - only use missing_knowledge if truly no relevant info exists
+- ALWAYS cite sources with line numbers
+- Use ONLY information from the documents
+- Be precise and factual
+- Do NOT fabricate information
+- Do NOT mention "According to the documents" or similar phrases – just provide a high-quality answer with citations
+- Do not be too strict on the question format; assume high-level answers are acceptable unless the question clearly asks for very specific details or requests depth beyond the documents
+
+# Additional Instructions
+Here are some additional instructions to follow about how to answer the question:
+${options.instructions || 'Provide a clear and concise answer based on the documents.'}`
+
+  const userPrompt = `<documents>
+${formattedDocs}
+</documents>
+
+Please answer the below question using the format specified above.
+Question to answer: "${question}"`
+
+  const { extracted } = await ctx.generateContent({
+    reasoningEffort: 'none',
+    systemPrompt,
+    stopSequences: [END],
+    messages: [
+      {
+        type: 'text',
+        role: 'user',
+        content: userPrompt,
+      },
+    ],
+    transform: (text) => {
+      text = text.slice(0, text.lastIndexOf(END.slice(0, -1))) // Remove anything after END
+      // Parse and validate response - errors will be caught and retried
+      return parseResponse(text || '', mappings)
+    },
+  })
+
+  return extracted
+}
+
+/**
+ * Parse LLM response into structured result
+ * @internal - Exported for testing purposes only
+ */
+export const parseResponse = <T>(response: string, mappings: LineMapping<T>[]): AnswerResult<T> => {
+  const text = response.trim()
+
+  const answersCount = (text.match(new RegExp(ANSWER_START, 'g')) || []).length
+
+  // Check response type
+  if (text.includes(AMBIGUOUS_START) || answersCount >= 2) {
+    return parseAmbiguousResponse(text, mappings)
+  } else if (text.includes(ANSWER_START)) {
+    return parseAnswerResponse(text, mappings)
+  } else if (text.includes(OUT_OF_TOPIC_START)) {
+    return parseOutOfTopicResponse(text)
+  } else if (text.includes(INVALID_QUESTION_START)) {
+    return parseInvalidQuestionResponse(text)
+  } else if (text.includes(MISSING_KNOWLEDGE_START)) {
+    return parseMissingKnowledgeResponse(text)
+  }
+
+  // Default to missing knowledge if format not recognized
+  return {
+    type: 'missing_knowledge',
+    reason: 'Unable to determine response type from the model output.',
+  }
+}
+
+/**
+ * Parse answer response
+ */
+const parseAnswerResponse = <T>(text: string, mappings: LineMapping<T>[]): AnswerResult<T> => {
+  // Match from answer start to end of string (END is a stop sequence and never appears)
+  const answerMatch = text.match(new RegExp(`${ANSWER_START}(.+)$`, 's'))
+  if (!answerMatch) {
+    return {
+      type: 'missing_knowledge',
+      reason: 'Could not extract answer from response.',
+    }
+  }
+
+  const answerText = answerMatch[1].trim()
+  const { cleanAnswer, citations } = parseCitations(answerText, mappings)
+
+  // Validate that citations are present
+  if (citations.length === 0) {
+    throw new Error(
+      'Answer must include citations using the ■XXX format (e.g., ■001, ■001-005). Every statement should be backed by a citation to the source documents.'
+    )
+  }
+
+  return {
+    type: 'answer',
+    answer: cleanAnswer,
+    citations,
+  }
+}
+
+/**
+ * Parse ambiguous response
+ */
+const parseAmbiguousResponse = <T>(text: string, mappings: LineMapping<T>[]): AmbiguousResponse<T> => {
+  // Extract ambiguity explanation
+  const ambiguityMatch = text.match(new RegExp(`${AMBIGUOUS_START}(.+?)■follow_up`, 's'))
+  const ambiguity = ambiguityMatch ? ambiguityMatch[1].trim() : 'The question has multiple interpretations.'
+
+  // Extract follow-up question
+  const followUpMatch = text.match(/■follow_up(.+?)■answer/s)
+  const follow_up = followUpMatch ? followUpMatch[1].trim() : 'Please clarify your question.'
+
+  // Extract all possible answers (match until next ■answer or end of string)
+  const answerPattern = /■answer(.+?)(?=■answer|$)/gs
+  const answers: AnswerWithCitations<T>[] = []
+  let match: RegExpExecArray | null
+
+  while ((match = answerPattern.exec(text)) !== null) {
+    const answerText = match[1].trim()
+    const { cleanAnswer, citations } = parseCitations(answerText, mappings)
+    answers.push({ answer: cleanAnswer, citations })
+  }
+
+  // Validate that each answer has citations
+  const answersWithoutCitations = answers.filter((a) => a.citations.length === 0)
+  if (answersWithoutCitations.length > 0) {
+    throw new Error(
+      'Each answer in an ambiguous response must include citations using the ■XXX format (e.g., ■001, ■001-005). Every statement should be backed by a citation to the source documents.'
+    )
+  }
+
+  return {
+    type: 'ambiguous',
+    ambiguity,
+    follow_up,
+    answers: answers.length >= 2 ? answers.slice(0, 3) : answers,
+  }
+}
+
+/**
+ * Parse out of topic response
+ */
+const parseOutOfTopicResponse = (text: string): OutOfTopicResponse => {
+  const reasonMatch = text.match(new RegExp(`${OUT_OF_TOPIC_START}(.+)$`, 's'))
+  const reason = reasonMatch ? reasonMatch[1].trim() : 'The question is not related to the provided documents.'
+
+  return {
+    type: 'out_of_topic',
+    reason,
+  }
+}
+
+/**
+ * Parse invalid question response
+ */
+const parseInvalidQuestionResponse = (text: string): InvalidQuestionResponse => {
+  const reasonMatch = text.match(new RegExp(`${INVALID_QUESTION_START}(.+)$`, 's'))
+  const reason = reasonMatch ? reasonMatch[1].trim() : 'The question is invalid or malformed.'
+
+  return {
+    type: 'invalid_question',
+    reason,
+  }
+}
+
+/**
+ * Parse missing knowledge response
+ */
+const parseMissingKnowledgeResponse = (text: string): MissingKnowledgeResponse => {
+  const reasonMatch = text.match(new RegExp(`${MISSING_KNOWLEDGE_START}(.+)$`, 's'))
+  const reason = reasonMatch
+    ? reasonMatch[1].trim()
+    : 'The documents do not contain sufficient information to answer the question.'
+
+  return {
+    type: 'missing_knowledge',
+    reason,
+  }
+}
+
+/**
+ * Answer generation with intelligent chunking and merging strategy
+ *
+ * Strategy for handling large document sets:
+ *
+ * 1. **Single-pass case**: If all documents fit in one LLM call, process directly
+ *
+ * 2. **Multi-pass case**: If documents exceed token limit:
+ *    a. Split documents into chunks that fit in token budget
+ *    b. Run question over EVERY chunk (ensures every document is considered at least once)
+ *    c. Collect all partial answers with their citations
+ *    d. If citations span multiple chunks:
+ *       - Extract only the cited lines from each chunk
+ *       - Create a new, smaller document set with just cited content
+ *       - Re-run the question with this refined set
+ *       - Merge into a unified answer
+ *    e. Repeat refinement until:
+ *       - All citations fit in one call (unified answer achieved), OR
+ *       - Max refinement passes reached (default: 3)
+ *
+ * This approach ensures:
+ * - Every document is evaluated at least once
+ * - Progressive refinement focuses on most relevant content
+ * - Typical convergence in 1-3 LLM calls total
+ * - Coherent final answer even from scattered information
+ */
+const answer = async <T>(
+  documents: T[],
+  question: string,
+  options: Options<T>,
+  ctx: ZaiContext
+): Promise<AnswerResult<T>> => {
+  ctx.controller.signal.throwIfAborted()
+
+  if (documents.length === 0) {
+    return {
+      type: 'missing_knowledge',
+      reason: 'No documents provided to answer the question.',
+    }
+  }
+
+  if (!question || question.trim().length === 0) {
+    return {
+      type: 'invalid_question',
+      reason: 'The question is empty or contains no content.',
+    }
+  }
+
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+  const TOTAL_MAX_TOKENS = Math.min(options.chunkLength, model.input.maxTokens - PROMPT_INPUT_BUFFER)
+
+  // Format all documents with line numbers
+  const { formatted: allFormattedDocs, mappings: allMappings } = formatDocumentsWithLineNumbers(documents)
+
+  const CHUNK_DOC_TOKENS = Math.floor(TOTAL_MAX_TOKENS * 0.6)
+  const totalDocTokens = tokenizer.count(allFormattedDocs)
+
+  // Check if we need to chunk
+  if (totalDocTokens <= CHUNK_DOC_TOKENS) {
+    // Single pass - all documents fit
+    return await processSingleChunk(allFormattedDocs, allMappings, question, options, ctx)
+  } else {
+    // Multi-pass chunking required
+    return await processMultipleChunks(documents, question, CHUNK_DOC_TOKENS, options, ctx)
+  }
+}
+
+Zai.prototype.answer = function <T>(
+  this: Zai,
+  documents: T[],
+  question: string,
+  _options?: Options<T>
+): Response<AnswerResult<T>, AnswerResult<T>> {
+  const parse = _Options.safeParse(_options ?? {})
+
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.answer',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  if (!parse.success) {
+    return Response.reject<AnswerResult<T>>(context, new Error(`Invalid options: ${parse.error.message}`))
+  }
+
+  return new Response<AnswerResult<T>, AnswerResult<T>>(
+    context,
+    answer(documents, question, parse.data, context),
+    (result) => result // No simplification - return full result
+  )
+}

packages/zai/src/operations/check.ts

@@ -0,0 +1,367 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER } from './constants'
+
+const _Example = z.object({
+  input: z.any(),
+  check: z.boolean(),
+  reason: z.string().optional(),
+  condition: z.string().optional(),
+})
+
+type Example = {
+  input: unknown
+  check: boolean
+  reason?: string
+  condition?: string
+}
+
+export type Options = {
+  /** Examples to check the condition against */
+  examples?: Array<Example>
+}
+
+const _Options = z.object({
+  examples: z.array(_Example).describe('Examples to check the condition against').default([]),
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Checks whether a condition is true for the given input, with an explanation.
+     *
+     * This operation evaluates natural language conditions against your input data,
+     * returning both a boolean result and a detailed explanation. Perfect for
+     * content moderation, sentiment analysis, quality checks, and business rule validation.
+     *
+     * @param input - The data to evaluate (text, object, or any value)
+     * @param condition - Natural language description of the condition to check
+     * @param options - Optional examples to guide the evaluation
+     * @returns Response with { value: boolean, explanation: string }, simplified to boolean when awaited
+     *
+     * @example Basic sentiment check
+     * ```typescript
+     * const review = "This product exceeded my expectations!"
+     * const isPositive = await zai.check(review, 'Is the sentiment positive?')
+     * // Result: true
+     *
+     * // Get full details
+     * const { value, explanation } = await zai.check(review, 'Is the sentiment positive?').result()
+     * // value: true
+     * // explanation: "The review expresses satisfaction and exceeded expectations..."
+     * ```
+     *
+     * @example Content moderation
+     * ```typescript
+     * const comment = "Great article! Very informative."
+     * const isSpam = await zai.check(comment, 'Is this spam or promotional content?')
+     * // Result: false
+     *
+     * const hasProfanity = await zai.check(comment, 'Does this contain profanity or offensive language?')
+     * // Result: false
+     * ```
+     *
+     * @example Business rules validation
+     * ```typescript
+     * const invoice = {
+     *   total: 1500,
+     *   items: ['laptop', 'mouse'],
+     *   customer: 'Enterprise Corp'
+     * }
+     *
+     * const needsApproval = await zai.check(
+     *   invoice,
+     *   'Does this invoice require manager approval? (over $1000 or enterprise customer)'
+     * )
+     * // Result: true
+     * ```
+     *
+     * @example With examples for consistency
+     * ```typescript
+     * const result = await zai.check(text, 'Is this a technical question?', {
+     *   examples: [
+     *     {
+     *       input: 'How do I deploy to production?',
+     *       check: true,
+     *       reason: 'Question about deployment process'
+     *     },
+     *     {
+     *       input: 'What time is the meeting?',
+     *       check: false,
+     *       reason: 'Not a technical question'
+     *     }
+     *   ]
+     * })
+     * ```
+     *
+     * @example Quality assurance
+     * ```typescript
+     * const code = "function add(a, b) { return a + b }"
+     * const hasDocumentation = await zai.check(code, 'Is this code properly documented?')
+     * const hasTests = await zai.check(code, 'Does this include unit tests?')
+     * const followsConventions = await zai.check(code, 'Does this follow naming conventions?')
+     * ```
+     */
+    check(
+      input: unknown,
+      condition: string,
+      options?: Options
+    ): Response<
+      {
+        /** Whether the condition is true or not */
+        value: boolean
+        /** The explanation of the decision */
+        explanation: string
+      },
+      boolean
+    >
+  }
+}
+
+const TRUE = '■TRUE■'
+const FALSE = '■FALSE■'
+const END = '■END■'
+
+const check = async (
+  input: unknown,
+  condition: string,
+  options: Options,
+  ctx: ZaiContext
+): Promise<{
+  value: boolean
+  explanation: string
+}> => {
+  ctx.controller.signal.throwIfAborted()
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+  const PROMPT_COMPONENT = Math.max(model.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.check'
+
+  const PROMPT_TOKENS = {
+    INPUT: Math.floor(0.5 * PROMPT_COMPONENT),
+    CONDITION: Math.floor(0.2 * PROMPT_COMPONENT),
+  }
+
+  // Truncate the input to fit the model's input size
+  const inputAsString = tokenizer.truncate(stringify(input), PROMPT_TOKENS.INPUT)
+  condition = tokenizer.truncate(condition, PROMPT_TOKENS.CONDITION)
+
+  // All tokens remaining after the input and condition are accounted can be used for examples
+  const EXAMPLES_TOKENS = PROMPT_COMPONENT - tokenizer.count(inputAsString) - tokenizer.count(condition)
+
+  const Key = fastHash(
+    JSON.stringify({
+      taskType,
+      taskId,
+      input: inputAsString,
+      condition,
+    })
+  )
+
+  const examples =
+    taskId && ctx.adapter
+      ? await ctx.adapter.getExamples<string, boolean>({
+          input: inputAsString,
+          taskType,
+          taskId,
+        })
+      : []
+
+  const exactMatch = examples.find((x) => x.key === Key)
+  if (exactMatch) {
+    return { explanation: exactMatch.explanation ?? '', value: exactMatch.output }
+  }
+
+  const defaultExamples = [
+    {
+      input: '50 Cent',
+      check: true,
+      reason: '50 Cent is widely recognized as a public personality.',
+      condition: 'Is the input a public personality?',
+    },
+    {
+      input: ['apple', 'banana', 'carrot', 'house'],
+      check: false,
+      reason:
+        'The list contains a house, which is not a fruit. Also, the list contains a carrot, which is a vegetable.',
+      condition: 'Is the input exclusively a list of fruits?',
+    },
+  ] satisfies Example[]
+
+  const userExamples = [
+    ...examples.map((e) => ({ input: e.input, check: e.output, reason: e.explanation })),
+    ...options.examples,
+  ]
+
+  let exampleId = 1
+
+  const formatInput = (input: string, condition: string) => {
+    const header = userExamples.length ? `Expert Example #${exampleId++}` : `Example of condition: "${condition}"`
+
+    return `
+${header}
+<|start_input|>
+${input.trim()}
+<|end_input|>
+`.trim()
+  }
+
+  const formatOutput = (answer: boolean, justification: string) => {
+    return `
+Analysis: ${justification}
+Final Answer: ${answer ? TRUE : FALSE}
+${END}
+`.trim()
+  }
+
+  const formatExample = (example: Example) => [
+    {
+      type: 'text' as const,
+      content: formatInput(stringify(example.input ?? null), example.condition ?? condition),
+      role: 'user' as const,
+    },
+    {
+      type: 'text' as const,
+      content: formatOutput(example.check, example.reason ?? ''),
+      role: 'assistant' as const,
+    },
+  ]
+
+  const allExamples = takeUntilTokens(
+    userExamples.length ? userExamples : defaultExamples,
+    EXAMPLES_TOKENS,
+    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(el.reason ?? '')
+  )
+    .map(formatExample)
+    .flat()
+
+  const specialInstructions = userExamples.length
+    ? `
+- You have been provided with examples from previous experts. Make sure to read them carefully before making your decision.
+- Make sure to refer to the examples provided by the experts to justify your decision (when applicable).
+- When in doubt, ground your decision on the examples provided by the experts instead of your own intuition.
+- When no example is similar to the input, make sure to provide a clear justification for your decision while inferring the decision-making process from the examples provided by the experts.
+`.trim()
+    : ''
+
+  const {
+    extracted: { finalAnswer, explanation },
+    meta,
+  } = await ctx.generateContent({
+    systemPrompt: `
+Check if the following condition is true or false for the given input. Before answering, make sure to read the input and the condition carefully.
+Justify your answer, then answer with either ${TRUE} or ${FALSE} at the very end, then add ${END} to finish the response.
+IMPORTANT: Make sure to answer with either ${TRUE} or ${FALSE} at the end of your response, but NOT both.
+---
+Expert Examples (#1 to #${exampleId - 1}):
+${specialInstructions}
+`.trim(),
+    stopSequences: [END],
+    messages: [
+      ...allExamples,
+      {
+        type: 'text',
+        content: `
+Considering the below input and above examples, is the following condition true or false?
+${formatInput(inputAsString, condition)}
+In your "Analysis", please refer to the Expert Examples # to justify your decision.`.trim(),
+        role: 'user',
+      },
+    ],
+    transform: (text) => {
+      const hasTrue = text.includes(TRUE)
+      const hasFalse = text.includes(FALSE)
+
+      if (!hasTrue && !hasFalse) {
+        throw new Error(`The model did not return a valid answer. The response was: ${text}`)
+      }
+
+      let finalAnswer: boolean
+      const explanation = text
+        .replace(TRUE, '')
+        .replace(FALSE, '')
+        .replace(END, '')
+        .replace('Final Answer:', '')
+        .replace('Analysis:', '')
+        .trim()
+
+      if (hasTrue && hasFalse) {
+        // If both TRUE and FALSE are present, we need to check which one was answered last
+        finalAnswer = text.lastIndexOf(TRUE) > text.lastIndexOf(FALSE)
+      } else {
+        finalAnswer = hasTrue
+      }
+
+      return { finalAnswer, explanation: explanation.trim() }
+    },
+  })
+
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
+      key: Key,
+      taskType,
+      taskId,
+      input: inputAsString,
+      instructions: condition,
+      metadata: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output,
+        },
+        latency: meta.latency,
+        model: ctx.modelId,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output,
+        },
+      },
+      output: finalAnswer,
+      explanation,
+    })
+  }
+
+  return {
+    value: finalAnswer,
+    explanation: explanation.trim(),
+  }
+}
+
+Zai.prototype.check = function (
+  this: Zai,
+  input: unknown,
+  condition: string,
+  _options: Options | undefined
+): Response<
+  {
+    value: boolean
+    explanation: string
+  },
+  boolean
+> {
+  const options = _Options.parse(_options ?? {}) as Options
+
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.check',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<
+    {
+      value: boolean
+      explanation: string
+    },
+    boolean
+  >(context, check(input, condition, options, context), (result) => result.value)
+}

packages/zai/src/operations/constants.ts

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

packages/zai/src/operations/errors.ts

@@ -0,0 +1,104 @@
+import z, { ZodError } from '@bpinternal/zui'
+
+export class JsonParsingError extends Error {
+  public constructor(
+    public json: unknown,
+    public error: Error
+  ) {
+    const message = JsonParsingError._formatError(json, error)
+    super(message)
+  }
+
+  private static _formatError(json: unknown, error: Error): string {
+    let errorMessage = 'Error parsing JSON:\n\n'
+    errorMessage += `---JSON---\n${json}\n\n`
+
+    if (z.is.zuiError(error)) {
+      errorMessage += '---Validation Errors---\n\n'
+      errorMessage += JsonParsingError._formatZodError(error)
+    } else {
+      errorMessage += '---Error---\n\n'
+      errorMessage += 'The JSON provided is not valid JSON.\n'
+      errorMessage += `Details: ${error.message}\n`
+    }
+
+    return errorMessage
+  }
+
+  private static _formatZodError(zodError: ZodError): string {
+    const issues = zodError.issues
+    if (issues.length === 0) {
+      return 'Unknown validation error\n'
+    }
+
+    let message = ''
+    for (let i = 0; i < issues.length; i++) {
+      const issue = issues[i]
+      const path = issue.path.length > 0 ? issue.path.join('.') : 'root'
+
+      message += `${i + 1}. Field: "${path}"\n`
+
+      switch (issue.code) {
+        case 'invalid_type':
+          message += `   Problem: Expected ${issue.expected}, but received ${issue.received}\n`
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'invalid_string':
+          if ('validation' in issue) {
+            message += `   Problem: Invalid ${issue.validation} format\n`
+          }
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'too_small':
+          if (issue.type === 'string') {
+            if (issue.exact) {
+              message += `   Problem: String must be exactly ${issue.minimum} characters\n`
+            } else {
+              message += `   Problem: String must be at least ${issue.minimum} characters\n`
+            }
+          } else if (issue.type === 'number') {
+            message += `   Problem: Number must be ${issue.inclusive ? 'at least' : 'greater than'} ${issue.minimum}\n`
+          } else if (issue.type === 'array') {
+            message += `   Problem: Array must contain ${issue.inclusive ? 'at least' : 'more than'} ${issue.minimum} items\n`
+          }
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'too_big':
+          if (issue.type === 'string') {
+            if (issue.exact) {
+              message += `   Problem: String must be exactly ${issue.maximum} characters\n`
+            } else {
+              message += `   Problem: String must be at most ${issue.maximum} characters\n`
+            }
+          } else if (issue.type === 'number') {
+            message += `   Problem: Number must be ${issue.inclusive ? 'at most' : 'less than'} ${issue.maximum}\n`
+          } else if (issue.type === 'array') {
+            message += `   Problem: Array must contain ${issue.inclusive ? 'at most' : 'fewer than'} ${issue.maximum} items\n`
+          }
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'invalid_enum_value':
+          message += `   Problem: Invalid value "${issue.received}"\n`
+          message += `   Allowed values: ${issue.options.map((o) => `"${o}"`).join(', ')}\n`
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'invalid_literal':
+          message += `   Problem: Expected the literal value "${issue.expected}", but received "${issue.received}"\n`
+          message += `   Message: ${issue.message}\n`
+          break
+        case 'invalid_union':
+          message += "   Problem: Value doesn't match any of the expected formats\n"
+          message += `   Message: ${issue.message}\n`
+          break
+        default:
+          message += `   Problem: ${issue.message}\n`
+      }
+
+      if (i < issues.length - 1) {
+        message += '\n'
+      }
+    }
+
+    return message
+  }
+}

packages/zai/src/operations/extract.ts

@@ -0,0 +1,509 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+
+import JSON5 from 'json5'
+import { jsonrepair } from 'jsonrepair'
+
+import { chunk, isArray } from 'lodash-es'
+import pLimit from 'p-limit'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER } from './constants'
+import { JsonParsingError } from './errors'
+
+export type Options = {
+  /** Instructions to guide the user on how to extract the data */
+  instructions?: string
+  /** The maximum number of tokens per chunk */
+  chunkLength?: number
+  /** Whether to strictly follow the schema or not */
+  strict?: boolean
+}
+
+const INVALID_SCHEMA_ERROR = 'zai.extract only accepts schemas created with @bpinternal/zui'
+
+const Options = z.object({
+  instructions: z.string().optional().describe('Instructions to guide the user on how to extract the data'),
+  chunkLength: z
+    .number()
+    .min(100)
+    .max(100_000)
+    .optional()
+    .describe('The maximum number of tokens per chunk')
+    .default(16_000),
+  strict: z.boolean().optional().default(true).describe('Whether to strictly follow the schema or not'),
+})
+
+type __Z<T> = { __type__: 'ZuiType'; _output: T }
+type OfType<O, T extends __Z<any> = __Z<O>> = T extends __Z<O> ? T : never
+type AnyObjectOrArray = Record<string, unknown> | Array<unknown>
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Extracts structured data from unstructured text using a Zod schema.
+     *
+     * This operation uses LLMs to intelligently parse text and extract information
+     * according to your schema. It handles large inputs automatically by chunking
+     * and supports both objects and arrays.
+     *
+     * @param input - The text or data to extract information from
+     * @param schema - Zod schema defining the structure to extract
+     * @param options - Optional configuration for extraction behavior
+     * @returns A Response promise that resolves to data matching your schema
+     *
+     * @example Extract a single object
+     * ```typescript
+     * import { z } from '@bpinternal/zui'
+     *
+     * const personSchema = z.object({
+     *   name: z.string(),
+     *   age: z.number(),
+     *   email: z.string().email()
+     * })
+     *
+     * const text = "Contact John Doe (35) at john@example.com"
+     * const person = await zai.extract(text, personSchema)
+     * // Result: { name: 'John Doe', age: 35, email: 'john@example.com' }
+     * ```
+     *
+     * @example Extract an array of items
+     * ```typescript
+     * const productSchema = z.array(z.object({
+     *   name: z.string(),
+     *   price: z.number()
+     * }))
+     *
+     * const text = "We have Apple ($1.50), Banana ($0.80), and Orange ($1.20)"
+     * const products = await zai.extract(text, productSchema)
+     * // Result: [
+     * //   { name: 'Apple', price: 1.50 },
+     * //   { name: 'Banana', price: 0.80 },
+     * //   { name: 'Orange', price: 1.20 }
+     * // ]
+     * ```
+     *
+     * @example With custom instructions
+     * ```typescript
+     * const result = await zai.extract(document, schema, {
+     *   instructions: 'Only extract confirmed information, skip uncertain data',
+     *   chunkLength: 10000, // Smaller chunks for better accuracy
+     *   strict: true // Enforce strict schema validation
+     * })
+     * ```
+     *
+     * @example Track usage and cost
+     * ```typescript
+     * const response = zai.extract(text, schema)
+     *
+     * // Monitor progress
+     * response.on('progress', (usage) => {
+     *   console.log(`Tokens used: ${usage.tokens.total}`)
+     *   console.log(`Cost so far: $${usage.cost.total}`)
+     * })
+     *
+     * // Get full results
+     * const { output, usage, elapsed } = await response.result()
+     * console.log(`Extraction took ${elapsed}ms and cost $${usage.cost.total}`)
+     * ```
+     */
+    extract<S extends OfType<any>>(input: unknown, schema: S, options?: Options): Response<S['_output']>
+  }
+}
+const SPECIAL_CHAR = '■'
+const START = '■json_start■'
+const END = '■json_end■'
+const NO_MORE = '■NO_MORE_ELEMENT■'
+const ZERO_ELEMENTS = '■ZERO_ELEMENTS■'
+
+const extract = async <S extends OfType<AnyObjectOrArray>>(
+  input: unknown,
+  _schema: S,
+  _options: Options | undefined,
+  ctx: ZaiContext
+): Promise<S['_output']> => {
+  ctx.controller.signal.throwIfAborted()
+
+  if (!z.is.zuiType(_schema)) {
+    throw new Error(INVALID_SCHEMA_ERROR)
+  }
+
+  let originalSchema: z.ZodType
+  try {
+    originalSchema = z.transforms.fromJSONSchema(z.transforms.toJSONSchema(_schema))
+  } catch {
+    // The above transformers arent the legacy ones. They are very strict and might fail on some schema types.
+    originalSchema = _schema
+  }
+
+  const options = Options.parse(_options ?? {})
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.extract'
+
+  const PROMPT_COMPONENT = Math.max(model.input.maxTokens - PROMPT_INPUT_BUFFER, 100)
+
+  const { isArrayOfObjects, isWrappedValue, objSchema } = _parsing.parseSchema(originalSchema, {
+    partial: !options.strict,
+  })
+
+  const schemaTypescript = objSchema.toTypescriptType({ declaration: false, treatDefaultAsOptional: true })
+  const schemaLength = tokenizer.count(schemaTypescript)
+
+  options.chunkLength = Math.min(options.chunkLength, model.input.maxTokens - PROMPT_INPUT_BUFFER - schemaLength)
+
+  const keys = Object.keys(objSchema.shape)
+
+  const inputAsString = stringify(input)
+
+  if (tokenizer.count(inputAsString) > options.chunkLength) {
+    const limit = pLimit(10) // Limit to 10 concurrent extraction operations
+    const tokens = tokenizer.split(inputAsString)
+    const chunks = chunk(tokens, options.chunkLength).map((x) => x.join(''))
+    const all = await Promise.allSettled(
+      chunks.map((chunk) =>
+        limit(() =>
+          extract(
+            chunk,
+            originalSchema,
+            {
+              ...options,
+              strict: false, // We don't want to fail on strict mode for sub-chunks
+            },
+            ctx
+          )
+        )
+      )
+    ).then((results) =>
+      results.filter((x) => x.status === 'fulfilled').map((x) => (x as PromiseFulfilledResult<S['_output']>).value)
+    )
+
+    ctx.controller.signal.throwIfAborted()
+
+    // We run this function recursively until all chunks are merged into a single output
+    const rows = all.map((x, idx) => `<part-${idx + 1}>\n${stringify(x, true)}\n</part-${idx + 1}>`).join('\n')
+    return extract(
+      `
+The result has been split into ${all.length} parts. Recursively merge the result into the final result.
+When merging arrays, take unique values.
+When merging conflictual (but defined) information, take the most reasonable and frequent value.
+Non-defined values are OK and normal. Don't delete fields because of null values. Focus on defined values.
+
+Here's the data:
+${rows}
+
+Merge it back into a final result.`.trim(),
+      originalSchema,
+      options,
+      ctx
+    )
+  }
+
+  const instructions: string[] = []
+
+  if (options.instructions) {
+    instructions.push(options.instructions)
+  }
+
+  const shape = `{ ${keys.map((key) => `"${key}": ...`).join(', ')} }`
+  const abbv = '{ ... }'
+
+  if (isArrayOfObjects) {
+    instructions.push('You may have multiple elements, or zero elements in the input.')
+    instructions.push('You must extract each element separately.')
+    instructions.push(`Each element must be a JSON object with exactly the format: ${START}${shape}${END}`)
+    instructions.push(`If there are no elements to extract, respond with ${ZERO_ELEMENTS}.`)
+    instructions.push(`When you are done extracting all elements, type "${NO_MORE}" to finish.`)
+    instructions.push(
+      `For example, if you have zero elements, the output should look like this: ${ZERO_ELEMENTS}${NO_MORE}`
+    )
+    instructions.push(
+      `For example, if you have two elements, the output should look like this: ${START}${abbv}${END}${START}${abbv}${END}${NO_MORE}`
+    )
+  } else {
+    instructions.push('You may have exactly one element in the input.')
+    instructions.push(`The element must be a JSON object with exactly the format: ${START}${shape}${END}`)
+  }
+
+  if (!options.strict) {
+    instructions.push('You may ignore any fields that are not present in the input. All keys are optional.')
+  }
+
+  // All tokens remaining after the input and condition are accounted can be used for examples
+  const EXAMPLES_TOKENS = PROMPT_COMPONENT - tokenizer.count(inputAsString) - tokenizer.count(instructions.join('\n'))
+
+  const Key = fastHash(
+    JSON.stringify({
+      taskType,
+      taskId,
+      input: inputAsString,
+      instructions: options.instructions,
+    })
+  )
+
+  const examples =
+    taskId && ctx.adapter
+      ? await ctx.adapter.getExamples<string, unknown>({
+          input: inputAsString,
+          taskType,
+          taskId,
+        })
+      : []
+
+  const exactMatch = examples.find((x) => x.key === Key)
+  if (exactMatch) {
+    return exactMatch.output as S['_output']
+  }
+
+  const defaultExample = isArrayOfObjects
+    ? {
+        input: `The story goes as follow.
+Once upon a time, there was a person named Alice who was 30 years old.
+Then, there was a person named Bob who was 25 years old.
+The end.`,
+        schema: 'Array<{ name: string, age: number }>',
+        instructions: 'Extract all people',
+        extracted: [
+          {
+            name: 'Alice',
+            age: 30,
+          },
+          {
+            name: 'Bob',
+            age: 25,
+          },
+        ],
+      }
+    : {
+        input: `The story goes as follow.
+Once upon a time, there was a person named Alice who was 30 years old.
+The end.`,
+        schema: '{ name: string, age: number }',
+        instructions: 'Extract the person',
+        extracted: { name: 'Alice', age: 30 },
+      }
+
+  const userExamples = examples.map((e) => ({
+    input: e.input,
+    extracted: e.output,
+    schema: schemaTypescript,
+    instructions: options.instructions,
+  }))
+
+  let exampleId = 1
+
+  const formatInput = (input: string, schema: string, instructions?: string) => {
+    const header = userExamples.length
+      ? `Expert Example #${exampleId++}`
+      : "Here's an example to help you understand the format:"
+
+    return `
+${header}
+
+<|start_schema|>
+${schema}
+<|end_schema|>
+
+<|start_instructions|>
+${instructions ?? 'No specific instructions, just follow the schema above.'}
+<|end_instructions|>
+
+<|start_input|>
+${input.trim()}
+<|end_input|>
+  `.trim()
+  }
+
+  const formatOutput = (extracted: unknown) => {
+    const extractedArr = isArray(extracted) ? extracted : [extracted]
+
+    return (
+      extractedArr
+        .map((x: string) =>
+          `
+${START}
+${JSON.stringify(x, null, 2)}
+${END}`.trim()
+        )
+        .join('\n') + NO_MORE
+    )
+  }
+
+  const formatExample = (example: { input?: unknown; schema: string; instructions?: string; extracted: unknown }) => [
+    {
+      type: 'text' as const,
+      content: formatInput(stringify(example.input ?? null), example.schema, example.instructions),
+      role: 'user' as const,
+    },
+    {
+      type: 'text' as const,
+      content: formatOutput(example.extracted),
+      role: 'assistant' as const,
+    },
+  ]
+
+  const allExamples = takeUntilTokens(
+    userExamples.length ? userExamples : [defaultExample],
+    EXAMPLES_TOKENS,
+    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.extracted))
+  )
+    .map(formatExample)
+    .flat()
+
+  const { meta, extracted } = await ctx.generateContent({
+    systemPrompt: `
+Extract the following information from the input:
+${schemaTypescript}
+====
+
+${instructions.map((x) => `• ${x}`).join('\n')}
+`.trim(),
+    stopSequences: [isArrayOfObjects ? NO_MORE : END],
+    messages: [
+      ...allExamples,
+      {
+        role: 'user',
+        type: 'text',
+        content: formatInput(inputAsString, schemaTypescript, options.instructions ?? ''),
+      },
+    ],
+    transform: (text): unknown[] =>
+      (text || '{}')
+        .split(START)
+        .filter((x) => x.trim().length > 0 && x.includes('}'))
+        .map((x) => {
+          try {
+            let json = x.slice(0, x.indexOf(END)).trim()
+
+            if (json.includes(SPECIAL_CHAR)) {
+              json = json.slice(0, json.indexOf(SPECIAL_CHAR)).trim()
+            }
+
+            const repairedJson = jsonrepair(json)
+            const parsedJson: unknown = JSON5.parse(repairedJson)
+            const safe = objSchema.safeParse(parsedJson)
+
+            if (safe.success) {
+              return safe.data
+            }
+
+            if (options.strict) {
+              throw new JsonParsingError(x, safe.error)
+            }
+
+            return parsedJson
+          } catch (error) {
+            throw new JsonParsingError(x, error instanceof Error ? error : new Error('Unknown error'))
+          }
+        })
+        .filter((x) => x !== null),
+  })
+
+  let final: unknown
+
+  if (isArrayOfObjects) {
+    final = extracted
+  } else if (extracted.length === 0) {
+    final = options.strict ? objSchema.parse({}) : {}
+  } else {
+    final = extracted[0]
+  }
+
+  if (isWrappedValue) {
+    if (Array.isArray(final)) {
+      final = final.map((x) => ('value' in x ? x.value : x))
+    } else if (typeof final === 'object' && final !== null) {
+      final = 'value' in final ? final.value : final
+    }
+  }
+
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
+      key: Key,
+      taskId: `zai/${taskId}`,
+      taskType,
+      instructions: options.instructions ?? 'No specific instructions',
+      input: inputAsString,
+      output: final,
+      metadata: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output,
+        },
+        latency: meta.latency,
+        model: ctx.modelId,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output,
+        },
+      },
+    })
+  }
+
+  return final as S['_output']
+}
+
+Zai.prototype.extract = function <S extends OfType<AnyObjectOrArray>>(
+  this: Zai,
+  input: unknown,
+  schema: S,
+  _options?: Options
+): Response<S['_output']> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.extract',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<S['_output']>(context, extract(input, schema, _options, context), (result) => result)
+}
+
+namespace _parsing {
+  const _getBaseSchema = (schema: z.ZodType): z.ZodType => (schema.naked ? schema.naked() : schema)
+
+  const _maybeWrapObjSchema = (s: z.ZodType): { schema: z.ZodObject; isWrapped: boolean } => {
+    if (z.is.zuiObject(s)) {
+      return { schema: s, isWrapped: false }
+    }
+    return {
+      schema: z.object({
+        value: s,
+      }),
+      isWrapped: true,
+    }
+  }
+
+  const _applySchemaOptions = (schema: z.ZodObject, opts: ParseSchemaOptions): z.ZodObject => {
+    if (!opts.partial) {
+      return schema
+    }
+
+    try {
+      return schema.partial()
+    } catch {
+      return schema
+    }
+  }
+
+  export type ParseSchemaOptions = { partial: boolean }
+  export const parseSchema = (originalSchema: z.ZodType, opts: ParseSchemaOptions) => {
+    const baseType = _getBaseSchema(originalSchema)
+
+    if (z.is.zuiArray(baseType)) {
+      const elementType = _getBaseSchema(baseType.element)
+      const { schema, isWrapped } = _maybeWrapObjSchema(elementType)
+      return { isArrayOfObjects: true, isWrappedValue: isWrapped, objSchema: _applySchemaOptions(schema, opts) }
+    }
+
+    const { schema, isWrapped: wrapped } = _maybeWrapObjSchema(baseType)
+    return { isArrayOfObjects: false, isWrappedValue: wrapped, objSchema: _applySchemaOptions(schema, opts) }
+  }
+}

packages/zai/src/operations/filter.ts

@@ -0,0 +1,370 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+
+import { clamp } from 'lodash-es'
+import pLimit from 'p-limit'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+
+type Example = {
+  input: unknown
+  filter: boolean
+  reason?: string
+}
+
+const _Example = z.object({
+  input: z.any(),
+  filter: z.boolean(),
+  reason: z.string().optional(),
+})
+
+export type Options = {
+  /** The maximum number of tokens per item */
+  tokensPerItem?: number
+  /** Examples to filter the condition against */
+  examples?: Array<Example>
+}
+
+const _Options = z.object({
+  tokensPerItem: z
+    .number()
+    .min(1)
+    .max(100_000)
+    .optional()
+    .describe('The maximum number of tokens per item')
+    .default(250),
+  examples: z.array(_Example).describe('Examples to filter the condition against').default([]),
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Filters array elements based on a natural language condition.
+     *
+     * This operation evaluates each element against a condition using LLMs,
+     * returning only elements that match. Handles large arrays automatically
+     * by processing in parallel chunks.
+     *
+     * @param input - Array of elements to filter
+     * @param condition - Natural language description of what to keep
+     * @param options - Configuration for token limits per item and examples
+     * @returns Response promise resolving to filtered array
+     *
+     * @example Filter positive reviews
+     * ```typescript
+     * const reviews = [
+     *   "Great product, love it!",
+     *   "Terrible quality, broke immediately",
+     *   "Amazing! Exceeded expectations",
+     *   "Worst purchase ever"
+     * ]
+     *
+     * const positive = await zai.filter(reviews, 'Keep only positive reviews')
+     * // Result: ["Great product, love it!", "Amazing! Exceeded expectations"]
+     * ```
+     *
+     * @example Filter technical questions
+     * ```typescript
+     * const questions = [
+     *   "How do I deploy to production?",
+     *   "What time is lunch?",
+     *   "Why is the API returning 500 errors?",
+     *   "Can you book the conference room?"
+     * ]
+     *
+     * const technical = await zai.filter(
+     *   questions,
+     *   'Keep only technical or engineering questions'
+     * )
+     * // Result: ["How do I deploy to production?", "Why is the API returning 500 errors?"]
+     * ```
+     *
+     * @example Filter with object array
+     * ```typescript
+     * const products = [
+     *   { name: 'Laptop', category: 'Electronics', inStock: true },
+     *   { name: 'Desk', category: 'Furniture', inStock: false },
+     *   { name: 'Mouse', category: 'Electronics', inStock: true },
+     *   { name: 'Chair', category: 'Furniture', inStock: true }
+     * ]
+     *
+     * const available = await zai.filter(
+     *   products,
+     *   'Keep only products that are in stock'
+     * )
+     * // Returns products where inStock === true
+     * ```
+     *
+     * @example Complex filtering logic
+     * ```typescript
+     * const emails = [...] // Array of email objects
+     *
+     * const urgent = await zai.filter(
+     *   emails,
+     *   'Keep emails that are urgent, from the CEO, or mention "critical bug"'
+     * )
+     * ```
+     *
+     * @example With examples for consistency
+     * ```typescript
+     * const filtered = await zai.filter(items, 'Keep valid entries', {
+     *   examples: [
+     *     {
+     *       input: { status: 'active', verified: true },
+     *       filter: true,
+     *       reason: 'Active and verified'
+     *     },
+     *     {
+     *       input: { status: 'pending', verified: false },
+     *       filter: false,
+     *       reason: 'Not yet verified'
+     *     }
+     *   ],
+     *   tokensPerItem: 100 // Limit tokens per item for performance
+     * })
+     * ```
+     */
+    filter<T>(input: Array<T>, condition: string, options?: Options): Response<Array<T>>
+  }
+}
+
+const END = '■END■'
+
+const filter = async <T>(
+  input: Array<T>,
+  condition: string,
+  _options: Options | undefined,
+  ctx: ZaiContext
+): Promise<Array<T>> => {
+  ctx.controller.signal.throwIfAborted()
+  const options = _Options.parse(_options ?? {}) as Options
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.filter'
+
+  const MAX_ITEMS_PER_CHUNK = 50
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
+  const TOKENS_EXAMPLES_MAX = Math.floor(Math.max(250, TOKENS_TOTAL_MAX * 0.5))
+  const TOKENS_CONDITION_MAX = clamp(TOKENS_TOTAL_MAX * 0.25, 250, tokenizer.count(condition))
+  const TOKENS_INPUT_ARRAY_MAX = TOKENS_TOTAL_MAX - TOKENS_EXAMPLES_MAX - TOKENS_CONDITION_MAX
+
+  condition = tokenizer.truncate(condition, TOKENS_CONDITION_MAX)
+
+  let chunks: Array<typeof input> = []
+  let currentChunk: typeof input = []
+  let currentChunkTokens = 0
+
+  for (const element of input) {
+    const elementAsString = tokenizer.truncate(stringify(element, false), options.tokensPerItem)
+    const elementTokens = tokenizer.count(elementAsString)
+
+    if (currentChunkTokens + elementTokens > TOKENS_INPUT_ARRAY_MAX || currentChunk.length >= MAX_ITEMS_PER_CHUNK) {
+      chunks.push(currentChunk)
+      currentChunk = []
+      currentChunkTokens = 0
+    }
+
+    currentChunk.push(element)
+    currentChunkTokens += elementTokens
+  }
+
+  if (currentChunk.length > 0) {
+    chunks.push(currentChunk)
+  }
+
+  chunks = chunks.filter((x) => x.length > 0)
+
+  // ■1:true■2:true■3:true
+
+  const formatInput = (input: Example[], condition: string) => {
+    return `
+Condition to check:
+${condition}
+
+Items (from ■0 to ■${input.length - 1})
+==============================
+${input.map((x, idx) => `■${idx} = ${stringify(x.input ?? null, false)}`).join('\n')}
+`.trim()
+  }
+
+  const formatExamples = (examples: Example[]) => {
+    return `
+${examples.map((x, idx) => `■${idx}:${!!x.filter ? 'true' : 'false'}`).join('')}
+${END}
+====
+Here's the reasoning behind each example:
+${examples.map((x, idx) => `■${idx}:${!!x.filter ? 'true' : 'false'}:${x.reason ?? 'No reason provided'}`).join('\n')}
+`.trim()
+  }
+
+  const genericExamples: Example[] = [
+    {
+      input: 'apple',
+      filter: true,
+      reason: 'Apples are fruits',
+    },
+    {
+      input: 'Apple Inc.',
+      filter: false,
+      reason: 'Apple Inc. is a company, not a fruit',
+    },
+    {
+      input: 'banana',
+      filter: true,
+      reason: 'Bananas are fruits',
+    },
+    {
+      input: 'potato',
+      filter: false,
+      reason: 'Potatoes are vegetables',
+    },
+  ]
+
+  const genericExamplesMessages = [
+    {
+      type: 'text' as const,
+      content: formatInput(genericExamples, 'is a fruit'),
+      role: 'user' as const,
+    },
+    {
+      type: 'text' as const,
+      content: formatExamples(genericExamples),
+      role: 'assistant' as const,
+    },
+  ]
+
+  const filterChunk = async (chunk: typeof input) => {
+    const examples =
+      taskId && ctx.adapter
+        ? await ctx.adapter
+            .getExamples<string, unknown>({
+              // The Table API can't search for a huge input string
+              input: JSON.stringify(chunk).slice(0, 1000),
+              taskType,
+              taskId,
+            })
+            .then((x) =>
+              x.map((y) => ({ filter: y.output as boolean, input: y.input, reason: y.explanation }) satisfies Example)
+            )
+        : []
+
+    const allExamples = takeUntilTokens([...examples, ...(options.examples ?? [])], TOKENS_EXAMPLES_MAX, (el) =>
+      tokenizer.count(stringify(el.input))
+    )
+
+    const exampleMessages = [
+      {
+        type: 'text' as const,
+        content: formatInput(allExamples, condition),
+        role: 'user' as const,
+      },
+      {
+        type: 'text' as const,
+        content: formatExamples(allExamples),
+        role: 'assistant' as const,
+      },
+    ]
+
+    const { extracted: partial, meta } = await ctx.generateContent({
+      systemPrompt: `
+You are given a list of items. Your task is to filter out the items that meet the condition below.
+You need to return the full list of items with the format:
+■x:true■y:false■z:true (where x, y, z are the indices of the items in the list)
+You need to start with "■0" and go up to the last index "■${chunk.length - 1}".
+If an item meets the condition, you should return ":true", otherwise ":false".
+
+IMPORTANT: Make sure to read the condition and the examples carefully before making your decision.
+The condition is: "${condition}"
+`.trim(),
+      stopSequences: [END],
+      messages: [
+        ...(exampleMessages.length ? exampleMessages : genericExamplesMessages),
+        {
+          type: 'text',
+          content: formatInput(
+            chunk.map((x) => ({ input: x }) as Example),
+            condition
+          ),
+          role: 'user',
+        },
+      ],
+      transform: (text) => {
+        const indices = text
+          .trim()
+          .split('■')
+          .filter((x) => x.length > 0)
+          .map((x) => {
+            const [idx, filter] = x.split(':')
+            return { idx: parseInt(idx?.trim() ?? ''), filter: filter?.toLowerCase().trim() === 'true' }
+          })
+
+        return chunk.filter((_, idx) => {
+          return indices.find((x) => x.idx === idx && x.filter) ?? false
+        })
+      },
+    })
+
+    if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+      const key = fastHash(
+        stringify({
+          taskId,
+          taskType,
+          input: JSON.stringify(chunk),
+          condition,
+        })
+      )
+
+      await ctx.adapter.saveExample({
+        key,
+        taskType,
+        taskId,
+        input: JSON.stringify(chunk),
+        output: partial,
+        instructions: condition,
+        metadata: {
+          cost: {
+            input: meta.cost.input,
+            output: meta.cost.output,
+          },
+          latency: meta.latency,
+          model: ctx.modelId,
+          tokens: {
+            input: meta.tokens.input,
+            output: meta.tokens.output,
+          },
+        },
+      })
+    }
+
+    return partial
+  }
+
+  const limit = pLimit(10) // Limit to 10 concurrent filtering operations
+  const filteredChunks = await Promise.all(chunks.map((chunk) => limit(() => filterChunk(chunk))))
+
+  return filteredChunks.flat()
+}
+
+Zai.prototype.filter = function <T>(
+  this: Zai,
+  input: Array<T>,
+  condition: string,
+  _options?: Options
+): Response<Array<T>> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.filter',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<Array<T>>(context, filter(input, condition, _options, context), (result) => result)
+}

packages/zai/src/operations/group.ts

@@ -0,0 +1,971 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+import { clamp } from 'lodash-es'
+import pLimit from 'p-limit'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+
+export type Group<T> = {
+  id: string
+  label: string
+  elements: T[]
+}
+
+type InitialGroup = {
+  id: string
+  label: string
+  elements?: unknown[]
+}
+
+const _InitialGroup = z.object({
+  id: z.string().min(1).max(100),
+  label: z.string().min(1).max(250),
+  elements: z.array(z.any()).optional().default([]),
+})
+
+export type Options = {
+  instructions?: string
+  tokensPerElement?: number
+  chunkLength?: number
+  initialGroups?: Array<InitialGroup>
+  maxGroups?: number
+  minElements?: number
+}
+
+const _Options = z.object({
+  instructions: z.string().optional(),
+  tokensPerElement: z.number().min(1).max(100_000).optional().default(250),
+  chunkLength: z.number().min(100).max(100_000).optional().default(16_000),
+  initialGroups: z.array(_InitialGroup).optional().default([]),
+  maxGroups: z.number().min(2).optional(),
+  minElements: z.number().min(1).optional(),
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Groups array items into categories based on semantic similarity or criteria.
+     *
+     * This operation intelligently categorizes items by analyzing their content and
+     * creating meaningful groups. It can discover natural groupings automatically or
+     * use predefined categories. Perfect for clustering, classification, and organization.
+     *
+     * @param input - Array of items to group
+     * @param options - Configuration for grouping behavior, instructions, and initial categories
+     * @param options.maxGroups - Maximum number of groups allowed (minimum 2). When set, groups are merged at the end until within limit.
+     * @param options.minElements - Minimum elements per group (minimum 1). Groups below this threshold have their elements redistributed via AI.
+     * @returns Response with groups array (simplified to Record<groupLabel, items[]>)
+     *
+     * @example Automatic grouping
+     * ```typescript
+     * const messages = [
+     *   "I can't log in to my account",
+     *   "How do I reset my password?",
+     *   "When will my order arrive?",
+     *   "The app keeps crashing",
+     *   "I haven't received my package",
+     *   "Error 500 on checkout"
+     * ]
+     *
+     * const groups = await zai.group(messages, {
+     *   instructions: 'Group by type of customer issue'
+     * })
+     * // Result (simplified):
+     * // {
+     * //   "Login Issues": ["I can't log in...", "How do I reset..."],
+     * //   "Shipping Questions": ["When will my order...", "I haven't received..."],
+     * //   "Technical Errors": ["The app keeps crashing", "Error 500..."]
+     * // }
+     *
+     * // Full result:
+     * const { output } = await zai.group(messages, { instructions: '...' }).result()
+     * // output: [
+     * //   { id: 'login_issues', label: 'Login Issues', elements: [...] },
+     * //   { id: 'shipping', label: 'Shipping Questions', elements: [...] },
+     * //   { id: 'errors', label: 'Technical Errors', elements: [...] }
+     * // ]
+     * ```
+     *
+     * @example With predefined categories
+     * ```typescript
+     * const articles = [
+     *   "How to build a React app",
+     *   "Python machine learning tutorial",
+     *   "Understanding Docker containers",
+     *   "Vue.js best practices",
+     *   "Deep learning with TensorFlow"
+     * ]
+     *
+     * const groups = await zai.group(articles, {
+     *   instructions: 'Categorize by technology',
+     *   initialGroups: [
+     *     { id: 'frontend', label: 'Frontend Development' },
+     *     { id: 'backend', label: 'Backend Development' },
+     *     { id: 'ml', label: 'Machine Learning' },
+     *     { id: 'devops', label: 'DevOps & Infrastructure' }
+     *   ]
+     * })
+     * // Groups articles into predefined categories
+     * ```
+     *
+     * @example Grouping products
+     * ```typescript
+     * const products = [
+     *   { name: 'Laptop', price: 999, category: 'Electronics' },
+     *   { name: 'Desk', price: 299, category: 'Furniture' },
+     *   { name: 'Mouse', price: 29, category: 'Electronics' },
+     *   { name: 'Chair', price: 199, category: 'Furniture' }
+     * ]
+     *
+     * const grouped = await zai.group(products, {
+     *   instructions: 'Group by price range: budget (< $100), mid-range ($100-$500), premium (> $500)'
+     * })
+     * // Result:
+     * // {
+     * //   "Budget": [Mouse],
+     * //   "Mid-range": [Chair, Desk],
+     * //   "Premium": [Laptop]
+     * // }
+     * ```
+     *
+     * @example Content categorization
+     * ```typescript
+     * const emails = [
+     *   { subject: 'Meeting tomorrow', body: '...', from: 'boss@company.com' },
+     *   { subject: 'Invoice #1234', body: '...', from: 'billing@vendor.com' },
+     *   { subject: 'Weekly report', body: '...', from: 'team@company.com' },
+     *   { subject: 'Payment received', body: '...', from: 'accounting@company.com' }
+     * ]
+     *
+     * const categorized = await zai.group(emails, {
+     *   instructions: 'Categorize by email type: work communication, financial, reports',
+     *   tokensPerElement: 300 // Allow more context per email
+     * })
+     * ```
+     *
+     * @example Customer feedback grouping
+     * ```typescript
+     * const feedback = [
+     *   "Love the new UI!",
+     *   "App is too slow",
+     *   "Great customer service",
+     *   "Confusing navigation",
+     *   "Fast shipping!",
+     *   "Hard to find features"
+     * ]
+     *
+     * const grouped = await zai.group(feedback, {
+     *   instructions: 'Group by aspect: UI/UX, Performance, Customer Service, Shipping'
+     * })
+     * ```
+     *
+     * @example Topic clustering for research
+     * ```typescript
+     * const papers = [
+     *   { title: 'Transformer Networks for NLP', abstract: '...' },
+     *   { title: 'CNN Image Classification', abstract: '...' },
+     *   { title: 'BERT Language Understanding', abstract: '...' },
+     *   { title: 'Object Detection with YOLO', abstract: '...' }
+     * ]
+     *
+     * const clusters = await zai.group(papers, {
+     *   instructions: 'Group by research area',
+     *   chunkLength: 10000 // Allow more tokens for detailed abstracts
+     * })
+     * // Result: Groups papers by topic (NLP, Computer Vision, etc.)
+     * ```
+     *
+     * @example With initial seed groups
+     * ```typescript
+     * const tasks = [
+     *   "Fix login bug",
+     *   "Update documentation",
+     *   "Add dark mode",
+     *   "Optimize database queries"
+     * ]
+     *
+     * const grouped = await zai.group(tasks, {
+     *   instructions: 'Categorize development tasks',
+     *   initialGroups: [
+     *     { id: 'bugs', label: 'Bug Fixes', elements: [] },
+     *     { id: 'features', label: 'New Features', elements: [] },
+     *     { id: 'docs', label: 'Documentation', elements: [] },
+     *     { id: 'performance', label: 'Performance', elements: [] }
+     *   ]
+     * })
+     * ```
+     */
+    /**
+     * @example Limiting number of groups
+     * ```typescript
+     * const items = ['apple', 'banana', 'carrot', 'chicken', 'rice', 'bread', 'salmon', 'milk']
+     *
+     * const groups = await zai.group(items, {
+     *   instructions: 'Group by food type',
+     *   maxGroups: 3 // At most 3 groups — smallest groups get merged if exceeded
+     * })
+     * // Guarantees no more than 3 groups in the result
+     * ```
+     */
+    group<T>(input: Array<T>, options?: Options): Response<Array<Group<T>>, Record<string, T[]>>
+  }
+}
+
+const END = '■END■'
+
+// Simplified data structures
+type GroupInfo = {
+  id: string
+  label: string
+  normalizedLabel: string
+}
+
+const normalizeLabel = (label: string): string => {
+  return label
+    .trim()
+    .toLowerCase()
+    .replace(/^(group|new group|new)\s*[-:]\s*/i, '')
+    .replace(/^(group|new group|new)\s+/i, '')
+    .trim()
+}
+
+const group = async <T>(input: Array<T>, _options: Options | undefined, ctx: ZaiContext): Promise<Array<Group<T>>> => {
+  ctx.controller.signal.throwIfAborted()
+
+  const options = _Options.parse(_options ?? {})
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.group'
+
+  if (input.length === 0) {
+    return []
+  }
+
+  // Simple data structures
+  const groups = new Map<string, GroupInfo>() // groupId -> GroupInfo
+  const groupElements = new Map<string, Set<number>>() // groupId -> Set of element indices
+  const elementGroups = new Map<number, Set<string>>() // elementIndex -> Set of groupIds seen/assigned
+  const labelToGroupId = new Map<string, string>() // normalized label -> groupId
+  let groupIdCounter = 0
+
+  // Initialize with provided groups
+  options.initialGroups.forEach((ig) => {
+    const normalized = normalizeLabel(ig.label)
+    groups.set(ig.id, { id: ig.id, label: ig.label, normalizedLabel: normalized })
+    groupElements.set(ig.id, new Set())
+    labelToGroupId.set(normalized, ig.id)
+  })
+
+  // Prepare elements
+  const elements = input.map((element, idx) => ({
+    element,
+    index: idx,
+    stringified: stringify(element, false),
+  }))
+
+  // Token budget
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
+  const TOKENS_INSTRUCTIONS_MAX = options.instructions
+    ? clamp(tokenizer.count(options.instructions), 100, TOKENS_TOTAL_MAX * 0.2)
+    : 0
+  const TOKENS_AVAILABLE = TOKENS_TOTAL_MAX - TOKENS_INSTRUCTIONS_MAX
+  const TOKENS_FOR_GROUPS_MAX = Math.floor(TOKENS_AVAILABLE * 0.4)
+  const TOKENS_FOR_ELEMENTS_MAX = Math.floor(TOKENS_AVAILABLE * 0.6)
+
+  // Chunk elements by token budget
+  const MAX_ELEMENTS_PER_CHUNK = 50
+  const elementChunks: number[][] = [] // Array of element indices
+  let currentChunk: number[] = []
+  let currentTokens = 0
+
+  for (const elem of elements) {
+    const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement)
+    const elemTokens = tokenizer.count(truncated)
+
+    if (
+      (currentTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX || currentChunk.length >= MAX_ELEMENTS_PER_CHUNK) &&
+      currentChunk.length > 0
+    ) {
+      elementChunks.push(currentChunk)
+      currentChunk = []
+      currentTokens = 0
+    }
+
+    currentChunk.push(elem.index)
+    currentTokens += elemTokens
+  }
+
+  if (currentChunk.length > 0) {
+    elementChunks.push(currentChunk)
+  }
+
+  // Helper to chunk groups
+  const getGroupChunks = (): string[][] => {
+    const allGroupIds = Array.from(groups.keys())
+    if (allGroupIds.length === 0) return [[]]
+
+    const chunks: string[][] = []
+    let currentChunk: string[] = []
+    let currentTokens = 0
+
+    for (const groupId of allGroupIds) {
+      const group = groups.get(groupId)!
+      const groupTokens = tokenizer.count(`${group.label}`) + 10
+
+      if (currentTokens + groupTokens > TOKENS_FOR_GROUPS_MAX && currentChunk.length > 0) {
+        chunks.push(currentChunk)
+        currentChunk = []
+        currentTokens = 0
+      }
+
+      currentChunk.push(groupId)
+      currentTokens += groupTokens
+    }
+
+    if (currentChunk.length > 0) {
+      chunks.push(currentChunk)
+    }
+
+    return chunks.length > 0 ? chunks : [[]]
+  }
+
+  // Process elements against groups and get assignments
+  const processChunk = async (
+    elementIndices: number[],
+    groupIds: string[]
+  ): Promise<Array<{ elementIndex: number; label: string }>> => {
+    // Get examples from adapter for active learning
+    const chunkElements = elementIndices.map((idx) => elements[idx].element)
+    const chunkInputStr = JSON.stringify(chunkElements)
+
+    const examples =
+      taskId && ctx.adapter
+        ? await ctx.adapter.getExamples<string, Array<{ elementIndex: number; label: string }>>({
+            input: chunkInputStr.slice(0, 1000), // Limit search string length
+            taskType,
+            taskId,
+          })
+        : []
+
+    // Check for exact match (cache hit)
+
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: chunkInputStr,
+        instructions: options.instructions ?? '',
+        groupIds: groupIds.join(','),
+      })
+    )
+
+    const exactMatch = examples.find((x) => x.key === key)
+    if (exactMatch && exactMatch.output) {
+      return exactMatch.output
+    }
+
+    const elementsText = elementIndices
+      .map((idx, i) => {
+        const elem = elements[idx]
+        const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement)
+        return `■${i}: ${truncated}■`
+      })
+      .join('\n')
+
+    const groupsList = groupIds.map((gid) => groups.get(gid)!.label)
+    const groupsText =
+      groupsList.length > 0
+        ? `**Existing Groups (prefer reusing these):**\n${groupsList.map((l) => `- ${l}`).join('\n')}\n\n`
+        : ''
+
+    // Format examples for few-shot learning
+    const exampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []
+
+    for (const example of examples.slice(0, 5)) {
+      try {
+        const exampleInput = JSON.parse(example.input)
+        const exampleElements = Array.isArray(exampleInput) ? exampleInput : [exampleInput]
+
+        // User message
+        const exampleElementsText = exampleElements
+          .map((el, i) => `■${i}: ${stringify(el, false).slice(0, 200)}■`)
+          .join('\n')
+
+        exampleMessages.push({
+          type: 'text',
+          role: 'user',
+          content: `Expert Example - Elements to group:
+${exampleElementsText}
+
+Group each element.`,
+        })
+
+        // Assistant message
+        const exampleOutput = example.output
+        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
+          const formattedAssignments = exampleOutput
+            .map((assignment) => `■${assignment.elementIndex}:${assignment.label}■`)
+            .join('\n')
+
+          exampleMessages.push({
+            type: 'text',
+            role: 'assistant',
+            content: `${formattedAssignments}\n${END}`,
+          })
+
+          if (example.explanation) {
+            exampleMessages.push({
+              type: 'text',
+              role: 'assistant',
+              content: `Reasoning: ${example.explanation}`,
+            })
+          }
+        }
+      } catch {
+        // Skip malformed examples
+      }
+    }
+
+    const systemPrompt = `You are grouping elements into cohesive groups.
+
+${options.instructions ? `**Instructions:** ${options.instructions}\n` : '**Instructions:** Group similar elements together.'}
+
+**Important:**
+- Each element gets exactly ONE group label
+- Use EXACT SAME label for similar items (case-sensitive)
+- Create new descriptive labels when needed
+
+**Output Format:**
+One line per element:
+■0:Group Label■
+■1:Group Label■
+${END}`.trim()
+
+    const userPrompt = `${groupsText}**Elements (■0 to ■${elementIndices.length - 1}):**
+${elementsText}
+
+**Task:** For each element, output one line with its group label.
+${END}`.trim()
+
+    const { extracted } = await ctx.generateContent({
+      systemPrompt,
+      stopSequences: [END],
+      messages: [...exampleMessages, { type: 'text', role: 'user', content: userPrompt }],
+      transform: (text) => {
+        const assignments: Array<{ elementIndex: number; label: string }> = []
+        const regex = /■(\d+):([^■]+)■/g
+        let match: RegExpExecArray | null
+
+        while ((match = regex.exec(text)) !== null) {
+          const idx = parseInt(match[1] ?? '', 10)
+          if (isNaN(idx) || idx < 0 || idx >= elementIndices.length) continue
+
+          const label = (match[2] ?? '').trim()
+          if (!label) continue
+
+          assignments.push({
+            elementIndex: elementIndices[idx],
+            label: label.slice(0, 250),
+          })
+        }
+
+        return assignments
+      },
+    })
+
+    return extracted
+  }
+
+  // Phase 1: Process all element chunks against current groups IN PARALLEL
+  const elementLimit = pLimit(10) // Separate limiter for element chunks
+  const groupLimit = pLimit(10) // Separate limiter for group chunks
+
+  // Collect all assignments from parallel processing
+  const allChunkResults = await Promise.all(
+    elementChunks.map((elementChunk) =>
+      elementLimit(async () => {
+        const groupChunks = getGroupChunks()
+
+        const allAssignments = await Promise.all(
+          groupChunks.map((groupChunk) => groupLimit(() => processChunk(elementChunk, groupChunk)))
+        )
+
+        return allAssignments.flat()
+      })
+    )
+  )
+
+  // Process all assignments sequentially to avoid race conditions
+  for (const assignments of allChunkResults) {
+    for (const { elementIndex, label } of assignments) {
+      const normalized = normalizeLabel(label)
+      let groupId = labelToGroupId.get(normalized)
+
+      if (!groupId) {
+        // Create new group
+        groupId = `group_${groupIdCounter++}`
+        groups.set(groupId, { id: groupId, label, normalizedLabel: normalized })
+        groupElements.set(groupId, new Set())
+        labelToGroupId.set(normalized, groupId)
+      }
+
+      // Add element to group
+      groupElements.get(groupId)!.add(elementIndex)
+
+      // Track that element saw this group
+      if (!elementGroups.has(elementIndex)) {
+        elementGroups.set(elementIndex, new Set())
+      }
+      elementGroups.get(elementIndex)!.add(groupId)
+    }
+  }
+
+  // Phase 2: Ensure all elements saw all groups (coverage guarantee)
+  const allGroupIds = Array.from(groups.keys())
+
+  if (allGroupIds.length > 0) {
+    const elementsNeedingReview: number[] = []
+
+    for (const elem of elements) {
+      const seenGroups = elementGroups.get(elem.index) ?? new Set()
+      const unseenCount = allGroupIds.filter((gid) => !seenGroups.has(gid)).length
+
+      if (unseenCount > 0) {
+        elementsNeedingReview.push(elem.index)
+      }
+    }
+
+    if (elementsNeedingReview.length > 0) {
+      // Chunk elements needing review
+      const reviewChunks: number[][] = []
+      let reviewChunk: number[] = []
+      let reviewTokens = 0
+
+      for (const elemIdx of elementsNeedingReview) {
+        const elem = elements[elemIdx]
+        const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement)
+        const elemTokens = tokenizer.count(truncated)
+
+        const shouldStartNewChunk =
+          (reviewTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX || reviewChunk.length >= MAX_ELEMENTS_PER_CHUNK) &&
+          reviewChunk.length > 0
+
+        if (shouldStartNewChunk) {
+          reviewChunks.push(reviewChunk)
+          reviewChunk = []
+          reviewTokens = 0
+        }
+
+        reviewChunk.push(elemIdx)
+        reviewTokens += elemTokens
+      }
+
+      if (reviewChunk.length > 0) {
+        reviewChunks.push(reviewChunk)
+      }
+
+      // Process review chunks IN PARALLEL
+      const reviewResults = await Promise.all(
+        reviewChunks.map((chunk) =>
+          elementLimit(async () => {
+            const groupChunks = getGroupChunks()
+
+            const allAssignments = await Promise.all(
+              groupChunks.map((groupChunk) => groupLimit(() => processChunk(chunk, groupChunk)))
+            )
+
+            return allAssignments.flat()
+          })
+        )
+      )
+
+      // Mark groups as seen and update assignments (sequential to avoid races)
+      const updateElementGroupAssignment = (elementIndex: number, label: string) => {
+        const normalized = normalizeLabel(label)
+        const groupId = labelToGroupId.get(normalized)
+        if (!groupId) return
+
+        // Add to group and mark as seen
+        groupElements.get(groupId)!.add(elementIndex)
+
+        // Initialize element groups if needed
+        const elemGroups = elementGroups.get(elementIndex) ?? new Set()
+        if (!elementGroups.has(elementIndex)) {
+          elementGroups.set(elementIndex, elemGroups)
+        }
+        elemGroups.add(groupId)
+      }
+
+      for (const assignments of reviewResults) {
+        for (const { elementIndex, label } of assignments) {
+          updateElementGroupAssignment(elementIndex, label)
+        }
+      }
+    }
+  }
+
+  // Phase 3: Resolve conflicts (elements in multiple groups)
+  for (const [elementIndex, groupSet] of elementGroups.entries()) {
+    if (groupSet.size > 1) {
+      // Element is in multiple groups, keep only the most common assignment
+      const groupIds = Array.from(groupSet)
+
+      // Remove from all groups
+      for (const gid of groupIds) {
+        groupElements.get(gid)?.delete(elementIndex)
+      }
+
+      // Re-assign to first group (or could use LLM to decide)
+      const finalGroupId = groupIds[0]
+      groupElements.get(finalGroupId)!.add(elementIndex)
+    }
+  }
+
+  // Phase 4: Merge groups if maxGroups is set (AI-driven)
+  if (options.maxGroups !== undefined) {
+    const nonEmptyGroupIds = () =>
+      Array.from(groupElements.entries())
+        .filter(([, s]) => s.size > 0)
+        .map(([id]) => id)
+
+    let currentIds = nonEmptyGroupIds()
+
+    if (currentIds.length > options.maxGroups) {
+      // Build a summary of each group: label + element count + sample elements
+      const groupSummaries = currentIds.map((gid, idx) => {
+        const info = groups.get(gid)!
+        const elemIndices = Array.from(groupElements.get(gid)!)
+        const sampleElements = elemIndices
+          .slice(0, 3)
+          .map((i) => tokenizer.truncate(elements[i].stringified, 60))
+          .join(', ')
+        return `■${idx}:${info.label} (${elemIndices.length} elements, e.g. ${sampleElements})■`
+      })
+
+      const mergeSystemPrompt = `You are consolidating groups into fewer, broader categories.
+
+${options.instructions ? `**Original instructions:** ${options.instructions}\n` : ''}
+**Task:** Merge ${currentIds.length} groups down to at most ${options.maxGroups} groups.
+Combine the most semantically related groups together. Give each merged group a new descriptive label.
+
+**Output Format:**
+For each input group (■0 to ■${currentIds.length - 1}), output which target label it maps to:
+■0:Merged Label■
+■1:Merged Label■
+${END}
+
+Use the EXACT SAME label for groups that should be merged together.`.trim()
+
+      const mergeUserPrompt = `**Current groups:**
+${groupSummaries.join('\n')}
+
+Merge into at most ${options.maxGroups} groups.
+${END}`.trim()
+
+      const { extracted: mergeAssignments } = await ctx.generateContent({
+        systemPrompt: mergeSystemPrompt,
+        stopSequences: [END],
+        messages: [{ type: 'text', role: 'user', content: mergeUserPrompt }],
+        transform: (text) => {
+          const assignments: Array<{ sourceIdx: number; label: string }> = []
+          const regex = /■(\d+):([^■]+)■/g
+          let match: RegExpExecArray | null
+
+          while ((match = regex.exec(text)) !== null) {
+            const idx = parseInt(match[1] ?? '', 10)
+            if (isNaN(idx) || idx < 0 || idx >= currentIds.length) continue
+
+            const label = (match[2] ?? '').trim()
+            if (!label) continue
+
+            assignments.push({ sourceIdx: idx, label: label.slice(0, 250) })
+          }
+
+          return assignments
+        },
+      })
+
+      // Build merge map: normalized merge label → list of source group IDs
+      const mergeMap = new Map<string, { label: string; sourceGroupIds: string[] }>()
+
+      for (const { sourceIdx, label } of mergeAssignments) {
+        const sourceGid = currentIds[sourceIdx]
+        if (!sourceGid) continue
+
+        const normalized = normalizeLabel(label)
+        if (!mergeMap.has(normalized)) {
+          mergeMap.set(normalized, { label, sourceGroupIds: [] })
+        }
+        mergeMap.get(normalized)!.sourceGroupIds.push(sourceGid)
+      }
+
+      // Apply merges: for each merge target, pick the first source group as the target
+      // and move all elements from other source groups into it
+      for (const [, { label, sourceGroupIds }] of mergeMap) {
+        if (sourceGroupIds.length <= 1) continue
+
+        const targetGid = sourceGroupIds[0]
+        const targetSet = groupElements.get(targetGid)!
+
+        // Update label on the target group
+        const targetInfo = groups.get(targetGid)!
+        targetInfo.label = label
+        targetInfo.normalizedLabel = normalizeLabel(label)
+
+        for (let i = 1; i < sourceGroupIds.length; i++) {
+          const sourceGid = sourceGroupIds[i]
+          const sourceSet = groupElements.get(sourceGid)!
+          sourceSet.forEach((elemIdx) => targetSet.add(elemIdx))
+          sourceSet.clear()
+        }
+      }
+
+      // Safety: if LLM still produced too many groups, fall back to merging smallest pairs
+      currentIds = nonEmptyGroupIds()
+      while (currentIds.length > options.maxGroups) {
+        currentIds.sort((a, b) => groupElements.get(a)!.size - groupElements.get(b)!.size)
+
+        const sourceSet = groupElements.get(currentIds[0])!
+        const targetSet = groupElements.get(currentIds[1])!
+        for (const elemIdx of sourceSet) {
+          targetSet.add(elemIdx)
+        }
+        sourceSet.clear()
+
+        currentIds = nonEmptyGroupIds()
+      }
+    }
+  }
+
+  // Phase 5: Redistribute undersized groups if minElements is set
+  // Reuses processChunk so orphans see the valid groups as available buckets
+  if (options.minElements !== undefined && options.minElements > 1) {
+    const getNonEmptyGroupIds = () =>
+      Array.from(groupElements.entries())
+        .filter(([, s]) => s.size > 0)
+        .map(([id]) => id)
+
+    // Collect orphan elements from all undersized groups
+    const orphanIndices: number[] = []
+
+    for (const gid of getNonEmptyGroupIds()) {
+      const elemSet = groupElements.get(gid)!
+      if (elemSet.size > 0 && elemSet.size < options.minElements) {
+        for (const idx of elemSet) {
+          orphanIndices.push(idx)
+        }
+        elemSet.clear()
+      }
+    }
+
+    if (orphanIndices.length > 0) {
+      // Valid groups = everything that's still non-empty (i.e. above minElements)
+      const validGroupIds = getNonEmptyGroupIds()
+
+      // Chunk orphans and run them through processChunk with only valid groups visible
+      const orphanChunks: number[][] = []
+      let currentOrphanChunk: number[] = []
+      let currentOrphanTokens = 0
+
+      for (const elemIdx of orphanIndices) {
+        const elem = elements[elemIdx]
+        const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement)
+        const elemTokens = tokenizer.count(truncated)
+
+        if (
+          (currentOrphanTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX ||
+            currentOrphanChunk.length >= MAX_ELEMENTS_PER_CHUNK) &&
+          currentOrphanChunk.length > 0
+        ) {
+          orphanChunks.push(currentOrphanChunk)
+          currentOrphanChunk = []
+          currentOrphanTokens = 0
+        }
+
+        currentOrphanChunk.push(elemIdx)
+        currentOrphanTokens += elemTokens
+      }
+
+      if (currentOrphanChunk.length > 0) {
+        orphanChunks.push(currentOrphanChunk)
+      }
+
+      // Process orphan chunks against valid groups (reuses the same processChunk as Phase 1)
+      const orphanResults = await Promise.all(
+        orphanChunks.map((chunk) =>
+          elementLimit(async () => {
+            // If there are valid groups, chunk them; otherwise pass empty so LLM creates new groups
+            const groupChunksForOrphans = validGroupIds.length > 0 ? getGroupChunks() : [[]]
+
+            const allAssignments = await Promise.all(
+              groupChunksForOrphans
+                .filter((gc) => gc.length === 0 || gc.some((gid) => validGroupIds.includes(gid)))
+                .map((groupChunk) => {
+                  // Only show valid groups (exclude the orphaned/undersized ones)
+                  const filteredGroupChunk = groupChunk.filter((gid) => validGroupIds.includes(gid))
+                  return groupLimit(() => processChunk(chunk, filteredGroupChunk))
+                })
+            )
+
+            return allAssignments.flat()
+          })
+        )
+      )
+
+      // Apply assignments
+      const flatAssignments = orphanResults.flat()
+      for (const { elementIndex, label } of flatAssignments) {
+        const normalized = normalizeLabel(label)
+        let groupId = labelToGroupId.get(normalized)
+
+        if (!groupId) {
+          groupId = `group_${groupIdCounter++}`
+          groups.set(groupId, { id: groupId, label, normalizedLabel: normalized })
+          groupElements.set(groupId, new Set())
+          labelToGroupId.set(normalized, groupId)
+        }
+        groupElements.get(groupId)!.add(elementIndex)
+      }
+
+      // Safety: any orphans the LLM missed get placed into the largest group
+      const isAssigned = (idx: number) => {
+        for (const [, elemSet] of groupElements) {
+          if (elemSet.has(idx)) return true
+        }
+        return false
+      }
+      const unassigned = orphanIndices.filter((idx) => !isAssigned(idx))
+      const placeIntoLargest = (indices: number[]) => {
+        const allNonEmpty = getNonEmptyGroupIds()
+        if (allNonEmpty.length === 0) return
+        const largestGid = allNonEmpty.reduce((a, b) =>
+          groupElements.get(a)!.size >= groupElements.get(b)!.size ? a : b
+        )
+        for (const idx of indices) {
+          groupElements.get(largestGid)!.add(idx)
+        }
+      }
+
+      if (unassigned.length > 0) {
+        placeIntoLargest(unassigned)
+      }
+
+      // Second pass: if any groups are still undersized after redistribution,
+      // merge their elements into the largest group
+      const mergeUndersizedGroups = () => {
+        const allNonEmpty = getNonEmptyGroupIds()
+        if (allNonEmpty.length <= 1) return false
+
+        const largestGid = allNonEmpty.reduce((a, b) =>
+          groupElements.get(a)!.size >= groupElements.get(b)!.size ? a : b
+        )
+        const targetSet = groupElements.get(largestGid)!
+        let merged = false
+
+        for (const gid of allNonEmpty) {
+          if (gid === largestGid) continue
+          const elemSet = groupElements.get(gid)!
+          if (elemSet.size > 0 && elemSet.size < options.minElements) {
+            elemSet.forEach((idx) => targetSet.add(idx))
+            elemSet.clear()
+            merged = true
+          }
+        }
+        return merged
+      }
+
+      while (mergeUndersizedGroups()) {
+        // keep merging until no undersized groups remain
+      }
+    }
+  }
+
+  // Build final result
+  const result: Array<Group<T>> = []
+
+  for (const [groupId, elementIndices] of groupElements.entries()) {
+    if (elementIndices.size > 0) {
+      const groupInfo = groups.get(groupId)!
+      result.push({
+        id: groupInfo.id,
+        label: groupInfo.label,
+        elements: Array.from(elementIndices).map((idx) => elements[idx].element),
+      })
+    }
+  }
+
+  // Save example for active learning
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: JSON.stringify(input),
+        instructions: options.instructions ?? '',
+      })
+    )
+
+    // Build output format for saving
+    const outputAssignments: Array<{ elementIndex: number; label: string }> = []
+    for (const [groupId, elementIndices] of groupElements.entries()) {
+      const groupInfo = groups.get(groupId)!
+      for (const idx of elementIndices) {
+        outputAssignments.push({
+          elementIndex: idx,
+          label: groupInfo.label,
+        })
+      }
+    }
+
+    // Note: We don't have direct access to usage metadata here since it's distributed
+    // across many parallel operations. We'll use default values.
+    await ctx.adapter.saveExample({
+      key,
+      taskType,
+      taskId,
+      input: JSON.stringify(input),
+      output: result,
+      instructions: options.instructions ?? '',
+      metadata: {
+        cost: { input: 0, output: 0 },
+        latency: 0,
+        model: ctx.modelId,
+        tokens: { input: 0, output: 0 },
+      },
+    })
+  }
+
+  return result
+}
+
+Zai.prototype.group = function <T>(
+  this: Zai,
+  input: Array<T>,
+  _options?: Options
+): Response<Array<Group<T>>, Record<string, T[]>> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.group',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<Array<Group<T>>, Record<string, T[]>>(context, group(input, _options, context), (result) => {
+    const merged: Record<string, T[]> = {}
+    result.forEach((group) => {
+      if (!merged[group.label]) {
+        merged[group.label] = []
+      }
+      merged[group.label].push(...group.elements)
+    })
+    return merged
+  })
+}

packages/zai/src/operations/label.ts

@@ -0,0 +1,566 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+
+import { chunk, clamp } from 'lodash-es'
+import pLimit from 'p-limit'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER } from './constants'
+
+type Label = keyof typeof LABELS
+const LABELS = {
+  ABSOLUTELY_NOT: 'ABSOLUTELY_NOT',
+  PROBABLY_NOT: 'PROBABLY_NOT',
+  AMBIGUOUS: 'AMBIGUOUS',
+  PROBABLY_YES: 'PROBABLY_YES',
+  ABSOLUTELY_YES: 'ABSOLUTELY_YES',
+} as const
+
+const ALL_LABELS = Object.values(LABELS).join(' | ')
+
+type Example<T extends string> = {
+  input: unknown
+  labels: Partial<Record<T, { label: Label; explanation?: string }>>
+}
+
+export type Options<T extends string> = {
+  /** Examples to help the user make a decision */
+  examples?: Array<Example<T>>
+  /** Instructions to guide the user on how to extract the data */
+  instructions?: string
+  /** The maximum number of tokens per chunk */
+  chunkLength?: number
+}
+
+const _Options = z.object({
+  examples: z
+    .array(
+      z.object({
+        input: z.any(),
+        labels: z.record(z.object({ label: z.enum(ALL_LABELS as never), explanation: z.string().optional() })),
+      })
+    )
+    .default([])
+    .describe('Examples to help the user make a decision'),
+  instructions: z.string().optional().describe('Instructions to guide the user on how to extract the data'),
+  chunkLength: z
+    .number()
+    .min(100)
+    .max(100_000)
+    .optional()
+    .describe('The maximum number of tokens per chunk')
+    .default(16_000),
+})
+
+type Labels<T extends string> = Record<T, string>
+
+const _Labels = z.record(z.string().min(1).max(250), z.string()).superRefine((labels, ctx) => {
+  const keys = Object.keys(labels)
+
+  for (const key of keys) {
+    if (key.length < 1 || key.length > 250) {
+      ctx.addIssue({ message: `The label key "${key}" must be between 1 and 250 characters long`, code: 'custom' })
+    }
+
+    if (keys.lastIndexOf(key) !== keys.indexOf(key)) {
+      ctx.addIssue({ message: `Duplicate label: ${labels[key]}`, code: 'custom' })
+    }
+
+    if (/[^a-zA-Z0-9_]/.test(key)) {
+      ctx.addIssue({
+        message: `The label key "${key}" must only contain alphanumeric characters and underscores`,
+        code: 'custom',
+      })
+    }
+  }
+
+  return true
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Tags input with multiple labels, each with explanation, value, and confidence level.
+     *
+     * This operation evaluates input against multiple categories simultaneously, providing
+     * nuanced confidence levels (ABSOLUTELY_YES, PROBABLY_YES, AMBIGUOUS, PROBABLY_NOT, ABSOLUTELY_NOT).
+     * Perfect for content classification, intent detection, and multi-criteria evaluation.
+     *
+     * @param input - The data to label (text, object, or any value)
+     * @param labels - Object mapping label keys to their descriptions
+     * @param options - Configuration for examples, instructions, and chunking
+     * @returns Response with detailed results per label (explanation, value, confidence), simplified to booleans
+     *
+     * @example Content moderation
+     * ```typescript
+     * const comment = "This is a great article! Click here for cheap products!"
+     *
+     * const result = await zai.label(comment, {
+     *   spam: 'Is this spam or promotional content?',
+     *   helpful: 'Is this comment helpful and constructive?',
+     *   profanity: 'Does this contain profanity or offensive language?'
+     * }).result()
+     *
+     * // result.output:
+     * // {
+     * //   spam: { value: true, confidence: 0.5, explanation: 'Contains promotional link...' },
+     * //   helpful: { value: true, confidence: 1, explanation: 'Positive feedback...' },
+     * //   profanity: { value: false, confidence: 1, explanation: 'No offensive language' }
+     * // }
+     *
+     * // Simplified (when awaited directly):
+     * const simple = await zai.label(comment, { spam: 'Is this spam?' })
+     * // simple: { spam: true }
+     * ```
+     *
+     * @example Intent classification
+     * ```typescript
+     * const message = "I can't log in to my account"
+     *
+     * const intents = await zai.label(message, {
+     *   technical_issue: 'Is this a technical problem?',
+     *   urgent: 'Does this require immediate attention?',
+     *   needs_human: 'Should this be escalated to a human agent?',
+     *   billing_related: 'Is this about billing or payments?'
+     * })
+     * // Returns boolean for each intent
+     * ```
+     *
+     * @example Sentiment analysis with nuance
+     * ```typescript
+     * const review = "The product is okay, but shipping was slow"
+     *
+     * const { output } = await zai.label(review, {
+     *   positive: 'Is the overall sentiment positive?',
+     *   negative: 'Is the overall sentiment negative?',
+     *   mixed: 'Does this express mixed feelings?'
+     * }).result()
+     *
+     * // Check confidence levels
+     * if (output.mixed.confidence > 0.5) {
+     *   console.log('Mixed sentiment detected:', output.mixed.explanation)
+     * }
+     * ```
+     *
+     * @example Product categorization
+     * ```typescript
+     * const product = {
+     *   name: 'Wireless Headphones',
+     *   description: 'Noise-canceling Bluetooth headphones for music lovers',
+     *   price: 199
+     * }
+     *
+     * const categories = await zai.label(product, {
+     *   electronics: 'Is this an electronic device?',
+     *   audio: 'Is this audio equipment?',
+     *   premium: 'Is this a premium/luxury product?',
+     *   portable: 'Is this portable/mobile?'
+     * })
+     * // Returns: { electronics: true, audio: true, premium: true, portable: true }
+     * ```
+     *
+     * @example With examples for consistency
+     * ```typescript
+     * const result = await zai.label(email, {
+     *   urgent: 'Requires immediate response',
+     *   complaint: 'Customer is dissatisfied'
+     * }, {
+     *   examples: [
+     *     {
+     *       input: 'ASAP! System is down!',
+     *       labels: {
+     *         urgent: { label: 'ABSOLUTELY_YES', explanation: 'Critical issue' },
+     *         complaint: { label: 'PROBABLY_YES', explanation: 'Implicit complaint' }
+     *       }
+     *     }
+     *   ],
+     *   instructions: 'Consider tone, urgency keywords, and context'
+     * })
+     * ```
+     *
+     * @example Understanding confidence levels
+     * ```typescript
+     * const { output } = await zai.label(text, labels).result()
+     *
+     * // Confidence values:
+     * // ABSOLUTELY_YES: confidence = 1.0, value = true
+     * // PROBABLY_YES: confidence = 0.5, value = true
+     * // AMBIGUOUS: confidence = 0, value = false
+     * // PROBABLY_NOT: confidence = 0.5, value = false
+     * // ABSOLUTELY_NOT: confidence = 1.0, value = false
+     *
+     * Object.entries(output).forEach(([key, result]) => {
+     *   if (result.value && result.confidence === 1) {
+     *     console.log(`Definitely ${key}:`, result.explanation)
+     *   } else if (result.value && result.confidence === 0.5) {
+     *     console.log(`Probably ${key}:`, result.explanation)
+     *   }
+     * })
+     * ```
+     */
+    label<T extends string>(
+      input: unknown,
+      labels: Labels<T>,
+      options?: Options<T>
+    ): Response<
+      {
+        [K in T]: {
+          explanation: string
+          value: boolean
+          confidence: number
+        }
+      },
+      { [K in T]: boolean }
+    >
+  }
+}
+
+const parseLabel = (label: string): Label => {
+  label = label.toUpperCase().replace(/\s+/g, '_').replace(/_{2,}/g, '_').trim()
+  if (label.includes('ABSOLUTELY') && label.includes('NOT')) {
+    return LABELS.ABSOLUTELY_NOT
+  } else if (label.includes('NOT')) {
+    return LABELS.PROBABLY_NOT
+  } else if (label.includes('AMBIGUOUS')) {
+    return LABELS.AMBIGUOUS
+  }
+  if (label.includes('YES')) {
+    return LABELS.PROBABLY_YES
+  } else if (label.includes('ABSOLUTELY') && label.includes('YES')) {
+    return LABELS.ABSOLUTELY_YES
+  }
+  return LABELS.AMBIGUOUS
+}
+
+const getConfidence = (label: Label) => {
+  switch (label) {
+    case LABELS.ABSOLUTELY_NOT:
+    case LABELS.ABSOLUTELY_YES:
+      return 1
+
+    case LABELS.PROBABLY_NOT:
+    case LABELS.PROBABLY_YES:
+      return 0.5
+    default:
+      return 0
+  }
+}
+
+const label = async <T extends string>(
+  input: unknown,
+  _labels: Labels<T>,
+  _options: Options<T> | undefined,
+  ctx: ZaiContext
+): Promise<{
+  [K in T]: {
+    explanation: string
+    value: boolean
+    confidence: number
+  }
+}> => {
+  ctx.controller.signal.throwIfAborted()
+  const options = _Options.parse(_options ?? {}) as unknown as Options<T>
+  const labels = _Labels.parse(_labels) as Labels<T>
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.label'
+
+  const TOTAL_MAX_TOKENS = clamp(options.chunkLength, 1000, model.input.maxTokens - PROMPT_INPUT_BUFFER)
+  const CHUNK_EXAMPLES_MAX_TOKENS = clamp(Math.floor(TOTAL_MAX_TOKENS * 0.5), 250, 10_000)
+  const CHUNK_INPUT_MAX_TOKENS = clamp(
+    TOTAL_MAX_TOKENS - CHUNK_EXAMPLES_MAX_TOKENS,
+    TOTAL_MAX_TOKENS * 0.5,
+    TOTAL_MAX_TOKENS
+  )
+
+  const inputAsString = stringify(input)
+
+  if (tokenizer.count(inputAsString) > CHUNK_INPUT_MAX_TOKENS) {
+    const limit = pLimit(10) // Limit to 10 concurrent labeling operations
+    const tokens = tokenizer.split(inputAsString)
+    const chunks = chunk(tokens, CHUNK_INPUT_MAX_TOKENS).map((x) => x.join(''))
+    const allLabels = await Promise.all(chunks.map((chunk) => limit(() => label(chunk, _labels, _options, ctx))))
+
+    // Merge all the labels together (those who are true will remain true)
+    return allLabels.reduce((acc, x) => {
+      Object.keys(x).forEach((key) => {
+        if (acc[key]?.value === true) {
+          acc[key] = acc[key]
+        } else if (x[key]?.value === true) {
+          acc[key] = x[key]
+        } else {
+          acc[key] = acc[key] || x[key]
+        }
+      })
+      return acc
+    }, {}) as {
+      [K in T]: {
+        explanation: string
+        value: boolean
+        confidence: number
+      }
+    }
+  }
+
+  const END = '■END■'
+
+  const Key = fastHash(
+    JSON.stringify({
+      taskType,
+      taskId,
+      input: inputAsString,
+      instructions: options.instructions ?? '',
+    })
+  )
+
+  const convertToAnswer = (mapping: { [K in T]: { explanation: string; label: Label } }) => {
+    return Object.keys(labels).reduce((acc, key) => {
+      acc[key] = {
+        explanation: mapping[key]?.explanation ?? '',
+        value: mapping[key]?.label === LABELS.ABSOLUTELY_YES || mapping[key]?.label === LABELS.PROBABLY_YES,
+        confidence: getConfidence(mapping[key]?.label),
+      }
+      return acc
+    }, {}) as {
+      [K in T]: {
+        explanation: string
+        value: boolean
+        confidence: number
+      }
+    }
+  }
+
+  const examples =
+    taskId && ctx.adapter
+      ? await ctx.adapter.getExamples<
+          string,
+          {
+            [K in T]: {
+              explanation: string
+              label: Label
+            }
+          }
+        >({
+          input: inputAsString,
+          taskType,
+          taskId,
+        })
+      : []
+
+  options.examples.forEach((example) => {
+    examples.push({
+      key: fastHash(JSON.stringify(example)),
+      input: stringify(example.input),
+      similarity: 1,
+      explanation: '',
+      output: example.labels as unknown as {
+        [K in T]: {
+          explanation: string
+          label: Label
+        }
+      },
+    })
+  })
+
+  const exactMatch = examples.find((x) => x.key === Key)
+  if (exactMatch) {
+    return convertToAnswer(exactMatch.output)
+  }
+
+  const allExamples = takeUntilTokens(
+    examples,
+    CHUNK_EXAMPLES_MAX_TOKENS,
+    (el) =>
+      tokenizer.count(stringify(el.input)) +
+      tokenizer.count(stringify(el.output)) +
+      tokenizer.count(el.explanation ?? '') +
+      100
+  )
+    .map((example, idx) => [
+      {
+        type: 'text' as const,
+        role: 'user' as const,
+        content: `
+Expert Example #${idx + 1}
+
+<|start_input|>
+${stringify(example.input)}
+<|end_input|>`.trim(),
+      },
+      {
+        type: 'text' as const,
+        role: 'assistant' as const,
+        content: `
+Expert Example #${idx + 1}
+============
+${Object.keys(example.output)
+  .map((key) =>
+    `
+■${key}:【${example.output[key]?.explanation}】:${example.output[key]?.label}■
+`.trim()
+  )
+  .join('\n')}
+${END}
+`.trim(),
+      },
+    ])
+    .flat()
+
+  const format = Object.keys(labels)
+    .map((key) => {
+      return `
+■${key}:【explanation (where "explanation" is answering the question "${labels[key]}")】:x■ (where x is ${ALL_LABELS})
+`.trim()
+    })
+    .join('\n\n')
+
+  const { extracted, meta } = await ctx.generateContent({
+    stopSequences: [END],
+    systemPrompt: `
+You need to tag the input with the following labels based on the question asked:
+${LABELS.ABSOLUTELY_NOT}: You are absolutely sure that the answer is "NO" to the question.
+${LABELS.PROBABLY_NOT}: You are leaning towards "NO" to the question.
+${LABELS.AMBIGUOUS}: You are unsure about the answer to the question.
+${LABELS.PROBABLY_YES}: You are leaning towards "YES" to the question.
+${LABELS.ABSOLUTELY_YES}: You are absolutely sure that the answer is "YES" to the question.
+
+You need to return a mapping of the labels, an explanation and the answer for each label following the format below:
+\`\`\`
+${format}
+${END}
+\`\`\`
+
+${options.instructions}
+
+===
+You should consider the Expert Examples below to help you make your decision.
+In your "Analysis", please refer to the Expert Examples # to justify your decision.
+`.trim(),
+    messages: [
+      ...allExamples,
+      {
+        type: 'text',
+        role: 'user',
+        content: `
+Input to tag:
+<|start_input|>
+${inputAsString}
+<|end_input|>
+
+Answer with this following format:
+\`\`\`
+${format}
+${END}
+\`\`\`
+
+Format cheatsheet:
+\`\`\`
+■label:【explanation】:x■
+\`\`\`
+
+Where \`x\` is one of the following: ${ALL_LABELS}
+
+Remember: In your \`explanation\`, please refer to the Expert Examples # (and quote them) that are relevant to ground your decision-making process.
+The Expert Examples are there to help you make your decision. They have been provided by experts in the field and their answers (and reasoning) are considered the ground truth and should be used as a reference to make your decision when applicable.
+For example, you can say: "According to Expert Example #1, ..."`.trim(),
+      },
+    ],
+    transform: (text) =>
+      Object.keys(labels).reduce((acc, key) => {
+        const match = text.match(new RegExp(`■${key}:【(.+)】:(\\w{2,})■`, 'i'))
+        if (match) {
+          const explanation = match[1].trim()
+          const label = parseLabel(match[2])
+          acc[key] = {
+            explanation,
+            label,
+          }
+        } else {
+          acc[key] = {
+            explanation: '',
+            label: LABELS.AMBIGUOUS,
+          }
+        }
+        return acc
+      }, {}) as {
+        [K in T]: {
+          explanation: string
+          label: Label
+        }
+      },
+  })
+
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
+      key: Key,
+      taskType,
+      taskId,
+      instructions: options.instructions ?? '',
+      metadata: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output,
+        },
+        latency: meta.latency,
+        model: ctx.modelId,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output,
+        },
+      },
+      input: inputAsString,
+      output: extracted,
+    })
+  }
+
+  return convertToAnswer(extracted)
+}
+
+Zai.prototype.label = function <T extends string>(
+  this: Zai,
+  input: unknown,
+  labels: Labels<T>,
+  _options?: Options<T>
+): Response<
+  {
+    [K in T]: {
+      explanation: string
+      value: boolean
+      confidence: number
+    }
+  },
+  { [K in T]: boolean }
+> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.label',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<
+    {
+      [K in T]: {
+        explanation: string
+        value: boolean
+        confidence: number
+      }
+    },
+    { [K in T]: boolean }
+  >(context, label(input, labels, _options, context), (result) =>
+    Object.keys(result).reduce(
+      (acc, key) => {
+        acc[key] = result[key].value
+        return acc
+      },
+      {} as { [K in T]: boolean }
+    )
+  )
+}

packages/zai/src/operations/patch.ts

@@ -0,0 +1,832 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+import pLimit from 'p-limit'
+
+import { ZaiContext } from '../context'
+import { isJsonFile, validateOrRepairJson } from '../json-utils'
+import { Micropatch } from '../micropatch'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+import { JsonParsingError } from './errors'
+
+/**
+ * Represents a file to be patched
+ */
+export type File = {
+  /** The file path (e.g., 'src/components/Button.tsx') */
+  path: string
+  /** The file name (e.g., 'Button.tsx') */
+  name: string
+  /** The file content */
+  content: string
+  /** The patch operations that were applied (only present in output) */
+  patch?: string
+}
+
+const _File = z.object({
+  path: z.string(),
+  name: z.string(),
+  content: z.string(),
+})
+
+type __Z<T> = { __type__: 'ZuiType'; _output: T }
+type OfType<O, T extends __Z<any> = __Z<O>> = T extends __Z<O> ? T : never
+
+export type Options = {
+  /**
+   * Maximum tokens per chunk when processing large files or many files.
+   * If a single file exceeds this limit, it will be split into chunks.
+   * If all files together exceed this limit, they will be processed in batches.
+   * If not specified, all files must fit in a single prompt.
+   */
+  maxTokensPerChunk?: number
+  /**
+   * Optional Zui schema to validate JSON files against after patching.
+   * Only applies to files detected as JSON (by .json extension).
+   * If the patched JSON doesn't match the schema, the LLM is re-prompted to fix it.
+   */
+  schema?: OfType<any>
+}
+
+const INVALID_SCHEMA_ERROR = 'zai.patch only accepts schemas created with @bpinternal/zui'
+
+const Options = z.object({
+  maxTokensPerChunk: z.number().optional(),
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Patches files based on natural language instructions using the micropatch protocol.
+     *
+     * This operation takes an array of files and instructions, then returns the modified files.
+     * It uses a token-efficient line-based patching protocol (micropatch) that allows precise
+     * modifications without regenerating entire files.
+     *
+     * @param files - Array of files to patch, each with path, name, and content
+     * @param instructions - Natural language instructions describing what changes to make
+     * @param options - Optional configuration for patch generation
+     * @returns Response promise resolving to array of patched files
+     *
+     * @example Simple text replacement
+     * ```typescript
+     * const files = [{
+     *   path: 'src/hello.ts',
+     *   name: 'hello.ts',
+     *   content: 'console.log("Hello World")'
+     * }]
+     *
+     * const patched = await zai.patch(
+     *   files,
+     *   'change the message to say "Hi World"'
+     * )
+     * // patched[0].content contains: console.log("Hi World")
+     * // patched[0].patch contains: ◼︎=1|console.log("Hi World")
+     * ```
+     *
+     * @example Adding documentation
+     * ```typescript
+     * const files = [{
+     *   path: 'src/utils.ts',
+     *   name: 'utils.ts',
+     *   content: 'export function add(a: number, b: number) {\n  return a + b\n}'
+     * }]
+     *
+     * const patched = await zai.patch(
+     *   files,
+     *   'add JSDoc comments to all exported functions'
+     * )
+     * ```
+     *
+     * @example Patching multiple files
+     * ```typescript
+     * const files = [
+     *   { path: 'package.json', name: 'package.json', content: '...' },
+     *   { path: 'config.json', name: 'config.json', content: '...' }
+     * ]
+     *
+     * const patched = await zai.patch(
+     *   files,
+     *   'update version to 2.0.0 in all config files'
+     * )
+     * ```
+     *
+     * @example Refactoring code
+     * ```typescript
+     * const files = [{
+     *   path: 'src/api.ts',
+     *   name: 'api.ts',
+     *   content: 'function fetchUser() { ... }'
+     * }]
+     *
+     * const patched = await zai.patch(
+     *   files,
+     *   'convert fetchUser to an async function and add error handling'
+     * )
+     * ```
+     *
+     * @example Removing code
+     * ```typescript
+     * const files = [{
+     *   path: 'src/legacy.ts',
+     *   name: 'legacy.ts',
+     *   content: 'const debug = true\nconsole.log("Debug mode")\nfunction main() {...}'
+     * }]
+     *
+     * const patched = await zai.patch(
+     *   files,
+     *   'remove all debug-related code'
+     * )
+     * ```
+     *
+     * @example Inspecting applied patches
+     * ```typescript
+     * const patched = await zai.patch(files, 'add error handling')
+     *
+     * // Check what patches were applied
+     * for (const file of patched) {
+     *   if (file.patch) {
+     *     console.log(`Patches for ${file.path}:`)
+     *     console.log(file.patch)
+     *   }
+     * }
+     * ```
+     */
+    patch(files: Array<File>, instructions: string, options?: Options): Response<Array<File>>
+  }
+}
+
+/**
+ * Represents a chunk of a file (partial file)
+ */
+type FileChunk = {
+  path: string
+  name: string
+  content: string
+  startLine: number // 1-based line number where this chunk starts in the original file
+  endLine: number // 1-based line number where this chunk ends in the original file
+  totalLines: number // Total lines in the complete file
+  isPartial: boolean // True if this is a chunk of a larger file
+}
+
+/**
+ * Represents a batch of files or file chunks to process together
+ */
+type ProcessingBatch = {
+  items: Array<FileChunk>
+  tokenCount: number
+}
+
+const patch = async (
+  files: Array<File>,
+  instructions: string,
+  _options: Options | undefined,
+  ctx: ZaiContext
+): Promise<Array<File>> => {
+  ctx.controller.signal.throwIfAborted()
+
+  const options: Options = { ...Options.parse(_options ?? {}), schema: _options?.schema }
+
+  let zSchema: z.ZodTypeAny | null = null
+  if (options.schema) {
+    if (!z.is.zuiType(options.schema)) {
+      throw new Error(INVALID_SCHEMA_ERROR)
+    }
+    zSchema = options.schema
+  }
+
+  if (files.length === 0) {
+    return []
+  }
+
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.patch'
+
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
+  const TOKENS_INSTRUCTIONS_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.2)
+  const TOKENS_FILES_MAX = TOKENS_TOTAL_MAX - TOKENS_INSTRUCTIONS_MAX
+
+  const truncatedInstructions = tokenizer.truncate(instructions, TOKENS_INSTRUCTIONS_MAX)
+
+  // Determine max tokens per chunk
+  const maxTokensPerChunk = options.maxTokensPerChunk ?? TOKENS_FILES_MAX
+
+  // Convert files to file chunks (initially full files)
+  const fileTokenCounts = files.map((file) => ({
+    file,
+    tokens: tokenizer.count(file.content),
+    lines: file.content.split(/\r?\n/).length,
+  }))
+
+  const totalInputTokens = fileTokenCounts.reduce((sum, f) => sum + f.tokens, 0)
+
+  /**
+   * Split a file into chunks by line ranges
+   */
+  const splitFileIntoChunks = (file: File, totalLines: number, fileTokens: number): Array<FileChunk> => {
+    const lines = file.content.split(/\r?\n/)
+    const tokensPerLine = fileTokens / totalLines
+    const linesPerChunk = Math.floor(maxTokensPerChunk / tokensPerLine)
+
+    if (linesPerChunk >= totalLines) {
+      // File fits in one chunk
+      return [
+        {
+          path: file.path,
+          name: file.name,
+          content: file.content,
+          startLine: 1,
+          endLine: totalLines,
+          totalLines,
+          isPartial: false,
+        },
+      ]
+    }
+
+    const chunks: Array<FileChunk> = []
+    for (let start = 0; start < totalLines; start += linesPerChunk) {
+      const end = Math.min(start + linesPerChunk, totalLines)
+      const chunkLines = lines.slice(start, end)
+      const chunkContent = chunkLines.join('\n')
+
+      chunks.push({
+        path: file.path,
+        name: file.name,
+        content: chunkContent,
+        startLine: start + 1,
+        endLine: end,
+        totalLines,
+        isPartial: true,
+      })
+    }
+
+    return chunks
+  }
+
+  /**
+   * Create batches of file chunks that fit within token limits
+   */
+  const createBatches = (chunks: Array<FileChunk>): Array<ProcessingBatch> => {
+    const batches: Array<ProcessingBatch> = []
+    let currentBatch: ProcessingBatch = { items: [], tokenCount: 0 }
+
+    for (const chunk of chunks) {
+      const chunkTokens = tokenizer.count(chunk.content)
+
+      if (currentBatch.tokenCount + chunkTokens > maxTokensPerChunk && currentBatch.items.length > 0) {
+        batches.push(currentBatch)
+        currentBatch = { items: [], tokenCount: 0 }
+      }
+
+      currentBatch.items.push(chunk)
+      currentBatch.tokenCount += chunkTokens
+    }
+
+    if (currentBatch.items.length > 0) {
+      batches.push(currentBatch)
+    }
+
+    return batches
+  }
+
+  /**
+   * Format file chunks using XML tags with numbered lines
+   */
+  const formatChunksForInput = (chunks: Array<FileChunk>): string => {
+    return chunks
+      .map((chunk) => {
+        const lines = chunk.content.split(/\r?\n/)
+
+        // Render with global line numbers
+        const numberedView = lines
+          .map((line, idx) => {
+            const lineNum = chunk.startLine + idx
+            return `${String(lineNum).padStart(3, '0')}|${line}`
+          })
+          .join('\n')
+
+        const partialNote = chunk.isPartial
+          ? ` (PARTIAL: lines ${chunk.startLine}-${chunk.endLine} of ${chunk.totalLines} total lines)`
+          : ''
+
+        return `<FILE path="${chunk.path}" name="${chunk.name}"${partialNote}>
+${numberedView}
+</FILE>`
+      })
+      .join('\n\n')
+  }
+
+  /**
+   * Parse XML output to extract patches per file
+   */
+  const parsePatchOutput = (output: string): Map<string, string> => {
+    const patchMap = new Map<string, string>()
+
+    // Match <FILE path="...">...</FILE> blocks
+    const fileBlockRegex = /<FILE[^>]*path="([^"]+)"[^>]*>([\s\S]*?)<\/FILE>/g
+    let match
+
+    while ((match = fileBlockRegex.exec(output)) !== null) {
+      const filePath = match[1]
+      const patchOps = match[2].trim()
+      patchMap.set(filePath, patchOps)
+    }
+
+    return patchMap
+  }
+
+  /**
+   * For JSON files, validate the patched content is valid JSON.
+   * Attempts repair with jsonrepair if invalid.
+   * Returns the (possibly repaired) content, or null if unrecoverable.
+   */
+  const ensureValidJson = (content: string): string | null => {
+    const result = validateOrRepairJson(content)
+    if (result.valid) {
+      return result.content
+    }
+    return null
+  }
+
+  /**
+   * Apply a patch to a file. If the file is JSON, validate and optionally repair the output.
+   * If repair fails, re-prompt the LLM to fix the JSON.
+   */
+  const applyPatchToFile = async (file: File, patchOps: string | undefined): Promise<File> => {
+    if (!patchOps || patchOps.trim().length === 0) {
+      return { ...file, patch: '' }
+    }
+
+    let patchedContent: string
+    try {
+      patchedContent = Micropatch.applyText(file.content, patchOps)
+    } catch (error) {
+      console.error(`Failed to apply patch to ${file.path}:`, error)
+      return { ...file, patch: `ERROR: ${error instanceof Error ? error.message : String(error)}` }
+    }
+
+    const jsonMode = isJsonFile(file.path, file.name)
+
+    if (!jsonMode) {
+      return { ...file, content: patchedContent, patch: patchOps }
+    }
+
+    // JSON mode: validate the original input is actually JSON
+    const originalParse = validateOrRepairJson(file.content)
+    if (!originalParse.valid) {
+      // Original file wasn't valid JSON either, skip JSON validation
+      return { ...file, content: patchedContent, patch: patchOps }
+    }
+
+    // Validate/repair the patched output
+    const validContent = ensureValidJson(patchedContent)
+    if (validContent === null) {
+      // Repair failed — give the LLM the broken result and ask it to fix the syntax
+      return retryJsonPatch(file, patchedContent)
+    }
+
+    // If a schema is provided, validate the JSON against it
+    if (zSchema) {
+      const parsed = JSON.parse(validContent)
+      const safe = zSchema.safeParse(parsed)
+      if (!safe.success) {
+        return retrySchemaValidation(file, validContent, safe.error)
+      }
+    }
+
+    return { ...file, content: validContent, patch: patchOps }
+  }
+
+  /**
+   * Re-prompt the LLM when a JSON file patch produced invalid JSON.
+   * Gives the LLM the broken result as a new file to patch, asking it to fix just the JSON syntax errors.
+   */
+  const retryJsonPatch = async (file: File, brokenContent: string): Promise<File> => {
+    // Get the parse error message
+    const parseResult = validateOrRepairJson(brokenContent)
+    const errorMessage = parseResult.error
+
+    // Format the broken content as a numbered file for the LLM to patch
+    const brokenChunk: FileChunk = {
+      path: file.path,
+      name: file.name,
+      content: brokenContent,
+      startLine: 1,
+      endLine: brokenContent.split(/\r?\n/).length,
+      totalLines: brokenContent.split(/\r?\n/).length,
+      isPartial: false,
+    }
+    const brokenFileInput = formatChunksForInput([brokenChunk])
+
+    const { extracted } = await ctx.generateContent({
+      systemPrompt: getMicropatchSystemPrompt(),
+      messages: [
+        {
+          type: 'text',
+          role: 'user',
+          content: [
+            'The following JSON file has syntax errors. The original instructions were:',
+            truncatedInstructions,
+            '',
+            'Here is the broken result:',
+            '',
+            brokenFileInput,
+            '',
+            `JSON parse error: ${errorMessage}`,
+            '',
+            'Fix ONLY the JSON syntax errors (missing commas, orphaned braces, etc). Do NOT change the data — only fix the syntax to make it valid JSON.',
+          ].join('\n'),
+        },
+      ],
+      stopSequences: [],
+      transform: (text) => (text ?? '').trim(),
+    })
+
+    const retryPatchMap = parsePatchOutput(extracted)
+    const retryPatchOps = retryPatchMap.get(file.path)
+
+    if (!retryPatchOps || retryPatchOps.trim().length === 0) {
+      throw new Error(`JSON patch retry for "${file.path}" produced no patches.\n\n${errorMessage}`)
+    }
+
+    try {
+      const retryContent = Micropatch.applyText(brokenContent, retryPatchOps)
+      const validContent = ensureValidJson(retryContent)
+      if (validContent !== null) {
+        if (zSchema) {
+          const parsed = JSON.parse(validContent)
+          const safe = zSchema.safeParse(parsed)
+          if (!safe.success) {
+            return retrySchemaValidation(file, validContent, safe.error)
+          }
+        }
+        return { ...file, content: validContent, patch: retryPatchOps }
+      }
+    } catch {
+      // Fall through to throw
+    }
+
+    throw new Error(`JSON patch retry for "${file.path}" still produced invalid JSON.\n\n${errorMessage}`)
+  }
+
+  /**
+   * Re-prompt the LLM when a JSON file's content is valid JSON but doesn't match the schema.
+   * Gives the LLM the content + the schema validation errors and asks it to fix the data.
+   */
+  const retrySchemaValidation = async (file: File, content: string, zodError: z.ZodError): Promise<File> => {
+    const errorMessage = new JsonParsingError(content, zodError).message
+
+    const contentChunk: FileChunk = {
+      path: file.path,
+      name: file.name,
+      content,
+      startLine: 1,
+      endLine: content.split(/\r?\n/).length,
+      totalLines: content.split(/\r?\n/).length,
+      isPartial: false,
+    }
+    const fileInput = formatChunksForInput([contentChunk])
+
+    const { extracted } = await ctx.generateContent({
+      systemPrompt: getMicropatchSystemPrompt(),
+      messages: [
+        {
+          type: 'text',
+          role: 'user',
+          content: [
+            'The following JSON file is valid JSON but does not match the expected schema.',
+            '',
+            fileInput,
+            '',
+            errorMessage,
+            '',
+            "Fix the JSON to match the schema. Only change what's needed to satisfy the validation errors above.",
+          ].join('\n'),
+        },
+      ],
+      stopSequences: [],
+      transform: (text) => (text ?? '').trim(),
+    })
+
+    const retryPatchMap = parsePatchOutput(extracted)
+    const retryPatchOps = retryPatchMap.get(file.path)
+
+    if (!retryPatchOps || retryPatchOps.trim().length === 0) {
+      throw new Error(`Schema validation retry for "${file.path}" produced no patches.\n\n${errorMessage}`)
+    }
+
+    try {
+      const retryContent = Micropatch.applyText(content, retryPatchOps)
+      const validContent = ensureValidJson(retryContent)
+      if (validContent !== null) {
+        if (zSchema) {
+          const parsed = JSON.parse(validContent)
+          const safe = zSchema.safeParse(parsed)
+          if (safe.success) {
+            return { ...file, content: validContent, patch: retryPatchOps }
+          }
+        } else {
+          return { ...file, content: validContent, patch: retryPatchOps }
+        }
+      }
+    } catch {
+      // Fall through to throw
+    }
+
+    throw new Error(`Schema validation retry for "${file.path}" still failed.\n\n${errorMessage}`)
+  }
+
+  /**
+   * Process a single batch of file chunks
+   */
+  const processBatch = async (batch: ProcessingBatch): Promise<Map<string, string>> => {
+    const chunksInput = formatChunksForInput(batch.items)
+    const { extracted } = await ctx.generateContent({
+      systemPrompt: getMicropatchSystemPrompt(),
+      messages: [
+        {
+          type: 'text',
+          role: 'user',
+          content: `
+Instructions: ${truncatedInstructions}
+
+${chunksInput}
+
+Generate patches for each file that needs modification:
+`.trim(),
+        },
+      ],
+      stopSequences: [],
+      transform: (text) => {
+        return text.trim()
+      },
+    })
+
+    return parsePatchOutput(extracted)
+  }
+
+  // Check if we need chunking
+  const needsChunking =
+    totalInputTokens > maxTokensPerChunk || fileTokenCounts.some((f) => f.tokens > maxTokensPerChunk)
+
+  if (!needsChunking) {
+    // Simple case: all files fit in one prompt (existing logic)
+    // Check for exact match in examples
+    const Key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        files: files.map((f) => ({ path: f.path, content: f.content })),
+        instructions: truncatedInstructions,
+      })
+    )
+
+    const tableExamples =
+      taskId && ctx.adapter
+        ? await ctx.adapter.getExamples<Array<File>, Array<File>>({
+            input: files,
+            taskId,
+            taskType,
+          })
+        : []
+
+    const exactMatch = tableExamples.find((x) => x.key === Key)
+    if (exactMatch) {
+      return exactMatch.output as Array<File>
+    }
+
+    // Process all files in one batch
+    const allChunks: Array<FileChunk> = fileTokenCounts.map(({ file }) => ({
+      path: file.path,
+      name: file.name,
+      content: file.content,
+      startLine: 1,
+      endLine: file.content.split(/\r?\n/).length,
+      totalLines: file.content.split(/\r?\n/).length,
+      isPartial: false,
+    }))
+
+    const batch: ProcessingBatch = { items: allChunks, tokenCount: totalInputTokens }
+    const patchMap = await processBatch(batch)
+
+    // Apply patches to each file (with JSON validation/repair for .json files)
+    const patchedFiles = await Promise.all(files.map((file) => applyPatchToFile(file, patchMap.get(file.path))))
+
+    // Save example for active learning
+    if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+      await ctx.adapter.saveExample({
+        key: Key,
+        taskType,
+        taskId,
+        input: files,
+        output: patchedFiles,
+        instructions: truncatedInstructions,
+        metadata: {
+          cost: {
+            input: ctx.usage.cost.input,
+            output: ctx.usage.cost.output,
+          },
+          latency: Date.now(),
+          model: ctx.modelId,
+          tokens: {
+            input: ctx.usage.tokens.input,
+            output: ctx.usage.tokens.output,
+          },
+        },
+      })
+    }
+
+    return patchedFiles
+  }
+
+  // Complex case: needs chunking
+  // Step 1: Split files that are too large
+  const allChunks: Array<FileChunk> = []
+  for (const { file, tokens, lines } of fileTokenCounts) {
+    const chunks = splitFileIntoChunks(file, lines, tokens)
+    allChunks.push(...chunks)
+  }
+
+  // Step 2: Create batches that fit within token limits
+  const batches = createBatches(allChunks)
+
+  // Step 3: Process batches in parallel
+  const limit = pLimit(10)
+  const batchResults = await Promise.all(batches.map((batch) => limit(() => processBatch(batch))))
+
+  // Step 4: Merge results - combine patches from all batches per file
+  const mergedPatches = new Map<string, string>()
+  for (const patchMap of batchResults) {
+    for (const [filePath, patchOps] of patchMap.entries()) {
+      const existing = mergedPatches.get(filePath) || ''
+      const combined = existing ? `${existing}\n${patchOps}` : patchOps
+      mergedPatches.set(filePath, combined)
+    }
+  }
+
+  // Step 5: Apply merged patches to original files (with JSON validation/repair for .json files)
+  const patchedFiles = await Promise.all(files.map((file) => applyPatchToFile(file, mergedPatches.get(file.path))))
+
+  return patchedFiles
+}
+
+/**
+ * Generate the system prompt that explains the micropatch protocol to the LLM
+ */
+function getMicropatchSystemPrompt(): string {
+  return `
+You are a code patching assistant. Your task is to generate precise line-based patches using the Micropatch protocol.
+
+## Input Format
+
+You will receive files in this XML format:
+
+\`\`\`
+<FILE path="src/hello.ts" name="hello.ts">
+001|const x = 1
+002|const y = 2
+003|console.log(x + y)
+</FILE>
+
+<FILE path="src/utils.ts" name="utils.ts">
+001|export function add(a, b) {
+002|  return a + b
+003|}
+</FILE>
+\`\`\`
+
+Each file has:
+- **path**: Full file path
+- **name**: File name
+- **Numbered lines**: Format is \`NNN|content\` where NNN is the ORIGINAL line number (1-based)
+
+## Output Format
+
+Generate patches for EACH file that needs modification using this EXACT XML format:
+
+\`\`\`
+<FILE path="src/hello.ts">
+◼︎=1|const a = 1
+◼︎=2|const b = 2
+◼︎=3|console.log(a + b)
+</FILE>
+
+<FILE path="src/utils.ts">
+◼︎<1|/**
+ * Adds two numbers
+ */
+</FILE>
+\`\`\`
+
+**CRITICAL RULES**:
+1. Each \`<FILE>\` tag MUST include the exact \`path\` attribute from the input
+2. Put patch operations for EACH file inside its own \`<FILE>...</FILE>\` block
+3. If a file doesn't need changes, omit its \`<FILE>\` block entirely
+4. DO NOT mix patches from different files
+5. DO NOT include line numbers or any text outside the patch operations
+
+## Micropatch Protocol
+
+The Micropatch protocol uses line numbers to reference ORIGINAL lines (before any edits).
+
+### Operations
+
+Each operation starts with the marker \`◼︎\` at the beginning of a line:
+
+1. **Insert BEFORE line**: \`◼︎<NNN|text\`
+   - Inserts \`text\` as a new line BEFORE original line NNN
+   - Example: \`◼︎<5|console.log('debug')\`
+
+2. **Insert AFTER line**: \`◼︎>NNN|text\`
+   - Inserts \`text\` as a new line AFTER original line NNN
+   - Example: \`◼︎>10|}\`
+
+3. **Replace single line**: \`◼︎=NNN|new text\`
+   - Replaces original line NNN with \`new text\`
+   - Can span multiple lines (continue until next ◼︎ or end)
+   - Example:
+     \`\`\`
+     ◼︎=7|function newName() {
+       return 42
+     }
+     \`\`\`
+
+4. **Replace range**: \`◼︎=NNN-MMM|replacement\`
+   - Replaces lines NNN through MMM with replacement text
+   - Example: \`◼︎=5-8|const combined = a + b + c + d\`
+
+5. **Delete single line**: \`◼︎-NNN\`
+   - Deletes original line NNN
+   - Example: \`◼︎-12\`
+
+6. **Delete range**: \`◼︎-NNN-MMM\`
+   - Deletes lines NNN through MMM inclusive
+   - Example: \`◼︎-5-10\`
+
+### Escaping
+
+- To include a literal \`◼︎\` in your text, use \`\\◼︎\`
+- No other escape sequences are recognized
+
+### Important Rules
+
+1. **Use ORIGINAL line numbers**: Always reference the line numbers shown in the input (001, 002, etc.)
+2. **One operation per line**: Each operation must start on a new line with \`◼︎\`
+3. **No explanations**: Output ONLY patch operations inside \`<FILE>\` tags
+4. **Precise operations**: Use the minimal set of operations to achieve the goal
+5. **Verify line numbers**: Double-check that line numbers match the input
+
+## Example
+
+**Input:**
+\`\`\`
+<FILE path="src/math.ts" name="math.ts">
+001|const x = 1
+002|const y = 2
+003|console.log(x + y)
+004|
+005|export { x, y }
+</FILE>
+\`\`\`
+
+**Task:** Change variable names from x,y to a,b
+
+**Output:**
+\`\`\`
+<FILE path="src/math.ts">
+◼︎=1|const a = 1
+◼︎=2|const b = 2
+◼︎=3|console.log(a + b)
+◼︎=5|export { a, b }
+</FILE>
+\`\`\`
+
+## Your Task
+
+Generate ONLY the \`<FILE>\` blocks with patch operations. Do not include explanations, comments, or any other text.
+`.trim()
+}
+
+Zai.prototype.patch = function (
+  this: Zai,
+  files: Array<File>,
+  instructions: string,
+  _options?: Options
+): Response<Array<File>> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.patch',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<Array<File>>(context, patch(files, instructions, _options, context), (result) => result)
+}

packages/zai/src/operations/rate.ts

@@ -0,0 +1,629 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+import pLimit from 'p-limit'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+
+// Rating scale constants
+const RATING_VALUES = {
+  very_bad: 1,
+  bad: 2,
+  average: 3,
+  good: 4,
+  very_good: 5,
+} as const
+
+type RatingLabel = keyof typeof RATING_VALUES
+
+// Evaluation criteria generated by LLM
+type EvaluationCriteria = Record<
+  string,
+  {
+    very_bad: string
+    bad: string
+    average: string
+    good: string
+    very_good: string
+  }
+>
+
+export type RatingInstructions = string | Record<string, string>
+
+export type Options = {
+  /** The maximum number of tokens per item */
+  tokensPerItem?: number
+  /** The maximum number of items to rate per chunk */
+  maxItemsPerChunk?: number
+}
+
+const _Options = z.object({
+  tokensPerItem: z
+    .number()
+    .min(1)
+    .max(100_000)
+    .optional()
+    .describe('The maximum number of tokens per item')
+    .default(250),
+  maxItemsPerChunk: z
+    .number()
+    .min(1)
+    .max(100)
+    .optional()
+    .describe('The maximum number of items to rate per chunk')
+    .default(50),
+})
+
+// Result types based on instructions type
+export type RatingResult<T extends RatingInstructions> = T extends string
+  ? {
+      [key: string]: number // criteria scores
+      total: number // sum of all criteria
+    }
+  : T extends Record<string, string>
+    ? {
+        [K in keyof T]: number // score for each criterion
+      } & {
+        total: number // sum of all criteria
+      }
+    : never
+
+export type SimplifiedRatingResult<T extends RatingInstructions> = T extends string ? number : RatingResult<T>
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Rates array items on a 1-5 scale based on single or multiple criteria.
+     *
+     * This operation evaluates each item and assigns numeric ratings (1-5) where:
+     * - 1 = Very Bad, 2 = Bad, 3 = Average, 4 = Good, 5 = Very Good
+     *
+     * Supports both simple single-criterion rating (string instructions) and
+     * multi-criteria rating (object with criterion → description mapping).
+     *
+     * @param input - Array of items to rate
+     * @param instructions - Single criterion (string) or multiple criteria (object)
+     * @param options - Configuration for chunking and tokens per item
+     * @returns Response with ratings array (simplified to numbers for single criterion)
+     *
+     * @example Single criterion rating
+     * ```typescript
+     * const reviews = [
+     *   "Amazing product! Best purchase ever!",
+     *   "It's okay, nothing special",
+     *   "Terrible quality, broke immediately"
+     * ]
+     *
+     * const ratings = await zai.rate(reviews, 'Rate the sentiment')
+     * // Result: [5, 3, 1] (simplified to numbers)
+     *
+     * // Get full details
+     * const { output } = await zai.rate(reviews, 'Rate the sentiment').result()
+     * // output[0]: { sentiment: 5, total: 5 }
+     * ```
+     *
+     * @example Multi-criteria rating
+     * ```typescript
+     * const essays = [
+     *   "... student essay text ...",
+     *   "... another essay ...",
+     * ]
+     *
+     * const ratings = await zai.rate(essays, {
+     *   grammar: 'Rate the grammar and spelling',
+     *   clarity: 'Rate how clear and well-organized the writing is',
+     *   argumentation: 'Rate the strength of arguments and evidence'
+     * })
+     *
+     * // Result: [
+     * //   { grammar: 4, clarity: 5, argumentation: 3, total: 12 },
+     * //   { grammar: 3, clarity: 4, argumentation: 4, total: 11 }
+     * // ]
+     * ```
+     *
+     * @example Rating customer support conversations
+     * ```typescript
+     * const conversations = [
+     *   { agent: 'John', messages: [...], duration: 300 },
+     *   { agent: 'Jane', messages: [...], duration: 180 }
+     * ]
+     *
+     * const ratings = await zai.rate(conversations, {
+     *   professionalism: 'How professional was the agent?',
+     *   helpfulness: 'How helpful was the agent in solving the issue?',
+     *   efficiency: 'How efficiently was the conversation handled?'
+     * })
+     *
+     * // Calculate average scores
+     * const avgProfessionalism = ratings.reduce((sum, r) => sum + r.professionalism, 0) / ratings.length
+     * ```
+     *
+     * @example Rating code quality
+     * ```typescript
+     * const codeSamples = [
+     *   "function foo() { return x + y }",
+     *   "const calculateTotal = (items) => items.reduce((sum, item) => sum + item.price, 0)",
+     * ]
+     *
+     * const quality = await zai.rate(codeSamples, {
+     *   readability: 'How readable and clear is the code?',
+     *   best_practices: 'Does it follow coding best practices?',
+     *   documentation: 'Is the code well-documented?'
+     * })
+     * ```
+     *
+     * @example Product review rating
+     * ```typescript
+     * const products = [
+     *   { name: 'Laptop', reviews: [...], avgStars: 4.2 },
+     *   { name: 'Mouse', reviews: [...], avgStars: 3.8 }
+     * ]
+     *
+     * const ratings = await zai.rate(
+     *   products,
+     *   'Rate overall product quality based on reviews and rating',
+     *   {
+     *     tokensPerItem: 500, // Allow more tokens for detailed reviews
+     *     maxItemsPerChunk: 20 // Process 20 products per chunk
+     *   }
+     * )
+     * ```
+     *
+     * @example Finding highest-rated items
+     * ```typescript
+     * const ratings = await zai.rate(items, {
+     *   quality: 'Product quality',
+     *   value: 'Value for money',
+     *   design: 'Design and aesthetics'
+     * })
+     *
+     * // Sort by total score
+     * const sorted = items
+     *   .map((item, idx) => ({ item, rating: ratings[idx] }))
+     *   .sort((a, b) => b.rating.total - a.rating.total)
+     *
+     * console.log('Top rated:', sorted[0].item)
+     * console.log('Score breakdown:', sorted[0].rating)
+     * ```
+     */
+    rate<T, I extends RatingInstructions>(
+      input: Array<T>,
+      instructions: I,
+      options?: Options
+    ): Response<Array<RatingResult<I>>, Array<SimplifiedRatingResult<I>>>
+  }
+}
+
+const END = '■END■'
+
+const rate = async <T, I extends RatingInstructions>(
+  input: Array<T>,
+  instructions: I,
+  _options: Options | undefined,
+  ctx: ZaiContext
+): Promise<Array<RatingResult<I>>> => {
+  ctx.controller.signal.throwIfAborted()
+  const options = _Options.parse(_options ?? {})
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  // Handle empty array
+  if (input.length === 0) {
+    return []
+  }
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.rate'
+
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
+
+  // Phase 1: Generate evaluation criteria
+  const isStringInstructions = typeof instructions === 'string'
+  const criteriaKeys: string[] = isStringInstructions
+    ? [] // Will be generated by LLM
+    : Object.keys(instructions as Record<string, string>)
+
+  const generateCriteriaPrompt = isStringInstructions
+    ? `Generate 3-5 evaluation criteria for: "${instructions}"
+
+For each criterion, provide 5 labels (very_bad, bad, average, good, very_good) with brief descriptions.
+
+Output format (JSON):
+{
+  "criterion1_name": {
+    "very_bad": "description",
+    "bad": "description",
+    "average": "description",
+    "good": "description",
+    "very_good": "description"
+  },
+  "criterion2_name": { ... }
+}
+
+Keep criterion names short (1-2 words, lowercase, use underscores).
+Keep descriptions brief (5-10 words each).`
+    : `For these evaluation criteria, provide 5 labels (very_bad, bad, average, good, very_good) with brief descriptions for each:
+
+${criteriaKeys.map((key) => `- ${key}: ${(instructions as Record<string, string>)[key]}`).join('\n')}
+
+Output format (JSON):
+{
+  "${criteriaKeys[0]}": {
+    "very_bad": "description",
+    "bad": "description",
+    "average": "description",
+    "good": "description",
+    "very_good": "description"
+  }
+  ${criteriaKeys.length > 1 ? '...' : ''}
+}
+
+Keep descriptions brief (5-10 words each).`
+
+  const { extracted: evaluationCriteria } = await ctx.generateContent({
+    systemPrompt: `You are creating evaluation criteria for rating items on a 1-5 scale.
+Each criterion must have exactly 5 labels: very_bad (1), bad (2), average (3), good (4), very_good (5).
+Output valid JSON only.`,
+    messages: [
+      {
+        type: 'text',
+        role: 'user',
+        content: generateCriteriaPrompt,
+      },
+    ],
+    transform: (text) => {
+      // Extract JSON from markdown code blocks if present
+      const jsonMatch = text.match(/```(?:json)?\s*(\{[\s\S]*?\})\s*```/) || text.match(/(\{[\s\S]*\})/)
+      if (!jsonMatch) {
+        throw new Error('Failed to parse evaluation criteria JSON')
+      }
+      return JSON.parse(jsonMatch[1]) as EvaluationCriteria
+    },
+  })
+
+  // Extract final criteria keys
+  const finalCriteriaKeys = Object.keys(evaluationCriteria)
+  if (finalCriteriaKeys.length === 0) {
+    throw new Error('No evaluation criteria generated')
+  }
+
+  // Phase 2: Chunk items
+  const TOKENS_CRITERIA_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.3)
+  const TOKENS_ITEMS_MAX = TOKENS_TOTAL_MAX - TOKENS_CRITERIA_MAX
+
+  let chunks: Array<typeof input> = []
+  let currentChunk: typeof input = []
+  let currentChunkTokens = 0
+
+  for (const element of input) {
+    const elementAsString = tokenizer.truncate(stringify(element, false), options.tokensPerItem)
+    const elementTokens = tokenizer.count(elementAsString)
+
+    if (currentChunkTokens + elementTokens > TOKENS_ITEMS_MAX || currentChunk.length >= options.maxItemsPerChunk) {
+      if (currentChunk.length > 0) {
+        chunks.push(currentChunk)
+      }
+      currentChunk = []
+      currentChunkTokens = 0
+    }
+
+    currentChunk.push(element)
+    currentChunkTokens += elementTokens
+  }
+
+  if (currentChunk.length > 0) {
+    chunks.push(currentChunk)
+  }
+
+  chunks = chunks.filter((x) => x.length > 0)
+
+  // Phase 3: Rate each chunk in parallel
+  type ChunkResult = {
+    ratings: Array<Record<string, number>>
+    meta: { cost: { input: number; output: number }; latency: number; tokens: { input: number; output: number } }
+  }
+
+  const rateChunk = async (chunk: typeof input): Promise<ChunkResult> => {
+    ctx.controller.signal.throwIfAborted()
+
+    // Get examples from adapter for active learning
+    const chunkInputStr = JSON.stringify(chunk)
+    const examples =
+      taskId && ctx.adapter
+        ? await ctx.adapter.getExamples<string, Array<Record<string, number>>>({
+            input: chunkInputStr.slice(0, 1000), // Limit search string length
+            taskType,
+            taskId,
+          })
+        : []
+
+    // Check for exact match (cache hit)
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: chunkInputStr,
+        instructions: stringify(instructions),
+      })
+    )
+    const exactMatch = examples.find((x) => x.key === key)
+    if (exactMatch && exactMatch.output) {
+      // Return cached result with zero cost
+      return {
+        ratings: exactMatch.output,
+        meta: { cost: { input: 0, output: 0 }, latency: 0, tokens: { input: 0, output: 0 } },
+      }
+    }
+
+    const formatCriteria = () => {
+      return finalCriteriaKeys
+        .map((key) => {
+          const labels = evaluationCriteria[key]
+          return `**${key}**:
+  - very_bad (1): ${labels?.very_bad}
+  - bad (2): ${labels?.bad}
+  - average (3): ${labels?.average}
+  - good (4): ${labels?.good}
+  - very_good (5): ${labels?.very_good}`
+        })
+        .join('\n\n')
+    }
+
+    const formatItems = (items: typeof chunk) => {
+      return items
+        .map((item, idx) => {
+          const itemStr = tokenizer.truncate(stringify(item, false), options.tokensPerItem)
+          return `■${idx}: ${itemStr}■`
+        })
+        .join('\n')
+    }
+
+    // Format examples for few-shot learning
+    const exampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []
+
+    for (const example of examples.slice(0, 5)) {
+      // User message with input
+      try {
+        const exampleInput = JSON.parse(example.input)
+        exampleMessages.push({
+          type: 'text',
+          role: 'user',
+          content: `Expert Example - Items to rate:
+${formatItems(Array.isArray(exampleInput) ? exampleInput : [exampleInput])}
+
+Rate each item on all criteria.`,
+        })
+
+        // Assistant message with ratings
+        const exampleOutput = example.output
+        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
+          const formattedRatings = exampleOutput
+            .map((rating, idx) => {
+              const pairs = finalCriteriaKeys
+                .map((key) => {
+                  const value = rating[key]
+                  if (typeof value === 'number') {
+                    // Convert number back to label
+                    const labelMap: Record<number, string> = {
+                      1: 'very_bad',
+                      2: 'bad',
+                      3: 'average',
+                      4: 'good',
+                      5: 'very_good',
+                    }
+                    return `${key}=${labelMap[value] || 'average'}`
+                  }
+                  return null
+                })
+                .filter(Boolean)
+                .join(';')
+              return `■${idx}:${pairs}■`
+            })
+            .join('\n')
+
+          exampleMessages.push({
+            type: 'text',
+            role: 'assistant',
+            content: `${formattedRatings}\n${END}`,
+          })
+
+          if (example.explanation) {
+            exampleMessages.push({
+              type: 'text',
+              role: 'assistant',
+              content: `Reasoning: ${example.explanation}`,
+            })
+          }
+        }
+      } catch {
+        // Skip malformed examples
+      }
+    }
+
+    const { extracted, meta } = await ctx.generateContent({
+      systemPrompt: `You are rating items based on evaluation criteria.
+
+Evaluation Criteria:
+${formatCriteria()}
+
+For each item, rate it on EACH criterion using one of these labels:
+very_bad, bad, average, good, very_good
+
+Output format:
+■0:criterion1=label;criterion2=label;criterion3=label■
+■1:criterion1=label;criterion2=label;criterion3=label■
+${END}
+
+IMPORTANT:
+- Rate every item (■0 to ■${chunk.length - 1})
+- Use exact criterion names: ${finalCriteriaKeys.join(', ')}
+- Use exact label names: very_bad, bad, average, good, very_good
+- Use semicolons (;) between criteria
+- Use equals (=) between criterion and label`,
+      stopSequences: [END],
+      messages: [
+        ...exampleMessages,
+        {
+          type: 'text',
+          role: 'user',
+          content: `Items to rate (■0 to ■${chunk.length - 1}):
+${formatItems(chunk)}
+
+Rate each item on all criteria.
+Output format: ■index:criterion1=label;criterion2=label■
+${END}`,
+        },
+      ],
+      transform: (text) => {
+        const results: Array<Record<string, number>> = []
+
+        // Parse ratings: ■0:affordability=good;quality=very_good■
+        const regex = /■(\d+):([^■]+)■/g
+        let match: RegExpExecArray | null
+
+        while ((match = regex.exec(text)) !== null) {
+          const idx = parseInt(match[1] ?? '', 10)
+          const ratingsStr = match[2] ?? ''
+
+          if (isNaN(idx) || idx < 0 || idx >= chunk.length) {
+            continue
+          }
+
+          const itemRatings: Record<string, number> = {}
+          let total = 0
+
+          // Parse criterion=label pairs
+          const pairs = ratingsStr.split(';').filter((x) => x.trim().length > 0)
+          for (const pair of pairs) {
+            const [criterion, label] = pair.split('=').map((x) => x.trim())
+            if (!criterion || !label) continue
+
+            // Convert label to number
+            const labelLower = label.toLowerCase().replace(/\s+/g, '_')
+            const ratingValue = RATING_VALUES[labelLower as RatingLabel] ?? 3 // default to average
+
+            itemRatings[criterion] = ratingValue
+            total += ratingValue
+          }
+
+          itemRatings.total = total
+          results[idx] = itemRatings
+        }
+
+        // Fill in missing results with defaults
+        for (let i = 0; i < chunk.length; i++) {
+          if (!results[i]) {
+            const defaultRatings: Record<string, number> = {}
+            let total = 0
+            for (const key of finalCriteriaKeys) {
+              defaultRatings[key] = 3 // average
+              total += 3
+            }
+            defaultRatings.total = total
+            results[i] = defaultRatings
+          }
+        }
+
+        return results
+      },
+    })
+
+    return { ratings: extracted, meta }
+  }
+
+  // Process chunks in parallel with p-limit
+  const limit = pLimit(10)
+  const chunkPromises = chunks.map((chunk) => limit(() => rateChunk(chunk)))
+
+  const ratedChunks = await Promise.all(chunkPromises)
+
+  // Phase 4: Flatten results and accumulate metadata
+  const allRatings = ratedChunks.flatMap((result) => result.ratings) as Array<RatingResult<I>>
+
+  // Accumulate metadata from all chunks
+  const totalMeta = ratedChunks.reduce(
+    (acc, result) => ({
+      cost: {
+        input: acc.cost.input + result.meta.cost.input,
+        output: acc.cost.output + result.meta.cost.output,
+      },
+      latency: Math.max(acc.latency, result.meta.latency), // Use max latency
+      tokens: {
+        input: acc.tokens.input + result.meta.tokens.input,
+        output: acc.tokens.output + result.meta.tokens.output,
+      },
+    }),
+    {
+      cost: { input: 0, output: 0 },
+      latency: 0,
+      tokens: { input: 0, output: 0 },
+    }
+  )
+
+  // Save example for active learning
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: JSON.stringify(input),
+        instructions: stringify(instructions),
+      })
+    )
+
+    await ctx.adapter.saveExample({
+      key,
+      taskType,
+      taskId,
+      input: JSON.stringify(input),
+      output: allRatings,
+      instructions: typeof instructions === 'string' ? instructions : JSON.stringify(instructions),
+      metadata: {
+        cost: {
+          input: totalMeta.cost.input,
+          output: totalMeta.cost.output,
+        },
+        latency: totalMeta.latency,
+        model: ctx.modelId,
+        tokens: {
+          input: totalMeta.tokens.input,
+          output: totalMeta.tokens.output,
+        },
+      },
+    })
+  }
+
+  return allRatings
+}
+
+Zai.prototype.rate = function <T, I extends RatingInstructions>(
+  this: Zai,
+  input: Array<T>,
+  instructions: I,
+  _options?: Options
+): Response<Array<RatingResult<I>>, Array<SimplifiedRatingResult<I>>> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.rate',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<Array<RatingResult<I>>, Array<SimplifiedRatingResult<I>>>(
+    context,
+    rate(input, instructions, _options, context),
+    (result) => {
+      // If instructions is a string, simplify to just the total number
+      if (typeof instructions === 'string') {
+        return result.map((r) => r.total as SimplifiedRatingResult<I>)
+      }
+      // Otherwise return the full result (including total)
+      return result as Array<SimplifiedRatingResult<I>>
+    }
+  )
+}

packages/zai/src/operations/rewrite.ts

@@ -0,0 +1,295 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify, takeUntilTokens } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+
+type Example = {
+  input: string
+  output: string
+  instructions?: string
+}
+
+const _Example = z.object({
+  input: z.string(),
+  output: z.string(),
+})
+
+export type Options = {
+  /** Examples to guide the rewriting */
+  examples?: Array<Example>
+  /** The maximum number of tokens to generate */
+  length?: number
+}
+
+const Options = z.object({
+  examples: z.array(_Example).default([]),
+  length: z.number().min(10).max(16_000).optional().describe('The maximum number of tokens to generate'),
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Rewrites text according to specific instructions while preserving the core meaning.
+     *
+     * This operation transforms text based on natural language instructions. Perfect for
+     * tone adjustment, format conversion, translation, simplification, and style changes.
+     *
+     * @param original - The text to rewrite
+     * @param prompt - Instructions describing how to transform the text
+     * @param options - Configuration for examples and length constraints
+     * @returns Response promise resolving to the rewritten text
+     *
+     * @example Change tone
+     * ```typescript
+     * const original = "Your request has been denied due to insufficient funds."
+     * const friendly = await zai.rewrite(
+     *   original,
+     *   'Make this sound more friendly and empathetic'
+     * )
+     * // Result: "We appreciate your interest, but unfortunately we're unable to proceed..."
+     * ```
+     *
+     * @example Simplify technical content
+     * ```typescript
+     * const technical = "The API implements RESTful architecture with OAuth 2.0 authentication..."
+     * const simple = await zai.rewrite(
+     *   technical,
+     *   'Explain this in simple terms for non-technical users'
+     * )
+     * ```
+     *
+     * @example Professional email conversion
+     * ```typescript
+     * const casual = "Hey, can you send me that report? Thanks!"
+     * const professional = await zai.rewrite(
+     *   casual,
+     *   'Rewrite this as a formal business email'
+     * )
+     * // Result: "Dear colleague, I would appreciate it if you could share the report..."
+     * ```
+     *
+     * @example Format conversion
+     * ```typescript
+     * const paragraph = "First do this. Then do that. Finally do this other thing."
+     * const bullets = await zai.rewrite(
+     *   paragraph,
+     *   'Convert this to a bulleted list'
+     * )
+     * // Result:
+     * // - First do this
+     * // - Then do that
+     * // - Finally do this other thing
+     * ```
+     *
+     * @example With examples for consistent style
+     * ```typescript
+     * const result = await zai.rewrite(
+     *   original,
+     *   'Rewrite in our brand voice',
+     *   {
+     *     examples: [
+     *       {
+     *         input: 'We offer many products.',
+     *         output: 'Discover our curated collection of innovative solutions.'
+     *       },
+     *       {
+     *         input: 'Contact us for help.',
+     *         output: "We're here to support your success. Let's connect!"
+     *       }
+     *     ],
+     *     length: 200 // Limit output length
+     *   }
+     * )
+     * ```
+     *
+     * @example Translation-like transformation
+     * ```typescript
+     * const code = "if (user.isActive && user.hasPermission) { allowAccess() }"
+     * const pseudocode = await zai.rewrite(
+     *   code,
+     *   'Convert this code to natural language pseudocode'
+     * )
+     * // Result: "If the user is active AND has permission, then allow access"
+     * ```
+     */
+    rewrite(original: string, prompt: string, options?: Options): Response<string>
+  }
+}
+
+const START = '■START■'
+const END = '■END■'
+
+const rewrite = async (
+  original: string,
+  prompt: string,
+  _options: Options | undefined,
+  ctx: ZaiContext
+): Promise<string> => {
+  ctx.controller.signal.throwIfAborted()
+  const options = Options.parse(_options ?? {}) as Options
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.rewrite'
+
+  const INPUT_COMPONENT_SIZE = Math.max(100, (model.input.maxTokens - PROMPT_INPUT_BUFFER) / 2)
+  prompt = tokenizer.truncate(prompt, INPUT_COMPONENT_SIZE)
+
+  const inputSize = tokenizer.count(original) + tokenizer.count(prompt)
+  const maxInputSize = model.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER
+  if (inputSize > maxInputSize) {
+    throw new Error(
+      `The input size is ${inputSize} tokens long, which is more than the maximum of ${maxInputSize} tokens for this model (${model.name} = ${model.input.maxTokens} tokens)`
+    )
+  }
+
+  const instructions: string[] = []
+
+  const originalSize = tokenizer.count(original)
+  const generationLength = options.length
+    ? Math.min(Math.max(options.length * 10, 500), model.output.maxTokens - PROMPT_OUTPUT_BUFFER)
+    : undefined
+
+  if (options.length && originalSize > options.length) {
+    instructions.push(`The original text is ${originalSize} tokens long – it should be less than ${options.length}`)
+    instructions.push(
+      `The text must be standalone and complete in less than ${options.length} tokens, so it has to be shortened to fit the length as well`
+    )
+  }
+
+  const format = (before: string, prompt: string) => {
+    return `
+Prompt: ${prompt}
+
+${START}
+${before}
+${END}
+`.trim()
+  }
+
+  const Key = fastHash(
+    stringify({
+      taskId,
+      taskType,
+      input: original,
+      prompt,
+    })
+  )
+
+  const formatExample = ({ input, output, instructions }: Example) => {
+    return [
+      { type: 'text' as const, role: 'user' as const, content: format(input, instructions || prompt) },
+      { type: 'text' as const, role: 'assistant' as const, content: `${START}${output}${END}` },
+    ]
+  }
+
+  const defaultExamples: Example[] = [
+    { input: 'Hello, how are you?', output: 'Bonjour, comment ça va?', instructions: 'translate to French' },
+    { input: '1\n2\n3', output: '3\n2\n1', instructions: 'reverse the order' },
+  ]
+
+  const tableExamples =
+    taskId && ctx.adapter
+      ? await ctx.adapter.getExamples<string, string>({
+          input: original,
+          taskId,
+          taskType,
+        })
+      : []
+
+  const exactMatch = tableExamples.find((x) => x.key === Key)
+  if (exactMatch) {
+    return exactMatch.output
+  }
+
+  const savedExamples: Example[] = [
+    ...tableExamples.map((x) => ({ input: x.input as string, output: x.output as string }) satisfies Example),
+    ...options.examples,
+  ]
+
+  const REMAINING_TOKENS = model.input.maxTokens - tokenizer.count(prompt) - PROMPT_INPUT_BUFFER
+  const examples = takeUntilTokens(
+    savedExamples.length ? savedExamples : defaultExamples,
+    REMAINING_TOKENS,
+    (el) => tokenizer.count(stringify(el.input)) + tokenizer.count(stringify(el.output))
+  )
+    .map(formatExample)
+    .flat()
+
+  const { extracted, meta } = await ctx.generateContent({
+    systemPrompt: `
+Rewrite the text between the ${START} and ${END} tags to match the user prompt.
+${instructions.map((x) => `• ${x}`).join('\n')}
+`.trim(),
+    messages: [...examples, { type: 'text', content: format(original, prompt), role: 'user' }],
+    stopSequences: [END],
+    ...(generationLength ? { maxTokens: generationLength } : {}),
+    transform: (text) => {
+      if (!text.trim().length) {
+        throw new Error('The model did not return a valid rewrite. The response was empty.')
+      }
+
+      return text
+    },
+  })
+
+  let result = extracted
+
+  if (result.includes(START)) {
+    result = result.slice(result.indexOf(START) + START.length)
+  }
+
+  if (result.includes(END)) {
+    result = result.slice(0, result.indexOf(END))
+  }
+
+  result = result.trim()
+
+  if (options.length && tokenizer.count(result) > options.length) {
+    result = tokenizer.truncate(result, options.length).trim()
+  }
+
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    await ctx.adapter.saveExample({
+      key: Key,
+      metadata: {
+        cost: {
+          input: meta.cost.input,
+          output: meta.cost.output,
+        },
+        latency: meta.latency,
+        model: ctx.modelId,
+        tokens: {
+          input: meta.tokens.input,
+          output: meta.tokens.output,
+        },
+      },
+      instructions: prompt,
+      input: original,
+      output: result,
+      taskType,
+      taskId,
+    })
+  }
+
+  return result
+}
+
+Zai.prototype.rewrite = function (this: Zai, original: string, prompt: string, _options?: Options): Response<string> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.rewrite',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<string>(context, rewrite(original, prompt, _options, context), (result) => result)
+}

packages/zai/src/operations/sort.ts

@@ -0,0 +1,811 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+import pLimit from 'p-limit'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+
+export type Options = {
+  /** The maximum number of tokens per item */
+  tokensPerItem?: number
+}
+
+const _Options = z.object({
+  tokensPerItem: z
+    .number()
+    .min(1)
+    .max(100_000)
+    .optional()
+    .describe('The maximum number of tokens per item')
+    .default(250),
+})
+
+// Evaluation criteria generated by LLM
+type SortingCriteria = Record<
+  string,
+  {
+    description: string
+    labels: string[] // Ordered array of labels from FIRST to LAST, e.g., ['very_slow', 'slow', 'medium', 'fast', 'very_fast']
+  }
+>
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Sorts array items based on natural language sorting criteria.
+     *
+     * This operation intelligently orders items according to your instructions, understanding
+     * complex sorting logic like priority, quality, chronology, or any custom criteria.
+     * Perfect for ranking, organizing, and prioritizing lists based on subjective or
+     * multi-faceted criteria.
+     *
+     * @param input - Array of items to sort
+     * @param instructions - Natural language description of how to sort (e.g., "by priority", "newest first")
+     * @param options - Configuration for tokens per item
+     * @returns Response resolving to the sorted array
+     *
+     * @example Sort by price
+     * ```typescript
+     * const products = [
+     *   { name: 'Laptop', price: 999 },
+     *   { name: 'Mouse', price: 29 },
+     *   { name: 'Keyboard', price: 79 }
+     * ]
+     *
+     * const sorted = await zai.sort(products, 'from least expensive to most expensive')
+     * // Result: [Mouse ($29), Keyboard ($79), Laptop ($999)]
+     * ```
+     *
+     * @example Sort by priority/urgency
+     * ```typescript
+     * const tasks = [
+     *   "Update documentation",
+     *   "Fix critical security bug",
+     *   "Add new feature",
+     *   "System is down - all users affected"
+     * ]
+     *
+     * const prioritized = await zai.sort(tasks, 'by urgency and impact, most urgent first')
+     * // Result: ["System is down...", "Fix critical security bug", "Add new feature", "Update documentation"]
+     * ```
+     *
+     * @example Sort by quality/rating
+     * ```typescript
+     * const reviews = [
+     *   "Product is okay",
+     *   "Absolutely amazing! Best purchase ever!",
+     *   "Terrible, broke immediately",
+     *   "Good quality for the price"
+     * ]
+     *
+     * const sorted = await zai.sort(reviews, 'by sentiment, most positive first')
+     * // Result: [amazing review, good review, okay review, terrible review]
+     * ```
+     *
+     * @example Sort by complexity
+     * ```typescript
+     * const problems = [
+     *   "Fix typo in README",
+     *   "Redesign entire authentication system",
+     *   "Update a dependency version",
+     *   "Implement new microservice architecture"
+     * ]
+     *
+     * const sorted = await zai.sort(problems, 'by complexity, simplest first')
+     * ```
+     *
+     * @example Sort by relevance to query
+     * ```typescript
+     * const documents = [
+     *   "Article about cats",
+     *   "Article about dogs and their training",
+     *   "Article about dog breeds",
+     *   "Article about fish"
+     * ]
+     *
+     * const sorted = await zai.sort(
+     *   documents,
+     *   'by relevance to "dog training", most relevant first'
+     * )
+     * // Result: [dogs and training, dog breeds, cats, fish]
+     * ```
+     *
+     * @example Sort candidates by fit
+     * ```typescript
+     * const candidates = [
+     *   { name: 'Alice', experience: 5, skills: ['React', 'Node'] },
+     *   { name: 'Bob', experience: 10, skills: ['Python', 'ML'] },
+     *   { name: 'Charlie', experience: 3, skills: ['React', 'TypeScript'] }
+     * ]
+     *
+     * const sorted = await zai.sort(
+     *   candidates,
+     *   'by fit for a senior React developer position, best fit first'
+     * )
+     * ```
+     *
+     * @example Sort chronologically
+     * ```typescript
+     * const events = [
+     *   "Started the project last month",
+     *   "Will launch next week",
+     *   "Met with client yesterday",
+     *   "Planning meeting tomorrow"
+     * ]
+     *
+     * const chronological = await zai.sort(events, 'in chronological order')
+     * // Understands relative time expressions
+     * ```
+     *
+     * @example With token limit per item
+     * ```typescript
+     * const sorted = await zai.sort(
+     *   longDocuments,
+     *   'by relevance to climate change research',
+     *   { tokensPerItem: 500 } // Allow 500 tokens per document
+     * )
+     * ```
+     */
+    sort<T>(input: Array<T>, instructions: string, options?: Options): Response<Array<T>, Array<T>>
+  }
+}
+
+const END = '■END■'
+
+const sort = async <T>(
+  input: Array<T>,
+  instructions: string,
+  _options: Options | undefined,
+  ctx: ZaiContext
+): Promise<Array<T>> => {
+  ctx.controller.signal.throwIfAborted()
+
+  const options = _Options.parse(_options ?? {})
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.sort'
+
+  // Handle empty or single element arrays
+  if (input.length === 0) {
+    return []
+  }
+
+  if (input.length === 1) {
+    return input
+  }
+
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
+
+  // Phase 1: Generate sorting criteria from instructions + sample items
+  const sampleSize = Math.min(5, input.length)
+  const sampleItems = input.slice(0, sampleSize)
+  const sampleItemsText = sampleItems.map((item, idx) => `■${idx}: ${stringify(item, false)}`).join('\n')
+
+  // Get examples from adapter for criteria generation
+  const criteriaInputStr = JSON.stringify({ instructions, sampleItems })
+  const criteriaExamples =
+    taskId && ctx.adapter
+      ? await ctx.adapter.getExamples<string, Array<T>>({
+          input: criteriaInputStr.slice(0, 1000),
+          taskType,
+          taskId,
+        })
+      : []
+
+  // Format examples for few-shot learning in criteria generation
+  const criteriaExampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []
+
+  for (const example of criteriaExamples.slice(0, 2)) {
+    try {
+      const exampleInput = JSON.parse(example.input)
+      const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput]
+
+      criteriaExampleMessages.push({
+        type: 'text',
+        role: 'user',
+        content: `Expert Example - Analyze sorting instruction: "${instructions}"\n\nSample items:\n${exampleItems
+          .slice(0, 3)
+          .map((el, i) => `■${i}: ${stringify(el, false).slice(0, 200)}`)
+          .join('\n')}\n\nGenerate sorting criteria.`,
+      })
+
+      // Try to infer criteria from the example if available
+      if (example.explanation) {
+        criteriaExampleMessages.push({
+          type: 'text',
+          role: 'assistant',
+          content: `${example.explanation}\n${END}`,
+        })
+      }
+    } catch {
+      // Skip malformed examples
+    }
+  }
+
+  const generateCriteriaPrompt = `Analyze this sorting instruction: "${instructions}"
+
+Sample items to be sorted:
+${sampleItemsText}
+
+Create 1-3 sorting criteria with ordered label arrays (3-10 labels each).
+
+**CRITICAL RULES**:
+1. Labels are single words, lowercase, no spaces, use underscores
+2. Labels are ordered from FIRST to LAST in sorted result
+3. If instruction says "from X to Y": first label represents X, last label represents Y
+4. If instruction says "prioritize" or "highest/lowest priority":
+   - First label = HIGHEST priority (top of todo list)
+   - Last label = LOWEST priority (bottom of todo list)
+
+Examples:
+
+"from slowest to fastest" → first=slowest, last=fastest
+■speed■
+very_slow;slow;medium;fast;very_fast
+■END■
+
+"from most dangerous to least dangerous" → first=most dangerous, last=least dangerous
+■danger■
+extremely_dangerous;very_dangerous;dangerous;moderate;slightly_dangerous;harmless
+■END■
+
+"from least urgent (spam) to most urgent (bills)" → first=spam, last=bills
+■urgency■
+spam;promotional;normal;important;urgent;critical
+■END■
+
+"prioritize: highest priority=open old tickets; lowest priority=closed" → first=high priority, last=low priority
+■status■
+open_old;open_recent;closed
+■age■
+oldest;old;recent;new
+■END■
+
+Output format:
+■criterion_name■
+label1;label2;label3;label4
+■END■
+
+Use 3-10 labels per criterion. Labels should be intuitive and match the domain.
+Keep criterion names short (1-2 words, lowercase, underscores).
+`
+
+  const { extracted: sortingCriteria } = await ctx.generateContent({
+    systemPrompt: `You are creating sorting criteria with ordered label arrays.
+
+CRITICAL: Output ordered labels from FIRST to LAST position in sorted result.
+- Labels are single words, lowercase, underscores only
+- 3-10 labels per criterion
+- Order matters: first label = appears first, last label = appears last`,
+    messages: [
+      ...criteriaExampleMessages,
+      {
+        type: 'text',
+        role: 'user',
+        content: generateCriteriaPrompt,
+      },
+    ],
+    transform: (text) => {
+      const criteria: SortingCriteria = {}
+      const criterionRegex = /■([^■]+)■\s*([^\n■]+)/g
+      let match: RegExpExecArray | null
+
+      while ((match = criterionRegex.exec(text)) !== null) {
+        const name = (match[1] ?? '').trim().toLowerCase()
+        const labelsStr = (match[2] ?? '').trim()
+
+        if (!name || name === 'end') continue
+
+        // Parse semicolon-separated labels
+        const labels = labelsStr
+          .split(';')
+          .map((l) => l.trim().toLowerCase().replace(/\s+/g, '_'))
+          .filter((l) => l.length > 0 && l.length < 50)
+
+        if (labels.length >= 3 && labels.length <= 10) {
+          criteria[name] = {
+            description: `${labels.length} ordered labels`,
+            labels,
+          }
+        }
+      }
+
+      if (Object.keys(criteria).length === 0) {
+        throw new Error(`Failed to parse sorting criteria. LLM output: ${text.slice(0, 500)}`)
+      }
+
+      return criteria
+    },
+  })
+
+  const criteriaKeys = Object.keys(sortingCriteria)
+  if (criteriaKeys.length === 0) {
+    throw new Error('No sorting criteria generated')
+  }
+
+  // Phase 2: Chunk items and score them in parallel
+  const TOKENS_CRITERIA_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.2)
+  const TOKENS_ITEMS_MAX = TOKENS_TOTAL_MAX - TOKENS_CRITERIA_MAX
+
+  const MAX_ITEMS_PER_CHUNK = 50
+
+  // Prepare elements with indices
+  const elements = input.map((element, idx) => ({
+    element,
+    index: idx,
+    stringified: stringify(element, false),
+  }))
+
+  // Chunk elements
+  const chunks: Array<typeof elements> = []
+  let currentChunk: typeof elements = []
+  let currentTokens = 0
+
+  for (const elem of elements) {
+    const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem)
+    const elemTokens = tokenizer.count(truncated)
+
+    if (
+      (currentTokens + elemTokens > TOKENS_ITEMS_MAX || currentChunk.length >= MAX_ITEMS_PER_CHUNK) &&
+      currentChunk.length > 0
+    ) {
+      chunks.push(currentChunk)
+      currentChunk = []
+      currentTokens = 0
+    }
+
+    currentChunk.push(elem)
+    currentTokens += elemTokens
+  }
+
+  if (currentChunk.length > 0) {
+    chunks.push(currentChunk)
+  }
+
+  // Phase 3: Score each chunk
+  type ItemScore = {
+    elementIndex: number
+    scores: Record<string, number>
+    totalScore: number
+  }
+
+  const scoreChunk = async (chunk: typeof elements): Promise<ItemScore[]> => {
+    ctx.controller.signal.throwIfAborted()
+
+    const chunkSize = chunk.length
+    const chunkInputStr = JSON.stringify(chunk.map((c) => c.element))
+
+    // Get examples from adapter for active learning
+    const examples =
+      taskId && ctx.adapter
+        ? await ctx.adapter.getExamples<string, ItemScore[]>({
+            input: chunkInputStr.slice(0, 1000),
+            taskType,
+            taskId,
+          })
+        : []
+
+    // Check for exact match (cache hit)
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: chunkInputStr,
+        instructions,
+      })
+    )
+
+    const exactMatch = examples.find((x) => x.key === key)
+    if (exactMatch && exactMatch.output) {
+      return exactMatch.output
+    }
+
+    const elementsText = chunk
+      .map((elem, i) => {
+        const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem)
+        return `■${i}: ${truncated}■`
+      })
+      .join('\n')
+
+    const criteriaText = criteriaKeys
+      .map((key) => {
+        const criterion = sortingCriteria[key]
+        const labelsText = criterion.labels.join(';')
+        return `**${key}**: ${labelsText}`
+      })
+      .join('\n')
+
+    // Format examples for few-shot learning
+    const exampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []
+
+    for (const example of examples.slice(0, 3)) {
+      try {
+        const exampleInput = JSON.parse(example.input)
+        const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput]
+
+        exampleMessages.push({
+          type: 'text',
+          role: 'user',
+          content: `Expert Example - Items to score:\n${exampleItems.map((el, i) => `■${i}: ${stringify(el, false).slice(0, 200)}■`).join('\n')}\n\nScore each item.`,
+        })
+
+        const exampleOutput = example.output
+        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
+          const formattedScores = exampleOutput
+            .map((score) => {
+              const pairs = criteriaKeys.map((key) => `${key}=${score.scores[key] ?? 0}`).join(';')
+              return `■${score.elementIndex}:${pairs}■`
+            })
+            .join('\n')
+
+          exampleMessages.push({
+            type: 'text',
+            role: 'assistant',
+            content: `${formattedScores}\n${END}`,
+          })
+
+          if (example.explanation) {
+            exampleMessages.push({
+              type: 'text',
+              role: 'assistant',
+              content: `Reasoning: ${example.explanation}`,
+            })
+          }
+        }
+      } catch {
+        // Skip malformed examples
+      }
+    }
+
+    const { extracted } = await ctx.generateContent({
+      systemPrompt: `You are ranking items for sorting using ordered label arrays.
+
+${criteriaText}
+
+Instructions: "${instructions}"
+
+SCORING RULES:
+- For each item and each criterion, assign ONE label from the ordered list
+- Labels are ordered: first label = appears FIRST in sorted result, last label = appears LAST
+- Choose the label that best describes each item
+
+Output format:
+■0:criterion1=label;criterion2=label■
+■1:criterion1=label;criterion2=label■
+${END}
+
+IMPORTANT:
+- Rank every item (■0 to ■${chunkSize - 1})
+- Use exact criterion names: ${criteriaKeys.join(', ')}
+- Use exact labels from the lists above (lowercase, underscores)
+- Use semicolons (;) between criteria
+- Use equals (=) between criterion and label`,
+      stopSequences: [END],
+      messages: [
+        ...exampleMessages,
+        {
+          type: 'text',
+          role: 'user',
+          content: `Items to rank (■0 to ■${chunkSize - 1}):\n${elementsText}\n\nRank each item using the labeled scales.\nOutput format: ■index:criterion1=label;criterion2=label■\n${END}`,
+        },
+      ],
+      transform: (text) => {
+        const results: ItemScore[] = []
+        const regex = /■(\d+):([^■]+)■/g
+        let match: RegExpExecArray | null
+
+        while ((match = regex.exec(text)) !== null) {
+          const idx = parseInt(match[1] ?? '', 10)
+          const labelsStr = match[2] ?? ''
+
+          if (isNaN(idx) || idx < 0 || idx >= chunkSize) continue
+
+          const scores: Record<string, number> = {}
+          let total = 0
+
+          const pairs = labelsStr.split(';').filter((x) => x.trim().length > 0)
+          for (const pair of pairs) {
+            const [criterion, labelStr] = pair.split('=').map((x) => x.trim().toLowerCase().replace(/\s+/g, '_'))
+            if (!criterion || !labelStr) continue
+
+            // Find the label index in the ordered array (index = score)
+            const labels = sortingCriteria[criterion]?.labels ?? []
+            const labelIndex = labels.findIndex((l) => l === labelStr)
+
+            if (labelIndex >= 0) {
+              // Use index as score (0 = first, higher = later)
+              scores[criterion] = labelIndex
+              total += labelIndex
+            } else {
+              // If label not found, use middle value
+              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5
+              scores[criterion] = middleIndex
+              total += middleIndex
+            }
+          }
+
+          results[idx] = {
+            elementIndex: chunk[idx].index,
+            scores,
+            totalScore: total,
+          }
+        }
+
+        // Fill in missing results with middle scores
+        for (let i = 0; i < chunkSize; i++) {
+          if (!results[i]) {
+            const scores: Record<string, number> = {}
+            let total = 0
+
+            for (const key of criteriaKeys) {
+              const labels = sortingCriteria[key]?.labels ?? []
+              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5
+              scores[key] = middleIndex
+              total += middleIndex
+            }
+
+            results[i] = {
+              elementIndex: chunk[i].index,
+              scores,
+              totalScore: total,
+            }
+          }
+        }
+
+        return results
+      },
+    })
+
+    return extracted
+  }
+
+  // Process all chunks in parallel
+  const limit = pLimit(10)
+  const chunkPromises = chunks.map((chunk) => limit(() => scoreChunk(chunk)))
+  const allScores = await Promise.all(chunkPromises)
+
+  // Phase 4: Merge scores from all chunks
+  // Build a map of elementIndex -> accumulated scores
+  const scoreMap = new Map<number, { scores: Record<string, number>; totalScore: number; tieBreakOrder?: number }>()
+
+  for (const chunkScores of allScores) {
+    for (const itemScore of chunkScores) {
+      const existing = scoreMap.get(itemScore.elementIndex)
+
+      if (existing) {
+        // Average the scores
+        for (const key of criteriaKeys) {
+          existing.scores[key] = (existing.scores[key] + (itemScore.scores[key] ?? 0)) / 2
+        }
+        existing.totalScore = (existing.totalScore + itemScore.totalScore) / 2
+      } else {
+        scoreMap.set(itemScore.elementIndex, {
+          scores: { ...itemScore.scores },
+          totalScore: itemScore.totalScore,
+        })
+      }
+    }
+  }
+
+  // Verify we have scores for all elements
+  if (scoreMap.size !== input.length) {
+    throw new Error(`Score map size mismatch: expected ${input.length}, got ${scoreMap.size}`)
+  }
+
+  // Phase 5: Identify ties
+  const scoreGroups = new Map<number, number[]>()
+  for (const [index, scoreData] of scoreMap.entries()) {
+    const roundedScore = Math.round(scoreData.totalScore * 100) // Round to avoid floating point issues
+    const group = scoreGroups.get(roundedScore) ?? []
+    group.push(index)
+    scoreGroups.set(roundedScore, group)
+  }
+
+  // Find groups with more than one item (ties)
+  const tiedGroups = Array.from(scoreGroups.values()).filter((group) => group.length > 1)
+
+  // Phase 6: Tie-breaking - process all tied groups IN PARALLEL
+  if (tiedGroups.length > 0) {
+    const tieBreakLimit = pLimit(10)
+
+    await Promise.all(
+      tiedGroups.map((tiedIndices) =>
+        tieBreakLimit(async () => {
+          if (tiedIndices.length <= 1) return
+
+          const tiedElements = tiedIndices.map((idx) => elements[idx])
+
+          // Re-score just these items for tie-breaking
+          const tieBreakText = tiedElements
+            .map((elem, i) => {
+              const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem)
+              return `■${i}: ${truncated}■`
+            })
+            .join('\n')
+
+          // Get examples from adapter for tie-breaking
+          const tieBreakInputStr = JSON.stringify(tiedElements.map((e) => e.element))
+          const tieBreakExamples =
+            taskId && ctx.adapter
+              ? await ctx.adapter.getExamples<string, Array<T>>({
+                  input: tieBreakInputStr.slice(0, 1000),
+                  taskType,
+                  taskId,
+                })
+              : []
+
+          // Format examples for few-shot learning in tie-breaking
+          const tieBreakExampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []
+
+          for (const example of tieBreakExamples.slice(0, 3)) {
+            try {
+              const exampleInput = JSON.parse(example.input)
+              const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput]
+
+              tieBreakExampleMessages.push({
+                type: 'text',
+                role: 'user',
+                content: `Expert Example - Items to order:\n${exampleItems.map((el, i) => `■${i}: ${stringify(el, false).slice(0, 200)}■`).join('\n')}\n\nOrder them from first to last.`,
+              })
+
+              if (Array.isArray(example.output) && example.output.length > 0) {
+                const formattedOrder = example.output.map((_, i) => `■${i}■`).join('\n')
+
+                tieBreakExampleMessages.push({
+                  type: 'text',
+                  role: 'assistant',
+                  content: `${formattedOrder}\n${END}`,
+                })
+
+                if (example.explanation) {
+                  tieBreakExampleMessages.push({
+                    type: 'text',
+                    role: 'assistant',
+                    content: `Reasoning: ${example.explanation}`,
+                  })
+                }
+              }
+            } catch {
+              // Skip malformed examples
+            }
+          }
+
+          const { extracted: tieBreakOrder } = await ctx.generateContent({
+            systemPrompt: `You are breaking a tie between items with identical total scores.
+
+Instructions: ${instructions}
+
+Criteria:
+${criteriaKeys
+  .map((key) => {
+    const labels = sortingCriteria[key].labels.join(';')
+    return `- ${key}: ${labels}`
+  })
+  .join('\n')}
+
+Order these ${tiedElements.length} items from FIRST to LAST based on the instructions.
+Earlier labels in each criterion should come FIRST.
+
+Output format:
+■original_index■
+■original_index■
+${END}
+
+Output the indices in the order they should appear (first item at top).`,
+            stopSequences: [END],
+            messages: [
+              ...tieBreakExampleMessages,
+              {
+                type: 'text',
+                role: 'user',
+                content: `Items with identical scores (need tie-breaking):\n${tieBreakText}\n\nOrder them from first to last.\nOutput format: ■index■ (one per line)\n${END}`,
+              },
+            ],
+            transform: (text) => {
+              const order: number[] = []
+              const regex = /■(\d+)■/g
+              let match: RegExpExecArray | null
+
+              while ((match = regex.exec(text)) !== null) {
+                const idx = parseInt(match[1] ?? '', 10)
+                if (!isNaN(idx) && idx >= 0 && idx < tiedElements.length) {
+                  order.push(idx)
+                }
+              }
+
+              // If not all items were ordered, append missing ones
+              for (let i = 0; i < tiedElements.length; i++) {
+                if (!order.includes(i)) {
+                  order.push(i)
+                }
+              }
+
+              return order
+            },
+          })
+
+          // Update scoreMap with tie-break order
+          for (let i = 0; i < tieBreakOrder.length; i++) {
+            const elementIndex = tiedElements[tieBreakOrder[i]].index
+            const scoreData = scoreMap.get(elementIndex)
+            if (scoreData) {
+              scoreData.tieBreakOrder = i
+            }
+          }
+        })
+      )
+    )
+  }
+
+  // Phase 7: Sort by total score, then by tie-break order
+  const sorted = Array.from(scoreMap.entries())
+    .sort((a, b) => {
+      // First sort by total score
+      const scoreDiff = a[1].totalScore - b[1].totalScore
+      if (scoreDiff !== 0) return scoreDiff
+
+      // If scores are equal, use tie-break order
+      const orderA = a[1].tieBreakOrder ?? 0
+      const orderB = b[1].tieBreakOrder ?? 0
+      return orderA - orderB
+    })
+    .map(([index]) => elements[index].element)
+
+  const result = sorted
+
+  // Save example for active learning
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: JSON.stringify(input),
+        instructions,
+      })
+    )
+
+    await ctx.adapter.saveExample({
+      key,
+      taskType,
+      taskId,
+      input: JSON.stringify(input),
+      output: result,
+      instructions,
+      metadata: {
+        cost: { input: 0, output: 0 },
+        latency: 0,
+        model: ctx.modelId,
+        tokens: { input: 0, output: 0 },
+      },
+    })
+  }
+
+  return result
+}
+
+Zai.prototype.sort = function <T>(
+  this: Zai,
+  input: Array<T>,
+  instructions: string,
+  _options?: Options
+): Response<Array<T>, Array<T>> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.sort',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<Array<T>, Array<T>>(
+    context,
+    sort(input, instructions, _options, context),
+    (result) => result // Simplified form is just the sorted array
+  )
+}

packages/zai/src/operations/summarize.ts

@@ -0,0 +1,313 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+
+import { chunk } from 'lodash-es'
+import pLimit from 'p-limit'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+
+import { getTokenizer } from '../tokenizer'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+
+export type Options = {
+  /** What should the text be summarized to? */
+  prompt?: string
+  /** How to format the example text */
+  format?: string
+  /** The length of the summary in tokens */
+  length?: number
+  /** How many times longer (than final length) are the intermediate summaries generated */
+  intermediateFactor?: number
+  /** The maximum number of iterations to perform */
+  maxIterations?: number
+  /** Sliding window options */
+  sliding?: {
+    window: number
+    overlap: number
+  }
+}
+
+const Options = z.object({
+  prompt: z
+    .string()
+    .describe('What should the text be summarized to?')
+    .default('New information, concepts and ideas that are deemed important'),
+  format: z
+    .string()
+    .describe('How to format the example text')
+    .default(
+      'A normal text with multiple sentences and paragraphs. Use markdown to format the text into sections. Use headings, lists, and other markdown features to make the text more readable. Do not include links, images, or other non-text elements.'
+    ),
+  length: z.number().min(10).max(100_000).describe('The length of the summary in tokens').default(250),
+  intermediateFactor: z
+    .number()
+    .min(1)
+    .max(10)
+    .describe('How many times longer (than final length) are the intermediate summaries generated')
+    .default(4),
+  maxIterations: z.number().min(1).default(100),
+  sliding: z
+    .object({
+      window: z.number().min(10).max(100_000),
+      overlap: z.number().min(0).max(100_000),
+    })
+    .describe('Sliding window options')
+    .default({ window: 50_000, overlap: 250 }),
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Summarizes text of any length to a target length using intelligent chunking strategies.
+     *
+     * This operation can handle documents from a few paragraphs to entire books. It uses
+     * two strategies based on document size:
+     * - **Sliding window**: For moderate documents, processes overlapping chunks iteratively
+     * - **Merge sort**: For very large documents, recursively summarizes and merges
+     *
+     * @param original - The text to summarize
+     * @param options - Configuration for length, focus, format, and chunking strategy
+     * @returns Response promise resolving to the summary text
+     *
+     * @example Basic summarization
+     * ```typescript
+     * const article = "Long article text here..."
+     * const summary = await zai.summarize(article, {
+     *   length: 100 // Target 100 tokens (~75 words)
+     * })
+     * ```
+     *
+     * @example Custom focus and format
+     * ```typescript
+     * const meetingNotes = "... detailed meeting transcript ..."
+     * const summary = await zai.summarize(meetingNotes, {
+     *   length: 200,
+     *   prompt: 'Key decisions, action items, and next steps',
+     *   format: 'Bullet points with clear sections for Decisions, Actions, and Next Steps'
+     * })
+     * ```
+     *
+     * @example Summarizing very large documents
+     * ```typescript
+     * const book = await readFile('large-book.txt', 'utf-8') // 100k+ tokens
+     * const summary = await zai.summarize(book, {
+     *   length: 500,
+     *   intermediateFactor: 4, // Intermediate summaries can be 4x target length
+     *   prompt: 'Main themes, key events, and character development'
+     * })
+     * // Automatically uses merge-sort strategy for efficiency
+     * ```
+     *
+     * @example Technical documentation summary
+     * ```typescript
+     * const docs = "... API documentation ..."
+     * const summary = await zai.summarize(docs, {
+     *   length: 300,
+     *   prompt: 'Core API endpoints, authentication methods, and rate limits',
+     *   format: 'Structured markdown with code examples where relevant'
+     * })
+     * ```
+     *
+     * @example Adjusting chunking strategy
+     * ```typescript
+     * const summary = await zai.summarize(document, {
+     *   length: 150,
+     *   sliding: {
+     *     window: 30000,  // Process 30k tokens at a time
+     *     overlap: 500    // 500 token overlap between windows
+     *   }
+     * })
+     * ```
+     *
+     * @example Progress tracking for long documents
+     * ```typescript
+     * const response = zai.summarize(veryLongDocument, { length: 400 })
+     *
+     * response.on('progress', (usage) => {
+     *   console.log(`Progress: ${Math.round(usage.requests.percentage * 100)}%`)
+     *   console.log(`Tokens used: ${usage.tokens.total}`)
+     * })
+     *
+     * const summary = await response
+     * ```
+     */
+    summarize(original: string, options?: Options): Response<string>
+  }
+}
+
+const START = '■START■'
+const END = '■END■'
+
+const summarize = async (original: string, options: Options, ctx: ZaiContext): Promise<string> => {
+  ctx.controller.signal.throwIfAborted()
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  const INPUT_COMPONENT_SIZE = Math.max(100, (model.input.maxTokens - PROMPT_INPUT_BUFFER) / 4)
+  options.prompt = tokenizer.truncate(options.prompt, INPUT_COMPONENT_SIZE)
+  options.format = tokenizer.truncate(options.format, INPUT_COMPONENT_SIZE)
+
+  const maxOutputSize = model.output.maxTokens - PROMPT_OUTPUT_BUFFER
+  if (options.length > maxOutputSize) {
+    throw new Error(
+      `The desired output length is ${maxOutputSize} tokens long, which is more than the maximum of ${model.output.maxTokens} tokens for this model (${model.name})`
+    )
+  }
+
+  // Ensure the sliding window is not bigger than the model input size
+  options.sliding.window = Math.min(options.sliding.window, model.input.maxTokens - PROMPT_INPUT_BUFFER)
+
+  // Ensure the overlap is not bigger than the window
+  // Most extreme case possible (all 3 same size)
+  // |ooooooooooooooooooo|wwwwwwwwwwwwwww|ooooooooooooooooooo|
+  // |<---- overlap ---->|<--  window -->|<---- overlap ---->|
+  options.sliding.overlap = Math.min(options.sliding.overlap, options.sliding.window - 3 * options.sliding.overlap)
+
+  const format = (summary: string, newText: string) => {
+    return `
+${START}
+${summary.length ? summary : '<summary still empty>'}
+${END}
+
+Please amend the summary between the ${START} and ${END} tags to accurately reflect the prompt and the additional text below.
+
+<|start_new_information|>
+${newText}
+<|new_information|>`.trim()
+  }
+
+  const tokens = tokenizer.split(original)
+  const parts = Math.ceil(tokens.length / (options.sliding.window - options.sliding.overlap))
+  let iteration = 0
+
+  // We split it recursively into smaller parts until we're at less than 4 window slides per part
+  // Then we use a merge strategy to combine the sub-chunks summaries
+  // This is basically a merge sort algorithm (but summary instead of sorting)
+  const N = 2 // This is the merge sort exponent
+  const useMergeSort = parts >= Math.pow(2, N)
+  const chunkSize = Math.ceil(tokens.length / (parts * N))
+
+  if (useMergeSort) {
+    const limit = pLimit(10) // Limit to 10 concurrent summarization operations
+    const chunks = chunk(tokens, chunkSize).map((x) => x.join(''))
+    const allSummaries = (await Promise.allSettled(chunks.map((chunk) => limit(() => summarize(chunk, options, ctx)))))
+      .filter((x) => x.status === 'fulfilled')
+      .map((x) => x.value)
+    return summarize(allSummaries.join('\n\n============\n\n'), options, ctx)
+  }
+
+  const summaries: string[] = []
+  let currentSummary = ''
+
+  for (let i = 0; i < tokens.length; i += options.sliding.window) {
+    const from = Math.max(0, i - options.sliding.overlap)
+    const to = Math.min(tokens.length, i + options.sliding.window + options.sliding.overlap)
+    const isFirst = i === 0
+    const isLast = to >= tokens.length
+
+    const slice = tokens.slice(from, to).join('')
+
+    if (iteration++ >= options.maxIterations) {
+      break
+    }
+
+    const instructions: string[] = [
+      `At each step, you will receive a part of the text to summarize. Make sure to reply with the new summary in the tags ${START} and ${END}.`,
+      'Summarize the text and make sure that the main points are included.',
+      'Ignore any unnecessary details and focus on the main points.',
+      'Use short and concise sentences to increase readability and information density.',
+      'When looking at the new information, focus on: ' + options.prompt,
+    ]
+
+    if (isFirst) {
+      instructions.push(
+        'The current summary is empty. You need to generate a summary that covers the main points of the text.'
+      )
+    }
+
+    let generationLength = options.length
+
+    if (!isLast) {
+      generationLength = Math.min(
+        tokenizer.count(currentSummary) + options.length * options.intermediateFactor,
+        maxOutputSize
+      )
+
+      instructions.push(
+        'You need to amend the summary to include the new information. Make sure the summary is complete and covers all the main points.'
+      )
+
+      instructions.push(`The current summary is ${currentSummary.length} tokens long.`)
+      instructions.push(`You can amend the summary to be up to ${generationLength} tokens long.`)
+    }
+
+    if (isLast) {
+      instructions.push(
+        'This is the last part you will have to summarize. Make sure the summary is complete and covers all the main points.'
+      )
+      instructions.push(
+        `The current summary is ${currentSummary.length} tokens long. You need to make sure it is ${options.length} tokens or less.`
+      )
+
+      if (currentSummary.length > options.length) {
+        instructions.push(
+          `The current summary is already too long, so you need to shorten it to ${options.length} tokens while also including the new information.`
+        )
+      }
+    }
+
+    let { extracted: result } = await ctx.generateContent({
+      systemPrompt: `
+You are summarizing a text. The text is split into ${parts} parts, and you are currently working on part ${iteration}.
+At every step, you will receive the current summary and a new part of the text. You need to amend the summary to include the new information (if needed).
+The summary needs to cover the main points of the text and must be concise.
+
+IMPORTANT INSTRUCTIONS:
+${instructions.map((x) => `- ${x.trim()}`).join('\n')}
+
+FORMAT OF THE SUMMARY:
+${options.format}
+`.trim(),
+      messages: [{ type: 'text', content: format(currentSummary, slice), role: 'user' }],
+      maxTokens: generationLength,
+      stopSequences: [END],
+      transform: (text) => {
+        if (!text.trim().length) {
+          throw new Error('The model did not return a valid summary. The response was empty.')
+        }
+
+        return text
+      },
+    })
+
+    if (result.includes(START)) {
+      result = result.slice(result.indexOf(START) + START.length)
+    }
+
+    if (result.includes('■')) {
+      // can happen if the model truncates the text before the entire END tag is written
+      result = result.slice(0, result.indexOf('■'))
+    }
+
+    summaries.push(result)
+    currentSummary = result
+  }
+
+  return currentSummary.trim()
+}
+
+Zai.prototype.summarize = function (this: Zai, original, _options): Response<string> {
+  const options = Options.parse(_options ?? {}) as Options
+
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'summarize',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<string, string>(context, summarize(original, options, context), (value) => value)
+}

packages/zai/src/operations/text.ts

@@ -0,0 +1,142 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+
+import { clamp } from 'lodash-es'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+
+export type Options = {
+  /** The maximum number of tokens to generate */
+  length?: number
+}
+
+const Options = z.object({
+  length: z.number().min(1).max(100_000).optional().describe('The maximum number of tokens to generate'),
+})
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Generates text content based on a natural language prompt.
+     *
+     * This operation creates original text content using LLMs with optional length constraints.
+     * Perfect for generating descriptions, emails, articles, creative content, and more.
+     *
+     * @param prompt - Natural language description of what text to generate
+     * @param options - Optional configuration for text length
+     * @returns Response promise resolving to the generated text
+     *
+     * @example Product description
+     * ```typescript
+     * const description = await zai.text(
+     *   'Write a compelling product description for eco-friendly bamboo toothbrushes'
+     * )
+     * ```
+     *
+     * @example With length constraint
+     * ```typescript
+     * const tagline = await zai.text(
+     *   'Create a catchy tagline for a fitness app',
+     *   { length: 10 } // ~10 tokens (7-8 words)
+     * )
+     * ```
+     *
+     * @example Email generation
+     * ```typescript
+     * const email = await zai.text(`
+     *   Write a professional email to a customer explaining
+     *   that their order will be delayed by 2 days due to weather.
+     *   Apologize and offer a 10% discount on their next purchase.
+     * `, { length: 150 })
+     * ```
+     *
+     * @example Blog post
+     * ```typescript
+     * const blogPost = await zai.text(`
+     *   Write an informative blog post about the benefits of meditation
+     *   for software developers. Include practical tips and scientific research.
+     * `, { length: 500 })
+     * ```
+     *
+     * @example Social media content
+     * ```typescript
+     * const tweet = await zai.text(
+     *   'Write an engaging tweet announcing our new AI-powered chatbot feature',
+     *   { length: 30 } // Twitter-friendly length
+     * )
+     * ```
+     */
+    text(prompt: string, options?: Options): Response<string>
+  }
+}
+
+const text = async (prompt: string, _options: Options | undefined, ctx: ZaiContext): Promise<string> => {
+  ctx.controller.signal.throwIfAborted()
+  const options = Options.parse(_options ?? {})
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  prompt = tokenizer.truncate(prompt, Math.max(model.input.maxTokens - PROMPT_INPUT_BUFFER, 100))
+
+  if (options.length) {
+    options.length = Math.min(model.output.maxTokens - PROMPT_OUTPUT_BUFFER, options.length)
+  }
+
+  const instructions: string[] = []
+  let chart = ''
+
+  if (options.length) {
+    const length = clamp(options.length * 0.75, 5, options.length)
+    instructions.push(`IMPORTANT: Length constraint: ${length} tokens/words`)
+    instructions.push(`The text must be standalone and complete in less than ${length} tokens/words`)
+  }
+
+  if (options.length && options.length <= 500) {
+    chart = `
+| Tokens      | Text Length (approximate)            |
+|-------------|--------------------------------------|
+| < 5 tokens  | 1-3 words                            |
+| 5-10 tokens | 3-6 words                            |
+| 10-20 tokens| 6-15 words                           |
+| 20-50 tokens| A short sentence (15-30 words)       |
+| 50-100 tokens| A medium sentence (30-70 words)     |
+| 100-200 tokens| A short paragraph (70-150 words)   |
+| 200-300 tokens| A medium paragraph (150-200 words) |
+| 300-500 tokens| A long paragraph (200-300 words)   |`.trim()
+  }
+
+  const { extracted } = await ctx.generateContent({
+    systemPrompt: `
+Generate a text that fulfills the user prompt below. Answer directly to the prompt, without any acknowledgements or fluff. Also, make sure the text is standalone and complete.
+${instructions.map((x) => `- ${x}`).join('\n')}
+${chart}
+`.trim(),
+    temperature: 0.7,
+    messages: [{ type: 'text', content: prompt, role: 'user' }],
+    transform: (text) => {
+      if (!text.trim().length) {
+        throw new Error('The model did not return a valid summary. The response was empty.')
+      }
+
+      return text
+    },
+  })
+
+  return extracted
+}
+
+Zai.prototype.text = function (this: Zai, prompt: string, _options?: Options): Response<string> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.text',
+    adapter: this.adapter,
+    memoizer: this._resolveMemoizer(),
+  })
+
+  return new Response<string>(context, text(prompt, _options, context), (result) => result)
+}

packages/zai/src/response.ts

@@ -0,0 +1,381 @@
+import { Usage, ZaiContext } from './context'
+import { EventEmitter } from './emitter'
+
+/**
+ * Event types emitted during operation execution.
+ *
+ * @property progress - Emitted periodically with usage statistics (tokens, cost, requests)
+ * @property complete - Emitted when operation completes successfully with the result
+ * @property error - Emitted when operation fails with the error
+ */
+export type ResponseEvents<TComplete> = {
+  /** Emitted during execution with updated usage statistics */
+  progress: Usage
+  /** Emitted when the operation completes with the full result */
+  complete: TComplete
+  /** Emitted when the operation fails with an error */
+  error: unknown
+}
+
+/**
+ * Promise-like wrapper for Zai operations with observability and control.
+ *
+ * Response provides a dual-value system:
+ * - **Simplified value**: When awaited directly, returns a simplified result (e.g., boolean for `check()`)
+ * - **Full result**: Via `.result()` method, returns `{ output, usage, elapsed }`
+ *
+ * All Zai operations return a Response instance, allowing you to:
+ * - Track progress and usage in real-time
+ * - Abort operations
+ * - Bind to external abort signals
+ * - Get detailed cost and performance metrics
+ *
+ * @template T - The full output type
+ * @template S - The simplified output type (defaults to T)
+ *
+ * @example Basic usage (simplified)
+ * ```typescript
+ * // Simplified result (boolean)
+ * const isPositive = await zai.check(review, 'Is this positive?')
+ * console.log(isPositive) // true or false
+ * ```
+ *
+ * @example Full result with usage
+ * ```typescript
+ * const response = zai.check(review, 'Is this positive?')
+ * const { output, usage, elapsed } = await response.result()
+ *
+ * console.log(output.value) // true/false
+ * console.log(output.explanation) // "The review expresses satisfaction..."
+ * console.log(usage.tokens.total) // 150
+ * console.log(usage.cost.total) // 0.002
+ * console.log(elapsed) // 1234 (ms)
+ * ```
+ *
+ * @example Progress tracking
+ * ```typescript
+ * const response = zai.summarize(longDocument, { length: 500 })
+ *
+ * response.on('progress', (usage) => {
+ *   console.log(`Progress: ${usage.requests.percentage * 100}%`)
+ *   console.log(`Tokens: ${usage.tokens.total}`)
+ *   console.log(`Cost: $${usage.cost.total}`)
+ * })
+ *
+ * const summary = await response
+ * ```
+ *
+ * @example Aborting operations
+ * ```typescript
+ * const response = zai.extract(hugeDocument, schema)
+ *
+ * // Abort after 5 seconds
+ * setTimeout(() => response.abort('Timeout'), 5000)
+ *
+ * try {
+ *   const result = await response
+ * } catch (error) {
+ *   console.log('Operation aborted:', error)
+ * }
+ * ```
+ *
+ * @example External abort signal
+ * ```typescript
+ * const controller = new AbortController()
+ * const response = zai.answer(documents, question).bindSignal(controller.signal)
+ *
+ * // User clicks cancel button
+ * cancelButton.onclick = () => controller.abort()
+ *
+ * const answer = await response
+ * ```
+ *
+ * @example Error handling
+ * ```typescript
+ * const response = zai.extract(text, schema)
+ *
+ * response.on('error', (error) => {
+ *   console.error('Operation failed:', error)
+ * })
+ *
+ * try {
+ *   const result = await response
+ * } catch (error) {
+ *   // Handle error
+ * }
+ * ```
+ */
+export class Response<T, S = T> implements PromiseLike<S> {
+  private _promise: Promise<T>
+  private _eventEmitter: EventEmitter<ResponseEvents<T>>
+  private _context: ZaiContext
+  private _elasped: number | null = null
+  private _simplify: (value: T) => S
+
+  public static reject<T>(context: ZaiContext, err: Error): Response<T, T> {
+    return new Response<T, T>(context, Promise.reject(err), (value) => value)
+  }
+
+  public constructor(context: ZaiContext, promise: Promise<T>, simplify: (value: T) => S) {
+    this._context = context
+    this._eventEmitter = new EventEmitter<ResponseEvents<T>>()
+    this._simplify = simplify
+    this._promise = promise.then(
+      (value) => {
+        this._elasped ||= this._context.elapsedTime
+        this._eventEmitter.emit('complete', value)
+        this._eventEmitter.clear()
+        this._context.clear()
+        return value
+      },
+      (reason) => {
+        this._elasped ||= this._context.elapsedTime
+        this._eventEmitter.emit('error', reason)
+        this._eventEmitter.clear()
+        this._context.clear()
+        throw reason
+      }
+    )
+
+    this._context.on('update', (usage) => {
+      this._eventEmitter.emit('progress', usage)
+    })
+  }
+
+  /**
+   * Subscribes to events emitted during operation execution.
+   *
+   * @param type - Event type: 'progress', 'complete', or 'error'
+   * @param listener - Callback function to handle the event
+   * @returns This Response instance for chaining
+   *
+   * @example Track progress
+   * ```typescript
+   * response.on('progress', (usage) => {
+   *   console.log(`${usage.requests.percentage * 100}% complete`)
+   *   console.log(`Cost: $${usage.cost.total}`)
+   * })
+   * ```
+   *
+   * @example Handle completion
+   * ```typescript
+   * response.on('complete', (result) => {
+   *   console.log('Operation completed:', result)
+   * })
+   * ```
+   *
+   * @example Handle errors
+   * ```typescript
+   * response.on('error', (error) => {
+   *   console.error('Operation failed:', error)
+   * })
+   * ```
+   */
+  public on<K extends keyof ResponseEvents<T>>(type: K, listener: (event: ResponseEvents<T>[K]) => void) {
+    this._eventEmitter.on(type, listener)
+    return this
+  }
+
+  /**
+   * Unsubscribes from events.
+   *
+   * @param type - Event type to unsubscribe from
+   * @param listener - The exact listener function to remove
+   * @returns This Response instance for chaining
+   *
+   * @example
+   * ```typescript
+   * const progressHandler = (usage) => console.log(usage.tokens.total)
+   * response.on('progress', progressHandler)
+   * // Later...
+   * response.off('progress', progressHandler)
+   * ```
+   */
+  public off<K extends keyof ResponseEvents<T>>(type: K, listener: (event: ResponseEvents<T>[K]) => void) {
+    this._eventEmitter.off(type, listener)
+    return this
+  }
+
+  /**
+   * Subscribes to an event for a single emission.
+   *
+   * The listener is automatically removed after being called once.
+   *
+   * @param type - Event type: 'progress', 'complete', or 'error'
+   * @param listener - Callback function to handle the event once
+   * @returns This Response instance for chaining
+   *
+   * @example
+   * ```typescript
+   * response.once('complete', (result) => {
+   *   console.log('Finished:', result)
+   * })
+   * ```
+   */
+  public once<K extends keyof ResponseEvents<T>>(type: K, listener: (event: ResponseEvents<T>[K]) => void) {
+    this._eventEmitter.once(type, listener)
+    return this
+  }
+
+  /**
+   * Binds an external AbortSignal to this operation.
+   *
+   * When the signal is aborted, the operation will be cancelled automatically.
+   * Useful for integrating with UI cancel buttons or request timeouts.
+   *
+   * @param signal - AbortSignal to bind
+   * @returns This Response instance for chaining
+   *
+   * @example With AbortController
+   * ```typescript
+   * const controller = new AbortController()
+   * const response = zai.extract(data, schema).bindSignal(controller.signal)
+   *
+   * // Cancel from elsewhere
+   * cancelButton.onclick = () => controller.abort()
+   * ```
+   *
+   * @example With timeout
+   * ```typescript
+   * const controller = new AbortController()
+   * setTimeout(() => controller.abort('Timeout'), 10000)
+   *
+   * const response = zai.answer(docs, question).bindSignal(controller.signal)
+   * ```
+   */
+  public bindSignal(signal: AbortSignal): this {
+    if (signal.aborted) {
+      this.abort(signal.reason)
+    }
+
+    const signalAbort = () => {
+      this.abort(signal.reason)
+    }
+
+    signal.addEventListener('abort', () => signalAbort())
+
+    void this.once('complete', () => signal.removeEventListener('abort', signalAbort))
+    void this.once('error', () => signal.removeEventListener('abort', signalAbort))
+
+    return this
+  }
+
+  /**
+   * Aborts the operation in progress.
+   *
+   * The operation will be cancelled and throw an abort error.
+   * Any partial results will not be returned.
+   *
+   * @param reason - Optional reason for aborting (string or Error)
+   *
+   * @example
+   * ```typescript
+   * const response = zai.extract(largeDocument, schema)
+   *
+   * // Abort after 5 seconds
+   * setTimeout(() => response.abort('Operation timeout'), 5000)
+   *
+   * try {
+   *   await response
+   * } catch (error) {
+   *   console.log('Aborted:', error)
+   * }
+   * ```
+   */
+  public abort(reason?: string | Error) {
+    this._context.controller.abort(reason)
+  }
+
+  /**
+   * Promise interface - allows awaiting the Response.
+   *
+   * When awaited, returns the simplified value (S).
+   * Use `.result()` for full output with usage statistics.
+   *
+   * @param onfulfilled - Success handler
+   * @param onrejected - Error handler
+   * @returns Promise resolving to simplified value
+   *
+   * @example
+   * ```typescript
+   * // Simplified value
+   * const isPositive = await zai.check(review, 'Is positive?')
+   * console.log(isPositive) // true
+   * ```
+   */
+  // oxlint-disable-next-line no-thenable
+  public then<TResult1 = S, TResult2 = never>(
+    onfulfilled?: ((value: S) => TResult1 | PromiseLike<TResult1>) | null,
+    onrejected?: ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null
+  ): PromiseLike<TResult1 | TResult2> {
+    return this._promise.then(
+      (value: T) => {
+        const simplified = this._simplify(value)
+        return onfulfilled ? onfulfilled(simplified) : simplified
+      },
+      (reason) => {
+        if (onrejected) {
+          return onrejected(reason)
+        }
+        throw reason
+      }
+    ) as PromiseLike<TResult1 | TResult2>
+  }
+
+  /**
+   * Promise interface - handles errors.
+   *
+   * @param onrejected - Error handler
+   * @returns Promise resolving to simplified value or error result
+   */
+  public catch<TResult = never>(
+    onrejected?: ((reason: unknown) => TResult | PromiseLike<TResult>) | null
+  ): PromiseLike<S | TResult> {
+    return this._promise.catch(onrejected) as PromiseLike<S | TResult>
+  }
+
+  /**
+   * Gets the full result with detailed usage statistics and timing.
+   *
+   * Unlike awaiting the Response directly (which returns simplified value),
+   * this method provides:
+   * - `output`: Full operation result (not simplified)
+   * - `usage`: Detailed token usage, cost, and request statistics
+   * - `elapsed`: Operation duration in milliseconds
+   *
+   * @returns Promise resolving to full result object
+   *
+   * @example
+   * ```typescript
+   * const { output, usage, elapsed } = await zai.check(text, condition).result()
+   *
+   * console.log(output.value) // true/false
+   * console.log(output.explanation) // "The text expresses..."
+   * console.log(usage.tokens.total) // 245
+   * console.log(usage.cost.total) // 0.0012
+   * console.log(elapsed) // 1523 (ms)
+   * ```
+   *
+   * @example Usage statistics breakdown
+   * ```typescript
+   * const { usage } = await response.result()
+   *
+   * console.log('Requests:', usage.requests.requests)
+   * console.log('Cached:', usage.requests.cached)
+   * console.log('Input tokens:', usage.tokens.input)
+   * console.log('Output tokens:', usage.tokens.output)
+   * console.log('Input cost:', usage.cost.input)
+   * console.log('Output cost:', usage.cost.output)
+   * console.log('Total cost:', usage.cost.total)
+   * ```
+   */
+  public async result(): Promise<{
+    output: T
+    usage: Usage
+    elapsed: number
+  }> {
+    const output = await this._promise
+    const usage = this._context.usage
+    return { output, usage, elapsed: this._elasped }
+  }
+}

packages/zai/src/thicktoken.d.ts

@@ -0,0 +1,7 @@
+// ThickToken v3 exposes its declarations through package exports, which this
+// package's current `moduleResolution: node` config cannot read. Keep this
+// shim limited to the API Zai imports until Zai can move to a modern resolver.
+declare module '@bpinternal/thicktoken/micro' {
+  export type TextTokenizer = import('./tokenizer').TextTokenizer
+  export const getWasmTokenizer: () => Promise<TextTokenizer>
+}

packages/zai/src/tokenizer.ts

@@ -0,0 +1,25 @@
+import { getWasmTokenizer } from '@bpinternal/thicktoken/micro'
+
+type TruncateMode = 'head' | 'tail' | 'middle'
+type CountOptions = {
+  approximate?: boolean
+}
+
+export type TextTokenizer = {
+  count(text: string, options?: CountOptions): number
+  truncate(text: string, maxTokens: number, mode?: TruncateMode): string
+  split(text: string): string[]
+}
+
+let tokenizer: TextTokenizer | null = null
+
+export async function getTokenizer(): Promise<TextTokenizer> {
+  if (!tokenizer) {
+    try {
+      tokenizer = await getWasmTokenizer()
+    } catch (err) {
+      throw new Error('Failed to initialize ThickToken tokenizer', { cause: err })
+    }
+  }
+  return tokenizer
+}

packages/zai/src/utils.ts

@@ -0,0 +1,45 @@
+export const stringify = (input: unknown, beautify = true) => {
+  return typeof input === 'string' && !!input.length
+    ? input
+    : input
+      ? JSON.stringify(input, beautify ? null : undefined, beautify ? 2 : undefined)
+      : '<input is null, false, undefined or empty>'
+}
+
+export 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
+}
+
+export const takeUntilTokens = <T>(arr: T[], tokens: number, count: (el: T) => number) => {
+  const result: T[] = []
+  let total = 0
+
+  for (const value of arr) {
+    const valueTokens = count(value)
+    if (total + valueTokens > tokens) {
+      break
+    }
+    total += valueTokens
+    result.push(value)
+  }
+
+  return result
+}
+
+export type GenerationMetadata = {
+  model: string
+  cost: {
+    input: number
+    output: number
+  }
+  latency: number
+  tokens: {
+    input: number
+    output: number
+  }
+}

packages/zai/src/zai.ts

@@ -0,0 +1,410 @@
+import { Client } from '@botpress/client'
+import { BotpressClientLike, Cognitive, Model, Models } from '@botpress/cognitive'
+
+import { z } from '@bpinternal/zui'
+
+import { Adapter } from './adapters/adapter'
+import { TableAdapter } from './adapters/botpress-table'
+import { MemoryAdapter } from './adapters/memory'
+import { type TextTokenizer, getTokenizer } from './tokenizer'
+
+/**
+ * A memoizer that caches the result of async operations by a unique key.
+ *
+ * When used with the Botpress ADK workflow `step` function, this enables
+ * Zai operations to resume where they left off if a workflow is interrupted.
+ *
+ */
+export type Memoizer = {
+  run: <T>(id: string, fn: () => Promise<T>) => Promise<T>
+}
+
+/**
+ * Active learning configuration for improving AI operations over time.
+ *
+ * When enabled, Zai stores successful operation results in a table and uses them as examples
+ * for future operations, improving accuracy and consistency.
+ *
+ * @example
+ * ```typescript
+ * const activeLearning = {
+ *   enable: true,
+ *   tableName: 'MyAppLearningTable',
+ *   taskId: 'sentiment-analysis'
+ * }
+ * ```
+ */
+type ActiveLearning = {
+  /** Whether to enable active learning for this Zai instance */
+  enable: boolean
+  /** Name of the Botpress table to store learning examples (must end with 'Table') */
+  tableName: string
+  /** Unique identifier for this learning task */
+  taskId: string
+}
+
+const _ActiveLearning = z.object({
+  enable: z.boolean().describe('Whether to enable active learning').default(false),
+  tableName: z
+    .string()
+    .regex(
+      /^[A-Za-z0-9_/-]{1,100}Table$/,
+      'Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes'
+    )
+    .describe('The name of the table to store active learning tasks')
+    .default('ActiveLearningTable'),
+  taskId: z
+    .string()
+    .regex(
+      /^[A-Za-z0-9_/-]{1,100}$/,
+      'Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes'
+    )
+    .describe('The ID of the task')
+    .default('default'),
+})
+
+/**
+ * Configuration options for creating a Zai instance.
+ *
+ * @example
+ * ```typescript
+ * import { Client } from '@botpress/client'
+ * import { Zai } from '@botpress/zai'
+ *
+ * const client = new Client({ token: 'your-token' })
+ * const config: ZaiConfig = {
+ *   client,
+ *   modelId: 'best', // Use the best available model
+ *   userId: 'user-123',
+ *   namespace: 'my-app',
+ *   activeLearning: {
+ *     enable: true,
+ *     tableName: 'MyLearningTable',
+ *     taskId: 'extraction'
+ *   }
+ * }
+ * const zai = new Zai(config)
+ * ```
+ */
+export type ZaiConfig = {
+  /** Botpress client or Cognitive client instance */
+  client: BotpressClientLike | Cognitive
+  /** Optional user ID for tracking and attribution */
+  userId?: string
+  /**
+   * Model to use: 'best' (default), 'fast', or specific model like 'openai:gpt-4'.
+   * An array can be provided to specify ordered fallback models if the primary
+   * model is unavailable. Server-side fallback is honored on the cognitive-v2
+   * path; the legacy integration path uses the first entry only.
+   */
+  modelId?: Models | Models[]
+  /** Active learning configuration to improve operations over time */
+  activeLearning?: ActiveLearning
+  /** Namespace for organizing tasks (default: 'zai') */
+  namespace?: string
+  /**
+   * Memoizer (or factory returning one) for caching cognitive call results.
+   *
+   * When provided, all LLM calls are wrapped in the memoizer, allowing results
+   * to be cached and replayed. This is useful for resuming workflow runs where
+   * Zai operations have already completed their cognitive calls.
+   *
+   * If a factory function is provided, it is called once per Zai operation invocation.
+   */
+  memoize?: Memoizer | (() => Memoizer)
+}
+
+const _ZaiConfig = z.object({
+  client: z.custom<BotpressClientLike | Cognitive>(),
+  userId: z.string().describe('The ID of the user consuming the API').optional(),
+  modelId: z
+    .custom<Models | Models[]>(
+      (value) => {
+        const isValidSingle = (v: unknown): v is string =>
+          typeof v === 'string' && (v === 'best' || v === 'fast' || v === 'auto' || v.includes(':'))
+
+        if (Array.isArray(value)) {
+          return value.length > 0 && value.every(isValidSingle)
+        }
+
+        return isValidSingle(value)
+      },
+      {
+        message: 'At least one model ID is invalid. Expected a model string or an array of model strings.',
+      }
+    )
+    .describe('The ID of the model you want to use, or an ordered list of fallback models')
+    .default('best' satisfies Models),
+  activeLearning: _ActiveLearning.default({ enable: false }),
+  namespace: z
+    .string()
+    .regex(
+      /^[A-Za-z0-9_/-]{1,100}$/,
+      'Namespace must be alphanumeric and contain only letters, numbers, underscores, hyphens and slashes'
+    )
+    .default('zai'),
+})
+
+/**
+ * Zai - A type-safe LLM utility library for production-ready AI operations.
+ *
+ * Zai provides high-level abstractions for common AI tasks with built-in features like:
+ * - Active learning (learns from successful operations)
+ * - Automatic chunking for large inputs
+ * - Retry logic with error recovery
+ * - Usage tracking (tokens, cost, latency)
+ * - Type-safe schema validation with Zod
+ *
+ * @example Basic usage
+ * ```typescript
+ * import { Client } from '@botpress/client'
+ * import { Zai } from '@botpress/zai'
+ * import { z } from '@bpinternal/zui'
+ *
+ * const client = new Client({ token: process.env.BOTPRESS_TOKEN })
+ * const zai = new Zai({ client })
+ *
+ * // Extract structured data
+ * const schema = z.object({
+ *   name: z.string(),
+ *   age: z.number()
+ * })
+ * const person = await zai.extract('John is 30 years old', schema)
+ * // Output: { name: 'John', age: 30 }
+ *
+ * // Check conditions
+ * const isPositive = await zai.check('I love this product!', 'Is the sentiment positive?')
+ * // Output: true
+ *
+ * // Summarize text
+ * const summary = await zai.summarize(longDocument, { length: 100 })
+ * ```
+ *
+ * @example With active learning
+ * ```typescript
+ * const zai = new Zai({
+ *   client,
+ *   activeLearning: {
+ *     enable: true,
+ *     tableName: 'SentimentTable',
+ *     taskId: 'product-reviews'
+ *   }
+ * })
+ *
+ * // Enable learning for specific task
+ * const result = await zai.learn('sentiment').check(review, 'Is this positive?')
+ * // Future calls will use approved examples for better accuracy
+ * ```
+ *
+ * @example Chaining configuration
+ * ```typescript
+ * // Use fast model for quick operations
+ * const fastZai = zai.with({ modelId: 'fast' })
+ * await fastZai.check(text, 'Is this spam?')
+ *
+ * // Use specific model
+ * const gpt4Zai = zai.with({ modelId: 'openai:gpt-4' })
+ * await gpt4Zai.extract(document, complexSchema)
+ * ```
+ */
+export class Zai {
+  protected static tokenizer: TextTokenizer = null!
+  protected client: Cognitive
+
+  private _originalConfig: ZaiConfig
+
+  private _userId: string | undefined
+
+  protected Model: Models | Models[]
+  protected ModelDetails: Model
+  protected namespace: string
+  protected adapter: Adapter
+  protected activeLearning: ActiveLearning
+  protected _memoize?: Memoizer | (() => Memoizer)
+
+  /**
+   * Creates a new Zai instance with the specified configuration.
+   *
+   * @param config - Configuration object containing client, model, and learning settings
+   *
+   * @example
+   * ```typescript
+   * import { Client } from '@botpress/client'
+   * import { Zai } from '@botpress/zai'
+   *
+   * const client = new Client({ token: 'your-token' })
+   * const zai = new Zai({
+   *   client,
+   *   modelId: 'best',
+   *   namespace: 'my-app',
+   *   userId: 'user-123'
+   * })
+   * ```
+   *
+   * @throws {Error} If the configuration is invalid (e.g., invalid modelId format)
+   */
+  public constructor(config: ZaiConfig) {
+    this._originalConfig = config
+    const parsed = _ZaiConfig.parse(config) as ZaiConfig
+
+    this.client = Cognitive.isCognitiveClient(parsed.client)
+      ? (parsed.client as unknown as Cognitive)
+      : new Cognitive({ client: parsed.client })
+
+    this.namespace = parsed.namespace
+    this._userId = parsed.userId
+    this.Model = parsed.modelId as Models | Models[]
+    this.activeLearning = parsed.activeLearning as ActiveLearning
+
+    this.adapter = parsed.activeLearning?.enable
+      ? new TableAdapter({
+          client: this.client.client as unknown as Client,
+          tableName: parsed.activeLearning.tableName,
+        })
+      : new MemoryAdapter([])
+
+    this._memoize = config.memoize
+  }
+
+  /** @internal */
+  protected async callModel(
+    props: Parameters<Cognitive['generateContent']>[0]
+  ): ReturnType<Cognitive['generateContent']> {
+    return this.client.generateContent({
+      reasoningEffort: 'none',
+      ...props,
+      model: this.Model as Required<Parameters<Cognitive['generateContent']>[0]>['model'],
+      userId: this._userId,
+    })
+  }
+
+  /** @internal */
+  protected _resolveMemoizer(): Memoizer | undefined {
+    if (!this._memoize) {
+      return undefined
+    }
+    return typeof this._memoize === 'function' ? this._memoize() : this._memoize
+  }
+
+  protected async getTokenizer() {
+    Zai.tokenizer ??= await getTokenizer()
+    return Zai.tokenizer
+  }
+
+  protected async fetchModelDetails(): Promise<void> {
+    if (!this.ModelDetails) {
+      // getModelDetails resolves a single model. When a fallback array is
+      // configured, we describe the primary model — fallbacks are only relevant
+      // at the request layer.
+      const primaryModel = Array.isArray(this.Model) ? this.Model[0] : this.Model
+      this.ModelDetails = await this.client.getModelDetails(primaryModel)
+    }
+  }
+
+  protected get taskId() {
+    if (!this.activeLearning.enable) {
+      return undefined
+    }
+
+    return `${this.namespace}/${this.activeLearning.taskId}`.replace(/\/+/g, '/')
+  }
+
+  /**
+   * Creates a new Zai instance with merged configuration options.
+   *
+   * This method allows you to create variations of your Zai instance with different
+   * settings without modifying the original. Useful for switching models, namespaces,
+   * or other configuration on a per-operation basis.
+   *
+   * @param options - Partial configuration to override the current settings
+   * @returns A new Zai instance with the merged configuration
+   *
+   * @example Switch to a faster model
+   * ```typescript
+   * const zai = new Zai({ client })
+   *
+   * // Use fast model for simple operations
+   * const fastZai = zai.with({ modelId: 'fast' })
+   * await fastZai.check(text, 'Is this spam?')
+   *
+   * // Use best model for complex operations
+   * const bestZai = zai.with({ modelId: 'best' })
+   * await bestZai.extract(document, complexSchema)
+   * ```
+   *
+   * @example Change namespace
+   * ```typescript
+   * const customerZai = zai.with({ namespace: 'customer-support' })
+   * const salesZai = zai.with({ namespace: 'sales' })
+   * ```
+   *
+   * @example Use specific model
+   * ```typescript
+   * const gpt4 = zai.with({ modelId: 'openai:gpt-4' })
+   * const claude = zai.with({ modelId: 'anthropic:claude-3-5-sonnet-20241022' })
+   * ```
+   */
+  public with(options: Partial<ZaiConfig>): Zai {
+    return new Zai({
+      ...this._originalConfig,
+      ...options,
+    })
+  }
+
+  /**
+   * Creates a new Zai instance with active learning enabled for a specific task.
+   *
+   * Active learning stores successful operation results and uses them as examples for
+   * future operations, improving accuracy and consistency over time. Each task ID
+   * maintains its own set of learned examples.
+   *
+   * @param taskId - Unique identifier for the learning task (alphanumeric, hyphens, underscores, slashes)
+   * @returns A new Zai instance with active learning enabled for the specified task
+   *
+   * @example Sentiment analysis with learning
+   * ```typescript
+   * const zai = new Zai({
+   *   client,
+   *   activeLearning: {
+   *     enable: false,
+   *     tableName: 'AppLearningTable',
+   *     taskId: 'default'
+   *   }
+   * })
+   *
+   * // Enable learning for sentiment analysis
+   * const sentimentZai = zai.learn('sentiment-analysis')
+   * const result = await sentimentZai.check(review, 'Is this review positive?')
+   *
+   * // Each successful call is stored and used to improve future calls
+   * ```
+   *
+   * @example Different tasks for different purposes
+   * ```typescript
+   * // Extract user info with learning
+   * const userExtractor = zai.learn('user-extraction')
+   * await userExtractor.extract(text, userSchema)
+   *
+   * // Extract product info with separate learning
+   * const productExtractor = zai.learn('product-extraction')
+   * await productExtractor.extract(text, productSchema)
+   *
+   * // Each task learns independently
+   * ```
+   *
+   * @example Combining with other configuration
+   * ```typescript
+   * // Use fast model + learning
+   * const fastLearner = zai.with({ modelId: 'fast' }).learn('quick-checks')
+   * await fastLearner.check(email, 'Is this spam?')
+   * ```
+   *
+   * @see {@link ZaiConfig.activeLearning} for configuration options
+   */
+  public learn(taskId: string) {
+    return new Zai({
+      ...this._originalConfig,
+      activeLearning: { ...this.activeLearning, taskId, enable: true },
+    })
+  }
+}

packages/zai/tsconfig.json

@@ -0,0 +1,12 @@
+{
+  "extends": "../../tsconfig.json",
+  "compilerOptions": {
+    "outDir": "dist",
+    "strict": false,
+    "paths": {
+      "@botpress/zai": ["./src/zai.ts"]
+    }
+  },
+  "exclude": ["node_modules", "dist"],
+  "include": ["src/**/*", "e2e/**/*", "*.ts"]
+}

packages/zai/tsdown.config.ts

@@ -0,0 +1,14 @@
+import { defineConfig } from 'tsdown'
+
+export default defineConfig({
+  entry: ['src/index.ts'],
+  dts: true,
+  outDir: 'dist',
+  platform: 'neutral',
+  clean: true,
+  // Keep the declaration entrypoint bundled so operation module augmentations
+  // like zai.extract/check/summarize are visible from @botpress/zai.
+  unbundle: false,
+  format: 'cjs',
+  target: undefined,
+})

packages/zai/vitest.config.ts

@@ -0,0 +1,12 @@
+import 'dotenv/config'
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    retry: 2, // because LLMs can fail
+    fileParallelism: false,
+    testTimeout: 60_000, // because LLMs can be slow
+    include: ['./e2e/**/*.test.ts'],
+    setupFiles: './vitest.setup.ts',
+  },
+})

packages/zai/vitest.setup.ts

@@ -0,0 +1,53 @@
+import 'dotenv/config'
+import { setupClient } from '@botpress/vai'
+import { beforeAll, afterEach, afterAll } from 'vitest'
+import { server } from './e2e/mocks/server'
+import { getClient } from './e2e/utils'
+
+globalThis.STUDIO = false
+
+// Setup MSW before all tests
+beforeAll(() => {
+  // Start MSW server
+  server.listen({ onUnhandledRequest: 'warn' })
+})
+
+// Reset MSW handlers after each test
+afterEach(() => {
+  server.resetHandlers()
+})
+
+// Clean up MSW after all tests
+afterAll(() => {
+  server.close()
+})
+
+beforeAll(async () => {
+  const token = process.env.CLOUD_PAT
+  if (!token) {
+    throw new Error('Missing CLOUD_PAT')
+  }
+
+  const botId = process.env.CLOUD_BOT_ID
+  if (!botId) {
+    throw new Error('Missing CLOUD_BOT_ID')
+  }
+
+  const client = getClient()
+
+  const { integration: openai } = await client.getPublicIntegration({
+    name: 'openai',
+    version: 'latest',
+  })
+
+  await client.updateBot({
+    id: botId,
+    integrations: {
+      [openai.id]: {
+        enabled: true,
+      },
+    },
+  })
+
+  setupClient(client)
+})