Refactor models and verification functions for improved validation and clarity

- Updated the DependentString model to utilize enumeration for order in dependent placeholders, enhancing readability.
- Refactored validation methods in NodeTemplate and GraphTemplate to ensure non-empty inputs and unique identifiers.
- Improved error handling in verification functions to check for string types in inputs and ensure proper validation of graph structures.
- Streamlined the RegisteredNode model's list_nodes_by_templates method to handle empty templates gracefully.

These changes enhance the robustness and maintainability of the models and their validation processes.

Author: NiveditJain

state-manager/app/models/db/graph_template_model.py

@@ -41,14 +41,12 @@ def _build_root_node(self) -> None:
         in_degree = {node.identifier: 0 for node in self.nodes}
 
         for node in self.nodes:
-            if node.next_nodes is None:
-                continue
-            for next_node in node.next_nodes:
-                in_degree[next_node] += 1
+            if node.next_nodes is not None:
+                for next_node in node.next_nodes:
+                    in_degree[next_node] += 1
 
-            if node.unites is None:
-                continue
-            in_degree[node.identifier] += 1
+            if node.unites is not None:
+                in_degree[node.identifier] += 1
 
         zero_in_degree_nodes = [node for node in self.nodes if in_degree[node.identifier] == 0]
         if len(zero_in_degree_nodes) != 1:
@@ -58,29 +56,31 @@ def _build_root_node(self) -> None:
     def _build_parents_by_identifier(self) -> None:
         try:
             root_node_identifier = self.get_root_node().identifier
-            self._parents_by_identifier = {
-                node.identifier: set() for node in self.nodes
-            }
 
-            visited = set()
+            visited = {}
+
+            self._parents_by_identifier = {}
+            for node in self.nodes:
+                self._parents_by_identifier[node.identifier] = set()
+                visited[node.identifier] = False
 
             def dfs(node_identifier: str, parents: set[str]) -> None:
                 assert self._parents_by_identifier is not None
 
                 self._parents_by_identifier[node_identifier] = parents | self._parents_by_identifier[node_identifier]
 
-                if node_identifier in visited:
+                if visited[node_identifier]:
                     return
 
-                visited.add(node_identifier)
+                visited[node_identifier] = True
 
                 node = self.get_node_by_identifier(node_identifier)
                 if node is None:
                     return
                 if node.next_nodes is None:
                     return
                 if node.unites is not None:
-                    self._parents_by_identifier[node.unites.identifier].add(node_identifier)
+                    self._parents_by_identifier[node_identifier].add(node.unites.identifier)
                 for next_node_identifier in node.next_nodes:
                     dfs(next_node_identifier, parents | {node_identifier})
 
@@ -162,18 +162,13 @@ def _validate_secret_value(cls, secret_value: str) -> None:
         except Exception:
             raise ValueError("Value is not valid URL-safe base64 encoded")
 
-    @model_validator(mode='after')
-    def validate_nodes(self) -> Self:
-        for node in self.nodes:
-            if node.namespace != self.namespace:
-                raise ValueError(f"Node namespace {node.namespace} does not match graph namespace {self.namespace}")
-        return self
-    
     @model_validator(mode='after')
     def validate_graph_is_connected(self) -> Self:
         errors = []
         root_node_identifier = self.get_root_node().identifier
         for node in self.nodes:
+            if node.identifier == root_node_identifier:
+                continue
             if root_node_identifier not in self.get_parents_by_identifier(node.identifier):
                 errors.append(f"Node {node.identifier} is not connected to the root node")
         if errors:
@@ -213,6 +208,10 @@ def verify_input_dependencies(self) -> Self:
         for node in self.nodes:
             for input_value in node.inputs.values():
                 try:
+                    if not isinstance(input_value, str):
+                        errors.append(f"Input {input_value} is not a string")
+                        continue
+
                     dependent_string = DependentString.create_dependent_string(input_value)
                     dependent_identifiers = set([identifier for identifier, _ in dependent_string.get_identifier_field()])
 
@@ -272,7 +271,7 @@ def get_parents_by_identifier(self, identifier: str) -> set[str]:
             self._build_parents_by_identifier()
 
         assert self._parents_by_identifier is not None
-        return self._parents_by_identifier[identifier]
+        return self._parents_by_identifier.get(identifier, set())
 
     @staticmethod
     async def get(namespace: str, graph_name: str) -> "GraphTemplate":

state-manager/app/models/db/registered_node.py

@@ -32,11 +32,13 @@ async def get_by_name_and_namespace(name: str, namespace: str) -> "RegisteredNod
 
     @staticmethod
     async def list_nodes_by_templates(templates: list[NodeTemplate]) -> list["RegisteredNode"]:
+        if len(templates) == 0:
+            return []
+        
         query = {
             "$or": [
                 {"name": node.node_name, "namespace": node.namespace}
                 for node in templates
             ]
         }
-        nodes = await RegisteredNode.find(query).to_list()
-        return [node for node in nodes if node is not None]
+        return await RegisteredNode.find(query).to_list()

state-manager/app/models/dependent_string.py

@@ -27,19 +27,17 @@ def create_dependent_string(syntax_string: str) -> "DependentString":
             return DependentString(head=syntax_string, dependents={})
 
         dependent_string = DependentString(head=splits[0], dependents={})
