Pipeline callbacks (#1729)

* Add pipeline_start and pipeline_end callbacks

* Collapse redundant callback/logger logic

* Remove redundant reporting config classes

* Remove a few out-of-date type ignores

* Semver

---------

Co-authored-by: Alonso Guevara <alonsog@microsoft.com>

Author: natoverse

.semversioner/next-release/minor-20250221001507100298.json

@@ -0,0 +1,4 @@
+{
+  "type": "minor",
+  "description": "Add pipeline_start and pipeline_end callbacks."
+}

graphrag/api/index.py

@@ -10,15 +10,17 @@
 
 import logging
 
-from graphrag.cache.noop_pipeline_cache import NoopPipelineCache
 from graphrag.callbacks.reporting import create_pipeline_reporter
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
-from graphrag.config.enums import CacheType, IndexingMethod
+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.typing import PipelineRunResult, WorkflowFunction
+from graphrag.index.run.utils import create_callback_chain
+from graphrag.index.typing import WorkflowFunction
 from graphrag.index.workflows.factory import PipelineFactory
 from graphrag.logger.base import ProgressLogger
+from graphrag.logger.null_progress import NullProgressLogger
 
 log = logging.getLogger(__name__)
 
@@ -51,36 +53,37 @@ async def build_index(
     list[PipelineRunResult]
         The list of pipeline run results
     """
-    pipeline_cache = (
-        NoopPipelineCache() if config.cache.type == CacheType.none is None else None
-    )
+    logger = progress_logger or NullProgressLogger()
     # create a pipeline reporter and add to any additional callbacks
-    # TODO: remove the type ignore once the new config engine has been refactored
     callbacks = callbacks or []
-    callbacks.append(create_pipeline_reporter(config.reporting, None))  # type: ignore
+    callbacks.append(create_pipeline_reporter(config.reporting, None))
+
+    workflow_callbacks = create_callback_chain(callbacks, logger)
+
     outputs: list[PipelineRunResult] = []
 
     if memory_profile:
         log.warning("New pipeline does not yet support memory profiling.")
 
     pipeline = PipelineFactory.create_pipeline(config, method)
 
+    workflow_callbacks.pipeline_start(pipeline.names())
+
     async for output in run_pipeline(
         pipeline,
         config,
-        cache=pipeline_cache,
-        callbacks=callbacks,
-        logger=progress_logger,
+        callbacks=workflow_callbacks,
+        logger=logger,
         is_update_run=is_update_run,
     ):
         outputs.append(output)
-        if progress_logger:
-            if output.errors and len(output.errors) > 0:
-                progress_logger.error(output.workflow)
-            else:
-                progress_logger.success(output.workflow)
-            progress_logger.info(str(output.result))
+        if output.errors and len(output.errors) > 0:
+            logger.error(output.workflow)
+        else:
+            logger.success(output.workflow)
+        logger.info(str(output.result))
 
+    workflow_callbacks.pipeline_end(outputs)
     return outputs
 
 

graphrag/api/query.py

@@ -435,10 +435,11 @@ def local_search_streaming(
     vector_store_args = {}
     for index, store in config.vector_store.items():
         vector_store_args[index] = store.model_dump()
-    logger.info(f"Vector Store Args: {redact(vector_store_args)}")  # type: ignore # noqa
+    msg = f"Vector Store Args: {redact(vector_store_args)}"
+    logger.info(msg)
 
     description_embedding_store = get_embedding_store(
-        config_args=vector_store_args,  # type: ignore
+        config_args=vector_store_args,
         embedding_name=entity_description_embedding,
     )
 
@@ -453,7 +454,7 @@ def local_search_streaming(
         entities=entities_,
         relationships=read_indexer_relationships(relationships),
         covariates={"claims": covariates_},
-        description_embedding_store=description_embedding_store,  # type: ignore
+        description_embedding_store=description_embedding_store,
         response_type=response_type,
         system_prompt=prompt,
         callbacks=callbacks,
@@ -789,15 +790,16 @@ def drift_search_streaming(
     vector_store_args = {}
     for index, store in config.vector_store.items():
         vector_store_args[index] = store.model_dump()
-    logger.info(f"Vector Store Args: {redact(vector_store_args)}")  # type: ignore # noqa
+    msg = f"Vector Store Args: {redact(vector_store_args)}"
+    logger.info(msg)
 
     description_embedding_store = get_embedding_store(
-        config_args=vector_store_args,  # type: ignore
+        config_args=vector_store_args,
         embedding_name=entity_description_embedding,
     )
 
     full_content_embedding_store = get_embedding_store(
-        config_args=vector_store_args,  # type: ignore
+        config_args=vector_store_args,
         embedding_name=community_full_content_embedding,
     )
 
@@ -815,7 +817,7 @@ def drift_search_streaming(
         text_units=read_indexer_text_units(text_units),
         entities=entities_,
         relationships=read_indexer_relationships(relationships),
-        description_embedding_store=description_embedding_store,  # type: ignore
+        description_embedding_store=description_embedding_store,
         local_system_prompt=prompt,
         reduce_system_prompt=reduce_prompt,
         response_type=response_type,
@@ -1104,10 +1106,11 @@ def basic_search_streaming(
     vector_store_args = {}
     for index, store in config.vector_store.items():
         vector_store_args[index] = store.model_dump()
-    logger.info(f"Vector Store Args: {redact(vector_store_args)}")  # type: ignore # noqa
+    msg = f"Vector Store Args: {redact(vector_store_args)}"
+    logger.info(msg)
 
     description_embedding_store = get_embedding_store(
-        config_args=vector_store_args,  # type: ignore
+        config_args=vector_store_args,
         embedding_name=text_unit_text_embedding,
     )
 

graphrag/callbacks/blob_workflow_callbacks.py

@@ -24,11 +24,11 @@ class BlobWorkflowCallbacks(NoopWorkflowCallbacks):
     def __init__(
         self,
         connection_string: str | None,
-        container_name: str,
+        container_name: str | None,
         blob_name: str = "",
         base_dir: str | None = None,
         storage_account_blob_url: str | None = None,
-    ):  # type: ignore
+    ):
         """Create a new instance of the BlobStorageReporter class."""
         if container_name is None:
             msg = "No container name provided for blob storage."

graphrag/callbacks/noop_workflow_callbacks.py

@@ -4,12 +4,19 @@
 """A no-op implementation of WorkflowCallbacks."""
 
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
+from graphrag.index.run.pipeline_run_result import PipelineRunResult
 from graphrag.logger.progress import Progress
 
 
 class NoopWorkflowCallbacks(WorkflowCallbacks):
     """A no-op implementation of WorkflowCallbacks."""
 
+    def pipeline_start(self, names: list[str]) -> None:
+        """Execute this callback when a the entire pipeline starts."""
+
+    def pipeline_end(self, results: list[PipelineRunResult]) -> None:
+        """Execute this callback when the entire pipeline ends."""
+
     def workflow_start(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow starts."""
 

graphrag/callbacks/reporting.py

@@ -1,110 +1,39 @@
 # Copyright (c) 2024 Microsoft Corporation.
 # Licensed under the MIT License
 
-"""A module containing 'PipelineReportingConfig', 'PipelineFileReportingConfig' and 'PipelineConsoleReportingConfig' models."""
+"""A module containing the pipeline reporter factory."""
 
 from __future__ import annotations
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Generic, Literal, TypeVar, cast
-
-from pydantic import BaseModel, Field
+from typing import TYPE_CHECKING
 
 from graphrag.callbacks.blob_workflow_callbacks import BlobWorkflowCallbacks
 from graphrag.callbacks.console_workflow_callbacks import ConsoleWorkflowCallbacks
 from graphrag.callbacks.file_workflow_callbacks import FileWorkflowCallbacks
 from graphrag.config.enums import ReportingType
+from graphrag.config.models.reporting_config import ReportingConfig
 
 if TYPE_CHECKING:
     from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
 
-T = TypeVar("T")
-
-
-class PipelineReportingConfig(BaseModel, Generic[T]):
-    """Represent the reporting configuration for the pipeline."""
-
-    type: T
-
-
-class PipelineFileReportingConfig(PipelineReportingConfig[Literal[ReportingType.file]]):
-    """Represent the file reporting configuration for the pipeline."""
-
-    type: Literal[ReportingType.file] = ReportingType.file
-    """The type of reporting."""
-
-    base_dir: str | None = Field(
-        description="The base directory for the reporting.", default=None
-    )
-    """The base directory for the reporting."""
-
-
-class PipelineConsoleReportingConfig(
-    PipelineReportingConfig[Literal[ReportingType.console]]
-):
-    """Represent the console reporting configuration for the pipeline."""
-
-    type: Literal[ReportingType.console] = ReportingType.console
-    """The type of reporting."""
-
-
-class PipelineBlobReportingConfig(PipelineReportingConfig[Literal[ReportingType.blob]]):
-    """Represents the blob reporting configuration for the pipeline."""
-
-    type: Literal[ReportingType.blob] = ReportingType.blob
-    """The type of reporting."""
-
-    connection_string: str | None = Field(
-        description="The blob reporting connection string for the reporting.",
-        default=None,
-    )
-    """The blob reporting connection string for the reporting."""
-
-    container_name: str = Field(
-        description="The container name for reporting", default=""
-    )
-    """The container name for reporting"""
-
-    storage_account_blob_url: str | None = Field(
-        description="The storage account blob url for reporting", default=None
-    )
-    """The storage account blob url for reporting"""
-
-    base_dir: str | None = Field(
-        description="The base directory for the reporting.", default=None
-    )
-    """The base directory for the reporting."""
-
-
-PipelineReportingConfigTypes = (
-    PipelineFileReportingConfig
-    | PipelineConsoleReportingConfig
-    | PipelineBlobReportingConfig
-)
-
 
 def create_pipeline_reporter(
-    config: PipelineReportingConfig | None, root_dir: str | None
+    config: ReportingConfig | None, root_dir: str | None
 ) -> WorkflowCallbacks:
     """Create a logger for the given pipeline config."""
-    config = config or PipelineFileReportingConfig(base_dir="logs")
-
+    config = config or ReportingConfig(base_dir="logs", type=ReportingType.file)
     match config.type:
         case ReportingType.file:
-            config = cast("PipelineFileReportingConfig", config)
             return FileWorkflowCallbacks(
                 str(Path(root_dir or "") / (config.base_dir or ""))
             )
         case ReportingType.console:
             return ConsoleWorkflowCallbacks()
         case ReportingType.blob:
-            config = cast("PipelineBlobReportingConfig", config)
             return BlobWorkflowCallbacks(
                 config.connection_string,
                 config.container_name,
                 base_dir=config.base_dir,
                 storage_account_blob_url=config.storage_account_blob_url,
             )
-        case _:
-            msg = f"Unknown reporting type: {config.type}"
-            raise ValueError(msg)

graphrag/callbacks/workflow_callbacks.py

@@ -5,6 +5,7 @@
 
 from typing import Protocol
 
+from graphrag.index.run.pipeline_run_result import PipelineRunResult
 from graphrag.logger.progress import Progress
 
 
@@ -15,6 +16,14 @@ class WorkflowCallbacks(Protocol):
     This base class is a "noop" implementation so that clients may implement just the callbacks they need.
     """
 
+    def pipeline_start(self, names: list[str]) -> None:
+        """Execute this callback to signal when the entire pipeline starts."""
+        ...
+
+    def pipeline_end(self, results: list[PipelineRunResult]) -> None:
+        """Execute this callback to signal when the entire pipeline ends."""
+        ...
+
     def workflow_start(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow starts."""
         ...

graphrag/callbacks/workflow_callbacks_manager.py

@@ -4,6 +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.logger.progress import Progress
 
 
@@ -20,6 +21,18 @@ def register(self, callbacks: WorkflowCallbacks) -> None:
         """Register a new WorkflowCallbacks type."""
         self._callbacks.append(callbacks)
 
+    def pipeline_start(self, names: list[str]) -> None:
+        """Execute this callback when a the entire pipeline starts."""
+        for callback in self._callbacks:
+            if hasattr(callback, "pipeline_start"):
+                callback.pipeline_start(names)
+
+    def pipeline_end(self, results: list[PipelineRunResult]) -> None:
+        """Execute this callback when the entire pipeline ends."""
+        for callback in self._callbacks:
+            if hasattr(callback, "pipeline_end"):
+                callback.pipeline_end(results)
+
     def workflow_start(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow starts."""
         for callback in self._callbacks:

graphrag/cli/query.py

@@ -527,7 +527,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()  # type: ignore
+    output_config = config.output.model_dump()
     storage_obj = StorageFactory().create_storage(
         storage_type=output_config["type"], kwargs=output_config
     )

graphrag/index/run/pipeline.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""A module containing the Pipeline class."""
+
+from collections.abc import Generator
+
+from graphrag.index.typing import Workflow
+
+
+class Pipeline:
+    """Encapsulates running workflows."""
+
+    def __init__(self, workflows: list[Workflow]):
+        self.workflows = workflows
+
+    def run(self) -> Generator[Workflow]:
+        """Return a Generator over the pipeline workflows."""
+        yield from self.workflows
+
+    def names(self) -> list[str]:
+        """Return the names of the workflows in the pipeline."""
+        return [name for name, _ in self.workflows]