from datetime import datetime import logging import os from typing import Any, Dict, List, Optional from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore from apscheduler.schedulers.asyncio import AsyncIOScheduler from apscheduler.triggers.cron import CronTrigger from apscheduler.triggers.date import DateTrigger from apscheduler.triggers.interval import IntervalTrigger logger = logging.getLogger(__name__) class WorkflowScheduler: """ Manages scheduled workflow executions using APScheduler. Persists jobs to a SQLite database. """ def __init__(self, db_url: Optional[str] = None): from core.config import get_config self.config = get_config() jobstores = {} # Use configured job store if self.config.scheduler.job_store_type == 'redis' and self.config.redis.enabled: try: from apscheduler.jobstores.redis import RedisJobStore jobstores['default'] = RedisJobStore( host=self.config.redis.host, port=self.config.redis.port, db=self.config.redis.db, password=self.config.redis.password ) logger.info("WorkflowScheduler using RedisJobStore") except Exception as e: logger.warning(f"Failed to initialize RedisJobStore: {e}. Falling back to SQLAlchemy.") jobstores['default'] = SQLAlchemyJobStore(url=db_url or self.config.scheduler.job_store_url) else: job_store_url = db_url or self.config.scheduler.job_store_url jobstores['default'] = SQLAlchemyJobStore(url=job_store_url) logger.info(f"WorkflowScheduler using SQLAlchemyJobStore") self.scheduler = AsyncIOScheduler( jobstores=jobstores, job_defaults={ 'misfire_grace_time': self.config.scheduler.misfire_grace_time, 'coalesce': self.config.scheduler.coalesce, 'max_instances': self.config.scheduler.max_instances } ) self.engine = None # Will be set later to avoid circular imports def start(self): """Start the scheduler""" if not self.scheduler.running: self.reschedule_system_pipelines() self.scheduler.start() logger.info("WorkflowScheduler started") def reschedule_system_pipelines(self): """Register or refresh System Pipelines (Memory Ingestion) based on settings""" try: from core.automation_settings import get_automation_settings settings = get_automation_settings().get_settings() pipeline_config = settings.get("pipelines", {}) from integrations.atom_finance_memory_pipeline import finance_pipeline from integrations.atom_projects_memory_pipeline import projects_pipeline from integrations.atom_sales_memory_pipeline import sales_pipeline pipelines = { 'sales': sales_pipeline, 'projects': projects_pipeline, 'finance': finance_pipeline } for name, pipeline in pipelines.items(): config = pipeline_config.get(name, {}) mode = config.get("mode", "scheduled") job_id = f"system_{name}_ingestion" if mode == "real_time": # For real-time, we use a high-frequency interval (e.g., 1 minute) trigger = IntervalTrigger(minutes=1) logger.info(f"Setting {name} pipeline to REAL-TIME (1m interval)") else: # Scheduled mode uses cron cron_expr = config.get("cron", "*/30 * * * *" if name != 'finance' else "0 * * * *") trigger = CronTrigger.from_crontab(cron_expr) logger.info(f"Setting {name} pipeline to SCHEDULED ({cron_expr})") self.scheduler.add_job( pipeline.run_pipeline, trigger, id=job_id, replace_existing=True ) logger.info("✓ System Memory Pipelines (Re)Scheduled") except Exception as e: logger.error(f"Error rescheduling system pipelines: {e}") def shutdown(self): """Shutdown the scheduler""" if self.scheduler.running: self.scheduler.shutdown() logger.info("WorkflowScheduler shutdown") def set_engine(self, engine): """Set the AutomationEngine instance""" self.engine = engine @staticmethod async def _execute_job(workflow_id: str, input_data: Dict[str, Any] = None): """Internal job function to execute a workflow""" logger.info(f"Executing scheduled workflow: {workflow_id}") try: # Instantiate engine on demand to ensure fresh state and avoid circular imports at module level from ai.automation_engine import AutomationEngine engine = AutomationEngine() # Load workflows from core.workflow_endpoints import load_workflows workflows = load_workflows() workflow_def = next((w for w in workflows if w.get('id') == workflow_id or w.get('workflow_id') == workflow_id), None) if workflow_def: # Execute with a special execution ID prefix execution_id = f"sched_{datetime.now().strftime('%Y%m%d%H%M%S')}_{workflow_id[:8]}" await engine.execute_workflow_definition(workflow_def, input_data or {}, execution_id=execution_id) logger.info(f"Scheduled execution {execution_id} completed") else: logger.error(f"Scheduled workflow {workflow_id} not found") except Exception as e: logger.error(f"Error executing scheduled workflow {workflow_id}: {e}") def schedule_workflow(self, workflow_id: str, trigger_type: str, trigger_config: Dict[str, Any], input_data: Dict[str, Any] = None) -> str: """ Schedule a workflow execution. Args: workflow_id: ID of the workflow to schedule trigger_type: 'cron', 'interval', or 'date' trigger_config: Configuration for the trigger (e.g. cron expression) input_data: Optional input data for the workflow Returns: job_id: The ID of the scheduled job """ job_id = f"job_{workflow_id}_{datetime.now().timestamp()}" trigger = None if trigger_type == 'cron': trigger = CronTrigger(**trigger_config) elif trigger_type == 'interval': trigger = IntervalTrigger(**trigger_config) elif trigger_type == 'date': trigger = DateTrigger(**trigger_config) else: raise ValueError(f"Unsupported trigger type: {trigger_type}") self.scheduler.add_job( self._execute_job, trigger=trigger, args=[workflow_id, input_data], id=job_id, replace_existing=True ) return job_id def schedule_workflow_cron(self, job_id: str, workflow_id: str, cron_expression: str): """Schedule a workflow using cron expression""" self.scheduler.add_job( self._execute_job, CronTrigger.from_crontab(cron_expression), args=[workflow_id], id=job_id, replace_existing=True ) logger.info(f"Scheduled cron job {job_id} for workflow {workflow_id}: {cron_expression}") return job_id def schedule_workflow_interval(self, job_id: str, workflow_id: str, interval_minutes: int): """Schedule a workflow using interval""" self.scheduler.add_job( self._execute_job, IntervalTrigger(minutes=interval_minutes), args=[workflow_id], id=job_id, replace_existing=True ) logger.info(f"Scheduled interval job {job_id} for workflow {workflow_id}: {interval_minutes}m") return job_id def schedule_workflow_once(self, job_id: str, workflow_id: str, run_date: str): """Schedule a workflow once at a specific date""" self.scheduler.add_job( self._execute_job, DateTrigger(run_date=run_date), args=[workflow_id], id=job_id, replace_existing=True ) logger.info(f"Scheduled one-time job {job_id} for workflow {workflow_id} at {run_date}") return job_id def remove_job(self, job_id: str) -> bool: """Remove a scheduled job""" try: self.scheduler.remove_job(job_id) logger.info(f"Removed job {job_id}") return True except Exception: return False logger.info(f"Scheduled workflow {workflow_id} with {trigger_type} trigger (Job ID: {job_id})") return job_id def remove_schedule(self, job_id: str): """Remove a scheduled job""" try: self.scheduler.remove_job(job_id) logger.info(f"Removed job {job_id}") except Exception as e: logger.error(f"Error removing job {job_id}: {e}") def list_jobs(self) -> List[Dict[str, Any]]: """List all scheduled jobs""" jobs = [] for job in self.scheduler.get_jobs(): jobs.append({ "id": job.id, "next_run_time": job.next_run_time.isoformat() if job.next_run_time else None, "trigger": str(job.trigger) }) return jobs # Global instance workflow_scheduler = WorkflowScheduler()