/* eslint-disable turbo/no-undeclared-env-vars */ import { DataType, MilvusClient } from "@zilliz/milvus2-sdk-node";
import { MetadataMode } from "../../Node.js";
import { VectorStoreBase } from "./types.js";
import { metadataDictToNode, nodeToMetadata } from "./utils.js";
export class MilvusVectorStore extends VectorStoreBase {
    storesText = true;
    isEmbeddingQuery;
    flatMetadata = true;
    milvusClient;
    collectionInitialized = false;
    collectionName;
    idKey;
    contentKey;
    metadataKey;
    embeddingKey;
    constructor(init){
        super(init?.embedModel);
        if (init?.milvusClient) {
            this.milvusClient = init.milvusClient;
        } else {
            const configOrAddress = init?.params?.configOrAddress ?? process.env.MILVUS_ADDRESS;
            const ssl = init?.params?.ssl ?? process.env.MILVUS_SSL === "true";
            const username = init?.params?.username ?? process.env.MILVUS_USERNAME;
            const password = init?.params?.password ?? process.env.MILVUS_PASSWORD;
            if (!configOrAddress) {
                throw new Error("Must specify MILVUS_ADDRESS via env variable.");
            }
            this.milvusClient = new MilvusClient(configOrAddress, ssl, username, password, init?.params?.channelOptions);
        }
        this.collectionName = init?.collection ?? "llamacollection";
        this.idKey = init?.idKey ?? "id";
        this.contentKey = init?.contentKey ?? "content";
        this.metadataKey = init?.metadataKey ?? "metadata";
        this.embeddingKey = init?.embeddingKey ?? "embedding";
    }
    client() {
        return this.milvusClient;
    }
    async createCollection() {
        await this.milvusClient.createCollection({
            collection_name: this.collectionName,
            fields: [
                {
                    name: this.idKey,
                    data_type: DataType.VarChar,
                    is_primary_key: true,
                    max_length: 200
                },
                {
                    name: this.embeddingKey,
                    data_type: DataType.FloatVector,
                    dim: 1536
                },
                {
                    name: this.contentKey,
                    data_type: DataType.VarChar,
                    max_length: 9000
                },
                {
                    name: this.metadataKey,
                    data_type: DataType.JSON
                }
            ]
        });
        await this.milvusClient.createIndex({
            collection_name: this.collectionName,
            field_name: this.embeddingKey
        });
    }
    async ensureCollection() {
        if (!this.collectionInitialized) {
            await this.milvusClient.connectPromise;
            // Check collection exists
            const isCollectionExist = await this.milvusClient.hasCollection({
                collection_name: this.collectionName
            });
            if (!isCollectionExist.value) {
                await this.createCollection();
            }
            await this.milvusClient.loadCollectionSync({
                collection_name: this.collectionName
            });
            this.collectionInitialized = true;
        }
    }
    async add(nodes) {
        await this.ensureCollection();
        const result = await this.milvusClient.insert({
            collection_name: this.collectionName,
            data: nodes.map((node)=>{
                const metadata = nodeToMetadata(node, true, this.contentKey, this.flatMetadata);
                const entry = {
                    [this.idKey]: node.id_,
                    [this.embeddingKey]: node.getEmbedding(),
                    [this.contentKey]: node.getContent(MetadataMode.NONE),
                    [this.metadataKey]: metadata
                };
                return entry;
            })
        });
        if (!result.IDs) {
            return [];
        }
        if ("int_id" in result.IDs) {
            return result.IDs.int_id.data.map((i)=>String(i));
        }
        return result.IDs.str_id.data.map((s)=>String(s));
    }
    async delete(refDocId, deleteOptions) {
        await this.ensureCollection();
        await this.milvusClient.delete({
            ids: [
                refDocId
            ],
            collection_name: this.collectionName,
            ...deleteOptions
        });
    }
    async query(query, _options) {
        await this.ensureCollection();
        const found = await this.milvusClient.search({
            collection_name: this.collectionName,
            limit: query.similarityTopK,
            vector: query.queryEmbedding
        });
        const nodes = [];
        const similarities = [];
        const ids = [];
        found.results.forEach((result)=>{
            const node = metadataDictToNode(result.metadata);
            node.setContent(result.content);
            nodes.push(node);
            similarities.push(result.score);
            ids.push(String(result.id));
        });
        return {
            nodes,
            similarities,
            ids
        };
    }
    async persist() {
    // no need to do anything
    }
}
