Update entity_relationship_generator.py

entity_relationship_generator.py

@@ -1,3 +1,650 @@
+# import graphviz
+# import json
+# from tempfile import NamedTemporaryFile
+# import os
+
+# def generate_entity_relationship_diagram(json_input: str, output_format: str) -> str:
+#     """
+#     Generates an Entity Relationship (ER) diagram from JSON input.
+
+#     Args:
+#         json_input (str): A JSON string describing the ER diagram structure.
+#                           It must follow the Expected JSON Format Example below.
+
+#     Expected JSON Format Example:
+#     {
+#       "entities": [
+#         {
+#           "name": "User",
+#           "type": "strong",
+#           "attributes": [
+#             {
+#               "name": "user_id",
+#               "type": "primary_key"
+#             },
+#             {
+#               "name": "username",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "email",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "password_hash",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "full_name",
+#               "type": "composite"
+#             },
+#             {
+#               "name": "phone_numbers",
+#               "type": "multivalued"
+#             },
+#             {
+#               "name": "age",
+#               "type": "derived"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "Product",
+#           "type": "strong",
+#           "attributes": [
+#             {
+#               "name": "product_id",
+#               "type": "primary_key"
+#             },
+#             {
+#               "name": "name",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "description",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "price",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "stock_quantity",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "tags",
+#               "type": "multivalued"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "Category",
+#           "type": "strong",
+#           "attributes": [
+#             {
+#               "name": "category_id",
+#               "type": "primary_key"
+#             },
+#             {
+#               "name": "name",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "description",
+#               "type": "regular"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "Order",
+#           "type": "strong",
+#           "attributes": [
+#             {
+#               "name": "order_id",
+#               "type": "primary_key"
+#             },
+#             {
+#               "name": "order_date",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "status",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "total_amount",
+#               "type": "derived"
+#             },
+#             {
+#               "name": "shipping_address",
+#               "type": "composite"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "OrderItem",
+#           "type": "weak",
+#           "attributes": [
+#             {
+#               "name": "line_number",
+#               "type": "partial_key"
+#             },
+#             {
+#               "name": "quantity",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "unit_price",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "subtotal",
+#               "type": "derived"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "Payment",
+#           "type": "strong",
+#           "attributes": [
+#             {
+#               "name": "payment_id",
+#               "type": "primary_key"
+#             },
+#             {
+#               "name": "amount",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "payment_method",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "payment_date",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "status",
+#               "type": "regular"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "Review",
+#           "type": "strong",
+#           "attributes": [
+#             {
+#               "name": "review_id",
+#               "type": "primary_key"
+#             },
+#             {
+#               "name": "rating",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "comment",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "review_date",
+#               "type": "regular"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "Vendor",
+#           "type": "strong",
+#           "attributes": [
+#             {
+#               "name": "vendor_id",
+#               "type": "primary_key"
+#             },
+#             {
+#               "name": "company_name",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "contact_person",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "contact_emails",
+#               "type": "multivalued"
+#             },
+#             {
+#               "name": "business_address",
+#               "type": "composite"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "ShoppingCart",
+#           "type": "strong",
+#           "attributes": [
+#             {
+#               "name": "cart_id",
+#               "type": "primary_key"
+#             },
+#             {
+#               "name": "created_date",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "last_updated",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "total_items",
+#               "type": "derived"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "CartItem",
+#           "type": "weak",
+#           "attributes": [
+#             {
+#               "name": "item_position",
+#               "type": "partial_key"
+#             },
+#             {
+#               "name": "quantity",
+#               "type": "regular"
+#             },
+#             {
+#               "name": "added_date",
+#               "type": "regular"
+#             }
+#           ]
+#         }
+#       ],
+#       "relationships": [
+#         {
+#           "name": "PlacesOrder",
+#           "type": "regular",
+#           "entities": ["User", "Order"],
+#           "cardinalities": {
+#             "User": "1",
+#             "Order": "M"
+#           },
+#           "attributes": []
+#         },
+#         {
+#           "name": "Contains",
+#           "type": "identifying",
+#           "entities": ["Order", "OrderItem"],
+#           "cardinalities": {
+#             "Order": "1",
+#             "OrderItem": "M"
+#           },
+#           "attributes": []
+#         },
+#         {
+#           "name": "OrdersProduct",
+#           "type": "regular",
+#           "entities": ["OrderItem", "Product"],
+#           "cardinalities": {
+#             "OrderItem": "M",
+#             "Product": "1"
+#           },
+#           "attributes": []
+#         },
+#         {
+#           "name": "BelongsTo",
+#           "type": "regular",
+#           "entities": ["Product", "Category"],
+#           "cardinalities": {
+#             "Product": "M",
+#             "Category": "1"
+#           },
+#           "attributes": []
+#         },
+#         {
+#           "name": "ProcessesPayment",
+#           "type": "regular",
+#           "entities": ["Order", "Payment"],
+#           "cardinalities": {
+#             "Order": "1",
+#             "Payment": "M"
+#           },
+#           "attributes": []
+#         },
+#         {
+#           "name": "WritesReview",
+#           "type": "regular",
+#           "entities": ["User", "Review"],
+#           "cardinalities": {
+#             "User": "1",
+#             "Review": "M"
+#           },
+#           "attributes": []
+#         },
+#         {
+#           "name": "ReviewsProduct",
+#           "type": "regular",
+#           "entities": ["Review", "Product"],
+#           "cardinalities": {
+#             "Review": "M",
+#             "Product": "1"
+#           },
+#           "attributes": []
+#         },
+#         {
+#           "name": "Supplies",
+#           "type": "regular",
+#           "entities": ["Vendor", "Product"],
+#           "cardinalities": {
+#             "Vendor": "M",
+#             "Product": "M"
+#           },
+#           "attributes": [
+#             {
+#               "name": "supply_price"
+#             },
+#             {
+#               "name": "lead_time"
+#             }
+#           ]
+#         },
+#         {
+#           "name": "HasCart",
+#           "type": "regular",
+#           "entities": ["User", "ShoppingCart"],
+#           "cardinalities": {
+#             "User": "1",
+#             "ShoppingCart": "1"
+#           },
+#           "attributes": []
+#         },
+#         {
+#           "name": "CartContains",
+#           "type": "identifying",
+#           "entities": ["ShoppingCart", "CartItem"],
+#           "cardinalities": {
+#             "ShoppingCart": "1",
+#             "CartItem": "M"
+#           },
+#           "attributes": []
+#         },
+#         {
+#           "name": "CartHasProduct",
+#           "type": "regular",
+#           "entities": ["CartItem", "Product"],
+#           "cardinalities": {
+#             "CartItem": "M",
+#             "Product": "1"
+#           },
+#           "attributes": []
+#         }
+#       ]
+#     }
+
+#     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 'entities' not in data:
+#             raise ValueError("Missing required field: entities")
+
+#         dot = graphviz.Graph(comment='ER Diagram', engine='neato')
+#         dot.attr(
+#             bgcolor='white',
+#             pad='1.5',
+#             overlap='false',
+#             splines='true',
+#             sep='+25',
+#             esep='+15'
+#         )
+#         dot.attr('node', fontname='Arial', fontsize='10', color='#404040')
+#         dot.attr('edge', fontname='Arial', fontsize='9', color='#4a4a4a')
+        
+#         entity_color = '#BEBEBE'        
+#         attribute_color = '#D4D4D4'     
+#         relationship_color = '#E8E8E8'  
+#         isa_color = '#F0F0F0'          
+        
+#         font_color = 'black'
+        
+#         entities = data.get('entities', [])
+#         relationships = data.get('relationships', [])
+        
+#         # Process entities with new styling
+#         for entity in entities:
+#             entity_name = entity.get('name')
+#             entity_type = entity.get('type', 'strong')
+#             attributes = entity.get('attributes', [])
+            
+#             if not entity_name:
+#                 continue
+            
+#             if entity_type == 'weak':
+#                 dot.node(
+#                     entity_name, 
+#                     entity_name, 
+#                     shape='box', 
+#                     style='filled,rounded',
+#                     fillcolor=entity_color,  
+#                     fontcolor=font_color,
+#                     color='#404040',
+#                     penwidth='3',
+#                     width='1.8',
+#                     height='0.8',
+#                     fontsize='12'
+#                 )
+#             else:
+#                 dot.node(
+#                     entity_name, 
+#                     entity_name, 
+#                     shape='box', 
+#                     style='filled,rounded',
+#                     fillcolor=entity_color,  
+#                     fontcolor=font_color,
+#                     color='#404040',
+#                     penwidth='1',
+#                     width='1.8',
+#                     height='0.8',
+#                     fontsize='12'
+#                 )
+            
+#             for i, attr in enumerate(attributes):
+#                 attr_name = attr.get('name', '')
+#                 attr_type = attr.get('type', 'regular')
+                
+#                 attr_id = f"{entity_name}_attr_{i}"
+                
+#                 if attr_type == 'primary_key':
+#                     dot.node(
+#                         attr_id, 
+#                         f'{attr_name} (PK)', 
+#                         shape='ellipse', 
+#                         style='filled,rounded',
+#                         fillcolor=attribute_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         width='1.2',
+#                         height='0.6',
+#                         fontsize='10'
+#                     )
+#                 elif attr_type == 'partial_key':
+#                     dot.node(
+#                         attr_id, 
+#                         f'{attr_name} (Partial)', 
+#                         shape='ellipse', 
+#                         style='filled,rounded,dashed',
+#                         fillcolor=attribute_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         width='1.2',
+#                         height='0.6',
+#                         fontsize='10'
+#                     )
+#                 elif attr_type == 'multivalued':
+#                     dot.node(
+#                         attr_id, 
+#                         attr_name, 
+#                         shape='ellipse', 
+#                         style='filled,rounded',
+#                         fillcolor=attribute_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         penwidth='3',
+#                         width='1.2',
+#                         height='0.6',
+#                         fontsize='10'
+#                     )
+#                 elif attr_type == 'derived':
+#                     dot.node(
+#                         attr_id, 
+#                         f'/{attr_name}/', 
+#                         shape='ellipse', 
+#                         style='filled,rounded,dashed',
+#                         fillcolor=attribute_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         width='1.2',
+#                         height='0.6',
+#                         fontsize='10'
+#                     )
+#                 elif attr_type == 'composite':
+#                     dot.node(
+#                         attr_id, 
+#                         attr_name, 
+#                         shape='ellipse', 
+#                         style='filled,rounded',
+#                         fillcolor=attribute_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         width='1.2',
+#                         height='0.6',
+#                         fontsize='10'
+#                     )
+#                 else:  
+#                     dot.node(
+#                         attr_id, 
+#                         attr_name, 
+#                         shape='ellipse', 
+#                         style='filled,rounded',
+#                         fillcolor=attribute_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         width='1.2',
+#                         height='0.6',
+#                         fontsize='10'
+#                     )
+                
+#                 dot.edge(entity_name, attr_id, color='#4a4a4a', len='1.5')
+        
+#         for relationship in relationships:
+#             rel_name = relationship.get('name')
+#             rel_type = relationship.get('type', 'regular')
+#             entities_involved = relationship.get('entities', [])
+#             cardinalities = relationship.get('cardinalities', {})
+#             rel_attributes = relationship.get('attributes', [])
+            
+#             if not rel_name:
+#                 continue
+            
+#             if rel_type == 'isa':
+#                 parent = relationship.get('parent')
+#                 children = relationship.get('children', [])
+                
+#                 if parent and children:
+#                     isa_id = f"isa_{rel_name}"
+                    
+#                     dot.node(
+#                         isa_id, 
+#                         'ISA', 
+#                         shape='triangle', 
+#                         style='filled,rounded',
+#                         fillcolor=isa_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         penwidth='2',
+#                         width='1.0',
+#                         height='0.8',
+#                         fontsize='10'
+#                     )
+                    
+#                     dot.edge(parent, isa_id, color='#4a4a4a', len='2.0')
+                    
+#                     for child in children:
+#                         dot.edge(isa_id, child, color='#4a4a4a', len='2.0')
+                
+#             elif len(entities_involved) >= 2:
+#                 if rel_type == 'identifying':
+#                     dot.node(
+#                         rel_name, 
+#                         rel_name, 
+#                         shape='diamond', 
+#                         style='filled,rounded',
+#                         fillcolor=relationship_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         penwidth='3',
+#                         width='1.8',
+#                         height='1.0',
+#                         fontsize='11'
+#                     )
+#                 else:
+#                     dot.node(
+#                         rel_name, 
+#                         rel_name, 
+#                         shape='diamond', 
+#                         style='filled,rounded',
+#                         fillcolor=relationship_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         penwidth='1',
+#                         width='1.8',
+#                         height='1.0',
+#                         fontsize='11'
+#                     )
+                
+#                 for j, attr in enumerate(rel_attributes):
+#                     attr_name = attr.get('name', '')
+#                     attr_id = f"{rel_name}_attr_{j}"
+                    
+#                     dot.node(
+#                         attr_id, 
+#                         attr_name, 
+#                         shape='ellipse', 
+#                         style='filled,rounded',
+#                         fillcolor=attribute_color,  
+#                         fontcolor=font_color,
+#                         color='#404040',
+#                         width='1.0',
+#                         height='0.5',
+#                         fontsize='9'
+#                     )
+#                     dot.edge(rel_name, attr_id, color='#4a4a4a', len='1.0')
+                
+#                 for entity in entities_involved:
+#                     cardinality = cardinalities.get(entity, '1')
+#                     dot.edge(
+#                         entity, 
+#                         rel_name, 
+#                         label=f' {cardinality} ',
+#                         color='#4a4a4a',
+#                         len='2.5',
+#                         fontcolor='#4a4a4a',
+#                         fontsize='10'
+#                     )
+
+#         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)}"
+
 import graphviz
 import json
 from tempfile import NamedTemporaryFile
