"use strict";
Object.defineProperty(exports, "__esModule", {
    value: true
});
Object.defineProperty(exports, "MongoDBAtlasVectorSearch", {
    enumerable: true,
    get: function() {
        return MongoDBAtlasVectorSearch;
    }
});
const _env = require("@llamaindex/env");
const _mongodb = require("mongodb");
const _Node = require("../../Node.js");
const _types = require("./types.js");
const _utils = require("./utils.js");
// Utility function to convert metadata filters to MongoDB filter
function toMongoDBFilter(standardFilters) {
    const filters = {};
    for (const filter of standardFilters?.filters ?? []){
        filters[filter.key] = filter.value;
    }
    return filters;
}
class MongoDBAtlasVectorSearch extends _types.VectorStoreBase {
    storesText = true;
    flatMetadata = true;
    /**
   * The used MongoClient. If not given, a new MongoClient is created based on the MONGODB_URI env variable.
   */ mongodbClient;
    /**
   * Name of the vector index. If invalid, Mongo will silently ignore this issue and return 0 results.
   *
   * Default: "default"
   */ indexName;
    /**
   * Name of the key containing the embedding vector.
   *
   * Default: "embedding"
   */ embeddingKey;
    /**
   * Name of the key containing the node id.
   *
   * Default: "id"
   */ idKey;
    /**
   * Name of the key containing the node text.
   *
   * Default: "text"
   */ textKey;
    /**
   * Name of the key containing the node metadata.
   *
   * Default: "metadata"
   */ metadataKey;
    /**
   * Options to pass to the insertMany function when adding nodes.
   */ insertOptions;
    /**
   * Function to determine the number of candidates to retrieve for a given query.
   * In case your results are not good, you might tune this value.
   *
   * {@link https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/|Run Vector Search Queries}
   *
   * {@link https://arxiv.org/abs/1603.09320|Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs}
   *
   *
   * Default: query.similarityTopK * 10
   */ numCandidates;
    collection;
    constructor(init){
        super(init.embedModel);
        if (init.mongodbClient) {
            this.mongodbClient = init.mongodbClient;
        } else {
            const mongoUri = (0, _env.getEnv)("MONGODB_URI");
            if (!mongoUri) {
                throw new Error("Must specify MONGODB_URI via env variable if not directly passing in client.");
            }
            this.mongodbClient = new _mongodb.MongoClient(mongoUri);
        }
        this.collection = this.mongodbClient.db(init.dbName ?? "default_db").collection(init.collectionName ?? "default_collection");
        this.indexName = init.indexName ?? "default";
        this.embeddingKey = init.embeddingKey ?? "embedding";
        this.idKey = init.idKey ?? "id";
        this.textKey = init.textKey ?? "text";
        this.metadataKey = init.metadataKey ?? "metadata";
        this.numCandidates = init.numCandidates ?? ((query)=>query.similarityTopK * 10);
        this.insertOptions = init.insertOptions;
    }
    /**
   * Add nodes to the vector store.
   *
   * @param nodes Nodes to add to the vector store
   * @returns List of node ids that were added
   */ async add(nodes) {
        if (!nodes || nodes.length === 0) {
            return [];
        }
        const dataToInsert = nodes.map((node)=>{
            const metadata = (0, _utils.nodeToMetadata)(node, true, this.textKey, this.flatMetadata);
            return {
                [this.idKey]: node.id_,
                [this.embeddingKey]: node.getEmbedding(),
                [this.textKey]: node.getContent(_Node.MetadataMode.NONE) || "",
                [this.metadataKey]: metadata
            };
        });
        console.debug("Inserting data into MongoDB: ", dataToInsert);
        const insertResult = await this.collection.insertMany(dataToInsert, this.insertOptions);
        console.debug("Result of insert: ", insertResult);
        return nodes.map((node)=>node.id_);
    }
    /**
   * Delete nodes from the vector store with the given redDocId.
   *
   * @param refDocId The refDocId of the nodes to delete
   * @param deleteOptions Options to pass to the deleteOne function
   */ async delete(refDocId, deleteOptions) {
        await this.collection.deleteMany({
            [`${this.metadataKey}.ref_doc_id`]: refDocId
        }, deleteOptions);
    }
    get client() {
        return this.mongodbClient;
    }
    /**
   * Perform a vector similarity search query.
   *
   * @param query The query to run
   * @returns List of nodes and their similarities
   */ async query(query, options) {
        const params = {
            queryVector: query.queryEmbedding,
            path: this.embeddingKey,
            numCandidates: this.numCandidates(query),
            limit: query.similarityTopK,
            index: this.indexName
        };
        if (query.filters) {
            params.filter = toMongoDBFilter(query.filters);
        }
        const queryField = {
            $vectorSearch: params
        };
        const pipeline = [
            queryField,
            {
                $project: {
                    score: {
                        $meta: "vectorSearchScore"
                    },
                    [this.embeddingKey]: 0
                }
            }
        ];
        console.debug("Running query pipeline: ", pipeline);
        const cursor = await this.collection.aggregate(pipeline);
        const nodes = [];
        const ids = [];
        const similarities = [];
        for await (const res of (await cursor)){
            const text = res[this.textKey];
            const score = res.score;
            const id = res[this.idKey];
            const metadata = res[this.metadataKey];
            const node = (0, _utils.metadataDictToNode)(metadata);
            node.setContent(text);
            ids.push(id);
            nodes.push(node);
            similarities.push(score);
        }
        const result = {
            nodes,
            similarities,
            ids
        };
        console.debug("Result of query (ids):", ids);
        return result;
    }
}
