chainlit-langchain

chainlit-langchain.py

from typing import Any, cast
import chainlit as cl
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_google_genai import ChatGoogleGenerativeAI


@cl.on_chat_start
async def on_chat_start():
    model = 'gemini-2.5-flash-lite'
    llm = ChatGoogleGenerativeAI(model=model)
    cl.user_session.set("runnable", llm)

@cl.on_message
async def on_message(message: cl.Message):
    runnable = cast(Runnable[Any, Any], 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.content)

    await msg.send()