LangchainCallbackHandler and DataLayer causing issues

[wizhippo]

Describe the bug
Using the cl.LangchainCallbackHandler and the official data layer logs have multiple reports

RuntimeError: Task <Task pending name='Task-154' coro=<ChainlitDataLayer.update_step() running at /home/wizhi/src/joy-chat/.venv/lib/python3.12/site-packages/chainlit/data/utils.py:25>> got Future <Future pending cb=[_chain_future.._call_check_cancel() at /usr/lib/python3.12/asyncio/futures.py:387]> attached to a different loop

To Reproduce
Steps to reproduce the behavior:

Run simple app: with datalayer enabled:

from dotenv import load_dotenv

load_dotenv()

import logging

logging.basicConfig(level=logging.INFO)

from typing import Dict, Optional
from langchain.chat_models import init_chat_model
from langchain_core.messages import AIMessageChunk
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.config import RunnableConfig
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.prebuilt import ToolNode, tools_condition

import chainlit as cl


@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b


tools = [multiply]
model = init_chat_model("openai:gpt-4o-mini", temperature=0)
model = model.bind_tools(tools)

system_prompt = """
Role:
    You are a specialized assistant.

Response Process:

1. Initial Search for Relevant Content:
    Tool: multiply
    Action: Use this tool multiply numbers
    Integration: If you find usable and pertinent information, incorporate it into your response.
    """

prompt_template = ChatPromptTemplate([
    ("system", system_prompt),
    MessagesPlaceholder("messages")
])


def should_continue(state: MessagesState):
    messages = state["messages"]
    last_message = messages[-1]
    if last_message.tool_calls:
        return "tools"
    return END


def call_model(state: MessagesState):
    messages = state['messages']
    messages = prompt_template.invoke({"messages": messages})
    response = model.invoke(messages)
    return {"messages": [response]}


workflow = StateGraph(state_schema=MessagesState)
tool_node = ToolNode(tools)

workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)

workflow.add_edge(START, "agent")
workflow.add_conditional_edges("agent", tools_condition)
workflow.add_edge("tools", "agent")

memory = MemorySaver()

app = workflow.compile(checkpointer=memory)


@cl.password_auth_callback
def auth_callback(username: str, password: str):
    # Fetch the user matching username from your database
    # and compare the hashed password with the value stored in the database
    if (username, password) == ("admin", "admin"):
        return cl.User(
            identifier="admin", metadata={"role": "admin", "provider": "credentials"}
        )
    else:
        return None


@cl.on_message
async def main(message: cl.Message):
    answer = cl.Message(content="")
    await answer.send()

    config: RunnableConfig = {
        "callbacks": [
            cl.LangchainCallbackHandler()
        ],
        "configurable": {"thread_id": cl.context.session.thread_id}
    }

    for msg, _ in app.stream(
            {"messages": [HumanMessage(content=message.content)]},
            config,
            stream_mode="messages",
    ):
        if isinstance(msg, AIMessageChunk):
            answer.content += msg.content  # type: ignore
            await answer.update()


if __name__ == "__main__":
    from chainlit.cli import run_chainlit

    run_chainlit(__file__)

Expected behavior
Expect no errors regarding the datastore

• OS: Ubunti 24.04
• Browser chrome

Additional context
Errors are not present if not using the LangchainCallbackHandler

[wizhippo]

Tools and node calls must be asyc too:

This works

from dotenv import load_dotenv

load_dotenv()

import logging

logging.basicConfig(level=logging.INFO)

from typing import Literal
from langchain_core.tools import tool
from langchain.chat_models import init_chat_model
from langgraph.prebuilt import ToolNode
from langchain.schema.runnable.config import RunnableConfig
from langgraph.graph import END, StateGraph, START
from langgraph.graph.message import MessagesState
from langchain_core.messages import SystemMessage, HumanMessage

import chainlit as cl


@tool
async def multiply(a: int, b: int) -> int:
    """Multiply two integers.

    Args:
        a: First integer
        b: Second integer
    """
    return a * b


tools = [multiply]
model = init_chat_model("openai:gpt-4o-mini", temperature=0, streaming=True)
model = model.bind_tools(tools)

final_model = init_chat_model("openai:gpt-4o-mini", temperature=0, streaming=True)
final_model = final_model.with_config(tags=["final_node"])


def should_continue(state: MessagesState) -> Literal["tools", "final"]:
    messages = state["messages"]
    last_message = messages[-1]
    if last_message.tool_calls:
        return "tools"
    return "final"


async def call_model(state: MessagesState):
    messages = state["messages"]
    response = await model.ainvoke(messages)
    return {"messages": [response]}


async def call_final_model(state: MessagesState):
    messages = state["messages"]
    last_ai_message = messages[-1]
    response = await final_model.ainvoke(
        [
            SystemMessage("Rewrite this in the voice of Al Roker"),
            HumanMessage(last_ai_message.content),
        ]
    )
    response.id = last_ai_message.id
    return {"messages": [response]}


tool_node = ToolNode(tools=tools)

builder = StateGraph(MessagesState)
builder.add_node("agent", call_model)
builder.add_node("tools", tool_node)
builder.add_node("final", call_final_model)
builder.add_edge(START, "agent")
builder.add_conditional_edges(
    "agent",
    should_continue,
)
builder.add_edge("tools", "agent")
builder.add_edge("final", END)

graph = builder.compile()


@cl.password_auth_callback
def auth_callback(username: str, password: str):
    if (username, password) == ("admin", "admin"):
        return cl.User(
            identifier="admin", metadata={"role": "admin", "provider": "credentials"}
        )
    else:
        return None


@cl.on_chat_start
async def on_chat_start():
    thread_id = cl.user_session.get("id")
    cl.user_session.set("thread_id", thread_id)


@cl.on_message
async def on_message(message: cl.Message):
    thread_id = cl.user_session.get("thread_id")
    config = {"configurable": {"thread_id": thread_id}}
    final_answer = cl.Message(content="")

    async for msg, metadata in graph.astream(
            {"messages": [HumanMessage(content=message.content)]},
            stream_mode="messages",
            config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()], **config)
    ):
        if (
                msg.content
                and not isinstance(msg, HumanMessage)
                and metadata["langgraph_node"] == "final"
        ):
            await final_answer.stream_token(msg.content)

    await final_answer.send()


@cl.set_starters
async def set_starters():
    return [
        cl.Starter(
            label="Multiply 4 * 5",
            message="Using the multiply tool, multiply 4 * 5",
            icon=None,
        )
    ]


if __name__ == "__main__":
    from chainlit.cli import run_chainlit

    run_chainlit(__file__)

[vio-p]

I have also encountered this problem when using the prebuilt ReAct agent from LangGraph, even without any tools. I have looked through the implementation but I can't figure out where the problem comes from, if it's on the LangGraph side or the Chainlit side. I have temporarily fixed it by implementing my own agent, thanks to your issue, but I would like to eventually use the prebuilt one as it has more features. @wizhippo Do you have any insights on this?

[wizhippo]

@vio-p My issue was I was using non-async tools. Once I made sure everything was properly async my this issue went away.

I have another issue with the data layer when sending a message vs updating message #1937 probably not related but is seams something is off in the data layer when it comes to async.

[Hmkhalla]

Hi, I'm still facing this issue, I don't think it has been solved