90 lines
3.2 KiB
TypeScript
90 lines
3.2 KiB
TypeScript
import "dotenv/config";
|
|
|
|
import { loadConfig } from "@voidbot/config";
|
|
import {
|
|
createTextEmbedder,
|
|
FileMessageArchiveRepository,
|
|
FileSourceDocumentArchiveRepository,
|
|
HistoryIngester,
|
|
SourceDocumentIngester,
|
|
createVectorStores,
|
|
} from "@voidbot/rag";
|
|
|
|
const REBUILD_CHUNK_BATCH_SIZE = 64;
|
|
|
|
async function main(): Promise<void> {
|
|
if (!process.env.DISCORD_OWNER_ID) {
|
|
process.env.DISCORD_OWNER_ID = "__rag_rebuild__";
|
|
}
|
|
|
|
const config = loadConfig();
|
|
const archiveRepository = new FileMessageArchiveRepository(config.ragArchivePath);
|
|
const sourceArchiveRepository = new FileSourceDocumentArchiveRepository(config.ragSourceArchivePath);
|
|
const embedder = createTextEmbedder({
|
|
backend: config.ragEmbeddingBackend,
|
|
hashDimensions: config.ragEmbeddingDimensions,
|
|
ollamaBaseUrl: config.ragOllamaBaseUrl,
|
|
ollamaModel: config.ragOllamaModel,
|
|
ollamaTimeoutMs: config.ragOllamaTimeoutMs,
|
|
queryInstruction: config.ragQueryInstruction,
|
|
});
|
|
const vectorStores = createVectorStores({
|
|
kind: config.vectorStore.kind,
|
|
historyPath: config.vectorStore.path,
|
|
sourceRoot: config.sourceVectorStoreRoot,
|
|
qdrant: config.qdrant,
|
|
historyEmbedder: embedder,
|
|
sourceEmbedder: embedder,
|
|
});
|
|
const vectorStore = vectorStores.history;
|
|
const sourceVectorStore = vectorStores.source;
|
|
const historyIngester = new HistoryIngester();
|
|
const sourceIngester = new SourceDocumentIngester();
|
|
const activeMessages = await archiveRepository.listAllActive();
|
|
const historyChunks = historyIngester.chunkMessages(activeMessages);
|
|
const sourceDocuments = await sourceArchiveRepository.listAll();
|
|
const sourceChunks = sourceIngester.chunkDocuments(sourceDocuments);
|
|
|
|
await vectorStore.clear();
|
|
await sourceVectorStore.clear();
|
|
|
|
for (let index = 0; index < historyChunks.length; index += REBUILD_CHUNK_BATCH_SIZE) {
|
|
const batch = historyChunks.slice(index, index + REBUILD_CHUNK_BATCH_SIZE);
|
|
|
|
if (batch.length === 0) {
|
|
continue;
|
|
}
|
|
|
|
await vectorStore.upsert(batch);
|
|
|
|
if ((index / REBUILD_CHUNK_BATCH_SIZE + 1) % 25 === 0 || index + batch.length >= historyChunks.length) {
|
|
console.log(`History rebuild progress: ${Math.min(index + batch.length, historyChunks.length)}/${historyChunks.length} chunks`);
|
|
}
|
|
}
|
|
|
|
for (let index = 0; index < sourceChunks.length; index += REBUILD_CHUNK_BATCH_SIZE) {
|
|
const batch = sourceChunks.slice(index, index + REBUILD_CHUNK_BATCH_SIZE);
|
|
|
|
if (batch.length === 0) {
|
|
continue;
|
|
}
|
|
|
|
await sourceVectorStore.upsert(batch);
|
|
|
|
if ((index / REBUILD_CHUNK_BATCH_SIZE + 1) % 25 === 0 || index + batch.length >= sourceChunks.length) {
|
|
console.log(`Source rebuild progress: ${Math.min(index + batch.length, sourceChunks.length)}/${sourceChunks.length} chunks`);
|
|
}
|
|
}
|
|
|
|
console.log(`Embedder: ${embedder.id}`);
|
|
console.log(`Archived active messages: ${activeMessages.length}`);
|
|
console.log(`Archived source documents: ${sourceDocuments.length}`);
|
|
console.log(`History chunks rebuilt: ${historyChunks.length}`);
|
|
console.log(`Source chunks rebuilt: ${sourceChunks.length}`);
|
|
console.log(`Indexed chunks rebuilt: ${historyChunks.length + sourceChunks.length}`);
|
|
}
|
|
|
|
void main().catch((error) => {
|
|
console.error(error);
|
|
process.exitCode = 1;
|
|
});
|