import { RunnableWithMessageHistory } from "@langchain/core/runnables"; 
import { ChatPromptTemplate, MessagesPlaceholder } from "@langchain/core/prompts"; 

export function createRephraseQuestionChain() {
    const REPHRASE_QUESTION_SYSTEM_TEMPLATE = `
    Given the following conversation and a follow up question,
    rephrase the follow up question to be a standalone question.
    `;

    // Rephrase the question to be a standalone question
    // along with passing the chat history
    const rephraseQuestionChainPrompt = ChatPromptTemplate.fromMessages([
        ["system", REPHRASE_QUESTION_SYSTEM_TEMPLATE],
        new MessagesPlaceholder("history"),
        ["human", "Rephrase the following question as a standalone question:\n{question}"],
    ]);

    // Runnable to rephrase the question
    const rephraseQuestionChain = RunnableSequence.from([
        rephraseQuestionChainPrompt,
        new ChatGoogleGenerativeAI({
            modelName: "gemini-pro",
            maxOutputTokens: 2048,
            apiKey: process.env.google_api_key
        }),
        new StringOutputParser(),
    ]);
    return rephraseQuestionChain;
}

const rephraseQuestionChain = createRephraseQuestionChain();

// make a prompt template for the response
const ANSWER_CHAIN_SYSTEM_TEMPLATE = `You are an experienced researcher,
expert at interpreting and answering questions based on provided sources.
Using the below provided context and chat history, 
answer the user's question to the best of your ability
using only the resources provided. Be verbose!

<context>
{context}
</context>`;

// Creat a prompt template from Messages
const answerGenerationChainPrompt = ChatPromptTemplate.fromMessages([
    ["system", ANSWER_CHAIN_SYSTEM_TEMPLATE],
    new MessagesPlaceholder("history"),
    [
        "human",
        `Now, answer this question using the previous context and chat history:

    {standalone_question}`
    ]
]);

// initialize the chain to retrieve the answer for the question
const conversationalRetrievalChain = RunnableSequence.from([
    RunnablePassthrough.assign({
        standalone_question: rephraseQuestionChain,
    }),
    RunnablePassthrough.assign({
        context: documentRetrievalChain,
    }),
    answerGenerationChainPrompt,
    new ChatGoogleGenerativeAI({
        modelName: "gemini-pro",
        maxOutputTokens: 2048,
        apiKey: process.env.google_api_key
    }),
]);

import { ChatMessageHistory } from "langchain/stores/message/in_memory";

// we should create a new history object per session
const messageHistories = {};
const getMessageHistoryForSession = (sessionId) => {
    if (messageHistories[sessionId] !== undefined) {
        return messageHistories[sessionId];
    }
    const newChatSessionHistory = new ChatMessageHistory();
    messageHistories[sessionId] = newChatSessionHistory;
    return newChatSessionHistory;
};
const finalRetrievalChain = new RunnableWithMessageHistory({
    runnable: conversationalRetrievalChain,
    getMessageHistory: getMessageHistoryForSession,
    inputMessagesKey: "question",
    historyMessagesKey: "history",
}).pipe(new StringOutputParser())

const result = await finalRetrievalChain.invoke({
  question: "What are the pre-requisites of this course?",
}, {
  configurable: { sessionId: "test" }
});

console.log(result)

//Input: Our loaded PDF has details about a Machine Learning Course by Andrew NG
//Output:
/* 
To enroll in this class, you must meet the following prerequisites:

* Basic probability and statistics. A course like Statistics 116 would be more than enough.

* Basic linear algebra. Courses like Math 51, 103, 113, or CS205 at Stanford are recommended. Matrix and vector multiplication, matrix inverse, and eigenvectors are concepts you should be familiar with.

* Programming skills. You should be able to write a simple computer program.

* Understanding of big O notation and data structures such as queues, stacks, and binary trees.

If you have taken these courses or have equivalent knowledge, you should be prepared for this class. However, it is important to note that while programming will not be a major focus of the course, some programming in MATLAB or Octave will be required.
*/