Update app.py

app.py

@@ -26,7 +26,7 @@ def stream_chat_with_rag(
 ):
     # print(f"Message: {message}")
     #answer = client.predict(question=question, api_name="/run_graph")
-    answer, fig = client.predict(
+    answer, sources = client.predict(
     	query= message,
 		election_year=year,
 		api_name="/process_query"
@@ -38,16 +38,32 @@ def stream_chat_with_rag(
     print(answer)
     history.append((message, response +"\n"+ answer))
     
-    # print("top works from API:")
-    # print(fig)
+
     
-    # fig_dict = json.loads(plotly_data['plot'])
+    # Render the figure
+
     
-    # # Render the figure
-    # fig = pio.from_json(json.dumps(fig_dict))
-    # fig.show() 
     return answer
 
+def topic_plot_gener(message: str, year: str):
+    fig = client.predict(
+    	query= message,
+		election_year=year,
+		api_name="/topics_plot_genera"
+    )
+        # print("top works from API:")
+    print(fig)
+    plot_base64 = fig
+
+    plot_bytes = base64.b64decode(plot_base64['plot'].split(',')[1])
+    img = plt.imread(BytesIO(plot_bytes), format='PNG')
+    plt.figure(figsize = (12, 6), dpi = 150)
+    plt.imshow(img)
+    plt.axis('off')
+    plt.show()
+
+    return plt.gcf()
+
 
 # def predict(message, history):
 #     history_langchain_format = []
@@ -228,26 +244,26 @@ with gr.Blocks(title="Reddit Election Analysis") as demo:
                             )
             
     gr.Markdown("## Top words of the relevant Q&A")
-    # with gr.Row():
-    #     output_plot = gr.Plot(
-    #         label="Topic Distribution",
-    #         container=True,  # Ensures the plot is contained within its area
-    #         elem_classes="topic-plot"  # Add a custom class for styling
-    #     )
+    with gr.Row():
+        topic_plot = gr.Plot(
+            label="Topic Distribution",
+            container=True,  # Ensures the plot is contained within its area
+            elem_classes="topic-plot"  # Add a custom class for styling
+        )
 
     # Add custom CSS to ensure proper plot sizing
     gr.HTML("""
         <style>
-            # .topic-plot {
-            #     min-height: 600px;
-            #     width: 100%;
-            #     margin: auto;
-            # }
             .heatmap-plot {
                 min-height: 400px;
                 width: 100%;
                 margin: auto;
             }
+            .topic-plot {
+                min-width: 600px;
+                height: 100%;
+                margin: auto;
+            }
         </style>
     """)
     # topics_df = gr.Dataframe(value=df, label="Data Input")
@@ -266,12 +282,12 @@ with gr.Blocks(title="Reddit Election Analysis") as demo:
         outputs = [time_series_fig, linePlot_status_text]
     )
 
-    # # Update both outputs when submit is clicked
-    # submit_btn.click(
-    #     fn=chat_function,
-    #     inputs=[query_input, year_selector],
-    #     outputs= chatbot
-    # )
+    # Update both outputs when submit is clicked
+    topic_btn.click(
+        fn= topic_plot_gener,
+        inputs=[query_input, year_selector],
+        outputs= topic_plot
+    )
 
 
 if __name__ == "__main__":