43 lines
1.6 KiB
JavaScript
43 lines
1.6 KiB
JavaScript
const { mockClient } = require('aws-sdk-client-mock');
|
|
const { BedrockRuntimeClient, CountTokensCommand } = require('@aws-sdk/client-bedrock-runtime');
|
|
const { countBook, estimateTokens } = require('../src/tokens');
|
|
|
|
const bedrock = mockClient(BedrockRuntimeClient);
|
|
|
|
describe('token counts', () => {
|
|
beforeEach(() => {
|
|
bedrock.reset();
|
|
});
|
|
|
|
it('falls back to chars/4 when CountTokens fails', async () => {
|
|
bedrock.on(CountTokensCommand).rejects(new Error('model does not support CountTokens'));
|
|
const client = new BedrockRuntimeClient({ region: 'ap-south-1' });
|
|
const marked = '[Page 1]\nHello world\n';
|
|
const result = await countBook({
|
|
client,
|
|
modelId: 'test-count',
|
|
markedText: marked,
|
|
chapters: [{ chapterNo: 1, startPage: 1, endPage: 1 }],
|
|
pages: [{ page: 1, text: 'Hello world' }],
|
|
});
|
|
expect(result.tokenCountMethod).toBe('estimate');
|
|
expect(result.tokenCount).toBe(estimateTokens(marked));
|
|
expect(result.chapters[0].tokenCountMethod).toBe('estimate');
|
|
});
|
|
|
|
it('uses Bedrock inputTokens when the call succeeds', async () => {
|
|
bedrock.on(CountTokensCommand).resolves({ inputTokens: 42 });
|
|
const client = new BedrockRuntimeClient({ region: 'ap-south-1' });
|
|
const result = await countBook({
|
|
client,
|
|
modelId: 'test-count',
|
|
markedText: 'Hello',
|
|
chapters: [{ chapterNo: 1, startPage: 1, endPage: 1 }],
|
|
pages: [{ page: 1, text: 'Hello' }],
|
|
});
|
|
expect(result.tokenCountMethod).toBe('bedrock');
|
|
expect(result.tokenCount).toBe(42);
|
|
expect(result.chapters[0].tokenCount).toBe(42);
|
|
});
|
|
});
|