Update app.py

app.py

@@ -1,6 +1,7 @@
 from langgraph.graph import StateGraph
 from typing import TypedDict, Annotated
 from langgraph.graph.message import add_messages
+from langchain_core.runnables.graph import MermaidDrawMethod
 
 class State(TypedDict):
     messages: Annotated[list[str], add_messages]
@@ -42,6 +43,7 @@ workflow.add_edge("collect", "process")
 
 # Set entry and finish points
 workflow.set_entry_point("collect")
+workflow.set_finish_point("process")
 
 app = workflow.compile()
 
@@ -49,4 +51,10 @@ app = workflow.compile()
 print("\nStarting workflow...")
 initial_state = State(messages=["Starting"], current_step="collect")
 final_state = app.invoke(initial_state)
-print(f"\nFinal messages: {final_state['messages']}")
+print(f"\nFinal messages: {final_state['messages']}")
+
+# Save the graph visualization as PNG
+png_data = app.get_graph().draw_mermaid_png(draw_method=MermaidDrawMethod.API)
+with open("workflow_graph.png", "wb") as f:
+    f.write(png_data)
+print("\nGraph visualization saved as 'workflow_graph.png'")