Update app.py

app.py

@@ -1,69 +1,59 @@
-import streamlit as st
+from streamlit import container, text_input, spinner
 from streamlit_chat import message
 from src.langchain_agent import init, init_agent
 
-def main():
-    
-    #initialise agent and streamlit page
-    init()
-    agent_executor = init_agent()
-    
-    # initialize message history
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-
-    # store agent in memory
-    if "clarina" not in st.session_state:
-        st.session_state.clarina = agent_executor
-
-    # store generated responses in memory
-    if 'generated' not in st.session_state:
-        st.session_state.generated = []
-    
-    # define function to generate response
-    def generate_response(user_input):
-        # handle user input
-        if user_input:
-            # save user input
-            st.session_state.messages.append(user_input) 
-            # get response from agent
-            with st.spinner("Thinking..."):
-                response = st.session_state.clarina.reverse_prompt_engineer(user_input)
-            # save response
-            st.session_state.messages.append(response) 
-            st.session_state.generated.append(response)    
+# Initialize agent and streamlit page
+init()
+agent_executor = init_agent()
+
+# Initialize session state variables
+st.session_state.update({
+    'messages': set(),
+    'clarina': agent_executor,
+    'generated': [],
+    'temp': ""
+})
+
+def generate_response(user_input):
+    # Handle user input
+    if user_input:
+        # Save user input
+        st.session_state.messages.add(user_input)
+        # Get response from agent
+        with spinner("Thinking..."):
+            response = st.session_state.clarina.reverse_prompt_engineer(user_input)
+        # Save response
+        st.session_state.messages.add(response)
+        st.session_state.generated.append(response)
 
-    # container for chat history
-    response_container = st.container()
-    # container for text box
-    container = st.container()
+def main():
+    # Container for chat history
+    response_container = container()
+    # Container for text box
+    container = container()
 
     with container:
-
-        # initialize session state to clear input text box after user enters input
-        if "temp" not in st.session_state:
-            st.session_state.temp = ""
-
         def clear_text():
-            """callback function to clear input text box"""
+            """Callback function to clear input text box"""
             st.session_state.temp = st.session_state.user_input
             st.session_state.user_input = ""
 
-        st.text_input("user input",key="user_input",placeholder = "Enter your code here", label_visibility="hidden",on_change=clear_text) # get user input
-        generate_response(st.session_state.temp) # generate response
+        text_input("user input", key="user_input", placeholder="Enter your code here", label_visibility="hidden", on_change=clear_text)  # Get user input
+        generate_response(st.session_state.user_input)  # Generate response
 
-    # display message history
+    # Display message history
     if st.session_state.generated:
-        
         with response_container:
-            messages = st.session_state.get('messages', []) 
-            for i, msg in enumerate(messages):
-                if i % 2 == 0:
-                    # display user input
-                    message(msg, is_user=True, key=str(i) + '_user')
-                else:
-                    # display response
-                    message(msg, is_user=False, key=str(i) + '_ai')    
+            messages = st.session_state.messages
+            i = 0
+            while i < len(messages):
+                # Display user input
+                message(list(messages)[i], is_user=True, key=str(i) + '_user')
+                i += 1
+                if i < len(messages):
+                    # Display response
+                    message(list(messages)[i], is_user=False, key=str(i) + '_ai')
+                i += 1
 
 if __name__ == '__main__':
-    main()
+    main()