251 lines
10 KiB
TypeScript
251 lines
10 KiB
TypeScript
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()
|
|
})
|
|
})
|