UNPKG

generator-begcode

Version:

Spring Boot + Angular/React/Vue in one handy generator

153 lines (152 loc) 5.76 kB
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\`\`\``, }); }