Claude/audit code generation t w pbv

project/block_manager/services/codegen/base.py

@@ -44,6 +44,16 @@ def topological_sort(nodes: List[Dict], edges: List[Dict]) -> List[Dict]:
             if in_degree[neighbor] == 0:
                 queue.append(neighbor)
 
+    # Cycle detection: if not all nodes were sorted, there's a cycle
+    if len(sorted_ids) != len(nodes):
+        # Find nodes that are still in the cycle (have non-zero in-degree)
+        cycle_nodes = [node_id for node_id, degree in in_degree.items() if degree > 0]
+        raise ValueError(
+            f"Graph contains a cycle. Neural networks must be acyclic (feedforward). "
+            f"Nodes involved in cycle: {', '.join(cycle_nodes[:5])}"
+            + (" and more..." if len(cycle_nodes) > 5 else "")
+        )
+
     # Return nodes in sorted order
     return [node_map[node_id] for node_id in sorted_ids if node_id in node_map]
 

project/block_manager/services/codegen/base_orchestrator.py

@@ -273,7 +273,7 @@ def _generate_config_file(self, nodes: List[Dict[str, Any]]) -> str:
         input_shape = self._extract_input_shape(nodes)
         layer_count = sum(
             1 for n in nodes
-            if get_node_type(n) not in ('input', 'output', 'dataloader')
+            if get_node_type(n) not in ('input', 'output', 'dataloader', 'loss')
         )
 
         if layer_count > 20:

project/block_manager/services/codegen/pytorch_group_generator.py

@@ -169,7 +169,7 @@ def _generate_internal_node_specs(
             node_type = get_node_type(node)
 
             # Skip special nodes
-            if node_type in ('input', 'output', 'dataloader'):
+            if node_type in ('input', 'output', 'dataloader', 'loss'):
                 continue
 
             node_id = node['id']
@@ -273,8 +273,8 @@ def _generate_forward_pass(
                         var_map[node_id] = var_name
                 continue
 
-            # Skip output and dataloader nodes (they don't produce code)
-            if node_type in ('output', 'dataloader'):
+            # Skip output, dataloader, and loss nodes (they don't produce code)
+            if node_type in ('output', 'dataloader', 'loss'):
                 continue
 
             # Get the spec for this node