Dark-Alex-17
be0c5462ed
Agent nodes gain an optional fallback target with llm-node parity: on any failure (retries exhausted, hard error, or structured-extraction failure) the failure chain is written into state as 'Agent node failed: <chain>' and the chain routes to the fallback node. Without a fallback the error propagates unchanged. Validator collects agent fallback edges for missing-target checks, reachability, and map-branch locality.
1653 lines
44 KiB
Rust
1653 lines
44 KiB
Rust
use crate::config::AgentVariable;
|
|
use anyhow::Result;
|
|
use indexmap::IndexMap;
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::Value;
|
|
use std::collections::HashMap;
|
|
use std::slice;
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct Graph {
|
|
pub name: String,
|
|
|
|
#[serde(default)]
|
|
pub description: String,
|
|
|
|
#[serde(default = "default_schema_version")]
|
|
pub version: String,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub model: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub temperature: Option<f64>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub top_p: Option<f64>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub reasoning_effort: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub max_concurrent_jobs: Option<usize>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub can_spawn_agents: Option<bool>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub max_concurrent_agents: Option<usize>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub max_agent_depth: Option<usize>,
|
|
|
|
#[serde(default)]
|
|
pub global_tools: Vec<String>,
|
|
|
|
#[serde(default)]
|
|
pub mcp_servers: Vec<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub skills_enabled: Option<bool>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub enabled_skills: Option<Vec<String>>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub inject_skill_instructions: Option<bool>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub skill_instructions: Option<String>,
|
|
|
|
#[serde(default)]
|
|
pub conversation_starters: Vec<String>,
|
|
|
|
#[serde(default)]
|
|
pub variables: Vec<AgentVariable>,
|
|
|
|
#[serde(default)]
|
|
pub settings: GraphSettings,
|
|
|
|
#[serde(default)]
|
|
pub initial_state: HashMap<String, Value>,
|
|
|
|
#[serde(default)]
|
|
pub reducers: HashMap<String, Reducer>,
|
|
|
|
pub start: String,
|
|
|
|
pub nodes: IndexMap<String, Node>,
|
|
}
|
|
|
|
impl Graph {
|
|
pub fn get_node(&self, id: &str) -> Option<&Node> {
|
|
self.nodes.get(id)
|
|
}
|
|
|
|
pub fn has_node(&self, id: &str) -> bool {
|
|
self.nodes.contains_key(id)
|
|
}
|
|
|
|
pub fn node_ids(&self) -> Vec<&str> {
|
|
self.nodes.keys().map(|s| s.as_str()).collect()
|
|
}
|
|
|
|
pub fn has_agent_node(&self) -> bool {
|
|
self.nodes
|
|
.values()
|
|
.any(|n| matches!(n.node_type, NodeType::Agent(_)))
|
|
}
|
|
}
|
|
|
|
fn default_schema_version() -> String {
|
|
super::GRAPH_SCHEMA_VERSION.to_string()
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct GraphSettings {
|
|
#[serde(default = "default_max_loop_iterations")]
|
|
pub max_loop_iterations: usize,
|
|
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub timeout: Option<u64>,
|
|
|
|
#[serde(default = "default_true")]
|
|
pub log_state_snapshots: bool,
|
|
|
|
#[serde(default = "default_true")]
|
|
pub validate_before_run: bool,
|
|
|
|
#[serde(default = "default_max_concurrency")]
|
|
pub max_concurrency: usize,
|
|
}
|
|
|
|
impl Default for GraphSettings {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_loop_iterations: default_max_loop_iterations(),
|
|
timeout: None,
|
|
log_state_snapshots: true,
|
|
validate_before_run: true,
|
|
max_concurrency: default_max_concurrency(),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn default_max_loop_iterations() -> usize {
|
|
super::DEFAULT_MAX_LOOP_ITERATIONS
|
|
}
|
|
|
|
fn default_true() -> bool {
|
|
true
|
|
}
|
|
|
|
fn default_max_concurrency() -> usize {
|
|
8
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct Node {
|
|
#[serde(default)]
|
|
pub id: String,
|
|
|
|
#[serde(default)]
|
|
pub description: String,
|
|
|
|
#[serde(flatten)]
|
|
pub node_type: NodeType,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub next: Option<NextTargets>,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
impl Node {
|
|
pub fn next_target(&self) -> Option<&str> {
|
|
match &self.next {
|
|
None => None,
|
|
Some(NextTargets::One(s)) => Some(s),
|
|
Some(NextTargets::Many(v)) if v.len() == 1 => Some(&v[0]),
|
|
Some(NextTargets::Many(_)) => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
#[serde(untagged)]
|
|
pub enum NextTargets {
|
|
One(String),
|
|
Many(Vec<String>),
|
|
}
|
|
|
|
impl NextTargets {
|
|
pub fn as_slice(&self) -> &[String] {
|
|
match self {
|
|
NextTargets::One(s) => slice::from_ref(s),
|
|
NextTargets::Many(v) => v.as_slice(),
|
|
}
|
|
}
|
|
|
|
pub fn is_fan_out(&self) -> bool {
|
|
matches!(self, NextTargets::Many(v) if v.len() > 1)
|
|
}
|
|
}
|
|
|
|
impl From<String> for NextTargets {
|
|
fn from(s: String) -> Self {
|
|
NextTargets::One(s)
|
|
}
|
|
}
|
|
|
|
impl From<&str> for NextTargets {
|
|
fn from(s: &str) -> Self {
|
|
NextTargets::One(s.to_string())
|
|
}
|
|
}
|
|
|
|
impl From<Vec<String>> for NextTargets {
|
|
fn from(v: Vec<String>) -> Self {
|
|
NextTargets::Many(v)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
#[serde(tag = "type", rename_all = "lowercase")]
|
|
pub enum NodeType {
|
|
Agent(AgentNode),
|
|
Script(ScriptNode),
|
|
Approval(ApprovalNode),
|
|
Input(InputNode),
|
|
Llm(LlmNode),
|
|
Rag(RagNode),
|
|
End(EndNode),
|
|
Map(MapNode),
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct AgentNode {
|
|
pub agent: String,
|
|
|
|
pub prompt: String,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub state_updates: Option<HashMap<String, String>>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub output_schema: Option<Value>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub timeout: Option<u64>,
|
|
|
|
#[serde(default = "default_max_attempts")]
|
|
pub max_attempts: u32,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub fallback: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub inputs: Option<HashMap<String, String>>,
|
|
|
|
#[serde(default)]
|
|
pub teammates: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct ScriptNode {
|
|
pub script: String,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub state_updates: Option<HashMap<String, String>>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub fallback: Option<String>,
|
|
|
|
#[serde(default = "default_script_timeout")]
|
|
pub timeout: u64,
|
|
}
|
|
|
|
fn default_script_timeout() -> u64 {
|
|
30
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct ApprovalNode {
|
|
pub question: String,
|
|
|
|
pub options: Vec<String>,
|
|
|
|
pub routes: HashMap<String, String>,
|
|
|
|
pub on_other: String,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub state_updates: Option<HashMap<String, String>>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct InputNode {
|
|
pub question: String,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub default: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub validation: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub state_updates: Option<HashMap<String, String>>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct LlmNode {
|
|
pub prompt: String,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub instructions: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub tools: Option<Vec<String>>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub mcp_tools: Option<IndexMap<String, Vec<String>>>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub model: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub temperature: Option<f64>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub top_p: Option<f64>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub reasoning_effort: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub fallback: Option<String>,
|
|
|
|
#[serde(default = "default_max_attempts")]
|
|
pub max_attempts: u32,
|
|
|
|
#[serde(default = "default_llm_max_iterations")]
|
|
pub max_iterations: u32,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub state_updates: Option<HashMap<String, String>>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub output_schema: Option<Value>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub timeout: Option<u64>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub skills_enabled: Option<bool>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub enabled_skills: Option<Vec<String>>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub inject_skill_instructions: Option<bool>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub skill_instructions: Option<String>,
|
|
}
|
|
|
|
fn default_max_attempts() -> u32 {
|
|
1
|
|
}
|
|
|
|
fn default_llm_max_iterations() -> u32 {
|
|
10
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct RagNode {
|
|
pub documents: Vec<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub query: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub top_k: Option<usize>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub embedding_model: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub chunk_size: Option<usize>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub chunk_overlap: Option<usize>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub reranker_model: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub batch_size: Option<usize>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub extractor_model: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub extractor_prompt: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub graph_hops: Option<usize>,
|
|
|
|
/// Storage driver for this node's knowledge base ("yaml", "duckdb"). `None`
|
|
/// means "yaml". Only honored when the knowledge base is first built;
|
|
/// changing it afterwards has no effect until the RAG is deleted and
|
|
/// re-initialized.
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub driver: Option<String>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub state_updates: Option<HashMap<String, String>>,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub timeout: Option<u64>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct EndNode {
|
|
#[serde(default)]
|
|
pub output: String,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub state_updates: Option<HashMap<String, String>>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
pub struct MapNode {
|
|
pub over: String,
|
|
|
|
#[serde(rename = "as")]
|
|
pub as_name: String,
|
|
|
|
pub branch: String,
|
|
|
|
#[serde(default = "default_map_output_key")]
|
|
pub output_key: String,
|
|
|
|
pub collect_into: String,
|
|
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
pub max_concurrency: Option<ConcurrencyCap>,
|
|
}
|
|
|
|
fn default_map_output_key() -> String {
|
|
"output".to_string()
|
|
}
|
|
|
|
/// A map node's per-item concurrency cap: either a literal integer or a
|
|
/// `{{template}}` resolved against graph state when the map runs.
|
|
#[derive(Debug, Clone, Serialize, PartialEq)]
|
|
#[serde(untagged)]
|
|
pub enum ConcurrencyCap {
|
|
Fixed(usize),
|
|
Template(String),
|
|
}
|
|
|
|
// Hand-written so a malformed value (`-1`, `4.5`, `true`) reports what was
|
|
// received and what is accepted, instead of serde's untagged-enum "did not
|
|
// match any variant".
|
|
impl<'de> Deserialize<'de> for ConcurrencyCap {
|
|
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
|
where
|
|
D: serde::Deserializer<'de>,
|
|
{
|
|
struct CapVisitor;
|
|
|
|
impl serde::de::Visitor<'_> for CapVisitor {
|
|
type Value = ConcurrencyCap;
|
|
|
|
fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
|
f.write_str("a positive integer or a `{{template}}` string for max_concurrency")
|
|
}
|
|
|
|
fn visit_u64<E: serde::de::Error>(self, v: u64) -> Result<Self::Value, E> {
|
|
usize::try_from(v)
|
|
.map(ConcurrencyCap::Fixed)
|
|
.map_err(|_| E::invalid_value(serde::de::Unexpected::Unsigned(v), &self))
|
|
}
|
|
|
|
fn visit_i64<E: serde::de::Error>(self, v: i64) -> Result<Self::Value, E> {
|
|
usize::try_from(v)
|
|
.map(ConcurrencyCap::Fixed)
|
|
.map_err(|_| E::invalid_value(serde::de::Unexpected::Signed(v), &self))
|
|
}
|
|
|
|
fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<Self::Value, E> {
|
|
Ok(ConcurrencyCap::Template(v.to_owned()))
|
|
}
|
|
}
|
|
|
|
deserializer.deserialize_any(CapVisitor)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
|
|
#[serde(rename_all = "lowercase")]
|
|
pub enum Reducer {
|
|
Append,
|
|
Extend,
|
|
Concat,
|
|
Sum,
|
|
Max,
|
|
Min,
|
|
Merge,
|
|
Overwrite,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct GraphState {
|
|
data: HashMap<String, Value>,
|
|
#[cfg(test)]
|
|
history: Vec<String>,
|
|
loop_counts: HashMap<String, usize>,
|
|
}
|
|
|
|
impl GraphState {
|
|
pub fn new(initial: HashMap<String, Value>) -> Self {
|
|
Self {
|
|
data: initial,
|
|
#[cfg(test)]
|
|
history: Vec::new(),
|
|
loop_counts: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
pub fn get(&self, key: &str) -> Option<&Value> {
|
|
self.data.get(key)
|
|
}
|
|
|
|
pub fn set(&mut self, key: String, value: Value) {
|
|
self.data.insert(key, value);
|
|
}
|
|
|
|
pub fn remove(&mut self, key: &str) -> Option<Value> {
|
|
self.data.remove(key)
|
|
}
|
|
|
|
pub fn merge(&mut self, json_obj: &serde_json::Map<String, Value>) {
|
|
for (key, value) in json_obj {
|
|
self.data.insert(key.clone(), value.clone());
|
|
}
|
|
}
|
|
|
|
pub fn data(&self) -> &HashMap<String, Value> {
|
|
&self.data
|
|
}
|
|
|
|
pub fn visit_node(&mut self, node_id: &str) {
|
|
#[cfg(test)]
|
|
self.history.push(node_id.to_string());
|
|
*self.loop_counts.entry(node_id.to_string()).or_insert(0) += 1;
|
|
}
|
|
|
|
pub fn loop_count(&self, node_id: &str) -> usize {
|
|
self.loop_counts.get(node_id).copied().unwrap_or(0)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub fn current_node(&self) -> Option<&str> {
|
|
self.history.last().map(|s| s.as_str())
|
|
}
|
|
|
|
pub fn to_json(&self) -> Result<String> {
|
|
serde_json::to_string(&self.data)
|
|
.map_err(|e| anyhow::anyhow!("Failed to serialize graph state: {}", e))
|
|
}
|
|
|
|
pub fn size_bytes(&self) -> usize {
|
|
self.to_json().map(|s| s.len()).unwrap_or(0)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use serde_json::json;
|
|
|
|
#[test]
|
|
fn deserializes_a_simple_graph() {
|
|
let yaml = r#"
|
|
name: test_graph
|
|
description: A test graph
|
|
version: "1.0"
|
|
start: node1
|
|
nodes:
|
|
node1:
|
|
id: node1
|
|
type: agent
|
|
agent: test_agent
|
|
prompt: "Hello {{name}}"
|
|
state_updates:
|
|
result: "{{output}}"
|
|
next: node2
|
|
node2:
|
|
id: node2
|
|
type: end
|
|
output: "{{result}}"
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
assert_eq!(graph.name, "test_graph");
|
|
assert_eq!(graph.start, "node1");
|
|
assert_eq!(graph.nodes.len(), 2);
|
|
assert!(graph.has_node("node1"));
|
|
assert!(graph.has_node("node2"));
|
|
assert!(!graph.has_node("missing"));
|
|
|
|
let node1 = graph.get_node("node1").unwrap();
|
|
assert!(matches!(node1.node_type, NodeType::Agent(_)));
|
|
|
|
let node2 = graph.get_node("node2").unwrap();
|
|
match &node2.node_type {
|
|
NodeType::End(end) => assert_eq!(end.output, "{{result}}"),
|
|
_ => panic!("expected End variant"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn agent_node_teammates_parses_and_defaults_false() {
|
|
let yaml = r#"
|
|
name: g
|
|
start: a
|
|
nodes:
|
|
a:
|
|
id: a
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
teammates: true
|
|
next: b
|
|
b:
|
|
id: b
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
next: e
|
|
e:
|
|
id: e
|
|
type: end
|
|
output: done
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
match &graph.get_node("a").unwrap().node_type {
|
|
NodeType::Agent(n) => assert!(n.teammates),
|
|
_ => panic!("expected Agent variant"),
|
|
}
|
|
match &graph.get_node("b").unwrap().node_type {
|
|
NodeType::Agent(n) => assert!(!n.teammates),
|
|
_ => panic!("expected Agent variant"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn agent_node_max_attempts_parses_and_defaults_to_one() {
|
|
let yaml = r#"
|
|
name: g
|
|
start: a
|
|
nodes:
|
|
a:
|
|
id: a
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
max_attempts: 2
|
|
next: b
|
|
b:
|
|
id: b
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
next: e
|
|
e:
|
|
id: e
|
|
type: end
|
|
output: done
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
match &graph.get_node("a").unwrap().node_type {
|
|
NodeType::Agent(n) => assert_eq!(n.max_attempts, 2),
|
|
_ => panic!("expected Agent variant"),
|
|
}
|
|
match &graph.get_node("b").unwrap().node_type {
|
|
NodeType::Agent(n) => assert_eq!(n.max_attempts, 1),
|
|
_ => panic!("expected Agent variant"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn agent_node_fallback_parses_and_defaults_to_none() {
|
|
let yaml = r#"
|
|
name: g
|
|
start: a
|
|
nodes:
|
|
a:
|
|
id: a
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
fallback: b
|
|
next: b
|
|
b:
|
|
id: b
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
next: e
|
|
e:
|
|
id: e
|
|
type: end
|
|
output: done
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
match &graph.get_node("a").unwrap().node_type {
|
|
NodeType::Agent(n) => assert_eq!(n.fallback.as_deref(), Some("b")),
|
|
_ => panic!("expected Agent variant"),
|
|
}
|
|
match &graph.get_node("b").unwrap().node_type {
|
|
NodeType::Agent(n) => assert!(n.fallback.is_none()),
|
|
_ => panic!("expected Agent variant"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn agent_node_inputs_parse_optionally_and_empty_map_is_some() {
|
|
let yaml = r#"
|
|
name: g
|
|
start: with_inputs
|
|
nodes:
|
|
with_inputs:
|
|
id: with_inputs
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
inputs:
|
|
width: "{{n}}"
|
|
note: "Repo: {{repo}}"
|
|
next: without_inputs
|
|
without_inputs:
|
|
id: without_inputs
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
next: empty_inputs
|
|
empty_inputs:
|
|
id: empty_inputs
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
inputs: {}
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
let inputs_of = |id: &str| match &graph.get_node(id).unwrap().node_type {
|
|
NodeType::Agent(n) => n.inputs.clone(),
|
|
_ => panic!("expected Agent variant"),
|
|
};
|
|
|
|
let with = inputs_of("with_inputs").expect("inputs should parse");
|
|
assert_eq!(with.len(), 2);
|
|
assert_eq!(with.get("width").map(String::as_str), Some("{{n}}"));
|
|
assert_eq!(with.get("note").map(String::as_str), Some("Repo: {{repo}}"));
|
|
assert_eq!(inputs_of("without_inputs"), None);
|
|
assert_eq!(inputs_of("empty_inputs"), Some(HashMap::new()));
|
|
}
|
|
|
|
#[test]
|
|
fn deserializes_mcp_tools_at_graph_and_node_level() {
|
|
let yaml = r#"
|
|
name: g
|
|
mcp_servers: [github]
|
|
mcp_tools:
|
|
github:
|
|
- get_*
|
|
- list_*
|
|
start: l
|
|
nodes:
|
|
l:
|
|
id: l
|
|
type: llm
|
|
prompt: hi
|
|
mcp_tools:
|
|
github:
|
|
- search_code
|
|
next: e
|
|
e:
|
|
id: e
|
|
type: end
|
|
output: done
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert_eq!(
|
|
graph.mcp_tools.as_ref().unwrap().get("github"),
|
|
Some(&vec!["get_*".to_string(), "list_*".to_string()])
|
|
);
|
|
match &graph.get_node("l").unwrap().node_type {
|
|
NodeType::Llm(llm) => {
|
|
assert_eq!(
|
|
llm.mcp_tools.as_ref().unwrap().get("github"),
|
|
Some(&vec!["search_code".to_string()])
|
|
);
|
|
}
|
|
_ => panic!("expected Llm variant"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn deserializes_every_node_type() {
|
|
let yaml = r#"
|
|
name: all_types
|
|
start: a
|
|
nodes:
|
|
a:
|
|
id: a
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
next: s
|
|
s:
|
|
id: s
|
|
type: script
|
|
script: scripts/decide.py
|
|
next: ap
|
|
ap:
|
|
id: ap
|
|
type: approval
|
|
question: ok?
|
|
options: [yes, no]
|
|
routes:
|
|
yes: i
|
|
no: e
|
|
on_other: e
|
|
i:
|
|
id: i
|
|
type: input
|
|
question: name?
|
|
state_updates:
|
|
name: "{{input}}"
|
|
next: e
|
|
e:
|
|
id: e
|
|
type: end
|
|
output: done
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
assert!(matches!(
|
|
graph.get_node("a").unwrap().node_type,
|
|
NodeType::Agent(_)
|
|
));
|
|
assert!(matches!(
|
|
graph.get_node("s").unwrap().node_type,
|
|
NodeType::Script(_)
|
|
));
|
|
assert!(matches!(
|
|
graph.get_node("ap").unwrap().node_type,
|
|
NodeType::Approval(_)
|
|
));
|
|
assert!(matches!(
|
|
graph.get_node("i").unwrap().node_type,
|
|
NodeType::Input(_)
|
|
));
|
|
assert!(matches!(
|
|
graph.get_node("e").unwrap().node_type,
|
|
NodeType::End(_)
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn graph_silently_ignores_enabled_macros_key() {
|
|
let yaml = "name: g\nenabled_macros: [\"x\"]\nstart: x\nnodes:\n x:\n id: x\n type: end\n output: ok\n";
|
|
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert_eq!(graph.name, "g");
|
|
assert_eq!(graph.start, "x");
|
|
}
|
|
|
|
#[test]
|
|
fn graph_settings_have_sensible_defaults() {
|
|
let yaml = "name: g\nstart: x\nnodes:\n x:\n id: x\n type: end\n output: ok\n";
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
assert_eq!(graph.version, super::super::GRAPH_SCHEMA_VERSION);
|
|
assert_eq!(
|
|
graph.settings.max_loop_iterations,
|
|
super::super::DEFAULT_MAX_LOOP_ITERATIONS
|
|
);
|
|
assert!(graph.settings.log_state_snapshots);
|
|
assert!(graph.settings.validate_before_run);
|
|
assert!(graph.settings.timeout.is_none());
|
|
assert!(graph.initial_state.is_empty());
|
|
assert_eq!(graph.description, "");
|
|
}
|
|
|
|
#[test]
|
|
fn input_node_with_all_fields() {
|
|
let yaml = r#"
|
|
id: get_key
|
|
type: input
|
|
question: "Enter your API key:"
|
|
default: "{{previous_api_key}}"
|
|
validation: "len(input) > 0"
|
|
state_updates:
|
|
api_key: "{{input}}"
|
|
next: configure
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
let next_target = node.next_target().map(str::to_string);
|
|
let input = match node.node_type {
|
|
NodeType::Input(i) => i,
|
|
_ => panic!("expected Input variant"),
|
|
};
|
|
assert_eq!(input.question, "Enter your API key:");
|
|
assert_eq!(input.default.as_deref(), Some("{{previous_api_key}}"));
|
|
assert_eq!(input.validation.as_deref(), Some("len(input) > 0"));
|
|
let updates = input.state_updates.unwrap();
|
|
assert_eq!(
|
|
updates.get("api_key").map(|s| s.as_str()),
|
|
Some("{{input}}")
|
|
);
|
|
assert_eq!(next_target.as_deref(), Some("configure"));
|
|
}
|
|
|
|
#[test]
|
|
fn input_node_with_minimal_fields() {
|
|
let yaml = r#"
|
|
id: ask
|
|
type: input
|
|
question: "Describe the feature:"
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
let input = match node.node_type {
|
|
NodeType::Input(i) => i,
|
|
_ => panic!("expected Input variant"),
|
|
};
|
|
assert_eq!(input.question, "Describe the feature:");
|
|
assert!(input.default.is_none());
|
|
assert!(input.validation.is_none());
|
|
assert!(input.state_updates.is_none());
|
|
assert!(node.next.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn script_node_defaults_timeout_to_30() {
|
|
let yaml = r#"
|
|
id: s
|
|
type: script
|
|
script: scripts/decide.py
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
let script = match node.node_type {
|
|
NodeType::Script(s) => s,
|
|
_ => panic!("expected Script variant"),
|
|
};
|
|
assert_eq!(script.timeout, 30);
|
|
assert!(script.fallback.is_none());
|
|
assert!(script.state_updates.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn approval_node_carries_routes() {
|
|
let yaml = r#"
|
|
id: approve
|
|
type: approval
|
|
question: "Approve {{filename}}?"
|
|
options: [approve, reject, edit]
|
|
routes:
|
|
approve: apply
|
|
reject: end_reject
|
|
edit: edit_loop
|
|
on_other: edit_loop
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
let approval = match node.node_type {
|
|
NodeType::Approval(a) => a,
|
|
_ => panic!("expected Approval variant"),
|
|
};
|
|
assert_eq!(approval.options.len(), 3);
|
|
assert_eq!(
|
|
approval.routes.get("approve").map(|s| s.as_str()),
|
|
Some("apply")
|
|
);
|
|
assert_eq!(
|
|
approval.routes.get("reject").map(|s| s.as_str()),
|
|
Some("end_reject")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn graph_state_basic_operations() {
|
|
let mut state = GraphState::new(HashMap::new());
|
|
state.set("key1".to_string(), json!("value1"));
|
|
assert_eq!(state.get("key1"), Some(&json!("value1")));
|
|
|
|
state.visit_node("node1");
|
|
state.visit_node("node2");
|
|
state.visit_node("node1");
|
|
|
|
assert_eq!(state.loop_count("node1"), 2);
|
|
assert_eq!(state.loop_count("node2"), 1);
|
|
assert_eq!(state.loop_count("never"), 0);
|
|
assert_eq!(state.history.len(), 3);
|
|
assert_eq!(state.current_node(), Some("node1"));
|
|
}
|
|
|
|
#[test]
|
|
fn graph_state_merge_overwrites_existing_keys() {
|
|
let mut state = GraphState::new(HashMap::new());
|
|
state.set("existing".to_string(), json!("value"));
|
|
state.set("kept".to_string(), json!("untouched"));
|
|
|
|
let mut obj = serde_json::Map::new();
|
|
obj.insert("new_key".to_string(), json!("new_value"));
|
|
obj.insert("count".to_string(), json!(42));
|
|
obj.insert("existing".to_string(), json!("replaced"));
|
|
|
|
state.merge(&obj);
|
|
|
|
assert_eq!(state.get("existing"), Some(&json!("replaced")));
|
|
assert_eq!(state.get("kept"), Some(&json!("untouched")));
|
|
assert_eq!(state.get("new_key"), Some(&json!("new_value")));
|
|
assert_eq!(state.get("count"), Some(&json!(42)));
|
|
}
|
|
|
|
#[test]
|
|
fn graph_state_serializes_to_json() {
|
|
let mut initial = HashMap::new();
|
|
initial.insert("k".to_string(), json!("v"));
|
|
let state = GraphState::new(initial);
|
|
let serialized = state.to_json().unwrap();
|
|
let parsed: Value = serde_json::from_str(&serialized).unwrap();
|
|
assert_eq!(parsed.get("k"), Some(&json!("v")));
|
|
assert!(state.size_bytes() > 0);
|
|
}
|
|
|
|
#[test]
|
|
fn graph_state_initial_values_are_seeded() {
|
|
let mut initial = HashMap::new();
|
|
initial.insert("user".to_string(), json!("alice"));
|
|
let state = GraphState::new(initial);
|
|
assert_eq!(state.get("user"), Some(&json!("alice")));
|
|
assert!(state.history.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn llm_node_with_all_fields() {
|
|
let yaml = r#"
|
|
id: classify
|
|
type: llm
|
|
instructions: "You are a classifier."
|
|
prompt: "Classify: {{input_text}}"
|
|
tools:
|
|
- read_query
|
|
- "mcp:pubmed-search"
|
|
model: anthropic:claude-3-5-haiku-20241022
|
|
temperature: 0.0
|
|
top_p: 0.5
|
|
fallback: skip_classify
|
|
max_attempts: 3
|
|
max_iterations: 5
|
|
state_updates:
|
|
category: "{{output}}"
|
|
timeout: 30
|
|
next: review
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
let next_target = node.next_target().map(str::to_string);
|
|
let llm = match node.node_type {
|
|
NodeType::Llm(l) => l,
|
|
_ => panic!("expected Llm variant"),
|
|
};
|
|
assert_eq!(llm.instructions.as_deref(), Some("You are a classifier."));
|
|
assert_eq!(llm.prompt, "Classify: {{input_text}}");
|
|
let tools = llm.tools.unwrap();
|
|
assert_eq!(tools, vec!["read_query", "mcp:pubmed-search"]);
|
|
assert_eq!(
|
|
llm.model.as_deref(),
|
|
Some("anthropic:claude-3-5-haiku-20241022")
|
|
);
|
|
assert_eq!(llm.temperature, Some(0.0));
|
|
assert_eq!(llm.top_p, Some(0.5));
|
|
assert_eq!(llm.fallback.as_deref(), Some("skip_classify"));
|
|
assert_eq!(llm.max_attempts, 3);
|
|
assert_eq!(llm.max_iterations, 5);
|
|
assert_eq!(llm.timeout, Some(30));
|
|
assert!(llm.state_updates.is_some());
|
|
assert_eq!(next_target.as_deref(), Some("review"));
|
|
}
|
|
|
|
#[test]
|
|
fn llm_node_minimal_fields_use_defaults() {
|
|
let yaml = r#"
|
|
id: pure_text
|
|
type: llm
|
|
instructions: "System."
|
|
prompt: "User."
|
|
next: done
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
let llm = match node.node_type {
|
|
NodeType::Llm(l) => l,
|
|
_ => panic!("expected Llm variant"),
|
|
};
|
|
assert_eq!(llm.instructions.as_deref(), Some("System."));
|
|
assert_eq!(llm.prompt, "User.");
|
|
assert!(llm.tools.is_none());
|
|
assert!(llm.model.is_none());
|
|
assert!(llm.fallback.is_none());
|
|
assert_eq!(llm.max_attempts, 1);
|
|
assert_eq!(llm.max_iterations, 10);
|
|
}
|
|
|
|
#[test]
|
|
fn llm_node_with_just_prompt_succeeds() {
|
|
let yaml = r#"
|
|
id: pure
|
|
type: llm
|
|
prompt: "User-only — no system prompt."
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
let llm = match node.node_type {
|
|
NodeType::Llm(l) => l,
|
|
_ => panic!("expected Llm variant"),
|
|
};
|
|
assert!(llm.instructions.is_none());
|
|
assert_eq!(llm.prompt, "User-only — no system prompt.");
|
|
}
|
|
|
|
#[test]
|
|
fn llm_node_missing_prompt_fails() {
|
|
let yaml = r#"
|
|
id: bad
|
|
type: llm
|
|
instructions: "System only — no user prompt."
|
|
"#;
|
|
let result: std::result::Result<Node, _> = serde_yaml::from_str(yaml);
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn graph_parses_agent_level_top_level_fields() {
|
|
let yaml = r#"
|
|
name: single_file
|
|
start: e
|
|
model: anthropic:claude-sonnet-4-6
|
|
temperature: 0.2
|
|
top_p: 0.9
|
|
global_tools:
|
|
- web_search_coyote.sh
|
|
mcp_servers:
|
|
- pubmed-search
|
|
conversation_starters:
|
|
- "Look up 2160-0"
|
|
nodes:
|
|
e:
|
|
id: e
|
|
type: end
|
|
output: done
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
assert_eq!(graph.model.as_deref(), Some("anthropic:claude-sonnet-4-6"));
|
|
assert_eq!(graph.temperature, Some(0.2));
|
|
assert_eq!(graph.top_p, Some(0.9));
|
|
assert_eq!(graph.global_tools, vec!["web_search_coyote.sh"]);
|
|
assert_eq!(graph.mcp_servers, vec!["pubmed-search"]);
|
|
assert_eq!(graph.conversation_starters, vec!["Look up 2160-0"]);
|
|
}
|
|
|
|
#[test]
|
|
fn graph_agent_level_fields_default_when_absent() {
|
|
let yaml = "name: g\nstart: x\nnodes:\n x:\n id: x\n type: end\n output: ok\n";
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
assert!(graph.model.is_none());
|
|
assert!(graph.temperature.is_none());
|
|
assert!(graph.top_p.is_none());
|
|
assert!(graph.max_concurrent_jobs.is_none());
|
|
assert!(graph.can_spawn_agents.is_none());
|
|
assert!(graph.max_concurrent_agents.is_none());
|
|
assert!(graph.max_agent_depth.is_none());
|
|
assert!(graph.global_tools.is_empty());
|
|
assert!(graph.mcp_servers.is_empty());
|
|
assert!(graph.conversation_starters.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn graph_parses_agent_orchestration_keys() {
|
|
let yaml = r#"
|
|
name: g
|
|
can_spawn_agents: true
|
|
max_concurrent_agents: 6
|
|
max_agent_depth: 2
|
|
start: x
|
|
nodes:
|
|
x:
|
|
id: x
|
|
type: end
|
|
output: ok
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
assert_eq!(graph.can_spawn_agents, Some(true));
|
|
assert_eq!(graph.max_concurrent_agents, Some(6));
|
|
assert_eq!(graph.max_agent_depth, Some(2));
|
|
}
|
|
|
|
#[test]
|
|
fn node_ids_lists_nodes_in_order() {
|
|
let yaml = r#"
|
|
name: g
|
|
start: first
|
|
nodes:
|
|
first:
|
|
id: first
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
next: last
|
|
last:
|
|
id: last
|
|
type: end
|
|
output: done
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
assert_eq!(graph.node_ids(), vec!["first", "last"]);
|
|
}
|
|
|
|
#[test]
|
|
fn has_agent_node_detects_agent_nodes() {
|
|
let with_agent = r#"
|
|
name: g
|
|
start: a
|
|
nodes:
|
|
a:
|
|
id: a
|
|
type: agent
|
|
agent: helper
|
|
prompt: hi
|
|
next: e
|
|
e:
|
|
id: e
|
|
type: end
|
|
output: done
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(with_agent).unwrap();
|
|
assert!(graph.has_agent_node());
|
|
}
|
|
|
|
#[test]
|
|
fn has_agent_node_false_without_agent_nodes() {
|
|
let yaml = "name: g\nstart: x\nnodes:\n x:\n id: x\n type: end\n output: ok\n";
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
assert!(!graph.has_agent_node());
|
|
}
|
|
|
|
#[test]
|
|
fn parses_static_fan_out_as_many_next_targets() {
|
|
let yaml = r#"
|
|
id: triage
|
|
type: llm
|
|
prompt: Classify
|
|
next: [retrieve_local, retrieve_web, retrieve_docs]
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
let targets = node.next.as_ref().expect("next should be present");
|
|
|
|
assert!(targets.is_fan_out());
|
|
assert_eq!(
|
|
targets.as_slice(),
|
|
&[
|
|
"retrieve_local".to_string(),
|
|
"retrieve_web".to_string(),
|
|
"retrieve_docs".to_string()
|
|
]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn parses_single_target_next_as_one_variant() {
|
|
let yaml = r#"
|
|
id: triage
|
|
type: llm
|
|
prompt: Classify
|
|
next: retrieve
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
let targets = node.next.as_ref().expect("next should be present");
|
|
|
|
assert!(!targets.is_fan_out());
|
|
assert_eq!(node.next_target(), Some("retrieve"));
|
|
}
|
|
|
|
#[test]
|
|
fn next_target_treats_many_of_one_as_single() {
|
|
let yaml = r#"
|
|
id: triage
|
|
type: llm
|
|
prompt: Classify
|
|
next: [retrieve]
|
|
"#;
|
|
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert_eq!(node.next_target(), Some("retrieve"));
|
|
}
|
|
|
|
#[test]
|
|
fn next_targets_round_trips_through_yaml_for_both_variants() {
|
|
let one: NextTargets = serde_yaml::from_str(r#""foo""#).unwrap();
|
|
let reparsed: NextTargets =
|
|
serde_yaml::from_str(&serde_yaml::to_string(&one).unwrap()).unwrap();
|
|
assert_eq!(reparsed.as_slice(), &["foo".to_string()]);
|
|
|
|
let many: NextTargets = serde_yaml::from_str("[a, b, c]").unwrap();
|
|
let reparsed: NextTargets =
|
|
serde_yaml::from_str(&serde_yaml::to_string(&many).unwrap()).unwrap();
|
|
assert_eq!(
|
|
reparsed.as_slice(),
|
|
&["a".to_string(), "b".to_string(), "c".to_string()]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn parses_reducers_block_with_all_builtins() {
|
|
let yaml = r#"
|
|
name: g
|
|
start: e
|
|
reducers:
|
|
sources: append
|
|
findings: extend
|
|
context: concat
|
|
cost_usd: sum
|
|
high_score: max
|
|
low_score: min
|
|
config: merge
|
|
forced: overwrite
|
|
nodes:
|
|
e:
|
|
type: end
|
|
output: ok
|
|
"#;
|
|
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert_eq!(graph.reducers.len(), 8);
|
|
assert_eq!(graph.reducers.get("sources"), Some(&Reducer::Append));
|
|
assert_eq!(graph.reducers.get("findings"), Some(&Reducer::Extend));
|
|
assert_eq!(graph.reducers.get("context"), Some(&Reducer::Concat));
|
|
assert_eq!(graph.reducers.get("cost_usd"), Some(&Reducer::Sum));
|
|
assert_eq!(graph.reducers.get("high_score"), Some(&Reducer::Max));
|
|
assert_eq!(graph.reducers.get("low_score"), Some(&Reducer::Min));
|
|
assert_eq!(graph.reducers.get("config"), Some(&Reducer::Merge));
|
|
assert_eq!(graph.reducers.get("forced"), Some(&Reducer::Overwrite));
|
|
}
|
|
|
|
#[test]
|
|
fn reducers_default_to_empty_when_block_absent() {
|
|
let yaml = "name: g\nstart: x\nnodes:\n x:\n type: end\n";
|
|
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert!(graph.reducers.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn max_concurrency_defaults_to_eight() {
|
|
let yaml = "name: g\nstart: x\nnodes:\n x:\n type: end\n";
|
|
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert_eq!(graph.settings.max_concurrency, 8);
|
|
}
|
|
|
|
#[test]
|
|
fn max_concurrency_can_be_overridden() {
|
|
let yaml = r#"
|
|
name: g
|
|
start: x
|
|
settings:
|
|
max_concurrency: 16
|
|
nodes:
|
|
x:
|
|
type: end
|
|
"#;
|
|
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert_eq!(graph.settings.max_concurrency, 16);
|
|
}
|
|
|
|
#[test]
|
|
fn parses_map_node_with_all_fields() {
|
|
let yaml = r#"
|
|
id: fan_out
|
|
type: map
|
|
over: "{{subjects}}"
|
|
as: subject
|
|
branch: research_subject
|
|
output_key: research_result
|
|
collect_into: research_results
|
|
max_concurrency: 5
|
|
next: rank
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
let map = match node.node_type {
|
|
NodeType::Map(m) => m,
|
|
_ => panic!("expected Map variant"),
|
|
};
|
|
|
|
assert_eq!(map.over, "{{subjects}}");
|
|
assert_eq!(map.as_name, "subject");
|
|
assert_eq!(map.branch, "research_subject");
|
|
assert_eq!(map.output_key, "research_result");
|
|
assert_eq!(map.collect_into, "research_results");
|
|
assert_eq!(map.max_concurrency, Some(ConcurrencyCap::Fixed(5)));
|
|
}
|
|
|
|
#[test]
|
|
fn map_node_uses_default_output_key_and_no_concurrency_cap() {
|
|
let yaml = r#"
|
|
id: fan_out
|
|
type: map
|
|
over: "{{items}}"
|
|
as: item
|
|
branch: process
|
|
collect_into: results
|
|
"#;
|
|
let node: Node = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
let map = match node.node_type {
|
|
NodeType::Map(m) => m,
|
|
_ => panic!("expected Map variant"),
|
|
};
|
|
|
|
assert_eq!(map.output_key, "output");
|
|
assert!(map.max_concurrency.is_none());
|
|
}
|
|
|
|
fn map_node_yaml_with_cap(cap: &str) -> String {
|
|
format!(
|
|
"id: fan_out\ntype: map\nover: \"{{{{items}}}}\"\nas: item\nbranch: process\n\
|
|
collect_into: results\nmax_concurrency: {cap}\n"
|
|
)
|
|
}
|
|
|
|
fn parse_map_cap(cap: &str) -> Option<ConcurrencyCap> {
|
|
let node: Node = serde_yaml::from_str(&map_node_yaml_with_cap(cap)).unwrap();
|
|
match node.node_type {
|
|
NodeType::Map(m) => m.max_concurrency,
|
|
_ => panic!("expected Map variant"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn map_max_concurrency_integer_parses_as_fixed() {
|
|
assert_eq!(parse_map_cap("4"), Some(ConcurrencyCap::Fixed(4)));
|
|
}
|
|
|
|
#[test]
|
|
fn map_max_concurrency_zero_parses_as_fixed_zero() {
|
|
assert_eq!(parse_map_cap("0"), Some(ConcurrencyCap::Fixed(0)));
|
|
}
|
|
|
|
#[test]
|
|
fn map_max_concurrency_string_parses_as_template() {
|
|
assert_eq!(
|
|
parse_map_cap("\"{{budget}}\""),
|
|
Some(ConcurrencyCap::Template("{{budget}}".into()))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn map_max_concurrency_quoted_integer_parses_as_template() {
|
|
assert_eq!(
|
|
parse_map_cap("\"4\""),
|
|
Some(ConcurrencyCap::Template("4".into()))
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn map_max_concurrency_malformed_values_name_accepted_forms() {
|
|
for bad in ["-1", "4.5", "true", "[4]", "{n: 4}"] {
|
|
let err = serde_yaml::from_str::<Node>(&map_node_yaml_with_cap(bad))
|
|
.expect_err(&format!("max_concurrency: {bad} should fail to parse"));
|
|
let msg = err.to_string();
|
|
assert!(
|
|
msg.contains("a positive integer or a `{{template}}` string for max_concurrency"),
|
|
"max_concurrency: {bad} produced an unhelpful error: {msg}"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn map_max_concurrency_negative_error_names_offending_value() {
|
|
let err = serde_yaml::from_str::<Node>(&map_node_yaml_with_cap("-1")).unwrap_err();
|
|
assert!(err.to_string().contains("-1"), "{err}");
|
|
}
|
|
|
|
#[test]
|
|
fn concurrency_cap_serializes_as_bare_scalar() {
|
|
assert_eq!(
|
|
serde_json::to_value(ConcurrencyCap::Fixed(4)).unwrap(),
|
|
serde_json::json!(4)
|
|
);
|
|
assert_eq!(
|
|
serde_json::to_value(ConcurrencyCap::Template("{{k}}".into())).unwrap(),
|
|
serde_json::json!("{{k}}")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn full_graph_with_all_new_phase_a_fields_parses() {
|
|
let yaml = r#"
|
|
name: deep_research
|
|
start: triage
|
|
settings:
|
|
max_concurrency: 4
|
|
reducers:
|
|
sources: append
|
|
cost_usd: sum
|
|
nodes:
|
|
triage:
|
|
type: llm
|
|
prompt: Classify
|
|
next: [retrieve_local, retrieve_web]
|
|
retrieve_local:
|
|
type: rag
|
|
documents: ["./docs"]
|
|
next: synthesize
|
|
retrieve_web:
|
|
type: llm
|
|
prompt: Search web
|
|
next: synthesize
|
|
synthesize:
|
|
type: end
|
|
output: done
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
assert_eq!(graph.settings.max_concurrency, 4);
|
|
assert_eq!(graph.reducers.len(), 2);
|
|
let triage = graph.get_node("triage").unwrap();
|
|
assert!(triage.next.as_ref().unwrap().is_fan_out());
|
|
assert_eq!(triage.next.as_ref().unwrap().as_slice().len(), 2);
|
|
}
|
|
|
|
fn rag_node_of(graph: &Graph, id: &str) -> RagNode {
|
|
match &graph.get_node(id).unwrap().node_type {
|
|
NodeType::Rag(r) => r.clone(),
|
|
other => panic!("expected a rag node, got {other:?}"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn rag_node_deserializes_an_explicit_driver() {
|
|
let yaml = r#"
|
|
name: kb
|
|
start: research
|
|
nodes:
|
|
research:
|
|
type: rag
|
|
documents: ["./docs"]
|
|
driver: duckdb
|
|
next: done
|
|
done:
|
|
type: end
|
|
output: ok
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert_eq!(
|
|
rag_node_of(&graph, "research").driver.as_deref(),
|
|
Some("duckdb")
|
|
);
|
|
}
|
|
|
|
/// Workflows written before drivers existed must keep parsing, and must keep
|
|
/// asking for nothing, so `RagInitConfig` resolves them to the yaml default.
|
|
#[test]
|
|
fn rag_node_without_a_driver_stays_unset() {
|
|
let yaml = r#"
|
|
name: kb
|
|
start: research
|
|
nodes:
|
|
research:
|
|
type: rag
|
|
documents: ["./docs"]
|
|
next: done
|
|
done:
|
|
type: end
|
|
output: ok
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert_eq!(rag_node_of(&graph, "research").driver, None);
|
|
}
|
|
|
|
#[test]
|
|
fn rag_node_driver_survives_a_serialize_round_trip() {
|
|
let yaml = r#"
|
|
name: kb
|
|
start: research
|
|
nodes:
|
|
research:
|
|
type: rag
|
|
documents: ["./docs"]
|
|
driver: duckdb
|
|
next: done
|
|
done:
|
|
type: end
|
|
output: ok
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
let reparsed: Graph =
|
|
serde_yaml::from_str(&serde_yaml::to_string(&graph).unwrap()).unwrap();
|
|
|
|
assert_eq!(
|
|
rag_node_of(&reparsed, "research").driver.as_deref(),
|
|
Some("duckdb")
|
|
);
|
|
}
|
|
|
|
/// `skip_serializing_if` must keep `driver:` out of graphs that never set it.
|
|
#[test]
|
|
fn rag_node_without_a_driver_omits_the_key_when_serialized() {
|
|
let yaml = r#"
|
|
name: kb
|
|
start: research
|
|
nodes:
|
|
research:
|
|
type: rag
|
|
documents: ["./docs"]
|
|
next: done
|
|
done:
|
|
type: end
|
|
output: ok
|
|
"#;
|
|
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
|
|
|
assert!(!serde_yaml::to_string(&graph).unwrap().contains("driver"));
|
|
}
|
|
|
|
#[test]
|
|
fn graph_state_remove_returns_previous_value_and_drops_key() {
|
|
let mut state = GraphState::new(HashMap::from([("k".to_string(), json!(1))]));
|
|
|
|
assert_eq!(state.remove("k"), Some(json!(1)));
|
|
assert_eq!(state.get("k"), None);
|
|
assert!(state.data().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn graph_state_remove_missing_key_is_none() {
|
|
let mut state = GraphState::new(HashMap::new());
|
|
|
|
assert_eq!(state.remove("k"), None);
|
|
}
|
|
}
|