Update version to 0.0.2b5 and introduce new models for graph execution

- Incremented the version in _version.py to 0.0.2b5.
- Added new models in models.py for GraphNode, Unites, RetryPolicy, and StoreConfig, enhancing the structure and validation of graph execution parameters.
- Refactored StateManager methods to utilize the new models, improving type safety and clarity in the trigger and upsert_graph functions.
- Updated trigger_graph to include a start_delay parameter for delayed execution, enhancing flexibility in graph triggering.

Author: NiveditJain

python-sdk/exospherehost/_version.py

@@ -1 +1 @@
-version = "0.0.2b4"
+version = "0.0.2b5"

python-sdk/exospherehost/models.py

@@ -0,0 +1,158 @@
+from pydantic import BaseModel, Field, field_validator
+from typing import Any, Optional, List
+from enum import Enum
+
+
+class UnitesStrategyEnum(str, Enum):
+    ALL_SUCCESS = "ALL_SUCCESS"
+    ALL_DONE = "ALL_DONE"
+
+
+class UnitesModel(BaseModel):
+    identifier: str = Field(..., description="Identifier of the node")
+    strategy: UnitesStrategyEnum = Field(default=UnitesStrategyEnum.ALL_SUCCESS, description="Strategy of the unites")
+
+
+class GraphNodeModel(BaseModel):
+    node_name: str = Field(..., description="Name of the node")
+    namespace: str = Field(..., description="Namespace of the node")
+    identifier: str = Field(..., description="Identifier of the node")
+    inputs: dict[str, Any] = Field(..., description="Inputs of the node")
+    next_nodes: Optional[List[str]] = Field(None, description="Next nodes to execute")
+    unites: Optional[UnitesModel] = Field(None, description="Unites of the node")
+
+    @field_validator('node_name')
+    @classmethod
+    def validate_node_name(cls, v: str) -> str:
+        trimmed_v = v.strip()
+        if trimmed_v == "" or trimmed_v is None:
+            raise ValueError("Node name cannot be empty")
+        return trimmed_v
+    
+    @field_validator('identifier')
+    @classmethod
+    def validate_identifier(cls, v: str) -> str:
+        trimmed_v = v.strip()
+        if trimmed_v == "" or trimmed_v is None:
+            raise ValueError("Node identifier cannot be empty")
+        elif trimmed_v == "store":
+            raise ValueError("Node identifier cannot be reserved word 'store'")
+        return trimmed_v
+    
+    @field_validator('next_nodes')
+    @classmethod
+    def validate_next_nodes(cls, v: Optional[List[str]]) -> Optional[List[str]]:
+        identifiers = set()
+        errors = []
+        trimmed_v = []
+
+        if v is not None:
+            for next_node_identifier in v:
+                trimmed_next_node_identifier = next_node_identifier.strip()
+                
+                if trimmed_next_node_identifier == "" or trimmed_next_node_identifier is None:
+                    errors.append("Next node identifier cannot be empty")
+                    continue
+
+                if trimmed_next_node_identifier in identifiers:
+                    errors.append(f"Next node identifier {trimmed_next_node_identifier} is not unique")
+                    continue
+
+                identifiers.add(trimmed_next_node_identifier)
+                trimmed_v.append(trimmed_next_node_identifier)
+        if errors:
+            raise ValueError("\n".join(errors))
+        return trimmed_v
+    
+    @field_validator('unites')
+    @classmethod
+    def validate_unites(cls, v: Optional[UnitesModel]) -> Optional[UnitesModel]:
+        trimmed_v = v
+        if v is not None:
+            trimmed_v = UnitesModel(identifier=v.identifier.strip(), strategy=v.strategy)
+            if trimmed_v.identifier == "" or trimmed_v.identifier is None:
+                raise ValueError("Unites identifier cannot be empty")
+        return trimmed_v
+
+
+class RetryStrategyEnum(str, Enum):
+    EXPONENTIAL = "EXPONENTIAL"
+    EXPONENTIAL_FULL_JITTER = "EXPONENTIAL_FULL_JITTER"
+    EXPONENTIAL_EQUAL_JITTER = "EXPONENTIAL_EQUAL_JITTER"
+
+    LINEAR = "LINEAR"
+    LINEAR_FULL_JITTER = "LINEAR_FULL_JITTER"
+    LINEAR_EQUAL_JITTER = "LINEAR_EQUAL_JITTER"
+
+    FIXED = "FIXED"
+    FIXED_FULL_JITTER = "FIXED_FULL_JITTER"
+    FIXED_EQUAL_JITTER = "FIXED_EQUAL_JITTER"
+
+
+class RetryPolicyModel(BaseModel):
+    max_retries: int = Field(default=3, description="The maximum number of retries", ge=0)
+    strategy: RetryStrategyEnum = Field(default=RetryStrategyEnum.EXPONENTIAL, description="The method of retry")
+    backoff_factor: int = Field(default=2000, description="The backoff factor in milliseconds (default: 2000 = 2 seconds)", gt=0)
+    exponent: int = Field(default=2, description="The exponent for the exponential retry strategy", gt=0)
+    max_delay: int | None = Field(default=None, description="The maximum delay in milliseconds (no default limit when None)", gt=0)
+
+
+class StoreConfigModel(BaseModel):
+    required_keys: list[str] = Field(default_factory=list, description="Required keys of the store")
+    default_values: dict[str, str] = Field(default_factory=dict, description="Default values of the store")
+
+    @field_validator("required_keys")
+    def validate_required_keys(cls, v: list[str]) -> list[str]:
+        errors = []
+        keys = set()
+        trimmed_keys = []
+        
+        for key in v:
+            trimmed_key = key.strip() if key is not None else ""
+            
+            if trimmed_key == "":
+                errors.append("Key cannot be empty or contain only whitespace")
+                continue
+                
+            if '.' in trimmed_key:
+                errors.append(f"Key '{trimmed_key}' cannot contain '.' character")
+                continue
+                
+            if trimmed_key in keys:
+                errors.append(f"Key '{trimmed_key}' is duplicated")
+                continue
+                
+            keys.add(trimmed_key)
+            trimmed_keys.append(trimmed_key)
+        
+        if len(errors) > 0:
+            raise ValueError("\n".join(errors))
+        return trimmed_keys
+
+    @field_validator("default_values")
+    def validate_default_values(cls, v: dict[str, str]) -> dict[str, str]:
+        errors = []
+        keys = set()
+        normalized_dict = {}
+        
+        for key, value in v.items():
+            trimmed_key = key.strip() if key is not None else ""
+            
+            if trimmed_key == "":
+                errors.append("Key cannot be empty or contain only whitespace")
+                continue
+                
+            if '.' in trimmed_key:
+                errors.append(f"Key '{trimmed_key}' cannot contain '.' character")
+                continue
+                
+            if trimmed_key in keys:
+                errors.append(f"Key '{trimmed_key}' is duplicated")
+                continue
+                
+            keys.add(trimmed_key)
+            normalized_dict[trimmed_key] = str(value)
+        
+        if len(errors) > 0:
+            raise ValueError("\n".join(errors))
+        return normalized_dict

