from langchain import PromptTemplate, LLMChain
import chainlit as cl
from custom_llm import CustomLLM
from langchain.prompts import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
)

template = """ Write a code for the following problem :
{question} 

Code:
"""


@cl.langchain_factory(use_async=False)
def factory():
    system_message_prompt = SystemMessagePromptTemplate.from_template(template)

    prompt = ChatPromptTemplate.from_messages([system_message_prompt])
    llm = CustomLLM()

    llm_chain = LLMChain(prompt=prompt, llm=llm, verbose=True,)

    cl.user_session.set("llm_chain", llm_chain)
    return llm_chain


@cl.langchain_rename   # This will be particularly useful when we want to customize this thing for production.
def rename(orig_author):
    rename_dict = {
        'LLMChain': 'Scooby'
    }
    return rename_dict.get(orig_author, orig_author)


@cl.on_chat_start
async def main():
    await cl.Message(content="Welcome to CodeSmith !! Let's start this").send()