# %%
import json
import os
import http.server
import socketserver
import threading
import webbrowser
import time
import socket
import importlib
import sys
from pathlib import Path
from typing import Any, Optional, Tuple, Union
from pocketflow import Flow
from async_flow import order_pipeline
def build_mermaid(start):
ids, visited, lines = {}, set(), ["graph LR"]
ctr = 1
def get_id(n):
nonlocal ctr
return (
ids[n] if n in ids else (ids.setdefault(n, f"N{ctr}"), (ctr := ctr + 1))[0]
)
def link(a, b, action=None):
if action:
lines.append(f" {a} -->|{action}| {b}")
else:
lines.append(f" {a} --> {b}")
def walk(node, parent=None, action=None):
if node in visited:
return parent and link(parent, get_id(node), action)
visited.add(node)
if isinstance(node, Flow):
node.start_node and parent and link(parent, get_id(node.start_node), action)
lines.append(
f"\n subgraph sub_flow_{get_id(node)}[{type(node).__name__}]"
)
node.start_node and walk(node.start_node)
for act, nxt in node.successors.items():
node.start_node and walk(nxt, get_id(node.start_node), act) or (
parent and link(parent, get_id(nxt), action)
) or walk(nxt, None, act)
lines.append(" end\n")
else:
lines.append(f" {(nid := get_id(node))}['{type(node).__name__}']")
parent and link(parent, nid, action)
[walk(nxt, nid, act) for act, nxt in node.successors.items()]
walk(start)
return "\n".join(lines)
def flow_to_json(start):
"""Convert a flow to JSON format suitable for D3.js visualization.
This function walks through the flow graph and builds a structure with:
- nodes: All non-Flow nodes with their group memberships
- links: Connections between nodes within the same group
- group_links: Connections between different groups (for inter-flow connections)
- flows: Flow information for group labeling
Returns:
dict: A JSON-serializable dictionary with 'nodes' and 'links' arrays.
"""
nodes = []
links = []
group_links = [] # For connections between groups (Flow to Flow)
ids = {}
node_types = {}
flow_nodes = {} # Keep track of flow nodes
ctr = 1
visited = set()
def get_id(n):
nonlocal ctr
if n not in ids:
ids[n] = ctr
node_types[ctr] = type(n).__name__
if isinstance(n, Flow):
flow_nodes[ctr] = n # Store flow reference
ctr += 1
return ids[n]
def walk(node, parent=None, group=None, parent_group=None, action=None):
"""Recursively walk the flow graph to build the visualization data.
Args:
node: Current node being processed
parent: ID of the parent node that connects to this node
group: Group (Flow) ID this node belongs to
parent_group: Group ID of the parent node
action: Action label on the edge from parent to this node
"""
node_id = get_id(node)
if (node_id, action) in visited:
return
visited.add((node_id, action))
# Add node if not already in nodes list and not a Flow
if not any(n["id"] == node_id for n in nodes) and not isinstance(node, Flow):
node_data = {
"id": node_id,
"name": node_types[node_id],
"group": group or 0, # Default group
}
nodes.append(node_data)
# Add link from parent if exists
if parent and not isinstance(node, Flow):
links.append(
{"source": parent, "target": node_id, "action": action or "default"}
)
# Process different types of nodes
if isinstance(node, Flow):
# This is a Flow node - it becomes a group container
flow_group = node_id # Use flow's ID as group for contained nodes
# Add a group-to-group link if this flow has a parent group
# This creates connections between nested flows
if parent_group is not None and parent_group != flow_group:
# Check if this link already exists
if not any(
l["source"] == parent_group and l["target"] == flow_group
for l in group_links
):
group_links.append(
{
"source": parent_group,
"target": flow_group,
"action": action or "default",
}
)
if node.start_node:
# Process the start node of this flow
walk(node.start_node, parent, flow_group, parent_group, action)
# Process successors of the flow's start node
for next_action, nxt in node.successors.items():
walk(
nxt,
get_id(node.start_node),
flow_group,
parent_group,
next_action,
)
else:
# Process successors for regular nodes
for next_action, nxt in node.successors.items():
if isinstance(nxt, Flow):
# This node connects to a flow - track the group relationship
flow_group_id = get_id(nxt)
walk(nxt, node_id, None, group, next_action)
else:
# Regular node-to-node connection
walk(nxt, node_id, group, parent_group, next_action)
# Start the traversal
walk(start)
# Post-processing: Generate group links based on node connections between different groups
# This ensures that when nodes in different groups are connected, we show a group-to-group
# link rather than a direct node-to-node link
node_groups = {n["id"]: n["group"] for n in nodes}
filtered_links = []
for link in links:
source_id = link["source"]
target_id = link["target"]
source_group = node_groups.get(source_id, 0)
target_group = node_groups.get(target_id, 0)
# If source and target are in different groups and both groups are valid
if source_group != target_group and source_group > 0 and target_group > 0:
# Add to group links if not already there
# This creates the dashed lines connecting group boxes
if not any(
gl["source"] == source_group and gl["target"] == target_group
for gl in group_links
):
group_links.append(
{
"source": source_group,
"target": target_group,
"action": link["action"],
}
)
# Skip adding this link to filtered_links - we don't want direct node connections across groups
else:
# Keep links within the same group
filtered_links.append(link)
return {
"nodes": nodes,
"links": filtered_links, # Use filtered links instead of all links
"group_links": group_links,
"flows": {str(k): v.__class__.__name__ for k, v in flow_nodes.items()},
}
def create_d3_visualization(
json_data,
output_dir="./viz",
filename="flow_viz",
html_title="PocketFlow Visualization",
):
"""Create a D3.js visualization from JSON data.
Args:
json_data: The JSON data for the visualization
output_dir: Directory to save the files
filename: Base filename (without extension)
html_title: Title for the HTML page
Returns:
str: Path to the HTML file
"""
# Create output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)
# Save JSON data to file
json_path = os.path.join(output_dir, f"{filename}.json")
with open(json_path, "w") as f:
json.dump(json_data, f, indent=2)
# Create HTML file with D3.js visualization
html_content = r"""
TITLE_PLACEHOLDER
"""
# Replace the placeholders with the actual values
html_content = html_content.replace("FILENAME_PLACEHOLDER", filename)
html_content = html_content.replace("TITLE_PLACEHOLDER", html_title)
# Write HTML to file
html_path = os.path.join(output_dir, f"{filename}.html")
with open(html_path, "w") as f:
f.write(html_content)
print(f"Visualization created at {html_path}")
return html_path
def find_free_port():
"""Find a free port on localhost."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
return s.getsockname()[1]
def start_http_server(directory, port=None):
"""Start an HTTP server in the given directory.
Args:
directory: Directory to serve files from
port: Port to use (finds a free port if None)
Returns:
tuple: (server_thread, port)
"""
if port is None:
port = find_free_port()
# Get the absolute path of the directory
directory = str(Path(directory).absolute())
# Change to the directory to serve files
os.chdir(directory)
# Create HTTP server
handler = http.server.SimpleHTTPRequestHandler
httpd = socketserver.TCPServer(("", port), handler)
# Start server in a separate thread
server_thread = threading.Thread(target=httpd.serve_forever)
server_thread.daemon = (
True # This makes the thread exit when the main program exits
)
server_thread.start()
print(f"Server started at http://localhost:{port}")
return server_thread, port
def serve_and_open_visualization(html_path, auto_open=True):
"""Serve the HTML file and open it in a browser.
Args:
html_path: Path to the HTML file
auto_open: Whether to automatically open the browser
Returns:
tuple: (server_thread, url)
"""
# Get the directory and filename
directory = os.path.dirname(os.path.abspath(html_path))
filename = os.path.basename(html_path)
# Start the server
server_thread, port = start_http_server(directory)
# Build the URL
url = f"http://localhost:{port}/{filename}"
# Open the URL in a browser
if auto_open:
print(f"Opening {url} in your browser...")
webbrowser.open(url)
else:
print(f"Visualization available at {url}")
return server_thread, url
def visualize_flow(
flow: Flow,
flow_name: str,
serve: bool = True,
auto_open: bool = True,
output_dir: str = "./viz",
html_title: Optional[str] = None,
) -> Union[str, Tuple[str, Any, str]]:
"""Helper function to visualize a flow with both mermaid and D3.js
Args:
flow: Flow object to visualize
flow_name: Name of the flow (used for filename and display)
serve: Whether to start a server for the visualization
auto_open: Whether to automatically open in browser
output_dir: Directory to save visualization files
html_title: Custom title for the HTML page (defaults to flow_name if None)
Returns:
str or tuple: Path to HTML file, or (path, server_thread, url) if serve=True
"""
print(f"\n--- {flow_name} Mermaid Diagram ---")
print(build_mermaid(start=flow))
print(f"\n--- {flow_name} D3.js Visualization ---")
json_data = flow_to_json(flow)
# Create the visualization
output_filename = f"{flow_name.lower().replace(' ', '_')}"
# Use flow_name as the HTML title if not specified
if html_title is None:
html_title = f"PocketFlow: {flow_name}"
html_path = create_d3_visualization(
json_data,
output_dir=output_dir,
filename=output_filename,
html_title=html_title,
)
# Serve and open if requested
if serve:
server_thread, url = serve_and_open_visualization(html_path, auto_open)
return html_path, server_thread, url
return html_path
def load_flow_from_module(module_path: str, flow_variable: str) -> Flow:
"""Dynamically load a flow from a module.
Args:
module_path: Path to the module (e.g., 'my_package.my_module')
flow_variable: Name of the flow variable in the module
Returns:
Flow: The loaded flow object
"""
try:
module = importlib.import_module(module_path)
return getattr(module, flow_variable)
except (ImportError, AttributeError) as e:
print(f"Error loading flow: {e}")
sys.exit(1)
# Example usage
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Visualize a PocketFlow flow")
parser.add_argument(
"--module", default="async_loop_flow", help="Module containing the flow"
)
parser.add_argument(
"--flow", default="order_pipeline", help="Flow variable name in the module"
)
parser.add_argument(
"--name", default="Flow Visualization", help="Name for the visualization"
)
parser.add_argument(
"--output-dir", default="./viz", help="Directory to save visualization files"
)
parser.add_argument("--no-serve", action="store_true", help="Don't start a server")
parser.add_argument(
"--no-open", action="store_true", help="Don't open browser automatically"
)
parser.add_argument("--title", help="Custom HTML title")
args = parser.parse_args()
# Load flow from the specified module
flow_obj = load_flow_from_module(args.module, args.flow)
# Visualize the flow
visualize_flow(
flow=flow_obj,
flow_name=args.name,
serve=not args.no_serve,
auto_open=not args.no_open,
output_dir=args.output_dir,
html_title=args.title,
)
# Keep server running if serving
if not args.no_serve:
try:
print("\nServer is running. Press Ctrl+C to stop...")
while True:
time.sleep(1)
except KeyboardInterrupt:
print("\nShutting down...")