Add multi-label group matching for Gremlin visualization (#746)

ChangeLog.md

@@ -4,6 +4,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 
 ## Upcoming
 - Autoformatter for gremlin queries
+- Support group matching for multi-label vertices in Gremlin visualization ([Link to Issue](https://github.com/aws/graph-notebook/issues/746))
 
 ## Release 5.2.0 (Mar 11, 2026)
 - Added %degreeDistribution magic command ([PR](https://github.com/aws/graph-notebook/pull/749)) 

src/graph_notebook/magics/graph_magic.py

@@ -1393,7 +1393,8 @@ def gremlin(self, line, cell, local_ns: dict = None):
                                             label_max_length=args.label_max_length,
                                             edge_label_max_length=args.edge_label_max_length,
                                             ignore_groups=args.ignore_groups,
-                                            using_http=using_http)
+                                            using_http=using_http,
+                                            vis_group_keys=list(self.graph_notebook_vis_options.get('groups', {}).keys()))
 
                         if using_http and 'path()' in cell and query_res and isinstance(query_res, list):
                             first_path = query_res[0]

src/graph_notebook/network/gremlin/GremlinNetwork.py

@@ -110,7 +110,7 @@ class GremlinNetwork(EventfulNetwork):
     def __init__(self, graph: MultiDiGraph = None, callbacks=None, label_max_length=DEFAULT_LABEL_MAX_LENGTH,
                  edge_label_max_length=DEFAULT_LABEL_MAX_LENGTH, group_by_property=None, display_property=None,
                  edge_display_property=None, tooltip_property=None, edge_tooltip_property=None, ignore_groups=False,
-                 group_by_depth=False, group_by_raw=False, using_http=False):
+                 group_by_depth=False, group_by_raw=False, using_http=False, vis_group_keys=None):
         if graph is None:
             graph = MultiDiGraph()
         if group_by_depth:
@@ -121,10 +121,29 @@ def __init__(self, graph: MultiDiGraph = None, callbacks=None, label_max_length=
             display_property = 'label' if using_http else T_LABEL
         if not edge_display_property:
             edge_display_property = 'label' if using_http else T_LABEL
+        self.vis_group_keys = set(vis_group_keys) if vis_group_keys else set()
         super().__init__(graph, callbacks, label_max_length, edge_label_max_length, group_by_property,
                          display_property, edge_display_property, tooltip_property, edge_tooltip_property,
                          ignore_groups, group_by_raw)
 
+    def _resolve_multilabel_group(self, label):
+        """For multi-label vertices (:: separated), check if any individual label
+        component matches a defined vis options group key. Returns the first match,
+        or the full label if no match is found."""
+        if not self.vis_group_keys or '::' not in label:
+            return label
+        if label in self.vis_group_keys:
+            return label
+        # Check if any vis group key has the same components (order-independent)
+        label_components = set(label.split('::'))
+        for key in self.vis_group_keys:
+            if '::' in key and set(key.split('::')) == label_components:
+                return key
+        for component in label.split('::'):
+            if component in self.vis_group_keys:
+                return component
+        return label
+
     def get_dict_element_property_value(self, element, k, temp_label, custom_property):
         property_value = None
         if isinstance(temp_label, list):
@@ -399,7 +418,7 @@ def add_vertex(self, v, path_index: int = -1):
                     # This sets the group key to the label if either "label" is passed in or
                     # T.label is set in order to handle the default case of grouping by label
                     # when no explicit key is specified
-                    group = v.label
+                    group = self._resolve_multilabel_group(v.label)
                 elif str(self.group_by_property) in [T_ID, 'id', '~id']:
                     group = v.id
                 elif self.group_by_property == DEPTH_GRP_KEY:
@@ -408,17 +427,26 @@ def add_vertex(self, v, path_index: int = -1):
                     group = DEFAULT_GRP
             else:  # handle dict format group_by
                 try:
+                    # For multi-label vertices, try the full label first, then individual components
+                    matched_label = None
                     if str(v.label) in self.group_by_property:
-                        if self.group_by_property[str(v.label)] == DEFAULT_RAW_GRP_KEY:
+                        matched_label = str(v.label)
+                    elif '::' in v.label:
+                        for component in v.label.split('::'):
+                            if component in self.group_by_property:
+                                matched_label = component
+                                break
+                    if matched_label is not None:
+                        if self.group_by_property[matched_label] == DEFAULT_RAW_GRP_KEY:
                             group = str(v)
-                        elif self.group_by_property[str(v.label)] in [T_LABEL, 'label']:
-                            group = v.label
-                        elif self.group_by_property[str(v.label)] in [T_ID, 'id', '~id']:
+                        elif self.group_by_property[matched_label] in [T_LABEL, 'label']:
+                            group = self._resolve_multilabel_group(v.label)
+                        elif self.group_by_property[matched_label] in [T_ID, 'id', '~id']:
                             group = v.id
-                        elif self.group_by_property[str(v.label)] == DEPTH_GRP_KEY:
+                        elif self.group_by_property[matched_label] == DEPTH_GRP_KEY:
                             group = depth_group
                         else:
-                            group = vertex_dict[self.group_by_property[str(v.label)]]
+                            group = vertex_dict[self.group_by_property[matched_label]]
                     else:
                         group = DEFAULT_GRP
                 except KeyError:
@@ -515,7 +543,10 @@ def add_vertex(self, v, path_index: int = -1):
                             group = depth_group
                             group_is_set = True
                         elif str(k) == self.group_by_property:
-                            group = str(v[k])
+                            if str(self.group_by_property) in ['label', T_LABEL] and isinstance(v[k], str) and '::' in v[k]:
+                                group = self._resolve_multilabel_group(v[k])
+                            else:
+                                group = str(v[k])
                             group_is_set = True
                 if not display_is_set:
                     label_property_raw_value = self.get_dict_element_property_value(v, k, label_raw,

src/graph_notebook/notebooks/01-Neptune-Database/02-Visualization/Grouping-and-Appearance-Customization-Gremlin.ipynb

@@ -1107,6 +1107,61 @@
     "g.V().hasId('10000').drop()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "multilabel-grouping-section",
+   "metadata": {},
+   "source": [
+    "### Grouping Multi-Label Vertices\n",
+    "\n",
+    "In Amazon Neptune, a vertex can have multiple labels by using the `::` separator when creating it. For example:\n",
+    "\n",
+    "```\n",
+    "g.addV('TAG::TAG_ANIMAL_CONTROL').property('name', 'animal_control')\n",
+    "```\n",
+    "\n",
+    "This creates a vertex with two labels: `TAG` and `TAG_ANIMAL_CONTROL`.\n",
+    "\n",
+    "When defining groups in your visualization options, you can target any individual label component. ",
+    "The graph notebook will check each label component against your defined group keys and use the first match. ",
+    "If no component matches, the full label string is used as the group.\n",
+    "\n",
+    "For example, with the following vis options, the multi-label vertex above will match the `TAG` group:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "multilabel-grouping-vis-options",
+   "metadata": {},
+   "source": [
+    "%%graph_notebook_vis_options\n",
+    "{\n",
+    "  \"groups\": {\n",
+    "    \"TAG\": {\n",
+    "        \"shape\": \"star\",\n",
+    "        \"color\": \"#00FF00\"\n",
+    "    },\n",
+    "    \"AGENT\": {\n",
+    "        \"shape\": \"triangle\",\n",
+    "        \"color\": \"#0000FF\"\n",
+    "    }\n",
+    "  }\n",
+    "}"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "id": "multilabel-grouping-note",
+   "metadata": {},
+   "source": [
+    "With this configuration:\n",
+    "- A vertex with label `TAG::TAG_ANIMAL_CONTROL` will match the `TAG` group (green star)\n",
+    "- A vertex with label `AGENT` will match the `AGENT` group (blue triangle)\n",
+    "- A vertex with label `FOO::BAR` (no matching group key) will use the default appearance"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "79cfb095",
@@ -1139,4 +1194,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

src/graph_notebook/widgets/package.json

@@ -61,7 +61,7 @@
     "stylelint-prettier": "^4.0.0",
     "ts-loader": "^9.4.4",
     "typescript": "~5.0.4",
-    "webpack": "^5.88.2",
+    "webpack": ">=5.88.2 <5.107.0",
     "webpack-cli": "^5.1.4",
     "webpack-dev-server": "^4.15.1",
     "@types/backbone": "^1.4.15"

test/unit/network/gremlin/test_gremlin_network.py

@@ -3014,6 +3014,126 @@ def test_add_path_with_pattern_groupby_depth(self):
         self.assertEqual(v1_data['group'], '__DEPTH-2__')
         self.assertEqual(v2_data['group'], '__DEPTH-4__')
 
+    # Multi-label group matching tests (Issue #746)
+
+    def test_multilabel_vertex_matches_first_label_component(self):
+        """Vertex with :: separator should match group key for first label component"""
+        vertex = Vertex(id='1', label='TAG::TAG_ANIMAL_CONTROL')
+        gn = GremlinNetwork(vis_group_keys=['TAG', 'AGENT'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG')
+
+    def test_multilabel_vertex_matches_second_label_component(self):
+        """Vertex with :: separator should match group key for second label component"""
+        vertex = Vertex(id='1', label='TAG::TAG_ANIMAL_CONTROL')
+        gn = GremlinNetwork(vis_group_keys=['TAG_ANIMAL_CONTROL', 'AGENT'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG_ANIMAL_CONTROL')
+
+    def test_multilabel_vertex_no_match_falls_back_to_full_label(self):
+        """Vertex with :: separator falls back to full label when no component matches"""
+        vertex = Vertex(id='1', label='TAG::TAG_ANIMAL_CONTROL')
+        gn = GremlinNetwork(vis_group_keys=['AGENT', 'PERSON'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG::TAG_ANIMAL_CONTROL')
+
+    def test_multilabel_vertex_no_vis_group_keys_uses_full_label(self):
+        """Without vis_group_keys, multi-label vertex uses full label as group"""
+        vertex = Vertex(id='1', label='TAG::TAG_ANIMAL_CONTROL')
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG::TAG_ANIMAL_CONTROL')
+
+    def test_single_label_vertex_unaffected_by_vis_group_keys(self):
+        """Single-label vertex behavior unchanged with vis_group_keys"""
+        vertex = Vertex(id='1', label='AGENT')
+        gn = GremlinNetwork(vis_group_keys=['AGENT', 'TAG'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'AGENT')
+
+    def test_multilabel_dict_vertex_matches_component(self):
+        """Dict vertex with :: string label should match group key for individual component"""
+        vertex = {
+            'id': '1',
+            'label': 'TAG_ANIMAL_CONTROL::TAG',
+            'name': 'animal_control'
+        }
+        gn = GremlinNetwork(using_http=True, vis_group_keys=['TAG', 'AGENT'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG')
+
+    def test_multilabel_dict_vertex_no_match_falls_back(self):
+        """Dict vertex with :: string label falls back to full label when no component matches"""
+        vertex = {
+            'id': '1',
+            'label': 'TAG_ANIMAL_CONTROL::TAG',
+            'name': 'animal_control'
+        }
+        gn = GremlinNetwork(using_http=True, vis_group_keys=['AGENT'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG_ANIMAL_CONTROL::TAG')
+
+    def test_multilabel_dict_format_groupby_matches_component(self):
+        """Dict format group_by should match multi-label vertex by component"""
+        vertex = Vertex(id='1', label='TAG::TAG_ANIMAL_CONTROL')
+        gn = GremlinNetwork(group_by_property='{"TAG":"label"}', vis_group_keys=['TAG'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG')
+
+    def test_multilabel_vertex_full_label_group_key_order_independent(self):
+        """Full multi-label group key should match regardless of component order"""
+        vertex = Vertex(id='1', label='TAG_ANIMAL_CONTROL::TAG')
+        gn = GremlinNetwork(vis_group_keys=['TAG::TAG_ANIMAL_CONTROL', 'AGENT'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG::TAG_ANIMAL_CONTROL')
+
+    def test_multilabel_dict_vertex_full_label_group_key_order_independent(self):
+        """Dict vertex: full multi-label group key should match regardless of component order"""
+        vertex = {
+            T.id: '1',
+            T.label: 'TAG_ANIMAL_CONTROL::TAG',
+            'name': 'animal_control'
+        }
+        gn = GremlinNetwork(vis_group_keys=['TAG::TAG_ANIMAL_CONTROL', 'AGENT'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG::TAG_ANIMAL_CONTROL')
+
+    def test_multilabel_full_label_key_takes_priority_over_component(self):
+        """Full multi-label group key should take priority over individual component match"""
+        vertex = Vertex(id='1', label='TAG_ANIMAL_CONTROL::TAG')
+        gn = GremlinNetwork(vis_group_keys=['TAG::TAG_ANIMAL_CONTROL', 'TAG', 'AGENT'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertNotEqual(node['group'], 'TAG')
+        self.assertEqual(node['group'], 'TAG::TAG_ANIMAL_CONTROL')
+
+    def test_multilabel_component_match_without_full_label_key(self):
+        """After reset, only current group keys should be used for matching"""
+        vertex = Vertex(id='1', label='TAG_ANIMAL_CONTROL::TAG')
+
+        # Before reset: full label key matches
+        gn = GremlinNetwork(vis_group_keys=['TAG::TAG_ANIMAL_CONTROL', 'TAG', 'AGENT'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertEqual(node['group'], 'TAG::TAG_ANIMAL_CONTROL')
+
+        # Simulate reset: new network with only component keys
+        gn = GremlinNetwork(vis_group_keys=['TAG', 'AGENT'])
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get('1')
+        self.assertNotEqual(node['group'], 'TAG::TAG_ANIMAL_CONTROL')
+        self.assertEqual(node['group'], 'TAG')
+
 
 if __name__ == '__main__':
     unittest.main()