Spaces:

dn-scribe
/

scribe-mcp

Sleeping

scribe-mcp / scribe_python_client /lineage_graph.py

fix: add demo-data and cleanup lineage_graph main

b3c3356 unverified over 1 year ago

10.4 kB

import json
import networkx as nx
# import plotly.graph_objects as go
from pyvis.network import Network
from pathlib import Path


def create_network(data, include_types, type_to_position):
    G = nx.DiGraph()
    for record in data:
        asset_type = record.get("asset_type")
        parent_type = record.get("parent_type")
        child = record.get("external_id")
        parent = record.get("parent_external_id")
        # Filter nodes based on allowed types.
        add_child = include_types is None or asset_type in include_types
        add_parent = include_types is None or parent_type in include_types

        if add_child and child:
            node_attrs = {"type": asset_type, "label": record.get("asset_name")}
            if record.get("uri"):
                node_attrs["uri"] = record.get("uri")
            G.add_node(child, **node_attrs)
        if parent and add_parent:
            node_attrs = {"type": parent_type, "label": record.get("parent_name")}
            if record.get("uri"):
                node_attrs["uri"] = record.get("uri")
            G.add_node(parent, **node_attrs)
        if child and parent and add_child and add_parent:
            G.add_edge(parent, child)

    # Assign 'layer' attribute for each node based on its type.
    for node in G.nodes():
        node_type = G.nodes[node].get('type')
        G.nodes[node]['layer'] = type_to_position.get(node_type, 1)
    return G

# def visualize_network_plotly(data, include_types, type_to_position):
#     """
#     Build and return a Plotly figure of a network graph.

#     Parameters:
#         data (list): JSON array of records.
#         include_types (set): Allowed asset types.
#         type_to_position (dict): Mapping of asset types to numeric layer positions.

#     Returns:
#         go.Figure: The Plotly figure with the network visualization.
#     """
#     # Create network graph

#     G = create_network(data, include_types, type_to_position)
#     # Compute positions using a multipartite layout (vertical alignment)
#     pos = nx.multipartite_layout(G, subset_key="layer", align="vertical")

#     # Build edge coordinates for Plotly
#     edge_x, edge_y = [], []
#     for source, target in G.edges():
#         x0, y0 = pos[source]
#         x1, y1 = pos[target]
#         edge_x.extend([x0, x1, None])
#         edge_y.extend([y0, y1, None])
#     edge_trace = go.Scatter(
#         x=edge_x, y=edge_y,
#         line=dict(width=1, color='#888'),
#         hoverinfo='none',
#         mode='lines'
#     )

#     # Build node coordinates and text lists.
#     node_x, node_y = [], []
#     node_text_trunc, node_text_hover, node_color = [], [], []
#     for node in G.nodes():
#         x, y = pos[node]
#         node_x.append(x)
#         node_y.append(y)
#         full_label = G.nodes[node].get('label', str(node))
#         truncated_label = full_label[0:30]  # truncated version for display
#         # Check for a "uri" attribute (lowercase) and create a clickable link if present.
#         uri = G.nodes[node].get('uri')
#         uri_link = f'<br><a href="{uri}" target="_blank">{uri}</a>' if uri else ""
#         node_text_trunc.append(f"{truncated_label}<br>Type: {G.nodes[node].get('type')}")
#         node_text_hover.append(f"{full_label}<br>Type: {G.nodes[node].get('type')}{uri_link}")
#         node_color.append(G.nodes[node]['layer'])

#     node_trace = go.Scatter(
#         x=node_x, y=node_y,
#         mode='markers+text',
#         textposition="bottom center",
#         hoverinfo='text',
#         marker=dict(
#             showscale=True,
#             colorscale='YlGnBu',
#             reversescale=True,
#             color=node_color,
#             size=15,
#             colorbar=dict(title='Layer')
#         ),
#         text=node_text_trunc,       # Display truncated text on the plot.
#         hovertext=node_text_hover   # Full text (with clickable URL if present) for hover.
#     )

