use serde_json::Value;
use crate::document::{
AgentNode, BranchCase, BranchCondition, BranchNode, Edge, FoldBody, FoldJoin, FoldNode,
GateNode, Graph, MapBody, MapNode, Node, SCHEMA_VERSION, ToolNode,
};
#[derive(Clone, Debug, Default)]
pub struct GraphBuilder {
nodes: Vec<Node>,
edges: Vec<Edge>,
}
impl GraphBuilder {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn agent(mut self, spec: AgentSpec) -> Self {
self.nodes.push(spec.into_node());
self
}
#[must_use]
pub fn tool(mut self, spec: ToolSpec) -> Self {
self.nodes.push(spec.into_node());
self
}
#[must_use]
pub fn gate(mut self, spec: GateSpec) -> Self {
self.nodes.push(spec.into_node());
self
}
#[must_use]
pub fn branch(mut self, spec: BranchSpec) -> Self {
self.nodes.push(spec.into_node());
self
}
#[must_use]
pub fn map(mut self, spec: MapSpec) -> Self {
self.nodes.push(spec.into_node());
self
}
#[must_use]
pub fn fold(mut self, spec: FoldSpec) -> Self {
self.nodes.push(spec.into_node());
self
}
#[must_use]
pub fn edge(mut self, from: impl Into<String>, to: impl Into<String>) -> Self {
self.edges.push(Edge {
from: from.into(),
to: to.into(),
label: None,
});
self
}
#[must_use]
pub fn labeled_edge(
mut self,
from: impl Into<String>,
to: impl Into<String>,
label: impl Into<String>,
) -> Self {
self.edges.push(Edge {
from: from.into(),
to: to.into(),
label: Some(label.into()),
});
self
}
#[must_use]
pub fn build(self) -> Graph {
Graph {
schema_version: SCHEMA_VERSION,
nodes: self.nodes,
edges: self.edges,
}
}
}
#[derive(Clone, Debug)]
pub struct AgentSpec {
id: String,
agent_hash: String,
name: Option<String>,
input_schema: Option<Value>,
output_schema: Option<Value>,
}
impl AgentSpec {
pub fn new(id: impl Into<String>, agent_hash: impl Into<String>) -> Self {
Self {
id: id.into(),
agent_hash: agent_hash.into(),
name: None,
input_schema: None,
output_schema: None,
}
}
#[must_use]
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
#[must_use]
pub fn input_schema(mut self, schema: Value) -> Self {
self.input_schema = Some(schema);
self
}
#[must_use]
pub fn output_schema(mut self, schema: Value) -> Self {
self.output_schema = Some(schema);
self
}
fn into_node(self) -> Node {
Node::Agent(AgentNode {
id: self.id,
agent_hash: self.agent_hash,
name: self.name,
input_schema: self.input_schema,
output_schema: self.output_schema,
})
}
}
#[derive(Clone, Debug)]
pub struct ToolSpec {
id: String,
tool: String,
name: Option<String>,
input: std::collections::BTreeMap<String, String>,
input_schema: Option<Value>,
output_schema: Option<Value>,
}
impl ToolSpec {
pub fn new(id: impl Into<String>, tool: impl Into<String>) -> Self {
Self {
id: id.into(),
tool: tool.into(),
name: None,
input: std::collections::BTreeMap::new(),
input_schema: None,
output_schema: None,
}
}
#[must_use]
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
#[must_use]
pub fn input(mut self, field: impl Into<String>, source: impl Into<String>) -> Self {
self.input.insert(field.into(), source.into());
self
}
#[must_use]
pub fn input_schema(mut self, schema: Value) -> Self {
self.input_schema = Some(schema);
self
}
#[must_use]
pub fn output_schema(mut self, schema: Value) -> Self {
self.output_schema = Some(schema);
self
}
fn into_node(self) -> Node {
Node::Tool(ToolNode {
id: self.id,
tool: self.tool,
name: self.name,
input: self.input,
input_schema: self.input_schema,
output_schema: self.output_schema,
})
}
}
#[derive(Clone, Debug)]
pub struct GateSpec {
id: String,
name: Option<String>,
prompt: Option<String>,
approval_schema: Value,
}
impl GateSpec {
pub fn new(id: impl Into<String>, approval_schema: Value) -> Self {
Self {
id: id.into(),
name: None,
prompt: None,
approval_schema,
}
}
#[must_use]
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
#[must_use]
pub fn prompt(mut self, prompt: impl Into<String>) -> Self {
self.prompt = Some(prompt.into());
self
}
fn into_node(self) -> Node {
Node::Gate(GateNode {
id: self.id,
name: self.name,
prompt: self.prompt,
approval_schema: self.approval_schema,
})
}
}
#[derive(Clone, Debug)]
pub struct BranchSpec {
id: String,
name: Option<String>,
on: Option<String>,
agent_hash: Option<String>,
cases: Vec<BranchCase>,
}
impl BranchSpec {
pub fn new(id: impl Into<String>) -> Self {
Self {
id: id.into(),
name: None,
on: None,
agent_hash: None,
cases: Vec::new(),
}
}
#[must_use]
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
#[must_use]
pub fn on(mut self, on: impl Into<String>) -> Self {
self.on = Some(on.into());
self
}
#[must_use]
pub fn agent_hash(mut self, agent_hash: impl Into<String>) -> Self {
self.agent_hash = Some(agent_hash.into());
self
}
#[must_use]
pub fn case(mut self, name: impl Into<String>, when: BranchCondition) -> Self {
self.cases.push(BranchCase {
name: name.into(),
when,
});
self
}
fn into_node(self) -> Node {
Node::Branch(BranchNode {
id: self.id,
name: self.name,
on: self.on,
agent_hash: self.agent_hash,
cases: self.cases,
})
}
}
#[derive(Clone, Debug)]
pub struct MapSpec {
id: String,
name: Option<String>,
over: String,
concurrency: u32,
body: MapBody,
output_schema: Option<Value>,
}
impl MapSpec {
pub fn new(
id: impl Into<String>,
over: impl Into<String>,
concurrency: u32,
body: MapBody,
) -> Self {
Self {
id: id.into(),
name: None,
over: over.into(),
concurrency,
body,
output_schema: None,
}
}
#[must_use]
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
#[must_use]
pub fn output_schema(mut self, schema: Value) -> Self {
self.output_schema = Some(schema);
self
}
fn into_node(self) -> Node {
Node::Map(MapNode {
id: self.id,
name: self.name,
over: self.over,
concurrency: self.concurrency,
body: self.body,
output_schema: self.output_schema,
})
}
}
#[derive(Clone, Debug)]
pub struct FoldSpec {
id: String,
name: Option<String>,
body: FoldBody,
max_iterations: u32,
stop_when: String,
join: FoldJoin,
accumulator_schema: Option<Value>,
}
impl FoldSpec {
pub fn new(
id: impl Into<String>,
body: FoldBody,
max_iterations: u32,
stop_when: impl Into<String>,
join: FoldJoin,
) -> Self {
Self {
id: id.into(),
name: None,
body,
max_iterations,
stop_when: stop_when.into(),
join,
accumulator_schema: None,
}
}
#[must_use]
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
#[must_use]
pub fn accumulator_schema(mut self, schema: Value) -> Self {
self.accumulator_schema = Some(schema);
self
}
fn into_node(self) -> Node {
Node::Fold(FoldNode {
id: self.id,
name: self.name,
body: self.body,
max_iterations: self.max_iterations,
stop_when: self.stop_when,
join: self.join,
accumulator_schema: self.accumulator_schema,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn draft_schema() -> Value {
json!({
"type": "object",
"properties": { "draft": { "type": "string" } },
"required": ["draft"]
})
}
fn canonical_flow() -> Graph {
GraphBuilder::new()
.agent(
AgentSpec::new("research", format!("sha256:{}", "1".repeat(64)))
.output_schema(draft_schema()),
)
.agent(
AgentSpec::new("review", format!("sha256:{}", "2".repeat(64)))
.input_schema(draft_schema())
.output_schema(draft_schema()),
)
.gate(
GateSpec::new(
"approve",
json!({
"type": "object",
"properties": { "approved": { "type": "boolean" } },
"required": ["approved"]
}),
)
.prompt("Approve this draft for publication?"),
)
.tool(
ToolSpec::new("publish", "http_post")
.input("body", "approve.draft")
.input("url", "config.publish_url"),
)
.edge("research", "review")
.edge("review", "approve")
.edge("approve", "publish")
.build()
}
#[test]
fn builds_the_canonical_document() {
let built = serde_json::to_value(canonical_flow()).expect("serialize built graph");
let canonical: Value = serde_json::from_str(include_str!(
"../../../examples/graphs/research-review-publish.json"
))
.expect("parse canonical fixture");
assert_eq!(
built, canonical,
"builder output must match the canonical fixture exactly"
);
}
fn fold_flow() -> Graph {
use crate::document::{FoldBody, FoldJoin};
let score_schema = json!({
"type": "object",
"properties": { "score": { "type": "number" } },
"required": ["score"]
});
GraphBuilder::new()
.agent(
AgentSpec::new("tailor", format!("sha256:{}", "3".repeat(64)))
.output_schema(score_schema.clone()),
)
.fold(
FoldSpec::new(
"refine",
FoldBody::Node("tailor".into()),
3,
"score >= 0.85",
FoldJoin::BestBy("score".into()),
)
.name("Refine to threshold")
.accumulator_schema(score_schema),
)
.build()
}
#[test]
fn builds_the_fold_document() {
let built = serde_json::to_value(fold_flow()).expect("serialize built graph");
let canonical: Value =
serde_json::from_str(include_str!("../../../examples/graphs/fold-refine.json"))
.expect("parse fold fixture");
assert_eq!(
built, canonical,
"builder output must match the fold fixture exactly"
);
}
#[test]
fn fold_document_validates() {
let summary = crate::validate(&fold_flow()).expect("fold flow is valid");
assert_eq!(summary.node_count, 2);
assert_eq!(summary.edge_count, 0);
}
#[test]
fn canonical_document_validates() {
let summary = crate::validate(&canonical_flow()).expect("canonical flow is valid");
assert_eq!(summary.node_count, 4);
assert_eq!(summary.edge_count, 3);
assert_eq!(summary.entry_nodes, vec!["research"]);
assert_eq!(summary.terminal_nodes, vec!["publish"]);
}
#[test]
fn branch_and_map_specs_build_expected_shapes() {
let graph = GraphBuilder::new()
.agent(AgentSpec::new(
"score",
format!("sha256:{}", "a".repeat(64)),
))
.branch(
BranchSpec::new("route")
.on("score.value")
.case("high", BranchCondition::Expression("score > 0.8".into()))
.case("ask", BranchCondition::ModelDecision),
)
.agent(AgentSpec::new(
"worker",
format!("sha256:{}", "b".repeat(64)),
))
.map(MapSpec::new(
"fanout",
"route.items",
4,
MapBody::Node("worker".into()),
))
.edge("score", "route")
.labeled_edge("route", "fanout", "high")
.build();
let value = serde_json::to_value(&graph).expect("serialize");
let map_payload = value["nodes"][3]["payload"].clone();
assert!(
map_payload.get("output_schema").is_none(),
"unset optional field must stay off the wire: {map_payload}"
);
assert_eq!(value["edges"][1]["label"], json!("high"));
}
#[test]
fn fold_spec_builds_expected_shape() {
use crate::document::{FoldBody, FoldJoin};
let graph = GraphBuilder::new()
.agent(AgentSpec::new(
"tailor",
format!("sha256:{}", "a".repeat(64)),
))
.fold(
FoldSpec::new(
"refine",
FoldBody::Node("tailor".into()),
3,
"score >= 0.85",
FoldJoin::BestBy("score".into()),
)
.name("Refine to threshold"),
)
.build();
let summary = crate::validate(&graph).expect("the fold flow is valid");
assert_eq!(summary.node_count, 2);
let value = serde_json::to_value(&graph).expect("serialize");
let payload = &value["nodes"][1]["payload"];
assert_eq!(value["nodes"][1]["kind"], json!("fold"));
assert_eq!(payload["max_iterations"], json!(3));
assert_eq!(payload["stop_when"], json!("score >= 0.85"));
assert_eq!(
payload["join"],
json!({"kind": "best_by", "value": "score"})
);
assert_eq!(payload["body"], json!({"kind": "node", "value": "tailor"}));
assert_eq!(payload["name"], json!("Refine to threshold"));
assert!(
payload.get("accumulator_schema").is_none(),
"unset optional field must stay off the wire: {payload}"
);
}
#[test]
fn every_spec_kind_accepts_a_display_name() {
let graph = GraphBuilder::new()
.agent(
AgentSpec::new("research", format!("sha256:{}", "1".repeat(64)))
.name("Research the topic"),
)
.tool(ToolSpec::new("publish", "http_post").name("Publish the draft"))
.gate(GateSpec::new("approve", json!({"type": "object"})).name("Approve the draft"))
.branch(
BranchSpec::new("route")
.name("Route on confidence")
.case("high", BranchCondition::Expression("score > 0.8".into())),
)
.map(
MapSpec::new("fanout", "route.items", 2, MapBody::Node("research".into()))
.name("Notify each watcher"),
)
.edge("research", "publish")
.build();
let summary = crate::validate(&graph).expect("named nodes still validate");
assert_eq!(summary.node_count, 5);
let value = serde_json::to_value(&graph).expect("serialize");
for (index, expected) in [
(0, "Research the topic"),
(1, "Publish the draft"),
(2, "Approve the draft"),
(3, "Route on confidence"),
(4, "Notify each watcher"),
] {
assert_eq!(
value["nodes"][index]["payload"]["name"],
json!(expected),
"node {index} carries its display name on the wire"
);
}
}
}