Pipeline viz

.github/workflows/pr-e2e-tests.yaml

@@ -47,6 +47,8 @@ jobs:
           - 7474:7474
 
     steps:
+      - name: Install graphviz package
+        run: sudo apt install graphviz graphviz-dev
       - name: Check out repository code
         uses: actions/checkout@v4
       - name: Docker Prune

.github/workflows/pr.yaml

@@ -8,6 +8,8 @@ jobs:
       matrix:
         python-version: [ '3.8', '3.9', '3.10', '3.11', '3.12' ]
     steps:
+      - name: Install graphviz package
+        run: sudo apt install graphviz graphviz-dev
       - name: Check out repository code
         uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}

.github/workflows/scheduled-e2e-tests.yaml

@@ -54,6 +54,8 @@ jobs:
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 
     steps:
+      - name: Install graphviz package
+        run: sudo apt install graphviz graphviz-dev
       - name: Check out repository code
         uses: actions/checkout@v4
       - name: Docker Prune

CHANGELOG.md

@@ -2,9 +2,16 @@
 
 ## Next
 
-## Fixed
+### Added
+- Ability to visualise pipeline with `my_pipeline.draw("pipeline.png")`
+
+### Fixed
 - Pipelines now return correct results when the same pipeline is run in parallel.
 
+### Changed
+- Pipeline run method now return a PipelineResult object.
+
+
 ## 0.5.0
 
 ### Added

docs/source/user_guide_pipeline.rst

@@ -96,3 +96,38 @@ The data flow is illustrated in the diagram below:
             asyncio.run(pipe.run({"a": {"number1": 10, "number2": 1}, "b": {"number1": 4, "number2": 42}}))
 
     The result will still be **15** because the user input `"number2": 42` is ignored.
+
+
+**********************
+Visualising a Pipeline
+**********************
+
+Pipelines can be visualized using the `draw` method:
+
+.. code:: python
+
+    import asyncio
+    from neo4j_genai.experimental.pipeline import Pipeline
+
+    pipe = Pipeline()
+    # ... define components and connections
+
+    pipe.draw("pipeline.png")
+
+Here is an example pipeline rendering:
+
+.. image:: images/pipeline_no_unused_outputs.png
+  :alt: Pipeline visualisation with hidden outputs if unused
+
+
+By default, output fields which are not mapped to any component are hidden. They
+can be added to the canvas by setting `hide_unused_outputs` to `False`:
+
+.. code:: python
+
+    pipe.draw("pipeline.png", hide_unused_outputs=False)
+
+Here is an example of final result:
+
+.. image:: images/pipeline_full.png
+  :alt: Pipeline visualisation

examples/pipeline/visualization.py

@@ -0,0 +1,58 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""This example illustrates how to visualize a Pipeline"""
+
+from neo4j_genai.experimental.pipeline import Component, Pipeline
+from neo4j_genai.experimental.pipeline.component import DataModel
+from pydantic import validate_call
+
+
+class IntDataModel(DataModel):
+    value: int
+    message: str
+
+
+class Addition(Component):
+    async def run(self, a: int, b: int) -> IntDataModel:
+        return IntDataModel(value=a + b, message="addition complete")
+
+
+class Duplicate(Component):
+    def __init__(self, factor: int = 2) -> None:
+        self.factor = factor
+
+    async def run(self, number: int) -> IntDataModel:
+        return IntDataModel(
+            value=number * self.factor, message=f"multiplication by {self.factor} done"
+        )
+
+
+class Save(Component):
+    @validate_call
+    async def run(self, number: IntDataModel) -> IntDataModel:
+        return IntDataModel(value=number.value, message="saved")
+
+
+if __name__ == "__main__":
+    pipe = Pipeline()
+    pipe.add_component(Duplicate(), "times_two")
+    pipe.add_component(Duplicate(factor=10), "times_ten")
+    pipe.add_component(Addition(), "addition")
+    pipe.add_component(Save(), "save")
+    pipe.connect("times_two", "addition", {"a": "times_two.value"})
+    pipe.connect("times_ten", "addition", {"b": "times_ten.value"})
+    pipe.connect("addition", "save", {"number": "addition"})
+    pipe.draw("graph.png")
+    pipe.draw("graph_full.png", hide_unused_outputs=False)

pyproject.toml

@@ -37,6 +37,10 @@ pinecone-client = {version = "^4.1.0", optional = true}
 types-mock = "^5.1.0.20240425"
 eval-type-backport = "^0.2.0"
 pypdf = "^4.3.1"
+pygraphviz = [
+    {version = "^1.13.0", python = ">=3.10,<4.0.0"},
+    {version = "^1.0.0", python = "<3.10"}
+]
 
 [tool.poetry.group.dev.dependencies]
 pylint = "^3.1.0"