#     # Create the Plotly figure.
#     fig = go.Figure(data=[edge_trace, node_trace],
#                     layout=go.Layout(
#                         title='<br>Integrated Network Graph Visualization',
#                         titlefont_size=16,
#                         dragmode='pan',
#                         showlegend=False,
#                         hovermode='closest',
#                         margin=dict(b=20, l=5, r=5, t=40),
#                         width=1000,
#                         height=800,
#                         xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
#                         yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)
#                     ))
#     return fig

def filter_network(data, include_types, include_fields):
    """
    Filter the data to only include records of the specified types
    and remove specified fields.
    Parameters:
        data (list): List of JSON records.
        include_types (set): Asset types to keep.
        include_fields (set): Fields to remove from records.
    Returns:
        list: Filtered data.
    """
    filtered_data = []
    
    for record in data:
        if record.get("asset_type") in include_types:
            # Create a copy to avoid modifying original data
            filtered_record = {key: value for key, value in record.items() if key in include_fields}
            filtered_data.append(filtered_record)
    
    return filtered_data

def visualize_network_pyvis(data, include_types, type_to_position, output_html = "graph.html"):

    width="100%" 
    height="1200px"
    layer_spacing=2.0
    node_spacing=2.0

    # import tempfile
    # output_html = tempfile.NamedTemporaryFile(delete=False)

    G = create_network(data, include_types, type_to_position)  # Create the network graph
    # 1. Assign hierarchy and determine types
    nx.set_node_attributes(
        G,
        {n: type_to_position.get(attrs["type"], 0) for n, attrs in G.nodes(data=True)},
        name="layer"
    )
    types = sorted({attrs["type"] for _, attrs in G.nodes(data=True)})

    # 2. Layout
    pos = nx.multipartite_layout(G, subset_key="layer", align="horizontal")
    pos = {n: (x * layer_spacing * 1000, y * node_spacing * 1000) for n, (x, y) in pos.items()}

    # 3. Assign colors
    base_colors = ["#e6194B", "#3cb44b", "#ffe119", "#4363d8", "#f58231",
                   "#911eb4", "#46f0f0", "#f032e6", "#bcf60c", "#fabebe"]
    color_map = {t: base_colors[i % len(base_colors)] for i, t in enumerate(types)}

    # 4. Build the graph
    net = Network(width=width, height=height, directed=True)
    net.toggle_physics(False)
    net.toggle_drag_nodes(True)
    net.show_buttons(filter_=["physics", "interaction", "clustering"])

    for node_id, attrs in G.nodes(data=True):
        x, y = pos[node_id]
        t = attrs["type"]
        c = color_map[t]
        net.add_node(node_id, label=attrs.get("label", str(node_id)), title=f"{attrs.get('label')}<br>Type: {t}<br>{attrs.get('uri','')}",
                     x=x, y=y, physics=False, group=t,
                     color={"background": c, "border": c, "highlight": {"background": c, "border": c}, "hover": {"background": c, "border": c}})
    for u, v in G.edges():
        net.add_edge(u, v)

    # 5. Inject clustering + double-click expand
    cluster_js = f"""
<script type="text/javascript">
window.addEventListener('load', () => {{
  const types = {types};
  const colors = {color_map};
  types.forEach(t => {{
    network.cluster({{
      joinCondition: opts => opts.group === t,
      clusterNodeProperties: {{
        id: 'cluster:' + t,
        label: 'Cluster: ' + t,
        shape: 'box',
        color: {{
          background: colors[t],
          border: colors[t],
          highlight: {{background: colors[t], border: colors[t]}},
          hover: {{background: colors[t], border: colors[t]}}
        }}
      }}
    }});
  }});
  network.on('doubleClick', params => {{
    if (params.nodes.length === 1) {{
      const id = params.nodes[0];
      if (network.isCluster(id)) {{
        network.openCluster(id);
      }}
    }}
  }});
}});
</script>
"""

    # 6. Inject a floating HTML legend panel
    legend_html = "<div id='legend' style='position:fixed; top:10px; right:10px; background:white; border:1px solid #ccc; padding:10px; font-family:Arial;'><strong>Legend</strong><br>"
    for t in types:
        c = color_map[t]
        legend_html += f"<div style='display:flex; align-items:center; margin-top:5px;'><span style='width:15px; height:15px; background:{c}; border:1px solid #000; display:inline-block; margin-right:5px;'></span>{t}</div>"
    legend_html += "</div>"
    
    # net.toggle_physics(True)


    # Combine HTML parts
    net.save_graph(output_html)
    with open(output_html, "r+", encoding="utf-8") as f:
        html = f.read().replace("</body>", cluster_js + legend_html + "\n</body>")
        f.seek(0); f.write(html); f.truncate()

    return output_html

