11import os
22import aiohttp
3+ import asyncio
4+ import time
5+
6+ from typing import Any
37from 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