Update app.py

app.py

@@ -9,7 +9,7 @@ from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import tools_condition
 from langgraph.prebuilt import ToolNode
 from langchain_core.messages import HumanMessage, SystemMessage
-from IPython.display import Image, display
+import tempfile
 
 # ------------------- Environment Variable Setup -------------------
 # Fetch API keys from environment variables
@@ -62,17 +62,17 @@ app_graph.add_conditional_edges("assistant", tools_condition)
 app_graph.add_edge("tools", "assistant")
 react_graph = app_graph.compile()
 
-# Display graph visualization
-st.header("LangGraph Workflow Visualization")
-try:
-    image_path = react_graph.get_graph(xray=True).draw_mermaid_png()
-    display(Image(image_path))
-    st.image(image_path, caption="LangGraph Workflow Visualization")
-except Exception as e:
-    st.error(f"Failed to display graph visualization: {e}")
+# Save graph visualization as an image
+with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
+    graph_image_path = tmpfile.name
+    react_graph.get_graph(xray=True).draw_mermaid_png(graph_image_path)
 
 # ------------------- Streamlit Interface -------------------
-#st.title("ReAct Agent with Arithmetic and Search")
+st.title("ReAct Agent with Arithmetic and Search")
+
+# Display the workflow graph
+st.header("LangGraph Workflow Visualization")
+st.image(graph_image_path, caption="LangGraph Workflow Visualization")
 
 # Prompt user for inputs
 user_question = st.text_area("Enter your question:",