import os import json from cognee.shared.logging_utils import get_logger from cognee.infrastructure.files.storage.LocalFileStorage import LocalFileStorage logger = get_logger() async def cognee_network_visualization(graph_data, destination_file_path: str = None): import networkx nodes_data, edges_data = graph_data G = networkx.DiGraph() nodes_list = [] color_map = { "Entity": "#5C10F4", "EntityType": "#A550FF", "DocumentChunk": "#0DFF00", "TextSummary": "#5C10F4", "TableRow": "#A550FF", "TableType": "#5C10F4", "ColumnValue": "#757470", "SchemaTable": "#A550FF", "DatabaseSchema": "#5C10F4", "SchemaRelationship": "#323332", "default": "#D8D8D8", } for node_id, node_info in nodes_data: node_info = node_info.copy() node_info["id"] = str(node_id) node_info["color"] = color_map.get(node_info.get("type", "default"), "#D3D3D3") node_info["name"] = node_info.get("name", str(node_id)) try: del node_info[ "updated_at" ] #:TODO: We should decide what properties to show on the nodes and edges, we dont necessarily need all. except KeyError: pass try: del node_info["created_at"] except KeyError: pass nodes_list.append(node_info) G.add_node(node_id, **node_info) edge_labels = {} links_list = [] for source, target, relation, edge_info in edges_data: source = str(source) target = str(target) G.add_edge(source, target) edge_labels[(source, target)] = relation # Extract edge metadata including all weights all_weights = {} primary_weight = None if edge_info: # Single weight (backward compatibility) if "weight" in edge_info: all_weights["default"] = edge_info["weight"] primary_weight = edge_info["weight"] # Multiple weights if "weights" in edge_info and isinstance(edge_info["weights"], dict): all_weights.update(edge_info["weights"]) # Use the first weight as primary for visual thickness if no default weight if primary_weight is None and edge_info["weights"]: primary_weight = next(iter(edge_info["weights"].values())) # Individual weight fields (weight_strength, weight_confidence, etc.) for key, value in edge_info.items(): if key.startswith("weight_") and isinstance(value, (int, float)): weight_name = key[7:] # Remove "weight_" prefix all_weights[weight_name] = value link_data = { "source": source, "target": target, "relation": relation, "weight": primary_weight, # Primary weight for backward compatibility "all_weights": all_weights, # All weights for display "relationship_type": edge_info.get("relationship_type") if edge_info else None, "edge_info": edge_info if edge_info else {}, } links_list.append(link_data) html_template = """

Hover a node or edge to inspect details

""" # Safely embed JSON inside