import sys

sys.path.append(".")

import chainlit as cl

from earnings_app import extract_information

@cl.author_rename
def rename(orig_author: str):
    diamond_char = u'\U0001F537'
    phrase = diamond_char + " Diamond Hands " + diamond_char
    rename_dict = {"RetrievalQA": phrase}
    return rename_dict.get(orig_author, orig_author)

@cl.on_chat_start
async def start():
    """
    This is called when the Chainlit chat is started!
    """
    chain = await extract_information()
    cl.user_session.set("chain", chain)
    await cl.Message("Welcome to the information extraction chat!").send()

@cl.on_message
async def main(message: cl.Message):
    """
    This is called when a message is received!
    """
    chain = cl.user_session.get("chain")
    res = await chain.achat(message.content)
    # res = await chain.aiinvoke({"input": message})
    # res = res["text"]
    out = str(res)
    await cl.Message(content=out).send()