-        order = 0
 
-        for split in splits[1:]:
+        for order, split in enumerate(splits[1:]):
             if "}}" not in split:
                 raise ValueError(f"Invalid syntax string placeholder {split} for: {syntax_string} '${{' not closed")
-            placeholder_content, tail = split.split("}}")
+            placeholder_content, tail = split.split("}}", 1)
 
             parts = [p.strip() for p in placeholder_content.split(".")]
             if len(parts) != 3 or parts[1] != "outputs":
                 raise ValueError(f"Invalid syntax string placeholder {placeholder_content} for: {syntax_string}")
 
             dependent_string.dependents[order] = Dependent(identifier=parts[0], field=parts[2], tail=tail)
-            order += 1
 
         return dependent_string
 

state-manager/app/models/node_template_model.py

@@ -22,16 +22,9 @@ def validate_node_name(cls, v: str) -> str:
             raise ValueError("Node name cannot be empty")
         return v
 
-    @field_validator('node_name')
-    @classmethod
-    def validate_node_name_unique(cls, v: str) -> str:
-        if v == "" or v is None:
-            raise ValueError("Node name cannot be empty")
-        return v
-    
     @field_validator('identifier')
     @classmethod
-    def validate_identifier_unique(cls, v: str) -> str:
+    def validate_identifier(cls, v: str) -> str:
         if v == "" or v is None:
             raise ValueError("Node identifier cannot be empty")
         return v
@@ -43,10 +36,15 @@ def validate_next_nodes(cls, v: Optional[List[str]]) -> Optional[List[str]]:
         errors = []
         if v is not None:
             for next_node_identifier in v:
+                
                 if next_node_identifier == "" or next_node_identifier is None:
                     errors.append("Next node identifier cannot be empty")
-                elif next_node_identifier in identifiers:
+                    continue
+
+                if next_node_identifier in identifiers:
                     errors.append(f"Next node identifier {next_node_identifier} is not unique")
+                    continue
+
                 identifiers.add(next_node_identifier)
         if errors:
             raise ValueError("\n".join(errors))
@@ -63,5 +61,7 @@ def validate_unites(cls, v: Optional[Unites]) -> Optional[Unites]:
     def get_dependent_strings(self) -> list[DependentString]:
         dependent_strings = []
         for input_value in self.inputs.values():
+            if not isinstance(input_value, str):
+                raise ValueError(f"Input {input_value} is not a string")
             dependent_strings.append(DependentString.create_dependent_string(input_value))
         return dependent_strings

state-manager/app/tasks/verify_graph.py

@@ -43,22 +43,22 @@ async def verify_secrets(graph_template: GraphTemplate, registered_nodes: list[R
 async def verify_inputs(graph_template: GraphTemplate, registered_nodes: list[RegisteredNode]) -> list[str]:
     errors = []
     look_up_table = {
-        (node.node_name, node.namespace): node
-        for node in registered_nodes
+        (rn.node_name, rn.namespace): rn
+        for rn in registered_nodes
     }
 
     for node in graph_template.nodes:
         if node.inputs is None:
             continue
 
-        if (node.node_name, node.namespace) not in look_up_table:
+        registered_node = look_up_table.get((node.node_name, node.namespace))
+        if registered_node is None:
             errors.append(f"Node {node.node_name} in namespace {node.namespace} does not exist")
             continue
 
-        registered_node = look_up_table[(node.node_name, node.namespace)]
-        registerd_node_input_model = create_model(registered_node.inputs_schema)
+        registered_node_input_model  = create_model(registered_node.inputs_schema)
 
-        for input_name, input_info in registerd_node_input_model.model_fields.items():
+        for input_name, input_info in registered_node_input_model.model_fields.items():
             if input_info.annotation is not str:
                 errors.append(f"Input {input_name} in node {node.node_name} in namespace {node.namespace} is not a string")
                 continue
@@ -75,7 +75,7 @@ async def verify_inputs(graph_template: GraphTemplate, registered_nodes: list[Re
                 temp_node = graph_template.get_node_by_identifier(identifier)
                 assert temp_node is not None
 
-                registered_node = look_up_table[(temp_node.node_name, temp_node.namespace)]
+                registered_node = look_up_table.get((temp_node.node_name, temp_node.namespace))
                 if registered_node is None:
                     errors.append(f"Node {temp_node.node_name} in namespace {temp_node.namespace} does not exist")
                     continue
@@ -100,7 +100,16 @@ async def verify_graph(graph_template: GraphTemplate):
             verify_secrets(graph_template, registered_nodes),
             verify_inputs(graph_template, registered_nodes)
         ]
-        errors.extend(await asyncio.gather(*basic_verify_tasks))
+        resultant_errors = await asyncio.gather(*basic_verify_tasks)
+
+        for error in resultant_errors:
+            errors.extend(error)
+        
+        if len(errors) > 0:
+            graph_template.validation_status = GraphTemplateValidationStatus.INVALID
+            graph_template.validation_errors = errors
+            await graph_template.save()
+            return
 
         graph_template.validation_status = GraphTemplateValidationStatus.VALID
         graph_template.validation_errors = None