import { globalsHelper } from "./GlobalsHelper.js";
import { defaultSummaryPrompt, messagesToHistoryStr } from "./Prompt.js";
import { OpenAI } from "./llm/openai.js";
import { extractText } from "./llm/utils.js";
/**
 * A ChatHistory is used to keep the state of back and forth chat messages
 */ export class ChatHistory {
}
export class SimpleChatHistory extends ChatHistory {
    messages;
    messagesBefore;
    constructor(init){
        super();
        this.messages = init?.messages ?? [];
        this.messagesBefore = this.messages.length;
    }
    addMessage(message) {
        this.messages.push(message);
    }
    async requestMessages(transientMessages) {
        return [
            ...transientMessages ?? [],
            ...this.messages
        ];
    }
    reset() {
        this.messages = [];
    }
    newMessages() {
        const newMessages = this.messages.slice(this.messagesBefore);
        this.messagesBefore = this.messages.length;
        return newMessages;
    }
}
export class SummaryChatHistory extends ChatHistory {
    /**
   * Tokenizer function that converts text to tokens,
   *  this is used to calculate the number of tokens in a message.
   */ tokenizer = globalsHelper.defaultTokenizer.encode;
    tokensToSummarize;
    messages;
    summaryPrompt;
    llm;
    messagesBefore;
    constructor(init){
        super();
        this.messages = init?.messages ?? [];
        this.messagesBefore = this.messages.length;
        this.summaryPrompt = init?.summaryPrompt ?? defaultSummaryPrompt;
        this.llm = init?.llm ?? new OpenAI();
        if (!this.llm.metadata.maxTokens) {
            throw new Error("LLM maxTokens is not set. Needed so the summarizer ensures the context window size of the LLM.");
        }
        this.tokensToSummarize = this.llm.metadata.contextWindow - this.llm.metadata.maxTokens;
        if (this.tokensToSummarize < this.llm.metadata.contextWindow * 0.25) {
            throw new Error("The number of tokens that trigger the summarize process are less than 25% of the context window. Try lowering maxTokens or use a model with a larger context window.");
        }
    }
    async summarize() {
        // get the conversation messages to create summary
        const messagesToSummarize = this.calcConversationMessages();
        let promptMessages;
        do {
            promptMessages = [
                {
                    content: this.summaryPrompt({
                        context: messagesToHistoryStr(messagesToSummarize)
                    }),
                    role: "user",
                    options: {}
                }
            ];
            // remove oldest message until the chat history is short enough for the context window
            messagesToSummarize.shift();
        }while (this.tokenizer(promptMessages[0].content).length > this.tokensToSummarize)
        const response = await this.llm.chat({
            messages: promptMessages
        });
        return {
            content: response.message.content,
            role: "memory"
        };
    }
    addMessage(message) {
        this.messages.push(message);
    }
    // Find last summary message
    getLastSummaryIndex() {
        const reversedMessages = this.messages.slice().reverse();
        const index = reversedMessages.findIndex((message)=>message.role === "memory");
        if (index === -1) {
            return null;
        }
        return this.messages.length - 1 - index;
    }
    getLastSummary() {
        const lastSummaryIndex = this.getLastSummaryIndex();
        return lastSummaryIndex ? this.messages[lastSummaryIndex] : null;
    }
    get systemMessages() {
        // get array of all system messages
        return this.messages.filter((message)=>message.role === "system");
    }
    get nonSystemMessages() {
        // get array of all non-system messages
        return this.messages.filter((message)=>message.role !== "system");
    }
    /**
   * Calculates the messages that describe the conversation so far.
   * If there's no memory, all non-system messages are used.
   * If there's a memory, uses all messages after the last summary message.
   */ calcConversationMessages(transformSummary) {
        const lastSummaryIndex = this.getLastSummaryIndex();
        if (!lastSummaryIndex) {
            // there's no memory, so just use all non-system messages
            return this.nonSystemMessages;
        } else {
            // there's a memory, so use all messages after the last summary message
            // and convert summary message so it can be send to the LLM
            const summaryMessage = transformSummary ? {
                content: `Summary of the conversation so far: ${this.messages[lastSummaryIndex].content}`,
                role: "system"
            } : this.messages[lastSummaryIndex];
            return [
                summaryMessage,
                ...this.messages.slice(lastSummaryIndex + 1)
            ];
        }
    }
    calcCurrentRequestMessages(transientMessages) {
        // TODO: check order: currently, we're sending:
        // system messages first, then transient messages and then the messages that describe the conversation so far
        return [
            ...this.systemMessages,
            ...transientMessages ? transientMessages : [],
            ...this.calcConversationMessages(true)
        ];
    }
    async requestMessages(transientMessages) {
        const requestMessages = this.calcCurrentRequestMessages(transientMessages);
        // get tokens of current request messages and the transient messages
        const tokens = requestMessages.reduce((count, message)=>count + this.tokenizer(extractText(message.content)).length, 0);
        if (tokens > this.tokensToSummarize) {
            // if there are too many tokens for the next request, call summarize
            const memoryMessage = await this.summarize();
            const lastMessage = this.messages.at(-1);
            if (lastMessage && lastMessage.role === "user") {
                // if last message is a user message, ensure that it's sent after the new memory message
                this.messages.pop();
                this.messages.push(memoryMessage);
                this.messages.push(lastMessage);
            } else {
                // otherwise just add the memory message
                this.messages.push(memoryMessage);
            }
            // TODO: we still might have too many tokens
            // e.g. too large system messages or transient messages
            // how should we deal with that?
            return this.calcCurrentRequestMessages(transientMessages);
        }
        return requestMessages;
    }
    reset() {
        this.messages = [];
    }
    newMessages() {
        const newMessages = this.messages.slice(this.messagesBefore);
        this.messagesBefore = this.messages.length;
        return newMessages;
    }
}
export function getHistory(chatHistory) {
    if (chatHistory instanceof ChatHistory) {
        return chatHistory;
    }
    return new SimpleChatHistory({
        messages: chatHistory
    });
}
