109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
from pydantic import BaseModel
|
|
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
|
from apscheduler.triggers.cron import CronTrigger
|
|
from app.services.session_manager import session_manager
|
|
from app.services.storage import storage
|
|
from app.core.logging import logger
|
|
|
|
SCHEDULE_KEY = "schedules"
|
|
|
|
|
|
class ScheduleTask(BaseModel):
|
|
id: str
|
|
name: str
|
|
cron: str
|
|
prompt: str
|
|
enabled: bool = True
|
|
|
|
|
|
class SchedulerService:
|
|
def __init__(self):
|
|
self.scheduler = AsyncIOScheduler()
|
|
self.scheduled_tasks: dict[str, ScheduleTask] = {}
|
|
|
|
def load(self):
|
|
data = storage.load(SCHEDULE_KEY, {})
|
|
if isinstance(data, dict):
|
|
for k, v in data.items():
|
|
self.scheduled_tasks[k] = ScheduleTask(**v)
|
|
logger.info(f"Loaded {len(self.scheduled_tasks)} scheduled tasks")
|
|
|
|
def save(self):
|
|
data = {k: v.model_dump() for k, v in self.scheduled_tasks.items()}
|
|
storage.save(SCHEDULE_KEY, data)
|
|
|
|
def start(self):
|
|
self.scheduler.start()
|
|
for task in self.scheduled_tasks.values():
|
|
if task.enabled:
|
|
trigger = CronTrigger.from_crontab(task.cron)
|
|
self.scheduler.add_job(
|
|
self._run_scheduled_task,
|
|
trigger,
|
|
args=[task.id],
|
|
id=task.id,
|
|
)
|
|
logger.info(f"Loaded schedule: {task.name} ({task.cron})")
|
|
logger.info("Scheduler started")
|
|
|
|
def shutdown(self):
|
|
self.scheduler.shutdown()
|
|
logger.info("Scheduler shutdown")
|
|
|
|
async def add_schedule(self, task: ScheduleTask) -> ScheduleTask:
|
|
self.scheduled_tasks[task.id] = task
|
|
self.save()
|
|
|
|
if task.enabled:
|
|
trigger = CronTrigger.from_crontab(task.cron)
|
|
self.scheduler.add_job(
|
|
self._run_scheduled_task,
|
|
trigger,
|
|
args=[task.id],
|
|
id=task.id,
|
|
)
|
|
logger.info(f"Added schedule: {task.name} ({task.cron})")
|
|
|
|
return task
|
|
|
|
async def remove_schedule(self, task_id: str) -> bool:
|
|
task = self.scheduled_tasks.pop(task_id, None)
|
|
if not task:
|
|
return False
|
|
|
|
try:
|
|
self.scheduler.remove_job(task_id)
|
|
except Exception:
|
|
pass
|
|
|
|
self.save()
|
|
logger.info(f"Removed schedule: {task.name}")
|
|
return True
|
|
|
|
async def list_schedules(self) -> list[ScheduleTask]:
|
|
return list(self.scheduled_tasks.values())
|
|
|
|
async def _run_scheduled_task(self, task_id: str):
|
|
task = self.scheduled_tasks.get(task_id)
|
|
if not task:
|
|
return
|
|
|
|
logger.info(f"Running scheduled task: {task.name}")
|
|
|
|
from app.models.session import CreateSessionRequest, SessionType
|
|
|
|
request = CreateSessionRequest(
|
|
type=SessionType.SCHEDULED,
|
|
prompt=task.prompt,
|
|
)
|
|
task_obj = await session_manager.create_task(request)
|
|
|
|
try:
|
|
await session_manager.execute_task_async(task_obj.id)
|
|
logger.info(f"Scheduled task {task.name} started")
|
|
except Exception as e:
|
|
logger.error(f"Scheduled task {task.name} failed to start: {e}")
|
|
|
|
|
|
scheduler_service = SchedulerService()
|