import chainlit as cl
from langchain_core.runnables import RunnableConfig
from rag import chain


@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("runnable", chain)


@cl.on_message
async def on_message(message: cl.Message):
    runnable = cl.user_session.get("runnable")

    msg = cl.Message(content="")

    async for chunk in runnable.astream(
        message.content,
        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
    ):
        await msg.stream_token(chunk)

    await msg.send()