Adding background task to create next state

state-manager/app/controller/executed_state.py

@@ -1,14 +1,15 @@
 from app.models.executed_models import ExecutedRequestModel, ExecutedResponseModel
 from bson import ObjectId
-from fastapi import HTTPException, status
+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
 
 logger = LogsManager().get_logger()
 
-async def executed_state(namespace_name: str, state_id: ObjectId, body: ExecutedRequestModel, x_exosphere_request_id: str) -> ExecutedResponseModel:
+async def executed_state(namespace_name: str, state_id: ObjectId, body: ExecutedRequestModel, x_exosphere_request_id: str, background_tasks: BackgroundTasks) -> ExecutedResponseModel:
 
     try:
         logger.info(f"Executed state {state_id} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
@@ -20,9 +21,36 @@ async def executed_state(namespace_name: str, state_id: ObjectId, body: Executed
         if state.status != StateStatusEnum.QUEUED:
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="State is not queued")
 
-        await State.find_one(State.id == state_id).set(
-            {"status": StateStatusEnum.EXECUTED, "outputs": body.outputs}
-        )
+        if len(body.outputs) == 0:
+            await State.find_one(State.id == state_id).set(
+                {"status": StateStatusEnum.EXECUTED, "outputs": {}, "parents": {**state.parents, state.identifier: ObjectId(state.id)}}
+            )
+
+            background_tasks.add_task(create_next_state, state)
+
+        else:
+            await State.find_one(State.id == state_id).set(
+                {"status": StateStatusEnum.EXECUTED, "outputs": body.outputs[0], "parents": {**state.parents, state.identifier: ObjectId(state.id)}}
+            )
+            background_tasks.add_task(create_next_state, state)
+
+            for output in body.outputs[1:]:
+                new_state = State(
+                    node_name=state.node_name,
+                    namespace_name=state.namespace_name,
+                    identifier=state.identifier,
+                    graph_name=state.graph_name,
+                    status=StateStatusEnum.CREATED,
+                    inputs=state.inputs,
+                    outputs=output,
+                    error=None,
+                    parents={
+                        **state.parents,
+                        state.identifier: ObjectId(state.id)
+                    }
+                )
+                await new_state.save()
+                background_tasks.add_task(create_next_state, new_state)
 
         return ExecutedResponseModel(status=StateStatusEnum.EXECUTED)
 

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

@@ -1,3 +1,4 @@
+from bson import ObjectId
 from .base import BaseDatabaseModel
 from ..state_status_enum import StateStatusEnum
 from pydantic import Field
@@ -13,4 +14,5 @@ class State(BaseDatabaseModel):
     status: StateStatusEnum = Field(..., description="Status of the state")
     inputs: dict[str, Any] = Field(..., description="Inputs of the state")
     outputs: dict[str, Any] = Field(..., description="Outputs of the state")
-    error: Optional[str] = Field(None, description="Error message")
+    error: Optional[str] = Field(None, description="Error message")
+    parents: dict[str, ObjectId] = Field(default_factory=dict, description="Parents of the state")

state-manager/app/routes.py

@@ -81,7 +81,7 @@ async def create_state(namespace_name: str, graph_name: str, body: CreateRequest
     response_description="State executed successfully",
     tags=["state"]
 )
-async def executed_state_route(namespace_name: str, state_id: str, body: ExecutedRequestModel, request: Request, api_key: str = Depends(check_api_key)):
+async def executed_state_route(namespace_name: str, state_id: str, body: ExecutedRequestModel, request: Request, background_tasks: BackgroundTasks, api_key: str = Depends(check_api_key)):
 
     x_exosphere_request_id = getattr(request.state, "x_exosphere_request_id", str(uuid4()))
 
@@ -91,7 +91,7 @@ async def executed_state_route(namespace_name: str, state_id: str, body: Execute
         logger.error(f"API key is invalid for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id)
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key")
 
-    return await executed_state(namespace_name, ObjectId(state_id), body, x_exosphere_request_id)
+    return await executed_state(namespace_name, ObjectId(state_id), body, x_exosphere_request_id, background_tasks)
 
 
 @router.post(

state-manager/app/tasks/create_next_state.py

@@ -0,0 +1,124 @@
+import asyncio
+import time
+
+from bson import ObjectId
+
+from app.models.db.state import State
+from app.models.db.graph_template_model import GraphTemplate
+from app.models.graph_template_validation_status import GraphTemplateValidationStatus
+from app.models.db.registered_node import RegisteredNode
+from app.models.state_status_enum import StateStatusEnum
+
+from json_schema_to_pydantic import create_model
+
+async def create_next_state(state: State):
+    graph_template = None
+
+    try:
+        start_time = time.time()
+        timeout_seconds = 300  # 5 minutes
+
+        while True:
+            graph_template = await GraphTemplate.find_one(GraphTemplate.name == state.graph_name, GraphTemplate.namespace == state.namespace_name)
+            if not graph_template:
+                raise Exception(f"Graph template {state.graph_name} not found")
+            if graph_template.validation_status == GraphTemplateValidationStatus.VALID:
+                break
+
+            # Check if we've exceeded the timeout
+            if time.time() - start_time > timeout_seconds:
+                raise Exception(f"Timeout waiting for graph template {state.graph_name} to become valid after {timeout_seconds} seconds")
+
+            await asyncio.sleep(1)
+
+        node_template = graph_template.get_node_by_identifier(state.identifier)
+        if not node_template:
+            raise Exception(f"Node template {state.identifier} not found")
+
+        next_node_identifier = node_template.next_nodes
+        if not next_node_identifier:
+            raise Exception(f"Node template {state.identifier} has no next nodes")
+
+        cache_states = {}         
+
+        for identifier in next_node_identifier:
+            next_node_template = graph_template.get_node_by_identifier(identifier)
+            if not next_node_template:
+                continue
+
+            registered_node = await RegisteredNode.find_one(RegisteredNode.name == next_node_template.node_name, RegisteredNode.namespace == next_node_template.namespace)
+
+            if not registered_node:
+                raise Exception(f"Registered node {next_node_template.node_name} not found")
+
+            next_node_input_model = create_model(registered_node.inputs_schema)
+            next_node_input_data = {}
+
+            for field_name, _ in next_node_input_model.model_fields.items():
+                temporary_input = next_node_template.inputs[field_name]
+                splits = temporary_input.split("${{")
+
+                if len(splits) == 0:
+                    next_node_input_data[field_name] = temporary_input
+                    continue
+
+                constructed_string = ""
+                for split in splits:
+                    if "}}" in split:
+                        placeholder_content = split.split("}}")[0]
+                        parts = [p.strip() for p in placeholder_content.split('.')]
+
+                        if len(parts) != 3 or parts[1] != 'outputs':
+                            raise Exception(f"Invalid input placeholder format: '{placeholder_content}' for field {field_name}")
+
+                        input_identifier = parts[0]
+                        input_field = parts[2]
+
+                        parent_id = state.parents.get(input_identifier)
+
+                        if not parent_id:
+                            raise Exception(f"Parent identifier '{input_identifier}' not found in state parents.")
+
+                        if parent_id not in cache_states:
+                            dependent_state = await State.get(ObjectId(parent_id))
+                            if not dependent_state:
+                                raise Exception(f"Dependent state {input_identifier} not found")
+                            cache_states[parent_id] = dependent_state
+                        else:
+                            dependent_state = cache_states[parent_id]
+
+                        if input_field not in dependent_state.outputs:
+                            raise Exception(f"Input field {input_field} not found in dependent state {input_identifier}")
+
+                        constructed_string += dependent_state.outputs[input_field] + split.split("}}")[1]
+
+                    else:
+                        constructed_string += split
+
+                next_node_input_data[field_name] = constructed_string
+
+            new_state = State(
+                node_name=next_node_template.node_name,
+                namespace_name=next_node_template.namespace,
+                identifier=next_node_template.identifier,
+                graph_name=state.graph_name,
+                status=StateStatusEnum.CREATED,
+                inputs=next_node_input_data,
+                outputs={},
+                error=None,
+                parents={
+                    **state.parents,
+                    next_node_template.identifier: ObjectId(state.id)
+                }
+            )
+
+            await new_state.save()
+
+        state.status = StateStatusEnum.SUCCESS
+        await state.save()
+
+    except Exception as e:
+        state.status = StateStatusEnum.ERRORED
+        state.error = str(e)
+        await state.save()
+        return