@@ -143,120 +790,6 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
               "type": "derived"
             }
           ]
-        },
-        {
-          "name": "Payment",
-          "type": "strong",
-          "attributes": [
-            {
-              "name": "payment_id",
-              "type": "primary_key"
-            },
-            {
-              "name": "amount",
-              "type": "regular"
-            },
-            {
-              "name": "payment_method",
-              "type": "regular"
-            },
-            {
-              "name": "payment_date",
-              "type": "regular"
-            },
-            {
-              "name": "status",
-              "type": "regular"
-            }
-          ]
-        },
-        {
-          "name": "Review",
-          "type": "strong",
-          "attributes": [
-            {
-              "name": "review_id",
-              "type": "primary_key"
-            },
-            {
-              "name": "rating",
-              "type": "regular"
-            },
-            {
-              "name": "comment",
-              "type": "regular"
-            },
-            {
-              "name": "review_date",
-              "type": "regular"
-            }
-          ]
-        },
-        {
-          "name": "Vendor",
-          "type": "strong",
-          "attributes": [
-            {
-              "name": "vendor_id",
-              "type": "primary_key"
-            },
-            {
-              "name": "company_name",
-              "type": "regular"
-            },
-            {
-              "name": "contact_person",
-              "type": "regular"
-            },
-            {
-              "name": "contact_emails",
-              "type": "multivalued"
-            },
-            {
-              "name": "business_address",
-              "type": "composite"
-            }
-          ]
-        },
-        {
-          "name": "ShoppingCart",
-          "type": "strong",
-          "attributes": [
-            {
-              "name": "cart_id",
-              "type": "primary_key"
-            },
-            {
-              "name": "created_date",
-              "type": "regular"
-            },
-            {
-              "name": "last_updated",
-              "type": "regular"
-            },
-            {
-              "name": "total_items",
-              "type": "derived"
-            }
-          ]
-        },
-        {
-          "name": "CartItem",
-          "type": "weak",
-          "attributes": [
-            {
-              "name": "item_position",
-              "type": "partial_key"
-            },
-            {
-              "name": "quantity",
-              "type": "regular"
-            },
-            {
-              "name": "added_date",
-              "type": "regular"
-            }
-          ]
         }
       ],
       "relationships": [
@@ -279,103 +812,6 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             "OrderItem": "M"
           },
           "attributes": []
