如何在Langgraph中正确使用StreamlitCallbackHandler？

问题描述 #

运行结合Langgraph多智能体和StreamlitCallbackHandler的代码时，出现以下报错：

2024-02-18 13:30:17.030 Thread 'ThreadPoolExecutor-5_0': missing ScriptRunContext
Error in StreamlitCallbackHandler.on_llm_start callback: NoSessionContext()
Error in StreamlitCallbackHandler.on_llm_end callback: RuntimeError('Current LLMThought is unexpectedly None!')
Error in StreamlitCallbackHandler.on_tool_end callback: RuntimeError('Current LLMThought is unexpectedly None!')
2024-02-18 13:30:18.630 Thread 'ThreadPoolExecutor-5_0': missing ScriptRunContext
Error in StreamlitCallbackHandler.on_llm_start callback: NoSessionContext()
Error in StreamlitCallbackHandler.on_llm_end callback: RuntimeError('Current LLMThought is unexpectedly None!')

核心代码片段如下：

#...Other Langgraph code from the example

workflow = StateGraph(AgentState)

workflow.add_node("Researcher", research_node)
workflow.add_node("Chart Generator", chart_node)
workflow.add_node("call_tool", tool_node)

workflow.add_conditional_edges(
    "Researcher",
    router,
    {"continue": "Chart Generator", "call_tool": "call_tool", "end": END},
)
workflow.add_conditional_edges(
    "Chart Generator",
    router,
    {"continue": "Researcher", "call_tool": "call_tool", "end": END},
)

workflow.add_conditional_edges(
    "call_tool",
    lambda x: x["sender"],
    {
        "Researcher": "Researcher",
        "Chart Generator": "Chart Generator",
    },
)
workflow.set_entry_point("Researcher")
graph = workflow.compile()

#... Other Streamlit configurations

with st.form(key="form"):
    user_input = st.text_input("Define the task")
    submit_clicked = st.form_submit_button("Execute")

output_container = st.empty()
if with_clear_container(submit_clicked):
    output_container = output_container.container()
    output_container.chat_message("user").write(user_input)

    answer_container = output_container.chat_message("assistant", avatar="🦜")
    st_callback = StreamlitCallbackHandler(answer_container)
    cfg = RunnableConfig()
    cfg["callbacks"] = [st_callback]
    cfg["recursion_limit"] = 100

    answer = graph.invoke({
            "messages": [
                HumanMessage(
                    content=user_input
                )
            ],
        }, cfg)

    answer_container.write(answer["content"])

报错原因 #

1. Streamlit会话上下文丢失：Langgraph在多线程环境中执行节点，而Streamlit的ScriptRunContext默认只在主线程中存在，子线程无法访问，导致回调无法正常渲染UI。
2. CallbackHandler适配问题：原生StreamlitCallbackHandler是为单Agent流程设计的，无法适配Langgraph多Agent节点切换的场景，导致LLMThought对象管理混乱。

解决方法 #

1. 传递Streamlit会话上下文到子线程 #

在每个Langgraph节点的执行逻辑中，手动推送Streamlit的会话上下文，确保子线程能访问到UI渲染所需的上下文：

首先导入相关工具：

from streamlit.runtime.scriptrunner import get_script_run_context, push_script_run_context

然后修改节点函数，比如research_node和chart_node：

def research_node(state: AgentState):
    # 获取主线程的会话上下文
    ctx = get_script_run_context()
    # 在当前子线程中推送上下文
    push_script_run_context(ctx)
    
    # 原节点逻辑...
    # 比如调用LLM、处理消息等
    return state

同理修改chart_node和tool_node，确保每个节点执行时都有会话上下文。

2. 改用适配多Agent的回调处理 #

原生StreamlitCallbackHandler不适合多Agent场景，推荐自定义适配Langgraph的回调类，简化UI更新逻辑：

from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import LLMResult

