import { AsyncTaskId, AIStreamParser, AIResult, AsyncFeatures } from '@isdk/ai-tool';
import { AITextGenerationOptions, AIOptions, AIModelQuantType, AILavaModelSettings, LLMProvider, AIModelParams } from '@isdk/ai-tool-llm';

interface LlamaModelOptions extends AITextGenerationOptions {
    /**
     * Stop generation when any of these tokens are generated in llama.cpp.
     * same as stop_words in ai-tool-llm
     */
    stop?: string[];
    /**
     * Limit the next token selection to the K most probable tokens (default: 40).
     */
    top_k?: number;
    /**
     * Limit the next token selection to a subset of tokens with a cumulative probability above a threshold P (default: 0.95).
     */
    top_p?: number;
    /**
     * The minimum probability for a token to be considered, relative to the probability of the most likely token (default: 0.05).
     */
    min_p?: number;
    /**
     * Specify the number of tokens from the prompt to retain when the context size is exceeded
     * and tokens need to be discarded. By default, this value is set to 0 (meaning no tokens
     * are kept). Use -1 to retain all tokens from the prompt.
     */
    n_keep?: number;
    /**
     * Enable tail free sampling with parameter z (default: 1.0, 1.0 = disabled).
     */
    tfs_z?: number;
    /**
     * Enable locally typical sampling with parameter p (default: 1.0, 1.0 = disabled).
     */
    typical_p?: number;
    /**
     * Control the repetition of token sequences in the generated text (default: 1.1).
     */
    repeat_penalty?: number;
    /**
     * Last n tokens to consider for penalizing repetition (default: 64, 0 = disabled, -1 = ctx-size).
     */
    repeat_last_n?: number;
    /**
     * Penalize newline tokens when applying the repeat penalty (default: true).
     */
    penalize_nl?: boolean;
    /**
     * Repeat alpha presence penalty (default: 0.0, 0.0 = disabled).
     */
    presence_penalty?: number;
    /**
     * Repeat alpha frequency penalty (default: 0.0, 0.0 = disabled).
     */
    frequency_penalty?: number;
    /**
     * This will replace the prompt for the purpose of the penalty evaluation.
     * Can be either null, a string or an array of numbers representing tokens
     * (default: null = use the original prompt).
     */
    penalty_prompt?: string | number[];
    /**
     * Enable Mirostat sampling, controlling perplexity during text generation
     * (default: 0, 0 = disabled, 1 = Mirostat, 2 = Mirostat 2.0).
     */
    mirostat?: number;
    /**
     * Set the Mirostat target entropy, parameter tau (default: 5.0).
     */
    mirostat_tau?: number;
    /**
     * Set the Mirostat learning rate, parameter eta (default: 0.1).
     */
    mirostat_eta?: number;
    /**
     * Set grammar for grammar-based sampling (default: no grammar)
     *
     * @see https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md
     */
    grammar?: string;
    /**
     * Set the random number generator (RNG) seed
     * (default: -1, -1 = random seed).
     */
    seed?: number;
    /**
     * Ignore end of stream token and continue generating (default: false).
     */
    ignore_eos?: boolean;
    /**
     * Modify the likelihood of a token appearing in the generated text completion.
     * For example, use "logit_bias": [[15043,1.0]] to increase the likelihood of the token
     * 'Hello', or "logit_bias": [[15043,-1.0]] to decrease its likelihood.
     * Setting the value to false, "logit_bias": [[15043,false]] ensures that the token Hello is
     * never produced (default: []).
     */
    logit_bias?: Array<[number, number | false]>;
    /**
     * If greater than 0, the response also contains the probabilities of top N tokens
     * for each generated token (default: 0)
     */
    n_probs?: number;
    /**
     * Save the prompt and generation for avoid reprocess entire prompt if a part of this isn't change (default: false)
     */
    cache_prompt?: boolean;
    /**
     * Assign the completion task to an specific slot.
     * If is -1 the task will be assigned to a Idle slot (default: -1)
     */
    slot_id?: number;
    /**
     * maximum tokens to predict, -1 means no limits.
     * It could not exceed the server's n_predict.
     */
    n_predict?: number;
}
type LlamaRopeScalingType = 'none' | 'linear' | 'yarn';
type LlamaCacheQuantType = 'f32' | 'f16' | 'q8_0' | 'q4_0' | 'q4_1' | 'q5_0' | 'q5_1';
interface LoraItemObject {
    name: string;
    scale?: number;
}
type LoraItem = LoraItemObject | string;
type LoraItems = LoraItem[];
declare enum LlmSizeType {
    Unknown = 0,
    _17M = 1,
    _22M = 2,
    _33M = 3,
    _109M = 4,
    _137M = 5,
    _335M = 6,
    _0_5B = 7,
    _1B = 8,
    _2B = 9,
    _3B = 10,
    _4B = 11,
    _7B = 12,
    _8B = 13,
    _13B = 14,
    _14B = 15,
    _15B = 16,
    _20B = 17,
    _30B = 18,
    _34B = 19,
    _35B = 20,
    _40B = 21,
    _65B = 22,
    _70B = 23,
    Small = 24,
    Medium = 25,
    Large = 26,
    XL = 27
}
declare enum LlmArch {
    Llama = 0,
    Falcon = 1,
    Baichuan = 2,
    Grok = 3,
    Gpt2 = 4,
    Gptj = 5,
    Gptneox = 6,
    Mpt = 7,
    Starcoder = 8,
    Refact = 9,
    Bert = 10,
    Nomic_bert = 11,
    Jina_bert_v2 = 12,
    Bloom = 13,
    Stablelm = 14,
    Qwen = 15,
    Qwen2 = 16,
    Qwen2moe = 17,
    Phi2 = 18,
    Phi3 = 19,
    Plamo = 20,
    Codeshell = 21,
    Orion = 22,
    Internlm2 = 23,
    Minicpm = 24,
    Gemma = 25,
    Starcoder2 = 26,
    Mamba = 27,
    Xverse = 28,
    Command_r = 29,
    Dbrx = 30,
    Olmo = 31,
    Arctic = 32,
    Deepseek2 = 33,
    Unknown = 34
}
declare enum LlamaVocabType {
    None = 0,// For models without vocab
    Spm = 1,// SentencePiece
    Bpe = 2,// Byte Pair Encoding
    Wpm = 3
}
interface LlamaLoadModelOptions extends AIOptions {
    model: string;
    mmproj?: string;
    alias?: string;
    gpu_layers?: number;
    /**
     * Specify the context window size of the model that you have loaded in your
     * Llama.cpp server.
    // Set the size of the prompt context. The default is 512, but LLaMA models were built with a context of 2048, which will provide better results for longer input/inference.
    // The size may differ in other models, for example, baichuan models were build with a context of 4096.
     */
    n_ctx?: number;
    embedding?: boolean;
    batch?: number;
    parallel?: number;
    threads?: number;
    n_threads_batch?: number;
    grp_attn_n?: number;
    grp_attn_w?: number;
    rope_scaling?: LlamaRopeScalingType;
    rope_freq_base?: number;
    rope_freq_scale?: number;
    yarn_ext_factor?: number;
    yarn_attn_factor?: number;
    yarn_beta_fast?: number;
    yarn_beta_slow?: number;
    use_mmap?: boolean;
    use_mlock?: boolean;
    cont_batching?: boolean;
    n_predict?: number;
    cache_type_k?: LlamaCacheQuantType;
    cache_type_v?: LlamaCacheQuantType;
    lora?: LoraItem | LoraItems;
    lora_base?: string;
}
declare const LlamaLoadModelOptionsKeys: (keyof LlamaLoadModelOptions)[];
declare enum LlamaCppServerState {
    LoadingModel = 0,// Server is starting up, model not fully loaded yet
    Ready = 1,// Server is ready and model is loaded
    Unloaded = 2,// model unloaded
    Error = 3
}
interface LlamaModelSettings {
    model: string;
    state?: LlamaCppServerState;
    arch?: LlmArch;
    size_type?: LlmSizeType;
    size?: number;
    quant?: AIModelQuantType;
    vocab_type?: LlamaVocabType;
    /**
     * the content length in train
     */
    n_ctx_train?: number;
    /**
     * the model params size in bytes
     */
    n_params?: number;
    /**
     * the quanted model size in bytes
     */
    n_vocab?: number;
    n_embd?: number;
    chat_template?: string;
    [name: string]: any;
}
interface LlamaLavaModelSettings extends AILavaModelSettings {
    n_embd?: number;
}
interface LLamaCppModelGenerationSettings extends LlamaModelOptions {
    model?: string;
    model_id?: string;
    samplers: string[];
    use_penalty_prompt_tokens: boolean;
}
interface LLamaCppModelTimings {
    predicted_ms: number;
    predicted_n: number;
    predicted_per_second: number;
    predicted_per_token_ms: number;
    prompt_ms: number;
    prompt_n: number;
    prompt_per_second: number;
    prompt_per_token_ms: number;
}
interface LLamaCppPartResult {
    content: string;
    multimodal?: boolean;
    slot_id: number;
    stop: boolean;
}
interface LLamaCppResult {
    content: string;
    generation_settings: LLamaCppModelGenerationSettings;
    model: string;
    prompt: string;
    slot_id: number;
    stop: boolean;
    stopped_eos: boolean;
    stopped_limit: boolean;
    stopped_word: boolean;
    stopping_word: string;
    timings: LLamaCppModelTimings;
    tokens_cached: number;
    tokens_evaluated: number;
    tokens_predicted: number;
    truncated: boolean;
    chatTemplateId?: {
        id: string;
        version?: string;
    };
    taskId?: AsyncTaskId;
}
type LlamaCppAIStreamParser<T = LLamaCppResult> = AIStreamParser<string, T>;
type LlamaCppAIResult = AIResult<string, LLamaCppResult>;
declare function llamaCppToAIResult(data: LLamaCppResult): LlamaCppAIResult;
declare function parseLlamaCppStream(opts?: any): LlamaCppAIStreamParser;
declare const AIOptionsMap: {
    max_tokens: string;
    stop_words: string;
    content_size: string;
    'response_format.type': string;
    'response_format.schema': string;
};
declare const ReverseAIOptionsMap: Record<string, string | string[]>;
declare const AILoadModelOptionsMap: {
    content_size: string;
    stop_words: string;
};
declare const ReverseAILoadModelOptions: Record<string, string | string[]>;
declare const AIModelSettingsMap: {
    n_params: string;
    n_ctx_train: string;
    model: string;
};

