import chainlit as cl
from extract_app import extract_information


@cl.on_chat_start
def start():
    """
    This is called when the Chainlit chat is started!
    """
    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!
    """
    # res = await extract_information().ainvoke({"input": message})["text"]
    res = await extract_information().ainvoke({"input": message})
    res = res["text"]
    
    # res = await cl.make_async(extract_information().invoke)(
    #     input=message.content, callbacks=[cl.LangchainCallbackHandler()]
    # )
    await cl.Message(content=res).send()