added the charts

app.py

@@ -13,13 +13,17 @@ import streamlit as st
 from dotenv import load_dotenv
 from sqlalchemy import create_engine, text, URL
 from htmlTemplates import css, bot_template, user_template
-
-import openai
+from openai import OpenAI
 import os
 import time
-
-from be import config, get_db_chain, outreach_sms_message, get_store_address
+from be import config, get_db_chain, outreach_sms_message, get_store_address, data_visualization
 from PIL import Image
+from langchain_core.messages import HumanMessage, AIMessage
+import matplotlib.pyplot as plt
+import re
+
+
+
 def conversation_agent(question):
     llm = ChatOpenAI(temperature=0.5, model="gpt-3.5-turbo")   # create LLM
     # search = DuckDuckGoSearchRun()
@@ -52,26 +56,41 @@ def conversation_agent(question):
         name="Get Store Address",
         func=get_store_address,
         description="Use this tool with store number to get the store address.  INPUT to this tool is Store number.  Do not use this tool if you don't have Store number as input"
+    ) 
+        # create the tool for finding the store details
+    data_visualization_tool = Tool(
+        name="Data Visualization Tool",
+        func=data_visualization,
+        description="ONLY USE THIS TOOL for data visualization, do not create on your own.  This tool is used to create python code to generate charts for data visualization with user question, column names, and pandas dataframe. DO NOT create sample dataframe, always use the provided dataframe in df.  Always make sure the query generation tool is executed and the response is shown to the user before executing this tool"
     ) 
     #tools.append(get_db_chain_tool)
     # List all the tools for the agent to use 
-    tools = [get_db_chain_tool, get_store_address_tool, outreach_sms_message_tool]
+    tools = [get_db_chain_tool, get_store_address_tool, outreach_sms_message_tool, data_visualization_tool]
     conversational_agent = initialize_agent(
         agent="conversational-react-description",
         tools=tools,
         llm=llm,
         verbose=True,
         max_iterations=10,
+        return_intermediate_steps=True,
         memory=st.session_state.memory,
         handle_parsing_errors=True
     )
-    response = conversational_agent.invoke(question)
+
+    #print(conversational_agent.agent.llm_chain.prompt.template)
+    #response = conversational_agent.invoke(question)
+    response = conversational_agent({"input": question})
     return response 
 
+client = OpenAI()
+
+
+
 def main():
+    st.set_option('deprecation.showPyplotGlobalUse', False)
     img = Image.open('assist_logo.jpg')
-    #user_avatar = Image.open('renga_profile.jpg')
-    #ai_avatar = Image.open('Designer.png')
+   # user_avatar = Image.open('renga_profile.jpg')
+   # ai_avatar = Image.open('Designer.png')
     load_dotenv()  # load env parameters
     st.set_page_config(page_title="Assist", page_icon=img)
     st.write(css, unsafe_allow_html=True)
@@ -88,43 +107,47 @@ def main():
         st.write("c. Send Text to Customers")
         st.write("d. Search websites and look up Product prices & other info")
         st.write("e. Generate charts for greater visualization")
-
     if "chat_history" not in st.session_state:
         st.session_state.memory = ConversationBufferMemory(memory_key="chat_history")
+        st.session_state.chat_history = []
     # ini chat history
     if "messages" not in st.session_state:
         st.session_state.messages = []
+       
+    for message in st.session_state.chat_history:
+        if isinstance(message, HumanMessage):
+            with st.chat_message("Human"):
+                st.markdown(message.content)
+        else:
+            with st.chat_message("AI"):
+                st.markdown(message.content)
 
     user_question = st.chat_input("Type your question")
-    if user_question:
-            
-            assistant_response = conversation_agent(user_question)
-            st.session_state.chat_history = assistant_response['chat_history']
-            chistory = assistant_response["chat_history"]
-            # Process the chat history 
-            messages = chistory.split("Human: ")
-            l = len(messages)       
-            # Step 3: Print the output
-            for i in range(1, l):
-                response = messages[i].strip()
-                response = response.split("AI:")
-            # Print the Human message from history
-                with st.chat_message("human"):
-                    st.markdown(response[0])
-            # Print the AI message from history
-                with st.chat_message("ai"):
-                    st.markdown(response[1])
-            #Print the last question from user
-            with st.chat_message("human"):
+
+    if user_question is not None:
+        #Print the last question from user
+          with st.chat_message("Human"):
                st.markdown(user_question)
-            #st.write(user_template.replace("{{MSG}}", user_question), unsafe_allow_html=True)
+        
+          with st.spinner("Processing"):  
+            st.session_state.chat_history.append(HumanMessage(user_question))
+
             # Print the last answer from user
-            assistant_response_output = assistant_response["output"]
-            with st.chat_message("ai"):
-                st.write(assistant_response_output)
-            #st.write(bot_template.replace("{{MSG}}", assistant_response_output), unsafe_allow_html=True)
+            
+            with st.chat_message("AI"):
+                assistant_response = conversation_agent(user_question)
+
+                assistant_response_output = assistant_response["output"]
+
+                st.markdown(assistant_response_output)
+              
+            st.session_state.chat_history.append(AIMessage(assistant_response_output))
+                 
 
 
 if __name__== '__main__':
     main()
 
+#if st.button:
+#    question = "Top 3 sales in store 4057"
+#    output = data_visualization(question)