python-sdk/exospherehost/statemanager.py

@@ -3,40 +3,7 @@
 import asyncio
 import time
 
-from typing import Any
-from pydantic import BaseModel
-
-
-class TriggerState(BaseModel):
-    """
-    Represents a trigger state for graph execution.
-    
-    A trigger state contains an identifier and a set of input parameters that
-    will be passed to the graph when it is triggered for execution.
-    
-    Attributes:
-        identifier (str): A unique identifier for this trigger state. This is used
-            to distinguish between different trigger states and may be used by the
-            graph to determine how to process the trigger.
-        inputs (dict[str, str]): A dictionary of input parameters that will be
-            passed to the graph. The keys are parameter names and values are
-            parameter values, both as strings.
-    
-    Example:
-        ```python
-        # Create a trigger state with identifier and inputs
-        trigger_state = TriggerState(
-            identifier="user-login",
-            inputs={
-                "user_id": "12345",
-                "session_token": "abc123def456",
-                "timestamp": "2024-01-15T10:30:00Z"
-            }
-        )
-        ```
-    """
-    identifier: str
-    inputs: dict[str, str]
+from .models import GraphNodeModel, RetryPolicyModel, StoreConfigModel
 
 
 class StateManager:
@@ -67,7 +34,7 @@ def _get_upsert_graph_endpoint(self, graph_name: str):
     def _get_get_graph_endpoint(self, graph_name: str):
         return f"{self._state_manager_uri}/{self._state_manager_version}/namespace/{self._namespace}/graph/{graph_name}"
 
