"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
    var ownKeys = function(o) {
        ownKeys = Object.getOwnPropertyNames || function (o) {
            var ar = [];
            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
            return ar;
        };
        return ownKeys(o);
    };
    return function (mod) {
        if (mod && mod.__esModule) return mod;
        var result = {};
        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
        __setModuleDefault(result, mod);
        return result;
    };
})();
Object.defineProperty(exports, "__esModule", { value: true });
exports.cosineSimilarity = void 0;
exports.loadEmbedding = loadEmbedding;
exports.getConfiguredEmbeddingMap = getConfiguredEmbeddingMap;
exports.getConfiguredEmbeddings = getConfiguredEmbeddings;
exports.detectDeletedEmbeddingFiles = detectDeletedEmbeddingFiles;
exports.embedSource = embedSource;
exports.embed = embed;
exports.isEmbeddingFile = isEmbeddingFile;
exports.embedJson = embedJson;
exports.embedKind = embedKind;
exports.handleAllKinds = handleAllKinds;
exports.saveEmbedding = saveEmbedding;
exports.pruneEmbedding = pruneEmbedding;
exports.pruneVector = pruneVector;
exports.pruneMetadata = pruneMetadata;
exports.queryEmbedding = queryEmbedding;
const glob_1 = require("glob");
const path = __importStar(require("path"));
const config_1 = require("./config");
const types_1 = require("./types");
const utils_1 = require("./utils");
Object.defineProperty(exports, "cosineSimilarity", { enumerable: true, get: function () { return utils_1.cosineSimilarity; } });
const ai_1 = require("./ai");
const services_1 = require("./services");
const hashes_1 = require("./hashes");
const conversion_1 = require("./conversion");
const clients_1 = require("./clients");
async function loadEmbedding(filePath) {
    if (await (0, utils_1.fileExists)(filePath)) {
        try {
            const parsed = JSON.parse(await (0, utils_1.readFile)(filePath, "utf8"));
            return parsed;
        }
        catch (e) {
            console.error("Error loading embedding file", filePath, e);
            return [];
        }
    }
    return [];
}
async function getConfiguredEmbeddingMap() {
    const config = await (0, config_1.getConfig)();
    const embedSources = config.embedSources || [];
    const files = Array.from(new Set(embedSources.map((s) => s.output)));
    const embeddings = {};
    for (const file of files) {
        if (!embeddings[file]) {
            embeddings[file] = [];
        }
        const fileEmbeddings = await loadEmbedding(file);
        embeddings[file].push(...fileEmbeddings);
    }
    return embeddings;
}
async function getConfiguredEmbeddings() {
    const config = await (0, config_1.getConfig)();
    const embedSources = config.embedSources || [];
    const files = Array.from(new Set(embedSources.map((s) => s.output)));
    const embeddings = [];
    for (const file of files) {
        const fileEmbeddings = await loadEmbedding(file);
        embeddings.push(...fileEmbeddings);
    }
    return embeddings;
}
function getChunkId(id, index, chunkSize) {
    const split = id.split("-");
    if (split.length > 1 &&
        Number.isInteger(parseInt(split[split.length - 1], 10))) {
        return id;
    }
    return chunkSize ? `${id}-${index}` : id;
}
function parseIdFromChunk(chunkName) {
    const split = chunkName.split("-");
    if (split.length > 1 &&
        Number.isInteger(parseInt(split[split.length - 1], 10))) {
        return split.slice(0, -1).join("-");
    }
    return chunkName;
}
async function detectDeletedEmbeddingFiles(source, ignorePattern, embeddings) {
    if (source.kind !== "file") {
        return;
    }
    const inputs = (await (0, glob_1.globSync)(source.input, {
        ignore: ignorePattern,
    }));
    for (const embedding of embeddings) {
        const id = parseIdFromChunk(embedding.id);
        const dirName = path.dirname(id);
        const exists = await (0, utils_1.fileExists)(path.resolve(id));
        const hasSomeFilesInDir = inputs.some((input) => input.startsWith(dirName));
        if (hasSomeFilesInDir && !exists) {
            console.log("Detected deleted embedding file", embedding.id);
            const index = embeddings.findIndex((e) => e.id === embedding.id);
            if (index !== -1) {
                embeddings.splice(index, 1);
            }
        }
    }
}
async function embedSource(model, source, ignorePattern) {
    if (!source.input) {
        console.log("Skipping", source.output, "with blank input property");
        return;
    }
    console.log("Embedding", source.input.slice(0, 100), "... to", source.output);
    let inputs = [];
    const kind = source.kind || "file";
    if (kind === "file") {
        inputs = await (0, glob_1.globSync)(source.input, { ignore: ignorePattern });
    }
    if (source.kind && inputs.length === 0) {
        inputs = [source.input];
    }
    console.log(`Checking ${inputs.length} files`);
    if (inputs.length > 1000) {
        console.error("Large number of files detected. This may take a while");
    }
    const embeddings = await loadEmbedding(source.output);
    let batch = [];
    let index = 0;
    for (const file of inputs) {
        index++;
        const shouldSave = batch.length > 20;
        if (shouldSave) {
            await Promise.all(batch);
            await saveEmbedding(source.output, embeddings);
            batch = [];
        }
        batch.push(embedKind(model, file, source, embeddings));
    }
    if (batch.length > 0) {
        await Promise.all(batch);
    }
    await detectDeletedEmbeddingFiles(source, ignorePattern, embeddings);
    await saveEmbedding(source.output, embeddings);
}
async function embed(model, id, text, metadata, embeddings, prompt, chunkSize, uploadMode, minLength) {
    let chunks = [text];
    if (!text) {
        console.log("Skipping", id, "with blank text");
        return [];
    }
    if (chunkSize) {
        chunks = await (0, ai_1.chunkText)(text, chunkSize);
    }
    const MAX_CHUNKS = 300;
    if (chunks.length <= MAX_CHUNKS) {
        chunks = chunks.slice(0, MAX_CHUNKS);
    }
    const dontPrune = [];
    const updates = new Array();
    for (let index = 0; index < chunks.length; index++) {
        const chunkId = getChunkId(id, index, chunkSize);
        let textOfChunk = chunks[index];
        const tooShort = minLength && textOfChunk.length < minLength;
        if (tooShort) {
            console.log("Skipping (too short)", chunkId);
            continue;
        }
        dontPrune.push(chunkId);
        const foundMany = embeddings.filter((e) => e.id === chunkId);
        const found = foundMany.length > 0 ? foundMany[0] : null;
        const alreadyEmbedded = found && found.text === textOfChunk;
        if (foundMany.length > 1) {
            console.log(`Removing ${foundMany.length} duplicates for`, chunkId);
            let toDelete = embeddings.findIndex((e) => e.id === chunkId);
            while (toDelete !== -1) {
                embeddings.splice(toDelete, 1);
                toDelete = embeddings.findIndex((e) => e.id === chunkId);
            }
            if (alreadyEmbedded) {
                embeddings.push({
                    ...found,
                });
                updates.push(chunkId);
            }
        }
        if (alreadyEmbedded) {
            continue;
        }
        if (prompt) {
            console.log("Summarizing", textOfChunk);
            textOfChunk = await (0, ai_1.summarizeTexts)([textOfChunk], prompt);
        }
        let vector = [];
        if (!uploadMode) {
            console.log("Embedding", chunkId);
            const providerEmbeddings = await clients_1.Clients.createEmbedding("", {
                input: textOfChunk,
                model: model || types_1.EmbeddingModels.openai.EmbeddingAda2,
            });
            vector = providerEmbeddings.data[0].embedding;
        }
        const embeddable = {
            id: chunkId,
            text: textOfChunk,
            vector,
            metadata: {
                ...metadata,
                ...(prompt && { text: ai_1.chunkText }),
            },
        };
        embeddings.push(embeddable);
        updates.push(chunkId);
    }
    pruneEmbedding(id, dontPrune, embeddings);
    return updates;
}
async function isEmbeddingFile(inputFile) {
    const filePath = path.join(__dirname, inputFile);
    const sourceJson = JSON.parse(await (0, utils_1.readFile)(inputFile, "utf8"));
    const isEmbedding = Array.isArray(sourceJson) &&
        sourceJson.every((e) => e.id && e.text && e.metadata);
    console.log(`Checking file ${inputFile} for embeddings: ${isEmbedding}`);
    return isEmbedding;
}
async function embedJson(model, inputFile, source) {
    const { prompt, output, uploadMode, chunkSize } = source;
    const filePath = path.join(__dirname, inputFile);
    const sourceJson = JSON.parse(await (0, utils_1.readFile)(inputFile, "utf8"));
    const embeddings = await loadEmbedding(output);
    const updates = [];
    let batch = [];
    for (const row of sourceJson) {
        if (embeddings.find((e) => e.id === row.id)) {
            continue;
        }
        console.log("Embedding", row.id);
        batch.push(embed(model, row.id, row.text, row.metadata, embeddings, prompt, chunkSize, uploadMode, source.minLength));
        let embedded = [];
        if (batch.length > 20) {
            embedded = (await Promise.all(batch)).flat();
            batch = [];
        }
        updates.push(...embedded);
    }
    await Promise.all(batch);
}
async function embedKind(model, id, source, embeddings = []) {
    const { prompt, output, uploadMode, chunkSize } = source;
    const stat = !source.kind && (await (0, utils_1.fileStat)(id));
    if (stat && stat.isDirectory()) {
        return;
    }
    if (id.endsWith(".json") && (await isEmbeddingFile(id))) {
        console.log("Embedding JSON", id);
        return embedJson(model, id, source);
    }
    const toEmbed = await handleAllKinds(id, source);
    const updates = [];
    for (const row of toEmbed) {
        const { id: rowId, text, metadata } = row;
        const embedded = await embed(model, rowId, text, metadata, embeddings, prompt, chunkSize, uploadMode);
        updates.push(...embedded);
    }
}
async function handleFileKind(filePath) {
    return [
        {
            id: filePath,
            text: await (0, conversion_1.convertToText)(filePath),
            metadata: {
                filepath: filePath,
                date: new Date().toISOString(),
            },
        },
    ];
}
async function handleTextKind(contents) {
    const hash = await (0, hashes_1.md5Hash)(contents);
    return [
        {
            id: hash,
            text: contents,
            metadata: {
                date: new Date().toISOString(),
            },
        },
    ];
}
async function handleAllKinds(id, source) {
    const { input, kind } = source;
    const contents = "";
    const ids = [];
    const { Plugins } = (0, services_1.services)();
    if (Plugins.isPlugin(kind)) {
        console.log("Embedding with plugin", kind);
        return Plugins.embed(kind, input);
    }
    switch (kind) {
        case "text":
            return handleTextKind(source.input);
        case "file":
        default:
            return handleFileKind(id);
    }
}
async function saveEmbedding(output, embeddings) {
    const fileString = "[" +
        embeddings
            .sort((a, b) => (a.id < b.id ? -1 : 1))
            .map((e) => JSON.stringify(e))
            .join(",") +
        "]";
    await (0, utils_1.writeFile)(output, fileString);
}
function pruneEmbedding(id, chunkIds, embeddings) {
    const relatedChunks = embeddings.filter((e) => e.id.startsWith(id));
    for (const chunk of relatedChunks) {
        if (!chunkIds.includes(chunk.id)) {
            console.log("Removing", chunk.id);
            const index = embeddings.findIndex((e) => e.id === chunk.id);
            embeddings.splice(index, 1);
        }
    }
    return embeddings;
}
function pruneVector(embeddings) {
    for (const entry of embeddings) {
        delete entry.vector;
    }
}
function pruneMetadata(embeddings, characterLimit = 5000) {
    for (const entry of embeddings) {
        for (const key of Object.keys(entry.metadata)) {
            if (JSON.stringify(entry.metadata[key]).length > characterLimit) {
                delete entry.metadata[key];
            }
        }
    }
}
async function queryEmbedding(query, embeddings, model = types_1.EmbeddingModels.openai.EmbeddingAda2) {
    const providerEmbeddings = await clients_1.Clients.createEmbedding("", {
        input: (0, utils_1.takeFirstNWords)(query, 5000).slice(0, 16000),
        model,
    });
    console.log(providerEmbeddings);
    const queryVector = providerEmbeddings.data[0].embedding;
    const results = new Array();
    for (const embedding of embeddings) {
        const similarity = (0, utils_1.cosineSimilarity)(embedding.vector, queryVector);
        results.push({
            ...embedding,
            similarity,
        });
    }
    return results.sort((a, b) => b.similarity - a.similarity);
}
//# sourceMappingURL=embeddings.js.map