wehub-resource-sync
115 lines
3.4 KiB
TypeScript
115 lines
3.4 KiB
TypeScript
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)
|
|
})
|
|
})
|