concept_map_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os
from graph_generator_utils import add_nodes_and_edges

def generate_concept_map(json_input: str, output_format: str) -> str:
    """
    Generates a concept map from JSON input.

    Args:
        json_input (str): A JSON string describing the concept map structure.
                          It must follow the Expected JSON Format Example below.

    Expected JSON Format Example:
    {
      "central_node": "Artificial Intelligence (AI)",
      "nodes": [
        {
          "id": "ml_fundamental",
          "label": "Machine Learning",
          "relationship": "is essential for",
          "subnodes": [
            {
              "id": "dl_branch",
              "label": "Deep Learning",
              "relationship": "for example",
              "subnodes": [
                {
                  "id": "cnn_example",
                  "label": "CNNs",
                  "relationship": "for example"
                },
                {
                  "id": "rnn_example",
                  "label": "RNNs",
                  "relationship": "for example"
                }
              ]
            },
            {
              "id": "rl_branch",
              "label": "Reinforcement Learning",
              "relationship": "for example",
              "subnodes": [
                {
                  "id": "qlearning_example",
                  "label": "Q-Learning",
                  "relationship": "example"
                },
                {
                  "id": "pg_example",
                  "label": "Policy Gradients",
                  "relationship": "example"
                }
              ]
            }
          ]
        },
        {
          "id": "ai_types",
          "label": "Types",
          "relationship": "formed by",
          "subnodes": [
            {
              "id": "agi_type",
              "label": "AGI",
              "relationship": "this is",
              "subnodes": [
                {
                  "id": "strong_ai",
                  "label": "Strong AI",
                  "relationship": "provoked by",
                  "subnodes": [
                    {
                      "id": "human_intel",
                      "label": "Human-level Intel.",
                      "relationship": "of"
                    }
                  ]
                }
              ]
            },
            {
              "id": "ani_type",
              "label": "ANI",
              "relationship": "this is",
              "subnodes": [
                {
                  "id": "weak_ai",
                  "label": "Weak AI",
                  "relationship": "provoked by",
                  "subnodes": [
                    {
                      "id": "narrow_tasks",
                      "label": "Narrow Tasks",
                      "relationship": "of"
                    }
                  ]
                }
              ]
            }
          ]
        },
        {
          "id": "ai_capabilities",
          "label": "Capabilities",
          "relationship": "change",
          "subnodes": [
            {
              "id": "data_proc",
              "label": "Data Processing",
              "relationship": "can",
              "subnodes": [
                {
                  "id": "big_data",
                  "label": "Big Data",
                  "relationship": "as",
                  "subnodes": [
                    {
                      "id": "analysis_example",
                      "label": "Data Analysis",
                      "relationship": "example"
                    },
                    {
                      "id": "prediction_example",
                      "label": "Prediction",
                      "relationship": "example"
                    }
                  ]
                }
              ]
            },
            {
              "id": "decision_making",
              "label": "Decision Making",
              "relationship": "can be",
              "subnodes": [
                {
                  "id": "automation",
                  "label": "Automation",
                  "relationship": "as",
                  "subnodes": [
                    {
                      "id": "robotics_example",
                      "label": "Robotics",
                      "relationship": "Example"},
                    {
                      "id": "autonomous_example",
                      "label": "Autonomous Vehicles",
                      "relationship": "of one"
                    }
                  ]
                }
              ]
            },
            {
              "id": "problem_solving",
              "label": "Problem Solving",
              "relationship": "can",
              "subnodes": [
                {
                  "id": "optimization",
                  "label": "Optimization",
                  "relationship": "as is",
                  "subnodes": [
                    {
                      "id": "algorithms_example",
                      "label": "Algorithms",
                      "relationship": "for example"
                    }
                  ]
                }
              ]
            }
          ]
        }
      ]
    }

    Returns:
        str: The filepath to the generated PNG image file.
    """
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        if 'central_node' not in data or 'nodes' not in data:
            raise ValueError("Missing required fields: central_node or nodes")

        dot = graphviz.Digraph(
            name='ConceptMap',
            format='png',
            graph_attr={
                'rankdir': 'TB',        # Top-to-Bottom layout (vertical hierarchy)
                'splines': 'ortho',     # Straight lines
                'bgcolor': 'white',     # White background
                'pad': '0.5'            # Padding around the graph
            }
        )
        
        base_color = '#19191a' # Hardcoded base color

        # Central node styling (rounded box, dark color)
        dot.node(
            'central',
            data['central_node'],
            shape='box',            # Rectangular shape
            style='filled,rounded', # Filled and rounded corners
            fillcolor=base_color,   # Darkest color
            fontcolor='white',      # White text for dark background
            fontsize='16'           # Larger font for central node
        )
        
        # Add child nodes and edges recursively starting from depth 1
        add_nodes_and_edges(dot, 'central', data.get('nodes', []), current_depth=1, base_color=base_color)

        with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
            dot.render(tmp.name, format=output_format, cleanup=True)
            return f"{tmp.name}.{output_format}"

    except json.JSONDecodeError:
        return "Error: Invalid JSON format"
    except Exception as e:
        return f"Error: {str(e)}"