Pipeline rerun + exception on cyclic graph (for now)

Author: stellasia

src/neo4j_genai/components/rag.py

@@ -36,6 +36,7 @@ async def process(self, prompt: str) -> dict[str, Any]:
     pipe.add_component("generate", LLM())
     pipe.connect("retrieve", "augment", {"context": "retrieve.context"})
     pipe.connect("augment", "generate", {"prompt": "augment.prompt"})
+    pipe.connect("generate", "retrieve", {"prompt": "augment.prompt"})
 
     query = "my question"
     print(

src/neo4j_genai/core/graph.py

@@ -70,3 +70,41 @@ def previous_edges(self, node: Node) -> list[Edge]:
 
     def __contains__(self, node: Node) -> bool:
         return node.name in self._nodes
+
+    def _check_stack(self, v, visited, stack):
+        visited[v.name] = True
+        stack[v.name] = True
+        for linked_vertices in self.next_edges(v):
+            next_node = linked_vertices.end
+            if not visited[next_node.name]:
+                if self._check_stack(next_node, visited, stack):
+                    return True
+            elif stack[next_node.name]:
+                return True
+        stack[v.name] = False
+        return False
+
+    def __is_cyclic(self) -> bool:
+        visited = {node: False for node in self._nodes}
+        stack = {node: False for node in self._nodes}
+        for node_name, node in self._nodes.items():
+            if not visited[node_name]:
+                if self._check_stack(node, visited, stack):
+                    return True
+        return False
+
+    def dfs(self, visited, node):
+        if node in visited:
+            return True
+        else:
+            for edge in self.next_edges(node):
+                neighbour = edge.end
+                if self.dfs(visited | {node}, neighbour):
+                    return True
+            return False
+
+    def is_cyclic(self):
+        for node in self._nodes.values():
+            if self.dfs(set(), node):
+                return True
+        return False

src/neo4j_genai/core/pipeline.py

@@ -81,6 +81,9 @@ def get_input_defs_from_parents(self) -> dict[str, str]:
                 input_defs.update(**prev_edge_data)
         return input_defs
 
+    def reinitialize(self) -> None:
+        self.status = RunStatus.SCHEDULED
+
     async def run(
         self, data: dict[str, Any], callback: Callable[[Any, Any, Any], Awaitable[Any]]
     ) -> None:
@@ -163,7 +166,6 @@ def __init__(self, store: Optional[Store] = None) -> None:
         super().__init__()
         self._store = store or InMemoryStore()
         self._final_results = InMemoryStore()
-        self.orchestrator = Orchestrator(self)
 
     def add_component(self, name: str, component: Component) -> None:
         task = TaskNode(name, component, self)
@@ -178,6 +180,8 @@ def connect(  # type: ignore
         start_node = self.get_node_by_name(start_component_name, raise_exception=True)
         end_node = self.get_node_by_name(end_component_name, raise_exception=True)
         super().connect(start_node, end_node, data={"input_defs": input_defs})
+        if self.is_cyclic():
+            raise Exception("Cyclic graph")
 
     def add_result_for_component(
         self, name: str, result: dict[str, Any], is_final: bool = False
@@ -208,6 +212,14 @@ def get_component_inputs(
                 component_inputs[input_def] = value
         return component_inputs
 
+    def reinitialize(self):
+        self._store.empty()
+        self._final_results.empty()
+        for task in self._nodes.values():
+            task.reinitialize()  # type: ignore
+
     async def run(self, data: dict[str, Any]) -> dict[str, Any]:
-        await self.orchestrator.run(data)
+        self.reinitialize()
+        orchestrator = Orchestrator(self)
+        await orchestrator.run(data)
         return self._final_results.all()

src/neo4j_genai/core/stores.py

@@ -42,6 +42,10 @@ def all(self) -> dict[str, Any]:
         """
         raise NotImplementedError()
 
+    def empty(self) -> None:
+        """Remove everything from store"""
+        raise NotImplementedError()
+
 
 class InMemoryStore(Store):
     """Simple in-memory store.
@@ -67,3 +71,6 @@ def find_all(self, pattern: str) -> list[Any]:
 
     def all(self) -> dict[str, Any]:
         return self._data
+
+    def empty(self) -> None:
+        self._data = {}