pub mod algo;
pub mod read;
pub mod write;
pub mod runner;
pub mod io;
pub use algo::{prepare_next_tasks, apply_writes};
pub use read::PregelNode;
pub use write::{ChannelWrite, ChannelWriteEntry};
pub use runner::PregelRunner;
pub use io::{map_input, map_command, read_channels, NULL_TASK_ID};
use std::collections::HashMap;
use std::sync::Arc;
use serde_json::Value as JsonValue;
use langgraph_checkpoint::config::RunnableConfig;
use crate::channels::Channel;
use crate::runnable::Runnable;
#[derive(Debug, Clone)]
pub struct PregelTask {
pub name: String,
pub input: JsonValue,
pub input_channels: Vec<String>,
pub triggers: Vec<String>,
pub id: String,
}
pub struct PregelExecutableTask {
pub name: String,
pub input: JsonValue,
pub proc: Arc<dyn Runnable>,
pub writes: Vec<(String, JsonValue)>,
pub config: RunnableConfig,
pub triggers: Vec<String>,
pub id: String,
}
#[derive(Debug, Clone, PartialEq)]
pub enum LoopStatus {
Input,
Pending,
Done,
InterruptBefore,
InterruptAfter,
OutOfSteps,
}
pub type ChannelVersions = HashMap<String, JsonValue>;
pub type TriggerToNodes = HashMap<String, Vec<String>>;
pub fn build_trigger_to_nodes(nodes: &HashMap<String, PregelNode>) -> TriggerToNodes {
let mut map: TriggerToNodes = HashMap::new();
for (name, node) in nodes {
for trigger in &node.triggers {
map.entry(trigger.clone()).or_default().push(name.clone());
}
}
map
}
pub fn channels_from_checkpoint(
specs: &HashMap<String, Box<dyn Channel>>,
checkpoint_channels: &HashMap<String, Option<JsonValue>>,
) -> HashMap<String, Box<dyn Channel>> {
let mut channels = HashMap::new();
for (key, spec) in specs {
let cp = checkpoint_channels.get(key).and_then(|v| v.as_ref());
channels.insert(key.clone(), spec.from_checkpoint(cp));
}
channels
}