Update app.py

app.py

@@ -81,84 +81,63 @@ class CustomHfAgent(Agent):
 
 st.title("Hugging Face Agent and tools")
 
-# Input field for the user's message
-message = st.text_input("Enter your message:", "")
-
-# 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")
-
 # Initialize the agent
-
 agent = CustomHfAgent(
     url_endpoint="https://api-inference.huggingface.co/models/bigcode/starcoder",
     token=os.environ['HF_token'],
     additional_tools=[],
-    input_params={"max_new_tokens": 192},  # Set the desired value
+    input_params={"max_new_tokens": 192},
 )
 
+# Display a welcome message
+with st.chat_message("assistant"):
+    st.markdown("Hello there! How can I assist you today?")
 
-# Define the callback function to handle the form submission
-def handle_submission():
-    # Get the user's message and the selected tools
-    #message = st.text_input("Enter your message:", "")
-    
-    #selected_tools = []
-    #selected_tools = [tool for idx, tool in enumerate(tools) if tool_checkboxes[idx]]
-    selected_tools = [tools[idx] for idx, checkbox in enumerate(tool_checkboxes) if checkbox]
-    #selected_tools = [tool for idx, checkbox in enumerate(tool_checkboxes) if checkbox]
+# Main chat loop
+while True:
+    # Input field for the user's message
+    user_message = st.text_input("User:", "")
 
-    print(selected_tools)
+    # Check if the user wants to end the conversation
+    if user_message.lower() in ["bye", "goodbye", "exit"]:
+        with st.chat_message("assistant"):
+            st.markdown("Goodbye! Have a great day.")
+        break
 
-    agent.tools = selected_tools
+    # Checkboxes for the tools to be used by the agent
+    tool_checkboxes = [st.checkbox(f"Use {tool.name} --- {tool.description} ") for tool in tools]
 
-    
-    # Run the agent with the user's message and selected tools
-    #response = agent.run(message)
-    response = agent.chat(message)
+    # Submit button
+    submit_button = st.button("Submit")
 
-    print("Agent Response\n {}".format(response))
-    
-    # Display the agent's response
-    if response is None:
-        st.warning("The agent's response is None.  Please try again. For Example: Generate an image of a boat in the water")
-    elif isinstance(response, Image.Image):
-        # Display the image response
-        st.image(response)
-    elif "audio" in response:
-        # Handle audio response (replace with your audio rendering code)
-        audio_data = base64.b64decode(response.split(",")[1])
-        audio = AudioSegment.from_file(io.BytesIO(audio_data))
-        st.audio(audio)
-    elif isinstance(response, AudioSegment):
-        # Handle audio response (replace with your audio rendering code)
-        st.audio(response)
-    elif isinstance(response, str):
-        # Display the text response
-        st.write(response)
-    elif "text" in response:
-        # Display the text response
-        st.write(response)
-    else:
-        # Handle unrecognized response type
-        st.warning("Unrecognized response type. Please try again. For Example: Generate an image of a boat in the water")
-
-
-
-# Add the callback function to the Streamlit app
-submit_button = st.button("Submit", on_click=handle_submission)
-
-# Define a callback function to handle the button click
-def ask_again():
-    # Reset the message input field
-    message_input.value = ""
-
-    # Run the agent again with an empty message
-    agent.run("")
-
-# Add the callback function to the button
-ask_again = st.button("Ask again", on_click=ask_again)
+    # 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]
+        agent.tools = selected_tools
+
+        response = agent.chat(user_message)
+
+        print("Agent Response\n {}".format(response))
+        
+        # Display the agent's response
+        with st.chat_message("assistant"):
+            if response is None:
+                st.warning("The agent's response is None. Please try again.")
+            elif isinstance(response, Image.Image):
+                st.image(response)
+            elif "audio" in response:
+                audio_data = base64.b64decode(response.split(",")[1])
+                audio = AudioSegment.from_file(io.BytesIO(audio_data))
+                st.audio(audio)
+            elif isinstance(response, AudioSegment):
+                st.audio(response)
+            elif isinstance(response, str):
+                st.markdown(response)
+            elif "text" in response:
+                st.markdown(response)
+            else:
+                st.warning("Unrecognized response type. Please try again.")
+
+    # Add the callback function to the Streamlit app
+    if submit_button:
+        handle_submission()