Introduce new models for graph execution and added Start_delay (#348)

* 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.

* Refactor __init__.py to streamline exports and include new models

- Removed TriggerState from exports and added new models: UnitesStrategyEnum, UnitesModel, GraphNodeModel, RetryStrategyEnum, RetryPolicyModel, and StoreConfigModel.
- Updated __all__ to reflect the changes in exported components, enhancing the module's structure and usability.

* Refactor tests to remove TriggerState and utilize GraphNodeModel

- Removed instances of TriggerState from tests, replacing them with a dictionary structure for state representation.
- Updated test assertions to use GraphNodeModel for graph node definitions, enhancing clarity and consistency in the test suite.
- Adjusted related test cases to ensure compatibility with the new state representation, improving overall test robustness.

* Add class method decorators to StoreConfigModel validators

- Updated the `validate_required_keys` and `validate_default_values` methods in `StoreConfigModel` to be class methods, enhancing their functionality and consistency with the model's design.
- Introduced a new test file for comprehensive testing of `GraphNodeModel`, `StoreConfigModel`, and `StateManager`, including validation checks and default behavior, improving test coverage and reliability.

* Remove unused import of asyncio in test_models_and_statemanager_new.py to streamline the test file.

Author: NiveditJain

python-sdk/exospherehost/__init__.py

@@ -37,9 +37,10 @@ async def execute(self, inputs: Inputs) -> Outputs:
 from ._version import version as __version__
 from .runtime import Runtime
 from .node.BaseNode import BaseNode
-from .statemanager import StateManager, TriggerState
+from .statemanager import StateManager
 from .signals import PruneSignal, ReQueueAfterSignal
+from .models import UnitesStrategyEnum, UnitesModel, GraphNodeModel, RetryStrategyEnum, RetryPolicyModel, StoreConfigModel
 
 VERSION = __version__
 
-__all__ = ["Runtime", "BaseNode", "StateManager", "TriggerState", "VERSION", "PruneSignal", "ReQueueAfterSignal"]
+__all__ = ["Runtime", "BaseNode", "StateManager", "VERSION", "PruneSignal", "ReQueueAfterSignal", "UnitesStrategyEnum", "UnitesModel", "GraphNodeModel", "RetryStrategyEnum", "RetryPolicyModel", "StoreConfigModel"]

python-sdk/exospherehost/_version.py

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

python-sdk/exospherehost/models.py

@@ -0,0 +1,160 @@
+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")
+    @classmethod
+    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")
+    @classmethod
+    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

python-sdk/tests/test_coverage_additions.py

@@ -43,7 +43,7 @@ async def test_statemanager_trigger_defaults(monkeypatch):
     # Verify it sent empty inputs/store when omitted
     mock_session.post.assert_called_once()
     _, kwargs = mock_session.post.call_args
-    assert kwargs["json"] == {"inputs": {}, "store": {}}
+    assert kwargs["json"] == {"inputs": {}, "store": {}, "start_delay": 0}
 
 
 class _DummyNode(BaseNode):

python-sdk/tests/test_integration.py

@@ -2,7 +2,7 @@
 import asyncio
 from unittest.mock import AsyncMock, patch, MagicMock
 from pydantic import BaseModel
-from exospherehost import Runtime, BaseNode, StateManager, TriggerState
+from exospherehost import Runtime, BaseNode, StateManager
 
 
 def create_mock_aiohttp_session():
@@ -205,21 +205,26 @@ async def test_state_manager_graph_lifecycle(self, mock_env_vars):
             sm = StateManager(namespace="test_namespace")
 
             # Test graph creation
+            from exospherehost.models import GraphNodeModel
             graph_nodes = [
-                {"name": "IntegrationTestNode", "type": "test"}
+                GraphNodeModel(
+                    node_name="IntegrationTestNode",
+                    namespace="test_namespace",
+                    identifier="IntegrationTestNode",
+                    inputs={"type": "test"},
+                    next_nodes=None,
+                    unites=None
+                )
             ]
             secrets = {"api_key": "test_key", "database_url": "db://test"}
 
             result = await sm.upsert_graph("test_graph", graph_nodes, secrets, validation_timeout=10, polling_interval=0.1) # type: ignore
             assert result["validation_status"] == "VALID"
 
             # Test graph triggering
-            trigger_state = TriggerState(
-                identifier="test_trigger",
-                inputs={"user_id": "123", "action": "login"}
-            )
+            trigger_state = {"identifier": "test_trigger", "inputs": {"user_id": "123", "action": "login"}}
 
-            trigger_result = await sm.trigger("test_graph", inputs=trigger_state.inputs)
+            trigger_result = await sm.trigger("test_graph", inputs=trigger_state["inputs"])
             assert trigger_result == {"status": "triggered"}
 
 
@@ -448,10 +453,10 @@ async def test_state_manager_error_propagation(self, mock_env_vars):
             mock_session_class.return_value = mock_session
 
             sm = StateManager(namespace="error_test")
-            trigger_state = TriggerState(identifier="test", inputs={"key": "value"})
+            trigger_state = {"identifier": "test", "inputs": {"key": "value"}}
 
             with pytest.raises(Exception, match="Failed to trigger state: 404 Graph not found"):
-                await sm.trigger("nonexistent_graph", inputs=trigger_state.inputs)
+                await sm.trigger("nonexistent_graph", inputs=trigger_state["inputs"])
 
 
 class TestConcurrencyIntegration: