generator-begcode
Version:
Spring Boot + Angular/React/Vue in one handy generator
99 lines (98 loc) • 3.9 kB
JavaScript
const calculateTotalTokens = (chunks) => {
return chunks.reduce((acc, chunk) => acc + chunk.tokens, 0);
};
const tryGetFunctionMessagePair = (chunk, chunkIdx, originalItems, chatLogs, chatLogType) => {
const msg = JSON.parse(chunk.json);
if ((msg.role === 'assistant' && 'function_call' in msg) || msg.role === 'function') {
const direction = msg.role === 'assistant' ? 1 : -1;
const search = msg.role === 'assistant' ? 'function' : 'assistant';
let funcChunkIdx;
let nextChunkIdx = chunkIdx + direction;
while (!funcChunkIdx) {
if (originalItems.length >= nextChunkIdx || nextChunkIdx < 0) {
break;
}
const msg = chatLogs.getMsg(chatLogType, originalItems[nextChunkIdx].msgIdx);
if (msg?.role === search) {
funcChunkIdx = nextChunkIdx;
}
else {
nextChunkIdx += direction;
}
}
if (funcChunkIdx) {
return funcChunkIdx;
}
}
return undefined;
};
export class MessageRecombiner {
static standard(tokenLimit, chatLogs, chatLogType, initChunks) {
return async (results, originalItems) => {
if (calculateTotalTokens(originalItems) < tokenLimit) {
return originalItems;
}
const chunks = [];
const chunksAdded = new Set();
const msgsAdded = {};
let tokenCount = 0;
const canAddTokens = (tokens) => tokenCount + tokens <= tokenLimit;
const addChunk = (chunkIdx, recursive = false) => {
const chunk = originalItems[chunkIdx];
if (chunksAdded.has(chunkIdx)) {
return true;
}
chunksAdded.add(chunkIdx);
if (!canAddTokens(chunk.tokens)) {
return false;
}
if (!msgsAdded[chunk.msgIdx]) {
msgsAdded[chunk.msgIdx] = 0;
}
if (msgsAdded[chunk.msgIdx] >= 4) {
return true;
}
msgsAdded[chunk.msgIdx] += 1;
chunks.push(chunk);
tokenCount += chunk.tokens;
const firstChunkIdx = originalItems.slice(0, chunkIdx).findIndex(item => item.msgIdx === chunk.msgIdx);
if (firstChunkIdx !== -1) {
addChunk(firstChunkIdx);
}
if (!recursive) {
const funcPairChunk = tryGetFunctionMessagePair(chunk, chunkIdx, originalItems, chatLogs, chatLogType);
if (funcPairChunk) {
return addChunk(funcPairChunk, true);
}
}
return true;
};
const iterator = await results();
if (chatLogType === 'persistent') {
const smallPersistentChunks = [];
let lastMsgIndex = -1;
for (const chunk of originalItems) {
if (chunk.msgIdx !== lastMsgIndex) {
smallPersistentChunks.push(chunk);
lastMsgIndex = chunk.msgIdx;
}
else {
smallPersistentChunks.pop();
break;
}
}
for (const chunk of smallPersistentChunks) {
addChunk(chunk.chunkIdx);
}
}
for (const chunk of initChunks) {
addChunk(chunk);
}
for await (const result of iterator) {
const chunkIdx = result.metadata().index;
addChunk(chunkIdx);
}
return chunks.sort((a, b) => a.chunkIdx - b.chunkIdx);
};
}
}