Update app.py

app.py

@@ -98,31 +98,16 @@ class CustomHfAgent(Agent):
                 return result[: -len(stop_seq)]
         return result
 
-st.title("Hugging Face Agent and tools")
-
-
-
-# Display a welcome message
-with st.chat_message("assistant"):
-    st.markdown("Hello there! How can I assist you today?")
- 
-# Input field for the user's message
-user_message = st.text_input("User:", key="user_input")
- 
-# Checkboxes for the tools to be used by the agent
-tool_checkboxes = [st.checkbox(f"Use {tool.name} --- {tool.description} ") for tool in tools]
-
-# Submit button
-submit_button = st.button("Submit")
-
+#################
+#################
 # Define the callback function to handle the form submission
 def handle_submission():
     selected_tools = [tools[idx] for idx, checkbox in enumerate(tool_checkboxes) if checkbox]
     print(selected_tools)
     # Initialize the agent
     agent = CustomHfAgent(
-        #url_endpoint="https://api-inference.huggingface.co/models/bigcode/starcoder", # mistralai/Mixtral-8x7B-v0.1
-        url_endpoint="https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-v0.1", # 
+        #url_endpoint="https://api-inference.huggingface.co/models/bigcode/starcoder", # mistralai/Mixtral-8x7B-v0.1 # bigscience/bloom
+        url_endpoint="https://api-inference.huggingface.co/models/bigscience/bloom", # 
         token=os.environ['HF_token'],
         additional_tools=selected_tools,
         input_params={"max_new_tokens": 192},
@@ -133,6 +118,41 @@ def handle_submission():
 
     print("Agent Response\n {}".format(response))
     
+    return response
+
+# Add the callback function to the Streamlit app
+#if submit_button:
+#    handle_submission()
+#################
+
+###### 
+        
+st.title("Hugging Face Agent and tools")
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# 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"])
+        
+# Display a welcome message
+with st.chat_message("assistant"):
+    st.markdown("Hello there! How can I assist you today?")
+ 
+# Input field for the user's message
+#user_message = st.chat_input("Enter message")
+
+# React to user input
+if user_message := st.chat_input("Enter message")
+    # Display user message in chat message container
+    st.chat_message("user").markdown(prompt)
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    response = handle_submission
     # Display the agent's response
     with st.chat_message("assistant"):
         if response is None:
@@ -154,7 +174,20 @@ def handle_submission():
             st.markdown(response)
         else:
             st.warning("Unrecognized response type. Please try again.")
+            
+    
+    # Display assistant response in chat message container
+    #with st.chat_message("assistant"):
+    #    st.markdown(response)
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})
+
+
+# Submit button
+#submit_button = st.button("Submit")
+
+# Checkboxes for the tools to be used by the agent
+tool_checkboxes = [st.checkbox(f"{tool.name} --- {tool.description} ") for tool in tools]
+
+
 
-# Add the callback function to the Streamlit app
-if submit_button:
-    handle_submission()