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'