Update app.py

app.py

@@ -103,14 +103,23 @@ userMessage2 = st.sidebar.text_area("Define User Message", value="This is the co
 st.sidebar.caption(f"Session ID: {genuuid()}")
 # Main chat interface
 st.header("Chat with the Agents")
-user_id = st.text_input("User ID:", key="user_id")
-user_input = st.text_input("Write your message here:", key="user_input")
 
-if 'history' not in st.session_state:
-    st.session_state.history = []
+# Initialize chat history in session state
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
-if st.button("Send"):
-    # Placeholder for processing the input and generating a response
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Collect user input
+if user_input := st.chat_input("Write your message here:"):
+    # Add user message to the chat history
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    st.chat_message("user").markdown(user_input)
+
+    # Prepare data for API call
     data = ChatRequestClient(
         user_id=user_id,
         user_input=user_input,
@@ -126,19 +135,23 @@ if st.button("Send"):
         tokens2=tokens2,
         temperature2=temp2
     )
+
+    # Call the API
     response = call_chat_api(data)
-    
-    
-    st.markdown(f"##### Time take: {format_elapsed_time(response['elapsed_time'])}")
-    st.markdown(f"##### Question Count : {response['count']} of {numberOfQuestions}")
-    
-    
-    # {"count":count, "user_id":user_id,"time_stamp":time_stamp, "elapsed_time":elapsed_time, "content":content}
 
-    st.session_state.history.append("You: " + user_input)
-                
-    st.session_state.history.append("Agent: " + response['content'])  # Using 'response' after it's defined
-    for message in st.session_state.history:
-        st.write(message)
+    # Process the API response
+    agent_message = response.get("content", "No response received from the agent.")
+    elapsed_time = response.get("elapsed_time", 0)
+    count = response.get("count", 0)
+
+    # Add agent response to the chat history
+    st.session_state.messages.append({"role": "assistant", "content": agent_message})
+    with st.chat_message("assistant"):
+        st.markdown(agent_message)
+
+    # Display additional metadata
+    st.markdown(f"##### Time taken: {format_elapsed_time(elapsed_time)} seconds")
+    st.markdown(f"##### Question Count: {count} of {numberOfQuestions}")
+