from . import Tool, LLMChain, llm
from langchain.memory import ConversationBufferMemory
from langchain_core.messages import SystemMessage
from langchain_core.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
)

system_message = "You are AI Assistant. You need to give crystal clear answer.\n"
template_messages = [
        SystemMessage(content=system_message),
        MessagesPlaceholder(variable_name="chat_history"),
        HumanMessagePromptTemplate.from_template("{text}"),
    ]
prompt_template = ChatPromptTemplate.from_messages(template_messages)
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

word_problem_chain = LLMChain(llm=llm, prompt=prompt_template, memory=memory)
word_problem_tool = Tool.from_function(name="Reasoning Tool", \
    func=word_problem_chain.run, \
    description="Useful for when you need to answer logic-based/reasoning \
                questions.",
)