declare const LlamaCppProviderName = "llamacppPro";
declare class LlamaCppProvider extends LLMProvider {
    rule: RegExp;
    asyncFeatures: AsyncFeatures;
    processModelOptions(model: string, value: any, options: AIOptions): Promise<AIOptions>;
    func({ model, value, options }: {
        model: string;
        value: any;
        options: AIOptions;
    }): Promise<any>;
    loadModel(model: LlamaLoadModelOptions): Promise<LlamaLoadModelOptions>;
    getModelInfo(modelName?: string): Promise<AIModelParams>;
    listModels(): Promise<string[] | undefined>;
}
declare const llamaCpp: LlamaCppProvider;

export { AILoadModelOptionsMap, AIModelSettingsMap, AIOptionsMap, type LLamaCppModelGenerationSettings, type LLamaCppModelTimings, type LLamaCppPartResult, type LLamaCppResult, type LlamaCacheQuantType, type LlamaCppAIResult, type LlamaCppAIStreamParser, LlamaCppProvider, LlamaCppProviderName, type LlamaLavaModelSettings, type LlamaLoadModelOptions, LlamaLoadModelOptionsKeys, type LlamaModelOptions, type LlamaModelSettings, type LlamaRopeScalingType, LlamaVocabType, LlmArch, LlmSizeType, type LoraItem, type LoraItemObject, type LoraItems, ReverseAILoadModelOptions, ReverseAIOptionsMap, llamaCpp, llamaCppToAIResult, parseLlamaCppStream };
