104 lines
2.8 KiB
JavaScript
104 lines
2.8 KiB
JavaScript
const { CountTokensCommand } = require('@aws-sdk/client-bedrock-runtime');
|
|
const { BookStatsSchema } = require('./schemas');
|
|
|
|
/**
|
|
* @param {string} text
|
|
*/
|
|
function estimateTokens(text) {
|
|
return Math.ceil(String(text || '').length / 4);
|
|
}
|
|
|
|
/**
|
|
* @param {{ send: Function }} client
|
|
* @param {string} modelId
|
|
* @param {string} text
|
|
*/
|
|
async function countWithBedrock(client, modelId, text) {
|
|
const response = await client.send(
|
|
new CountTokensCommand({
|
|
modelId,
|
|
input: {
|
|
converse: {
|
|
messages: [{ role: 'user', content: [{ text }] }],
|
|
},
|
|
},
|
|
}),
|
|
);
|
|
if (typeof response.inputTokens !== 'number') {
|
|
throw new Error('CountTokens did not return inputTokens');
|
|
}
|
|
return response.inputTokens;
|
|
}
|
|
|
|
/**
|
|
* Count tokens for the marked book and each chapter.
|
|
* Falls back to chars/4 when CountTokens is missing or fails.
|
|
* @param {object} input
|
|
*/
|
|
async function countBook(input) {
|
|
const { client, modelId, markedText, chapters, pages } = input;
|
|
let tokenCountMethod = 'bedrock';
|
|
let tokenCount;
|
|
try {
|
|
if (!client || !modelId) throw new Error('CountTokens client is not configured');
|
|
tokenCount = await countWithBedrock(client, modelId, markedText);
|
|
} catch {
|
|
tokenCountMethod = 'estimate';
|
|
tokenCount = estimateTokens(markedText);
|
|
}
|
|
|
|
const chapterCounts = [];
|
|
for (const chapter of chapters) {
|
|
const text = pages
|
|
.filter((page) => page.page >= chapter.startPage && page.page <= chapter.endPage)
|
|
.map((page) => page.text)
|
|
.join('\n');
|
|
if (tokenCountMethod === 'bedrock') {
|
|
try {
|
|
chapterCounts.push({
|
|
chapterNo: chapter.chapterNo,
|
|
tokenCount: await countWithBedrock(client, modelId, text),
|
|
tokenCountMethod: 'bedrock',
|
|
});
|
|
continue;
|
|
} catch {
|
|
// This chapter falls back on its own.
|
|
}
|
|
}
|
|
chapterCounts.push({
|
|
chapterNo: chapter.chapterNo,
|
|
tokenCount: estimateTokens(text),
|
|
tokenCountMethod: 'estimate',
|
|
});
|
|
}
|
|
|
|
return { tokenCount, tokenCountMethod, chapters: chapterCounts };
|
|
}
|
|
|
|
/**
|
|
* @param {object} input
|
|
*/
|
|
function buildStats(input) {
|
|
const { pages, markedText, tokenResult, cleaning, llm, timingsMs } = input;
|
|
const stats = {
|
|
pageCount: pages.length,
|
|
charCount: markedText.length,
|
|
wordCount: markedText.split(/\s+/).filter(Boolean).length,
|
|
tokenCount: tokenResult.tokenCount,
|
|
tokenCountMethod: tokenResult.tokenCountMethod,
|
|
chapters: tokenResult.chapters,
|
|
cleaning: {
|
|
removedLineCount: cleaning.removedLineCount || 0,
|
|
hyphenFixCount: cleaning.hyphenFixCount || 0,
|
|
},
|
|
llm: {
|
|
inputTokens: llm.inputTokens || 0,
|
|
outputTokens: llm.outputTokens || 0,
|
|
},
|
|
timingsMs,
|
|
};
|
|
return BookStatsSchema.parse(stats);
|
|
}
|
|
|
|
module.exports = { estimateTokens, countWithBedrock, countBook, buildStats };
|