Merge pull request #14 from Dark-Alex-17/feat/rag-driver-abstraction-v3
feat(rag): pluggable RagProvider abstraction with DuckDB and Qdrant drivers
This commit is contained in:
@@ -36,6 +36,17 @@ jobs:
|
||||
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
|
||||
- name: Cache DuckDB Extensions
|
||||
id: duckdb-extensions
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.duckdb/extensions
|
||||
key: duckdb-ext-${{ matrix.os }}-${{ hashFiles('Cargo.lock') }}
|
||||
|
||||
- name: Install DuckDB Extensions
|
||||
if: steps.duckdb-extensions.outputs.cache-hit != 'true'
|
||||
run: cargo test --all duckdb
|
||||
|
||||
- name: Test
|
||||
run: cargo test --all
|
||||
|
||||
|
||||
@@ -5,3 +5,4 @@
|
||||
.idea/
|
||||
/coyote.iml
|
||||
/.idea/
|
||||
.coyote
|
||||
|
||||
Generated
+566
-141
File diff suppressed because it is too large
Load Diff
+2
-1
@@ -17,8 +17,9 @@ exclude = [".github", "CONTRIBUTING.md"]
|
||||
anyhow = "1.0.69"
|
||||
bytes = "1.4.0"
|
||||
clap = { version = "4.5.40", features = ["cargo", "derive", "wrap_help"] }
|
||||
comfy-table = { version = "7.2.2", features = ["custom_styling"] }
|
||||
comfy-table = { version = "7.1.4", features = ["custom_styling"] }
|
||||
dirs = "6.0.0"
|
||||
duckdb = { version = "1.10505.0", features = ["bundled"] }
|
||||
dunce = "1.0.5"
|
||||
futures-util = "0.3.29"
|
||||
inquire = "0.9.4"
|
||||
|
||||
+96
-14
@@ -4,6 +4,7 @@ use crate::{
|
||||
client::Model,
|
||||
config::memory,
|
||||
function::{Functions, run_llm_function},
|
||||
graph, rag,
|
||||
};
|
||||
|
||||
use super::rag_cache::RagKey;
|
||||
@@ -12,6 +13,7 @@ use crate::config::prompts::{
|
||||
DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS,
|
||||
DEFAULT_USER_INTERACTION_INSTRUCTIONS,
|
||||
};
|
||||
use crate::graph::types::RagNode;
|
||||
use crate::graph::{Graph, GraphParser, NodeType};
|
||||
use crate::rag::RagInitConfig;
|
||||
use crate::vault::SECRET_RE;
|
||||
@@ -146,11 +148,18 @@ impl Agent {
|
||||
let rag = if rag_path.exists() {
|
||||
let key = RagKey::Agent(name.to_string());
|
||||
let app_clone = app.clone();
|
||||
let vault_clone = app_state.vault.clone();
|
||||
let rag_path_clone = rag_path.clone();
|
||||
let rag = app_state
|
||||
.rag_cache
|
||||
.load_with(key, || async move {
|
||||
Rag::load(&app_clone, DEFAULT_AGENT_NAME, &rag_path_clone)
|
||||
Rag::load_async(
|
||||
&app_clone,
|
||||
&vault_clone,
|
||||
DEFAULT_AGENT_NAME,
|
||||
&rag_path_clone,
|
||||
)
|
||||
.await
|
||||
})
|
||||
.await?;
|
||||
Some(rag)
|
||||
@@ -171,7 +180,15 @@ impl Agent {
|
||||
let rag = app_state
|
||||
.rag_cache
|
||||
.load_with(key, || async move {
|
||||
Rag::init(&app_clone, "rag", &rag_path_clone, &document_paths, abort).await
|
||||
Rag::init(
|
||||
&app_clone,
|
||||
"rag",
|
||||
&rag_path_clone,
|
||||
&document_paths,
|
||||
abort,
|
||||
true,
|
||||
)
|
||||
.await
|
||||
})
|
||||
.await?;
|
||||
Some(rag)
|
||||
@@ -937,6 +954,30 @@ fn resolve_document_paths(
|
||||
Ok(document_paths)
|
||||
}
|
||||
|
||||
/// How a graph rag node describes the knowledge base it wants built.
|
||||
///
|
||||
/// `driver` is forwarded as-is: `None` means the node did not ask for one, which
|
||||
/// `RagInitConfig` resolves to yaml, so workflows written before drivers existed
|
||||
/// keep their current storage.
|
||||
///
|
||||
/// Every field is now named explicitly, so adding one to `RagInitConfig` breaks
|
||||
/// this literal. That is deliberate: the new field then gets a decision about
|
||||
/// whether a rag node can drive it, instead of silently taking its default.
|
||||
fn rag_init_config(rag_node: &RagNode) -> RagInitConfig {
|
||||
RagInitConfig {
|
||||
embedding_model: rag_node.embedding_model.clone(),
|
||||
chunk_size: rag_node.chunk_size,
|
||||
chunk_overlap: rag_node.chunk_overlap,
|
||||
reranker_model: rag_node.reranker_model.clone(),
|
||||
top_k: rag_node.top_k,
|
||||
batch_size: rag_node.batch_size,
|
||||
extractor_model: rag_node.extractor_model.clone(),
|
||||
extractor_prompt: rag_node.extractor_prompt.clone(),
|
||||
graph_hops: rag_node.graph_hops,
|
||||
driver: rag_node.driver.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn init_graph_rags(
|
||||
app: &AppConfig,
|
||||
@@ -964,26 +1005,28 @@ async fn init_graph_rags(
|
||||
};
|
||||
let rag = if rag_path.exists() {
|
||||
let app_clone = app.clone();
|
||||
let vault_clone = app_state.vault.clone();
|
||||
let path_clone = rag_path.clone();
|
||||
let name_clone = node_id.clone();
|
||||
app_state
|
||||
.rag_cache
|
||||
.load_with(key, || async move {
|
||||
Rag::load(&app_clone, &name_clone, &path_clone)
|
||||
Rag::load_async(&app_clone, &vault_clone, &name_clone, &path_clone).await
|
||||
})
|
||||
.await?
|
||||
} else {
|
||||
let config = RagInitConfig {
|
||||
embedding_model: rag_node.embedding_model.clone(),
|
||||
chunk_size: rag_node.chunk_size,
|
||||
chunk_overlap: rag_node.chunk_overlap,
|
||||
reranker_model: rag_node.reranker_model.clone(),
|
||||
top_k: rag_node.top_k,
|
||||
batch_size: rag_node.batch_size,
|
||||
extractor_model: rag_node.extractor_model.clone(),
|
||||
extractor_prompt: rag_node.extractor_prompt.clone(),
|
||||
graph_hops: rag_node.graph_hops,
|
||||
};
|
||||
// Checked before anything is built: an unknown driver would otherwise
|
||||
// fall through `Rag::create`'s catch-all to a yaml store, embed every
|
||||
// document, and persist the bogus driver string. The RAG would then be
|
||||
// rejected on every subsequent load, leaving the agent unstartable.
|
||||
// Graph validation catches this too, but it is skipped when
|
||||
// `validate_before_run` is off, so this guard is the load-bearing one.
|
||||
if let Some(driver) = &rag_node.driver
|
||||
&& let Some(message) = graph::validator::rag_driver_error(driver)
|
||||
{
|
||||
bail!("rag node '{node_id}': {message}");
|
||||
}
|
||||
let mut config = rag_init_config(rag_node);
|
||||
let fully_specified = config.embedding_model.is_some()
|
||||
&& config.chunk_size.is_some()
|
||||
&& config.chunk_overlap.is_some();
|
||||
@@ -1009,6 +1052,10 @@ async fn init_graph_rags(
|
||||
initialized. RAG initialization is required for this agent."
|
||||
);
|
||||
}
|
||||
|
||||
if config.driver.is_none() {
|
||||
config.driver = Some(rag::select_rag_driver()?);
|
||||
}
|
||||
}
|
||||
|
||||
let document_paths =
|
||||
@@ -1317,4 +1364,39 @@ version: "1.0"
|
||||
|
||||
assert_eq!(meta.description, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rag_init_config_forwards_an_explicit_driver() {
|
||||
let node: RagNode =
|
||||
serde_yaml::from_str("documents: [\"./docs\"]\ndriver: duckdb\n").unwrap();
|
||||
|
||||
assert_eq!(rag_init_config(&node).driver.as_deref(), Some("duckdb"));
|
||||
}
|
||||
|
||||
/// A node that names no driver must forward `None`, which `RagInitConfig`
|
||||
/// documents as "yaml". Existing workflows therefore keep their yaml store.
|
||||
#[test]
|
||||
fn rag_init_config_leaves_the_driver_unset_by_default() {
|
||||
let node: RagNode = serde_yaml::from_str("documents: [\"./docs\"]\n").unwrap();
|
||||
|
||||
assert_eq!(rag_init_config(&node).driver, None);
|
||||
}
|
||||
|
||||
/// The driver must ride alongside the rest of the node's settings, not
|
||||
/// replace them.
|
||||
#[test]
|
||||
fn rag_init_config_forwards_the_other_settings_too() {
|
||||
let node: RagNode = serde_yaml::from_str(
|
||||
"documents: [\"./docs\"]\ndriver: duckdb\nchunk_size: 512\nchunk_overlap: 64\ntop_k: 7\nembedding_model: some:model\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let config = rag_init_config(&node);
|
||||
|
||||
assert_eq!(config.driver.as_deref(), Some("duckdb"));
|
||||
assert_eq!(config.chunk_size, Some(512));
|
||||
assert_eq!(config.chunk_overlap, Some(64));
|
||||
assert_eq!(config.top_k, Some(7));
|
||||
assert_eq!(config.embedding_model.as_deref(), Some("some:model"));
|
||||
}
|
||||
}
|
||||
|
||||
+101
-1
@@ -16,7 +16,7 @@ use anyhow::{Context, Result, anyhow, bail};
|
||||
use log::LevelFilter;
|
||||
use std::collections::HashSet;
|
||||
use std::env;
|
||||
use std::fs::{read_dir, read_to_string};
|
||||
use std::fs::{read_dir, read_to_string, remove_file};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
pub fn config_dir() -> PathBuf {
|
||||
@@ -414,6 +414,10 @@ pub fn list_rags() -> Vec<String> {
|
||||
for entry in rd.flatten() {
|
||||
let name = entry.file_name();
|
||||
if let Some(name) = name.to_string_lossy().strip_suffix(".yaml") {
|
||||
if is_rag_sidecar_name(name) {
|
||||
continue;
|
||||
}
|
||||
|
||||
names.push(name.to_string());
|
||||
}
|
||||
}
|
||||
@@ -424,6 +428,34 @@ pub fn list_rags() -> Vec<String> {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_rag_sidecar_name(name: &str) -> bool {
|
||||
name.ends_with(".sbx-mixin")
|
||||
}
|
||||
|
||||
pub(crate) fn remove_rag_sidecars(dir: &Path, name: &str) -> Result<()> {
|
||||
let duckdb_path = dir.join(format!("{name}.duckdb"));
|
||||
if duckdb_path.exists() {
|
||||
let _ = remove_file(&duckdb_path);
|
||||
}
|
||||
let wal_path = dir.join(format!("{name}.duckdb.wal"));
|
||||
if wal_path.exists() {
|
||||
let _ = remove_file(&wal_path);
|
||||
}
|
||||
let mixin_path = dir.join(format!("{name}.sbx-mixin.yaml"));
|
||||
if mixin_path.exists() {
|
||||
remove_file(&mixin_path).with_context(|| {
|
||||
format!(
|
||||
"Failed to remove the sandbox mixin for RAG '{name}' at '{}'. \
|
||||
The RAG was NOT deleted so you can retry; this host remains \
|
||||
whitelisted in the sandbox until the file is removed.",
|
||||
mixin_path.display()
|
||||
)
|
||||
})?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn list_macros() -> Vec<String> {
|
||||
list_file_names(macros_dir(), ".yaml")
|
||||
}
|
||||
@@ -846,4 +878,72 @@ mod tests {
|
||||
}
|
||||
let _ = fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
fn sidecar_temp_dir(label: &str) -> PathBuf {
|
||||
let unique = time::SystemTime::now()
|
||||
.duration_since(time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
let root = env::temp_dir().join(format!("coyote-{label}-test-{unique}"));
|
||||
fs::create_dir_all(&root).unwrap();
|
||||
root
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_rag_sidecar_name_accepts_dotted_rag_names() {
|
||||
assert!(!is_rag_sidecar_name("v2.docs"));
|
||||
assert!(!is_rag_sidecar_name("myrag"));
|
||||
assert!(is_rag_sidecar_name("myrag.sbx-mixin"));
|
||||
assert!(is_rag_sidecar_name("v2.docs.sbx-mixin"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_rag_sidecars_removes_duckdb_wal_and_mixin() {
|
||||
let root = sidecar_temp_dir("rag-sidecars-both");
|
||||
let duckdb = root.join("docs.duckdb");
|
||||
let wal = root.join("docs.duckdb.wal");
|
||||
let mixin = root.join("docs.sbx-mixin.yaml");
|
||||
fs::write(&duckdb, "db").unwrap();
|
||||
fs::write(&wal, "wal").unwrap();
|
||||
fs::write(&mixin, "mixin").unwrap();
|
||||
|
||||
remove_rag_sidecars(&root, "docs").unwrap();
|
||||
|
||||
assert!(!duckdb.exists(), "the .duckdb sidecar must be removed");
|
||||
assert!(!wal.exists(), "the .duckdb.wal sidecar must be removed");
|
||||
assert!(
|
||||
!mixin.exists(),
|
||||
"the .sbx-mixin.yaml sidecar must be removed"
|
||||
);
|
||||
let _ = fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_rag_sidecars_is_ok_when_absent() {
|
||||
let root = sidecar_temp_dir("rag-sidecars-absent");
|
||||
assert!(remove_rag_sidecars(&root, "docs").is_ok());
|
||||
let _ = fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_rag_sidecars_runs_before_yaml_unlink() {
|
||||
let root = sidecar_temp_dir("rag-sidecars-order");
|
||||
let yaml = root.join("docs.yaml");
|
||||
fs::write(&yaml, "rag").unwrap();
|
||||
let mixin = root.join("docs.sbx-mixin.yaml");
|
||||
fs::create_dir_all(&mixin).unwrap();
|
||||
fs::write(mixin.join("blocker"), "x").unwrap();
|
||||
|
||||
let err = remove_rag_sidecars(&root, "docs").unwrap_err();
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("Failed to remove the sandbox mixin"),
|
||||
"got: {err}"
|
||||
);
|
||||
assert!(
|
||||
yaml.exists(),
|
||||
"the .yaml must survive a sidecar-removal failure so the delete is retryable"
|
||||
);
|
||||
let _ = fs::remove_dir_all(&root);
|
||||
}
|
||||
}
|
||||
|
||||
+175
-20
@@ -142,6 +142,7 @@ pub struct RequestContext {
|
||||
pub role: Option<Role>,
|
||||
pub session: Option<Session>,
|
||||
pub rag: Option<Arc<Rag>>,
|
||||
pub rag_key: Option<RagKey>,
|
||||
pub agent: Option<Agent>,
|
||||
|
||||
pub last_message: Option<LastMessage>,
|
||||
@@ -176,6 +177,7 @@ impl RequestContext {
|
||||
role: None,
|
||||
session: None,
|
||||
rag: None,
|
||||
rag_key: None,
|
||||
agent: None,
|
||||
last_message: None,
|
||||
tool_scope: ToolScope::default(),
|
||||
@@ -229,6 +231,7 @@ impl RequestContext {
|
||||
role: None,
|
||||
session: None,
|
||||
rag: None,
|
||||
rag_key: None,
|
||||
agent: None,
|
||||
last_message: None,
|
||||
tool_scope: ToolScope {
|
||||
@@ -277,6 +280,7 @@ impl RequestContext {
|
||||
role: self.role.clone(),
|
||||
session: self.session.clone(),
|
||||
rag: self.rag.clone(),
|
||||
rag_key: self.rag_key.clone(),
|
||||
agent: self.agent.clone(),
|
||||
last_message: self.last_message.clone(),
|
||||
tool_scope: self.tool_scope.clone(),
|
||||
@@ -315,6 +319,7 @@ impl RequestContext {
|
||||
role: None,
|
||||
session: None,
|
||||
rag: None,
|
||||
rag_key: None,
|
||||
agent: None,
|
||||
last_message: None,
|
||||
tool_scope: ToolScope {
|
||||
@@ -2554,6 +2559,14 @@ impl RequestContext {
|
||||
match file_ext {
|
||||
Some(file_ext) => {
|
||||
if let Some(name) = name.to_string_lossy().strip_suffix(file_ext) {
|
||||
// Sidecars are not independently deletable assets.
|
||||
// Guarded on `kind == "rag"` because this scan is shared
|
||||
// by all six kinds, and `session`/`macro` also use
|
||||
// `.yaml`. The helper lives in paths.rs beside
|
||||
// list_rags() so both filters cannot drift apart.
|
||||
if kind == "rag" && paths::is_rag_sidecar_name(name) {
|
||||
continue;
|
||||
}
|
||||
names.push(name.to_string());
|
||||
}
|
||||
}
|
||||
@@ -2590,6 +2603,13 @@ impl RequestContext {
|
||||
match file_ext {
|
||||
Some(ext) => {
|
||||
let path = dir.join(format!("{name}{ext}"));
|
||||
// Sidecars FIRST. If this fails, the .yaml is still on disk, the
|
||||
// RAG is still listed, and the user can retry. Unlinking the .yaml
|
||||
// first would make the deletion unretryable while leaving an
|
||||
// orphaned mixin whitelisting a host in every sandbox launch.
|
||||
if kind == "rag" {
|
||||
paths::remove_rag_sidecars(&dir, &name)?;
|
||||
}
|
||||
remove_file(&path).with_context(|| {
|
||||
format!("Failed to delete {kind} at '{}'", path.display())
|
||||
})?;
|
||||
@@ -2773,7 +2793,12 @@ impl RequestContext {
|
||||
}
|
||||
}
|
||||
"rag_top_k" => {
|
||||
let value = value.parse().with_context(|| "Invalid value")?;
|
||||
let value: usize = value.parse().with_context(|| "Invalid value")?;
|
||||
if value == 0 {
|
||||
bail!(
|
||||
"rag_top_k must be >= 1; a top_k of 0 makes every query return no results."
|
||||
);
|
||||
}
|
||||
if !self.set_rag_top_k(value)? {
|
||||
self.update_app_config(|app| app.rag_top_k = value);
|
||||
}
|
||||
@@ -3724,6 +3749,14 @@ impl RequestContext {
|
||||
.then(|| Arc::new(RwLock::new(Supervisor::new(max_concurrent, max_depth))));
|
||||
|
||||
self.rag = agent.rag();
|
||||
// Keep `rag_key` in lockstep with `rag`. Agent RAGs are cached under
|
||||
// `RagKey::Agent(<agent name>)` (see `Agent::init`), so mirror that key exactly;
|
||||
// leaving the previous key in place would let `.rebuild rag` invalidate an
|
||||
// unrelated RAG's cache entry, and leaving it `None` would invalidate nothing.
|
||||
self.rag_key = self
|
||||
.rag
|
||||
.is_some()
|
||||
.then(|| RagKey::Agent(agent.name().to_string()));
|
||||
self.agent = Some(agent);
|
||||
self.supervisor = supervisor;
|
||||
self.inbox = None;
|
||||
@@ -3772,6 +3805,11 @@ impl RequestContext {
|
||||
self.pending_agents_guardrail_count = 0;
|
||||
self.todo_list = TodoList::default();
|
||||
self.rag.take();
|
||||
// Cleared alongside `rag` so the pair never disagrees: an agent RAG is
|
||||
// cached under `RagKey::Agent(<agent name>)`, and leaving that key behind
|
||||
// would outlive the RAG it names. Latent rather than live today only
|
||||
// because `rebuild_rag`/`edit_rag_docs` bail on `rag.is_none()` first.
|
||||
self.rag_key = None;
|
||||
self.discontinuous_last_message();
|
||||
}
|
||||
Ok(())
|
||||
@@ -4079,10 +4117,11 @@ impl RequestContext {
|
||||
}
|
||||
|
||||
let app = self.app.config.clone();
|
||||
let vault = self.app.vault.clone();
|
||||
let rag_cache = self.rag_cache();
|
||||
let working_mode = self.working_mode;
|
||||
|
||||
let rag: Arc<Rag> = match rag {
|
||||
let (rag, rag_key): (Arc<Rag>, Option<RagKey>) = match rag {
|
||||
None => {
|
||||
let rag_path = self.rag_file(super::TEMP_RAG_NAME);
|
||||
if rag_path.exists() {
|
||||
@@ -4090,15 +4129,29 @@ impl RequestContext {
|
||||
format!("Failed to cleanup previous '{}' rag", super::TEMP_RAG_NAME)
|
||||
})?;
|
||||
}
|
||||
Arc::new(Rag::init(&app, super::TEMP_RAG_NAME, &rag_path, &[], abort_signal).await?)
|
||||
(
|
||||
Arc::new(
|
||||
Rag::init(
|
||||
&app,
|
||||
super::TEMP_RAG_NAME,
|
||||
&rag_path,
|
||||
&[],
|
||||
abort_signal,
|
||||
false,
|
||||
)
|
||||
.await?,
|
||||
),
|
||||
None,
|
||||
)
|
||||
}
|
||||
Some(name) => {
|
||||
let rag_path = self.rag_file(name);
|
||||
let key = RagKey::Named(name.to_string());
|
||||
|
||||
rag_cache
|
||||
.load_with(key, || {
|
||||
let loaded = rag_cache
|
||||
.load_with(key.clone(), || {
|
||||
let app = app.clone();
|
||||
let vault = vault.clone();
|
||||
let rag_path = rag_path.clone();
|
||||
let abort_signal = abort_signal.clone();
|
||||
async move {
|
||||
@@ -4106,16 +4159,39 @@ impl RequestContext {
|
||||
if working_mode.is_cmd() {
|
||||
bail!("Unknown RAG '{name}'");
|
||||
}
|
||||
Rag::init(&app, name, &rag_path, &[], abort_signal.clone()).await
|
||||
Rag::init(&app, name, &rag_path, &[], abort_signal.clone(), true)
|
||||
.await
|
||||
} else {
|
||||
Rag::load(&app, name, &rag_path)
|
||||
Rag::load_async(&app, &vault, name, &rag_path).await
|
||||
}
|
||||
}
|
||||
})
|
||||
.await?
|
||||
.await?;
|
||||
(loaded, Some(key))
|
||||
}
|
||||
};
|
||||
self.rag = Some(rag);
|
||||
self.rag_key = rag_key;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn attach_rag(&mut self, name: &str) -> Result<()> {
|
||||
let rag_path = self.rag_file(name);
|
||||
if rag_path.exists() {
|
||||
bail!(
|
||||
"RAG '{name}' already exists at '{}'. \
|
||||
Use a different name, or delete the existing file first.",
|
||||
rag_path.display()
|
||||
);
|
||||
}
|
||||
let app = self.app.config.as_ref();
|
||||
let vault = self.app.vault.clone();
|
||||
let rag = Rag::attach(app, &vault, name, &rag_path).await?;
|
||||
let rag = Arc::new(rag);
|
||||
let key = RagKey::Named(name.to_string());
|
||||
self.rag_cache().insert(key.clone(), &rag);
|
||||
self.rag = Some(rag);
|
||||
self.rag_key = Some(key);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -4125,6 +4201,12 @@ impl RequestContext {
|
||||
None => bail!("No RAG"),
|
||||
};
|
||||
|
||||
if rag.is_attached() {
|
||||
bail!(
|
||||
"Cannot edit documents on an attached RAG; Coyote does not own its source documents."
|
||||
);
|
||||
}
|
||||
|
||||
let document_paths = rag.document_paths();
|
||||
let temp_file = temp_file(&format!("-rag-{}", rag.name()), ".txt");
|
||||
tokio::fs::write(&temp_file, &document_paths.join("\n"))
|
||||
@@ -4150,14 +4232,17 @@ impl RequestContext {
|
||||
bail!("No changes")
|
||||
}
|
||||
|
||||
let key = if self.agent.is_some() {
|
||||
RagKey::Agent(rag.name().to_string())
|
||||
} else {
|
||||
RagKey::Named(rag.name().to_string())
|
||||
};
|
||||
if let Some(key) = self.rag_key.clone() {
|
||||
self.rag_cache().invalidate(&key);
|
||||
}
|
||||
|
||||
rag.refresh_document_paths(&new_document_paths, false, &self.app.config, abort_signal)
|
||||
rag.refresh_document_paths(
|
||||
&new_document_paths,
|
||||
false,
|
||||
false,
|
||||
&self.app.config,
|
||||
abort_signal,
|
||||
)
|
||||
.await?;
|
||||
self.rag = Some(Arc::new(rag));
|
||||
Ok(())
|
||||
@@ -4169,15 +4254,25 @@ impl RequestContext {
|
||||
None => bail!("No RAG"),
|
||||
};
|
||||
|
||||
let key = if self.agent.is_some() {
|
||||
RagKey::Agent(rag.name().to_string())
|
||||
} else {
|
||||
RagKey::Named(rag.name().to_string())
|
||||
};
|
||||
if rag.is_attached() {
|
||||
bail!(
|
||||
"Cannot rebuild an attached RAG; Coyote does not own its source documents. \
|
||||
Re-index from the system that originally created '{}'.",
|
||||
rag.name()
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(key) = self.rag_key.clone() {
|
||||
self.rag_cache().invalidate(&key);
|
||||
}
|
||||
|
||||
let document_paths = rag.document_paths().to_vec();
|
||||
rag.refresh_document_paths(&document_paths, true, &self.app.config, abort_signal)
|
||||
println!(
|
||||
"Rebuilding re-embeds every document ({} files). \
|
||||
This will call the embedding API and may take a while.",
|
||||
rag.file_count()
|
||||
);
|
||||
rag.refresh_document_paths(&document_paths, true, true, &self.app.config, abort_signal)
|
||||
.await?;
|
||||
self.rag = Some(Arc::new(rag));
|
||||
Ok(())
|
||||
@@ -4583,6 +4678,44 @@ mod tests {
|
||||
|
||||
assert!(ctx.agent.is_none());
|
||||
assert!(ctx.rag.is_none());
|
||||
assert_eq!(ctx.rag_key, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn use_agent_does_not_carry_stale_rag_key() {
|
||||
let _guard = TestConfigDirGuard::new();
|
||||
let mut ctx = create_test_ctx();
|
||||
let app = ctx.app.config.clone();
|
||||
let agent_name = format!(
|
||||
"test_agent_{}",
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos()
|
||||
);
|
||||
let agent_dir = paths::agent_data_dir(&agent_name);
|
||||
create_dir_all(&agent_dir).unwrap();
|
||||
write(
|
||||
agent_dir.join("config.yaml"),
|
||||
format!("name: {agent_name}\ninstructions: hi\n"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
ctx.rag_key = Some(RagKey::Named("docs".to_string()));
|
||||
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap()
|
||||
.block_on(async {
|
||||
ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal())
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
assert!(ctx.rag.is_none());
|
||||
assert_eq!(ctx.rag_key, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -5970,6 +6103,28 @@ mod tests {
|
||||
assert!(paths::list_rags().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn list_rags_skips_sbx_mixin_sidecars() {
|
||||
let _guard = TestConfigDirGuard::new();
|
||||
let rags_dir = paths::rags_dir();
|
||||
create_dir_all(&rags_dir).unwrap();
|
||||
write(rags_dir.join("docs.yaml"), "embedding_model: test").unwrap();
|
||||
write(rags_dir.join("docs.sbx-mixin.yaml"), "kind: mixin").unwrap();
|
||||
write(rags_dir.join("v2.docs.yaml"), "embedding_model: test").unwrap();
|
||||
|
||||
let names = paths::list_rags();
|
||||
assert!(names.contains(&"docs".to_string()));
|
||||
assert!(
|
||||
names.contains(&"v2.docs".to_string()),
|
||||
"a dotted RAG name must still be listed: {names:?}"
|
||||
);
|
||||
assert!(
|
||||
!names.contains(&"docs.sbx-mixin".to_string()),
|
||||
"the sandbox mixin sidecar must not appear as a RAG: {names:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial]
|
||||
fn use_agent_errors_when_already_in_session() {
|
||||
|
||||
@@ -367,6 +367,13 @@ pub struct RagNode {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub graph_hops: Option<usize>,
|
||||
|
||||
/// Storage driver for this node's knowledge base ("yaml", "duckdb"). `None`
|
||||
/// means "yaml". Only honored when the knowledge base is first built;
|
||||
/// changing it afterwards has no effect until the RAG is deleted and
|
||||
/// re-initialized.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub driver: Option<String>,
|
||||
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub state_updates: Option<HashMap<String, String>>,
|
||||
|
||||
@@ -1152,4 +1159,100 @@ nodes:
|
||||
assert!(triage.next.as_ref().unwrap().is_fan_out());
|
||||
assert_eq!(triage.next.as_ref().unwrap().as_slice().len(), 2);
|
||||
}
|
||||
|
||||
fn rag_node_of(graph: &Graph, id: &str) -> RagNode {
|
||||
match &graph.get_node(id).unwrap().node_type {
|
||||
NodeType::Rag(r) => r.clone(),
|
||||
other => panic!("expected a rag node, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rag_node_deserializes_an_explicit_driver() {
|
||||
let yaml = r#"
|
||||
name: kb
|
||||
start: research
|
||||
nodes:
|
||||
research:
|
||||
type: rag
|
||||
documents: ["./docs"]
|
||||
driver: duckdb
|
||||
next: done
|
||||
done:
|
||||
type: end
|
||||
output: ok
|
||||
"#;
|
||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
rag_node_of(&graph, "research").driver.as_deref(),
|
||||
Some("duckdb")
|
||||
);
|
||||
}
|
||||
|
||||
/// Workflows written before drivers existed must keep parsing, and must keep
|
||||
/// asking for nothing, so `RagInitConfig` resolves them to the yaml default.
|
||||
#[test]
|
||||
fn rag_node_without_a_driver_stays_unset() {
|
||||
let yaml = r#"
|
||||
name: kb
|
||||
start: research
|
||||
nodes:
|
||||
research:
|
||||
type: rag
|
||||
documents: ["./docs"]
|
||||
next: done
|
||||
done:
|
||||
type: end
|
||||
output: ok
|
||||
"#;
|
||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
||||
|
||||
assert_eq!(rag_node_of(&graph, "research").driver, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rag_node_driver_survives_a_serialize_round_trip() {
|
||||
let yaml = r#"
|
||||
name: kb
|
||||
start: research
|
||||
nodes:
|
||||
research:
|
||||
type: rag
|
||||
documents: ["./docs"]
|
||||
driver: duckdb
|
||||
next: done
|
||||
done:
|
||||
type: end
|
||||
output: ok
|
||||
"#;
|
||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
||||
let reparsed: Graph =
|
||||
serde_yaml::from_str(&serde_yaml::to_string(&graph).unwrap()).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
rag_node_of(&reparsed, "research").driver.as_deref(),
|
||||
Some("duckdb")
|
||||
);
|
||||
}
|
||||
|
||||
/// `skip_serializing_if` must keep `driver:` out of graphs that never set it.
|
||||
#[test]
|
||||
fn rag_node_without_a_driver_omits_the_key_when_serialized() {
|
||||
let yaml = r#"
|
||||
name: kb
|
||||
start: research
|
||||
nodes:
|
||||
research:
|
||||
type: rag
|
||||
documents: ["./docs"]
|
||||
next: done
|
||||
done:
|
||||
type: end
|
||||
output: ok
|
||||
"#;
|
||||
let graph: Graph = serde_yaml::from_str(yaml).unwrap();
|
||||
|
||||
assert!(!serde_yaml::to_string(&graph).unwrap().contains("driver"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ use super::state::template_root_keys;
|
||||
use super::types::{Graph, Node, NodeType};
|
||||
use crate::client::{Model, ModelType};
|
||||
use crate::config::{Agent, AppConfig, paths};
|
||||
use crate::rag::{GraphRagConfig, RagData};
|
||||
use anyhow::{Result, bail};
|
||||
use std::collections::{BTreeMap, HashSet, VecDeque};
|
||||
use std::path::PathBuf;
|
||||
@@ -96,6 +97,51 @@ pub struct GraphValidator {
|
||||
skill_exists: fn(&str) -> bool,
|
||||
}
|
||||
|
||||
/// A minimal `RagData` whose only interesting field is `driver`. The numeric
|
||||
/// arguments are the smallest values that satisfy `validate()`'s unrelated
|
||||
/// floors (top_k >= 1, and chunk_size >= 1 with chunk_overlap < chunk_size for
|
||||
/// a non-attached RAG). `RagData::new` sets `attached: false`, which is the
|
||||
/// correct shape here: a graph rag node always builds its own local knowledge
|
||||
/// base from `documents` and can never be attached.
|
||||
fn rag_driver_probe(driver: &str) -> RagData {
|
||||
let mut data = RagData::new(
|
||||
String::new(),
|
||||
1,
|
||||
0,
|
||||
None,
|
||||
1,
|
||||
None,
|
||||
GraphRagConfig::default(),
|
||||
);
|
||||
data.driver = driver.to_string();
|
||||
data
|
||||
}
|
||||
|
||||
/// `Some(message)` when `driver` is one that `RagData::validate()` would reject.
|
||||
///
|
||||
/// The set of valid drivers is defined in exactly one place, `RagData::validate()`,
|
||||
/// so this asks that function rather than restating the list here.
|
||||
///
|
||||
/// Fails open on purpose: the first probe below uses the default driver, which is
|
||||
/// valid by definition. If even that one is rejected, `validate()` has grown a
|
||||
/// precondition the probe fixture no longer satisfies, and every verdict from here
|
||||
/// would be a false positive that rejects working graphs. In that case we decline
|
||||
/// to judge and leave enforcement to RAG construction. The
|
||||
/// `rag_driver_probe_fixture_is_accepted` test turns that silent degradation into a
|
||||
/// loud failure. Both `validate()` calls are load-bearing; neither is redundant.
|
||||
pub(crate) fn rag_driver_error(driver: &str) -> Option<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 {
|
||||
pub fn new(base_dir: impl Into<PathBuf>) -> Self {
|
||||
Self {
|
||||
@@ -154,6 +200,11 @@ impl GraphValidator {
|
||||
not be written to state",
|
||||
));
|
||||
}
|
||||
if let Some(driver) = &r.driver
|
||||
&& let Some(message) = rag_driver_error(driver)
|
||||
{
|
||||
result.error(ValidationError::with_node(node_id, message));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1031,6 +1082,7 @@ mod tests {
|
||||
extractor_model: None,
|
||||
extractor_prompt: None,
|
||||
graph_hops: None,
|
||||
driver: None,
|
||||
state_updates,
|
||||
timeout: None,
|
||||
}),
|
||||
@@ -1385,6 +1437,55 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// Guards the fail-open branch in `rag_driver_error`. If this fails,
|
||||
/// `RagData::validate()` grew a precondition the probe fixture no longer
|
||||
/// satisfies and rag-node driver validation has silently switched itself off.
|
||||
/// Repair the fixture in `rag_driver_probe`; do not delete this test.
|
||||
#[test]
|
||||
fn rag_driver_probe_fixture_is_accepted() {
|
||||
let default_driver = RagData::default().driver;
|
||||
assert!(
|
||||
rag_driver_probe(&default_driver).validate().is_ok(),
|
||||
"probe fixture rejected for the default driver '{default_driver}'"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rag_driver_error_defers_to_ragdata_validate() {
|
||||
assert_eq!(rag_driver_error("yaml"), None);
|
||||
assert_eq!(rag_driver_error("duckdb"), None);
|
||||
|
||||
let message = rag_driver_error("duckdbb").expect("unknown driver must be rejected");
|
||||
assert!(message.contains("duckdbb"), "got: {message}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rag_node_with_unknown_driver_errors_naming_the_node() {
|
||||
let mut node = rag_node("kb", &["./docs"], true);
|
||||
if let NodeType::Rag(ref mut r) = node.node_type {
|
||||
r.driver = Some("postgres".into());
|
||||
}
|
||||
let graph = graph_with(vec![("kb", node), ("end", end_node("end"))], "kb");
|
||||
|
||||
let result = validator().validate(&graph);
|
||||
|
||||
assert!(!result.is_valid());
|
||||
let err = result.into_result().unwrap_err().to_string();
|
||||
assert!(err.contains("[kb]"), "must name the node: {err}");
|
||||
assert!(err.contains("postgres"), "must name the driver: {err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rag_node_with_duckdb_driver_produces_no_findings() {
|
||||
let mut node = rag_node("kb", &["./docs"], true);
|
||||
if let NodeType::Rag(ref mut r) = node.node_type {
|
||||
r.driver = Some("duckdb".into());
|
||||
}
|
||||
let graph = graph_with(vec![("kb", node), ("end", end_node("end"))], "kb");
|
||||
|
||||
assert!(validator().validate(&graph).is_valid());
|
||||
}
|
||||
|
||||
fn agent_node(id: &str, agent: &str, next: Option<&str>) -> Node {
|
||||
Node {
|
||||
id: id.into(),
|
||||
|
||||
+1532
-84
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,76 @@
|
||||
use super::{DocumentId, RagData};
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
|
||||
/// Abstracts where RAG vector data is stored and queried.
|
||||
///
|
||||
/// The Rag orchestrator owns: embeddings, chunking, BM25 keyword search, graph RAG,
|
||||
/// entity extraction, RRF merging. Providers own: vector storage and content retrieval.
|
||||
#[async_trait]
|
||||
pub trait RagProvider: Send + Sync {
|
||||
/// Vector similarity search. Returns (DocumentId, score) sorted by score desc.
|
||||
/// `embedding` is a single query vector from Coyote's embedding model.
|
||||
async fn vector_search(
|
||||
&self,
|
||||
embedding: &[f32],
|
||||
top_k: usize,
|
||||
min_score: f32,
|
||||
) -> Result<Vec<(DocumentId, f32)>>;
|
||||
|
||||
/// Resolve document IDs to their page content.
|
||||
///
|
||||
/// **Ordering contract:** implementations MUST return results in the same
|
||||
/// relative order as the input `ids` slice. `hybrid_search` passes an
|
||||
/// RRF-ranked list and feeds the result straight to the LLM. A provider
|
||||
/// that returns rows in storage order (e.g. Qdrant `get_points`, DuckDB
|
||||
/// `WHERE id IN (...)`) would silently discard the ranking. Implementations
|
||||
/// that query an unordered backend must re-sort by input position before
|
||||
/// returning.
|
||||
///
|
||||
/// Returns only IDs that were found; callers must handle partial returns
|
||||
/// (a missing ID is skipped, not an error).
|
||||
/// YamlProvider: reads from an in-memory content map built from data.files.
|
||||
/// DuckDbProvider: queries the documents table by id.
|
||||
/// QdrantProvider: fetches payload from the remote collection.
|
||||
async fn fetch_content(&self, ids: &[DocumentId]) -> Result<Vec<(DocumentId, String)>>;
|
||||
|
||||
/// Rebuild internal indexes from freshly updated RagData.
|
||||
/// Called once at the end of every sync_documents pass.
|
||||
///
|
||||
/// `full_rebuild` mirrors `sync_documents`' `refresh` parameter:
|
||||
/// - `true`: a full re-index (`.rebuild rag`, `--rebuild-rag`, initial build).
|
||||
/// Destructive strategies (wipe-then-reindex) are permitted.
|
||||
/// - `false`: an incremental change (`.edit rag-docs` adding/removing a file).
|
||||
/// Implementations MUST NOT wipe existing state; upsert only.
|
||||
///
|
||||
/// The parameter is part of the signature from the outset so it is fixed
|
||||
/// while there is exactly one implementor. Yaml/DuckDb ignore it,
|
||||
/// rebuilding their local state wholesale is fast and always correct.
|
||||
/// Only a remote provider is destructive enough to care.
|
||||
async fn rebuild_indexes(&mut self, data: &RagData, full_rebuild: bool) -> Result<()>;
|
||||
|
||||
/// Keyword / full-text search. Returns (DocumentId, BM25-style score) sorted desc.
|
||||
///
|
||||
/// Default impl returns `Ok(vec![])`. Callers fall back to `Rag.bm25` (local in-memory
|
||||
/// BM25 built from `data.files`).
|
||||
///
|
||||
/// Callers check `has_native_keyword_search()` before deciding which path to take:
|
||||
/// - true → call this method; skip `Rag.bm25`
|
||||
/// - false → call `Rag.keyword_search()` which uses `Rag.bm25` (sync, infallible)
|
||||
async fn keyword_search(&self, query: &str, top_k: usize) -> Result<Vec<(DocumentId, f32)>> {
|
||||
let _ = (query, top_k);
|
||||
|
||||
Ok(vec![])
|
||||
}
|
||||
|
||||
/// Returns true if this provider implements a native keyword-search index.
|
||||
/// When false, `Rag.hybrid_search` uses the local `Rag.bm25` field instead.
|
||||
fn has_native_keyword_search(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Deep-clone the provider with fresh indexes derived from `data`.
|
||||
/// Required because Box<dyn RagProvider> is not Clone.
|
||||
/// Called by Rag's Clone impl (which clones before mutating in rebuild_rag/edit_rag_docs).
|
||||
fn duplicate(&self, data: &RagData) -> Box<dyn RagProvider>;
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,9 @@
|
||||
mod yaml;
|
||||
pub use self::yaml::YamlProvider;
|
||||
|
||||
mod duckdb;
|
||||
pub use self::duckdb::DuckDbProvider;
|
||||
pub(crate) use self::duckdb::duckdb_path_from_yaml;
|
||||
|
||||
mod qdrant;
|
||||
pub use self::qdrant::QdrantProvider;
|
||||
@@ -0,0 +1,828 @@
|
||||
use crate::rag::provider::RagProvider;
|
||||
use crate::rag::{DocumentId, RagData};
|
||||
|
||||
use anyhow::{Context, Result, bail};
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::RwLock;
|
||||
use reqwest::header::{HeaderMap, HeaderValue};
|
||||
use reqwest::{Client, Response, StatusCode};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use url::{Host, Url};
|
||||
|
||||
/// 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)| min_score <= 0.0 || *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.
|
||||
///
|
||||
/// `body` is the raw response text. Two shapes have to be tolerated:
|
||||
/// * application-level errors carry `{"status": {"error": "..."}, "time": 0.0}`,
|
||||
/// while successful responses carry a bare string `{"status": "ok", ...}` — so
|
||||
/// `status` is string-or-object and a struct with `status: String` fails to
|
||||
/// parse every error body;
|
||||
/// * routing-level 404s (a wrong HTTP verb) return an EMPTY body with no JSON at
|
||||
/// all, which without the length check surfaces as "EOF while parsing a value"
|
||||
/// instead of the actual 404.
|
||||
fn format_error_body(status: StatusCode, body: &str) -> String {
|
||||
if body.is_empty() {
|
||||
return format!("HTTP {status} (empty body — check the HTTP verb and path)");
|
||||
}
|
||||
serde_json::from_str::<Value>(body)
|
||||
.ok()
|
||||
.and_then(|v| v["status"]["error"].as_str().map(str::to_string))
|
||||
.unwrap_or_else(|| format!("HTTP {status}: {body}"))
|
||||
}
|
||||
|
||||
/// Read the vector dimension out of a parsed `GET /collections/{name}` response.
|
||||
fn vector_dimension_from_collection(body: &Value) -> Result<u64> {
|
||||
let params = &body["result"]["config"]["params"];
|
||||
params["vectors"]["size"]
|
||||
.as_u64()
|
||||
.or_else(|| {
|
||||
params["vectors"]
|
||||
.as_object()
|
||||
.and_then(|m| m.values().next())
|
||||
.and_then(|v| v["size"].as_u64())
|
||||
})
|
||||
.context("Could not determine vector dimension from collection config")
|
||||
}
|
||||
|
||||
/// True if a parsed `GET /collections/{name}` response describes a NAMED
|
||||
/// (multi-vector) collection.
|
||||
///
|
||||
/// `vector_search` posts an unnamed vector, which a named-vector collection
|
||||
/// rejects with HTTP 400 on every query, so attaching one yields a RAG that is
|
||||
/// silently 100% broken. A named collection holding a SINGLE vector is
|
||||
/// structurally a map, identical in kind to the multi-named case, and rejects
|
||||
/// the same way; testing for a numeric `size` directly under `vectors` catches
|
||||
/// it, whereas counting keys (`len() > 1`) would wrongly accept it.
|
||||
fn is_multi_vector_config(body: &Value) -> bool {
|
||||
body["result"]["config"]["params"]["vectors"]["size"]
|
||||
.as_u64()
|
||||
.is_none()
|
||||
}
|
||||
|
||||
/// Query-only client for an external Qdrant collection.
|
||||
///
|
||||
/// Attach-only: this provider never writes to the remote collection. Coyote does
|
||||
/// not own the data, and `rebuild_indexes` refuses rather than pretending to.
|
||||
pub struct QdrantProvider {
|
||||
client: Client,
|
||||
base_url: String,
|
||||
collection: String,
|
||||
point_ids: Arc<RwLock<PointIdInterner>>,
|
||||
}
|
||||
|
||||
impl QdrantProvider {
|
||||
fn skips_proxy(base_url: &str) -> bool {
|
||||
let Ok(url) = Url::parse(base_url) else {
|
||||
return false;
|
||||
};
|
||||
match url.host() {
|
||||
Some(Host::Domain(name)) => {
|
||||
name == "localhost" || name.ends_with(".localhost") || name.ends_with(".local")
|
||||
}
|
||||
Some(Host::Ipv4(ip)) => ip.is_loopback() || ip.is_private() || ip.is_link_local(),
|
||||
// No stable is_unique_local, so fc00::/7 is matched directly.
|
||||
Some(Host::Ipv6(ip)) => ip.is_loopback() || ip.segments()[0] & 0xfe00 == 0xfc00,
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_client(base_url: &str, api_key: Option<&str>) -> Result<Client> {
|
||||
let mut headers = HeaderMap::new();
|
||||
if let Some(key) = api_key {
|
||||
let mut value =
|
||||
HeaderValue::from_str(key).context("api-key header value is not valid ASCII")?;
|
||||
value.set_sensitive(true);
|
||||
headers.insert("api-key", value);
|
||||
}
|
||||
let mut builder = Client::builder().default_headers(headers);
|
||||
if Self::skips_proxy(base_url) {
|
||||
builder = builder.no_proxy();
|
||||
}
|
||||
builder.build().context("Failed to build reqwest client")
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_base_url(host: &str) -> String {
|
||||
if host.starts_with("http://") || host.starts_with("https://") {
|
||||
host.to_string()
|
||||
} else {
|
||||
format!("http://{host}")
|
||||
}
|
||||
}
|
||||
|
||||
async fn error_message(resp: Response) -> String {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
format_error_body(status, &body)
|
||||
}
|
||||
|
||||
/// Shared `GET /collections/{name}` fetch. Both the dimension and the
|
||||
/// multi-vector probe discriminate on this same response.
|
||||
async fn fetch_collection(
|
||||
host: &str,
|
||||
collection: &str,
|
||||
api_key: Option<&str>,
|
||||
) -> Result<Value> {
|
||||
let base_url = Self::normalize_base_url(host);
|
||||
let client = Self::make_client(&base_url, api_key)?;
|
||||
let resp = client
|
||||
.get(format!("{base_url}/collections/{collection}"))
|
||||
.send()
|
||||
.await
|
||||
.with_context(|| format!("Failed to connect to {host}"))?;
|
||||
if !resp.status().is_success() {
|
||||
bail!(
|
||||
"Failed to read collection '{collection}': {}",
|
||||
Self::error_message(resp).await
|
||||
);
|
||||
}
|
||||
|
||||
Ok(resp.json().await?)
|
||||
}
|
||||
|
||||
pub async fn new(host: &str, collection: &str, api_key: Option<&str>) -> Result<Self> {
|
||||
let base_url = Self::normalize_base_url(host);
|
||||
let client = Self::make_client(&base_url, api_key)?;
|
||||
let resp = client
|
||||
.get(format!("{base_url}/collections/{collection}"))
|
||||
.send()
|
||||
.await
|
||||
.with_context(|| format!("Failed to connect to {host}"))?;
|
||||
if !resp.status().is_success() {
|
||||
bail!(
|
||||
"Collection '{collection}' not accessible at {host}: {}",
|
||||
Self::error_message(resp).await
|
||||
);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
client,
|
||||
base_url,
|
||||
collection: collection.to_string(),
|
||||
point_ids: Arc::default(),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn list_collections(host: &str, api_key: Option<&str>) -> Result<Vec<String>> {
|
||||
let base_url = Self::normalize_base_url(host);
|
||||
let client = Self::make_client(&base_url, api_key)?;
|
||||
let resp = client
|
||||
.get(format!("{base_url}/collections"))
|
||||
.send()
|
||||
.await
|
||||
.with_context(|| format!("Failed to connect to {host}"))?;
|
||||
if !resp.status().is_success() {
|
||||
bail!(
|
||||
"Failed to list collections: {}",
|
||||
Self::error_message(resp).await
|
||||
);
|
||||
}
|
||||
|
||||
let body: Value = resp.json().await?;
|
||||
let names = body["result"]["collections"]
|
||||
.as_array()
|
||||
.context("Unexpected /collections response shape")?
|
||||
.iter()
|
||||
.filter_map(|v| v["name"].as_str().map(str::to_string))
|
||||
.collect();
|
||||
|
||||
Ok(names)
|
||||
}
|
||||
|
||||
pub async fn get_vector_dimension(
|
||||
host: &str,
|
||||
collection: &str,
|
||||
api_key: Option<&str>,
|
||||
) -> Result<u64> {
|
||||
let body = Self::fetch_collection(host, collection, api_key).await?;
|
||||
|
||||
vector_dimension_from_collection(&body)
|
||||
}
|
||||
|
||||
pub async fn is_multi_vector(
|
||||
host: &str,
|
||||
collection: &str,
|
||||
api_key: Option<&str>,
|
||||
) -> Result<bool> {
|
||||
let body = Self::fetch_collection(host, collection, api_key).await?;
|
||||
|
||||
Ok(is_multi_vector_config(&body))
|
||||
}
|
||||
|
||||
pub async fn sample_point_id(
|
||||
host: &str,
|
||||
collection: &str,
|
||||
api_key: Option<&str>,
|
||||
) -> Result<Option<String>> {
|
||||
let base_url = Self::normalize_base_url(host);
|
||||
let client = Self::make_client(&base_url, api_key)?;
|
||||
let url = format!("{base_url}/collections/{collection}/points/scroll");
|
||||
let body = serde_json::json!({ "limit": 1, "with_payload": false });
|
||||
|
||||
let resp = client
|
||||
.post(&url)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.with_context(|| format!("Failed to connect to {host}"))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
bail!(
|
||||
"Failed to sample a point from '{collection}': {}",
|
||||
Self::error_message(resp).await
|
||||
);
|
||||
}
|
||||
|
||||
let data: Value = resp.json().await?;
|
||||
let id_val = data["result"]["points"]
|
||||
.as_array()
|
||||
.and_then(|pts| pts.first())
|
||||
.map(|pt| pt["id"].to_string());
|
||||
|
||||
Ok(id_val)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RagProvider for QdrantProvider {
|
||||
async fn vector_search(
|
||||
&self,
|
||||
embedding: &[f32],
|
||||
top_k: usize,
|
||||
min_score: f32,
|
||||
) -> Result<Vec<(DocumentId, f32)>> {
|
||||
let url = format!(
|
||||
"{}/collections/{}/points/search",
|
||||
self.base_url, self.collection
|
||||
);
|
||||
// `score_threshold` is deliberately NOT sent. It is metric-aware: on Cosine
|
||||
// collections 0.0 means "no floor" as expected, but Euclid collections score
|
||||
// by negative distance, where 0.0 filters everything out. The attach wizard
|
||||
// does not pin the distance metric, so filter locally instead; i.e. where a
|
||||
// 0.0 floor is correctly treated as "no floor" (see `parse_search_hits`).
|
||||
let body = serde_json::json!({
|
||||
"vector": embedding,
|
||||
"limit": top_k,
|
||||
"with_payload": false,
|
||||
});
|
||||
let resp = self.client.post(&url).json(&body).send().await?;
|
||||
if !resp.status().is_success() {
|
||||
bail!(
|
||||
"Qdrant search on '{}' failed: {}",
|
||||
self.collection,
|
||||
Self::error_message(resp).await
|
||||
);
|
||||
}
|
||||
let data: Value = resp.json().await?;
|
||||
// The interner is what lets a UUID-keyed collection work: a string id gets
|
||||
// a synthetic handle here and the original is replayed by `fetch_content`.
|
||||
let mut interner = self.point_ids.write();
|
||||
|
||||
parse_search_hits(&mut interner, &data, min_score)
|
||||
}
|
||||
|
||||
async fn fetch_content(&self, ids: &[DocumentId]) -> Result<Vec<(DocumentId, String)>> {
|
||||
if ids.is_empty() {
|
||||
return Ok(vec![]);
|
||||
}
|
||||
let url = format!("{}/collections/{}/points", self.base_url, self.collection);
|
||||
// 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!({
|
||||
"ids": id_list,
|
||||
"with_payload": true,
|
||||
});
|
||||
|
||||
let resp = self.client.post(&url).json(&body).send().await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
bail!(
|
||||
"Qdrant point fetch on '{}' failed: {}",
|
||||
self.collection,
|
||||
Self::error_message(resp).await
|
||||
);
|
||||
}
|
||||
let data: Value = resp.json().await?;
|
||||
let mut rows = {
|
||||
let mut interner = self.point_ids.write();
|
||||
parse_points(&mut interner, &data)?
|
||||
};
|
||||
// `/points` does not guarantee response order matches request order, and the
|
||||
// caller's RRF ranking is carried by that order. Restore it.
|
||||
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));
|
||||
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
async fn rebuild_indexes(&mut self, data: &RagData, _full_rebuild: bool) -> Result<()> {
|
||||
// Both arms refuse. A silent `Ok(())` would make `.rebuild rag` and
|
||||
// `.edit rag-docs` look like they worked while writing nothing to the
|
||||
// remote, leaving the user believing the collection was updated.
|
||||
if data.attached {
|
||||
bail!(
|
||||
"This RAG is attached to an external Qdrant collection. Coyote does not own \
|
||||
its documents and cannot rebuild it. Manage the collection directly, or \
|
||||
create a Coyote-owned RAG with `.rag <name>`."
|
||||
);
|
||||
}
|
||||
bail!("Writing to Qdrant is not supported yet (attach-only).");
|
||||
}
|
||||
|
||||
fn duplicate(&self, _data: &RagData) -> Box<dyn RagProvider> {
|
||||
// Cloning the client shares the connection pool and the injected api-key
|
||||
// header. Sharing is correct: both handles address the same remote
|
||||
// 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 {
|
||||
client: self.client.clone(),
|
||||
base_url: self.base_url.clone(),
|
||||
collection: self.collection.clone(),
|
||||
point_ids: Arc::clone(&self.point_ids),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn error_message_reads_the_object_status_envelope() {
|
||||
let body =
|
||||
r#"{"status": {"error": "Wrong input: Not existing vector name error:"}, "time": 0.0}"#;
|
||||
let msg = format_error_body(StatusCode::BAD_REQUEST, body);
|
||||
assert!(msg.contains("Not existing vector name"), "got: {msg}");
|
||||
assert!(
|
||||
!msg.contains("EOF"),
|
||||
"must not fall through to a parse error"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_message_survives_the_string_status_and_the_empty_body() {
|
||||
let ok = format_error_body(StatusCode::OK, r#"{"status": "ok", "time": 0.0}"#);
|
||||
assert!(
|
||||
ok.contains("200"),
|
||||
"no `status.error` present → fall back to status+body: {ok}"
|
||||
);
|
||||
|
||||
let empty = format_error_body(StatusCode::NOT_FOUND, "");
|
||||
|
||||
assert!(empty.contains("empty body"), "got: {empty}");
|
||||
assert!(
|
||||
empty.contains("verb"),
|
||||
"the message must point at the likely cause: {empty}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vector_dimension_handles_both_collection_shapes() {
|
||||
let unnamed = serde_json::json!({
|
||||
"result": {"config": {"params": {"vectors": {"size": 1536, "distance": "Cosine"}}}}
|
||||
});
|
||||
assert_eq!(vector_dimension_from_collection(&unnamed).unwrap(), 1536);
|
||||
|
||||
let named = serde_json::json!({
|
||||
"result": {"config": {"params": {"vectors": {"text": {"size": 768, "distance": "Cosine"}}}}}
|
||||
});
|
||||
assert_eq!(vector_dimension_from_collection(&named).unwrap(), 768);
|
||||
|
||||
let junk = serde_json::json!({"result": {"config": {"params": {}}}});
|
||||
assert!(vector_dimension_from_collection(&junk).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_multi_vector_rejects_the_named_single_collection() {
|
||||
let unnamed = serde_json::json!({
|
||||
"result": {"config": {"params": {"vectors": {"size": 1536, "distance": "Cosine"}}}}
|
||||
});
|
||||
assert!(!is_multi_vector_config(&unnamed));
|
||||
|
||||
let named_single = serde_json::json!({
|
||||
"result": {"config": {"params": {"vectors": {"text": {"size": 1536}}}}}
|
||||
});
|
||||
assert!(
|
||||
is_multi_vector_config(&named_single),
|
||||
"named-single must be rejected too"
|
||||
);
|
||||
|
||||
let named_multi = serde_json::json!({
|
||||
"result": {"config": {"params": {"vectors": {"text": {"size": 1536}, "image": {"size": 512}}}}}
|
||||
});
|
||||
assert!(is_multi_vector_config(&named_multi));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_base_url_only_adds_a_scheme_when_missing() {
|
||||
assert_eq!(
|
||||
QdrantProvider::normalize_base_url("qdrant.example.com:6333"),
|
||||
"http://qdrant.example.com:6333"
|
||||
);
|
||||
assert_eq!(
|
||||
QdrantProvider::normalize_base_url("https://xyz.cloud.qdrant.io"),
|
||||
"https://xyz.cloud.qdrant.io"
|
||||
);
|
||||
assert_eq!(
|
||||
QdrantProvider::normalize_base_url("http://localhost:6333"),
|
||||
"http://localhost:6333"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rebuild_indexes_refuses_for_attached_and_unattached_alike() {
|
||||
let mut provider = QdrantProvider {
|
||||
client: Client::new(),
|
||||
base_url: "http://localhost:6333".to_string(),
|
||||
collection: "c".to_string(),
|
||||
point_ids: Arc::default(),
|
||||
};
|
||||
|
||||
let attached = RagData {
|
||||
driver: "qdrant".to_string(),
|
||||
attached: true,
|
||||
..Default::default()
|
||||
};
|
||||
let err = provider
|
||||
.rebuild_indexes(&attached, true)
|
||||
.await
|
||||
.expect_err("an attached qdrant RAG must never report a successful rebuild");
|
||||
assert!(err.to_string().contains("cannot rebuild"), "got: {err}");
|
||||
|
||||
// `attached: false` is reserved for the (unimplemented) write path. It must
|
||||
// also refuse: silently succeeding would run a full paid embedding pass and
|
||||
// then discard every vector.
|
||||
let owned = RagData {
|
||||
driver: "qdrant".to_string(),
|
||||
attached: false,
|
||||
..Default::default()
|
||||
};
|
||||
let err = provider
|
||||
.rebuild_indexes(&owned, true)
|
||||
.await
|
||||
.expect_err("writing to qdrant is unimplemented and must fail loudly");
|
||||
assert!(err.to_string().contains("not supported yet"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_content_short_circuits_on_an_empty_id_list() {
|
||||
let provider = QdrantProvider {
|
||||
client: Client::new(),
|
||||
base_url: "http://127.0.0.1:1".to_string(),
|
||||
collection: "c".to_string(),
|
||||
point_ids: Arc::default(),
|
||||
};
|
||||
|
||||
assert!(provider.fetch_content(&[]).await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_and_private_hosts_skip_the_proxy() {
|
||||
for host in [
|
||||
"http://localhost:6333",
|
||||
"http://127.0.0.1:6333",
|
||||
"http://192.168.0.56:6333",
|
||||
"http://10.1.2.3:6333",
|
||||
"http://172.16.4.5:6333",
|
||||
"http://qdrant.local:6333",
|
||||
"http://[::1]:6333",
|
||||
] {
|
||||
assert!(
|
||||
QdrantProvider::skips_proxy(host),
|
||||
"{host} should not be proxied"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn public_hosts_still_honour_the_environment() {
|
||||
for host in [
|
||||
"https://qdrant.example.com",
|
||||
"http://8.8.8.8:6333",
|
||||
"https://xyz.eu-central.aws.cloud.qdrant.io:6333",
|
||||
"http://172.32.0.1:6333",
|
||||
] {
|
||||
assert!(
|
||||
!QdrantProvider::skips_proxy(host),
|
||||
"{host} must keep the environment's proxy"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Euclid collections score by NEGATIVE distance, so the 0.0 the caller
|
||||
/// passes must mean "no floor". Filtering on it drops every hit — the exact
|
||||
/// bug that keeps Qdrant's own `score_threshold` off the wire.
|
||||
#[test]
|
||||
fn a_zero_floor_keeps_negative_euclid_scores() {
|
||||
let mut interner = PointIdInterner::default();
|
||||
let search = serde_json::json!({
|
||||
"result": [
|
||||
{"id": 1, "score": -0.12},
|
||||
{"id": 2, "score": -8.5},
|
||||
]
|
||||
});
|
||||
|
||||
let hits = parse_search_hits(&mut interner, &search, 0.0).unwrap();
|
||||
|
||||
assert_eq!(hits.len(), 2, "a 0.0 floor must not drop negative scores");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_positive_floor_still_filters() {
|
||||
let mut interner = PointIdInterner::default();
|
||||
let search = serde_json::json!({
|
||||
"result": [
|
||||
{"id": 1, "score": 0.9},
|
||||
{"id": 2, "score": 0.2},
|
||||
]
|
||||
});
|
||||
|
||||
let hits = parse_search_hits(&mut interner, &search, 0.5).unwrap();
|
||||
|
||||
assert_eq!(hits.len(), 1);
|
||||
assert_eq!(hits[0].0, DocumentId(1));
|
||||
}
|
||||
|
||||
/// 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]
|
||||
#[ignore]
|
||||
async fn qdrant_list_collections_requires_running_instance() {
|
||||
let collections = QdrantProvider::list_collections("http://localhost:6333", None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!collections.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn qdrant_vector_search_returns_results() {
|
||||
let provider = QdrantProvider::new("http://localhost:6333", "test-collection", None)
|
||||
.await
|
||||
.unwrap();
|
||||
let embedding = vec![0.0f32; 1536];
|
||||
|
||||
let results = provider.vector_search(&embedding, 5, 0.0).await.unwrap();
|
||||
|
||||
assert!(results.len() <= 5);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,219 @@
|
||||
use crate::rag::provider::RagProvider;
|
||||
use crate::rag::{DocumentId, RagData};
|
||||
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use hnsw_rs::prelude::*;
|
||||
use indexmap::IndexMap;
|
||||
|
||||
pub struct YamlProvider {
|
||||
hnsw: Hnsw<'static, f32, DistCosine>,
|
||||
content_map: IndexMap<DocumentId, String>,
|
||||
}
|
||||
|
||||
impl YamlProvider {
|
||||
pub fn from_data(data: &RagData) -> Self {
|
||||
Self {
|
||||
hnsw: data.build_hnsw(),
|
||||
content_map: Self::build_content_map(data),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_content_map(data: &RagData) -> IndexMap<DocumentId, String> {
|
||||
data.iter_documents()
|
||||
.map(|(id, doc)| (id, doc.page_content.clone()))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RagProvider for YamlProvider {
|
||||
async fn vector_search(
|
||||
&self,
|
||||
embedding: &[f32],
|
||||
top_k: usize,
|
||||
min_score: f32,
|
||||
) -> Result<Vec<(DocumentId, f32)>> {
|
||||
let results = self
|
||||
.hnsw
|
||||
.parallel_search(&[embedding.to_vec()], top_k, 30)
|
||||
.into_iter()
|
||||
.flat_map(|list| {
|
||||
list.into_iter().filter_map(|v| {
|
||||
let score = 1.0 - v.distance;
|
||||
if score > min_score {
|
||||
Some((DocumentId(v.d_id), score))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
async fn fetch_content(&self, ids: &[DocumentId]) -> Result<Vec<(DocumentId, String)>> {
|
||||
Ok(ids
|
||||
.iter()
|
||||
.filter_map(|id| self.content_map.get(id).map(|text| (*id, text.clone())))
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn rebuild_indexes(&mut self, data: &RagData, _full_rebuild: bool) -> Result<()> {
|
||||
self.hnsw = data.build_hnsw();
|
||||
|
||||
self.content_map = Self::build_content_map(data);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn duplicate(&self, data: &RagData) -> Box<dyn RagProvider> {
|
||||
Box::new(YamlProvider::from_data(data))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod provider_tests {
|
||||
use super::*;
|
||||
use crate::rag::{RagDocument, RagFile};
|
||||
|
||||
fn minimal_rag_data() -> RagData {
|
||||
RagData {
|
||||
embedding_model: "text-embedding-3-small".to_string(),
|
||||
chunk_size: 1024,
|
||||
chunk_overlap: 50,
|
||||
top_k: 5,
|
||||
driver: "yaml".to_string(),
|
||||
attached: false,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Two files, one chunk each, with vectors, the minimum needed to exercise
|
||||
/// `build_content_map` and the `fetch_content` ordering contract.
|
||||
/// `DocumentId::new(f, d)` packs (file_index, document_index); `RagData::add`
|
||||
/// is the real insertion path but a direct literal is sufficient and avoids
|
||||
/// the embedding pipeline.
|
||||
fn populated_rag_data() -> RagData {
|
||||
let mut data = minimal_rag_data();
|
||||
// `files` must be populated: build_content_map iterates data.iter_documents(),
|
||||
// which enumerates `files`. Populating `vectors` alone would produce an EMPTY
|
||||
// content map, and every assertion below would vacuously pass on a broken impl.
|
||||
// The vectors inserted at the end are for the HNSW side only.
|
||||
data.files.insert(
|
||||
0,
|
||||
RagFile {
|
||||
hash: "h0".to_string(),
|
||||
path: "/tmp/a.md".to_string(),
|
||||
documents: vec![RagDocument {
|
||||
page_content: "alpha".to_string(),
|
||||
metadata: Default::default(),
|
||||
}],
|
||||
},
|
||||
);
|
||||
data.files.insert(
|
||||
1,
|
||||
RagFile {
|
||||
hash: "h1".to_string(),
|
||||
path: "/tmp/b.md".to_string(),
|
||||
documents: vec![RagDocument {
|
||||
page_content: "beta".to_string(),
|
||||
metadata: Default::default(),
|
||||
}],
|
||||
},
|
||||
);
|
||||
data.vectors
|
||||
.insert(DocumentId::new(0, 0), vec![1.0, 0.0, 0.0]);
|
||||
data.vectors
|
||||
.insert(DocumentId::new(1, 0), vec![0.0, 1.0, 0.0]);
|
||||
data
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn yaml_provider_empty_data_returns_nothing() {
|
||||
let data = minimal_rag_data();
|
||||
let provider = YamlProvider::from_data(&data);
|
||||
|
||||
let results = provider.fetch_content(&[]).await.unwrap();
|
||||
|
||||
assert!(results.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn yaml_provider_fetch_content_preserves_input_order() {
|
||||
let data = populated_rag_data();
|
||||
let provider = YamlProvider::from_data(&data);
|
||||
|
||||
let a = DocumentId::new(0, 0);
|
||||
let b = DocumentId::new(1, 0);
|
||||
|
||||
let forward = provider.fetch_content(&[a, b]).await.unwrap();
|
||||
assert_eq!(forward.len(), 2, "both documents must resolve");
|
||||
assert_eq!(forward[0].1, "alpha");
|
||||
assert_eq!(forward[1].1, "beta");
|
||||
|
||||
let reversed = provider.fetch_content(&[b, a]).await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
reversed[0].1, "beta",
|
||||
"fetch_content must honor input order"
|
||||
);
|
||||
assert_eq!(reversed[1].1, "alpha");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn yaml_provider_fetch_content_skips_missing_ids() {
|
||||
let data = populated_rag_data();
|
||||
let provider = YamlProvider::from_data(&data);
|
||||
|
||||
let a = DocumentId::new(0, 0);
|
||||
let missing = DocumentId::new(99, 0);
|
||||
let b = DocumentId::new(1, 0);
|
||||
|
||||
let out = provider.fetch_content(&[a, missing, b]).await.unwrap();
|
||||
assert_eq!(out.len(), 2, "missing id is skipped, not an error");
|
||||
assert_eq!(out[0].1, "alpha");
|
||||
assert_eq!(out[1].1, "beta");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn yaml_provider_duplicate_returns_equivalent_content() {
|
||||
let data = populated_rag_data();
|
||||
let provider = YamlProvider::from_data(&data);
|
||||
let dup = provider.duplicate(&data);
|
||||
let ids = [DocumentId::new(0, 0), DocumentId::new(1, 0)];
|
||||
|
||||
let r1 = provider.fetch_content(&ids).await.unwrap();
|
||||
let r2 = dup.fetch_content(&ids).await.unwrap();
|
||||
|
||||
assert_eq!(r1.len(), 2, "fixture must resolve both documents");
|
||||
assert_eq!(
|
||||
r1, r2,
|
||||
"duplicate must resolve the same content as the original"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn yaml_provider_content_is_keyed_on_files_not_vectors() {
|
||||
let mut data = populated_rag_data();
|
||||
let orphan = DocumentId::new(9, 0);
|
||||
data.vectors.insert(orphan, vec![0.0, 0.0, 1.0]);
|
||||
|
||||
let provider = YamlProvider::from_data(&data);
|
||||
|
||||
let out = provider.fetch_content(&[orphan]).await.unwrap();
|
||||
assert!(
|
||||
out.is_empty(),
|
||||
"an id present only in `vectors` must not resolve to content"
|
||||
);
|
||||
|
||||
let real = provider
|
||||
.fetch_content(&[DocumentId::new(0, 0), DocumentId::new(1, 0)])
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(real.len(), 2, "file-backed documents must still resolve");
|
||||
assert_eq!(real[0].1, "alpha");
|
||||
assert_eq!(real[1].1, "beta");
|
||||
}
|
||||
}
|
||||
@@ -1749,6 +1749,86 @@ std::error::Error>> {
|
||||
);
|
||||
}
|
||||
|
||||
fn strip_ansi(text: &str) -> String {
|
||||
let mut out = String::with_capacity(text.len());
|
||||
let mut chars = text.chars();
|
||||
while let Some(c) = chars.next() {
|
||||
if c != '\u{1b}' {
|
||||
out.push(c);
|
||||
continue;
|
||||
}
|
||||
if chars.next() == Some('[') {
|
||||
for c in chars.by_ref() {
|
||||
if ('\u{40}'..='\u{7e}').contains(&c) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_ansi_removes_sgr_and_keeps_text() {
|
||||
assert_eq!(strip_ansi("\x1b[1mbold\x1b[0m"), "bold");
|
||||
assert_eq!(strip_ansi("\x1b[38;5;120mx\x1b[39m"), "x");
|
||||
assert_eq!(strip_ansi("plain"), "plain");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_table_pads_columns_by_display_width() {
|
||||
use unicode_width::UnicodeWidthStr;
|
||||
|
||||
const WRAP_WIDTH: u16 = 80;
|
||||
|
||||
let options = RenderOptions::default();
|
||||
let mut render = MarkdownRender::init(options).unwrap();
|
||||
render.wrap_width = Some(WRAP_WIDTH);
|
||||
|
||||
let header = vec![
|
||||
"**Setting**".into(),
|
||||
"*Default*".into(),
|
||||
"`Description`".into(),
|
||||
];
|
||||
let alignments = vec![
|
||||
CellAlignment::Left,
|
||||
CellAlignment::Right,
|
||||
CellAlignment::Center,
|
||||
];
|
||||
let rows = vec![
|
||||
vec![
|
||||
"**temperature**".into(),
|
||||
"`0.7`".into(),
|
||||
"Controls how *random* the sampled reply is allowed to be".into(),
|
||||
],
|
||||
vec![
|
||||
"**top_p**".into(),
|
||||
"`1.0`".into(),
|
||||
"Nucleus sampling cutoff, applied **after** temperature".into(),
|
||||
],
|
||||
];
|
||||
|
||||
let output = render.render_table(header, alignments, rows);
|
||||
|
||||
assert!(
|
||||
output.contains('\u{1b}'),
|
||||
"fixture must actually contain ANSI escapes: {output:?}",
|
||||
);
|
||||
|
||||
let widths: Vec<usize> = output
|
||||
.lines()
|
||||
.map(|line| strip_ansi(line).width())
|
||||
.collect();
|
||||
assert!(!widths.is_empty(), "table rendered no lines");
|
||||
|
||||
for (index, width) in widths.iter().enumerate() {
|
||||
assert_eq!(
|
||||
*width, WRAP_WIDTH as usize,
|
||||
"line {index} display width; all widths were {widths:?} in output:\n{output}",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_machine_renders_full_table_and_flushes_on_paragraph() {
|
||||
let options = RenderOptions::default();
|
||||
|
||||
+17
-4
@@ -53,7 +53,7 @@ pub const DEFAULT_CONTINUATION_PROMPT: &str = indoc! {"
|
||||
4. Continue with the next pending item now. Call tools immediately."
|
||||
};
|
||||
|
||||
static REPL_COMMANDS: LazyLock<[ReplCommand; 59]> = LazyLock::new(|| {
|
||||
static REPL_COMMANDS: LazyLock<[ReplCommand; 60]> = LazyLock::new(|| {
|
||||
[
|
||||
ReplCommand::new(".help", "Show this help guide", AssertState::pass()),
|
||||
ReplCommand::new(".info", "Show system info", AssertState::pass()),
|
||||
@@ -217,6 +217,11 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 59]> = LazyLock::new(|| {
|
||||
"Initialize or access RAG",
|
||||
AssertState::False(StateFlags::AGENT),
|
||||
),
|
||||
ReplCommand::new(
|
||||
".rag attach",
|
||||
"Attach to a pre-existing external RAG",
|
||||
AssertState::False(StateFlags::AGENT),
|
||||
),
|
||||
ReplCommand::new(
|
||||
".edit rag-docs",
|
||||
"Add or remove documents from an existing RAG",
|
||||
@@ -884,9 +889,17 @@ pub async fn run_repl_command(
|
||||
let version = args.map(|s| s.trim().to_string());
|
||||
task::spawn_blocking(move || config::run_self_update(version, false)).await??;
|
||||
}
|
||||
".rag" => {
|
||||
".rag" => match split_first_arg(args) {
|
||||
Some(("attach", rest)) => match rest {
|
||||
Some(name) if !name.trim().is_empty() => {
|
||||
ctx.attach_rag(name.trim()).await?;
|
||||
}
|
||||
_ => println!("Usage: .rag attach <name>"),
|
||||
},
|
||||
_ => {
|
||||
ctx.use_rag(args, abort_signal.clone()).await?;
|
||||
}
|
||||
},
|
||||
".agent" => match split_first_arg(args) {
|
||||
Some((agent_name, args)) => {
|
||||
let (new_args, _) = split_args_text(args.unwrap_or_default(), cfg!(windows));
|
||||
@@ -1711,8 +1724,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repl_commands_has_59_entries() {
|
||||
assert_eq!(REPL_COMMANDS.len(), 59);
|
||||
fn repl_commands_has_60_entries() {
|
||||
assert_eq!(REPL_COMMANDS.len(), 60);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -387,7 +387,7 @@ pub(crate) fn collect_server_allow_entries(
|
||||
out.into_iter().collect()
|
||||
}
|
||||
|
||||
fn allow_entry_for_url(raw: &str) -> Option<String> {
|
||||
pub(crate) fn allow_entry_for_url(raw: &str) -> Option<String> {
|
||||
let url = Url::parse(raw).ok()?;
|
||||
let scheme = url.scheme();
|
||||
if scheme != "https" && scheme != "http" {
|
||||
@@ -426,8 +426,8 @@ fn placeholders(text: &str) -> Result<Vec<PlaceholderMatch>> {
|
||||
struct CredentialsMixin {
|
||||
schema_version: &'static str,
|
||||
kind: &'static str,
|
||||
name: &'static str,
|
||||
description: &'static str,
|
||||
name: String,
|
||||
description: String,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
credentials: Vec<CredentialEntry>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
@@ -435,20 +435,20 @@ struct CredentialsMixin {
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct CredentialEntry {
|
||||
service: String,
|
||||
description: String,
|
||||
pub(crate) struct CredentialEntry {
|
||||
pub service: String,
|
||||
pub description: String,
|
||||
#[serde(rename = "apiKey")]
|
||||
api_key: ApiKey,
|
||||
pub api_key: ApiKey,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ApiKey {
|
||||
name: String,
|
||||
proxy_managed: bool,
|
||||
pub(crate) struct ApiKey {
|
||||
pub name: String,
|
||||
pub proxy_managed: bool,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
inject: Vec<InjectRule>,
|
||||
pub inject: Vec<InjectRule>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
@@ -461,25 +461,39 @@ struct Network {
|
||||
allow: Vec<String>,
|
||||
}
|
||||
|
||||
pub(crate) fn render_mixin_yaml(
|
||||
credentials: &[CredentialSpec],
|
||||
server_allow_entries: &[String],
|
||||
pub(crate) fn render_mixin_document(
|
||||
name: &str,
|
||||
description: &str,
|
||||
credentials: Vec<CredentialEntry>,
|
||||
extra_allow_entries: &[String],
|
||||
) -> Result<String> {
|
||||
let mut allow: BTreeSet<String> = credentials
|
||||
.iter()
|
||||
.flat_map(|c| c.inject.iter().map(|r| r.domain.clone()))
|
||||
.flat_map(|c| c.api_key.inject.iter().map(|r| r.domain.clone()))
|
||||
.collect();
|
||||
allow.extend(server_allow_entries.iter().cloned());
|
||||
allow.extend(extra_allow_entries.iter().cloned());
|
||||
|
||||
let mixin = CredentialsMixin {
|
||||
schema_version: "2",
|
||||
kind: "mixin",
|
||||
name: MCP_MIXIN_NAME,
|
||||
description: "Auto-generated by Coyote at launch: allows network egress to the user's \
|
||||
remote MCP servers and declares their credentials so Docker Sandboxes \
|
||||
binds them (bindings are approved on first interactive run). Values are \
|
||||
pre-seeded from Coyote's vault via `sbx secret set`.",
|
||||
credentials: credentials
|
||||
name: name.to_string(),
|
||||
description: description.to_string(),
|
||||
credentials,
|
||||
permissions: (!allow.is_empty()).then(|| Permissions {
|
||||
network: Network {
|
||||
allow: allow.into_iter().collect(),
|
||||
},
|
||||
}),
|
||||
};
|
||||
|
||||
serde_yaml::to_string(&mixin).context("Failed to serialize generated sandbox mixin")
|
||||
}
|
||||
|
||||
pub(crate) fn render_mixin_yaml(
|
||||
credentials: &[CredentialSpec],
|
||||
server_allow_entries: &[String],
|
||||
) -> Result<String> {
|
||||
let entries = credentials
|
||||
.iter()
|
||||
.map(|c| CredentialEntry {
|
||||
service: c.service_id.clone(),
|
||||
@@ -494,15 +508,17 @@ pub(crate) fn render_mixin_yaml(
|
||||
inject: c.inject.clone(),
|
||||
},
|
||||
})
|
||||
.collect(),
|
||||
permissions: (!allow.is_empty()).then(|| Permissions {
|
||||
network: Network {
|
||||
allow: allow.into_iter().collect(),
|
||||
},
|
||||
}),
|
||||
};
|
||||
.collect();
|
||||
|
||||
serde_yaml::to_string(&mixin).context("Failed to serialize generated MCP credentials mixin")
|
||||
render_mixin_document(
|
||||
MCP_MIXIN_NAME,
|
||||
"Auto-generated by Coyote at launch: allows network egress to the user's remote MCP \
|
||||
servers and declares their credentials so Docker Sandboxes binds them (bindings are \
|
||||
approved on first interactive run). Values are pre-seeded from Coyote's vault via \
|
||||
`sbx secret set`.",
|
||||
entries,
|
||||
server_allow_entries,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
+210
-14
@@ -10,6 +10,7 @@ use sha2::{Digest, Sha256};
|
||||
use crate::config::paths;
|
||||
|
||||
const SBX_MIXIN_FILE_NAME: &str = "sbx-mixin.yaml";
|
||||
const SBX_MIXIN_FILE_SUFFIX: &str = ".sbx-mixin.yaml";
|
||||
const KIT_SPEC_FILE_NAME: &str = "spec.yaml";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -66,10 +67,16 @@ pub fn discover() -> Result<Vec<DiscoveredMixin>> {
|
||||
push_if_exists(&mut out, paths::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)?);
|
||||
}
|
||||
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)?);
|
||||
}
|
||||
for path in collect_mixins(&paths::rags_dir(), &[ScanMode::Flat]) {
|
||||
out.push(read_mixin(path)?);
|
||||
}
|
||||
|
||||
@@ -156,7 +163,73 @@ 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();
|
||||
|
||||
if modes.contains(&ScanMode::Flat) {
|
||||
result.extend(suffixed_mixins_in(dir));
|
||||
}
|
||||
|
||||
let named = modes.contains(&ScanMode::SubdirNamed);
|
||||
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() {
|
||||
result.push(candidate);
|
||||
}
|
||||
}
|
||||
if subdir_flat {
|
||||
result.extend(suffixed_mixins_in(&subdir));
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn suffixed_mixins_in(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_file()).unwrap_or(false))
|
||||
.filter(|e| {
|
||||
e.file_name()
|
||||
.to_str()
|
||||
.is_some_and(|n| n.ends_with(SBX_MIXIN_FILE_SUFFIX))
|
||||
})
|
||||
.collect();
|
||||
entries.sort_by_key(|e| e.file_name());
|
||||
|
||||
result.extend(entries.into_iter().map(|e| e.path()));
|
||||
result
|
||||
}
|
||||
|
||||
fn subdirs_of(dir: &Path) -> Vec<PathBuf> {
|
||||
let mut result = Vec::new();
|
||||
let Ok(rd) = read_dir(dir) else { return result };
|
||||
|
||||
@@ -166,13 +239,7 @@ fn collect_subdir_mixins(dir: &Path) -> Vec<PathBuf> {
|
||||
.collect();
|
||||
entries.sort_by_key(|e| e.file_name());
|
||||
|
||||
for entry in entries {
|
||||
let candidate = entry.path().join(SBX_MIXIN_FILE_NAME);
|
||||
if candidate.exists() {
|
||||
result.push(candidate);
|
||||
}
|
||||
}
|
||||
|
||||
result.extend(entries.into_iter().map(|e| e.path()));
|
||||
result
|
||||
}
|
||||
|
||||
@@ -192,6 +259,13 @@ mod tests {
|
||||
root
|
||||
}
|
||||
|
||||
fn file_names(paths: &[PathBuf]) -> Vec<&str> {
|
||||
paths
|
||||
.iter()
|
||||
.map(|p| p.file_name().unwrap().to_str().unwrap())
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summarize_counts_installs_and_domains() {
|
||||
let root = unique_root("sbx-mixin-counts");
|
||||
@@ -275,7 +349,7 @@ network:
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collect_subdir_mixins_sorts_and_skips_missing() {
|
||||
fn subdir_named_scan_sorts_and_skips_missing() {
|
||||
let root = unique_root("sbx-mixin-subdirs");
|
||||
for name in ["zebra", "apple", "no-mixin", "mango"] {
|
||||
let dir = root.join(name);
|
||||
@@ -285,7 +359,7 @@ network:
|
||||
}
|
||||
}
|
||||
|
||||
let found = collect_subdir_mixins(&root);
|
||||
let found = collect_mixins(&root, &[ScanMode::SubdirNamed]);
|
||||
let names: Vec<String> = found
|
||||
.iter()
|
||||
.map(|p| {
|
||||
@@ -303,9 +377,9 @@ network:
|
||||
}
|
||||
|
||||
#[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 found = collect_subdir_mixins(&absent);
|
||||
let found = collect_mixins(&absent, &[ScanMode::SubdirNamed]);
|
||||
assert!(found.is_empty());
|
||||
}
|
||||
|
||||
@@ -483,4 +557,126 @@ network:
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flat_scan_matches_rag_sidecars_by_suffix() {
|
||||
let root = unique_root("flat-mixins");
|
||||
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("company-docs.yaml"), "driver: qdrant\n").unwrap();
|
||||
fs::write(root.join("notes.yaml"), "driver: yaml\n").unwrap();
|
||||
fs::create_dir_all(root.join("decoy.sbx-mixin.yaml")).unwrap();
|
||||
|
||||
let found = collect_mixins(&root, &[ScanMode::Flat]);
|
||||
assert_eq!(
|
||||
file_names(&found),
|
||||
vec!["alpha.sbx-mixin.yaml", "company-docs.sbx-mixin.yaml"]
|
||||
);
|
||||
|
||||
let _ = fs::remove_dir_all(&root);
|
||||
}
|
||||
|
||||
/// Every scan site in `discover()` picks its modes assuming each mode owns
|
||||
/// exactly one layout and nothing else. `agents_data_dir()` requests two
|
||||
/// modes at once, so an overlap would collect the same file twice and
|
||||
/// `create_sandbox` would pass it as two `--kit` flags.
|
||||
#[test]
|
||||
fn each_scan_mode_owns_exactly_one_layout() {
|
||||
let root = unique_root("scan-mode-ownership");
|
||||
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_eq!(collect_mixins(&root, &[ScanMode::Flat]), vec![flat.clone()]);
|
||||
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);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flat_scan_tolerates_a_missing_directory() {
|
||||
let root = unique_root("flat-missing");
|
||||
let absent = root.join("nope");
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
+145
-1
@@ -10,7 +10,7 @@ use std::path::{Path, PathBuf};
|
||||
use std::process::{Command, Stdio};
|
||||
use which::which;
|
||||
|
||||
mod mcp_credentials;
|
||||
pub(crate) mod mcp_credentials;
|
||||
mod mixins;
|
||||
|
||||
pub(crate) use mcp_credentials::sandbox_secret_env_var;
|
||||
@@ -19,6 +19,7 @@ use crate::config::AppConfig;
|
||||
use crate::config::Config;
|
||||
use crate::config::VAULT_DATA_FILE_NAME;
|
||||
use crate::config::paths;
|
||||
use crate::rag::RagData;
|
||||
use crate::sandbox::mcp_credentials::MCP_MIXIN_NAME;
|
||||
use crate::sandbox::mixins::DiscoveredMixin;
|
||||
use crate::utils::run_command_with_output;
|
||||
@@ -53,6 +54,10 @@ pub fn launch(name: Option<String>, fresh: bool) -> Result<()> {
|
||||
let vault = Vault::init(&bootstrap)?;
|
||||
let registered = sbx_registered_services()?;
|
||||
inject_llm_secret(&config_content, &vault, ®istered)?;
|
||||
if !fresh {
|
||||
inject_rag_secrets(&vault, ®istered)?;
|
||||
}
|
||||
|
||||
let credentials_mixin = if fresh {
|
||||
None
|
||||
} else {
|
||||
@@ -309,6 +314,87 @@ fn inject_mcp_secrets(vault: &Vault, registered: &HashSet<String>) -> Result<Opt
|
||||
)?))
|
||||
}
|
||||
|
||||
fn inject_rag_secrets(vault: &Vault, registered: &HashSet<String>) -> Result<()> {
|
||||
let rags_dir = paths::rags_dir();
|
||||
if !rags_dir.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
for entry in fs::read_dir(&rags_dir)?.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("yaml") {
|
||||
continue;
|
||||
}
|
||||
let stem = match path.file_stem().and_then(|s| s.to_str()) {
|
||||
Some(s) if !paths::is_rag_sidecar_name(s) => s.to_string(),
|
||||
_ => continue,
|
||||
};
|
||||
let Ok(raw) = fs::read_to_string(&path) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(data) = serde_yaml::from_str::<RagData>(&raw) else {
|
||||
continue;
|
||||
};
|
||||
if !data.attached {
|
||||
continue;
|
||||
}
|
||||
let secret_names = driver_config_secret_names(&data);
|
||||
let Some((primary, extra)) = secret_names.split_first() else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let service_id = mcp_credentials::secret_service_id(&stem);
|
||||
if !service_id.is_empty() && !registered.contains(&service_id) {
|
||||
bind_rag_secret(vault, &service_id, primary, &stem)?;
|
||||
}
|
||||
|
||||
for name in extra {
|
||||
let id = mcp_credentials::secret_service_id(name);
|
||||
if !id.is_empty() && !registered.contains(&id) {
|
||||
bind_rag_secret(vault, &id, name, &stem)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn driver_config_secret_names(data: &RagData) -> Vec<String> {
|
||||
let mut names: Vec<String> = Vec::new();
|
||||
for value in data.driver_config.values() {
|
||||
let trimmed = value.trim();
|
||||
let Ok(Some(caps)) = SECRET_RE.captures(trimmed) else {
|
||||
continue;
|
||||
};
|
||||
if caps.get(0).map(|m| m.as_str()) != Some(trimmed) {
|
||||
continue;
|
||||
}
|
||||
let Some(name) = caps.get(1).map(|m| m.as_str().trim()) else {
|
||||
continue;
|
||||
};
|
||||
if !name.is_empty() && !names.iter().any(|n| n == name) {
|
||||
names.push(name.to_string());
|
||||
}
|
||||
}
|
||||
names
|
||||
}
|
||||
|
||||
fn bind_rag_secret(vault: &Vault, service_id: &str, secret_name: &str, stem: &str) -> Result<()> {
|
||||
match vault.get_secret(secret_name, false) {
|
||||
Ok(secret_value) => {
|
||||
sbx_secret_set(service_id, &secret_value)
|
||||
.context("Failed to register RAG secret with sbx")?;
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"Warning: could not load secret '{secret_name}' for RAG '{stem}': {e}. \
|
||||
Queries to this RAG will fail inside the sandbox. \
|
||||
Run `coyote --add-secret {secret_name}` to fix."
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn provider_to_sbx_service(provider_type: &str, client_name: Option<&str>) -> String {
|
||||
match provider_type {
|
||||
"claude" => "anthropic".to_string(),
|
||||
@@ -588,6 +674,64 @@ fn chown_agent_recursive(sandbox: &str, path: &str) -> Result<()> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn rag_with(driver_config: &[(&str, &str)]) -> RagData {
|
||||
let mut data = RagData::new("m".into(), 1024, 50, None, 5, None, Default::default());
|
||||
data.driver = "qdrant".to_string();
|
||||
data.attached = true;
|
||||
for (k, v) in driver_config {
|
||||
data.driver_config.insert(k.to_string(), v.to_string());
|
||||
}
|
||||
data
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn secret_names_are_found_whatever_the_field_is_called() {
|
||||
let data = rag_with(&[
|
||||
("host", "qdrant.example.com:6333"),
|
||||
("collection", "docs"),
|
||||
("token", "{{SOME_TOKEN}}"),
|
||||
]);
|
||||
|
||||
assert_eq!(driver_config_secret_names(&data), vec!["SOME_TOKEN"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_literal_credential_is_not_treated_as_a_secret_name() {
|
||||
let data = rag_with(&[("api_key", "sk-a-real-looking-key")]);
|
||||
|
||||
assert!(driver_config_secret_names(&data).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plain_values_are_never_mistaken_for_secrets() {
|
||||
let data = rag_with(&[("host", "localhost:6333"), ("collection", "docs")]);
|
||||
|
||||
assert!(driver_config_secret_names(&data).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_partial_placeholder_is_not_a_credential() {
|
||||
let data = rag_with(&[("api_key", "Bearer {{KEY}}")]);
|
||||
|
||||
assert!(driver_config_secret_names(&data).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn several_secrets_are_all_found_and_deduped() {
|
||||
let data = rag_with(&[
|
||||
("api_key", "{{QDRANT_KEY}}"),
|
||||
("host", "localhost:6333"),
|
||||
("token", "{{ OTHER_TOKEN }}"),
|
||||
("fallback_key", "{{QDRANT_KEY}}"),
|
||||
]);
|
||||
|
||||
assert_eq!(
|
||||
driver_config_secret_names(&data),
|
||||
vec!["QDRANT_KEY", "OTHER_TOKEN"],
|
||||
"order follows driver_config, and a repeat is not registered twice"
|
||||
);
|
||||
}
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user