13
13
from autogpt_libs .utils .cache import thread_cached
14
14
from dotenv import load_dotenv
15
15
from prisma .enums import NotificationType
16
- from pydantic import BaseModel
16
+ from pydantic import BaseModel , ValidationError
17
17
from sqlalchemy import MetaData , create_engine
18
18
19
19
from backend .data .block import BlockInput
@@ -72,7 +72,7 @@ def get_notification_client():
72
72
73
73
74
74
def execute_graph (** kwargs ):
75
- args = ExecutionJobArgs (** kwargs )
75
+ args = GraphExecutionJobArgs (** kwargs )
76
76
try :
77
77
log (f"Executing recurring job for graph #{ args .graph_id } " )
78
78
execution_utils .add_graph_execution (
@@ -140,22 +140,24 @@ class Jobstores(Enum):
140
140
WEEKLY_NOTIFICATIONS = "weekly_notifications"
141
141
142
142
143
- class ExecutionJobArgs (BaseModel ):
143
+ class GraphExecutionJobArgs (BaseModel ):
144
144
graph_id : str
145
145
input_data : BlockInput
146
146
user_id : str
147
147
graph_version : int
148
148
cron : str
149
149
150
150
151
- class ExecutionJobInfo ( ExecutionJobArgs ):
151
+ class GraphExecutionJobInfo ( GraphExecutionJobArgs ):
152
152
id : str
153
153
name : str
154
154
next_run_time : str
155
155
156
156
@staticmethod
157
- def from_db (job_args : ExecutionJobArgs , job_obj : JobObj ) -> "ExecutionJobInfo" :
158
- return ExecutionJobInfo (
157
+ def from_db (
158
+ job_args : GraphExecutionJobArgs , job_obj : JobObj
159
+ ) -> "GraphExecutionJobInfo" :
160
+ return GraphExecutionJobInfo (
159
161
id = job_obj .id ,
160
162
name = job_obj .name ,
161
163
next_run_time = job_obj .next_run_time .isoformat (),
@@ -269,15 +271,15 @@ def cleanup(self):
269
271
self .scheduler .shutdown (wait = False )
270
272
271
273
@expose
272
- def add_execution_schedule (
274
+ def add_graph_execution_schedule (
273
275
self ,
274
276
graph_id : str ,
275
277
graph_version : int ,
276
278
cron : str ,
277
279
input_data : BlockInput ,
278
280
user_id : str ,
279
- ) -> ExecutionJobInfo :
280
- job_args = ExecutionJobArgs (
281
+ ) -> GraphExecutionJobInfo :
282
+ job_args = GraphExecutionJobArgs (
281
283
graph_id = graph_id ,
282
284
input_data = input_data ,
283
285
user_id = user_id ,
@@ -292,40 +294,46 @@ def add_execution_schedule(
292
294
jobstore = Jobstores .EXECUTION .value ,
293
295
)
294
296
log (f"Added job { job .id } with cron schedule '{ cron } ' input data: { input_data } " )
295
- return ExecutionJobInfo .from_db (job_args , job )
297
+ return GraphExecutionJobInfo .from_db (job_args , job )
296
298
297
299
@expose
298
- def delete_schedule (self , schedule_id : str , user_id : str ) -> ExecutionJobInfo :
300
+ def delete_graph_execution_schedule (
301
+ self , schedule_id : str , user_id : str
302
+ ) -> GraphExecutionJobInfo :
299
303
job = self .scheduler .get_job (schedule_id , jobstore = Jobstores .EXECUTION .value )
300
304
if not job :
301
305
log (f"Job { schedule_id } not found." )
302
306
raise ValueError (f"Job #{ schedule_id } not found." )
303
307
304
- job_args = ExecutionJobArgs (** job .kwargs )
308
+ job_args = GraphExecutionJobArgs (** job .kwargs )
305
309
if job_args .user_id != user_id :
306
310
raise ValueError ("User ID does not match the job's user ID." )
307
311
308
312
log (f"Deleting job { schedule_id } " )
309
313
job .remove ()
310
314
311
- return ExecutionJobInfo .from_db (job_args , job )
315
+ return GraphExecutionJobInfo .from_db (job_args , job )
312
316
313
317
@expose
314
- def get_execution_schedules (
318
+ def get_graph_execution_schedules (
315
319
self , graph_id : str | None = None , user_id : str | None = None
316
- ) -> list [ExecutionJobInfo ]:
320
+ ) -> list [GraphExecutionJobInfo ]:
321
+ jobs : list [JobObj ] = self .scheduler .get_jobs (jobstore = Jobstores .EXECUTION .value )
317
322
schedules = []
318
- for job in self . scheduler . get_jobs ( jobstore = Jobstores . EXECUTION . value ) :
319
- logger .info (
323
+ for job in jobs :
324
+ logger .debug (
320
325
f"Found job { job .id } with cron schedule { job .trigger } and args { job .kwargs } "
321
326
)
322
- job_args = ExecutionJobArgs (** job .kwargs )
327
+ try :
328
+ job_args = GraphExecutionJobArgs .model_validate (job .kwargs )
329
+ except ValidationError :
330
+ continue
323
331
if (
324
332
job .next_run_time is not None
325
333
and (graph_id is None or job_args .graph_id == graph_id )
326
334
and (user_id is None or job_args .user_id == user_id )
327
335
):
328
- schedules .append (ExecutionJobInfo .from_db (job_args , job ))
336
+ schedules .append (GraphExecutionJobInfo .from_db (job_args , job ))
329
337
return schedules
330
338
331
339
@expose
@@ -346,6 +354,6 @@ class SchedulerClient(AppServiceClient):
346
354
def get_service_type (cls ):
347
355
return Scheduler
348
356
349
- add_execution_schedule = endpoint_to_async (Scheduler .add_execution_schedule )
350
- delete_schedule = endpoint_to_async (Scheduler .delete_schedule )
351
- get_execution_schedules = endpoint_to_async (Scheduler .get_execution_schedules )
357
+ add_execution_schedule = endpoint_to_async (Scheduler .add_graph_execution_schedule )
358
+ delete_schedule = endpoint_to_async (Scheduler .delete_graph_execution_schedule )
359
+ get_execution_schedules = endpoint_to_async (Scheduler .get_graph_execution_schedules )
0 commit comments