Add RunContext to doc and improve protocol description (#320)

* Add RunContext to doc and improve protocol description

* Ruff

* Rename for clarity

* Add RunContext to doc and improve protocol description

* Ruff

* Rename for clarity

* Update docs/source/types.rst

Co-authored-by: Alex Thomas <alexthomas93@users.noreply.github.com>

---------

Co-authored-by: Alex Thomas <alexthomas93@users.noreply.github.com>

Author: stellasia

docs/source/types.rst

@@ -177,3 +177,16 @@ EventCallbackProtocol
 
 .. autoclass:: neo4j_graphrag.experimental.pipeline.notification.EventCallbackProtocol
     :members: __call__
+
+
+TaskProgressCallbackProtocol
+============================
+
+.. autoclass:: neo4j_graphrag.experimental.pipeline.types.context.TaskProgressNotifierProtocol
+    :members: __call__
+
+
+RunContext
+==========
+
+.. autoclass:: neo4j_graphrag.experimental.pipeline.types.context.RunContext

src/neo4j_graphrag/experimental/pipeline/notification.py

@@ -74,6 +74,11 @@ class TaskEvent(Event):
 
 
 class EventCallbackProtocol(Protocol):
+    """This protocol is used to send events about pipeline progress
+    from the pipeline itself. It will receive either PipelineEvent or
+    TaskEvent depending on the event type.
+    """
+
     def __call__(self, event: Event) -> Awaitable[None]: ...
 
 

src/neo4j_graphrag/experimental/pipeline/types/context.py

@@ -19,7 +19,13 @@
 
 
 @runtime_checkable
-class TaskProgressCallbackProtocol(Protocol):
+class TaskProgressNotifierProtocol(Protocol):
+    """This protocol is used to send events from the component to the
+    Pipeline callback protocol.
+    The event sent to the callback will be of type :ref:`TaskEvent`,
+    with `event_type=TASK_PROGRESS`.
+    """
+
     def __call__(self, message: str, data: dict[str, Any]) -> Awaitable[None]: ...
 
 
@@ -28,7 +34,7 @@ class RunContext(BaseModel):
 
     run_id: str
     task_name: str
-    notifier: Optional[TaskProgressCallbackProtocol] = None
+    notifier: Optional[TaskProgressNotifierProtocol] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)