awslabs · KilianTrunk · Aug 17, 2026
diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/chunk_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/chunk_graph_builder.py
@@ -5,7 +5,7 @@
 from typing import Any
 
 from graphrag_toolkit.lexical_graph.config import GraphRAGConfig
-from graphrag_toolkit.lexical_graph.storage.graph import GraphStore
+from graphrag_toolkit.lexical_graph.storage.graph import GraphQueryOperation, GraphStore
 from graphrag_toolkit.lexical_graph.storage.chunk_store_factory import ChunkStoreFactory
 from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder
 
@@ -112,7 +112,13 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any):
 
                 query_c = '\n'.join(statements_c)
 
-                graph_client.execute_query_with_retry(query_c, self._to_params(properties_c), max_attempts=5, max_wait=7)
+                graph_client.execute_query_with_retry(
+                    query_c,
+                    self._to_params(properties_c),
+                    max_attempts=5,
+                    max_wait=7,
+                    operation=GraphQueryOperation.UPSERT_CHUNK,
+                )
 
             source_info = node.relationships.get(NodeRelationship.SOURCE, None)
 
@@ -135,7 +141,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any):
 
                 query_s = '\n'.join(statements_s)
 
-                graph_client.execute_query_with_retry(query_s, self._to_params(properties_s), max_attempts=5, max_wait=7)
+                graph_client.execute_query_with_retry(query_s, self._to_params(properties_s), max_attempts=5, max_wait=7, operation=GraphQueryOperation.LINK_CHUNK_SOURCE)
 
             else:
                 logger.warning(f'source_id missing from chunk node [node_id: {chunk_id}]')
@@ -153,12 +159,13 @@ def insert_chunk_to_chunk_relationship(node_id:str, relationship_type:str):
 
                 properties_c2c = {
                     'chunk_id': chunk_id,
-                    'target_id': node_id
+                    'target_id': node_id,
+                    '_relationship_type': relationship_type,
                 }
 
                 query_c2c = '\n'.join(statements_c2c)
 
-                graph_client.execute_query_with_retry(query_c2c, self._to_params(properties_c2c), max_attempts=5, max_wait=7)
+                graph_client.execute_query_with_retry(query_c2c, self._to_params(properties_c2c), max_attempts=5, max_wait=7, operation=GraphQueryOperation.LINK_CHUNKS)
 
 
             for node_relationship,relationship_info in node.relationships.items():

diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_graph_builder.py
@@ -5,7 +5,7 @@
 from typing import Any
 
 from graphrag_toolkit.lexical_graph.indexing.model import Fact, Entity
-from graphrag_toolkit.lexical_graph.storage.graph import GraphStore
+from graphrag_toolkit.lexical_graph.storage.graph import GraphQueryOperation, GraphStore
 from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import search_string_from, label_from, new_query_var, escape_cypher_label
 from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder
 from graphrag_toolkit.lexical_graph.indexing.constants import DEFAULT_CLASSIFICATION, LOCAL_ENTITY_CLASSIFICATION
@@ -98,7 +98,7 @@ def insert_for_entity(entity:Entity):
 
                 query = '\n'.join(statements)
 
-                graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7)
+                graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7, operation=GraphQueryOperation.UPSERT_ENTITY)
 
             insert_for_entity(fact.subject)
 
@@ -125,7 +125,17 @@ def insert_domain_entity(entity:Entity):
                     e_label = escape_cypher_label(label_from(entity.classification or DEFAULT_CLASSIFICATION))
                     e_comment = f'// awsqid:{e_id}-{e_label}'.replace('\r', ' ').replace('\n', ' ')
                     query_e = f"UNWIND $params AS params MERGE ({e_var}:`__Entity__`{{{graph_client.node_id('entityId')}: params.entityId}}) SET {e_var} :`{e_label}` {e_comment}"
-                    graph_client.execute_query_with_retry(query_e, self._to_params({'entityId': e_id}), max_attempts=5, max_wait=7)
+                    properties = self._to_params({
+                        'entityId': e_id,
+                        '_classification': entity.classification or DEFAULT_CLASSIFICATION,
+                    })
+                    graph_client.execute_query_with_retry(
+                        query_e,
+                        properties,
+                        max_attempts=5,
+                        max_wait=7,
+                        operation=GraphQueryOperation.ADD_ENTITY_TYPE,
+                    )
 
                 insert_domain_entity(fact.subject)
 
@@ -136,4 +146,4 @@ def insert_domain_entity(entity:Entity):
                     insert_domain_entity(fact.complement)
 
         else:
-            logger.warning(f'fact_id missing from fact node [node_id: {node.node_id}]')
+            logger.warning(f'fact_id missing from fact node [node_id: {node.node_id}]')
diff --git a/...-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_relation_graph_builder.py b/...-graph/src/graphrag_toolkit/lexical_graph/indexing/build/entity_relation_graph_builder.py
@@ -5,7 +5,7 @@
 from typing import Any
 
 from graphrag_toolkit.lexical_graph.indexing.model import Fact
-from graphrag_toolkit.lexical_graph.storage.graph import GraphStore
+from graphrag_toolkit.lexical_graph.storage.graph import GraphQueryOperation, GraphStore
 from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import relationship_name_from, new_query_var
 from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder
 from graphrag_toolkit.lexical_graph.indexing.utils.fact_utils import string_complement_to_entity
@@ -90,7 +90,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any):
 
                 query = '\n'.join(statements)
 
-                graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7)
+                graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7, operation=GraphQueryOperation.LINK_ENTITIES)
 
                 # if include_domain_labels:
 
@@ -137,7 +137,7 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any):
 
                 query = '\n'.join(statements)
 
-                graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7)
+                graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7, operation=GraphQueryOperation.LINK_ENTITIES)
 
                 # if include_domain_labels:
 

diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/fact_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/fact_graph_builder.py
@@ -5,7 +5,7 @@
 from typing import Any, Optional
 
 from graphrag_toolkit.lexical_graph.indexing.model import Fact
-from graphrag_toolkit.lexical_graph.storage.graph import GraphStore, Query, QueryTree
+from graphrag_toolkit.lexical_graph.storage.graph import GraphQueryOperation, GraphStore, Query, QueryTree
 from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder
 from graphrag_toolkit.lexical_graph.indexing.constants import LOCAL_ENTITY_CLASSIFICATION
 from graphrag_toolkit.lexical_graph.indexing.utils.fact_utils import string_complement_to_entity
@@ -82,15 +82,26 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any):
                 'MERGE (fact)-[:`__SUPPORTS__`]->(statement)'
             ]
 
+            subject_literal = None
+            if fact.subject.classification == LOCAL_ENTITY_CLASSIFICATION and not include_local_entities:
+                subject_literal = fact.subject.value
+
+            object_literal = None
+            if not fact.object and fact.complement and not include_local_entities:
+                object_literal = fact.complement.value
+
             properties = {
                 'statement_id': fact.statementId,
                 'fact_id': fact.factId,
-                'fact': node.text
+                'fact': node.text,
+                '_predicate': fact.predicate.value,
+                '_subject_literal': subject_literal,
+                '_object_literal': object_literal
             }
 
             query = '\n'.join(statements)
 
-            graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7)
+            graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=5, max_wait=7, operation=GraphQueryOperation.UPSERT_FACT)
 
             def insert_entity_fact_relationship(relationship_type:str, entity_id:Optional[str]=None):
 
