from langchain.prompts import (
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
    ChatPromptTemplate,
    MessagesPlaceholder
)
from langchain.chains import ConversationChain

class Chain:
    def __init__(self, llm, history=None):
        self.llm = llm
        # self.chain = self.get_conversational_chain()
        if history is not None:
            self.history = history

    def run_conversational_chain(self, prompt_template):
        
        ans = self.llm.invoke(prompt_template).content

        return ans
    
    def get_chain_with_history(self):
        system_msg_template = SystemMessagePromptTemplate.from_template(template="""Answer the question as truthfully as possible using the provided context, 
        and if the answer is not contained within the text below, say 'I don't know'""")
        human_msg_template = HumanMessagePromptTemplate.from_template(template="{input}")
        prompt_template = ChatPromptTemplate.from_messages([system_msg_template, MessagesPlaceholder(variable_name="history"), human_msg_template])
        conversation = ConversationChain(memory=self.history, prompt=prompt_template, llm=self.llm, verbose=True)
        return conversation