Enhance StateManager with Upsert and Get Graph functions

python-sdk/exospherehost/_version.py

@@ -1 +1 @@
-version = "0.0.7b6"
+version = "0.0.7b7"

python-sdk/exospherehost/statemanager.py

@@ -1,5 +1,9 @@
 import os
 import aiohttp
+import asyncio
+import time
+
+from typing import Any
 from pydantic import BaseModel
 
 
@@ -56,6 +60,12 @@ def _set_config_from_env(self):
 
     def _get_trigger_state_endpoint(self, graph_name: str):
         return f"{self._state_manager_uri}/{self._state_manager_version}/namespace/{self._namespace}/graph/{graph_name}/trigger"
+
+    def _get_upsert_graph_endpoint(self, graph_name: str):
+        return f"{self._state_manager_uri}/{self._state_manager_version}/namespace/{self._namespace}/graph/{graph_name}"
+
+    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, state: TriggerState | None = None, states: list[TriggerState] | None = None):
         """
@@ -121,3 +131,97 @@ async def trigger(self, graph_name: str, state: TriggerState | None = None, stat
                 if response.status != 200:
                     raise Exception(f"Failed to trigger state: {response.status} {await response.text()}")
                 return await response.json()
+
+    async def get_graph(self, graph_name: str):
+        """
+        Retrieve information about a specific graph from the state manager.
+
+        This method fetches the current state and configuration of a graph,
+        including its validation status, nodes, and any validation errors.
+
+        Args:
+            graph_name (str): The name of the graph to retrieve.
+
+        Returns:
+            dict: The JSON response from the state manager API containing the
+                graph information, including validation status, nodes, and errors.
+
+        Raises:
+            Exception: If the API request fails with a non-200 status code. The exception
+                message includes the HTTP status code and response text for debugging.
+
+        Example:
+            ```python
+            # Get information about a specific graph
+            graph_info = await state_manager.get_graph("my-workflow-graph")
+            print(f"Graph status: {graph_info['validation_status']}")
+            ```
+        """
+        endpoint = self._get_get_graph_endpoint(graph_name)
+        headers = {
+            "x-api-key": self._key
+        }
+        async with aiohttp.ClientSession() as session:
+            async with session.get(endpoint, headers=headers) as response: # type: ignore
+                if response.status != 200:
+                    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], validation_timeout: int = 60, polling_interval: int = 1):
+        """
+        Create or update a graph in the state manager with validation.
+
+        This method sends a graph definition to the state manager API for creation
+        or update. After submission, it polls the API to wait for graph validation
+        to complete, ensuring the graph is properly configured before returning.
+
+        Args:
+            graph_name (str): The name of the graph to create or update.
+            graph_nodes (list[dict[str, Any]]): A list of node definitions that make up
+                the graph. Each node should contain the necessary configuration for
+                the graph execution engine.
+            secrets (dict[str, str]): A dictionary of secret values that will be
+                available to the graph during execution. Keys are secret names and
+                values are the secret values.
+            validation_timeout (int, optional): Maximum time in seconds to wait for
+                graph validation to complete. Defaults to 60.
+            polling_interval (int, optional): Time in seconds between validation
+                status checks. Defaults to 1.
+
+        Returns:
+            dict: The JSON response from the state manager API containing the
+                validated graph information.
+
+        Raises:
+            Exception: If the API request fails with a non-201 status code, if graph
+                validation times out, or if validation fails. The exception message
+                includes relevant error details for debugging.
+        """
+        endpoint = self._get_upsert_graph_endpoint(graph_name)
+        headers = {
+            "x-api-key": self._key
+        }
+        body = {
+            "secrets": secrets,
+            "nodes": graph_nodes
+        }
+        async with aiohttp.ClientSession() as session:
+            async with session.put(endpoint, json=body, headers=headers) as response: # type: ignore
+                if response.status not in [200, 201]:
+                    raise Exception(f"Failed to upsert graph: {response.status} {await response.text()}")
+                graph = await response.json()
+
+        validation_state = graph["validation_status"]
+
+        start_time = time.monotonic()
+        while validation_state == "PENDING":
+            if time.monotonic() - start_time > validation_timeout:
+                raise Exception(f"Graph validation check timed out after {validation_timeout} seconds")
+            await asyncio.sleep(polling_interval)
+            graph = await self.get_graph(graph_name)
+            validation_state = graph["validation_status"]
+
+        if validation_state != "VALID":
+            raise Exception(f"Graph validation failed: {graph['validation_status']} and errors: {graph['validation_errors']}")
+
+        return graph