fix state management

Author: andresmor-ms

graphrag/index/workflows/__init__.py

@@ -42,6 +42,9 @@
 from .prune_graph import (
     run_workflow as run_prune_graph,
 )
+from .update_clean_state import (
+    run_workflow as run_update_clean_state,
+)
 from .update_communities import (
     run_workflow as run_update_communities,
 )
@@ -85,4 +88,5 @@
     "update_communities": run_update_communities,
     "update_covariates": run_update_covariates,
     "update_text_units": run_update_text_units,
+    "update_clean_state": run_update_clean_state,
 })

graphrag/index/workflows/factory.py

@@ -53,6 +53,7 @@ def _get_workflows_list(
         "update_communities",
         "update_community_reports",
         "update_text_embeddings",
+        "update_clean_state",
     ]
     if config.workflows:
         return config.workflows + (update_workflows if is_update_run else [])

graphrag/index/workflows/update_clean_state.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing run_workflow method definition."""
+
+import logging
+
+from graphrag.config.models.graph_rag_config import GraphRagConfig
+from graphrag.index.typing.context import PipelineRunContext
+from graphrag.index.typing.workflow import WorkflowFunctionOutput
+
+logger = logging.getLogger(__name__)
+
+
+async def run_workflow(
+    config: GraphRagConfig,
+    context: PipelineRunContext,
+) -> WorkflowFunctionOutput:
+    """Clean the state after the update."""
+    logger.info("Cleaning State")
+    keys_to_delete = [key_name for key_name in context.state if key_name.startswith("incremental_update_")]
+    
+    for key_name in keys_to_delete:
+        del context.state[key_name]
+
+    return WorkflowFunctionOutput(result=None)