-        },
-        {
-          "name": "OrdersProduct",
-          "type": "regular",
-          "entities": ["OrderItem", "Product"],
-          "cardinalities": {
-            "OrderItem": "M",
-            "Product": "1"
-          },
-          "attributes": []
-        },
-        {
-          "name": "BelongsTo",
-          "type": "regular",
-          "entities": ["Product", "Category"],
-          "cardinalities": {
-            "Product": "M",
-            "Category": "1"
-          },
-          "attributes": []
-        },
-        {
-          "name": "ProcessesPayment",
-          "type": "regular",
-          "entities": ["Order", "Payment"],
-          "cardinalities": {
-            "Order": "1",
-            "Payment": "M"
-          },
-          "attributes": []
-        },
-        {
-          "name": "WritesReview",
-          "type": "regular",
-          "entities": ["User", "Review"],
-          "cardinalities": {
-            "User": "1",
-            "Review": "M"
-          },
-          "attributes": []
-        },
-        {
-          "name": "ReviewsProduct",
-          "type": "regular",
-          "entities": ["Review", "Product"],
-          "cardinalities": {
-            "Review": "M",
-            "Product": "1"
-          },
-          "attributes": []
-        },
-        {
-          "name": "Supplies",
-          "type": "regular",
-          "entities": ["Vendor", "Product"],
-          "cardinalities": {
-            "Vendor": "M",
-            "Product": "M"
-          },
-          "attributes": [
-            {
-              "name": "supply_price"
-            },
-            {
-              "name": "lead_time"
-            }
-          ]
-        },
-        {
-          "name": "HasCart",
-          "type": "regular",
-          "entities": ["User", "ShoppingCart"],
-          "cardinalities": {
-            "User": "1",
-            "ShoppingCart": "1"
-          },
-          "attributes": []
-        },
-        {
-          "name": "CartContains",
-          "type": "identifying",
-          "entities": ["ShoppingCart", "CartItem"],
-          "cardinalities": {
-            "ShoppingCart": "1",
-            "CartItem": "M"
-          },
-          "attributes": []
-        },
-        {
-          "name": "CartHasProduct",
-          "type": "regular",
-          "entities": ["CartItem", "Product"],
-          "cardinalities": {
-            "CartItem": "M",
-            "Product": "1"
-          },
-          "attributes": []
         }
       ]
     }
