added bulk_create_next

state-manager/app/controller/executed_state.py

@@ -1,13 +1,12 @@
 from beanie import PydanticObjectId
-from beanie.operators import In
 from app.models.executed_models import ExecutedRequestModel, ExecutedResponseModel
 
 from fastapi import HTTPException, status, BackgroundTasks
 
 from app.models.db.state import State
 from app.models.state_status_enum import StateStatusEnum
 from app.singletons.logs_manager import LogsManager
-from app.tasks.create_next_state import create_next_state
+from app.tasks.create_next_states import create_next_states
 
 logger = LogsManager().get_logger()
 
@@ -23,19 +22,20 @@ async def executed_state(namespace_name: str, state_id: PydanticObjectId, body:
         if state.status != StateStatusEnum.QUEUED:
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is not queued")
 
+        next_state_ids = []
         if len(body.outputs) == 0:
             state.status = StateStatusEnum.EXECUTED
             state.outputs = {}
             await state.save()
 
-            background_tasks.add_task(create_next_state, state)
+            next_state_ids.append(state.id)
 
         else:            
             state.outputs = body.outputs[0]
             state.status = StateStatusEnum.EXECUTED
 
             await state.save()
-            background_tasks.add_task(create_next_state, state)
+            next_state_ids.append(state.id)
 
             new_states = []
             for output in body.outputs[1:]:
@@ -54,16 +54,9 @@ async def executed_state(namespace_name: str, state_id: PydanticObjectId, body:
 
             if len(new_states) > 0:
                 inserted_ids = (await State.insert_many(new_states)).inserted_ids
+                next_state_ids.extend(inserted_ids)
 
-                inserted_states = await State.find(
-                    In(State.id, inserted_ids)
-                ).to_list()
-
-                if len(inserted_states) != len(new_states):
-                    raise RuntimeError(f"Failed to insert all new states. Expected {len(new_states)} states, but only {len(inserted_states)} were inserted")
-
-                for inserted_state in inserted_states:
-                    background_tasks.add_task(create_next_state, inserted_state)
+        background_tasks.add_task(create_next_states, next_state_ids, state.identifier, state.namespace_name, state.graph_name, state.parents)
 
         return ExecutedResponseModel(status=StateStatusEnum.EXECUTED)
 

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

@@ -1,7 +1,9 @@
 import base64
+import time
+import asyncio
 
 from .base import BaseDatabaseModel
-from pydantic import Field, field_validator
+from pydantic import Field, field_validator, PrivateAttr
 from typing import Optional, List
 from ..graph_template_validation_status import GraphTemplateValidationStatus
 from ..node_template_model import NodeTemplate
@@ -17,6 +19,7 @@ class GraphTemplate(BaseDatabaseModel):
     validation_status: GraphTemplateValidationStatus = Field(..., description="Validation status of the graph")
     validation_errors: Optional[List[str]] = Field(None, description="Validation errors of the graph")
     secrets: Dict[str, str] = Field(default_factory=dict, description="Secrets of the graph")
+    _node_by_identifier: Dict[str, NodeTemplate] | None = PrivateAttr(default=None)
 
     class Settings:
         indexes = [
@@ -27,12 +30,18 @@ class Settings:
             )
         ]
 
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def _build_node_by_identifier(self) -> None:
+        self._node_by_identifier = {node.identifier: node for node in self.nodes}
+
     def get_node_by_identifier(self, identifier: str) -> NodeTemplate | None:
         """Get a node by its identifier using O(1) dictionary lookup."""
-        for node in self.nodes:
-            if node.identifier == identifier:
-                return node
-        return None
+        if self._node_by_identifier is None:
+            self._build_node_by_identifier()
+
+        return self._node_by_identifier.get(identifier) # type: ignore
 
     @field_validator('secrets')
     @classmethod
@@ -78,4 +87,40 @@ def get_secret(self, secret_name: str) -> str | None:
             return None
         if secret_name not in self.secrets:
             return None
-        return get_encrypter().decrypt(self.secrets[secret_name])
+        return get_encrypter().decrypt(self.secrets[secret_name])
+
+    def is_valid(self) -> bool:
+        return self.validation_status == GraphTemplateValidationStatus.VALID
+
+    def is_validating(self) -> bool:
+        return self.validation_status in (GraphTemplateValidationStatus.ONGOING, GraphTemplateValidationStatus.PENDING)
+
+    @staticmethod
+    async def get(namespace: str, graph_name: str) -> "GraphTemplate":
+        graph_template = await GraphTemplate.find_one(GraphTemplate.namespace == namespace, GraphTemplate.name == graph_name)
+        if not graph_template:
+            raise ValueError(f"Graph template not found for namespace: {namespace} and graph name: {graph_name}")
+        return graph_template
+
+    @staticmethod
+    async def get_valid(namespace: str, graph_name: str, polling_interval: float = 1.0, timeout: float = 300.0) -> "GraphTemplate":
+        # Validate polling_interval and timeout
+        if polling_interval <= 0:
+            raise ValueError("polling_interval must be positive")
+        if timeout <= 0:
+            raise ValueError("timeout must be positive")
+
+        # Coerce polling_interval to a sensible minimum
+        if polling_interval < 0.1:
+            polling_interval = 0.1
+
+        start_time = time.monotonic()
+        while time.monotonic() - start_time < timeout:
+            graph_template = await GraphTemplate.get(namespace, graph_name)
+            if graph_template.is_valid():
+                return graph_template
+            if graph_template.is_validating():
+                await asyncio.sleep(polling_interval)
+            else:
+                raise ValueError(f"Graph template is in a non-validating state: {graph_template.validation_status.value} for namespace: {namespace} and graph name: {graph_name}")
+        raise ValueError(f"Graph template is not valid for namespace: {namespace} and graph name: {graph_name} after {timeout} seconds")

state-manager/app/models/state_status_enum.py

@@ -11,4 +11,5 @@ class StateStatusEnum(str, Enum):
     TIMEDOUT = 'TIMEDOUT'
     ERRORED = 'ERRORED'
     CANCELLED = 'CANCELLED'
-    SUCCESS = 'SUCCESS'
+    SUCCESS = 'SUCCESS'
+    NEXT_CREATED_ERROR = 'NEXT_CREATED_ERROR'