convert_output_to_mermaid.py

import json


class MermaidConverter:
    def __init__(self):
        # Define style mappings
        self.style_colors = [
            {'entity_type':'location', 'node_color':'#264653', 'text_color': '#fff'},
            {'entity_type':'organization', 'node_color':'#287271', 'text_color': '#fff'},
            {'entity_type':'person', 'node_color':'#2a9d8f', 'text_color': '#fff'},
            {'entity_type':'event', 'node_color':'#8ab17d', 'text_color': '#000'},
            {'entity_type':'object', 'node_color':'#e9c46a', 'text_color': '#000'},
            {'entity_type':'date', 'node_color':'#f4a261', 'text_color': '#000'},
            {'entity_type':'skill', 'node_color':'#e76f51', 'text_color': '#000'}
        ]
        
        # Node name sanitization mapping
        self.special_chars = str.maketrans("çğıöşüÇĞİÖŞÜ ", "cgioscCGIOSU_")

    def sanitize_node_name(self, name):
        """Convert a string into a valid Mermaid node identifier."""
        return name.translate(self.special_chars)

    def generate_style_definitions(self):
        """Generate Mermaid style class definitions."""
        styles = []
        for style in self.style_colors:
            entity_type = style['entity_type']
            node_color = style['node_color']
            text_color = style['text_color']

            styles.append(
                f"classDef {entity_type} fill:{node_color},stroke:#333,stroke-width:2px,color:{text_color};"
            )
        return "\n    ".join(styles)

    def generate_nodes(self, data):
        """Generate Mermaid node definitions."""
        nodes = []
        node_mappings = {}  # Store sanitized names for relationships
        
        for node in data['nodes']:
            name = node['name']
            node_type = node['type'].lower()
            sanitized_name = self.sanitize_node_name(name)
            node_mappings[name] = sanitized_name
            
            # Create node definition with label showing name and type
            nodes.append(f'{sanitized_name}(["{name}\\n({node_type})"])')
            
        return nodes, node_mappings

    def generate_relationships(self, data, node_mappings):
        """Generate Mermaid relationship definitions."""
        relationships = []
        
        for rel in data['relationships']:
            source = node_mappings[rel['source']]
            target = node_mappings[rel['target']]
            relationship = rel['relationship']
            
            relationships.append(
                f'{source} -->|"{relationship}"| {target}'
            )
            
        return relationships

    def generate_style_applications(self, data, node_mappings):
        """Generate Mermaid style class applications."""
        # Group nodes by type
        types = {}
        for node in data['nodes']:
            node_type = node['type'].lower()
            if node_type not in types:
                types[node_type] = []
            types[node_type].append(node_mappings[node['name']])
        
        # Generate style applications
        style_apps = []
        for node_type, nodes in types.items():
            style_apps.append(f"class {','.join(nodes)} {node_type};")
            
        return style_apps

    def convert(self, data):
        """Convert the data dictionary to Mermaid diagram format."""
        # Generate nodes and get mappings
        nodes, node_mappings = self.generate_nodes(data)
        
        # Generate other components
        relationships = self.generate_relationships(data, node_mappings)
        style_applications = self.generate_style_applications(data, node_mappings)
        
        # Construct the complete diagram
        diagram = [
            "%%{",
            "  init: {",
            "    'theme': 'base',",
            "    'themeVariables': {",
            "      'primaryTextColor': '#000',",
            "      'lineColor': '#232323',",
            "      'secondaryColor': '#ffffff',",
            "      'tertiaryColor': '#fff'",
            "    }",
            "  }",
            "}%%",
            "graph LR",
            "    %% Style definitions",
            f"    {self.generate_style_definitions()}",
            "",
            "    %% Nodes",
            "    " + "\n    ".join(nodes),
            "",
            "    %% Relationships",
            "    " + "\n    ".join(relationships),
            "",
            "    %% Applying styles",
            "    " + "\n    ".join(style_applications)
        ]
        
        return "\n".join(diagram)

# Example usage
if __name__ == "__main__":

    with open('knowledge_graph.json', 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    converter = MermaidConverter()
    mermaid_diagram = converter.convert(data)

    with open('knowledge_graph.txt', 'w', encoding='utf-8') as f:
        f.write(mermaid_diagram)