def create_lineage_graph(lineage_data, type="pyvis", output_file="graph.html"):

    all_types = {record.get("asset_type") for record in lineage_data if record.get("asset_type")}
    # Exclude unwanted types.
    exclude_types = {"user", "secret", "org_secret", "commit", "actor", "build_machine"}
    types_to_position = {
            "organization": 1,
            "repo": 2,
            "branch": 3,
            "workflow": 4,
            "workflow_run": 5,
            "image": 6,
            "pod": 7,
            "namespace": 8

    }
    include_types = all_types - exclude_types
    # Generate and display the network visualization.
    if type == "plotly":
        # fig = visualize_network_plotly(lineage_data, include_types, types_to_position)
        return None
    else:
        output_file = Path(output_file)
        output_file.touch(exist_ok=True)
        output_file = output_file.resolve()
        output_file = str(output_file)
        file = visualize_network_pyvis(lineage_data, include_types, types_to_position, output_html=output_file)
        return file

if __name__ == "__main__":
    # Open the JSON file and load data.
    filepath = "demo-data/astro-lineage.json"
    with open(filepath, "r") as f:
        data = json.load(f)
    

    filename = create_lineage_graph(data, output_file="demo-data/lineage_graph.html")
    print(f"Network graph saved to {filename}")
    print(f"Open by browsing file:///{filename}")

	import json
	import networkx as nx
	# import plotly.graph_objects as go
	from pyvis.network import Network
	from pathlib import Path


	def create_network(data, include_types, type_to_position):
	G = nx.DiGraph()
	for record in data:
	asset_type = record.get("asset_type")
	parent_type = record.get("parent_type")
	child = record.get("external_id")
	parent = record.get("parent_external_id")
	# Filter nodes based on allowed types.
	add_child = include_types is None or asset_type in include_types
	add_parent = include_types is None or parent_type in include_types

	if add_child and child:
	node_attrs = {"type": asset_type, "label": record.get("asset_name")}
	if record.get("uri"):
	node_attrs["uri"] = record.get("uri")
	G.add_node(child, **node_attrs)
	if parent and add_parent:
	node_attrs = {"type": parent_type, "label": record.get("parent_name")}
	if record.get("uri"):
	node_attrs["uri"] = record.get("uri")
	G.add_node(parent, **node_attrs)
	if child and parent and add_child and add_parent:
	G.add_edge(parent, child)

	# Assign 'layer' attribute for each node based on its type.
	for node in G.nodes():
	node_type = G.nodes[node].get('type')
	G.nodes[node]['layer'] = type_to_position.get(node_type, 1)
	return G

	# def visualize_network_plotly(data, include_types, type_to_position):
	# """
	# Build and return a Plotly figure of a network graph.

	# Parameters:
	# data (list): JSON array of records.
	# include_types (set): Allowed asset types.
	# type_to_position (dict): Mapping of asset types to numeric layer positions.

	# Returns:
	# go.Figure: The Plotly figure with the network visualization.
	# """
	# # Create network graph

	# G = create_network(data, include_types, type_to_position)
	# # Compute positions using a multipartite layout (vertical alignment)
	# pos = nx.multipartite_layout(G, subset_key="layer", align="vertical")

	# # Build edge coordinates for Plotly
	# edge_x, edge_y = [], []
	# for source, target in G.edges():
	# x0, y0 = pos[source]
	# x1, y1 = pos[target]
	# edge_x.extend([x0, x1, None])
	# edge_y.extend([y0, y1, None])
	# edge_trace = go.Scatter(
	# x=edge_x, y=edge_y,
	# line=dict(width=1, color='#888'),
	# hoverinfo='none',
	# mode='lines'
	# )

	# # Build node coordinates and text lists.
	# node_x, node_y = [], []
	# node_text_trunc, node_text_hover, node_color = [], [], []
	# for node in G.nodes():
	# x, y = pos[node]
	# node_x.append(x)
	# node_y.append(y)
	# full_label = G.nodes[node].get('label', str(node))
	# truncated_label = full_label[0:30] # truncated version for display
	# # Check for a "uri" attribute (lowercase) and create a clickable link if present.
	# uri = G.nodes[node].get('uri')
	# uri_link = f'<br><a href="{uri}" target="_blank">{uri}</a>' if uri else ""
	# node_text_trunc.append(f"{truncated_label}<br>Type: {G.nodes[node].get('type')}")
	# node_text_hover.append(f"{full_label}<br>Type: {G.nodes[node].get('type')}{uri_link}")
	# node_color.append(G.nodes[node]['layer'])

	# node_trace = go.Scatter(
	# x=node_x, y=node_y,
	# mode='markers+text',
	# textposition="bottom center",
	# hoverinfo='text',
	# marker=dict(
	# showscale=True,
	# colorscale='YlGnBu',
	# reversescale=True,
	# color=node_color,
	# size=15,
	# colorbar=dict(title='Layer')
	# ),
	# text=node_text_trunc, # Display truncated text on the plot.
	# hovertext=node_text_hover # Full text (with clickable URL if present) for hover.
	# )

	# # Create the Plotly figure.
	# fig = go.Figure(data=[edge_trace, node_trace],
	# layout=go.Layout(
	# title='<br>Integrated Network Graph Visualization',
	# titlefont_size=16,
	# dragmode='pan',
	# showlegend=False,
	# hovermode='closest',
	# margin=dict(b=20, l=5, r=5, t=40),
	# width=1000,
	# height=800,
	# xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
	# yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)
	# ))
	# return fig

	def filter_network(data, include_types, include_fields):
	"""
	Filter the data to only include records of the specified types
	and remove specified fields.

	Parameters:
	data (list): List of JSON records.
	include_types (set): Asset types to keep.
	include_fields (set): Fields to remove from records.

	Returns:
	list: Filtered data.
	"""
	filtered_data = []

	for record in data:
	if record.get("asset_type") in include_types:
	# Create a copy to avoid modifying original data
	filtered_record = {key: value for key, value in record.items() if key in include_fields}
	filtered_data.append(filtered_record)

	return filtered_data

	def visualize_network_pyvis(data, include_types, type_to_position, output_html = "graph.html"):

	width="100%"
	height="1200px"
	layer_spacing=2.0
	node_spacing=2.0

	# import tempfile
	# output_html = tempfile.NamedTemporaryFile(delete=False)

	G = create_network(data, include_types, type_to_position) # Create the network graph
	# 1. Assign hierarchy and determine types
	nx.set_node_attributes(
	G,
	{n: type_to_position.get(attrs["type"], 0) for n, attrs in G.nodes(data=True)},
	name="layer"
	)
	types = sorted({attrs["type"] for _, attrs in G.nodes(data=True)})

	# 2. Layout
	pos = nx.multipartite_layout(G, subset_key="layer", align="horizontal")
	pos = {n: (x * layer_spacing * 1000, y * node_spacing * 1000) for n, (x, y) in pos.items()}

	# 3. Assign colors
	base_colors = ["#e6194B", "#3cb44b", "#ffe119", "#4363d8", "#f58231",
	"#911eb4", "#46f0f0", "#f032e6", "#bcf60c", "#fabebe"]
	color_map = {t: base_colors[i % len(base_colors)] for i, t in enumerate(types)}

	# 4. Build the graph
	net = Network(width=width, height=height, directed=True)
	net.toggle_physics(False)
	net.toggle_drag_nodes(True)
	net.show_buttons(filter_=["physics", "interaction", "clustering"])

	for node_id, attrs in G.nodes(data=True):
	x, y = pos[node_id]
	t = attrs["type"]
	c = color_map[t]
	net.add_node(node_id, label=attrs.get("label", str(node_id)), title=f"{attrs.get('label')}<br>Type: {t}<br>{attrs.get('uri','')}",
	x=x, y=y, physics=False, group=t,
	color={"background": c, "border": c, "highlight": {"background": c, "border": c}, "hover": {"background": c, "border": c}})
	for u, v in G.edges():
	net.add_edge(u, v)

	# 5. Inject clustering + double-click expand
	cluster_js = f"""
	<script type="text/javascript">
	window.addEventListener('load', () => {{
	const types = {types};
	const colors = {color_map};
	types.forEach(t => {{
	network.cluster({{
	joinCondition: opts => opts.group === t,
	clusterNodeProperties: {{
	id: 'cluster:' + t,
	label: 'Cluster: ' + t,
	shape: 'box',
	color: {{
	background: colors[t],
	border: colors[t],
	highlight: {{background: colors[t], border: colors[t]}},
	hover: {{background: colors[t], border: colors[t]}}
	}}
	}}
	}});
	}});
	network.on('doubleClick', params => {{
	if (params.nodes.length === 1) {{
	const id = params.nodes[0];
	if (network.isCluster(id)) {{
	network.openCluster(id);
	}}
	}}
	}});
	}});
	</script>
	"""

	# 6. Inject a floating HTML legend panel
	legend_html = "<div id='legend' style='position:fixed; top:10px; right:10px; background:white; border:1px solid #ccc; padding:10px; font-family:Arial;'><strong>Legend</strong><br>"
	for t in types:
	c = color_map[t]
	legend_html += f"<div style='display:flex; align-items:center; margin-top:5px;'><span style='width:15px; height:15px; background:{c}; border:1px solid #000; display:inline-block; margin-right:5px;'></span>{t}</div>"
	legend_html += "</div>"

	# net.toggle_physics(True)


	# Combine HTML parts
	net.save_graph(output_html)
	with open(output_html, "r+", encoding="utf-8") as f:
	html = f.read().replace("</body>", cluster_js + legend_html + "\n</body>")
	f.seek(0); f.write(html); f.truncate()

	return output_html

	def create_lineage_graph(lineage_data, type="pyvis", output_file="graph.html"):

	all_types = {record.get("asset_type") for record in lineage_data if record.get("asset_type")}
	# Exclude unwanted types.
	exclude_types = {"user", "secret", "org_secret", "commit", "actor", "build_machine"}
	types_to_position = {
	"organization": 1,
	"repo": 2,
	"branch": 3,
	"workflow": 4,
	"workflow_run": 5,
	"image": 6,
	"pod": 7,
	"namespace": 8

	}
	include_types = all_types - exclude_types
	# Generate and display the network visualization.
	if type == "plotly":
	# fig = visualize_network_plotly(lineage_data, include_types, types_to_position)
	return None
	else:
	output_file = Path(output_file)
	output_file.touch(exist_ok=True)
	output_file = output_file.resolve()
	output_file = str(output_file)
	file = visualize_network_pyvis(lineage_data, include_types, types_to_position, output_html=output_file)
	return file

	if __name__ == "__main__":
	# Open the JSON file and load data.
	filepath = "demo-data/astro-lineage.json"
	with open(filepath, "r") as f:
	data = json.load(f)


	filename = create_lineage_graph(data, output_file="demo-data/lineage_graph.html")
	print(f"Network graph saved to {filename}")
	print(f"Open by browsing file:///{filename}")