-    async def trigger(self, graph_name: str, inputs: dict[str, str] | None = None, store: dict[str, str] | None = None):
+    async def trigger(self, graph_name: str, inputs: dict[str, str] | None = None, store: dict[str, str] | None = None, start_delay: int = 0):
         """
         Trigger execution of a graph.
         
@@ -82,7 +49,8 @@ async def trigger(self, graph_name: str, inputs: dict[str, str] | None = None, s
                 graph. Strings only.
             store (dict[str, str] | None): Optional key-value store that will be merged
                 into the graph-level store before execution (beta).
-        
+            start_delay (int): Optional delay in milliseconds before the graph starts execution.
+
         Returns:
             dict: JSON payload returned by the state-manager API.
         
@@ -108,6 +76,7 @@ async def trigger(self, graph_name: str, inputs: dict[str, str] | None = None, s
             store = {}
 
         body = {
+            "start_delay": start_delay,
             "inputs": inputs,
             "store": store
         }
@@ -156,7 +125,7 @@ async def get_graph(self, graph_name: str):
                     raise Exception(f"Failed to get graph: {response.status} {await response.text()}")
                 return await response.json()
 
-    async def upsert_graph(self, graph_name: str, graph_nodes: list[dict[str, Any]], secrets: dict[str, str], retry_policy: dict[str, Any] | None = None, store_config: dict[str, Any] | None = None, validation_timeout: int = 60, polling_interval: int = 1):
+    async def upsert_graph(self, graph_name: str, graph_nodes: list[GraphNodeModel], secrets: dict[str, str], retry_policy: RetryPolicyModel | None = None, store_config: StoreConfigModel | None = None, validation_timeout: int = 60, polling_interval: int = 1):
         """
         Create or update a graph definition.
 
@@ -169,10 +138,10 @@ async def upsert_graph(self, graph_name: str, graph_nodes: list[dict[str, Any]],
         
         Args:
             graph_name (str): Graph identifier.
-            graph_nodes (list[dict[str, Any]]): Graph node list.
+            graph_nodes (list[GraphNodeModel]): List of graph node models defining the workflow.
             secrets (dict[str, str]): Secrets available to all nodes.
-            retry_policy (dict[str, Any] | None): Optional per-node retry policy.
-            store_config (dict[str, Any] | None): Beta configuration for the
+            retry_policy (RetryPolicyModel | None): Optional per-node retry policy configuration.
+            store_config (StoreConfigModel | None): Beta configuration for the
                 graph-level store (schema is subject to change).
             validation_timeout (int): Seconds to wait for validation (default 60).
             polling_interval (int): Polling interval in seconds (default 1).
@@ -189,13 +158,13 @@ async def upsert_graph(self, graph_name: str, graph_nodes: list[dict[str, Any]],
         }
         body = {
             "secrets": secrets,
-            "nodes": graph_nodes
+            "nodes": [node.model_dump() for node in graph_nodes]
         }
 
         if retry_policy is not None:
-            body["retry_policy"] = retry_policy
+            body["retry_policy"] = retry_policy.model_dump()
         if store_config is not None:
-            body["store_config"] = store_config
+            body["store_config"] = store_config.model_dump()
 
         async with aiohttp.ClientSession() as session:
             async with session.put(endpoint, json=body, headers=headers) as response: # type: ignore

state-manager/app/controller/trigger_graph.py

@@ -9,7 +9,9 @@
 from app.models.db.graph_template_model import GraphTemplate
 from app.models.node_template_model import NodeTemplate
 from app.models.dependent_string import DependentString
+
 import uuid
+import time
 
 logger = LogsManager().get_logger()
 
@@ -96,6 +98,7 @@ async def trigger_graph(namespace_name: str, graph_name: str, body: TriggerGraph
             graph_name=graph_name,
             run_id=run_id,
             status=StateStatusEnum.CREATED,
+            enqueue_after=int(time.time() * 1000) + body.start_delay,
             inputs=inputs,
             outputs={},
             error=None

state-manager/app/models/trigger_model.py

@@ -4,6 +4,7 @@
 class TriggerGraphRequestModel(BaseModel):
     store: dict[str, str] = Field(default_factory=dict, description="Store for the runtime")
     inputs: dict[str, str] = Field(default_factory=dict, description="Inputs for the graph execution")
+    start_delay: int = Field(default=0, ge=0, description="Start delay in milliseconds")
 
 class TriggerGraphResponseModel(BaseModel):
     status: StateStatusEnum = Field(..., description="Status of the states")