app.py

import chainlit as cl

from modules.main_pipeline import Inference
from modules.utils import apply_decorator_to_methods, log_execution_time
from setting import CFG


@cl.on_chat_start
async def factory():
    if CFG.track_execution_time:
        apply_decorator_to_methods(Inference, log_execution_time)

    inference_model = Inference(
        llm_config=dict(model=CFG.model)
    )
    cl.user_session.set("inference_model", inference_model)


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

    assistant_messages = inference_model(message.content)
    await cl.Message(content=assistant_messages).send()