use crate::retry::RetryPolicy;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Duration;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct GraphConfig {
pub thread_id: String,
pub checkpoint_id: Option<String>,
pub recursion_limit: u32,
pub max_concurrency: Option<usize>,
pub configurable: HashMap<String, serde_json::Value>,
#[serde(skip)]
pub max_execution_time: Option<Duration>,
#[serde(skip)]
pub retry_policy: Option<RetryPolicy>,
}
impl Default for GraphConfig {
fn default() -> Self {
Self {
thread_id: uuid::Uuid::new_v4().to_string(),
checkpoint_id: None,
recursion_limit: 25,
max_concurrency: None,
configurable: HashMap::new(),
max_execution_time: None,
retry_policy: None,
}
}
}
impl GraphConfig {
#[must_use = "builder methods return the modified config"]
pub fn with_thread_id(mut self, id: impl Into<String>) -> Self {
self.thread_id = id.into();
self
}
#[must_use = "builder methods return the modified config"]
pub fn with_recursion_limit(mut self, limit: u32) -> Self {
self.recursion_limit = limit;
self
}
#[must_use = "builder methods return the modified config"]
pub fn with_max_concurrency(mut self, n: usize) -> Self {
self.max_concurrency = Some(n);
self
}
#[must_use = "builder methods return the modified config"]
pub fn with_checkpoint_id(mut self, id: impl Into<String>) -> Self {
self.checkpoint_id = Some(id.into());
self
}
#[must_use = "builder methods return the modified config"]
pub fn with_max_execution_time(mut self, duration: Duration) -> Self {
self.max_execution_time = Some(duration);
self
}
#[must_use = "builder methods return the modified config"]
pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
self.retry_policy = Some(policy);
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_values() {
let config = GraphConfig::default();
assert_eq!(config.recursion_limit, 25);
assert!(config.max_concurrency.is_none());
assert!(config.checkpoint_id.is_none());
assert!(!config.thread_id.is_empty());
assert!(config.configurable.is_empty());
assert!(config.max_execution_time.is_none());
assert!(config.retry_policy.is_none());
}
#[test]
fn test_builder_methods() {
let config = GraphConfig::default()
.with_thread_id("test-thread")
.with_recursion_limit(100)
.with_max_concurrency(4)
.with_checkpoint_id("cp-123");
assert_eq!(config.thread_id, "test-thread");
assert_eq!(config.recursion_limit, 100);
assert_eq!(config.max_concurrency, Some(4));
assert_eq!(config.checkpoint_id.as_deref(), Some("cp-123"));
}
#[test]
fn test_with_retry_policy() {
let policy = RetryPolicy {
max_attempts: 5,
..Default::default()
};
let config = GraphConfig::default().with_retry_policy(policy);
assert!(config.retry_policy.is_some());
assert_eq!(config.retry_policy.unwrap().max_attempts, 5);
}
#[test]
fn test_unique_thread_ids() {
let a = GraphConfig::default();
let b = GraphConfig::default();
assert_ne!(a.thread_id, b.thread_id);
}
}