fix: a bug where the SubGrapWorker did not correctly setup the graph and distances were not computed which let to failure when multiple JoinedTaskWorkers were waiting on the same provenance

Author: provos

src/planai/graph_task.py

@@ -50,7 +50,8 @@ def consume_work(self, task: output_type):
                 old_task = task.get_private_state(PRIVATE_STATE_KEY)
                 if old_task is None:
                     raise ValueError(
-                        f"No task provenance found for {PRIVATE_STATE_KEY}"
+                        f"No task state found for {PRIVATE_STATE_KEY}: {task._private_state} "
+                        f"(provenance: {task._provenance})",
                     )
                 assert isinstance(task, Task)
 
@@ -75,7 +76,7 @@ def consume_work(self, task: output_type):
                 with graph_task.lock:
                     if provenance not in graph_task._state:
                         raise ValueError(
-                            f"Task {provenance} does not have any associated state."
+                            f"Task {provenance} does not have any associated state: {graph_task._state}"
                         )
                     logging.debug(
                         "Subgraph is removing provenance for %s in %s",
@@ -92,6 +93,9 @@ def consume_work(self, task: output_type):
         instance = AdapterSinkWorker()
         self.graph.add_workers(instance)
         self.graph.set_dependency(self.exit_worker, instance)
+        self.graph.set_entry(self.entry_worker)
+        self.graph.finalize()  # compute the worker distances
+        self.graph.init_workers()
 
     def get_task_class(self) -> Type[Task]:
         # usually the entry task gets dynamically determined from consume_work but we are overriding it here
@@ -101,7 +105,6 @@ def init(self):
         # we need to install the graph dispatcher into the sub-graph
         assert self._graph is not None
         self.graph._dispatcher = self._graph._dispatcher
-        self.graph.init_workers()
 
     def consume_work(self, task: Task):
         new_task = task.copy_public()
@@ -143,7 +146,9 @@ def notify(self, prefix: str):
         self.graph.unwatch(prefix, self)
         with self.lock:
             if prefix not in self._state:
-                raise ValueError(f"Task {prefix} does not have any associated state.")
+                raise ValueError(
+                    f"Task {prefix} does not have any associated state: {self._state}"
+                )
             task, remove_provenance = self._state.pop(prefix)
 
         if remove_provenance:

tests/planai/patterns/test_search_fetch.py

@@ -223,6 +223,42 @@ def test_search_fetch_graph_workflow(self):
             self.assertTrue(page.title.startswith("Test Result"))
             self.assertEqual(page.content, self.mock_page_content)
 
+    def test_search_fetch_worker_distances(self):
+        # Create graph and get references to workers
+        graph, search_executor, exit_worker = create_search_fetch_graph(
+            llm=self.mock_llm, name="TestSearchFetch"
+        )
+
+        # Compute distances
+        graph.set_entry(search_executor)
+        graph.finalize()
+
+        # Validate specific distances from InitialTaskWorker
+        initial_distances = graph._worker_distances["InitialTaskWorker"]
+        expected_order = {
+            "SearchExecutor": 1,
+            "SearchResultSplitter": 2,
+            "PageFetcher": 3,
+            "PageRelevanceFilter": 4,
+            "PageAnalysisConsumer": 4,
+            "PageConsolidator": 5,
+        }
+        for worker_name, expected_distance in expected_order.items():
+            self.assertEqual(
+                initial_distances[worker_name],
+                expected_distance,
+                f"Expected {worker_name} to be at distance {expected_distance}, "
+                f"got {initial_distances[worker_name]}",
+            )
+
+        # Validate specific distance between PageAnalysisConsumer and PageConsolidator
+        analysis_consumer_distances = graph._worker_distances["PageAnalysisConsumer"]
+        self.assertEqual(
+            analysis_consumer_distances["PageConsolidator"],
+            1,
+            "Expected PageConsolidator to be at distance 1 from PageAnalysisConsumer",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

tests/planai/test_graph_task.py

@@ -558,6 +558,51 @@ def abort_thread():
         self.assertEqual(dispatcher.work_queue.qsize(), 0)
         self.assertEqual(dispatcher.active_tasks, 0)
 
+    def test_graph_task_entry_setup(self):
+        # Create subgraph
+        subgraph = Graph(name="SubGraph")
+        subgraph_worker = SubGraphHandler()
+        subgraph.add_workers(subgraph_worker)
+
+        # Create GraphTask
+        _ = SubGraphWorker(
+            graph=subgraph, entry_worker=subgraph_worker, exit_worker=subgraph_worker
+        )
+
+        # Verify that entry worker has been set correctly
+        self.assertIn(subgraph_worker, subgraph.dependencies[subgraph._initial_worker])
+        self.assertEqual(len(subgraph.dependencies[subgraph._initial_worker]), 1)
+
+    def test_graph_task_distances(self):
+        # Create subgraph with multiple workers
+        subgraph = Graph(name="SubGraph")
+
+        class Worker1(TaskWorker):
+            output_types: List[Type[Task]] = [SubGraphTask]
+
+            def consume_work(self, task: SubGraphTask):
+                self.publish_work(task, input_task=task)
+
+        class Worker2(TaskWorker):
+            output_types: List[Type[Task]] = [SubGraphTask]
+
+            def consume_work(self, task: SubGraphTask):
+                self.publish_work(task, input_task=task)
+
+        worker1 = Worker1()
+        worker2 = Worker2()
+        subgraph.add_workers(worker1, worker2)
+        subgraph.set_dependency(worker1, worker2)
+
+        # Create GraphTask
+        _ = SubGraphWorker(graph=subgraph, entry_worker=worker1, exit_worker=worker2)
+
+        # Verify subgraph distances
+        sub_distances = subgraph._worker_distances["InitialTaskWorker"]
+        self.assertEqual(sub_distances["Worker1"], 1)
+        self.assertEqual(sub_distances["Worker2"], 2)
+        self.assertEqual(sub_distances["AdapterSinkWorker"], 3)
+
 
 class StatusNotifyingWorker(TaskWorker):
     output_types: List[Type[Task]] = [FinalTask]