Compare commits
11
Commits
ecda258d3a
...
d6c114fe58
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d6c114fe58
|
||
|
|
912e00a627
|
||
|
|
e006e29ff1
|
||
|
|
c0067d387c
|
||
|
|
118c346345
|
||
|
|
dc677a2529
|
||
|
|
5e2b9c98ad
|
||
|
|
c458ca93a9
|
||
|
|
860566bf50
|
||
|
|
7f90710427
|
||
|
|
93a934439b
|
+72
-15
@@ -12,6 +12,7 @@ use crate::config::prompts::{
|
|||||||
DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS,
|
DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS,
|
||||||
DEFAULT_USER_INTERACTION_INSTRUCTIONS,
|
DEFAULT_USER_INTERACTION_INSTRUCTIONS,
|
||||||
};
|
};
|
||||||
|
use crate::graph::types::RagNode;
|
||||||
use crate::graph::{Graph, GraphParser, NodeType};
|
use crate::graph::{Graph, GraphParser, NodeType};
|
||||||
use crate::rag::RagInitConfig;
|
use crate::rag::RagInitConfig;
|
||||||
use crate::vault::SECRET_RE;
|
use crate::vault::SECRET_RE;
|
||||||
@@ -952,6 +953,30 @@ fn resolve_document_paths(
|
|||||||
Ok(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)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
async fn init_graph_rags(
|
async fn init_graph_rags(
|
||||||
app: &AppConfig,
|
app: &AppConfig,
|
||||||
@@ -989,21 +1014,18 @@ async fn init_graph_rags(
|
|||||||
})
|
})
|
||||||
.await?
|
.await?
|
||||||
} else {
|
} else {
|
||||||
let config = RagInitConfig {
|
// Checked before anything is built: an unknown driver would otherwise
|
||||||
embedding_model: rag_node.embedding_model.clone(),
|
// fall through `Rag::create`'s catch-all to a yaml store, embed every
|
||||||
chunk_size: rag_node.chunk_size,
|
// document, and persist the bogus driver string. The RAG would then be
|
||||||
chunk_overlap: rag_node.chunk_overlap,
|
// rejected on every subsequent load, leaving the agent unstartable.
|
||||||
reranker_model: rag_node.reranker_model.clone(),
|
// Graph validation catches this too, but it is skipped when
|
||||||
top_k: rag_node.top_k,
|
// `validate_before_run` is off, so this guard is the load-bearing one.
|
||||||
batch_size: rag_node.batch_size,
|
if let Some(driver) = &rag_node.driver
|
||||||
extractor_model: rag_node.extractor_model.clone(),
|
&& let Some(message) = crate::graph::validator::rag_driver_error(driver)
|
||||||
extractor_prompt: rag_node.extractor_prompt.clone(),
|
{
|
||||||
graph_hops: rag_node.graph_hops,
|
bail!("rag node '{node_id}': {message}");
|
||||||
// Graph-node RAGs are yaml-only: `RagNode` has no `driver` field, so
|
}
|
||||||
// there is nothing to forward. The rest-pattern also keeps this literal
|
let config = rag_init_config(rag_node);
|
||||||
// from breaking on future `RagInitConfig` additions.
|
|
||||||
..Default::default()
|
|
||||||
};
|
|
||||||
let fully_specified = config.embedding_model.is_some()
|
let fully_specified = config.embedding_model.is_some()
|
||||||
&& config.chunk_size.is_some()
|
&& config.chunk_size.is_some()
|
||||||
&& config.chunk_overlap.is_some();
|
&& config.chunk_overlap.is_some();
|
||||||
@@ -1337,4 +1359,39 @@ version: "1.0"
|
|||||||
|
|
||||||
assert_eq!(meta.description, "");
|
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"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -367,6 +367,13 @@ pub struct RagNode {
|
|||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub graph_hops: Option<usize>,
|
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")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub state_updates: Option<HashMap<String, String>>,
|
pub state_updates: Option<HashMap<String, String>>,
|
||||||
|
|
||||||
@@ -1152,4 +1159,100 @@ nodes:
|
|||||||
assert!(triage.next.as_ref().unwrap().is_fan_out());
|
assert!(triage.next.as_ref().unwrap().is_fan_out());
|
||||||
assert_eq!(triage.next.as_ref().unwrap().as_slice().len(), 2);
|
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"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ use super::state::template_root_keys;
|
|||||||
use super::types::{Graph, Node, NodeType};
|
use super::types::{Graph, Node, NodeType};
|
||||||
use crate::client::{Model, ModelType};
|
use crate::client::{Model, ModelType};
|
||||||
use crate::config::{Agent, AppConfig, paths};
|
use crate::config::{Agent, AppConfig, paths};
|
||||||
|
use crate::rag::{GraphRagConfig, RagData};
|
||||||
use anyhow::{Result, bail};
|
use anyhow::{Result, bail};
|
||||||
use std::collections::{BTreeMap, HashSet, VecDeque};
|
use std::collections::{BTreeMap, HashSet, VecDeque};
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
@@ -96,6 +97,51 @@ pub struct GraphValidator {
|
|||||||
skill_exists: fn(&str) -> bool,
|
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<String> {
|
||||||
|
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 {
|
impl GraphValidator {
|
||||||
pub fn new(base_dir: impl Into<PathBuf>) -> Self {
|
pub fn new(base_dir: impl Into<PathBuf>) -> Self {
|
||||||
Self {
|
Self {
|
||||||
@@ -154,6 +200,11 @@ impl GraphValidator {
|
|||||||
not be written to state",
|
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_model: None,
|
||||||
extractor_prompt: None,
|
extractor_prompt: None,
|
||||||
graph_hops: None,
|
graph_hops: None,
|
||||||
|
driver: None,
|
||||||
state_updates,
|
state_updates,
|
||||||
timeout: None,
|
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 {
|
fn agent_node(id: &str, agent: &str, next: Option<&str>) -> Node {
|
||||||
Node {
|
Node {
|
||||||
id: id.into(),
|
id: id.into(),
|
||||||
|
|||||||
+292
-60
@@ -18,6 +18,7 @@ use crate::vault::{Vault, interpolate_secrets};
|
|||||||
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
use bm25::{Language, SearchEngine, SearchEngineBuilder};
|
use bm25::{Language, SearchEngine, SearchEngineBuilder};
|
||||||
|
use gman::SecretError;
|
||||||
use hnsw_rs::prelude::*;
|
use hnsw_rs::prelude::*;
|
||||||
use indexmap::{IndexMap, IndexSet};
|
use indexmap::{IndexMap, IndexSet};
|
||||||
use inquire::{Confirm, Select, Text, required, validator::Validation};
|
use inquire::{Confirm, Select, Text, required, validator::Validation};
|
||||||
@@ -273,15 +274,17 @@ impl Rag {
|
|||||||
let driver = if prompt_for_driver {
|
let driver = if prompt_for_driver {
|
||||||
let options = vec![
|
let options = vec![
|
||||||
"yaml — portable, in-memory HNSW; usable from several Coyote processes at once (default)",
|
"yaml — portable, in-memory HNSW; usable from several Coyote processes at once (default)",
|
||||||
"duckdb — persistent on-disk store; vectors and content survive restarts; HNSW approximate search. Can only be open in ONE Coyote process at a time, and its driver cannot be changed later without recreating the RAG",
|
"duckdb — persistent on-disk store; vectors and content survive restarts; HNSW approximate search.",
|
||||||
];
|
];
|
||||||
let sel = Select::new("RAG storage driver:", options)
|
let sel = Select::new("RAG storage driver:", options)
|
||||||
.with_starting_cursor(0)
|
.with_starting_cursor(0)
|
||||||
.prompt()?;
|
.prompt()?;
|
||||||
if sel.starts_with("duckdb") {
|
if sel.starts_with("duckdb") {
|
||||||
println!(
|
println!(
|
||||||
"Note: a duckdb RAG can only be open in one Coyote process at a time, \
|
"Note: several Coyote processes can query a duckdb RAG at the same time, \
|
||||||
and changing its driver later means deleting and recreating the RAG."
|
but while one process is ingesting or rebuilding it the others cannot \
|
||||||
|
read it until that finishes. Changing its driver later means deleting \
|
||||||
|
and recreating the RAG."
|
||||||
);
|
);
|
||||||
"duckdb"
|
"duckdb"
|
||||||
} else {
|
} else {
|
||||||
@@ -355,38 +358,27 @@ impl Rag {
|
|||||||
let raw_content = fs::read_to_string(path).with_context(err)?;
|
let raw_content = fs::read_to_string(path).with_context(err)?;
|
||||||
|
|
||||||
// Parsed WITHOUT secret interpolation, so `driver_config` keeps its
|
// Parsed WITHOUT secret interpolation, so `driver_config` keeps its
|
||||||
// `{{...}}` placeholders. Interpolating here would bake the resolved API
|
// `{{...}}` placeholders in `self.data`. Resolution happens below, into a
|
||||||
// key into `self.data`, which `save()` then writes back to disk in
|
// function-local copy only — see `resolve_driver_config` for why the
|
||||||
// plaintext.
|
// resolved values must never travel back into `data`.
|
||||||
let data: RagData = serde_yaml::from_str(&raw_content).with_context(err)?;
|
let data: RagData = serde_yaml::from_str(&raw_content).with_context(err)?;
|
||||||
|
|
||||||
data.validate().with_context(err)?;
|
data.validate().with_context(err)?;
|
||||||
|
|
||||||
match data.driver.as_str() {
|
match data.driver.as_str() {
|
||||||
"qdrant" => {
|
"qdrant" => {
|
||||||
let host = data
|
let driver_config = resolve_driver_config(&data.driver_config, vault, name)?;
|
||||||
.driver_config
|
let host = driver_config
|
||||||
.get("host")
|
.get("host")
|
||||||
.context("qdrant driver requires 'host' in driver_config")?
|
.context("qdrant driver requires 'host' in driver_config")?
|
||||||
.clone();
|
.clone();
|
||||||
let collection = data
|
let collection = driver_config
|
||||||
.driver_config
|
|
||||||
.get("collection")
|
.get("collection")
|
||||||
.context("qdrant driver requires 'collection' in driver_config")?
|
.context("qdrant driver requires 'collection' in driver_config")?
|
||||||
.clone();
|
.clone();
|
||||||
|
let api_key = driver_config.get("api_key").map(String::as_str);
|
||||||
|
|
||||||
let api_key: Option<String> = match data.driver_config.get("api_key") {
|
let provider = QdrantProvider::new(&host, &collection, api_key).await?;
|
||||||
Some(placeholder) => {
|
|
||||||
let (resolved, _) =
|
|
||||||
interpolate_secrets(placeholder, vault).with_context(|| {
|
|
||||||
format!("Failed to resolve api_key secret for RAG '{name}'")
|
|
||||||
})?;
|
|
||||||
Some(resolved)
|
|
||||||
}
|
|
||||||
None => None,
|
|
||||||
};
|
|
||||||
|
|
||||||
let provider = QdrantProvider::new(&host, &collection, api_key.as_deref()).await?;
|
|
||||||
let embedding_model =
|
let embedding_model =
|
||||||
Model::retrieve_model(app, &data.embedding_model, ModelType::Embedding)?;
|
Model::retrieve_model(app, &data.embedding_model, ModelType::Embedding)?;
|
||||||
Ok(Rag {
|
Ok(Rag {
|
||||||
@@ -440,12 +432,7 @@ impl Rag {
|
|||||||
.with_default("QDRANT_API_KEY")
|
.with_default("QDRANT_API_KEY")
|
||||||
.with_validator(required!("This field is required"))
|
.with_validator(required!("This field is required"))
|
||||||
.prompt()?;
|
.prompt()?;
|
||||||
let resolved = vault.get_secret(&secret_name, false).with_context(|| {
|
let resolved = resolve_or_create_api_key_secret(vault, &secret_name)?;
|
||||||
format!(
|
|
||||||
"Secret '{secret_name}' not found in vault. \
|
|
||||||
Run `coyote --add-secret {secret_name}` first."
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
Some((secret_name, resolved))
|
Some((secret_name, resolved))
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
@@ -468,18 +455,28 @@ impl Rag {
|
|||||||
|
|
||||||
let collection = Select::new("Select collection:", collections).prompt()?;
|
let collection = Select::new("Select collection:", collections).prompt()?;
|
||||||
|
|
||||||
// Point IDs are read with `as_u64()`, which yields None for a JSON string.
|
let sample_id = QdrantProvider::sample_point_id(&host, &collection, api_key).await?;
|
||||||
// A UUID-keyed collection would therefore return zero hits with no error,
|
|
||||||
// so refuse it here instead of attaching something silently broken.
|
// `None` means the scroll came back with no points at all: the collection
|
||||||
if let Some(raw_id) = QdrantProvider::sample_point_id(&host, &collection, api_key).await?
|
// is empty. Attaching is not necessarily wrong — another tool may be about
|
||||||
&& raw_id.starts_with('"')
|
// to fill it — but accepting it silently yields a RAG that answers every
|
||||||
{
|
// query with nothing and never explains why, and none of the checks below
|
||||||
bail!(
|
// can tell that apart from a misconfiguration. Ask, defaulting to no, so it
|
||||||
"Collection '{collection}' uses string (UUID) point IDs. \
|
// cannot happen by accident. (`attach` already refuses to run
|
||||||
Coyote requires integer point IDs. Rebuild the collection with integer IDs \
|
// non-interactively, so there is no unattended path through this prompt.)
|
||||||
(e.g. LangChain: pass ids=list(range(len(docs))) to add_documents())."
|
if sample_id.is_none() {
|
||||||
|
println!(
|
||||||
|
"⚠️ Collection '{collection}' contains no points. Queries will return \
|
||||||
|
nothing until something writes to it."
|
||||||
);
|
);
|
||||||
|
let attach_anyway = Confirm::new("Attach to this empty collection anyway?")
|
||||||
|
.with_default(false)
|
||||||
|
.prompt()?;
|
||||||
|
if !attach_anyway {
|
||||||
|
bail!("Collection '{collection}' is empty; nothing to attach to.");
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
println!("ℹ This collection must store document text in a 'page_content' payload field.");
|
println!("ℹ This collection must store document text in a 'page_content' payload field.");
|
||||||
|
|
||||||
let dim = QdrantProvider::get_vector_dimension(&host, &collection, api_key)
|
let dim = QdrantProvider::get_vector_dimension(&host, &collection, api_key)
|
||||||
@@ -496,13 +493,7 @@ impl Rag {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
if dim > 0 {
|
if dim > 0 {
|
||||||
let candidates = embedding_model_candidates_for_dimension(dim);
|
println!("Collection uses {dim}-dim vectors.");
|
||||||
if !candidates.is_empty() {
|
|
||||||
println!(
|
|
||||||
"Collection uses {dim}-dim vectors. Likely models: {}",
|
|
||||||
candidates.join(", ")
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
println!(
|
println!(
|
||||||
"⚠️ If the embedding model doesn't match what built this collection, \
|
"⚠️ If the embedding model doesn't match what built this collection, \
|
||||||
@@ -1839,21 +1830,6 @@ fn driver_auth_header(driver: &str) -> (&'static str, &'static str) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Embedding models known to produce a given vector dimension, used to hint the
|
|
||||||
/// user toward a model compatible with the collection they just picked.
|
|
||||||
fn embedding_model_candidates_for_dimension(dim: u64) -> Vec<&'static str> {
|
|
||||||
match dim {
|
|
||||||
1536 => vec!["text-embedding-3-small", "text-embedding-ada-002"],
|
|
||||||
3072 => vec!["text-embedding-3-large"],
|
|
||||||
768 => vec!["nomic-embed-text", "all-minilm-l6-v2"],
|
|
||||||
1024 => vec![
|
|
||||||
"text-embedding-3-small (matryoshka-1024)",
|
|
||||||
"jina-embeddings-v2-base",
|
|
||||||
],
|
|
||||||
_ => vec![],
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn select_embedding_model(models: &[&Model]) -> Result<String> {
|
fn select_embedding_model(models: &[&Model]) -> Result<String> {
|
||||||
let max_width = models.iter().map(|v| v.id().len()).max().unwrap_or(0);
|
let max_width = models.iter().map(|v| v.id().len()).max().unwrap_or(0);
|
||||||
let models: Vec<_> = models
|
let models: Vec<_> = models
|
||||||
@@ -2121,6 +2097,123 @@ fn embedding_dim_for_model(model_id: &str) -> usize {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// True only for "the vault does not hold this key".
|
||||||
|
///
|
||||||
|
/// Everything else — an auth failure, a provider outage, or the vault being
|
||||||
|
/// disabled because Coyote is running inside a sandbox — must NOT be treated as
|
||||||
|
/// a missing secret. Offering to create one in those cases would prompt for a
|
||||||
|
/// value that cannot be stored and bury the real reason.
|
||||||
|
fn is_missing_secret(err: &anyhow::Error) -> bool {
|
||||||
|
matches!(
|
||||||
|
err.downcast_ref::<SecretError>(),
|
||||||
|
Some(SecretError::NotFound { .. })
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reads `secret_name` out of the vault, offering to create it in place when the
|
||||||
|
/// vault simply does not hold it yet.
|
||||||
|
///
|
||||||
|
/// Sending the user off to run `coyote --add-secret` mid-wizard discarded every
|
||||||
|
/// answer they had already given. `Vault::add_secret` does the masked prompt,
|
||||||
|
/// the provider write and the confirmation line, so this defers to it rather
|
||||||
|
/// than collecting or storing the value itself.
|
||||||
|
fn resolve_or_create_api_key_secret(vault: &Vault, secret_name: &str) -> Result<String> {
|
||||||
|
let read_err = match vault.get_secret(secret_name, false) {
|
||||||
|
Ok(secret) => return Ok(secret),
|
||||||
|
Err(err) => err,
|
||||||
|
};
|
||||||
|
if !is_missing_secret(&read_err) {
|
||||||
|
return Err(read_err)
|
||||||
|
.with_context(|| format!("Cannot read secret '{secret_name}' from the vault"));
|
||||||
|
}
|
||||||
|
|
||||||
|
let create = Confirm::new(&format!(
|
||||||
|
"Secret '{secret_name}' is not in the vault. Create it now?"
|
||||||
|
))
|
||||||
|
.with_default(true)
|
||||||
|
.prompt()?;
|
||||||
|
if !create {
|
||||||
|
bail!(
|
||||||
|
"This instance needs an API key, so '{secret_name}' has to exist before \
|
||||||
|
attaching. Add it with `coyote --add-secret {secret_name}` and re-run, or \
|
||||||
|
re-run and answer 'no' when asked whether the instance requires an API key."
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
vault
|
||||||
|
.add_secret(secret_name)
|
||||||
|
.with_context(|| format!("Failed to store secret '{secret_name}' in the vault"))?;
|
||||||
|
vault
|
||||||
|
.get_secret(secret_name, false)
|
||||||
|
.with_context(|| format!("Secret '{secret_name}' is unreadable after being stored"))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolves `{{SECRET}}` placeholders in every `driver_config` value against the
|
||||||
|
/// vault, returning a DETACHED copy.
|
||||||
|
///
|
||||||
|
/// Three properties this must preserve, each of which has already bitten:
|
||||||
|
///
|
||||||
|
/// 1. The resolved values never go back into `RagData`. `Rag::save()`
|
||||||
|
/// serializes `self.data`, and `.set rag_top_k`, `.set rag_reranker_model`
|
||||||
|
/// and every post-sync save call it — so a resolved credential parked in
|
||||||
|
/// `data.driver_config` gets written to the RAG's YAML file in plaintext the
|
||||||
|
/// next time the user changes any setting.
|
||||||
|
/// 2. The literal `{{NAME}}` text survives in `data` and on disk. Sandbox
|
||||||
|
/// credential provisioning parses that placeholder back out of the file to
|
||||||
|
/// learn which vault secret to bind into the sandbox; resolve it away and
|
||||||
|
/// provisioning silently finds nothing to register.
|
||||||
|
/// 3. Only `driver_config` is interpolated, never the whole file. The rest of a
|
||||||
|
/// RAG file is ingested document text and vectors — where `{{...}}` is
|
||||||
|
/// ordinary content (Jinja, Mustache, Vue, Go templates) that would be read
|
||||||
|
/// as a secret reference, blanked to `""`, and persisted on the next save.
|
||||||
|
/// `driver_config` is small and is the only place credentials live.
|
||||||
|
fn resolve_driver_config(
|
||||||
|
driver_config: &IndexMap<String, String>,
|
||||||
|
vault: &Vault,
|
||||||
|
rag_name: &str,
|
||||||
|
) -> Result<IndexMap<String, String>> {
|
||||||
|
resolve_driver_config_with(driver_config, rag_name, |value| {
|
||||||
|
interpolate_secrets(value, vault)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Interpolation core, taking the resolver as an argument so it can be exercised
|
||||||
|
/// without a vault. Mirrors `interpolate_secrets` / `interpolate_secrets_with`.
|
||||||
|
fn resolve_driver_config_with<F>(
|
||||||
|
driver_config: &IndexMap<String, String>,
|
||||||
|
rag_name: &str,
|
||||||
|
mut interpolate: F,
|
||||||
|
) -> Result<IndexMap<String, String>>
|
||||||
|
where
|
||||||
|
F: FnMut(&str) -> Result<(String, Vec<String>)>,
|
||||||
|
{
|
||||||
|
let mut resolved = IndexMap::with_capacity(driver_config.len());
|
||||||
|
let mut missing: Vec<String> = Vec::new();
|
||||||
|
for (key, value) in driver_config {
|
||||||
|
let (value, value_missing) = interpolate(value).with_context(|| {
|
||||||
|
format!("Failed to resolve '{key}' in driver_config for RAG '{rag_name}'")
|
||||||
|
})?;
|
||||||
|
missing.extend(value_missing);
|
||||||
|
resolved.insert(key.clone(), value);
|
||||||
|
}
|
||||||
|
|
||||||
|
// A secret the vault does not hold is NOT an error inside
|
||||||
|
// `interpolate_secrets`: it substitutes the empty string and only reports the
|
||||||
|
// name. Accepting that ships an empty credential, and the user sees an
|
||||||
|
// unexplained 401 from the server instead of the typo they made.
|
||||||
|
if !missing.is_empty() {
|
||||||
|
missing.sort();
|
||||||
|
missing.dedup();
|
||||||
|
bail!(
|
||||||
|
"RAG '{rag_name}' references secrets that are missing from the vault: {}. \
|
||||||
|
Add them with `coyote --add-secret <name>`, then try again.",
|
||||||
|
missing.join(", ")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(resolved)
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -2171,6 +2264,145 @@ mod tests {
|
|||||||
assert!(yaml.contains("{{QDRANT_API_KEY}}"));
|
assert!(yaml.contains("{{QDRANT_API_KEY}}"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const FAKE_SECRET: &str = "sk-live-fake-value-for-tests";
|
||||||
|
|
||||||
|
fn attached_qdrant_data() -> RagData {
|
||||||
|
let mut data = RagData {
|
||||||
|
driver: "qdrant".to_string(),
|
||||||
|
attached: true,
|
||||||
|
embedding_model: "text-embedding-3-small".to_string(),
|
||||||
|
top_k: 5,
|
||||||
|
..Default::default()
|
||||||
|
};
|
||||||
|
data.driver_config
|
||||||
|
.insert("host".into(), "localhost:6333".into());
|
||||||
|
data.driver_config.insert("collection".into(), "c".into());
|
||||||
|
data.driver_config
|
||||||
|
.insert("api_key".into(), "{{QDRANT_API_KEY}}".into());
|
||||||
|
data
|
||||||
|
}
|
||||||
|
|
||||||
|
/// THE invariant behind `resolve_driver_config` returning a detached copy.
|
||||||
|
///
|
||||||
|
/// `save()` serializes `self.data`, and `.set rag_top_k`, `.set
|
||||||
|
/// rag_reranker_model` and every post-sync save call it. If load ever bakes
|
||||||
|
/// the resolved credential into `data.driver_config`, the next trivial
|
||||||
|
/// setting change writes the user's plaintext API key into the RAG's YAML
|
||||||
|
/// file. The literal placeholder must also survive, because sandbox
|
||||||
|
/// credential provisioning parses it back off disk.
|
||||||
|
#[test]
|
||||||
|
fn a_save_after_load_writes_the_placeholder_not_the_resolved_secret() {
|
||||||
|
let dir = TempDir::new("driver-config-secret");
|
||||||
|
let path = dir.path.join("kb.yaml");
|
||||||
|
let data = attached_qdrant_data();
|
||||||
|
|
||||||
|
// Exactly what `load_async` does with the parsed data.
|
||||||
|
let resolved = resolve_driver_config_with(&data.driver_config, "kb", |value| {
|
||||||
|
Ok((value.replace("{{QDRANT_API_KEY}}", FAKE_SECRET), vec![]))
|
||||||
|
})
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
resolved["api_key"], FAKE_SECRET,
|
||||||
|
"the live client still has to receive the real key"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
data.driver_config["api_key"], "{{QDRANT_API_KEY}}",
|
||||||
|
"resolution must not mutate the RagData that save() serializes"
|
||||||
|
);
|
||||||
|
|
||||||
|
let rag = Rag {
|
||||||
|
app_config: Arc::new(AppConfig::default()),
|
||||||
|
name: "kb".to_string(),
|
||||||
|
path: path.display().to_string(),
|
||||||
|
embedding_model: Model::new("openai", "text-embedding-3-small"),
|
||||||
|
bm25: data.build_bm25(),
|
||||||
|
provider: Box::new(YamlProvider::from_data(&data)),
|
||||||
|
node_to_docs: IndexMap::new(),
|
||||||
|
data,
|
||||||
|
last_sources: RwLock::new(None),
|
||||||
|
};
|
||||||
|
assert!(rag.save().unwrap());
|
||||||
|
|
||||||
|
let on_disk = fs::read_to_string(&path).unwrap();
|
||||||
|
assert!(
|
||||||
|
on_disk.contains("{{QDRANT_API_KEY}}"),
|
||||||
|
"sandbox provisioning parses this placeholder back off disk: {on_disk}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!on_disk.contains(FAKE_SECRET),
|
||||||
|
"a save after load leaked the plaintext secret to {}",
|
||||||
|
path.display()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Every value is interpolated, not just `api_key` — a credential-bearing
|
||||||
|
/// field added later must not ship its raw placeholder to the server.
|
||||||
|
#[test]
|
||||||
|
fn resolution_covers_every_driver_config_value() {
|
||||||
|
let mut driver_config = IndexMap::new();
|
||||||
|
driver_config.insert("host".to_string(), "{{QDRANT_HOST}}".to_string());
|
||||||
|
driver_config.insert("collection".to_string(), "c".to_string());
|
||||||
|
driver_config.insert("api_key".to_string(), "{{QDRANT_API_KEY}}".to_string());
|
||||||
|
|
||||||
|
let resolved = resolve_driver_config_with(&driver_config, "kb", |value| {
|
||||||
|
let out = value
|
||||||
|
.replace("{{QDRANT_HOST}}", "qdrant.internal:6333")
|
||||||
|
.replace("{{QDRANT_API_KEY}}", FAKE_SECRET);
|
||||||
|
Ok((out, vec![]))
|
||||||
|
})
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(resolved["host"], "qdrant.internal:6333");
|
||||||
|
assert_eq!(resolved["collection"], "c");
|
||||||
|
assert_eq!(resolved["api_key"], FAKE_SECRET);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Missing secrets are reported together, deduplicated, and name the RAG.
|
||||||
|
#[test]
|
||||||
|
fn missing_secrets_fail_the_load_instead_of_resolving_to_empty() {
|
||||||
|
let mut driver_config = IndexMap::new();
|
||||||
|
driver_config.insert("host".to_string(), "{{QDRANT_HOST}}".to_string());
|
||||||
|
driver_config.insert("api_key".to_string(), "{{QDRANT_API_KEY}}".to_string());
|
||||||
|
|
||||||
|
let err = resolve_driver_config_with(&driver_config, "kb", |value| {
|
||||||
|
// What `interpolate_secrets` really does for an absent secret: blank it
|
||||||
|
// out and report the name rather than returning Err.
|
||||||
|
Ok((
|
||||||
|
String::new(),
|
||||||
|
vec![value.trim_matches(['{', '}']).to_string()],
|
||||||
|
))
|
||||||
|
})
|
||||||
|
.expect_err("an empty API key must not be accepted as a successful load");
|
||||||
|
|
||||||
|
let msg = err.to_string();
|
||||||
|
assert!(msg.contains("kb"), "the RAG must be named: {msg}");
|
||||||
|
assert!(msg.contains("QDRANT_HOST"), "got: {msg}");
|
||||||
|
assert!(msg.contains("QDRANT_API_KEY"), "got: {msg}");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Only a genuine NotFound may trigger the attach wizard's "create it now?"
|
||||||
|
/// offer. The vault is disabled wholesale inside a sandbox, where creating a
|
||||||
|
/// secret is impossible — misreading that as "missing" would prompt for a
|
||||||
|
/// value that cannot be stored and hide why.
|
||||||
|
#[test]
|
||||||
|
fn only_a_not_found_error_counts_as_a_missing_secret() {
|
||||||
|
let not_found = anyhow::Error::new(SecretError::NotFound {
|
||||||
|
key: "QDRANT_API_KEY".to_string(),
|
||||||
|
provider: "local",
|
||||||
|
});
|
||||||
|
assert!(is_missing_secret(¬_found));
|
||||||
|
|
||||||
|
let auth_failed = anyhow::Error::new(SecretError::AuthFailed {
|
||||||
|
provider: "local",
|
||||||
|
source: anyhow!("bad vault password"),
|
||||||
|
});
|
||||||
|
assert!(!is_missing_secret(&auth_failed));
|
||||||
|
|
||||||
|
// What `Vault::get_secret` returns in sandbox mode: a plain anyhow error.
|
||||||
|
let sandboxed = anyhow!("Vault management is disabled in sandbox mode.");
|
||||||
|
assert!(!is_missing_secret(&sandboxed));
|
||||||
|
}
|
||||||
|
|
||||||
/// A qdrant RAG's vectors MUST survive serialization.
|
/// A qdrant RAG's vectors MUST survive serialization.
|
||||||
///
|
///
|
||||||
/// `save()` omits vectors only for `driver == "duckdb"`. Qdrant must not join
|
/// `save()` omits vectors only for `driver == "duckdb"`. Qdrant must not join
|
||||||
|
|||||||
+561
-51
@@ -4,8 +4,8 @@ use std::collections::HashMap;
|
|||||||
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use duckdb::Connection;
|
|
||||||
use duckdb::types::Value;
|
use duckdb::types::Value;
|
||||||
|
use duckdb::{AccessMode, Config, Connection};
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
use log::warn;
|
use log::warn;
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
@@ -17,8 +17,12 @@ use std::sync::{Arc, Mutex, MutexGuard};
|
|||||||
/// `~/.duckdb/extensions/...`. Two threads installing the same extension at once
|
/// `~/.duckdb/extensions/...`. Two threads installing the same extension at once
|
||||||
/// both perform that move; on Windows the loser's move targets a file the winner
|
/// both perform that move; on Windows the loser's move targets a file the winner
|
||||||
/// already holds open and fails with "Access is denied", where POSIX would let the
|
/// already holds open and fails with "Access is denied", where POSIX would let the
|
||||||
/// replacement through. Guards nothing but the install step, so it is never held
|
/// replacement through. Guards nothing but the install step.
|
||||||
/// across a `DuckDbProvider::conn` guard and cannot invert lock order.
|
///
|
||||||
|
/// Lock order is `DuckDbProvider::conn` -> INSTALL_LOCK, never the reverse:
|
||||||
|
/// `ensure_writable` reopens the connection, and so may install, while holding the
|
||||||
|
/// `ConnHandle` guard, whereas nothing ever acquires a `ConnHandle` guard while
|
||||||
|
/// holding this lock. The cycle that would deadlock cannot form.
|
||||||
static INSTALL_LOCK: Mutex<()> = Mutex::new(());
|
static INSTALL_LOCK: Mutex<()> = Mutex::new(());
|
||||||
|
|
||||||
/// Derive the DuckDB sidecar path from a RAG's YAML path: `docs.yaml` -> `docs.duckdb`.
|
/// Derive the DuckDB sidecar path from a RAG's YAML path: `docs.yaml` -> `docs.duckdb`.
|
||||||
@@ -26,9 +30,44 @@ pub(crate) fn duckdb_path_from_yaml(yaml_path: &Path) -> PathBuf {
|
|||||||
yaml_path.with_extension("duckdb")
|
yaml_path.with_extension("duckdb")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// The shared connection together with the access mode it was opened with.
|
||||||
|
///
|
||||||
|
/// `conn` is an `Option` only so that an upgrade can DROP the read-only connection
|
||||||
|
/// before asking DuckDB for a read-write one. It is `Some` at every point an outside
|
||||||
|
/// caller can observe, and is never left `None` on a path that returns `Ok`.
|
||||||
|
struct ConnHandle {
|
||||||
|
conn: Option<Connection>,
|
||||||
|
/// True when `conn` was opened READ_WRITE. This lives behind the same mutex as the
|
||||||
|
/// connection itself rather than next to it in `DuckDbProvider`, so that a
|
||||||
|
/// `duplicate()` clone sharing the `Arc` observes an upgrade performed through any
|
||||||
|
/// other handle instead of keeping its own stale copy of the mode.
|
||||||
|
writable: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ConnHandle {
|
||||||
|
fn conn(&self) -> Result<&Connection> {
|
||||||
|
self.conn.as_ref().ok_or_else(Self::lost)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn conn_mut(&mut self) -> Result<&mut Connection> {
|
||||||
|
self.conn.as_mut().ok_or_else(Self::lost)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Only reachable when a read-write upgrade failed AND reopening read-only failed
|
||||||
|
/// too. Returning an error beats panicking inside a locked scope, which would
|
||||||
|
/// poison the mutex for the remaining life of the process.
|
||||||
|
fn lost() -> anyhow::Error {
|
||||||
|
anyhow!(
|
||||||
|
"The DuckDB connection was lost: upgrading it to read-write failed and the \
|
||||||
|
store could not be reopened read-only afterwards. Another process is \
|
||||||
|
holding the file; retry once it has released it."
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub struct DuckDbProvider {
|
pub struct DuckDbProvider {
|
||||||
path: PathBuf,
|
path: PathBuf,
|
||||||
conn: Arc<Mutex<Connection>>,
|
conn: Arc<Mutex<ConnHandle>>,
|
||||||
/// Embedding dimension; fixed at open time because the `FLOAT[N]` column depends on it.
|
/// Embedding dimension; fixed at open time because the `FLOAT[N]` column depends on it.
|
||||||
dim: usize,
|
dim: usize,
|
||||||
/// True once an FTS index has been built on `documents`. Until then
|
/// True once an FTS index has been built on `documents`. Until then
|
||||||
@@ -40,28 +79,140 @@ pub struct DuckDbProvider {
|
|||||||
impl DuckDbProvider {
|
impl DuckDbProvider {
|
||||||
/// Open (or create) the DuckDB file. `dim` is the embedding vector dimension,
|
/// Open (or create) the DuckDB file. `dim` is the embedding vector dimension,
|
||||||
/// supplied by the caller who knows the model.
|
/// supplied by the caller who knows the model.
|
||||||
|
///
|
||||||
|
/// Opens READ-ONLY whenever the file already carries a complete schema, so that any
|
||||||
|
/// number of Coyote processes can query the same RAG at the same time. DuckDB allows
|
||||||
|
/// many concurrent readers XOR exactly one writer, so the exclusive read-write handle
|
||||||
|
/// is taken only when there is actually something to write: when the store has to be
|
||||||
|
/// created or initialized here, or lazily through `ensure_writable` on the rebuild
|
||||||
|
/// path.
|
||||||
pub fn open(db_path: &Path, dim: usize) -> Result<Self> {
|
pub fn open(db_path: &Path, dim: usize) -> Result<Self> {
|
||||||
let conn = Connection::open(db_path).with_context(|| {
|
let (conn, writable) = Self::open_for_workload(db_path, dim)?;
|
||||||
|
// A reopened file may already carry a live FTS index from a previous session,
|
||||||
|
// in which case keyword search works immediately.
|
||||||
|
let fts_exists = Self::probe_fts_index(&conn);
|
||||||
|
Ok(Self {
|
||||||
|
path: db_path.to_path_buf(),
|
||||||
|
conn: Arc::new(Mutex::new(ConnHandle {
|
||||||
|
conn: Some(conn),
|
||||||
|
writable,
|
||||||
|
})),
|
||||||
|
dim,
|
||||||
|
fts_ready: AtomicBool::new(fts_exists),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pick the weakest access mode that can serve this store, returning the connection
|
||||||
|
/// and whether it came back writable.
|
||||||
|
fn open_for_workload(db_path: &Path, dim: usize) -> Result<(Connection, bool)> {
|
||||||
|
if db_path.exists()
|
||||||
|
&& let Ok(conn) = Self::open_read_only(db_path)
|
||||||
|
&& Self::store_is_initialized(&conn)
|
||||||
|
{
|
||||||
|
return Ok((conn, false));
|
||||||
|
}
|
||||||
|
// Three cases land here: the file does not exist yet, it could not be opened
|
||||||
|
// read-only (another process holds it read-write), or it carries no usable
|
||||||
|
// schema. All of them need a read-write handle, and the read-write attempt is
|
||||||
|
// also what produces the actionable lock error for the middle case.
|
||||||
|
let conn = Self::open_read_write(db_path, dim)?;
|
||||||
|
Ok((conn, true))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Is this file already a fully initialized Coyote store?
|
||||||
|
///
|
||||||
|
/// This gate decides whether a read-only open is viable, so it must be exact: every
|
||||||
|
/// statement in `init_schema` is rejected outright on a read-only handle, INCLUDING
|
||||||
|
/// `CREATE TABLE IF NOT EXISTS` against a table that already exists, which DuckDB
|
||||||
|
/// refuses rather than treating as a no-op. Anything missing therefore forces a
|
||||||
|
/// read-write open. The HNSW index is part of the check because a store whose tables
|
||||||
|
/// survived but whose index did not would otherwise be opened read-only and silently
|
||||||
|
/// serve every `vector_search` from a full scan.
|
||||||
|
fn store_is_initialized(conn: &Connection) -> bool {
|
||||||
|
let tables: i64 = conn
|
||||||
|
.query_row(
|
||||||
|
"SELECT count(*) FROM duckdb_tables() \
|
||||||
|
WHERE table_name IN ('vectors', 'documents')",
|
||||||
|
[],
|
||||||
|
|r| r.get(0),
|
||||||
|
)
|
||||||
|
.unwrap_or(0);
|
||||||
|
if tables < 2 {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
conn.query_row(
|
||||||
|
"SELECT count(*) FROM duckdb_indexes() WHERE index_name = 'hnsw_idx'",
|
||||||
|
[],
|
||||||
|
|r| r.get::<_, i64>(0),
|
||||||
|
)
|
||||||
|
.map(|n| n > 0)
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Open the store read-only. Many processes may hold such a handle at once.
|
||||||
|
fn open_read_only(db_path: &Path) -> Result<Connection> {
|
||||||
|
let config = Config::default()
|
||||||
|
.access_mode(AccessMode::ReadOnly)
|
||||||
|
.context("Failed to build a read-only DuckDB configuration")?;
|
||||||
|
let conn = Connection::open_with_flags(db_path, config).with_context(|| {
|
||||||
format!(
|
format!(
|
||||||
"Failed to open the DuckDB store at '{}'. If another Coyote process (or \
|
"Failed to open the DuckDB store at '{}' read-only",
|
||||||
another window) has this RAG open, close it and retry — a duckdb RAG can \
|
|
||||||
only be open in ONE process at a time. Unlike the yaml driver, its data \
|
|
||||||
lives in a single file with an exclusive lock.",
|
|
||||||
db_path.display()
|
db_path.display()
|
||||||
)
|
)
|
||||||
})?;
|
})?;
|
||||||
// Statement order is load-bearing. `hnsw_enable_experimental_persistence` is
|
Self::establish_session(&conn)?;
|
||||||
// registered BY the vss extension, so setting it before `LOAD vss` fails with
|
Ok(conn)
|
||||||
// "Setting with name ... is not in the catalog, but it exists in the vss
|
}
|
||||||
// extension". Omitting it entirely makes CREATE INDEX ... USING HNSW on a
|
|
||||||
// file-backed database fail with "HNSW index persistence is not yet supported
|
/// Open the store read-write and make sure its schema exists. Exactly one process
|
||||||
// by default". ensure vss (installing it if missing) -> ensure fts -> SET ->
|
/// may hold such a handle, and no reader from another process may hold it meanwhile.
|
||||||
// CREATE INDEX.
|
fn open_read_write(db_path: &Path, dim: usize) -> Result<Connection> {
|
||||||
Self::ensure_extension(&conn, "vss")?;
|
let config = Config::default()
|
||||||
Self::ensure_extension(&conn, "fts")?;
|
.access_mode(AccessMode::ReadWrite)
|
||||||
|
.context("Failed to build a read-write DuckDB configuration")?;
|
||||||
|
let conn = Connection::open_with_flags(db_path, config).with_context(|| {
|
||||||
|
format!(
|
||||||
|
"Failed to open the DuckDB store at '{}' for writing. Another Coyote \
|
||||||
|
process (or another window) has this RAG open: a duckdb RAG supports MANY \
|
||||||
|
concurrent READERS, but only ONE writer at a time, and a writer excludes \
|
||||||
|
readers in other processes. Close that process, or wait for its sync to \
|
||||||
|
finish, and retry.",
|
||||||
|
db_path.display()
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
Self::establish_session(&conn)?;
|
||||||
|
Self::init_schema(&conn, dim)?;
|
||||||
|
Ok(conn)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Install the per-connection session state that every connection needs, whatever
|
||||||
|
/// its access mode.
|
||||||
|
///
|
||||||
|
/// Extension `LOAD`s and `SET` are per-CONNECTION, not per-database: a connection
|
||||||
|
/// opened later — an upgrade, in particular — starts with none of this and must run
|
||||||
|
/// it again. None of these statements write to the database, so they all succeed on
|
||||||
|
/// a read-only handle.
|
||||||
|
///
|
||||||
|
/// Statement order is load-bearing. `hnsw_enable_experimental_persistence` is
|
||||||
|
/// registered BY the vss extension, so setting it before `LOAD vss` fails with
|
||||||
|
/// "Setting with name ... is not in the catalog, but it exists in the vss
|
||||||
|
/// extension". ensure vss (installing it if missing) -> ensure fts -> SET.
|
||||||
|
fn establish_session(conn: &Connection) -> Result<()> {
|
||||||
|
Self::ensure_extension(conn, "vss")?;
|
||||||
|
Self::ensure_extension(conn, "fts")?;
|
||||||
|
conn.execute_batch("SET hnsw_enable_experimental_persistence = true;")
|
||||||
|
.context("Failed to enable DuckDB HNSW index persistence")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create the tables and the vector index. Every statement here WRITES, so this only
|
||||||
|
/// ever runs on a read-write connection.
|
||||||
|
///
|
||||||
|
/// Must be preceded by `establish_session`: without the `SET` it performs, a
|
||||||
|
/// CREATE INDEX ... USING HNSW on a file-backed database fails with "HNSW index
|
||||||
|
/// persistence is not yet supported by default".
|
||||||
|
fn init_schema(conn: &Connection, dim: usize) -> Result<()> {
|
||||||
conn.execute_batch(&format!(
|
conn.execute_batch(&format!(
|
||||||
"SET hnsw_enable_experimental_persistence = true;
|
"CREATE TABLE IF NOT EXISTS vectors (
|
||||||
CREATE TABLE IF NOT EXISTS vectors (
|
|
||||||
doc_id UBIGINT PRIMARY KEY,
|
doc_id UBIGINT PRIMARY KEY,
|
||||||
embedding FLOAT[{dim}]
|
embedding FLOAT[{dim}]
|
||||||
);
|
);
|
||||||
@@ -73,16 +224,44 @@ impl DuckDbProvider {
|
|||||||
page_content TEXT NOT NULL
|
page_content TEXT NOT NULL
|
||||||
);"
|
);"
|
||||||
))
|
))
|
||||||
.context("Failed to initialize DuckDB schema")?;
|
.context("Failed to initialize DuckDB schema")
|
||||||
// A reopened file may already carry a live FTS index from a previous session,
|
}
|
||||||
// in which case keyword search works immediately.
|
|
||||||
let fts_exists = Self::probe_fts_index(&conn);
|
/// Guarantee the shared connection is read-write, upgrading it in place if it is not.
|
||||||
Ok(Self {
|
/// EVERY write path must call this before touching the store.
|
||||||
path: db_path.to_path_buf(),
|
///
|
||||||
conn: Arc::new(Mutex::new(conn)),
|
/// The upgrade replaces the `Connection` INSIDE the shared `Arc<Mutex<..>>`, so
|
||||||
dim,
|
/// `duplicate()` clones, which share that `Arc`, see it too. The read-only connection
|
||||||
fts_ready: AtomicBool::new(fts_exists),
|
/// is dropped before the read-write open because DuckDB tracks the file lock per
|
||||||
})
|
/// database instance and the old handle still holds one.
|
||||||
|
///
|
||||||
|
/// On failure the store is reopened read-only so that queries keep working, and the
|
||||||
|
/// error is propagated so the caller aborts instead of writing. A failed upgrade must
|
||||||
|
/// leave the provider degraded, never bricked, and never silently read-only-with-a-
|
||||||
|
/// caller-that-thinks-it-wrote.
|
||||||
|
fn ensure_writable(&self) -> Result<()> {
|
||||||
|
let mut handle = self.lock_conn()?;
|
||||||
|
if handle.writable {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
drop(handle.conn.take());
|
||||||
|
match Self::open_read_write(&self.path, self.dim) {
|
||||||
|
Ok(conn) => {
|
||||||
|
handle.conn = Some(conn);
|
||||||
|
handle.writable = true;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
handle.conn = Self::open_read_only(&self.path).ok();
|
||||||
|
Err(e.context(format!(
|
||||||
|
"Cannot write to the DuckDB RAG at '{}': it is open read-only and could \
|
||||||
|
not be upgraded to read-write, because another Coyote process has this \
|
||||||
|
RAG open. NOTHING WAS WRITTEN. Close the other process, or wait for it \
|
||||||
|
to finish, and retry.",
|
||||||
|
self.path.display()
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Make a DuckDB extension available on `conn`, installing it if this machine does
|
/// Make a DuckDB extension available on `conn`, installing it if this machine does
|
||||||
@@ -136,8 +315,10 @@ impl DuckDbProvider {
|
|||||||
/// `rebuild_indexes` does `CREATE OR REPLACE TABLE` and writes exactly what it is
|
/// `rebuild_indexes` does `CREATE OR REPLACE TABLE` and writes exactly what it is
|
||||||
/// given, so a thinned map is committed as the new truth on the next sync.
|
/// given, so a thinned map is committed as the new truth on the next sync.
|
||||||
pub(crate) fn read_all_vectors(&self) -> Result<IndexMap<DocumentId, Vec<f32>>> {
|
pub(crate) fn read_all_vectors(&self) -> Result<IndexMap<DocumentId, Vec<f32>>> {
|
||||||
let conn = self.lock_conn()?;
|
let handle = self.lock_conn()?;
|
||||||
let mut stmt = conn.prepare("SELECT doc_id, embedding FROM vectors")?;
|
let mut stmt = handle
|
||||||
|
.conn()?
|
||||||
|
.prepare("SELECT doc_id, embedding FROM vectors")?;
|
||||||
let raw: Vec<(u64, Vec<f32>)> = stmt
|
let raw: Vec<(u64, Vec<f32>)> = stmt
|
||||||
.query_map([], |row| {
|
.query_map([], |row| {
|
||||||
let id: u64 = row.get(0)?;
|
let id: u64 = row.get(0)?;
|
||||||
@@ -226,7 +407,7 @@ impl DuckDbProvider {
|
|||||||
/// Never `.lock().unwrap()` here: a panic anywhere inside a locked scope poisons the
|
/// Never `.lock().unwrap()` here: a panic anywhere inside a locked scope poisons the
|
||||||
/// mutex permanently, and an unwrap would then turn every subsequent RAG query into
|
/// mutex permanently, and an unwrap would then turn every subsequent RAG query into
|
||||||
/// a panic for the remaining life of the process.
|
/// a panic for the remaining life of the process.
|
||||||
fn lock_conn(&self) -> Result<MutexGuard<'_, Connection>> {
|
fn lock_conn(&self) -> Result<MutexGuard<'_, ConnHandle>> {
|
||||||
self.conn
|
self.conn
|
||||||
.lock()
|
.lock()
|
||||||
.map_err(|e| anyhow!("DuckDB connection mutex was poisoned: {e}"))
|
.map_err(|e| anyhow!("DuckDB connection mutex was poisoned: {e}"))
|
||||||
@@ -253,7 +434,7 @@ impl RagProvider for DuckDbProvider {
|
|||||||
.collect::<Vec<_>>()
|
.collect::<Vec<_>>()
|
||||||
.join(", ");
|
.join(", ");
|
||||||
let dim = self.dim;
|
let dim = self.dim;
|
||||||
let conn = self.lock_conn()?;
|
let handle = self.lock_conn()?;
|
||||||
// array_cosine_distance requires a FLOAT[N] ARRAY, not the LIST type FLOAT[].
|
// array_cosine_distance requires a FLOAT[N] ARRAY, not the LIST type FLOAT[].
|
||||||
// ORDER BY distance ASC is required for the planner to use hnsw_idx; the
|
// ORDER BY distance ASC is required for the planner to use hnsw_idx; the
|
||||||
// similarity form (DESC) does NOT trigger the ANN index. Distance is converted
|
// similarity form (DESC) does NOT trigger the ANN index. Distance is converted
|
||||||
@@ -263,7 +444,7 @@ impl RagProvider for DuckDbProvider {
|
|||||||
array_cosine_distance(embedding, [{vals}]::FLOAT[{dim}]) AS distance \
|
array_cosine_distance(embedding, [{vals}]::FLOAT[{dim}]) AS distance \
|
||||||
FROM vectors ORDER BY distance ASC LIMIT {top_k}"
|
FROM vectors ORDER BY distance ASC LIMIT {top_k}"
|
||||||
);
|
);
|
||||||
let mut stmt = conn.prepare(&sql)?;
|
let mut stmt = handle.conn()?.prepare(&sql)?;
|
||||||
let results = stmt
|
let results = stmt
|
||||||
.query_map([], |row| {
|
.query_map([], |row| {
|
||||||
let id: u64 = row.get(0)?;
|
let id: u64 = row.get(0)?;
|
||||||
@@ -300,12 +481,12 @@ impl RagProvider for DuckDbProvider {
|
|||||||
if ids.is_empty() {
|
if ids.is_empty() {
|
||||||
return Ok(vec![]);
|
return Ok(vec![]);
|
||||||
}
|
}
|
||||||
let conn = self.lock_conn()?;
|
let handle = self.lock_conn()?;
|
||||||
let placeholders = ids.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
|
let placeholders = ids.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
|
||||||
let sql =
|
let sql =
|
||||||
format!("SELECT doc_id, page_content FROM documents WHERE doc_id IN ({placeholders})");
|
format!("SELECT doc_id, page_content FROM documents WHERE doc_id IN ({placeholders})");
|
||||||
let params: Vec<Value> = ids.iter().map(|id| Value::UBigInt(id.0 as u64)).collect();
|
let params: Vec<Value> = ids.iter().map(|id| Value::UBigInt(id.0 as u64)).collect();
|
||||||
let mut stmt = conn.prepare(&sql)?;
|
let mut stmt = handle.conn()?.prepare(&sql)?;
|
||||||
let mut rows: Vec<(DocumentId, String)> = stmt
|
let mut rows: Vec<(DocumentId, String)> = stmt
|
||||||
.query_map(duckdb::params_from_iter(params.iter()), |row| {
|
.query_map(duckdb::params_from_iter(params.iter()), |row| {
|
||||||
let id: u64 = row.get(0)?;
|
let id: u64 = row.get(0)?;
|
||||||
@@ -345,8 +526,10 @@ impl RagProvider for DuckDbProvider {
|
|||||||
// Scoped: the guard MUST be dropped before `lock_conn()` is taken again
|
// Scoped: the guard MUST be dropped before `lock_conn()` is taken again
|
||||||
// below. `Mutex` is not reentrant; holding both self-deadlocks at
|
// below. `Mutex` is not reentrant; holding both self-deadlocks at
|
||||||
// runtime, with no compile error.
|
// runtime, with no compile error.
|
||||||
let conn = self.lock_conn()?;
|
let handle = self.lock_conn()?;
|
||||||
conn.query_row("SELECT count(*) FROM vectors", [], |r| r.get(0))
|
handle
|
||||||
|
.conn()?
|
||||||
|
.query_row("SELECT count(*) FROM vectors", [], |r| r.get(0))
|
||||||
.context("Failed to count existing vectors before rebuild")?
|
.context("Failed to count existing vectors before rebuild")?
|
||||||
};
|
};
|
||||||
if existing > 0 {
|
if existing > 0 {
|
||||||
@@ -372,7 +555,16 @@ impl RagProvider for DuckDbProvider {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
let dim = self.dim;
|
let dim = self.dim;
|
||||||
let mut conn = self.lock_conn()?;
|
// THE write path. Everything above this line only reads, so the upgrade happens
|
||||||
|
// here, after both guards have had their say: a rebuild that is going to be
|
||||||
|
// refused must not first take the exclusive lock away from other processes.
|
||||||
|
//
|
||||||
|
// This is also the point that makes a silently-dropped write impossible. If the
|
||||||
|
// upgrade fails, `?` aborts the rebuild before a single statement is issued and
|
||||||
|
// the caller gets the error. Nothing below can run on a read-only connection.
|
||||||
|
self.ensure_writable()?;
|
||||||
|
let mut handle = self.lock_conn()?;
|
||||||
|
let conn = handle.conn_mut()?;
|
||||||
let tx = conn
|
let tx = conn
|
||||||
.transaction()
|
.transaction()
|
||||||
.context("Failed to begin DuckDB transaction")?;
|
.context("Failed to begin DuckDB transaction")?;
|
||||||
@@ -476,9 +668,9 @@ impl RagProvider for DuckDbProvider {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn keyword_search(&self, query: &str, top_k: usize) -> Result<Vec<(DocumentId, f32)>> {
|
async fn keyword_search(&self, query: &str, top_k: usize) -> Result<Vec<(DocumentId, f32)>> {
|
||||||
let conn = self.lock_conn()?;
|
let handle = self.lock_conn()?;
|
||||||
// match_bm25 returns NULL for non-matching rows; WHERE filters them out.
|
// match_bm25 returns NULL for non-matching rows; WHERE filters them out.
|
||||||
let mut stmt = conn.prepare(
|
let mut stmt = handle.conn()?.prepare(
|
||||||
"SELECT doc_id, fts_main_documents.match_bm25(doc_id, ?) AS score
|
"SELECT doc_id, fts_main_documents.match_bm25(doc_id, ?) AS score
|
||||||
FROM documents
|
FROM documents
|
||||||
WHERE score IS NOT NULL
|
WHERE score IS NOT NULL
|
||||||
@@ -533,6 +725,10 @@ impl RagProvider for DuckDbProvider {
|
|||||||
// disk and a rebuild through one handle is immediately visible to the other.
|
// disk and a rebuild through one handle is immediately visible to the other.
|
||||||
// That is unavoidable for any on-disk store and is handled by the discipline
|
// That is unavoidable for any on-disk store and is handled by the discipline
|
||||||
// documented on `Rag`'s Clone impl, the pre-clone instance must be discarded.
|
// documented on `Rag`'s Clone impl, the pre-clone instance must be discarded.
|
||||||
|
//
|
||||||
|
// Sharing the Arc also shares the ACCESS MODE, which lives inside the ConnHandle
|
||||||
|
// rather than beside it: when one handle upgrades itself to read-write, every
|
||||||
|
// clone is upgraded with it and none is left holding a stale "read-only" belief.
|
||||||
Box::new(DuckDbProvider {
|
Box::new(DuckDbProvider {
|
||||||
path: self.path.clone(),
|
path: self.path.clone(),
|
||||||
conn: Arc::clone(&self.conn),
|
conn: Arc::clone(&self.conn),
|
||||||
@@ -547,6 +743,7 @@ mod tests {
|
|||||||
use super::*;
|
use super::*;
|
||||||
use crate::rag::provider::RagProvider;
|
use crate::rag::provider::RagProvider;
|
||||||
use crate::rag::{RagDocument, RagFile};
|
use crate::rag::{RagDocument, RagFile};
|
||||||
|
use std::process::Command;
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
use std::{env, fs};
|
use std::{env, fs};
|
||||||
|
|
||||||
@@ -625,7 +822,8 @@ mod tests {
|
|||||||
async fn open_creates_schema() {
|
async fn open_creates_schema() {
|
||||||
let db = TempDb::new("schema");
|
let db = TempDb::new("schema");
|
||||||
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
let conn = provider.conn.lock().unwrap();
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let conn = handle.conn().unwrap();
|
||||||
|
|
||||||
let v: i64 = conn
|
let v: i64 = conn
|
||||||
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
|
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
|
||||||
@@ -643,7 +841,8 @@ mod tests {
|
|||||||
let db = TempDb::new("vsearch");
|
let db = TempDb::new("vsearch");
|
||||||
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
{
|
{
|
||||||
let conn = provider.conn.lock().unwrap();
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let conn = handle.conn().unwrap();
|
||||||
// The ::FLOAT[3] cast is REQUIRED: a bare [0.1, 0.2, 0.3] literal infers
|
// The ::FLOAT[3] cast is REQUIRED: a bare [0.1, 0.2, 0.3] literal infers
|
||||||
// DOUBLE[], which does not match the FLOAT[N] ARRAY column type.
|
// DOUBLE[], which does not match the FLOAT[N] ARRAY column type.
|
||||||
conn.execute(
|
conn.execute(
|
||||||
@@ -668,7 +867,8 @@ mod tests {
|
|||||||
let db = TempDb::new("fetch");
|
let db = TempDb::new("fetch");
|
||||||
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
{
|
{
|
||||||
let conn = provider.conn.lock().unwrap();
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let conn = handle.conn().unwrap();
|
||||||
conn.execute(
|
conn.execute(
|
||||||
"INSERT INTO documents (doc_id, page_content) VALUES (42, 'hello world')",
|
"INSERT INTO documents (doc_id, page_content) VALUES (42, 'hello world')",
|
||||||
[],
|
[],
|
||||||
@@ -712,7 +912,8 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
.expect("second rebuild must not violate the primary key constraint");
|
.expect("second rebuild must not violate the primary key constraint");
|
||||||
|
|
||||||
let conn = provider.conn.lock().unwrap();
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let conn = handle.conn().unwrap();
|
||||||
let count: i64 = conn
|
let count: i64 = conn
|
||||||
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
|
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -744,7 +945,8 @@ mod tests {
|
|||||||
reloaded.vectors.insert(DocumentId(2), vec![0.7, 0.8, 0.9]);
|
reloaded.vectors.insert(DocumentId(2), vec![0.7, 0.8, 0.9]);
|
||||||
provider.rebuild_indexes(&reloaded, false).await.unwrap();
|
provider.rebuild_indexes(&reloaded, false).await.unwrap();
|
||||||
|
|
||||||
let conn = provider.conn.lock().unwrap();
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let conn = handle.conn().unwrap();
|
||||||
let count: i64 = conn
|
let count: i64 = conn
|
||||||
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
|
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -768,7 +970,8 @@ mod tests {
|
|||||||
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
|
data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]);
|
||||||
provider.rebuild_indexes(&data, true).await.unwrap();
|
provider.rebuild_indexes(&data, true).await.unwrap();
|
||||||
{
|
{
|
||||||
let conn = provider.conn.lock().unwrap();
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let conn = handle.conn().unwrap();
|
||||||
let docs: i64 = conn
|
let docs: i64 = conn
|
||||||
.query_row("SELECT COUNT(*) FROM documents", [], |r| r.get(0))
|
.query_row("SELECT COUNT(*) FROM documents", [], |r| r.get(0))
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -849,7 +1052,8 @@ mod tests {
|
|||||||
"got: {err}"
|
"got: {err}"
|
||||||
);
|
);
|
||||||
|
|
||||||
let conn = provider.conn.lock().unwrap();
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let conn = handle.conn().unwrap();
|
||||||
let count: i64 = conn
|
let count: i64 = conn
|
||||||
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
|
.query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0))
|
||||||
.unwrap();
|
.unwrap();
|
||||||
@@ -877,7 +1081,8 @@ mod tests {
|
|||||||
let dup = provider.duplicate(&minimal_rag_data());
|
let dup = provider.duplicate(&minimal_rag_data());
|
||||||
|
|
||||||
{
|
{
|
||||||
let conn = provider.conn.lock().unwrap();
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let conn = handle.conn().unwrap();
|
||||||
conn.execute(
|
conn.execute(
|
||||||
"INSERT INTO documents (doc_id, page_content) VALUES (7, 'shared row')",
|
"INSERT INTO documents (doc_id, page_content) VALUES (7, 'shared row')",
|
||||||
[],
|
[],
|
||||||
@@ -900,7 +1105,8 @@ mod tests {
|
|||||||
let db = TempDb::new("nonfinite");
|
let db = TempDb::new("nonfinite");
|
||||||
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
{
|
{
|
||||||
let conn = provider.conn.lock().unwrap();
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let conn = handle.conn().unwrap();
|
||||||
conn.execute(
|
conn.execute(
|
||||||
"INSERT INTO vectors (doc_id, embedding) VALUES (0, [0.1, 0.2, 0.3]::FLOAT[3])",
|
"INSERT INTO vectors (doc_id, embedding) VALUES (0, [0.1, 0.2, 0.3]::FLOAT[3])",
|
||||||
[],
|
[],
|
||||||
@@ -923,6 +1129,310 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Was the shared connection opened read-write? Reads the flag that lives inside the
|
||||||
|
/// shared handle, which is the same one `ensure_writable` flips.
|
||||||
|
fn is_writable(provider: &DuckDbProvider) -> bool {
|
||||||
|
provider.conn.lock().unwrap().writable
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a fully initialized store, then let the read-write handle go so the file is
|
||||||
|
/// unlocked for the next opener.
|
||||||
|
async fn seed_store(path: &Path) {
|
||||||
|
let mut provider = DuckDbProvider::open(path, 3).unwrap();
|
||||||
|
assert!(is_writable(&provider), "a fresh file must open read-write");
|
||||||
|
let mut data = populated_rag_data();
|
||||||
|
data.vectors
|
||||||
|
.insert(DocumentId::new(0, 0), vec![0.1, 0.2, 0.3]);
|
||||||
|
provider.rebuild_indexes(&data, true).await.unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a_fresh_file_is_opened_read_write() {
|
||||||
|
let db = TempDb::new("freshrw");
|
||||||
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
is_writable(&provider),
|
||||||
|
"the schema has to be created, which writes, so a missing file must open \
|
||||||
|
read-write"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn an_initialized_store_is_reopened_read_only() {
|
||||||
|
let db = TempDb::new("reopenro");
|
||||||
|
seed_store(&db.path).await;
|
||||||
|
|
||||||
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
!is_writable(&provider),
|
||||||
|
"a store that needs no schema work must open read-only, so that other Coyote \
|
||||||
|
processes can query it at the same time"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a_file_without_the_coyote_schema_is_opened_read_write() {
|
||||||
|
let db = TempDb::new("noschema");
|
||||||
|
{
|
||||||
|
// A valid DuckDB file that is not one of ours. Opening it read-only would
|
||||||
|
// strand it forever: the init batch is refused on a read-only handle.
|
||||||
|
let conn = Connection::open(&db.path).unwrap();
|
||||||
|
conn.execute_batch("CREATE TABLE unrelated (x INTEGER);")
|
||||||
|
.unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
is_writable(&provider),
|
||||||
|
"a schema-less file must open read-write"
|
||||||
|
);
|
||||||
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let tables: i64 = handle
|
||||||
|
.conn()
|
||||||
|
.unwrap()
|
||||||
|
.query_row(
|
||||||
|
"SELECT count(*) FROM duckdb_tables() \
|
||||||
|
WHERE table_name IN ('vectors', 'documents')",
|
||||||
|
[],
|
||||||
|
|r| r.get(0),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(tables, 2, "the schema must have been created");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a_read_only_store_still_serves_vector_and_keyword_search() {
|
||||||
|
let db = TempDb::new("rosearch");
|
||||||
|
seed_store(&db.path).await;
|
||||||
|
|
||||||
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
assert!(!is_writable(&provider), "precondition: opened read-only");
|
||||||
|
|
||||||
|
let hits = provider
|
||||||
|
.vector_search(&[0.1, 0.2, 0.3], 5, 0.0)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
hits.len(),
|
||||||
|
1,
|
||||||
|
"the persisted HNSW index must be queryable read-only"
|
||||||
|
);
|
||||||
|
assert!(hits[0].1 > 0.99);
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
provider.has_native_keyword_search(),
|
||||||
|
"the FTS index built by the previous session must still be detected on a \
|
||||||
|
read-only handle"
|
||||||
|
);
|
||||||
|
let kw = provider.keyword_search("alpha", 5).await.unwrap();
|
||||||
|
assert_eq!(kw.len(), 1, "keyword search must work read-only");
|
||||||
|
|
||||||
|
let docs = provider
|
||||||
|
.fetch_content(&[DocumentId::new(0, 0)])
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(docs[0].1, "alpha keyword");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a_read_only_handle_refuses_a_direct_write() {
|
||||||
|
let db = TempDb::new("rorefuse");
|
||||||
|
seed_store(&db.path).await;
|
||||||
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
assert!(!is_writable(&provider), "precondition: opened read-only");
|
||||||
|
|
||||||
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
let err = handle
|
||||||
|
.conn()
|
||||||
|
.unwrap()
|
||||||
|
.execute(
|
||||||
|
"INSERT INTO documents (doc_id, page_content) VALUES (99, 'nope')",
|
||||||
|
[],
|
||||||
|
)
|
||||||
|
.unwrap_err();
|
||||||
|
|
||||||
|
// The backstop, not the primary defence: `rebuild_indexes` upgrades first and
|
||||||
|
// never reaches a write on a read-only handle. It matters anyway because DuckDB
|
||||||
|
// lets `transaction()` open and `commit()` return Ok on a read-only connection,
|
||||||
|
// so a write that slipped through would look like it had succeeded.
|
||||||
|
assert!(
|
||||||
|
err.to_string().contains("read-only mode"),
|
||||||
|
"a read-only handle must refuse writes loudly; got: {err}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn rebuild_indexes_upgrades_a_read_only_connection() {
|
||||||
|
let db = TempDb::new("upgrade");
|
||||||
|
seed_store(&db.path).await;
|
||||||
|
|
||||||
|
let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
assert!(!is_writable(&provider), "precondition: opened read-only");
|
||||||
|
|
||||||
|
let mut data = populated_rag_data();
|
||||||
|
data.vectors = provider.read_all_vectors().unwrap();
|
||||||
|
data.vectors
|
||||||
|
.insert(DocumentId::new(1, 0), vec![0.4, 0.5, 0.6]);
|
||||||
|
provider.rebuild_indexes(&data, false).await.unwrap();
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
is_writable(&provider),
|
||||||
|
"the write path must have upgraded the connection in place"
|
||||||
|
);
|
||||||
|
|
||||||
|
// The upgraded connection is a NEW connection, so the per-connection session
|
||||||
|
// state has to have been re-established on it. Without the re-run `SET`, the
|
||||||
|
// CREATE INDEX ... USING HNSW inside rebuild_indexes would already have failed.
|
||||||
|
let persisted: String = {
|
||||||
|
let handle = provider.conn.lock().unwrap();
|
||||||
|
handle
|
||||||
|
.conn()
|
||||||
|
.unwrap()
|
||||||
|
.query_row(
|
||||||
|
"SELECT CAST(current_setting('hnsw_enable_experimental_persistence') \
|
||||||
|
AS VARCHAR)",
|
||||||
|
[],
|
||||||
|
|r| r.get(0),
|
||||||
|
)
|
||||||
|
.unwrap()
|
||||||
|
};
|
||||||
|
assert_eq!(
|
||||||
|
persisted, "true",
|
||||||
|
"the upgraded connection must re-run the SET; session state does not carry \
|
||||||
|
over from the dropped read-only connection"
|
||||||
|
);
|
||||||
|
|
||||||
|
drop(provider);
|
||||||
|
let reopened = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
let all = reopened.read_all_vectors().unwrap();
|
||||||
|
assert_eq!(all.len(), 2, "the upgraded write must have reached disk");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn an_upgrade_is_visible_through_duplicate_clones() {
|
||||||
|
let db = TempDb::new("upgradedup");
|
||||||
|
seed_store(&db.path).await;
|
||||||
|
|
||||||
|
let mut provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
assert!(!is_writable(&provider), "precondition: opened read-only");
|
||||||
|
let dup = provider.duplicate(&minimal_rag_data());
|
||||||
|
|
||||||
|
let mut data = populated_rag_data();
|
||||||
|
data.vectors = provider.read_all_vectors().unwrap();
|
||||||
|
provider.rebuild_indexes(&data, false).await.unwrap();
|
||||||
|
|
||||||
|
// `duplicate()` shares the Arc, and the access mode lives inside it, so the clone
|
||||||
|
// must observe the upgrade rather than keep believing it is read-only.
|
||||||
|
let via_dup = dup.fetch_content(&[DocumentId::new(0, 0)]).await.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
via_dup.len(),
|
||||||
|
1,
|
||||||
|
"the clone must still read after an upgrade"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
is_writable(&provider),
|
||||||
|
"the shared handle must report writable to every clone"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Two REAL OS processes reading one store at the same time.
|
||||||
|
///
|
||||||
|
/// Ignored by default because it re-executes the test binary as a child process,
|
||||||
|
/// which is heavier and more environment-dependent than the rest of the suite. Run it
|
||||||
|
/// with:
|
||||||
|
/// cargo test --all -- --ignored duckdb_store_is_shared_across_processes
|
||||||
|
///
|
||||||
|
/// It cannot be written as an ordinary in-process test: DuckDB keeps ONE database
|
||||||
|
/// instance per process, so a second open in the same process bypasses the file lock
|
||||||
|
/// entirely (a read-write open succeeds even while this process holds a read-only
|
||||||
|
/// one). Only separate processes exercise the lock this feature exists to avoid.
|
||||||
|
#[tokio::test]
|
||||||
|
#[ignore = "spawns a second OS process; run explicitly with --ignored"]
|
||||||
|
async fn duckdb_store_is_shared_across_processes() {
|
||||||
|
const CHILD_DB: &str = "COYOTE_DUCKDB_MULTIPROC_DB";
|
||||||
|
const CHILD_EXPECT: &str = "COYOTE_DUCKDB_MULTIPROC_EXPECT";
|
||||||
|
const TEST_NAME: &str =
|
||||||
|
"rag::providers::duckdb::tests::duckdb_store_is_shared_across_processes";
|
||||||
|
|
||||||
|
if let Ok(path) = env::var(CHILD_DB) {
|
||||||
|
let expect = env::var(CHILD_EXPECT).unwrap_or_default();
|
||||||
|
let opened = DuckDbProvider::open(Path::new(&path), 3);
|
||||||
|
match expect.as_str() {
|
||||||
|
"readable" => {
|
||||||
|
let provider = opened.expect(
|
||||||
|
"a second process must be able to open a store that another \
|
||||||
|
process holds READ-ONLY",
|
||||||
|
);
|
||||||
|
assert!(!is_writable(&provider), "the child must land read-only");
|
||||||
|
let hits = provider
|
||||||
|
.vector_search(&[0.1, 0.2, 0.3], 5, 0.0)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(hits.len(), 1, "the child must read the seeded vector");
|
||||||
|
}
|
||||||
|
"blocked" => {
|
||||||
|
let err = opened.err().expect(
|
||||||
|
"a second process must NOT be able to open a store that another \
|
||||||
|
process holds READ-WRITE",
|
||||||
|
);
|
||||||
|
let msg = format!("{err:#}");
|
||||||
|
assert!(
|
||||||
|
msg.contains("concurrent READERS") && msg.contains("only ONE writer"),
|
||||||
|
"the lock error must explain the reader/writer rule; got: {msg}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
other => panic!("unknown child expectation {other:?}"),
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let db = TempDb::new("multiproc");
|
||||||
|
seed_store(&db.path).await;
|
||||||
|
|
||||||
|
let run_child = |expect: &str| {
|
||||||
|
Command::new(env::current_exe().unwrap())
|
||||||
|
.args(["--exact", "--ignored", "--nocapture", TEST_NAME])
|
||||||
|
.env(CHILD_DB, &db.path)
|
||||||
|
.env(CHILD_EXPECT, expect)
|
||||||
|
.output()
|
||||||
|
.expect("failed to spawn the child test process")
|
||||||
|
};
|
||||||
|
|
||||||
|
// Phase 1: this process holds a READ-ONLY handle. The child must get one too.
|
||||||
|
let provider = DuckDbProvider::open(&db.path, 3).unwrap();
|
||||||
|
assert!(!is_writable(&provider), "precondition: parent is read-only");
|
||||||
|
let out = run_child("readable");
|
||||||
|
assert!(
|
||||||
|
out.status.success(),
|
||||||
|
"child could not share the read-only store:\n{}\n{}",
|
||||||
|
String::from_utf8_lossy(&out.stdout),
|
||||||
|
String::from_utf8_lossy(&out.stderr)
|
||||||
|
);
|
||||||
|
let hits = provider
|
||||||
|
.vector_search(&[0.1, 0.2, 0.3], 5, 0.0)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
hits.len(),
|
||||||
|
1,
|
||||||
|
"the parent must still read after the child ran"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Phase 2: upgrade this process to READ-WRITE. The child must now be refused,
|
||||||
|
// with the message that explains why.
|
||||||
|
provider.ensure_writable().unwrap();
|
||||||
|
let out = run_child("blocked");
|
||||||
|
assert!(
|
||||||
|
out.status.success(),
|
||||||
|
"a writer must exclude other processes, with an actionable error:\n{}\n{}",
|
||||||
|
String::from_utf8_lossy(&out.stdout),
|
||||||
|
String::from_utf8_lossy(&out.stderr)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn duckdb_path_from_yaml_swaps_extension() {
|
fn duckdb_path_from_yaml_swaps_extension() {
|
||||||
let p = duckdb_path_from_yaml(Path::new("/tmp/rags/docs.yaml"));
|
let p = duckdb_path_from_yaml(Path::new("/tmp/rags/docs.yaml"));
|
||||||
|
|||||||
+277
-26
@@ -3,10 +3,124 @@ use crate::rag::{DocumentId, RagData};
|
|||||||
|
|
||||||
use anyhow::{Context, Result, bail};
|
use anyhow::{Context, Result, bail};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
use parking_lot::RwLock;
|
||||||
use reqwest::header::{HeaderMap, HeaderValue};
|
use reqwest::header::{HeaderMap, HeaderValue};
|
||||||
use reqwest::{Client, Response, StatusCode};
|
use reqwest::{Client, Response, StatusCode};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
/// Marks a `DocumentId` that stands in for a point id Coyote cannot carry
|
||||||
|
/// directly. Qdrant accepts UUID strings as point ids, and that is what
|
||||||
|
/// LangChain writes by default.
|
||||||
|
///
|
||||||
|
/// `DocumentId` packs `(file_index, document_index)` into one `usize` with the
|
||||||
|
/// file index in the high half, so this bit is only reachable at a file index of
|
||||||
|
/// 2^31. Nothing local gets near that, and an attached RAG builds no local index
|
||||||
|
/// at all — `data.files` and `data.vectors` stay empty and every
|
||||||
|
/// `DocumentId::split` caller early-returns on `data.attached`. Along the
|
||||||
|
/// attached path the id is an opaque key carried through RRF, which is what
|
||||||
|
/// makes a synthetic one safe here and nowhere else.
|
||||||
|
const SYNTHETIC_ID_TAG: usize = 1 << (usize::BITS - 1);
|
||||||
|
|
||||||
|
/// Two-way map between a raw Qdrant point id and the `DocumentId` the retrieval
|
||||||
|
/// pipeline sees.
|
||||||
|
///
|
||||||
|
/// Only ids that cannot survive the round trip are interned. A plain `u64` that
|
||||||
|
/// fits below the tag keeps mapping to itself, so integer-keyed collections
|
||||||
|
/// behave exactly as they did before this map existed.
|
||||||
|
#[derive(Default)]
|
||||||
|
struct PointIdInterner {
|
||||||
|
handles: HashMap<String, DocumentId>,
|
||||||
|
raw: HashMap<DocumentId, Value>,
|
||||||
|
next: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PointIdInterner {
|
||||||
|
/// The `DocumentId` for a raw point id, minting a handle if one is needed.
|
||||||
|
///
|
||||||
|
/// `None` only for a missing id, which is a malformed response.
|
||||||
|
fn document_id(&mut self, raw: &Value) -> Option<DocumentId> {
|
||||||
|
if raw.is_null() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
// The pre-existing integer path, unchanged. `try_from` rather than `as`
|
||||||
|
// so a value too wide for the target's `usize` is interned instead of
|
||||||
|
// silently truncated into a different point.
|
||||||
|
if let Some(n) = raw.as_u64()
|
||||||
|
&& let Ok(n) = usize::try_from(n)
|
||||||
|
&& n & SYNTHETIC_ID_TAG == 0
|
||||||
|
{
|
||||||
|
return Some(DocumentId(n));
|
||||||
|
}
|
||||||
|
Some(self.intern(raw))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn intern(&mut self, raw: &Value) -> DocumentId {
|
||||||
|
// Keyed on the JSON rendering, so the string "1" and the integer 1 are
|
||||||
|
// not conflated into one point.
|
||||||
|
let key = raw.to_string();
|
||||||
|
if let Some(handle) = self.handles.get(&key) {
|
||||||
|
return *handle;
|
||||||
|
}
|
||||||
|
let handle = DocumentId(SYNTHETIC_ID_TAG | self.next);
|
||||||
|
self.next += 1;
|
||||||
|
self.handles.insert(key, handle);
|
||||||
|
self.raw.insert(handle, raw.clone());
|
||||||
|
handle
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The original id for a handle, or `None` when the id was never interned —
|
||||||
|
/// i.e. it is a plain integer that is already its own id.
|
||||||
|
fn raw_id(&self, handle: DocumentId) -> Option<&Value> {
|
||||||
|
self.raw.get(&handle)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Builds the `ids` array for an outbound `/points` fetch. Every entry is the
|
||||||
|
/// id Qdrant issued, integer or string; a synthetic handle must never leave
|
||||||
|
/// this process.
|
||||||
|
fn outbound_ids(&self, ids: &[DocumentId]) -> Vec<Value> {
|
||||||
|
ids.iter()
|
||||||
|
.map(|id| match self.raw_id(*id) {
|
||||||
|
Some(raw) => raw.clone(),
|
||||||
|
None => Value::from(id.0 as u64),
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_search_hits(
|
||||||
|
interner: &mut PointIdInterner,
|
||||||
|
body: &Value,
|
||||||
|
min_score: f32,
|
||||||
|
) -> Result<Vec<(DocumentId, f32)>> {
|
||||||
|
let hits = body["result"]
|
||||||
|
.as_array()
|
||||||
|
.context("Unexpected /points/search response shape")?;
|
||||||
|
|
||||||
|
Ok(hits
|
||||||
|
.iter()
|
||||||
|
.filter_map(|pt| {
|
||||||
|
let score = pt["score"].as_f64()? as f32;
|
||||||
|
Some((interner.document_id(&pt["id"])?, score))
|
||||||
|
})
|
||||||
|
.filter(|(_, score)| *score > min_score)
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_points(interner: &mut PointIdInterner, body: &Value) -> Result<Vec<(DocumentId, String)>> {
|
||||||
|
let points = body["result"]
|
||||||
|
.as_array()
|
||||||
|
.context("Unexpected /points response shape")?;
|
||||||
|
|
||||||
|
Ok(points
|
||||||
|
.iter()
|
||||||
|
.filter_map(|pt| {
|
||||||
|
let text = pt["payload"]["page_content"].as_str()?.to_string();
|
||||||
|
Some((interner.document_id(&pt["id"])?, text))
|
||||||
|
})
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
/// Render Qdrant's error envelope into a human-readable message.
|
/// Render Qdrant's error envelope into a human-readable message.
|
||||||
///
|
///
|
||||||
@@ -65,6 +179,7 @@ pub struct QdrantProvider {
|
|||||||
client: Client,
|
client: Client,
|
||||||
base_url: String,
|
base_url: String,
|
||||||
collection: String,
|
collection: String,
|
||||||
|
point_ids: Arc<RwLock<PointIdInterner>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl QdrantProvider {
|
impl QdrantProvider {
|
||||||
@@ -139,6 +254,7 @@ impl QdrantProvider {
|
|||||||
client,
|
client,
|
||||||
base_url,
|
base_url,
|
||||||
collection: collection.to_string(),
|
collection: collection.to_string(),
|
||||||
|
point_ids: Arc::default(),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -252,22 +368,11 @@ impl RagProvider for QdrantProvider {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
let data: Value = resp.json().await?;
|
let data: Value = resp.json().await?;
|
||||||
let results = data["result"]
|
// The interner is what lets a UUID-keyed collection work: a string id gets
|
||||||
.as_array()
|
// a synthetic handle here and the original is replayed by `fetch_content`.
|
||||||
.context("Unexpected /points/search response shape")?
|
let mut interner = self.point_ids.write();
|
||||||
.iter()
|
|
||||||
.filter_map(|pt| {
|
|
||||||
// String (UUID) IDs yield None here and are dropped. The attach
|
|
||||||
// wizard rejects such collections up front so this cannot silently
|
|
||||||
// become "zero results, no error".
|
|
||||||
let id = pt["id"].as_u64()? as usize;
|
|
||||||
let score = pt["score"].as_f64()? as f32;
|
|
||||||
Some((DocumentId(id), score))
|
|
||||||
})
|
|
||||||
.filter(|(_, score)| *score > min_score)
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
Ok(results)
|
parse_search_hits(&mut interner, &data, min_score)
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn fetch_content(&self, ids: &[DocumentId]) -> Result<Vec<(DocumentId, String)>> {
|
async fn fetch_content(&self, ids: &[DocumentId]) -> Result<Vec<(DocumentId, String)>> {
|
||||||
@@ -275,7 +380,8 @@ impl RagProvider for QdrantProvider {
|
|||||||
return Ok(vec![]);
|
return Ok(vec![]);
|
||||||
}
|
}
|
||||||
let url = format!("{}/collections/{}/points", self.base_url, self.collection);
|
let url = format!("{}/collections/{}/points", self.base_url, self.collection);
|
||||||
let id_list: Vec<u64> = ids.iter().map(|d| d.0 as u64).collect();
|
// Qdrant is asked for the ids it issued, never for a synthetic handle.
|
||||||
|
let id_list = self.point_ids.read().outbound_ids(ids);
|
||||||
let body = serde_json::json!({
|
let body = serde_json::json!({
|
||||||
"ids": id_list,
|
"ids": id_list,
|
||||||
"with_payload": true,
|
"with_payload": true,
|
||||||
@@ -291,16 +397,10 @@ impl RagProvider for QdrantProvider {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
let data: Value = resp.json().await?;
|
let data: Value = resp.json().await?;
|
||||||
let mut rows: Vec<(DocumentId, String)> = data["result"]
|
let mut rows = {
|
||||||
.as_array()
|
let mut interner = self.point_ids.write();
|
||||||
.context("Unexpected /points response shape")?
|
parse_points(&mut interner, &data)?
|
||||||
.iter()
|
};
|
||||||
.filter_map(|pt| {
|
|
||||||
let id = pt["id"].as_u64()? as usize;
|
|
||||||
let text = pt["payload"]["page_content"].as_str()?.to_string();
|
|
||||||
Some((DocumentId(id), text))
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
// `/points` does not guarantee response order matches request order, and the
|
// `/points` does not guarantee response order matches request order, and the
|
||||||
// caller's RRF ranking is carried by that order. Restore it.
|
// caller's RRF ranking is carried by that order. Restore it.
|
||||||
let position: HashMap<DocumentId, usize> =
|
let position: HashMap<DocumentId, usize> =
|
||||||
@@ -328,10 +428,17 @@ impl RagProvider for QdrantProvider {
|
|||||||
// Cloning the client shares the connection pool and the injected api-key
|
// Cloning the client shares the connection pool and the injected api-key
|
||||||
// header. Sharing is correct: both handles address the same remote
|
// header. Sharing is correct: both handles address the same remote
|
||||||
// collection, and neither of them writes to it.
|
// collection, and neither of them writes to it.
|
||||||
|
//
|
||||||
|
// The point-id map is shared for the same reason, and because it MUST be:
|
||||||
|
// `Rag::clone()` hands the clone `DocumentId`s that the original minted,
|
||||||
|
// so a fresh map would resolve them to nothing and `fetch_content` would
|
||||||
|
// ask Qdrant for a synthetic handle — zero results, no error. Resetting it
|
||||||
|
// would also re-mint handles for ids the original still holds.
|
||||||
Box::new(Self {
|
Box::new(Self {
|
||||||
client: self.client.clone(),
|
client: self.client.clone(),
|
||||||
base_url: self.base_url.clone(),
|
base_url: self.base_url.clone(),
|
||||||
collection: self.collection.clone(),
|
collection: self.collection.clone(),
|
||||||
|
point_ids: Arc::clone(&self.point_ids),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -428,6 +535,7 @@ mod tests {
|
|||||||
client: Client::new(),
|
client: Client::new(),
|
||||||
base_url: "http://localhost:6333".to_string(),
|
base_url: "http://localhost:6333".to_string(),
|
||||||
collection: "c".to_string(),
|
collection: "c".to_string(),
|
||||||
|
point_ids: Arc::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
let attached = RagData {
|
let attached = RagData {
|
||||||
@@ -462,11 +570,154 @@ mod tests {
|
|||||||
client: Client::new(),
|
client: Client::new(),
|
||||||
base_url: "http://127.0.0.1:1".to_string(),
|
base_url: "http://127.0.0.1:1".to_string(),
|
||||||
collection: "c".to_string(),
|
collection: "c".to_string(),
|
||||||
|
point_ids: Arc::default(),
|
||||||
};
|
};
|
||||||
|
|
||||||
assert!(provider.fetch_content(&[]).await.unwrap().is_empty());
|
assert!(provider.fetch_content(&[]).await.unwrap().is_empty());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// A UUID-keyed collection has to survive the whole `vector_search` →
|
||||||
|
/// `fetch_content` round trip, and the fetch must ask Qdrant for the ORIGINAL
|
||||||
|
/// string id. Parsing ids with `as_u64()` used to drop these hits inside a
|
||||||
|
/// `filter_map`, i.e. zero results and no error.
|
||||||
|
#[test]
|
||||||
|
fn uuid_point_ids_round_trip_and_are_requested_verbatim() {
|
||||||
|
let mut interner = PointIdInterner::default();
|
||||||
|
let first_uuid = "3f1b0c2e-1111-4000-8000-000000000001";
|
||||||
|
let second_uuid = "3f1b0c2e-2222-4000-8000-000000000002";
|
||||||
|
|
||||||
|
let search = serde_json::json!({
|
||||||
|
"result": [
|
||||||
|
{"id": first_uuid, "score": 0.91},
|
||||||
|
{"id": second_uuid, "score": 0.42},
|
||||||
|
]
|
||||||
|
});
|
||||||
|
let hits = parse_search_hits(&mut interner, &search, 0.0).unwrap();
|
||||||
|
assert_eq!(hits.len(), 2, "string ids must not be silently dropped");
|
||||||
|
|
||||||
|
let ids: Vec<DocumentId> = hits.iter().map(|(id, _)| *id).collect();
|
||||||
|
assert_eq!(
|
||||||
|
interner.outbound_ids(&ids),
|
||||||
|
vec![Value::from(first_uuid), Value::from(second_uuid)],
|
||||||
|
"the fetch must send the ids Qdrant issued, not the handles"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Qdrant may answer /points in any order; the handles still map back and
|
||||||
|
// the caller's RRF ranking is recoverable.
|
||||||
|
let points = serde_json::json!({
|
||||||
|
"result": [
|
||||||
|
{"id": second_uuid, "payload": {"page_content": "second"}},
|
||||||
|
{"id": first_uuid, "payload": {"page_content": "first"}},
|
||||||
|
]
|
||||||
|
});
|
||||||
|
let mut rows = parse_points(&mut interner, &points).unwrap();
|
||||||
|
let position: HashMap<DocumentId, usize> =
|
||||||
|
ids.iter().enumerate().map(|(i, id)| (*id, i)).collect();
|
||||||
|
rows.sort_by_key(|(id, _)| position.get(id).copied().unwrap_or(usize::MAX));
|
||||||
|
assert_eq!(
|
||||||
|
rows,
|
||||||
|
vec![
|
||||||
|
(ids[0], "first".to_string()),
|
||||||
|
(ids[1], "second".to_string())
|
||||||
|
]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Integer-keyed collections must be untouched by the interner: the id maps to
|
||||||
|
/// itself on the way in and goes back out as the same integer.
|
||||||
|
#[test]
|
||||||
|
fn integer_point_ids_are_passed_through_untouched() {
|
||||||
|
let mut interner = PointIdInterner::default();
|
||||||
|
let search = serde_json::json!({
|
||||||
|
"result": [{"id": 7, "score": 0.9}, {"id": 0, "score": 0.5}]
|
||||||
|
});
|
||||||
|
|
||||||
|
let hits = parse_search_hits(&mut interner, &search, 0.0).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
hits,
|
||||||
|
vec![(DocumentId(7), 0.9_f32), (DocumentId(0), 0.5_f32)]
|
||||||
|
);
|
||||||
|
|
||||||
|
let ids: Vec<DocumentId> = hits.iter().map(|(id, _)| *id).collect();
|
||||||
|
assert_eq!(
|
||||||
|
interner.outbound_ids(&ids),
|
||||||
|
vec![Value::from(7_u64), Value::from(0_u64)],
|
||||||
|
"integer ids must not be regressed into synthetic handles"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
interner.raw_id(DocumentId(7)).is_none(),
|
||||||
|
"a plain integer id is its own id and needs no map entry"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Synthetic handles are stable per point id and live in a range no packed
|
||||||
|
/// `DocumentId` can reach.
|
||||||
|
#[test]
|
||||||
|
fn synthetic_handles_are_stable_and_never_collide_with_packed_ids() {
|
||||||
|
let mut interner = PointIdInterner::default();
|
||||||
|
let uuid = Value::from("9d2f0a11-3333-4000-8000-00000000000a");
|
||||||
|
|
||||||
|
let handle = interner.document_id(&uuid).unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
interner.document_id(&uuid).unwrap(),
|
||||||
|
handle,
|
||||||
|
"the same point id must keep the same handle across queries"
|
||||||
|
);
|
||||||
|
assert_ne!(
|
||||||
|
interner.document_id(&Value::from("other")).unwrap(),
|
||||||
|
handle,
|
||||||
|
"distinct point ids must not share a handle"
|
||||||
|
);
|
||||||
|
assert_ne!(handle.0 & SYNTHETIC_ID_TAG, 0, "a handle carries the tag");
|
||||||
|
|
||||||
|
// A packed (file_index, document_index) never sets the tag bit: it is the
|
||||||
|
// top bit of the file index, which would take 2^31 indexed files.
|
||||||
|
for (file_index, document_index) in [(0, 0), (1, 0), (0, 4242), (1_000_000, 999)] {
|
||||||
|
assert_eq!(
|
||||||
|
DocumentId::new(file_index, document_index).0 & SYNTHETIC_ID_TAG,
|
||||||
|
0,
|
||||||
|
"packed ({file_index}, {document_index}) must stay out of the handle range"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
// The one integer id that WOULD land on the tag is interned instead of
|
||||||
|
// being handed back as itself, so it cannot alias a handle.
|
||||||
|
let collides = Value::from(SYNTHETIC_ID_TAG as u64);
|
||||||
|
let interned = interner.document_id(&collides).unwrap();
|
||||||
|
assert_eq!(interner.raw_id(interned), Some(&collides));
|
||||||
|
assert_eq!(
|
||||||
|
interner.outbound_ids(&[interned]),
|
||||||
|
vec![collides],
|
||||||
|
"the original integer must still be what Qdrant is asked for"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `duplicate()` shares the map rather than resetting it: `Rag::clone()` hands
|
||||||
|
/// the clone `DocumentId`s the original minted, and a fresh map would turn
|
||||||
|
/// those into requests for a synthetic handle — zero results, no error.
|
||||||
|
#[test]
|
||||||
|
fn duplicate_shares_the_point_id_map() {
|
||||||
|
let provider = QdrantProvider {
|
||||||
|
client: Client::new(),
|
||||||
|
base_url: "http://127.0.0.1:1".to_string(),
|
||||||
|
collection: "c".to_string(),
|
||||||
|
point_ids: Arc::default(),
|
||||||
|
};
|
||||||
|
let uuid = Value::from("c0ffee00-4444-4000-8000-000000000007");
|
||||||
|
let handle = provider.point_ids.write().document_id(&uuid).unwrap();
|
||||||
|
|
||||||
|
let dup = provider.duplicate(&RagData {
|
||||||
|
driver: "qdrant".to_string(),
|
||||||
|
attached: true,
|
||||||
|
..Default::default()
|
||||||
|
});
|
||||||
|
// Downcasting is not available through `dyn RagProvider`, so go via the
|
||||||
|
// shared Arc: the clone must observe the original's interning.
|
||||||
|
assert_eq!(Arc::strong_count(&provider.point_ids), 2);
|
||||||
|
assert_eq!(provider.point_ids.read().raw_id(handle), Some(&uuid));
|
||||||
|
drop(dup);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
#[ignore]
|
#[ignore]
|
||||||
async fn qdrant_list_collections_requires_running_instance() {
|
async fn qdrant_list_collections_requires_running_instance() {
|
||||||
|
|||||||
+162
-38
@@ -67,13 +67,16 @@ pub fn discover() -> Result<Vec<DiscoveredMixin>> {
|
|||||||
push_if_exists(&mut out, paths::sbx_mixin_file())?;
|
push_if_exists(&mut out, paths::sbx_mixin_file())?;
|
||||||
push_if_exists(&mut out, paths::global_tools_sbx_mixin_file())?;
|
push_if_exists(&mut out, paths::global_tools_sbx_mixin_file())?;
|
||||||
|
|
||||||
for path in collect_subdir_mixins(&paths::functions_dir()) {
|
for path in collect_mixins(&paths::functions_dir(), &[ScanMode::SubdirNamed]) {
|
||||||
out.push(read_mixin(path)?);
|
out.push(read_mixin(path)?);
|
||||||
}
|
}
|
||||||
for path in collect_subdir_mixins(&paths::agents_data_dir()) {
|
for path in collect_mixins(
|
||||||
|
&paths::agents_data_dir(),
|
||||||
|
&[ScanMode::SubdirNamed, ScanMode::SubdirFlat],
|
||||||
|
) {
|
||||||
out.push(read_mixin(path)?);
|
out.push(read_mixin(path)?);
|
||||||
}
|
}
|
||||||
for path in collect_flat_mixins(&paths::rags_dir()) {
|
for path in collect_mixins(&paths::rags_dir(), &[ScanMode::Flat]) {
|
||||||
out.push(read_mixin(path)?);
|
out.push(read_mixin(path)?);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -160,27 +163,54 @@ fn read_mixin(path: PathBuf) -> Result<DiscoveredMixin> {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn collect_subdir_mixins(dir: &Path) -> Vec<PathBuf> {
|
/// One on-disk layout a mixin scan can look for. A scan takes a set of these,
|
||||||
|
/// and each mode contributes only the shape it names.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
enum ScanMode {
|
||||||
|
/// `<dir>/*.sbx-mixin.yaml`
|
||||||
|
Flat,
|
||||||
|
/// `<dir>/*/sbx-mixin.yaml`
|
||||||
|
SubdirNamed,
|
||||||
|
/// `<dir>/*/*.sbx-mixin.yaml`
|
||||||
|
SubdirFlat,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Collects mixin paths under `dir` for every requested layout. Missing or
|
||||||
|
/// unreadable directories yield nothing rather than an error — these paths are
|
||||||
|
/// all optional on disk.
|
||||||
|
///
|
||||||
|
/// Order is deterministic: flat matches first (sorted by file name), then each
|
||||||
|
/// subdirectory in sorted order, contributing its named mixin before its
|
||||||
|
/// suffixed ones.
|
||||||
|
fn collect_mixins(dir: &Path, modes: &[ScanMode]) -> Vec<PathBuf> {
|
||||||
let mut result = Vec::new();
|
let mut result = Vec::new();
|
||||||
let Ok(rd) = read_dir(dir) else { return result };
|
|
||||||
|
|
||||||
let mut entries: Vec<_> = rd
|
if modes.contains(&ScanMode::Flat) {
|
||||||
.flatten()
|
result.extend(suffixed_mixins_in(dir));
|
||||||
.filter(|e| e.file_type().map(|t| t.is_dir()).unwrap_or(false))
|
}
|
||||||
.collect();
|
|
||||||
entries.sort_by_key(|e| e.file_name());
|
|
||||||
|
|
||||||
for entry in entries {
|
let named = modes.contains(&ScanMode::SubdirNamed);
|
||||||
let candidate = entry.path().join(SBX_MIXIN_FILE_NAME);
|
let subdir_flat = modes.contains(&ScanMode::SubdirFlat);
|
||||||
|
if !named && !subdir_flat {
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
for subdir in subdirs_of(dir) {
|
||||||
|
if named {
|
||||||
|
let candidate = subdir.join(SBX_MIXIN_FILE_NAME);
|
||||||
if candidate.exists() {
|
if candidate.exists() {
|
||||||
result.push(candidate);
|
result.push(candidate);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if subdir_flat {
|
||||||
|
result.extend(suffixed_mixins_in(&subdir));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
result
|
result
|
||||||
}
|
}
|
||||||
|
|
||||||
fn collect_flat_mixins(dir: &Path) -> Vec<PathBuf> {
|
fn suffixed_mixins_in(dir: &Path) -> Vec<PathBuf> {
|
||||||
let mut result = Vec::new();
|
let mut result = Vec::new();
|
||||||
let Ok(rd) = read_dir(dir) else { return result };
|
let Ok(rd) = read_dir(dir) else { return result };
|
||||||
|
|
||||||
@@ -195,10 +225,21 @@ fn collect_flat_mixins(dir: &Path) -> Vec<PathBuf> {
|
|||||||
.collect();
|
.collect();
|
||||||
entries.sort_by_key(|e| e.file_name());
|
entries.sort_by_key(|e| e.file_name());
|
||||||
|
|
||||||
for entry in entries {
|
result.extend(entries.into_iter().map(|e| e.path()));
|
||||||
result.push(entry.path());
|
result
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn subdirs_of(dir: &Path) -> Vec<PathBuf> {
|
||||||
|
let mut result = Vec::new();
|
||||||
|
let Ok(rd) = read_dir(dir) else { return result };
|
||||||
|
|
||||||
|
let mut entries: Vec<_> = rd
|
||||||
|
.flatten()
|
||||||
|
.filter(|e| e.file_type().map(|t| t.is_dir()).unwrap_or(false))
|
||||||
|
.collect();
|
||||||
|
entries.sort_by_key(|e| e.file_name());
|
||||||
|
|
||||||
|
result.extend(entries.into_iter().map(|e| e.path()));
|
||||||
result
|
result
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -218,6 +259,13 @@ mod tests {
|
|||||||
root
|
root
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn file_names(paths: &[PathBuf]) -> Vec<&str> {
|
||||||
|
paths
|
||||||
|
.iter()
|
||||||
|
.map(|p| p.file_name().unwrap().to_str().unwrap())
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn summarize_counts_installs_and_domains() {
|
fn summarize_counts_installs_and_domains() {
|
||||||
let root = unique_root("sbx-mixin-counts");
|
let root = unique_root("sbx-mixin-counts");
|
||||||
@@ -301,7 +349,7 @@ network:
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn collect_subdir_mixins_sorts_and_skips_missing() {
|
fn subdir_named_scan_sorts_and_skips_missing() {
|
||||||
let root = unique_root("sbx-mixin-subdirs");
|
let root = unique_root("sbx-mixin-subdirs");
|
||||||
for name in ["zebra", "apple", "no-mixin", "mango"] {
|
for name in ["zebra", "apple", "no-mixin", "mango"] {
|
||||||
let dir = root.join(name);
|
let dir = root.join(name);
|
||||||
@@ -311,7 +359,7 @@ network:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let found = collect_subdir_mixins(&root);
|
let found = collect_mixins(&root, &[ScanMode::SubdirNamed]);
|
||||||
let names: Vec<String> = found
|
let names: Vec<String> = found
|
||||||
.iter()
|
.iter()
|
||||||
.map(|p| {
|
.map(|p| {
|
||||||
@@ -329,9 +377,9 @@ network:
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn collect_subdir_mixins_returns_empty_for_missing_dir() {
|
fn subdir_named_scan_returns_empty_for_missing_dir() {
|
||||||
let absent = env::temp_dir().join("coyote-definitely-not-here-xyz");
|
let absent = env::temp_dir().join("coyote-definitely-not-here-xyz");
|
||||||
let found = collect_subdir_mixins(&absent);
|
let found = collect_mixins(&absent, &[ScanMode::SubdirNamed]);
|
||||||
assert!(found.is_empty());
|
assert!(found.is_empty());
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -511,7 +559,7 @@ network:
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn collect_flat_mixins_matches_rag_sidecars_by_suffix() {
|
fn flat_scan_matches_rag_sidecars_by_suffix() {
|
||||||
let root = unique_root("flat-mixins");
|
let root = unique_root("flat-mixins");
|
||||||
fs::write(root.join("company-docs.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
|
fs::write(root.join("company-docs.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
|
||||||
fs::write(root.join("alpha.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
|
fs::write(root.join("alpha.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
|
||||||
@@ -519,39 +567,115 @@ network:
|
|||||||
fs::write(root.join("notes.yaml"), "driver: yaml\n").unwrap();
|
fs::write(root.join("notes.yaml"), "driver: yaml\n").unwrap();
|
||||||
fs::create_dir_all(root.join("decoy.sbx-mixin.yaml")).unwrap();
|
fs::create_dir_all(root.join("decoy.sbx-mixin.yaml")).unwrap();
|
||||||
|
|
||||||
let found = collect_flat_mixins(&root);
|
let found = collect_mixins(&root, &[ScanMode::Flat]);
|
||||||
let names: Vec<_> = found
|
|
||||||
.iter()
|
|
||||||
.map(|p| p.file_name().unwrap().to_str().unwrap())
|
|
||||||
.collect();
|
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
names,
|
file_names(&found),
|
||||||
vec!["alpha.sbx-mixin.yaml", "company-docs.sbx-mixin.yaml"]
|
vec!["alpha.sbx-mixin.yaml", "company-docs.sbx-mixin.yaml"]
|
||||||
);
|
);
|
||||||
|
|
||||||
let _ = fs::remove_dir_all(&root);
|
let _ = fs::remove_dir_all(&root);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Why `collect_flat_mixins` had to be written: the existing collector walks
|
/// Every scan site in `discover()` picks its modes assuming each mode owns
|
||||||
/// SUBDIRECTORIES for a file named exactly `sbx-mixin.yaml`, so it cannot see
|
/// exactly one layout and nothing else. `agents_data_dir()` requests two
|
||||||
/// a flat sidecar. If this ever starts finding them, the new collector is
|
/// modes at once, so an overlap would collect the same file twice and
|
||||||
/// redundant — but until then, removing it silently drops every RAG mixin.
|
/// `create_sandbox` would pass it as two `--kit` flags.
|
||||||
#[test]
|
#[test]
|
||||||
fn collect_subdir_mixins_cannot_see_flat_rag_sidecars() {
|
fn each_scan_mode_owns_exactly_one_layout() {
|
||||||
let root = unique_root("flat-vs-subdir");
|
let root = unique_root("scan-mode-ownership");
|
||||||
fs::write(root.join("company-docs.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
|
let agent = root.join("researcher");
|
||||||
|
fs::create_dir_all(&agent).unwrap();
|
||||||
|
let flat = root.join("company-docs.sbx-mixin.yaml");
|
||||||
|
let subdir_named = agent.join("sbx-mixin.yaml");
|
||||||
|
let subdir_flat = agent.join("handbook.sbx-mixin.yaml");
|
||||||
|
for path in [&flat, &subdir_named, &subdir_flat] {
|
||||||
|
fs::write(path, "kind: mixin\n").unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
assert!(collect_subdir_mixins(&root).is_empty());
|
assert_eq!(collect_mixins(&root, &[ScanMode::Flat]), vec![flat.clone()]);
|
||||||
assert_eq!(collect_flat_mixins(&root).len(), 1);
|
assert_eq!(
|
||||||
|
collect_mixins(&root, &[ScanMode::SubdirNamed]),
|
||||||
|
vec![subdir_named.clone()]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
collect_mixins(&root, &[ScanMode::SubdirFlat]),
|
||||||
|
vec![subdir_flat.clone()]
|
||||||
|
);
|
||||||
|
|
||||||
|
let all = collect_mixins(
|
||||||
|
&root,
|
||||||
|
&[ScanMode::Flat, ScanMode::SubdirNamed, ScanMode::SubdirFlat],
|
||||||
|
);
|
||||||
|
assert_eq!(all, vec![flat, subdir_named, subdir_flat]);
|
||||||
|
|
||||||
|
let mut deduped = all.clone();
|
||||||
|
deduped.sort();
|
||||||
|
deduped.dedup();
|
||||||
|
assert_eq!(
|
||||||
|
deduped.len(),
|
||||||
|
all.len(),
|
||||||
|
"no mixin may be collected twice: {all:?}"
|
||||||
|
);
|
||||||
|
|
||||||
let _ = fs::remove_dir_all(&root);
|
let _ = fs::remove_dir_all(&root);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn collect_flat_mixins_tolerates_a_missing_directory() {
|
fn flat_scan_tolerates_a_missing_directory() {
|
||||||
let root = unique_root("flat-missing");
|
let root = unique_root("flat-missing");
|
||||||
let absent = root.join("nope");
|
let absent = root.join("nope");
|
||||||
assert!(collect_flat_mixins(&absent).is_empty());
|
assert!(collect_mixins(&absent, &[ScanMode::Flat]).is_empty());
|
||||||
|
|
||||||
|
let _ = fs::remove_dir_all(&root);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `generate_rag_sbx_mixin` writes an agent-scoped RAG sidecar next to the
|
||||||
|
/// rag yaml, at `<agents>/<agent>/<rag>.sbx-mixin.yaml`. Before `SubdirFlat`
|
||||||
|
/// existed, nothing scanned that shape and attaching a Qdrant RAG from
|
||||||
|
/// inside an agent produced no network allow rule and no credential.
|
||||||
|
#[test]
|
||||||
|
fn agent_scoped_rag_sidecar_is_discovered() {
|
||||||
|
let root = unique_root("agent-scoped-rag");
|
||||||
|
let agent = root.join("researcher");
|
||||||
|
fs::create_dir_all(&agent).unwrap();
|
||||||
|
fs::write(agent.join("company-docs.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
|
||||||
|
fs::write(agent.join("company-docs.yaml"), "driver: qdrant\n").unwrap();
|
||||||
|
|
||||||
|
let found = collect_mixins(&root, &[ScanMode::SubdirNamed, ScanMode::SubdirFlat]);
|
||||||
|
assert_eq!(found, vec![agent.join("company-docs.sbx-mixin.yaml")]);
|
||||||
|
|
||||||
|
let _ = fs::remove_dir_all(&root);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn agent_level_mixin_and_rag_sidecars_are_both_discovered() {
|
||||||
|
let root = unique_root("agent-both-shapes");
|
||||||
|
let agent = root.join("researcher");
|
||||||
|
fs::create_dir_all(&agent).unwrap();
|
||||||
|
fs::write(agent.join("sbx-mixin.yaml"), "kind: mixin\n").unwrap();
|
||||||
|
fs::write(agent.join("zebra.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
|
||||||
|
fs::write(agent.join("alpha.sbx-mixin.yaml"), "kind: mixin\n").unwrap();
|
||||||
|
|
||||||
|
let found = collect_mixins(&root, &[ScanMode::SubdirNamed, ScanMode::SubdirFlat]);
|
||||||
|
assert_eq!(
|
||||||
|
file_names(&found),
|
||||||
|
vec![
|
||||||
|
"sbx-mixin.yaml",
|
||||||
|
"alpha.sbx-mixin.yaml",
|
||||||
|
"zebra.sbx-mixin.yaml"
|
||||||
|
]
|
||||||
|
);
|
||||||
|
|
||||||
|
let _ = fs::remove_dir_all(&root);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn subdir_flat_scan_ignores_a_directory_named_like_a_mixin() {
|
||||||
|
let root = unique_root("subdir-flat-decoy");
|
||||||
|
let agent = root.join("researcher");
|
||||||
|
fs::create_dir_all(agent.join("decoy.sbx-mixin.yaml")).unwrap();
|
||||||
|
|
||||||
|
assert!(collect_mixins(&root, &[ScanMode::SubdirFlat]).is_empty());
|
||||||
|
|
||||||
let _ = fs::remove_dir_all(&root);
|
let _ = fs::remove_dir_all(&root);
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user