Move components test in the example folder - add some tests

Author: stellasia

examples/pipeline/__init__.py

@@ -1,3 +1,17 @@
+#  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.
 from __future__ import annotations
 
 import asyncio

src/neo4j_genai/components/kg_builder.py renamed to examples/pipeline/kg_builder.py

@@ -1,3 +1,17 @@
+#  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.
 from __future__ import annotations
 
 import asyncio

src/neo4j_genai/components/rag.py renamed to examples/pipeline/rag.py

@@ -61,6 +61,7 @@ scipy = [
     {version = "^1", python = "<3.12"},
     {version = "^1.7.0", python = ">=3.12"}
 ]
+pytest-asyncio = "^0.23.8"
 
 [tool.poetry.extras]
 external_clients = ["weaviate-client", "pinecone-client"]

poetry.lock

@@ -1,3 +1,17 @@
+#  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.
 """
 Basic graph structure for Pipeline
 """

pyproject.toml

@@ -1,3 +1,17 @@
+#  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.
 """
 Pipeline implementation.
 
@@ -118,26 +132,26 @@ def save_results(self, node: Node, res: RunResult) -> None:
             node.name, res.result or {}, is_final=node.is_leaf()
         )
 
-    async def run_node(self, node: TaskNode, data: dict[str, Any]) -> None:
-        await node.run(data, callback=self.on_task_complete)
+    async def run_task(self, task: TaskNode, data: dict[str, Any]) -> None:
+        await task.run(data, callback=self.on_task_complete)
 
     async def on_task_complete(
-        self, node: TaskNode, data: dict[str, Any], res: RunResult
+        self, task: TaskNode, data: dict[str, Any], res: RunResult
     ) -> None:
-        self.save_results(node, res)
-        await asyncio.gather(*[self.run_node(n, data) for n in self.next(node)])
+        self.pipeline.on_task_complete(task, res)
+        await asyncio.gather(*[self.run_task(n, data) for n in self.next(task)])
 
-    def check_dependencies_complete(self, node: TaskNode) -> None:
-        dependencies = self.pipeline.previous_edges(node)
+    def check_dependencies_complete(self, task: TaskNode) -> None:
+        dependencies = self.pipeline.previous_edges(task)
         for d in dependencies:
             if d.start.status != RunStatus.DONE:  # type: ignore
                 logger.warning(
-                    f"Missing dependency {d.start.name} for {node.name} (status: {d.start.status})"  # type: ignore
+                    f"Missing dependency {d.start.name} for {task.name} (status: {d.start.status})"  # type: ignore
                 )
                 raise MissingDependencyError()
 
-    def next(self, node: TaskNode) -> Generator[TaskNode, None, None]:
-        possible_nexts = self.pipeline.next_edges(node)
+    def next(self, task: TaskNode) -> Generator[TaskNode, None, None]:
+        possible_nexts = self.pipeline.next_edges(task)
         for next_edge in possible_nexts:
             next_node = next_edge.end
             # check status
@@ -185,6 +199,13 @@ def connect(  # type: ignore
         if self.is_cyclic():
             raise Exception("Cyclic graph")
 
+    def on_task_complete(self, node, result) -> None:
+        self.add_result_for_component(
+            node.name,
+            result.result,
+            is_final=node.is_leaf()
+        )
+
     def add_result_for_component(
         self, name: str, result: dict[str, Any], is_final: bool = False
     ) -> None:

src/neo4j_genai/core/graph.py

@@ -1,3 +1,20 @@
+#  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.
+"""Result store interface
+and in-memory store implementation.
+"""
 from __future__ import annotations
 
 import abc

src/neo4j_genai/core/pipeline.py

@@ -0,0 +1,62 @@
+import pytest
+
+from neo4j_genai.core.graph import Node, Graph
+from neo4j_genai.core.pipeline import Orchestrator, Component, Pipeline, RunStatus
+
+
+@pytest.fixture(scope="function")
+def component():
+    return Component()
+
+
+@pytest.fixture(scope="function")
+def pipeline_branch(component):
+    pipe = Pipeline()
+    pipe.add_component("a", component)
+    pipe.add_component("b", component)
+    pipe.add_component("c", component)
+    pipe.connect("a", "b")
+    pipe.connect("a", "c")
+    return pipe
+
+
+@pytest.fixture(scope="function")
+def pipeline_aggregation(component):
+    pipe = Pipeline()
+    pipe.add_component("a", component)
+    pipe.add_component("b", component)
+    pipe.add_component("c", component)
+    pipe.connect("a", "b")
+    pipe.connect("a", "c")
+    return pipe
+
+
+def test_orchestrator_branch(pipeline_branch):
+    orchestrator = Orchestrator(pipeline=pipeline_branch)
+    node_a = pipeline_branch.get_node_by_name("a")
+    node_a.status = RunStatus.DONE
+    next_tasks = orchestrator.next(node_a)
+    next_task_names = [n.name for n in next_tasks]
+    assert next_task_names == ["b", "c"]
+
+
+def test_orchestrator_aggregation(pipeline_aggregation):
+    orchestrator = Orchestrator(pipeline=pipeline_aggregation)
+    node_a = pipeline_aggregation.get_node_by_name("a")
+    node_a.status = RunStatus.DONE
+    node_b = pipeline_aggregation.get_node_by_name("b")
+    node_b.status = RunStatus.DONE
+    next_tasks = orchestrator.next(node_a)
+    next_task_names = [n.name for n in next_tasks]
+    assert next_task_names == ["c"]
+
+
+def test_orchestrator_aggregation_waiting(pipeline_aggregation):
+    orchestrator = Orchestrator(pipeline=pipeline_aggregation)
+    node_a = pipeline_aggregation.get_node_by_name("a")
+    node_a.status = RunStatus.DONE
+    node_b = pipeline_aggregation.get_node_by_name("a")
+    node_b.status = RunStatus.UNKNOWN
+    next_tasks = orchestrator.next(node_a)
+    next_task_names = [n.name for n in next_tasks]
+    assert next_task_names == []

src/neo4j_genai/core/stores.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+from typing import Any
+from unittest.mock import AsyncMock
+
+import pytest
+
+from neo4j_genai.core.pipeline import Component, Pipeline
+
+
+@pytest.fixture(scope="function")
+def component_multiply():
+    class ComponentMultiply(Component):
+        def __init__(self, r: float = 2.0) -> None:
+            self.r = r
+
+        async def run(self, number: float):
+            return {"product": number * self.r}
+
+    return ComponentMultiply()
+
+
+@pytest.fixture(scope="function")
+def component_add():
+    class ComponentAdd(Component):
+        async def run(self, number1: float, number2: float):
+            return {"sum": number1 + number2}
+
+    return ComponentAdd()
+
+
+@pytest.mark.asyncio
+async def test_simple_pipeline_two_components():
+    pipe = Pipeline()
+    component_a = AsyncMock(spec=Component)
+    component_a.run = AsyncMock(return_value={})
+    component_b = AsyncMock(spec=Component)
+    component_b.run = AsyncMock(return_value={"result": "result"})
+    pipe.add_component("a", component_a)
+    pipe.add_component("b", component_b)
+    pipe.connect("a", "b", {})
+    res = await pipe.run({})
+    assert component_a.run.called_one_with({})
+    assert component_b.run.called_one_with({})
+    assert "b" in res
+    assert res["b"] == {"result": "result"}
+
+
+@pytest.mark.asyncio
+async def test_simple_pipeline_two_components_parameter_propagation():
+    pipe = Pipeline()
+    component_a = AsyncMock(spec=Component)
+    component_a.run = AsyncMock(return_value={"product": 20})
+    component_b = AsyncMock(spec=Component)
+    component_b.run = AsyncMock(return_value={"sum": 54})
+    pipe.add_component("a", component_a)
+    pipe.add_component("b", component_b)
+    # first component output product goes to second component input number1
+    pipe.connect("a", "b", {
+        "number1": "a.product",
+    })
+    res = await pipe.run({"a": {}, "b": {"number2": 1}})
+    assert component_a.run.called_one_with({})
+    assert component_b.run.called_one_with({"number1": 20, "number2": 1})
+    assert res == {"b": {"sum": 54}}
+
+
+@pytest.mark.asyncio
+async def test_pipeline_branches():
+    pipe = Pipeline()
+    component_a = AsyncMock(spec=Component)
+    component_a.run = AsyncMock(return_value={})
+    component_b = AsyncMock(spec=Component)
+    component_b.run = AsyncMock(return_value={})
+    component_c = AsyncMock(spec=Component)
+    component_c.run = AsyncMock(return_value={})
+
+    pipe.add_component("a", component_a)
+    pipe.add_component("b", component_b)
+    pipe.add_component("c", component_c)
+    pipe.connect("a", "b")
+    pipe.connect("a", "c")
+    res = await pipe.run({})
+    assert "b" in res
+    assert "c" in res
+
+
+@pytest.mark.asyncio
+async def test_pipeline_aggregation():
+    pipe = Pipeline()
+    component_a = AsyncMock(spec=Component)
+    component_a.run = AsyncMock(return_value={})
+    component_b = AsyncMock(spec=Component)
+    component_b.run = AsyncMock(return_value={})
+    component_c = AsyncMock(spec=Component)
+    component_c.run = AsyncMock(return_value={})
+
+    pipe.add_component("a", component_a)
+    pipe.add_component("b", component_b)
+    pipe.add_component("c", component_c)
+    pipe.connect("a", "c")
+    pipe.connect("b", "c")
+    res = await pipe.run({})
+    assert "c" in res