Context property bag ("state") (#1774)

* Add pipeline state property bag to run context

* Move state creation out of context util

* Move callbacks into PipelineRunContext

* Semver

* Rename state.json to context.json to avoid confusion with stats.json

* Expand smoke test row count

* Add util to create storage and cache

Author: natoverse

.semversioner/next-release/patch-20250228005421960221.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Adds general-purpose pipeline run state object."
+}

graphrag/api/index.py

@@ -14,10 +14,10 @@
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
 from graphrag.config.enums import IndexingMethod
 from graphrag.config.models.graph_rag_config import GraphRagConfig
-from graphrag.index.run.pipeline_run_result import PipelineRunResult
 from graphrag.index.run.run_pipeline import run_pipeline
 from graphrag.index.run.utils import create_callback_chain
-from graphrag.index.typing import WorkflowFunction
+from graphrag.index.typing.pipeline_run_result import PipelineRunResult
+from graphrag.index.typing.workflow import WorkflowFunction
 from graphrag.index.workflows.factory import PipelineFactory
 from graphrag.logger.base import ProgressLogger
 from graphrag.logger.null_progress import NullProgressLogger

graphrag/callbacks/noop_workflow_callbacks.py

@@ -4,7 +4,7 @@
 """A no-op implementation of WorkflowCallbacks."""
 
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
-from graphrag.index.run.pipeline_run_result import PipelineRunResult
+from graphrag.index.typing.pipeline_run_result import PipelineRunResult
 from graphrag.logger.progress import Progress
 
 

graphrag/callbacks/workflow_callbacks.py

@@ -5,7 +5,7 @@
 
 from typing import Protocol
 
-from graphrag.index.run.pipeline_run_result import PipelineRunResult
+from graphrag.index.typing.pipeline_run_result import PipelineRunResult
 from graphrag.logger.progress import Progress
 
 

graphrag/callbacks/workflow_callbacks_manager.py

@@ -4,7 +4,7 @@
 """A module containing the WorkflowCallbacks registry."""
 
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
-from graphrag.index.run.pipeline_run_result import PipelineRunResult
+from graphrag.index.typing.pipeline_run_result import PipelineRunResult
 from graphrag.logger.progress import Progress
 
 

graphrag/cli/query.py

@@ -13,7 +13,7 @@
 from graphrag.config.load_config import load_config
 from graphrag.config.models.graph_rag_config import GraphRagConfig
 from graphrag.logger.print_progress import PrintProgressLogger
-from graphrag.storage.factory import StorageFactory
+from graphrag.utils.api import create_storage_from_config
 from graphrag.utils.storage import load_table_from_storage, storage_has_table
 
 if TYPE_CHECKING:
@@ -497,10 +497,7 @@ def _resolve_output_files(
         dataframe_dict["num_indexes"] = len(config.outputs)
         dataframe_dict["index_names"] = config.outputs.keys()
         for output in config.outputs.values():
-            output_config = output.model_dump()
-            storage_obj = StorageFactory().create_storage(
-                storage_type=output_config["type"], kwargs=output_config
-            )
+            storage_obj = create_storage_from_config(output)
             for name in output_list:
                 if name not in dataframe_dict:
                     dataframe_dict[name] = []
@@ -527,10 +524,7 @@ def _resolve_output_files(
         return dataframe_dict
     # Loading output files for single-index search
     dataframe_dict["multi-index"] = False
-    output_config = config.output.model_dump()
-    storage_obj = StorageFactory().create_storage(
-        storage_type=output_config["type"], kwargs=output_config
-    )
+    storage_obj = create_storage_from_config(config.output)
     for name in output_list:
         df_value = asyncio.run(load_table_from_storage(name=name, storage=storage_obj))
         dataframe_dict[name] = df_value

graphrag/index/context.py

@@ -13,7 +13,7 @@
 from graphrag.index.operations.build_noun_graph.np_extractors.base import (
     BaseNounPhraseExtractor,
 )
-from graphrag.index.run.derive_from_rows import derive_from_rows
+from graphrag.index.utils.derive_from_rows import derive_from_rows
 from graphrag.index.utils.hashing import gen_sha512_hash
 
 

graphrag/index/operations/build_noun_graph/build_noun_graph.py

@@ -11,7 +11,7 @@
 import tiktoken
 
 from graphrag.config.defaults import ENCODING_MODEL, graphrag_config_defaults
-from graphrag.index.typing import ErrorHandlerFn
+from graphrag.index.typing.error_handler import ErrorHandlerFn
 from graphrag.language_model.protocol.base import ChatModel
 from graphrag.prompts.index.extract_claims import (
     CONTINUE_PROMPT,

graphrag/index/operations/extract_covariates/claim_extractor.py

@@ -20,7 +20,7 @@
     Covariate,
     CovariateExtractionResult,
 )
-from graphrag.index.run.derive_from_rows import derive_from_rows
+from graphrag.index.utils.derive_from_rows import derive_from_rows
 from graphrag.language_model.manager import ModelManager
 
 log = logging.getLogger(__name__)