Files
agent/examples/stream_modes/graph_visualize.py
2025-08-30 23:28:22 +08:00

68 lines
1.5 KiB
Python

"""
Demonstrate compiling a LangGraph and visualizing its structure.
Requires no external viz packages for ASCII. To export PNG, install:
uv pip install -e '.[viz]'
and ensure Graphviz is present on your system.
"""
from typing import TypedDict, Annotated
import operator
from langgraph.graph import StateGraph, START, END
class S(TypedDict):
a: int
b: int
c: int
def n1(s: S) -> S:
return {"a": s.get("a", 0) + 1, "b": s.get("b", 0), "c": s.get("c", 0)}
def n2(s: S) -> S:
return {"a": s.get("a", 0), "b": s.get("b", 0) + 2, "c": s.get("c", 0)}
def n3(s: S) -> S:
return {"a": s.get("a", 0), "b": s.get("b", 0), "c": s.get("c", 0) + 3}
def build_app():
g = StateGraph(S)
g.add_node("n1", n1)
g.add_node("n2", n2)
g.add_node("n3", n3)
g.add_edge(START, "n1")
g.add_edge("n1", "n2")
g.add_edge("n2", "n3")
g.add_edge("n3", END)
return g.compile()
def main():
app = build_app()
graph = app.get_graph()
# Try ASCII render if available
draw_ascii = getattr(graph, "draw_ascii", None)
if callable(draw_ascii):
print("=== Graph ASCII ===")
print(draw_ascii())
# Try PNG export if available and pydot/graphviz installed
draw_png = getattr(graph, "draw_png", None)
if callable(draw_png):
try:
out = "graph_demo.png"
draw_png(out)
print(f"Saved PNG: {out}")
except Exception as e:
print(f"PNG export unavailable: {e}")
if __name__ == "__main__":
main()