Update app.py

app.py

@@ -8,9 +8,8 @@ from langgraph.graph import StateGraph, END
 from typing import TypedDict, Annotated, Sequence
 from langchain_core.messages import BaseMessage
 import operator
-import pygraphviz as pgv
-from PIL import Image
-import tempfile
+import networkx as nx
+import matplotlib.pyplot as plt
 
 # Set API keys and validate credentials
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
@@ -89,13 +88,20 @@ graph.add_edge("tool", END)
 graph.set_entry_point("agent")
 compiled_app = graph.compile()
 
-# Function to render graph
-def render_graph(graph):
-    dot_string = graph.get_graph().to_string()
-    G = pgv.AGraph(string=dot_string)
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
-        G.draw(temp_file.name, prog="dot", format="png")
-        return Image.open(temp_file.name)
+# Function to render graph with NetworkX
+def render_graph_nx(graph):
+    G = nx.DiGraph()
+    G.add_edge("agent", "tool", label="invoke tool")
+    G.add_edge("agent", "end", label="end condition")
+    G.add_edge("tool", "end", label="finish")
+
+    pos = nx.spring_layout(G, seed=42)
+    plt.figure(figsize=(8, 6))
+    nx.draw(G, pos, with_labels=True, node_color="lightblue", node_size=3000, font_size=10, font_weight="bold")
+    edge_labels = nx.get_edge_attributes(G, "label")
+    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=9)
+    plt.title("Workflow Graph")
+    st.pyplot(plt)
 
 # Streamlit UI
 st.title("LLM Tool Workflow Demo")
@@ -124,5 +130,4 @@ if st.button("Run Workflow"):
 
     # Display Graph
     st.subheader("Workflow Graph")
-    graph_image = render_graph(graph)
-    st.image(graph_image, caption="Workflow Graph", use_column_width=True)
+    render_graph_nx(graph)