feat: let workflow rag nodes select a RAG driver

`RagNode` gains an optional `driver`, forwarded into `RagInitConfig` so a
graph node can build its knowledge base on duckdb instead of yaml. Nodes
that name no driver forward `None`, which still resolves to yaml, so
existing workflows are unaffected.

An unknown driver is rejected up front rather than at construction time.
`Rag::create` dispatches unknown drivers to its yaml catch-all, so a typo
would otherwise embed every document and persist the bogus string, after
which every subsequent load fails validation and the agent cannot start.
The check asks `RagData::validate()` through a probe value instead of
restating the list of valid drivers, so the two cannot drift.
This commit is contained in:
2026-08-11 13:46:59 -06:00
parent 7f90710427
commit 860566bf50
3 changed files with 276 additions and 15 deletions
+103
View File
@@ -367,6 +367,13 @@ pub struct RagNode {
#[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>>,
@@ -1152,4 +1159,100 @@ nodes:
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"));
}
}