StreamlitCallbackHandlerがLangGraphなどMulti-Threadで使えない

Example code

import streamlit as st
from langchain.callbacks.streamlit import StreamlitCallbackHandler

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langgraph.graph import END, MessageGraph


handler = StreamlitCallbackHandler(st.container())

model = ChatOpenAI(temperature=0)

graph = MessageGraph()

graph.add_node("oracle", model)
graph.add_edge("oracle", END)

graph.set_entry_point("oracle")

runnable = graph.compile()

runnable.invoke(HumanMessage("What is 1 + 1?"), {"callbacks": [handler]})

Error

2024-05-24 11:50:06.661 Thread 'ThreadPoolExecutor-1_0': missing ScriptRunContext
Error in StreamlitCallbackHandler.on_llm_start callback: NoSessionContext()
Error in StreamlitCallbackHandler.on_llm_end callback: RuntimeError('Current LLMThought is unexpectedly None!')

Workaround?

試してみたが、一部動いたケースもあったが、全てのケースに対してカバーはできなそうだった。

Related Issues

1. [StreamlitCallbackHandler] - Not compatible with LangGraph #101 in langchain-ai/langgraph
2. Feature Request: Add support for multi-threading/multi-processing in Streamlit #8490 in streamlit/streamlit
3. Stream Chat LLM Token By Token is not working #78