From ac6ef359596def361fb70046f8d430c8af72e212 Mon Sep 17 00:00:00 2001 From: estelle Date: Tue, 1 Apr 2025 15:55:28 +0200 Subject: [PATCH 1/7] Add RunContext to doc and improve protocol description --- docs/source/types.rst | 13 +++++++++++++ .../experimental/pipeline/notification.py | 4 ++++ .../experimental/pipeline/types/context.py | 4 ++++ 3 files changed, 21 insertions(+) diff --git a/docs/source/types.rst b/docs/source/types.rst index 4b4ae7155..b826c3c5f 100644 --- a/docs/source/types.rst +++ b/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.TaskProgressCallbackProtocol + :members: __call__ + + +RunContext +========== + +.. autoclass:: neo4j_graphrag.experimental.pipeline.types.context.RunContext diff --git a/src/neo4j_graphrag/experimental/pipeline/notification.py b/src/neo4j_graphrag/experimental/pipeline/notification.py index e9cb63cc6..5e3dc03c8 100644 --- a/src/neo4j_graphrag/experimental/pipeline/notification.py +++ b/src/neo4j_graphrag/experimental/pipeline/notification.py @@ -68,6 +68,10 @@ 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]: ... diff --git a/src/neo4j_graphrag/experimental/pipeline/types/context.py b/src/neo4j_graphrag/experimental/pipeline/types/context.py index f0b4caf97..da52c8e3d 100644 --- a/src/neo4j_graphrag/experimental/pipeline/types/context.py +++ b/src/neo4j_graphrag/experimental/pipeline/types/context.py @@ -20,6 +20,10 @@ @runtime_checkable class TaskProgressCallbackProtocol(Protocol): + """This protocol is used to send events from within the component. + The final event will be of type TaskEvent, but event type and task name + are predefined and can not be changed by the user calling this protocol. + """ def __call__(self, message: str, data: dict[str, Any]) -> Awaitable[None]: ... From 50b9a40066b43cfbf78983f451c11ab30354901d Mon Sep 17 00:00:00 2001 From: estelle Date: Tue, 1 Apr 2025 15:58:36 +0200 Subject: [PATCH 2/7] Ruff --- src/neo4j_graphrag/experimental/pipeline/notification.py | 1 + src/neo4j_graphrag/experimental/pipeline/types/context.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/neo4j_graphrag/experimental/pipeline/notification.py b/src/neo4j_graphrag/experimental/pipeline/notification.py index 5e3dc03c8..822e62919 100644 --- a/src/neo4j_graphrag/experimental/pipeline/notification.py +++ b/src/neo4j_graphrag/experimental/pipeline/notification.py @@ -72,6 +72,7 @@ class EventCallbackProtocol(Protocol): from the pipeline itself. It will receive either PipelineEvent or TaskEvent depending on the event type. """ + def __call__(self, event: Event) -> Awaitable[None]: ... diff --git a/src/neo4j_graphrag/experimental/pipeline/types/context.py b/src/neo4j_graphrag/experimental/pipeline/types/context.py index da52c8e3d..566af355c 100644 --- a/src/neo4j_graphrag/experimental/pipeline/types/context.py +++ b/src/neo4j_graphrag/experimental/pipeline/types/context.py @@ -24,6 +24,7 @@ class TaskProgressCallbackProtocol(Protocol): The final event will be of type TaskEvent, but event type and task name are predefined and can not be changed by the user calling this protocol. """ + def __call__(self, message: str, data: dict[str, Any]) -> Awaitable[None]: ... From d67e1b3b484bb5d0d13fccb785b0d71d23db1407 Mon Sep 17 00:00:00 2001 From: estelle Date: Wed, 2 Apr 2025 19:33:44 +0200 Subject: [PATCH 3/7] Rename for clarity --- .../experimental/pipeline/types/context.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/neo4j_graphrag/experimental/pipeline/types/context.py b/src/neo4j_graphrag/experimental/pipeline/types/context.py index 566af355c..d641aae38 100644 --- a/src/neo4j_graphrag/experimental/pipeline/types/context.py +++ b/src/neo4j_graphrag/experimental/pipeline/types/context.py @@ -19,10 +19,11 @@ @runtime_checkable -class TaskProgressCallbackProtocol(Protocol): - """This protocol is used to send events from within the component. - The final event will be of type TaskEvent, but event type and task name - are predefined and can not be changed by the user calling this 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]: ... @@ -33,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) From abea591f825cdadcac21b289905ecf01cd87e0d8 Mon Sep 17 00:00:00 2001 From: estelle Date: Tue, 1 Apr 2025 15:55:28 +0200 Subject: [PATCH 4/7] Add RunContext to doc and improve protocol description --- docs/source/types.rst | 13 +++++++++++++ .../experimental/pipeline/notification.py | 4 ++++ .../experimental/pipeline/types/context.py | 4 ++++ 3 files changed, 21 insertions(+) diff --git a/docs/source/types.rst b/docs/source/types.rst index 4b4ae7155..b826c3c5f 100644 --- a/docs/source/types.rst +++ b/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.TaskProgressCallbackProtocol + :members: __call__ + + +RunContext +========== + +.. autoclass:: neo4j_graphrag.experimental.pipeline.types.context.RunContext diff --git a/src/neo4j_graphrag/experimental/pipeline/notification.py b/src/neo4j_graphrag/experimental/pipeline/notification.py index cc69a69e2..bfb3dbec6 100644 --- a/src/neo4j_graphrag/experimental/pipeline/notification.py +++ b/src/neo4j_graphrag/experimental/pipeline/notification.py @@ -74,6 +74,10 @@ 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]: ... diff --git a/src/neo4j_graphrag/experimental/pipeline/types/context.py b/src/neo4j_graphrag/experimental/pipeline/types/context.py index f0b4caf97..da52c8e3d 100644 --- a/src/neo4j_graphrag/experimental/pipeline/types/context.py +++ b/src/neo4j_graphrag/experimental/pipeline/types/context.py @@ -20,6 +20,10 @@ @runtime_checkable class TaskProgressCallbackProtocol(Protocol): + """This protocol is used to send events from within the component. + The final event will be of type TaskEvent, but event type and task name + are predefined and can not be changed by the user calling this protocol. + """ def __call__(self, message: str, data: dict[str, Any]) -> Awaitable[None]: ... From 0729a2098b33e88b00e4ee2a4af2702e0e4c3794 Mon Sep 17 00:00:00 2001 From: estelle Date: Tue, 1 Apr 2025 15:58:36 +0200 Subject: [PATCH 5/7] Ruff --- src/neo4j_graphrag/experimental/pipeline/notification.py | 1 + src/neo4j_graphrag/experimental/pipeline/types/context.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/neo4j_graphrag/experimental/pipeline/notification.py b/src/neo4j_graphrag/experimental/pipeline/notification.py index bfb3dbec6..94665fe2c 100644 --- a/src/neo4j_graphrag/experimental/pipeline/notification.py +++ b/src/neo4j_graphrag/experimental/pipeline/notification.py @@ -78,6 +78,7 @@ class EventCallbackProtocol(Protocol): from the pipeline itself. It will receive either PipelineEvent or TaskEvent depending on the event type. """ + def __call__(self, event: Event) -> Awaitable[None]: ... diff --git a/src/neo4j_graphrag/experimental/pipeline/types/context.py b/src/neo4j_graphrag/experimental/pipeline/types/context.py index da52c8e3d..566af355c 100644 --- a/src/neo4j_graphrag/experimental/pipeline/types/context.py +++ b/src/neo4j_graphrag/experimental/pipeline/types/context.py @@ -24,6 +24,7 @@ class TaskProgressCallbackProtocol(Protocol): The final event will be of type TaskEvent, but event type and task name are predefined and can not be changed by the user calling this protocol. """ + def __call__(self, message: str, data: dict[str, Any]) -> Awaitable[None]: ... From 23dbcca22fa79976cb3c006966f5f61c8acbdac4 Mon Sep 17 00:00:00 2001 From: estelle Date: Wed, 2 Apr 2025 19:33:44 +0200 Subject: [PATCH 6/7] Rename for clarity --- .../experimental/pipeline/types/context.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/neo4j_graphrag/experimental/pipeline/types/context.py b/src/neo4j_graphrag/experimental/pipeline/types/context.py index 566af355c..d641aae38 100644 --- a/src/neo4j_graphrag/experimental/pipeline/types/context.py +++ b/src/neo4j_graphrag/experimental/pipeline/types/context.py @@ -19,10 +19,11 @@ @runtime_checkable -class TaskProgressCallbackProtocol(Protocol): - """This protocol is used to send events from within the component. - The final event will be of type TaskEvent, but event type and task name - are predefined and can not be changed by the user calling this 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]: ... @@ -33,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) From a634da2cfc5dcf0ac7d243f9d99f30b0998dc1ff Mon Sep 17 00:00:00 2001 From: Estelle Scifo Date: Tue, 15 Apr 2025 12:55:44 +0200 Subject: [PATCH 7/7] Update docs/source/types.rst Co-authored-by: Alex Thomas --- docs/source/types.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/types.rst b/docs/source/types.rst index b826c3c5f..73afcffb7 100644 --- a/docs/source/types.rst +++ b/docs/source/types.rst @@ -182,7 +182,7 @@ EventCallbackProtocol TaskProgressCallbackProtocol ============================ -.. autoclass:: neo4j_graphrag.experimental.pipeline.types.context.TaskProgressCallbackProtocol +.. autoclass:: neo4j_graphrag.experimental.pipeline.types.context.TaskProgressNotifierProtocol :members: __call__