118 lines
3.5 KiB
Rust
118 lines
3.5 KiB
Rust
use crate::agent::Agent;
|
|
use crate::entities::task::Entity as Task;
|
|
use crate::entities::task_run;
|
|
use dashmap::DashMap;
|
|
use sea_orm::{DatabaseConnection, EntityTrait, Set};
|
|
use tokio_cron_scheduler::{Job, JobScheduler};
|
|
use uuid::Uuid;
|
|
|
|
pub struct Scheduler {
|
|
scheduler: JobScheduler,
|
|
db: DatabaseConnection,
|
|
tasks_to_jobs: DashMap<Uuid, Uuid>,
|
|
zen_api_key: Option<String>,
|
|
tavily_api_key: Option<String>,
|
|
}
|
|
|
|
impl Scheduler {
|
|
pub async fn new(
|
|
db: DatabaseConnection,
|
|
zen_api_key: Option<String>,
|
|
tavily_api_key: Option<String>,
|
|
) -> Result<Self, Box<dyn std::error::Error>> {
|
|
let scheduler = JobScheduler::new().await?;
|
|
scheduler.start().await?;
|
|
Ok(Self {
|
|
scheduler,
|
|
db,
|
|
tasks_to_jobs: DashMap::new(),
|
|
zen_api_key,
|
|
tavily_api_key,
|
|
})
|
|
}
|
|
|
|
pub async fn add_task_job(
|
|
&self,
|
|
task_id: Uuid,
|
|
cron_expr: &str,
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
|
// Remove existing job if any
|
|
if let Some((_, old_job_id)) = self.tasks_to_jobs.remove(&task_id) {
|
|
let _ = self.scheduler.remove(&old_job_id).await;
|
|
}
|
|
|
|
let db = self.db.clone();
|
|
let zen_key = self.zen_api_key.clone();
|
|
let tavily_key = self.tavily_api_key.clone();
|
|
|
|
let job = Job::new_async(cron_expr, move |_uuid, _l| {
|
|
let db = db.clone();
|
|
let zen_key = zen_key.clone();
|
|
let tavily_key = tavily_key.clone();
|
|
Box::pin(async move {
|
|
if let Err(e) = Self::run_task(db, zen_key, tavily_key, task_id).await {
|
|
eprintln!("Error in scheduled task {}: {}", task_id, e);
|
|
}
|
|
})
|
|
})?;
|
|
|
|
let job_id = self.scheduler.add(job).await?;
|
|
self.tasks_to_jobs.insert(task_id, job_id);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn remove_task_job(&self, task_id: Uuid) -> Result<(), Box<dyn std::error::Error>> {
|
|
if let Some((_, job_id)) = self.tasks_to_jobs.remove(&task_id) {
|
|
self.scheduler.remove(&job_id).await?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn run_task(
|
|
db: DatabaseConnection,
|
|
zen_key: Option<String>,
|
|
tavily_key: Option<String>,
|
|
task_id: Uuid,
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
|
let task = Task::find_by_id(task_id)
|
|
.one(&db)
|
|
.await?
|
|
.ok_or("Task not found")?;
|
|
|
|
// Create a new run entry
|
|
let run_id = Uuid::new_v4();
|
|
let run = task_run::ActiveModel {
|
|
id: Set(run_id),
|
|
task_id: Set(task_id),
|
|
status: Set("running".to_string()),
|
|
logs: Set(String::new()),
|
|
answer: Set(None),
|
|
created_at: Set(chrono::Utc::now().into()),
|
|
};
|
|
|
|
use sea_orm::ActiveModelTrait;
|
|
run.insert(&db).await?;
|
|
|
|
// Start agent in background
|
|
let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone())?;
|
|
|
|
tokio::spawn(async move {
|
|
let (logs, answer) = match agent.run().await {
|
|
Ok(res) => res,
|
|
Err(e) => (format!("Scheduled run failed: {}", e), None),
|
|
};
|
|
|
|
let run_complete = task_run::ActiveModel {
|
|
id: Set(run_id),
|
|
status: Set("completed".to_string()),
|
|
logs: Set(logs),
|
|
answer: Set(answer),
|
|
..Default::default()
|
|
};
|
|
let _ = run_complete.update(&db).await;
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
}
|