@@ -107,10 +118,11 @@ def insert_entity_fact_relationship(relationship_type:str, entity_id:Optional[st
                 if entity_id:
                     properties_e2f['fact_id'] = fact.factId
                     properties_e2f['entity_id'] = entity_id
+                    properties_e2f['_relationship_type'] = relationship_type
 
                 query_e2f = '\n'.join(statements_e2f)
 
-                graph_client.execute_query_with_retry(query_e2f, self._to_params(properties_e2f), max_attempts=5, max_wait=7)
+                graph_client.execute_query_with_retry(query_e2f, self._to_params(properties_e2f), max_attempts=5, max_wait=7, operation=GraphQueryOperation.LINK_FACT_ENTITY)
 
 
             insert_entity_fact_relationship('subject')

diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_batch_client.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_batch_client.py
@@ -46,6 +46,7 @@ def __init__(self, graph_client:GraphStore, batch_writes_enabled:bool, batch_wri
         self.batch_writes_enabled = batch_writes_enabled
         self.batch_write_size = batch_write_size
         self.batches:Dict[str, List] = {}
+        self.batch_operations:Dict[str, Any] = {}
         self.query_trees:Dict[str, QueryTree] = {}
         self.all_nodes = []
         self.parameterless_queries:Dict[str, str] = {}
@@ -138,8 +139,9 @@ def execute_query_with_retry(self, query:QueryTree, properties:Dict[str, Any], *
                 if properties:
                     if query not in self.batches:
                         self.batches[query] = []
+                        self.batch_operations[query] = kwargs.get('operation')
                     self.batches[query].extend(properties['params'])
-                else:
+                elif kwargs.get('operation') is None:
                     self._add_parameterless_query(query)
             elif isinstance(query, QueryTree):
                 properties = properties or {'params':[]}
@@ -224,7 +226,7 @@ def _apply_batch_query(self, query, parameters):
                 'params': p
             }
             try:
-                self.graph_client.execute_query_with_retry(query, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT)
+                self._execute_batch(query, params)
             except Exception as e:
                 logger.debug(f'Batch failed - queuing for retry: [query: {query}, params: {params}]')
                 retry_batches.append((query, params))
@@ -233,20 +235,29 @@ def _apply_batch_query(self, query, parameters):
 
         for (query, params) in retry_batches:
             try:
-                self.graph_client.execute_query_with_retry(query, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT)
+                self._execute_batch(query, params)
             except Exception as e:
                 logger.debug(f'Retry batch failed - queuing for return: [query: {query}, params: {params}]')
                 failed_batches.append((query, params))
 
         return failed_batches
+
+    def _execute_batch(self, query, parameters):
+        return self.graph_client.execute_query_with_retry(
+            query,
+            parameters,
+            max_attempts=BATCH_MAX_ATTEMPTS,
+            max_wait=BATCH_MAX_WAIT,
+            operation=self.batch_operations.get(query),
+        )
 
     def _retry_failed_batches(self, failed_batches):
 
         last_chance_batches = []
 
         for (query, params) in failed_batches:
             try:
-                self.graph_client.execute_query_with_retry(query, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT)
+                self._execute_batch(query, params)
             except Exception as e:
                 logger.debug(f'Retry failed batch failed - queuing for individual writes retry: [query: {query}, params: {params}]')
                 last_chance_batches.append((query, params))
@@ -259,7 +270,7 @@ def _retry_failed_batches(self, failed_batches):
                     'params': [p]
                 }
                 try:
-                    self.graph_client.execute_query_with_retry(query, single_params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT)
+                    self._execute_batch(query, single_params)
                 except Exception as e:
                     logger.error(f'Failed single write: [query: {query}, params: {params}, error: {str(e)}]') 
                     raise e    
@@ -269,7 +280,7 @@ def _apply_batch_query_tree(self, query_tree_id, parameters):
 
         query_tree = self.query_trees[query_tree_id]
 
-        def graph_store_op(q, p):
+        def graph_store_op(q, p, **kwargs):
 
             all_params = p['params']
 
@@ -284,7 +295,7 @@ def graph_store_op(q, p):
                     'params': chunk
                 }
 
-                results = self.graph_client.execute_query_with_retry(q, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT)
+                results = self.graph_client.execute_query_with_retry(q, params, max_attempts=BATCH_MAX_ATTEMPTS, max_wait=BATCH_MAX_WAIT, **kwargs)
 
                 for r in results:
                     yield r

diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_summary_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/graph_summary_builder.py
@@ -5,7 +5,7 @@
 from typing import Any
 
 from graphrag_toolkit.lexical_graph.indexing.model import Fact
-from graphrag_toolkit.lexical_graph.storage.graph import GraphStore
+from graphrag_toolkit.lexical_graph.storage.graph import GraphQueryOperation, GraphStore
 from graphrag_toolkit.lexical_graph.storage.graph.graph_utils import label_from, relationship_name_from
 from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder
 from graphrag_toolkit.lexical_graph.indexing.constants import DEFAULT_CLASSIFICATION, LOCAL_ENTITY_CLASSIFICATION
@@ -116,7 +116,7 @@ def build(self, node:BaseNode, graph_client:GraphStore, **kwargs:Any):
 
                 query = '\n'.join(statements)
 
-                graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=10, max_wait=10)
+                graph_client.execute_query_with_retry(query, self._to_params(properties), max_attempts=10, max_wait=10, operation=GraphQueryOperation.UPDATE_GRAPH_SUMMARY)
 
         else:
             logger.warning(f'fact_id missing from fact node [node_id: {node.node_id}]')
