Skip to content

Commit fc94542

Browse files
authored
Enhance StateManager with Upsert and Get Graph functions (#215)
* Enhance StateManager with new graph management methods - Added methods to retrieve and upsert graphs in the StateManager class. - Implemented asynchronous functionality for graph retrieval and validation. - Updated dependencies in the Python SDK, including pytest and pytest-cov for testing. - Introduced new packages: colorama, coverage, pygments, and pytest-cov with their respective versions. - Improved documentation for new methods and their usage examples. * 0.0.7b6 -> 0.0.7b7 * fixing review comments
1 parent 5b7076d commit fc94542

3 files changed

Lines changed: 328 additions & 78 deletions

File tree

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
version = "0.0.7b6"
1+
version = "0.0.7b7"

python-sdk/exospherehost/statemanager.py

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
import os
22
import aiohttp
3+
import asyncio
4+
import time
5+
6+
from typing import Any
37
from pydantic import BaseModel
48

59

@@ -56,6 +60,12 @@ def _set_config_from_env(self):
5660

5761
def _get_trigger_state_endpoint(self, graph_name: str):
5862
return f"{self._state_manager_uri}/{self._state_manager_version}/namespace/{self._namespace}/graph/{graph_name}/trigger"
63+
64+
def _get_upsert_graph_endpoint(self, graph_name: str):
65+
return f"{self._state_manager_uri}/{self._state_manager_version}/namespace/{self._namespace}/graph/{graph_name}"
66+
67+
def _get_get_graph_endpoint(self, graph_name: str):
68+
return f"{self._state_manager_uri}/{self._state_manager_version}/namespace/{self._namespace}/graph/{graph_name}"
5969

6070
async def trigger(self, graph_name: str, state: TriggerState | None = None, states: list[TriggerState] | None = None):
6171
"""
@@ -121,3 +131,97 @@ async def trigger(self, graph_name: str, state: TriggerState | None = None, stat
121131
if response.status != 200:
122132
raise Exception(f"Failed to trigger state: {response.status} {await response.text()}")
123133
return await response.json()
134+
135+
async def get_graph(self, graph_name: str):
136+
"""
137+
Retrieve information about a specific graph from the state manager.
138+
139+
This method fetches the current state and configuration of a graph,
140+
including its validation status, nodes, and any validation errors.
141+
142+
Args:
143+
graph_name (str): The name of the graph to retrieve.
144+
145+
Returns:
146+
dict: The JSON response from the state manager API containing the
147+
graph information, including validation status, nodes, and errors.
148+
149+
Raises:
150+
Exception: If the API request fails with a non-200 status code. The exception
151+
message includes the HTTP status code and response text for debugging.
152+
153+
Example:
154+
```python
155+
# Get information about a specific graph
156+
graph_info = await state_manager.get_graph("my-workflow-graph")
157+
print(f"Graph status: {graph_info['validation_status']}")
158+
```
159+
"""
160+
endpoint = self._get_get_graph_endpoint(graph_name)
161+
headers = {
162+
"x-api-key": self._key
163+
}
164+
async with aiohttp.ClientSession() as session:
165+
async with session.get(endpoint, headers=headers) as response: # type: ignore
166+
if response.status != 200:
167+
raise Exception(f"Failed to get graph: {response.status} {await response.text()}")
168+
return await response.json()
169+
170+
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):
171+
"""
172+
Create or update a graph in the state manager with validation.
173+
174+
This method sends a graph definition to the state manager API for creation
175+
or update. After submission, it polls the API to wait for graph validation
176+
to complete, ensuring the graph is properly configured before returning.
177+
178+
Args:
179+
graph_name (str): The name of the graph to create or update.
180+
graph_nodes (list[dict[str, Any]]): A list of node definitions that make up
181+
the graph. Each node should contain the necessary configuration for
182+
the graph execution engine.
183+
secrets (dict[str, str]): A dictionary of secret values that will be
184+
available to the graph during execution. Keys are secret names and
185+
values are the secret values.
186+
validation_timeout (int, optional): Maximum time in seconds to wait for
187+
graph validation to complete. Defaults to 60.
188+
polling_interval (int, optional): Time in seconds between validation
189+
status checks. Defaults to 1.
190+
191+
Returns:
192+
dict: The JSON response from the state manager API containing the
193+
validated graph information.
194+
195+
Raises:
196+
Exception: If the API request fails with a non-201 status code, if graph
197+
validation times out, or if validation fails. The exception message
198+
includes relevant error details for debugging.
199+
"""
200+
endpoint = self._get_upsert_graph_endpoint(graph_name)
201+
headers = {
202+
"x-api-key": self._key
203+
}
204+
body = {
205+
"secrets": secrets,
206+
"nodes": graph_nodes
207+
}
208+
async with aiohttp.ClientSession() as session:
209+
async with session.put(endpoint, json=body, headers=headers) as response: # type: ignore
210+
if response.status not in [200, 201]:
211+
raise Exception(f"Failed to upsert graph: {response.status} {await response.text()}")
212+
graph = await response.json()
213+
214+
validation_state = graph["validation_status"]
215+
216+
start_time = time.monotonic()
217+
while validation_state == "PENDING":
218+
if time.monotonic() - start_time > validation_timeout:
219+
raise Exception(f"Graph validation check timed out after {validation_timeout} seconds")
220+
await asyncio.sleep(polling_interval)
221+
graph = await self.get_graph(graph_name)
222+
validation_state = graph["validation_status"]
223+
224+
if validation_state != "VALID":
225+
raise Exception(f"Graph validation failed: {graph['validation_status']} and errors: {graph['validation_errors']}")
226+
227+
return graph

0 commit comments

Comments
 (0)