Add chart or graph to streamlit session state

Using Streamlit
,
1

Hi all, I am new to streamlit and LLM, I am trying to build a basic chatbot which would output the intent (either plot_intent/ other_intent/ greeting intent). So, using the below code I am able to see historical text data, I am also able to plot the line chart but as soon as I type the next sentence, the chart vanishes but text is available in history. Could you please let me know if we can save the chart in session state as well. I would like to keep the chart in chat history.

import streamlit as st
from langchain_core.messages import AIMessage, HumanMessage
from langchain_openai import OpenAI, ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.prompts import (
    ChatPromptTemplate,
    FewShotChatMessagePromptTemplate,
)
import matplotlib.pyplot as plt
import matplotlib
import datetime
import pandas as pd
matplotlib.use('Agg')

if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
    
def process_data(data_path):
    data = pd.read_csv(data_path)
    return data


def plot_data(datapath, datetime_column):
    df = process_data(datapath)
    df[datetime_column] = pd.to_datetime(df[datetime_column], format='mixed')
    df = df.set_index(datetime_column)
    # Add the chart data to the chat history
    st.session_state.chat_history.append(('assistant', df))
    # Create a line chart
    return st.line_chart(df)


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

def get_prompt(input, output):
    examples = [
    {"input": "How many pens are there?", "output": "other_intent"},
    {"input": "Compare top 5 contibutors to WHO", "output": "plot_intent"},
    {"input": "What is the average sales citiwise per stores in teh 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"}
    ]

    # This is a prompt template used to format each individual example.
    example_prompt = ChatPromptTemplate.from_messages(
        [
            ("human", "{input}"),
            ("ai", "{output}"),
        ]
    )
    few_shot_prompt = FewShotChatMessagePromptTemplate(
        example_prompt=example_prompt,
        examples=examples,
    )

    print(few_shot_prompt.format())

    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, output):
    prompt = get_prompt(input, output)
    openai_api_key = "abcdedf"
    llm = ChatOpenAI(temperature=0, openai_api_key=openai_api_key)
    chain = prompt | llm | StrOutputParser()
    response = chain.stream({
        "input": input
        })
    st.session_state.chat_history.append(('user', input))
    st.session_state.chat_history.append(('assistant', response))
    return response

# session state
if "output" not in st.session_state:
    st.session_state.output = [
        AIMessage(content="Hello, I am an Intent Classifier Assistant. How can I help you?"),
    ]
# conversation
for message in st.session_state.output:
    if isinstance(message, AIMessage):
        with st.chat_message("AI"):
            st.write(message.content)
    elif isinstance(message, HumanMessage):
        with st.chat_message("Human"):
            st.write(message.content)

# user input
input = st.chat_input("Type your message here...")
if input is not None and input != "":
    #The below line keeps user input history on chat boxes for that session
    st.session_state.output.append(HumanMessage(content=input))

    with st.chat_message("Human"):
        st.markdown(input)

    with st.chat_message("AI"):
        response = st.write_stream(get_intent(input, st.session_state.output))
        # print("THIS IS THE RESPONSE::: " + response, type(response), len(response))
    
        if "plot_intent" in response:
            plotter = plot_data(datapath="./data/Electric_Production.csv", datetime_column="DATE")
    #Below line keeps answer history from the chatbot in chat box history
    st.session_state.output.append(AIMessage(content=response))