基于LLM的Streamlit聊天机器人图表无法持久化问题

你遇到的图表消失问题，核心原因是Streamlit每次用户交互都会重新执行整个脚本，而当前代码里仅把文本消息存在st.session_state.output中，图表没有被持久化存储。每次新输入触发脚本重跑时，之前生成的图表不会被重新渲染，自然就消失了。

需要扩展会话状态的消息结构，让它能同时保存文本和图表数据，然后在遍历历史消息时，针对性渲染文本或图表。具体步骤如下：

1. 重构会话状态的消息格式 #

把原来仅存AIMessage/HumanMessage的列表，改成包含消息类型、角色、内容的字典结构，支持两种类型：

• text：存储普通对话文本
• chart：存储图表对应的DataFrame

2. 修改历史消息渲染逻辑 #

遍历会话状态的消息时，根据类型判断是渲染文本还是图表。

3. 调整图表生成和存储逻辑 #

生成图表时，把图表数据以chart类型的消息存入会话状态，而不是临时渲染后就丢弃。

import streamlit as st
from langchain_core.messages import AIMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
import pandas as pd

# 初始化会话状态：支持文本和图表消息
if "chat_messages" not in st.session_state:
    st.session_state.chat_messages = [
        {"type": "text", "role": "assistant", "content": "Hello, I am an Intent Classifier Assistant. How can I help you?"}
    ]

def process_data(data_path):
    return pd.read_csv(data_path)

def prepare_plot_data(datapath, datetime_column):
    df = process_data(datapath)
    df[datetime_column] = pd.to_datetime(df[datetime_column], format='mixed')
    return df.set_index(datetime_column)

# app config
st.set_page_config(page_title="Intent Classifier Assistant", page_icon="🤖")
st.title("Intent Classifier")

def get_prompt(input_text):
    examples = [
        {"input": "How many pens are there?", "output": "other_intent"},
        {"input": "Compare top 5 contributors to WHO", "output": "plot_intent"},
        {"input": "What is the average sales citiwise per stores in the US", "output": "plot_intent"},
        {"input": "What is the average cost of bike in Indian small towns", "output": "other_intent"},
        {"input": "What is the maximum cost of bike in Indian small towns", "output": "other_intent"},
        {"input": "What is the max cost of bike in Indian small towns", "output": "other_intent"},
        {"input": "What is the min cost of bike in Indian small towns", "output": "other_intent"},
        {"input": "What is the mean cost of bike in Indian small towns", "output": "other_intent"},
        {"input": "What is the maximum cost of bike in Asian small towns, provide countrywise details", "output": "plot_intent"}
    ]

    example_prompt = ChatPromptTemplate.from_messages(
        [("human", "{input}"), ("ai", "{output}")]
    )
    few_shot_prompt = FewShotChatMessagePromptTemplate(
        example_prompt=example_prompt, examples=examples
    )

    final_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are an AI assistant that classifies user intents based on the questions asked. You have to answer one intent among: greeting_intent, plot_intent , other_intent."),
            few_shot_prompt,
            ("human", "{input}"),
        ]
    )
    return final_prompt

def get_intent(input_text):
    prompt = get_prompt(input_text)
    openai_api_key = "abcdedf"  # 替换为你的实际API密钥
    llm = ChatOpenAI(temperature=0, openai_api_key=openai_api_key)
    chain = prompt | llm | StrOutputParser()
    return chain.stream({"input": input_text})

# 渲染历史消息
for msg in st.session_state.chat_messages:
    with st.chat_message(msg["role"]):
        if msg["type"] == "text":
            st.write(msg["content"])
        elif msg["type"] == "chart":
            st.line_chart(msg["content"])

# 用户输入处理
user_input = st.chat_input("Type your message here...")
if user_input:
    # 添加用户文本消息到会话状态
    st.session_state.chat_messages.append({
        "type": "text", "role": "user", "content": user_input
    })

    with st.chat_message("user"):
        st.markdown(user_input)

    with st.chat_message("assistant"):
        # 获取意图分类结果
        intent_response = st.write_stream(get_intent(user_input))
        # 添加AI文本消息到会话状态
        st.session_state.chat_messages.append({
            "type": "text", "role": "assistant", "content": intent_response
        })
        
        # 如果是绘图意图，生成并存储图表数据
        if "plot_intent" in intent_response:
            chart_df = prepare_plot_data("./data/Electric_Production.csv", "DATE")
            # 添加图表消息到会话状态
            st.session_state.chat_messages.append({
                "type": "chart", "role": "assistant", "content": chart_df
            })
            # 渲染当前图表
            st.line_chart(chart_df)

• 新增chat_messages会话状态，用字典存储消息类型、角色和内容，支持文本和图表两种类型
• 修改prepare_plot_data函数，只负责处理数据返回DataFrame，不直接渲染图表
• 历史消息渲染时，根据消息类型判断是显示文本还是绘制图表
• 用户输入处理时，把文本和图表消息都存入会话状态，确保每次脚本重跑都能重新渲染所有历史内容

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