77 lines
2.3 KiB
JavaScript
77 lines
2.3 KiB
JavaScript
const { mockClient } = require('aws-sdk-client-mock');
|
|
const { BedrockRuntimeClient, ConverseCommand } = require('@aws-sdk/client-bedrock-runtime');
|
|
const { complete } = require('../src/llm');
|
|
const { ChapterSummarySchema } = require('../src/schemas');
|
|
|
|
const bedrock = mockClient(BedrockRuntimeClient);
|
|
|
|
function words(count) {
|
|
return Array.from({ length: count }, (_, index) => `word${index}`).join(' ');
|
|
}
|
|
|
|
function summary(overrides = {}) {
|
|
return {
|
|
chapterNo: 1,
|
|
title: 'Harbor',
|
|
pageStart: 1,
|
|
pageEnd: 3,
|
|
oneLine: 'Ships leave the harbor.',
|
|
summary: words(120),
|
|
keyPoints: ['One', 'Two', 'Three'],
|
|
entities: { people: ['Ada'], places: ['Harbor'], terms: ['Tide'] },
|
|
questionsAnswered: ['Who leaves?', 'Where?', 'When?'],
|
|
...overrides,
|
|
};
|
|
}
|
|
|
|
function toolResponse(input, name = 'chapter_summary') {
|
|
return {
|
|
output: { message: { content: [{ toolUse: { name, input } }] } },
|
|
usage: { inputTokens: 11, outputTokens: 7 },
|
|
};
|
|
}
|
|
|
|
describe('llm repair', () => {
|
|
let client;
|
|
|
|
beforeEach(() => {
|
|
bedrock.reset();
|
|
client = new BedrockRuntimeClient({ region: 'ap-south-1' });
|
|
});
|
|
|
|
it('retries once when the first tool input fails validation', async () => {
|
|
bedrock
|
|
.on(ConverseCommand)
|
|
.resolvesOnce(toolResponse({ ...summary(), summary: 'too short' }))
|
|
.resolvesOnce(toolResponse(summary()));
|
|
const result = await complete({
|
|
client,
|
|
modelId: 'test-model',
|
|
system: 'system',
|
|
user: '<book-text>\nHello\n</book-text>',
|
|
schema: ChapterSummarySchema,
|
|
toolName: 'chapter_summary',
|
|
jsonSchema: { type: 'object' },
|
|
});
|
|
expect(result.data.summary.split(/\s+/)).toHaveLength(120);
|
|
expect(result.usage.inputTokens).toBe(22);
|
|
expect(bedrock.commandCalls(ConverseCommand)).toHaveLength(2);
|
|
});
|
|
|
|
it('throws after the repair attempt is still invalid', async () => {
|
|
bedrock.on(ConverseCommand).resolves(toolResponse({ ...summary(), summary: 'still short' }));
|
|
await expect(
|
|
complete({
|
|
client,
|
|
modelId: 'test-model',
|
|
system: 'system',
|
|
user: 'text',
|
|
schema: ChapterSummarySchema,
|
|
toolName: 'chapter_summary',
|
|
jsonSchema: { type: 'object' },
|
|
}),
|
|
).rejects.toThrow(/after repair/);
|
|
expect(bedrock.commandCalls(ConverseCommand)).toHaveLength(2);
|
|
});
|
|
});
|