"use strict";

Object.defineProperty(exports, "__esModule", {
  value: true
});
exports.LlamaContext = void 0;
Object.defineProperty(exports, "SchemaGrammarConverter", {
  enumerable: true,
  get: function () {
    return _grammar.SchemaGrammarConverter;
  }
});
Object.defineProperty(exports, "convertJsonSchemaToGrammar", {
  enumerable: true,
  get: function () {
    return _grammar.convertJsonSchemaToGrammar;
  }
});
exports.initLlama = initLlama;
exports.releaseAllLlama = releaseAllLlama;
exports.setContextLimit = setContextLimit;
var _reactNative = require("react-native");
var _NativeRNLlama = _interopRequireDefault(require("./NativeRNLlama"));
var _grammar = require("./grammar");
var _chat = require("./chat");
function _interopRequireDefault(obj) { return obj && obj.__esModule ? obj : { default: obj }; }
const EVENT_ON_TOKEN = '@RNLlama_onToken';
let EventEmitter;
if (_reactNative.Platform.OS === 'ios') {
  // @ts-ignore
  EventEmitter = new _reactNative.NativeEventEmitter(_NativeRNLlama.default);
}
if (_reactNative.Platform.OS === 'android') {
  EventEmitter = _reactNative.DeviceEventEmitter;
}
class LlamaContext {
  gpu = false;
  reasonNoGPU = '';
  model = {};
  constructor(_ref) {
    let {
      contextId,
      gpu,
      reasonNoGPU,
      model
    } = _ref;
    this.id = contextId;
    this.gpu = gpu;
    this.reasonNoGPU = reasonNoGPU;
    this.model = model;
  }

  /**
   * Load cached prompt & completion state from a file.
   */
  async loadSession(filepath) {
    let path = filepath;
    if (path.startsWith('file://')) path = path.slice(7);
    return _NativeRNLlama.default.loadSession(this.id, path);
  }

  /**
   * Save current cached prompt & completion state to a file.
   */
  async saveSession(filepath, options) {
    return _NativeRNLlama.default.saveSession(this.id, filepath, (options === null || options === void 0 ? void 0 : options.tokenSize) || -1);
  }
  async getFormattedChat(messages) {
    var _this$model;
    const chat = (0, _chat.formatChat)(messages);
    return _NativeRNLlama.default.getFormattedChat(this.id, chat, (_this$model = this.model) !== null && _this$model !== void 0 && _this$model.isChatTemplateSupported ? undefined : 'chatml');
  }
  async completion(params, callback) {
    let finalPrompt = params.prompt;
    if (params.messages) {
      // messages always win
      finalPrompt = await this.getFormattedChat(params.messages);
    }
    let tokenListener = callback && EventEmitter.addListener(EVENT_ON_TOKEN, evt => {
      const {
        contextId,
        tokenResult
      } = evt;
      if (contextId !== this.id) return;
      callback(tokenResult);
    });
    if (!finalPrompt) throw new Error('Prompt is required');
    const promise = _NativeRNLlama.default.completion(this.id, {
      ...params,
      prompt: finalPrompt,
      emit_partial_completion: !!callback
    });
    return promise.then(completionResult => {
      var _tokenListener;
      (_tokenListener = tokenListener) === null || _tokenListener === void 0 ? void 0 : _tokenListener.remove();
      tokenListener = null;
      return completionResult;
    }).catch(err => {
      var _tokenListener2;
      (_tokenListener2 = tokenListener) === null || _tokenListener2 === void 0 ? void 0 : _tokenListener2.remove();
      tokenListener = null;
      throw err;
    });
  }
  stopCompletion() {
    return _NativeRNLlama.default.stopCompletion(this.id);
  }
  tokenize(text) {
    return _NativeRNLlama.default.tokenize(this.id, text);
  }
  detokenize(tokens) {
    return _NativeRNLlama.default.detokenize(this.id, tokens);
  }
  embedding(text) {
    return _NativeRNLlama.default.embedding(this.id, text);
  }
  async bench(pp, tg, pl, nr) {
    const result = await _NativeRNLlama.default.bench(this.id, pp, tg, pl, nr);
    const [modelDesc, modelSize, modelNParams, ppAvg, ppStd, tgAvg, tgStd] = JSON.parse(result);
    return {
      modelDesc,
      modelSize,
      modelNParams,
      ppAvg,
      ppStd,
      tgAvg,
      tgStd
    };
  }
  async release() {
    return _NativeRNLlama.default.releaseContext(this.id);
  }
}
exports.LlamaContext = LlamaContext;
async function setContextLimit(limit) {
  return _NativeRNLlama.default.setContextLimit(limit);
}
async function initLlama(_ref2) {
  let {
    model,
    is_model_asset: isModelAsset,
    ...rest
  } = _ref2;
  let path = model;
  if (path.startsWith('file://')) path = path.slice(7);
  const {
    contextId,
    gpu,
    reasonNoGPU,
    model: modelDetails
  } = await _NativeRNLlama.default.initContext({
    model: path,
    is_model_asset: !!isModelAsset,
    ...rest
  });
  return new LlamaContext({
    contextId,
    gpu,
    reasonNoGPU,
    model: modelDetails
  });
}
async function releaseAllLlama() {
  return _NativeRNLlama.default.releaseAllContexts();
}
//# sourceMappingURL=index.js.map