use serde::{Deserialize, Serialize};
use std::time::Duration;
pub type JobId = String;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum JobState {
Pending,
Queued,
Running,
Succeeded,
Failed,
Retrying,
DeadLettered,
Cancelled,
}
impl JobState {
pub fn is_terminal(self) -> bool {
matches!(
self,
JobState::Succeeded | JobState::DeadLettered | JobState::Cancelled
)
}
pub fn label(self) -> &'static str {
match self {
JobState::Pending => "pending",
JobState::Queued => "queued",
JobState::Running => "running",
JobState::Succeeded => "succeeded",
JobState::Failed => "failed",
JobState::Retrying => "retrying",
JobState::DeadLettered => "dead-lettered",
JobState::Cancelled => "cancelled",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct JobResult {
pub text: String,
pub input_tokens: u64,
pub output_tokens: u64,
}
impl JobResult {
pub fn new(text: impl Into<String>) -> Self {
Self {
text: text.into(),
input_tokens: 0,
output_tokens: 0,
}
}
pub fn with_tokens(mut self, input: u64, output: u64) -> Self {
self.input_tokens = input;
self.output_tokens = output;
self
}
pub fn total_tokens(&self) -> u64 {
self.input_tokens + self.output_tokens
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Backoff {
None,
Fixed(Duration),
Exponential {
base: Duration,
factor: u32,
max: Duration,
},
}
impl Backoff {
pub fn delay(self, attempt: u32) -> Duration {
match self {
Backoff::None => Duration::ZERO,
Backoff::Fixed(d) => d,
Backoff::Exponential { base, factor, max } => {
let mult = factor.saturating_pow(attempt.saturating_sub(1));
base.saturating_mul(mult).min(max)
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct RetryPolicy {
pub max_attempts: u32,
pub backoff: Backoff,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_attempts: 1,
backoff: Backoff::None,
}
}
}
impl RetryPolicy {
pub fn once() -> Self {
Self::default()
}
pub fn new(max_attempts: u32, backoff: Backoff) -> Self {
Self {
max_attempts: max_attempts.max(1),
backoff,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Job {
pub id: JobId,
pub deps: Vec<JobId>,
pub prompt: String,
pub retry: RetryPolicy,
pub state: JobState,
pub result: Option<JobResult>,
pub attempts: u32,
#[serde(default)]
pub visits: u32,
pub last_error: Option<String>,
#[serde(default)]
pub prior_attempt: Option<String>,
#[serde(default)]
pub feedback: Option<String>,
}
impl Job {
pub fn new(id: impl Into<JobId>, prompt: impl Into<String>) -> Self {
Self {
id: id.into(),
deps: Vec::new(),
prompt: prompt.into(),
retry: RetryPolicy::default(),
state: JobState::Pending,
result: None,
attempts: 0,
visits: 0,
last_error: None,
prior_attempt: None,
feedback: None,
}
}
pub fn with_deps<I, S>(mut self, deps: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<JobId>,
{
self.deps = deps.into_iter().map(Into::into).collect();
self
}
pub fn with_retry(mut self, retry: RetryPolicy) -> Self {
self.retry = retry;
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn exponential_backoff_grows_and_caps() {
let b = Backoff::Exponential {
base: Duration::from_secs(1),
factor: 2,
max: Duration::from_secs(10),
};
assert_eq!(b.delay(1), Duration::from_secs(1)); assert_eq!(b.delay(2), Duration::from_secs(2)); assert_eq!(b.delay(3), Duration::from_secs(4)); assert_eq!(b.delay(5), Duration::from_secs(10)); }
#[test]
fn retry_policy_min_one_attempt() {
assert_eq!(RetryPolicy::new(0, Backoff::None).max_attempts, 1);
}
#[test]
fn terminal_states() {
assert!(JobState::Succeeded.is_terminal());
assert!(JobState::DeadLettered.is_terminal());
assert!(!JobState::Pending.is_terminal());
assert!(!JobState::Running.is_terminal());
}
}