Add new param to cli for new pipeline

Author: natoverse

graphrag/api/index.py

@@ -8,6 +8,8 @@
 Backwards compatibility is not guaranteed at this time.
 """
 
+import logging
+
 from datashaper import WorkflowCallbacks
 
 from graphrag.cache.noop_pipeline_cache import NoopPipelineCache
@@ -20,6 +22,8 @@
 from graphrag.index.typing import PipelineRunResult
 from graphrag.logger.base import ProgressLogger
 
+log = logging.getLogger(__name__)
+
 
 async def build_index(
     config: GraphRagConfig,
@@ -28,7 +32,7 @@ async def build_index(
     memory_profile: bool = False,
     callbacks: list[WorkflowCallbacks] | None = None,
     progress_logger: ProgressLogger | None = None,
-    use_new_pipeline: bool = False,
+    new_pipeline: bool = False,
 ) -> list[PipelineRunResult]:
     """Run the pipeline with the given configuration.
 
@@ -67,7 +71,8 @@ async def build_index(
     callbacks.append(create_pipeline_reporter(config.reporting, None))  # type: ignore
     outputs: list[PipelineRunResult] = []
 
-    if use_new_pipeline:
+    if new_pipeline:
+        log.info("RUNNING NEW WORKFLOWS WITHOUT DATASHAPER")
         async for output in run_workflows(
             config,
             cache=pipeline_cache,
@@ -83,6 +88,7 @@ async def build_index(
                     progress_logger.success(output.workflow)
                 progress_logger.info(str(output.result))
     else:
+        log.info("RUNNING ORIGINAL PIPELINE")
         pipeline_config = create_pipeline_config(config)
         async for output in run_pipeline_with_config(
             pipeline_config,

graphrag/cli/index.py

@@ -74,6 +74,7 @@ def index_cli(
     dry_run: bool,
     skip_validation: bool,
     output_dir: Path | None,
+    new_pipeline: bool,
 ):
     """Run the pipeline with the given config."""
     config = load_config(root_dir, config_filepath)
@@ -88,6 +89,7 @@ def index_cli(
         dry_run=dry_run,
         skip_validation=skip_validation,
         output_dir=output_dir,
+        new_pipeline=new_pipeline,
     )
 
 
@@ -124,6 +126,7 @@ def update_cli(
         dry_run=False,
         skip_validation=skip_validation,
         output_dir=output_dir,
+        new_pipeline=False,
     )
 
 
@@ -137,6 +140,7 @@ def _run_index(
     dry_run,
     skip_validation,
     output_dir,
+    new_pipeline,
 ):
     progress_logger = LoggerFactory().create_logger(logger)
     info, error, success = _logger(progress_logger)
@@ -182,7 +186,7 @@ def _run_index(
             is_resume_run=bool(resume),
             memory_profile=memprofile,
             progress_logger=progress_logger,
-            use_new_pipeline=True,
+            new_pipeline=new_pipeline,
         )
     )
     encountered_errors = any(

graphrag/cli/main.py

@@ -170,6 +170,7 @@ def _index_cli(
             resolve_path=True,
         ),
     ] = None,
+    new: Annotated[bool, typer.Option(help="Use the new pipeline.")] = False,
 ):
     """Build a knowledge graph index."""
     from graphrag.cli.index import index_cli
@@ -185,6 +186,7 @@ def _index_cli(
         dry_run=dry_run,
         skip_validation=skip_validation,
         output_dir=output,
+        new_pipeline=new,
     )
 
 

graphrag/index/run/run_workflows.py

@@ -58,7 +58,6 @@ async def run_workflows(
     run_id: str | None = None,
 ) -> AsyncIterable[PipelineRunResult]:
     """Run all workflows using a simplified pipeline."""
-    log.info("RUNNING NEW WORKFLOWS WITHOUT DATASHAPER")
     start_time = time.time()
 
     run_id = run_id or time.strftime("%Y%m%d-%H%M%S")