Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import requests
-from IPython.display import Image, display
 from langchain.embeddings import SentenceTransformerEmbeddings
 from langchain.vectorstores import FAISS
 from langchain_community.chat_models.huggingface import ChatHuggingFace
@@ -32,17 +31,21 @@ messages = [
 ]
 
 def handle_message(message: str, mode: str):
+    result_text, result_image = "", None
     
     if not message.strip():
-        return "Enter a valid message."
+        return "Enter a valid message.", None
+    
     if mode == "Chat-Message":
-        return chat_message(message)
+        result_text = chat_message(message)
     elif mode == "Web-Search":
-        return web_search(message)
+        result_text = web_search(message)
     elif mode == "Chart-Generator":
-        return chart_generator(message)
+        result_text, result_image = chart_generator(message)
     else:
-        return "Select a valid mode."
+        result_text = "Select a valid mode."
+    
+    return result_text, result_image
 
 def chat_message(message: str):
     global messages
@@ -106,16 +109,14 @@ def chart_generator(message: str):
         if len(messages) >= 6:
             messages = messages[-6:]
 
-        display(Image(url=chart_url, width=500, height=300))
-
-        return f"IT-Assistant: {response.content}"
+        return f"IT-Assistant: {response.content}", chart_url
     else:
-        return f"Can't generate this image. Please provide valid chart details."
+        return f"Can't generate this image. Please provide valid chart details.", None
         
 demo = gr.Interface(
     fn=handle_message,
     inputs=["text", gr.Radio(["Chat-Message", "Web-Search", "Chart-Generator"], label="mode", info="Choose a mode and enter your message, then click submit to interact.")],
-    outputs="text",
+    outputs=[gr.outputs.Textbox(label="Response"), gr.outputs.Image(label="Chart", type="auto")],
     title="IT Assistant")
 
 demo.launch()