import Handlebars from "handlebars";
import { logDevData } from "../util/devdata.js";
import { DevDataSqliteDb } from "../util/devdataSqlite.js";
import { fetchwithRequestOptions } from "../util/fetchWithOptions.js";
import mergeJson from "../util/merge.js";
import { Telemetry } from "../util/posthog.js";
import { withExponentialBackoff } from "../util/withExponentialBackoff.js";
import { autodetectPromptTemplates, autodetectTemplateFunction, autodetectTemplateType, modelSupportsImages, } from "./autodetect.js";
import { CONTEXT_LENGTH_FOR_MODEL, DEFAULT_ARGS, DEFAULT_CONTEXT_LENGTH, DEFAULT_MAX_TOKENS, } from "./constants.js";
import { compileChatMessages, countTokens, pruneRawPromptFromTop, stripImages, } from "./countTokens.js";
import CompletionOptionsForModels from "./templates/options.js";
export class BaseLLM {
    static providerName;
    static defaultOptions = undefined;
    get providerName() {
        return this.constructor.providerName;
    }
    supportsImages() {
        return modelSupportsImages(this.providerName, this.model, this.title);
    }
    supportsCompletions() {
        if (this.providerName === "openai") {
            if (this.apiBase?.includes("api.groq.com") ||
                this.apiBase?.includes(":1337") ||
                this._llmOptions.useLegacyCompletionsEndpoint?.valueOf() === false) {
                // Jan + Groq don't support completions : (
                return false;
            }
        }
        if (this.providerName === "groq") {
            return false;
        }
        return true;
    }
    supportsPrefill() {
        return ["ollama", "anthropic"].includes(this.providerName);
    }
    uniqueId;
    model;
    title;
    systemMessage;
    contextLength;
    completionOptions;
    requestOptions;
    template;
    promptTemplates;
    templateMessages;
    writeLog;
    llmRequestHook;
    apiKey;
    apiBase;
    engine;
    apiVersion;
    apiType;
    region;
    projectId;
    _llmOptions;
    constructor(_options) {
        this._llmOptions = _options;
        // Set default options
        const options = {
            title: this.constructor.providerName,
            ...this.constructor.defaultOptions,
            ..._options,
        };
        const templateType = options.template ?? autodetectTemplateType(options.model);
        this.title = options.title;
        this.uniqueId = options.uniqueId ?? "None";
        this.model = options.model;
        this.systemMessage = options.systemMessage;
        this.contextLength = options.contextLength ?? DEFAULT_CONTEXT_LENGTH;
        this.completionOptions = {
            ...options.completionOptions,
            model: options.model || "gpt-4",
            maxTokens: options.completionOptions?.maxTokens ?? DEFAULT_MAX_TOKENS,
        };
        if (CompletionOptionsForModels[options.model]) {
            this.completionOptions = mergeJson(this.completionOptions, CompletionOptionsForModels[options.model] ?? {});
        }
        this.requestOptions = options.requestOptions;
        this.promptTemplates = {
            ...autodetectPromptTemplates(options.model, templateType),
            ...options.promptTemplates,
        };
        this.templateMessages =
            options.templateMessages ??
                autodetectTemplateFunction(options.model, this.providerName, options.template);
        this.writeLog = options.writeLog;
        this.llmRequestHook = options.llmRequestHook;
        this.apiKey = options.apiKey;
        this.apiBase = options.apiBase;
        if (this.apiBase && !this.apiBase.endsWith("/")) {
            this.apiBase = `${this.apiBase}/`;
        }
        this.engine = options.engine;
        this.apiVersion = options.apiVersion;
        this.apiType = options.apiType;
        this.region = options.region;
        this.projectId = options.projectId;
    }
    listModels() {
        return Promise.resolve([]);
    }
    _compileChatMessages(options, messages, functions) {
        let contextLength = this.contextLength;
        if (options.model !== this.model &&
            options.model in CONTEXT_LENGTH_FOR_MODEL) {
            contextLength =
                CONTEXT_LENGTH_FOR_MODEL[options.model] || DEFAULT_CONTEXT_LENGTH;
        }
        return compileChatMessages(options.model, messages, contextLength, options.maxTokens ?? DEFAULT_MAX_TOKENS, this.supportsImages(), undefined, functions, this.systemMessage);
    }
    _getSystemMessage() {
        // TODO: Merge with config system message
        return this.systemMessage;
    }
    _templatePromptLikeMessages(prompt) {
        if (!this.templateMessages) {
            return prompt;
        }
        const msgs = [{ role: "user", content: prompt }];
        const systemMessage = this._getSystemMessage();
        if (systemMessage) {
            msgs.unshift({ role: "system", content: systemMessage });
        }
        return this.templateMessages(msgs);
    }
    _compileLogMessage(prompt, completionOptions) {
        const dict = { contextLength: this.contextLength, ...completionOptions };
        const settings = Object.entries(dict)
            .map(([key, value]) => `${key}: ${value}`)
            .join("\n");
        return `Settings:
${settings}

############################################

${prompt}`;
    }
    _logTokensGenerated(model, prompt, completion) {
        let promptTokens = this.countTokens(prompt);
        let generatedTokens = this.countTokens(completion);
        Telemetry.capture("tokens_generated", {
            model: model,
            provider: this.providerName,
            promptTokens: promptTokens,
            generatedTokens: generatedTokens,
        });
        DevDataSqliteDb.logTokensGenerated(model, this.providerName, promptTokens, generatedTokens);
        logDevData("tokens_generated", {
            model: model,
            provider: this.providerName,
            promptTokens: promptTokens,
            generatedTokens: generatedTokens,
        });
    }
    fetch(url, init) {
        // Custom Node.js fetch
        const customFetch = async (input, init) => {
            try {
                const resp = await fetchwithRequestOptions(new URL(input), { ...init }, { ...this.requestOptions });
                if (!resp.ok) {
                    let text = await resp.text();
                    if (resp.status === 404 && !resp.url.includes("/v1")) {
                        if (text.includes("try pulling it first")) {
                            const model = JSON.parse(text).error.split(" ")[1].slice(1, -1);
                            text = `The model "${model}" was not found. To download it, run \`ollama run ${model}\`.`;
                        }
                        else if (text.includes("/api/chat")) {
                            text =
                                "The /api/chat endpoint was not found. This may mean that you are using an older version of Ollama that does not support /api/chat. Upgrading to the latest version will solve the issue.";
                        }
                        else {
                            text =
                                "This may mean that you forgot to add '/v1' to the end of your 'apiBase' in config.json.";
                        }
                    }
                    throw new Error(`HTTP ${resp.status} ${resp.statusText} from ${resp.url}\n\n${text}`);
                }
                return resp;
            }
            catch (e) {
                if (e.code === "ECONNREFUSED" &&
                    e.message.includes("http://127.0.0.1:11434")) {
                    throw new Error("Failed to connect to local Ollama instance. To start Ollama, first download it at https://ollama.ai.");
                }
                throw new Error(`${e.message}\n\nCode: ${e.code}\nError number: ${e.errno}\nSyscall: ${e.erroredSysCall}\nType: ${e.type}\n\n${e.stack}`);
            }
        };
        return withExponentialBackoff(() => customFetch(url, init), 5, 0.5);
    }
    _parseCompletionOptions(options) {
        const log = options.log ?? true;
        const raw = options.raw ?? false;
        options.log = undefined;
        const completionOptions = mergeJson(this.completionOptions, options);
        return { completionOptions, log, raw };
    }
    _formatChatMessages(messages) {
        const msgsCopy = messages ? messages.map((msg) => ({ ...msg })) : [];
        let formatted = "";
        for (const msg of msgsCopy) {
            if ("content" in msg && Array.isArray(msg.content)) {
                const content = stripImages(msg.content);
                msg.content = content;
            }
            formatted += `<${msg.role}>\n${msg.content || ""}\n\n`;
        }
        return formatted;
    }
    async *streamComplete(_prompt, options = {}) {
        const { completionOptions, log, raw } = this._parseCompletionOptions(options);
        let prompt = pruneRawPromptFromTop(completionOptions.model, this.contextLength, _prompt, completionOptions.maxTokens ?? DEFAULT_MAX_TOKENS);
        if (!raw) {
            prompt = this._templatePromptLikeMessages(prompt);
        }
        if (log) {
            if (this.writeLog) {
                await this.writeLog(this._compileLogMessage(prompt, completionOptions));
            }
            if (this.llmRequestHook) {
                this.llmRequestHook(completionOptions.model, prompt);
            }
        }
        let completion = "";
        for await (const chunk of this._streamComplete(prompt, completionOptions)) {
            completion += chunk;
            yield chunk;
        }
        this._logTokensGenerated(completionOptions.model, prompt, completion);
        if (log && this.writeLog) {
            await this.writeLog(`Completion:\n\n${completion}\n\n`);
        }
        return { prompt, completion, completionOptions };
    }
    async complete(_prompt, options = {}) {
        const { completionOptions, log, raw } = this._parseCompletionOptions(options);
        let prompt = pruneRawPromptFromTop(completionOptions.model, this.contextLength, _prompt, completionOptions.maxTokens ?? DEFAULT_MAX_TOKENS);
        if (!raw) {
            prompt = this._templatePromptLikeMessages(prompt);
        }
        if (log) {
            if (this.writeLog) {
                await this.writeLog(this._compileLogMessage(prompt, completionOptions));
            }
            if (this.llmRequestHook) {
                this.llmRequestHook(completionOptions.model, prompt);
            }
        }
        const completion = await this._complete(prompt, completionOptions);
        this._logTokensGenerated(completionOptions.model, prompt, completion);
        if (log && this.writeLog) {
            await this.writeLog(`Completion:\n\n${completion}\n\n`);
        }
        return completion;
    }
    async chat(messages, options = {}) {
        let completion = "";
        for await (const chunk of this.streamChat(messages, options)) {
            completion += chunk.content;
        }
        return { role: "assistant", content: completion };
    }
    async *streamChat(_messages, options = {}) {
        const { completionOptions, log, raw } = this._parseCompletionOptions(options);
        const messages = this._compileChatMessages(completionOptions, _messages);
        const prompt = this.templateMessages
            ? this.templateMessages(messages)
            : this._formatChatMessages(messages);
        if (log) {
            if (this.writeLog) {
                await this.writeLog(this._compileLogMessage(prompt, completionOptions));
            }
            if (this.llmRequestHook) {
                this.llmRequestHook(completionOptions.model, prompt);
            }
        }
        let completion = "";
        try {
            if (this.templateMessages) {
                for await (const chunk of this._streamComplete(prompt, completionOptions)) {
                    completion += chunk;
                    yield { role: "assistant", content: chunk };
                }
            }
            else {
                for await (const chunk of this._streamChat(messages, completionOptions)) {
                    completion += chunk.content;
                    yield chunk;
                }
            }
        }
        catch (error) {
            console.log(error);
            throw error;
        }
        this._logTokensGenerated(completionOptions.model, prompt, completion);
        if (log && this.writeLog) {
            await this.writeLog(`Completion:\n\n${completion}\n\n`);
        }
        return {
            prompt,
            completion,
            completionOptions,
        };
    }
    // biome-ignore lint/correctness/useYield: Purposefully not implemented
    async *_streamComplete(prompt, options) {
        throw new Error("Not implemented");
    }
    async *_streamChat(messages, options) {
        if (!this.templateMessages) {
            throw new Error("You must either implement templateMessages or _streamChat");
        }
        for await (const chunk of this._streamComplete(this.templateMessages(messages), options)) {
            yield { role: "assistant", content: chunk };
        }
    }
    async _complete(prompt, options) {
        let completion = "";
        for await (const chunk of this._streamComplete(prompt, options)) {
            completion += chunk;
        }
        return completion;
    }
    countTokens(text) {
        return countTokens(text, this.model);
    }
    collectArgs(options) {
        return {
            ...DEFAULT_ARGS,
            // model: this.model,
            ...options,
        };
    }
    renderPromptTemplate(template, history, otherData, canPutWordsInModelsMouth = false) {
        if (typeof template === "string") {
            const data = {
                history: history,
                ...otherData,
            };
            if (history.length > 0 && history[0].role === "system") {
                data.system_message = history.shift().content;
            }
            const compiledTemplate = Handlebars.compile(template);
            return compiledTemplate(data);
        }
        const rendered = template(history, {
            ...otherData,
            supportsCompletions: this.supportsCompletions() ? "true" : "false",
            supportsPrefill: this.supportsPrefill() ? "true" : "false",
        });
        if (typeof rendered !== "string" &&
            rendered[rendered.length - 1]?.role === "assistant" &&
            !canPutWordsInModelsMouth) {
            // Some providers don't allow you to put words in the model's mouth
            // So we have to manually compile the prompt template and use
            // raw /completions, not /chat/completions
            const templateMessages = autodetectTemplateFunction(this.model, this.providerName, autodetectTemplateType(this.model));
            return templateMessages(rendered);
        }
        return rendered;
    }
}
