diff --git a/src/config/agent.rs b/src/config/agent.rs index a6c051a..569ec48 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -12,6 +12,7 @@ use crate::config::prompts::{ DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS, DEFAULT_USER_INTERACTION_INSTRUCTIONS, }; +use crate::graph::types::RagNode; use crate::graph::{Graph, GraphParser, NodeType}; use crate::rag::RagInitConfig; use crate::vault::SECRET_RE; @@ -952,6 +953,30 @@ fn resolve_document_paths( Ok(document_paths) } +/// How a graph rag node describes the knowledge base it wants built. +/// +/// `driver` is forwarded as-is: `None` means the node did not ask for one, which +/// `RagInitConfig` resolves to yaml, so workflows written before drivers existed +/// keep their current storage. +/// +/// Every field is now named explicitly, so adding one to `RagInitConfig` breaks +/// this literal. That is deliberate: the new field then gets a decision about +/// whether a rag node can drive it, instead of silently taking its default. +fn rag_init_config(rag_node: &RagNode) -> RagInitConfig { + RagInitConfig { + embedding_model: rag_node.embedding_model.clone(), + chunk_size: rag_node.chunk_size, + chunk_overlap: rag_node.chunk_overlap, + reranker_model: rag_node.reranker_model.clone(), + top_k: rag_node.top_k, + batch_size: rag_node.batch_size, + extractor_model: rag_node.extractor_model.clone(), + extractor_prompt: rag_node.extractor_prompt.clone(), + graph_hops: rag_node.graph_hops, + driver: rag_node.driver.clone(), + } +} + #[allow(clippy::too_many_arguments)] async fn init_graph_rags( app: &AppConfig, @@ -989,21 +1014,18 @@ async fn init_graph_rags( }) .await? } else { - let config = RagInitConfig { - embedding_model: rag_node.embedding_model.clone(), - chunk_size: rag_node.chunk_size, - chunk_overlap: rag_node.chunk_overlap, - reranker_model: rag_node.reranker_model.clone(), - top_k: rag_node.top_k, - batch_size: rag_node.batch_size, - extractor_model: rag_node.extractor_model.clone(), - extractor_prompt: rag_node.extractor_prompt.clone(), - graph_hops: rag_node.graph_hops, - // Graph-node RAGs are yaml-only: `RagNode` has no `driver` field, so - // there is nothing to forward. The rest-pattern also keeps this literal - // from breaking on future `RagInitConfig` additions. - ..Default::default() - }; + // Checked before anything is built: an unknown driver would otherwise + // fall through `Rag::create`'s catch-all to a yaml store, embed every + // document, and persist the bogus driver string. The RAG would then be + // rejected on every subsequent load, leaving the agent unstartable. + // Graph validation catches this too, but it is skipped when + // `validate_before_run` is off, so this guard is the load-bearing one. + if let Some(driver) = &rag_node.driver + && let Some(message) = crate::graph::validator::rag_driver_error(driver) + { + bail!("rag node '{node_id}': {message}"); + } + let config = rag_init_config(rag_node); let fully_specified = config.embedding_model.is_some() && config.chunk_size.is_some() && config.chunk_overlap.is_some(); @@ -1337,4 +1359,39 @@ version: "1.0" assert_eq!(meta.description, ""); } + + #[test] + fn rag_init_config_forwards_an_explicit_driver() { + let node: RagNode = + serde_yaml::from_str("documents: [\"./docs\"]\ndriver: duckdb\n").unwrap(); + + assert_eq!(rag_init_config(&node).driver.as_deref(), Some("duckdb")); + } + + /// A node that names no driver must forward `None`, which `RagInitConfig` + /// documents as "yaml". Existing workflows therefore keep their yaml store. + #[test] + fn rag_init_config_leaves_the_driver_unset_by_default() { + let node: RagNode = serde_yaml::from_str("documents: [\"./docs\"]\n").unwrap(); + + assert_eq!(rag_init_config(&node).driver, None); + } + + /// The driver must ride alongside the rest of the node's settings, not + /// replace them. + #[test] + fn rag_init_config_forwards_the_other_settings_too() { + let node: RagNode = serde_yaml::from_str( + "documents: [\"./docs\"]\ndriver: duckdb\nchunk_size: 512\nchunk_overlap: 64\ntop_k: 7\nembedding_model: some:model\n", + ) + .unwrap(); + + let config = rag_init_config(&node); + + assert_eq!(config.driver.as_deref(), Some("duckdb")); + assert_eq!(config.chunk_size, Some(512)); + assert_eq!(config.chunk_overlap, Some(64)); + assert_eq!(config.top_k, Some(7)); + assert_eq!(config.embedding_model.as_deref(), Some("some:model")); + } } diff --git a/src/graph/types.rs b/src/graph/types.rs index 020badf..fcd3292 100644 --- a/src/graph/types.rs +++ b/src/graph/types.rs @@ -367,6 +367,13 @@ pub struct RagNode { #[serde(default, skip_serializing_if = "Option::is_none")] pub graph_hops: Option, + /// 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, + #[serde(default, skip_serializing_if = "Option::is_none")] pub state_updates: Option>, @@ -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")); + } } diff --git a/src/graph/validator.rs b/src/graph/validator.rs index f655c8e..fe438f7 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -2,6 +2,7 @@ use super::state::template_root_keys; use super::types::{Graph, Node, NodeType}; use crate::client::{Model, ModelType}; use crate::config::{Agent, AppConfig, paths}; +use crate::rag::{GraphRagConfig, RagData}; use anyhow::{Result, bail}; use std::collections::{BTreeMap, HashSet, VecDeque}; use std::path::PathBuf; @@ -96,6 +97,51 @@ pub struct GraphValidator { skill_exists: fn(&str) -> bool, } +/// A minimal `RagData` whose only interesting field is `driver`. The numeric +/// arguments are the smallest values that satisfy `validate()`'s unrelated +/// floors (top_k >= 1, and chunk_size >= 1 with chunk_overlap < chunk_size for +/// a non-attached RAG). `RagData::new` sets `attached: false`, which is the +/// correct shape here: a graph rag node always builds its own local knowledge +/// base from `documents` and can never be attached. +fn rag_driver_probe(driver: &str) -> RagData { + let mut data = RagData::new( + String::new(), + 1, + 0, + None, + 1, + None, + GraphRagConfig::default(), + ); + data.driver = driver.to_string(); + data +} + +/// `Some(message)` when `driver` is one that `RagData::validate()` would reject. +/// +/// The set of valid drivers is defined in exactly one place, `RagData::validate()`, +/// so this asks that function rather than restating the list here. +/// +/// Fails open on purpose: the first probe below uses the default driver, which is +/// valid by definition. If even that one is rejected, `validate()` has grown a +/// precondition the probe fixture no longer satisfies, and every verdict from here +/// would be a false positive that rejects working graphs. In that case we decline +/// to judge and leave enforcement to RAG construction. The +/// `rag_driver_probe_fixture_is_accepted` test turns that silent degradation into a +/// loud failure. Both `validate()` calls are load-bearing; neither is redundant. +pub(crate) fn rag_driver_error(driver: &str) -> Option { + if rag_driver_probe(&RagData::default().driver) + .validate() + .is_err() + { + return None; + } + rag_driver_probe(driver) + .validate() + .err() + .map(|err| err.to_string()) +} + impl GraphValidator { pub fn new(base_dir: impl Into) -> Self { Self { @@ -154,6 +200,11 @@ impl GraphValidator { not be written to state", )); } + if let Some(driver) = &r.driver + && let Some(message) = rag_driver_error(driver) + { + result.error(ValidationError::with_node(node_id, message)); + } } } } @@ -1031,6 +1082,7 @@ mod tests { extractor_model: None, extractor_prompt: None, graph_hops: None, + driver: None, state_updates, timeout: None, }), @@ -1385,6 +1437,55 @@ mod tests { ); } + /// Guards the fail-open branch in `rag_driver_error`. If this fails, + /// `RagData::validate()` grew a precondition the probe fixture no longer + /// satisfies and rag-node driver validation has silently switched itself off. + /// Repair the fixture in `rag_driver_probe`; do not delete this test. + #[test] + fn rag_driver_probe_fixture_is_accepted() { + let default_driver = RagData::default().driver; + assert!( + rag_driver_probe(&default_driver).validate().is_ok(), + "probe fixture rejected for the default driver '{default_driver}'" + ); + } + + #[test] + fn rag_driver_error_defers_to_ragdata_validate() { + assert_eq!(rag_driver_error("yaml"), None); + assert_eq!(rag_driver_error("duckdb"), None); + + let message = rag_driver_error("duckdbb").expect("unknown driver must be rejected"); + assert!(message.contains("duckdbb"), "got: {message}"); + } + + #[test] + fn rag_node_with_unknown_driver_errors_naming_the_node() { + let mut node = rag_node("kb", &["./docs"], true); + if let NodeType::Rag(ref mut r) = node.node_type { + r.driver = Some("postgres".into()); + } + let graph = graph_with(vec![("kb", node), ("end", end_node("end"))], "kb"); + + let result = validator().validate(&graph); + + assert!(!result.is_valid()); + let err = result.into_result().unwrap_err().to_string(); + assert!(err.contains("[kb]"), "must name the node: {err}"); + assert!(err.contains("postgres"), "must name the driver: {err}"); + } + + #[test] + fn rag_node_with_duckdb_driver_produces_no_findings() { + let mut node = rag_node("kb", &["./docs"], true); + if let NodeType::Rag(ref mut r) = node.node_type { + r.driver = Some("duckdb".into()); + } + let graph = graph_with(vec![("kb", node), ("end", end_node("end"))], "kb"); + + assert!(validator().validate(&graph).is_valid()); + } + fn agent_node(id: &str, agent: &str, next: Option<&str>) -> Node { Node { id: id.into(),