"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.EmbeddingService = void 0;
const path = __importStar(require("path"));
const glob_1 = require("glob");
const types_1 = require("../types");
const utils_1 = require("../utils");
const ai_1 = require("../ai");
const clients_1 = require("../clients");
class EmbeddingService {
    context;
    config;
    embeddingClient;
    plugins;
    constructor(context) {
        this.context = context;
        this.config = context.config;
        this.embeddingClient = context.embeddingClient;
        this.plugins = context.plugins;
    }
    setEmbeddingClient(client) {
        this.embeddingClient = client;
    }
    createEmbedding({ provider, model, input }) {
        if (!this.embeddingClient) {
            return clients_1.Clients.createEmbedding(provider, { model, input });
        }
        return this.embeddingClient.createEmbedding({ model, input });
    }
    async loadEmbedding(filePath) {
        if (await (0, utils_1.fileExists)(filePath)) {
            return JSON.parse(await (0, utils_1.readFile)(filePath, "utf8"));
        }
        return [];
    }
    async getConfiguredEmbeddingMap() {
        const files = Array.from(new Set(this.config.embedSources.map((s) => s.output)));
        const embeddings = {};
        for (const file of files) {
            if (!embeddings[file]) {
                embeddings[file] = [];
            }
            const fileEmbeddings = await this.loadEmbedding(file);
            embeddings[file].push(...fileEmbeddings);
        }
        return embeddings;
    }
    async getConfiguredEmbeddings() {
        const files = Array.from(new Set(this.config.embedSources.map((s) => s.output)));
        const embeddings = [];
        for (const file of files) {
            const fileEmbeddings = await this.loadEmbedding(file);
            embeddings.push(...fileEmbeddings);
        }
        return embeddings;
    }
    getChunkId(id, index, chunkSize) {
        return chunkSize ? `${id}-${index}` : id;
    }
    async embedSource(model, source, ignorePattern) {
        if (!source.input) {
            console.log("Skipping", source.output, "with blank input property");
            return;
        }
        console.log("Embedding", source.input, "to", source.output);
        let files = await (0, glob_1.globSync)(source.input, { ignore: ignorePattern });
        if (source.kind && files.length === 0) {
            files = [source.input];
        }
        console.log(`Found ${files.length} files`);
        if (files.length > 100) {
            console.error("woah there, that's a lot of files. I'm not going to embed that many");
        }
        console.log(files);
        const embeddings = await this.loadEmbedding(source.output);
        let batch = [];
        let index = 0;
        for (const file of files) {
            index++;
            const shouldSave = batch.length > 20 || index === files.length;
            if (shouldSave) {
                await Promise.all(batch);
                batch = [];
            }
            batch.push(this.embedKind(model, file, source, embeddings, shouldSave));
        }
        if (batch.length > 0) {
            await Promise.all(batch);
        }
        await this.saveEmbedding(source.output, embeddings);
    }
    async saveEmbedding(output, embeddings) {
        const fileString = "[" + embeddings.map((e) => JSON.stringify(e)).join(",") + "]";
        await (0, utils_1.writeFile)(output, fileString);
    }
    async embedKind(model, id, source, embeddings = [], save = true) {
        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 this.isEmbeddingFile(id))) {
            console.log("Embedding JSON", id);
            return this.embedJson(model, id, source);
        }
        const toEmbed = await this.handleAllKinds(id, source);
        const updates = [];
        for (const row of toEmbed) {
            const { id: rowId, text, metadata } = row;
            const embedded = await this.embed(model, rowId, text, metadata, embeddings, {
                prompt,
                chunkSize,
                uploadMode,
            });
            updates.push(...embedded);
        }
        if (save && updates.length > 0) {
            await this.saveEmbedding(output, embeddings);
        }
    }
    async handleAllKinds(id, source) {
        const { input, kind } = source;
        const contents = "";
        const ids = [];
        if (this.plugins.isPlugin(kind)) {
            console.log("Embedding with plugin", kind);
            return this.plugins.embed(kind, input);
        }
        switch (kind) {
            case "text":
                return this.handleTextKind(source.input);
            case "file":
            default:
                return this.handleFileKind(id);
        }
    }
    async handleTextKind(contents) {
        return [{ id: "", text: contents, metadata: {} }];
    }
    async handleFileKind(filePath) {
        const contents = await (0, utils_1.readFile)(filePath, "utf8");
        return [{ id: filePath, text: contents, metadata: { filePath } }];
    }
    async embed(model, id, text, metadata, embeddings, options = {}) {
        const { prompt, chunkSize, uploadMode, minLength } = options;
        let chunks = [text];
        if (chunkSize) {
            chunks = await (0, ai_1.chunkText)(text, chunkSize);
        }
        const MAX_CHUNKS = 300;
        chunks = chunks.slice(0, MAX_CHUNKS);
        const dontPrune = [];
        const updates = new Array();
        for (let index = 0; index < chunks.length; index++) {
            const chunkId = this.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 alreadyEmbedded = embeddings.find((e) => e.id === chunkId && e.text === textOfChunk);
            if (alreadyEmbedded) {
                console.log("Skipping", chunkId);
                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 this.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: chunks[index] }),
                },
            };
            embeddings.push(embeddable);
            updates.push(chunkId);
        }
        this.pruneEmbedding(id, dontPrune, embeddings);
        return updates;
    }
    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;
    }
    async queryEmbedding(query, embeddings, model = types_1.EmbeddingModels.openai.EmbeddingAda2) {
        const providerEmbeddings = await this.createEmbedding({
            input: query,
            model,
        });
        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);
    }
    pruneVector(embeddings) {
        for (const entry of embeddings) {
            delete entry.vector;
        }
    }
    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 isEmbeddingFile(inputFile) {
        const filePath = path.join(process.cwd(), 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 embedJson(model, inputFile, source) {
        const { prompt, output, uploadMode, chunkSize, minLength } = source;
        const filePath = path.join(process.cwd(), inputFile);
        const sourceJson = JSON.parse(await (0, utils_1.readFile)(inputFile, "utf8"));
        const embeddings = await this.loadEmbedding(output);
        let updates = [];
        let batch = [];
        for (const row of sourceJson) {
            if (embeddings.find((e) => e.id === row.id)) {
                continue;
            }
            console.log("Embedding", row.id);
            batch.push(this.embed(model, row.id, row.text, row.metadata, embeddings, {
                prompt,
                chunkSize,
                uploadMode,
                minLength,
            }));
            let embedded = [];
            if (batch.length > 20) {
                embedded = (await Promise.all(batch)).flat();
                batch = [];
            }
            updates.push(...embedded);
            if (updates.length > 20) {
                await this.saveEmbedding(output, embeddings);
                updates = [];
            }
        }
        if (batch.length > 0) {
            const embedded = (await Promise.all(batch)).flat();
            updates.push(...embedded);
        }
        if (updates.length > 0) {
            await this.saveEmbedding(output, embeddings);
        }
    }
}
exports.EmbeddingService = EmbeddingService;
//# sourceMappingURL=EmbeddingService.js.map