repomix
Version:
A tool to pack repository contents to single file for AI consumption
43 lines (42 loc) • 1.58 kB
JavaScript
import { logger, setLogLevelByWorkerData } from '../../../shared/logger.js';
import { freeTokenCounters, getTokenCounter } from '../tokenCounterFactory.js';
setLogLevelByWorkerData();
export const countTokens = async (task) => {
const processStartAt = process.hrtime.bigint();
try {
const counter = await getTokenCounter(task.encoding);
const tokenCount = counter.countTokens(task.content, task.path);
logger.trace(`Counted tokens. Count: ${tokenCount}. Took: ${getProcessDuration(processStartAt)}ms`);
return tokenCount;
}
catch (error) {
logger.error('Error in token counting worker:', error);
throw error;
}
};
const countTokensBatch = async (task) => {
const processStartAt = process.hrtime.bigint();
try {
const counter = await getTokenCounter(task.encoding);
const results = task.items.map((item) => counter.countTokens(item.content, item.path));
logger.trace(`Counted tokens for ${task.items.length} items. Took: ${getProcessDuration(processStartAt)}ms`);
return results;
}
catch (error) {
logger.error('Error in batch token counting worker:', error);
throw error;
}
};
const getProcessDuration = (startTime) => {
const endTime = process.hrtime.bigint();
return (Number(endTime - startTime) / 1e6).toFixed(2);
};
export default async (task) => {
if ('items' in task) {
return countTokensBatch(task);
}
return countTokens(task);
};
export const onWorkerTermination = async () => {
freeTokenCounters();
};