class LanggraphStreamlitCallback(BaseCallbackHandler):
    def __init__(self, container):
        self.container = container
        self.current_agent = None

    def on_llm_start(self, serialized, prompts, **kwargs):
        # 记录当前执行的Agent
        self.current_agent = kwargs.get("tags", ["Unknown"])[0]
        self.container.markdown(f"**{self.current_agent}** 正在思考...")

    def on_llm_end(self, response: LLMResult, **kwargs):
        # 输出LLM结果
        if response.generations:
            content = response.generations[0][0].text
            self.container.markdown(f"**{self.current_agent}** 输出：{content}")

    def on_tool_start(self, serialized, input_str, **kwargs):
        # 输出工具调用信息
        tool_name = serialized.get("name", "Unknown Tool")
        self.container.markdown(f"**{self.current_agent}** 调用工具：{tool_name}，输入：{input_str}")

    def on_tool_end(self, output, **kwargs):
        # 输出工具返回结果
        self.container.markdown(f"工具返回：{output}")

然后在代码中使用这个自定义回调：

answer_container = output_container.chat_message("assistant", avatar="🦜")
st_callback = LanggraphStreamlitCallback(answer_container)
cfg = RunnableConfig()
cfg["callbacks"] = [st_callback]
cfg["recursion_limit"] = 100

3. 调整Langgraph的执行模式 #

如果不需要多线程执行，可以强制Langgraph使用同步模式，避免上下文丢失：

# 调用时使用同步执行
answer = graph.invoke(
    {"messages": [HumanMessage(content=user_input)]},
    cfg,
    sync=True  # 强制同步执行
)

修改后核心代码示例 #

from streamlit.runtime.scriptrunner import get_script_run_context, push_script_run_context
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import LLMResult

# 自定义回调类
class LanggraphStreamlitCallback(BaseCallbackHandler):
    def __init__(self, container):
        self.container = container
        self.current_agent = None

    def on_llm_start(self, serialized, prompts, **kwargs):
        self.current_agent = kwargs.get("tags", ["Unknown"])[0]
        self.container.markdown(f"**{self.current_agent}** 正在思考...")

    def on_llm_end(self, response: LLMResult, **kwargs):
        if response.generations:
            content = response.generations[0][0].text
            self.container.markdown(f"**{self.current_agent}** 输出：{content}")

    def on_tool_start(self, serialized, input_str, **kwargs):
        tool_name = serialized.get("name", "Unknown Tool")
        self.container.markdown(f"**{self.current_agent}** 调用工具：{tool_name}，输入：{input_str}")

    def on_tool_end(self, output, **kwargs):
        self.container.markdown(f"工具返回：{output}")

# 修改节点函数，添加上下文传递
def research_node(state: AgentState):
    ctx = get_script_run_context()
    push_script_run_context(ctx)
    # 原节点逻辑
    # ...
    state["sender"] = "Researcher"
    return state

def chart_node(state: AgentState):
    ctx = get_script_run_context()
    push_script_run_context(ctx)
    # 原节点逻辑
    # ...
    state["sender"] = "Chart Generator"
    return state

def tool_node(state: AgentState):
    ctx = get_script_run_context()
    push_script_run_context(ctx)
    # 原工具调用逻辑
    # ...
    return state

# ... Langgraph workflow定义部分不变 ...

# Streamlit部分
with st.form(key="form"):
    user_input = st.text_input("Define the task")
    submit_clicked = st.form_submit_button("Execute")

output_container = st.empty()
if with_clear_container(submit_clicked):
    output_container = output_container.container()
    output_container.chat_message("user").write(user_input)

    answer_container = output_container.chat_message("assistant", avatar="🦜")
    st_callback = LanggraphStreamlitCallback(answer_container)
    cfg = RunnableConfig()
    cfg["callbacks"] = [st_callback]
    cfg["recursion_limit"] = 100

    answer = graph.invoke(
        {"messages": [HumanMessage(content=user_input)]},
        cfg,
        sync=True
    )

    answer_container.write(f"最终结果：{answer['messages'][-1].content}")

内容的提问来源于stack exchange，提问作者James