Update app.py

app.py

@@ -1,19 +1,49 @@
 import streamlit as st
 import os
-import requests
 import base64
 import io
-import time
 from PIL import Image
 from pydub import AudioSegment
 import IPython
 import soundfile as sf
+import requests
+import pandas as pd  # If you're working with DataFrames
+import matplotlib.figure  # If you're using matplotlib figures
+
+# For Altair charts
+import streamlit.graphics_altair
+# For Bokeh charts
+import streamlit.graphics_bokeh
+
+# For Plotly charts
+import streamlit.graphics_plotly
+
+# For Pydeck charts
+import streamlit.graphics_pydeck
+# For Vega-Lite charts
+import streamlit.graphics_vega_lite
+
+
+import time  
 from transformers import load_tool, Agent
 import torch
 
+
 class ToolLoader:
     def __init__(self, tool_names):
-        self.tools = [load_tool(tool_name) for tool_name in tool_names]
+        self.tools = self.load_tools(tool_names)
+
+    def load_tools(self, tool_names):
+        loaded_tools = []
+        for tool_name in tool_names:
+            try:
+                tool = load_tool(tool_name)
+                loaded_tools.append(tool)
+            except Exception as e:
+                print(f"Error loading tool '{tool_name}': {e}")
+                # Handle the error as needed, e.g., continue with other tools or take corrective action
+
+        return loaded_tools
 
 class CustomHfAgent(Agent):
     def __init__(self, url_endpoint, token, chat_prompt_template=None, run_prompt_template=None, additional_tools=None, input_params=None):
@@ -56,6 +86,9 @@ def load_tools(tool_names):
 tool_names = [
     "Chris4K/random-character-tool",
     "Chris4K/text-generation-tool",
+    "Chris4K/sentiment-tool",
+    "Chris4K/EmojifyTextTool",
+    
     # Add other tool names as needed
 ]
 
@@ -100,22 +133,43 @@ if user_message := st.chat_input("Enter message"):
 
     with st.chat_message("assistant"):
         if response is None:
-            st.warning("The agent's response is None. Please try again.")
-        elif "emojified_text" in response:
-            st.markdown(f"Emojified Text: {response['emojified_text']}")
+            st.warning("The agent's response is None. Please try again. Generate an image of a flying horse.")
         elif isinstance(response, Image.Image):
             st.image(response)
-        elif "audio" in str(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 isinstance(response, int):
             st.markdown(response)
+        elif isinstance(response, str):
+            if "emojified_text" in response:
+                st.markdown(f"{response['emojified_text']}")
+            else:
+                st.markdown(response)
+        elif isinstance(response, list):
+            for item in response:
+                st.markdown(item)  # Assuming the list contains strings
+        elif isinstance(response, pd.DataFrame):
+            st.dataframe(response)
+        elif isinstance(response, pd.Series):
+            st.table(response.iloc[0:10])
+        elif isinstance(response, dict):
+            st.json(response)
+        elif isinstance(response, streamlit.graphics_altair.AltairChart):
+            st.altair_chart(response)
+        elif isinstance(response, streamlit.graphics_bokeh.BokehChart):
+            st.bokeh_chart(response)
+        elif isinstance(response, streamlit.graphics_graphviz.GraphvizChart):
+            st.graphviz_chart(response)
+        elif isinstance(response, streamlit.graphics_plotly.PlotlyChart):
+            st.plotly_chart(response)
+        elif isinstance(response, streamlit.graphics_pydeck.PydeckChart):
+            st.pydeck_chart(response)
+        elif isinstance(response, matplotlib.figure.Figure):
+            st.pyplot(response)
+        elif isinstance(response, streamlit.graphics_vega_lite.VegaLiteChart):
+            st.vega_lite_chart(response)
         else:
-            st.warning("Unrecognized response type. Please try again.")
+            st.warning("Unrecognized response type. Please try again. e.g. Generate an image of a flying horse.")
+
             
     st.session_state.messages.append({"role": "assistant", "content": response})