"use strict";

Object.defineProperty(exports, "__esModule", {
  value: true
});
exports.sortSourceFields = exports.sortModels = exports.isTextExpansionModel = exports.isSupportedMLModel = exports.getModelDisplayTitle = exports.getMLType = exports.RECOMMENDED_FIELDS = exports.NLP_DISPLAY_TITLES = exports.NLP_CONFIG_KEYS = void 0;
var _i18n = require("@kbn/i18n");
var _mlTrainedModelsUtils = require("@kbn/ml-trained-models-utils");
/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0; you may not use this file except in compliance with the Elastic License
 * 2.0.
 */

const NLP_CONFIG_KEYS = Object.values(_mlTrainedModelsUtils.SUPPORTED_PYTORCH_TASKS);
exports.NLP_CONFIG_KEYS = NLP_CONFIG_KEYS;
const RECOMMENDED_FIELDS = ['body', 'body_content', 'title'];
exports.RECOMMENDED_FIELDS = RECOMMENDED_FIELDS;
const NLP_DISPLAY_TITLES = {
  fill_mask: _i18n.i18n.translate('xpack.enterpriseSearch.content.ml_inference.fill_mask', {
    defaultMessage: 'Fill Mask'
  }),
  lang_ident: _i18n.i18n.translate('xpack.enterpriseSearch.content.ml_inference.lang_ident', {
    defaultMessage: 'Language Identification'
  }),
  ner: _i18n.i18n.translate('xpack.enterpriseSearch.content.ml_inference.ner', {
    defaultMessage: 'Named Entity Recognition'
  }),
  question_answering: _i18n.i18n.translate('xpack.enterpriseSearch.content.ml_inference.question_answering', {
    defaultMessage: 'Named Entity Recognition'
  }),
  text_classification: _i18n.i18n.translate('xpack.enterpriseSearch.content.ml_inference.text_classification', {
    defaultMessage: 'Text Classification'
  }),
  text_embedding: _i18n.i18n.translate('xpack.enterpriseSearch.content.ml_inference.text_embedding', {
    defaultMessage: 'Dense Vector Text Embedding'
  }),
  text_expansion: _i18n.i18n.translate('xpack.enterpriseSearch.content.ml_inference.text_expansion', {
    defaultMessage: 'ELSER Text Expansion'
  }),
  zero_shot_classification: _i18n.i18n.translate('xpack.enterpriseSearch.content.ml_inference.zero_shot_classification', {
    defaultMessage: 'Zero-Shot Text Classification'
  })
};
exports.NLP_DISPLAY_TITLES = NLP_DISPLAY_TITLES;
const isSupportedMLModel = model => {
  return Object.keys(model.inference_config || {}).some(key => NLP_CONFIG_KEYS.includes(key)) || model.model_type === _mlTrainedModelsUtils.TRAINED_MODEL_TYPE.LANG_IDENT;
};
exports.isSupportedMLModel = isSupportedMLModel;
const sortSourceFields = (a, b) => {
  const promoteA = RECOMMENDED_FIELDS.includes(a);
  const promoteB = RECOMMENDED_FIELDS.includes(b);
  if (promoteA && promoteB) {
    return RECOMMENDED_FIELDS.indexOf(a) > RECOMMENDED_FIELDS.indexOf(b) ? 1 : -1;
  } else if (promoteA) {
    return -1;
  } else if (promoteB) {
    return 1;
  }
  return a.localeCompare(b);
};
exports.sortSourceFields = sortSourceFields;
const getMLType = modelTypes => {
  var _modelTypes$;
  for (const type of modelTypes) {
    if (NLP_CONFIG_KEYS.includes(type)) {
      return type;
    }
  }
  if (modelTypes !== null && modelTypes !== void 0 && modelTypes.includes(_mlTrainedModelsUtils.TRAINED_MODEL_TYPE.LANG_IDENT)) return _mlTrainedModelsUtils.TRAINED_MODEL_TYPE.LANG_IDENT;
  return (_modelTypes$ = modelTypes === null || modelTypes === void 0 ? void 0 : modelTypes[0]) !== null && _modelTypes$ !== void 0 ? _modelTypes$ : '';
};
exports.getMLType = getMLType;
const getModelDisplayTitle = type => NLP_DISPLAY_TITLES[type];
exports.getModelDisplayTitle = getModelDisplayTitle;
const isTextExpansionModel = model => {
  var _model$inference_conf;
  return Boolean((_model$inference_conf = model.inference_config) === null || _model$inference_conf === void 0 ? void 0 : _model$inference_conf.text_expansion);
};

/**
 * Sort function for displaying a list of models. Promotes text_expansion models and sorts the rest by model ID.
 */
exports.isTextExpansionModel = isTextExpansionModel;
const sortModels = (m1, m2) => isTextExpansionModel(m1) ? -1 : isTextExpansionModel(m2) ? 1 : m1.model_id.localeCompare(m2.model_id);
exports.sortModels = sortModels;