-
Notifications
You must be signed in to change notification settings - Fork 3.5k
Expand file tree
/
Copy path__init__.py
More file actions
239 lines (186 loc) · 7.85 KB
/
__init__.py
File metadata and controls
239 lines (186 loc) · 7.85 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
"""
Sim Studio SDK for Python
Official Python SDK for Sim Studio, allowing you to execute workflows programmatically.
"""
from typing import Any, Dict, Optional
from dataclasses import dataclass
import requests
__version__ = "0.1.0"
__all__ = ["SimStudioClient", "SimStudioError", "WorkflowExecutionResult", "WorkflowStatus"]
@dataclass
class WorkflowExecutionResult:
"""Result of a workflow execution."""
success: bool
output: Optional[Any] = None
error: Optional[str] = None
logs: Optional[list] = None
metadata: Optional[Dict[str, Any]] = None
trace_spans: Optional[list] = None
total_duration: Optional[float] = None
@dataclass
class WorkflowStatus:
"""Status of a workflow."""
is_deployed: bool
deployed_at: Optional[str] = None
is_published: bool = False
needs_redeployment: bool = False
class SimStudioError(Exception):
"""Exception raised for Sim Studio API errors."""
def __init__(self, message: str, code: Optional[str] = None, status: Optional[int] = None):
super().__init__(message)
self.code = code
self.status = status
class SimStudioClient:
"""
Sim Studio API client for executing workflows programmatically.
Args:
api_key: Your Sim Studio API key
base_url: Base URL for the Sim Studio API (defaults to https://simstudio.ai)
"""
def __init__(self, api_key: str, base_url: str = "https://simstudio.ai"):
self.api_key = api_key
self.base_url = base_url.rstrip('/')
self._session = requests.Session()
self._session.headers.update({
'X-API-Key': self.api_key,
'Content-Type': 'application/json',
})
def execute_workflow(
self,
workflow_id: str,
input_data: Optional[Dict[str, Any]] = None,
timeout: float = 30.0
) -> WorkflowExecutionResult:
"""
Execute a workflow with optional input data.
Args:
workflow_id: The ID of the workflow to execute
input_data: Input data to pass to the workflow
timeout: Timeout in seconds (default: 30.0)
Returns:
WorkflowExecutionResult object containing the execution result
Raises:
SimStudioError: If the workflow execution fails
"""
url = f"{self.base_url}/api/workflows/{workflow_id}/execute"
try:
response = self._session.post(
url,
json=input_data or {},
timeout=timeout
)
if not response.ok:
try:
error_data = response.json()
error_message = error_data.get('error', f'HTTP {response.status_code}: {response.reason}')
error_code = error_data.get('code')
except (ValueError, KeyError):
error_message = f'HTTP {response.status_code}: {response.reason}'
error_code = None
raise SimStudioError(error_message, error_code, response.status_code)
result_data = response.json()
return WorkflowExecutionResult(
success=result_data['success'],
output=result_data.get('output'),
error=result_data.get('error'),
logs=result_data.get('logs'),
metadata=result_data.get('metadata'),
trace_spans=result_data.get('traceSpans'),
total_duration=result_data.get('totalDuration')
)
except requests.Timeout:
raise SimStudioError(f'Workflow execution timed out after {timeout} seconds', 'TIMEOUT')
except requests.RequestException as e:
raise SimStudioError(f'Failed to execute workflow: {str(e)}', 'EXECUTION_ERROR')
def get_workflow_status(self, workflow_id: str) -> WorkflowStatus:
"""
Get the status of a workflow (deployment status, etc.).
Args:
workflow_id: The ID of the workflow
Returns:
WorkflowStatus object containing the workflow status
Raises:
SimStudioError: If getting the status fails
"""
url = f"{self.base_url}/api/workflows/{workflow_id}/status"
try:
response = self._session.get(url)
if not response.ok:
try:
error_data = response.json()
error_message = error_data.get('error', f'HTTP {response.status_code}: {response.reason}')
error_code = error_data.get('code')
except (ValueError, KeyError):
error_message = f'HTTP {response.status_code}: {response.reason}'
error_code = None
raise SimStudioError(error_message, error_code, response.status_code)
status_data = response.json()
return WorkflowStatus(
is_deployed=status_data.get('isDeployed', False),
deployed_at=status_data.get('deployedAt'),
is_published=status_data.get('isPublished', False),
needs_redeployment=status_data.get('needsRedeployment', False)
)
except requests.RequestException as e:
raise SimStudioError(f'Failed to get workflow status: {str(e)}', 'STATUS_ERROR')
def validate_workflow(self, workflow_id: str) -> bool:
"""
Validate that a workflow is ready for execution.
Args:
workflow_id: The ID of the workflow
Returns:
True if the workflow is deployed and ready, False otherwise
"""
try:
status = self.get_workflow_status(workflow_id)
return status.is_deployed
except SimStudioError:
return False
def execute_workflow_sync(
self,
workflow_id: str,
input_data: Optional[Dict[str, Any]] = None,
timeout: float = 30.0
) -> WorkflowExecutionResult:
"""
Execute a workflow and poll for completion (useful for long-running workflows).
Note: Currently, the API is synchronous, so this method just calls execute_workflow.
In the future, if async execution is added, this method can be enhanced.
Args:
workflow_id: The ID of the workflow to execute
input_data: Input data to pass to the workflow
timeout: Timeout for the initial request in seconds
Returns:
WorkflowExecutionResult object containing the execution result
Raises:
SimStudioError: If the workflow execution fails
"""
# For now, the API is synchronous, so we just execute directly
# In the future, if async execution is added, this method can be enhanced
return self.execute_workflow(workflow_id, input_data, timeout)
def set_api_key(self, api_key: str) -> None:
"""
Update the API key.
Args:
api_key: New API key
"""
self.api_key = api_key
self._session.headers.update({'X-API-Key': api_key})
def set_base_url(self, base_url: str) -> None:
"""
Update the base URL.
Args:
base_url: New base URL
"""
self.base_url = base_url.rstrip('/')
def close(self) -> None:
"""Close the underlying HTTP session."""
self._session.close()
def __enter__(self):
"""Context manager entry."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Context manager exit."""
self.close()
# For backward compatibility
Client = SimStudioClient