import { stripImages } from "../countTokens.js";
import { BaseLLM } from "../index.js";
import { streamResponse } from "../stream.js";
class Gemini extends BaseLLM {
    static providerName = "gemini";
    static defaultOptions = {
        model: "gemini-pro",
        apiBase: "https://generativelanguage.googleapis.com/v1beta/",
    };
    async *_streamComplete(prompt, options) {
        for await (const message of this._streamChat([{ content: prompt, role: "user" }], options)) {
            yield stripImages(message.content);
        }
    }
    removeSystemMessage(messages) {
        const msgs = [...messages];
        if (msgs[0]?.role === "system") {
            const sysMsg = msgs.shift()?.content;
            // @ts-ignore
            if (msgs[0]?.role === "user") {
                msgs[0].content = `System message - follow these instructions in every response: ${sysMsg}\n\n---\n\n${msgs[0].content}`;
            }
        }
        return msgs;
    }
    async *_streamChat(messages, options) {
        const convertedMsgs = this.removeSystemMessage(messages);
        if (options.model.includes("gemini")) {
            for await (const message of this.streamChatGemini(convertedMsgs, options)) {
                yield message;
            }
        }
        else {
            for await (const message of this.streamChatBison(convertedMsgs, options)) {
                yield message;
            }
        }
    }
    _continuePartToGeminiPart(part) {
        return part.type === "text"
            ? {
                text: part.text,
            }
            : {
                inlineData: {
                    mimeType: "image/jpeg",
                    data: part.imageUrl?.url.split(",")[1],
                },
            };
    }
    async *streamChatGemini(messages, options) {
        const apiURL = new URL(`models/${options.model}:streamGenerateContent?key=${this.apiKey}`, this.apiBase);
        const body = {
            contents: messages.map((msg) => {
                return {
                    role: msg.role === "assistant" ? "model" : "user",
                    parts: typeof msg.content === "string"
                        ? [{ text: msg.content }]
                        : msg.content.map(this._continuePartToGeminiPart),
                };
            }),
        };
        const response = await this.fetch(apiURL, {
            method: "POST",
            body: JSON.stringify(body),
        });
        let buffer = "";
        for await (const chunk of streamResponse(response)) {
            buffer += chunk;
            if (buffer.startsWith("[")) {
                buffer = buffer.slice(1);
            }
            if (buffer.endsWith("]")) {
                buffer = buffer.slice(0, -1);
            }
            if (buffer.startsWith(",")) {
                buffer = buffer.slice(1);
            }
            const parts = buffer.split("\n,");
            let foundIncomplete = false;
            for (let i = 0; i < parts.length; i++) {
                const part = parts[i];
                let data;
                try {
                    data = JSON.parse(part);
                }
                catch (e) {
                    foundIncomplete = true;
                    continue;
                }
                if (data.error) {
                    throw new Error(data.error.message);
                }
                // Incrementally stream the content to make it smoother
                const content = data.candidates[0].content.parts[0].text;
                const words = content.split(/(\s+)/);
                const delaySeconds = Math.min(4.0 / (words.length + 1), 0.1);
                while (words.length > 0) {
                    const wordsToYield = Math.min(3, words.length);
                    yield {
                        role: "assistant",
                        content: words.splice(0, wordsToYield).join(""),
                    };
                    await delay(delaySeconds);
                }
            }
            if (foundIncomplete) {
                buffer = parts[parts.length - 1];
            }
            else {
                buffer = "";
            }
        }
    }
    async *streamChatBison(messages, options) {
        const msgList = [];
        for (const message of messages) {
            msgList.push({ content: message.content });
        }
        const apiURL = new URL(`models/${options.model}:generateMessage?key=${this.apiKey}`, this.apiBase);
        const body = { prompt: { messages: msgList } };
        const response = await this.fetch(apiURL, {
            method: "POST",
            body: JSON.stringify(body),
        });
        const data = await response.json();
        yield { role: "assistant", content: data.candidates[0].content };
    }
}
async function delay(seconds) {
    return new Promise((resolve) => setTimeout(resolve, seconds * 1000));
}
export default Gemini;
