107 lines
4.6 KiB
JavaScript
107 lines
4.6 KiB
JavaScript
const { mockClient } = require('aws-sdk-client-mock');
|
|
const { BedrockRuntimeClient, ConverseStreamCommand } = require('@aws-sdk/client-bedrock-runtime');
|
|
const { createEngine } = require('../src/engine');
|
|
const { historyMessages, prefixBeforeCheckpoint } = require('../src/prompt/build');
|
|
const { MemorySessionStore } = require('../src/sessions');
|
|
const { BookCache } = require('../src/bookCache');
|
|
const { sentenceStream } = require('../src/sentenceStream');
|
|
const { chatConfig, memoryStorage, textStream, usage } = require('./fixture');
|
|
|
|
const bedrock = mockClient(BedrockRuntimeClient);
|
|
|
|
async function collect(iterable) {
|
|
const events = [];
|
|
for await (const event of iterable) events.push(event);
|
|
return events;
|
|
}
|
|
|
|
describe('engine', () => {
|
|
let client;
|
|
|
|
beforeEach(() => {
|
|
bedrock.reset();
|
|
client = new BedrockRuntimeClient({ region: 'ap-south-1' });
|
|
});
|
|
|
|
it('persists an interrupted turn', async () => {
|
|
const abort = new AbortController();
|
|
bedrock.on(ConverseStreamCommand).callsFake(async () => ({
|
|
stream: (async function* () {
|
|
yield { contentBlockDelta: { contentBlockIndex: 0, delta: { text: 'Hello' } } };
|
|
abort.abort();
|
|
yield { contentBlockDelta: { contentBlockIndex: 0, delta: { text: ' more' } } };
|
|
yield { metadata: { usage: usage({ outputTokens: 2 }) } };
|
|
})(),
|
|
}));
|
|
const sessions = new MemorySessionStore();
|
|
const storage = memoryStorage('3');
|
|
const engine = createEngine({
|
|
config: chatConfig(),
|
|
logger: { info() {}, warn() {} },
|
|
storage,
|
|
books: { async get() { return { bookId: 3, tokenCount: 10 }; } },
|
|
sessions,
|
|
bedrock: client,
|
|
opensearch: null,
|
|
bookCache: new BookCache(storage),
|
|
});
|
|
const session = await engine.startSession({ bookId: 3, channel: 'chat', modeOverride: 'full' });
|
|
const events = await collect(engine.answer({ sessionId: session.sessionId, userText: 'Hi', signal: abort.signal }));
|
|
expect(events.find((event) => event.type === 'final').interrupted).toBe(true);
|
|
const saved = await sessions.get(session.sessionId);
|
|
expect(saved.history[1].text).toBe('Hello');
|
|
expect(saved.history[1].interrupted).toBe(true);
|
|
expect(saved.history[1].rawText).toBeUndefined();
|
|
});
|
|
|
|
it('sends markers back on the next turn and leaves the cached prefix unchanged', async () => {
|
|
const inputs = [];
|
|
bedrock.on(ConverseStreamCommand).callsFake(async (input) => {
|
|
inputs.push(input);
|
|
return {
|
|
stream: textStream('It is listed [[p:1]].', usage({ outputTokens: 6 })),
|
|
};
|
|
});
|
|
const sessions = new MemorySessionStore();
|
|
const storage = memoryStorage('3');
|
|
const engine = createEngine({
|
|
config: chatConfig(),
|
|
logger: { info() {}, warn() {} },
|
|
storage,
|
|
books: { async get() { return { bookId: 3, tokenCount: 10, status: 'completed' }; } },
|
|
sessions,
|
|
bedrock: client,
|
|
opensearch: null,
|
|
bookCache: new BookCache(storage),
|
|
});
|
|
const session = await engine.startSession({ bookId: 3, channel: 'chat', modeOverride: 'full' });
|
|
const first = await collect(engine.answer({ sessionId: session.sessionId, userText: 'Where?' }));
|
|
expect(first.filter((event) => event.type === 'text_delta').map((event) => event.text).join('')).not.toContain('[[');
|
|
expect(first.find((event) => event.type === 'final').rawText).toBeUndefined();
|
|
const saved = await sessions.get(session.sessionId);
|
|
expect(saved.history[1].rawText).toContain('[[p:1]]');
|
|
expect(saved.history[1].text).toBe('It is listed.');
|
|
const replay = historyMessages({ history: saved.history });
|
|
expect(replay[1].content[0].text).toContain('[[p:1]]');
|
|
await collect(engine.answer({ sessionId: session.sessionId, userText: 'And the name?' }));
|
|
const prefix = (input) => JSON.stringify(prefixBeforeCheckpoint(input.system));
|
|
expect(prefix(inputs[0])).toBe(prefix(inputs[1]));
|
|
const secondAssistant = inputs[1].messages.find((message) => message.role === 'assistant');
|
|
expect(secondAssistant.content[0].text).toContain('[[p:1]]');
|
|
});
|
|
});
|
|
|
|
describe('sentenceStream', () => {
|
|
it('keeps abbreviations and decimals intact and flushes the tail', async () => {
|
|
async function* deltas() {
|
|
yield { type: 'text_delta', text: 'Dr. Ada left. ' };
|
|
yield { type: 'text_delta', text: 'The value is 3.14 today. Next' };
|
|
}
|
|
const sentences = [];
|
|
for await (const event of sentenceStream(deltas())) {
|
|
if (event.type === 'sentence') sentences.push(event.text);
|
|
}
|
|
expect(sentences).toEqual(['Dr. Ada left.', 'The value is 3.14 today.', 'Next']);
|
|
});
|
|
});
|