diff --git a/.../src/graphrag_toolkit/lexical_graph/indexing/build/local_entity_rewrites_graph_builder.py b/.../src/graphrag_toolkit/lexical_graph/indexing/build/local_entity_rewrites_graph_builder.py
@@ -5,7 +5,7 @@
 from typing import Any
 
 from graphrag_toolkit.lexical_graph.indexing.model import Fact
-from graphrag_toolkit.lexical_graph.storage.graph import GraphStore, Query, QueryTree
+from graphrag_toolkit.lexical_graph.storage.graph import GraphQueryOperation, GraphStore, Query, QueryTree
 from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder
 from graphrag_toolkit.lexical_graph.indexing.utils.fact_utils import string_complement_to_entity
 from graphrag_toolkit.lexical_graph.indexing.constants import LOCAL_ENTITY_CLASSIFICATION
@@ -43,7 +43,8 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any):
                 WHERE {graph_client.node_id('n.entityId')} = params.n_id AND {graph_client.node_id('c.entityId')} = params.c_id
                 MERGE (s)-[:`__RELATION__`{{value:r.value}}]->(n)
                 MERGE (n)-[:`__OBJECT__`]->(f)
-                """
+                """,
+                operation=GraphQueryOperation.COPY_COMPLEMENT_RELATIONSHIPS,
             )
 
             delete_complement_relationships = Query(
@@ -54,7 +55,8 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any):
                 DELETE r1
                 DELETE r2
                 DETACH DELETE c
-                """
+                """,
+                operation=GraphQueryOperation.DELETE_COMPLEMENT,
             )
 
             if fact.subject:
@@ -72,7 +74,8 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any):
                     child_queries=[
                         copy_complement_relationships_to_subject, 
                         delete_complement_relationships
-                    ]
+                    ],
+                    operation=GraphQueryOperation.FIND_COMPLEMENTS,
                 )
 
                 params = {
@@ -98,7 +101,8 @@ def build(self, node:BaseNode, graph_client: GraphStore, **kwargs:Any):
                     child_queries=[
                         copy_complement_relationships_to_subject, 
                         delete_complement_relationships
-                    ]
+                    ],
+                    operation=GraphQueryOperation.FIND_SUBJECTS,
                 )
 
                 params = {

diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/source_graph_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/build/source_graph_builder.py
@@ -4,7 +4,7 @@
 import logging
 from typing import Any
 
-from graphrag_toolkit.lexical_graph.storage.graph import GraphStore
+from graphrag_toolkit.lexical_graph.storage.graph import GraphQueryOperation, GraphStore
 from graphrag_toolkit.lexical_graph.indexing.build.graph_builder import GraphBuilder
 from graphrag_toolkit.lexical_graph.versioning import VALID_FROM, VALID_TO, VERSION_INDEPENDENT_ID_FIELDS
 from graphrag_toolkit.lexical_graph.versioning import EXTRACT_TIMESTAMP, BUILD_TIMESTAMP, PREV_VERSIONS
@@ -111,7 +111,8 @@ def format_assigment(key):
 
             query = '\n'.join(statements)
 
-            graph_client.execute_query_with_retry(query, self._to_params(clean_metadata))
+            clean_metadata['_source_id'] = source_id
+            graph_client.execute_query_with_retry(query, self._to_params(clean_metadata), operation=GraphQueryOperation.UPSERT_SOURCE)
 
             # prev_source_ids = source_metadata.get('prev_versions', [])