@tonytruong/chatbot-ai-lib/dist/src/modules/embedding/application/EmbeddingService.js

AI-powered healthcare automation, document parsing, OpenAI, embeddings, RAG, vector DB, Facebook OAuth.

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.getCollection = getCollection;
exports.saveTexts = saveTexts;
exports.querySimilar = querySimilar;
exports.embedChunks = embedChunks;
const tslib_1 = require("tslib");
const chromadb_1 = require("chromadb");
const openai_1 = require("../../../config/openai");
const openai_2 = tslib_1.__importDefault(require("openai"));
const getEmbedder = (apiKey, model) => new chromadb_1.OpenAIEmbeddingFunction({
    openai_api_key: apiKey,
    openai_model: model,
});
const getClient = (chromaUrl) => new chromadb_1.ChromaClient({ path: chromaUrl });
const COLLECTION_NAME = "documents";
const DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small";
function getCollection(options) {
    return tslib_1.__awaiter(this, void 0, void 0, function* () {
        const apiKey = (options === null || options === void 0 ? void 0 : options.apiKey) || openai_1.OPENAI_API_KEY;
        const model = (options === null || options === void 0 ? void 0 : options.model) || DEFAULT_EMBEDDING_MODEL;
        const chromaUrl = (options === null || options === void 0 ? void 0 : options.chromaUrl) || "http://localhost:8000";
        const embedder = getEmbedder(apiKey, model);
        const client = getClient(chromaUrl);
        let collection;
        try {
            collection = yield client.getCollection({ name: COLLECTION_NAME });
        }
        catch (e) {
            collection = yield client.createCollection({
                name: COLLECTION_NAME,
                embeddingFunction: embedder,
            });
        }
        return collection;
    });
}
function saveTexts(texts_1) {
    return tslib_1.__awaiter(this, arguments, void 0, function* (texts, metadatas = [], options) {
        const collection = yield getCollection(options);
        const ids = texts.map((_, i) => `doc-${Date.now()}-${i}`);
        yield collection.add({
            ids,
            documents: texts,
            metadatas,
        });
    });
}
function querySimilar(queryText_1) {
    return tslib_1.__awaiter(this, arguments, void 0, function* (queryText, topK = 5, options) {
        const collection = yield getCollection(options);
        const results = yield collection.query({
            queryTexts: [queryText],
            nResults: topK,
        });
        return results;
    });
}
function embedChunks(chunks, options) {
    return tslib_1.__awaiter(this, void 0, void 0, function* () {
        const model = (options === null || options === void 0 ? void 0 : options.model) || DEFAULT_EMBEDDING_MODEL;
        const apiKey = (options === null || options === void 0 ? void 0 : options.apiKey) || openai_1.OPENAI_API_KEY;
        const openai = new openai_2.default({ apiKey });
        const embeddings = [];
        for (const chunk of chunks) {
            const response = yield openai.embeddings.create({
                model,
                input: chunk,
            });
            embeddings.push(response.data[0].embedding);
        }
        return embeddings;
    });
}
//# sourceMappingURL=EmbeddingService.js.map