Update app.py

app.py

@@ -50,10 +50,8 @@ image_transformation = load_tool("Chris4K/image-transformation")
 latent_upscaler_tool = load_tool("Chris4K/latent-upscaler-tool")  
 
 tools = [random_character_tool, text_generation_tool, sentiment_tool, token_counter_tool, most_downloaded_model, word_counter_tool, sentence_counter_tool, emojify_text_tool , namedEntityRecognitionTool, sourcecode_retriever_tool, text_to_image, text_to_video, image_transformation, latent_upscaler_tool ]
-
-
-
-# Define the custom HfAgent class
+ 
+# Define the custom HfAgent class with token and input_params for e.g max_new_token
 class CustomHfAgent(Agent):
     def __init__(
         self, url_endpoint, token=os.environ['HF_token'], chat_prompt_template=None, run_prompt_template=None, additional_tools=None, input_params=None
@@ -109,52 +107,44 @@ agent = CustomHfAgent(
 # 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]
 
-# Main chat loop
-while True:
-    # Input field for the user's message
-    user_message = st.text_input("User:", key="user_input")
-
-    # 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
+# Submit button
+submit_button = st.button("Submit")
 
-    # Checkboxes for the tools to be used by the agent
-    tool_checkboxes = [st.checkbox(f"Use {tool.name} --- {tool.description} ") for tool in tools]
+# 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
 
-    # Submit button
-    submit_button = st.button("Submit")
+    response = agent.chat(user_message)
 
-    # 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()
+    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()