@@ -404,17 +840,31 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
         dot.attr('node', fontname='Arial', fontsize='10', color='#404040')
         dot.attr('edge', fontname='Arial', fontsize='9', color='#4a4a4a')
         
-        entity_color = '#BEBEBE'        
-        attribute_color = '#D4D4D4'     
-        relationship_color = '#E8E8E8'  
-        isa_color = '#F0F0F0'          
+        entity_colors = {
+            'strong': '#BEBEBE',
+            'weak': '#FFF9C4'
+        }
+        
+        attribute_colors = {
+            'primary_key': '#A8E6CF',
+            'partial_key': '#FFB3BA',
+            'multivalued': '#B8D4F1',
+            'derived': '#F0F8FF',
+            'composite': '#E8E8E8',
+            'regular': '#D4D4D4'
+        }
+        
+        relationship_colors = {
+            'regular': '#E8E8E8',
+            'identifying': '#B8D4F1',
+            'isa': '#F0F0F0'
+        }
         
         font_color = 'black'
         
         entities = data.get('entities', [])
         relationships = data.get('relationships', [])
         
-        # Process entities with new styling
         for entity in entities:
             entity_name = entity.get('name')
             entity_type = entity.get('type', 'strong')
@@ -423,13 +873,15 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             if not entity_name:
                 continue
             
