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