generator-begcode
Version:
Spring Boot + Angular/React/Vue in one handy generator
153 lines (152 loc) • 5.76 kB
JavaScript
import { AgentVariables, PriorityContainer } from '../../index.js';
import { Rag } from '../../rag/Rag.js';
import { MessageRecombiner } from '../../chunking/MessageRecombiner.js';
export class ContextualizedChat {
_rawChat;
_chunker;
_variables;
_chunks = {
persistent: [],
temporary: [],
};
_rags;
_functionCallResultFirstChunks = {};
_lastTwoMsgs = {
persistent: new PriorityContainer(2, (a, b) => b - a),
temporary: new PriorityContainer(2, (a, b) => b - a),
};
constructor(_context, _rawChat, _chunker, _variables) {
this._rawChat = _rawChat;
this._chunker = _chunker;
this._variables = _variables;
this._rags = {
persistent: Rag.standard(_context).selector(x => x.json),
temporary: Rag.standard(_context).selector(x => x.json),
};
}
get rawChat() {
return this._rawChat;
}
async contextualize(contextVector, tokenLimits) {
await this._processNewMessages();
const persistentLargeChunks = await this._rags.persistent
.query(contextVector)
.recombine(MessageRecombiner.standard(tokenLimits.persistent, this._rawChat.chatLogs, 'persistent', this._lastTwoMsgs.persistent.getItems()));
const temporaryChunks = await this._rags.temporary
.query(contextVector)
.recombine(MessageRecombiner.standard(tokenLimits.temporary, this._rawChat.chatLogs, 'temporary', this._lastTwoMsgs.temporary.getItems()));
const sorted = {
persistent: persistentLargeChunks.map(x => JSON.parse(x.json)),
temporary: temporaryChunks.map(x => JSON.parse(x.json)),
};
sorted.persistent = postProcessMessages(sorted.persistent);
sorted.temporary = postProcessMessages(sorted.temporary);
const chat = this._rawChat.cloneEmpty();
await chat.persistent(sorted.persistent);
await chat.temporary(sorted.temporary);
return chat;
}
async _processNewMessages() {
this._processNewMessagesByType('persistent');
this._processNewMessagesByType('temporary');
}
_processNewMessagesByType(type) {
const messages = this._rawChat.chatLogs.get(type).msgs;
if (messages.length === 0) {
return;
}
const lastProcessedIdx = getLastProcessedMessageIndex(this._chunks[type]);
if (lastProcessedIdx === messages.length - 1) {
return;
}
for (let i = lastProcessedIdx + 1; i < messages.length; ++i) {
const message = messages[i];
this._processNewMessage(message, i, type);
}
}
_processNewMessage(message, msgIdx, type) {
const newChunks = [];
const varName = message.content || '';
let isVariable = false;
if (AgentVariables.hasSyntax(varName)) {
const data = this._variables.get(varName);
if (data) {
message.content = data;
isVariable = true;
}
}
if (this._chunker.shouldChunk(message)) {
if (isVariable) {
newChunks.push(...this._chunker.chunk(message).map((chunk, index) => variableChunkText(chunk, index, varName)));
}
else {
newChunks.push(...this._chunker.chunk(message));
}
}
else if (isVariable) {
newChunks.push(variableChunkText(message, 0, varName));
}
else {
newChunks.push(JSON.stringify(message));
}
const chunks = this._chunks[type];
const startChunkIdx = chunks.length;
for (const _ of newChunks) {
chunks.push({ msgIdx });
}
this._rags[type].addItems(newChunks.map((chunk, index) => ({
chunkIdx: startChunkIdx + index,
msgIdx,
tokens: this._rawChat.tokenizer.encode(chunk).length,
json: chunk,
})));
if (message.role === 'function' || 'function_call' in message) {
this._functionCallResultFirstChunks[startChunkIdx] = {
text: newChunks[0],
metadata: { index: startChunkIdx },
};
}
this._lastTwoMsgs[type].addItem(startChunkIdx);
}
}
function getLastProcessedMessageIndex(chunks) {
const lastIdx = chunks.length - 1;
if (lastIdx < 0) {
return -1;
}
return chunks[lastIdx].msgIdx;
}
function postProcessMessages(messages) {
const result = [];
const varChunkPrefix = 'Variable "${';
const varNameRegex = /\$\{([^}]+)\}/;
let prevVarName;
for (const message of messages) {
let varName;
if (message.content?.startsWith(varChunkPrefix)) {
const match = message.content.match(varNameRegex);
varName = (match && match[1]) || undefined;
}
if (varName && prevVarName === varName) {
const lastMessage = result.at(-1);
lastMessage.content += `\n\n${message.content}`;
result[result.length - 1] = lastMessage;
}
else {
result.push(message);
}
if (prevVarName && !varName) {
const lastMessage = result.at(-1);
lastMessage.content += `\nThe above function result was too large, so it was stored in the variable "\${${prevVarName}}".`;
}
prevVarName = varName;
}
return result;
}
function variableChunkText(chunk, index, varName) {
const message = typeof chunk === 'string' ? JSON.parse(chunk) : chunk;
return JSON.stringify({
...message,
content: `Variable "${varName}" chunk #${index}\n\`\`\`\n${message.content}\n\`\`\``,
});
}