+            entity_color = entity_colors.get(entity_type, entity_colors['strong'])
+            
             if entity_type == 'weak':
                 dot.node(
                     entity_name, 
                     entity_name, 
                     shape='box', 
                     style='filled,rounded',
-                    fillcolor=entity_color,  
+                    fillcolor=entity_color,
                     fontcolor=font_color,
                     color='#404040',
                     penwidth='3',
@@ -443,7 +895,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                     entity_name, 
                     shape='box', 
                     style='filled,rounded',
-                    fillcolor=entity_color,  
+                    fillcolor=entity_color,
                     fontcolor=font_color,
                     color='#404040',
                     penwidth='1',
@@ -457,6 +909,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                 attr_type = attr.get('type', 'regular')
                 
                 attr_id = f"{entity_name}_attr_{i}"
+                attr_color = attribute_colors.get(attr_type, attribute_colors['regular'])
                 
                 if attr_type == 'primary_key':
                     dot.node(
@@ -464,7 +917,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         f'{attr_name} (PK)', 
                         shape='ellipse', 
                         style='filled,rounded',
-                        fillcolor=attribute_color,  
+                        fillcolor=attr_color,
                         fontcolor=font_color,
                         color='#404040',
                         width='1.2',
@@ -477,7 +930,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         f'{attr_name} (Partial)', 
                         shape='ellipse', 
                         style='filled,rounded,dashed',
-                        fillcolor=attribute_color,  
+                        fillcolor=attr_color,
                         fontcolor=font_color,
                         color='#404040',
                         width='1.2',
@@ -490,7 +943,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         attr_name, 
                         shape='ellipse', 
                         style='filled,rounded',
-                        fillcolor=attribute_color,  
+                        fillcolor=attr_color,
                         fontcolor=font_color,
                         color='#404040',
                         penwidth='3',
@@ -504,7 +957,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         f'/{attr_name}/', 
                         shape='ellipse', 
                         style='filled,rounded,dashed',
-                        fillcolor=attribute_color,  
+                        fillcolor=attr_color,
                         fontcolor=font_color,
                         color='#404040',
                         width='1.2',
@@ -517,20 +970,20 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         attr_name, 
                         shape='ellipse', 
                         style='filled,rounded',
-                        fillcolor=attribute_color,  
+                        fillcolor=attr_color,
                         fontcolor=font_color,
                         color='#404040',
                         width='1.2',
                         height='0.6',
                         fontsize='10'
                     )
-                else:  
+                else:
                     dot.node(
                         attr_id, 
                         attr_name, 
                         shape='ellipse', 
                         style='filled,rounded',
-                        fillcolor=attribute_color,  
+                        fillcolor=attr_color,
                         fontcolor=font_color,
                         color='#404040',
                         width='1.2',
@@ -556,13 +1009,14 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                 
                 if parent and children:
                     isa_id = f"isa_{rel_name}"
+                    isa_color = relationship_colors.get('isa', relationship_colors['regular'])
                     
                     dot.node(
                         isa_id, 
                         'ISA', 
                         shape='triangle', 
                         style='filled,rounded',
-                        fillcolor=isa_color,  
+                        fillcolor=isa_color,
                         fontcolor=font_color,
                         color='#404040',
                         penwidth='2',
@@ -577,13 +1031,15 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         dot.edge(isa_id, child, color='#4a4a4a', len='2.0')
                 
             elif len(entities_involved) >= 2:
+                rel_color = relationship_colors.get(rel_type, relationship_colors['regular'])
+                
                 if rel_type == 'identifying':
                     dot.node(
                         rel_name, 
                         rel_name, 
                         shape='diamond', 
                         style='filled,rounded',
-                        fillcolor=relationship_color,  
+                        fillcolor=rel_color,
                         fontcolor=font_color,
                         color='#404040',
                         penwidth='3',
@@ -597,7 +1053,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         rel_name, 
                         shape='diamond', 
                         style='filled,rounded',
-                        fillcolor=relationship_color,  
+                        fillcolor=rel_color,
                         fontcolor=font_color,
                         color='#404040',
                         penwidth='1',
@@ -609,13 +1065,14 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                 for j, attr in enumerate(rel_attributes):
                     attr_name = attr.get('name', '')
                     attr_id = f"{rel_name}_attr_{j}"
+                    attr_color = attribute_colors['regular']
                     
                     dot.node(
                         attr_id, 
                         attr_name, 
                         shape='ellipse', 
                         style='filled,rounded',
-                        fillcolor=attribute_color,  
+                        fillcolor=attr_color,
                         fontcolor=font_color,
                         color='#404040',
                         width='1.0',