87 lines
3.6 KiB
JavaScript
87 lines
3.6 KiB
JavaScript
const { mockClient } = require('aws-sdk-client-mock');
|
|
const { BedrockRuntimeClient, ConverseStreamCommand } = require('@aws-sdk/client-bedrock-runtime');
|
|
const { runFullContext } = require('../src/strategies/fullContext');
|
|
const { parseCitations } = require('../src/citations');
|
|
const { chatConfig, sampleBook, 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('full context', () => {
|
|
let client;
|
|
|
|
beforeEach(() => {
|
|
bedrock.reset();
|
|
client = new BedrockRuntimeClient({ region: 'ap-south-1' });
|
|
});
|
|
|
|
it('reports cache reads on the second turn', async () => {
|
|
let call = 0;
|
|
bedrock.on(ConverseStreamCommand).callsFake(async () => {
|
|
call += 1;
|
|
return {
|
|
stream: textStream('In chapter 1 (p. 2).', usage({
|
|
cacheReadInputTokens: call === 1 ? 0 : 8000,
|
|
cacheWriteInputTokens: call === 1 ? 8000 : 0,
|
|
})),
|
|
};
|
|
});
|
|
const session = { channel: 'chat', history: [], rollingSummary: '', lastCacheTouchAt: null, keepaliveCount: 0, sessionId: 's' };
|
|
const ctx = {
|
|
config: chatConfig(),
|
|
logger: { info() {}, warn() {} },
|
|
bedrock: client,
|
|
book: sampleBook(),
|
|
session,
|
|
userText: 'Where is that?',
|
|
};
|
|
const first = await collect(runFullContext(ctx));
|
|
session.history = [{ role: 'user', text: 'Hi' }, { role: 'assistant', text: 'Hello' }];
|
|
session.lastCacheTouchAt = new Date().toISOString();
|
|
const second = await collect(runFullContext({ ...ctx, userText: 'And then?' }));
|
|
const finalEvent = second.find((event) => event.type === 'final');
|
|
expect(finalEvent.usage.cacheReadTokens).toBe(8000);
|
|
expect(first.find((event) => event.type === 'final').usage.cacheWriteTokens).toBe(8000);
|
|
expect(finalEvent.usage.estimatedCostUsd).toBeGreaterThan(0);
|
|
});
|
|
|
|
it('drops citations for pages that are not in the book', () => {
|
|
const items = parseCitations('See (p. 2) and (p. 99).', sampleBook());
|
|
expect(items.map((item) => item.page)).toEqual([2]);
|
|
});
|
|
|
|
it('strips markers before text deltas and emits a citation chip', async () => {
|
|
bedrock.on(ConverseStreamCommand).resolves({
|
|
stream: (async function* () {
|
|
yield { contentBlockDelta: { contentBlockIndex: 0, delta: { text: 'The dose is 50 mg ' } } };
|
|
yield { contentBlockDelta: { contentBlockIndex: 0, delta: { text: '[[p:2-4]]' } } };
|
|
yield { contentBlockDelta: { contentBlockIndex: 0, delta: { text: '.' } } };
|
|
yield { messageStop: { stopReason: 'end_turn' } };
|
|
yield { metadata: { usage: usage({ outputTokens: 12 }) } };
|
|
})(),
|
|
});
|
|
const events = await collect(runFullContext({
|
|
config: chatConfig(),
|
|
logger: { info() {}, warn() {} },
|
|
bedrock: client,
|
|
book: sampleBook(),
|
|
session: { channel: 'chat', history: [], rollingSummary: '', lastCacheTouchAt: null, keepaliveCount: 0, sessionId: 's' },
|
|
userText: 'What was the dose?',
|
|
}));
|
|
const spoken = events.filter((event) => event.type === 'text_delta').map((event) => event.text).join('');
|
|
expect(spoken).toBe('The dose is 50 mg.');
|
|
expect(spoken).not.toContain('[[');
|
|
const citations = events.find((event) => event.type === 'citations');
|
|
expect(citations.items[0]).toMatchObject({ page: 2, pageEnd: 4, chapterNo: 1 });
|
|
expect(citations.items[0].verified).toBeUndefined();
|
|
const finalEvent = events.find((event) => event.type === 'final');
|
|
expect(finalEvent.text).toBe('The dose is 50 mg.');
|
|
expect(finalEvent.rawText).toContain('[[p:2-4]]');
|
|
});
|
|
});
|