Update app.py

app.py

@@ -1,4 +1,6 @@
 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
@@ -34,14 +36,16 @@ def handle_message(message: str, mode: str):
     # Check if message is empty
     if not message.strip():
         return "Enter a valid message."
-    if mode == "Chat":
-        return chat_mode(message)
+    if mode == "Chat-Message":
+        return chat_message(message)
     elif mode == "Web-Search":
         return web_search(message)
+    elif mode == "Chart-Generator":
+        return chart_generator(message)
     else:
         return "Select a valid mode."
 
-def chat_mode(message: str):
+def chat_message(message: str):
     global messages
     prompt = HumanMessage(content=message)
     messages.append(prompt)
@@ -49,7 +53,7 @@ def chat_mode(message: str):
     messages.append(response.content)
     if len(messages) >= 6:
         messages = messages[-6:]
-    return f"You: {message}\n\nIT-Assistant: {response.content}"
+    return f"IT-Assistant: {response.content}"
 
 def web_search(message: str):
     global messages
@@ -72,11 +76,41 @@ def web_search(message: str):
     messages.append(response.content)
     if len(messages) >= 6:
         messages = messages[-6:]
-    return f"You: {message}\n\nIT-Assistant: {response.content}"
+    return f"IT-Assistant: {response.content}"
 
+def chart_generator(message: str):
+    # Construct the full chart URL
+    chart_url = f"https://quickchart.io/natural/{message}"
+
+    # Send a GET request to the chart URL
+    response = requests.get(chart_url)
+
+    # Check if the request was successful (status code 200)
+    if response.status_code == 200:
+        # Add a description to the message
+        message_with_description = f"Describe and analyse the content of this chart: {message}"
+
+        # Create a HumanMessage object with the message including the description
+        prompt = HumanMessage(content=message_with_description)
+        messages.append(prompt)
+
+        # Invoke the chat model with the updated messages
+        response = chat_model.invoke(messages)
+        messages.append(response.content)
+        if len(messages) >= 6:
+            messages = messages[-6:]
+
+        # Display the chart image with a smaller size
+        display(Image(url=chart_url, width=500, height=300))
+
+        # Return the message and response
+        return f"IT-Assistant: {response.content}"
+    else:
+        return f"Can't generate this image. Please provide valid chart details."
+        
 demo = gr.Interface(
     fn=handle_message,
-    inputs=["text", gr.Radio(["Chat", "Web-Search"], label="mode", info="Choose a mode and enter your message, then click submit to interact.")],
+    inputs=["text", gr.Radio(["Chat", "Web-Search", "Chart-Generator"], label="mode", info="Choose a mode and enter your message, then click submit to interact.")],
     outputs="text",
     title="IT Assistant")
 demo.launch()