import os
import json
import networkx
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):
nodes_data, edges_data = graph_data
G = networkx.DiGraph()
nodes_list = []
color_map = {
"Entity": "#f47710",
"EntityType": "#6510f4",
"DocumentChunk": "#801212",
"TextSummary": "#1077f4",
"TableRow": "#f47710",
"TableType": "#6510f4",
"ColumnValue": "#13613a",
"default": "#D3D3D3",
}
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 = """
"""
html_content = html_template.replace("{nodes}", json.dumps(nodes_list))
html_content = html_content.replace("{links}", json.dumps(links_list))
if not destination_file_path:
home_dir = os.path.expanduser("~")
destination_file_path = os.path.join(home_dir, "graph_visualization.html")
dir_path = os.path.dirname(destination_file_path)
file_path = os.path.basename(destination_file_path)
file_storage = LocalFileStorage(dir_path)
file_storage.store(file_path, html_content, overwrite=True)
logger.info(f"Graph visualization saved as {destination_file_path}")
return html_content