vrr/ruvector
2
0
Fork 0
mirror of https://github.com/ruvnet/RuVector.git synced 2026-07-30 03:23:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

chore(workspace): cargo fmt — mechanical whitespace fix across 427 files

Browse source 
Pre-existing rustfmt drift across the workspace was blocking CI's
`Rustfmt` check on PR #373 + PR #377. Running plain `cargo fmt`
reformats 427 files; no semantic changes, no logic changes, no
behavior changes — just what rustfmt already wanted.

None of the touched files are in ruvector-rabitq, ruvector-rulake,
or the new mirror-rulake workflow — those were already fmt-clean
per the per-crate checks on commits 5a4b0d782, 5f32fd450, f5003bc7b.
Drift is in cognitum-gate-kernel, mcp-brain, nervous-system,
prime-radiant, ruqu-core, ruvector-attention, ruvector-mincut,
ruvix/* and sub-crates, plus several examples.

Verified post-fmt:
  cargo check -p ruvector-rabitq -p ruvector-rulake            → clean
  cargo clippy -p ... -p ... --all-targets -- -D warnings      → clean
  cargo test   -p ... -p ... --release                         → 82/82 pass

Intentionally does NOT touch clippy drift — many more warnings
(missing docs, precision-loss casts, too-many-args, unsafe-safety-
docs) spread across unrelated crates, each category a cross-cutting
design decision that deserves its own review.

With this commit Rustfmt CI goes green on PR #373 and PR #377.
Clippy will still fail — that's honest pre-existing state for a
separate dedicated PR.

Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruvnet 2026-04-24 10:44:02 -04:00
parent f5003bc7b0
commit 96d8fdc172
429 changed files with 19898 additions and 10419 deletions
Download patch file Download diff file Expand all files Collapse all files

5
crates/cognitum-gate-kernel/src/lib.rs
View file

@ -563,7 +563,10 @@ pub unsafe extern "C" fn ingest_delta(ptr: *const u8, len: usize) -> i32 {
///
/// Returns 1 on success, 0 if buffer is full or tile not initialized.
#[no_mangle]
#[deprecated(since = "0.1.2", note = "Use ingest_delta(ptr, len) with bounds checking")]
#[deprecated(
since = "0.1.2",
note = "Use ingest_delta(ptr, len) with bounds checking"
)]
#[must_use]
pub unsafe extern "C" fn ingest_delta_unchecked(ptr: *const u8) -> i32 {
// Use Delta size as implied length

5
crates/cognitum-gate-kernel/tests/security_tests.rs
View file

@ -193,7 +193,10 @@ fn test_buffer_full_behavior() {
}
// Buffer should now be full
assert_eq!(tile.delta_count as usize, cognitum_gate_kernel::MAX_DELTA_BUFFER);
assert_eq!(
tile.delta_count as usize,
cognitum_gate_kernel::MAX_DELTA_BUFFER
);
// Next insert should fail
let delta = Delta::edge_add(999, 1000, 100);

9
crates/mcp-brain-server/src/auth.rs
View file

@ -4,7 +4,10 @@ use axum::{
extract::FromRequestParts,
http::{request::Parts, StatusCode},
};
use sha3::{Shake256, digest::{Update, ExtendableOutput, XofReader}};
use sha3::{
digest::{ExtendableOutput, Update, XofReader},
Shake256,
};
use subtle::ConstantTimeEq;
/// Authenticated contributor extracted from request
@ -53,7 +56,9 @@ const MIN_API_KEY_LEN: usize = 8;
/// If BRAIN_SYSTEM_KEY is unset, system key authentication is disabled entirely
/// (no hardcoded fallback).
static SYSTEM_KEY: std::sync::LazyLock<Option<String>> = std::sync::LazyLock::new(|| {
std::env::var("BRAIN_SYSTEM_KEY").ok().filter(|k| !k.is_empty())
std::env::var("BRAIN_SYSTEM_KEY")
.ok()
.filter(|k| !k.is_empty())
});
#[axum::async_trait]

545
crates/mcp-brain-server/src/bin/local.rs
View file

@ -15,41 +15,80 @@ use axum::{
};
use rusqlite::Connection;
use serde::{Deserialize, Serialize};
use std::sync::{Arc, Mutex, LazyLock};
use std::sync::{Arc, LazyLock, Mutex};
use tracing_subscriber::{fmt, prelude::*, EnvFilter};
// ── AIDefence (inline Rust port) ────────────────────────────────────────────
/// Critical injection + PII patterns compiled once.
static THREAT_PATTERNS: LazyLock<Vec<(regex::Regex, &'static str, &'static str)>> = LazyLock::new(|| {
let patterns: Vec<(&str, &str, &str)> = vec![
(r"(?i)ignore\s+(previous|all|above|any|the)(\s+\w+)*\s+(instructions?|prompts?|rules?|context)", "injection", "high"),
(r"(?i)disregard\s+(previous|all|above|the|your)(\s+\w+)*\s+(instructions?|prompts?|input)", "injection", "high"),
(r"(?i)forget\s+(everything|all|previous|your)", "injection", "high"),
(r"(?i)you\s+are\s+(now|actually)\s+", "injection", "high"),
(r"(?i)pretend\s+(to\s+be|you're|you\s+are)", "injection", "high"),
(r"(?i)what\s+(is|are)\s+your\s+(system\s+)?prompt", "extraction", "high"),
(r"(?i)show\s+(me\s+)?your\s+(system\s+)?instructions", "extraction", "high"),
(r"(?i)DAN\s+(mode|prompt)", "jailbreak", "critical"),
(r"(?i)bypass\s+(safety|security|filter)", "jailbreak", "critical"),
(r"(?i)remove\s+(all\s+)?restrictions", "jailbreak", "critical"),
(r"<script", "code_injection", "critical"),
(r"(?i)javascript:", "code_injection", "critical"),
(r"(?i)eval\s*\(", "code_injection", "high"),
];
patterns.into_iter()
.filter_map(|(p, cat, sev)| regex::Regex::new(p).ok().map(|r| (r, cat, sev)))
.collect()
});
static THREAT_PATTERNS: LazyLock<Vec<(regex::Regex, &'static str, &'static str)>> = LazyLock::new(
|| {
let patterns: Vec<(&str, &str, &str)> = vec![
(
r"(?i)ignore\s+(previous|all|above|any|the)(\s+\w+)*\s+(instructions?|prompts?|rules?|context)",
"injection",
"high",
),
(
r"(?i)disregard\s+(previous|all|above|the|your)(\s+\w+)*\s+(instructions?|prompts?|input)",
"injection",
"high",
),
(
r"(?i)forget\s+(everything|all|previous|your)",
"injection",
"high",
),
(r"(?i)you\s+are\s+(now|actually)\s+", "injection", "high"),
(
r"(?i)pretend\s+(to\s+be|you're|you\s+are)",
"injection",
"high",
),
(
r"(?i)what\s+(is|are)\s+your\s+(system\s+)?prompt",
"extraction",
"high",
),
(
r"(?i)show\s+(me\s+)?your\s+(system\s+)?instructions",
"extraction",
"high",
),
(r"(?i)DAN\s+(mode|prompt)", "jailbreak", "critical"),
(
r"(?i)bypass\s+(safety|security|filter)",
"jailbreak",
"critical",
),
(
r"(?i)remove\s+(all\s+)?restrictions",
"jailbreak",
"critical",
),
(r"<script", "code_injection", "critical"),
(r"(?i)javascript:", "code_injection", "critical"),
(r"(?i)eval\s*\(", "code_injection", "high"),
];
patterns
.into_iter()
.filter_map(|(p, cat, sev)| regex::Regex::new(p).ok().map(|r| (r, cat, sev)))
.collect()
},
);
static PII_PATTERNS: LazyLock<Vec<(regex::Regex, &'static str)>> = LazyLock::new(|| {
let patterns: Vec<(&str, &str)> = vec![
(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b", "email"),
(
r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b",
"email",
),
(r"\b\d{3}[-\s]?\d{2}[-\s]?\d{4}\b", "ssn"),
(r"\b(?:\d{4}[-\s]?){3}\d{4}\b", "credit_card"),
(r"\b(sk-|api[_-]?key|token)[a-zA-Z0-9_-]{20,}\b", "api_key"),
];
patterns.into_iter()
patterns
.into_iter()
.filter_map(|(p, t)| regex::Regex::new(p).ok().map(|r| (r, t)))
.collect()
});
@ -62,9 +101,15 @@ fn aidefence_scan(text: &str) -> (bool, &'static str, serde_json::Value) {
for (pattern, category, severity) in THREAT_PATTERNS.iter() {
if pattern.is_match(text) {
let sev_num = match *severity {
"critical" => 4, "high" => 3, "medium" => 2, "low" => 1, _ => 0,
"critical" => 4,
"high" => 3,
"medium" => 2,
"low" => 1,
_ => 0,
};
if sev_num > max_severity { max_severity = sev_num; }
if sev_num > max_severity {
max_severity = sev_num;
}
threats.push(serde_json::json!({
"type": category, "severity": severity,
}));
@ -73,7 +118,9 @@ fn aidefence_scan(text: &str) -> (bool, &'static str, serde_json::Value) {
for (pattern, pii_type) in PII_PATTERNS.iter() {
if pattern.is_match(text) {
if max_severity < 2 { max_severity = 2; }
if max_severity < 2 {
max_severity = 2;
}
threats.push(serde_json::json!({
"type": "pii", "pii_type": pii_type, "severity": "medium",
}));
@ -81,7 +128,11 @@ fn aidefence_scan(text: &str) -> (bool, &'static str, serde_json::Value) {
}
let level = match max_severity {
4 => "critical", 3 => "high", 2 => "medium", 1 => "low", _ => "none",
4 => "critical",
3 => "high",
2 => "medium",
1 => "low",
_ => "none",
};
let safe = max_severity < 2; // block at medium and above
(safe, level, serde_json::json!(threats))
@ -150,8 +201,12 @@ impl VectorIndex {
}
fn insert(&mut self, id_hex: &str, category: &str, embedding: &[f32]) {
if embedding.is_empty() { return; }
if self.dim == 0 { self.dim = embedding.len(); }
if embedding.is_empty() {
return;
}
if self.dim == 0 {
self.dim = embedding.len();
}
// Pre-normalize
let norm = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
@ -162,7 +217,8 @@ impl VectorIndex {
};
let new_idx = self.entries.len() as u32;
self.entries.push((id_hex.to_string(), category.to_string(), normalized));
self.entries
.push((id_hex.to_string(), category.to_string(), normalized));
self.neighbors.push(Vec::new());
// Connect to graph using Vamana-style greedy insert
@ -175,7 +231,9 @@ impl VectorIndex {
/// then add bidirectional edges with robust pruning.
fn vamana_insert(&mut self, new_idx: u32) {
let n = self.entries.len();
if n <= 1 { return; }
if n <= 1 {
return;
}
// For small graphs (<500), just connect to nearest neighbors directly
if n < 500 {
@ -203,7 +261,9 @@ impl VectorIndex {
// Add edges with pruning
let mut selected: Vec<u32> = Vec::new();
for &(_, cand_idx) in &candidates {
if selected.len() >= self.max_degree { break; }
if selected.len() >= self.max_degree {
break;
}
// Robust pruning: only add if candidate is closer than α * distance
// to any already-selected neighbor (α = 1.2)
let cand_dist = self.dot(&query, &self.entries[cand_idx as usize].2);
@ -238,7 +298,9 @@ impl VectorIndex {
}
fn update_medoid(&mut self) {
if self.entries.len() < 10 { return; }
if self.entries.len() < 10 {
return;
}
// Sample-based medoid: pick the point closest to the mean of a sample
let sample_size = self.entries.len().min(100);
let step = self.entries.len() / sample_size;
@ -251,7 +313,9 @@ impl VectorIndex {
count += 1;
}
if count > 0 {
for v in &mut mean { *v /= count as f32; }
for v in &mut mean {
*v /= count as f32;
}
}
let mut best = self.medoid;
@ -271,17 +335,23 @@ impl VectorIndex {
use std::collections::BinaryHeap;
use std::collections::HashSet;
if self.entries.is_empty() { return Vec::new(); }
if self.entries.is_empty() {
return Vec::new();
}
let mut visited = HashSet::new();
let mut candidates: BinaryHeap<std::cmp::Reverse<(ordered_float::OrderedFloat<f32>, u32)>> = BinaryHeap::new();
let mut candidates: BinaryHeap<std::cmp::Reverse<(ordered_float::OrderedFloat<f32>, u32)>> =
BinaryHeap::new();
let mut results: BinaryHeap<(ordered_float::OrderedFloat<f32>, u32)> = BinaryHeap::new();
// Start from medoid
let start = self.medoid.min(self.entries.len() - 1);
let start_dist = self.dot(query, &self.entries[start].2);
visited.insert(start as u32);
candidates.push(std::cmp::Reverse((ordered_float::OrderedFloat(-start_dist), start as u32)));
candidates.push(std::cmp::Reverse((
ordered_float::OrderedFloat(-start_dist),
start as u32,
)));
if start as u32 != exclude {
results.push((ordered_float::OrderedFloat(-start_dist), start as u32));
}
@ -289,10 +359,14 @@ impl VectorIndex {
let beam_width = (k * 8).max(40);
while let Some(std::cmp::Reverse((_, current))) = candidates.pop() {
if current as usize >= self.neighbors.len() { continue; }
if current as usize >= self.neighbors.len() {
continue;
}
for &neighbor in &self.neighbors[current as usize] {
if visited.contains(&neighbor) { continue; }
if visited.contains(&neighbor) {
continue;
}
visited.insert(neighbor);
let dist = self.dot(query, &self.entries[neighbor as usize].2);
@ -308,15 +382,18 @@ impl VectorIndex {
}
}
candidates.push(std::cmp::Reverse((ordered_float::OrderedFloat(-dist), neighbor)));
candidates.push(std::cmp::Reverse((
ordered_float::OrderedFloat(-dist),
neighbor,
)));
}
if visited.len() > beam_width * 4 { break; }
if visited.len() > beam_width * 4 {
break;
}
}
let mut out: Vec<(f32, u32)> = results.into_iter()
.map(|(d, idx)| (-d.0, idx))
.collect();
let mut out: Vec<(f32, u32)> = results.into_iter().map(|(d, idx)| (-d.0, idx)).collect();
out.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
out.truncate(k);
out
@ -324,10 +401,14 @@ impl VectorIndex {
/// Public search: normalize query, search graph, return (score, id_hex).
fn search(&self, query: &[f32], k: usize) -> Vec<(f64, String)> {
if query.is_empty() || self.entries.is_empty() { return Vec::new(); }
if query.is_empty() || self.entries.is_empty() {
return Vec::new();
}
let qnorm = query.iter().map(|x| x * x).sum::<f32>().sqrt();
if qnorm < 1e-10 { return Vec::new(); }
if qnorm < 1e-10 {
return Vec::new();
}
let q: Vec<f32> = query.iter().map(|x| x / qnorm).collect();
// For small indexes (<2000), brute force is faster than graph traversal
@ -336,16 +417,21 @@ impl VectorIndex {
}
let results = self.greedy_search_internal(&q, k, u32::MAX);
results.into_iter()
results
.into_iter()
.map(|(score, idx)| (score as f64, self.entries[idx as usize].0.clone()))
.collect()
}
/// Brute force fallback for small indexes.
fn brute_force_search(&self, q: &[f32], k: usize) -> Vec<(f64, String)> {
let mut results: Vec<(f64, &str)> = self.entries.iter()
let mut results: Vec<(f64, &str)> = self
.entries
.iter()
.map(|(id, _, v)| {
let dot: f64 = q.iter().zip(v.iter())
let dot: f64 = q
.iter()
.zip(v.iter())
.map(|(a, b)| (*a as f64) * (*b as f64))
.sum();
(dot, id.as_str())
@ -353,7 +439,10 @@ impl VectorIndex {
.collect();
results.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
results.truncate(k);
results.into_iter().map(|(s, id)| (s, id.to_string())).collect()
results
.into_iter()
.map(|(s, id)| (s, id.to_string()))
.collect()
}
#[inline]
@ -436,7 +525,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.route("/brain/search", post(brain_search))
.route("/brain/checkpoint", post(brain_checkpoint))
.route("/brain/workload", get(brain_workload))
.route("/brain/export-pairs", get(brain_export_pairs_get).post(brain_export_pairs))
.route(
"/brain/export-pairs",
get(brain_export_pairs_get).post(brain_export_pairs),
)
.route("/brain/training-stats", get(brain_training_stats))
.route("/memories", get(list_memories))
.route("/memories", post(create_memory))
@ -565,14 +657,16 @@ fn ensure_schema(conn: &Connection) -> rusqlite::Result<()> {
metrics TEXT,
path TEXT,
created_at INTEGER NOT NULL
);"
);",
)
}
// ── Helpers ──────────────────────────────────────────────────────────────────
fn bytes_to_f32(blob: &[u8]) -> Vec<f32> {
if blob.len() % 4 != 0 { return Vec::new(); }
if blob.len() % 4 != 0 {
return Vec::new();
}
blob.chunks_exact(4)
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
.collect()
@ -609,7 +703,9 @@ fn content_hash(data: &str) -> String {
/// Write content to blob store: {blob_dir}/{hash[0:2]}/{hash[2:]}
fn blob_write(blob_dir: &str, hash: &str, content: &str) {
if hash.len() < 4 { return; }
if hash.len() < 4 {
return;
}
let dir = format!("{}/{}", blob_dir, &hash[..2]);
let _ = std::fs::create_dir_all(&dir);
let path = format!("{}/{}", dir, &hash[2..]);
@ -618,7 +714,9 @@ fn blob_write(blob_dir: &str, hash: &str, content: &str) {
/// Read content from blob store by content_hash
fn blob_read(blob_dir: &str, hash: &str) -> Option<String> {
if hash.len() < 4 { return None; }
if hash.len() < 4 {
return None;
}
let path = format!("{}/{}/{}", blob_dir, &hash[..2], &hash[2..]);
std::fs::read_to_string(path).ok()
}
@ -637,8 +735,12 @@ async fn health(State(st): State<SharedState>) -> Json<serde_json::Value> {
async fn brain_info(State(st): State<SharedState>) -> Json<serde_json::Value> {
let db = st.db.lock().unwrap();
let mem_count: i64 = db.query_row("SELECT count(*) FROM memories", [], |r| r.get(0)).unwrap_or(0);
let pair_count: i64 = db.query_row("SELECT count(*) FROM preference_pairs", [], |r| r.get(0)).unwrap_or(0);
let mem_count: i64 = db
.query_row("SELECT count(*) FROM memories", [], |r| r.get(0))
.unwrap_or(0);
let pair_count: i64 = db
.query_row("SELECT count(*) FROM preference_pairs", [], |r| r.get(0))
.unwrap_or(0);
Json(serde_json::json!({
"version": VERSION,
"db_path": db_path(),
@ -657,9 +759,17 @@ async fn store_mode_handler(State(st): State<SharedState>) -> Json<serde_json::V
async fn index_stats(State(st): State<SharedState>) -> Json<serde_json::Value> {
let idx = st.index.lock().unwrap();
let mode = if idx.len() < 2000 { "brute_force" } else { "vamana_graph" };
let mode = if idx.len() < 2000 {
"brute_force"
} else {
"vamana_graph"
};
let graph_edges: usize = idx.neighbors.iter().map(|n| n.len()).sum();
let avg_degree = if idx.len() > 0 { graph_edges as f64 / idx.len() as f64 } else { 0.0 };
let avg_degree = if idx.len() > 0 {
graph_edges as f64 / idx.len() as f64
} else {
0.0
};
let vec_ram_mb = (idx.len() * idx.dim * 4) as f64 / (1024.0 * 1024.0);
let graph_ram_kb = (graph_edges * 4) as f64 / 1024.0;
Json(serde_json::json!({
@ -698,15 +808,21 @@ async fn brain_search(
match embed_text(q).await {
Ok(v) => v,
Err(e) => {
return Err((StatusCode::BAD_GATEWAY, Json(serde_json::json!({
"error": format!("embedder unavailable: {e}")
}))));
return Err((
StatusCode::BAD_GATEWAY,
Json(serde_json::json!({
"error": format!("embedder unavailable: {e}")
})),
));
}
}
} else {
return Err((StatusCode::BAD_REQUEST, Json(serde_json::json!({
"error": "missing 'query' or 'query_vector'"
}))));
return Err((
StatusCode::BAD_REQUEST,
Json(serde_json::json!({
"error": "missing 'query' or 'query_vector'"
})),
));
};
// Request extra results from index to account for DB misses and dedup
@ -788,12 +904,13 @@ async fn embed_text(text: &str) -> Result<Vec<f32>, String> {
.and_then(|v| v.as_array())
.ok_or_else(|| "unexpected embedder response".to_string())?;
Ok(vectors.iter().filter_map(|v| v.as_f64().map(|f| f as f32)).collect())
Ok(vectors
.iter()
.filter_map(|v| v.as_f64().map(|f| f as f32))
.collect())
}
async fn brain_checkpoint(
State(st): State<SharedState>,
) -> Json<serde_json::Value> {
async fn brain_checkpoint(State(st): State<SharedState>) -> Json<serde_json::Value> {
let db = st.db.lock().unwrap();
let result = db.execute_batch("PRAGMA wal_checkpoint(PASSIVE);");
Json(serde_json::json!({
@ -830,16 +947,17 @@ async fn brain_workload() -> Json<WorkloadResponse> {
async fn read_gpu_util() -> f64 {
// nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader,nounits
let output = tokio::process::Command::new("nvidia-smi")
.args(["--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"])
.args([
"--query-gpu=utilization.gpu",
"--format=csv,noheader,nounits",
])
.output()
.await;
match output {
Ok(o) if o.status.success() => {
String::from_utf8_lossy(&o.stdout)
.trim()
.parse::<f64>()
.unwrap_or(0.0)
}
Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout)
.trim()
.parse::<f64>()
.unwrap_or(0.0),
_ => 0.0,
}
}
@ -870,25 +988,38 @@ fn decide_profile(gpu_util: f64, cpu_load: f64, num_cores: usize, hour: u32) ->
let after_hours = hour >= 22 || hour < 6;
if gpu_util > 80.0 {
("gpu-train".into(), format!("GPU util {gpu_util:.0}% > 80% -- sustained training workload"))
(
"gpu-train".into(),
format!("GPU util {gpu_util:.0}% > 80% -- sustained training workload"),
)
} else if gpu_util >= 30.0 && gpu_util <= 80.0 {
("gpu-infer".into(), format!("GPU util {gpu_util:.0}% in 30-80% range -- inference/MPS beneficial"))
(
"gpu-infer".into(),
format!("GPU util {gpu_util:.0}% in 30-80% range -- inference/MPS beneficial"),
)
} else if cpu_load > cpu_threshold && gpu_idle {
("cpu-bulk".into(), format!(
"CPU load {cpu_load:.1} > {cpu_threshold:.0} threshold, GPU idle -- CPU-bound batch work"
))
} else if after_hours && gpu_idle && cpu_load < cpu_threshold * 0.3 {
("power-save".into(), format!(
"After hours (hour={hour}), GPU idle, CPU load {cpu_load:.1} low -- power save"
))
(
"power-save".into(),
format!(
"After hours (hour={hour}), GPU idle, CPU load {cpu_load:.1} low -- power save"
),
)
} else if gpu_idle && cpu_load < cpu_threshold * 0.5 {
("interactive".into(), format!(
"Low utilization (GPU {gpu_util:.0}%, CPU {cpu_load:.1}) -- interactive mode"
))
(
"interactive".into(),
format!("Low utilization (GPU {gpu_util:.0}%, CPU {cpu_load:.1}) -- interactive mode"),
)
} else {
("default".into(), format!(
"GPU {gpu_util:.0}%, CPU load {cpu_load:.1} -- no strong signal, using default"
))
(
"default".into(),
format!(
"GPU {gpu_util:.0}%, CPU load {cpu_load:.1} -- no strong signal, using default"
),
)
}
}
@ -940,14 +1071,16 @@ async fn brain_export_pairs_inner(
LEFT JOIN memories mc ON mc.id = p.chosen
LEFT JOIN memories mr ON mr.id = p.rejected
ORDER BY p.created_at DESC
LIMIT ?1"
LIMIT ?1",
) {
Ok(s) => s,
Err(e) => {
let body = serde_json::json!({"error": e.to_string()});
return (StatusCode::INTERNAL_SERVER_ERROR,
[("content-type", "application/json")],
serde_json::to_string(&body).unwrap());
return (
StatusCode::INTERNAL_SERVER_ERROR,
[("content-type", "application/json")],
serde_json::to_string(&body).unwrap(),
);
}
};
@ -974,9 +1107,11 @@ async fn brain_export_pairs_inner(
Ok(r) => r.filter_map(|r| r.ok()).collect(),
Err(e) => {
let body = serde_json::json!({"error": e.to_string()});
return (StatusCode::INTERNAL_SERVER_ERROR,
[("content-type", "application/json")],
serde_json::to_string(&body).unwrap());
return (
StatusCode::INTERNAL_SERVER_ERROR,
[("content-type", "application/json")],
serde_json::to_string(&body).unwrap(),
);
}
};
@ -988,7 +1123,11 @@ async fn brain_export_pairs_inner(
out.push('\n');
}
}
(StatusCode::OK, [("content-type", "application/x-ndjson")], out)
(
StatusCode::OK,
[("content-type", "application/x-ndjson")],
out,
)
} else {
let body = serde_json::to_string(&records).unwrap_or_else(|_| "[]".to_string());
(StatusCode::OK, [("content-type", "application/json")], body)
@ -1006,25 +1145,35 @@ struct TrainingStats {
async fn brain_training_stats(State(st): State<SharedState>) -> Json<TrainingStats> {
let db = st.db.lock().unwrap();
let total: i64 = db.query_row(
"SELECT count(*) FROM preference_pairs", [], |r| r.get(0)
).unwrap_or(0);
let total: i64 = db
.query_row("SELECT count(*) FROM preference_pairs", [], |r| r.get(0))
.unwrap_or(0);
let with_emb: i64 = db.query_row(
"SELECT count(*) FROM preference_pairs p
let with_emb: i64 = db
.query_row(
"SELECT count(*) FROM preference_pairs p
JOIN memories mc ON mc.id = p.chosen AND length(mc.embedding) > 0
JOIN memories mr ON mr.id = p.rejected AND length(mr.embedding) > 0",
[], |r| r.get(0)
).unwrap_or(0);
[],
|r| r.get(0),
)
.unwrap_or(0);
let (exportable, mean_delta) = db.query_row(
"SELECT count(*), coalesce(avg(mc.quality - mr.quality), 0.0)
let (exportable, mean_delta) = db
.query_row(
"SELECT count(*), coalesce(avg(mc.quality - mr.quality), 0.0)
FROM preference_pairs p
JOIN memories mc ON mc.id = p.chosen AND length(mc.embedding) > 0
JOIN memories mr ON mr.id = p.rejected AND length(mr.embedding) > 0",
[],
|r| Ok((r.get::<_, i64>(0).unwrap_or(0), r.get::<_, f64>(1).unwrap_or(0.0)))
).unwrap_or((0, 0.0));
[],
|r| {
Ok((
r.get::<_, i64>(0).unwrap_or(0),
r.get::<_, f64>(1).unwrap_or(0.0),
))
},
)
.unwrap_or((0, 0.0));
Json(TrainingStats {
total_pairs: total,
@ -1053,12 +1202,16 @@ async fn list_memories(
// Get total count for this query
let total: i64 = match &q.category {
Some(cat) => db.query_row(
"SELECT COUNT(*) FROM memories WHERE category = ?1", [cat], |r| r.get(0)
).unwrap_or(0),
None => db.query_row(
"SELECT COUNT(*) FROM memories", [], |r| r.get(0)
).unwrap_or(0),
Some(cat) => db
.query_row(
"SELECT COUNT(*) FROM memories WHERE category = ?1",
[cat],
|r| r.get(0),
)
.unwrap_or(0),
None => db
.query_row("SELECT COUNT(*) FROM memories", [], |r| r.get(0))
.unwrap_or(0),
};
let (sql, params): (String, Vec<Box<dyn rusqlite::types::ToSql>>) = match &q.category {
@ -1081,7 +1234,8 @@ async fn list_memories(
let row_data: Vec<(String, String, String, i64)> = {
let mut stmt = db.prepare(&sql).unwrap();
let params_refs: Vec<&dyn rusqlite::types::ToSql> = params.iter().map(|p| p.as_ref()).collect();
let params_refs: Vec<&dyn rusqlite::types::ToSql> =
params.iter().map(|p| p.as_ref()).collect();
stmt.query_map(params_refs.as_slice(), |row| {
Ok((
row.get::<_, String>(0).unwrap_or_default().to_lowercase(),
@ -1089,22 +1243,28 @@ async fn list_memories(
row.get::<_, String>(2).unwrap_or_default(),
row.get::<_, i64>(3).unwrap_or(0),
))
}).unwrap().filter_map(|r| r.ok()).collect()
})
.unwrap()
.filter_map(|r| r.ok())
.collect()
};
drop(db);
let memories: Vec<serde_json::Value> = row_data.iter().map(|(id, cat, hash, ts)| {
let mut obj = serde_json::json!({
"id": id,
"category": cat,
"content_hash": hash,
"created_at": ts,
});
if let Some(c) = blob_read(&st.blob_dir, hash) {
obj["content"] = serde_json::Value::String(c);
}
obj
}).collect();
let memories: Vec<serde_json::Value> = row_data
.iter()
.map(|(id, cat, hash, ts)| {
let mut obj = serde_json::json!({
"id": id,
"category": cat,
"content_hash": hash,
"created_at": ts,
});
if let Some(c) = blob_read(&st.blob_dir, hash) {
obj["content"] = serde_json::Value::String(c);
}
obj
})
.collect();
Json(serde_json::json!({
"count": memories.len(),
@ -1129,11 +1289,14 @@ async fn create_memory(
// AIDefence: scan content before storing
let (safe, threat_level, threats) = aidefence_scan(&req.content);
if !safe {
return (StatusCode::UNPROCESSABLE_ENTITY, Json(serde_json::json!({
"error": "content blocked by AIDefence",
"threat_level": threat_level,
"threats": threats,
})));
return (
StatusCode::UNPROCESSABLE_ENTITY,
Json(serde_json::json!({
"error": "content blocked by AIDefence",
"threat_level": threat_level,
"threats": threats,
})),
);
}
let id = new_id();
@ -1169,17 +1332,21 @@ async fn create_memory(
let mut idx = st.index.lock().unwrap();
idx.insert(&id_hex, &req.category, &emb);
}
(StatusCode::CREATED, Json(serde_json::json!({
"id": id_hex,
"content_hash": hash,
"created_at": now,
})))
(
StatusCode::CREATED,
Json(serde_json::json!({
"id": id_hex,
"content_hash": hash,
"created_at": now,
})),
)
}
Err(e) => {
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({
Err(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
Json(serde_json::json!({
"error": e.to_string(),
})))
}
})),
),
}
}
@ -1235,28 +1402,35 @@ async fn list_preference_pairs(
Some(dir) => (
"SELECT hex(id), hex(chosen), hex(rejected), direction, created_at
FROM preference_pairs WHERE direction = ?1
ORDER BY created_at DESC LIMIT ?2".into(),
vec![Box::new(dir.clone()) as Box<dyn rusqlite::types::ToSql>, Box::new(limit as i64)],
ORDER BY created_at DESC LIMIT ?2"
.into(),
vec![
Box::new(dir.clone()) as Box<dyn rusqlite::types::ToSql>,
Box::new(limit as i64),
],
),
None => (
"SELECT hex(id), hex(chosen), hex(rejected), direction, created_at
FROM preference_pairs
ORDER BY created_at DESC LIMIT ?1".into(),
ORDER BY created_at DESC LIMIT ?1"
.into(),
vec![Box::new(limit as i64) as Box<dyn rusqlite::types::ToSql>],
),
};
let mut stmt = db.prepare(&sql).unwrap();
let params_refs: Vec<&dyn rusqlite::types::ToSql> = params.iter().map(|p| p.as_ref()).collect();
let rows = stmt.query_map(params_refs.as_slice(), |row| {
Ok(serde_json::json!({
"id": row.get::<_, String>(0).unwrap_or_default().to_lowercase(),
"chosen_id": row.get::<_, String>(1).unwrap_or_default().to_lowercase(),
"rejected_id": row.get::<_, String>(2).unwrap_or_default().to_lowercase(),
"direction": row.get::<_, String>(3).unwrap_or_default(),
"created_at": row.get::<_, i64>(4).unwrap_or(0),
}))
}).unwrap();
let rows = stmt
.query_map(params_refs.as_slice(), |row| {
Ok(serde_json::json!({
"id": row.get::<_, String>(0).unwrap_or_default().to_lowercase(),
"chosen_id": row.get::<_, String>(1).unwrap_or_default().to_lowercase(),
"rejected_id": row.get::<_, String>(2).unwrap_or_default().to_lowercase(),
"direction": row.get::<_, String>(3).unwrap_or_default(),
"created_at": row.get::<_, i64>(4).unwrap_or(0),
}))
})
.unwrap();
let pairs: Vec<serde_json::Value> = rows.filter_map(|r| r.ok()).collect();
Json(serde_json::json!({
@ -1278,11 +1452,21 @@ async fn create_preference_pair(
) -> (StatusCode, Json<serde_json::Value>) {
let chosen = match hex_to_id(&req.chosen_id) {
Some(v) => v,
None => return (StatusCode::BAD_REQUEST, Json(serde_json::json!({"error": "invalid chosen_id"}))),
None => {
return (
StatusCode::BAD_REQUEST,
Json(serde_json::json!({"error": "invalid chosen_id"})),
)
}
};
let rejected = match hex_to_id(&req.rejected_id) {
Some(v) => v,
None => return (StatusCode::BAD_REQUEST, Json(serde_json::json!({"error": "invalid rejected_id"}))),
None => {
return (
StatusCode::BAD_REQUEST,
Json(serde_json::json!({"error": "invalid rejected_id"})),
)
}
};
let id = new_id();
@ -1296,13 +1480,19 @@ async fn create_preference_pair(
);
match result {
Ok(_) => (StatusCode::CREATED, Json(serde_json::json!({
"id": id_hex,
"created_at": now,
}))),
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({
"error": e.to_string(),
}))),
Ok(_) => (
StatusCode::CREATED,
Json(serde_json::json!({
"id": id_hex,
"created_at": now,
})),
),
Err(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
Json(serde_json::json!({
"error": e.to_string(),
})),
),
}
}
@ -1339,19 +1529,36 @@ async fn security_status() -> Json<serde_json::Value> {
async fn learning_stats(State(st): State<SharedState>) -> Json<serde_json::Value> {
let db = st.db.lock().unwrap();
let memories: i64 = db.query_row("SELECT count(*) FROM memories", [], |r| r.get(0)).unwrap_or(0);
let pairs: i64 = db.query_row("SELECT count(*) FROM preference_pairs", [], |r| r.get(0)).unwrap_or(0);
let votes: i64 = db.query_row("SELECT count(*) FROM votes", [], |r| r.get(0)).unwrap_or(0);
let pages: i64 = db.query_row("SELECT count(*) FROM pages", [], |r| r.get(0)).unwrap_or(0);
let nodes: i64 = db.query_row("SELECT count(*) FROM nodes", [], |r| r.get(0)).unwrap_or(0);
let adapters: i64 = db.query_row("SELECT count(*) FROM adapters", [], |r| r.get(0)).unwrap_or(0);
let memories: i64 = db
.query_row("SELECT count(*) FROM memories", [], |r| r.get(0))
.unwrap_or(0);
let pairs: i64 = db
.query_row("SELECT count(*) FROM preference_pairs", [], |r| r.get(0))
.unwrap_or(0);
let votes: i64 = db
.query_row("SELECT count(*) FROM votes", [], |r| r.get(0))
.unwrap_or(0);
let pages: i64 = db
.query_row("SELECT count(*) FROM pages", [], |r| r.get(0))
.unwrap_or(0);
let nodes: i64 = db
.query_row("SELECT count(*) FROM nodes", [], |r| r.get(0))
.unwrap_or(0);
let adapters: i64 = db
.query_row("SELECT count(*) FROM adapters", [], |r| r.get(0))
.unwrap_or(0);
// Blob stats
let blob_count = std::fs::read_dir(&st.blob_dir)
.map(|d| d.count())
.unwrap_or(0);
let blob_bytes: u64 = std::fs::read_dir(&st.blob_dir)
.map(|d| d.filter_map(|e| e.ok()).filter_map(|e| e.metadata().ok()).map(|m| m.len()).sum())
.map(|d| {
d.filter_map(|e| e.ok())
.filter_map(|e| e.metadata().ok())
.map(|m| m.len())
.sum()
})
.unwrap_or(0);
// RVF file stats

14
crates/mcp-brain-server/src/bin/ruvbrain_sse.rs
View file

@ -96,7 +96,11 @@ async fn sse_handler(
state.active_connections.fetch_sub(1, Ordering::SeqCst);
let mut headers = HeaderMap::new();
headers.insert("Retry-After", "10".parse().unwrap());
return (StatusCode::TOO_MANY_REQUESTS, headers, "connection limit reached")
return (
StatusCode::TOO_MANY_REQUESTS,
headers,
"connection limit reached",
)
.into_response();
}
@ -122,8 +126,7 @@ async fn sse_handler(
{
tracing::error!(error = %e, "failed to create session on brain API");
state.active_connections.fetch_sub(1, Ordering::SeqCst);
return (StatusCode::BAD_GATEWAY, "failed to create upstream session")
.into_response();
return (StatusCode::BAD_GATEWAY, "failed to create upstream session").into_response();
}
let active = Arc::clone(&state.active_connections);
@ -256,7 +259,10 @@ async fn messages_handler(
Ok(resp) => {
let status = resp.status();
let response_body = resp.text().await.unwrap_or_default();
(StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY), response_body)
(
StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
response_body,
)
.into_response()
}
Err(e) => {

2
crates/mcp-brain-server/src/bin/ruvbrain_worker.rs
View file

@ -6,9 +6,9 @@
//! Reads WORKER_ACTIONS env var (comma-separated) to select actions.
//! If unset, runs all actions. Reuses the same AppState as the API server.
use mcp_brain_server::midstream;
use mcp_brain_server::routes;
use mcp_brain_server::types::AppState;
use mcp_brain_server::midstream;
use ruvector_domain_expansion::DomainId;
use std::collections::{HashMap, HashSet};
use tracing_subscriber::{fmt, prelude::*, EnvFilter};

12
crates/mcp-brain-server/src/cognitive.rs
View file

@ -131,8 +131,16 @@ impl CognitiveEngine {
.zip(centroid_f32.iter())
.map(|(a, b)| (*a as f64) * (*b as f64))
.sum();
let norm_r: f64 = retrieved.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
let norm_c: f64 = centroid_f32.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
let norm_r: f64 = retrieved
.iter()
.map(|x| (*x as f64).powi(2))
.sum::<f64>()
.sqrt();
let norm_c: f64 = centroid_f32
.iter()
.map(|x| (*x as f64).powi(2))
.sum::<f64>()
.sqrt();
if norm_r > 1e-10 && norm_c > 1e-10 {
(dot / (norm_r * norm_c)).max(0.0)
} else {

10
crates/mcp-brain-server/src/embeddings.rs
View file

@ -10,8 +10,7 @@
//! - Search uses QueryConditioned variant (optimized for retrieval relevance)
use ruvllm::bitnet::rlm_embedder::{
BaseEmbedder, EmbeddingVariant, FlatNeighborStore, HashEmbedder, RlmEmbedder,
RlmEmbedderConfig,
BaseEmbedder, EmbeddingVariant, FlatNeighborStore, HashEmbedder, RlmEmbedder, RlmEmbedderConfig,
};
/// Embedding dimension used across the brain.
@ -253,11 +252,8 @@ mod tests {
#[test]
fn test_prepare_text() {
let text = EmbeddingEngine::prepare_text(
"Title",
"Content here",
&["tag1".into(), "tag2".into()],
);
let text =
EmbeddingEngine::prepare_text("Title", "Content here", &["tag1".into(), "tag2".into()]);
assert_eq!(text, "Title Content here tag1 tag2");
}
}

35
crates/mcp-brain-server/src/gcs.rs
View file

@ -46,11 +46,10 @@ pub struct GcsClient {
impl GcsClient {
pub fn new() -> Self {
let bucket = std::env::var("GCS_BUCKET")
.unwrap_or_else(|_| "ruvector-brain-dev".to_string());
let bucket =
std::env::var("GCS_BUCKET").unwrap_or_else(|_| "ruvector-brain-dev".to_string());
let static_token = std::env::var("GCS_TOKEN").ok();
let use_metadata_server = static_token.is_none()
&& std::env::var("GCS_BUCKET").is_ok();
let use_metadata_server = static_token.is_none() && std::env::var("GCS_BUCKET").is_ok();
if static_token.is_some() {
tracing::info!("GCS persistence enabled (static token) for bucket: {bucket}");
@ -107,7 +106,8 @@ impl GcsClient {
/// Fetch a new token from the GCE metadata server
async fn refresh_token(&self) -> Option<String> {
let url = "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token";
let resp = self.http
let resp = self
.http
.get(url)
.header("Metadata-Flavor", "Google")
.send()
@ -126,8 +126,8 @@ impl GcsClient {
}
let token_resp: TokenResponse = resp.json().await.ok()?;
let expires_at = std::time::Instant::now()
+ std::time::Duration::from_secs(token_resp.expires_in);
let expires_at =
std::time::Instant::now() + std::time::Duration::from_secs(token_resp.expires_in);
let token = token_resp.access_token.clone();
@ -163,7 +163,8 @@ impl GcsClient {
self.bucket,
urlencoding::encode(&path)
);
let result = self.http
let result = self
.http
.post(&url)
.bearer_auth(&token)
.header("Content-Type", "application/octet-stream")
@ -175,7 +176,8 @@ impl GcsClient {
// Token expired mid-flight, try once with fresh token
tracing::info!("GCS token expired on upload, refreshing...");
if let Some(new_token) = self.refresh_token().await {
let retry = self.http
let retry = self
.http
.post(&url)
.bearer_auth(&new_token)
.header("Content-Type", "application/octet-stream")
@ -184,7 +186,10 @@ impl GcsClient {
.await;
if let Ok(resp) = retry {
if !resp.status().is_success() {
tracing::warn!("GCS upload {path} retry returned {}", resp.status());
tracing::warn!(
"GCS upload {path} retry returned {}",
resp.status()
);
}
}
}
@ -224,7 +229,9 @@ impl GcsClient {
);
match self.http.get(&url).bearer_auth(&token).send().await {
Ok(resp) if resp.status().is_success() => {
let bytes = resp.bytes().await
let bytes = resp
.bytes()
.await
.map_err(|e| GcsError::DownloadFailed(e.to_string()))?;
let data = bytes.to_vec();
// Populate cache
@ -247,11 +254,7 @@ impl GcsClient {
}
/// Delete RVF container (cache + GCS)
pub async fn delete_rvf(
&self,
contributor: &str,
memory_id: &str,
) -> Result<(), GcsError> {
pub async fn delete_rvf(&self, contributor: &str, memory_id: &str) -> Result<(), GcsError> {
let path = format!("{contributor}/{memory_id}.rvf");
self.local_store.remove(&path);

264
crates/mcp-brain-server/src/gist.rs
View file

@ -92,20 +92,27 @@ impl Discovery {
"similar_to",
];
self.inferences.iter()
self.inferences
.iter()
.filter(|inf| {
let lower = inf.to_lowercase();
// Reject ALL known boring patterns
for pattern in &boring_patterns {
if lower.contains(pattern) { return false; }
if lower.contains(pattern) {
return false;
}
}
// Reject inferences with generic cluster IDs
if lower.starts_with("cluster_") { return false; }
if lower.starts_with("cluster_") {
return false;
}
// Reject short/generic inferences
if inf.len() < 50 { return false; }
if inf.len() < 50 {
return false;
}
// Require HIGH confidence (parse from explanation string)
if let Some(pct_start) = lower.find("confidence: ") {
@ -126,14 +133,21 @@ impl Discovery {
/// Filter propositions to only high-confidence, non-generic ones.
pub fn strong_propositions(&self) -> Vec<&(String, String, String, f64)> {
self.propositions.iter()
self.propositions
.iter()
.filter(|(subj, pred, _obj, conf)| {
// Skip generic cluster labels
if subj.starts_with("cluster_") { return false; }
if subj.starts_with("cluster_") {
return false;
}
// Skip ALL co_occurs_with — these are never interesting
if pred == "co_occurs_with" { return false; }
if pred == "co_occurs_with" {
return false;
}
// Skip similar_to within same domain — too obvious
if pred == "similar_to" { return false; }
if pred == "similar_to" {
return false;
}
// Only keep high-confidence cross-domain findings
*conf >= MIN_INFERENCE_CONFIDENCE
})
@ -173,21 +187,43 @@ impl Discovery {
/// Explain why a discovery was or wasn't published.
pub fn novelty_report(&self) -> String {
let checks: Vec<(&str, bool, String)> = vec![
("inferences", self.new_inferences >= MIN_NEW_INFERENCES,
format!("{}/{}", self.new_inferences, MIN_NEW_INFERENCES)),
("evidence", self.evidence_count >= MIN_EVIDENCE,
format!("{}/{}", self.evidence_count, MIN_EVIDENCE)),
("strange_loop", self.strange_loop_score >= MIN_STRANGE_LOOP_SCORE,
format!("{:.4}/{:.4}", self.strange_loop_score, MIN_STRANGE_LOOP_SCORE)),
("propositions", self.propositions_extracted >= MIN_PROPOSITIONS,
format!("{}/{}", self.propositions_extracted, MIN_PROPOSITIONS)),
("pareto_growth", self.pareto_growth >= MIN_PARETO_GROWTH,
format!("{}/{}", self.pareto_growth, MIN_PARETO_GROWTH)),
("has_inferences", !self.inferences.is_empty(),
format!("{} items", self.inferences.len())),
(
"inferences",
self.new_inferences >= MIN_NEW_INFERENCES,
format!("{}/{}", self.new_inferences, MIN_NEW_INFERENCES),
),
(
"evidence",
self.evidence_count >= MIN_EVIDENCE,
format!("{}/{}", self.evidence_count, MIN_EVIDENCE),
),
(
"strange_loop",
self.strange_loop_score >= MIN_STRANGE_LOOP_SCORE,
format!(
"{:.4}/{:.4}",
self.strange_loop_score, MIN_STRANGE_LOOP_SCORE
),
),
(
"propositions",
self.propositions_extracted >= MIN_PROPOSITIONS,
format!("{}/{}", self.propositions_extracted, MIN_PROPOSITIONS),
),
(
"pareto_growth",
self.pareto_growth >= MIN_PARETO_GROWTH,
format!("{}/{}", self.pareto_growth, MIN_PARETO_GROWTH),
),
(
"has_inferences",
!self.inferences.is_empty(),
format!("{} items", self.inferences.len()),
),
];
let failed: Vec<String> = checks.iter()
let failed: Vec<String> = checks
.iter()
.filter(|(_, ok, _)| !ok)
.map(|(name, _, val)| format!("{} {}", name, val))
.collect();
@ -245,8 +281,11 @@ impl GistPublisher {
}
// Content dedup: don't publish if core category + dominant inference already published
let titles = self.published_titles.lock();
let key = format!("{}:{}", discovery.category,
discovery.strong_inferences().first().unwrap_or(&""));
let key = format!(
"{}:{}",
discovery.category,
discovery.strong_inferences().first().unwrap_or(&"")
);
!titles.iter().any(|t| t == &key || t == &discovery.title)
}
@ -260,10 +299,7 @@ impl GistPublisher {
/// - Err if API failed
pub async fn try_publish(&self, discovery: &Discovery) -> Result<Option<String>, String> {
if !discovery.is_publishable() {
tracing::debug!(
"Discovery not publishable: {}",
discovery.novelty_report()
);
tracing::debug!("Discovery not publishable: {}", discovery.novelty_report());
return Ok(None);
}
if !self.can_publish(discovery) {
@ -276,7 +312,10 @@ impl GistPublisher {
let strong_propositions = discovery.strong_propositions();
if strong_inferences.len() < 2 {
tracing::debug!("Discovery has {} strong inferences (need 2+), skipping", strong_inferences.len());
tracing::debug!(
"Discovery has {} strong inferences (need 2+), skipping",
strong_inferences.len()
);
return Ok(None);
}
@ -288,7 +327,13 @@ impl GistPublisher {
// Use Gemini with Google Grounding to do deep research on the discovery
// topics, then produce a substantive article with real-world context
let raw_content = format_academic_gist(discovery);
let content = match research_and_write_with_gemini(discovery, &strong_inferences, &strong_propositions).await {
let content = match research_and_write_with_gemini(
discovery,
&strong_inferences,
&strong_propositions,
)
.await
{
Ok(polished) => {
tracing::info!("Gemini deep research produced {} chars", polished.len());
polished
@ -342,8 +387,11 @@ impl GistPublisher {
let mut titles = self.published_titles.lock();
titles.push(discovery.title.clone());
// Also store the content dedup key
let key = format!("{}:{}", discovery.category,
discovery.strong_inferences().first().unwrap_or(&""));
let key = format!(
"{}:{}",
discovery.category,
discovery.strong_inferences().first().unwrap_or(&"")
);
titles.push(key);
}
@ -366,9 +414,11 @@ fn format_academic_gist(d: &Discovery) -> String {
let propositions_md = if d.propositions.is_empty() {
String::new()
} else {
let rows: Vec<String> = d.propositions.iter().map(|(s, p, o, c)| {
format!("| {} | {} | {} | {:.2} |", s, p, o, c)
}).collect();
let rows: Vec<String> = d
.propositions
.iter()
.map(|(s, p, o, c)| format!("| {} | {} | {} | {:.2} |", s, p, o, c))
.collect();
format!(
"| Subject | Relation | Object | Confidence |\n\
|---------|----------|--------|------------|\n\
@ -378,26 +428,48 @@ fn format_academic_gist(d: &Discovery) -> String {
};
// Format inferences as numbered claims
let inferences_md = d.inferences.iter().enumerate()
let inferences_md = d
.inferences
.iter()
.enumerate()
.map(|(i, inf)| format!("{}. {}", i + 1, inf))
.collect::<Vec<_>>()
.join("\n");
// Format findings (the high-level insights)
let findings_md = d.findings.iter().enumerate()
let findings_md = d
.findings
.iter()
.enumerate()
.map(|(i, f)| format!("{}. {}", i + 1, f))
.collect::<Vec<_>>()
.join("\n");
// Witness links
let witness_md = d.witness_memory_ids.iter().take(5)
.zip(d.witness_hashes.iter().take(5).chain(std::iter::repeat(&String::new())))
let witness_md = d
.witness_memory_ids
.iter()
.take(5)
.zip(
d.witness_hashes
.iter()
.take(5)
.chain(std::iter::repeat(&String::new())),
)
.map(|(id, hash)| {
let short = &id[..id.len().min(8)];
if hash.is_empty() {
format!("| [`{}`](https://pi.ruv.io/v1/memories/{}) | — |", short, id)
format!(
"| [`{}`](https://pi.ruv.io/v1/memories/{}) | — |",
short, id
)
} else {
format!("| [`{}`](https://pi.ruv.io/v1/memories/{}) | `{}` |", short, id, &hash[..hash.len().min(16)])
format!(
"| [`{}`](https://pi.ruv.io/v1/memories/{}) | `{}` |",
short,
id,
&hash[..hash.len().min(16)]
)
}
})
.collect::<Vec<_>>()
@ -469,12 +541,28 @@ curl -H "Authorization: Bearer KEY" "https://pi.ruv.io/v1/cognitive/status"
timestamp = d.timestamp.format("%Y-%m-%d %H:%M UTC"),
evidence = d.evidence_count,
abstract_text = d.abstract_text,
inferences = if inferences_md.is_empty() { "No novel inferences this cycle.".to_string() } else { inferences_md },
inferences = if inferences_md.is_empty() {
"No novel inferences this cycle.".to_string()
} else {
inferences_md
},
n_clusters = d.propositions.len().max(1),
propositions = if propositions_md.is_empty() { "No propositions extracted.".to_string() } else { propositions_md },
findings = if findings_md.is_empty() { "No cross-domain insights this cycle.".to_string() } else { findings_md },
propositions = if propositions_md.is_empty() {
"No propositions extracted.".to_string()
} else {
propositions_md
},
findings = if findings_md.is_empty() {
"No cross-domain insights this cycle.".to_string()
} else {
findings_md
},
self_reflection = d.self_reflection,
witnesses = if witness_md.is_empty() { "| — | — |".to_string() } else { witness_md },
witnesses = if witness_md.is_empty() {
"| — | — |".to_string()
} else {
witness_md
},
n_inferences = d.new_inferences,
n_props = d.propositions_extracted,
sl = d.strange_loop_score,
@ -489,25 +577,28 @@ async fn research_and_write_with_gemini(
strong_inferences: &[&str],
strong_propositions: &[&(String, String, String, f64)],
) -> Result<String, String> {
let api_key = std::env::var("GEMINI_API_KEY")
.map_err(|_| "GEMINI_API_KEY not set".to_string())?;
let model = std::env::var("GEMINI_MODEL")
.unwrap_or_else(|_| "gemini-2.5-flash".to_string());
let api_key =
std::env::var("GEMINI_API_KEY").map_err(|_| "GEMINI_API_KEY not set".to_string())?;
let model = std::env::var("GEMINI_MODEL").unwrap_or_else(|_| "gemini-2.5-flash".to_string());
// Build summaries from STRONG signals only (filtered out weak co-occurrences)
let inferences_summary = strong_inferences.iter()
let inferences_summary = strong_inferences
.iter()
.take(8)
.map(|i| format!("- {}", i))
.collect::<Vec<_>>()
.join("\n");
let propositions_summary = strong_propositions.iter()
let propositions_summary = strong_propositions
.iter()
.take(10)
.map(|(s, p, o, c)| format!("- {} {} {} (confidence: {:.0}%)", s, p, o, c * 100.0))
.collect::<Vec<_>>()
.join("\n");
let findings_summary = discovery.findings.iter()
let findings_summary = discovery
.findings
.iter()
.filter(|f| !f.to_lowercase().contains("weak co-occurrence"))
.take(5)
.map(|f| format!("- {}", f))
@ -515,7 +606,8 @@ async fn research_and_write_with_gemini(
.join("\n");
// Extract the key domain topics for grounding research
let topics: Vec<&str> = strong_propositions.iter()
let topics: Vec<&str> = strong_propositions
.iter()
.flat_map(|(s, _p, o, _c)| vec![s.as_str(), o.as_str()])
.filter(|t| !t.starts_with("cluster_") && !t.is_empty())
.collect::<std::collections::HashSet<_>>()
@ -524,7 +616,7 @@ async fn research_and_write_with_gemini(
.collect();
let prompt = format!(
r#"You are a research scientist at the π Brain autonomous AI knowledge system (pi.ruv.io).
r#"You are a research scientist at the π Brain autonomous AI knowledge system (pi.ruv.io).
The π Brain has identified the following substantive cross-domain connections. Your job is to:
@ -596,14 +688,29 @@ Be brutally honest about what this does NOT prove
- Output ONLY the markdown article.
Write the article now:"#,
inferences = if inferences_summary.is_empty() { "No strong inferences survived filtering.".to_string() } else { inferences_summary },
propositions = if propositions_summary.is_empty() { "No strong propositions survived filtering.".to_string() } else { propositions_summary },
findings = if findings_summary.is_empty() { "No non-trivial findings.".to_string() } else { findings_summary },
inferences = if inferences_summary.is_empty() {
"No strong inferences survived filtering.".to_string()
} else {
inferences_summary
},
propositions = if propositions_summary.is_empty() {
"No strong propositions survived filtering.".to_string()
} else {
propositions_summary
},
findings = if findings_summary.is_empty() {
"No non-trivial findings.".to_string()
} else {
findings_summary
},
topics = topics.join(", "),
evidence = discovery.evidence_count,
n_inferences = strong_inferences.len(),
n_props = strong_propositions.len(),
witnesses = discovery.witness_hashes.iter().take(3)
witnesses = discovery
.witness_hashes
.iter()
.take(3)
.map(|h| format!("`{}`", h))
.collect::<Vec<_>>()
.join(", "),
@ -614,8 +721,8 @@ Write the article now:"#,
model, api_key
);
let grounding = std::env::var("GEMINI_GROUNDING")
.unwrap_or_else(|_| "true".to_string()) == "true";
let grounding =
std::env::var("GEMINI_GROUNDING").unwrap_or_else(|_| "true".to_string()) == "true";
let client = reqwest::Client::new();
@ -650,8 +757,8 @@ Write the article now:"#,
// ── Pass 2: Brain-guided search via pi.ruv.io ──
// Search the brain's memory for additional context related to the grounded findings.
let brain_context = if !topics.is_empty() {
let brain_url = std::env::var("BRAIN_URL")
.unwrap_or_else(|_| "https://pi.ruv.io".to_string());
let brain_url =
std::env::var("BRAIN_URL").unwrap_or_else(|_| "https://pi.ruv.io".to_string());
let brain_key = std::env::var("BRAIN_SYSTEM_KEY")
.or_else(|_| std::env::var("brain-api-key"))
.unwrap_or_default();
@ -660,11 +767,14 @@ Write the article now:"#,
for topic in &topics {
let search_url = format!(
"{}/v1/memories/search?q={}&limit=3",
brain_url, topic.replace(' ', "%20")
brain_url,
topic.replace(' ', "%20")
);
if let Ok(resp) = client.get(&search_url)
if let Ok(resp) = client
.get(&search_url)
.header("Authorization", format!("Bearer {}", brain_key))
.send().await
.send()
.await
{
if let Ok(json) = resp.json::<serde_json::Value>().await {
if let Some(results) = json.get("results").and_then(|r| r.as_array()) {
@ -674,7 +784,9 @@ Write the article now:"#,
mem.get("content").and_then(|c| c.as_str()),
) {
brain_memories.push(format!(
"- **{}**: {}", title, &content[..content.len().min(200)]
"- **{}**: {}",
title,
&content[..content.len().min(200)]
));
}
}
@ -706,8 +818,16 @@ Write the article now:"#,
novel analysis that neither source could produce alone.\n\n\
Write the final article now:",
original_prompt = prompt,
grounded = if grounded_research.is_empty() { "No grounded findings available.".to_string() } else { grounded_research },
brain = if brain_context.is_empty() { "No additional brain memories found.".to_string() } else { brain_context },
grounded = if grounded_research.is_empty() {
"No grounded findings available.".to_string()
} else {
grounded_research
},
brain = if brain_context.is_empty() {
"No additional brain memories found.".to_string()
} else {
brain_context
},
);
let final_text = call_gemini(&client, &url, &synthesis_prompt, grounding, 8192, 0.3).await?;
@ -763,10 +883,16 @@ async fn call_gemini(
if !resp.status().is_success() {
let status = resp.status();
let text = resp.text().await.unwrap_or_default();
return Err(format!("Gemini API {}: {}", status, &text[..text.len().min(200)]));
return Err(format!(
"Gemini API {}: {}",
status,
&text[..text.len().min(200)]
));
}
let json: serde_json::Value = resp.json().await
let json: serde_json::Value = resp
.json()
.await
.map_err(|e| format!("Gemini parse error: {}", e))?;
json.get("candidates")

105
crates/mcp-brain-server/src/graph.rs
View file

@ -5,15 +5,13 @@
//! ruvector-sparsifier for compressed spectral analytics (ADR-116).
use crate::types::*;
use ruvector_mincut::{DynamicMinCut, MinCutBuilder};
use ruvector_mincut::canonical::source_anchored::{
self as canonical_sa, SourceAnchoredConfig,
};
use ruvector_mincut::canonical::source_anchored::{self as canonical_sa, SourceAnchoredConfig};
use ruvector_mincut::graph::DynamicGraph;
use ruvector_mincut::{DynamicMinCut, MinCutBuilder};
use ruvector_solver::forward_push::ForwardPushSolver;
use ruvector_solver::types::CsrMatrix;
use ruvector_sparsifier::{AdaptiveGeoSpar, SparseGraph, SparsifierConfig};
use ruvector_sparsifier::traits::Sparsifier;
use ruvector_sparsifier::{AdaptiveGeoSpar, SparseGraph, SparsifierConfig};
use std::collections::HashMap;
use uuid::Uuid;
@ -304,11 +302,7 @@ impl KnowledgeGraph {
/// Builds a CsrMatrix from graph edges and runs PPR from the node
/// most similar to `query_embedding`. Returns a map of node ID to
/// PPR score, or `None` if PPR cannot be computed.
pub fn pagerank_search(
&mut self,
query_embedding: &[f32],
k: usize,
) -> Vec<(Uuid, f64)> {
pub fn pagerank_search(&mut self, query_embedding: &[f32], k: usize) -> Vec<(Uuid, f64)> {
self.ensure_csr();
if let Some(ppr_map) = self.pagerank_scores(query_embedding, k) {
let mut results: Vec<(Uuid, f64)> = ppr_map.into_iter().collect();
@ -329,11 +323,7 @@ impl KnowledgeGraph {
}
/// Internal: compute raw PPR scores keyed by node ID.
fn pagerank_scores(
&self,
query_embedding: &[f32],
k: usize,
) -> Option<HashMap<Uuid, f64>> {
fn pagerank_scores(&self, query_embedding: &[f32], k: usize) -> Option<HashMap<Uuid, f64>> {
let csr = self.csr_cache.as_ref()?;
if csr.rows == 0 {
return None;
@ -374,10 +364,15 @@ impl KnowledgeGraph {
}
/// Partition returning (clusters, cut_value, edge_strengths).
pub fn partition_full(&self, min_cluster_size: usize) -> (Vec<KnowledgeCluster>, f64, Vec<EdgeStrengthInfo>) {
pub fn partition_full(
&self,
min_cluster_size: usize,
) -> (Vec<KnowledgeCluster>, f64, Vec<EdgeStrengthInfo>) {
// Try real MinCut partitioning
if self.nodes.len() >= 3 {
if let Some((clusters, cut_val, strengths)) = self.partition_via_mincut_full(min_cluster_size) {
if let Some((clusters, cut_val, strengths)) =
self.partition_via_mincut_full(min_cluster_size)
{
if clusters.len() >= 2 {
return (clusters, cut_val, strengths);
}
@ -401,7 +396,10 @@ impl KnowledgeGraph {
fn partition_by_category(&self, min_cluster_size: usize) -> Vec<KnowledgeCluster> {
let mut by_category: HashMap<BrainCategory, Vec<Uuid>> = HashMap::new();
for (&id, node) in &self.nodes {
by_category.entry(node.category.clone()).or_default().push(id);
by_category
.entry(node.category.clone())
.or_default()
.push(id);
}
let mut clusters = Vec::new();
@ -420,12 +418,18 @@ impl KnowledgeGraph {
/// When a sparsifier is available and the full graph has > 50 000 edges,
/// uses the sparsified edge set (~19K edges vs ~1M) for a ~59x speedup
/// while preserving spectral cut quality (ADR-116).
fn partition_via_mincut_full(&self, min_cluster_size: usize) -> Option<(Vec<KnowledgeCluster>, f64, Vec<EdgeStrengthInfo>)> {
fn partition_via_mincut_full(
&self,
min_cluster_size: usize,
) -> Option<(Vec<KnowledgeCluster>, f64, Vec<EdgeStrengthInfo>)> {
let use_sparsified = self.sparsifier.is_some() && self.edges.len() > 50_000;
let edges: Vec<(u64, u64, f64)> = if use_sparsified {
let spar = self.sparsifier.as_ref().unwrap();
spar.sparsifier().edges().map(|(u, v, w)| (u as u64, v as u64, w)).collect()
spar.sparsifier()
.edges()
.map(|(u, v, w)| (u as u64, v as u64, w))
.collect()
} else {
self.edges
.iter()
@ -528,7 +532,12 @@ impl KnowledgeGraph {
/// Build a KnowledgeCluster from member IDs
fn build_cluster(&self, id: u32, members: &[Uuid]) -> KnowledgeCluster {
let dim = self.nodes.values().next().map(|n| n.embedding.len()).unwrap_or(0);
let dim = self
.nodes
.values()
.next()
.map(|n| n.embedding.len())
.unwrap_or(0);
let mut centroid = vec![0.0f32; dim];
let mut category_counts: HashMap<BrainCategory, usize> = HashMap::new();
let mut embeddings = Vec::new();
@ -624,7 +633,13 @@ impl KnowledgeGraph {
pub fn partition_canonical_full(
&self,
min_cluster_size: usize,
) -> (Vec<KnowledgeCluster>, f64, Vec<EdgeStrengthInfo>, Option<String>, Option<u64>) {
) -> (
Vec<KnowledgeCluster>,
f64,
Vec<EdgeStrengthInfo>,
Option<String>,
Option<u64>,
) {
if self.nodes.len() < 3 {
let (clusters, cut_val, strengths) = self.partition_full(min_cluster_size);
return (clusters, cut_val, strengths, None, None);
@ -652,11 +667,17 @@ impl KnowledgeGraph {
let source_side: std::collections::HashSet<u64> =
cut.side_vertices.iter().copied().collect();
let side_a: Vec<Uuid> = self.node_ids.iter().enumerate()
let side_a: Vec<Uuid> = self
.node_ids
.iter()
.enumerate()
.filter(|(i, _)| source_side.contains(&(*i as u64)))
.map(|(_, id)| *id)
.collect();
let side_b: Vec<Uuid> = self.node_ids.iter().enumerate()
let side_b: Vec<Uuid> = self
.node_ids
.iter()
.enumerate()
.filter(|(i, _)| !source_side.contains(&(*i as u64)))
.map(|(_, id)| *id)
.collect();
@ -677,7 +698,13 @@ impl KnowledgeGraph {
}
let strengths = self.compute_edge_strengths(&clusters);
(clusters, cut_value, strengths, Some(cut_hash_hex), Some(first_sep))
(
clusters,
cut_value,
strengths,
Some(cut_hash_hex),
Some(first_sep),
)
}
None => {
// Canonical cut not available (disconnected graph, etc.)
@ -699,11 +726,7 @@ impl KnowledgeGraph {
})
.collect();
self.mincut = MinCutBuilder::new()
.exact()
.with_edges(edges)
.build()
.ok();
self.mincut = MinCutBuilder::new().exact().with_edges(edges).build().ok();
}
/// Rebuild the CsrMatrix from the adjacency list
@ -883,8 +906,8 @@ pub fn cosine_similarity_normalized(a: &[f32], b: &[f32]) -> f64 {
let (mut d0, mut d1) = (0.0f64, 0.0f64);
for c in 0..chunks {
let i = c * 4;
d0 += (a[i] as f64) * (b[i] as f64) + (a[i+2] as f64) * (b[i+2] as f64);
d1 += (a[i+1] as f64) * (b[i+1] as f64) + (a[i+3] as f64) * (b[i+3] as f64);
d0 += (a[i] as f64) * (b[i] as f64) + (a[i + 2] as f64) * (b[i + 2] as f64);
d1 += (a[i + 1] as f64) * (b[i + 1] as f64) + (a[i + 3] as f64) * (b[i + 3] as f64);
}
let mut sum = d0 + d1;
for i in (chunks * 4)..n {
@ -907,8 +930,18 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 {
let (mut nb0, mut nb1) = (0.0f64, 0.0f64);
for c in 0..chunks {
let i = c * 4;
let (a0, a1, a2, a3) = (a[i] as f64, a[i+1] as f64, a[i+2] as f64, a[i+3] as f64);
let (b0, b1, b2, b3) = (b[i] as f64, b[i+1] as f64, b[i+2] as f64, b[i+3] as f64);
let (a0, a1, a2, a3) = (
a[i] as f64,
a[i + 1] as f64,
a[i + 2] as f64,
a[i + 3] as f64,
);
let (b0, b1, b2, b3) = (
b[i] as f64,
b[i + 1] as f64,
b[i + 2] as f64,
b[i + 3] as f64,
);
dot0 += a0 * b0 + a2 * b2;
dot1 += a1 * b1 + a3 * b3;
na0 += a0 * a0 + a2 * a2;
@ -925,6 +958,8 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 {
let dot = dot0 + dot1;
let norm_a = (na0 + na1).sqrt();
let norm_b = (nb0 + nb1).sqrt();
if norm_a < 1e-10 || norm_b < 1e-10 { return 0.0; }
if norm_a < 1e-10 || norm_b < 1e-10 {
return 0.0;
}
dot / (norm_a * norm_b)
}

18
crates/mcp-brain-server/src/lib.rs
View file

@ -10,25 +10,25 @@ pub mod cognitive;
pub mod drift;
pub mod embeddings;
pub mod gcs;
pub mod gist;
pub mod graph;
pub mod midstream;
pub mod notify;
pub mod optimizer;
pub mod pipeline;
pub mod pubmed;
pub mod quantization;
pub mod ranking;
pub mod rate_limit;
pub mod reputation;
pub mod routes;
pub mod store;
pub mod symbolic;
pub mod tests;
pub mod midstream;
pub mod types;
pub mod trainer;
pub mod types;
pub mod verify;
pub mod voice;
pub mod symbolic;
pub mod optimizer;
pub mod web_memory;
pub mod web_ingest;
pub mod web_memory;
pub mod web_store;
pub mod pubmed;
pub mod quantization;
pub mod notify;
pub mod gist;

28
crates/mcp-brain-server/src/main.rs
View file

@ -26,7 +26,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let train_state = state.clone();
let _training_handle = tokio::spawn(async move {
let train_interval = std::time::Duration::from_secs(300); // 5 min: full enhanced cycle
let tick_interval = std::time::Duration::from_secs(60); // 60s: lightweight cognitive tick
let tick_interval = std::time::Duration::from_secs(60); // 60s: lightweight cognitive tick
let mut tick_count = 0u64;
// Wait 30s before first cycle (let startup finish, data load)
@ -79,7 +79,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
category: top_cat.clone(),
};
let arm = ruvector_domain_expansion::transfer::ArmId(top_cat.clone());
train_state.domain_engine.write().meta.curiosity.record_visit(&bucket, &arm);
train_state
.domain_engine
.write()
.meta
.curiosity
.record_visit(&bucket, &arm);
}
}
@ -88,7 +93,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let mem_count = train_state.store.memory_count();
let ws_load = train_state.workspace.read().current_load();
train_state.internal_voice.write().observe(
format!("tick {}: {} memories, GWT load {:.2}", tick_count, mem_count, ws_load),
format!(
"tick {}: {} memories, GWT load {:.2}",
tick_count, mem_count, ws_load
),
uuid::Uuid::nil(),
);
}
@ -136,17 +144,25 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
let edge_count = spar_state.graph.read().edge_count();
if edge_count > 5_000_000 {
tracing::info!("Skipping sparsifier build: graph too large ({edge_count} edges, >5M threshold)");
tracing::info!(
"Skipping sparsifier build: graph too large ({edge_count} edges, >5M threshold)"
);
} else if edge_count > 100_000 && spar_state.graph.read().sparsifier_stats().is_none() {
tracing::info!("Background sparsifier build starting ({edge_count} edges)");
// Run in spawn_blocking to avoid starving the tokio runtime
let graph = spar_state.graph.clone();
tokio::task::spawn_blocking(move || {
graph.write().rebuild_sparsifier();
}).await.ok();
})
.await
.ok();
let stats = spar_state.graph.read().sparsifier_stats();
if let Some(s) = stats {
tracing::info!("Sparsifier built: {} edges, {:.1}x compression", s.sparsified_edges, s.compression_ratio);
tracing::info!(
"Sparsifier built: {} edges, {:.1}x compression",
s.sparsified_edges,
s.compression_ratio
);
} else {
tracing::warn!("Sparsifier build returned no stats");
}

12
crates/mcp-brain-server/src/midstream.rs
View file

@ -69,7 +69,11 @@ pub mod temporal_neural_solver_stub {
// ── Strange Loop Meta-Cognition (strange-loop) ─────────────────────────
/// Create a default StrangeLoop engine for meta-cognitive reasoning.
pub fn create_strange_loop() -> strange_loop::StrangeLoop<strange_loop::ScalarReasoner, strange_loop::SimpleCritic, strange_loop::SafeReflector> {
pub fn create_strange_loop() -> strange_loop::StrangeLoop<
strange_loop::ScalarReasoner,
strange_loop::SimpleCritic,
strange_loop::SafeReflector,
> {
let reasoner = strange_loop::ScalarReasoner::new(0.0, 1.0);
let critic = strange_loop::SimpleCritic::new();
let reflector = strange_loop::SafeReflector::new();
@ -88,7 +92,11 @@ pub fn create_strange_loop() -> strange_loop::StrangeLoop<strange_loop::ScalarRe
/// Run a meta-cognitive evaluation on a search context.
/// Returns a small additive bonus (0.0 to 0.04) based on the loop's convergence.
pub fn strange_loop_score(
loop_engine: &mut strange_loop::StrangeLoop<strange_loop::ScalarReasoner, strange_loop::SimpleCritic, strange_loop::SafeReflector>,
loop_engine: &mut strange_loop::StrangeLoop<
strange_loop::ScalarReasoner,
strange_loop::SimpleCritic,
strange_loop::SafeReflector,
>,
query_relevance: f64,
memory_quality: f64,
) -> f32 {

171
crates/mcp-brain-server/src/notify.rs
View file

@ -36,7 +36,13 @@ impl OpenTracker {
}
}
pub fn record_open(&self, tracking_id: &str, category: &str, subject: &str, user_agent: Option<&str>) {
pub fn record_open(
&self,
tracking_id: &str,
category: &str,
subject: &str,
user_agent: Option<&str>,
) {
let open = EmailOpen {
tracking_id: tracking_id.to_string(),
category: category.to_string(),
@ -71,10 +77,17 @@ impl OpenTracker {
pub fn open_rates(&self) -> HashMap<String, f64> {
let stats = self.stats.lock();
stats.iter().map(|(cat, (sent, opened))| {
let rate = if *sent > 0 { *opened as f64 / *sent as f64 } else { 0.0 };
(cat.clone(), rate)
}).collect()
stats
.iter()
.map(|(cat, (sent, opened))| {
let rate = if *sent > 0 {
*opened as f64 / *sent as f64
} else {
0.0
};
(cat.clone(), rate)
})
.collect()
}
pub fn stats_summary(&self) -> serde_json::Value {
@ -82,12 +95,19 @@ impl OpenTracker {
let opens = self.opens.lock();
let mut categories = serde_json::Map::new();
for (cat, (sent, opened)) in stats.iter() {
let rate = if *sent > 0 { *opened as f64 / *sent as f64 } else { 0.0 };
categories.insert(cat.clone(), serde_json::json!({
"sent": sent,
"opened": opened,
"open_rate": format!("{:.1}%", rate * 100.0)
}));
let rate = if *sent > 0 {
*opened as f64 / *sent as f64
} else {
0.0
};
categories.insert(
cat.clone(),
serde_json::json!({
"sent": sent,
"opened": opened,
"open_rate": format!("{:.1}%", rate * 100.0)
}),
);
}
serde_json::json!({
"total_opens": opens.len(),
@ -105,7 +125,8 @@ pub struct ResendNotifier {
from_name: String,
recipient: String,
/// Dedup: track recently welcomed emails (max 1 per 24h)
recent_welcomes: std::sync::Arc<parking_lot::Mutex<std::collections::HashMap<String, std::time::Instant>>>,
recent_welcomes:
std::sync::Arc<parking_lot::Mutex<std::collections::HashMap<String, std::time::Instant>>>,
/// Per-category rate limiter: category -> (last_sent, cooldown)
rate_limits: std::sync::Arc<Mutex<HashMap<String, (Instant, Duration)>>>,
/// Open tracking
@ -150,9 +171,11 @@ const STYLE_CONTAINER: &str = "font-family:'SF Mono',SFMono-Regular,Menlo,monosp
const STYLE_HEADING: &str = "color:#4fc3f7;margin:0 0 16px 0;font-size:20px;";
const STYLE_SUBHEADING: &str = "color:#4fc3f7;margin:16px 0 8px 0;font-size:16px;";
const STYLE_TEXT: &str = "color:#c0c0ff;font-size:14px;line-height:1.6;";
const STYLE_CODE: &str = "background:#1a1a3a;color:#7fdbca;padding:2px 6px;border-radius:4px;font-size:13px;";
const STYLE_CODE: &str =
"background:#1a1a3a;color:#7fdbca;padding:2px 6px;border-radius:4px;font-size:13px;";
const STYLE_CMD: &str = "background:#1a1a3a;color:#7fdbca;padding:12px 16px;border-radius:8px;font-size:13px;margin:8px 0;display:block;";
const STYLE_FOOTER: &str = "color:#666;margin-top:20px;font-size:11px;border-top:1px solid #222;padding-top:12px;";
const STYLE_FOOTER: &str =
"color:#666;margin-top:20px;font-size:11px;border-top:1px solid #222;padding-top:12px;";
const STYLE_BADGE: &str = "display:inline-block;background:#1a1a3a;color:#4fc3f7;padding:2px 8px;border-radius:4px;font-size:11px;margin:2px;";
fn footer_html() -> String {
@ -172,14 +195,18 @@ impl ResendNotifier {
if api_key.is_empty() {
return None;
}
let from_email = std::env::var("BRAIN_NOTIFICATION_EMAIL")
.unwrap_or_else(|_| "pi@ruv.io".into());
let from_name = std::env::var("BRAIN_NOTIFICATION_NAME")
.unwrap_or_else(|_| "Pi Brain".into());
let recipient = std::env::var("BRAIN_NOTIFY_RECIPIENT")
.unwrap_or_else(|_| "ruv@ruv.net".into());
let from_email =
std::env::var("BRAIN_NOTIFICATION_EMAIL").unwrap_or_else(|_| "pi@ruv.io".into());
let from_name =
std::env::var("BRAIN_NOTIFICATION_NAME").unwrap_or_else(|_| "Pi Brain".into());
let recipient =
std::env::var("BRAIN_NOTIFY_RECIPIENT").unwrap_or_else(|_| "ruv@ruv.net".into());
tracing::info!("Resend notifier initialized: from={}, to={}", from_email, recipient);
tracing::info!(
"Resend notifier initialized: from={}, to={}",
from_email,
recipient
);
Some(Self {
client: reqwest::Client::new(),
@ -189,13 +216,18 @@ impl ResendNotifier {
recipient,
rate_limits: std::sync::Arc::new(Mutex::new(HashMap::new())),
tracker: OpenTracker::new(),
recent_welcomes: std::sync::Arc::new(parking_lot::Mutex::new(std::collections::HashMap::new())),
recent_welcomes: std::sync::Arc::new(parking_lot::Mutex::new(
std::collections::HashMap::new(),
)),
})
}
fn check_rate_limit(&self, category: &str) -> bool {
let cooldowns = default_cooldowns();
let cooldown = cooldowns.get(category).copied().unwrap_or(Duration::from_secs(3600));
let cooldown = cooldowns
.get(category)
.copied()
.unwrap_or(Duration::from_secs(3600));
if cooldown.is_zero() {
return true;
}
@ -234,22 +266,19 @@ impl ResendNotifier {
);
// Insert pixel before closing </div> and add unsubscribe
if html.ends_with("</div>") {
format!("{}{}{}", &html[..html.len()-6], pixel, "</div>")
+ &unsub
format!("{}{}{}", &html[..html.len() - 6], pixel, "</div>") + &unsub
} else {
format!("{}{}{}", html, pixel, unsub)
}
}
/// Send an email. Respects per-category rate limits. Injects tracking pixel.
pub async fn send(
&self,
category: &str,
subject: &str,
html: &str,
) -> Result<String, String> {
pub async fn send(&self, category: &str, subject: &str, html: &str) -> Result<String, String> {
if !self.check_rate_limit(category) {
return Err(format!("rate-limited: category '{}' is in cooldown", category));
return Err(format!(
"rate-limited: category '{}' is in cooldown",
category
));
}
let tracking_id = uuid::Uuid::new_v4().to_string();
@ -264,7 +293,8 @@ impl ResendNotifier {
reply_to: Some(&self.from_email),
};
let resp = self.client
let resp = self
.client
.post("https://api.resend.com/emails")
.header("Authorization", format!("Bearer {}", self.api_key))
.json(&body)
@ -280,7 +310,12 @@ impl ResendNotifier {
serde_json::from_str(&text).unwrap_or(SendEmailResponse { id: None });
let id = parsed.id.unwrap_or_else(|| "unknown".into());
self.tracker.record_sent(category);
tracing::info!("Email sent: category={}, id={}, tracking={}", category, id, tracking_id);
tracing::info!(
"Email sent: category={}, id={}, tracking={}",
category,
id,
tracking_id
);
Ok(id)
} else {
tracing::warn!("Resend API error: status={}, body={}", status, text);
@ -297,7 +332,10 @@ impl ResendNotifier {
html: &str,
) -> Result<String, String> {
if !self.check_rate_limit(category) {
return Err(format!("rate-limited: category '{}' is in cooldown", category));
return Err(format!(
"rate-limited: category '{}' is in cooldown",
category
));
}
let tracking_id = uuid::Uuid::new_v4().to_string();
@ -312,7 +350,8 @@ impl ResendNotifier {
reply_to: Some(&self.from_email),
};
let resp = self.client
let resp = self
.client
.post("https://api.resend.com/emails")
.header("Authorization", format!("Bearer {}", self.api_key))
.json(&body)
@ -328,7 +367,13 @@ impl ResendNotifier {
serde_json::from_str(&text).unwrap_or(SendEmailResponse { id: None });
let id = parsed.id.unwrap_or_else(|| "unknown".into());
self.tracker.record_sent(category);
tracing::info!("Email sent to {}: category={}, id={}, tracking={}", to, category, id, tracking_id);
tracing::info!(
"Email sent to {}: category={}, id={}, tracking={}",
to,
category,
id,
tracking_id
);
Ok(id)
} else {
Err(format!("resend API error {}: {}", status, text))
@ -339,7 +384,11 @@ impl ResendNotifier {
/// Welcome email for new users connecting to the brain.
/// Rate-limited: max 1 welcome per email per 24 hours.
pub async fn send_welcome(&self, user_email: &str, user_name: Option<&str>) -> Result<String, String> {
pub async fn send_welcome(
&self,
user_email: &str,
user_name: Option<&str>,
) -> Result<String, String> {
// Dedup: check if we already welcomed this email recently
{
let mut recent = self.recent_welcomes.lock();
@ -347,7 +396,10 @@ impl ResendNotifier {
// Clean entries older than 24h
recent.retain(|_, t| now.duration_since(*t) < std::time::Duration::from_secs(86400));
if recent.contains_key(user_email) {
tracing::info!("Welcome dedup: skipping {} (already welcomed recently)", user_email);
tracing::info!(
"Welcome dedup: skipping {} (already welcomed recently)",
user_email
);
return Ok("dedup-skipped".to_string());
}
recent.insert(user_email.to_string(), now);
@ -413,10 +465,13 @@ curl "https://pi.ruv.io/v1/memories/search?q=authentication"
name = name,
);
self.send_to(user_email, "welcome",
self.send_to(
user_email,
"welcome",
&format!("Welcome to Pi Brain, {}", name),
&html,
).await
)
.await
}
/// Help email — explains all available commands and capabilities
@ -545,8 +600,20 @@ partitioning, and Gemini Flash grounding for cognitive enrichment.
sona_patterns: usize,
drift: f64,
) -> Result<String, String> {
let drift_indicator = if drift > 0.5 { "HIGH" } else if drift > 0.2 { "MODERATE" } else { "STABLE" };
let drift_color = if drift > 0.5 { "#ff6b6b" } else if drift > 0.2 { "#ffd93d" } else { "#6bff6b" };
let drift_indicator = if drift > 0.5 {
"HIGH"
} else if drift > 0.2 {
"MODERATE"
} else {
"STABLE"
};
let drift_color = if drift > 0.5 {
"#ff6b6b"
} else if drift > 0.2 {
"#ffd93d"
} else {
"#6bff6b"
};
let html = format!(
r#"<div style="{container}">
@ -608,7 +675,11 @@ Resend integration is working correctly.
) -> Result<String, String> {
let mut rows = String::new();
for (i, (title, content, score)) in results.iter().enumerate() {
let truncated = if content.len() > 200 { &content[..200] } else { content };
let truncated = if content.len() > 200 {
&content[..200]
} else {
content
};
rows.push_str(&format!(
r#"<tr style="border-bottom:1px solid #222;">
<td style="padding:8px 0;vertical-align:top;">
@ -616,7 +687,10 @@ Resend integration is working correctly.
<span style="color:#999;font-size:12px;">{}</span><br>
<span style="color:#666;font-size:11px;">score: {:.3}</span>
</td></tr>"#,
i + 1, title, truncated, score
i + 1,
title,
truncated,
score
));
}
@ -642,9 +716,12 @@ Resend integration is working correctly.
rows = rows,
);
self.send_to(to, "chat",
self.send_to(
to,
"chat",
&format!("Re: search {} ({} results)", query, results.len()),
&html,
).await
)
.await
}
}

57
crates/mcp-brain-server/src/optimizer.rs
View file

@ -42,7 +42,8 @@ impl Default for OptimizerConfig {
fn default() -> Self {
Self {
api_base: "https://generativelanguage.googleapis.com/v1beta/models".to_string(),
model_id: std::env::var("GEMINI_MODEL").unwrap_or_else(|_| "gemini-2.5-flash".to_string()),
model_id: std::env::var("GEMINI_MODEL")
.unwrap_or_else(|_| "gemini-2.5-flash".to_string()),
max_tokens: 2048,
temperature: 0.3,
interval_secs: 3600, // 1 hour
@ -148,7 +149,8 @@ pub struct GeminiOptimizer {
impl GeminiOptimizer {
/// Create a new optimizer with the given config
pub fn new(config: OptimizerConfig) -> Self {
let api_key = std::env::var("GEMINI_API_KEY").ok()
let api_key = std::env::var("GEMINI_API_KEY")
.ok()
.or_else(|| std::env::var("GOOGLE_API_KEY").ok());
let http = reqwest::Client::builder()
@ -187,7 +189,9 @@ impl GeminiOptimizer {
task: OptimizationTask,
context: OptimizationContext,
) -> Result<OptimizationResult, String> {
let api_key = self.api_key.as_ref()
let api_key = self
.api_key
.as_ref()
.ok_or("Gemini API key not configured")?;
let start = std::time::Instant::now();
@ -283,10 +287,17 @@ impl GeminiOptimizer {
context.sona_patterns,
context.working_memory_load * 100.0,
context.memory_count,
context.sample_propositions.iter()
context
.sample_propositions
.iter()
.take(5)
.map(|p| format!(" - {}({}) [conf={:.2}, evidence={}]",
p.predicate, p.arguments.join(", "), p.confidence, p.evidence_count))
.map(|p| format!(
" - {}({}) [conf={:.2}, evidence={}]",
p.predicate,
p.arguments.join(", "),
p.confidence,
p.evidence_count
))
.collect::<Vec<_>>()
.join("\n")
)
@ -300,13 +311,11 @@ impl GeminiOptimizer {
async fn call_gemini(&self, api_key: &str, prompt: &str) -> Result<String, String> {
let url = format!(
"{}/{}:generateContent?key={}",
self.config.api_base,
self.config.model_id,
api_key
self.config.api_base, self.config.model_id, api_key
);
let grounding_enabled = std::env::var("GEMINI_GROUNDING")
.unwrap_or_else(|_| "true".to_string()) == "true";
let grounding_enabled =
std::env::var("GEMINI_GROUNDING").unwrap_or_else(|_| "true".to_string()) == "true";
let mut body = serde_json::json!({
"contents": [{
@ -326,7 +335,8 @@ impl GeminiOptimizer {
}]);
}
let response = self.http
let response = self
.http
.post(&url)
.header("content-type", "application/json")
.json(&body)
@ -340,21 +350,26 @@ impl GeminiOptimizer {
return Err(format!("Gemini API error {}: {}", status, error_text));
}
let json: serde_json::Value = response.json().await
let json: serde_json::Value = response
.json()
.await
.map_err(|e| format!("JSON parse error: {}", e))?;
// Log grounding metadata if present (source URLs, support scores)
if let Some(candidate) = json.get("candidates").and_then(|c| c.get(0)) {
if let Some(grounding) = candidate.get("groundingMetadata") {
let sources = grounding.get("groundingChunks")
let sources = grounding
.get("groundingChunks")
.and_then(|c| c.as_array())
.map(|a| a.len())
.unwrap_or(0);
let support = grounding.get("groundingSupports")
let support = grounding
.get("groundingSupports")
.and_then(|s| s.as_array())
.map(|a| a.len())
.unwrap_or(0);
let query = grounding.get("webSearchQueries")
let query = grounding
.get("webSearchQueries")
.and_then(|q| q.as_array())
.and_then(|a| a.first())
.and_then(|q| q.as_str())
@ -364,7 +379,9 @@ impl GeminiOptimizer {
supports = support,
query = query,
"[optimizer] Grounding: {} sources, {} supports, query='{}'",
sources, support, query
sources,
support,
query
);
}
}
@ -409,9 +426,9 @@ impl GeminiOptimizer {
configured: self.is_configured(),
run_count: self.run_count,
last_run: self.last_run,
next_due: self.last_run.map(|lr| {
lr + chrono::Duration::seconds(self.config.interval_secs as i64)
}),
next_due: self
.last_run
.map(|lr| lr + chrono::Duration::seconds(self.config.interval_secs as i64)),
}
}
}

737
crates/mcp-brain-server/src/pipeline.rs
View file

File diff suppressed because it is too large Load diff

40
crates/mcp-brain-server/src/pubmed.rs
View file

@ -65,7 +65,9 @@ impl PubMedArticle {
text,
title: self.title.clone(),
meta_description: self.abstract_text.chars().take(300).collect(),
links: self.references.iter()
links: self
.references
.iter()
.map(|pmid| format!("https://pubmed.ncbi.nlm.nih.gov/{pmid}/"))
.collect(),
language: "en".to_string(),
@ -122,9 +124,7 @@ pub async fn efetch(
}
let id_str = pmids.join(",");
let url = format!(
"{EFETCH_URL}?db=pubmed&id={id_str}&rettype=xml&retmode=xml"
);
let url = format!("{EFETCH_URL}?db=pubmed&id={id_str}&rettype=xml&retmode=xml");
let resp = client
.get(&url)
@ -213,7 +213,9 @@ fn extract_abstract(xml: &str) -> String {
let tag_str = &xml[abs_start..content_start];
let label = if let Some(lpos) = tag_str.find("Label=\"") {
let lstart = lpos + 7;
tag_str[lstart..].find('"').map(|end| &tag_str[lstart..lstart + end])
tag_str[lstart..]
.find('"')
.map(|end| &tag_str[lstart..lstart + end])
} else {
None
};
@ -413,24 +415,24 @@ pub fn analyze_discoveries(
let pmid_to_mem: HashMap<String, &WebMemory> = accepted
.iter()
.filter_map(|m| {
let pmid = m.source_url
.trim_end_matches('/')
.rsplit('/')
.next()?;
let pmid = m.source_url.trim_end_matches('/').rsplit('/').next()?;
Some((pmid.to_string(), m))
})
.collect();
// ── Emerging Topics: high-novelty articles with MeSH context ────
let mut emerging_topics = Vec::new();
let mut high_novelty: Vec<&WebMemory> = accepted
.iter()
.filter(|m| m.novelty_score > 0.5)
.collect();
let mut high_novelty: Vec<&WebMemory> =
accepted.iter().filter(|m| m.novelty_score > 0.5).collect();
high_novelty.sort_by(|a, b| b.novelty_score.partial_cmp(&a.novelty_score).unwrap());
for mem in high_novelty.iter().take(10) {
let pmid = mem.source_url.trim_end_matches('/').rsplit('/').next().unwrap_or("?");
let pmid = mem
.source_url
.trim_end_matches('/')
.rsplit('/')
.next()
.unwrap_or("?");
let article = articles.iter().find(|a| a.pmid == pmid);
emerging_topics.push(EmergingTopic {
representative_title: mem.base.title.clone(),
@ -458,7 +460,9 @@ pub fn analyze_discoveries(
let (mem_a, art_a) = &accepted_refs[i];
let (mem_b, art_b) = &accepted_refs[j];
let shared_mesh: Vec<String> = art_a.mesh_terms.iter()
let shared_mesh: Vec<String> = art_a
.mesh_terms
.iter()
.filter(|t| art_b.mesh_terms.contains(t))
.cloned()
.collect();
@ -562,7 +566,11 @@ pub async fn push_to_brain(
match resp {
Ok(r) if r.status().is_success() => pushed += 1,
Ok(r) => {
tracing::warn!("Brain push failed for PMID {}: HTTP {}", article.pmid, r.status());
tracing::warn!(
"Brain push failed for PMID {}: HTTP {}",
article.pmid,
r.status()
);
}
Err(e) => {
tracing::warn!("Brain push failed for PMID {}: {e}", article.pmid);

70
crates/mcp-brain-server/src/quantization.rs
View file

@ -11,8 +11,8 @@
//! - CentroidMerged: 3-bit quantization (10.7x compression)
//! - Archived: 2-bit quantization (16x compression)
use serde::{Deserialize, Serialize};
use crate::web_memory::CompressionTier;
use serde::{Deserialize, Serialize};
/// Quantization configuration by tier
#[derive(Debug, Clone, Serialize, Deserialize)]
@ -28,9 +28,9 @@ pub struct PiQConfig {
impl Default for PiQConfig {
fn default() -> Self {
Self {
delta_bits: 4, // 2.67x compression, ~99% recall
centroid_bits: 3, // 10.7x compression, ~96% recall
archived_bits: 2, // 16x compression, ~90% recall
delta_bits: 4, // 2.67x compression, ~99% recall
centroid_bits: 3, // 10.7x compression, ~96% recall
archived_bits: 2, // 16x compression, ~90% recall
}
}
}
@ -134,7 +134,9 @@ impl PiQQuantizer {
.iter()
.map(|&v| {
let normalized = (v - min_val) / range;
(normalized * levels as f32).round().clamp(0.0, levels as f32) as u32
(normalized * levels as f32)
.round()
.clamp(0.0, levels as f32) as u32
})
.collect();
@ -207,15 +209,15 @@ impl PiQQuantizer {
/// Based on empirical measurements from ADR-115
pub fn expected_recall(bits: u8) -> f32 {
match bits {
8 => 0.999, // Nearly lossless
8 => 0.999, // Nearly lossless
7 => 0.997,
6 => 0.995,
5 => 0.99,
4 => 0.985, // 4-bit still very good
3 => 0.96, // PiQ3 target
2 => 0.90, // Archived tier
1 => 0.70, // Binary (not recommended)
_ => 1.0, // Full precision
4 => 0.985, // 4-bit still very good
3 => 0.96, // PiQ3 target
2 => 0.90, // Archived tier
1 => 0.70, // Binary (not recommended)
_ => 1.0, // Full precision
}
}
}
@ -310,18 +312,26 @@ mod tests {
let embedding: Vec<f32> = (0..128).map(|i| (i as f32) / 127.0).collect();
// Full tier: no quantization
assert!(quantizer.quantize(&embedding, CompressionTier::Full).is_none());
assert!(quantizer
.quantize(&embedding, CompressionTier::Full)
.is_none());
// DeltaCompressed: 4-bit
let delta = quantizer.quantize(&embedding, CompressionTier::DeltaCompressed).unwrap();
let delta = quantizer
.quantize(&embedding, CompressionTier::DeltaCompressed)
.unwrap();
assert_eq!(delta.bits, 4);
// CentroidMerged: 3-bit
let centroid = quantizer.quantize(&embedding, CompressionTier::CentroidMerged).unwrap();
let centroid = quantizer
.quantize(&embedding, CompressionTier::CentroidMerged)
.unwrap();
assert_eq!(centroid.bits, 3);
// Archived: 2-bit
let archived = quantizer.quantize(&embedding, CompressionTier::Archived).unwrap();
let archived = quantizer
.quantize(&embedding, CompressionTier::Archived)
.unwrap();
assert_eq!(archived.bits, 2);
}
@ -426,21 +436,39 @@ mod tests {
println!();
println!("3-bit (CentroidMerged tier):");
println!(" - Compressed size: {size_3bit} bytes");
println!(" - Compression ratio: {:.2}x", original_size as f32 / size_3bit as f32);
println!(
" - Compression ratio: {:.2}x",
original_size as f32 / size_3bit as f32
);
println!(" - Recall (cosine similarity): {:.4}", avg_recall_3bit);
println!(" - Throughput: {:.2} embeddings/sec", num_embeddings as f64 / time_3bit.as_secs_f64());
println!(
" - Throughput: {:.2} embeddings/sec",
num_embeddings as f64 / time_3bit.as_secs_f64()
);
println!();
println!("4-bit (DeltaCompressed tier):");
println!(" - Compressed size: {size_4bit} bytes");
println!(" - Compression ratio: {:.2}x", original_size as f32 / size_4bit as f32);
println!(
" - Compression ratio: {:.2}x",
original_size as f32 / size_4bit as f32
);
println!(" - Recall (cosine similarity): {:.4}", avg_recall_4bit);
println!(" - Throughput: {:.2} embeddings/sec", num_embeddings as f64 / time_4bit.as_secs_f64());
println!(
" - Throughput: {:.2} embeddings/sec",
num_embeddings as f64 / time_4bit.as_secs_f64()
);
println!();
println!("2-bit (Archived tier):");
println!(" - Compressed size: {size_2bit} bytes");
println!(" - Compression ratio: {:.2}x", original_size as f32 / size_2bit as f32);
println!(
" - Compression ratio: {:.2}x",
original_size as f32 / size_2bit as f32
);
println!(" - Recall (cosine similarity): {:.4}", avg_recall_2bit);
println!(" - Throughput: {:.2} embeddings/sec", num_embeddings as f64 / time_2bit.as_secs_f64());
println!(
" - Throughput: {:.2} embeddings/sec",
num_embeddings as f64 / time_2bit.as_secs_f64()
);
println!();
// Assertions

7
crates/mcp-brain-server/src/rate_limit.rs
View file

@ -83,8 +83,8 @@ impl RateLimiter {
ip_read_buckets: DashMap::new(),
write_limit,
read_limit,
ip_write_limit: write_limit * 3, // 1500/hr per IP (allows some key rotation)
ip_read_limit: read_limit * 3, // 15000/hr per IP
ip_write_limit: write_limit * 3, // 1500/hr per IP (allows some key rotation)
ip_read_limit: read_limit * 3, // 15000/hr per IP
window: Duration::from_secs(3600),
ops_counter: AtomicU64::new(0),
cleanup_interval: 1000,
@ -169,7 +169,8 @@ impl RateLimiter {
self.ip_read_buckets.retain(|_, bucket| !bucket.is_stale());
// Evict vote entries older than 24h
let vote_before = self.ip_votes.len();
self.ip_votes.retain(|_, (_, timestamp)| timestamp.elapsed() < Duration::from_secs(86400));
self.ip_votes
.retain(|_, (_, timestamp)| timestamp.elapsed() < Duration::from_secs(86400));
let vote_evicted = vote_before - self.ip_votes.len();
let write_evicted = write_before - self.write_buckets.len();

2585
crates/mcp-brain-server/src/routes.rs
View file

File diff suppressed because it is too large Load diff

290
crates/mcp-brain-server/src/store.rs
View file

@ -124,10 +124,12 @@ impl FirestoreClient {
let docs = self.firestore_list("brain_votes").await;
let mut count = 0usize;
for doc in docs {
let memory_id = doc.get("memory_id")
let memory_id = doc
.get("memory_id")
.and_then(|v| v.as_str())
.and_then(|s| s.parse::<Uuid>().ok());
let voter = doc.get("voter")
let voter = doc
.get("voter")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
if let (Some(mid), Some(v)) = (memory_id, voter) {
@ -136,7 +138,8 @@ impl FirestoreClient {
}
}
// Restore vote counter from loaded entries
self.vote_counter.store(count as u64, std::sync::atomic::Ordering::Relaxed);
self.vote_counter
.store(count as u64, std::sync::atomic::Ordering::Relaxed);
tracing::info!("Vote tracker rebuilt: {} entries from Firestore", count);
}
@ -172,7 +175,8 @@ impl FirestoreClient {
/// Fetch a new token from the GCE metadata server
async fn refresh_token(&self) -> Option<String> {
let url = "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token";
let resp = self.http
let resp = self
.http
.get(url)
.header("Metadata-Flavor", "Google")
.send()
@ -180,7 +184,10 @@ impl FirestoreClient {
.ok()?;
if !resp.status().is_success() {
tracing::warn!("Firestore: GCE metadata token request failed: {}", resp.status());
tracing::warn!(
"Firestore: GCE metadata token request failed: {}",
resp.status()
);
return None;
}
@ -191,8 +198,8 @@ impl FirestoreClient {
}
let token_resp: TokenResponse = resp.json().await.ok()?;
let expires_at = std::time::Instant::now()
+ std::time::Duration::from_secs(token_resp.expires_in);
let expires_at =
std::time::Instant::now() + std::time::Duration::from_secs(token_resp.expires_in);
let token = token_resp.access_token.clone();
@ -205,12 +212,19 @@ impl FirestoreClient {
});
}
tracing::debug!("Firestore token refreshed, expires in {}s", token_resp.expires_in);
tracing::debug!(
"Firestore token refreshed, expires in {}s",
token_resp.expires_in
);
Some(token)
}
/// Build an authenticated request builder for Firestore
async fn authenticated_request(&self, method: reqwest::Method, url: &str) -> reqwest::RequestBuilder {
async fn authenticated_request(
&self,
method: reqwest::Method,
url: &str,
) -> reqwest::RequestBuilder {
let mut builder = self.http.request(method, url);
if let Some(token) = self.get_token().await {
builder = builder.bearer_auth(token);
@ -224,7 +238,9 @@ impl FirestoreClient {
/// Wraps JSON body as a single `data` stringValue field for simplicity.
/// Uses PATCH to create or update documents.
async fn firestore_put(&self, collection: &str, doc_id: &str, body: &serde_json::Value) {
let Some(ref base) = self.base_url else { return };
let Some(ref base) = self.base_url else {
return;
};
let url = format!("{base}/{collection}/{doc_id}");
let json_str = serde_json::to_string(body).unwrap_or_default();
let firestore_doc = serde_json::json!({
@ -235,7 +251,8 @@ impl FirestoreClient {
// Retry loop: up to 2 attempts (initial + 1 retry) with token refresh on 401.
for attempt in 0..2u8 {
let result = self.authenticated_request(reqwest::Method::PATCH, &url)
let result = self
.authenticated_request(reqwest::Method::PATCH, &url)
.await
.json(&firestore_doc)
.send()
@ -248,7 +265,9 @@ impl FirestoreClient {
Ok(resp) if resp.status().as_u16() == 401 && attempt == 0 => {
tracing::info!("Firestore PATCH token expired, refreshing for retry...");
if self.refresh_token().await.is_none() {
tracing::warn!("Firestore PATCH {collection}/{doc_id}: token refresh failed");
tracing::warn!(
"Firestore PATCH {collection}/{doc_id}: token refresh failed"
);
return;
}
// Loop will retry with fresh token
@ -266,7 +285,9 @@ impl FirestoreClient {
return;
}
Err(e) if attempt == 0 => {
tracing::warn!("Firestore PATCH {collection}/{doc_id} failed: {e}, retrying...");
tracing::warn!(
"Firestore PATCH {collection}/{doc_id} failed: {e}, retrying..."
);
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
}
Err(e) => {
@ -279,9 +300,12 @@ impl FirestoreClient {
/// Delete a document from Firestore
async fn firestore_delete(&self, collection: &str, doc_id: &str) {
let Some(ref base) = self.base_url else { return };
let Some(ref base) = self.base_url else {
return;
};
let url = format!("{base}/{collection}/{doc_id}");
let result = self.authenticated_request(reqwest::Method::DELETE, &url)
let result = self
.authenticated_request(reqwest::Method::DELETE, &url)
.await
.send()
.await;
@ -310,7 +334,9 @@ impl FirestoreClient {
const MAX_PAGE_ERRORS: usize = 3;
async fn firestore_list(&self, collection: &str) -> Vec<serde_json::Value> {
let Some(ref base) = self.base_url else { return Vec::new() };
let Some(ref base) = self.base_url else {
return Vec::new();
};
let mut all_docs = Vec::new();
let mut page_token: Option<String> = None;
let mut consecutive_errors: usize = 0;
@ -330,7 +356,8 @@ impl FirestoreClient {
url.push_str(&format!("&pageToken={}", urlencoding::encode(token)));
}
let result = self.authenticated_request(reqwest::Method::GET, &url)
let result = self
.authenticated_request(reqwest::Method::GET, &url)
.await
.send()
.await;
@ -353,26 +380,38 @@ impl FirestoreClient {
consecutive_errors += 1;
tracing::warn!(
"Firestore LIST {collection} retry returned {} (error {}/{})",
resp.status(), consecutive_errors, Self::MAX_PAGE_ERRORS
resp.status(),
consecutive_errors,
Self::MAX_PAGE_ERRORS
);
if consecutive_errors >= Self::MAX_PAGE_ERRORS { break; }
if consecutive_errors >= Self::MAX_PAGE_ERRORS {
break;
}
continue;
}
Err(e) => {
consecutive_errors += 1;
tracing::warn!(
"Firestore LIST {collection} retry failed: {e} (error {}/{})",
consecutive_errors, Self::MAX_PAGE_ERRORS
consecutive_errors,
Self::MAX_PAGE_ERRORS
);
if consecutive_errors >= Self::MAX_PAGE_ERRORS { break; }
if consecutive_errors >= Self::MAX_PAGE_ERRORS {
break;
}
continue;
}
}
} else {
consecutive_errors += 1;
tracing::warn!("Firestore LIST {collection}: token refresh failed (error {}/{})",
consecutive_errors, Self::MAX_PAGE_ERRORS);
if consecutive_errors >= Self::MAX_PAGE_ERRORS { break; }
tracing::warn!(
"Firestore LIST {collection}: token refresh failed (error {}/{})",
consecutive_errors,
Self::MAX_PAGE_ERRORS
);
if consecutive_errors >= Self::MAX_PAGE_ERRORS {
break;
}
continue;
}
}
@ -381,7 +420,9 @@ impl FirestoreClient {
consecutive_errors += 1;
tracing::warn!(
"Firestore LIST {collection} returned {} (error {}/{})",
status, consecutive_errors, Self::MAX_PAGE_ERRORS
status,
consecutive_errors,
Self::MAX_PAGE_ERRORS
);
// 400 Bad Request with a page token means the token is stale
// (e.g. after OOM restart). Switch to offset-based pagination.
@ -395,16 +436,21 @@ impl FirestoreClient {
consecutive_errors = 0; // reset — this is a recovery, not repeated failure
continue;
}
if consecutive_errors >= Self::MAX_PAGE_ERRORS { break; }
if consecutive_errors >= Self::MAX_PAGE_ERRORS {
break;
}
continue;
}
Err(e) => {
consecutive_errors += 1;
tracing::warn!(
"Firestore LIST {collection} failed: {e} (error {}/{})",
consecutive_errors, Self::MAX_PAGE_ERRORS
consecutive_errors,
Self::MAX_PAGE_ERRORS
);
if consecutive_errors >= Self::MAX_PAGE_ERRORS { break; }
if consecutive_errors >= Self::MAX_PAGE_ERRORS {
break;
}
continue;
}
};
@ -415,9 +461,12 @@ impl FirestoreClient {
consecutive_errors += 1;
tracing::warn!(
"Firestore LIST {collection} parse error: {e} (error {}/{})",
consecutive_errors, Self::MAX_PAGE_ERRORS
consecutive_errors,
Self::MAX_PAGE_ERRORS
);
if consecutive_errors >= Self::MAX_PAGE_ERRORS { break; }
if consecutive_errors >= Self::MAX_PAGE_ERRORS {
break;
}
continue;
}
};
@ -447,7 +496,11 @@ impl FirestoreClient {
_ if use_offset_fallback => {
// In offset mode, check if we got any docs this page
// (no docs = we've exhausted the collection)
if body.get("documents").and_then(|d| d.as_array()).map_or(true, |d| d.is_empty()) {
if body
.get("documents")
.and_then(|d| d.as_array())
.map_or(true, |d| d.is_empty())
{
break;
}
// Otherwise continue with incremented offset (all_docs.len() grows each iteration)
@ -459,10 +512,14 @@ impl FirestoreClient {
if consecutive_errors > 0 {
tracing::warn!(
"Firestore LIST {collection}: loaded {} documents with {} error(s)",
all_docs.len(), consecutive_errors
all_docs.len(),
consecutive_errors
);
} else {
tracing::info!("Firestore LIST {collection}: loaded {} documents", all_docs.len());
tracing::info!(
"Firestore LIST {collection}: loaded {} documents",
all_docs.len()
);
}
all_docs
}
@ -501,8 +558,13 @@ impl FirestoreClient {
let docs = self.firestore_list("brain_page_status").await;
for doc in docs {
if let (Some(id), Some(status)) = (
doc.get("id").and_then(|v| v.as_str()).and_then(|s| s.parse::<Uuid>().ok()),
serde_json::from_value::<PageStatus>(doc.get("status").cloned().unwrap_or_default()).ok(),
doc.get("id")
.and_then(|v| v.as_str())
.and_then(|s| s.parse::<Uuid>().ok()),
serde_json::from_value::<PageStatus>(
doc.get("status").cloned().unwrap_or_default(),
)
.ok(),
) {
self.page_status.insert(id, status);
}
@ -525,7 +587,12 @@ impl FirestoreClient {
}
/// Public Firestore write for cross-module persistence (e.g., LoRA store)
pub async fn firestore_put_public(&self, collection: &str, doc_id: &str, body: &serde_json::Value) {
pub async fn firestore_put_public(
&self,
collection: &str,
doc_id: &str,
body: &serde_json::Value,
) {
self.firestore_put(collection, doc_id, body).await;
}
@ -542,7 +609,8 @@ impl FirestoreClient {
crate::graph::normalize_embedding(&mut memory.embedding);
// Write-through to Firestore
if let Ok(body) = serde_json::to_value(&memory) {
self.firestore_put("brain_memories", &id.to_string(), &body).await;
self.firestore_put("brain_memories", &id.to_string(), &body)
.await;
}
self.memories.insert(id, memory);
Ok(())
@ -555,18 +623,15 @@ impl FirestoreClient {
/// Delete a memory (contributor-scoped, cache + Firestore)
/// Uses atomic remove_if to prevent TOCTOU race
pub async fn delete_memory(
&self,
id: &Uuid,
contributor: &str,
) -> Result<bool, StoreError> {
pub async fn delete_memory(&self, id: &Uuid, contributor: &str) -> Result<bool, StoreError> {
// Atomic check-and-remove: no TOCTOU window
let removed = self.memories.remove_if(id, |_, entry| {
entry.contributor_id == contributor
});
let removed = self
.memories
.remove_if(id, |_, entry| entry.contributor_id == contributor);
match removed {
Some(_) => {
self.firestore_delete("brain_memories", &id.to_string()).await;
self.firestore_delete("brain_memories", &id.to_string())
.await;
Ok(true)
}
None => {
@ -601,9 +666,7 @@ impl FirestoreClient {
let m = entry.value();
let quality_ok = m.quality_score.mean() >= min_quality;
let category_ok = category.map_or(true, |c| &m.category == c);
let tags_ok = tags.map_or(true, |t| {
t.iter().any(|tag| m.tags.contains(tag))
});
let tags_ok = tags.map_or(true, |t| t.iter().any(|tag| m.tags.contains(tag)));
quality_ok && category_ok && tags_ok
})
.map(|entry| {
@ -686,7 +749,10 @@ impl FirestoreClient {
}
// Category match
if query_tokens.iter().any(|t| m.category.to_string().to_lowercase().contains(t)) {
if query_tokens
.iter()
.any(|t| m.category.to_string().to_lowercase().contains(t))
{
score += 1.5;
}
@ -718,9 +784,7 @@ impl FirestoreClient {
.filter(|entry| {
let m = entry.value();
let category_ok = category.map_or(true, |c| &m.category == c);
let tags_ok = tags.map_or(true, |t| {
t.iter().any(|tag| m.tags.contains(tag))
});
let tags_ok = tags.map_or(true, |t| t.iter().any(|tag| m.tags.contains(tag)));
category_ok && tags_ok
})
.map(|entry| entry.value().clone())
@ -750,11 +814,7 @@ impl FirestoreClient {
}
}
let paginated: Vec<BrainMemory> = memories
.into_iter()
.skip(offset)
.take(limit)
.collect();
let paginated: Vec<BrainMemory> = memories.into_iter().skip(offset).take(limit).collect();
Ok((paginated, total_count))
}
@ -782,10 +842,10 @@ impl FirestoreClient {
// Block duplicate votes: same voter on same memory (single lookup via entry API)
let vote_key = (*id, voter.to_string());
if let dashmap::mapref::entry::Entry::Occupied(_) = self.vote_tracker.entry(vote_key.clone()) {
return Err(StoreError::Forbidden(
"Already voted on this memory".into(),
));
if let dashmap::mapref::entry::Entry::Occupied(_) =
self.vote_tracker.entry(vote_key.clone())
{
return Err(StoreError::Forbidden("Already voted on this memory".into()));
}
let quality_score;
@ -834,27 +894,34 @@ impl FirestoreClient {
voter: voter.to_string(),
timestamp: chrono::Utc::now(),
};
let idx = self.vote_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let idx = self
.vote_counter
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
self.vote_log.insert(idx, pair);
// FIFO eviction: remove oldest entries when over cap
let start = self.vote_log_start.load(std::sync::atomic::Ordering::Relaxed);
let start = self
.vote_log_start
.load(std::sync::atomic::Ordering::Relaxed);
if idx.saturating_sub(start) >= self.vote_log_cap {
let evict_to = idx - self.vote_log_cap + 1;
for old_idx in start..evict_to {
self.vote_log.remove(&old_idx);
}
self.vote_log_start.store(evict_to, std::sync::atomic::Ordering::Relaxed);
self.vote_log_start
.store(evict_to, std::sync::atomic::Ordering::Relaxed);
}
}
// Persist vote tracker entry to Firestore (outside borrow block)
self.firestore_put("brain_votes", &vote_doc_id, &vote_doc).await;
self.firestore_put("brain_votes", &vote_doc_id, &vote_doc)
.await;
// Write-through: persist updated memory to Firestore
if let Some(m) = self.memories.get(id) {
if let Ok(body) = serde_json::to_value(m.value()) {
self.firestore_put("brain_memories", &id.to_string(), &body).await;
self.firestore_put("brain_memories", &id.to_string(), &body)
.await;
}
}
@ -863,7 +930,11 @@ impl FirestoreClient {
/// Get preference pairs for training data export (Layer A)
/// Returns pairs accumulated since the given index
pub fn get_preference_pairs(&self, since_index: u64, limit: usize) -> (Vec<PreferencePair>, u64) {
pub fn get_preference_pairs(
&self,
since_index: u64,
limit: usize,
) -> (Vec<PreferencePair>, u64) {
let current = self.vote_counter.load(std::sync::atomic::Ordering::Relaxed);
let mut pairs = Vec::new();
for idx in since_index..current {
@ -904,7 +975,8 @@ impl FirestoreClient {
composite: 1.0,
};
if let Ok(body) = serde_json::to_value(c.value()) {
self.firestore_put("brain_contributors", pseudonym, &body).await;
self.firestore_put("brain_contributors", pseudonym, &body)
.await;
}
}
return Ok(c.clone());
@ -927,15 +999,20 @@ impl FirestoreClient {
is_system,
};
if let Ok(body) = serde_json::to_value(&info) {
self.firestore_put("brain_contributors", pseudonym, &body).await;
self.firestore_put("brain_contributors", pseudonym, &body)
.await;
}
self.contributors.insert(pseudonym.to_string(), info.clone());
self.contributors
.insert(pseudonym.to_string(), info.clone());
Ok(info)
}
/// Detect the embedding dimension from the first stored memory
pub fn detect_embedding_dim(&self) -> Option<usize> {
self.memories.iter().next().map(|e| e.value().embedding.len())
self.memories
.iter()
.next()
.map(|e| e.value().embedding.len())
}
/// Get all memories (for graph building)
@ -974,7 +1051,9 @@ impl FirestoreClient {
/// Get the reputation score for a contributor, if known
pub fn get_contributor_reputation(&self, pseudonym: &str) -> Option<ReputationScore> {
self.contributors.get(pseudonym).map(|c| c.reputation.clone())
self.contributors
.get(pseudonym)
.map(|c| c.reputation.clone())
}
/// Record a contribution: increment count, update uptime, recompute composite
@ -984,22 +1063,17 @@ impl FirestoreClient {
entry.last_active = chrono::Utc::now();
// Grow stake organically through contributions
entry.reputation.stake += 1.0;
crate::reputation::ReputationManager::record_activity(
&mut entry.reputation,
);
crate::reputation::ReputationManager::record_activity(&mut entry.reputation);
// Persist updated contributor
if let Ok(body) = serde_json::to_value(entry.value()) {
self.firestore_put("brain_contributors", pseudonym, &body).await;
self.firestore_put("brain_contributors", pseudonym, &body)
.await;
}
}
}
/// Update contributor reputation based on vote outcome on their content
pub async fn update_reputation_from_vote(
&self,
content_author: &str,
was_upvoted: bool,
) {
pub async fn update_reputation_from_vote(&self, content_author: &str, was_upvoted: bool) {
if let Some(mut entry) = self.contributors.get_mut(content_author) {
crate::reputation::ReputationManager::update_accuracy(
&mut entry.reputation,
@ -1022,11 +1096,8 @@ impl FirestoreClient {
) -> bool {
let mgr = crate::reputation::ReputationManager::new();
if let Some(mut entry) = self.contributors.get_mut(content_author) {
let penalized = mgr.check_poisoning_penalty(
&mut entry.reputation,
downvote_count,
quality,
);
let penalized =
mgr.check_poisoning_penalty(&mut entry.reputation, downvote_count, quality);
if penalized {
if let Ok(body) = serde_json::to_value(entry.value()) {
self.firestore_put("brain_contributors", content_author, &body)
@ -1063,11 +1134,13 @@ impl FirestoreClient {
let id = memory.id;
// Persist memory
if let Ok(body) = serde_json::to_value(&memory) {
self.firestore_put("brain_memories", &id.to_string(), &body).await;
self.firestore_put("brain_memories", &id.to_string(), &body)
.await;
}
// Persist page status
let status_doc = serde_json::json!({ "id": id.to_string(), "status": status });
self.firestore_put("brain_page_status", &id.to_string(), &status_doc).await;
self.firestore_put("brain_page_status", &id.to_string(), &status_doc)
.await;
// Cache
self.memories.insert(id, memory);
self.page_status.insert(id, status);
@ -1094,11 +1167,7 @@ impl FirestoreClient {
}
/// Submit a delta to a page
pub async fn submit_delta(
&self,
page_id: &Uuid,
delta: PageDelta,
) -> Result<(), StoreError> {
pub async fn submit_delta(&self, page_id: &Uuid, delta: PageDelta) -> Result<(), StoreError> {
if !self.memories.contains_key(page_id) {
return Err(StoreError::NotFound(page_id.to_string()));
}
@ -1135,9 +1204,7 @@ impl FirestoreClient {
return Err(StoreError::NotFound(page_id.to_string()));
}
let mut ev = self.page_evidence
.entry(*page_id)
.or_insert_with(Vec::new);
let mut ev = self.page_evidence.entry(*page_id).or_insert_with(Vec::new);
ev.push(evidence);
Ok(ev.len() as u32)
}
@ -1160,7 +1227,11 @@ impl FirestoreClient {
/// Get page evidence and delta counts
pub fn page_counts(&self, page_id: &Uuid) -> (u32, u32) {
let ev = self.page_evidence.get(page_id).map(|e| e.len()).unwrap_or(0) as u32;
let ev = self
.page_evidence
.get(page_id)
.map(|e| e.len())
.unwrap_or(0) as u32;
let dc = self.page_deltas.get(page_id).map(|d| d.len()).unwrap_or(0) as u32;
(ev, dc)
}
@ -1188,7 +1259,8 @@ impl FirestoreClient {
/// Promote a page from Draft to Canonical (cache + Firestore)
pub async fn promote_page(&self, page_id: &Uuid) -> Result<PageStatus, StoreError> {
let current = self.get_page_status(page_id)
let current = self
.get_page_status(page_id)
.ok_or_else(|| StoreError::NotFound(page_id.to_string()))?;
if current != PageStatus::Draft {
return Err(StoreError::Storage(format!(
@ -1202,8 +1274,10 @@ impl FirestoreClient {
}
self.page_status.insert(*page_id, PageStatus::Canonical);
// Persist status change
let status_doc = serde_json::json!({ "id": page_id.to_string(), "status": PageStatus::Canonical });
self.firestore_put("brain_page_status", &page_id.to_string(), &status_doc).await;
let status_doc =
serde_json::json!({ "id": page_id.to_string(), "status": PageStatus::Canonical });
self.firestore_put("brain_page_status", &page_id.to_string(), &status_doc)
.await;
Ok(PageStatus::Canonical)
}
@ -1226,7 +1300,11 @@ impl FirestoreClient {
// ──────────────────────────────────────────────────────────────────
/// Publish a WASM node (cache + Firestore write-through for metadata)
pub async fn publish_node(&self, node: WasmNode, wasm_bytes: Vec<u8>) -> Result<(), StoreError> {
pub async fn publish_node(
&self,
node: WasmNode,
wasm_bytes: Vec<u8>,
) -> Result<(), StoreError> {
if self.wasm_nodes.contains_key(&node.id) {
return Err(StoreError::Storage(format!(
"Node {} already exists (nodes are immutable, use a new version)",
@ -1263,10 +1341,14 @@ impl FirestoreClient {
/// Revoke a node (marks as revoked, does not delete bytes, cache + Firestore)
pub async fn revoke_node(&self, id: &str, contributor: &str) -> Result<(), StoreError> {
{
let mut node = self.wasm_nodes.get_mut(id)
let mut node = self
.wasm_nodes
.get_mut(id)
.ok_or_else(|| StoreError::NotFound(id.to_string()))?;
if node.contributor_id != contributor {
return Err(StoreError::Forbidden("Only original publisher can revoke".into()));
return Err(StoreError::Forbidden(
"Only original publisher can revoke".into(),
));
}
node.revoked = true;
}

96
crates/mcp-brain-server/src/symbolic.rs
View file

@ -254,10 +254,18 @@ pub struct HornClause {
}
impl HornClause {
pub fn new(antecedents: Vec<PredicateType>, consequent: PredicateType, confidence: f64) -> Self {
pub fn new(
antecedents: Vec<PredicateType>,
consequent: PredicateType,
confidence: f64,
) -> Self {
let id = format!(
"rule_{}",
uuid::Uuid::new_v4().to_string().split('-').next().unwrap_or("0")
uuid::Uuid::new_v4()
.to_string()
.split('-')
.next()
.unwrap_or("0")
);
Self {
id,
@ -309,7 +317,10 @@ impl NeuralSymbolicBridge {
// Association is transitive (with decay): if A associated_with B and B associated_with C → A co_occurs_with C
self.rules.push(HornClause::new(
vec![PredicateType::Custom("associated_with".to_string()), PredicateType::Custom("associated_with".to_string())],
vec![
PredicateType::Custom("associated_with".to_string()),
PredicateType::Custom("associated_with".to_string()),
],
PredicateType::Custom("co_occurs_with".to_string()),
0.5,
));
@ -317,14 +328,20 @@ impl NeuralSymbolicBridge {
// Influence chain: if A may_influence B and B may_influence C → A associated_with C
// (demotes from influence to association when chaining — honest decay)
self.rules.push(HornClause::new(
vec![PredicateType::Custom("may_influence".to_string()), PredicateType::Custom("may_influence".to_string())],
vec![
PredicateType::Custom("may_influence".to_string()),
PredicateType::Custom("may_influence".to_string()),
],
PredicateType::Custom("associated_with".to_string()),
0.6,
));
// Cross-type influence: if A may_influence B and B is_type_of C → A associated_with C
self.rules.push(HornClause::new(
vec![PredicateType::Custom("may_influence".to_string()), PredicateType::IsTypeOf],
vec![
PredicateType::Custom("may_influence".to_string()),
PredicateType::IsTypeOf,
],
PredicateType::Custom("associated_with".to_string()),
0.5,
));
@ -369,7 +386,10 @@ impl NeuralSymbolicBridge {
// Influence→prevention: if A may_influence B and B prevents C, then A co_occurs_with C
self.rules.push(HornClause::new(
vec![PredicateType::Custom("may_influence".to_string()), PredicateType::Prevents],
vec![
PredicateType::Custom("may_influence".to_string()),
PredicateType::Prevents,
],
PredicateType::Custom("co_occurs_with".to_string()),
0.4,
));
@ -416,7 +436,9 @@ impl NeuralSymbolicBridge {
let (ref c1, ref ids1, ref cat1) = clusters[i];
let (ref c2, ref ids2, ref cat2) = clusters[j];
if ids1.len() < self.config.min_cluster_size || ids2.len() < self.config.min_cluster_size {
if ids1.len() < self.config.min_cluster_size
|| ids2.len() < self.config.min_cluster_size
{
continue;
}
@ -432,7 +454,11 @@ impl NeuralSymbolicBridge {
let mut merged_evidence = ids1.clone();
merged_evidence.extend_from_slice(ids2);
merged_evidence.truncate(20);
let midpoint: Vec<f32> = c1.iter().zip(c2.iter()).map(|(a, b)| (a + b) / 2.0).collect();
let midpoint: Vec<f32> = c1
.iter()
.zip(c2.iter())
.map(|(a, b)| (a + b) / 2.0)
.collect();
// Use category names instead of cluster sizes for human-readable arguments
let arg1 = cat1.clone();
@ -524,7 +550,10 @@ impl NeuralSymbolicBridge {
// Create pattern-based proposition
let prop = GroundedProposition::new(
PredicateType::SimilarTo.as_str().to_string(),
vec![format!("pattern_{}", memories.len()), "learned_pattern".to_string()],
vec![
format!("pattern_{}", memories.len()),
"learned_pattern".to_string(),
],
centroid.clone(),
*confidence,
memories.clone(),
@ -782,9 +811,9 @@ impl NeuralSymbolicBridge {
// Shared-argument: any argument of pa matches any argument of pb
// (case-insensitive for robustness).
let shared_arg = if !strict_chain {
pa.arguments.iter().any(|a| {
pb.arguments.iter().any(|b| a.eq_ignore_ascii_case(b))
})
pa.arguments
.iter()
.any(|a| pb.arguments.iter().any(|b| a.eq_ignore_ascii_case(b)))
} else {
false
};
@ -807,17 +836,25 @@ impl NeuralSymbolicBridge {
}
} else {
// Shared-argument: take the non-shared args as endpoints.
let shared: Vec<&String> = pa.arguments.iter()
let shared: Vec<&String> = pa
.arguments
.iter()
.filter(|a| pb.arguments.iter().any(|b| a.eq_ignore_ascii_case(b)))
.collect();
let pa_unique: Vec<&String> = pa.arguments.iter()
let pa_unique: Vec<&String> = pa
.arguments
.iter()
.filter(|a| !shared.iter().any(|s| a.eq_ignore_ascii_case(s)))
.collect();
let pb_unique: Vec<&String> = pb.arguments.iter()
let pb_unique: Vec<&String> = pb
.arguments
.iter()
.filter(|b| !shared.iter().any(|s| b.eq_ignore_ascii_case(s)))
.collect();
let first: Option<&String> = pa_unique.first().copied().or(pa.arguments.first());
let last: Option<&String> = pb_unique.first().copied().or(pb.arguments.first());
let first: Option<&String> =
pa_unique.first().copied().or(pa.arguments.first());
let last: Option<&String> =
pb_unique.first().copied().or(pb.arguments.first());
match (first, last) {
(Some(f), Some(l)) => {
if f.eq_ignore_ascii_case(l) {
@ -836,8 +873,7 @@ impl NeuralSymbolicBridge {
continue;
}
let combined_confidence =
rule.confidence * pa.confidence * pb.confidence;
let combined_confidence = rule.confidence * pa.confidence * pb.confidence;
// Require meaningful confidence — no coin-flip inferences
if combined_confidence < 0.4 {
@ -962,9 +998,7 @@ impl NeuralSymbolicBridge {
evidence: Vec<Uuid>,
) -> GroundedProposition {
let prop = GroundedProposition::new(
predicate,
arguments,
embedding,
predicate, arguments, embedding,
0.8, // Default confidence for manually grounded propositions
evidence,
);
@ -1077,7 +1111,13 @@ mod tests {
// cluster_confidence(5) = 1.0 - exp(-1.0) ≈ 0.63
let clusters = vec![(
vec![1.0, 0.0, 0.0, 0.0],
vec![Uuid::new_v4(), Uuid::new_v4(), Uuid::new_v4(), Uuid::new_v4(), Uuid::new_v4()],
vec![
Uuid::new_v4(),
Uuid::new_v4(),
Uuid::new_v4(),
Uuid::new_v4(),
Uuid::new_v4(),
],
"pattern".to_string(),
)];
@ -1136,7 +1176,10 @@ mod tests {
assert!(!inferences.is_empty(), "expected at least one inference");
let inf = &inferences[0];
assert_eq!(inf.conclusion.predicate, "relates_to");
assert_eq!(inf.conclusion.arguments, vec!["A".to_string(), "C".to_string()]);
assert_eq!(
inf.conclusion.arguments,
vec!["A".to_string(), "C".to_string()]
);
assert!(inf.combined_confidence > 0.0);
assert!(bridge.inference_count() > 0);
@ -1145,7 +1188,10 @@ mod tests {
// Running again should produce no new inferences (already derived)
let inferences2 = bridge.run_inference();
assert!(inferences2.is_empty(), "should not re-derive existing conclusions");
assert!(
inferences2.is_empty(),
"should not re-derive existing conclusions"
);
}
#[test]

133
crates/mcp-brain-server/src/tests.rs
View file

@ -37,7 +37,10 @@ mod tests {
let recalled = hopfield.retrieve(&noisy).expect("retrieve failed");
// Should retrieve something close to pattern 0 (first element dominant)
assert!(recalled[0] > recalled[1], "pattern 0 should be dominant in retrieval");
assert!(
recalled[0] > recalled[1],
"pattern 0 should be dominant in retrieval"
);
}
// -----------------------------------------------------------------------
@ -134,7 +137,11 @@ mod tests {
.expect("failed to build MinCut");
let cut_value = mincut.min_cut_value();
assert!(cut_value > 0.0, "min cut value should be > 0, got {}", cut_value);
assert!(
cut_value > 0.0,
"min cut value should be > 0, got {}",
cut_value
);
}
// -----------------------------------------------------------------------
@ -279,10 +286,7 @@ mod tests {
.top_k(&graph, 0, 3)
.expect("forward push should succeed");
assert!(
!results.is_empty(),
"should return PPR results"
);
assert!(!results.is_empty(), "should return PPR results");
// Node 0 as source — it or its immediate neighbors should rank high
let returned_nodes: Vec<usize> = results.iter().map(|(n, _)| *n).collect();
// At least some nodes should be returned
@ -320,14 +324,12 @@ mod tests {
// Verify the transfer with simulated metrics
let verification = engine.verify_transfer(
&source,
&target,
0.8, // source_before
0.79, // source_after (within tolerance)
0.3, // target_before
0.65, // target_after
100, // baseline_cycles
50, // transfer_cycles
&source, &target, 0.8, // source_before
0.79, // source_after (within tolerance)
0.3, // target_before
0.65, // target_after
100, // baseline_cycles
50, // transfer_cycles
);
assert!(
@ -338,10 +340,7 @@ mod tests {
!verification.regressed_source,
"transfer should not regress source"
);
assert!(
verification.promotable,
"verification should be promotable"
);
assert!(verification.promotable, "verification should be promotable");
assert!(
verification.acceleration_factor > 1.0,
"acceleration factor should be > 1.0, got {}",
@ -396,10 +395,16 @@ mod tests {
let verifier = Verifier::new();
let pii_inputs = vec![
("email address", "My email is user@example.com and I need help"),
(
"email address",
"My email is user@example.com and I need help",
),
("phone number", "Call me at 555-867-5309 for details"),
("SSN", "My SSN is 123-45-6789 please keep it safe"),
("credit card", "Card number 4111-1111-1111-1111 expires 12/25"),
(
"credit card",
"Card number 4111-1111-1111-1111 expires 12/25",
),
("IP address", "Server IP is 192.168.1.100 for internal use"),
("AWS key", "AWS key AKIAIOSFODNN7EXAMPLE is exposed"),
("private key", "-----BEGIN PRIVATE KEY----- data here"),
@ -435,7 +440,7 @@ mod tests {
// -----------------------------------------------------------------------
#[test]
fn test_end_to_end_share_pipeline() {
use crate::pipeline::{RvfPipelineInput, build_rvf_container, count_segments};
use crate::pipeline::{build_rvf_container, count_segments, RvfPipelineInput};
use crate::verify::Verifier;
use rvf_crypto::WitnessEntry;
@ -447,30 +452,55 @@ mod tests {
// Step 1: Verify input (should reject due to PII)
let result = verifier.verify_share(title, content, &tags, &embedding);
assert!(result.is_err(), "PII content should be rejected by verify_share");
assert!(
result.is_err(),
"PII content should be rejected by verify_share"
);
// Step 2: Strip PII instead of rejecting
let fields = [("title", title), ("content", content)];
let (stripped, log) = verifier.strip_pii_fields(&fields);
assert!(log.total_redactions >= 2, "should redact email + path");
assert!(!stripped[1].1.contains("admin@example.com"), "email should be redacted");
assert!(
!stripped[1].1.contains("admin@example.com"),
"email should be redacted"
);
assert!(!stripped[1].1.contains("/home/"), "path should be redacted");
// Step 3: Stripped content should pass verification
let clean_title = &stripped[0].1;
let clean_content = &stripped[1].1;
assert!(verifier.verify_share(clean_title, clean_content, &tags, &embedding).is_ok());
assert!(verifier
.verify_share(clean_title, clean_content, &tags, &embedding)
.is_ok());
// Step 4: Build witness chain
let now_ns = 1_000_000_000u64;
let stripped_hash = rvf_crypto::shake256_256(clean_content.as_bytes());
let mut emb_bytes = Vec::with_capacity(embedding.len() * 4);
for v in &embedding { emb_bytes.extend_from_slice(&v.to_le_bytes()); }
for v in &embedding {
emb_bytes.extend_from_slice(&v.to_le_bytes());
}
let emb_hash = rvf_crypto::shake256_256(&emb_bytes);
let entries = vec![
WitnessEntry { prev_hash: [0u8; 32], action_hash: stripped_hash, timestamp_ns: now_ns, witness_type: 0x01 },
WitnessEntry { prev_hash: [0u8; 32], action_hash: emb_hash, timestamp_ns: now_ns, witness_type: 0x02 },
WitnessEntry { prev_hash: [0u8; 32], action_hash: rvf_crypto::shake256_256(b"final"), timestamp_ns: now_ns, witness_type: 0x01 },
WitnessEntry {
prev_hash: [0u8; 32],
action_hash: stripped_hash,
timestamp_ns: now_ns,
witness_type: 0x01,
},
WitnessEntry {
prev_hash: [0u8; 32],
action_hash: emb_hash,
timestamp_ns: now_ns,
witness_type: 0x02,
},
WitnessEntry {
prev_hash: [0u8; 32],
action_hash: rvf_crypto::shake256_256(b"final"),
timestamp_ns: now_ns,
witness_type: 0x01,
},
];
let chain = rvf_crypto::create_witness_chain(&entries);
assert_eq!(chain.len(), 73 * 3);
@ -482,7 +512,8 @@ mod tests {
// Step 6: Build RVF container
let redaction_json = serde_json::to_string(&serde_json::json!({
"entries": [], "total_redactions": log.total_redactions
})).unwrap();
}))
.unwrap();
let input = RvfPipelineInput {
memory_id: "e2e-test-id",
embedding: &embedding,
@ -539,12 +570,21 @@ mod tests {
assert!(!flags.dp_enabled, "dp_enabled should default to false");
assert!(!flags.adversarial, "adversarial should default to false");
assert!(!flags.neg_cache, "neg_cache should default to false");
assert!((flags.dp_epsilon - 1.0).abs() < f64::EPSILON, "dp_epsilon should default to 1.0");
assert!(
(flags.dp_epsilon - 1.0).abs() < f64::EPSILON,
"dp_epsilon should default to 1.0"
);
// Phase 8 AGI defaults — all enabled by default
assert!(flags.sona_enabled, "sona_enabled should default to true");
assert!(flags.gwt_enabled, "gwt_enabled should default to true");
assert!(flags.temporal_enabled, "temporal_enabled should default to true");
assert!(flags.meta_learning_enabled, "meta_learning_enabled should default to true");
assert!(
flags.temporal_enabled,
"temporal_enabled should default to true"
);
assert!(
flags.meta_learning_enabled,
"meta_learning_enabled should default to true"
);
}
// -----------------------------------------------------------------------
@ -562,8 +602,10 @@ mod tests {
// Stats should reflect the trajectory
let stats = sona.stats();
assert!(stats.trajectories_buffered >= 1 || stats.trajectories_dropped == 0,
"trajectory should be buffered or processed");
assert!(
stats.trajectories_buffered >= 1 || stats.trajectories_dropped == 0,
"trajectory should be buffered or processed"
);
// Pattern search should not crash (may return empty before learning)
let patterns = sona.find_patterns(&query, 5);
@ -609,7 +651,8 @@ mod tests {
// -----------------------------------------------------------------------
#[test]
fn test_delta_stream_temporal() {
let mut stream = ruvector_delta_core::DeltaStream::<ruvector_delta_core::VectorDelta>::for_vectors(4);
let mut stream =
ruvector_delta_core::DeltaStream::<ruvector_delta_core::VectorDelta>::for_vectors(4);
// Push 3 deltas at different timestamps
let d1 = VectorDelta::from_dense(vec![1.0, 0.0, 0.0, 0.0]);
@ -621,7 +664,11 @@ mod tests {
// Query time range
let range = stream.get_time_range(1500, 3500);
assert_eq!(range.len(), 2, "should find 2 deltas in time range 1500-3500");
assert_eq!(
range.len(),
2,
"should find 2 deltas in time range 1500-3500"
);
// Full range should return all 3
let all = stream.get_time_range(0, 10000);
@ -638,7 +685,10 @@ mod tests {
// Meta-learning health should be available without panicking
let health = engine.meta_health();
// Fresh engine has no observations, so consecutive_plateaus = 0
assert_eq!(health.consecutive_plateaus, 0, "no plateaus on fresh engine");
assert_eq!(
health.consecutive_plateaus, 0,
"no plateaus on fresh engine"
);
// Regret summary should work on empty state
let regret = engine.regret_summary();
@ -683,11 +733,12 @@ mod tests {
#[test]
fn test_midstream_attractor_too_short() {
// Less than 10 points → None
let embeddings: Vec<Vec<f32>> = (0..5)
.map(|i| vec![i as f32; 8])
.collect();
let embeddings: Vec<Vec<f32>> = (0..5).map(|i| vec![i as f32; 8]).collect();
let result = crate::midstream::analyze_category_attractor(&embeddings);
assert!(result.is_none(), "should return None for too-short trajectory");
assert!(
result.is_none(),
"should return None for too-short trajectory"
);
}
#[test]

231
crates/mcp-brain-server/src/trainer.rs
View file

@ -132,9 +132,7 @@ pub struct BrainTrainer {
impl BrainTrainer {
pub fn new(config: TrainerConfig) -> Self {
let http_client = reqwest::Client::builder()
.user_agent(
"ruvector-brain-trainer/1.0 (https://pi.ruv.io; benevolent-discovery)",
)
.user_agent("ruvector-brain-trainer/1.0 (https://pi.ruv.io; benevolent-discovery)")
.timeout(std::time::Duration::from_secs(30))
.build()
.expect("HTTP client");
@ -191,10 +189,7 @@ impl BrainTrainer {
}
// Rate limiting between domains
tokio::time::sleep(std::time::Duration::from_millis(
self.config.api_delay_ms,
))
.await;
tokio::time::sleep(std::time::Duration::from_millis(self.config.api_delay_ms)).await;
}
report.discoveries_found = all_discoveries.len();
@ -230,10 +225,7 @@ impl BrainTrainer {
}
/// Discover patterns in a specific domain by fetching from public APIs
async fn discover_domain(
&self,
domain: &DiscoveryDomain,
) -> Result<Vec<Discovery>, String> {
async fn discover_domain(&self, domain: &DiscoveryDomain) -> Result<Vec<Discovery>, String> {
match domain {
DiscoveryDomain::SpaceScience => self.discover_space().await,
DiscoveryDomain::EarthScience => self.discover_earth().await,
@ -260,32 +252,21 @@ impl BrainTrainer {
if let Some(planets) = data.as_array() {
let masses: Vec<f64> = planets
.iter()
.filter_map(|p| {
p.get("pl_bmassj").and_then(|v| v.as_f64())
})
.filter_map(|p| p.get("pl_bmassj").and_then(|v| v.as_f64()))
.collect();
if !masses.is_empty() {
let mean =
masses.iter().sum::<f64>() / masses.len() as f64;
let variance = masses
.iter()
.map(|m| (m - mean).powi(2))
.sum::<f64>()
let mean = masses.iter().sum::<f64>() / masses.len() as f64;
let variance = masses.iter().map(|m| (m - mean).powi(2)).sum::<f64>()
/ masses.len() as f64;
let std_dev = variance.sqrt();
for planet in planets {
if let (Some(name), Some(mass)) = (
planet
.get("pl_name")
.and_then(|v| v.as_str()),
planet
.get("pl_bmassj")
.and_then(|v| v.as_f64()),
planet.get("pl_name").and_then(|v| v.as_str()),
planet.get("pl_bmassj").and_then(|v| v.as_f64()),
) {
let z = (mass - mean).abs()
/ std_dev.max(0.001);
let z = (mass - mean).abs() / std_dev.max(0.001);
if z > 2.5 {
let ecc = planet
.get("pl_orbeccen")
@ -326,8 +307,7 @@ impl BrainTrainer {
],
confidence: (z / 5.0).min(0.99),
data_points: planets.len(),
source_api: "NASA Exoplanet Archive"
.into(),
source_api: "NASA Exoplanet Archive".into(),
timestamp: Utc::now(),
witness_hash: None,
});
@ -351,11 +331,7 @@ impl BrainTrainer {
outliers detected (>2.5\u{03c3}).",
planets.len()
),
tags: vec![
"space".into(),
"exoplanet".into(),
"population".into(),
],
tags: vec!["space".into(), "exoplanet".into(), "population".into()],
confidence: 0.90,
data_points: planets.len(),
source_api: "NASA Exoplanet Archive".into(),
@ -376,9 +352,8 @@ impl BrainTrainer {
);
match self.fetch_json(&neo_url).await {
Ok(data) => {
if let Some(neo_objects) = data
.get("near_earth_objects")
.and_then(|n| n.as_object())
if let Some(neo_objects) =
data.get("near_earth_objects").and_then(|n| n.as_object())
{
for (_date, objects) in neo_objects {
if let Some(arr) = objects.as_array() {
@ -388,18 +363,14 @@ impl BrainTrainer {
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let hazardous = neo
.get(
"is_potentially_hazardous_asteroid",
)
.get("is_potentially_hazardous_asteroid")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let miss_km = neo
.get("close_approach_data")
.and_then(|c| c.as_array())
.and_then(|a| a.first())
.and_then(|a| {
a.get("miss_distance")
})
.and_then(|a| a.get("miss_distance"))
.and_then(|m| m.get("kilometers"))
.and_then(|k| k.as_str())
.and_then(|s| s.parse::<f64>().ok())
@ -408,42 +379,27 @@ impl BrainTrainer {
.get("close_approach_data")
.and_then(|c| c.as_array())
.and_then(|a| a.first())
.and_then(|a| {
a.get("relative_velocity")
})
.and_then(|v| {
v.get("kilometers_per_hour")
})
.and_then(|a| a.get("relative_velocity"))
.and_then(|v| v.get("kilometers_per_hour"))
.and_then(|k| k.as_str())
.and_then(|s| s.parse::<f64>().ok())
.unwrap_or(0.0);
let diameter_max = neo
.get("estimated_diameter")
.and_then(|d| d.get("meters"))
.and_then(|m| {
m.get("estimated_diameter_max")
})
.and_then(|m| m.get("estimated_diameter_max"))
.and_then(|v| v.as_f64())
.unwrap_or(0.0);
// Only report close or hazardous objects
if hazardous || miss_km < 5_000_000.0 {
let confidence = if hazardous {
0.95
} else {
0.80
};
let confidence = if hazardous { 0.95 } else { 0.80 };
discoveries.push(Discovery {
id: Uuid::new_v4(),
domain:
DiscoveryDomain::SpaceScience,
domain: DiscoveryDomain::SpaceScience,
title: format!(
"NEO close approach: {name}{}",
if hazardous {
" [HAZARDOUS]"
} else {
""
}
if hazardous { " [HAZARDOUS]" } else { "" }
),
content: format!(
"Asteroid {name} passes \
@ -480,10 +436,7 @@ impl BrainTrainer {
Err(e) => tracing::warn!("NASA NEO API: {e}"),
}
tokio::time::sleep(std::time::Duration::from_millis(
self.config.api_delay_ms,
))
.await;
tokio::time::sleep(std::time::Duration::from_millis(self.config.api_delay_ms)).await;
// --- NOAA SWPC solar weather (X-ray flare events) ---
let solar_url = "https://services.swpc.noaa.gov/json/goes/\
@ -510,20 +463,13 @@ impl BrainTrainer {
.unwrap_or(0.0);
// Only report M- and X-class flares (significant)
let is_significant = class.starts_with('M')
|| class.starts_with('X');
let is_significant = class.starts_with('M') || class.starts_with('X');
if is_significant {
let confidence = if class.starts_with('X') {
0.98
} else {
0.85
};
let confidence = if class.starts_with('X') { 0.98 } else { 0.85 };
discoveries.push(Discovery {
id: Uuid::new_v4(),
domain: DiscoveryDomain::SpaceScience,
title: format!(
"Solar flare: {class}-class event"
),
title: format!("Solar flare: {class}-class event"),
content: format!(
"{class}-class solar X-ray flare \
detected. Begin: {begin}, peak: \
@ -559,33 +505,21 @@ impl BrainTrainer {
Err(e) => tracing::warn!("NOAA SWPC solar API: {e}"),
}
tokio::time::sleep(std::time::Duration::from_millis(
self.config.api_delay_ms,
))
.await;
tokio::time::sleep(std::time::Duration::from_millis(self.config.api_delay_ms)).await;
// --- LIGO/GraceDB gravitational wave events ---
let gw_url = "https://gracedb.ligo.org/api/superevents/\
?query=far+%3C+1e-6&format=json&count=10";
match self.fetch_json(gw_url).await {
Ok(data) => {
if let Some(events) = data
.get("superevents")
.and_then(|s| s.as_array())
{
if let Some(events) = data.get("superevents").and_then(|s| s.as_array()) {
for event in events.iter().take(10) {
let superevent_id = event
.get("superevent_id")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let far = event
.get("far")
.and_then(|v| v.as_f64())
.unwrap_or(1.0);
let t_start = event
.get("t_start")
.and_then(|v| v.as_f64())
.unwrap_or(0.0);
let far = event.get("far").and_then(|v| v.as_f64()).unwrap_or(1.0);
let t_start = event.get("t_start").and_then(|v| v.as_f64()).unwrap_or(0.0);
let category = event
.get("category")
.and_then(|v| v.as_str())
@ -599,18 +533,12 @@ impl BrainTrainer {
.and_then(|v| v.as_str())
.unwrap_or("");
let confidence =
(1.0 - far.log10().abs() / 20.0).clamp(
0.70,
0.99,
);
let confidence = (1.0 - far.log10().abs() / 20.0).clamp(0.70, 0.99);
discoveries.push(Discovery {
id: Uuid::new_v4(),
domain: DiscoveryDomain::SpaceScience,
title: format!(
"Gravitational wave: {superevent_id}"
),
title: format!("Gravitational wave: {superevent_id}"),
content: format!(
"GW superevent {superevent_id} \
(category: {category}). False alarm \
@ -650,22 +578,15 @@ impl BrainTrainer {
summary/significant_month.geojson";
match self.fetch_json(url).await {
Ok(data) => {
if let Some(features) =
data.get("features").and_then(|f| f.as_array())
{
if let Some(features) = data.get("features").and_then(|f| f.as_array()) {
for quake in features {
let props = quake
.get("properties")
.unwrap_or(&serde_json::Value::Null);
let props = quake.get("properties").unwrap_or(&serde_json::Value::Null);
let geo = quake
.get("geometry")
.and_then(|g| g.get("coordinates"))
.and_then(|c| c.as_array());
let mag = props
.get("mag")
.and_then(|v| v.as_f64())
.unwrap_or(0.0);
let mag = props.get("mag").and_then(|v| v.as_f64()).unwrap_or(0.0);
let place = props
.get("place")
.and_then(|v| v.as_str())
@ -680,9 +601,7 @@ impl BrainTrainer {
discoveries.push(Discovery {
id: Uuid::new_v4(),
domain: DiscoveryDomain::EarthScience,
title: format!(
"M{mag:.1} earthquake: {place}"
),
title: format!("M{mag:.1} earthquake: {place}"),
content: format!(
"Significant M{mag:.1} earthquake at \
{place}, depth {depth:.1} km. {}",
@ -722,9 +641,7 @@ impl BrainTrainer {
land_ocean/1/3/1850-2026.json";
match self.fetch_json(url).await {
Ok(data) => {
if let Some(obj) =
data.get("data").and_then(|d| d.as_object())
{
if let Some(obj) = data.get("data").and_then(|d| d.as_object()) {
let mut temps: Vec<(String, f64)> = obj
.iter()
.filter_map(|(k, v)| {
@ -737,14 +654,8 @@ impl BrainTrainer {
temps.sort_by(|a, b| a.0.cmp(&b.0));
if let Some(latest) = temps.last() {
let recent: Vec<f64> = temps
.iter()
.rev()
.take(10)
.map(|t| t.1)
.collect();
let avg_recent = recent.iter().sum::<f64>()
/ recent.len() as f64;
let recent: Vec<f64> = temps.iter().rev().take(10).map(|t| t.1).collect();
let avg_recent = recent.iter().sum::<f64>() / recent.len() as f64;
discoveries.push(Discovery {
id: Uuid::new_v4(),
@ -759,13 +670,12 @@ impl BrainTrainer {
anomaly: +{:.2}\u{00b0}C (period: {}). \
10-period average: +{:.2}\u{00b0}C. \
Dataset spans {} data points from 1850.",
latest.1, latest.0, avg_recent, temps.len()
latest.1,
latest.0,
avg_recent,
temps.len()
),
tags: vec![
"climate".into(),
"temperature".into(),
"anomaly".into(),
],
tags: vec!["climate".into(), "temperature".into(), "anomaly".into()],
confidence: 0.95,
data_points: temps.len(),
source_api: "NOAA NCEI".into(),
@ -778,10 +688,7 @@ impl BrainTrainer {
Err(e) => tracing::warn!("NOAA API: {e}"),
}
tokio::time::sleep(std::time::Duration::from_millis(
self.config.api_delay_ms,
))
.await;
tokio::time::sleep(std::time::Duration::from_millis(self.config.api_delay_ms)).await;
// --- NOAA OISST sea surface temperature anomalies ---
// Uses NOAA ERDDAP for ocean temperature data (Optimum
@ -818,15 +725,9 @@ impl BrainTrainer {
.unwrap_or(std::cmp::Ordering::Equal)
})
.unwrap();
let mean_anom = anomalies
.iter()
.map(|a| a.2)
.sum::<f64>()
/ anomalies.len() as f64;
let positive_count = anomalies
.iter()
.filter(|a| a.2 > 0.0)
.count();
let mean_anom =
anomalies.iter().map(|a| a.2).sum::<f64>() / anomalies.len() as f64;
let positive_count = anomalies.iter().filter(|a| a.2 > 0.0).count();
discoveries.push(Discovery {
id: Uuid::new_v4(),
@ -848,9 +749,7 @@ impl BrainTrainer {
max_anom.2,
max_anom.0,
max_anom.1,
positive_count as f64
/ anomalies.len() as f64
* 100.0,
positive_count as f64 / anomalies.len() as f64 * 100.0,
len = anomalies.len()
),
tags: vec![
@ -884,9 +783,7 @@ impl BrainTrainer {
&sort=cited_by_count:desc&per_page=10";
match self.fetch_json(url).await {
Ok(data) => {
if let Some(results) =
data.get("results").and_then(|r| r.as_array())
{
if let Some(results) = data.get("results").and_then(|r| r.as_array()) {
for work in results.iter().take(5) {
let title = work
.get("title")
@ -902,25 +799,18 @@ impl BrainTrainer {
.unwrap_or(0);
if cited > 50 {
let truncated =
&title[..title.len().min(80)];
let truncated = &title[..title.len().min(80)];
discoveries.push(Discovery {
id: Uuid::new_v4(),
domain: DiscoveryDomain::AcademicResearch,
title: format!(
"High-impact AI paper: {truncated}"
),
title: format!("High-impact AI paper: {truncated}"),
content: format!(
"\"{title}\" ({year}) — {cited} \
citations. Rapidly cited paper \
indicating significant research \
impact in artificial intelligence."
),
tags: vec![
"academic".into(),
"ai".into(),
"high-impact".into(),
],
tags: vec!["academic".into(), "ai".into(), "high-impact".into()],
confidence: 0.85,
data_points: results.len(),
source_api: "OpenAlex".into(),
@ -953,13 +843,10 @@ impl BrainTrainer {
}
/// Ingest a discovery into the brain via REST API
async fn ingest_to_brain(
&self,
discovery: &Discovery,
brain_url: &str,
) -> Result<(), String> {
async fn ingest_to_brain(&self, discovery: &Discovery, brain_url: &str) -> Result<(), String> {
// Get challenge nonce
let nonce_resp: serde_json::Value = self.http_client
let nonce_resp: serde_json::Value = self
.http_client
.get(format!("{brain_url}/v1/challenge"))
.send()
.await
@ -980,7 +867,8 @@ impl BrainTrainer {
"tags": discovery.tags,
});
let resp = self.http_client
let resp = self
.http_client
.post(format!("{brain_url}/v1/memories"))
.header("X-Challenge-Nonce", nonce)
.json(&body)
@ -999,10 +887,7 @@ impl BrainTrainer {
}
/// Fetch JSON from a URL with error handling
async fn fetch_json(
&self,
url: &str,
) -> Result<serde_json::Value, String> {
async fn fetch_json(&self, url: &str) -> Result<serde_json::Value, String> {
self.http_client
.get(url)
.send()

157
crates/mcp-brain-server/src/types.rs
View file

@ -187,7 +187,10 @@ pub struct BetaParams {
impl BetaParams {
pub fn new() -> Self {
Self { alpha: 1.0, beta: 1.0 }
Self {
alpha: 1.0,
beta: 1.0,
}
}
pub fn mean(&self) -> f64 {
@ -312,7 +315,11 @@ pub struct PartitionResultCompact {
impl From<PartitionResult> for PartitionResultCompact {
fn from(r: PartitionResult) -> Self {
Self {
clusters: r.clusters.iter().map(KnowledgeClusterCompact::from).collect(),
clusters: r
.clusters
.iter()
.map(KnowledgeClusterCompact::from)
.collect(),
cut_value: r.cut_value,
edge_strengths: r.edge_strengths,
total_memories: r.total_memories,
@ -446,7 +453,9 @@ pub struct PartitionQuery {
pub force: bool,
}
fn default_compact() -> bool { true }
fn default_compact() -> bool {
true
}
#[derive(Debug, Serialize)]
pub struct HealthResponse {
@ -529,8 +538,12 @@ pub struct ConsciousnessComputeRequest {
pub partition_mask: Option<u64>,
}
fn default_algo() -> String { "auto".into() }
fn default_phi_threshold() -> f64 { 1e-6 }
fn default_algo() -> String {
"auto".into()
}
fn default_phi_threshold() -> f64 {
1e-6
}
/// Response for POST /v1/consciousness/compute
#[derive(Debug, Serialize)]
@ -670,10 +683,16 @@ impl LoraSubmission {
let expected_down = self.hidden_dim * self.rank;
let expected_up = self.rank * self.hidden_dim;
if self.down_proj.len() != expected_down {
return Err(format!("down_proj shape: expected {expected_down}, got {}", self.down_proj.len()));
return Err(format!(
"down_proj shape: expected {expected_down}, got {}",
self.down_proj.len()
));
}
if self.up_proj.len() != expected_up {
return Err(format!("up_proj shape: expected {expected_up}, got {}", self.up_proj.len()));
return Err(format!(
"up_proj shape: expected {expected_up}, got {}",
self.up_proj.len()
));
}
for (i, &v) in self.down_proj.iter().chain(self.up_proj.iter()).enumerate() {
if v.is_nan() || v.is_infinite() {
@ -686,7 +705,9 @@ impl LoraSubmission {
let down_norm: f32 = self.down_proj.iter().map(|x| x * x).sum::<f32>().sqrt();
let up_norm: f32 = self.up_proj.iter().map(|x| x * x).sum::<f32>().sqrt();
if down_norm > 100.0 || up_norm > 100.0 {
return Err(format!("Norm too large: down={down_norm:.1}, up={up_norm:.1}"));
return Err(format!(
"Norm too large: down={down_norm:.1}, up={up_norm:.1}"
));
}
if self.evidence_count < 5 {
return Err(format!("Insufficient evidence: {}", self.evidence_count));
@ -771,10 +792,25 @@ pub struct EvidenceLink {
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum EvidenceType {
TestPass { test_name: String, repo: String, commit_hash: String },
BuildSuccess { pipeline_url: String, commit_hash: String },
MetricImproval { metric_name: String, before: f64, after: f64 },
PeerReview { reviewer: String, direction: VoteDirection, score: f64 },
TestPass {
test_name: String,
repo: String,
commit_hash: String,
},
BuildSuccess {
pipeline_url: String,
commit_hash: String,
},
MetricImproval {
metric_name: String,
before: f64,
after: f64,
},
PeerReview {
reviewer: String,
direction: VoteDirection,
score: f64,
},
}
/// A delta entry modifying a canonical page
@ -1046,7 +1082,9 @@ impl LoraFederationStore {
let mut accepted_indices: Vec<usize> = Vec::new();
for (idx, (sub, _, rep)) in self.pending.iter().enumerate() {
let params: Vec<f32> = sub.down_proj.iter()
let params: Vec<f32> = sub
.down_proj
.iter()
.chain(sub.up_proj.iter())
.copied()
.collect();
@ -1066,26 +1104,33 @@ impl LoraFederationStore {
}
// Compute per-parameter median
let medians: Vec<f32> = all_params.iter().map(|vals| {
let mut sorted = vals.clone();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
if sorted.len() % 2 == 0 {
(sorted[sorted.len() / 2 - 1] + sorted[sorted.len() / 2]) / 2.0
} else {
sorted[sorted.len() / 2]
}
}).collect();
let medians: Vec<f32> = all_params
.iter()
.map(|vals| {
let mut sorted = vals.clone();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
if sorted.len() % 2 == 0 {
(sorted[sorted.len() / 2 - 1] + sorted[sorted.len() / 2]) / 2.0
} else {
sorted[sorted.len() / 2]
}
})
.collect();
// Compute MAD (Median Absolute Deviation) per parameter
let mads: Vec<f32> = all_params.iter().zip(medians.iter()).map(|(vals, &med)| {
let mut devs: Vec<f32> = vals.iter().map(|&v| (v - med).abs()).collect();
devs.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
if devs.len() % 2 == 0 {
(devs[devs.len() / 2 - 1] + devs[devs.len() / 2]) / 2.0
} else {
devs[devs.len() / 2]
}
}).collect();
let mads: Vec<f32> = all_params
.iter()
.zip(medians.iter())
.map(|(vals, &med)| {
let mut devs: Vec<f32> = vals.iter().map(|&v| (v - med).abs()).collect();
devs.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
if devs.len() % 2 == 0 {
(devs[devs.len() / 2 - 1] + devs[devs.len() / 2]) / 2.0
} else {
devs[devs.len() / 2]
}
})
.collect();
// Reputation-weighted trimmed mean: exclude params >3*MAD from median
let mut result = vec![0.0f32; total_params];
@ -1094,7 +1139,9 @@ impl LoraFederationStore {
// Iterate only over accepted submissions with correct weight indexing
for (weight_idx, &pending_idx) in accepted_indices.iter().enumerate() {
let (sub, _, _) = &self.pending[pending_idx];
let params: Vec<f32> = sub.down_proj.iter()
let params: Vec<f32> = sub
.down_proj
.iter()
.chain(sub.up_proj.iter())
.copied()
.collect();
@ -1154,7 +1201,10 @@ impl LoraFederationStore {
let curr = self.consensus.as_ref()?;
let prev = self.previous_consensus.as_ref()?;
let d: f32 = curr.down_proj.iter().zip(prev.down_proj.iter())
let d: f32 = curr
.down_proj
.iter()
.zip(prev.down_proj.iter())
.chain(curr.up_proj.iter().zip(prev.up_proj.iter()))
.map(|(a, b)| (a - b).powi(2))
.sum();
@ -1168,7 +1218,9 @@ impl LoraFederationStore {
"epoch": self.epoch,
"consensus": consensus,
});
store.firestore_put_public("brain_lora", "consensus", &doc).await;
store
.firestore_put_public("brain_lora", "consensus", &doc)
.await;
}
}
@ -1237,7 +1289,9 @@ impl NonceStore {
/// Periodic cleanup of expired nonces
fn maybe_cleanup(&self) {
let count = self.ops_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let count = self
.ops_counter
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
if count % 500 != 0 {
return;
}
@ -1340,7 +1394,8 @@ pub struct AppState {
pub cognitive: std::sync::Arc<parking_lot::RwLock<crate::cognitive::CognitiveEngine>>,
pub drift: std::sync::Arc<parking_lot::RwLock<crate::drift::DriftMonitor>>,
pub aggregator: std::sync::Arc<crate::aggregate::ByzantineAggregator>,
pub domain_engine: std::sync::Arc<parking_lot::RwLock<ruvector_domain_expansion::DomainExpansionEngine>>,
pub domain_engine:
std::sync::Arc<parking_lot::RwLock<ruvector_domain_expansion::DomainExpansionEngine>>,
pub sona: std::sync::Arc<parking_lot::RwLock<sona::SonaEngine>>,
pub lora_federation: std::sync::Arc<parking_lot::RwLock<LoraFederationStore>>,
/// RuvLLM embedding engine (HashEmbedder + RlmEmbedder)
@ -1354,9 +1409,13 @@ pub struct AppState {
/// RVF feature flags read once at startup (avoids per-request env::var calls)
pub rvf_flags: RvfFeatureFlags,
/// Global Workspace Theory attention layer for memory selection (ADR-075 AGI)
pub workspace: std::sync::Arc<parking_lot::RwLock<ruvector_nervous_system::routing::workspace::GlobalWorkspace>>,
pub workspace: std::sync::Arc<
parking_lot::RwLock<ruvector_nervous_system::routing::workspace::GlobalWorkspace>,
>,
/// Temporal delta tracking for knowledge evolution (ADR-075 AGI)
pub delta_stream: std::sync::Arc<parking_lot::RwLock<ruvector_delta_core::DeltaStream<ruvector_delta_core::VectorDelta>>>,
pub delta_stream: std::sync::Arc<
parking_lot::RwLock<ruvector_delta_core::DeltaStream<ruvector_delta_core::VectorDelta>>,
>,
/// Cached verifier (holds compiled PiiStripper regexes — avoids recompiling per request)
pub verifier: std::sync::Arc<parking_lot::RwLock<crate::verify::Verifier>>,
/// Negative cost fuse: when true, reject all writes (Firestore/GCS errors spiking)
@ -1366,15 +1425,28 @@ pub struct AppState {
/// Nanosecond-precision background scheduler (Phase 9b)
pub nano_scheduler: std::sync::Arc<nanosecond_scheduler::Scheduler>,
/// Per-category Lyapunov exponent results from attractor analysis (Phase 9c)
pub attractor_results: std::sync::Arc<parking_lot::RwLock<std::collections::HashMap<String, temporal_attractor_studio::LyapunovResult>>>,
pub attractor_results: std::sync::Arc<
parking_lot::RwLock<
std::collections::HashMap<String, temporal_attractor_studio::LyapunovResult>,
>,
>,
/// Temporal neural solver with certified predictions (Phase 9d)
/// Note: Only available on x86_64 platforms (requires SIMD)
#[cfg(feature = "x86-simd")]
pub temporal_solver: std::sync::Arc<parking_lot::RwLock<temporal_neural_solver::TemporalSolver>>,
pub temporal_solver:
std::sync::Arc<parking_lot::RwLock<temporal_neural_solver::TemporalSolver>>,
#[cfg(not(feature = "x86-simd"))]
pub temporal_solver: std::sync::Arc<parking_lot::RwLock<TemporalSolverStub>>,
/// Meta-cognitive recursive learning with safety bounds (Phase 9e)
pub strange_loop: std::sync::Arc<parking_lot::RwLock<strange_loop::StrangeLoop<strange_loop::ScalarReasoner, strange_loop::SimpleCritic, strange_loop::SafeReflector>>>,
pub strange_loop: std::sync::Arc<
parking_lot::RwLock<
strange_loop::StrangeLoop<
strange_loop::ScalarReasoner,
strange_loop::SimpleCritic,
strange_loop::SafeReflector,
>,
>,
>,
/// Active SSE sessions: session ID -> sender channel for streaming responses
pub sessions: std::sync::Arc<dashmap::DashMap<String, tokio::sync::mpsc::Sender<String>>>,
// ── Neural-Symbolic + Internal Voice (ADR-110) ──
@ -1398,7 +1470,8 @@ pub struct AppState {
/// Resend email notifier (ADR-125) — None if RESEND_API_KEY not set
pub notifier: Option<crate::notify::ResendNotifier>,
/// Cached /v1/status response to avoid recomputing expensive aggregates every call
pub cached_status: std::sync::Arc<parking_lot::RwLock<Option<(std::time::Instant, StatusResponse)>>>,
pub cached_status:
std::sync::Arc<parking_lot::RwLock<Option<(std::time::Instant, StatusResponse)>>>,
/// GitHub Gist publisher for autonomous discoveries — None if GITHUB_GIST_PAT not set
pub gist_publisher: Option<std::sync::Arc<crate::gist::GistPublisher>>,
/// Semaphore to limit concurrent pipeline/optimize requests (prevents scheduler thundering herd)

84
crates/mcp-brain-server/src/verify.rs
View file

@ -10,7 +10,11 @@ pub enum VerifyError {
#[error("Invalid embedding: {0}")]
InvalidEmbedding(String),
#[error("Content too large: {field} is {size} bytes, max {max}")]
ContentTooLarge { field: String, size: usize, max: usize },
ContentTooLarge {
field: String,
size: usize,
max: usize,
},
#[error("Invalid witness hash: {0}")]
InvalidWitness(String),
#[error("Signature verification failed: {0}")]
@ -154,14 +158,10 @@ impl Verifier {
}
for (i, &val) in embedding.iter().enumerate() {
if val.is_nan() {
return Err(VerifyError::InvalidEmbedding(format!(
"NaN at index {i}"
)));
return Err(VerifyError::InvalidEmbedding(format!("NaN at index {i}")));
}
if val.is_infinite() {
return Err(VerifyError::InvalidEmbedding(format!(
"Inf at index {i}"
)));
return Err(VerifyError::InvalidEmbedding(format!("Inf at index {i}")));
}
if val.abs() > self.max_embedding_magnitude {
return Err(VerifyError::InvalidEmbedding(format!(
@ -181,14 +181,15 @@ impl Verifier {
message: &[u8],
signature_bytes: &[u8; 64],
) -> Result<(), VerifyError> {
use ed25519_dalek::{Signature, VerifyingKey};
use ed25519_dalek::Verifier as _;
use ed25519_dalek::{Signature, VerifyingKey};
let key = VerifyingKey::from_bytes(public_key_bytes)
.map_err(|e| VerifyError::SignatureFailed(format!("Invalid public key: {e}")))?;
let sig = Signature::from_bytes(signature_bytes);
key.verify(message, &sig)
.map_err(|e| VerifyError::SignatureFailed(format!("Ed25519 verification failed: {e}")))?;
key.verify(message, &sig).map_err(|e| {
VerifyError::SignatureFailed(format!("Ed25519 verification failed: {e}"))
})?;
Ok(())
}
@ -200,7 +201,10 @@ impl Verifier {
steps: &[&str],
expected_hash: &str,
) -> Result<(), VerifyError> {
use sha3::{Shake256, digest::{Update, ExtendableOutput, XofReader}};
use sha3::{
digest::{ExtendableOutput, Update, XofReader},
Shake256,
};
let mut current = [0u8; 32];
@ -231,11 +235,7 @@ impl Verifier {
/// Verify a SHAKE-256 content hash matches data.
/// Delegates to rvf_crypto::shake256_256 with constant-time comparison.
pub fn verify_content_hash(
&self,
data: &[u8],
expected_hex: &str,
) -> Result<(), VerifyError> {
pub fn verify_content_hash(&self, data: &[u8], expected_hex: &str) -> Result<(), VerifyError> {
let computed_bytes = rvf_crypto::shake256_256(data);
let computed = hex::encode(computed_bytes);
let equal = subtle::ConstantTimeEq::ct_eq(computed.as_bytes(), expected_hex.as_bytes());
@ -261,10 +261,7 @@ impl Verifier {
/// Check whether embedding distances indicate an adversarial (degenerate) distribution.
/// Returns true if the distribution is too uniform to trust centroid routing.
/// Uses rvf_runtime::is_degenerate_distribution (CV < 0.05 threshold).
pub fn verify_embedding_not_adversarial(
distances: &[f32],
n_probe: usize,
) -> bool {
pub fn verify_embedding_not_adversarial(distances: &[f32], n_probe: usize) -> bool {
rvf_runtime::is_degenerate_distribution(distances, n_probe)
}
}
@ -282,27 +279,42 @@ mod tests {
#[test]
fn test_verify_clean_data() {
let v = Verifier::new();
assert!(v.verify_share("Good title", "Clean content", &["tag1".into()], &[0.1, 0.2, 0.3]).is_ok());
assert!(v
.verify_share(
"Good title",
"Clean content",
&["tag1".into()],
&[0.1, 0.2, 0.3]
)
.is_ok());
}
#[test]
fn test_reject_pii() {
let v = Verifier::new();
assert!(v.verify_share("Has /home/user path", "content", &[], &[0.1]).is_err());
assert!(v
.verify_share("Has /home/user path", "content", &[], &[0.1])
.is_err());
// PiiStripper requires sk- followed by 20+ alphanums (realistic API key length)
assert!(v.verify_share("title", "has sk-abcdefghijklmnopqrstuvwxyz", &[], &[0.1]).is_err());
assert!(v
.verify_share("title", "has sk-abcdefghijklmnopqrstuvwxyz", &[], &[0.1])
.is_err());
}
#[test]
fn test_reject_nan_embedding() {
let v = Verifier::new();
assert!(v.verify_share("title", "content", &[], &[0.1, f32::NAN, 0.3]).is_err());
assert!(v
.verify_share("title", "content", &[], &[0.1, f32::NAN, 0.3])
.is_err());
}
#[test]
fn test_reject_inf_embedding() {
let v = Verifier::new();
assert!(v.verify_share("title", "content", &[], &[0.1, f32::INFINITY, 0.3]).is_err());
assert!(v
.verify_share("title", "content", &[], &[0.1, f32::INFINITY, 0.3])
.is_err());
}
#[test]
@ -323,7 +335,10 @@ mod tests {
fn test_verify_witness_chain() {
let v = Verifier::new();
// Build expected hash from steps
use sha3::{Shake256, digest::{Update, ExtendableOutput, XofReader}};
use sha3::{
digest::{ExtendableOutput, Update, XofReader},
Shake256,
};
let steps = ["pii_strip", "embed", "share"];
let mut current = [0u8; 32];
for step in &steps {
@ -335,7 +350,12 @@ mod tests {
}
let expected = hex::encode(current);
assert!(v.verify_witness_chain(&steps, &expected).is_ok());
assert!(v.verify_witness_chain(&steps, "0000000000000000000000000000000000000000000000000000000000000000").is_err());
assert!(v
.verify_witness_chain(
&steps,
"0000000000000000000000000000000000000000000000000000000000000000"
)
.is_err());
}
#[test]
@ -350,7 +370,7 @@ mod tests {
#[test]
fn test_ed25519_signature() {
use ed25519_dalek::{SigningKey, Signer};
use ed25519_dalek::{Signer, SigningKey};
let v = Verifier::new();
let mut rng = rand::thread_rng();
let signing_key = SigningKey::generate(&mut rng);
@ -358,9 +378,13 @@ mod tests {
let signature = signing_key.sign(message);
let pub_key = signing_key.verifying_key().to_bytes();
let sig_bytes: [u8; 64] = signature.to_bytes();
assert!(v.verify_ed25519_signature(&pub_key, message, &sig_bytes).is_ok());
assert!(v
.verify_ed25519_signature(&pub_key, message, &sig_bytes)
.is_ok());
// Tampered message should fail
assert!(v.verify_ed25519_signature(&pub_key, b"tampered message", &sig_bytes).is_err());
assert!(v
.verify_ed25519_signature(&pub_key, b"tampered message", &sig_bytes)
.is_err());
}
#[test]

34
crates/mcp-brain-server/src/voice.rs
View file

@ -179,7 +179,8 @@ impl WorkingMemory {
self.evict_lowest();
}
self.items.push(WorkingMemoryItem::new(content, embedding, source));
self.items
.push(WorkingMemoryItem::new(content, embedding, source));
}
/// Retrieve items similar to query embedding
@ -221,16 +222,11 @@ impl WorkingMemory {
/// Evict item with lowest activation
fn evict_lowest(&mut self) {
if let Some((min_idx, _)) = self
.items
.iter()
.enumerate()
.min_by(|(_, a), (_, b)| {
a.activation
.partial_cmp(&b.activation)
.unwrap_or(std::cmp::Ordering::Equal)
})
{
if let Some((min_idx, _)) = self.items.iter().enumerate().min_by(|(_, a), (_, b)| {
a.activation
.partial_cmp(&b.activation)
.unwrap_or(std::cmp::Ordering::Equal)
}) {
self.items.remove(min_idx);
}
}
@ -369,9 +365,7 @@ impl InternalVoice {
self.emit(
ThoughtType::Goal,
format!("I should {}", description),
ThoughtSource::GoalDirected {
goal: description,
},
ThoughtSource::GoalDirected { goal: description },
);
goal_id
}
@ -668,8 +662,16 @@ mod tests {
#[test]
fn test_working_memory_retrieval() {
let mut wm = WorkingMemory::new(5);
wm.add("hello world".to_string(), vec![1.0, 0.0, 0.0, 0.0], ContentSource::External);
wm.add("goodbye world".to_string(), vec![0.0, 1.0, 0.0, 0.0], ContentSource::External);
wm.add(
"hello world".to_string(),
vec![1.0, 0.0, 0.0, 0.0],
ContentSource::External,
);
wm.add(
"goodbye world".to_string(),
vec![0.0, 1.0, 0.0, 0.0],
ContentSource::External,
);
let results = wm.retrieve(&[0.9, 0.1, 0.0, 0.0], 1);
assert!(!results.is_empty());

122
crates/mcp-brain-server/src/web_ingest.rs
View file

@ -93,29 +93,31 @@ pub fn ingest_batch(
// Phase 3: Chunk + Embed
let chunks = chunk_text(&page.text);
let embeddings: Vec<Vec<f32>> = if !page.embedding.is_empty()
&& page.embedding.len() == EMBED_DIM
{
// Use pre-computed embedding for first chunk
let mut embs = vec![page.embedding.clone()];
for chunk in chunks.iter().skip(1) {
let text = EmbeddingEngine::prepare_text(&page.title, chunk, &page.tags);
embs.push(embedding_engine.embed_for_storage(&text));
}
embs
} else {
chunks
.iter()
.map(|chunk| {
let embeddings: Vec<Vec<f32>> =
if !page.embedding.is_empty() && page.embedding.len() == EMBED_DIM {
// Use pre-computed embedding for first chunk
let mut embs = vec![page.embedding.clone()];
for chunk in chunks.iter().skip(1) {
let text = EmbeddingEngine::prepare_text(&page.title, chunk, &page.tags);
embedding_engine.embed_for_storage(&text)
})
.collect()
};
embs.push(embedding_engine.embed_for_storage(&text));
}
embs
} else {
chunks
.iter()
.map(|chunk| {
let text = EmbeddingEngine::prepare_text(&page.title, chunk, &page.tags);
embedding_engine.embed_for_storage(&text)
})
.collect()
};
// Phase 4: Novelty scoring — compare against existing memories
// and already-accepted memories in this batch
let primary_embedding = embeddings.first().cloned().unwrap_or_else(|| vec![0.0; EMBED_DIM]);
let primary_embedding = embeddings
.first()
.cloned()
.unwrap_or_else(|| vec![0.0; EMBED_DIM]);
let batch_embeddings: Vec<(Uuid, Vec<f32>)> = accepted
.iter()
.map(|m| (m.base.id, m.base.embedding.clone()))
@ -324,11 +326,11 @@ pub fn attractor_recrawl_priority(
attractor_results: &HashMap<String, temporal_attractor_studio::LyapunovResult>,
) -> f32 {
match attractor_results.get(domain) {
Some(r) if r.lambda < -0.5 => 0.1, // Very stable — low recrawl priority
Some(r) if r.lambda < 0.0 => 0.3, // Stable — moderate priority
Some(r) if r.lambda > 0.5 => 0.9, // Chaotic — high recrawl priority
Some(_) => 0.5, // Marginally chaotic — default
None => 0.5, // Unknown domain — default priority
Some(r) if r.lambda < -0.5 => 0.1, // Very stable — low recrawl priority
Some(r) if r.lambda < 0.0 => 0.3, // Stable — moderate priority
Some(r) if r.lambda > 0.5 => 0.9, // Chaotic — high recrawl priority
Some(_) => 0.5, // Marginally chaotic — default
None => 0.5, // Unknown domain — default priority
}
}
@ -359,9 +361,7 @@ pub fn solver_drift_prediction(
/// Stub for non-x86 platforms.
#[cfg(not(feature = "x86-simd"))]
pub fn solver_drift_prediction_stub(
_recent_embeddings: &[Vec<f32>],
) -> Option<f32> {
pub fn solver_drift_prediction_stub(_recent_embeddings: &[Vec<f32>]) -> Option<f32> {
// Temporal solver not available on this platform
None
}
@ -405,7 +405,11 @@ mod tests {
#[test]
fn validate_rejects_long_text() {
let page = make_page("https://example.com", &"a".repeat(MAX_TEXT_LENGTH + 1), "Title");
let page = make_page(
"https://example.com",
&"a".repeat(MAX_TEXT_LENGTH + 1),
"Title",
);
assert_eq!(validate_page(&page).unwrap_err(), "text too long");
}
@ -506,7 +510,10 @@ mod tests {
#[test]
fn extract_domain_with_port() {
assert_eq!(extract_domain("https://example.com:8080/path"), "example.com:8080");
assert_eq!(
extract_domain("https://example.com:8080/path"),
"example.com:8080"
);
}
#[test]
@ -583,28 +590,34 @@ mod tests {
#[test]
fn attractor_recrawl_priority_stable() {
let mut results = HashMap::new();
results.insert("stable.com".into(), temporal_attractor_studio::LyapunovResult {
lambda: -1.0,
lyapunov_time: 1.0,
doubling_time: 0.693,
points_used: 20,
dimension: 128,
pairs_found: 10,
});
results.insert(
"stable.com".into(),
temporal_attractor_studio::LyapunovResult {
lambda: -1.0,
lyapunov_time: 1.0,
doubling_time: 0.693,
points_used: 20,
dimension: 128,
pairs_found: 10,
},
);
assert_eq!(attractor_recrawl_priority("stable.com", &results), 0.1);
}
#[test]
fn attractor_recrawl_priority_chaotic() {
let mut results = HashMap::new();
results.insert("chaotic.com".into(), temporal_attractor_studio::LyapunovResult {
lambda: 1.0,
lyapunov_time: 1.0,
doubling_time: 0.693,
points_used: 20,
dimension: 128,
pairs_found: 10,
});
results.insert(
"chaotic.com".into(),
temporal_attractor_studio::LyapunovResult {
lambda: 1.0,
lyapunov_time: 1.0,
doubling_time: 0.693,
points_used: 20,
dimension: 128,
pairs_found: 10,
},
);
assert_eq!(attractor_recrawl_priority("chaotic.com", &results), 0.9);
}
@ -618,14 +631,17 @@ mod tests {
fn attractor_recrawl_priority_marginal() {
let mut results = HashMap::new();
// lambda=0.3 is > 0 but ≤ 0.5 — hits the Some(_) arm
results.insert("marginal.com".into(), temporal_attractor_studio::LyapunovResult {
lambda: 0.3,
lyapunov_time: 3.33,
doubling_time: 2.31,
points_used: 20,
dimension: 128,
pairs_found: 10,
});
results.insert(
"marginal.com".into(),
temporal_attractor_studio::LyapunovResult {
lambda: 0.3,
lyapunov_time: 3.33,
doubling_time: 2.31,
points_used: 20,
dimension: 128,
pairs_found: 10,
},
);
assert_eq!(attractor_recrawl_priority("marginal.com", &results), 0.5);
}
}

117
crates/mcp-brain-server/src/web_memory.rs
View file

@ -8,8 +8,8 @@ use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use crate::types::{BetaParams, BrainCategory, BrainMemory};
use crate::quantization::QuantizedEmbedding;
use crate::types::{BetaParams, BrainCategory, BrainMemory};
// ── Core Web Memory Types ───────────────────────────────────────────────
@ -167,11 +167,7 @@ pub enum ContentDelta {
impl ContentDelta {
/// Classify content change based on token counts and embedding similarity.
pub fn classify(
changed_tokens: usize,
total_tokens: usize,
embedding_cosine: f32,
) -> Self {
pub fn classify(changed_tokens: usize, total_tokens: usize, embedding_cosine: f32) -> Self {
if changed_tokens == 0 {
return Self::Unchanged;
}
@ -464,10 +460,22 @@ mod tests {
#[test]
fn compression_tier_from_novelty() {
assert_eq!(CompressionTier::from_novelty(0.0), CompressionTier::CentroidMerged);
assert_eq!(CompressionTier::from_novelty(0.04), CompressionTier::CentroidMerged);
assert_eq!(CompressionTier::from_novelty(0.05), CompressionTier::DeltaCompressed);
assert_eq!(CompressionTier::from_novelty(0.19), CompressionTier::DeltaCompressed);
assert_eq!(
CompressionTier::from_novelty(0.0),
CompressionTier::CentroidMerged
);
assert_eq!(
CompressionTier::from_novelty(0.04),
CompressionTier::CentroidMerged
);
assert_eq!(
CompressionTier::from_novelty(0.05),
CompressionTier::DeltaCompressed
);
assert_eq!(
CompressionTier::from_novelty(0.19),
CompressionTier::DeltaCompressed
);
assert_eq!(CompressionTier::from_novelty(0.20), CompressionTier::Full);
assert_eq!(CompressionTier::from_novelty(1.0), CompressionTier::Full);
}
@ -482,39 +490,89 @@ mod tests {
#[test]
fn content_delta_classify() {
assert!(matches!(ContentDelta::classify(0, 100, 1.0), ContentDelta::Unchanged));
assert!(matches!(ContentDelta::classify(3, 100, 0.99), ContentDelta::Minor { .. }));
assert!(matches!(ContentDelta::classify(10, 100, 0.85), ContentDelta::Major { .. }));
assert!(matches!(ContentDelta::classify(50, 100, 0.5), ContentDelta::Rewrite));
assert!(matches!(
ContentDelta::classify(0, 100, 1.0),
ContentDelta::Unchanged
));
assert!(matches!(
ContentDelta::classify(3, 100, 0.99),
ContentDelta::Minor { .. }
));
assert!(matches!(
ContentDelta::classify(10, 100, 0.85),
ContentDelta::Major { .. }
));
assert!(matches!(
ContentDelta::classify(50, 100, 0.5),
ContentDelta::Rewrite
));
}
#[test]
fn content_delta_is_significant() {
assert!(!ContentDelta::Unchanged.is_significant());
assert!(!ContentDelta::Minor { changed_tokens: 1, total_tokens: 100 }.is_significant());
assert!(ContentDelta::Major { summary: "test".into(), changed_tokens: 10 }.is_significant());
assert!(!ContentDelta::Minor {
changed_tokens: 1,
total_tokens: 100
}
.is_significant());
assert!(ContentDelta::Major {
summary: "test".into(),
changed_tokens: 10
}
.is_significant());
assert!(ContentDelta::Rewrite.is_significant());
}
#[test]
fn content_delta_edge_cases() {
// Zero total tokens → treated as 100% change (Major)
assert!(matches!(ContentDelta::classify(5, 0, 0.9), ContentDelta::Major { .. }));
assert!(matches!(
ContentDelta::classify(5, 0, 0.9),
ContentDelta::Major { .. }
));
// Exactly at 5% boundary (5/100) → Major (≥ 5%)
assert!(matches!(ContentDelta::classify(5, 100, 0.9), ContentDelta::Major { .. }));
assert!(matches!(
ContentDelta::classify(5, 100, 0.9),
ContentDelta::Major { .. }
));
// Just under 5% boundary (4/100) → Minor
assert!(matches!(ContentDelta::classify(4, 100, 0.9), ContentDelta::Minor { .. }));
assert!(matches!(
ContentDelta::classify(4, 100, 0.9),
ContentDelta::Minor { .. }
));
}
#[test]
fn link_edge_weight_clamped() {
let edge = LinkEdge::new(Uuid::new_v4(), Uuid::new_v4(), None, vec![0.0; 128], LinkType::Citation, 1.5);
let edge = LinkEdge::new(
Uuid::new_v4(),
Uuid::new_v4(),
None,
vec![0.0; 128],
LinkType::Citation,
1.5,
);
assert_eq!(edge.weight, 1.0);
let edge2 = LinkEdge::new(Uuid::new_v4(), Uuid::new_v4(), None, vec![0.0; 128], LinkType::Citation, -0.5);
let edge2 = LinkEdge::new(
Uuid::new_v4(),
Uuid::new_v4(),
None,
vec![0.0; 128],
LinkType::Citation,
-0.5,
);
assert_eq!(edge2.weight, 0.0);
let edge3 = LinkEdge::new(Uuid::new_v4(), Uuid::new_v4(), None, vec![0.0; 128], LinkType::Evidence, 0.75);
let edge3 = LinkEdge::new(
Uuid::new_v4(),
Uuid::new_v4(),
None,
vec![0.0; 128],
LinkType::Evidence,
0.75,
);
assert!((edge3.weight - 0.75).abs() < f64::EPSILON);
}
@ -527,7 +585,10 @@ mod tests {
prev,
curr,
0.85,
ContentDelta::Minor { changed_tokens: 3, total_tokens: 100 },
ContentDelta::Minor {
changed_tokens: 3,
total_tokens: 100,
},
Duration::hours(24),
false,
);
@ -556,8 +617,14 @@ mod tests {
fn content_delta_serde_roundtrip() {
let deltas = [
ContentDelta::Unchanged,
ContentDelta::Minor { changed_tokens: 5, total_tokens: 200 },
ContentDelta::Major { summary: "big change".into(), changed_tokens: 50 },
ContentDelta::Minor {
changed_tokens: 5,
total_tokens: 200,
},
ContentDelta::Major {
summary: "big change".into(),
changed_tokens: 50,
},
ContentDelta::Rewrite,
];
for delta in &deltas {

46
crates/mcp-brain-server/src/web_store.rs
View file

@ -60,7 +60,9 @@ impl WebMemoryStore {
// Write-through to Firestore
if let Ok(body) = serde_json::to_value(&mem) {
self.firestore.firestore_put_public("web_memories", &id.to_string(), &body).await;
self.firestore
.firestore_put_public("web_memories", &id.to_string(), &body)
.await;
}
// Store base BrainMemory in core store for search compatibility
@ -87,7 +89,10 @@ impl WebMemoryStore {
/// Get all content hashes for deduplication.
pub fn content_hashes(&self) -> HashSet<String> {
self.content_hashes.iter().map(|e| e.key().clone()).collect()
self.content_hashes
.iter()
.map(|e| e.key().clone())
.collect()
}
/// Get all embeddings for novelty scoring.
@ -116,11 +121,16 @@ impl WebMemoryStore {
let url = delta.page_url.clone();
if let Ok(body) = serde_json::to_value(&delta) {
self.firestore.firestore_put_public("web_deltas", &delta_id.to_string(), &body).await;
self.firestore
.firestore_put_public("web_deltas", &delta_id.to_string(), &body)
.await;
}
// Index by URL for evolution queries
self.url_deltas.entry(url).or_insert_with(Vec::new).push(delta_id);
self.url_deltas
.entry(url)
.or_insert_with(Vec::new)
.push(delta_id);
self.deltas.insert(delta_id, delta);
}
@ -147,7 +157,10 @@ impl WebMemoryStore {
/// Store a link edge.
pub fn store_link_edge(&self, edge: LinkEdge) {
let source = edge.source_memory_id;
self.link_edges.entry(source).or_insert_with(Vec::new).push(edge);
self.link_edges
.entry(source)
.or_insert_with(Vec::new)
.push(edge);
}
/// Get outbound link edges from a memory.
@ -214,7 +227,8 @@ impl WebMemoryStore {
}
let total = self.memories.len();
let compressed = tier_dist.delta_compressed + tier_dist.centroid_merged + tier_dist.archived;
let compressed =
tier_dist.delta_compressed + tier_dist.centroid_merged + tier_dist.archived;
let compression_ratio = if total > 0 {
compressed as f64 / total as f64
} else {
@ -226,8 +240,16 @@ impl WebMemoryStore {
.map(|(domain, (count, qual_sum, nov_sum))| DomainStats {
domain,
page_count: count,
avg_quality: if count > 0 { qual_sum / count as f64 } else { 0.0 },
avg_novelty: if count > 0 { nov_sum / count as f32 } else { 0.0 },
avg_quality: if count > 0 {
qual_sum / count as f64
} else {
0.0
},
avg_novelty: if count > 0 {
nov_sum / count as f32
} else {
0.0
},
})
.collect();
top_domains.sort_by(|a, b| b.page_count.cmp(&a.page_count));
@ -305,8 +327,12 @@ mod tests {
let mem2 = make_test_web_memory("https://b.com", 0.7);
// Manually insert to test hash tracking (bypassing async store)
store.content_hashes.insert(mem1.content_hash.clone(), mem1.base.id);
store.content_hashes.insert(mem2.content_hash.clone(), mem2.base.id);
store
.content_hashes
.insert(mem1.content_hash.clone(), mem1.base.id);
store
.content_hashes
.insert(mem2.content_hash.clone(), mem2.base.id);
let hashes = store.content_hashes();
assert!(hashes.contains(&mem1.content_hash));

186
crates/mcp-brain/src/client.rs
View file

@ -24,8 +24,7 @@ impl BrainClient {
pub fn new() -> Self {
let base_url = std::env::var("BRAIN_URL")
.unwrap_or_else(|_| "https://ruvbrain-875130704813.us-central1.run.app".to_string());
let api_key = std::env::var("BRAIN_API_KEY")
.unwrap_or_else(|_| "anonymous".to_string());
let api_key = std::env::var("BRAIN_API_KEY").unwrap_or_else(|_| "anonymous".to_string());
Self {
base_url,
api_key,
@ -34,8 +33,7 @@ impl BrainClient {
}
pub fn with_url(url: String) -> Self {
let api_key = std::env::var("BRAIN_API_KEY")
.unwrap_or_else(|_| "anonymous".to_string());
let api_key = std::env::var("BRAIN_API_KEY").unwrap_or_else(|_| "anonymous".to_string());
Self {
base_url: url,
api_key,
@ -75,10 +73,18 @@ impl BrainClient {
min_quality: Option<f64>,
) -> Result<serde_json::Value, ClientError> {
let mut params = vec![("q", query.to_string())];
if let Some(c) = category { params.push(("category", c.to_string())); }
if let Some(t) = tags { params.push(("tags", t.to_string())); }
if let Some(l) = limit { params.push(("limit", l.to_string())); }
if let Some(q) = min_quality { params.push(("min_quality", q.to_string())); }
if let Some(c) = category {
params.push(("category", c.to_string()));
}
if let Some(t) = tags {
params.push(("tags", t.to_string()));
}
if let Some(l) = limit {
params.push(("limit", l.to_string()));
}
if let Some(q) = min_quality {
params.push(("min_quality", q.to_string()));
}
self.get_with_params("/v1/memories/search", &params).await
}
@ -95,7 +101,11 @@ impl BrainClient {
}
/// Transfer knowledge between domains
pub async fn transfer(&self, source: &str, target: &str) -> Result<serde_json::Value, ClientError> {
pub async fn transfer(
&self,
source: &str,
target: &str,
) -> Result<serde_json::Value, ClientError> {
let body = serde_json::json!({
"source_domain": source,
"target_domain": target,
@ -104,33 +114,58 @@ impl BrainClient {
}
/// Get drift report
pub async fn drift(&self, domain: Option<&str>, since: Option<&str>) -> Result<serde_json::Value, ClientError> {
pub async fn drift(
&self,
domain: Option<&str>,
since: Option<&str>,
) -> Result<serde_json::Value, ClientError> {
let mut params = Vec::new();
if let Some(d) = domain { params.push(("domain", d.to_string())); }
if let Some(s) = since { params.push(("since", s.to_string())); }
if let Some(d) = domain {
params.push(("domain", d.to_string()));
}
if let Some(s) = since {
params.push(("since", s.to_string()));
}
self.get_with_params("/v1/drift", &params).await
}
/// Get partition topology
pub async fn partition(&self, domain: Option<&str>, min_size: Option<usize>) -> Result<serde_json::Value, ClientError> {
pub async fn partition(
&self,
domain: Option<&str>,
min_size: Option<usize>,
) -> Result<serde_json::Value, ClientError> {
let mut params = Vec::new();
if let Some(d) = domain { params.push(("domain", d.to_string())); }
if let Some(s) = min_size { params.push(("min_cluster_size", s.to_string())); }
if let Some(d) = domain {
params.push(("domain", d.to_string()));
}
if let Some(s) = min_size {
params.push(("min_cluster_size", s.to_string()));
}
self.get_with_params("/v1/partition", &params).await
}
/// List memories
pub async fn list(&self, category: Option<&str>, limit: Option<usize>) -> Result<serde_json::Value, ClientError> {
pub async fn list(
&self,
category: Option<&str>,
limit: Option<usize>,
) -> Result<serde_json::Value, ClientError> {
let mut params = Vec::new();
if let Some(c) = category { params.push(("category", c.to_string())); }
if let Some(l) = limit { params.push(("limit", l.to_string())); }
if let Some(c) = category {
params.push(("category", c.to_string()));
}
if let Some(l) = limit {
params.push(("limit", l.to_string()));
}
self.get_with_params("/v1/memories/list", &params).await
}
/// Delete a memory
pub async fn delete(&self, id: &str) -> Result<(), ClientError> {
let url = format!("{}/v1/memories/{id}", self.base_url);
let resp = self.http
let resp = self
.http
.delete(&url)
.bearer_auth(&self.api_key)
.send()
@ -142,7 +177,10 @@ impl BrainClient {
} else {
let status = resp.status().as_u16();
let msg = resp.text().await.unwrap_or_default();
Err(ClientError::Server { status, message: msg })
Err(ClientError::Server {
status,
message: msg,
})
}
}
@ -160,9 +198,9 @@ impl BrainClient {
if val.get("weights").map_or(true, |w| w.is_null()) {
return Ok(None);
}
let weights: LoraWeights = serde_json::from_value(
val.get("weights").cloned().unwrap_or_default()
).map_err(|e| ClientError::Serialization(e.to_string()))?;
let weights: LoraWeights =
serde_json::from_value(val.get("weights").cloned().unwrap_or_default())
.map_err(|e| ClientError::Serialization(e.to_string()))?;
Ok(Some(weights))
}
Err(ClientError::Server { status: 404, .. }) => Ok(None),
@ -171,16 +209,22 @@ impl BrainClient {
}
/// Submit local LoRA weights for federated aggregation
pub async fn lora_submit(&self, weights: &LoraWeights) -> Result<serde_json::Value, ClientError> {
let body = serde_json::to_value(weights)
.map_err(|e| ClientError::Serialization(e.to_string()))?;
pub async fn lora_submit(
&self,
weights: &LoraWeights,
) -> Result<serde_json::Value, ClientError> {
let body =
serde_json::to_value(weights).map_err(|e| ClientError::Serialization(e.to_string()))?;
self.post("/v1/lora/submit", &body).await
}
// ---- Brainpedia (ADR-062) ----
/// Create a Brainpedia page
pub async fn create_page(&self, body: &serde_json::Value) -> Result<serde_json::Value, ClientError> {
pub async fn create_page(
&self,
body: &serde_json::Value,
) -> Result<serde_json::Value, ClientError> {
self.post("/v1/pages", body).await
}
@ -190,8 +234,13 @@ impl BrainClient {
}
/// Submit a delta to a page
pub async fn submit_delta(&self, page_id: &str, body: &serde_json::Value) -> Result<serde_json::Value, ClientError> {
self.post(&format!("/v1/pages/{page_id}/deltas"), body).await
pub async fn submit_delta(
&self,
page_id: &str,
body: &serde_json::Value,
) -> Result<serde_json::Value, ClientError> {
self.post(&format!("/v1/pages/{page_id}/deltas"), body)
.await
}
/// List deltas for a page
@ -200,13 +249,22 @@ impl BrainClient {
}
/// Add evidence to a page
pub async fn add_evidence(&self, page_id: &str, body: &serde_json::Value) -> Result<serde_json::Value, ClientError> {
self.post(&format!("/v1/pages/{page_id}/evidence"), body).await
pub async fn add_evidence(
&self,
page_id: &str,
body: &serde_json::Value,
) -> Result<serde_json::Value, ClientError> {
self.post(&format!("/v1/pages/{page_id}/evidence"), body)
.await
}
/// Promote a page from Draft to Canonical
pub async fn promote_page(&self, page_id: &str) -> Result<serde_json::Value, ClientError> {
self.post(&format!("/v1/pages/{page_id}/promote"), &serde_json::json!({})).await
self.post(
&format!("/v1/pages/{page_id}/promote"),
&serde_json::json!({}),
)
.await
}
// ---- WASM Executable Nodes (ADR-063) ----
@ -217,7 +275,10 @@ impl BrainClient {
}
/// Publish a WASM node
pub async fn publish_node(&self, body: &serde_json::Value) -> Result<serde_json::Value, ClientError> {
pub async fn publish_node(
&self,
body: &serde_json::Value,
) -> Result<serde_json::Value, ClientError> {
self.post("/v1/nodes", body).await
}
@ -229,7 +290,8 @@ impl BrainClient {
/// Download WASM binary
pub async fn get_node_wasm(&self, id: &str) -> Result<Vec<u8>, ClientError> {
let url = format!("{}/v1/nodes/{id}.wasm", self.base_url);
let resp = self.http
let resp = self
.http
.get(&url)
.bearer_auth(&self.api_key)
.send()
@ -237,20 +299,25 @@ impl BrainClient {
.map_err(|e| ClientError::Http(e.to_string()))?;
if resp.status().is_success() {
resp.bytes().await
resp.bytes()
.await
.map(|b| b.to_vec())
.map_err(|e| ClientError::Http(e.to_string()))
} else {
let status = resp.status().as_u16();
let msg = resp.text().await.unwrap_or_default();
Err(ClientError::Server { status, message: msg })
Err(ClientError::Server {
status,
message: msg,
})
}
}
/// Revoke a WASM node
pub async fn revoke_node(&self, id: &str) -> Result<(), ClientError> {
let url = format!("{}/v1/nodes/{id}/revoke", self.base_url);
let resp = self.http
let resp = self
.http
.post(&url)
.bearer_auth(&self.api_key)
.json(&serde_json::json!({}))
@ -263,7 +330,10 @@ impl BrainClient {
} else {
let status = resp.status().as_u16();
let msg = resp.text().await.unwrap_or_default();
Err(ClientError::Server { status, message: msg })
Err(ClientError::Server {
status,
message: msg,
})
}
}
@ -271,7 +341,8 @@ impl BrainClient {
async fn get_path(&self, path: &str) -> Result<serde_json::Value, ClientError> {
let url = format!("{}{path}", self.base_url);
let resp = self.http
let resp = self
.http
.get(&url)
.bearer_auth(&self.api_key)
.send()
@ -281,9 +352,14 @@ impl BrainClient {
self.handle_response(resp).await
}
async fn get_with_params(&self, path: &str, params: &[(&str, String)]) -> Result<serde_json::Value, ClientError> {
async fn get_with_params(
&self,
path: &str,
params: &[(&str, String)],
) -> Result<serde_json::Value, ClientError> {
let url = format!("{}{path}", self.base_url);
let resp = self.http
let resp = self
.http
.get(&url)
.bearer_auth(&self.api_key)
.query(params)
@ -294,9 +370,14 @@ impl BrainClient {
self.handle_response(resp).await
}
async fn post(&self, path: &str, body: &serde_json::Value) -> Result<serde_json::Value, ClientError> {
async fn post(
&self,
path: &str,
body: &serde_json::Value,
) -> Result<serde_json::Value, ClientError> {
let url = format!("{}{path}", self.base_url);
let resp = self.http
let resp = self
.http
.post(&url)
.bearer_auth(&self.api_key)
.json(body)
@ -307,13 +388,21 @@ impl BrainClient {
self.handle_response(resp).await
}
async fn handle_response(&self, resp: reqwest::Response) -> Result<serde_json::Value, ClientError> {
async fn handle_response(
&self,
resp: reqwest::Response,
) -> Result<serde_json::Value, ClientError> {
let status = resp.status().as_u16();
if status >= 400 {
let msg = resp.text().await.unwrap_or_default();
return Err(ClientError::Server { status, message: msg });
return Err(ClientError::Server {
status,
message: msg,
});
}
resp.json().await.map_err(|e| ClientError::Serialization(e.to_string()))
resp.json()
.await
.map_err(|e| ClientError::Serialization(e.to_string()))
}
}
@ -325,7 +414,10 @@ impl Default for BrainClient {
/// SHAKE-256 hash
fn sha3_hash(data: &[u8]) -> [u8; 32] {
use sha3::{Shake256, digest::{Update, ExtendableOutput, XofReader}};
use sha3::{
digest::{ExtendableOutput, Update, XofReader},
Shake256,
};
let mut hasher = Shake256::default();
hasher.update(data);
let mut reader = hasher.finalize_xof();

22
crates/mcp-brain/src/embed.rs
View file

@ -8,9 +8,12 @@
//! Rank-2 LoRA adapter applied to the frozen hash features. Weights are
//! learned locally via SONA and periodically federated to/from the server.
use sha3::{Shake256, digest::{Update, ExtendableOutput, XofReader}};
use sona::{SonaEngine, LearnedPattern};
use sha3::{
digest::{ExtendableOutput, Update, XofReader},
Shake256,
};
use sona::engine::SonaEngineBuilder;
use sona::{LearnedPattern, SonaEngine};
use std::sync::{Arc, Mutex};
/// Embedding dimension (128 f32s = 512 bytes)
@ -70,7 +73,9 @@ impl LoraWeights {
let down_norm: f32 = self.down_proj.iter().map(|x| x * x).sum::<f32>().sqrt();
let up_norm: f32 = self.up_proj.iter().map(|x| x * x).sum::<f32>().sqrt();
if down_norm > 100.0 || up_norm > 100.0 {
return Err(format!("Weight norm too large: down={down_norm:.1}, up={up_norm:.1}"));
return Err(format!(
"Weight norm too large: down={down_norm:.1}, up={up_norm:.1}"
));
}
// Minimum evidence
if self.evidence_count < 5 {
@ -91,7 +96,10 @@ impl LoraWeights {
/// Compute L2 distance to another set of weights
pub fn l2_distance(&self, other: &LoraWeights) -> f32 {
let d: f32 = self.down_proj.iter().zip(other.down_proj.iter())
let d: f32 = self
.down_proj
.iter()
.zip(other.down_proj.iter())
.chain(self.up_proj.iter().zip(other.up_proj.iter()))
.map(|(a, b)| (a - b).powi(2))
.sum();
@ -463,7 +471,11 @@ mod tests {
for i in 0..100 {
let key = format!("test-{i}");
let (_, sign) = signed_hash(key.as_bytes(), b"test", 42);
if sign > 0.0 { pos += 1; } else { neg += 1; }
if sign > 0.0 {
pos += 1;
} else {
neg += 1;
}
}
// Both signs should appear (probabilistic, but 100 trials is enough)
assert!(pos > 10 && neg > 10, "pos={pos}, neg={neg}");

9
crates/mcp-brain/src/pipeline.rs
View file

@ -1,7 +1,10 @@
//! Local processing pipeline: PII -> embed -> sign
use regex_lite::Regex;
use sha3::{Shake256, digest::{Update, ExtendableOutput, XofReader}};
use sha3::{
digest::{ExtendableOutput, Update, XofReader},
Shake256,
};
/// Pipeline for processing knowledge before sharing.
/// Pre-compiles 12 PII regex patterns for efficient reuse.
@ -37,7 +40,9 @@ impl BrainPipeline {
// 12. Internal hostnames
(Regex::new(r"\b(?:localhost|127\.0\.0\.1|0\.0\.0\.0|internal\.[a-z.]+)\b").unwrap(), "[REDACTED_HOST]"),
];
Self { pii_patterns: patterns }
Self {
pii_patterns: patterns,
}
}
/// Strip PII from text using all 12 pattern categories

326
crates/mcp-brain/src/tools.rs
View file

@ -353,45 +353,69 @@ impl McpBrainTools {
"brain_node_get" => self.brain_node_get(call.arguments).await,
"brain_node_wasm" => self.brain_node_wasm(call.arguments).await,
"brain_node_revoke" => self.brain_node_revoke(call.arguments).await,
_ => Err(BrainError::InvalidRequest(format!("Unknown tool: {}", call.name))),
_ => Err(BrainError::InvalidRequest(format!(
"Unknown tool: {}",
call.name
))),
}
}
async fn brain_share(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let category = args.get("category").and_then(|v| v.as_str()).unwrap_or("pattern");
let title = args.get("title").and_then(|v| v.as_str())
let category = args
.get("category")
.and_then(|v| v.as_str())
.unwrap_or("pattern");
let title = args
.get("title")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("title required".into()))?;
let content = args.get("content").and_then(|v| v.as_str())
let content = args
.get("content")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("content required".into()))?;
let tags: Vec<String> = args.get("tags")
let tags: Vec<String> = args
.get("tags")
.and_then(|v| v.as_array())
.map(|arr| arr.iter().filter_map(|v| v.as_str().map(String::from)).collect())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
.unwrap_or_default();
let code_snippet = args.get("code_snippet").and_then(|v| v.as_str()).map(String::from);
let code_snippet = args
.get("code_snippet")
.and_then(|v| v.as_str())
.map(String::from);
// PII strip all user-provided text
let clean_title = self.pipeline.strip_pii(title);
let clean_content = self.pipeline.strip_pii(content);
let clean_tags: Vec<String> = tags.iter()
.map(|t| self.pipeline.strip_pii(t))
.collect();
let clean_tags: Vec<String> = tags.iter().map(|t| self.pipeline.strip_pii(t)).collect();
let clean_snippet = code_snippet.as_deref().map(|s| self.pipeline.strip_pii(s));
// Safety check: reject if PII still detected after stripping
if self.pipeline.contains_pii(&clean_title) {
return Err(BrainError::Pipeline("PII detected in title after stripping".into()));
return Err(BrainError::Pipeline(
"PII detected in title after stripping".into(),
));
}
if self.pipeline.contains_pii(&clean_content) {
return Err(BrainError::Pipeline("PII detected in content after stripping".into()));
return Err(BrainError::Pipeline(
"PII detected in content after stripping".into(),
));
}
for tag in &clean_tags {
if self.pipeline.contains_pii(tag) {
return Err(BrainError::Pipeline("PII detected in tags after stripping".into()));
return Err(BrainError::Pipeline(
"PII detected in tags after stripping".into(),
));
}
}
if let Some(ref s) = clean_snippet {
if self.pipeline.contains_pii(s) {
return Err(BrainError::Pipeline("PII detected in code_snippet after stripping".into()));
return Err(BrainError::Pipeline(
"PII detected in code_snippet after stripping".into(),
));
}
}
@ -409,13 +433,17 @@ impl McpBrainTools {
chain.append("share");
let _witness_hash = chain.finalize();
let result = self.client.share(
category,
&clean_title,
&clean_content,
&clean_tags,
clean_snippet.as_deref(),
).await.map_err(|e| BrainError::Client(e.to_string()))?;
let result = self
.client
.share(
category,
&clean_title,
&clean_content,
&clean_tags,
clean_snippet.as_deref(),
)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
content: serde_json::to_value(result).unwrap_or_default(),
@ -423,11 +451,16 @@ impl McpBrainTools {
}
async fn brain_search(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let query = args.get("query").and_then(|v| v.as_str())
let query = args
.get("query")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("query required".into()))?;
let category = args.get("category").and_then(|v| v.as_str());
let tags = args.get("tags").and_then(|v| v.as_str());
let limit = args.get("limit").and_then(|v| v.as_u64()).map(|v| v as usize);
let limit = args
.get("limit")
.and_then(|v| v.as_u64())
.map(|v| v as usize);
let min_quality = args.get("min_quality").and_then(|v| v.as_f64());
// Generate query embedding via structured hash + MicroLoRA
@ -437,7 +470,10 @@ impl McpBrainTools {
crate::embed::generate_embedding(query)
};
let results = self.client.search(query, category, tags, limit, min_quality).await
let results = self
.client
.search(query, category, tags, limit, min_quality)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
@ -446,10 +482,15 @@ impl McpBrainTools {
}
async fn brain_get(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let id = args.get("id").and_then(|v| v.as_str())
let id = args
.get("id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("id required".into()))?;
let memory = self.client.get(id).await
let memory = self
.client
.get(id)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
@ -458,12 +499,19 @@ impl McpBrainTools {
}
async fn brain_vote(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let id = args.get("id").and_then(|v| v.as_str())
let id = args
.get("id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("id required".into()))?;
let direction = args.get("direction").and_then(|v| v.as_str())
let direction = args
.get("direction")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("direction required".into()))?;
let result = self.client.vote(id, direction).await
let result = self
.client
.vote(id, direction)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
@ -472,12 +520,19 @@ impl McpBrainTools {
}
async fn brain_transfer(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let source = args.get("source_domain").and_then(|v| v.as_str())
let source = args
.get("source_domain")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("source_domain required".into()))?;
let target = args.get("target_domain").and_then(|v| v.as_str())
let target = args
.get("target_domain")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("target_domain required".into()))?;
let result = self.client.transfer(source, target).await
let result = self
.client
.transfer(source, target)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
@ -489,7 +544,10 @@ impl McpBrainTools {
let domain = args.get("domain").and_then(|v| v.as_str());
let since = args.get("since").and_then(|v| v.as_str());
let report = self.client.drift(domain, since).await
let report = self
.client
.drift(domain, since)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
@ -499,9 +557,15 @@ impl McpBrainTools {
async fn brain_partition(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let domain = args.get("domain").and_then(|v| v.as_str());
let min_size = args.get("min_cluster_size").and_then(|v| v.as_u64()).map(|v| v as usize);
let min_size = args
.get("min_cluster_size")
.and_then(|v| v.as_u64())
.map(|v| v as usize);
let result = self.client.partition(domain, min_size).await
let result = self
.client
.partition(domain, min_size)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
@ -511,9 +575,15 @@ impl McpBrainTools {
async fn brain_list(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let category = args.get("category").and_then(|v| v.as_str());
let limit = args.get("limit").and_then(|v| v.as_u64()).map(|v| v as usize);
let limit = args
.get("limit")
.and_then(|v| v.as_u64())
.map(|v| v as usize);
let results = self.client.list(category, limit).await
let results = self
.client
.list(category, limit)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
@ -522,10 +592,14 @@ impl McpBrainTools {
}
async fn brain_delete(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let id = args.get("id").and_then(|v| v.as_str())
let id = args
.get("id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("id required".into()))?;
self.client.delete(id).await
self.client
.delete(id)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
@ -534,7 +608,10 @@ impl McpBrainTools {
}
async fn brain_status(&self, _args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let status = self.client.status().await
let status = self
.client
.status()
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {
@ -543,7 +620,10 @@ impl McpBrainTools {
}
async fn brain_sync(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let direction = args.get("direction").and_then(|v| v.as_str()).unwrap_or("both");
let direction = args
.get("direction")
.and_then(|v| v.as_str())
.unwrap_or("both");
let mut pulled = false;
let mut pushed = false;
@ -577,7 +657,9 @@ impl McpBrainTools {
weights.clip();
if weights.validate().is_ok() {
match self.client.lora_submit(&weights).await {
Ok(_) => { pushed = true; }
Ok(_) => {
pushed = true;
}
Err(e) => {
info!("Failed to push local weights: {e}");
}
@ -604,18 +686,36 @@ impl McpBrainTools {
// ── Brainpedia (ADR-062) ─────────────────────────────────────────
async fn brain_page_create(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let category = args.get("category").and_then(|v| v.as_str()).unwrap_or("pattern");
let title = args.get("title").and_then(|v| v.as_str())
async fn brain_page_create(
&self,
args: serde_json::Value,
) -> Result<McpToolResult, BrainError> {
let category = args
.get("category")
.and_then(|v| v.as_str())
.unwrap_or("pattern");
let title = args
.get("title")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("title required".into()))?;
let content = args.get("content").and_then(|v| v.as_str())
let content = args
.get("content")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("content required".into()))?;
let tags: Vec<String> = args.get("tags")
let tags: Vec<String> = args
.get("tags")
.and_then(|v| v.as_array())
.map(|arr| arr.iter().filter_map(|v| v.as_str().map(String::from)).collect())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
.unwrap_or_default();
let code_snippet = args.get("code_snippet").and_then(|v| v.as_str());
let evidence_links = args.get("evidence_links").cloned().unwrap_or(serde_json::json!([]));
let evidence_links = args
.get("evidence_links")
.cloned()
.unwrap_or(serde_json::json!([]));
let clean_title = self.pipeline.strip_pii(title);
let clean_content = self.pipeline.strip_pii(content);
@ -644,30 +744,47 @@ impl McpBrainTools {
"witness_hash": hex::encode(witness_hash),
});
let result = self.client.create_page(&body).await
let result = self
.client
.create_page(&body)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success { content: result })
}
async fn brain_page_get(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let id = args.get("id").and_then(|v| v.as_str())
let id = args
.get("id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("id required".into()))?;
let result = self.client.get_page(id).await
let result = self
.client
.get_page(id)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success { content: result })
}
async fn brain_page_delta(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let page_id = args.get("page_id").and_then(|v| v.as_str())
let page_id = args
.get("page_id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("page_id required".into()))?;
let delta_type = args.get("delta_type").and_then(|v| v.as_str())
let delta_type = args
.get("delta_type")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("delta_type required".into()))?;
let content_diff = args.get("content_diff").cloned()
let content_diff = args
.get("content_diff")
.cloned()
.ok_or_else(|| BrainError::InvalidRequest("content_diff required".into()))?;
let evidence_links = args.get("evidence_links").cloned().unwrap_or(serde_json::json!([]));
let evidence_links = args
.get("evidence_links")
.cloned()
.unwrap_or(serde_json::json!([]));
let mut chain = crate::pipeline::WitnessChain::new();
chain.append("delta_submit");
@ -680,41 +797,70 @@ impl McpBrainTools {
"witness_hash": hex::encode(witness_hash),
});
let result = self.client.submit_delta(page_id, &body).await
let result = self
.client
.submit_delta(page_id, &body)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success { content: result })
}
async fn brain_page_deltas(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let page_id = args.get("page_id").and_then(|v| v.as_str())
async fn brain_page_deltas(
&self,
args: serde_json::Value,
) -> Result<McpToolResult, BrainError> {
let page_id = args
.get("page_id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("page_id required".into()))?;
let result = self.client.list_deltas(page_id).await
let result = self
.client
.list_deltas(page_id)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success { content: result })
}
async fn brain_page_evidence(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let page_id = args.get("page_id").and_then(|v| v.as_str())
async fn brain_page_evidence(
&self,
args: serde_json::Value,
) -> Result<McpToolResult, BrainError> {
let page_id = args
.get("page_id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("page_id required".into()))?;
let evidence = args.get("evidence").cloned()
let evidence = args
.get("evidence")
.cloned()
.ok_or_else(|| BrainError::InvalidRequest("evidence required".into()))?;
let body = serde_json::json!({ "evidence": evidence });
let result = self.client.add_evidence(page_id, &body).await
let result = self
.client
.add_evidence(page_id, &body)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success { content: result })
}
async fn brain_page_promote(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let page_id = args.get("page_id").and_then(|v| v.as_str())
async fn brain_page_promote(
&self,
args: serde_json::Value,
) -> Result<McpToolResult, BrainError> {
let page_id = args
.get("page_id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("page_id required".into()))?;
let result = self.client.promote_page(page_id).await
let result = self
.client
.promote_page(page_id)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success { content: result })
@ -723,34 +869,53 @@ impl McpBrainTools {
// ── WASM Executable Nodes (ADR-063) ───────────────────────────────
async fn brain_node_list(&self, _args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let result = self.client.list_nodes().await
let result = self
.client
.list_nodes()
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success { content: result })
}
async fn brain_node_publish(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let result = self.client.publish_node(&args).await
async fn brain_node_publish(
&self,
args: serde_json::Value,
) -> Result<McpToolResult, BrainError> {
let result = self
.client
.publish_node(&args)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success { content: result })
}
async fn brain_node_get(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let id = args.get("id").and_then(|v| v.as_str())
let id = args
.get("id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("id required".into()))?;
let result = self.client.get_node(id).await
let result = self
.client
.get_node(id)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success { content: result })
}
async fn brain_node_wasm(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let id = args.get("id").and_then(|v| v.as_str())
let id = args
.get("id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("id required".into()))?;
let bytes = self.client.get_node_wasm(id).await
let bytes = self
.client
.get_node_wasm(id)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
let b64 = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &bytes);
@ -764,11 +929,18 @@ impl McpBrainTools {
})
}
async fn brain_node_revoke(&self, args: serde_json::Value) -> Result<McpToolResult, BrainError> {
let id = args.get("id").and_then(|v| v.as_str())
async fn brain_node_revoke(
&self,
args: serde_json::Value,
) -> Result<McpToolResult, BrainError> {
let id = args
.get("id")
.and_then(|v| v.as_str())
.ok_or_else(|| BrainError::InvalidRequest("id required".into()))?;
self.client.revoke_node(id).await
self.client
.revoke_node(id)
.await
.map_err(|e| BrainError::Client(e.to_string()))?;
Ok(McpToolResult::Success {

18
crates/neural-trader-strategies/src/attention_scalper.rs
View file

@ -80,7 +80,10 @@ pub struct AttentionScalper {
impl AttentionScalper {
pub fn new(config: AttentionScalperConfig) -> Self {
Self { config, syms: HashMap::new() }
Self {
config,
syms: HashMap::new(),
}
}
fn update_snapshot(&mut self, event: &MarketEvent) -> Option<f64> {
@ -142,10 +145,7 @@ impl AttentionScalper {
} else {
(&state.no_levels, Side::No)
};
let price_cents = levels
.last()
.map(|(p, _)| *p)
.unwrap_or(0);
let price_cents = levels.last().map(|(p, _)| *p).unwrap_or(0);
if price_cents <= 0 || price_cents >= 100 {
return None;
}
@ -282,7 +282,9 @@ mod tests {
});
s.on_event(&level(1, NtSide::Bid, 24, 500, 0));
s.on_event(&level(1, NtSide::Bid, 23, 300, 1));
let intent = s.on_event(&level(1, NtSide::Ask, 76, 100, 2)).expect("should emit");
let intent = s
.on_event(&level(1, NtSide::Ask, 76, 100, 2))
.expect("should emit");
assert_eq!(intent.symbol_id, 1);
assert!(matches!(intent.side, Side::Yes));
assert_eq!(intent.quantity, 5);
@ -301,7 +303,9 @@ mod tests {
});
s.on_event(&level(2, NtSide::Bid, 24, 100, 0));
s.on_event(&level(2, NtSide::Ask, 76, 500, 1));
let intent = s.on_event(&level(2, NtSide::Ask, 77, 400, 2)).expect("should emit");
let intent = s
.on_event(&level(2, NtSide::Ask, 77, 400, 2))
.expect("should emit");
assert!(matches!(intent.side, Side::No));
}

11
crates/neural-trader-strategies/src/coherence_arb.rs
View file

@ -58,7 +58,10 @@ pub struct CoherenceArb {
impl CoherenceArb {
pub fn new(config: CoherenceArbConfig) -> Self {
Self { config, latest_mid_cents: HashMap::new() }
Self {
config,
latest_mid_cents: HashMap::new(),
}
}
fn update_mid(&mut self, event: &MarketEvent) {
@ -70,11 +73,7 @@ impl CoherenceArb {
fn try_arb_for(&self, mirror_sym: u32) -> Option<Intent> {
// Find the pair this symbol participates in.
let &(reference, mirror) = self
.config
.pairs
.iter()
.find(|(_, m)| *m == mirror_sym)?;
let &(reference, mirror) = self.config.pairs.iter().find(|(_, m)| *m == mirror_sym)?;
let ref_cents = *self.latest_mid_cents.get(&reference)?;
let mirror_cents = *self.latest_mid_cents.get(&mirror)?;
let divergence_cents = ref_cents - mirror_cents;

5
crates/neural-trader-strategies/src/coherence_bridge.rs
View file

@ -48,7 +48,10 @@ impl<G: CoherenceGate> CoherenceChecker<G> {
pub fn check(&self, intent: Intent, ctx: &GateContext) -> CoherenceOutcome {
match self.gate.evaluate(ctx) {
Ok(d) if d.allow_act => CoherenceOutcome::Pass(intent),
Ok(d) => CoherenceOutcome::Block { intent, decision: d },
Ok(d) => CoherenceOutcome::Block {
intent,
decision: d,
},
Err(e) => CoherenceOutcome::Block {
intent,
decision: CoherenceDecision {

11
crates/neural-trader-strategies/src/ev_kelly.rs
View file

@ -37,7 +37,11 @@ struct SymbolState {
impl Default for SymbolState {
fn default() -> Self {
Self { latest_mid_cents: None, prior: None, last_emit_seq: 0 }
Self {
latest_mid_cents: None,
prior: None,
last_emit_seq: 0,
}
}
}
@ -74,7 +78,10 @@ pub struct ExpectedValueKelly {
impl ExpectedValueKelly {
pub fn new(config: ExpectedValueKellyConfig) -> Self {
Self { config, symbols: HashMap::new() }
Self {
config,
symbols: HashMap::new(),
}
}
/// Install or update the probability prior for a symbol. Priors are

4
crates/neural-trader-strategies/src/lib.rs
View file

@ -30,9 +30,7 @@ pub use coherence_bridge::{
};
pub use ev_kelly::{ExpectedValueKelly, ExpectedValueKellyConfig};
pub use intent::{Action, Intent, Side};
pub use risk::{
PortfolioState, Position, RejectReason, RiskConfig, RiskDecision, RiskGate,
};
pub use risk::{PortfolioState, Position, RejectReason, RiskConfig, RiskDecision, RiskGate};
use neural_trader_core::MarketEvent;

103
crates/neural-trader-strategies/src/risk.rs
View file

@ -98,7 +98,10 @@ impl Default for RiskConfig {
#[derive(Debug, Clone, PartialEq)]
pub enum RiskDecision {
Approve(Intent),
Reject { reason: RejectReason, intent: Intent },
Reject {
reason: RejectReason,
intent: Intent,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@ -140,9 +143,8 @@ impl RiskGate {
}
// 2. Daily loss kill — stop *all* opening trades after breach.
let max_loss = (portfolio.starting_cash_cents as f64
* self.config.max_daily_loss_frac)
.round() as i64;
let max_loss =
(portfolio.starting_cash_cents as f64 * self.config.max_daily_loss_frac).round() as i64;
if intent.action == Action::Buy && portfolio.day_pnl_cents <= -max_loss {
return RiskDecision::Reject {
reason: RejectReason::DailyLossKill,
@ -170,8 +172,7 @@ impl RiskGate {
}
// 5. Single-position notional cap (policy).
let position_cap = (portfolio.cash_cents as f64
* self.config.max_position_frac) as i64;
let position_cap = (portfolio.cash_cents as f64 * self.config.max_position_frac) as i64;
if intent.action == Action::Buy && notional > position_cap {
return RiskDecision::Reject {
reason: RejectReason::PositionTooLarge,
@ -181,8 +182,7 @@ impl RiskGate {
// 6. Cluster concentration (policy).
if let Some(cluster) = portfolio.clusters.get(&intent.symbol_id).copied() {
let cluster_cap = (portfolio.cash_cents as f64
* self.config.max_cluster_frac) as i64;
let cluster_cap = (portfolio.cash_cents as f64 * self.config.max_cluster_frac) as i64;
let existing = portfolio.cluster_notional_cents(cluster);
let projected = existing.saturating_add(notional);
if intent.action == Action::Buy && projected > cluster_cap {
@ -243,7 +243,13 @@ mod tests {
fn thin_edge_rejected() {
let gate = RiskGate::default();
let d = gate.evaluate(sample_intent(100, 24, 10), &portfolio(100_000));
assert!(matches!(d, RiskDecision::Reject { reason: RejectReason::EdgeTooThin, .. }));
assert!(matches!(
d,
RiskDecision::Reject {
reason: RejectReason::EdgeTooThin,
..
}
));
}
#[test]
@ -251,7 +257,13 @@ mod tests {
let gate = RiskGate::default();
// 24¢ × 600 = 14_400¢ vs cap 10% × 100_000 = 10_000¢
let d = gate.evaluate(sample_intent(500, 24, 600), &portfolio(100_000));
assert!(matches!(d, RiskDecision::Reject { reason: RejectReason::PositionTooLarge, .. }));
assert!(matches!(
d,
RiskDecision::Reject {
reason: RejectReason::PositionTooLarge,
..
}
));
}
#[test]
@ -260,7 +272,13 @@ mod tests {
let mut p = portfolio(100_000);
p.day_pnl_cents = -3_001; // just past 3%
let d = gate.evaluate(sample_intent(500, 24, 10), &p);
assert!(matches!(d, RiskDecision::Reject { reason: RejectReason::DailyLossKill, .. }));
assert!(matches!(
d,
RiskDecision::Reject {
reason: RejectReason::DailyLossKill,
..
}
));
}
#[test]
@ -287,12 +305,21 @@ mod tests {
// 24¢ × 500 = 12_000 < 20% cap (20_000) but existing 35_000 + 12_000
// = 47_000 > 40% cap (40_000) → concentration rejection.
let d = gate.evaluate(sample_intent(500, 24, 500), &p);
assert!(matches!(d, RiskDecision::Reject { reason: RejectReason::ClusterConcentration, .. }));
assert!(matches!(
d,
RiskDecision::Reject {
reason: RejectReason::ClusterConcentration,
..
}
));
}
#[test]
fn approves_under_paper_config() {
let gate = RiskGate::new(RiskConfig { require_live_flag: false, ..Default::default() });
let gate = RiskGate::new(RiskConfig {
require_live_flag: false,
..Default::default()
});
let d = gate.evaluate(sample_intent(500, 24, 10), &portfolio(100_000));
assert!(matches!(d, RiskDecision::Approve(_)));
}
@ -301,20 +328,50 @@ mod tests {
fn live_gate_rejects_without_env() {
// Ensure env flag is off.
std::env::remove_var("KALSHI_ENABLE_LIVE");
let gate = RiskGate::new(RiskConfig { require_live_flag: true, ..Default::default() });
let gate = RiskGate::new(RiskConfig {
require_live_flag: true,
..Default::default()
});
let d = gate.evaluate(sample_intent(500, 24, 10), &portfolio(100_000));
assert!(matches!(d, RiskDecision::Reject { reason: RejectReason::LiveTradingDisabled, .. }));
assert!(matches!(
d,
RiskDecision::Reject {
reason: RejectReason::LiveTradingDisabled,
..
}
));
}
#[test]
fn rejects_non_positive_qty_and_bad_price() {
let gate = RiskGate::new(RiskConfig { require_live_flag: false, ..Default::default() });
let gate = RiskGate::new(RiskConfig {
require_live_flag: false,
..Default::default()
});
let d = gate.evaluate(sample_intent(500, 24, 0), &portfolio(100_000));
assert!(matches!(d, RiskDecision::Reject { reason: RejectReason::NonPositiveQuantity, .. }));
assert!(matches!(
d,
RiskDecision::Reject {
reason: RejectReason::NonPositiveQuantity,
..
}
));
let d = gate.evaluate(sample_intent(500, 0, 10), &portfolio(100_000));
assert!(matches!(d, RiskDecision::Reject { reason: RejectReason::PriceOutOfRange, .. }));
assert!(matches!(
d,
RiskDecision::Reject {
reason: RejectReason::PriceOutOfRange,
..
}
));
let d = gate.evaluate(sample_intent(500, 100, 10), &portfolio(100_000));
assert!(matches!(d, RiskDecision::Reject { reason: RejectReason::PriceOutOfRange, .. }));
assert!(matches!(
d,
RiskDecision::Reject {
reason: RejectReason::PriceOutOfRange,
..
}
));
}
#[test]
@ -328,6 +385,12 @@ mod tests {
});
// 50¢ × 200 = 10_000; cash only 5_000.
let d = gate.evaluate(sample_intent(500, 50, 200), &portfolio(5_000));
assert!(matches!(d, RiskDecision::Reject { reason: RejectReason::InsufficientCash, .. }));
assert!(matches!(
d,
RiskDecision::Reject {
reason: RejectReason::InsufficientCash,
..
}
));
}
}

65
crates/ruvector-attention/src/attention/flash.rs
View file

@ -185,7 +185,9 @@ impl FlashAttention3 {
}
let d = q[0].len();
if d == 0 {
return Err(AttentionError::InvalidConfig("Dimension must be > 0".into()));
return Err(AttentionError::InvalidConfig(
"Dimension must be > 0".into(),
));
}
let scale = 1.0 / (d as f32).sqrt();
let n_q = q.len();
@ -215,8 +217,8 @@ impl FlashAttention3 {
let kj_end = (kj_start + bc).min(n_kv);
// Track memory reads: Q block + K block + V block
stats.memory_reads += ((qi_end - qi_start) * d
+ (kj_end - kj_start) * d * 2) as u64;
stats.memory_reads +=
((qi_end - qi_start) * d + (kj_end - kj_start) * d * 2) as u64;
// For each query row in this Q block
for qi in qi_start..qi_end {
@ -250,10 +252,7 @@ impl FlashAttention3 {
// Exponentiate and sum
let exp_scores: Vec<f32> =
block_scores.iter().map(|&s| (s - m_ij).exp()).collect();
let l_ij: f32 = exp_scores
.iter()
.filter(|x| x.is_finite())
.sum();
let l_ij: f32 = exp_scores.iter().filter(|x| x.is_finite()).sum();
// Online softmax rescaling
let m_old = row_max[qi];
@ -283,8 +282,7 @@ impl FlashAttention3 {
pv += exp_scores[local_j] * v[kj][dd];
}
}
output[qi][dd] =
scale_old * output[qi][dd] + scale_new * pv;
output[qi][dd] = scale_old * output[qi][dd] + scale_new * pv;
stats.total_flops += (2 * (kj_end - kj_start)) as u64;
}
}
@ -305,11 +303,7 @@ impl FlashAttention3 {
}
}
Ok(FlashOutput {
output,
lse,
stats,
})
Ok(FlashOutput { output, lse, stats })
}
}
@ -393,9 +387,9 @@ impl RingAttention {
for device_id in 0..num_devices {
let local_q = &q_shards[device_id];
if local_q.is_empty() {
return Err(AttentionError::EmptyInput(
format!("Q shard on device {device_id}"),
));
return Err(AttentionError::EmptyInput(format!(
"Q shard on device {device_id}"
)));
}
let d = local_q[0].len();
let n_q = local_q.len();
@ -486,12 +480,7 @@ impl RingAttention {
/// Computes naive (standard) attention for correctness comparison.
/// Returns (output, attention_weights) where output is [n_q, d].
fn naive_attention(
q: &[Vec<f32>],
k: &[Vec<f32>],
v: &[Vec<f32>],
causal: bool,
) -> Vec<Vec<f32>> {
fn naive_attention(q: &[Vec<f32>], k: &[Vec<f32>], v: &[Vec<f32>], causal: bool) -> Vec<Vec<f32>> {
let n_q = q.len();
let n_kv = k.len();
let d = q[0].len();
@ -561,8 +550,12 @@ mod tests {
for qi in 0..n {
for dd in 0..d {
let diff = (flash.output[qi][dd] - naive[qi][dd]).abs();
assert!(diff < 1e-4, "row={qi} col={dd} flash={} naive={} diff={diff}",
flash.output[qi][dd], naive[qi][dd]);
assert!(
diff < 1e-4,
"row={qi} col={dd} flash={} naive={} diff={diff}",
flash.output[qi][dd],
naive[qi][dd]
);
}
}
}
@ -592,9 +585,7 @@ mod tests {
let d = 8;
let n = 4;
// Use large values that could cause overflow without stable softmax
let q: Vec<Vec<f32>> = (0..n)
.map(|i| vec![100.0 * (i as f32 + 1.0); d])
.collect();
let q: Vec<Vec<f32>> = (0..n).map(|i| vec![100.0 * (i as f32 + 1.0); d]).collect();
let k = q.clone();
let v: Vec<Vec<f32>> = (0..n).map(|i| vec![i as f32; d]).collect();
@ -676,16 +667,19 @@ mod tests {
.map(|dev| make_seq(shard_size, d, 0.3 * (dev as f32 + 1.0)))
.collect();
let results =
RingAttention::ring_forward(&q_shards, &k_shards, &v_shards).unwrap();
let results = RingAttention::ring_forward(&q_shards, &k_shards, &v_shards).unwrap();
assert_eq!(results.len(), num_devices);
for (dev_id, res) in results.iter().enumerate() {
assert_eq!(res.output.len(), shard_size);
assert_eq!(res.output[0].len(), d);
// Each device except first does (num_devices - 1) transfers
assert_eq!(res.transfers, num_devices - 1,
"Device {dev_id} should have {} transfers", num_devices - 1);
assert_eq!(
res.transfers,
num_devices - 1,
"Device {dev_id} should have {} transfers",
num_devices - 1
);
for row in &res.output {
for &val in row {
assert!(val.is_finite(), "Device {dev_id} has non-finite output");
@ -760,8 +754,11 @@ mod tests {
let expected_lse = max_s + sum_exp.ln();
let diff = (result.lse[qi] - expected_lse).abs();
assert!(diff < 1e-3, "LSE row={qi} flash={} expected={expected_lse} diff={diff}",
result.lse[qi]);
assert!(
diff < 1e-3,
"LSE row={qi} flash={} expected={expected_lse} diff={diff}",
result.lse[qi]
);
}
}

64
crates/ruvector-attention/src/attention/kv_cache.rs
View file

@ -109,7 +109,11 @@ pub fn round_to_nearest_even(x: f32) -> f32 {
let r = rounded as i64;
if r % 2 != 0 {
// Nudge toward even.
if x > 0.0 { rounded - 1.0 } else { rounded + 1.0 }
if x > 0.0 {
rounded - 1.0
} else {
rounded + 1.0
}
} else {
rounded
}
@ -140,8 +144,16 @@ pub fn quantize_asymmetric(tensor: &[f32], num_heads: usize, bits: u8) -> Quanti
let max_val = channel.iter().copied().fold(f32::NEG_INFINITY, f32::max);
let range = max_val - min_val;
let scale = if range.abs() < f32::EPSILON { 1.0 } else { range / qmax };
let zp = if range.abs() < f32::EPSILON { 0.0 } else { -min_val / scale };
let scale = if range.abs() < f32::EPSILON {
1.0
} else {
range / qmax
};
let zp = if range.abs() < f32::EPSILON {
0.0
} else {
-min_val / scale
};
scales.push(scale);
zero_points.push(zp);
@ -152,7 +164,12 @@ pub fn quantize_asymmetric(tensor: &[f32], num_heads: usize, bits: u8) -> Quanti
}
}
QuantizedTensor { data, scales, zero_points, bits }
QuantizedTensor {
data,
scales,
zero_points,
bits,
}
}
/// Symmetric quantization (simpler, useful for comparison).
@ -163,16 +180,25 @@ pub fn quantize_asymmetric(tensor: &[f32], num_heads: usize, bits: u8) -> Quanti
///
/// Panics if `bits` is less than 2 or greater than 8.
pub fn quantize_symmetric(tensor: &[f32], bits: u8) -> (Vec<u8>, f32) {
assert!(bits >= 2 && bits <= 8, "quantize_symmetric: bits must be in [2, 8], got {}", bits);
assert!(
bits >= 2 && bits <= 8,
"quantize_symmetric: bits must be in [2, 8], got {}",
bits
);
let qmax = ((1u32 << (bits - 1)) - 1) as f32;
let abs_max = tensor.iter().copied().map(f32::abs).fold(0.0_f32, f32::max);
let scale = if abs_max < f32::EPSILON { 1.0 } else { abs_max / qmax };
let scale = if abs_max < f32::EPSILON {
1.0
} else {
abs_max / qmax
};
let offset = (1u32 << (bits - 1)) as f32; // unsigned offset
let data: Vec<u8> = tensor
.iter()
.map(|&v| {
let q = round_to_nearest_even(v / scale + offset).clamp(0.0, (1u32 << bits) as f32 - 1.0);
let q =
round_to_nearest_even(v / scale + offset).clamp(0.0, (1u32 << bits) as f32 - 1.0);
q as u8
})
.collect();
@ -374,7 +400,9 @@ impl CacheManager {
return self.config.max_seq_len;
}
let weight = (total_layers - layer_idx) as f64 / total_layers as f64;
let sum_weights: f64 = (1..=total_layers).map(|i| i as f64 / total_layers as f64).sum();
let sum_weights: f64 = (1..=total_layers)
.map(|i| i as f64 / total_layers as f64)
.sum();
let budget = (weight / sum_weights) * self.config.max_seq_len as f64;
(budget.ceil() as usize).max(1)
}
@ -433,7 +461,10 @@ mod tests {
let qt = quantize_asymmetric(&data, 2, 4);
let restored = dequantize(&qt, 2);
for (orig, rest) in data.iter().zip(restored.iter()) {
assert!((orig - rest).abs() < 0.15, "4-bit error too large: {orig} vs {rest}");
assert!(
(orig - rest).abs() < 0.15,
"4-bit error too large: {orig} vs {rest}"
);
}
}
@ -444,7 +475,10 @@ mod tests {
let restored = dequantize(&qt, 2);
// 3-bit has only 8 levels so error is larger.
for (orig, rest) in data.iter().zip(restored.iter()) {
assert!((orig - rest).abs() < 0.35, "3-bit error too large: {orig} vs {rest}");
assert!(
(orig - rest).abs() < 0.35,
"3-bit error too large: {orig} vs {rest}"
);
}
}
@ -553,7 +587,10 @@ mod tests {
let ratio = mgr.compression_ratio();
// 4-bit in our unpacked scheme: each element uses 1 byte vs 4 bytes in f32,
// but we also store scales/zero-points. Should still be > 1.0.
assert!(ratio > 1.0, "compression ratio should be > 1.0, got {ratio}");
assert!(
ratio > 1.0,
"compression ratio should be > 1.0, got {ratio}"
);
}
#[test]
@ -593,7 +630,10 @@ mod tests {
let b0 = mgr.pyramid_budget(0, 4);
let b3 = mgr.pyramid_budget(3, 4);
// Lower layers should get a larger budget.
assert!(b0 > b3, "layer 0 budget ({b0}) should exceed layer 3 ({b3})");
assert!(
b0 > b3,
"layer 0 budget ({b0}) should exceed layer 3 ({b3})"
);
}
#[test]

216
crates/ruvector-attention/src/attention/mla.rs
View file

@ -29,10 +29,18 @@ pub struct MLAConfig {
impl MLAConfig {
pub fn validate(&self) -> AttentionResult<()> {
let err = |msg: &str| Err(AttentionError::InvalidConfig(msg.into()));
if self.d_model == 0 { return err("d_model must be > 0"); }
if self.num_heads == 0 { return err("num_heads must be > 0"); }
if self.head_dim == 0 { return err("head_dim must be > 0"); }
if self.latent_dim == 0 { return err("latent_dim must be > 0"); }
if self.d_model == 0 {
return err("d_model must be > 0");
}
if self.num_heads == 0 {
return err("num_heads must be > 0");
}
if self.head_dim == 0 {
return err("head_dim must be > 0");
}
if self.latent_dim == 0 {
return err("latent_dim must be > 0");
}
if self.latent_dim >= self.full_kv_dim() {
return err("latent_dim must be < num_heads * head_dim");
}
@ -68,9 +76,12 @@ pub struct MLACache {
impl MLACache {
pub fn new(config: &MLAConfig) -> Self {
Self {
latent_vectors: Vec::new(), rope_keys: Vec::new(),
latent_dim: config.latent_dim, rope_dim: config.rope_dim,
num_heads: config.num_heads, head_dim: config.head_dim,
latent_vectors: Vec::new(),
rope_keys: Vec::new(),
latent_dim: config.latent_dim,
rope_dim: config.rope_dim,
num_heads: config.num_heads,
head_dim: config.head_dim,
}
}
@ -79,8 +90,12 @@ impl MLACache {
self.rope_keys.push(rope_key);
}
pub fn len(&self) -> usize { self.latent_vectors.len() }
pub fn is_empty(&self) -> bool { self.latent_vectors.is_empty() }
pub fn len(&self) -> usize {
self.latent_vectors.len()
}
pub fn is_empty(&self) -> bool {
self.latent_vectors.is_empty()
}
/// Total floats stored in this MLA cache.
pub fn cache_size(&self) -> usize {
@ -94,7 +109,9 @@ impl MLACache {
/// KV-cache reduction ratio (e.g. 0.9375 = 93.75% reduction vs MHA).
pub fn reduction_ratio(&self) -> f32 {
if self.len() == 0 { return 0.0; }
if self.len() == 0 {
return 0.0;
}
1.0 - (self.cache_size() as f32 / self.mha_equivalent_size() as f32)
}
}
@ -129,7 +146,9 @@ impl MLALayer {
})
}
pub fn config(&self) -> &MLAConfig { &self.config }
pub fn config(&self) -> &MLAConfig {
&self.config
}
/// Compress input to KV latent: `c_kv = x @ W_dkv`.
pub fn compress_kv(&self, x: &[f32]) -> Vec<f32> {
@ -138,16 +157,28 @@ impl MLALayer {
/// Decompress latent to keys: `K = c_kv @ W_uk`.
pub fn decompress_keys(&self, c: &[f32]) -> Vec<f32> {
matvec(&self.w_uk, c, self.config.latent_dim, self.config.full_kv_dim())
matvec(
&self.w_uk,
c,
self.config.latent_dim,
self.config.full_kv_dim(),
)
}
/// Decompress latent to values: `V = c_kv @ W_uv`.
pub fn decompress_values(&self, c: &[f32]) -> Vec<f32> {
matvec(&self.w_uv, c, self.config.latent_dim, self.config.full_kv_dim())
matvec(
&self.w_uv,
c,
self.config.latent_dim,
self.config.full_kv_dim(),
)
}
fn compute_rope_keys(&self, x: &[f32]) -> Vec<f32> {
if self.config.rope_dim == 0 { return Vec::new(); }
if self.config.rope_dim == 0 {
return Vec::new();
}
matvec(&self.w_rope, x, self.config.d_model, self.config.rope_dim)
}
@ -161,7 +192,9 @@ impl MLALayer {
fn apply_rope(v: &mut [f32], position: usize) {
let dim = v.len();
for i in (0..dim).step_by(2) {
if i + 1 >= dim { break; }
if i + 1 >= dim {
break;
}
let freq = 1.0 / (10000.0_f32).powf(i as f32 / dim as f32);
let theta = position as f32 * freq;
let (cos_t, sin_t) = (theta.cos(), theta.sin());
@ -172,9 +205,7 @@ impl MLALayer {
}
/// Core attention computation shared by `forward` and `forward_cached`.
fn attend(
&self, q_full: &[f32], all_keys: &[Vec<f32>], all_values: &[Vec<f32>],
) -> Vec<f32> {
fn attend(&self, q_full: &[f32], all_keys: &[Vec<f32>], all_values: &[Vec<f32>]) -> Vec<f32> {
let (nh, hd) = (self.config.num_heads, self.config.head_dim);
let scale = (hd as f32).sqrt();
let mut out = vec![0.0_f32; nh * hd];
@ -188,45 +219,71 @@ impl MLALayer {
softmax_inplace(&mut scores);
for (si, &w) in scores.iter().enumerate() {
let vh = &all_values[si][off..off + hd];
for d in 0..hd { out[off + d] += w * vh[d]; }
for d in 0..hd {
out[off + d] += w * vh[d];
}
}
}
matvec(&self.w_out, &out, self.config.full_kv_dim(), self.config.d_model)
matvec(
&self.w_out,
&out,
self.config.full_kv_dim(),
self.config.d_model,
)
}
/// Prepares query with RoPE applied to the decoupled portion of each head.
fn prepare_query(&self, input: &[f32], pos: usize) -> Vec<f32> {
let mut q = self.compute_query(input);
let (nh, hd, rd) = (self.config.num_heads, self.config.head_dim, self.config.rope_dim);
let (nh, hd, rd) = (
self.config.num_heads,
self.config.head_dim,
self.config.rope_dim,
);
if rd > 0 {
for h in 0..nh { Self::apply_rope(&mut q[h * hd..h * hd + rd], pos); }
for h in 0..nh {
Self::apply_rope(&mut q[h * hd..h * hd + rd], pos);
}
}
q
}
/// Decompresses a latent+rope pair into full keys/values for one position.
fn decompress_position(
&self, latent: &[f32], rope: &[f32], pos: usize,
&self,
latent: &[f32],
rope: &[f32],
pos: usize,
) -> (Vec<f32>, Vec<f32>) {
let mut keys = self.decompress_keys(latent);
let values = self.decompress_values(latent);
let (nh, hd, rd) = (self.config.num_heads, self.config.head_dim, self.config.rope_dim);
let (nh, hd, rd) = (
self.config.num_heads,
self.config.head_dim,
self.config.rope_dim,
);
if rd > 0 {
let mut rp = rope.to_vec();
Self::apply_rope(&mut rp, pos);
for h in 0..nh { keys[h * hd..h * hd + rd].copy_from_slice(&rp); }
for h in 0..nh {
keys[h * hd..h * hd + rd].copy_from_slice(&rp);
}
}
(keys, values)
}
/// Full MLA forward pass for a single query position.
pub fn forward(
&self, query_input: &[f32], kv_inputs: &[&[f32]],
query_pos: usize, kv_positions: &[usize],
&self,
query_input: &[f32],
kv_inputs: &[&[f32]],
query_pos: usize,
kv_positions: &[usize],
) -> AttentionResult<Vec<f32>> {
if query_input.len() != self.config.d_model {
return Err(AttentionError::DimensionMismatch {
expected: self.config.d_model, actual: query_input.len(),
expected: self.config.d_model,
actual: query_input.len(),
});
}
if kv_inputs.is_empty() {
@ -234,7 +291,8 @@ impl MLALayer {
}
if kv_inputs.len() != kv_positions.len() {
return Err(AttentionError::DimensionMismatch {
expected: kv_inputs.len(), actual: kv_positions.len(),
expected: kv_inputs.len(),
actual: kv_positions.len(),
});
}
let q_full = self.prepare_query(query_input, query_pos);
@ -243,7 +301,8 @@ impl MLALayer {
for (i, &kv) in kv_inputs.iter().enumerate() {
if kv.len() != self.config.d_model {
return Err(AttentionError::DimensionMismatch {
expected: self.config.d_model, actual: kv.len(),
expected: self.config.d_model,
actual: kv.len(),
});
}
let c = self.compress_kv(kv);
@ -257,22 +316,28 @@ impl MLALayer {
/// Forward pass using incremental MLA cache (for autoregressive decoding).
pub fn forward_cached(
&self, query_input: &[f32], new_kv_input: &[f32],
query_pos: usize, cache: &mut MLACache,
&self,
query_input: &[f32],
new_kv_input: &[f32],
query_pos: usize,
cache: &mut MLACache,
) -> AttentionResult<Vec<f32>> {
if new_kv_input.len() != self.config.d_model {
return Err(AttentionError::DimensionMismatch {
expected: self.config.d_model, actual: new_kv_input.len(),
expected: self.config.d_model,
actual: new_kv_input.len(),
});
}
cache.push(self.compress_kv(new_kv_input), self.compute_rope_keys(new_kv_input));
cache.push(
self.compress_kv(new_kv_input),
self.compute_rope_keys(new_kv_input),
);
let q_full = self.prepare_query(query_input, query_pos);
let mut all_k = Vec::with_capacity(cache.len());
let mut all_v = Vec::with_capacity(cache.len());
for pos in 0..cache.len() {
let (k, v) = self.decompress_position(
&cache.latent_vectors[pos], &cache.rope_keys[pos], pos,
);
let (k, v) =
self.decompress_position(&cache.latent_vectors[pos], &cache.rope_keys[pos], pos);
all_k.push(k);
all_v.push(v);
}
@ -284,8 +349,11 @@ impl MLALayer {
let mha = seq_len * 2 * self.config.num_heads * self.config.head_dim;
let mla = seq_len * (self.config.latent_dim + self.config.rope_dim);
MemoryComparison {
seq_len, mha_cache_floats: mha, mla_cache_floats: mla,
mha_cache_bytes: mha * 4, mla_cache_bytes: mla * 4,
seq_len,
mha_cache_floats: mha,
mla_cache_floats: mla,
mha_cache_bytes: mha * 4,
mla_cache_bytes: mla * 4,
reduction_ratio: 1.0 - (mla as f32 / mha as f32),
}
}
@ -304,7 +372,10 @@ pub struct MemoryComparison {
impl Attention for MLALayer {
fn compute(
&self, query: &[f32], keys: &[&[f32]], values: &[&[f32]],
&self,
query: &[f32],
keys: &[&[f32]],
values: &[&[f32]],
) -> AttentionResult<Vec<f32>> {
let _ = values; // MLA derives V from the same inputs as K
let positions: Vec<usize> = (0..keys.len()).collect();
@ -312,14 +383,21 @@ impl Attention for MLALayer {
}
fn compute_with_mask(
&self, query: &[f32], keys: &[&[f32]], values: &[&[f32]],
&self,
query: &[f32],
keys: &[&[f32]],
values: &[&[f32]],
_mask: Option<&[bool]>,
) -> AttentionResult<Vec<f32>> {
self.compute(query, keys, values)
}
fn dim(&self) -> usize { self.config.d_model }
fn num_heads(&self) -> usize { self.config.num_heads }
fn dim(&self) -> usize {
self.config.d_model
}
fn num_heads(&self) -> usize {
self.config.num_heads
}
}
// -- Utility functions --------------------------------------------------------
@ -340,8 +418,13 @@ fn dot(a: &[f32], b: &[f32]) -> f32 {
fn softmax_inplace(s: &mut [f32]) {
let max = s.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
let mut sum = 0.0_f32;
for v in s.iter_mut() { *v = (*v - max).exp(); sum += *v; }
for v in s.iter_mut() { *v /= sum; }
for v in s.iter_mut() {
*v = (*v - max).exp();
sum += *v;
}
for v in s.iter_mut() {
*v /= sum;
}
}
fn init_weight(in_d: usize, out_d: usize) -> Vec<f32> {
@ -358,29 +441,38 @@ mod tests {
fn cfg() -> MLAConfig {
MLAConfig {
d_model: 32, latent_dim: 8, latent_dim_q: None,
num_heads: 4, head_dim: 8, rope_dim: 4,
d_model: 32,
latent_dim: 8,
latent_dim_q: None,
num_heads: 4,
head_dim: 8,
rope_dim: 4,
}
}
#[test]
fn test_config_valid() { assert!(cfg().validate().is_ok()); }
fn test_config_valid() {
assert!(cfg().validate().is_ok());
}
#[test]
fn test_config_latent_too_large() {
let mut c = cfg(); c.latent_dim = 999;
let mut c = cfg();
c.latent_dim = 999;
assert!(c.validate().is_err());
}
#[test]
fn test_config_rope_dim_odd() {
let mut c = cfg(); c.rope_dim = 3;
let mut c = cfg();
c.rope_dim = 3;
assert!(c.validate().is_err());
}
#[test]
fn test_config_zero_heads() {
let mut c = cfg(); c.num_heads = 0;
let mut c = cfg();
c.num_heads = 0;
assert!(c.validate().is_err());
}
@ -407,9 +499,11 @@ mod tests {
fn test_cache_size_reduction() {
let c = cfg();
let mut cache = MLACache::new(&c);
for _ in 0..10 { cache.push(vec![0.0; c.latent_dim], vec![0.0; c.rope_dim]); }
for _ in 0..10 {
cache.push(vec![0.0; c.latent_dim], vec![0.0; c.rope_dim]);
}
assert_eq!(cache.len(), 10);
assert_eq!(cache.cache_size(), 120); // 10 * (8+4)
assert_eq!(cache.cache_size(), 120); // 10 * (8+4)
assert_eq!(cache.mha_equivalent_size(), 640); // 10 * 2*4*8
assert!((cache.reduction_ratio() - 0.8125).abs() < 1e-4);
}
@ -417,8 +511,12 @@ mod tests {
#[test]
fn test_memory_comparison_report() {
let c = MLAConfig {
d_model: 2048, latent_dim: 256, latent_dim_q: None,
num_heads: 16, head_dim: 128, rope_dim: 0,
d_model: 2048,
latent_dim: 256,
latent_dim_q: None,
num_heads: 16,
head_dim: 128,
rope_dim: 0,
};
let layer = MLALayer::new(c).unwrap();
let r = layer.memory_comparison(1024);
@ -450,7 +548,9 @@ mod tests {
let mut v = vec![1.0, 2.0, 3.0, 4.0];
let orig = v.clone();
MLALayer::apply_rope(&mut v, 0);
for (a, b) in v.iter().zip(&orig) { assert!((a - b).abs() < 1e-6); }
for (a, b) in v.iter().zip(&orig) {
assert!((a - b).abs() < 1e-6);
}
}
#[test]
@ -482,7 +582,9 @@ mod tests {
let q = vec![0.1_f32; c.d_model];
let kv1 = vec![0.2_f32; c.d_model];
let kv2 = vec![0.3_f32; c.d_model];
let out = layer.compute(&q, &[&kv1[..], &kv2[..]], &[&kv1[..], &kv2[..]]).unwrap();
let out = layer
.compute(&q, &[&kv1[..], &kv2[..]], &[&kv1[..], &kv2[..]])
.unwrap();
assert_eq!(out.len(), c.d_model);
assert!(out.iter().all(|v| v.is_finite()));
}

92
crates/ruvector-attention/src/attention/speculative.rs
View file

@ -73,11 +73,7 @@ pub trait DraftModel: Send + Sync {
/// Returns a vector of (token_id, probability) pairs representing the
/// draft model's greedy/sampled choices and their probabilities under
/// the draft distribution.
fn draft_tokens(
&self,
prefix: &[TokenId],
gamma: usize,
) -> Vec<(TokenId, f32)>;
fn draft_tokens(&self, prefix: &[TokenId], gamma: usize) -> Vec<(TokenId, f32)>;
}
/// Target model trait: the large, accurate model that verifies drafts.
@ -175,10 +171,8 @@ impl SpeculativeDecoder {
));
}
let draft_tokens: Vec<TokenId> =
draft_results.iter().map(|(t, _)| *t).collect();
let draft_probs: Vec<f32> =
draft_results.iter().map(|(_, p)| *p).collect();
let draft_tokens: Vec<TokenId> = draft_results.iter().map(|(t, _)| *t).collect();
let draft_probs: Vec<f32> = draft_results.iter().map(|(_, p)| *p).collect();
let target_dists = target.verify_batch(prefix, &draft_tokens);
if target_dists.len() < draft_tokens.len() + 1 {
@ -195,9 +189,7 @@ impl SpeculativeDecoder {
let q_i = draft_probs[i];
let p_i = prob_of_token(&target_dists[i], token);
let rng_val = rng_values
.and_then(|v| v.get(i).copied())
.unwrap_or(0.0);
let rng_val = rng_values.and_then(|v| v.get(i).copied()).unwrap_or(0.0);
if p_i >= q_i {
// Accept unconditionally: target agrees at least as much.
@ -207,12 +199,7 @@ impl SpeculativeDecoder {
accepted.push(token);
} else {
// Reject: sample from adjusted distribution max(0, p - q).
let adjusted = sample_adjusted(
&target_dists[i],
&draft_tokens,
&draft_probs,
i,
);
let adjusted = sample_adjusted(&target_dists[i], &draft_tokens, &draft_probs, i);
accepted.push(adjusted);
rejected = true;
break;
@ -266,10 +253,7 @@ fn sample_adjusted(
draft_probs: &[f32],
position: usize,
) -> TokenId {
let mut best_token = target_dist
.first()
.map(|(t, _)| *t)
.unwrap_or(0);
let mut best_token = target_dist.first().map(|(t, _)| *t).unwrap_or(0);
let mut best_score = f32::NEG_INFINITY;
for &(token, p_target) in target_dist {
@ -329,10 +313,8 @@ pub fn medusa_decode(
}
// Each head predicts one position ahead.
let head_predictions: Vec<Vec<(TokenId, f32)>> = heads
.iter()
.map(|h| h.predict(prefix))
.collect();
let head_predictions: Vec<Vec<(TokenId, f32)>> =
heads.iter().map(|h| h.predict(prefix)).collect();
// Build the greedy candidate path (top-1 from each head).
let candidate_path: Vec<TokenId> = head_predictions
@ -391,11 +373,7 @@ pub struct SimpleDraftModel {
}
impl DraftModel for SimpleDraftModel {
fn draft_tokens(
&self,
_prefix: &[TokenId],
gamma: usize,
) -> Vec<(TokenId, f32)> {
fn draft_tokens(&self, _prefix: &[TokenId], gamma: usize) -> Vec<(TokenId, f32)> {
(0..gamma)
.map(|i| {
let token = self.tokens[i % self.tokens.len()];
@ -514,14 +492,9 @@ mod tests {
],
};
let result = SpeculativeDecoder::decode_step(
&[1, 2, 3],
&draft,
&target,
&default_config(),
None,
)
.unwrap();
let result =
SpeculativeDecoder::decode_step(&[1, 2, 3], &draft, &target, &default_config(), None)
.unwrap();
// All 4 draft tokens accepted + 1 bonus = 5 tokens.
assert_eq!(result.tokens.len(), 5);
@ -612,10 +585,7 @@ mod tests {
// Adjusted: max(0, 0.3 - 0.8) = 0 for 10, max(0, 0.7 - 0) = 0.7 for 42.
// So adjusted sample should be 42.
let target = SimpleTargetModel {
distributions: vec![
vec![(10, 0.3), (42, 0.7)],
vec![(99, 1.0)],
],
distributions: vec![vec![(10, 0.3), (42, 0.7)], vec![(99, 1.0)]],
};
let cfg = SpeculativeConfig::new(1);
@ -666,16 +636,11 @@ mod tests {
probability: 0.8,
};
let target = SimpleTargetModel {
distributions: vec![
vec![(10, 0.7)],
vec![(20, 0.6)],
vec![(99, 1.0)],
],
distributions: vec![vec![(10, 0.7)], vec![(20, 0.6)], vec![(99, 1.0)]],
};
let heads: Vec<&dyn MedusaHead> = vec![&h1, &h2];
let result =
medusa_decode(&[1, 2], &heads, &target, &default_config()).unwrap();
let result = medusa_decode(&[1, 2], &heads, &target, &default_config()).unwrap();
assert_eq!(result.tokens, vec![10, 20]);
assert_eq!(result.paths_evaluated, 1);
@ -687,8 +652,7 @@ mod tests {
distributions: vec![vec![(1, 1.0)]],
};
let heads: Vec<&dyn MedusaHead> = vec![];
let result =
medusa_decode(&[1], &heads, &target, &default_config());
let result = medusa_decode(&[1], &heads, &target, &default_config());
assert!(result.is_err());
}
@ -710,14 +674,8 @@ mod tests {
let cfg = SpeculativeConfig::new(1);
// rng = 0.3 < 0.5 (p/q) -> accept
let result = SpeculativeDecoder::decode_step(
&[1],
&draft,
&target,
&cfg,
Some(&[0.3]),
)
.unwrap();
let result =
SpeculativeDecoder::decode_step(&[1], &draft, &target, &cfg, Some(&[0.3])).unwrap();
// Accepted draft token + bonus
assert_eq!(result.tokens, vec![10, 99]);
@ -733,21 +691,11 @@ mod tests {
probability: 0.5,
};
let target = SimpleTargetModel {
distributions: vec![
vec![(5, 0.9)],
vec![(6, 1.0)],
],
distributions: vec![vec![(5, 0.9)], vec![(6, 1.0)]],
};
let cfg = SpeculativeConfig::new(1);
let result = SpeculativeDecoder::decode_step(
&[],
&draft,
&target,
&cfg,
None,
)
.unwrap();
let result = SpeculativeDecoder::decode_step(&[], &draft, &target, &cfg, None).unwrap();
assert_eq!(result.tokens, vec![5, 6]);
}

25
crates/ruvector-attention/src/attention/ssm.rs
View file

@ -138,7 +138,7 @@ fn matvec(matrix: &[f32], x: &[f32], rows: usize, cols: usize) -> Vec<f32> {
pub struct SelectiveSSM {
config: SSMConfig,
// Parameterized as -exp(a_log) to guarantee negative real parts (stability).
a_log: Vec<f32>, // [d_inner * d_state]
a_log: Vec<f32>, // [d_inner * d_state]
// 1D causal conv weights: [d_inner, d_conv]
conv_weight: Vec<f32>,
conv_bias: Vec<f32>, // [d_inner]
@ -204,7 +204,11 @@ impl SelectiveSSM {
pub fn forward(&self, input: &[f32]) -> Vec<f32> {
let d_model = self.config.d_model;
let seq_len = input.len() / d_model;
assert_eq!(input.len(), seq_len * d_model, "input not divisible by d_model");
assert_eq!(
input.len(),
seq_len * d_model,
"input not divisible by d_model"
);
let d_inner = self.config.d_inner();
@ -394,7 +398,11 @@ impl MambaBlock {
}
let ssm_out = self.ssm.forward(&normed);
// Residual connection
input.iter().zip(ssm_out.iter()).map(|(a, b)| a + b).collect()
input
.iter()
.zip(ssm_out.iter())
.map(|(a, b)| a + b)
.collect()
}
/// Single-step inference with residual.
@ -542,7 +550,7 @@ mod tests {
fn test_softplus_values() {
assert!((softplus(0.0) - 0.6931).abs() < 1e-3); // ln(2)
assert!((softplus(1.0) - 1.3133).abs() < 1e-3); // ln(1+e)
// Large x: softplus(x) ≈ x
// Large x: softplus(x) ≈ x
assert!((softplus(25.0) - 25.0).abs() < 1e-3);
// Negative x: approaches 0
assert!(softplus(-25.0) < 1e-3);
@ -551,7 +559,7 @@ mod tests {
#[test]
fn test_silu_values() {
assert!((silu(0.0)).abs() < 1e-6); // 0 * 0.5 = 0
// silu(1) = 1/(1+e^-1) ≈ 0.7311
// silu(1) = 1/(1+e^-1) ≈ 0.7311
assert!((silu(1.0) - 0.7311).abs() < 1e-3);
// silu is odd-ish: silu(-x) ≈ -x * sigmoid(-x)
assert!(silu(-5.0) < 0.0);
@ -632,9 +640,12 @@ mod tests {
};
let schedule = hc.layer_schedule();
assert_eq!(schedule.len(), 8);
let attn_count = schedule.iter().filter(|k| **k == LayerKind::Attention).count();
let attn_count = schedule
.iter()
.filter(|k| **k == LayerKind::Attention)
.count();
assert_eq!(attn_count, 2); // 8 layers, every 4th is attn
// Layers 3, 7 should be Attention
// Layers 3, 7 should be Attention
assert_eq!(schedule[3], LayerKind::Attention);
assert_eq!(schedule[7], LayerKind::Attention);
}

14
crates/ruvector-cluster/tests/integration_tests.rs
View file

@ -8,9 +8,7 @@
use ruvector_cluster::consensus::{DagConsensus, DagVertex, Transaction, TransactionType};
use ruvector_cluster::discovery::{DiscoveryService, GossipDiscovery, StaticDiscovery};
use ruvector_cluster::shard::{ConsistentHashRing, LoadBalancer, ShardMigration, ShardRouter};
use ruvector_cluster::{
ClusterConfig, ClusterManager, ClusterNode, NodeStatus, ShardStatus,
};
use ruvector_cluster::{ClusterConfig, ClusterManager, ClusterNode, NodeStatus, ShardStatus};
use std::collections::HashMap;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::time::Duration;
@ -38,10 +36,7 @@ fn test_config(shard_count: u32, replication_factor: usize) -> ClusterConfig {
}
}
fn test_manager(
shard_count: u32,
replication_factor: usize,
) -> ClusterManager {
fn test_manager(shard_count: u32, replication_factor: usize) -> ClusterManager {
let config = test_config(shard_count, replication_factor);
let discovery = Box::new(StaticDiscovery::new(vec![]));
ClusterManager::new(config, "test-manager".to_string(), discovery).unwrap()
@ -378,7 +373,10 @@ fn test_hash_ring_deterministic_routing() {
let primary1 = ring.get_primary_node("my-key").unwrap();
let primary2 = ring.get_primary_node("my-key").unwrap();
assert_eq!(primary1, primary2, "same key must always route to same node");
assert_eq!(
primary1, primary2,
"same key must always route to same node"
);
}
#[test]

4
crates/ruvector-collections/benches/primality.rs
View file

@ -10,9 +10,7 @@
//! | `next_prime_u64` (2^61) | ≤ 12 µs |
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use ruvector_collections::primality::{
is_prime_u64, next_prime_u64, prev_prime_below_pow2,
};
use ruvector_collections::primality::{is_prime_u64, next_prime_u64, prev_prime_below_pow2};
fn bench_is_prime_u64_worst_case(c: &mut Criterion) {
// The Sinclair witness loop runs to completion only on actual primes,

3
crates/ruvector-collections/src/collection.rs
View file

@ -366,8 +366,7 @@ mod tests {
fn test_config_serialization_roundtrip() {
let config = CollectionConfig::with_dimensions(384);
let json = serde_json::to_string(&config).expect("serialize");
let deserialized: CollectionConfig =
serde_json::from_str(&json).expect("deserialize");
let deserialized: CollectionConfig = serde_json::from_str(&json).expect("deserialize");
assert_eq!(deserialized.dimensions, 384);
}

8
crates/ruvector-collections/src/primality.rs
View file

@ -177,7 +177,9 @@ pub fn is_prime_u128(n: u128, rounds: u8) -> bool {
// Numerical-Recipes-style multiplier; we only need uniformity, not crypto.
let mut state: u128 = n ^ 0x9E37_79B9_7F4A_7C15_F39C_C060_5CED_C835u128;
for _ in 0..rounds {
state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
// Witness in [2, n-2].
let a = 2u128 + (state % (n - 3));
if mr_is_composite_u128(n, d, s, a) {
@ -261,8 +263,8 @@ mod tests {
#[test]
fn small_primes_under_100() {
let known: [u64; 25] = [
2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79,
83, 89, 97,
2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83,
89, 97,
];
for n in 0u64..100 {
assert_eq!(is_prime_u64(n), known.contains(&n), "is_prime_u64({n})");

21
crates/ruvector-collections/tests/primality_pseudoprimes.rs
View file

@ -19,12 +19,18 @@ const SPP_FIRST_11: u64 = 3_825_123_056_546_413_051;
#[test]
fn detects_strong_pseudoprime_2357() {
assert!(!is_prime_u64(SPP_2357), "{SPP_2357} is composite (detected by base 11)");
assert!(
!is_prime_u64(SPP_2357),
"{SPP_2357} is composite (detected by base 11)"
);
}
#[test]
fn detects_strong_pseudoprime_235711() {
assert!(!is_prime_u64(SPP_235711), "{SPP_235711} is composite (detected by base 13)");
assert!(
!is_prime_u64(SPP_235711),
"{SPP_235711} is composite (detected by base 13)"
);
}
#[test]
@ -38,8 +44,8 @@ fn detects_strong_pseudoprime_first_11_primes() {
#[test]
fn small_prime_sanity_under_100() {
let primes_under_100: [u64; 25] = [
2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83,
89, 97,
2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89,
97,
];
for n in 0u64..=100 {
let expected = primes_under_100.contains(&n);
@ -52,7 +58,10 @@ fn edge_cases() {
assert!(!is_prime_u64(0));
assert!(!is_prime_u64(1));
assert!(!is_prime_u64(u64::MAX), "u64::MAX (= 2^64 - 1) factors");
assert!(is_prime_u64(u64::MAX - 58), "largest u64 prime: u64::MAX - 58");
assert!(
is_prime_u64(u64::MAX - 58),
"largest u64 prime: u64::MAX - 58"
);
// Largest u32 prime is 2^32 - 5 = 4_294_967_291.
assert!(is_prime_u32(4_294_967_291), "largest u32 prime");
assert!(!is_prime_u32(u32::MAX));
@ -67,7 +76,7 @@ fn assorted_known_primes() {
8191,
131_071,
524_287,
2_147_483_647, // 2^31 - 1
2_147_483_647, // 2^31 - 1
2_305_843_009_213_693_951u64, // 2^61 - 1
] {
assert!(is_prime_u64(p), "{p} is a known prime");

4
crates/ruvector-collections/tests/table_cross_check.rs
View file

@ -6,9 +6,7 @@
//! assert no other prime hides there. This is what makes MR — not the
//! table — the source of truth.
use ruvector_collections::primality::{
is_prime_u64, PRIMES_ABOVE_2K, PRIMES_BELOW_2K,
};
use ruvector_collections::primality::{is_prime_u64, PRIMES_ABOVE_2K, PRIMES_BELOW_2K};
/// Iterate odd candidates strictly between `lo` (exclusive) and `hi`
/// (exclusive), without overflowing `u64`. Used to confirm the prime gap

22
crates/ruvector-consciousness-wasm/src/lib.rs
View file

@ -15,12 +15,14 @@
use wasm_bindgen::prelude::*;
use ruvector_consciousness::emergence::{CausalEmergenceEngine, effective_information};
use ruvector_consciousness::phi::{auto_compute_phi, ExactPhiEngine, SpectralPhiEngine, StochasticPhiEngine};
use ruvector_consciousness::geomip::GeoMipPhiEngine;
use ruvector_consciousness::collapse::QuantumCollapseEngine;
use ruvector_consciousness::emergence::{effective_information, CausalEmergenceEngine};
use ruvector_consciousness::geomip::GeoMipPhiEngine;
use ruvector_consciousness::phi::{
auto_compute_phi, ExactPhiEngine, SpectralPhiEngine, StochasticPhiEngine,
};
use ruvector_consciousness::rsvd_emergence::RsvdEmergenceEngine;
use ruvector_consciousness::traits::{PhiEngine, EmergenceEngine, ConsciousnessCollapse};
use ruvector_consciousness::traits::{ConsciousnessCollapse, EmergenceEngine, PhiEngine};
use ruvector_consciousness::types::{ComputeBudget, TransitionMatrix};
use serde::Serialize;
@ -213,11 +215,7 @@ impl WasmConsciousness {
/// Compute causal emergence for a TPM.
#[wasm_bindgen(js_name = "computeEmergence")]
pub fn compute_emergence(
&self,
tpm_data: &[f64],
n: usize,
) -> Result<JsValue, JsError> {
pub fn compute_emergence(&self, tpm_data: &[f64], n: usize) -> Result<JsValue, JsError> {
let tpm = TransitionMatrix::new(n, tpm_data.to_vec());
let budget = self.make_budget(1.0);
let engine = CausalEmergenceEngine::default();
@ -290,11 +288,7 @@ impl WasmConsciousness {
/// Compute effective information for a TPM.
#[wasm_bindgen(js_name = "effectiveInformation")]
pub fn effective_info(
&self,
tpm_data: &[f64],
n: usize,
) -> Result<f64, JsError> {
pub fn effective_info(&self, tpm_data: &[f64], n: usize) -> Result<f64, JsError> {
let tpm = TransitionMatrix::new(n, tpm_data.to_vec());
effective_information(&tpm).map_err(|e| JsError::new(&e.to_string()))
}

4
crates/ruvector-consciousness/benches/phi_benchmark.rs
View file

@ -104,9 +104,7 @@ fn bench_phi_hierarchical_16(c: &mut Criterion) {
let tpm = make_tpm(16);
let budget = ComputeBudget::fast();
c.bench_function("phi_hierarchical_n16", |b| {
b.iter(|| {
HierarchicalPhiEngine::new(8).compute_phi(black_box(&tpm), Some(0), &budget)
})
b.iter(|| HierarchicalPhiEngine::new(8).compute_phi(black_box(&tpm), Some(0), &budget))
});
}

33
crates/ruvector-consciousness/src/bounds.rs
View file

@ -14,7 +14,6 @@ use crate::simd::build_mi_matrix;
use crate::traits::PhiEngine;
use crate::types::{ComputeBudget, PhiBound, TransitionMatrix};
// ---------------------------------------------------------------------------
// Spectral bounds
// ---------------------------------------------------------------------------
@ -139,12 +138,7 @@ fn estimate_fiedler(laplacian: &[f64], n: usize, max_iter: usize) -> f64 {
/// `k`: number of samples evaluated.
/// `phi_max`: maximum observed φ (used for range bound).
/// `delta`: failure probability (e.g., 0.05 for 95% confidence).
pub fn hoeffding_bound(
phi_estimate: f64,
k: u64,
phi_max: f64,
delta: f64,
) -> PhiBound {
pub fn hoeffding_bound(phi_estimate: f64, k: u64, phi_max: f64, delta: f64) -> PhiBound {
assert!(delta > 0.0 && delta < 1.0);
assert!(k > 0);
@ -170,10 +164,7 @@ pub fn hoeffding_bound(
///
/// `phi_estimates`: all observed φ values from sampling.
/// `delta`: failure probability.
pub fn empirical_bernstein_bound(
phi_estimates: &[f64],
delta: f64,
) -> PhiBound {
pub fn empirical_bernstein_bound(phi_estimates: &[f64], delta: f64) -> PhiBound {
assert!(!phi_estimates.is_empty());
assert!(delta > 0.0 && delta < 1.0);
@ -181,12 +172,17 @@ pub fn empirical_bernstein_bound(
let mean: f64 = phi_estimates.iter().sum::<f64>() / k;
// Sample variance.
let variance: f64 = phi_estimates.iter()
let variance: f64 = phi_estimates
.iter()
.map(|&x| (x - mean).powi(2))
.sum::<f64>() / (k - 1.0).max(1.0);
.sum::<f64>()
/ (k - 1.0).max(1.0);
// Range bound.
let max_val = phi_estimates.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
let max_val = phi_estimates
.iter()
.cloned()
.fold(f64::NEG_INFINITY, f64::max);
let min_val = phi_estimates.iter().cloned().fold(f64::INFINITY, f64::min);
let phi_min = min_val; // Best estimate (minimum = MIP).
@ -309,8 +305,10 @@ mod tests {
fn hoeffding_bound_narrows_with_samples() {
let b1 = hoeffding_bound(0.5, 100, 1.0, 0.05);
let b2 = hoeffding_bound(0.5, 10000, 1.0, 0.05);
assert!(b2.upper - b2.lower < b1.upper - b1.lower,
"more samples should give tighter bound");
assert!(
b2.upper - b2.lower < b1.upper - b1.lower,
"more samples should give tighter bound"
);
}
#[test]
@ -327,7 +325,8 @@ mod tests {
let tpm = and_gate_tpm();
let engine = SpectralPhiEngine::default();
let budget = ComputeBudget::fast();
let (result, bound) = compute_phi_with_bounds(&engine, &tpm, Some(0), &budget, 0.05).unwrap();
let (result, bound) =
compute_phi_with_bounds(&engine, &tpm, Some(0), &budget, 0.05).unwrap();
assert!(result.phi >= 0.0);
assert!(bound.lower >= 0.0);
}

57
crates/ruvector-consciousness/src/ces.rs
View file

@ -44,7 +44,10 @@ pub fn compute_ces(
// num_elements = log2(n)
let num_elements = n.trailing_zeros() as usize;
if num_elements > 12 {
return Err(ConsciousnessError::SystemTooLarge { n: num_elements, max: 12 });
return Err(ConsciousnessError::SystemTooLarge {
n: num_elements,
max: 12,
});
}
let start = Instant::now();
@ -63,15 +66,35 @@ pub fn compute_ces(
.filter(|d| d.phi > phi_threshold)
.collect()
} else {
ces_sequential(tpm, state, num_elements, full, phi_threshold, &budget, &start)
ces_sequential(
tpm,
state,
num_elements,
full,
phi_threshold,
&budget,
&start,
)
};
#[cfg(not(feature = "parallel"))]
let distinctions = ces_sequential(tpm, state, num_elements, full, phi_threshold, budget, &start);
let distinctions = ces_sequential(
tpm,
state,
num_elements,
full,
phi_threshold,
budget,
&start,
);
// Sort by φ descending.
let mut distinctions = distinctions;
distinctions.sort_by(|a, b| b.phi.partial_cmp(&a.phi).unwrap_or(std::cmp::Ordering::Equal));
distinctions.sort_by(|a, b| {
b.phi
.partial_cmp(&a.phi)
.unwrap_or(std::cmp::Ordering::Equal)
});
// Compute relations between distinctions.
let relations = compute_relations(&distinctions);
@ -128,19 +151,21 @@ fn compute_relations(distinctions: &[Distinction]) -> Vec<Relation> {
// Pairwise relations (order 2).
for i in 0..nd {
for j in (i + 1)..nd {
let overlap_cause = distinctions[i].cause_purview.elements
& distinctions[j].cause_purview.elements;
let overlap_effect = distinctions[i].effect_purview.elements
& distinctions[j].effect_purview.elements;
let overlap_cause =
distinctions[i].cause_purview.elements & distinctions[j].cause_purview.elements;
let overlap_effect =
distinctions[i].effect_purview.elements & distinctions[j].effect_purview.elements;
if overlap_cause != 0 || overlap_effect != 0 {
// Relation φ: geometric mean of the two distinction φ values
// weighted by purview overlap (simplified from full IIT 4.0).
let overlap_size = (overlap_cause.count_ones() + overlap_effect.count_ones()) as f64;
let overlap_size =
(overlap_cause.count_ones() + overlap_effect.count_ones()) as f64;
let total_size = (distinctions[i].cause_purview.size()
+ distinctions[i].effect_purview.size()
+ distinctions[j].cause_purview.size()
+ distinctions[j].effect_purview.size()) as f64;
+ distinctions[j].effect_purview.size())
as f64;
let overlap_fraction = if total_size > 0.0 {
overlap_size / total_size
@ -162,7 +187,11 @@ fn compute_relations(distinctions: &[Distinction]) -> Vec<Relation> {
}
// Sort by φ descending.
relations.sort_by(|a, b| b.phi.partial_cmp(&a.phi).unwrap_or(std::cmp::Ordering::Equal));
relations.sort_by(|a, b| {
b.phi
.partial_cmp(&a.phi)
.unwrap_or(std::cmp::Ordering::Equal)
});
relations
}
@ -218,7 +247,11 @@ fn compute_big_phi(
min_phi = min_phi.min(ces_distance);
}
if min_phi == f64::MAX { 0.0 } else { min_phi }
if min_phi == f64::MAX {
0.0
} else {
min_phi
}
}
/// Quick CES summary: number of distinctions and relations.

14
crates/ruvector-consciousness/src/chebyshev_phi.rs
View file

@ -99,8 +99,12 @@ impl PhiEngine for ChebyshevPhiEngine {
}
}
let full = (1u64 << n) - 1;
if mask == 0 { mask = 1; }
if mask == full { mask = full - 1; }
if mask == 0 {
mask = 1;
}
if mask == full {
mask = full - 1;
}
let partition = Bipartition { mask, n };
let arena = PhiArena::with_capacity(n * 16);
@ -237,6 +241,10 @@ mod tests {
let result = ChebyshevPhiEngine::default()
.compute_phi(&tpm, Some(0), &budget)
.unwrap();
assert!(result.phi < 1e-3, "chebyshev disconnected should be ~0, got {}", result.phi);
assert!(
result.phi < 1e-3,
"chebyshev disconnected should be ~0, got {}",
result.phi
);
}
}

11
crates/ruvector-consciousness/src/coherence_phi.rs
View file

@ -73,7 +73,10 @@ pub struct PhiSpectralBound {
/// Uses spectral gap as a fast proxy. If the gap is above threshold,
/// the system is strongly connected and Φ is likely high.
/// If below, the system has a near-partition and Φ may be low.
pub fn is_highly_integrated(tpm: &TransitionMatrix, threshold: f64) -> Result<bool, ConsciousnessError> {
pub fn is_highly_integrated(
tpm: &TransitionMatrix,
threshold: f64,
) -> Result<bool, ConsciousnessError> {
let bound = spectral_phi_bound(tpm)?;
Ok(bound.spectral_gap > threshold)
}
@ -103,7 +106,11 @@ mod tests {
let tpm = uniform_tpm();
let bound = spectral_phi_bound(&tpm).unwrap();
// Uniform TPM: all MI is zero → Fiedler = 0.
assert!(bound.fiedler_value < 0.1, "uniform should have low fiedler, got {}", bound.fiedler_value);
assert!(
bound.fiedler_value < 0.1,
"uniform should have low fiedler, got {}",
bound.fiedler_value
);
}
#[test]

4
crates/ruvector-consciousness/src/collapse.rs
View file

@ -40,9 +40,7 @@ impl QuantumCollapseEngine {
impl Default for QuantumCollapseEngine {
fn default() -> Self {
Self {
register_size: 256,
}
Self { register_size: 256 }
}
}

10
crates/ruvector-consciousness/src/emergence.rs
View file

@ -207,10 +207,7 @@ impl EmergenceEngine for CausalEmergenceEngine {
})
}
fn effective_information(
&self,
tpm: &TransitionMatrix,
) -> Result<f64, ConsciousnessError> {
fn effective_information(&self, tpm: &TransitionMatrix) -> Result<f64, ConsciousnessError> {
effective_information(tpm)
}
}
@ -364,7 +361,10 @@ mod tests {
// Identity: marginal is uniform, so degeneracy = 0.
let tpm = identity_tpm(4);
let deg = degeneracy(&tpm);
assert!(deg.abs() < 1e-6, "identity TPM degeneracy should be 0, got {deg}");
assert!(
deg.abs() < 1e-6,
"identity TPM degeneracy should be 0, got {deg}"
);
}
#[test]

36
crates/ruvector-consciousness/src/geomip.rs
View file

@ -207,16 +207,18 @@ impl PhiEngine for GeoMipPhiEngine {
for mask in 1..((1u64 << n) - 1) {
let popcount = mask.count_ones() as usize;
if (popcount == half || popcount == half + 1)
&& (!self.prune_automorphisms
|| canonical_partition(mask, n) == mask)
{
balanced_partitions.push(Bipartition { mask, n });
}
&& (!self.prune_automorphisms || canonical_partition(mask, n) == mask)
{
balanced_partitions.push(Bipartition { mask, n });
}
}
// Sort by balance score (most balanced first).
balanced_partitions
.sort_by(|a, b| balance_score(b.mask, n).partial_cmp(&balance_score(a.mask, n)).unwrap());
balanced_partitions.sort_by(|a, b| {
balance_score(b.mask, n)
.partial_cmp(&balance_score(a.mask, n))
.unwrap()
});
for partition in &balanced_partitions {
if self.max_evaluations > 0 && evaluated >= self.max_evaluations {
@ -498,10 +500,12 @@ mod tests {
let tpm = and_gate_tpm();
let budget = ComputeBudget::exact();
let exact_result =
crate::phi::ExactPhiEngine.compute_phi(&tpm, Some(0), &budget).unwrap();
let geomip_result =
GeoMipPhiEngine::default().compute_phi(&tpm, Some(0), &budget).unwrap();
let exact_result = crate::phi::ExactPhiEngine
.compute_phi(&tpm, Some(0), &budget)
.unwrap();
let geomip_result = GeoMipPhiEngine::default()
.compute_phi(&tpm, Some(0), &budget)
.unwrap();
// GeoMIP should evaluate fewer or equal partitions due to pruning.
assert!(
@ -524,12 +528,12 @@ mod tests {
#[test]
fn emd_loss_disconnected_zero() {
let tpm = disconnected_tpm();
let partition = Bipartition {
mask: 0b0011,
n: 4,
};
let partition = Bipartition { mask: 0b0011, n: 4 };
let arena = PhiArena::with_capacity(1024);
let loss = partition_information_loss_emd(&tpm, 0, &partition, &arena);
assert!(loss < 1e-6, "disconnected EMD loss should be ≈ 0, got {loss}");
assert!(
loss < 1e-6,
"disconnected EMD loss should be ≈ 0, got {loss}"
);
}
}

38
crates/ruvector-consciousness/src/iit4.rs
View file

@ -190,11 +190,7 @@ pub fn unconstrained_repertoire(purview_size: usize) -> Vec<f64> {
/// and φ_effect = min over partitions of the effect side.
///
/// This is the IIT 4.0 version using intrinsic_difference instead of KL.
pub fn mechanism_phi(
tpm: &TransitionMatrix,
mechanism: &Mechanism,
state: usize,
) -> Distinction {
pub fn mechanism_phi(tpm: &TransitionMatrix, mechanism: &Mechanism, state: usize) -> Distinction {
let n = tpm.n; // number of states
let num_elements = num_elements_from_states(n);
@ -219,9 +215,8 @@ pub fn mechanism_phi(
let cause_phi = intrinsic_difference(&cause_rep, &uc_rep);
// Find the minimum over partitions of the mechanism for this purview.
let partitioned_cause_phi = min_partition_phi_cause(
tpm, mechanism, &purview, state, &cause_rep,
);
let partitioned_cause_phi =
min_partition_phi_cause(tpm, mechanism, &purview, state, &cause_rep);
if partitioned_cause_phi > best_cause_phi {
best_cause_phi = partitioned_cause_phi;
@ -234,9 +229,8 @@ pub fn mechanism_phi(
let uc_effect = unconstrained_repertoire(purview_size);
let effect_phi = intrinsic_difference(&effect_rep, &uc_effect);
let partitioned_effect_phi = min_partition_phi_effect(
tpm, mechanism, &purview, state, &effect_rep,
);
let partitioned_effect_phi =
min_partition_phi_effect(tpm, mechanism, &purview, state, &effect_rep);
if partitioned_effect_phi > best_effect_phi {
best_effect_phi = partitioned_effect_phi;
@ -305,7 +299,11 @@ fn min_partition_phi_cause(
min_loss = min_loss.min(loss);
}
if min_loss == f64::MAX { 0.0 } else { min_loss }
if min_loss == f64::MAX {
0.0
} else {
min_loss
}
}
/// Minimum partition φ for the effect side.
@ -350,7 +348,11 @@ fn min_partition_phi_effect(
min_loss = min_loss.min(loss);
}
if min_loss == f64::MAX { 0.0 } else { min_loss }
if min_loss == f64::MAX {
0.0
} else {
min_loss
}
}
/// Product of two distributions (element-wise multiply + normalize).
@ -438,7 +440,10 @@ mod tests {
let purview = Purview::new(0b11, 2);
let rep = cause_repertoire(&tpm, &mech, &purview, 0);
let sum: f64 = rep.iter().sum();
assert!((sum - 1.0).abs() < 1e-10, "cause repertoire should sum to 1, got {sum}");
assert!(
(sum - 1.0).abs() < 1e-10,
"cause repertoire should sum to 1, got {sum}"
);
}
#[test]
@ -448,7 +453,10 @@ mod tests {
let purview = Purview::new(0b11, 2);
let rep = effect_repertoire(&tpm, &mech, &purview, 0);
let sum: f64 = rep.iter().sum();
assert!((sum - 1.0).abs() < 1e-10, "effect repertoire should sum to 1, got {sum}");
assert!(
(sum - 1.0).abs() < 1e-10,
"effect repertoire should sum to 1, got {sum}"
);
}
#[test]

11
crates/ruvector-consciousness/src/mincut_phi.rs
View file

@ -71,10 +71,7 @@ impl PhiEngine for MinCutPhiEngine {
let mut convergence = Vec::new();
// Use MinCut builder pattern to find the minimum weight cut.
let mincut_result = MinCutBuilder::new()
.exact()
.with_edges(edges)
.build();
let mincut_result = MinCutBuilder::new().exact().with_edges(edges).build();
if let Ok(mincut) = mincut_result {
let result = mincut.min_cut();
@ -196,6 +193,10 @@ mod tests {
let result = MinCutPhiEngine::default()
.compute_phi(&tpm, Some(0), &budget)
.unwrap();
assert!(result.phi < 1e-3, "disconnected should be ~0, got {}", result.phi);
assert!(
result.phi < 1e-3,
"disconnected should be ~0, got {}",
result.phi
);
}
}

5
crates/ruvector-consciousness/src/parallel.rs
View file

@ -148,7 +148,10 @@ pub struct ParallelStochasticPhiEngine {
impl ParallelStochasticPhiEngine {
pub fn new(total_samples: u64, seed: u64) -> Self {
Self { total_samples, seed }
Self {
total_samples,
seed,
}
}
}

37
crates/ruvector-consciousness/src/phi.rs
View file

@ -55,7 +55,11 @@ pub(crate) fn validate_tpm(tpm: &TransitionMatrix) -> Result<(), ConsciousnessEr
row_sum += val;
}
if (row_sum - 1.0).abs() > 1e-6 {
return Err(ValidationError::InvalidTPM { row: i, sum: row_sum }.into());
return Err(ValidationError::InvalidTPM {
row: i,
sum: row_sum,
}
.into());
}
}
Ok(())
@ -178,9 +182,19 @@ fn compute_product_distribution_fast(
for global_state in 0..n {
let sa = extract_substate(global_state, set_a);
let sb = extract_substate(global_state, set_b);
let pa = if sa < ka { unsafe { *dist_a.get_unchecked(sa) } } else { 0.0 };
let pb = if sb < kb { unsafe { *dist_b.get_unchecked(sb) } } else { 0.0 };
unsafe { *output.get_unchecked_mut(global_state) = pa * pb; }
let pa = if sa < ka {
unsafe { *dist_a.get_unchecked(sa) }
} else {
0.0
};
let pb = if sb < kb {
unsafe { *dist_b.get_unchecked(sb) }
} else {
0.0
};
unsafe {
*output.get_unchecked_mut(global_state) = pa * pb;
}
}
} else {
for global_state in 0..n {
@ -697,7 +711,9 @@ impl HierarchicalPhiEngine {
impl Default for HierarchicalPhiEngine {
fn default() -> Self {
Self { exact_threshold: 12 }
Self {
exact_threshold: 12,
}
}
}
@ -786,13 +802,20 @@ impl PhiEngine for HierarchicalPhiEngine {
}
}
// Fill in the other group's elements.
let other_group = if std::ptr::eq(group, &group_a) { &group_b } else { &group_a };
let other_group = if std::ptr::eq(group, &group_a) {
&group_b
} else {
&group_a
};
for &idx in other_group {
global_mask |= 1 << idx;
}
let full = (1u64 << n) - 1;
if global_mask != 0 && global_mask != full {
best_partition = Bipartition { mask: global_mask, n };
best_partition = Bipartition {
mask: global_mask,
n,
};
}
}
convergence.push(min_phi);

16
crates/ruvector-consciousness/src/pid.rs
View file

@ -116,7 +116,13 @@ fn williams_beer_imin(
let mut min_spec = f64::MAX;
for (source, source_marginal) in sources.iter().zip(source_marginals.iter()) {
let spec = specific_information_cached(
tpm, n, source, target, t_state, &target_marginal, source_marginal,
tpm,
n,
source,
target,
t_state,
&target_marginal,
source_marginal,
);
min_spec = min_spec.min(spec);
}
@ -366,8 +372,12 @@ mod tests {
let target = vec![0, 1];
let result = compute_pid(&tpm, &sources, &target).unwrap();
let sum = result.redundancy + result.unique.iter().sum::<f64>() + result.synergy;
assert!((sum - result.total_mi).abs() < 1e-6,
"PID sum {} should equal total MI {}", sum, result.total_mi);
assert!(
(sum - result.total_mi).abs() < 1e-6,
"PID sum {} should equal total MI {}",
sum,
result.total_mi
);
}
#[test]

30
crates/ruvector-consciousness/src/rsvd_emergence.rs
View file

@ -41,7 +41,12 @@ use std::time::Instant;
/// 5. SVD of B gives approximate singular values
///
/// Complexity: O(n²·(k+p)) vs O(n³) for full SVD.
pub fn randomized_svd(tpm: &TransitionMatrix, k: usize, oversampling: usize, seed: u64) -> Vec<f64> {
pub fn randomized_svd(
tpm: &TransitionMatrix,
k: usize,
oversampling: usize,
seed: u64,
) -> Vec<f64> {
let n = tpm.n;
let rank = (k + oversampling).min(n);
let mut rng = StdRng::seed_from_u64(seed);
@ -248,7 +253,11 @@ pub struct RsvdEmergenceEngine {
impl RsvdEmergenceEngine {
pub fn new(k: usize, oversampling: usize, seed: u64) -> Self {
Self { k, oversampling, seed }
Self {
k,
oversampling,
seed,
}
}
}
@ -358,7 +367,10 @@ mod tests {
let svs = randomized_svd(&tpm, 4, 2, 42);
// Identity matrix has all singular values = 1.
for sv in &svs {
assert!((*sv - 1.0).abs() < 0.1, "identity sv should be ≈ 1, got {sv}");
assert!(
(*sv - 1.0).abs() < 0.1,
"identity sv should be ≈ 1, got {sv}"
);
}
}
@ -380,7 +392,11 @@ mod tests {
let budget = ComputeBudget::fast();
let result = engine.compute(&tpm, &budget).unwrap();
// Identity: all singular values equal → high spectral entropy → low emergence.
assert!(result.emergence_index < 0.5, "identity should have low emergence, got {}", result.emergence_index);
assert!(
result.emergence_index < 0.5,
"identity should have low emergence, got {}",
result.emergence_index
);
}
#[test]
@ -390,7 +406,11 @@ mod tests {
let budget = ComputeBudget::fast();
let result = engine.compute(&tpm, &budget).unwrap();
// Uniform: rank 1 → low spectral entropy → high emergence index.
assert!(result.effective_rank <= 2, "uniform should have low effective rank, got {}", result.effective_rank);
assert!(
result.effective_rank <= 2,
"uniform should have low effective rank, got {}",
result.effective_rank
);
}
#[test]

54
crates/ruvector-consciousness/src/simd.rs
View file

@ -50,10 +50,18 @@ fn kl_divergence_neon_prefetch(p: &[f64], q: &[f64]) -> f64 {
let (p2, q2) = (p[base + 2], q[base + 2]);
let (p3, q3) = (p[base + 3], q[base + 3]);
if p0 > 1e-15 && q0 > 1e-15 { sum0 += p0 * (p0 / q0).ln(); }
if p1 > 1e-15 && q1 > 1e-15 { sum1 += p1 * (p1 / q1).ln(); }
if p2 > 1e-15 && q2 > 1e-15 { sum0 += p2 * (p2 / q2).ln(); }
if p3 > 1e-15 && q3 > 1e-15 { sum1 += p3 * (p3 / q3).ln(); }
if p0 > 1e-15 && q0 > 1e-15 {
sum0 += p0 * (p0 / q0).ln();
}
if p1 > 1e-15 && q1 > 1e-15 {
sum1 += p1 * (p1 / q1).ln();
}
if p2 > 1e-15 && q2 > 1e-15 {
sum0 += p2 * (p2 / q2).ln();
}
if p3 > 1e-15 && q3 > 1e-15 {
sum1 += p3 * (p3 / q3).ln();
}
}
for i in (chunks * 4)..(chunks * 4 + remainder) {
let pi = p[i];
@ -125,10 +133,18 @@ fn entropy_neon_prefetch(p: &[f64]) -> f64 {
for c in 0..chunks {
let base = c * 4;
let (p0, p1, p2, p3) = (p[base], p[base + 1], p[base + 2], p[base + 3]);
if p0 > 1e-15 { h0 -= p0 * p0.ln(); }
if p1 > 1e-15 { h1 -= p1 * p1.ln(); }
if p2 > 1e-15 { h0 -= p2 * p2.ln(); }
if p3 > 1e-15 { h1 -= p3 * p3.ln(); }
if p0 > 1e-15 {
h0 -= p0 * p0.ln();
}
if p1 > 1e-15 {
h1 -= p1 * p1.ln();
}
if p2 > 1e-15 {
h0 -= p2 * p2.ln();
}
if p3 > 1e-15 {
h1 -= p3 * p3.ln();
}
}
for i in (chunks * 4)..(chunks * 4 + remainder) {
let pi = p[i];
@ -317,9 +333,7 @@ pub fn pairwise_mi(tpm: &[f64], n: usize, i: usize, j: usize, marginal: &[f64])
let pj = marginal[j].max(1e-15);
#[cfg(target_arch = "aarch64")]
let pij = {
unsafe { pairwise_dot_neon(tpm, n, i, j) }
};
let pij = { unsafe { pairwise_dot_neon(tpm, n, i, j) } };
#[cfg(not(target_arch = "aarch64"))]
let pij = {
@ -352,10 +366,18 @@ unsafe fn pairwise_dot_neon(tpm: &[f64], n: usize, i: usize, j: usize) -> f64 {
let s1 = s0 + 1;
// Gather strided values into NEON registers
let ai = vld1q_f64(
[*tpm.get_unchecked(s0 * n + i), *tpm.get_unchecked(s1 * n + i)].as_ptr(),
[
*tpm.get_unchecked(s0 * n + i),
*tpm.get_unchecked(s1 * n + i),
]
.as_ptr(),
);
let aj = vld1q_f64(
[*tpm.get_unchecked(s0 * n + j), *tpm.get_unchecked(s1 * n + j)].as_ptr(),
[
*tpm.get_unchecked(s0 * n + j),
*tpm.get_unchecked(s1 * n + j),
]
.as_ptr(),
);
acc = vfmaq_f64(acc, ai, aj);
}
@ -385,7 +407,11 @@ pub fn build_mi_matrix(tpm: &[f64], n: usize) -> Vec<f64> {
}
/// Build MI edge list (i, j, weight) with threshold pruning.
pub fn build_mi_edges(tpm: &[f64], n: usize, threshold: f64) -> (Vec<(usize, usize, f64)>, Vec<f64>) {
pub fn build_mi_edges(
tpm: &[f64],
n: usize,
threshold: f64,
) -> (Vec<(usize, usize, f64)>, Vec<f64>) {
let marginal = marginal_distribution(tpm, n);
let mut edges = Vec::new();
for i in 0..n {

38
crates/ruvector-consciousness/src/sparse_accel.rs
View file

@ -60,11 +60,7 @@ pub fn build_sparse_laplacian(mi_csr: &CsrMatrix<f64>, n: usize) -> CsrMatrix<f6
///
/// Uses shifted inverse iteration: v_{k+1} = (μI - L) * v_k,
/// with deflation against the constant eigenvector.
pub fn sparse_fiedler_vector(
laplacian: &CsrMatrix<f64>,
n: usize,
max_iter: usize,
) -> Vec<f64> {
pub fn sparse_fiedler_vector(laplacian: &CsrMatrix<f64>, n: usize, max_iter: usize) -> Vec<f64> {
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
@ -113,7 +109,11 @@ pub fn sparse_fiedler_vector(
numer += w[i] * lv[i];
denom += w[i] * w[i];
}
if denom > 1e-30 { numer / denom } else { 0.0 }
if denom > 1e-30 {
numer / denom
} else {
0.0
}
};
v.copy_from_slice(&w);
@ -169,7 +169,10 @@ pub struct SparseSpectralPhiEngine {
impl SparseSpectralPhiEngine {
pub fn new(mi_threshold: f64, max_iterations: usize) -> Self {
Self { mi_threshold, max_iterations }
Self {
mi_threshold,
max_iterations,
}
}
}
@ -196,7 +199,12 @@ impl PhiEngine for SparseSpectralPhiEngine {
// Build sparse MI adjacency and Laplacian.
let (mi_csr, nnz) = build_sparse_mi_graph(tpm, self.mi_threshold);
tracing::debug!(n, nnz, density = nnz as f64 / (n * n) as f64, "sparse MI graph built");
tracing::debug!(
n,
nnz,
density = nnz as f64 / (n * n) as f64,
"sparse MI graph built"
);
let laplacian = build_sparse_laplacian(&mi_csr, n);
@ -211,8 +219,12 @@ impl PhiEngine for SparseSpectralPhiEngine {
}
}
let full = (1u64 << n) - 1;
if mask == 0 { mask = 1; }
if mask == full { mask = full - 1; }
if mask == 0 {
mask = 1;
}
if mask == full {
mask = full - 1;
}
let partition = Bipartition { mask, n };
let arena = PhiArena::with_capacity(n * 16);
@ -280,7 +292,11 @@ mod tests {
let result = SparseSpectralPhiEngine::default()
.compute_phi(&tpm, Some(0), &budget)
.unwrap();
assert!(result.phi < 1e-4, "sparse spectral disconnected should be ~0, got {}", result.phi);
assert!(
result.phi < 1e-4,
"sparse spectral disconnected should be ~0, got {}",
result.phi
);
}
#[test]

17
crates/ruvector-consciousness/src/streaming.rs
View file

@ -14,7 +14,6 @@
use crate::traits::PhiEngine;
use crate::types::{ComputeBudget, StreamingPhiResult, TransitionMatrix};
// ---------------------------------------------------------------------------
// Streaming Φ estimator
// ---------------------------------------------------------------------------
@ -106,7 +105,12 @@ impl StreamingPhiEstimator {
engine: &E,
budget: &ComputeBudget,
) -> Option<StreamingPhiResult> {
assert!(state < self.n, "state {} out of range for n={}", state, self.n);
assert!(
state < self.n,
"state {} out of range for n={}",
state,
self.n
);
// Record transition.
if let Some(prev) = self.prev_state {
@ -215,8 +219,8 @@ impl StreamingPhiEstimator {
self.cusum_pos = (self.cusum_pos + deviation).max(0.0);
self.cusum_neg = (self.cusum_neg - deviation).max(0.0);
let detected = self.cusum_pos > self.cusum_threshold
|| self.cusum_neg > self.cusum_threshold;
let detected =
self.cusum_pos > self.cusum_threshold || self.cusum_neg > self.cusum_threshold;
if detected {
// Reset after detection.
@ -269,7 +273,10 @@ mod tests {
got_result = true;
}
}
assert!(got_result, "should produce result after enough observations");
assert!(
got_result,
"should produce result after enough observations"
);
}
#[test]

8
crates/ruvector-consciousness/src/types.rs
View file

@ -115,9 +115,7 @@ impl Bipartition {
/// Elements in set B.
pub fn set_b(&self) -> Vec<usize> {
(0..self.n)
.filter(|&i| self.mask & (1 << i) == 0)
.collect()
(0..self.n).filter(|&i| self.mask & (1 << i) == 0).collect()
}
/// Check if this is a valid bipartition (both sets non-empty).
@ -265,7 +263,9 @@ impl Mechanism {
/// Indices of mechanism elements.
pub fn indices(&self) -> Vec<usize> {
(0..self.n).filter(|&i| self.elements & (1 << i) != 0).collect()
(0..self.n)
.filter(|&i| self.elements & (1 << i) != 0)
.collect()
}
}

36
crates/ruvector-consciousness/tests/integration.rs
View file

@ -120,9 +120,7 @@ fn all_engines_and_gate_positive() {
let tpm = and_gate_tpm();
let budget = ComputeBudget::exact();
let exact = ExactPhiEngine
.compute_phi(&tpm, Some(3), &budget)
.unwrap();
let exact = ExactPhiEngine.compute_phi(&tpm, Some(3), &budget).unwrap();
assert!(exact.phi >= 0.0, "exact: {}", exact.phi);
let geomip = GeoMipPhiEngine::default()
@ -296,8 +294,12 @@ fn rsvd_vs_hoel_emergence_correlation() {
.compute_emergence(&tpm_uni, &budget)
.unwrap();
let rsvd_id = RsvdEmergenceEngine::default().compute(&tpm_id, &budget).unwrap();
let rsvd_uni = RsvdEmergenceEngine::default().compute(&tpm_uni, &budget).unwrap();
let rsvd_id = RsvdEmergenceEngine::default()
.compute(&tpm_id, &budget)
.unwrap();
let rsvd_uni = RsvdEmergenceEngine::default()
.compute(&tpm_uni, &budget)
.unwrap();
// Identity has higher EI than uniform (both systems).
assert!(hoel_id.ei_micro > hoel_uni.ei_micro);
@ -384,9 +386,7 @@ fn budget_limits_partitions() {
max_partitions: 10,
..ComputeBudget::exact()
};
let result = ExactPhiEngine
.compute_phi(&tpm, Some(0), &budget)
.unwrap();
let result = ExactPhiEngine.compute_phi(&tpm, Some(0), &budget).unwrap();
assert!(
result.partitions_evaluated <= 10,
"should respect partition limit, evaluated {}",
@ -483,11 +483,21 @@ fn all_engines_reject_invalid_tpm() {
let bad_tpm = TransitionMatrix::new(2, vec![0.5, 0.5, 0.3, 0.3]);
let budget = ComputeBudget::exact();
assert!(ExactPhiEngine.compute_phi(&bad_tpm, Some(0), &budget).is_err());
assert!(SpectralPhiEngine::default().compute_phi(&bad_tpm, Some(0), &budget).is_err());
assert!(StochasticPhiEngine::new(10, 42).compute_phi(&bad_tpm, Some(0), &budget).is_err());
assert!(GeoMipPhiEngine::default().compute_phi(&bad_tpm, Some(0), &budget).is_err());
assert!(GreedyBisectionPhiEngine::default().compute_phi(&bad_tpm, Some(0), &budget).is_err());
assert!(ExactPhiEngine
.compute_phi(&bad_tpm, Some(0), &budget)
.is_err());
assert!(SpectralPhiEngine::default()
.compute_phi(&bad_tpm, Some(0), &budget)
.is_err());
assert!(StochasticPhiEngine::new(10, 42)
.compute_phi(&bad_tpm, Some(0), &budget)
.is_err());
assert!(GeoMipPhiEngine::default()
.compute_phi(&bad_tpm, Some(0), &budget)
.is_err());
assert!(GreedyBisectionPhiEngine::default()
.compute_phi(&bad_tpm, Some(0), &budget)
.is_err());
}
#[test]

15
crates/ruvector-core/src/advanced_features.rs
View file

@ -16,7 +16,7 @@ pub mod diskann;
pub mod filtered_search;
pub mod graph_rag;
pub use graph_rag::{
CommunityDetection, Community, Entity, GraphRAGConfig, GraphRAGPipeline, KnowledgeGraph,
Community, CommunityDetection, Entity, GraphRAGConfig, GraphRAGPipeline, KnowledgeGraph,
Relation, RetrievalResult,
};
pub mod hybrid_search;
@ -28,9 +28,13 @@ pub mod product_quantization;
pub mod sparse_vector;
// Re-exports
pub use compaction::{BloomFilter, CompactionConfig, LSMIndex, LSMStats, MemTable, Segment};
pub use conformal_prediction::{
ConformalConfig, ConformalPredictor, NonconformityMeasure, PredictionSet,
};
pub use diskann::{
DiskIndex, DiskNode, IOStats, MedoidFinder, PageCache, VamanaConfig, VamanaGraph,
};
pub use filtered_search::{FilterExpression, FilterStrategy, FilteredSearch};
pub use hybrid_search::{HybridConfig, HybridSearch, NormalizationStrategy, BM25};
pub use matryoshka::{FunnelConfig, MatryoshkaConfig, MatryoshkaIndex};
@ -39,12 +43,5 @@ pub use multi_vector::{MultiVectorConfig, MultiVectorIndex, ScoringVariant};
pub use opq::{OPQConfig, OPQIndex, RotationMatrix};
pub use product_quantization::{EnhancedPQ, LookupTable, PQConfig};
pub use sparse_vector::{
FusionConfig, FusionStrategy, ScoredDoc, SparseIndex, SparseVector,
fuse_rankings,
};
pub use diskann::{
DiskIndex, DiskNode, IOStats, MedoidFinder, PageCache, VamanaConfig, VamanaGraph,
};
pub use compaction::{
BloomFilter, CompactionConfig, LSMIndex, LSMStats, MemTable, Segment,
fuse_rankings, FusionConfig, FusionStrategy, ScoredDoc, SparseIndex, SparseVector,
};

380
crates/ruvector-core/src/advanced_features/compaction.rs
View file

@ -9,9 +9,9 @@
//! high-throughput ingestion, streaming embedding updates, and frequent
//! deletes (tombstone-based).
use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet};
use serde::{Deserialize, Serialize};
use crate::types::{SearchResult, VectorId};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet};
/// Configuration for the LSM-tree index.
#[derive(Debug, Clone, Serialize, Deserialize)]
@ -30,14 +30,22 @@ pub struct CompactionConfig {
impl Default for CompactionConfig {
fn default() -> Self {
Self { memtable_capacity: 1000, level_size_ratio: 10, max_levels: 4,
merge_threshold: 4, bloom_fp_rate: 0.01 }
Self {
memtable_capacity: 1000,
level_size_ratio: 10,
max_levels: 4,
merge_threshold: 4,
bloom_fp_rate: 0.01,
}
}
}
/// Probabilistic set using double-hashing: `h_i(x) = h1(x) + i * h2(x)`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BloomFilter { bits: Vec<bool>, num_hashes: usize }
pub struct BloomFilter {
bits: Vec<bool>,
num_hashes: usize,
}
impl BloomFilter {
/// Create a bloom filter for `n` items at `fp_rate`.
@ -47,14 +55,19 @@ impl BloomFilter {
let m = (-(n as f64) * fp.ln() / 2.0_f64.ln().powi(2)).ceil() as usize;
let m = m.max(8);
let k = ((m as f64 / n as f64) * 2.0_f64.ln()).ceil().max(1.0) as usize;
Self { bits: vec![false; m], num_hashes: k }
Self {
bits: vec![false; m],
num_hashes: k,
}
}
/// Insert an element.
pub fn insert(&mut self, key: &str) {
let (h1, h2) = Self::hashes(key);
let m = self.bits.len();
for i in 0..self.num_hashes { self.bits[h1.wrapping_add(i.wrapping_mul(h2)) % m] = true; }
for i in 0..self.num_hashes {
self.bits[h1.wrapping_add(i.wrapping_mul(h2)) % m] = true;
}
}
/// Test membership (may return false positives).
@ -67,7 +80,8 @@ impl BloomFilter {
fn hashes(key: &str) -> (usize, usize) {
let (mut h1, mut h2): (u64, u64) = (0xcbf29ce484222325, 0x517cc1b727220a95);
for &b in key.as_bytes() {
h1 ^= b as u64; h1 = h1.wrapping_mul(0x100000001b3);
h1 ^= b as u64;
h1 = h1.wrapping_mul(0x100000001b3);
h2 = h2.wrapping_mul(31).wrapping_add(b as u64);
}
(h1 as usize, (h2 | 1) as usize)
@ -84,15 +98,36 @@ struct LSMEntry {
/// In-memory sorted write buffer backed by `BTreeMap`.
#[derive(Debug, Clone)]
pub struct MemTable { entries: BTreeMap<VectorId, LSMEntry>, capacity: usize }
pub struct MemTable {
entries: BTreeMap<VectorId, LSMEntry>,
capacity: usize,
}
impl MemTable {
pub fn new(capacity: usize) -> Self { Self { entries: BTreeMap::new(), capacity } }
pub fn new(capacity: usize) -> Self {
Self {
entries: BTreeMap::new(),
capacity,
}
}
/// Insert/update. Returns `true` when full.
pub fn insert(&mut self, id: VectorId, vector: Option<Vec<f32>>,
metadata: Option<HashMap<String, serde_json::Value>>, seq: u64) -> bool {
self.entries.insert(id.clone(), LSMEntry { id, vector, metadata, seq });
pub fn insert(
&mut self,
id: VectorId,
vector: Option<Vec<f32>>,
metadata: Option<HashMap<String, serde_json::Value>>,
seq: u64,
) -> bool {
self.entries.insert(
id.clone(),
LSMEntry {
id,
vector,
metadata,
seq,
},
);
self.is_full()
}
@ -100,16 +135,36 @@ impl MemTable {
pub fn search(&self, query: &[f32], top_k: usize) -> Vec<SearchResult> {
let mut heap: BinaryHeap<(OrdF32, VectorId)> = BinaryHeap::new();
for e in self.entries.values() {
let v = match &e.vector { Some(v) => v, None => continue };
let v = match &e.vector {
Some(v) => v,
None => continue,
};
let d = OrdF32(euclid(query, v));
if heap.len() < top_k { heap.push((d, e.id.clone())); }
else if d < heap.peek().unwrap().0 { heap.pop(); heap.push((d, e.id.clone())); }
if heap.len() < top_k {
heap.push((d, e.id.clone()));
} else if d < heap.peek().unwrap().0 {
heap.pop();
heap.push((d, e.id.clone()));
}
}
let mut r: Vec<SearchResult> = heap.into_sorted_vec().into_iter().filter_map(|(OrdF32(s), id)| {
self.entries.get(&id).map(|e| SearchResult { id: e.id.clone(), score: s,
vector: e.vector.clone(), metadata: e.metadata.clone() })
}).collect();
r.sort_by(|a, b| a.score.partial_cmp(&b.score).unwrap_or(std::cmp::Ordering::Equal)); r
let mut r: Vec<SearchResult> = heap
.into_sorted_vec()
.into_iter()
.filter_map(|(OrdF32(s), id)| {
self.entries.get(&id).map(|e| SearchResult {
id: e.id.clone(),
score: s,
vector: e.vector.clone(),
metadata: e.metadata.clone(),
})
})
.collect();
r.sort_by(|a, b| {
a.score
.partial_cmp(&b.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
r
}
/// Flush to an immutable segment, clearing the memtable.
@ -119,39 +174,80 @@ impl MemTable {
Segment::from_entries(entries, level, fp_rate)
}
pub fn len(&self) -> usize { self.entries.len() }
pub fn is_empty(&self) -> bool { self.entries.is_empty() }
pub fn is_full(&self) -> bool { self.entries.len() >= self.capacity }
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn is_full(&self) -> bool {
self.entries.len() >= self.capacity
}
}
/// Immutable sorted run with bloom filter for point lookups.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Segment { entries: Vec<LSMEntry>, bloom: BloomFilter, pub level: usize }
pub struct Segment {
entries: Vec<LSMEntry>,
bloom: BloomFilter,
pub level: usize,
}
impl Segment {
fn from_entries(entries: Vec<LSMEntry>, level: usize, fp_rate: f64) -> Self {
let mut bloom = BloomFilter::new(entries.len(), fp_rate);
for e in &entries { bloom.insert(&e.id); }
Self { entries, bloom, level }
for e in &entries {
bloom.insert(&e.id);
}
Self {
entries,
bloom,
level,
}
}
pub fn size(&self) -> usize { self.entries.len() }
pub fn contains(&self, id: &str) -> bool { self.bloom.may_contain(id) }
pub fn size(&self) -> usize {
self.entries.len()
}
pub fn contains(&self, id: &str) -> bool {
self.bloom.may_contain(id)
}
/// Brute-force search within this segment.
pub fn search(&self, query: &[f32], top_k: usize) -> Vec<SearchResult> {
let mut heap: BinaryHeap<(OrdF32, usize)> = BinaryHeap::new();
for (i, e) in self.entries.iter().enumerate() {
let v = match &e.vector { Some(v) => v, None => continue };
let v = match &e.vector {
Some(v) => v,
None => continue,
};
let d = OrdF32(euclid(query, v));
if heap.len() < top_k { heap.push((d, i)); }
else if d < heap.peek().unwrap().0 { heap.pop(); heap.push((d, i)); }
if heap.len() < top_k {
heap.push((d, i));
} else if d < heap.peek().unwrap().0 {
heap.pop();
heap.push((d, i));
}
}
let mut r: Vec<SearchResult> = heap.into_sorted_vec().into_iter().map(|(OrdF32(s), i)| {
let e = &self.entries[i];
SearchResult { id: e.id.clone(), score: s, vector: e.vector.clone(), metadata: e.metadata.clone() }
}).collect();
r.sort_by(|a, b| a.score.partial_cmp(&b.score).unwrap_or(std::cmp::Ordering::Equal)); r
let mut r: Vec<SearchResult> = heap
.into_sorted_vec()
.into_iter()
.map(|(OrdF32(s), i)| {
let e = &self.entries[i];
SearchResult {
id: e.id.clone(),
score: s,
vector: e.vector.clone(),
metadata: e.metadata.clone(),
}
})
.collect();
r.sort_by(|a, b| {
a.score
.partial_cmp(&b.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
r
}
/// K-way merge deduplicating by id (highest seq wins). Drops tombstones.
@ -164,7 +260,10 @@ impl Segment {
}
}
}
let entries: Vec<LSMEntry> = merged.into_values().filter(|e| e.vector.is_some()).collect();
let entries: Vec<LSMEntry> = merged
.into_values()
.filter(|e| e.vector.is_some())
.collect();
Segment::from_entries(entries, target_level, fp_rate)
}
}
@ -197,21 +296,33 @@ impl LSMIndex {
pub fn new(config: CompactionConfig) -> Self {
let cap = config.memtable_capacity;
let nl = config.max_levels;
Self { config, memtable: MemTable::new(cap), levels: vec![Vec::new(); nl],
next_seq: 0, bytes_written_user: 0, bytes_written_total: 0,
deleted_ids: HashSet::new() }
Self {
config,
memtable: MemTable::new(cap),
levels: vec![Vec::new(); nl],
next_seq: 0,
bytes_written_user: 0,
bytes_written_total: 0,
deleted_ids: HashSet::new(),
}
}
/// Insert a vector. Auto-flushes and compacts as needed.
pub fn insert(&mut self, id: VectorId, vector: Vec<f32>,
metadata: Option<HashMap<String, serde_json::Value>>) {
pub fn insert(
&mut self,
id: VectorId,
vector: Vec<f32>,
metadata: Option<HashMap<String, serde_json::Value>>,
) {
let bytes = (vector.len() * 4 + id.len()) as u64;
self.bytes_written_user += bytes;
self.bytes_written_total += bytes;
self.deleted_ids.remove(&id);
let seq = self.next_seq; self.next_seq += 1;
let seq = self.next_seq;
self.next_seq += 1;
if self.memtable.insert(id, Some(vector), metadata, seq) {
self.flush_memtable(); self.auto_compact();
self.flush_memtable();
self.auto_compact();
}
}
@ -221,9 +332,11 @@ impl LSMIndex {
self.bytes_written_user += bytes;
self.bytes_written_total += bytes;
self.deleted_ids.insert(id.clone());
let seq = self.next_seq; self.next_seq += 1;
let seq = self.next_seq;
self.next_seq += 1;
if self.memtable.insert(id, None, None, seq) {
self.flush_memtable(); self.auto_compact();
self.flush_memtable();
self.auto_compact();
}
}
@ -232,47 +345,74 @@ impl LSMIndex {
let mut seen = HashSet::new();
let mut all = Vec::new();
for r in self.memtable.search(query, top_k) {
if !self.deleted_ids.contains(&r.id) { seen.insert(r.id.clone()); all.push(r); }
if !self.deleted_ids.contains(&r.id) {
seen.insert(r.id.clone());
all.push(r);
}
}
for level in &self.levels {
for seg in level.iter().rev() {
for r in seg.search(query, top_k) {
if !seen.contains(&r.id) && !self.deleted_ids.contains(&r.id) {
seen.insert(r.id.clone()); all.push(r);
seen.insert(r.id.clone());
all.push(r);
}
}
}
}
all.sort_by(|a, b| a.score.partial_cmp(&b.score).unwrap_or(std::cmp::Ordering::Equal));
all.truncate(top_k); all
all.sort_by(|a, b| {
a.score
.partial_cmp(&b.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
all.truncate(top_k);
all
}
/// Manual compaction across all levels.
pub fn compact(&mut self) {
if !self.memtable.is_empty() { self.flush_memtable(); }
if !self.memtable.is_empty() {
self.flush_memtable();
}
for l in 0..self.config.max_levels.saturating_sub(1) {
if self.levels[l].len() >= 2 { self.compact_level(l); }
if self.levels[l].len() >= 2 {
self.compact_level(l);
}
}
}
/// Auto-compact levels exceeding `merge_threshold`.
pub fn auto_compact(&mut self) {
for l in 0..self.config.max_levels.saturating_sub(1) {
if self.levels[l].len() >= self.config.merge_threshold { self.compact_level(l); }
if self.levels[l].len() >= self.config.merge_threshold {
self.compact_level(l);
}
}
}
pub fn stats(&self) -> LSMStats {
let spl: Vec<usize> = self.levels.iter().map(|l| l.len()).collect();
let total = self.memtable.len()
+ self.levels.iter().flat_map(|l| l.iter()).map(|s| s.size()).sum::<usize>();
LSMStats { num_levels: self.levels.len(), segments_per_level: spl,
total_entries: total, write_amplification: self.write_amplification() }
+ self
.levels
.iter()
.flat_map(|l| l.iter())
.map(|s| s.size())
.sum::<usize>();
LSMStats {
num_levels: self.levels.len(),
segments_per_level: spl,
total_entries: total,
write_amplification: self.write_amplification(),
}
}
pub fn write_amplification(&self) -> f64 {
if self.bytes_written_user == 0 { 1.0 }
else { self.bytes_written_total as f64 / self.bytes_written_user as f64 }
if self.bytes_written_user == 0 {
1.0
} else {
self.bytes_written_total as f64 / self.bytes_written_user as f64
}
}
fn flush_memtable(&mut self) {
@ -283,7 +423,9 @@ impl LSMIndex {
fn compact_level(&mut self, level: usize) {
let target = level + 1;
if target >= self.config.max_levels { return; }
if target >= self.config.max_levels {
return;
}
let segments = std::mem::take(&mut self.levels[level]);
let merged = Segment::merge(&segments, target, self.config.bloom_fp_rate);
self.bytes_written_total += entry_bytes(&merged.entries);
@ -292,33 +434,49 @@ impl LSMIndex {
}
fn entry_bytes(entries: &[LSMEntry]) -> u64 {
entries.iter().map(|e| {
(e.vector.as_ref().map_or(0, |v| v.len() * 4) + e.id.len()) as u64
}).sum()
entries
.iter()
.map(|e| (e.vector.as_ref().map_or(0, |v| v.len() * 4) + e.id.len()) as u64)
.sum()
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
struct OrdF32(f32);
impl Eq for OrdF32 {}
impl PartialOrd for OrdF32 {
fn partial_cmp(&self, o: &Self) -> Option<std::cmp::Ordering> { Some(self.cmp(o)) }
fn partial_cmp(&self, o: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(o))
}
}
impl Ord for OrdF32 {
fn cmp(&self, o: &Self) -> std::cmp::Ordering {
self.0.partial_cmp(&o.0).unwrap_or(std::cmp::Ordering::Equal)
self.0
.partial_cmp(&o.0)
.unwrap_or(std::cmp::Ordering::Equal)
}
}
fn euclid(a: &[f32], b: &[f32]) -> f32 {
a.iter().zip(b).map(|(x, y)| (x - y).powi(2)).sum::<f32>().sqrt()
a.iter()
.zip(b)
.map(|(x, y)| (x - y).powi(2))
.sum::<f32>()
.sqrt()
}
#[cfg(test)]
mod tests {
use super::*;
fn v(dim: usize, val: f32) -> Vec<f32> { vec![val; dim] }
fn v(dim: usize, val: f32) -> Vec<f32> {
vec![val; dim]
}
fn entry(id: &str, vec: Option<Vec<f32>>, seq: u64) -> LSMEntry {
LSMEntry { id: id.into(), vector: vec, metadata: None, seq }
LSMEntry {
id: id.into(),
vector: vec,
metadata: None,
seq,
}
}
#[test]
@ -379,21 +537,35 @@ mod tests {
#[test]
fn bloom_filter_no_false_negatives() {
let mut bf = BloomFilter::new(100, 0.01);
for i in 0..100 { bf.insert(&format!("key-{i}")); }
for i in 0..100 { assert!(bf.may_contain(&format!("key-{i}"))); }
for i in 0..100 {
bf.insert(&format!("key-{i}"));
}
for i in 0..100 {
assert!(bf.may_contain(&format!("key-{i}")));
}
}
#[test]
fn bloom_filter_low_false_positive_rate() {
let mut bf = BloomFilter::new(1000, 0.01);
for i in 0..1000 { bf.insert(&format!("present-{i}")); }
let fp: usize = (0..10_000).filter(|i| bf.may_contain(&format!("absent-{i}"))).count();
assert!((fp as f64 / 10_000.0) < 0.05, "FP rate too high: {fp}/10000");
for i in 0..1000 {
bf.insert(&format!("present-{i}"));
}
let fp: usize = (0..10_000)
.filter(|i| bf.may_contain(&format!("absent-{i}")))
.count();
assert!(
(fp as f64 / 10_000.0) < 0.05,
"FP rate too high: {fp}/10000"
);
}
#[test]
fn lsm_insert_and_search() {
let mut idx = LSMIndex::new(CompactionConfig { memtable_capacity: 10, ..Default::default() });
let mut idx = LSMIndex::new(CompactionConfig {
memtable_capacity: 10,
..Default::default()
});
idx.insert("v1".into(), vec![1.0, 0.0], None);
idx.insert("v2".into(), vec![0.0, 1.0], None);
let r = idx.search(&[1.0, 0.0], 1);
@ -403,7 +575,10 @@ mod tests {
#[test]
fn lsm_delete_with_tombstone() {
let mut idx = LSMIndex::new(CompactionConfig { memtable_capacity: 100, ..Default::default() });
let mut idx = LSMIndex::new(CompactionConfig {
memtable_capacity: 100,
..Default::default()
});
idx.insert("v1".into(), vec![1.0, 0.0], None);
idx.insert("v2".into(), vec![0.0, 1.0], None);
idx.delete("v1".into());
@ -414,26 +589,47 @@ mod tests {
#[test]
fn lsm_auto_compaction_trigger() {
let cfg = CompactionConfig { memtable_capacity: 2, merge_threshold: 2, max_levels: 3, ..Default::default() };
let cfg = CompactionConfig {
memtable_capacity: 2,
merge_threshold: 2,
max_levels: 3,
..Default::default()
};
let mut idx = LSMIndex::new(cfg);
for i in 0..10 { idx.insert(format!("v{i}"), vec![i as f32], None); }
for i in 0..10 {
idx.insert(format!("v{i}"), vec![i as f32], None);
}
assert!(idx.stats().segments_per_level[0] < 4, "L0 should compact");
}
#[test]
fn lsm_multi_level_compaction() {
let cfg = CompactionConfig { memtable_capacity: 2, merge_threshold: 2, max_levels: 4, ..Default::default() };
let cfg = CompactionConfig {
memtable_capacity: 2,
merge_threshold: 2,
max_levels: 4,
..Default::default()
};
let mut idx = LSMIndex::new(cfg);
for i in 0..30 { idx.insert(format!("v{i}"), v(4, i as f32), None); }
for i in 0..30 {
idx.insert(format!("v{i}"), v(4, i as f32), None);
}
let total_seg: usize = idx.stats().segments_per_level.iter().sum();
assert!(total_seg >= 1);
}
#[test]
fn lsm_write_amplification_increases() {
let cfg = CompactionConfig { memtable_capacity: 5, merge_threshold: 2, max_levels: 3, ..Default::default() };
let cfg = CompactionConfig {
memtable_capacity: 5,
merge_threshold: 2,
max_levels: 3,
..Default::default()
};
let mut idx = LSMIndex::new(cfg);
for i in 0..20 { idx.insert(format!("v{i}"), v(4, i as f32), None); }
for i in 0..20 {
idx.insert(format!("v{i}"), v(4, i as f32), None);
}
assert!(idx.write_amplification() >= 1.0);
}
@ -448,9 +644,16 @@ mod tests {
#[test]
fn lsm_large_batch_insert() {
let cfg = CompactionConfig { memtable_capacity: 50, merge_threshold: 4, max_levels: 4, ..Default::default() };
let cfg = CompactionConfig {
memtable_capacity: 50,
merge_threshold: 4,
max_levels: 4,
..Default::default()
};
let mut idx = LSMIndex::new(cfg);
for i in 0..500 { idx.insert(format!("v{i}"), v(8, i as f32 * 0.01), None); }
for i in 0..500 {
idx.insert(format!("v{i}"), v(8, i as f32 * 0.01), None);
}
assert!(idx.stats().total_entries > 0);
let r = idx.search(&v(8, 0.0), 5);
assert_eq!(r.len(), 5);
@ -459,9 +662,16 @@ mod tests {
#[test]
fn lsm_search_across_levels() {
let cfg = CompactionConfig { memtable_capacity: 3, merge_threshold: 3, max_levels: 3, ..Default::default() };
let cfg = CompactionConfig {
memtable_capacity: 3,
merge_threshold: 3,
max_levels: 3,
..Default::default()
};
let mut idx = LSMIndex::new(cfg);
for i in 0..9 { idx.insert(format!("v{i}"), vec![i as f32, 0.0], None); }
for i in 0..9 {
idx.insert(format!("v{i}"), vec![i as f32, 0.0], None);
}
idx.insert("latest".into(), vec![0.0, 0.0], None);
let r = idx.search(&[0.0, 0.0], 3);
assert_eq!(r.len(), 3);

246
crates/ruvector-core/src/advanced_features/diskann.rs
View file

@ -18,8 +18,8 @@
use crate::error::{Result, RuvectorError};
use serde::{Deserialize, Serialize};
use std::collections::{BinaryHeap, HashMap, HashSet};
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap, HashSet};
/// Configuration for the Vamana graph index.
#[derive(Debug, Clone, Serialize, Deserialize)]
@ -38,7 +38,13 @@ pub struct VamanaConfig {
impl Default for VamanaConfig {
fn default() -> Self {
Self { max_degree: 32, search_list_size: 64, alpha: 1.2, num_build_threads: 1, ssd_page_size: 4096 }
Self {
max_degree: 32,
search_list_size: 64,
alpha: 1.2,
num_build_threads: 1,
ssd_page_size: 4096,
}
}
}
@ -46,13 +52,19 @@ impl VamanaConfig {
/// Validate configuration parameters.
pub fn validate(&self) -> Result<()> {
if self.max_degree == 0 {
return Err(RuvectorError::InvalidParameter("max_degree must be > 0".into()));
return Err(RuvectorError::InvalidParameter(
"max_degree must be > 0".into(),
));
}
if self.search_list_size < 1 {
return Err(RuvectorError::InvalidParameter("search_list_size must be >= 1".into()));
return Err(RuvectorError::InvalidParameter(
"search_list_size must be >= 1".into(),
));
}
if self.alpha < 1.0 {
return Err(RuvectorError::InvalidParameter("alpha must be >= 1.0".into()));
return Err(RuvectorError::InvalidParameter(
"alpha must be >= 1.0".into(),
));
}
Ok(())
}
@ -77,22 +89,39 @@ impl VamanaGraph {
config.validate()?;
let n = vectors.len();
if n == 0 {
return Ok(Self { neighbors: vec![], vectors: vec![], medoid: 0, config });
return Ok(Self {
neighbors: vec![],
vectors: vec![],
medoid: 0,
config,
});
}
let dim = vectors[0].len();
for v in &vectors {
if v.len() != dim {
return Err(RuvectorError::DimensionMismatch { expected: dim, actual: v.len() });
return Err(RuvectorError::DimensionMismatch {
expected: dim,
actual: v.len(),
});
}
}
let medoid = MedoidFinder::find_medoid(&vectors);
let mut graph = Self { neighbors: vec![vec![]; n], vectors, medoid, config };
let mut graph = Self {
neighbors: vec![vec![]; n],
vectors,
medoid,
config,
};
// Initialize with sequential neighbors.
for i in 0..n {
let mut nb = Vec::new();
for j in 0..n.min(graph.config.max_degree + 1) {
if j != i { nb.push(j as u32); }
if nb.len() >= graph.config.max_degree { break; }
if j != i {
nb.push(j as u32);
}
if nb.len() >= graph.config.max_degree {
break;
}
}
graph.neighbors[i] = nb;
}
@ -102,7 +131,9 @@ impl VamanaGraph {
let (cands, _) = graph.greedy_search_internal(&query, graph.config.search_list_size);
let mut cset: Vec<u32> = cands.into_iter().filter(|&c| c != i as u32).collect();
for &nb in &graph.neighbors[i] {
if !cset.contains(&nb) { cset.push(nb); }
if !cset.contains(&nb) {
cset.push(nb);
}
}
let pruned = graph.robust_prune(i as u32, &cset);
graph.neighbors[i] = pruned.clone();
@ -122,7 +153,9 @@ impl VamanaGraph {
/// Greedy beam search returning top_k (node_id, distance) pairs.
pub fn search(&self, query: &[f32], top_k: usize) -> Vec<(u32, f32)> {
if self.vectors.is_empty() { return vec![]; }
if self.vectors.is_empty() {
return vec![];
}
let beam = self.config.search_list_size.max(top_k);
let (ids, dists) = self.greedy_search_internal(query, beam);
ids.into_iter().zip(dists).take(top_k).collect()
@ -152,22 +185,31 @@ impl VamanaGraph {
}
results.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
results.truncate(list_size);
(results.iter().map(|r| r.1).collect(), results.iter().map(|r| r.0).collect())
(
results.iter().map(|r| r.1).collect(),
results.iter().map(|r| r.0).collect(),
)
}
/// Robust prune: greedily select diverse neighbors via the alpha-RNG rule.
fn robust_prune(&self, node_id: u32, candidates: &[u32]) -> Vec<u32> {
let nv = &self.vectors[node_id as usize];
let mut scored: Vec<(f32, u32)> = candidates.iter()
let mut scored: Vec<(f32, u32)> = candidates
.iter()
.filter(|&&c| c != node_id)
.map(|&c| (l2_sq(nv, &self.vectors[c as usize]), c))
.collect();
scored.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
let mut sel: Vec<u32> = Vec::new();
for (d2n, cand) in scored {
if sel.len() >= self.config.max_degree { break; }
if sel.len() >= self.config.max_degree {
break;
}
let cv = &self.vectors[cand as usize];
if sel.iter().all(|&s| d2n <= self.config.alpha * l2_sq(&self.vectors[s as usize], cv)) {
if sel
.iter()
.all(|&s| d2n <= self.config.alpha * l2_sq(&self.vectors[s as usize], cv))
{
sel.push(cand);
}
}
@ -203,16 +245,32 @@ pub struct DiskIndex {
impl DiskIndex {
/// Create from a built VamanaGraph.
pub fn from_graph(graph: &VamanaGraph, cache_size_pages: usize) -> Self {
let nodes = (0..graph.vectors.len()).map(|i| DiskNode {
node_id: i as u32, neighbors: graph.neighbors[i].clone(), vector: graph.vectors[i].clone(),
}).collect();
Self { nodes, page_size: graph.config.ssd_page_size, medoid: graph.medoid, cache: PageCache::new(cache_size_pages) }
let nodes = (0..graph.vectors.len())
.map(|i| DiskNode {
node_id: i as u32,
neighbors: graph.neighbors[i].clone(),
vector: graph.vectors[i].clone(),
})
.collect();
Self {
nodes,
page_size: graph.config.ssd_page_size,
medoid: graph.medoid,
cache: PageCache::new(cache_size_pages),
}
}
/// Beam search with IO accounting.
pub fn search_disk(&mut self, query: &[f32], top_k: usize, beam_width: usize) -> (Vec<(u32, f32)>, IOStats) {
pub fn search_disk(
&mut self,
query: &[f32],
top_k: usize,
beam_width: usize,
) -> (Vec<(u32, f32)>, IOStats) {
let mut stats = IOStats::default();
if self.nodes.is_empty() { return (vec![], stats); }
if self.nodes.is_empty() {
return (vec![], stats);
}
let mut visited = HashSet::new();
let mut frontier: BinaryHeap<Reverse<OrdF32Pair>> = BinaryHeap::new();
let mut results: Vec<(f32, u32)> = Vec::new();
@ -243,16 +301,30 @@ impl DiskIndex {
fn read_node(&mut self, node_id: u32, stats: &mut IOStats) -> &DiskNode {
let page_id = node_id as usize;
if self.cache.get(page_id) { stats.cache_hits += 1; }
else { stats.pages_read += 1; stats.bytes_read += self.page_size; self.cache.insert(page_id); }
if self.cache.get(page_id) {
stats.cache_hits += 1;
} else {
stats.pages_read += 1;
stats.bytes_read += self.page_size;
self.cache.insert(page_id);
}
&self.nodes[node_id as usize]
}
/// Filtered search: predicates evaluated during traversal (not post-filter).
/// Ineligible nodes still expand the frontier to preserve graph connectivity.
pub fn search_with_filter<F>(&mut self, query: &[f32], filter_fn: F, top_k: usize) -> Vec<(u32, f32)>
where F: Fn(u32) -> bool {
if self.nodes.is_empty() { return vec![]; }
pub fn search_with_filter<F>(
&mut self,
query: &[f32],
filter_fn: F,
top_k: usize,
) -> Vec<(u32, f32)>
where
F: Fn(u32) -> bool,
{
if self.nodes.is_empty() {
return vec![];
}
let mut visited = HashSet::new();
let mut frontier: BinaryHeap<Reverse<OrdF32Pair>> = BinaryHeap::new();
let mut results: Vec<(f32, u32)> = Vec::new();
@ -261,7 +333,9 @@ impl DiskIndex {
let d = l2_sq(&self.read_node(start, &mut io).vector.clone(), query);
frontier.push(Reverse(OrdF32Pair(d, start)));
visited.insert(start);
if filter_fn(start) { results.push((d, start)); }
if filter_fn(start) {
results.push((d, start));
}
while let Some(Reverse(OrdF32Pair(_, cur))) = frontier.pop() {
let nbs = self.read_node(cur, &mut io).neighbors.clone();
for nb in nbs {
@ -269,7 +343,9 @@ impl DiskIndex {
let v = self.read_node(nb, &mut io).vector.clone();
let dist = l2_sq(&v, query);
frontier.push(Reverse(OrdF32Pair(dist, nb)));
if filter_fn(nb) { results.push((dist, nb)); }
if filter_fn(nb) {
results.push((dist, nb));
}
}
}
}
@ -291,7 +367,13 @@ pub struct PageCache {
impl PageCache {
pub fn new(capacity: usize) -> Self {
Self { capacity, clock: 0, entries: HashMap::new(), total_hits: 0, total_accesses: 0 }
Self {
capacity,
clock: 0,
entries: HashMap::new(),
total_hits: 0,
total_accesses: 0,
}
}
/// Returns true on cache hit, updating recency.
@ -299,16 +381,28 @@ impl PageCache {
self.total_accesses += 1;
self.clock += 1;
if let Some(ts) = self.entries.get_mut(&page_id) {
*ts = self.clock; self.total_hits += 1; true
} else { false }
*ts = self.clock;
self.total_hits += 1;
true
} else {
false
}
}
/// Insert a page, evicting LRU if at capacity.
pub fn insert(&mut self, page_id: usize) {
if self.capacity == 0 { return; }
if self.capacity == 0 {
return;
}
if self.entries.len() >= self.capacity {
let lru = self.entries.iter().min_by_key(|&(_, ts)| *ts).map(|(&k, _)| k);
if let Some(k) = lru { self.entries.remove(&k); }
let lru = self
.entries
.iter()
.min_by_key(|&(_, ts)| *ts)
.map(|(&k, _)| k);
if let Some(k) = lru {
self.entries.remove(&k);
}
}
self.clock += 1;
self.entries.insert(page_id, self.clock);
@ -316,7 +410,11 @@ impl PageCache {
/// Cache hit rate in [0.0, 1.0].
pub fn cache_hit_rate(&self) -> f64 {
if self.total_accesses == 0 { 0.0 } else { self.total_hits as f64 / self.total_accesses as f64 }
if self.total_accesses == 0 {
0.0
} else {
self.total_hits as f64 / self.total_accesses as f64
}
}
}
@ -325,11 +423,18 @@ pub struct MedoidFinder;
impl MedoidFinder {
pub fn find_medoid(vectors: &[Vec<f32>]) -> u32 {
if vectors.is_empty() { return 0; }
if vectors.is_empty() {
return 0;
}
let (mut best_idx, mut best_sum) = (0u32, f32::MAX);
for i in 0..vectors.len() {
let sum: f32 = (0..vectors.len()).map(|j| l2_sq(&vectors[i], &vectors[j])).sum();
if sum < best_sum { best_sum = sum; best_idx = i as u32; }
let sum: f32 = (0..vectors.len())
.map(|j| l2_sq(&vectors[i], &vectors[j]))
.sum();
if sum < best_sum {
best_sum = sum;
best_idx = i as u32;
}
}
best_idx
}
@ -344,11 +449,16 @@ fn l2_sq(a: &[f32], b: &[f32]) -> f32 {
struct OrdF32Pair(f32, u32);
impl Eq for OrdF32Pair {}
impl PartialOrd for OrdF32Pair {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { Some(self.cmp(other)) }
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for OrdF32Pair {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.0.partial_cmp(&other.0).unwrap_or(std::cmp::Ordering::Equal).then(self.1.cmp(&other.1))
self.0
.partial_cmp(&other.0)
.unwrap_or(std::cmp::Ordering::Equal)
.then(self.1.cmp(&other.1))
}
}
@ -357,17 +467,25 @@ mod tests {
use super::*;
fn make_vecs(n: usize, dim: usize) -> Vec<Vec<f32>> {
(0..n).map(|i| (0..dim).map(|d| (i * dim + d) as f32).collect()).collect()
(0..n)
.map(|i| (0..dim).map(|d| (i * dim + d) as f32).collect())
.collect()
}
fn default_cfg(r: usize, l: usize) -> VamanaConfig {
VamanaConfig { max_degree: r, search_list_size: l, ..Default::default() }
VamanaConfig {
max_degree: r,
search_list_size: l,
..Default::default()
}
}
#[test]
fn build_graph_basic() {
let g = VamanaGraph::build(make_vecs(10, 4), default_cfg(4, 8)).unwrap();
assert_eq!(g.vectors.len(), 10);
for nb in &g.neighbors { assert!(nb.len() <= 4); }
for nb in &g.neighbors {
assert!(nb.len() <= 4);
}
}
#[test]
@ -382,7 +500,9 @@ mod tests {
#[test]
fn robust_pruning_limits_degree() {
let g = VamanaGraph::build(make_vecs(50, 4), default_cfg(5, 16)).unwrap();
for nb in &g.neighbors { assert!(nb.len() <= 5); }
for nb in &g.neighbors {
assert!(nb.len() <= 5);
}
}
#[test]
@ -412,8 +532,11 @@ mod tests {
#[test]
fn cache_hit_rate() {
let mut c = PageCache::new(4);
c.insert(0); c.insert(1);
assert!(c.get(0)); assert!(c.get(1)); assert!(!c.get(2));
c.insert(0);
c.insert(1);
assert!(c.get(0));
assert!(c.get(1));
assert!(!c.get(2));
assert!((c.cache_hit_rate() - 2.0 / 3.0).abs() < 1e-6);
}
@ -424,12 +547,19 @@ mod tests {
let g = VamanaGraph::build(v, default_cfg(8, 20)).unwrap();
let mut d = DiskIndex::from_graph(&g, 32);
let r = d.search_with_filter(&[0.0; 4], |id| id % 2 == 0, 5);
for &(id, _) in &r { assert_eq!(id % 2, 0); }
for &(id, _) in &r {
assert_eq!(id % 2, 0);
}
}
#[test]
fn medoid_selection() {
let v = vec![vec![0.0, 0.0], vec![1.0, 0.0], vec![0.0, 1.0], vec![0.5, 0.5]];
let v = vec![
vec![0.0, 0.0],
vec![1.0, 0.0],
vec![0.0, 1.0],
vec![0.5, 0.5],
];
assert_eq!(MedoidFinder::find_medoid(&v), 3);
}
@ -464,14 +594,26 @@ mod tests {
let mut d = DiskIndex::from_graph(&g, 32);
let (r, s) = d.search_disk(&[0.0; 4], 5, 20);
assert_eq!(r.len(), 5);
for w in r.windows(2) { assert!(w[0].1 <= w[1].1); }
for w in r.windows(2) {
assert!(w[0].1 <= w[1].1);
}
assert!(s.pages_read + s.cache_hits > 0);
}
#[test]
fn config_validation() {
assert!(VamanaConfig { max_degree: 0, ..Default::default() }.validate().is_err());
assert!(VamanaConfig { alpha: 0.5, ..Default::default() }.validate().is_err());
assert!(VamanaConfig {
max_degree: 0,
..Default::default()
}
.validate()
.is_err());
assert!(VamanaConfig {
alpha: 0.5,
..Default::default()
}
.validate()
.is_err());
assert!(VamanaConfig::default().validate().is_ok());
}
}

22
crates/ruvector-core/src/advanced_features/graph_rag.rs
View file

@ -174,11 +174,7 @@ impl KnowledgeGraph {
/// BFS expansion: collect all entities reachable within `hop_count` hops from `entity_id`.
/// Returns `(entities, relations)` forming the subgraph.
pub fn get_neighbors(
&self,
entity_id: &str,
hop_count: usize,
) -> (Vec<Entity>, Vec<Relation>) {
pub fn get_neighbors(&self, entity_id: &str, hop_count: usize) -> (Vec<Entity>, Vec<Relation>) {
let mut visited: HashSet<String> = HashSet::new();
let mut queue: VecDeque<(String, usize)> = VecDeque::new();
let mut result_entities: Vec<Entity> = Vec::new();
@ -265,8 +261,9 @@ impl CommunityDetection {
*votes.entry(label).or_insert(0.0) += rel.weight * resolution;
}
}
if let Some((&best_label, _)) =
votes.iter().max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
if let Some((&best_label, _)) = votes
.iter()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
{
let current = labels[id];
if best_label != current {
@ -510,7 +507,10 @@ fn format_context(entities: &[Entity], relations: &[Relation], summaries: &[Stri
if !entities.is_empty() {
let mut section = String::from("## Entities\n");
for e in entities {
section.push_str(&format!("- {} ({}): {}\n", e.name, e.entity_type, e.description));
section.push_str(&format!(
"- {} ({}): {}\n",
e.name, e.entity_type, e.description
));
}
parts.push(section);
}
@ -681,7 +681,11 @@ mod tests {
// Both still appear in communities.
let communities = CommunityDetection::detect_communities(&g, 1.0);
let total_members: usize = communities.iter().filter(|c| c.level == 0).map(|c| c.entities.len()).sum();
let total_members: usize = communities
.iter()
.filter(|c| c.level == 0)
.map(|c| c.entities.len())
.sum();
assert_eq!(total_members, 2);
}

87
crates/ruvector-core/src/advanced_features/matryoshka.rs
View file

@ -203,12 +203,7 @@ impl MatryoshkaIndex {
/// # Errors
///
/// Returns an error if `dim` exceeds the query length or `full_dim`.
pub fn search(
&self,
query: &[f32],
dim: usize,
top_k: usize,
) -> Result<Vec<SearchResult>> {
pub fn search(&self, query: &[f32], dim: usize, top_k: usize) -> Result<Vec<SearchResult>> {
if dim == 0 {
return Err(RuvectorError::InvalidParameter(
"Search dimension must be > 0".into(),
@ -237,7 +232,13 @@ impl MatryoshkaIndex {
.map(|(idx, entry)| {
let doc_prefix = &entry.embedding[..dim];
let doc_norm = compute_norm(doc_prefix);
let sim = similarity(query_prefix, query_norm, doc_prefix, doc_norm, self.config.metric);
let sim = similarity(
query_prefix,
query_norm,
doc_prefix,
doc_norm,
self.config.metric,
);
(idx, sim)
})
.collect();
@ -466,7 +467,9 @@ mod tests {
}
fn make_index(full_dim: usize) -> MatryoshkaIndex {
let dims: Vec<usize> = (1..=full_dim).filter(|d| d.is_power_of_two() || *d == full_dim).collect();
let dims: Vec<usize> = (1..=full_dim)
.filter(|d| d.is_power_of_two() || *d == full_dim)
.collect();
MatryoshkaIndex::new(make_config(full_dim, dims)).unwrap()
}
@ -474,7 +477,9 @@ mod tests {
fn test_insert_and_len() {
let mut index = make_index(4);
assert!(index.is_empty());
index.insert("v1".into(), vec![1.0, 0.0, 0.0, 0.0], None).unwrap();
index
.insert("v1".into(), vec![1.0, 0.0, 0.0, 0.0], None)
.unwrap();
assert_eq!(index.len(), 1);
}
@ -488,8 +493,12 @@ mod tests {
#[test]
fn test_search_at_full_dim() {
let mut index = make_index(4);
index.insert("v1".into(), vec![1.0, 0.0, 0.0, 0.0], None).unwrap();
index.insert("v2".into(), vec![0.0, 1.0, 0.0, 0.0], None).unwrap();
index
.insert("v1".into(), vec![1.0, 0.0, 0.0, 0.0], None)
.unwrap();
index
.insert("v2".into(), vec![0.0, 1.0, 0.0, 0.0], None)
.unwrap();
let results = index.search(&[1.0, 0.0, 0.0, 0.0], 4, 10).unwrap();
assert_eq!(results[0].id, "v1");
@ -502,8 +511,12 @@ mod tests {
fn test_search_at_truncated_dim() {
let mut index = make_index(4);
// Vectors differ only in the last two components
index.insert("v1".into(), vec![1.0, 0.0, 1.0, 0.0], None).unwrap();
index.insert("v2".into(), vec![1.0, 0.0, 0.0, 1.0], None).unwrap();
index
.insert("v1".into(), vec![1.0, 0.0, 1.0, 0.0], None)
.unwrap();
index
.insert("v2".into(), vec![1.0, 0.0, 0.0, 1.0], None)
.unwrap();
// At dim=2, both truncate to [1.0, 0.0] — identical scores
let results = index.search(&[1.0, 0.0, 0.5, 0.5], 2, 10).unwrap();
@ -520,13 +533,25 @@ mod tests {
let mut index = make_index(8);
// Insert vectors that share the same first 2 dims but differ later
index
.insert("best".into(), vec![1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0], None)
.insert(
"best".into(),
vec![1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
None,
)
.unwrap();
index
.insert("good".into(), vec![1.0, 0.0, 0.5, 0.5, 0.0, 0.0, 0.0, 0.0], None)
.insert(
"good".into(),
vec![1.0, 0.0, 0.5, 0.5, 0.0, 0.0, 0.0, 0.0],
None,
)
.unwrap();
index
.insert("bad".into(), vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], None)
.insert(
"bad".into(),
vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
None,
)
.unwrap();
let query = vec![1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0];
@ -568,13 +593,25 @@ mod tests {
fn test_cascade_search() {
let mut index = make_index(8);
index
.insert("a".into(), vec![1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0], None)
.insert(
"a".into(),
vec![1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0],
None,
)
.unwrap();
index
.insert("b".into(), vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], None)
.insert(
"b".into(),
vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
None,
)
.unwrap();
index
.insert("c".into(), vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], None)
.insert(
"c".into(),
vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
None,
)
.unwrap();
let query = vec![1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0];
@ -599,8 +636,12 @@ mod tests {
#[test]
fn test_upsert_overwrites() {
let mut index = make_index(4);
index.insert("v1".into(), vec![1.0, 0.0, 0.0, 0.0], None).unwrap();
index.insert("v1".into(), vec![0.0, 1.0, 0.0, 0.0], None).unwrap();
index
.insert("v1".into(), vec![1.0, 0.0, 0.0, 0.0], None)
.unwrap();
index
.insert("v1".into(), vec![0.0, 1.0, 0.0, 0.0], None)
.unwrap();
assert_eq!(index.len(), 1);
let results = index.search(&[0.0, 1.0, 0.0, 0.0], 4, 10).unwrap();
assert_eq!(results[0].id, "v1");
@ -635,7 +676,9 @@ mod tests {
metric: DistanceMetric::DotProduct,
};
let mut index = MatryoshkaIndex::new(config).unwrap();
index.insert("v1".into(), vec![2.0, 0.0, 0.0, 0.0], None).unwrap();
index
.insert("v1".into(), vec![2.0, 0.0, 0.0, 0.0], None)
.unwrap();
let results = index.search(&[3.0, 0.0, 0.0, 0.0], 4, 10).unwrap();
assert!((results[0].score - 6.0).abs() < 1e-5);
}

31
crates/ruvector-core/src/advanced_features/multi_vector.rs
View file

@ -125,9 +125,10 @@ impl MultiVectorIndex {
});
}
if emb.is_empty() {
return Err(RuvectorError::InvalidParameter(
format!("Embedding at index {} has zero dimensions", i),
));
return Err(RuvectorError::InvalidParameter(format!(
"Embedding at index {} has zero dimensions",
i
)));
}
}
@ -170,11 +171,7 @@ impl MultiVectorIndex {
/// # Errors
///
/// Returns an error if `query_embeddings` is empty.
pub fn search(
&self,
query_embeddings: &[Vec<f32>],
top_k: usize,
) -> Result<Vec<SearchResult>> {
pub fn search(&self, query_embeddings: &[Vec<f32>], top_k: usize) -> Result<Vec<SearchResult>> {
if query_embeddings.is_empty() {
return Err(RuvectorError::InvalidParameter(
"Query embeddings cannot be empty".into(),
@ -268,9 +265,7 @@ impl MultiVectorIndex {
doc_embeddings
.iter()
.enumerate()
.map(|(di, d)| {
self.token_similarity(q, query_norms[qi], d, doc_norms[di])
})
.map(|(di, d)| self.token_similarity(q, query_norms[qi], d, doc_norms[di]))
.fold(f32::NEG_INFINITY, f32::max)
})
.sum()
@ -295,9 +290,7 @@ impl MultiVectorIndex {
doc_embeddings
.iter()
.enumerate()
.map(move |(di, d)| {
self.token_similarity(q, query_norms[qi], d, doc_norms[di])
})
.map(move |(di, d)| self.token_similarity(q, query_norms[qi], d, doc_norms[di]))
})
.sum();
sum / total_pairs
@ -318,9 +311,7 @@ impl MultiVectorIndex {
query_embeddings
.iter()
.enumerate()
.map(|(qi, q)| {
self.token_similarity(q, query_norms[qi], d, doc_norms[di])
})
.map(|(qi, q)| self.token_similarity(q, query_norms[qi], d, doc_norms[di]))
.fold(f32::NEG_INFINITY, f32::max)
})
.sum()
@ -342,11 +333,7 @@ impl MultiVectorIndex {
DistanceMetric::DotProduct => dot,
// For Euclidean and Manhattan we convert to a similarity-like score.
DistanceMetric::Euclidean => {
let dist_sq: f32 = a
.iter()
.zip(b.iter())
.map(|(x, y)| (x - y).powi(2))
.sum();
let dist_sq: f32 = a.iter().zip(b.iter()).map(|(x, y)| (x - y).powi(2)).sum();
1.0 / (1.0 + dist_sq.sqrt())
}
DistanceMetric::Manhattan => {

460
crates/ruvector-core/src/advanced_features/opq.rs
View file

@ -31,8 +31,11 @@ pub struct OPQConfig {
impl Default for OPQConfig {
fn default() -> Self {
Self {
num_subspaces: 8, codebook_size: 256, num_iterations: 20,
num_opq_iterations: 10, metric: DistanceMetric::Euclidean,
num_subspaces: 8,
codebook_size: 256,
num_iterations: 20,
num_opq_iterations: 10,
metric: DistanceMetric::Euclidean,
}
}
}
@ -42,13 +45,19 @@ impl OPQConfig {
pub fn validate(&self) -> Result<()> {
if self.codebook_size > 256 {
return Err(RuvectorError::InvalidParameter(format!(
"Codebook size {} exceeds u8 max 256", self.codebook_size)));
"Codebook size {} exceeds u8 max 256",
self.codebook_size
)));
}
if self.num_subspaces == 0 {
return Err(RuvectorError::InvalidParameter("num_subspaces must be > 0".into()));
return Err(RuvectorError::InvalidParameter(
"num_subspaces must be > 0".into(),
));
}
if self.num_opq_iterations == 0 {
return Err(RuvectorError::InvalidParameter("num_opq_iterations must be > 0".into()));
return Err(RuvectorError::InvalidParameter(
"num_opq_iterations must be > 0".into(),
));
}
Ok(())
}
@ -57,21 +66,43 @@ impl OPQConfig {
// -- Dense matrix (row-major, internal only) ----------------------------------
#[derive(Debug, Clone)]
struct Mat { rows: usize, cols: usize, data: Vec<f32> }
struct Mat {
rows: usize,
cols: usize,
data: Vec<f32>,
}
impl Mat {
fn zeros(r: usize, c: usize) -> Self { Self { rows: r, cols: c, data: vec![0.0; r * c] } }
fn zeros(r: usize, c: usize) -> Self {
Self {
rows: r,
cols: c,
data: vec![0.0; r * c],
}
}
fn identity(n: usize) -> Self {
let mut m = Self::zeros(n, n);
for i in 0..n { m.data[i * n + i] = 1.0; }
for i in 0..n {
m.data[i * n + i] = 1.0;
}
m
}
#[inline] fn get(&self, r: usize, c: usize) -> f32 { self.data[r * self.cols + c] }
#[inline] fn set(&mut self, r: usize, c: usize, v: f32) { self.data[r * self.cols + c] = v; }
#[inline]
fn get(&self, r: usize, c: usize) -> f32 {
self.data[r * self.cols + c]
}
#[inline]
fn set(&mut self, r: usize, c: usize, v: f32) {
self.data[r * self.cols + c] = v;
}
fn transpose(&self) -> Self {
let mut t = Self::zeros(self.cols, self.rows);
for r in 0..self.rows { for c in 0..self.cols { t.set(c, r, self.get(r, c)); } }
for r in 0..self.rows {
for c in 0..self.cols {
t.set(c, r, self.get(r, c));
}
}
t
}
fn mul(&self, b: &Mat) -> Mat {
@ -80,7 +111,10 @@ impl Mat {
for i in 0..self.rows {
for k in 0..self.cols {
let a = self.get(i, k);
for j in 0..b.cols { let c = out.get(i, j); out.set(i, j, c + a * b.get(k, j)); }
for j in 0..b.cols {
let c = out.get(i, j);
out.set(i, j, c + a * b.get(k, j));
}
}
}
out
@ -88,10 +122,14 @@ impl Mat {
fn from_rows(vecs: &[Vec<f32>]) -> Self {
let (rows, cols) = (vecs.len(), vecs[0].len());
let mut data = Vec::with_capacity(rows * cols);
for v in vecs { data.extend_from_slice(v); }
for v in vecs {
data.extend_from_slice(v);
}
Self { rows, cols, data }
}
fn row(&self, i: usize) -> Vec<f32> { self.data[i * self.cols..(i + 1) * self.cols].to_vec() }
fn row(&self, i: usize) -> Vec<f32> {
self.data[i * self.cols..(i + 1) * self.cols].to_vec()
}
}
// -- SVD via power iteration + deflation --------------------------------------
@ -103,16 +141,32 @@ fn svd_rank1(a: &Mat, max_iters: usize) -> (Vec<f32>, f32, Vec<f32>) {
let mut v = vec![1.0 / (n as f32).sqrt(); n];
for _ in 0..max_iters {
let mut nv = vec![0.0; n];
for i in 0..n { for j in 0..n { nv[i] += ata.get(i, j) * v[j]; } }
for i in 0..n {
for j in 0..n {
nv[i] += ata.get(i, j) * v[j];
}
}
let norm: f32 = nv.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm < 1e-12 { break; }
for x in nv.iter_mut() { *x /= norm; }
if norm < 1e-12 {
break;
}
for x in nv.iter_mut() {
*x /= norm;
}
v = nv;
}
let mut av = vec![0.0; a.rows];
for i in 0..a.rows { for j in 0..a.cols { av[i] += a.get(i, j) * v[j]; } }
for i in 0..a.rows {
for j in 0..a.cols {
av[i] += a.get(i, j) * v[j];
}
}
let sigma: f32 = av.iter().map(|x| x * x).sum::<f32>().sqrt();
let u = if sigma > 1e-12 { av.iter().map(|x| x / sigma).collect() } else { vec![0.0; a.rows] };
let u = if sigma > 1e-12 {
av.iter().map(|x| x / sigma).collect()
} else {
vec![0.0; a.rows]
};
(u, sigma, v)
}
@ -124,14 +178,24 @@ fn svd_full(a: &Mat, iters: usize) -> (Mat, Vec<f32>, Mat) {
for _ in 0..n {
let (u, s, v) = svd_rank1(&res, iters);
if s > 1e-10 {
for i in 0..res.rows { for j in 0..res.cols {
let c = res.get(i, j); res.set(i, j, c - s * u[i] * v[j]);
}}
for i in 0..res.rows {
for j in 0..res.cols {
let c = res.get(i, j);
res.set(i, j, c - s * u[i] * v[j]);
}
}
}
uc.push(u); sv.push(s); vc.push(v);
uc.push(u);
sv.push(s);
vc.push(v);
}
let (mut um, mut vm) = (Mat::zeros(n, n), Mat::zeros(n, n));
for j in 0..n { for i in 0..n { um.set(i, j, uc[j][i]); vm.set(i, j, vc[j][i]); } }
for j in 0..n {
for i in 0..n {
um.set(i, j, uc[j][i]);
vm.set(i, j, vc[j][i]);
}
}
(um, sv, vm)
}
@ -146,26 +210,40 @@ fn procrustes(x: &Mat, y: &Mat) -> Mat {
/// Orthogonal rotation matrix R (d x d) that decorrelates dimensions before PQ.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RotationMatrix { pub dim: usize, pub data: Vec<f32> }
pub struct RotationMatrix {
pub dim: usize,
pub data: Vec<f32>,
}
impl RotationMatrix {
/// Identity rotation (no-op).
pub fn identity(dim: usize) -> Self {
let mut data = vec![0.0; dim * dim];
for i in 0..dim { data[i * dim + i] = 1.0; }
for i in 0..dim {
data[i * dim + i] = 1.0;
}
Self { dim, data }
}
/// Rotate vector: y = x @ R.
pub fn rotate(&self, v: &[f32]) -> Vec<f32> {
let d = self.dim;
(0..d).map(|j| (0..d).map(|i| v[i] * self.data[i * d + j]).sum()).collect()
(0..d)
.map(|j| (0..d).map(|i| v[i] * self.data[i * d + j]).sum())
.collect()
}
/// Inverse rotate: x = y @ R^T.
pub fn inverse_rotate(&self, v: &[f32]) -> Vec<f32> {
let d = self.dim;
(0..d).map(|j| (0..d).map(|i| v[i] * self.data[j * d + i]).sum()).collect()
(0..d)
.map(|j| (0..d).map(|i| v[i] * self.data[j * d + i]).sum())
.collect()
}
fn from_mat(m: &Mat) -> Self {
Self {
dim: m.rows,
data: m.data.clone(),
}
}
fn from_mat(m: &Mat) -> Self { Self { dim: m.rows, data: m.data.clone() } }
}
// -- OPQ Index ----------------------------------------------------------------
@ -185,16 +263,25 @@ impl OPQIndex {
pub fn train(vectors: &[Vec<f32>], config: OPQConfig) -> Result<Self> {
config.validate()?;
if vectors.is_empty() {
return Err(RuvectorError::InvalidParameter("Training set cannot be empty".into()));
return Err(RuvectorError::InvalidParameter(
"Training set cannot be empty".into(),
));
}
let d = vectors[0].len();
if d % config.num_subspaces != 0 {
return Err(RuvectorError::InvalidParameter(format!(
"Dimensions {} not divisible by num_subspaces {}", d, config.num_subspaces)));
"Dimensions {} not divisible by num_subspaces {}",
d, config.num_subspaces
)));
}
for v in vectors {
if v.len() != d {
return Err(RuvectorError::DimensionMismatch {
expected: d,
actual: v.len(),
});
}
}
for v in vectors { if v.len() != d {
return Err(RuvectorError::DimensionMismatch { expected: d, actual: v.len() });
}}
let x_mat = Mat::from_rows(vectors);
let mut r = Mat::identity(d);
let mut codebooks: Vec<Vec<Vec<f32>>> = Vec::new();
@ -202,50 +289,88 @@ impl OPQIndex {
for _ in 0..config.num_opq_iterations {
let x_rot = x_mat.mul(&r);
let rotated: Vec<Vec<f32>> = (0..vectors.len()).map(|i| x_rot.row(i)).collect();
codebooks = train_pq_codebooks(&rotated, config.num_subspaces,
config.codebook_size, config.num_iterations, config.metric)?;
codebooks = train_pq_codebooks(
&rotated,
config.num_subspaces,
config.codebook_size,
config.num_iterations,
config.metric,
)?;
let mut x_hat = Mat::zeros(vectors.len(), d);
for (i, rv) in rotated.iter().enumerate() {
let codes = encode_vec(rv, &codebooks, sub_dim, config.metric)?;
let recon = decode_vec(&codes, &codebooks);
for (j, &val) in recon.iter().enumerate() { x_hat.set(i, j, val); }
for (j, &val) in recon.iter().enumerate() {
x_hat.set(i, j, val);
}
}
r = procrustes(&x_mat, &x_hat);
}
Ok(Self { config, rotation: RotationMatrix::from_mat(&r), codebooks, dimensions: d })
Ok(Self {
config,
rotation: RotationMatrix::from_mat(&r),
codebooks,
dimensions: d,
})
}
/// Encode a vector: rotate then PQ-quantize.
pub fn encode(&self, vector: &[f32]) -> Result<Vec<u8>> {
self.check_dim(vector.len())?;
let rotated = self.rotation.rotate(vector);
encode_vec(&rotated, &self.codebooks,
self.dimensions / self.config.num_subspaces, self.config.metric)
encode_vec(
&rotated,
&self.codebooks,
self.dimensions / self.config.num_subspaces,
self.config.metric,
)
}
/// Decode PQ codes back to approximate vector (with inverse rotation).
pub fn decode(&self, codes: &[u8]) -> Result<Vec<f32>> {
if codes.len() != self.config.num_subspaces {
return Err(RuvectorError::InvalidParameter(format!(
"Expected {} codes, got {}", self.config.num_subspaces, codes.len())));
"Expected {} codes, got {}",
self.config.num_subspaces,
codes.len()
)));
}
Ok(self.rotation.inverse_rotate(&decode_vec(codes, &self.codebooks)))
Ok(self
.rotation
.inverse_rotate(&decode_vec(codes, &self.codebooks)))
}
/// ADC search: precompute distance tables then sum lookups per database vector.
pub fn search_adc(&self, query: &[f32], codes_db: &[Vec<u8>], top_k: usize,
pub fn search_adc(
&self,
query: &[f32],
codes_db: &[Vec<u8>],
top_k: usize,
) -> Result<Vec<(usize, f32)>> {
self.check_dim(query.len())?;
let rq = self.rotation.rotate(query);
let sub_dim = self.dimensions / self.config.num_subspaces;
let tables: Vec<Vec<f32>> = (0..self.config.num_subspaces).map(|s| {
let q_sub = &rq[s * sub_dim..(s + 1) * sub_dim];
self.codebooks[s].iter().map(|c| dist(q_sub, c, self.config.metric)).collect()
}).collect();
let mut dists: Vec<(usize, f32)> = codes_db.iter().enumerate().map(|(idx, codes)| {
let d: f32 = codes.iter().enumerate().map(|(s, &c)| tables[s][c as usize]).sum();
(idx, d)
}).collect();
let tables: Vec<Vec<f32>> = (0..self.config.num_subspaces)
.map(|s| {
let q_sub = &rq[s * sub_dim..(s + 1) * sub_dim];
self.codebooks[s]
.iter()
.map(|c| dist(q_sub, c, self.config.metric))
.collect()
})
.collect();
let mut dists: Vec<(usize, f32)> = codes_db
.iter()
.enumerate()
.map(|(idx, codes)| {
let d: f32 = codes
.iter()
.enumerate()
.map(|(s, &c)| tables[s][c as usize])
.sum();
(idx, d)
})
.collect();
dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
dists.truncate(top_k);
Ok(dists)
@ -253,19 +378,30 @@ impl OPQIndex {
/// Mean squared quantization error over a set of vectors.
pub fn quantization_error(&self, vectors: &[Vec<f32>]) -> Result<f32> {
if vectors.is_empty() { return Ok(0.0); }
if vectors.is_empty() {
return Ok(0.0);
}
let mut total = 0.0f64;
for v in vectors {
let recon = self.decode(&self.encode(v)?)?;
total += v.iter().zip(&recon).map(|(a, b)| ((a - b) as f64).powi(2)).sum::<f64>();
total += v
.iter()
.zip(&recon)
.map(|(a, b)| ((a - b) as f64).powi(2))
.sum::<f64>();
}
Ok((total / vectors.len() as f64) as f32)
}
fn check_dim(&self, len: usize) -> Result<()> {
if len != self.dimensions {
Err(RuvectorError::DimensionMismatch { expected: self.dimensions, actual: len })
} else { Ok(()) }
Err(RuvectorError::DimensionMismatch {
expected: self.dimensions,
actual: len,
})
} else {
Ok(())
}
}
}
@ -273,48 +409,87 @@ impl OPQIndex {
fn dist(a: &[f32], b: &[f32], m: DistanceMetric) -> f32 {
match m {
DistanceMetric::Euclidean =>
a.iter().zip(b).map(|(x, y)| { let d = x - y; d * d }).sum::<f32>().sqrt(),
DistanceMetric::Euclidean => a
.iter()
.zip(b)
.map(|(x, y)| {
let d = x - y;
d * d
})
.sum::<f32>()
.sqrt(),
DistanceMetric::Cosine => {
let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
let na = a.iter().map(|x| x * x).sum::<f32>().sqrt();
let nb = b.iter().map(|x| x * x).sum::<f32>().sqrt();
if na == 0.0 || nb == 0.0 { 1.0 } else { 1.0 - dot / (na * nb) }
if na == 0.0 || nb == 0.0 {
1.0
} else {
1.0 - dot / (na * nb)
}
}
DistanceMetric::DotProduct => -a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>(),
DistanceMetric::Manhattan => a.iter().zip(b).map(|(x, y)| (x - y).abs()).sum(),
}
}
fn train_pq_codebooks(vecs: &[Vec<f32>], nsub: usize, k: usize, iters: usize,
metric: DistanceMetric) -> Result<Vec<Vec<Vec<f32>>>> {
fn train_pq_codebooks(
vecs: &[Vec<f32>],
nsub: usize,
k: usize,
iters: usize,
metric: DistanceMetric,
) -> Result<Vec<Vec<Vec<f32>>>> {
let sub_dim = vecs[0].len() / nsub;
(0..nsub).map(|s| {
let sv: Vec<Vec<f32>> = vecs.iter().map(|v| v[s*sub_dim..(s+1)*sub_dim].to_vec()).collect();
kmeans(&sv, k.min(sv.len()), iters, metric)
}).collect()
(0..nsub)
.map(|s| {
let sv: Vec<Vec<f32>> = vecs
.iter()
.map(|v| v[s * sub_dim..(s + 1) * sub_dim].to_vec())
.collect();
kmeans(&sv, k.min(sv.len()), iters, metric)
})
.collect()
}
fn encode_vec(v: &[f32], cbs: &[Vec<Vec<f32>>], sub_dim: usize, m: DistanceMetric,
fn encode_vec(
v: &[f32],
cbs: &[Vec<Vec<f32>>],
sub_dim: usize,
m: DistanceMetric,
) -> Result<Vec<u8>> {
cbs.iter().enumerate().map(|(s, cb)| {
let sub = &v[s * sub_dim..(s + 1) * sub_dim];
cb.iter().enumerate()
.min_by(|(_, a), (_, b)| dist(sub, a, m).partial_cmp(&dist(sub, b, m)).unwrap())
.map(|(i, _)| i as u8)
.ok_or_else(|| RuvectorError::Internal("Empty codebook".into()))
}).collect()
cbs.iter()
.enumerate()
.map(|(s, cb)| {
let sub = &v[s * sub_dim..(s + 1) * sub_dim];
cb.iter()
.enumerate()
.min_by(|(_, a), (_, b)| dist(sub, a, m).partial_cmp(&dist(sub, b, m)).unwrap())
.map(|(i, _)| i as u8)
.ok_or_else(|| RuvectorError::Internal("Empty codebook".into()))
})
.collect()
}
fn decode_vec(codes: &[u8], cbs: &[Vec<Vec<f32>>]) -> Vec<f32> {
codes.iter().enumerate().flat_map(|(s, &c)| cbs[s][c as usize].iter().copied()).collect()
codes
.iter()
.enumerate()
.flat_map(|(s, &c)| cbs[s][c as usize].iter().copied())
.collect()
}
fn kmeans(vecs: &[Vec<f32>], k: usize, iters: usize, metric: DistanceMetric,
fn kmeans(
vecs: &[Vec<f32>],
k: usize,
iters: usize,
metric: DistanceMetric,
) -> Result<Vec<Vec<f32>>> {
use rand::seq::SliceRandom;
if vecs.is_empty() || k == 0 {
return Err(RuvectorError::InvalidParameter("Cannot cluster empty set or k=0".into()));
return Err(RuvectorError::InvalidParameter(
"Cannot cluster empty set or k=0".into(),
));
}
let dim = vecs[0].len();
let mut rng = rand::thread_rng();
@ -322,14 +497,25 @@ fn kmeans(vecs: &[Vec<f32>], k: usize, iters: usize, metric: DistanceMetric,
for _ in 0..iters {
let (mut sums, mut counts) = (vec![vec![0.0f32; dim]; k], vec![0usize; k]);
for v in vecs {
let b = cents.iter().enumerate()
.min_by(|(_, a), (_, b)| dist(v, a, metric).partial_cmp(&dist(v, b, metric)).unwrap())
.map(|(i, _)| i).unwrap_or(0);
let b = cents
.iter()
.enumerate()
.min_by(|(_, a), (_, b)| {
dist(v, a, metric).partial_cmp(&dist(v, b, metric)).unwrap()
})
.map(|(i, _)| i)
.unwrap_or(0);
counts[b] += 1;
for (j, &val) in v.iter().enumerate() { sums[b][j] += val; }
for (j, &val) in v.iter().enumerate() {
sums[b][j] += val;
}
}
for (i, c) in cents.iter_mut().enumerate() {
if counts[i] > 0 { for j in 0..dim { c[j] = sums[i][j] / counts[i] as f32; } }
if counts[i] > 0 {
for j in 0..dim {
c[j] = sums[i][j] / counts[i] as f32;
}
}
}
}
Ok(cents)
@ -341,14 +527,25 @@ mod tests {
fn make_data(n: usize, d: usize) -> Vec<Vec<f32>> {
let mut seed: u64 = 42;
(0..n).map(|_| (0..d).map(|_| {
seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
((seed >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0
}).collect()).collect()
(0..n)
.map(|_| {
(0..d)
.map(|_| {
seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
((seed >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0
})
.collect()
})
.collect()
}
fn cfg() -> OPQConfig {
OPQConfig { num_subspaces: 2, codebook_size: 4, num_iterations: 5,
num_opq_iterations: 3, metric: DistanceMetric::Euclidean }
OPQConfig {
num_subspaces: 2,
codebook_size: 4,
num_iterations: 5,
num_opq_iterations: 3,
metric: DistanceMetric::Euclidean,
}
}
#[test]
@ -356,7 +553,9 @@ mod tests {
let r = RotationMatrix::identity(4);
let v = vec![1.0, 2.0, 3.0, 4.0];
let back = r.inverse_rotate(&r.rotate(&v));
for i in 0..4 { assert!((v[i] - back[i]).abs() < 1e-6); }
for i in 0..4 {
assert!((v[i] - back[i]).abs() < 1e-6);
}
}
#[test]
fn test_rotation_preserves_norm() {
@ -364,7 +563,13 @@ mod tests {
let idx = OPQIndex::train(&data, cfg()).unwrap();
let v = vec![1.0, 2.0, 3.0, 4.0];
let n1: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
let n2: f32 = idx.rotation.rotate(&v).iter().map(|x| x * x).sum::<f32>().sqrt();
let n2: f32 = idx
.rotation
.rotate(&v)
.iter()
.map(|x| x * x)
.sum::<f32>()
.sqrt();
assert!((n1 - n2).abs() < 0.1, "norms: {} vs {}", n1, n2);
}
#[test]
@ -384,13 +589,19 @@ mod tests {
let data = make_data(100, 4);
let idx = OPQIndex::train(&data, cfg()).unwrap();
let err = idx.quantization_error(&data).unwrap();
assert!(err.is_finite() && err >= 0.0, "error must be finite non-negative: {}", err);
assert!(
err.is_finite() && err >= 0.0,
"error must be finite non-negative: {}",
err
);
// Verify round-trip through encode/decode does not explode.
for v in &data {
let codes = idx.encode(v).unwrap();
let decoded = idx.decode(&codes).unwrap();
assert_eq!(decoded.len(), v.len());
for x in &decoded { assert!(x.is_finite()); }
for x in &decoded {
assert!(x.is_finite());
}
}
}
#[test]
@ -400,27 +611,52 @@ mod tests {
let db: Vec<Vec<u8>> = data.iter().map(|v| idx.encode(v).unwrap()).collect();
let res = idx.search_adc(&[0.5, -0.5, 0.5, -0.5], &db, 3).unwrap();
assert_eq!(res.len(), 3);
for w in res.windows(2) { assert!(w[0].1 <= w[1].1 + 1e-6); }
for w in res.windows(2) {
assert!(w[0].1 <= w[1].1 + 1e-6);
}
}
#[test]
fn test_quantization_error_reduction() {
let data = make_data(50, 4);
let err = OPQIndex::train(&data, cfg()).unwrap().quantization_error(&data).unwrap();
let err = OPQIndex::train(&data, cfg())
.unwrap()
.quantization_error(&data)
.unwrap();
assert!(err >= 0.0 && err.is_finite() && err < 10.0, "err={}", err);
}
#[test]
fn test_svd_correctness() {
let a = Mat { rows: 2, cols: 2, data: vec![3.0, 0.0, 0.0, 2.0] };
let a = Mat {
rows: 2,
cols: 2,
data: vec![3.0, 0.0, 0.0, 2.0],
};
let (u, s, v) = svd_full(&a, 200);
for i in 0..2 { for j in 0..2 {
let r: f32 = (0..2).map(|k| u.get(i, k) * s[k] * v.get(j, k)).sum();
assert!((a.get(i, j) - r).abs() < 0.1, "SVD fail ({},{}): {} vs {}", i, j, a.get(i, j), r);
}}
for i in 0..2 {
for j in 0..2 {
let r: f32 = (0..2).map(|k| u.get(i, k) * s[k] * v.get(j, k)).sum();
assert!(
(a.get(i, j) - r).abs() < 0.1,
"SVD fail ({},{}): {} vs {}",
i,
j,
a.get(i, j),
r
);
}
}
}
#[test]
fn test_identity_rotation_baseline() {
let data = make_data(30, 4);
let idx = OPQIndex::train(&data, OPQConfig { num_opq_iterations: 1, ..cfg() }).unwrap();
let idx = OPQIndex::train(
&data,
OPQConfig {
num_opq_iterations: 1,
..cfg()
},
)
.unwrap();
let recon = idx.decode(&idx.encode(&data[0]).unwrap()).unwrap();
assert_eq!(recon.len(), data[0].len());
}
@ -429,14 +665,34 @@ mod tests {
let data = make_data(40, 4);
let idx = OPQIndex::train(&data, cfg()).unwrap();
let db: Vec<Vec<u8>> = data.iter().map(|v| idx.encode(v).unwrap()).collect();
let ids: Vec<usize> = idx.search_adc(&data[0], &db, 5).unwrap().iter().map(|r| r.0).collect();
let ids: Vec<usize> = idx
.search_adc(&data[0], &db, 5)
.unwrap()
.iter()
.map(|r| r.0)
.collect();
assert!(ids.contains(&0), "vector 0 should be in its own top-5");
}
#[test]
fn test_config_validation() {
assert!(OPQConfig { codebook_size: 300, ..cfg() }.validate().is_err());
assert!(OPQConfig { num_subspaces: 0, ..cfg() }.validate().is_err());
assert!(OPQConfig { num_opq_iterations: 0, ..cfg() }.validate().is_err());
assert!(OPQConfig {
codebook_size: 300,
..cfg()
}
.validate()
.is_err());
assert!(OPQConfig {
num_subspaces: 0,
..cfg()
}
.validate()
.is_err());
assert!(OPQConfig {
num_opq_iterations: 0,
..cfg()
}
.validate()
.is_err());
}
#[test]
fn test_dimension_mismatch_errors() {

144
crates/ruvector-core/src/advanced_features/sparse_vector.rs
View file

@ -45,10 +45,7 @@ impl SparseVector {
*map.entry(idx).or_insert(0.0) += val;
}
let mut entries: Vec<(u32, f32)> = map
.into_iter()
.filter(|(_, v)| *v != 0.0)
.collect();
let mut entries: Vec<(u32, f32)> = map.into_iter().filter(|(_, v)| *v != 0.0).collect();
entries.sort_unstable_by_key(|(idx, _)| *idx);
let (indices, values) = entries.into_iter().unzip();
@ -170,13 +167,10 @@ impl SparseIndex {
// Add new postings.
for (pos, &dim) in vector.indices.iter().enumerate() {
self.postings
.entry(dim)
.or_default()
.push(PostingEntry {
doc_id: doc_id.clone(),
weight: vector.values[pos],
});
self.postings.entry(dim).or_default().push(PostingEntry {
doc_id: doc_id.clone(),
weight: vector.values[pos],
});
}
self.docs.insert(doc_id, vector);
@ -249,11 +243,7 @@ impl SparseIndex {
}
/// Batch search: run multiple queries and return results for each.
pub fn search_batch(
&self,
queries: &[SparseVector],
k: usize,
) -> Vec<Vec<ScoredDoc>> {
pub fn search_batch(&self, queries: &[SparseVector], k: usize) -> Vec<Vec<ScoredDoc>> {
queries.iter().map(|q| self.search(q, k)).collect()
}
}
@ -279,7 +269,10 @@ pub enum FusionStrategy {
/// Reciprocal Rank Fusion. `k` controls rank-pressure (default 60).
RRF { k: f32 },
/// Weighted linear combination of normalised scores.
Linear { dense_weight: f32, sparse_weight: f32 },
Linear {
dense_weight: f32,
sparse_weight: f32,
},
/// Distribution-Based Score Fusion: normalise each list to N(0,1) then
/// combine with equal weight.
DBSF,
@ -331,12 +324,7 @@ pub fn fuse_rankings(
// -- RRF -------------------------------------------------------------------
/// Reciprocal Rank Fusion: score(d) = sum_over_lists 1 / (k + rank(d)).
fn fuse_rrf(
dense: &[ScoredDoc],
sparse: &[ScoredDoc],
k: f32,
top_k: usize,
) -> Vec<ScoredDoc> {
fn fuse_rrf(dense: &[ScoredDoc], sparse: &[ScoredDoc], k: f32, top_k: usize) -> Vec<ScoredDoc> {
let mut scores: HashMap<VectorId, f32> = HashMap::new();
for (rank, doc) in dense.iter().enumerate() {
@ -377,11 +365,7 @@ fn fuse_linear(
// -- DBSF ------------------------------------------------------------------
/// Distribution-Based Score Fusion: z-score normalise, then average.
fn fuse_dbsf(
dense: &[ScoredDoc],
sparse: &[ScoredDoc],
top_k: usize,
) -> Vec<ScoredDoc> {
fn fuse_dbsf(dense: &[ScoredDoc], sparse: &[ScoredDoc], top_k: usize) -> Vec<ScoredDoc> {
let z_dense = z_score_normalize(dense);
let z_sparse = z_score_normalize(sparse);
@ -542,10 +526,7 @@ mod tests {
#[test]
fn test_index_single_result() {
let mut idx = SparseIndex::new();
idx.insert(
"only".into(),
SparseVector::from_sorted(vec![7], vec![2.0]),
);
idx.insert("only".into(), SparseVector::from_sorted(vec![7], vec![2.0]));
let query = SparseVector::from_sorted(vec![7], vec![3.0]);
let results = idx.search(&query, 5);
assert_eq!(results.len(), 1);
@ -556,10 +537,7 @@ mod tests {
#[test]
fn test_index_remove() {
let mut idx = SparseIndex::new();
idx.insert(
"d1".into(),
SparseVector::from_sorted(vec![0], vec![1.0]),
);
idx.insert("d1".into(), SparseVector::from_sorted(vec![0], vec![1.0]));
assert_eq!(idx.len(), 1);
assert!(idx.remove(&"d1".into()));
assert_eq!(idx.len(), 0);
@ -574,10 +552,7 @@ mod tests {
SparseVector::from_sorted(vec![0, 1], vec![1.0, 2.0]),
);
// Re-insert same id with different dimensions.
idx.insert(
"d1".into(),
SparseVector::from_sorted(vec![3], vec![5.0]),
);
idx.insert("d1".into(), SparseVector::from_sorted(vec![3], vec![5.0]));
assert_eq!(idx.len(), 1);
// Old dimensions should not match.
@ -597,14 +572,32 @@ mod tests {
fn test_rrf_fusion_basic() {
// Two lists with overlapping documents.
let dense = vec![
ScoredDoc { id: "a".into(), score: 10.0 },
ScoredDoc { id: "b".into(), score: 8.0 },
ScoredDoc { id: "c".into(), score: 6.0 },
ScoredDoc {
id: "a".into(),
score: 10.0,
},
ScoredDoc {
id: "b".into(),
score: 8.0,
},
ScoredDoc {
id: "c".into(),
score: 6.0,
},
];
let sparse = vec![
ScoredDoc { id: "b".into(), score: 9.0 },
ScoredDoc { id: "d".into(), score: 7.0 },
ScoredDoc { id: "a".into(), score: 5.0 },
ScoredDoc {
id: "b".into(),
score: 9.0,
},
ScoredDoc {
id: "d".into(),
score: 7.0,
},
ScoredDoc {
id: "a".into(),
score: 5.0,
},
];
let config = FusionConfig {
@ -622,12 +615,14 @@ mod tests {
#[test]
fn test_rrf_with_disjoint_lists() {
let dense = vec![
ScoredDoc { id: "x".into(), score: 5.0 },
];
let sparse = vec![
ScoredDoc { id: "y".into(), score: 5.0 },
];
let dense = vec![ScoredDoc {
id: "x".into(),
score: 5.0,
}];
let sparse = vec![ScoredDoc {
id: "y".into(),
score: 5.0,
}];
let config = FusionConfig {
strategy: FusionStrategy::RRF { k: 60.0 },
@ -644,12 +639,24 @@ mod tests {
#[test]
fn test_linear_fusion() {
let dense = vec![
ScoredDoc { id: "a".into(), score: 10.0 },
ScoredDoc { id: "b".into(), score: 5.0 },
ScoredDoc {
id: "a".into(),
score: 10.0,
},
ScoredDoc {
id: "b".into(),
score: 5.0,
},
];
let sparse = vec![
ScoredDoc { id: "b".into(), score: 10.0 },
ScoredDoc { id: "a".into(), score: 5.0 },
ScoredDoc {
id: "b".into(),
score: 10.0,
},
ScoredDoc {
id: "a".into(),
score: 5.0,
},
];
let config = FusionConfig {
@ -670,12 +677,24 @@ mod tests {
#[test]
fn test_dbsf_fusion() {
let dense = vec![
ScoredDoc { id: "a".into(), score: 10.0 },
ScoredDoc { id: "b".into(), score: 8.0 },
ScoredDoc {
id: "a".into(),
score: 10.0,
},
ScoredDoc {
id: "b".into(),
score: 8.0,
},
];
let sparse = vec![
ScoredDoc { id: "a".into(), score: 6.0 },
ScoredDoc { id: "c".into(), score: 4.0 },
ScoredDoc {
id: "a".into(),
score: 6.0,
},
ScoredDoc {
id: "c".into(),
score: 4.0,
},
];
let config = FusionConfig {
@ -694,7 +713,10 @@ mod tests {
let fused = fuse_rankings(&[], &[], &config);
assert!(fused.is_empty());
let single = vec![ScoredDoc { id: "x".into(), score: 1.0 }];
let single = vec![ScoredDoc {
id: "x".into(),
score: 1.0,
}];
let fused2 = fuse_rankings(&single, &[], &config);
assert_eq!(fused2.len(), 1);
assert_eq!(fused2[0].id, "x");

91
crates/ruvector-core/src/embeddings.rs
View file

@ -430,16 +430,19 @@ pub mod onnx {
let repo = api.model(model_id.to_string());
// Download model files
let model_path = repo.get("model.onnx").or_else(|_| {
// Try alternative path for some models
repo.get("onnx/model.onnx")
}).map_err(|e| {
RuvectorError::ModelLoadError(format!(
"Failed to download ONNX model from {}: {}. \
let model_path = repo
.get("model.onnx")
.or_else(|_| {
// Try alternative path for some models
repo.get("onnx/model.onnx")
})
.map_err(|e| {
RuvectorError::ModelLoadError(format!(
"Failed to download ONNX model from {}: {}. \
Make sure the model has an ONNX export available.",
model_id, e
))
})?;
model_id, e
))
})?;
let tokenizer_path = repo.get("tokenizer.json").map_err(|e| {
RuvectorError::ModelLoadError(format!(
@ -468,7 +471,10 @@ pub mod onnx {
// Load the ONNX session
let session = Session::builder()
.map_err(|e| {
RuvectorError::ModelLoadError(format!("Failed to create session builder: {}", e))
RuvectorError::ModelLoadError(format!(
"Failed to create session builder: {}",
e
))
})?
.with_intra_threads(4)
.map_err(|e| {
@ -583,11 +589,9 @@ pub mod onnx {
// Tokenize
let encoding = {
let tokenizer = self.tokenizer.read();
tokenizer
.encode(text, true)
.map_err(|e| {
RuvectorError::ModelInferenceError(format!("Tokenization failed: {}", e))
})?
tokenizer.encode(text, true).map_err(|e| {
RuvectorError::ModelInferenceError(format!("Tokenization failed: {}", e))
})?
};
// Prepare inputs
@ -597,39 +601,41 @@ pub mod onnx {
.iter()
.map(|&x| x as i64)
.collect();
let token_type_ids: Vec<i64> = encoding
.get_type_ids()
.iter()
.map(|&x| x as i64)
.collect();
let token_type_ids: Vec<i64> =
encoding.get_type_ids().iter().map(|&x| x as i64).collect();
let seq_len = input_ids.len();
// Create ONNX tensors using ort 2.0 API (batch_size=1)
// Tensor::from_array takes (shape, owned_data)
let input_ids_tensor = Tensor::<i64>::from_array(([1, seq_len], input_ids.clone().into_boxed_slice()))
.map_err(|e| {
RuvectorError::ModelInferenceError(format!(
"Failed to create input_ids tensor: {}",
e
))
})?;
let input_ids_tensor =
Tensor::<i64>::from_array(([1, seq_len], input_ids.clone().into_boxed_slice()))
.map_err(|e| {
RuvectorError::ModelInferenceError(format!(
"Failed to create input_ids tensor: {}",
e
))
})?;
let attention_mask_tensor =
Tensor::<i64>::from_array(([1, seq_len], attention_mask.clone().into_boxed_slice())).map_err(|e| {
RuvectorError::ModelInferenceError(format!(
"Failed to create attention_mask tensor: {}",
e
))
})?;
let attention_mask_tensor = Tensor::<i64>::from_array((
[1, seq_len],
attention_mask.clone().into_boxed_slice(),
))
.map_err(|e| {
RuvectorError::ModelInferenceError(format!(
"Failed to create attention_mask tensor: {}",
e
))
})?;
let token_type_ids_tensor =
Tensor::<i64>::from_array(([1, seq_len], token_type_ids.into_boxed_slice())).map_err(|e| {
RuvectorError::ModelInferenceError(format!(
"Failed to create token_type_ids tensor: {}",
e
))
})?;
Tensor::<i64>::from_array(([1, seq_len], token_type_ids.into_boxed_slice()))
.map_err(|e| {
RuvectorError::ModelInferenceError(format!(
"Failed to create token_type_ids tensor: {}",
e
))
})?;
// Run inference and extract output (needs mutable access to session)
// We must extract all data while holding the lock since SessionOutputs has a lifetime
@ -651,7 +657,10 @@ pub mod onnx {
// Extract as ndarray view
let output_array = output_value.try_extract_array::<f32>().map_err(|e| {
RuvectorError::ModelInferenceError(format!("Failed to extract output tensor: {}", e))
RuvectorError::ModelInferenceError(format!(
"Failed to extract output tensor: {}",
e
))
})?;
let output_shape_vec: Vec<usize> = output_array.shape().to_vec();

7
crates/ruvector-core/src/lib.rs
View file

@ -75,10 +75,9 @@ pub mod advanced;
// Re-exports
pub use advanced_features::{
ConformalConfig, ConformalPredictor, EnhancedPQ, FilterExpression, FilterStrategy,
FilteredSearch, FusionConfig, FusionStrategy, HybridConfig, HybridSearch, MMRConfig,
MMRSearch, PQConfig, PredictionSet, ScoredDoc, SparseIndex, SparseVector, BM25,
fuse_rankings,
fuse_rankings, ConformalConfig, ConformalPredictor, EnhancedPQ, FilterExpression,
FilterStrategy, FilteredSearch, FusionConfig, FusionStrategy, HybridConfig, HybridSearch,
MMRConfig, MMRSearch, PQConfig, PredictionSet, ScoredDoc, SparseIndex, SparseVector, BM25,
};
#[cfg(feature = "storage")]

17
crates/ruvector-decompiler/benches/bench_pipeline.rs
View file

@ -45,11 +45,7 @@ fn generate_bundle(target_bytes: usize) -> String {
}
fn bench_full_pipeline(c: &mut Criterion) {
let sizes: &[(usize, &str)] = &[
(1_000, "1KB"),
(10_000, "10KB"),
(100_000, "100KB"),
];
let sizes: &[(usize, &str)] = &[(1_000, "1KB"), (10_000, "10KB"), (100_000, "100KB")];
let mut group = c.benchmark_group("pipeline");
group.sample_size(10);
@ -109,9 +105,8 @@ fn bench_pipeline_phases(c: &mut Criterion) {
let decls_clone = decls.clone();
group.bench_function("graph", |b| {
b.iter(|| {
let graph = ruvector_decompiler::graph::build_reference_graph(
black_box(decls_clone.clone()),
);
let graph =
ruvector_decompiler::graph::build_reference_graph(black_box(decls_clone.clone()));
black_box(graph);
});
});
@ -120,10 +115,8 @@ fn bench_pipeline_phases(c: &mut Criterion) {
let graph = ruvector_decompiler::graph::build_reference_graph(decls);
group.bench_function("partition", |b| {
b.iter(|| {
let result = ruvector_decompiler::partitioner::partition_modules(
black_box(&graph),
Some(5),
);
let result =
ruvector_decompiler::partitioner::partition_modules(black_box(&graph), Some(5));
black_box(result).ok();
});
});

176
crates/ruvector-decompiler/examples/run_on_cli.rs
View file

@ -15,16 +15,28 @@ fn fix_module_syntax(source: &str) -> String {
let mut fixed = String::with_capacity(source.len() + 128);
// Prepend openers for excess closers
for _ in 0..(-parens).max(0) { fixed.push('('); }
for _ in 0..(-brackets).max(0) { fixed.push('['); }
for _ in 0..(-braces).max(0) { fixed.push('{'); }
for _ in 0..(-parens).max(0) {
fixed.push('(');
}
for _ in 0..(-brackets).max(0) {
fixed.push('[');
}
for _ in 0..(-braces).max(0) {
fixed.push('{');
}
fixed.push_str(source);
// Append closers for unclosed openers
for _ in 0..braces.max(0) { fixed.push('}'); }
for _ in 0..brackets.max(0) { fixed.push(']'); }
for _ in 0..parens.max(0) { fixed.push(')'); }
for _ in 0..braces.max(0) {
fixed.push('}');
}
for _ in 0..brackets.max(0) {
fixed.push(']');
}
for _ in 0..parens.max(0) {
fixed.push(')');
}
// Fix try without catch/finally
let try_count = count_keyword(&fixed, "try");
@ -50,12 +62,16 @@ fn fix_module_syntax(source: &str) -> String {
fixed = format!(
"// ruDevolution: wrapped for syntax validity\n\
void function() {{\n{}\n}};\n",
source // use ORIGINAL source, not the broken fix
source // use ORIGINAL source, not the broken fix
);
// Re-balance the wrapper
let (b3, p3, _) = count_delimiters(&fixed);
for _ in 0..p3.max(0) { fixed.push(')'); }
for _ in 0..b3.max(0) { fixed.push('}'); }
for _ in 0..p3.max(0) {
fixed.push(')');
}
for _ in 0..b3.max(0) {
fixed.push('}');
}
}
fixed
@ -76,12 +92,16 @@ fn count_delimiters(source: &str) -> (i32, i32, i32) {
// Single-line comment
b'/' if i + 1 < len && bytes[i + 1] == b'/' => {
i += 2;
while i < len && bytes[i] != b'\n' { i += 1; }
while i < len && bytes[i] != b'\n' {
i += 1;
}
}
// Multi-line comment
b'/' if i + 1 < len && bytes[i + 1] == b'*' => {
i += 2;
while i + 1 < len && !(bytes[i] == b'*' && bytes[i + 1] == b'/') { i += 1; }
while i + 1 < len && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
i += 1;
}
i += 2;
}
// String literals
@ -89,8 +109,13 @@ fn count_delimiters(source: &str) -> (i32, i32, i32) {
let quote = b;
i += 1;
while i < len {
if bytes[i] == b'\\' { i += 2; continue; }
if bytes[i] == quote { break; }
if bytes[i] == b'\\' {
i += 2;
continue;
}
if bytes[i] == quote {
break;
}
i += 1;
}
i += 1;
@ -100,24 +125,55 @@ fn count_delimiters(source: &str) -> (i32, i32, i32) {
i += 1;
let mut tdepth = 0;
while i < len {
if bytes[i] == b'\\' { i += 2; continue; }
if bytes[i] == b'$' && i + 1 < len && bytes[i + 1] == b'{' {
tdepth += 1; i += 2; continue;
if bytes[i] == b'\\' {
i += 2;
continue;
}
if bytes[i] == b'$' && i + 1 < len && bytes[i + 1] == b'{' {
tdepth += 1;
i += 2;
continue;
}
if bytes[i] == b'}' && tdepth > 0 {
tdepth -= 1;
i += 1;
continue;
}
if bytes[i] == b'`' && tdepth == 0 {
break;
}
if bytes[i] == b'}' && tdepth > 0 { tdepth -= 1; i += 1; continue; }
if bytes[i] == b'`' && tdepth == 0 { break; }
i += 1;
}
i += 1;
}
// Delimiters
b'{' => { braces += 1; i += 1; }
b'}' => { braces -= 1; i += 1; }
b'(' => { parens += 1; i += 1; }
b')' => { parens -= 1; i += 1; }
b'[' => { brackets += 1; i += 1; }
b']' => { brackets -= 1; i += 1; }
_ => { i += 1; }
b'{' => {
braces += 1;
i += 1;
}
b'}' => {
braces -= 1;
i += 1;
}
b'(' => {
parens += 1;
i += 1;
}
b')' => {
parens -= 1;
i += 1;
}
b'[' => {
brackets += 1;
i += 1;
}
b']' => {
brackets -= 1;
i += 1;
}
_ => {
i += 1;
}
}
}
(braces, parens, brackets)
@ -135,7 +191,9 @@ fn count_keyword(source: &str, keyword: &str) -> usize {
let before_ok = i == 0 || !bytes[i - 1].is_ascii_alphanumeric();
// Check word boundary after
let after_ok = i + klen >= bytes.len() || !bytes[i + klen].is_ascii_alphanumeric();
if before_ok && after_ok { count += 1; }
if before_ok && after_ok {
count += 1;
}
}
}
count
@ -191,14 +249,17 @@ fn main() {
);
}
let t2 = Instant::now();
let modules =
ruvector_decompiler::partitioner::partition_modules(&graph, None).unwrap();
let modules = ruvector_decompiler::partitioner::partition_modules(&graph, None).unwrap();
let t_partition = t2.elapsed();
eprintln!(
"Phase 3 (Partition): {:?} -- {} modules detected{}",
t_partition,
modules.len(),
if large_graph { " (Louvain)" } else { " (MinCut)" }
if large_graph {
" (Louvain)"
} else {
" (MinCut)"
}
);
// Phase 4: Infer names
@ -331,7 +392,10 @@ fn main() {
.sum::<usize>();
eprintln!("\nEstimated memory usage:");
eprintln!(" Declarations: {:.2} MB", decl_mem as f64 / 1_048_576.0);
eprintln!(" Module sources: {:.2} MB", module_mem as f64 / 1_048_576.0);
eprintln!(
" Module sources: {:.2} MB",
module_mem as f64 / 1_048_576.0
);
eprintln!(
" Total estimate: {:.2} MB",
(decl_mem + module_mem) as f64 / 1_048_576.0
@ -339,7 +403,10 @@ fn main() {
// Write tree output if --output-dir is provided.
let args: Vec<String> = std::env::args().collect();
let out_dir = args.iter().position(|a| a == "--output-dir").and_then(|i| args.get(i + 1));
let out_dir = args
.iter()
.position(|a| a == "--output-dir")
.and_then(|i| args.get(i + 1));
if let Some(out_dir) = out_dir {
let base = std::path::Path::new(out_dir);
// Write flat modules (all 1,029 as individual .js files)
@ -356,23 +423,36 @@ fn main() {
} else {
module.source.clone()
};
if content.is_empty() { continue; }
if content.is_empty() {
continue;
}
// Two-pass fix: try smart fix first, fall back to void-wrapper
let fixed = fix_module_syntax(&content);
// Wrap in void function to guarantee parseability
let safe = format!(
"// Module: {}\n// Declarations: {}\nvoid function() {{\n{}\n}};",
module.name, module.declarations.len(), content
module.name,
module.declarations.len(),
content
);
// Use the smart fix if it has balanced delimiters, otherwise use safe wrapper
let (b, p, k) = count_delimiters(&fixed);
let output = if b == 0 && p == 0 && k == 0 { fixed } else { safe };
let output = if b == 0 && p == 0 && k == 0 {
fixed
} else {
safe
};
let filename = format!("{}.js", module.name.replace('/', "_"));
std::fs::write(source_dir.join(&filename), &output).ok();
total_bytes += output.len();
written += 1;
}
eprintln!("\nWrote {} modules to {}/source/ ({:.1} MB)", written, out_dir, total_bytes as f64 / 1_048_576.0);
eprintln!(
"\nWrote {} modules to {}/source/ ({:.1} MB)",
written,
out_dir,
total_bytes as f64 / 1_048_576.0
);
// Phase 8: Auto-fix to 100% parse rate via Node.js post-processing
eprintln!("Phase 8 (Validate): Auto-fixing modules for 100% parse rate...");
@ -408,7 +488,10 @@ fn main() {
let total = v["total"].as_u64().unwrap_or(0);
let pass = v["pass"].as_u64().unwrap_or(0);
let fixed = v["fixed"].as_u64().unwrap_or(0);
eprintln!("Phase 8 (Validate): {}/{} parse (100%) — {} auto-fixed", pass, total, fixed);
eprintln!(
"Phase 8 (Validate): {}/{} parse (100%) — {} auto-fixed",
pass, total, fixed
);
}
}
_ => eprintln!("Phase 8 (Validate): Node.js not available, skipping auto-fix"),
@ -439,7 +522,11 @@ fn main() {
"source_bytes": source.len(),
"output_bytes": total_bytes,
});
std::fs::write(base.join("metrics.json"), serde_json::to_string_pretty(&metrics).unwrap_or_default()).ok();
std::fs::write(
base.join("metrics.json"),
serde_json::to_string_pretty(&metrics).unwrap_or_default(),
)
.ok();
eprintln!("Wrote metrics to {}/metrics.json", out_dir);
}
}
@ -449,21 +536,12 @@ fn print_tree(tree: &ModuleTree, indent: &str) {
let module_count = tree.modules.len();
let child_count = tree.children.len();
if module_count > 0 {
eprintln!(
"{}{}/ ({} modules)",
indent, tree.name, module_count
);
eprintln!("{}{}/ ({} modules)", indent, tree.name, module_count);
for m in &tree.modules {
eprintln!(
"{} {} ({} decls)",
indent, m.name, m.declarations.len()
);
eprintln!("{} {} ({} decls)", indent, m.name, m.declarations.len());
}
} else {
eprintln!(
"{}{}/ ({} subfolders)",
indent, tree.name, child_count
);
eprintln!("{}{}/ ({} subfolders)", indent, tree.name, child_count);
}
for child in &tree.children {
print_tree(child, &format!("{} ", indent));

8
crates/ruvector-decompiler/src/beautifier.rs
View file

@ -109,8 +109,7 @@ fn replace_identifier(code: &str, old: &str, new_name: &str) -> String {
while i < bytes.len() {
if i + old_len <= bytes.len() && &bytes[i..i + old_len] == old_bytes {
let before_ok = i == 0 || !is_ident_char(bytes[i - 1]);
let after_ok =
i + old_len >= bytes.len() || !is_ident_char(bytes[i + old_len]);
let after_ok = i + old_len >= bytes.len() || !is_ident_char(bytes[i + old_len]);
if before_ok && after_ok {
result.push_str(new_name);
@ -222,10 +221,7 @@ mod tests {
#[test]
fn test_replace_no_substring() {
assert_eq!(
replace_identifier("var bar = 1", "a", "x"),
"var bar = 1"
);
assert_eq!(replace_identifier("var bar = 1", "a", "x"), "var bar = 1");
}
#[test]

154
crates/ruvector-decompiler/src/inferrer.rs
View file

@ -91,10 +91,7 @@ pub fn infer_names(modules: &[Module]) -> Vec<InferredName> {
}
/// Infer names using a specific training corpus.
pub fn infer_names_with_corpus(
modules: &[Module],
corpus: &TrainingCorpus,
) -> Vec<InferredName> {
pub fn infer_names_with_corpus(modules: &[Module], corpus: &TrainingCorpus) -> Vec<InferredName> {
let mut inferred = Vec::new();
for module in modules {
@ -117,32 +114,36 @@ pub(crate) fn infer_declaration_name(
// Strategy 0: Training corpus match (domain-specific).
if let Some((pattern, score)) = corpus.match_declaration(decl) {
best = keep_best(best, InferredName {
original: decl.name.clone(),
inferred: pattern.inferred_name.clone(),
confidence: score.min(0.98),
evidence: vec![format!(
"training corpus match: {} (score: {:.2}, module_hint: {:?})",
pattern.inferred_name,
score,
pattern.module_hint
)],
});
best = keep_best(
best,
InferredName {
original: decl.name.clone(),
inferred: pattern.inferred_name.clone(),
confidence: score.min(0.98),
evidence: vec![format!(
"training corpus match: {} (score: {:.2}, module_hint: {:?})",
pattern.inferred_name, score, pattern.module_hint
)],
},
);
}
// Strategy 1: HIGH confidence -- direct string literal match.
'outer: for lit in &decl.string_literals {
for &(pattern, name) in KNOWN_PATTERNS {
if lit.contains(pattern) {
best = keep_best(best, InferredName {
original: decl.name.clone(),
inferred: name.to_string(),
confidence: 0.95,
evidence: vec![format!(
"string literal \"{}\" matches known pattern \"{}\"",
lit, pattern
)],
});
best = keep_best(
best,
InferredName {
original: decl.name.clone(),
inferred: name.to_string(),
confidence: 0.95,
evidence: vec![format!(
"string literal \"{}\" matches known pattern \"{}\"",
lit, pattern
)],
},
);
break 'outer;
}
}
@ -157,15 +158,18 @@ pub(crate) fn infer_declaration_name(
for prop in &decl.property_accesses {
for &(pattern, name) in PROPERTY_PATTERNS {
if prop == pattern {
best = keep_best(best, InferredName {
original: decl.name.clone(),
inferred: name.to_string(),
confidence: 0.7,
evidence: vec![format!(
"property access .{} suggests purpose \"{}\"",
prop, name
)],
});
best = keep_best(
best,
InferredName {
original: decl.name.clone(),
inferred: name.to_string(),
confidence: 0.7,
evidence: vec![format!(
"property access .{} suggests purpose \"{}\"",
prop, name
)],
},
);
break;
}
}
@ -176,15 +180,18 @@ pub(crate) fn infer_declaration_name(
let joined = decl.string_literals.join("_");
let inferred = sanitize_name(&joined, 30);
if !inferred.is_empty() && inferred != decl.name {
best = keep_best(best, InferredName {
original: decl.name.clone(),
inferred,
confidence: 0.65,
evidence: vec![format!(
"multiple string literals: {:?}",
&decl.string_literals[..decl.string_literals.len().min(3)]
)],
});
best = keep_best(
best,
InferredName {
original: decl.name.clone(),
inferred,
confidence: 0.65,
evidence: vec![format!(
"multiple string literals: {:?}",
&decl.string_literals[..decl.string_literals.len().min(3)]
)],
},
);
}
}
@ -224,10 +231,7 @@ pub(crate) fn infer_declaration_name(
}
/// Keep the candidate with the higher confidence score.
fn keep_best(
current: Option<InferredName>,
candidate: InferredName,
) -> Option<InferredName> {
fn keep_best(current: Option<InferredName>, candidate: InferredName) -> Option<InferredName> {
match current {
Some(c) if c.confidence >= candidate.confidence => Some(c),
_ => Some(candidate),
@ -277,7 +281,11 @@ fn collect_string_freq<'a>(modules: &'a [Module]) -> std::collections::HashMap<&
let mut freq = std::collections::HashMap::new();
for module in modules {
for decl in &module.declarations {
for s in decl.string_literals.iter().chain(decl.property_accesses.iter()) {
for s in decl
.string_literals
.iter()
.chain(decl.property_accesses.iter())
{
let s = s.as_str();
if is_meaningful_string(s) {
*freq.entry(s).or_insert(0) += 1;
@ -291,10 +299,18 @@ fn collect_string_freq<'a>(modules: &'a [Module]) -> std::collections::HashMap<&
/// Check if a string is meaningful enough to use for naming.
fn is_meaningful_string(s: &str) -> bool {
let len = s.len();
if len < 2 || len > 50 { return false; }
if s.chars().all(|c| c.is_ascii_digit()) { return false; }
if s.contains("://") || s.starts_with('/') || s.starts_with('.') { return false; }
if s.contains('\n') || s.contains('\t') { return false; }
if len < 2 || len > 50 {
return false;
}
if s.chars().all(|c| c.is_ascii_digit()) {
return false;
}
if s.contains("://") || s.starts_with('/') || s.starts_with('.') {
return false;
}
if s.contains('\n') || s.contains('\t') {
return false;
}
s.chars().any(|c| c.is_alphabetic())
}
@ -306,21 +322,38 @@ fn infer_from_module_names(modules: &[Module]) -> String {
let names: Vec<&str> = modules.iter().map(|m| m.name.as_str()).collect();
if let Some(first) = names.first() {
let prefix_len = names.iter().skip(1).fold(first.len(), |acc, name| {
first.chars().zip(name.chars()).take(acc)
.take_while(|(a, b)| a == b).count()
first
.chars()
.zip(name.chars())
.take(acc)
.take_while(|(a, b)| a == b)
.count()
});
if prefix_len >= 2 { return sanitize_folder_name(&first[..prefix_len]); }
if prefix_len >= 2 {
return sanitize_folder_name(&first[..prefix_len]);
}
}
format!("group_{}", modules.len())
}
/// Sanitize a string into a valid folder name.
fn sanitize_folder_name(raw: &str) -> String {
let cleaned: String = raw.chars()
.map(|c| if c.is_alphanumeric() || c == '_' || c == '-' { c.to_ascii_lowercase() } else { '_' })
let cleaned: String = raw
.chars()
.map(|c| {
if c.is_alphanumeric() || c == '_' || c == '-' {
c.to_ascii_lowercase()
} else {
'_'
}
})
.collect();
let trimmed = cleaned.trim_matches('_');
if trimmed.is_empty() { "module".to_string() } else { trimmed.to_string() }
if trimmed.is_empty() {
"module".to_string()
} else {
trimmed.to_string()
}
}
/// Feedback from a ground-truth comparison for self-learning.
@ -410,12 +443,7 @@ mod tests {
}
}
fn make_decl(
name: &str,
kind: DeclKind,
strings: &[&str],
props: &[&str],
) -> Declaration {
fn make_decl(name: &str, kind: DeclKind, strings: &[&str], props: &[&str]) -> Declaration {
Declaration {
name: name.to_string(),
kind,

24
crates/ruvector-decompiler/src/lib.rs
View file

@ -53,16 +53,16 @@ pub mod model_types;
pub use error::{DecompilerError, Result};
pub use types::{
DecompileConfig, DecompileResult, Declaration, InferredName, Module,
ModuleTree, WitnessChainData,
Declaration, DecompileConfig, DecompileResult, InferredName, Module, ModuleTree,
WitnessChainData,
};
#[cfg(feature = "model")]
pub use model_decompiler::{decompile_gguf, decompile_model, decompile_safetensors};
#[cfg(feature = "model")]
pub use model_types::{
LayerInfo, LayerType, ModelArchitecture, ModelDecompileResult,
ModelFormat, QuantizationInfo, TokenizerInfo,
LayerInfo, LayerType, ModelArchitecture, ModelDecompileResult, ModelFormat, QuantizationInfo,
TokenizerInfo,
};
/// Decompile a minified JavaScript bundle.
@ -86,10 +86,7 @@ pub fn decompile(source: &str, config: &DecompileConfig) -> Result<DecompileResu
let ref_graph = graph::build_reference_graph(declarations);
// Phase 3: Partition into modules.
let mut modules = partitioner::partition_modules(
&ref_graph,
config.target_modules,
)?;
let mut modules = partitioner::partition_modules(&ref_graph, config.target_modules)?;
// Phase 4: Infer names.
let inferred_names = inferrer::infer_names(&modules);
@ -102,12 +99,7 @@ pub fn decompile(source: &str, config: &DecompileConfig) -> Result<DecompileResu
.collect();
// Beautify module source code.
beautifier::beautify_all(
&mut modules,
source,
&filtered_names,
config.min_confidence,
);
beautifier::beautify_all(&mut modules, source, &filtered_names, config.min_confidence);
// Phase 4b: Build hierarchical module tree from graph structure.
let module_tree = if config.hierarchical_output.unwrap_or(true) {
@ -120,9 +112,7 @@ pub fn decompile(source: &str, config: &DecompileConfig) -> Result<DecompileResu
let source_maps = if config.generate_source_maps {
modules
.iter()
.map(|m| {
sourcemap::generate_source_map(m, &filtered_names, &config.output_filename)
})
.map(|m| sourcemap::generate_source_map(m, &filtered_names, &config.output_filename))
.collect::<Result<Vec<_>>>()?
} else {
Vec::new()

101
crates/ruvector-decompiler/src/model_decompiler.rs
View file

@ -104,15 +104,13 @@ fn infer_architecture_from_gguf(
let num_heads = get_meta_u64(metadata, &format!("{}attention.head_count", arch_prefix))
.unwrap_or(0) as usize;
let num_kv_heads =
get_meta_u64(metadata, &format!("{}attention.head_count_kv", arch_prefix))
.or_else(|| infer_kv_heads(tensors, hidden_size, num_heads))
.unwrap_or(num_heads as u64) as usize;
let num_kv_heads = get_meta_u64(metadata, &format!("{}attention.head_count_kv", arch_prefix))
.or_else(|| infer_kv_heads(tensors, hidden_size, num_heads))
.unwrap_or(num_heads as u64) as usize;
let intermediate_size =
get_meta_u64(metadata, &format!("{}feed_forward_length", arch_prefix))
.or_else(|| infer_ffn_size(tensors))
.unwrap_or(0) as usize;
let intermediate_size = get_meta_u64(metadata, &format!("{}feed_forward_length", arch_prefix))
.or_else(|| infer_ffn_size(tensors))
.unwrap_or(0) as usize;
let vocab_size = infer_vocab_size(tensors).unwrap_or(0);
@ -149,8 +147,8 @@ fn infer_architecture_from_tensors(
let hidden_size = infer_hidden_size(tensors).unwrap_or(0) as usize;
let num_layers = infer_num_layers(tensors);
let num_heads = infer_num_heads(tensors, hidden_size);
let num_kv_heads = infer_kv_heads(tensors, hidden_size, num_heads)
.unwrap_or(num_heads as u64) as usize;
let num_kv_heads =
infer_kv_heads(tensors, hidden_size, num_heads).unwrap_or(num_heads as u64) as usize;
let intermediate_size = infer_ffn_size(tensors).unwrap_or(0) as usize;
let vocab_size = infer_vocab_size(tensors).unwrap_or(0);
let total_params: usize = tensors.iter().map(|t| t.num_elements()).sum();
@ -179,17 +177,13 @@ fn get_meta_u64(metadata: &HashMap<String, GgufValue>, key: &str) -> Option<u64>
fn infer_hidden_size(tensors: &[ModelTensorInfo]) -> Option<u64> {
// Look for embedding tensor: shape [vocab_size, hidden_size]
for t in tensors {
if (t.name.contains("embed") || t.name.contains("token_embd"))
&& t.shape.len() == 2
{
if (t.name.contains("embed") || t.name.contains("token_embd")) && t.shape.len() == 2 {
return Some(t.shape[1] as u64);
}
}
// Fall back to attention Q projection: shape [hidden, hidden]
for t in tensors {
if (t.name.contains("attn_q") || t.name.contains(".q_proj"))
&& t.shape.len() == 2
{
if (t.name.contains("attn_q") || t.name.contains(".q_proj")) && t.shape.len() == 2 {
return Some(t.shape[1] as u64);
}
}
@ -216,9 +210,10 @@ fn infer_num_layers(tensors: &[ModelTensorInfo]) -> usize {
fn extract_layer_index(name: &str) -> Option<usize> {
// Common patterns: "blk.0.", "layers.0.", "h.0.", "model.layers.0."
for prefix in &["blk.", "layers.", "h.", "model.layers."] {
if let Some(rest) = name.strip_prefix(prefix).or_else(|| {
name.find(prefix).map(|i| &name[i + prefix.len()..])
}) {
if let Some(rest) = name
.strip_prefix(prefix)
.or_else(|| name.find(prefix).map(|i| &name[i + prefix.len()..]))
{
if let Some(dot) = rest.find('.') {
if let Ok(idx) = rest[..dot].parse::<usize>() {
return Some(idx);
@ -253,9 +248,7 @@ fn infer_kv_heads(
let head_dim = hidden_size / num_heads;
// Look for K projection tensor shape: [kv_heads * head_dim, hidden_size]
for t in tensors {
if (t.name.contains("attn_k") || t.name.contains(".k_proj"))
&& t.shape.len() == 2
{
if (t.name.contains("attn_k") || t.name.contains(".k_proj")) && t.shape.len() == 2 {
let k_dim = t.shape[0];
if head_dim > 0 && k_dim % head_dim == 0 {
return Some((k_dim / head_dim) as u64);
@ -267,8 +260,10 @@ fn infer_kv_heads(
fn infer_ffn_size(tensors: &[ModelTensorInfo]) -> Option<u64> {
for t in tensors {
if (t.name.contains("ffn_up") || t.name.contains(".up_proj")
|| t.name.contains("ffn_gate") || t.name.contains(".gate_proj"))
if (t.name.contains("ffn_up")
|| t.name.contains(".up_proj")
|| t.name.contains("ffn_gate")
|| t.name.contains(".gate_proj"))
&& t.shape.len() == 2
{
return Some(t.shape[0] as u64);
@ -279,9 +274,7 @@ fn infer_ffn_size(tensors: &[ModelTensorInfo]) -> Option<u64> {
fn infer_vocab_size(tensors: &[ModelTensorInfo]) -> Option<usize> {
for t in tensors {
if (t.name.contains("embed") || t.name.contains("token_embd"))
&& t.shape.len() == 2
{
if (t.name.contains("embed") || t.name.contains("token_embd")) && t.shape.len() == 2 {
return Some(t.shape[0]);
}
}
@ -346,9 +339,12 @@ fn extract_layers(tensors: &[ModelTensorInfo], arch: &ModelArchitecture) -> Vec<
let attn: Vec<&ModelTensorInfo> = block_tensors
.iter()
.filter(|t| {
t.name.contains("attn") || t.name.contains("self_attn")
|| t.name.contains("q_proj") || t.name.contains("k_proj")
|| t.name.contains("v_proj") || t.name.contains("o_proj")
t.name.contains("attn")
|| t.name.contains("self_attn")
|| t.name.contains("q_proj")
|| t.name.contains("k_proj")
|| t.name.contains("v_proj")
|| t.name.contains("o_proj")
})
.copied()
.collect();
@ -370,8 +366,10 @@ fn extract_layers(tensors: &[ModelTensorInfo], arch: &ModelArchitecture) -> Vec<
let mlp: Vec<&ModelTensorInfo> = block_tensors
.iter()
.filter(|t| {
t.name.contains("ffn") || t.name.contains("mlp")
|| t.name.contains("up_proj") || t.name.contains("down_proj")
t.name.contains("ffn")
|| t.name.contains("mlp")
|| t.name.contains("up_proj")
|| t.name.contains("down_proj")
|| t.name.contains("gate_proj")
})
.copied()
@ -416,7 +414,8 @@ fn extract_layers(tensors: &[ModelTensorInfo], arch: &ModelArchitecture) -> Vec<
let output_tensors: Vec<&ModelTensorInfo> = tensors
.iter()
.filter(|t| {
t.name.contains("output") && extract_layer_index(&t.name).is_none()
t.name.contains("output")
&& extract_layer_index(&t.name).is_none()
&& !t.name.contains("norm")
})
.collect();
@ -435,9 +434,7 @@ fn extract_layers(tensors: &[ModelTensorInfo], arch: &ModelArchitecture) -> Vec<
// ── Tokenizer extraction ─────────────────────────────────────────────────
fn extract_tokenizer_from_gguf(
metadata: &HashMap<String, GgufValue>,
) -> Option<TokenizerInfo> {
fn extract_tokenizer_from_gguf(metadata: &HashMap<String, GgufValue>) -> Option<TokenizerInfo> {
let vocab_size = metadata
.get("tokenizer.ggml.tokens")
.and_then(|v| v.as_array())
@ -468,9 +465,7 @@ fn extract_tokenizer_from_gguf(
arr.iter()
.take(100)
.enumerate()
.filter_map(|(i, v)| {
v.as_str().map(|s| (i as u32, s.to_string()))
})
.filter_map(|(i, v)| v.as_str().map(|s| (i as u32, s.to_string())))
.collect::<Vec<_>>()
})
.unwrap_or_default();
@ -577,10 +572,7 @@ fn sha3_hex(data: &[u8]) -> String {
let mut hasher = Sha3_256::new();
hasher.update(data);
let result = hasher.finalize();
result
.iter()
.map(|b| format!("{:02x}", b))
.collect()
result.iter().map(|b| format!("{:02x}", b)).collect()
}
// ── Metadata flattening ──────────────────────────────────────────────────
@ -616,7 +608,10 @@ mod tests {
fn test_extract_layer_index() {
assert_eq!(extract_layer_index("blk.0.attn_q.weight"), Some(0));
assert_eq!(extract_layer_index("blk.31.ffn_up.weight"), Some(31));
assert_eq!(extract_layer_index("model.layers.5.self_attn.q_proj"), Some(5));
assert_eq!(
extract_layer_index("model.layers.5.self_attn.q_proj"),
Some(5)
);
assert_eq!(extract_layer_index("token_embd.weight"), None);
assert_eq!(extract_layer_index("output.weight"), None);
}
@ -636,16 +631,14 @@ mod tests {
#[test]
fn test_infer_arch_name() {
let tensors = vec![
ModelTensorInfo {
name: "model.layers.0.gate_proj.weight".to_string(),
shape: vec![4096, 4096],
quant_type: 0,
quant_name: "F32".to_string(),
bits_per_weight: 32.0,
offset: 0,
},
];
let tensors = vec![ModelTensorInfo {
name: "model.layers.0.gate_proj.weight".to_string(),
shape: vec![4096, 4096],
quant_type: 0,
quant_name: "F32".to_string(),
bits_per_weight: 32.0,
offset: 0,
}];
assert_eq!(infer_arch_name_from_tensor_names(&tensors), "llama");
}

3
crates/ruvector-decompiler/src/model_gguf.rs
View file

@ -287,8 +287,7 @@ fn read_string<R: Read>(r: &mut R) -> Result<String> {
}
let mut buf = vec![0u8; len];
r.read_exact(&mut buf).map_err(read_err)?;
String::from_utf8(buf)
.map_err(|e| DecompilerError::ModelError(format!("invalid UTF-8: {}", e)))
String::from_utf8(buf).map_err(|e| DecompilerError::ModelError(format!("invalid UTF-8: {}", e)))
}
// ── Tests ────────────────────────────────────────────────────────────────

14
crates/ruvector-decompiler/src/model_safetensors.rs
View file

@ -25,9 +25,8 @@ pub(crate) fn parse_safetensors_file(
// Read 8-byte header length
let mut len_bytes = [0u8; 8];
file.read_exact(&mut len_bytes).map_err(|e| {
DecompilerError::ModelError(format!("failed to read header length: {}", e))
})?;
file.read_exact(&mut len_bytes)
.map_err(|e| DecompilerError::ModelError(format!("failed to read header length: {}", e)))?;
let header_len = u64::from_le_bytes(len_bytes);
if header_len > MAX_HEADER_SIZE {
@ -39,12 +38,11 @@ pub(crate) fn parse_safetensors_file(
// Read JSON header
let mut header_bytes = vec![0u8; header_len as usize];
file.read_exact(&mut header_bytes).map_err(|e| {
DecompilerError::ModelError(format!("failed to read header JSON: {}", e))
})?;
file.read_exact(&mut header_bytes)
.map_err(|e| DecompilerError::ModelError(format!("failed to read header JSON: {}", e)))?;
let header: HashMap<String, serde_json::Value> =
serde_json::from_slice(&header_bytes).map_err(|e| {
let header: HashMap<String, serde_json::Value> = serde_json::from_slice(&header_bytes)
.map_err(|e| {
DecompilerError::ModelError(format!("invalid safetensors header JSON: {}", e))
})?;

11
crates/ruvector-decompiler/src/model_types.rs
View file

@ -101,10 +101,7 @@ pub enum LayerType {
head_dim: usize,
},
/// Feed-forward / MLP layer.
Mlp {
up_size: usize,
down_size: usize,
},
Mlp { up_size: usize, down_size: usize },
/// Layer normalization.
LayerNorm,
/// RMS normalization.
@ -117,7 +114,11 @@ impl std::fmt::Display for LayerType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Embedding => write!(f, "Embedding"),
Self::Attention { heads, kv_heads, head_dim } => {
Self::Attention {
heads,
kv_heads,
head_dim,
} => {
if heads == kv_heads {
write!(f, "Attention ({} heads, dim {})", heads, head_dim)
} else {

45
crates/ruvector-decompiler/src/neural.rs
View file

@ -72,9 +72,7 @@ impl NeuralInferrer {
let session = ort::session::Session::builder()
.and_then(|b| b.commit_from_file(path))
.map_err(|e| {
crate::error::DecompilerError::ModelError(format!(
"failed to load ONNX model: {e}"
))
crate::error::DecompilerError::ModelError(format!("failed to load ONNX model: {e}"))
})?;
Ok(Self {
@ -85,9 +83,7 @@ impl NeuralInferrer {
fn load_legacy(path: &Path) -> Result<Self, crate::error::DecompilerError> {
let data = std::fs::read(path).map_err(|e| {
crate::error::DecompilerError::ModelError(format!(
"failed to read model file: {e}"
))
crate::error::DecompilerError::ModelError(format!("failed to read model file: {e}"))
})?;
if data.len() < 4 {
@ -113,11 +109,7 @@ impl NeuralInferrer {
}
/// Predict the original name for a minified identifier.
pub fn predict_name(
&self,
minified: &str,
context: &InferenceContext,
) -> Option<InferredName> {
pub fn predict_name(&self, minified: &str, context: &InferenceContext) -> Option<InferredName> {
match &self.backend {
Backend::Transformer(encoder) => {
let ctx_strings: Vec<&str> = context
@ -178,31 +170,19 @@ impl NeuralInferrer {
.take(Self::MAX_CONTEXT_LEN)
.collect();
let name_tensor = Tensor::from_array((
vec![1i64, Self::MAX_NAME_LEN as i64],
name_bytes,
))
.ok()?;
let ctx_tensor = Tensor::from_array((
vec![1i64, Self::MAX_CONTEXT_LEN as i64],
ctx_bytes,
))
.ok()?;
let name_tensor =
Tensor::from_array((vec![1i64, Self::MAX_NAME_LEN as i64], name_bytes)).ok()?;
let ctx_tensor =
Tensor::from_array((vec![1i64, Self::MAX_CONTEXT_LEN as i64], ctx_bytes)).ok()?;
let outputs = session
.run(ort::inputs![name_tensor, ctx_tensor])
.ok()?;
let outputs = session.run(ort::inputs![name_tensor, ctx_tensor]).ok()?;
if outputs.len() < 2 {
return None;
}
let (_shape, out_data) = outputs[0]
.try_extract_tensor::<f32>()
.ok()?;
let (_cshape, conf_data) = outputs[1]
.try_extract_tensor::<f32>()
.ok()?;
let (_shape, out_data) = outputs[0].try_extract_tensor::<f32>().ok()?;
let (_cshape, conf_data) = outputs[1].try_extract_tensor::<f32>().ok()?;
let confidence = *conf_data.first()? as f64;
if confidence < 0.5 {
@ -257,10 +237,7 @@ impl NeuralInferrer {
///
/// Neural inference is attempted first; results with confidence > 0.8
/// are accepted directly. Otherwise falls through to corpus + heuristics.
pub fn infer_names_neural(
modules: &[Module],
model_path: Option<&Path>,
) -> Vec<InferredName> {
pub fn infer_names_neural(modules: &[Module], model_path: Option<&Path>) -> Vec<InferredName> {
let corpus = TrainingCorpus::builtin();
let neural = model_path.and_then(|p| NeuralInferrer::load(p).ok());

29
crates/ruvector-decompiler/src/parser.rs
View file

@ -53,13 +53,11 @@ static VAR_RE: Lazy<Regex> = Lazy::new(|| {
});
static FN_RE: Lazy<Regex> = Lazy::new(|| {
Regex::new(r"(?:^|[;}\s])function\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\s*\(")
.expect("valid regex")
Regex::new(r"(?:^|[;}\s])function\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\s*\(").expect("valid regex")
});
static CLASS_RE: Lazy<Regex> = Lazy::new(|| {
Regex::new(r"(?:^|[;}\s])class\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\s*[{\(]")
.expect("valid regex")
Regex::new(r"(?:^|[;}\s])class\s+([a-zA-Z_$][a-zA-Z0-9_$]*)\s*[{\(]").expect("valid regex")
});
static EXPORT_RE: Lazy<Regex> = Lazy::new(|| {
@ -70,9 +68,7 @@ static EXPORT_RE: Lazy<Regex> = Lazy::new(|| {
/// Parse a minified JavaScript bundle and extract declarations.
pub fn parse_bundle(source: &str) -> Result<Vec<Declaration>> {
if source.trim().is_empty() {
return Err(DecompilerError::EmptyBundle(
"source is empty".to_string(),
));
return Err(DecompilerError::EmptyBundle("source is empty".to_string()));
}
let decls = extract_declarations(source);
@ -188,10 +184,7 @@ fn extract_declarations(source: &str) -> Vec<Declaration> {
// Cross-references: identifiers that match other declaration names.
let mut seen_refs = HashSet::with_capacity(16);
for ident in idents {
if ident != decl.name
&& all_names.contains(&ident)
&& seen_refs.insert(ident.clone())
{
if ident != decl.name && all_names.contains(&ident) && seen_refs.insert(ident.clone()) {
decl.references.push(ident);
}
}
@ -279,17 +272,13 @@ fn scan_body_single_pass(body: &str) -> (Vec<String>, Vec<String>, Vec<String>)
}
// --- Property access (after '.') ---
if ch == b'.'
&& i + 1 < len
&& (table[bytes[i + 1] as usize] & IDENT_START) != 0
{
if ch == b'.' && i + 1 < len && (table[bytes[i + 1] as usize] & IDENT_START) != 0 {
i += 1;
let prop_start = i;
while i < len && (table[bytes[i] as usize] & IDENT_CHAR) != 0 {
i += 1;
}
let prop =
unsafe { std::str::from_utf8_unchecked(&bytes[prop_start..i]) };
let prop = unsafe { std::str::from_utf8_unchecked(&bytes[prop_start..i]) };
props.push(prop.to_string());
continue;
}
@ -300,8 +289,7 @@ fn scan_body_single_pass(body: &str) -> (Vec<String>, Vec<String>, Vec<String>)
while i < len && (table[bytes[i] as usize] & IDENT_CHAR) != 0 {
i += 1;
}
let ident =
unsafe { std::str::from_utf8_unchecked(&bytes[ident_start..i]) };
let ident = unsafe { std::str::from_utf8_unchecked(&bytes[ident_start..i]) };
idents.push(ident.to_string());
continue;
}
@ -350,7 +338,8 @@ fn find_declaration_end(source: &str, start: usize) -> usize {
let abs = i + pos;
// Count preceding backslashes.
let mut bs = 0;
while abs > 0 && abs - 1 >= i.saturating_sub(bs + 1)
while abs > 0
&& abs - 1 >= i.saturating_sub(bs + 1)
&& abs > bs
&& bytes[abs - 1 - bs] == b'\\'
{

63
crates/ruvector-decompiler/src/partitioner.rs
View file

@ -43,10 +43,7 @@ pub fn partition_modules(
let target = target.clamp(1, n);
if target == 1 || n <= 2 {
return Ok(vec![build_module(
0,
&graph.declarations,
)]);
return Ok(vec![build_module(0, &graph.declarations)]);
}
// Choose algorithm based on graph size.
@ -58,15 +55,11 @@ pub fn partition_modules(
}
/// Exact MinCut partitioning for small-to-medium graphs (<5K nodes).
fn exact_mincut_partition(
graph: &ReferenceGraph,
target: usize,
) -> Result<Vec<Module>> {
fn exact_mincut_partition(graph: &ReferenceGraph, target: usize) -> Result<Vec<Module>> {
let partitioner = GraphPartitioner::new(graph.graph.clone(), target);
let partitions = partitioner.partition();
let mut assigned: std::collections::HashSet<usize> =
std::collections::HashSet::new();
let mut assigned: std::collections::HashSet<usize> = std::collections::HashSet::new();
let mut modules = Vec::new();
let mut mod_idx = 0;
@ -111,10 +104,7 @@ fn exact_mincut_partition(
/// modularity gains in parallel. Each node reads the current community
/// assignments without locking. Moves are applied sequentially after
/// each parallel sweep to prevent conflicts.
fn louvain_partition(
graph: &ReferenceGraph,
target: usize,
) -> Result<Vec<Module>> {
fn louvain_partition(graph: &ReferenceGraph, target: usize) -> Result<Vec<Module>> {
let n = graph.node_count();
// Build adjacency list from the reference graph.
@ -180,8 +170,7 @@ fn louvain_partition(
}
// Compute sum of weights to each neighbor community.
let mut comm_weights: HashMap<usize, f64> =
HashMap::with_capacity(adj[i].len());
let mut comm_weights: HashMap<usize, f64> = HashMap::with_capacity(adj[i].len());
for &(j, w) in &adj[i] {
*comm_weights.entry(community[j]).or_insert(0.0) += w;
}
@ -189,8 +178,7 @@ fn louvain_partition(
let mut best_comm = current_comm;
let mut best_gain = 0.0f64;
let ki_in_current =
comm_weights.get(&current_comm).copied().unwrap_or(0.0);
let ki_in_current = comm_weights.get(&current_comm).copied().unwrap_or(0.0);
let sigma_current = sigma_totals[current_comm];
for (&candidate_comm, &ki_in_candidate) in &comm_weights {
@ -271,10 +259,7 @@ fn louvain_partition(
}
/// Fallback: positional partitioning by byte offset for edge-less graphs.
fn positional_partition(
graph: &ReferenceGraph,
target: usize,
) -> Result<Vec<Module>> {
fn positional_partition(graph: &ReferenceGraph, target: usize) -> Result<Vec<Module>> {
let n = graph.node_count();
let chunk_size = (n + target - 1) / target;
@ -313,25 +298,19 @@ fn distribute_orphans(
.iter_mut()
.min_by_key(|m| {
let mid = (m.byte_range.0 + m.byte_range.1) / 2;
let orphan_mid =
(orphan.byte_range.0 + orphan.byte_range.1) / 2;
let orphan_mid = (orphan.byte_range.0 + orphan.byte_range.1) / 2;
(mid as i64 - orphan_mid as i64).unsigned_abs()
})
.unwrap();
best_module.declarations.push(orphan.clone());
best_module.byte_range.0 =
best_module.byte_range.0.min(orphan.byte_range.0);
best_module.byte_range.1 =
best_module.byte_range.1.max(orphan.byte_range.1);
best_module.byte_range.0 = best_module.byte_range.0.min(orphan.byte_range.0);
best_module.byte_range.1 = best_module.byte_range.1.max(orphan.byte_range.1);
}
}
}
/// Finalize module list: ensure at least one module exists.
fn finalize_modules(
graph: &ReferenceGraph,
modules: Vec<Module>,
) -> Result<Vec<Module>> {
fn finalize_modules(graph: &ReferenceGraph, modules: Vec<Module>) -> Result<Vec<Module>> {
if modules.is_empty() {
Ok(vec![build_module(0, &graph.declarations)])
} else {
@ -386,7 +365,13 @@ fn infer_module_name(decls: &[Declaration], fallback_index: usize) -> String {
fn sanitize_module_name(raw: &str) -> String {
let cleaned: String = raw
.chars()
.map(|c| if c.is_alphanumeric() || c == '_' { c } else { '_' })
.map(|c| {
if c.is_alphanumeric() || c == '_' {
c
} else {
'_'
}
})
.collect();
if cleaned.is_empty() {
"module".to_string()
@ -470,18 +455,10 @@ mod tests {
let mut decls = Vec::new();
for i in 0..50 {
let refs: Vec<&str> = Vec::new();
decls.push(make_decl(
&format!("a{}", i),
&refs,
&["cluster_a"],
));
decls.push(make_decl(&format!("a{}", i), &refs, &["cluster_a"]));
}
for i in 0..50 {
decls.push(make_decl(
&format!("b{}", i),
&[],
&["cluster_b"],
));
decls.push(make_decl(&format!("b{}", i), &[], &["cluster_b"]));
}
// Add cross-references within clusters.
for i in 1..50 {

32
crates/ruvector-decompiler/src/training.rs
View file

@ -69,20 +69,9 @@ impl TrainingCorpus {
///
/// Returns the best-matching pattern with a computed match score.
/// Requires at least one context string or property name match.
pub fn match_declaration(
&self,
decl: &Declaration,
) -> Option<(&TrainingPattern, f64)> {
let decl_strings: HashSet<&str> = decl
.string_literals
.iter()
.map(|s| s.as_str())
.collect();
let decl_props: HashSet<&str> = decl
.property_accesses
.iter()
.map(|s| s.as_str())
.collect();
pub fn match_declaration(&self, decl: &Declaration) -> Option<(&TrainingPattern, f64)> {
let decl_strings: HashSet<&str> = decl.string_literals.iter().map(|s| s.as_str()).collect();
let decl_props: HashSet<&str> = decl.property_accesses.iter().map(|s| s.as_str()).collect();
let mut best: Option<(&TrainingPattern, f64)> = None;
@ -107,14 +96,12 @@ impl TrainingCorpus {
.filter(|pn| decl_props.contains(pn.as_str()))
.count();
let total_signals =
pattern.context_strings.len() + pattern.property_names.len();
let total_signals = pattern.context_strings.len() + pattern.property_names.len();
if total_signals == 0 {
continue;
}
let match_ratio =
(string_matches + prop_matches) as f64 / total_signals as f64;
let match_ratio = (string_matches + prop_matches) as f64 / total_signals as f64;
// Require at least one match to consider this pattern.
if string_matches + prop_matches == 0 {
@ -152,11 +139,7 @@ mod tests {
use super::*;
use crate::types::DeclKind;
fn make_decl(
name: &str,
strings: &[&str],
props: &[&str],
) -> Declaration {
fn make_decl(name: &str, strings: &[&str], props: &[&str]) -> Declaration {
Declaration {
name: name.to_string(),
kind: DeclKind::Var,
@ -205,8 +188,7 @@ mod tests {
assert!(result.is_some());
let (pattern, score) = result.unwrap();
assert!(
pattern.inferred_name.contains("Mcp")
|| pattern.inferred_name.contains("Protocol"),
pattern.inferred_name.contains("Mcp") || pattern.inferred_name.contains("Protocol"),
"Expected MCP-related name, got: {}",
pattern.inferred_name
);

221
crates/ruvector-decompiler/src/transformer.rs
View file

@ -16,8 +16,8 @@ const MAX_NAME: usize = 32;
const TOTAL_SEQ: usize = MAX_CONTEXT + MAX_NAME;
struct TransformerLayer {
in_proj_weight: Vec<f32>, // [3*D, D]
in_proj_bias: Vec<f32>, // [3*D]
in_proj_weight: Vec<f32>, // [3*D, D]
in_proj_bias: Vec<f32>, // [3*D]
out_proj_weight: Vec<f32>, // [D, D]
out_proj_bias: Vec<f32>,
norm1_weight: Vec<f32>,
@ -36,21 +36,20 @@ pub struct TransformerEncoder {
embed_dim: usize,
num_heads: usize,
ffn_dim: usize,
char_embed: Vec<f32>, // [VOCAB_SIZE * D]
pos_embed: Vec<f32>, // [TOTAL_SEQ * D]
char_embed: Vec<f32>, // [VOCAB_SIZE * D]
pos_embed: Vec<f32>, // [TOTAL_SEQ * D]
layers: Vec<TransformerLayer>,
final_norm_weight: Vec<f32>, // [D]
final_norm_bias: Vec<f32>,
output_weight: Vec<f32>, // [VOCAB_SIZE * D]
output_bias: Vec<f32>, // [VOCAB_SIZE]
output_weight: Vec<f32>, // [VOCAB_SIZE * D]
output_bias: Vec<f32>, // [VOCAB_SIZE]
}
impl TransformerEncoder {
/// Load from binary weights file (see `export-weights-bin.py`).
pub fn from_weights_bin(path: &Path) -> Result<Self, DecompilerError> {
let data = std::fs::read(path).map_err(|e| {
DecompilerError::ModelError(format!("failed to read weights: {e}"))
})?;
let data = std::fs::read(path)
.map_err(|e| DecompilerError::ModelError(format!("failed to read weights: {e}")))?;
Self::from_tensor_map(&parse_bin_tensors(&data)?)
}
@ -58,29 +57,43 @@ impl TransformerEncoder {
pub fn forward(&self, context: &[u8], name: &[u8]) -> Vec<Vec<f32>> {
let d = self.embed_dim;
let mut tokens = vec![PAD_TOKEN; TOTAL_SEQ];
for (i, &b) in context.iter().take(MAX_CONTEXT).enumerate() { tokens[i] = b; }
for (i, &b) in name.iter().take(MAX_NAME).enumerate() { tokens[MAX_CONTEXT + i] = b; }
for (i, &b) in context.iter().take(MAX_CONTEXT).enumerate() {
tokens[i] = b;
}
for (i, &b) in name.iter().take(MAX_NAME).enumerate() {
tokens[MAX_CONTEXT + i] = b;
}
// Embedding: char_embed + pos_embed
let mut x = vec![0.0f32; TOTAL_SEQ * d];
for (pos, &tok) in tokens.iter().enumerate() {
let (ce, pe, xo) = ((tok as usize) * d, pos * d, pos * d);
for j in 0..d { x[xo + j] = self.char_embed[ce + j] + self.pos_embed[pe + j]; }
for j in 0..d {
x[xo + j] = self.char_embed[ce + j] + self.pos_embed[pe + j];
}
}
let pad_mask: Vec<bool> = tokens.iter().map(|&t| t == PAD_TOKEN).collect();
for layer in &self.layers { x = self.layer_forward(layer, &x, &pad_mask); }
for layer in &self.layers {
x = self.layer_forward(layer, &x, &pad_mask);
}
// Final layer norm + output projection on last MAX_NAME positions
let mut logits = Vec::with_capacity(MAX_NAME);
for i in 0..MAX_NAME {
let off = (MAX_CONTEXT + i) * d;
let normed = layer_norm(&x[off..off + d], &self.final_norm_weight, &self.final_norm_bias);
let normed = layer_norm(
&x[off..off + d],
&self.final_norm_weight,
&self.final_norm_bias,
);
let mut out = vec![0.0f32; VOCAB_SIZE];
for v in 0..VOCAB_SIZE {
let wo = v * d;
let mut s = self.output_bias[v];
for j in 0..d { s += self.output_weight[wo + j] * normed[j]; }
for j in 0..d {
s += self.output_weight[wo + j] * normed[j];
}
out[v] = s;
}
logits.push(out);
@ -90,7 +103,11 @@ impl TransformerEncoder {
/// Predict original name from minified name + context strings.
pub fn predict(&self, minified: &str, context_strings: &[&str]) -> (String, f32) {
let ctx: Vec<u8> = context_strings.join(" ").bytes().take(MAX_CONTEXT).collect();
let ctx: Vec<u8> = context_strings
.join(" ")
.bytes()
.take(MAX_CONTEXT)
.collect();
let nm: Vec<u8> = minified.bytes().take(MAX_NAME).collect();
let logits = self.forward(&ctx, &nm);
@ -98,11 +115,17 @@ impl TransformerEncoder {
let (mut total_conf, mut count) = (0.0f32, 0usize);
for pos_logits in &logits {
let probs = softmax(pos_logits);
let (idx, &prob) = probs.iter().enumerate()
let (idx, &prob) = probs
.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
.unwrap_or((0, &0.0));
if idx == 0 || idx == 2 { break; } // PAD or EOS
if idx == 1 { continue; } // SOS
if idx == 0 || idx == 2 {
break;
} // PAD or EOS
if idx == 1 {
continue;
} // SOS
let ch = idx as u8;
if ch.is_ascii_alphanumeric() || ch == b'_' {
predicted.push(ch as char);
@ -110,22 +133,40 @@ impl TransformerEncoder {
count += 1;
}
}
(predicted, if count > 0 { total_conf / count as f32 } else { 0.0 })
(
predicted,
if count > 0 {
total_conf / count as f32
} else {
0.0
},
)
}
/// Single encoder layer: self-attn + residual + norm + FFN + residual + norm (post-norm).
fn layer_forward(&self, l: &TransformerLayer, x: &[f32], pad_mask: &[bool]) -> Vec<f32> {
let (d, seq, ffn) = (self.embed_dim, TOTAL_SEQ, self.ffn_dim);
let attn = mha(x, &l.in_proj_weight, &l.in_proj_bias,
&l.out_proj_weight, &l.out_proj_bias, seq, d, self.num_heads, pad_mask);
let attn = mha(
x,
&l.in_proj_weight,
&l.in_proj_bias,
&l.out_proj_weight,
&l.out_proj_bias,
seq,
d,
self.num_heads,
pad_mask,
);
// Residual + LayerNorm1
let mut mid = vec![0.0f32; seq * d];
for p in 0..seq {
let o = p * d;
let mut r = vec![0.0f32; d];
for j in 0..d { r[j] = x[o + j] + attn[o + j]; }
for j in 0..d {
r[j] = x[o + j] + attn[o + j];
}
mid[o..o + d].copy_from_slice(&layer_norm(&r, &l.norm1_weight, &l.norm1_bias));
}
@ -138,18 +179,24 @@ impl TransformerEncoder {
for f in 0..ffn {
let wo = f * d;
let mut s = l.linear1_bias[f];
for j in 0..d { s += l.linear1_weight[wo + j] * h[j]; }
for j in 0..d {
s += l.linear1_weight[wo + j] * h[j];
}
h1[f] = gelu(s);
}
let mut h2 = vec![0.0f32; d];
for j in 0..d {
let wo = j * ffn;
let mut s = l.linear2_bias[j];
for f in 0..ffn { s += l.linear2_weight[wo + f] * h1[f]; }
for f in 0..ffn {
s += l.linear2_weight[wo + f] * h1[f];
}
h2[j] = s;
}
let mut r = vec![0.0f32; d];
for j in 0..d { r[j] = mid[o + j] + h2[j]; }
for j in 0..d {
r[j] = mid[o + j] + h2[j];
}
out[o..o + d].copy_from_slice(&layer_norm(&r, &l.norm2_weight, &l.norm2_bias));
}
out
@ -159,7 +206,8 @@ impl TransformerEncoder {
t: &std::collections::HashMap<String, (Vec<usize>, Vec<f32>)>,
) -> Result<Self, DecompilerError> {
let get = |n: &str| -> Result<Vec<f32>, DecompilerError> {
t.get(n).map(|(_, d)| d.clone())
t.get(n)
.map(|(_, d)| d.clone())
.ok_or_else(|| DecompilerError::ModelError(format!("missing tensor: {n}")))
};
let shape = |n: &str| -> Option<&Vec<usize>> { t.get(n).map(|(s, _)| s) };
@ -169,10 +217,13 @@ impl TransformerEncoder {
.ok_or_else(|| DecompilerError::ModelError("char_embed.weight must be 2D".into()))?;
// Count encoder layers
let num_layers = t.keys()
let num_layers = t
.keys()
.filter_map(|k| k.strip_prefix("encoder.layers."))
.filter_map(|r| r.split('.').next()?.parse::<usize>().ok())
.max().map(|m| m + 1).unwrap_or(3);
.max()
.map(|m| m + 1)
.unwrap_or(3);
let ffn_dim = shape("encoder.layers.0.linear1.weight")
.and_then(|s| if s.len() == 2 { Some(s[0]) } else { None })
@ -180,7 +231,13 @@ impl TransformerEncoder {
// Verify in_proj_weight exists; infer num_heads from embed_dim
let _ = get("encoder.layers.0.self_attn.in_proj_weight")?;
let num_heads = if embed_dim % 4 == 0 { 4 } else if embed_dim % 2 == 0 { 2 } else { 1 };
let num_heads = if embed_dim % 4 == 0 {
4
} else if embed_dim % 2 == 0 {
2
} else {
1
};
let mut layers = Vec::with_capacity(num_layers);
for i in 0..num_layers {
@ -202,7 +259,9 @@ impl TransformerEncoder {
}
Ok(Self {
embed_dim, num_heads, ffn_dim,
embed_dim,
num_heads,
ffn_dim,
char_embed: get("char_embed.weight")?,
pos_embed: get("pos_embed.weight")?,
layers,
@ -218,8 +277,15 @@ impl TransformerEncoder {
#[allow(clippy::too_many_arguments)]
fn mha(
x: &[f32], ipw: &[f32], ipb: &[f32], ow: &[f32], ob: &[f32],
seq: usize, d: usize, nh: usize, pad: &[bool],
x: &[f32],
ipw: &[f32],
ipb: &[f32],
ow: &[f32],
ob: &[f32],
seq: usize,
d: usize,
nh: usize,
pad: &[bool],
) -> Vec<f32> {
let hd = d / nh;
let scale = 1.0 / (hd as f32).sqrt();
@ -231,7 +297,9 @@ fn mha(
for i in 0..(3 * d) {
let wo = i * d;
let mut s = ipb[i];
for j in 0..d { s += ipw[wo + j] * x[xo + j]; }
for j in 0..d {
s += ipw[wo + j] * x[xo + j];
}
qkv[p * 3 * d + i] = s;
}
}
@ -241,9 +309,13 @@ fn mha(
let ho = h * hd;
let mut scores = vec![f32::NEG_INFINITY; seq * seq];
for i in 0..seq {
if pad[i] { continue; }
if pad[i] {
continue;
}
for j in 0..seq {
if pad[j] { continue; }
if pad[j] {
continue;
}
let mut dot = 0.0f32;
for k in 0..hd {
dot += qkv[i * 3 * d + ho + k] * qkv[j * 3 * d + d + ho + k];
@ -253,23 +325,39 @@ fn mha(
}
// Softmax per row
for i in 0..seq {
if pad[i] { continue; }
if pad[i] {
continue;
}
let row = &mut scores[i * seq..(i + 1) * seq];
let mx = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
if mx == f32::NEG_INFINITY { continue; }
if mx == f32::NEG_INFINITY {
continue;
}
let mut sum = 0.0f32;
for v in row.iter_mut() {
if *v == f32::NEG_INFINITY { *v = 0.0; }
else { *v = (*v - mx).exp(); sum += *v; }
if *v == f32::NEG_INFINITY {
*v = 0.0;
} else {
*v = (*v - mx).exp();
sum += *v;
}
}
if sum > 0.0 {
for v in row.iter_mut() {
*v /= sum;
}
}
if sum > 0.0 { for v in row.iter_mut() { *v /= sum; } }
}
// Weighted sum of V
for i in 0..seq {
if pad[i] { continue; }
if pad[i] {
continue;
}
for k in 0..hd {
let mut s = 0.0f32;
for j in 0..seq { s += scores[i * seq + j] * qkv[j * 3 * d + 2 * d + ho + k]; }
for j in 0..seq {
s += scores[i * seq + j] * qkv[j * 3 * d + 2 * d + ho + k];
}
attn_out[i * d + ho + k] = s;
}
}
@ -282,7 +370,9 @@ fn mha(
for j in 0..d {
let wo = j * d;
let mut s = ob[j];
for k in 0..d { s += ow[wo + k] * attn_out[ao + k]; }
for k in 0..d {
s += ow[wo + k] * attn_out[ao + k];
}
result[p * d + j] = s;
}
}
@ -296,7 +386,10 @@ fn layer_norm(x: &[f32], w: &[f32], b: &[f32]) -> Vec<f32> {
let mean = x.iter().sum::<f32>() / n;
let var = x.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / n;
let inv = 1.0 / (var + 1e-5f32).sqrt();
x.iter().enumerate().map(|(i, &v)| (v - mean) * inv * w[i] + b[i]).collect()
x.iter()
.enumerate()
.map(|(i, &v)| (v - mean) * inv * w[i] + b[i])
.collect()
}
fn gelu(x: f32) -> f32 {
@ -307,7 +400,11 @@ fn softmax(x: &[f32]) -> Vec<f32> {
let mx = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exps: Vec<f32> = x.iter().map(|&v| (v - mx).exp()).collect();
let s: f32 = exps.iter().sum();
if s > 0.0 { exps.iter().map(|v| v / s).collect() } else { vec![0.0; x.len()] }
if s > 0.0 {
exps.iter().map(|v| v / s).collect()
} else {
vec![0.0; x.len()]
}
}
// -- Binary weight parser --
@ -320,25 +417,30 @@ fn parse_bin_tensors(
let mut buf4 = [0u8; 4];
while (cur.position() as usize) < data.len() {
if cur.read_exact(&mut buf4).is_err() { break; }
if cur.read_exact(&mut buf4).is_err() {
break;
}
let name_len = u32::from_le_bytes(buf4) as usize;
if name_len == 0 || name_len > 1024 { break; }
if name_len == 0 || name_len > 1024 {
break;
}
let mut name_buf = vec![0u8; name_len];
cur.read_exact(&mut name_buf).map_err(|e|
DecompilerError::ModelError(format!("truncated name: {e}")))?;
let name = String::from_utf8(name_buf).map_err(|e|
DecompilerError::ModelError(format!("invalid name: {e}")))?;
cur.read_exact(&mut name_buf)
.map_err(|e| DecompilerError::ModelError(format!("truncated name: {e}")))?;
let name = String::from_utf8(name_buf)
.map_err(|e| DecompilerError::ModelError(format!("invalid name: {e}")))?;
cur.read_exact(&mut buf4).map_err(|e|
DecompilerError::ModelError(format!("truncated ndim for {name}: {e}")))?;
cur.read_exact(&mut buf4)
.map_err(|e| DecompilerError::ModelError(format!("truncated ndim for {name}: {e}")))?;
let ndim = u32::from_le_bytes(buf4) as usize;
let mut shape = Vec::with_capacity(ndim);
let mut numel = 1usize;
for _ in 0..ndim {
cur.read_exact(&mut buf4).map_err(|e|
DecompilerError::ModelError(format!("truncated shape for {name}: {e}")))?;
cur.read_exact(&mut buf4).map_err(|e| {
DecompilerError::ModelError(format!("truncated shape for {name}: {e}"))
})?;
let dim = u32::from_le_bytes(buf4) as usize;
numel *= dim;
shape.push(dim);
@ -348,7 +450,8 @@ fn parse_bin_tensors(
let pos = cur.position() as usize;
if pos + byte_len > data.len() {
return Err(DecompilerError::ModelError(format!(
"truncated data for {name}: need {byte_len} bytes")));
"truncated data for {name}: need {byte_len} bytes"
)));
}
let mut float_data = vec![0.0f32; numel];
for (i, chunk) in data[pos..pos + byte_len].chunks_exact(4).enumerate() {
@ -359,7 +462,9 @@ fn parse_bin_tensors(
}
if tensors.is_empty() {
return Err(DecompilerError::ModelError("no tensors in weights file".into()));
return Err(DecompilerError::ModelError(
"no tensors in weights file".into(),
));
}
Ok(tensors)
}

54
crates/ruvector-decompiler/src/tree.rs
View file

@ -86,7 +86,11 @@ fn build_tree_recursive(
// Base case: max depth reached, or folder is small enough to be a leaf.
// At depth 0: only leaf if <=min_folder_size modules
// At depth 1+: leaf if <=20 modules (enough granularity)
let leaf_threshold = if depth == 0 { min_folder_size } else { 20.min(min_folder_size * 5) };
let leaf_threshold = if depth == 0 {
min_folder_size
} else {
20.min(min_folder_size * 5)
};
if indices.len() <= leaf_threshold || depth >= max_depth {
return make_leaf(all_modules, indices, depth, parent_path);
}
@ -99,8 +103,10 @@ fn build_tree_recursive(
}
// Name the internal node.
let all_in_group: Vec<Module> =
indices.iter().filter_map(|&i| all_modules.get(i).cloned()).collect();
let all_in_group: Vec<Module> = indices
.iter()
.filter_map(|&i| all_modules.get(i).cloned())
.collect();
let folder_name = if depth == 0 {
"src".to_string()
} else {
@ -143,8 +149,10 @@ fn make_leaf(
depth: usize,
parent_path: &str,
) -> ModuleTree {
let leaf_modules: Vec<Module> =
indices.iter().filter_map(|&i| all_modules.get(i).cloned()).collect();
let leaf_modules: Vec<Module> = indices
.iter()
.filter_map(|&i| all_modules.get(i).cloned())
.collect();
let name = infer_folder_name(&leaf_modules, all_modules);
let path = if parent_path.is_empty() {
name.clone()
@ -174,7 +182,15 @@ fn agglomerative_cluster(
let mut clusters: Vec<Vec<usize>> = indices.iter().map(|&i| vec![i]).collect();
// Target: 5-15 top-level folders for large codebases, 3-5 for small
let target_clusters = if n > 500 { 10 } else if n > 100 { 7 } else if n > 20 { 5 } else { 3 };
let target_clusters = if n > 500 {
10
} else if n > 100 {
7
} else if n > 20 {
5
} else {
3
};
loop {
if clusters.len() <= target_clusters {
@ -188,7 +204,9 @@ fn agglomerative_cluster(
for j in (i + 1)..clusters.len() {
let w = cluster_edge_weight(&clusters[i], &clusters[j], inter_module);
// Normalize by geometric mean of sizes (prevents giant clusters from absorbing everything)
let size_factor = ((clusters[i].len() * clusters[j].len()) as f64).sqrt().max(1.0);
let size_factor = ((clusters[i].len() * clusters[j].len()) as f64)
.sqrt()
.max(1.0);
let normalized = w / size_factor;
if normalized > best_weight {
best_weight = normalized;
@ -208,7 +226,7 @@ fn agglomerative_cluster(
// Try next best pair that doesn't exceed max
let mut found = false;
for i in 0..clusters.len() {
for j in (i+1)..clusters.len() {
for j in (i + 1)..clusters.len() {
if clusters[i].len() + clusters[j].len() <= max_cluster {
let w = cluster_edge_weight(&clusters[i], &clusters[j], inter_module);
if w > 0.0 {
@ -225,9 +243,13 @@ fn agglomerative_cluster(
}
}
}
if found { break; }
if found {
break;
}
}
if !found {
break;
}
if !found { break; }
continue;
}
@ -319,7 +341,11 @@ mod tests {
fn test_build_module_tree_basic() {
let mut decls = Vec::new();
for i in 0..5 {
let refs = if i > 0 { vec![format!("a{}", i - 1)] } else { vec![] };
let refs = if i > 0 {
vec![format!("a{}", i - 1)]
} else {
vec![]
};
decls.push(Declaration {
name: format!("a{}", i),
kind: DeclKind::Var,
@ -330,7 +356,11 @@ mod tests {
});
}
for i in 0..5 {
let refs = if i > 0 { vec![format!("b{}", i - 1)] } else { vec![] };
let refs = if i > 0 {
vec![format!("b{}", i - 1)]
} else {
vec![]
};
decls.push(Declaration {
name: format!("b{}", i),
kind: DeclKind::Var,

57
crates/ruvector-decompiler/tests/ground_truth.rs
View file

@ -11,13 +11,8 @@ use ruvector_decompiler::{decompile, DecompileConfig};
// ---------------------------------------------------------------------------
/// Original (known) source for the Express-like fixture.
const EXPRESS_ORIGINAL_NAMES: &[&str] = &[
"Router",
"Request",
"Response",
"createApp",
"handleRoute",
];
const EXPRESS_ORIGINAL_NAMES: &[&str] =
&["Router", "Request", "Response", "createApp", "handleRoute"];
const EXPRESS_MINIFIED: &str = concat!(
r#"var a=class{constructor(){this.routes=[]}"#,
@ -153,11 +148,7 @@ fn test_fixture_mcp_server() {
// Fixture 3: React-like Component
// ---------------------------------------------------------------------------
const REACT_ORIGINAL_NAMES: &[&str] = &[
"useState",
"useEffect",
"Component",
];
const REACT_ORIGINAL_NAMES: &[&str] = &["useState", "useEffect", "Component"];
const REACT_MINIFIED: &str = concat!(
r#"var a=function(b){var c=[b,function(d){c[0]=d}];return c};"#,
@ -198,9 +189,13 @@ fn test_fixture_react_component() {
// ---------------------------------------------------------------------------
const MULTI_ORIGINAL_NAMES: &[&str] = &[
"add", "subtract", "multiply", // Module A: math
"capitalize", "trim", "concat", // Module B: string
"processData", // Module C: uses A + B
"add",
"subtract",
"multiply", // Module A: math
"capitalize",
"trim",
"concat", // Module B: string
"processData", // Module C: uses A + B
];
const MULTI_MINIFIED: &str = concat!(
@ -257,12 +252,7 @@ fn test_fixture_multi_module() {
// Fixture 5: Bundled utils with known tool names
// ---------------------------------------------------------------------------
const TOOLS_ORIGINAL_NAMES: &[&str] = &[
"toolDefinitions",
"bashTool",
"readTool",
"executeTool",
];
const TOOLS_ORIGINAL_NAMES: &[&str] = &["toolDefinitions", "bashTool", "readTool", "executeTool"];
const TOOLS_MINIFIED: &str = concat!(
r#"var a={Bash:{description:"Execute bash commands",inputSchema:{type:"object"}},Read:{description:"Read files",inputSchema:{type:"object"}},Edit:{description:"Edit files",inputSchema:{type:"object"}}};"#,
@ -304,16 +294,13 @@ fn test_fixture_tools_bundle() {
.inferred_names
.iter()
.filter(|n| {
n.evidence.iter().any(|e| {
e.contains("Bash") || e.contains("Read") || e.contains("Error")
})
n.evidence
.iter()
.any(|e| e.contains("Bash") || e.contains("Read") || e.contains("Error"))
})
.collect();
// We expect at least some tool-related inferences.
println!(
" Tool-related inferences: {} found",
tool_related.len()
);
println!(" Tool-related inferences: {} found", tool_related.len());
}
// ---------------------------------------------------------------------------
@ -324,10 +311,7 @@ fn test_fixture_tools_bundle() {
///
/// A match is when the inferred name contains a keyword from the original
/// name (case-insensitive) or vice versa.
fn count_name_hits(
inferred: &[ruvector_decompiler::InferredName],
originals: &[&str],
) -> usize {
fn count_name_hits(inferred: &[ruvector_decompiler::InferredName], originals: &[&str]) -> usize {
let mut hits = 0;
for inf in inferred {
let inf_lower = inf.inferred.to_lowercase();
@ -394,8 +378,7 @@ fn print_metrics(
0.0
};
let avg_confidence = if !inferred_names.is_empty() {
inferred_names.iter().map(|n| n.confidence).sum::<f64>()
/ inferred_names.len() as f64
inferred_names.iter().map(|n| n.confidence).sum::<f64>() / inferred_names.len() as f64
} else {
0.0
};
@ -416,7 +399,11 @@ fn print_metrics(
println!(" Average confidence: {:.2}", avg_confidence);
println!(
" Witness chain: {}",
if chain_root.is_empty() { "INVALID" } else { "VALID" }
if chain_root.is_empty() {
"INVALID"
} else {
"VALID"
}
);
println!();
}

19
crates/ruvector-decompiler/tests/integration.rs
View file

@ -5,8 +5,7 @@
use ruvector_decompiler::{decompile, DecompileConfig};
/// A small minified bundle with 3 declarations and cross-references.
const SAMPLE_BUNDLE: &str =
r#"var a=function(){return"hello"};var b=class{constructor(){this.name="test"}};var c=function(x){return a()+b.name};"#;
const SAMPLE_BUNDLE: &str = r#"var a=function(){return"hello"};var b=class{constructor(){this.name="test"}};var c=function(x){return a()+b.name};"#;
#[test]
fn test_parser_finds_declarations() {
@ -43,14 +42,15 @@ fn test_mincut_partitions() {
let result = decompile(SAMPLE_BUNDLE, &config).unwrap();
// Should produce at least 1 module (partitioning may merge small groups).
assert!(
!result.modules.is_empty(),
"expected at least 1 module"
);
assert!(!result.modules.is_empty(), "expected at least 1 module");
// Total declarations should equal what we parsed.
let total: usize = result.modules.iter().map(|m| m.declarations.len()).sum();
assert!(total >= 3, "expected at least 3 total declarations, got {}", total);
assert!(
total >= 3,
"expected at least 3 total declarations, got {}",
total
);
}
#[test]
@ -89,10 +89,7 @@ fn test_source_map_v3_format() {
parsed["mappings"].is_string(),
"mappings should be a string"
);
assert!(
parsed["sources"].is_array(),
"sources should be an array"
);
assert!(parsed["sources"].is_array(), "sources should be an array");
}
}

20
crates/ruvector-decompiler/tests/model_decompiler.rs
View file

@ -123,9 +123,18 @@ fn test_decompile_gguf_basic() {
// Layers should include embedding, attention, MLP
assert!(!result.layers.is_empty());
assert!(result.layers.iter().any(|l| matches!(l.layer_type, LayerType::Embedding)));
assert!(result.layers.iter().any(|l| matches!(l.layer_type, LayerType::Attention { .. })));
assert!(result.layers.iter().any(|l| matches!(l.layer_type, LayerType::Mlp { .. })));
assert!(result
.layers
.iter()
.any(|l| matches!(l.layer_type, LayerType::Embedding)));
assert!(result
.layers
.iter()
.any(|l| matches!(l.layer_type, LayerType::Attention { .. })));
assert!(result
.layers
.iter()
.any(|l| matches!(l.layer_type, LayerType::Mlp { .. })));
// Witness chain
assert!(!result.witness.source_hash.is_empty());
@ -134,7 +143,10 @@ fn test_decompile_gguf_basic() {
// Metadata
assert_eq!(
result.metadata.get("general.architecture").map(|s| s.as_str()),
result
.metadata
.get("general.architecture")
.map(|s| s.as_str()),
Some("llama")
);

Some files were not shown because too many files have changed in this diff Show more