feat: TaskScheduler - backend mechanics
Rafael Uzarowski committed
Mar 21, 2025 at 18:23 UTC
132f792e98e94f57261a223f4b6d8d95a15ae4ff
10 files changed
+424
-5
docker/run/Dockerfile
+3
@@ -30,5 +30,8 @@ EXPOSE 22 80
30
31
RUN chmod +x /exe/initialize.sh /exe/run_A0.sh /exe/run_searxng.sh
32
33
+# This is important or cron will not execute the file
34
+RUN chmod 0644 /etc/cron.d/scheduler_tick
35
+
36
# initialize runtime and switch to supervisord
37
CMD ["/exe/initialize.sh", "$BRANCH"]
docker/run/DockerfileKali
+3
@@ -31,5 +31,8 @@ EXPOSE 22 80
31
32
RUN chmod +x /exe/initialize.sh /exe/run_A0.sh /exe/run_searxng.sh
33
34
+# This is important or cron will not execute the file
35
+RUN chmod 0644 /etc/cron.d/scheduler_tick
36
+
37
# initialize runtime and switch to supervisord
38
CMD ["/exe/initialize.sh", "$BRANCH"]
docker/run/fs/etc/cron.d/scheduler_tick
new
+5
@@ -0,0 +1,5 @@
1
+# TaskScheduler CronJob
2
+
3
+* * * * * root curl http://127.0.0.1:80/scheduler_tick
4
+
5
+# END
docker/run/fs/etc/supervisor/conf.d/supervisord.conf
+5
-1
@@ -7,6 +7,10 @@ pidfile=/var/run/supervisord.pid
7
exitcodes=0,2
8
directory=/
9
10
+[unix_http_server]
11
+file=/var/run/supervisor.sock
12
+chmod=0777
13
+
14
[rpcinterface:supervisor]
15
supervisor.rpcinterface_factory = supervisor.rpcinterface:make_main_rpcinterface
16
@@ -15,7 +19,7 @@ serverurl=unix:///var/run/supervisor.sock
19
20
[program:run_sshd]
21
command=/usr/sbin/sshd -D
18
-environment=LC_ALL=en_US.UTF-8
22
+environment=
23
stopwaitsecs=1
24
stdout_logfile=/dev/stdout
25
stdout_logfile_maxbytes=0
python/api/chat_remove.py
-2
@@ -1,6 +1,4 @@
1
from python.helpers.api import ApiHandler, Input, Output, Request, Response
2
-
3
-
2
from agent import AgentContext
3
from python.helpers import persist_chat
4
python/api/scheduler_adhoc.py_
new
+5
@@ -0,0 +1,5 @@
1
+## JSON Body
2
+# api_key
3
+# task_id
4
+# params
5
+# callback
python/api/scheduler_tick.py
new
+17
@@ -0,0 +1,17 @@
1
+from datetime import datetime
2
+
3
+from python.helpers.api import ApiHandler, Input, Output, Request
4
+from python.helpers.print_style import PrintStyle
5
+from python.helpers.task_scheduler import TaskScheduler
6
+
7
+
8
+class SchedulerTick(ApiHandler):
9
+ @classmethod
10
+ def requires_loopback(cls):
11
+ return True
12
+
13
+ async def process(self, input: Input, request: Request) -> Output:
14
+ # timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
15
+ # PrintStyle().print(f"Scheduler tick - API: {timestamp}")
16
+ await TaskScheduler.get().tick()
17
+ return {"scheduler": "tick"}
python/helpers/persist_chat.py
+2
-2
@@ -26,7 +26,7 @@ def save_tmp_chat(context: AgentContext):
26
27
def load_tmp_chats():
28
_convert_v080_chats()
29
- folders = files.list_files("tmp/chats/", "*")
29
+ folders = files.list_files(CHATS_FOLDER, "*")
30
json_files = []
31
for folder in folders:
32
json_files.append(_get_chat_file_path(folder))
@@ -48,7 +48,7 @@ def _get_chat_file_path(ctxid: str):
48
49
50
def _convert_v080_chats():
51
- json_files = files.list_files("tmp/chats", "*.json")
51
+ json_files = files.list_files(CHATS_FOLDER, "*.json")
52
for file in json_files:
53
path = files.get_abs_path(CHATS_FOLDER, file)
54
name = file.rstrip(".json")
python/helpers/task_scheduler.py
new
+383
@@ -0,0 +1,383 @@
1
+import uuid
2
+import random
3
+import os
4
+from datetime import datetime, timezone
5
+import threading
6
+import asyncio
7
+import nest_asyncio
8
+nest_asyncio.apply()
9
+
10
+from typing import Union, Literal, Optional
11
+
12
+from crontab import CronTab
13
+from pydantic import BaseModel, Field, PrivateAttr
14
+from python.helpers.files import get_abs_path, exists, write_file, read_file, make_dirs
15
+from agent import Agent, AgentContext, UserMessage
16
+from initialize import initialize
17
+from python.helpers.persist_chat import export_json_chat, load_json_chats
18
+from python.helpers.print_style import PrintStyle
19
+from python.helpers.defer import DeferredTask
20
+from python.helpers.persist_chat import CHATS_FOLDER
21
+from python.helpers import errors
22
+
23
+SCHEDULER_FOLDER = "memory/scheduler"
24
+TASKS_FOLDER = CHATS_FOLDER
25
+
26
+
27
+class TaskSchedule(BaseModel):
28
+ minute: str
29
+ hour: str
30
+ day: str
31
+ month: str
32
+ weekday: str
33
+
34
+ def to_crontab(self) -> str:
35
+ return f"{self.minute} {self.hour} {self.day} {self.month} {self.weekday}"
36
+
37
+
38
+class AdHocTask(BaseModel):
39
+ uuid: str = Field(default_factory=lambda: str(uuid.uuid4()))
40
+ state: Literal["idle", "running"] = Field(default="idle")
41
+ name: str = Field()
42
+ system_prompt: str
43
+ prompt: str
44
+ attachments: list[str] = Field(default_factory=list)
45
+ token: str = Field(default_factory=lambda: str(random.randint(1000000000000000000, 9999999999999999999)))
46
+ created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
47
+ updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
48
+ last_run: datetime | None = None
49
+ last_result: str | None = None
50
+
51
+ # lock: Optional[threading.Lock] = Field(exclude=True, default=threading.Lock())
52
+
53
+ @classmethod
54
+ def create(
55
+ cls,
56
+ name: str,
57
+ system_prompt: str,
58
+ prompt: str,
59
+ token: str,
60
+ attachments: list[str] = list()
61
+ ):
62
+ return cls(name=name,
63
+ system_prompt=system_prompt,
64
+ prompt=prompt,
65
+ attachments=attachments,
66
+ token=token)
67
+
68
+ def __init__(self, *args, **kwargs):
69
+ super().__init__(*args, **kwargs)
70
+ self._lock = threading.Lock()
71
+
72
+ def update(self,
73
+ name: str | None = None,
74
+ state: Literal["idle", "running"] | None = None,
75
+ system_prompt: str | None = None,
76
+ prompt: str | None = None,
77
+ attachments: list[str] | None = None,
78
+ last_run: datetime | None = None,
79
+ last_result: str | None = None):
80
+ with self._lock:
81
+ if name is not None:
82
+ self.name = name
83
+ self.updated_at = datetime.now(timezone.utc)
84
+ if state is not None:
85
+ self.state = state
86
+ self.updated_at = datetime.now(timezone.utc)
87
+ if system_prompt is not None:
88
+ self.system_prompt = system_prompt
89
+ self.updated_at = datetime.now(timezone.utc)
90
+ if prompt is not None:
91
+ self.prompt = prompt
92
+ self.updated_at = datetime.now(timezone.utc)
93
+ if attachments is not None:
94
+ self.attachments = attachments
95
+ self.updated_at = datetime.now(timezone.utc)
96
+ if last_run is not None:
97
+ self.last_run = last_run
98
+ self.updated_at = datetime.now(timezone.utc)
99
+ if last_result is not None:
100
+ self.last_result = last_result
101
+ self.updated_at = datetime.now(timezone.utc)
102
+
103
+ def check_schedule(self) -> bool:
104
+ with self._lock:
105
+ return False
106
+
107
+
108
+class ScheduledTask(BaseModel):
109
+ uuid: str = Field(default_factory=lambda: str(uuid.uuid4()))
110
+ state: Literal["idle", "running"] = Field(default="idle")
111
+ name: str
112
+ schedule: TaskSchedule
113
+ system_prompt: str
114
+ prompt: str
115
+ attachments: list[str] = Field(default_factory=list)
116
+ created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
117
+ updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
118
+ last_run: datetime | None = None
119
+ last_result: str | None = None
120
+
121
+ # lock: Optional[threading.Lock] = Field(exclude=True, default=threading.Lock())
122
+
123
+ @classmethod
124
+ def create(
125
+ cls,
126
+ name: str,
127
+ system_prompt: str,
128
+ prompt: str,
129
+ schedule: TaskSchedule,
130
+ attachments: list[str] = []
131
+ ):
132
+ return cls(name=name,
133
+ system_prompt=system_prompt,
134
+ prompt=prompt,
135
+ attachments=attachments,
136
+ schedule=schedule)
137
+
138
+ def __init__(self, *args, **kwargs):
139
+ super().__init__(*args, **kwargs)
140
+ self._lock = threading.Lock()
141
+
142
+ def update(self,
143
+ name: str | None = None,
144
+ state: Literal["idle", "running"] | None = None,
145
+ system_prompt: str | None = None,
146
+ prompt: str | None = None,
147
+ attachments: list[str] | None = None,
148
+ schedule: TaskSchedule | None = None,
149
+ last_run: datetime | None = None,
150
+ last_result: str | None = None):
151
+ with self._lock:
152
+ if name is not None:
153
+ self.name = name
154
+ self.updated_at = datetime.now(timezone.utc)
155
+ if state is not None:
156
+ self.state = state
157
+ self.updated_at = datetime.now(timezone.utc)
158
+ if system_prompt is not None:
159
+ self.system_prompt = system_prompt
160
+ self.updated_at = datetime.now(timezone.utc)
161
+ if prompt is not None:
162
+ self.prompt = prompt
163
+ self.updated_at = datetime.now(timezone.utc)
164
+ if attachments is not None:
165
+ self.attachments = attachments
166
+ self.updated_at = datetime.now(timezone.utc)
167
+ if schedule is not None:
168
+ self.schedule = schedule
169
+ self.updated_at = datetime.now(timezone.utc)
170
+ if last_run is not None:
171
+ self.last_run = last_run
172
+ self.updated_at = datetime.now(timezone.utc)
173
+ if last_result is not None:
174
+ self.last_result = last_result
175
+ self.updated_at = datetime.now(timezone.utc)
176
+
177
+ def check_schedule(self, frequency_seconds: float = 60.0) -> bool:
178
+ with self._lock:
179
+ crontab = CronTab(crontab=self.schedule.to_crontab())
180
+ next_run: float | None = crontab.next(now=datetime.now(timezone.utc), return_datetime=False)
181
+ if next_run is None:
182
+ return False
183
+ return next_run < frequency_seconds
184
+
185
+ def run(self):
186
+ pass
187
+
188
+
189
+class SchedulerTaskList(BaseModel):
190
+ tasks: list[Union[ScheduledTask, AdHocTask]]
191
+
192
+ # lock: threading.Lock = Field(exclude=True, default=threading.Lock())
193
+
194
+ @classmethod
195
+ def get(cls) -> "SchedulerTaskList":
196
+ path = get_abs_path(SCHEDULER_FOLDER, "tasks.json")
197
+ if not exists(path):
198
+ make_dirs(path)
199
+ instance = asyncio.run(cls(tasks=[]).save())
200
+ else:
201
+ instance = cls.model_validate_json(read_file(path))
202
+ return instance
203
+
204
+ def __init__(self, *args, **kwargs):
205
+ super().__init__(*args, **kwargs)
206
+ self._lock = threading.Lock()
207
+
208
+ async def add_task(self, task: Union[ScheduledTask, AdHocTask]) -> "SchedulerTaskList":
209
+ with self._lock:
210
+ self.tasks.append(task)
211
+ await self.save()
212
+ return self
213
+
214
+ async def save(self) -> "SchedulerTaskList":
215
+ with self._lock:
216
+ path = get_abs_path(SCHEDULER_FOLDER, "tasks.json")
217
+ if not exists(path):
218
+ make_dirs(path)
219
+ write_file(path, self.model_dump_json())
220
+ return self
221
+
222
+ def get_due_tasks(self) -> list[Union[ScheduledTask, AdHocTask]]:
223
+ with self._lock:
224
+ return [task for task in self.tasks if task.check_schedule()]
225
+
226
+ def get_task_by_uuid(self, task_uuid: str) -> Union[ScheduledTask, AdHocTask] | None:
227
+ with self._lock:
228
+ return next((task for task in self.tasks if task.uuid == task_uuid), None)
229
+
230
+ def get_task_by_name(self, name: str) -> Union[ScheduledTask, AdHocTask] | None:
231
+ with self._lock:
232
+ return next((task for task in self.tasks if task.name == name), None)
233
+
234
+ async def remove_task_by_uuid(self, task_uuid: str) -> "SchedulerTaskList":
235
+ with self._lock:
236
+ self.tasks = [task for task in self.tasks if task.uuid != task_uuid]
237
+ await self.save()
238
+ return self
239
+
240
+ async def remove_task_by_name(self, name: str) -> "SchedulerTaskList":
241
+ with self._lock:
242
+ self.tasks = [task for task in self.tasks if task.name != name]
243
+ await self.save()
244
+ return self
245
+
246
+
247
+class TaskScheduler:
248
+
249
+ _tasks: SchedulerTaskList
250
+ _printer: PrintStyle
251
+
252
+ @classmethod
253
+ def get(cls) -> "TaskScheduler":
254
+ return cls()
255
+
256
+ def __init__(self):
257
+ self._tasks = SchedulerTaskList.get()
258
+ self._printer = PrintStyle(italic=True, font_color="green", padding=False)
259
+
260
+ async def tick(self):
261
+ for task in self._tasks.get_due_tasks():
262
+ await self._run_task(task)
263
+
264
+ async def run_task_by_uuid(self, task_uuid: str):
265
+ task = self._tasks.get_task_by_uuid(task_uuid)
266
+ if task is None:
267
+ raise ValueError(f"Task with UUID {task_uuid} not found")
268
+ await self._run_task(task)
269
+
270
+ async def run_task_by_name(self, name: str):
271
+ task = self._tasks.get_task_by_name(name)
272
+ if task is None:
273
+ raise ValueError(f"Task with name {name} not found")
274
+ await self._run_task(task)
275
+
276
+ async def __new_context(self, task: Union[ScheduledTask, AdHocTask]) -> AgentContext:
277
+ config = initialize()
278
+ context: AgentContext = AgentContext(config)
279
+ context.id = task.uuid
280
+ chat_json = export_json_chat(context)
281
+ chat_file = get_abs_path(TASKS_FOLDER, task.uuid, "chat.json")
282
+ make_dirs(chat_file)
283
+ write_file(chat_file, chat_json)
284
+ return context
285
+
286
+ async def _get_chat_context(self, task: Union[ScheduledTask, AdHocTask]) -> AgentContext:
287
+ chat_file = get_abs_path(TASKS_FOLDER, task.uuid, "chat.json")
288
+ if exists(chat_file):
289
+ chat = read_file(chat_file)
290
+ context = AgentContext.get(load_json_chats([chat])[0])
291
+ if isinstance(context, AgentContext):
292
+ self._printer.print(
293
+ f"Scheduler Task {task.name} loaded from chat {task.uuid}"
294
+ )
295
+ return context
296
+ else:
297
+ self._printer.print(
298
+ f"Scheduler Task {task.name} loaded from chat {task.uuid} but failed to load context"
299
+ )
300
+ return await self.__new_context(task)
301
+ else:
302
+ self._printer.print(
303
+ f"Scheduler Task {task.name} loaded from chat {task.uuid} but chat file not found"
304
+ )
305
+ return await self.__new_context(task)
306
+
307
+ async def _persist_chat(self, task: Union[ScheduledTask, AdHocTask], context: AgentContext):
308
+ chat_json = export_json_chat(context)
309
+ chat_file = get_abs_path(TASKS_FOLDER, task.uuid, "chat.json")
310
+ make_dirs(chat_file)
311
+ write_file(chat_file, chat_json)
312
+
313
+ async def _run_task(self, task: Union[ScheduledTask, AdHocTask]):
314
+
315
+ async def _run_task_wrapper(task: Union[ScheduledTask, AdHocTask]):
316
+ if task.state == "running":
317
+ self._printer.print(f"Scheduler Task {task.name} already running")
318
+ return
319
+
320
+ try:
321
+ self._printer.print(f"Scheduler Task {task.name} started")
322
+
323
+ task.update(state="running")
324
+ await self._tasks.save()
325
+
326
+ context = await self._get_chat_context(task)
327
+ agent = Agent(0, context.config, context)
328
+
329
+ # Prepare attachment filenames for logging
330
+ attachment_filenames = []
331
+ if task.attachments:
332
+ for attachment in task.attachments:
333
+ if os.path.exists(attachment):
334
+ attachment_filenames.append(os.path.basename(attachment))
335
+
336
+ self._printer.print("User message:")
337
+ self._printer.print(f"> {task.prompt}")
338
+ if attachment_filenames:
339
+ self._printer.print("Attachments:")
340
+ for filename in attachment_filenames:
341
+ self._printer.print(f"- {filename}")
342
+
343
+ # Log the message with message_id and attachments
344
+ context.log.log(
345
+ type="user",
346
+ heading="User message",
347
+ content=task.prompt,
348
+ kvps={"attachments": attachment_filenames},
349
+ id=str(uuid.uuid4()),
350
+ )
351
+
352
+ agent.hist_add_user_message(
353
+ UserMessage(
354
+ message=task.prompt,
355
+ attachments=[]))
356
+
357
+ await self._persist_chat(task, context)
358
+
359
+ result = await agent.monologue()
360
+ task.update(last_result="SUCCESS: " + result)
361
+
362
+ self._printer.print(f"Scheduler Task {task.name} completed: {result}")
363
+
364
+ await self._persist_chat(task, context)
365
+
366
+ except Exception as e:
367
+ self._printer.print(f"Scheduler Task {task.name} failed: {e}")
368
+ task.update(last_result=f"ERROR: {str(e)}")
369
+ if agent:
370
+ agent.handle_critical_exception(e)
371
+
372
+ finally:
373
+ try:
374
+ task.update(
375
+ state="idle",
376
+ last_run=datetime.now(timezone.utc)
377
+ )
378
+ await self._tasks.save()
379
+ except Exception as e:
380
+ self._printer.print(f"Scheduler Task {task.name} failed to save: {e}")
381
+
382
+ deferred_task = DeferredTask(thread_name=self.__class__.__name__)
383
+ deferred_task.start_task(_run_task_wrapper, task)
requirements.txt
+1
@@ -29,3 +29,4 @@ tiktoken==0.8.0
29
unstructured==0.15.13
30
unstructured-client==0.25.9
31
webcolors==24.6.0
32
+crontab==1.0.1