Pipeline status notifier (#254)

* WIP: notifier

* Add callback protocol

* Call notifier + add example

* ruff

* ruff

* import annotation to make Protocol | None work

* mypy

* Add unit test

* Move types to types.py

* Doc and changelog

* ruff & co

* ruff

Author: stellasia

CHANGELOG.md

@@ -2,10 +2,13 @@
 
 ## Next
 
+### Added
+
+- Ability to add event listener to get notifications about Pipeline progress.
+
 ### Changed
 - Changed the default behaviour of `FixedSizeSplitter` to avoid words cut-off in the chunks whenever it is possible.
 
-
 ## 1.4.2
 
 ### Fixed

docs/source/conf.py

@@ -45,6 +45,7 @@
     "sphinx.ext.napoleon",
     "sphinx.ext.viewcode",
     "sphinx.ext.autosectionlabel",
+    "enum_tools.autoenum",
 ]
 
 # The suffix(es) of source filenames.

docs/source/types.rst

@@ -153,3 +153,27 @@ ParamFromEnvConfig
 ==================
 
 .. autoclass:: neo4j_graphrag.experimental.pipeline.config.param_resolver.ParamFromEnvConfig
+
+
+EventType
+=========
+
+.. autoenum:: neo4j_graphrag.experimental.pipeline.types.EventType
+
+
+PipelineEvent
+==============
+
+.. autoclass:: neo4j_graphrag.experimental.pipeline.types.PipelineEvent
+
+TaskEvent
+==============
+
+.. autoclass:: neo4j_graphrag.experimental.pipeline.types.TaskEvent
+
+
+EventCallbackProtocol
+=====================
+
+.. autoclass:: neo4j_graphrag.experimental.pipeline.types.EventCallbackProtocol
+    :members: __call__

docs/source/user_guide_pipeline.rst

@@ -106,7 +106,6 @@ Pipelines can be visualized using the `draw` method:
 
 .. code:: python
 
-    import asyncio
     from neo4j_graphrag.experimental.pipeline import Pipeline
 
     pipe = Pipeline()
@@ -131,3 +130,45 @@ Here is an example of final result:
 
 .. image:: images/pipeline_full.png
   :alt: Pipeline visualisation
+
+
+************************
+Adding an Event Callback
+************************
+
+It is possible to add a callback to receive notification about pipeline progress:
+
+- `PIPELINE_STARTED`, when pipeline starts
+- `PIPELINE_FINISHED`, when pipeline ends
+- `TASK_STARTED`, when a task starts
+- `TASK_FINISHED`, when a task ends
+
+
+See :ref:`pipelineevent` and :ref:`taskevent` to see what is sent in each event type.
+
+.. code:: python
+
+    import asyncio
+    import logging
+
+    from neo4j_graphrag.experimental.pipeline import Pipeline
+    from neo4j_graphrag.experimental.pipeline.types import Event
+
+    logger = logging.getLogger(__name__)
+    logging.basicConfig()
+    logger.setLevel(logging.WARNING)
+
+
+    async def event_handler(event: Event) -> None:
+        """Function can do anything about the event,
+        here we're just logging it if it's a pipeline-level event.
+        """
+        if event.event_type.is_pipeline_event:
+            logger.warning(event)
+
+    pipeline = Pipeline(
+        callback=event_handler,
+    )
+    # ... add components, connect them as usual
+
+    await pipeline.run(...)

examples/README.md

@@ -101,6 +101,7 @@ are listed in [the last section of this file](#customize).
 - [Process multiple documents](./customize/build_graph/pipeline/kg_builder_two_documents_entity_resolution.py)
 - [Export lexical graph creation into another pipeline](./customize/build_graph/pipeline/text_to_lexical_graph_to_entity_graph_two_pipelines.py)
 - [Build pipeline from config file](customize/build_graph/pipeline/from_config_files/pipeline_from_config_file.py)
+- [Add event listener to get notification about Pipeline progress](./customize/build_graph/pipeline/pipeline_with_notifications.py)
 
 
 #### Components

examples/customize/build_graph/pipeline/pipeline_with_notifications.py

@@ -0,0 +1,93 @@
+"""This example demonstrates how to use event callback to receive notifications
+about the pipeline progress.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+
+import neo4j
+from neo4j_graphrag.experimental.components.kg_writer import Neo4jWriter
+from neo4j_graphrag.experimental.components.lexical_graph import LexicalGraphBuilder
+from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
+    FixedSizeSplitter,
+)
+from neo4j_graphrag.experimental.pipeline import Pipeline
+from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult
+from neo4j_graphrag.experimental.pipeline.types import Event
+
+logger = logging.getLogger(__name__)
+logging.basicConfig()
+logger.setLevel(logging.WARNING)
+
+
+async def event_handler(event: Event) -> None:
+    """Function can do anything about the event,
+    here we're just logging it if it's a pipeline-level event.
+    """
+    if event.event_type.is_pipeline_event:
+        logger.warning(event)
+
+
+async def main(neo4j_driver: neo4j.Driver) -> PipelineResult:
+    """This is where we define and run the Lexical Graph builder pipeline, instantiating
+    a few components:
+
+    - Text Splitter: to split the text into manageable chunks of fixed size
+    - Chunk Embedder: to embed the chunks' text
+    - Lexical Graph Builder: to build the lexical graph, ie creating the chunk nodes and relationships between them
+    - KG writer: save the lexical graph to Neo4j
+    """
+    pipe = Pipeline(
+        callback=event_handler,
+    )
+    # define the components
+    pipe.add_component(
+        FixedSizeSplitter(chunk_size=300, chunk_overlap=10, approximate=False),
+        "splitter",
+    )
+    pipe.add_component(
+        LexicalGraphBuilder(),
+        "lexical_graph_builder",
+    )
+    pipe.add_component(Neo4jWriter(neo4j_driver), "writer")
+    # define the execution order of component
+    # and how the output of previous components must be used
+    pipe.connect(
+        "splitter", "lexical_graph_builder", input_config={"text_chunks": "splitter"}
+    )
+    pipe.connect(
+        "lexical_graph_builder",
+        "writer",
+        input_config={
+            "graph": "lexical_graph_builder.graph",
+            "lexical_graph_config": "lexical_graph_builder.config",
+        },
+    )
+    # user input:
+    # the initial text
+    # and the list of entities and relations we are looking for
+    pipe_inputs = {
+        "splitter": {
+            "text": """Albert Einstein was a German physicist born in 1879 who
+            wrote many groundbreaking papers especially about general relativity
+            and quantum mechanics. He worked for many different institutions, including
+            the University of Bern in Switzerland and the University of Oxford."""
+        },
+        "lexical_graph_builder": {
+            "document_info": {
+                # 'path' can be anything
+                "path": "example/pipeline_with_notifications"
+            },
+        },
+    }
+    # run the pipeline
+    return await pipe.run(pipe_inputs)
+
+
+if __name__ == "__main__":
+    with neo4j.GraphDatabase.driver(
+        "bolt://localhost:7687", auth=("neo4j", "password")
+    ) as driver:
+        print(asyncio.run(main(driver)))