checkpoint: matrix-agent-validated (2026-04-25)

Architectural snapshot of the lakehouse codebase at the point where the
full matrix-driven agent loop with Mem0 versioning + deletion was
validated end-to-end.

WHAT THIS REPO IS
A clean single-commit snapshot of the lakehouse code. Heavy test data
(.parquet datasets, vector indexes) excluded — see REPLICATION.md for
regen path. Full lakehouse history at git.agentview.dev/profit/lakehouse.

WHAT WAS PROVEN
- Vector retrieval across multi-corpora matrix (chicago_permits + entity
  briefs + sec_tickers + distilled procedural + llm_team runs)
- Observer hand-review (cloud + heuristic fallback) gating each candidate
- Local-model agent loop (qwen3.5:latest) with tool use + scratchpad
- Playbook seal on success → next-iter retrieval surfaces it as preamble
- Mem0 versioning + deletion in pathway_memory:
    * UPSERT: ADD on new workflow, UPDATE bumps replay_count on identical
    * REVISE: chains versions, parent.superseded_at + superseded_by stamped
    * RETIRE: marks specific trace retired with reason, excluded from retrieval
    * HISTORY: walks chain root→tip, cycle-safe

KEY DIRECTORIES
- crates/vectord/src/pathway_memory.rs — Mem0 ops live here
- crates/vectord/src/playbook_memory.rs — original Mem0 reference
- tests/agent_test/ — local-model agent harness + PRD + session archives
- scripts/dump_raw_corpus.sh — MinIO bucket dump (raw test corpus)
- scripts/vectorize_raw_corpus.ts — corpus → vector indexes
- scripts/analyze_chicago_contracts.ts — real inference pipeline
- scripts/seal_agent_playbook.ts — Mem0 upsert from agent traces

Replication: see REPLICATION.md for Debian 13 clean install + cloud-only
adaptation (no local Ollama).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
profit
2026-04-25 19:43:27 -05:00
co-authored by Claude Opus 4.7
commit ac01fffd9a
325 changed files with 82720 additions and 0 deletions
+29
View File
@@ -0,0 +1,29 @@
[package]
name = "vectord"
version = "0.1.0"
edition = "2024"
[dependencies]
shared = { path = "../shared" }
storaged = { path = "../storaged" }
aibridge = { path = "../aibridge" }
catalogd = { path = "../catalogd" }
queryd = { path = "../queryd" }
futures = { workspace = true }
# ADR-019 firewall — vectord-lance owns its own Arrow 57 / Lance 4 deps.
# Public API uses only std types so no version conflict propagates here.
vectord-lance = { path = "../vectord-lance" }
tokio = { workspace = true }
axum = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tracing = { workspace = true }
bytes = { workspace = true }
object_store = { workspace = true }
parquet = { workspace = true }
arrow = { workspace = true }
chrono = { workspace = true }
instant-distance = { workspace = true }
uuid = { workspace = true }
sha2 = { workspace = true }
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
+118
View File
@@ -0,0 +1,118 @@
//! Profile activation tracking (Phase 41 PRD).
//!
//! Phase 41 PRD called out `crates/vectord/src/activation.rs` with
//! `ActivationTracker` + background-job pattern. The activation
//! handler itself lives in `service.rs::activate_profile` (200+ lines
//! of warm-up + bucket binding that's wired to VectorState); this
//! module provides the type the PRD named and a single-flight guard
//! that satisfies the PRD gate "refuse new activation if one is
//! pending/running."
//!
//! Handler extraction (moving the body of `activate_profile` here)
//! is deliberately NOT in this commit — it's a module-structure
//! refactor, not a semantic change. When that lands, the inline
//! `tokio::spawn` in `service.rs` moves into `ActivationTracker::start`
//! and the HTTP handler shrinks to ~20 lines of validate + start +
//! respond-202.
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
/// Tracks in-flight profile activations. The PRD's "single-flight guard"
/// lives here: callers check `is_pending` before starting a new activation
/// and register via `mark_pending` if they proceed. On completion, they
/// call `mark_complete` so the next caller can start.
///
/// Per-profile granularity — activating profile A doesn't block B.
#[derive(Clone, Default)]
pub struct ActivationTracker {
pending: Arc<RwLock<HashMap<String, String>>>, // profile_id → job_id
}
impl ActivationTracker {
pub fn new() -> Self {
Self::default()
}
/// Check if a profile has an activation already running. Returns the
/// in-flight job_id if so. Safe to call without holding a lock.
pub async fn is_pending(&self, profile_id: &str) -> Option<String> {
self.pending.read().await.get(profile_id).cloned()
}
/// Register a new activation as pending. Returns false if an
/// activation is already running for the same profile (caller should
/// return 409 Conflict or surface the existing job_id). Returns true
/// on successful registration.
pub async fn mark_pending(&self, profile_id: &str, job_id: &str) -> bool {
let mut guard = self.pending.write().await;
if guard.contains_key(profile_id) {
return false;
}
guard.insert(profile_id.to_string(), job_id.to_string());
true
}
/// Remove the pending marker when activation finishes (success OR
/// failure — both free the slot for the next caller).
pub async fn mark_complete(&self, profile_id: &str) {
self.pending.write().await.remove(profile_id);
}
/// How many activations are currently in-flight across all profiles.
pub async fn in_flight_count(&self) -> usize {
self.pending.read().await.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn empty_tracker_has_no_pending() {
let t = ActivationTracker::new();
assert_eq!(t.in_flight_count().await, 0);
assert!(t.is_pending("any-profile").await.is_none());
}
#[tokio::test]
async fn mark_pending_registers_the_job() {
let t = ActivationTracker::new();
assert!(t.mark_pending("profile-A", "job-1").await);
assert_eq!(t.in_flight_count().await, 1);
assert_eq!(t.is_pending("profile-A").await, Some("job-1".into()));
}
#[tokio::test]
async fn single_flight_guard_refuses_second_activation_same_profile() {
// PRD Phase 41 gate: "refuse new activation if one is
// pending/running." Same profile twice → second call returns
// false, caller must surface the in-flight job_id.
let t = ActivationTracker::new();
assert!(t.mark_pending("profile-A", "job-1").await);
assert!(!t.mark_pending("profile-A", "job-2").await);
// Still the first job — second registration didn't overwrite.
assert_eq!(t.is_pending("profile-A").await, Some("job-1".into()));
}
#[tokio::test]
async fn different_profiles_dont_block_each_other() {
// Per-profile granularity — activating A doesn't block B.
let t = ActivationTracker::new();
assert!(t.mark_pending("profile-A", "job-1").await);
assert!(t.mark_pending("profile-B", "job-2").await);
assert_eq!(t.in_flight_count().await, 2);
}
#[tokio::test]
async fn mark_complete_frees_the_slot() {
let t = ActivationTracker::new();
t.mark_pending("profile-A", "job-1").await;
t.mark_complete("profile-A").await;
assert_eq!(t.in_flight_count().await, 0);
// Next activation can now proceed.
assert!(t.mark_pending("profile-A", "job-2").await);
}
}
+789
View File
@@ -0,0 +1,789 @@
//! Phase 16.2 + 16.5 — The autotune agent.
//!
//! A long-running tokio task that watches the trial journal and
//! autonomously proposes + runs new HNSW configs. Distinct from
//! `autotune::run_autotune` which is synchronous (one HTTP call, grid
//! of trials, done). The agent is the continuous version: it sleeps,
//! wakes on triggers, proposes configs based on prior trial history,
//! runs them one at a time, and auto-promotes when it finds an
//! improvement.
//!
//! Design invariants:
//! - Trials are data (ADR-018). Every proposal reads the journal; every
//! attempt appends to it. The journal is the agent's memory.
//! - One trial at a time. Bounded Ollama load — the agent never fires
//! multiple parallel embeddings and respects `cooldown_between_trials_secs`.
//! - Rate-limited. `max_trials_per_hour` is a hard ceiling so a
//! misbehaving proposal function can't saturate the system.
//! - Never promotes below `min_recall`. Same safety gate as
//! `run_autotune` — we will not make the index worse.
//! - Triggered OR periodic. Ingest enqueues a `DatasetAppended` event
//! when a new batch lands; the agent also wakes periodically to keep
//! exploring even when nothing changed externally.
//! - Graceful shutdown via the `stop_tx` signal — the handle's Drop
//! doesn't force-kill, but `stop()` requests a clean exit after the
//! current trial.
use chrono::{DateTime, Utc};
use object_store::ObjectStore;
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock, mpsc, oneshot};
use aibridge::client::AiClient;
use catalogd::registry::Registry as CatalogRegistry;
use crate::embedding_cache::EmbeddingCache;
use crate::harness;
use crate::hnsw::HnswStore;
use crate::index_registry::IndexRegistry;
use crate::promotion::{PromotionEntry, PromotionRegistry};
use crate::trial::{HnswConfig, Trial, TrialJournal, TrialMetrics};
// -------- Public-facing types --------
/// Runtime configuration for the agent. Mirrored in shared::config under
/// `[agent]`. Defaults are conservative — designed to tune slowly in the
/// background without fighting real workloads for GPU time.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentConfig {
/// Master switch. When false, `spawn` returns a handle but the loop
/// doesn't run.
pub enabled: bool,
/// Periodic wake-up — even if the trigger queue is empty, every N
/// seconds the agent picks an index with trials and proposes one
/// more config. Keeps exploration alive on idle indexes.
pub cycle_interval_secs: u64,
/// Minimum gap between two trials on the SAME index. Prevents the
/// agent from hammering Ollama when a hot index has many pending
/// triggers in a row.
pub cooldown_between_trials_secs: u64,
/// Below this recall, a proposal is never promoted — even if it
/// beats the champion on latency.
pub min_recall: f32,
/// Budget cap: hard ceiling on trials per hour across all indexes.
/// When hit, the agent idles until the hour window rolls.
pub max_trials_per_hour: u32,
}
impl Default for AgentConfig {
fn default() -> Self {
Self {
enabled: false, // opt-in — don't auto-start until J turns it on
cycle_interval_secs: 60,
cooldown_between_trials_secs: 30,
min_recall: 0.9,
max_trials_per_hour: 30,
}
}
}
/// What caused the agent to look at a particular index. Recorded in the
/// trial's note field so we can tell "new data arrived" trials from
/// "periodic exploration" trials in the journal.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum TriggerReason {
/// Ingest just appended to a dataset that has attached HNSW indexes.
DatasetAppended { dataset: String },
/// A human or another agent hit `/agent/enqueue/{index}`.
Manual,
/// Periodic wake — no external event, just keep exploring.
Periodic,
}
/// One unit of work for the agent.
#[derive(Debug, Clone)]
pub struct TriggerEvent {
pub index_name: String,
pub reason: TriggerReason,
pub enqueued_at: DateTime<Utc>,
}
impl TriggerEvent {
pub fn manual(index_name: impl Into<String>) -> Self {
Self { index_name: index_name.into(), reason: TriggerReason::Manual, enqueued_at: Utc::now() }
}
pub fn dataset_appended(index_name: impl Into<String>, dataset: impl Into<String>) -> Self {
Self {
index_name: index_name.into(),
reason: TriggerReason::DatasetAppended { dataset: dataset.into() },
enqueued_at: Utc::now(),
}
}
pub fn periodic(index_name: impl Into<String>) -> Self {
Self { index_name: index_name.into(), reason: TriggerReason::Periodic, enqueued_at: Utc::now() }
}
}
/// Observable snapshot of the agent's state — what `/agent/status` returns.
#[derive(Debug, Clone, Serialize)]
pub struct AgentStatus {
pub running: bool,
pub config: AgentConfig,
pub queue_depth: usize,
pub trials_run: u64,
pub promotions: u64,
pub trials_in_last_hour: u32,
pub last_event: Option<AgentEvent>,
pub started_at: Option<DateTime<Utc>>,
}
/// Last thing that happened — useful for "why didn't it do anything?" debugging.
#[derive(Debug, Clone, Serialize)]
pub struct AgentEvent {
pub at: DateTime<Utc>,
pub kind: String, // "trial_completed" | "promoted" | "skipped_rate_limit" | etc
pub index_name: Option<String>,
pub detail: String,
}
/// Handle returned by `spawn`. Holds the trigger sender + shared status +
/// stop signal.
#[derive(Clone)]
pub struct AgentHandle {
trigger_tx: mpsc::Sender<TriggerEvent>,
inner: Arc<AgentInner>,
}
struct AgentInner {
status: RwLock<AgentStatus>,
stop_tx: Mutex<Option<oneshot::Sender<()>>>,
queue_len: Mutex<usize>, // mirror of the channel capacity — for status reporting
recent_trials: Mutex<VecDeque<DateTime<Utc>>>, // ring of recent trial timestamps for rate limit
}
impl AgentHandle {
/// Enqueue a trigger. Returns Err if the agent isn't running or the
/// queue is full (backpressure — dropping events is correct here
/// since periodic exploration will pick up the slack).
pub async fn enqueue(&self, event: TriggerEvent) -> Result<(), String> {
self.trigger_tx.try_send(event).map_err(|e| format!("enqueue: {e}"))?;
let mut guard = self.inner.queue_len.lock().await;
*guard = guard.saturating_add(1);
// Update queue_depth in status for observability.
let mut s = self.inner.status.write().await;
s.queue_depth = *guard;
Ok(())
}
pub async fn status(&self) -> AgentStatus {
let mut s = self.inner.status.read().await.clone();
// Refresh rate-limit window from ring buffer.
let cutoff = Utc::now() - chrono::Duration::hours(1);
let ring = self.inner.recent_trials.lock().await;
s.trials_in_last_hour = ring.iter().filter(|t| **t >= cutoff).count() as u32;
s
}
/// Request a graceful stop. Returns immediately — the loop exits
/// after its current trial completes.
pub async fn stop(&self) -> bool {
let mut guard = self.inner.stop_tx.lock().await;
if let Some(tx) = guard.take() {
let _ = tx.send(());
true
} else {
false
}
}
}
// -------- Agent state holder --------
/// Everything the agent needs to run a trial. Mirrors the fields of
/// `VectorState` the agent actually uses. Kept separate so the service
/// layer builds it explicitly — no clone of unneeded state.
#[derive(Clone)]
pub struct AgentDeps {
pub store: Arc<dyn ObjectStore>,
pub ai_client: AiClient,
pub catalog: CatalogRegistry,
pub index_registry: IndexRegistry,
pub hnsw_store: HnswStore,
pub embedding_cache: EmbeddingCache,
pub trial_journal: TrialJournal,
pub promotion_registry: PromotionRegistry,
pub harness_store: crate::harness::HarnessStore,
}
// -------- Spawn --------
/// Start the agent loop in a background tokio task. Returns a handle
/// the caller uses to enqueue events and read status.
pub fn spawn(config: AgentConfig, deps: AgentDeps) -> AgentHandle {
let (trigger_tx, trigger_rx) = mpsc::channel::<TriggerEvent>(256);
let (stop_tx, stop_rx) = oneshot::channel::<()>();
let status = AgentStatus {
running: config.enabled,
config: config.clone(),
queue_depth: 0,
trials_run: 0,
promotions: 0,
trials_in_last_hour: 0,
last_event: None,
started_at: if config.enabled { Some(Utc::now()) } else { None },
};
let inner = Arc::new(AgentInner {
status: RwLock::new(status),
stop_tx: Mutex::new(Some(stop_tx)),
queue_len: Mutex::new(0),
recent_trials: Mutex::new(VecDeque::with_capacity(64)),
});
if config.enabled {
tracing::info!(
"autotune agent started (cycle={}s, cooldown={}s, cap={}/hr, min_recall={})",
config.cycle_interval_secs, config.cooldown_between_trials_secs,
config.max_trials_per_hour, config.min_recall,
);
let loop_inner = inner.clone();
let loop_deps = deps.clone();
let loop_config = config.clone();
tokio::spawn(async move {
run_loop(loop_config, loop_deps, trigger_rx, stop_rx, loop_inner).await;
});
} else {
// Agent disabled — still drain the channel so sends don't back up.
tokio::spawn(async move {
let mut rx = trigger_rx;
while rx.recv().await.is_some() {}
});
tracing::info!("autotune agent configured but disabled (set [agent].enabled=true)");
}
AgentHandle { trigger_tx, inner }
}
// -------- Main loop --------
async fn run_loop(
config: AgentConfig,
deps: AgentDeps,
mut trigger_rx: mpsc::Receiver<TriggerEvent>,
mut stop_rx: oneshot::Receiver<()>,
inner: Arc<AgentInner>,
) {
let mut periodic = tokio::time::interval(std::time::Duration::from_secs(config.cycle_interval_secs));
// First tick fires immediately — skip it so we don't double-fire on startup.
periodic.tick().await;
loop {
let event = tokio::select! {
_ = &mut stop_rx => {
tracing::info!("autotune agent: stop signal received");
let mut s = inner.status.write().await;
s.running = false;
return;
}
maybe = trigger_rx.recv() => match maybe {
Some(ev) => {
let mut guard = inner.queue_len.lock().await;
*guard = guard.saturating_sub(1);
let mut s = inner.status.write().await;
s.queue_depth = *guard;
ev
}
None => {
tracing::info!("autotune agent: trigger channel closed");
return;
}
},
_ = periodic.tick() => {
// Periodic wake — pick an index with existing trials.
// If nothing's been tuned yet there's nothing to propose.
match pick_periodic_target(&deps).await {
Some(idx) => TriggerEvent::periodic(idx),
None => continue,
}
}
};
// Rate limit: count recent trials, skip if over budget.
if over_rate_limit(&inner, config.max_trials_per_hour).await {
record_event(&inner, "skipped_rate_limit", Some(&event.index_name),
format!("hit cap of {}/hour", config.max_trials_per_hour)).await;
continue;
}
// Per-index cooldown.
if cooling_down(&inner, &event.index_name, config.cooldown_between_trials_secs).await {
record_event(&inner, "skipped_cooldown", Some(&event.index_name),
format!("last trial too recent (<{}s)", config.cooldown_between_trials_secs)).await;
continue;
}
// Run one trial.
match run_one_cycle(&event, &deps, config.min_recall).await {
Ok(outcome) => {
mark_recent_trial(&inner).await;
{
let mut s = inner.status.write().await;
s.trials_run += 1;
if outcome.promoted { s.promotions += 1; }
}
record_event(&inner, if outcome.promoted { "promoted" } else { "trial_completed" },
Some(&event.index_name),
format!("config=ec{}/es{} recall={:.3} p50={:.0}us {}",
outcome.trial.config.ef_construction,
outcome.trial.config.ef_search,
outcome.trial.metrics.recall_at_k,
outcome.trial.metrics.search_latency_p50_us,
if outcome.promoted { "★ PROMOTED" } else { "" })).await;
}
Err(e) => {
record_event(&inner, "trial_error", Some(&event.index_name), e).await;
}
}
}
}
/// Result of one cycle — ran a trial, maybe promoted it.
struct CycleOutcome {
trial: Trial,
promoted: bool,
}
/// Core cycle: propose → build → bench → record → maybe promote.
async fn run_one_cycle(
event: &TriggerEvent,
deps: &AgentDeps,
min_recall: f32,
) -> Result<CycleOutcome, String> {
// Read history.
let history = deps.trial_journal.list(&event.index_name).await
.map_err(|e| format!("read journal: {e}"))?;
// Bootstrap mode: first ever visit to this index. Auto-generate a
// harness (if missing) and seed with the default config — the Phase 15
// known-good ec=80/es=30. Subsequent visits go through the proposer.
let is_bootstrap = history.is_empty();
// Current champion (if any) is the promoted config.
let champion = deps.promotion_registry.get_current(&event.index_name).await;
let champion_trial = champion.as_ref().and_then(|p| {
history.iter().find(|t| t.id == p.trial_id).cloned()
});
let (next_config, harness_name) = if is_bootstrap {
let name = ensure_auto_harness(&event.index_name, deps).await?;
(HnswConfig::default(), name)
} else {
let Some(cfg) = propose_next_config(&history, champion_trial.as_ref()) else {
return Err("proposer returned None — search space exhausted".into());
};
// Validate bounds defensively.
if !(10..=400).contains(&cfg.ef_construction) {
return Err(format!("proposed ef_construction={} out of bounds", cfg.ef_construction));
}
if !(10..=200).contains(&cfg.ef_search) {
return Err(format!("proposed ef_search={} out of bounds", cfg.ef_search));
}
// A future refinement: remember per-index "canonical harness" on
// the index metadata. For now: latest wins.
let hname = history.last().unwrap().eval_set.clone();
(cfg, hname)
};
let mut harness_set = deps.harness_store.load_for_index(&event.index_name, &harness_name).await
.map_err(|e| format!("load harness '{harness_name}': {e}"))?;
let embeddings = deps.embedding_cache.get_or_load(&event.index_name).await
.map_err(|e| format!("embeddings: {e}"))?;
if !harness_set.ground_truth_built {
harness::compute_ground_truth(&mut harness_set, &embeddings, &deps.ai_client).await
.map_err(|e| format!("ground truth: {e}"))?;
deps.harness_store.save(&harness_set).await.ok();
}
// Build + bench.
let trial_id = Trial::new_id();
let slot = format!("{}__{}", event.index_name, trial_id);
let build = deps.hnsw_store
.build_index_with_config(&slot, (*embeddings).clone(), &next_config)
.await?;
let query_vectors: Vec<Vec<f32>> = harness_set.queries
.iter().filter_map(|q| q.query_embedding.clone()).collect();
let bench = deps.hnsw_store.bench_search(&slot, &query_vectors, harness_set.k).await?;
let mut recalls = Vec::with_capacity(harness_set.queries.len());
for (q, hits) in harness_set.queries.iter().zip(bench.retrieved.iter()) {
if let Some(gt) = &q.ground_truth {
recalls.push(harness::recall_at_k(hits, gt, harness_set.k));
}
}
let mean_recall = if recalls.is_empty() { 0.0 } else {
recalls.iter().sum::<f32>() / recalls.len() as f32
};
let mut lats = bench.latencies_us.clone();
lats.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let p = |pct: f32| -> f32 {
if lats.is_empty() { return 0.0; }
let idx = ((lats.len() as f32 - 1.0) * pct).round() as usize;
lats[idx.min(lats.len() - 1)]
};
let dims = embeddings.first().map(|e| e.vector.len()).unwrap_or(0);
let memory_bytes = (embeddings.len() * dims * std::mem::size_of::<f32>() + embeddings.len() * 128) as u64;
let note = {
let base = match &event.reason {
TriggerReason::DatasetAppended { dataset } => format!("agent: dataset_appended({dataset})"),
TriggerReason::Manual => "agent: manual".to_string(),
TriggerReason::Periodic => "agent: periodic".to_string(),
};
if is_bootstrap { format!("{base} bootstrap") } else { base }
};
let trial = Trial {
id: trial_id,
index_name: event.index_name.clone(),
eval_set: harness_set.name.clone(),
config: next_config.clone(),
metrics: TrialMetrics {
build_time_secs: build.build_time_secs,
search_latency_p50_us: p(0.50),
search_latency_p95_us: p(0.95),
search_latency_p99_us: p(0.99),
recall_at_k: mean_recall,
memory_bytes,
vectors: build.vectors,
eval_queries: harness_set.queries.len(),
brute_force_latency_us: 0.0,
},
created_at: Utc::now(),
note: Some(note),
};
deps.trial_journal.append(&trial).await.ok();
deps.hnsw_store.drop(&slot).await;
// Promotion decision: the new trial must meet recall gate AND beat
// the current champion (higher recall OR same recall + lower p50).
let promoted = if trial.metrics.recall_at_k < min_recall {
false
} else {
let beats = match &champion_trial {
None => true, // no champion yet — anything passing the gate wins
Some(c) => beats_champion(&trial, c),
};
if beats {
let entry = PromotionEntry {
config: trial.config.clone(),
trial_id: trial.id.clone(),
promoted_at: Utc::now(),
promoted_by: "agent".to_string(),
note: Some(format!(
"auto-promote: recall={:.3} p50={:.0}us (was {:.3}/{:.0}us)",
trial.metrics.recall_at_k, trial.metrics.search_latency_p50_us,
champion_trial.as_ref().map(|t| t.metrics.recall_at_k).unwrap_or(0.0),
champion_trial.as_ref().map(|t| t.metrics.search_latency_p50_us).unwrap_or(0.0),
)),
};
deps.promotion_registry.promote(&event.index_name, entry).await.is_ok()
} else {
false
}
};
Ok(CycleOutcome { trial, promoted })
}
/// Champion-beat test: strictly higher recall, OR equal recall with
/// lower p50. Same rule as autotune::pick_winner — kept consistent so the
/// agent and the synchronous autotune agree on what "better" means.
fn beats_champion(candidate: &Trial, champion: &Trial) -> bool {
if candidate.metrics.recall_at_k > champion.metrics.recall_at_k {
return true;
}
if (candidate.metrics.recall_at_k - champion.metrics.recall_at_k).abs() < 1e-4
&& candidate.metrics.search_latency_p50_us < champion.metrics.search_latency_p50_us {
return true;
}
false
}
/// Propose the next HnswConfig given trial history and the current
/// champion.
///
/// ============================================================
/// J: THIS IS YOURS TO IMPLEMENT
/// ============================================================
///
/// Inputs:
/// - `history`: every trial ever run on this index, oldest first
/// - `champion`: the currently-promoted trial, if any
///
/// Output:
/// - `Some(HnswConfig)` with the config to try next
/// - `None` if you decide there's nothing worth trying (agent logs
/// "search space exhausted" and moves on)
///
/// Hard bounds the caller enforces: ef_construction ∈ [10, 400],
/// ef_search ∈ [10, 200]. Stay inside those — configs outside get
/// rejected and count as a wasted cycle.
///
/// Design options (pick one, or mix):
///
/// 1. LOCAL REFINEMENT (exploit-heavy):
/// Sample near champion ± small delta. Converges fast, risks local
/// minima. Good for "we know roughly where the sweet spot is."
///
/// 2. ε-GREEDY (mixed):
/// With prob ε, random sample from full bounds (explore). Otherwise
/// refinement around champion (exploit). ε=0.2 is a reasonable start.
/// Good for long-running tune with no prior knowledge.
///
/// 3. COARSE→FINE (annealed):
/// First N trials: wide random. Then shrink the neighborhood around
/// champion as more trials accumulate. Mimics simulated annealing.
///
/// 4. DEDUP-AWARE:
/// Whatever strategy, skip configs already in history. Prevents the
/// agent from re-running the same (ec, es) pair twice.
///
/// A simple starter implementation is provided below (local refinement
/// + dedup). Replace with your preferred strategy.
pub fn propose_next_config(history: &[Trial], champion: Option<&Trial>) -> Option<HnswConfig> {
// ε-greedy around the champion, dedup-aware.
//
// - With probability ε (≈0.25), sample a random config from the full
// bounds. Keeps exploration alive so we don't get stuck hill-climbing
// one axis.
// - Otherwise: perturb the champion symmetrically on BOTH axes (not
// just +20 / +40 like the starter did). Prefers small moves first
// so recall stays near the current level.
// - Always skip configs already in history — no point re-running.
// - Deterministic per-history: RNG is seeded from history length so
// the same journal state always proposes the same next config.
// Makes tests + offline replay reproducible.
let base = champion
.map(|t| t.config.clone())
.or_else(|| history.last().map(|t| t.config.clone()))
.unwrap_or_default();
let tried = |ec: usize, es: usize| -> bool {
history.iter().any(|t|
t.config.ef_construction == ec && t.config.ef_search == es
)
};
let clamp = |ec: i32, es: i32| -> (usize, usize) {
(ec.clamp(10, 400) as usize, es.clamp(10, 200) as usize)
};
// Tiny xorshift — no rand crate dep. Seeded from history length so the
// proposer is deterministic for a given journal state.
let mut rng = (history.len() as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
let mut next_u = || {
rng ^= rng << 13; rng ^= rng >> 7; rng ^= rng << 17;
rng
};
for _attempt in 0..32 {
let explore = (next_u() % 100) < 25; // ε = 0.25
let (ec, es) = if explore {
clamp(10 + (next_u() % 391) as i32, 10 + (next_u() % 191) as i32)
} else {
// Symmetric perturbation — signed steps on both axes.
let dec = [-40, -20, -10, 10, 20, 40][(next_u() % 6) as usize];
let des = [-20, -10, -5, 5, 10, 20][(next_u() % 6) as usize];
clamp(base.ef_construction as i32 + dec, base.ef_search as i32 + des)
};
if !tried(ec, es) {
return Some(HnswConfig { ef_construction: ec, ef_search: es, seed: Some(42) });
}
}
None // 32 attempts all landed on duplicates — likely saturated
}
// -------- Helpers --------
/// Minimum vectors for an index to be worth auto-tuning. Below this,
/// HNSW's win over brute-force is too small to matter and trial budget
/// is better spent on bigger indexes.
const AUTOTUNE_MIN_VECTORS: usize = 1_000;
/// Find an index to poke on a periodic wake. Strategy: least-recently-tuned
/// wins — pick the index (from `IndexRegistry`, not just promoted ones)
/// whose most-recent trial is oldest, so trial budget spreads across the
/// whole portfolio. Never-trialed indexes sort first (None < Some), which
/// is exactly what we want: bootstrap them on their first visit.
///
/// Why not "most recently promoted" (the original strategy): a converged
/// index like `threat_intel_v1` sits at recall=1.0 on a minimal graph and
/// can't be improved — but it was always the freshest promotion, so the
/// agent burned every trial on it while 40+ other indexes got zero
/// attention. Rotating by last-trial-time fixes that without any explicit
/// convergence detection.
///
/// Why IndexRegistry instead of PromotionRegistry: on a fresh system only
/// a handful of indexes have ever been promoted, so promotion-based
/// picking starves new ones. IndexRegistry lists every live index the
/// platform knows about; filtering by `AUTOTUNE_MIN_VECTORS` keeps the
/// proposer off indexes where the graph config doesn't matter.
///
/// Cost: N extra journal reads per periodic tick, where N = eligible
/// indexes. Fine at dozens; if the portfolio grows into thousands, cache
/// `last_trial_at` on the IndexMeta.
async fn pick_periodic_target(deps: &AgentDeps) -> Option<String> {
let candidates: Vec<String> = deps
.index_registry
.list(None, None)
.await
.into_iter()
.filter(|m| m.chunk_count >= AUTOTUNE_MIN_VECTORS)
.map(|m| m.index_name)
.collect();
if candidates.is_empty() { return None; }
let mut scored: Vec<(String, Option<DateTime<Utc>>)> = Vec::with_capacity(candidates.len());
for name in candidates {
let last_trial_at = deps
.trial_journal
.list(&name)
.await
.ok()
.and_then(|trials| trials.into_iter().map(|t| t.created_at).max());
scored.push((name, last_trial_at));
}
scored.into_iter().min_by_key(|(_, last)| *last).map(|(n, _)| n)
}
/// On the first visit to an index, load or synthesize an eval harness so
/// bootstrap trials have something to measure recall against. Returns the
/// harness name. Synthetic harnesses sample 20 chunks and use their first
/// 200 chars as self-queries — the ground truth is computed by brute-force
/// cosine, so recall numbers are real (a config that finds the chunk in
/// top-k counts as a hit).
async fn ensure_auto_harness(index_name: &str, deps: &AgentDeps) -> Result<String, String> {
let name = format!("{index_name}_auto");
if deps.harness_store.load_for_index(index_name, &name).await.is_ok() {
return Ok(name);
}
let embeddings = deps
.embedding_cache
.get_or_load(index_name)
.await
.map_err(|e| format!("load embeddings for auto-harness: {e}"))?;
if embeddings.is_empty() {
return Err(format!("index '{index_name}' has no embeddings — cannot bootstrap"));
}
let mut eval = harness::synthetic_from_chunks(&name, index_name, &embeddings, 20, 10);
harness::compute_ground_truth(&mut eval, &embeddings, &deps.ai_client)
.await
.map_err(|e| format!("auto-harness ground truth: {e}"))?;
deps.harness_store.save(&eval).await
.map_err(|e| format!("save auto-harness: {e}"))?;
tracing::info!(
"agent: bootstrapped harness '{}' for index '{}' (20 self-queries, k=10)",
name, index_name
);
Ok(name)
}
async fn over_rate_limit(inner: &Arc<AgentInner>, cap: u32) -> bool {
let cutoff = Utc::now() - chrono::Duration::hours(1);
let ring = inner.recent_trials.lock().await;
ring.iter().filter(|t| **t >= cutoff).count() as u32 >= cap
}
async fn cooling_down(inner: &Arc<AgentInner>, _index: &str, cooldown_secs: u64) -> bool {
// Minimal impl: gate on global most-recent trial rather than per-index.
// Per-index cooldown would be easy to add — keep a HashMap<String, DateTime>
// — but for Phase 16.2 MVP, global is fine. Ollama is the shared resource.
let ring = inner.recent_trials.lock().await;
if let Some(last) = ring.back() {
let since = Utc::now().signed_duration_since(*last);
return since < chrono::Duration::seconds(cooldown_secs as i64);
}
false
}
async fn mark_recent_trial(inner: &Arc<AgentInner>) {
let mut ring = inner.recent_trials.lock().await;
ring.push_back(Utc::now());
// Keep bounded.
while ring.len() > 256 {
ring.pop_front();
}
}
async fn record_event(
inner: &Arc<AgentInner>,
kind: &str,
index: Option<&str>,
detail: String,
) {
tracing::info!("agent: {} {}{}", kind,
index.map(|i| format!("[{i}] ")).unwrap_or_default(), detail);
let mut s = inner.status.write().await;
s.last_event = Some(AgentEvent {
at: Utc::now(),
kind: kind.to_string(),
index_name: index.map(String::from),
detail,
});
}
#[cfg(test)]
mod tests {
use super::*;
fn mk_trial(ec: usize, es: usize, recall: f32, p50: f32) -> Trial {
Trial {
id: format!("t-{ec}-{es}"),
index_name: "test".into(),
eval_set: "eval".into(),
config: HnswConfig { ef_construction: ec, ef_search: es, seed: Some(42) },
metrics: TrialMetrics {
build_time_secs: 1.0,
search_latency_p50_us: p50,
search_latency_p95_us: p50 * 1.5,
search_latency_p99_us: p50 * 2.0,
recall_at_k: recall,
memory_bytes: 0, vectors: 1000, eval_queries: 10,
brute_force_latency_us: 100.0,
},
created_at: Utc::now(),
note: None,
}
}
#[test]
fn propose_skips_duplicates() {
let hist = vec![
mk_trial(80, 30, 1.0, 500.0),
mk_trial(100, 30, 1.0, 520.0), // ec+20
];
let next = propose_next_config(&hist, Some(&hist[0])).unwrap();
// ec+20 is taken, so the proposer should skip it.
assert!(next.ef_construction != 100 || next.ef_search != 30);
}
#[test]
fn beats_champion_strict_recall() {
let champ = mk_trial(80, 30, 0.95, 500.0);
let better_recall = mk_trial(80, 30, 0.99, 600.0);
let worse_recall = mk_trial(80, 30, 0.90, 100.0);
assert!(beats_champion(&better_recall, &champ));
assert!(!beats_champion(&worse_recall, &champ));
}
#[test]
fn beats_champion_same_recall_lower_latency() {
let champ = mk_trial(80, 30, 1.0, 500.0);
let faster = mk_trial(60, 30, 1.0, 400.0);
let slower = mk_trial(60, 30, 1.0, 600.0);
assert!(beats_champion(&faster, &champ));
assert!(!beats_champion(&slower, &champ));
}
}
+332
View File
@@ -0,0 +1,332 @@
//! Phase 16: Autotune — the agent loop for HNSW tuning.
//!
//! Given an index and a harness, run a grid of (ef_construction,
//! ef_search) trials, pick the Pareto winner, and promote it.
//!
//! This is deliberately NOT a sophisticated optimizer — random-grid
//! search is good enough when the config space is small (a dozen
//! combinations) and the work per trial is seconds, not hours. Think
//! of it as a "smart operator" rather than Bayesian optimization.
//!
//! Score: we want high recall and low p50 latency. The function below
//! picks the Pareto-dominating trial, breaking ties by prefer-lower-
//! build-time (cheaper to rebuild on data changes).
//!
//! Safeguards (PRD risk mitigation):
//! - Never promote a trial with recall < `min_recall` (default 0.9)
//! - Hard config bounds: ef_construction ∈ [10, 400], ef_search ∈ [10, 200]
//! - Concurrency cap: one autotune per index at a time (tracked by
//! JobTracker — caller is responsible for not firing in parallel
//! for the same index)
use chrono::Utc;
use object_store::ObjectStore;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use aibridge::client::AiClient;
use catalogd::registry::Registry as CatalogRegistry;
use crate::embedding_cache::EmbeddingCache;
use crate::harness::{self, HarnessStore};
use crate::hnsw::HnswStore;
use crate::index_registry::IndexRegistry;
use crate::jobs::JobTracker;
use crate::promotion::{PromotionEntry, PromotionRegistry};
use crate::trial::{HnswConfig, Trial, TrialJournal, TrialMetrics};
#[derive(Debug, Clone, Deserialize)]
pub struct AutotuneRequest {
pub index_name: String,
pub harness: String,
/// Minimum recall a trial must achieve to be eligible for promotion.
/// Default 0.9 — below that, the tuning is not producing useful
/// answers. Never promote a bad index.
#[serde(default = "default_min_recall")]
pub min_recall: f32,
/// Override the default config grid (rarely needed). When omitted,
/// a small curated grid runs.
#[serde(default)]
pub grid: Option<Vec<HnswConfig>>,
}
fn default_min_recall() -> f32 { 0.9 }
#[derive(Debug, Clone, Serialize)]
pub struct AutotuneResult {
pub index_name: String,
pub trials_run: usize,
pub trials_eligible: usize,
pub winner: Option<WinnerReport>,
pub promoted: bool,
pub duration_secs: f32,
pub skipped: Vec<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct WinnerReport {
pub trial_id: String,
pub config: HnswConfig,
pub recall_at_k: f32,
pub search_latency_p50_us: f32,
pub build_time_secs: f32,
}
/// Default exploration grid. Covers the practical range of settings
/// from "fast build, OK recall" to "thorough build, great recall".
fn default_grid() -> Vec<HnswConfig> {
vec![
HnswConfig { ef_construction: 20, ef_search: 30, seed: Some(42) },
HnswConfig { ef_construction: 40, ef_search: 30, seed: Some(42) },
HnswConfig { ef_construction: 80, ef_search: 30, seed: Some(42) },
HnswConfig { ef_construction: 80, ef_search: 60, seed: Some(42) },
HnswConfig { ef_construction: 160, ef_search: 30, seed: Some(42) },
]
}
/// Config-space bounds. Rejects anything outside these before running a
/// trial — defends against an agent getting creative with absurd values.
fn sanitize_grid(grid: Vec<HnswConfig>) -> (Vec<HnswConfig>, Vec<String>) {
let mut ok = Vec::new();
let mut rejected = Vec::new();
for cfg in grid {
if !(10..=400).contains(&cfg.ef_construction) {
rejected.push(format!(
"ef_construction={} out of bounds [10,400]", cfg.ef_construction,
));
continue;
}
if !(10..=200).contains(&cfg.ef_search) {
rejected.push(format!(
"ef_search={} out of bounds [10,200]", cfg.ef_search,
));
continue;
}
ok.push(cfg);
}
(ok, rejected)
}
/// Pareto winner: highest recall, then lowest p50 latency among
/// eligible trials (recall >= min_recall).
fn pick_winner(trials: &[Trial], min_recall: f32) -> Option<&Trial> {
trials
.iter()
.filter(|t| t.metrics.recall_at_k >= min_recall)
.min_by(|a, b| {
// Compare by (−recall, p50_latency) so minimum is the
// best trial: higher recall wins, then lower latency.
let a_key = (-a.metrics.recall_at_k, a.metrics.search_latency_p50_us);
let b_key = (-b.metrics.recall_at_k, b.metrics.search_latency_p50_us);
a_key.partial_cmp(&b_key).unwrap_or(std::cmp::Ordering::Equal)
})
}
/// Run the autotune loop. Synchronous — caller can wrap in tokio::spawn
/// to make it fire-and-forget. Uses the existing trial flow so every
/// autotune step lands in the trial journal.
#[allow(clippy::too_many_arguments)]
pub async fn run_autotune(
req: AutotuneRequest,
_store: &Arc<dyn ObjectStore>,
catalog: &CatalogRegistry,
ai_client: &AiClient,
embedding_cache: &EmbeddingCache,
hnsw_store: &HnswStore,
index_registry: &IndexRegistry,
trial_journal: &TrialJournal,
promotion_registry: &PromotionRegistry,
harness_store: &HarnessStore,
_job_tracker: &JobTracker,
) -> Result<AutotuneResult, String> {
let t0 = std::time::Instant::now();
// Sanity: the index has to exist.
if index_registry.get(&req.index_name).await.is_none() {
return Err(format!("index not found: {}", req.index_name));
}
let _ = catalog; // reserved for future audit emission
// Load the harness once, compute ground truth once. Harness resolves
// to the index's bucket via HarnessStore, with primary as fallback for
// pre-federation evals.
let mut harness_set = harness_store
.load_for_index(&req.index_name, &req.harness)
.await
.map_err(|e| format!("load harness: {e}"))?;
let embeddings = embedding_cache
.get_or_load(&req.index_name)
.await
.map_err(|e| format!("embeddings: {e}"))?;
if !harness_set.ground_truth_built {
harness::compute_ground_truth(&mut harness_set, &embeddings, ai_client)
.await
.map_err(|e| format!("ground truth: {e}"))?;
harness_store.save(&harness_set).await.ok();
}
let (grid, rejected) = sanitize_grid(req.grid.clone().unwrap_or_else(default_grid));
if grid.is_empty() {
return Err("all configs rejected by bounds check".into());
}
tracing::info!(
"autotune '{}': {} trial configs, min_recall={}",
req.index_name, grid.len(), req.min_recall,
);
// Run each trial through the existing pipeline so every run lands
// in the trial journal and contributes to the long-term tuning
// memory. This is the llms3.com-style "trials are data" pattern.
let mut run_trials: Vec<Trial> = Vec::with_capacity(grid.len());
for config in &grid {
let trial = run_single_trial(
&req.index_name,
&harness_set,
&embeddings,
config,
hnsw_store,
trial_journal,
ai_client,
).await;
match trial {
Ok(t) => run_trials.push(t),
Err(e) => tracing::warn!("autotune: trial ec={} es={} failed: {e}",
config.ef_construction, config.ef_search),
}
}
let trials_eligible = run_trials
.iter()
.filter(|t| t.metrics.recall_at_k >= req.min_recall)
.count();
let winner = pick_winner(&run_trials, req.min_recall);
let mut promoted = false;
let winner_report = if let Some(w) = winner {
let entry = PromotionEntry {
config: w.config.clone(),
trial_id: w.id.clone(),
promoted_at: Utc::now(),
promoted_by: "autotune".to_string(),
note: Some(format!(
"autotune winner: recall={:.3} p50={:.0}us",
w.metrics.recall_at_k, w.metrics.search_latency_p50_us,
)),
};
if let Err(e) = promotion_registry.promote(&req.index_name, entry).await {
tracing::warn!("autotune: promote failed: {e}");
} else {
promoted = true;
}
Some(WinnerReport {
trial_id: w.id.clone(),
config: w.config.clone(),
recall_at_k: w.metrics.recall_at_k,
search_latency_p50_us: w.metrics.search_latency_p50_us,
build_time_secs: w.metrics.build_time_secs,
})
} else {
tracing::warn!(
"autotune '{}': no trial met min_recall={}; not promoting",
req.index_name, req.min_recall,
);
None
};
Ok(AutotuneResult {
index_name: req.index_name,
trials_run: run_trials.len(),
trials_eligible,
winner: winner_report,
promoted,
duration_secs: t0.elapsed().as_secs_f32(),
skipped: rejected,
})
}
/// Run one trial and journal it. Returns the recorded Trial.
#[allow(clippy::too_many_arguments)]
async fn run_single_trial(
index_name: &str,
harness_set: &harness::EvalSet,
embeddings: &Arc<Vec<crate::store::StoredEmbedding>>,
config: &HnswConfig,
hnsw_store: &HnswStore,
trial_journal: &TrialJournal,
_ai_client: &AiClient,
) -> Result<Trial, String> {
let trial_id = Trial::new_id();
let slot = format!("{}__{}", index_name, trial_id);
let build = hnsw_store
.build_index_with_config(&slot, (**embeddings).clone(), config)
.await?;
let query_vectors: Vec<Vec<f32>> = harness_set
.queries
.iter()
.filter_map(|q| q.query_embedding.clone())
.collect();
let bench = hnsw_store
.bench_search(&slot, &query_vectors, harness_set.k)
.await?;
let mut recalls = Vec::with_capacity(harness_set.queries.len());
for (q, hits) in harness_set.queries.iter().zip(bench.retrieved.iter()) {
if let Some(gt) = &q.ground_truth {
recalls.push(harness::recall_at_k(hits, gt, harness_set.k));
}
}
let mean_recall = if recalls.is_empty() {
0.0
} else {
recalls.iter().sum::<f32>() / recalls.len() as f32
};
let mut lats = bench.latencies_us.clone();
lats.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let p = |pct: f32| -> f32 {
if lats.is_empty() { return 0.0; }
let idx = ((lats.len() as f32 - 1.0) * pct).round() as usize;
lats[idx.min(lats.len() - 1)]
};
let brute_us = if let Some(qv) = query_vectors.first() {
let t = std::time::Instant::now();
let _ = harness::brute_force_top_k(qv, embeddings, harness_set.k);
t.elapsed().as_micros() as f32
} else {
0.0
};
let dims = embeddings.first().map(|e| e.vector.len()).unwrap_or(0);
let memory_bytes =
(embeddings.len() * dims * std::mem::size_of::<f32>() + embeddings.len() * 128) as u64;
let trial = Trial {
id: trial_id.clone(),
index_name: index_name.to_string(),
eval_set: harness_set.name.clone(),
config: config.clone(),
metrics: TrialMetrics {
build_time_secs: build.build_time_secs,
search_latency_p50_us: p(0.50),
search_latency_p95_us: p(0.95),
search_latency_p99_us: p(0.99),
recall_at_k: mean_recall,
memory_bytes,
vectors: build.vectors,
eval_queries: harness_set.queries.len(),
brute_force_latency_us: brute_us,
},
created_at: Utc::now(),
note: Some("autotune".to_string()),
};
let _ = trial_journal.append(&trial).await;
let _ = hnsw_store.drop(&slot).await;
Ok(trial)
}
+224
View File
@@ -0,0 +1,224 @@
/// Text chunking strategies for embedding.
/// Chunks need to be small enough for the embedding model (typically <512 tokens)
/// but large enough to carry meaning.
/// A chunk of text with metadata pointing back to its source.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TextChunk {
/// Source identifier (dataset name, filename, etc.)
pub source: String,
/// Row or document ID within the source
pub doc_id: String,
/// Chunk index within the document (0, 1, 2, ...)
pub chunk_idx: u32,
/// The actual text content
pub text: String,
}
/// Round `idx` UP to the nearest UTF-8 char boundary in `s`.
/// Needed because `&str[a..b]` panics if `a` or `b` lands mid-character —
/// and any text imported from Postgres/user data will contain multi-byte
/// chars like `\u{202f}` (narrow no-break space), em-dashes, etc.
fn ceil_char_boundary(s: &str, mut idx: usize) -> usize {
if idx >= s.len() { return s.len(); }
while idx <= s.len() && !s.is_char_boundary(idx) {
idx += 1;
}
idx
}
/// Split text into overlapping chunks.
/// - `chunk_size`: target bytes per chunk (UTF-8 boundary-respecting)
/// - `overlap`: bytes of overlap between consecutive chunks
///
/// Safety contract:
/// - Never slices at a byte that isn't a UTF-8 char boundary
/// - Always advances forward (monotonic `start`) — no infinite loops on
/// pathological inputs
/// - Bounded iteration — caps total chunks at `text.len()` to prevent
/// runaway allocation on degenerate text/parameter combinations
pub fn chunk_text(
text: &str,
source: &str,
doc_id: &str,
chunk_size: usize,
overlap: usize,
) -> Vec<TextChunk> {
let text = text.trim();
if text.is_empty() {
return vec![];
}
// Short text — single chunk
if text.len() <= chunk_size {
return vec![TextChunk {
source: source.to_string(),
doc_id: doc_id.to_string(),
chunk_idx: 0,
text: text.to_string(),
}];
}
let mut chunks = Vec::new();
let mut start = 0;
let mut idx = 0u32;
// Safety cap: a well-formed chunk_size advances `start` by at least
// `chunk_size - overlap` each iteration. Even in degenerate cases we
// should never emit more chunks than there are bytes of text.
let max_chunks = text.len();
let mut iterations = 0;
while start < text.len() {
iterations += 1;
if iterations > max_chunks {
tracing::warn!(
"chunker: exceeded safety cap ({}) on doc '{}' of length {} — stopping",
max_chunks, doc_id, text.len(),
);
break;
}
// UTF-8 safe upper bound: move forward to a char boundary.
let end = ceil_char_boundary(text, (start + chunk_size).min(text.len()));
let chunk_text_slice = &text[start..end];
let actual_end = if end < text.len() {
// Look for last sentence boundary inside the slice.
if let Some(pos) = chunk_text_slice.rfind(". ") {
start + pos + 2
} else if let Some(pos) = chunk_text_slice.rfind('\n') {
start + pos + 1
} else if let Some(pos) = chunk_text_slice.rfind(' ') {
start + pos + 1
} else {
end
}
} else {
end
};
// All the rfind'd positions are at char boundaries (they match
// ASCII byte patterns), so actual_end should already be safe.
// Defensive belt-and-suspenders: round up anyway.
let actual_end = ceil_char_boundary(text, actual_end);
let chunk = text[start..actual_end].trim();
if !chunk.is_empty() {
chunks.push(TextChunk {
source: source.to_string(),
doc_id: doc_id.to_string(),
chunk_idx: idx,
text: chunk.to_string(),
});
idx += 1;
}
// Advance with overlap — must strictly progress to avoid infinite
// loops on degenerate inputs (e.g. chunk_size ≤ overlap, or
// boundary-finding returning a position at `start`).
if actual_end >= text.len() {
break;
}
let tentative = if actual_end > overlap { actual_end - overlap } else { actual_end };
let next_start = ceil_char_boundary(text, tentative);
if next_start <= start {
// No progress — force forward one char to guarantee termination.
let mut forced = start + 1;
while forced < text.len() && !text.is_char_boundary(forced) {
forced += 1;
}
start = forced;
} else {
start = next_start;
}
}
chunks
}
/// Chunk a dataset's text column. Returns all chunks from all rows.
pub fn chunk_column(
source: &str,
doc_ids: &[String],
texts: &[String],
chunk_size: usize,
overlap: usize,
) -> Vec<TextChunk> {
let mut all_chunks = Vec::new();
for (doc_id, text) in doc_ids.iter().zip(texts.iter()) {
let chunks = chunk_text(text, source, doc_id, chunk_size, overlap);
all_chunks.extend(chunks);
}
all_chunks
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn short_text_single_chunk() {
let chunks = chunk_text("Hello world", "test", "1", 500, 50);
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].text, "Hello world");
}
#[test]
fn long_text_multiple_chunks() {
let text = "First sentence. Second sentence. Third sentence. Fourth sentence. Fifth sentence. Sixth sentence. Seventh sentence. Eighth sentence.";
let chunks = chunk_text(text, "test", "1", 50, 10);
assert!(chunks.len() > 1);
// All text should be covered
for chunk in &chunks {
assert!(!chunk.text.is_empty());
}
}
#[test]
fn chunk_preserves_source() {
let chunks = chunk_text("Some text here", "candidates", "CAND-001", 500, 50);
assert_eq!(chunks[0].source, "candidates");
assert_eq!(chunks[0].doc_id, "CAND-001");
}
#[test]
fn handles_multibyte_utf8_at_chunk_boundary() {
// Regression: multi-byte unicode (narrow no-break space \u{202f},
// 3 bytes) landing inside a chunk-size window used to panic with
// "byte index N is not a char boundary". Found on pg-imported
// data that contained editorial unicode punctuation.
let text = "**Claim 1**\u{202f}\u{2013} *Each large language model has its own personality \
and way of interacting.*\u{202f}\u{2013}\u{202f}Verdict:\u{202f}VERIFIED. \
Reasoning: widely accepted observation that different models exhibit distinct response \
styles stemming from variations in training data.".repeat(5);
let chunks = chunk_text(&text, "test", "unicode-1", 500, 50);
// Must not panic, must produce at least one chunk.
assert!(!chunks.is_empty());
// All chunks must be valid UTF-8 substrings of the original.
for chunk in &chunks {
assert!(!chunk.text.is_empty());
assert!(text.contains(&chunk.text));
}
}
#[test]
fn no_infinite_loop_on_no_spaces() {
// Long text with no spaces, sentences, or newlines — chunker must
// still advance via the hard `end` fallback.
let text = "a".repeat(5000);
let chunks = chunk_text(&text, "test", "noSpace", 500, 50);
assert!(!chunks.is_empty());
// Expected ~5000/450 chunks (with overlap=50 → step=450)
assert!(chunks.len() < 50, "chunk count {} seems excessive", chunks.len());
}
#[test]
fn no_infinite_loop_on_degenerate_params() {
// chunk_size equals overlap — previously would have made zero
// forward progress. Our safety forced-advance must catch this.
let text = "a b c d e f g h i j k l m n o p q r s t u v w x y z".repeat(20);
let chunks = chunk_text(&text, "test", "degen", 100, 100);
assert!(!chunks.is_empty());
assert!(chunks.len() <= text.len(), "runaway chunk count {}", chunks.len());
}
}
+166
View File
@@ -0,0 +1,166 @@
// Phase 45 slice 3 — context7 bridge client for doc-drift detection.
//
// Calls the Bun context7 bridge on :3900 (mcp-server/context7_bridge.ts)
// which itself wraps context7's public API. For each DocRef on a
// playbook, queries /docs/:tool/diff?since=<snippet_hash> and parses
// the response.
//
// Kept deliberately thin: no caching here (bridge caches for 5 min),
// no retry logic (transient fail = report as "unknown" drift status,
// don't flag). The handler decides whether to flag the playbook based
// on whether ANY tool came back drifted=true.
use serde::Deserialize;
use std::time::Duration;
use crate::playbook_memory::DocRef;
const DEFAULT_BRIDGE_URL: &str = "http://localhost:3900";
const CALL_TIMEOUT_SECS: u64 = 15;
#[derive(Debug, Clone)]
pub struct DriftCheck {
pub tool: String,
pub version_seen: String,
pub outcome: DriftOutcome,
}
#[derive(Debug, Clone)]
pub enum DriftOutcome {
/// Bridge reports drift (current snippet_hash != previous).
Drifted { current_snippet_hash: String, source_url: Option<String> },
/// Bridge reports no drift — the stored snippet_hash still matches
/// the current context7 docs.
Unchanged,
/// Bridge unreachable, 404 on the tool, or returned unparseable
/// data. We deliberately don't flag on these — a down bridge
/// shouldn't silently mark every playbook drift-flagged.
Unknown { reason: String },
}
#[derive(Debug, Clone, Deserialize)]
struct BridgeDiffResponse {
drifted: bool,
current_snippet_hash: Option<String>,
source_url: Option<String>,
}
pub struct DriftCheckerConfig {
pub bridge_url: String,
}
impl Default for DriftCheckerConfig {
fn default() -> Self {
Self {
bridge_url: std::env::var("LH_BRIDGE_URL")
.unwrap_or_else(|_| DEFAULT_BRIDGE_URL.to_string()),
}
}
}
/// For every doc_ref, ask the bridge whether it drifted against the
/// recorded snippet_hash. Returns per-tool outcomes.
pub async fn check_all_refs(
cfg: &DriftCheckerConfig,
doc_refs: &[DocRef],
) -> Vec<DriftCheck> {
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(CALL_TIMEOUT_SECS))
.build()
.expect("reqwest client build");
let mut out = Vec::with_capacity(doc_refs.len());
for r in doc_refs {
let hash = r.snippet_hash.as_deref().unwrap_or("");
if hash.is_empty() {
// No hash to compare against — can't detect drift. Report
// unknown so the caller isn't forced to flag.
out.push(DriftCheck {
tool: r.tool.clone(),
version_seen: r.version_seen.clone(),
outcome: DriftOutcome::Unknown {
reason: "no snippet_hash recorded on doc_ref".into(),
},
});
continue;
}
let url = format!(
"{}/docs/{}/diff?since={}",
cfg.bridge_url.trim_end_matches('/'),
urlencoding_minimal(&r.tool),
urlencoding_minimal(hash),
);
let outcome = match client.get(&url).send().await {
Err(e) => DriftOutcome::Unknown {
reason: format!("bridge unreachable: {}", e),
},
Ok(resp) => {
if resp.status() == reqwest::StatusCode::NOT_FOUND {
DriftOutcome::Unknown {
reason: format!("bridge 404 — no context7 library for tool '{}'", r.tool),
}
} else if !resp.status().is_success() {
let status = resp.status();
DriftOutcome::Unknown {
reason: format!("bridge {}: {}", status, resp.text().await.unwrap_or_default()),
}
} else {
match resp.json::<BridgeDiffResponse>().await {
Err(e) => DriftOutcome::Unknown {
reason: format!("bridge response parse: {}", e),
},
Ok(body) => {
if body.drifted {
DriftOutcome::Drifted {
current_snippet_hash: body.current_snippet_hash.unwrap_or_default(),
source_url: body.source_url,
}
} else {
DriftOutcome::Unchanged
}
}
}
}
}
};
out.push(DriftCheck {
tool: r.tool.clone(),
version_seen: r.version_seen.clone(),
outcome,
});
}
out
}
/// Minimal URL path+query encoder for ASCII tool names. We're not
/// pulling in the `urlencoding` crate for this single use — tool
/// names are short alphanumeric (docker/terraform/react), snippet
/// hashes are hex. Handles space + `/` + `?` + `&` defensively so
/// unusual tool names don't corrupt the URL.
fn urlencoding_minimal(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for b in s.bytes() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => out.push(b as char),
_ => out.push_str(&format!("%{:02X}", b)),
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn urlencoding_handles_ascii_safe_chars_passthrough() {
assert_eq!(urlencoding_minimal("docker"), "docker");
assert_eq!(urlencoding_minimal("next.js"), "next.js");
}
#[test]
fn urlencoding_encodes_slash_and_space() {
assert_eq!(urlencoding_minimal("foo/bar"), "foo%2Fbar");
assert_eq!(urlencoding_minimal("foo bar"), "foo%20bar");
}
}
+98
View File
@@ -0,0 +1,98 @@
/// In-memory cache for StoredEmbedding vectors.
///
/// Rationale: loading 100K embeddings from Parquet takes ~2-5s. When an AI agent
/// iterates on HNSW parameters, each trial would repeat that cost. The cache
/// pins embeddings in memory so trials reuse them.
///
/// This is a pure performance layer — the Parquet file is still the source of
/// truth (ADR-008). Eviction is safe; worst case is one slow reload.
use object_store::ObjectStore;
use serde::Serialize;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::store::{self, StoredEmbedding};
#[derive(Clone)]
pub struct EmbeddingCache {
store: Arc<dyn ObjectStore>,
cache: Arc<RwLock<HashMap<String, Arc<Vec<StoredEmbedding>>>>>,
}
#[derive(Debug, Clone, Serialize)]
pub struct CacheEntry {
pub index_name: String,
pub vectors: usize,
pub dimensions: usize,
pub memory_bytes: u64,
}
#[derive(Debug, Clone, Serialize)]
pub struct CacheStats {
pub entries: Vec<CacheEntry>,
pub total_memory_bytes: u64,
}
impl EmbeddingCache {
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
Self {
store,
cache: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Return cached embeddings, loading from object storage on first request.
pub async fn get_or_load(
&self,
index_name: &str,
) -> Result<Arc<Vec<StoredEmbedding>>, String> {
if let Some(cached) = self.cache.read().await.get(index_name) {
return Ok(cached.clone());
}
// Load under a write lock so concurrent callers only hit disk once.
let mut guard = self.cache.write().await;
if let Some(cached) = guard.get(index_name) {
return Ok(cached.clone());
}
tracing::info!("embedding_cache: loading '{index_name}' from object storage");
let t0 = std::time::Instant::now();
let loaded = store::load_embeddings(&self.store, index_name).await?;
let n = loaded.len();
let arc = Arc::new(loaded);
guard.insert(index_name.to_string(), arc.clone());
tracing::info!(
"embedding_cache: loaded '{index_name}' — {n} vectors in {:.2}s",
t0.elapsed().as_secs_f32()
);
Ok(arc)
}
pub async fn evict(&self, index_name: &str) -> bool {
self.cache.write().await.remove(index_name).is_some()
}
pub async fn stats(&self) -> CacheStats {
let cache = self.cache.read().await;
let mut entries = Vec::with_capacity(cache.len());
let mut total: u64 = 0;
for (name, embs) in cache.iter() {
let dims = embs.first().map(|e| e.vector.len()).unwrap_or(0);
// Rough estimate: vector data + chunk_text + metadata overhead.
let vec_bytes = (embs.len() * dims * std::mem::size_of::<f32>()) as u64;
let text_bytes: u64 = embs.iter().map(|e| e.chunk_text.len() as u64).sum();
let overhead = (embs.len() * 128) as u64; // strings + struct overhead
let mem = vec_bytes + text_bytes + overhead;
total += mem;
entries.push(CacheEntry {
index_name: name.clone(),
vectors: embs.len(),
dimensions: dims,
memory_bytes: mem,
});
}
CacheStats { entries, total_memory_bytes: total }
}
}
+387
View File
@@ -0,0 +1,387 @@
/// Eval harness for HNSW tuning.
///
/// An EvalSet is a named collection of queries against a specific vector index.
/// Each query has a query_text and (optionally pre-computed) ground_truth
/// doc_ids — the "correct" top-k results that brute-force cosine returns.
/// HNSW trials are scored by how well they recreate that top-k (recall@k).
///
/// Storage: `_hnsw_evals/{name}.json` as a single JSON document. Small.
use chrono::{DateTime, Utc};
use object_store::ObjectStore;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::sync::Arc;
use aibridge::client::{AiClient, EmbedRequest};
use storaged::ops;
use storaged::registry::BucketRegistry;
use crate::index_registry::IndexRegistry;
use crate::store::StoredEmbedding;
/// A single eval query with optional pre-computed ground truth.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EvalQuery {
pub id: String,
pub query_text: String,
/// Ordered list of doc_ids that brute-force returns as top-k.
/// `None` means the ground truth hasn't been computed yet; `compute_ground_truth`
/// fills it in against the current embedding set.
#[serde(default)]
pub ground_truth: Option<Vec<String>>,
/// Optional pre-computed query embedding. Filling this avoids re-embedding
/// the same query on every trial.
#[serde(default)]
pub query_embedding: Option<Vec<f32>>,
}
/// A named eval set tied to a specific vector index.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EvalSet {
pub name: String,
pub index_name: String,
pub k: usize, // top-k used for recall calculation
pub queries: Vec<EvalQuery>,
pub created_at: DateTime<Utc>,
pub ground_truth_built: bool,
}
impl EvalSet {
pub fn new(name: &str, index_name: &str, k: usize) -> Self {
Self {
name: name.to_string(),
index_name: index_name.to_string(),
k,
queries: Vec::new(),
created_at: Utc::now(),
ground_truth_built: false,
}
}
pub fn storage_key(&self) -> String {
format!("_hnsw_evals/{}.json", self.name)
}
pub async fn save(&self, store: &Arc<dyn ObjectStore>) -> Result<(), String> {
let json = serde_json::to_vec_pretty(self).map_err(|e| e.to_string())?;
ops::put(store, &self.storage_key(), json.into()).await
}
pub async fn load(store: &Arc<dyn ObjectStore>, name: &str) -> Result<Self, String> {
let key = format!("_hnsw_evals/{}.json", name);
let data = ops::get(store, &key).await?;
serde_json::from_slice(&data).map_err(|e| format!("parse eval set: {e}"))
}
pub async fn list(store: &Arc<dyn ObjectStore>) -> Result<Vec<String>, String> {
let keys = ops::list(store, Some("_hnsw_evals/")).await?;
Ok(keys
.into_iter()
.filter(|k| k.ends_with(".json"))
.map(|k| {
k.trim_start_matches("_hnsw_evals/")
.trim_end_matches(".json")
.to_string()
})
.collect())
}
}
/// Federation-aware wrapper around EvalSet persistence. Mirrors the
/// `TrialJournal` / `PromotionRegistry` pattern: harness files colocate
/// with their index's bucket (looked up via `IndexMeta.bucket`), falling
/// back to `primary` for indexes the registry has never seen. Legacy
/// harnesses predating federation remain discoverable — lookups transparently
/// try the resolved bucket first, then `primary` as a fallback. Cross-bucket
/// listing dedupes across every registered bucket so `GET /hnsw/evals`
/// returns a complete picture.
#[derive(Clone)]
pub struct HarnessStore {
buckets: Arc<BucketRegistry>,
index_registry: IndexRegistry,
}
impl HarnessStore {
pub fn new(buckets: Arc<BucketRegistry>, index_registry: IndexRegistry) -> Self {
Self { buckets, index_registry }
}
/// Resolve which bucket holds this index's eval artifacts. Indexes the
/// registry has never heard of fall through to `primary`.
async fn bucket_for_index(&self, index_name: &str) -> String {
self.index_registry
.get(index_name)
.await
.map(|m| m.bucket)
.unwrap_or_else(|| "primary".to_string())
}
/// Save to the bucket that owns `eval.index_name`. Writes under the
/// standard `_hnsw_evals/{name}.json` prefix of the resolved bucket.
pub async fn save(&self, eval: &EvalSet) -> Result<(), String> {
let bucket = self.bucket_for_index(&eval.index_name).await;
let store = self.buckets.get(&bucket)?;
eval.save(&store).await
}
/// Load a harness by name, given the index it belongs to. Tries the
/// index's bucket first; if the file is absent AND the resolved bucket
/// isn't `primary`, falls through to `primary` so pre-federation evals
/// remain reachable without migration.
pub async fn load_for_index(
&self,
index_name: &str,
harness_name: &str,
) -> Result<EvalSet, String> {
let bucket = self.bucket_for_index(index_name).await;
let primary_store = self.buckets.get("primary")?;
let store = self.buckets.get(&bucket)?;
match EvalSet::load(&store, harness_name).await {
Ok(e) => Ok(e),
Err(e) if bucket != "primary" => EvalSet::load(&primary_store, harness_name)
.await
.map_err(|primary_err| format!("{bucket}: {e}; primary fallback: {primary_err}")),
Err(e) => Err(e),
}
}
/// Find a harness by name without knowing which index it belongs to —
/// used by `GET /hnsw/evals/{name}`. Scans every registered bucket;
/// first hit wins. Primary is searched first so pre-federation evals
/// with the same name as a federated one resolve deterministically.
pub async fn get_any(&self, harness_name: &str) -> Result<EvalSet, String> {
let bucket_infos = self.buckets.list().await;
let mut ordered: Vec<String> = bucket_infos.iter().map(|b| b.name.clone()).collect();
ordered.sort_by_key(|n| if n == "primary" { 0 } else { 1 });
let mut last_err = None;
for b in ordered {
let store = match self.buckets.get(&b) {
Ok(s) => s,
Err(e) => { last_err = Some(e); continue; }
};
match EvalSet::load(&store, harness_name).await {
Ok(e) => return Ok(e),
Err(e) => { last_err = Some(e); }
}
}
Err(last_err.unwrap_or_else(|| format!("no buckets registered for eval '{harness_name}'")))
}
/// Union of every harness name across every registered bucket.
/// Duplicates (same name in multiple buckets — pathological but
/// possible after a manual migration) are collapsed.
pub async fn list_all(&self) -> Vec<String> {
let mut all: HashSet<String> = HashSet::new();
for b in self.buckets.list().await {
let store = match self.buckets.get(&b.name) {
Ok(s) => s,
Err(_) => continue,
};
if let Ok(names) = EvalSet::list(&store).await {
all.extend(names);
}
}
let mut out: Vec<String> = all.into_iter().collect();
out.sort();
out
}
}
/// Cosine similarity for two same-length f32 slices.
fn cosine(a: &[f32], b: &[f32]) -> f32 {
let mut dot = 0.0f32;
let mut na = 0.0f32;
let mut nb = 0.0f32;
for i in 0..a.len() {
dot += a[i] * b[i];
na += a[i] * a[i];
nb += b[i] * b[i];
}
if na == 0.0 || nb == 0.0 {
return 0.0;
}
dot / (na.sqrt() * nb.sqrt())
}
/// Brute-force top-k search for a single query against all embeddings.
/// This is the ground-truth oracle that HNSW trials must approximate.
pub fn brute_force_top_k(
query: &[f32],
embeddings: &[StoredEmbedding],
k: usize,
) -> Vec<String> {
let mut scored: Vec<(f32, usize)> = embeddings
.iter()
.enumerate()
.map(|(i, e)| (cosine(query, &e.vector), i))
.collect();
// Partial sort — we only need top-k.
scored.sort_unstable_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
scored
.into_iter()
.take(k)
.map(|(_, i)| embeddings[i].doc_id.clone())
.collect()
}
/// Recall@k — fraction of ground-truth doc_ids that appear in predicted.
pub fn recall_at_k(predicted: &[String], ground_truth: &[String], k: usize) -> f32 {
if ground_truth.is_empty() || k == 0 {
return 0.0;
}
// Set-intersection recall@k. Previous implementation counted duplicates
// in `predicted` (a corpus with repeated chunks — e.g. cached LLM
// responses — returns the same doc_id multiple times via HNSW), which
// inflated recall above 1.0 and poisoned promotion decisions.
let gt_set: std::collections::HashSet<&String> =
ground_truth.iter().take(k).collect();
let pred_set: std::collections::HashSet<&String> =
predicted.iter().take(k).collect();
let hits = pred_set.intersection(&gt_set).count();
hits as f32 / gt_set.len() as f32
}
/// Populate query_embedding and ground_truth for every query that lacks them.
pub async fn compute_ground_truth(
eval: &mut EvalSet,
embeddings: &[StoredEmbedding],
ai_client: &AiClient,
) -> Result<(), String> {
let need_embed: Vec<(usize, String)> = eval
.queries
.iter()
.enumerate()
.filter(|(_, q)| q.query_embedding.is_none())
.map(|(i, q)| (i, q.query_text.clone()))
.collect();
if !need_embed.is_empty() {
// Embed in one batch to keep things simple; for very large eval sets
// we'd batch this in chunks of 32.
let texts: Vec<String> = need_embed.iter().map(|(_, t)| t.clone()).collect();
let resp = ai_client
.embed(EmbedRequest { texts, model: None })
.await
.map_err(|e| format!("embed queries: {e}"))?;
for ((idx, _), vec) in need_embed.iter().zip(resp.embeddings.iter()) {
let v: Vec<f32> = vec.iter().map(|&x| x as f32).collect();
eval.queries[*idx].query_embedding = Some(v);
}
}
for q in eval.queries.iter_mut() {
if q.ground_truth.is_some() {
continue;
}
let emb = q.query_embedding.as_ref().ok_or("missing embedding")?;
q.ground_truth = Some(brute_force_top_k(emb, embeddings, eval.k));
}
eval.ground_truth_built = true;
Ok(())
}
/// Auto-generate a synthetic eval set by sampling every Nth chunk's text as
/// its own query. Useful for a quick-start eval when the user doesn't have
/// real natural-language queries yet.
pub fn synthetic_from_chunks(
eval_name: &str,
index_name: &str,
embeddings: &[StoredEmbedding],
sample_count: usize,
k: usize,
) -> EvalSet {
let n = embeddings.len();
let sample_count = sample_count.min(n);
let stride = (n / sample_count.max(1)).max(1);
let mut queries = Vec::with_capacity(sample_count);
for i in 0..sample_count {
let idx = (i * stride).min(n - 1);
let chunk = &embeddings[idx];
// Use the first ~200 chars of the chunk as the "query" — it should find
// itself and nearby chunks as top results.
let query_text: String = chunk.chunk_text.chars().take(200).collect();
queries.push(EvalQuery {
id: format!("syn-{}", i),
query_text,
ground_truth: None,
query_embedding: None,
});
}
EvalSet {
name: eval_name.to_string(),
index_name: index_name.to_string(),
k,
queries,
created_at: Utc::now(),
ground_truth_built: false,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn s(v: &[&str]) -> Vec<String> { v.iter().map(|x| x.to_string()).collect() }
#[test]
fn recall_empty_ground_truth_is_zero() {
assert_eq!(recall_at_k(&s(&["a", "b"]), &[], 10), 0.0);
}
#[test]
fn recall_k_zero_is_zero() {
assert_eq!(recall_at_k(&s(&["a"]), &s(&["a"]), 0), 0.0);
}
#[test]
fn recall_perfect_match_equals_one() {
let pred = s(&["a", "b", "c"]);
let gt = s(&["a", "b", "c"]);
assert!((recall_at_k(&pred, &gt, 3) - 1.0).abs() < 1e-6);
}
#[test]
fn recall_half_match() {
let pred = s(&["a", "b", "x", "y"]);
let gt = s(&["a", "b", "c", "d"]);
assert!((recall_at_k(&pred, &gt, 4) - 0.5).abs() < 1e-6);
}
#[test]
fn recall_duplicates_in_predicted_do_not_inflate() {
// Regression guard: the previous implementation counted each
// duplicate in `predicted` separately, which could push recall
// above 1.0 on corpora with repeated chunks (cached responses etc).
// Set-intersection semantics keep it bounded in [0, 1].
let pred = s(&["a", "a", "a", "a"]);
let gt = s(&["a", "b", "c", "d"]);
let r = recall_at_k(&pred, &gt, 4);
assert!(r <= 1.0, "recall {r} must not exceed 1.0");
// One unique match out of four in gt = 0.25.
assert!((r - 0.25).abs() < 1e-6);
}
#[test]
fn recall_duplicates_in_ground_truth_handled() {
// gt with dupes reduces effective |gt|; matching one hits all.
let pred = s(&["x"]);
let gt = s(&["x", "x", "x"]);
let r = recall_at_k(&pred, &gt, 3);
assert!(r <= 1.0);
assert!((r - 1.0).abs() < 1e-6); // |pred ∩ gt| / |gt_set| = 1/1
}
#[test]
fn recall_respects_k_bound() {
// k=2 means only the first 2 of pred and gt count toward the set.
let pred = s(&["a", "b", "c", "d"]);
let gt = s(&["a", "b", "c", "d"]);
let r = recall_at_k(&pred, &gt, 2);
assert!((r - 1.0).abs() < 1e-6);
}
}
+229
View File
@@ -0,0 +1,229 @@
/// HNSW (Hierarchical Navigable Small World) index for fast approximate nearest neighbor search.
/// Wraps instant-distance to provide <50ms search over 100K+ vectors.
/// Falls back to brute-force for small datasets.
use instant_distance::{Builder, HnswMap, Search, Point};
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::store::StoredEmbedding;
use crate::trial::HnswConfig;
/// A vector point for HNSW — wraps f32 slice with cosine distance.
#[derive(Clone)]
struct VecPoint(Vec<f32>);
impl Point for VecPoint {
fn distance(&self, other: &Self) -> f32 {
// Cosine distance = 1 - cosine_similarity
let dot: f32 = self.0.iter().zip(other.0.iter()).map(|(a, b)| a * b).sum();
let norm_a: f32 = self.0.iter().map(|x| x * x).sum::<f32>().sqrt();
let norm_b: f32 = other.0.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm_a == 0.0 || norm_b == 0.0 {
return 1.0;
}
1.0 - (dot / (norm_a * norm_b))
}
}
/// HNSW search result.
#[derive(Debug, Clone, serde::Serialize)]
pub struct HnswResult {
pub source: String,
pub doc_id: String,
pub chunk_idx: u32,
pub chunk_text: String,
pub score: f32, // cosine similarity (1.0 = identical)
}
/// An HNSW index built from stored embeddings.
pub struct HnswIndex {
map: HnswMap<VecPoint, usize>, // value is index into metadata vec
metadata: Vec<EmbeddingMeta>,
}
#[derive(Clone)]
struct EmbeddingMeta {
source: String,
doc_id: String,
chunk_idx: u32,
chunk_text: String,
}
/// Shared, thread-safe index store.
#[derive(Clone)]
pub struct HnswStore {
indexes: Arc<RwLock<std::collections::HashMap<String, Arc<HnswIndex>>>>,
}
impl HnswStore {
pub fn new() -> Self {
Self {
indexes: Arc::new(RwLock::new(std::collections::HashMap::new())),
}
}
/// Build an HNSW index from stored embeddings with default config.
pub async fn build_index(
&self,
index_name: &str,
embeddings: Vec<StoredEmbedding>,
) -> Result<BuildStats, String> {
self.build_index_with_config(index_name, embeddings, &HnswConfig::default()).await
}
/// Build an HNSW index from stored embeddings with explicit config.
/// Used by the trial system — each trial calls this with different params.
pub async fn build_index_with_config(
&self,
index_name: &str,
embeddings: Vec<StoredEmbedding>,
config: &HnswConfig,
) -> Result<BuildStats, String> {
let n = embeddings.len();
if n == 0 {
return Err("no embeddings to index".into());
}
tracing::info!(
"building HNSW '{}': {} vectors, ef_construction={} ef_search={} seed={:?}",
index_name, n, config.ef_construction, config.ef_search, config.seed,
);
let start = std::time::Instant::now();
// Separate points and metadata
let mut points = Vec::with_capacity(n);
let mut metadata = Vec::with_capacity(n);
let mut values = Vec::with_capacity(n);
for (i, emb) in embeddings.into_iter().enumerate() {
points.push(VecPoint(emb.vector));
metadata.push(EmbeddingMeta {
source: emb.source,
doc_id: emb.doc_id,
chunk_idx: emb.chunk_idx,
chunk_text: emb.chunk_text,
});
values.push(i);
}
// Build HNSW — the expensive part
let mut builder = Builder::default()
.ef_construction(config.ef_construction)
.ef_search(config.ef_search);
if let Some(seed) = config.seed {
builder = builder.seed(seed);
}
let map = builder.build(points, values);
let build_time = start.elapsed().as_secs_f32();
tracing::info!("HNSW '{}' built: {} vectors in {:.1}s", index_name, n, build_time);
let index = Arc::new(HnswIndex { map, metadata });
self.indexes.write().await.insert(index_name.to_string(), index);
Ok(BuildStats {
index_name: index_name.to_string(),
vectors: n,
build_time_secs: build_time,
})
}
/// Run a batch of search queries and return raw per-query latencies in microseconds.
/// Also returns the retrieved doc_ids per query (for recall calculation).
pub async fn bench_search(
&self,
index_name: &str,
query_vectors: &[Vec<f32>],
top_k: usize,
) -> Result<BenchResult, String> {
let indexes = self.indexes.read().await;
let index = indexes
.get(index_name)
.ok_or_else(|| format!("HNSW index not found: {index_name}"))?
.clone();
drop(indexes);
let mut latencies_us = Vec::with_capacity(query_vectors.len());
let mut retrieved: Vec<Vec<String>> = Vec::with_capacity(query_vectors.len());
for qv in query_vectors {
let query_point = VecPoint(qv.clone());
let t0 = std::time::Instant::now();
let mut search = Search::default();
let results = index.map.search(&query_point, &mut search);
let ids: Vec<String> = results
.take(top_k)
.map(|item| {
let meta_idx = *item.value;
index.metadata[meta_idx].doc_id.clone()
})
.collect();
latencies_us.push(t0.elapsed().as_micros() as f32);
retrieved.push(ids);
}
Ok(BenchResult { latencies_us, retrieved })
}
/// Search an HNSW index. Returns approximate nearest neighbors.
pub async fn search(
&self,
index_name: &str,
query: &[f32],
top_k: usize,
) -> Result<Vec<HnswResult>, String> {
let indexes = self.indexes.read().await;
let index = indexes.get(index_name)
.ok_or_else(|| format!("HNSW index not found: {index_name}"))?;
let query_point = VecPoint(query.to_vec());
let mut search = Search::default();
let results = index.map.search(&query_point, &mut search);
let mut out = Vec::with_capacity(top_k);
for item in results.take(top_k) {
let meta_idx = *item.value;
let meta = &index.metadata[meta_idx];
// Convert distance back to similarity
let similarity = 1.0 - item.distance;
out.push(HnswResult {
source: meta.source.clone(),
doc_id: meta.doc_id.clone(),
chunk_idx: meta.chunk_idx,
chunk_text: meta.chunk_text.clone(),
score: similarity,
});
}
Ok(out)
}
/// Check if an index exists.
pub async fn has_index(&self, name: &str) -> bool {
self.indexes.read().await.contains_key(name)
}
/// List all loaded indexes.
pub async fn list(&self) -> Vec<String> {
self.indexes.read().await.keys().cloned().collect()
}
/// Drop an index from memory.
pub async fn drop(&self, name: &str) -> bool {
self.indexes.write().await.remove(name).is_some()
}
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct BuildStats {
pub index_name: String,
pub vectors: usize,
pub build_time_secs: f32,
}
#[derive(Debug, Clone)]
pub struct BenchResult {
pub latencies_us: Vec<f32>,
pub retrieved: Vec<Vec<String>>,
}
+281
View File
@@ -0,0 +1,281 @@
/// Vector index registry — tracks all indexes with model versioning.
/// Each index knows which model created it, enabling:
/// - Multi-version indexes (same data, different models, coexist)
/// - Incremental re-embed (only new/changed docs on model upgrade)
/// - A/B search comparison between model versions
use chrono::{DateTime, Utc};
use object_store::ObjectStore;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use storaged::ops;
/// Metadata for a vector index.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexMeta {
pub index_name: String,
pub source: String, // dataset this was built from
pub model_name: String, // "nomic-embed-text"
pub model_version: String, // "latest" or specific version
pub dimensions: u32, // 768
pub chunk_count: usize,
pub doc_count: usize,
pub chunk_size: usize,
pub overlap: usize,
pub storage_key: String, // "vectors/resumes_v1_nomic.parquet"
pub created_at: DateTime<Utc>,
pub build_time_secs: f32,
pub chunks_per_sec: f32,
/// Federation layer 2: which bucket holds this index's artifacts
/// (trial journal + promotion file). Defaults to "primary" for
/// pre-federation indexes — the serde default keeps old metadata
/// files readable without migration.
#[serde(default = "default_bucket")]
pub bucket: String,
/// ADR-019: which physical backend stores this index. `Parquet`
/// means storage_key points at our binary-blob Parquet file;
/// `Lance` means it points at a Lance dataset directory.
#[serde(default)]
pub vector_backend: shared::types::VectorBackend,
/// ADR-020: prefix prepended to doc_ids during embedding. If set,
/// hybrid search strips this prefix to match against SQL primary keys.
/// None = doc_ids ARE the raw primary keys (no stripping needed).
/// Existing indexes: "W-", "CAND-", "W500K-", etc.
#[serde(default)]
pub id_prefix: Option<String>,
/// PRD 11.3 — when this index was last searched against. `None` =
/// never used since registration (or pre-field-existed metadata).
/// Incremental re-embed walks this to skip cold indexes.
/// Scrum iter 11 flagged the missing field as a UnitMismatch
/// because callers were reading `created_at` as a proxy for
/// liveness, which conflated "built" with "used."
#[serde(default)]
pub last_used: Option<DateTime<Utc>>,
/// PRD 11.3 — SHA-256 of (sorted source file list + chunk_size +
/// overlap + model_version). Lets incremental re-embed detect
/// "no change since last build" without scanning the source
/// Parquet. None = signature not computed yet (pre-existing
/// indexes before this field landed).
#[serde(default)]
pub build_signature: Option<String>,
}
fn default_bucket() -> String { "primary".to_string() }
/// Registry of all vector indexes.
#[derive(Clone)]
pub struct IndexRegistry {
indexes: Arc<RwLock<HashMap<String, IndexMeta>>>,
store: Arc<dyn ObjectStore>,
}
impl IndexRegistry {
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
Self {
indexes: Arc::new(RwLock::new(HashMap::new())),
store,
}
}
/// Rebuild from persisted index metadata on startup.
pub async fn rebuild(&self) -> Result<usize, String> {
let keys = ops::list(&self.store, Some("vectors/meta/")).await?;
let mut reg = self.indexes.write().await;
reg.clear();
for key in &keys {
if !key.ends_with(".json") { continue; }
let data = ops::get(&self.store, key).await?;
match serde_json::from_slice::<IndexMeta>(&data) {
Ok(meta) => { reg.insert(meta.index_name.clone(), meta); }
Err(e) => tracing::warn!("failed to load index meta {key}: {e}"),
}
}
let count = reg.len();
if count > 0 {
tracing::info!("loaded {count} vector index metadata entries");
}
Ok(count)
}
/// Register a new index.
pub async fn register(&self, meta: IndexMeta) -> Result<(), String> {
let key = format!("vectors/meta/{}.json", meta.index_name);
let json = serde_json::to_vec_pretty(&meta).map_err(|e| e.to_string())?;
ops::put(&self.store, &key, json.into()).await?;
self.indexes.write().await.insert(meta.index_name.clone(), meta);
Ok(())
}
/// Get metadata for an index.
pub async fn get(&self, index_name: &str) -> Option<IndexMeta> {
self.indexes.read().await.get(index_name).cloned()
}
/// List all indexes, optionally filtered by source or model.
pub async fn list(&self, source: Option<&str>, model: Option<&str>) -> Vec<IndexMeta> {
self.indexes.read().await.values()
.filter(|m| source.map_or(true, |s| m.source == s))
.filter(|m| model.map_or(true, |mo| m.model_name == mo))
.cloned()
.collect()
}
/// Find all versions of an index for a given source dataset.
/// Returns indexes sorted by creation time (newest first).
pub async fn versions_for_source(&self, source: &str) -> Vec<IndexMeta> {
let mut versions: Vec<IndexMeta> = self.indexes.read().await.values()
.filter(|m| m.source == source)
.cloned()
.collect();
versions.sort_by(|a, b| b.created_at.cmp(&a.created_at));
versions
}
/// Delete an index (metadata only — vector Parquet stays for safety).
pub async fn delete(&self, index_name: &str) -> Result<(), String> {
let key = format!("vectors/meta/{index_name}.json");
ops::delete(&self.store, &key).await?;
self.indexes.write().await.remove(index_name);
Ok(())
}
/// Stamp `last_used = now()` on an index. Search handlers call this
/// on every hit so incremental re-embed (PRD 11.3) can tell live
/// indexes from cold ones. Silently no-ops if the index is unknown
/// — callers get best-effort behavior, not a 500 on a missing row.
pub async fn touch_used(&self, index_name: &str) {
if let Some(m) = self.indexes.write().await.get_mut(index_name) {
m.last_used = Some(Utc::now());
}
}
}
/// Compute a stable build_signature for PRD 11.3 incremental re-embed.
/// Hashes (sorted source file list, chunk_size, overlap, model_version)
/// so a caller can ask "has anything we built from changed?" without
/// re-scanning the source parquet. Same inputs always produce the
/// same hash.
pub fn compute_build_signature(
source_files: &[impl AsRef<str>],
chunk_size: usize,
overlap: usize,
model_version: &str,
) -> String {
use sha2::{Digest, Sha256};
let mut sorted: Vec<&str> = source_files.iter().map(|s| s.as_ref()).collect();
sorted.sort();
let mut hasher = Sha256::new();
for f in &sorted {
hasher.update(f.as_bytes());
hasher.update(b"\n");
}
hasher.update(chunk_size.to_le_bytes());
hasher.update(overlap.to_le_bytes());
hasher.update(model_version.as_bytes());
format!("{:x}", hasher.finalize())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_signature_is_deterministic() {
let sig1 = compute_build_signature(&["a.parquet", "b.parquet"], 800, 80, "v1");
let sig2 = compute_build_signature(&["a.parquet", "b.parquet"], 800, 80, "v1");
assert_eq!(sig1, sig2, "same inputs → same hash");
}
#[test]
fn build_signature_order_invariant() {
// Files get sorted internally so caller's order doesn't matter.
let sig_a = compute_build_signature(&["a.parquet", "b.parquet"], 800, 80, "v1");
let sig_b = compute_build_signature(&["b.parquet", "a.parquet"], 800, 80, "v1");
assert_eq!(sig_a, sig_b, "file list order must not affect hash");
}
#[test]
fn build_signature_changes_on_chunk_param() {
let sig_a = compute_build_signature(&["a.parquet"], 800, 80, "v1");
let sig_b = compute_build_signature(&["a.parquet"], 900, 80, "v1");
assert_ne!(sig_a, sig_b, "chunk_size change → different hash");
}
#[test]
fn build_signature_changes_on_model_version() {
let sig_a = compute_build_signature(&["a.parquet"], 800, 80, "v1");
let sig_b = compute_build_signature(&["a.parquet"], 800, 80, "v2");
assert_ne!(sig_a, sig_b, "model version change → different hash");
}
#[tokio::test]
async fn touch_used_updates_last_used() {
use object_store::memory::InMemory;
let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
let reg = IndexRegistry::new(store);
let meta = IndexMeta {
index_name: "test".into(),
source: "s".into(),
model_name: "m".into(),
model_version: "v1".into(),
dimensions: 768,
chunk_count: 0,
doc_count: 0,
chunk_size: 800,
overlap: 80,
storage_key: "k".into(),
created_at: Utc::now(),
build_time_secs: 0.0,
chunks_per_sec: 0.0,
bucket: "primary".into(),
vector_backend: Default::default(),
id_prefix: None,
last_used: None,
build_signature: None,
};
reg.register(meta).await.unwrap();
assert!(reg.get("test").await.unwrap().last_used.is_none());
reg.touch_used("test").await;
assert!(reg.get("test").await.unwrap().last_used.is_some());
}
#[tokio::test]
async fn touch_used_is_noop_on_missing_index() {
use object_store::memory::InMemory;
let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
let reg = IndexRegistry::new(store);
// No panic — unknown index just doesn't get touched.
reg.touch_used("nonexistent").await;
}
#[test]
fn index_meta_deserializes_without_new_fields_backcompat() {
// Pre-field-existence metadata files on disk must still load.
// Critical — we have ~40 .json meta files under vectors/meta/
// that predate these fields.
let json = r#"{
"index_name": "resumes_v1",
"source": "resumes",
"model_name": "nomic-embed-text",
"model_version": "latest",
"dimensions": 768,
"chunk_count": 100,
"doc_count": 10,
"chunk_size": 800,
"overlap": 80,
"storage_key": "vectors/resumes_v1.parquet",
"created_at": "2026-04-20T00:00:00Z",
"build_time_secs": 1.0,
"chunks_per_sec": 100.0
}"#;
let meta: IndexMeta = serde_json::from_str(json).expect("must deserialize pre-field meta");
assert!(meta.last_used.is_none());
assert!(meta.build_signature.is_none());
assert_eq!(meta.bucket, "primary");
}
}
+148
View File
@@ -0,0 +1,148 @@
/// Background job system for long-running embedding tasks.
/// POST /vectors/index returns a job_id immediately.
/// GET /vectors/jobs/{id} returns progress.
/// Embedding runs in background via tokio::spawn.
use serde::Serialize;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum JobType {
Embed,
ProfileActivation,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum JobStatus {
Running,
Completed,
Failed,
}
#[derive(Debug, Clone, Serialize)]
pub struct Job {
pub id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub job_type: Option<JobType>,
pub status: JobStatus,
pub index_name: Option<String>,
pub profile_id: Option<String>,
pub total_chunks: usize,
#[serde(rename = "processed")]
pub processed_alias: usize,
pub progress_pct: f32,
pub storage_key: Option<String>,
pub error: Option<String>,
pub started_at: String,
pub completed_at: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub result: Option<serde_json::Value>,
}
impl Job {
pub fn new_embed(index_name: &str, total_chunks: usize) -> Self {
Self {
id: format!("job-{}", chrono::Utc::now().timestamp_millis()),
job_type: Some(JobType::Embed),
status: JobStatus::Running,
index_name: Some(index_name.to_string()),
profile_id: None,
total_chunks,
processed_alias: 0,
progress_pct: 0.0,
storage_key: None,
error: None,
started_at: chrono::Utc::now().to_rfc3339(),
completed_at: None,
result: None,
}
}
pub fn new_profile_activation(profile_id: &str) -> Self {
Self {
id: format!("job-{}", chrono::Utc::now().timestamp_millis()),
job_type: Some(JobType::ProfileActivation),
status: JobStatus::Running,
index_name: None,
profile_id: Some(profile_id.to_string()),
total_chunks: 0,
processed_alias: 0,
progress_pct: 0.0,
storage_key: None,
error: None,
started_at: chrono::Utc::now().to_rfc3339(),
completed_at: None,
result: None,
}
}
}
#[derive(Clone)]
pub struct JobTracker {
jobs: Arc<RwLock<HashMap<String, Job>>>,
}
impl JobTracker {
pub fn new() -> Self {
Self {
jobs: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn create_embed(&self, index_name: &str, total_chunks: usize) -> String {
let job = Job::new_embed(index_name, total_chunks);
let id = job.id.clone();
self.jobs.write().await.insert(id.clone(), job);
id
}
pub async fn create_profile_activation(&self, profile_id: &str) -> String {
let job = Job::new_profile_activation(profile_id);
let id = job.id.clone();
self.jobs.write().await.insert(id.clone(), job);
id
}
pub async fn update_embed_progress(&self, id: &str, embedded: usize, _rate: f32) {
let mut jobs = self.jobs.write().await;
if let Some(job) = jobs.get_mut(id) {
job.processed_alias = embedded;
job.progress_pct = if job.total_chunks > 0 {
(embedded as f32 / job.total_chunks as f32) * 100.0
} else {
0.0
};
}
}
pub async fn complete(&self, id: &str, result: Option<serde_json::Value>) {
let mut jobs = self.jobs.write().await;
if let Some(job) = jobs.get_mut(id) {
job.status = JobStatus::Completed;
job.progress_pct = 100.0;
job.completed_at = Some(chrono::Utc::now().to_rfc3339());
job.result = result;
}
}
pub async fn fail(&self, id: &str, error: String) {
let mut jobs = self.jobs.write().await;
if let Some(job) = jobs.get_mut(id) {
job.status = JobStatus::Failed;
job.error = Some(error);
job.completed_at = Some(chrono::Utc::now().to_rfc3339());
}
}
pub async fn get(&self, id: &str) -> Option<Job> {
self.jobs.read().await.get(id).cloned()
}
pub async fn list(&self) -> Vec<Job> {
self.jobs.read().await.values().cloned().collect()
}
}
+131
View File
@@ -0,0 +1,131 @@
//! ADR-019 hybrid: routing layer for the Lance vector backend.
//!
//! Holds a `LanceRegistry` which maps an index name to its
//! `LanceVectorStore` instance, lazy-creating on first touch. Path
//! resolution: the index's bucket gives us a bucket root; we append
//! `lance/{index_name}` and use that as the dataset URI. For local
//! buckets that means a directory under the bucket's root.
//!
//! S3 buckets: in principle Lance accepts s3://... URIs, but the
//! firewall crate (`vectord-lance`) doesn't pull the S3 feature in by
//! default. When we promote profile buckets onto S3 in a future phase,
//! enable that feature and update `lance_uri_for` accordingly.
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::RwLock;
use storaged::registry::BucketRegistry;
use vectord_lance::LanceVectorStore;
use crate::index_registry::IndexRegistry;
/// Convert a bucket+index pair into the URI Lance should use as the
/// dataset path. Supports both local (filesystem) and S3 buckets.
///
/// **Local buckets:** path resolution mirrors lakehouse.toml's convention.
/// Returns an absolute filesystem path.
///
/// **S3 buckets:** returns `s3://{s3_bucket}/lance/{index_name}`. Lance's
/// internal object_store crate reads `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY`
/// / `AWS_ENDPOINT` from environment (or the S3 feature's default chain).
/// For MinIO: set `AWS_ENDPOINT=http://localhost:9000` and
/// `AWS_ALLOW_HTTP=true` before starting the gateway.
///
/// Refuses unknown buckets so a typo doesn't silently land Lance data
/// in a directory / prefix the rest of the system can't see.
pub fn lance_uri_for(
buckets: &BucketRegistry,
bucket: &str,
index_name: &str,
) -> Result<String, String> {
if !buckets.contains(bucket) {
return Err(format!("bucket '{bucket}' not registered"));
}
// Check if this bucket is S3-backed by looking for a bucket config
// with backend="s3". BucketRegistry exposes backend type through
// the list() info, but that's async. The simpler signal: if the
// bucket name matches one we know is S3 (configured via lakehouse.toml
// with backend="s3"), use the s3:// URI scheme.
//
// For the synchronous path, we check a naming convention: buckets
// whose name starts with "s3:" are treated as S3 targets. The rest
// of the name is the S3 bucket name. Convention-based, explicit,
// no async needed.
//
// Additionally, any bucket registered with backend="s3" in the
// config will have its BucketConfig.bucket field set — that's the
// actual S3 bucket name. We can't access BucketConfig synchronously
// from the current registry API, so for now the naming convention
// is the primary signal.
if bucket.starts_with("s3:") {
let s3_bucket = &bucket["s3:".len()..];
return Ok(format!("s3://{s3_bucket}/lance/{index_name}"));
}
// Local path resolution.
let root: PathBuf = match bucket {
"primary" => PathBuf::from("./data"),
"rescue" => PathBuf::from("./data/_rescue"),
"testing" => PathBuf::from("./data/_testing"),
b if b.starts_with("profile:") => {
let safe = b.replace(':', "_");
PathBuf::from(format!("./data/_profiles/{safe}"))
}
b => PathBuf::from(format!("./data/_buckets/{}", b.replace(':', "_"))),
};
let _ = std::fs::create_dir_all(root.join("lance"));
let abs = match std::fs::canonicalize(&root) {
Ok(p) => p.join("lance").join(index_name),
Err(_) => root.join("lance").join(index_name),
};
Ok(abs.to_string_lossy().to_string())
}
/// Lookup-by-index registry of `LanceVectorStore` handles. One handle
/// per index, lazy-created on first call. Cheap to keep alive — it's
/// just a path string + handle to Lance's metadata cache.
#[derive(Clone)]
pub struct LanceRegistry {
buckets: Arc<BucketRegistry>,
indexes: IndexRegistry,
stores: Arc<RwLock<HashMap<String, LanceVectorStore>>>,
}
impl LanceRegistry {
pub fn new(buckets: Arc<BucketRegistry>, indexes: IndexRegistry) -> Self {
Self {
buckets,
indexes,
stores: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Get (or lazy-create) the LanceVectorStore for an index.
/// Resolves bucket → URI via the registry; backend stays whatever
/// IndexMeta says (no enforcement here — caller decides whether to
/// route to Lance based on profile/index backend choice).
pub async fn store_for(&self, index_name: &str) -> Result<LanceVectorStore, String> {
if let Some(s) = self.stores.read().await.get(index_name) {
return Ok(s.clone());
}
let bucket = self.indexes
.get(index_name).await
.map(|m| m.bucket)
.unwrap_or_else(|| "primary".to_string());
let uri = lance_uri_for(&self.buckets, &bucket, index_name)?;
let store = LanceVectorStore::new(uri);
self.stores.write().await.insert(index_name.to_string(), store.clone());
Ok(store)
}
/// For freshly-created indexes that don't have IndexMeta yet — caller
/// supplies the bucket explicitly.
pub async fn store_for_new(&self, index_name: &str, bucket: &str) -> Result<LanceVectorStore, String> {
let uri = lance_uri_for(&self.buckets, bucket, index_name)?;
let store = LanceVectorStore::new(uri);
self.stores.write().await.insert(index_name.to_string(), store.clone());
Ok(store)
}
}
+21
View File
@@ -0,0 +1,21 @@
pub mod agent;
pub mod autotune;
pub mod lance_backend;
pub mod chunker;
pub mod embedding_cache;
pub mod harness;
pub mod hnsw;
pub mod index_registry;
pub mod jobs;
pub mod activation;
pub mod playbook_memory;
pub mod pathway_memory;
pub mod doc_drift;
pub mod promotion;
pub mod refresh;
pub mod store;
pub mod search;
pub mod rag;
pub mod supervisor;
pub mod service;
pub mod trial;
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+227
View File
@@ -0,0 +1,227 @@
//! Phase 16: Promoted HNSW configs — the "active generation" pointer.
//!
//! An index's HNSW config used at build time normally defaults to the
//! system-wide default (`HnswConfig::default()`). An operator or the
//! autotune agent can *promote* a specific trial's config — subsequent
//! HNSW builds against that index use the promoted config instead.
//!
//! Every promotion is history-tracked so `rollback` can revert. The
//! history file lives at `primary://_hnsw_promotions/{index_name}.json`
//! and is small (< few KB) so we rewrite it on every promotion rather
//! than append-log.
//!
//! Not included here:
//! - Atomic graph rebuild on promote — promotion only updates the sticky
//! default. Next activation (or search that triggers lazy build) picks
//! up the new config. That's "zero-downtime swap after build" which is
//! what ADR-019 actually claimed; an instant-swap requires a
//! pre-built graph pool which we don't have yet.
//! - Agent loop — lives in `vectord::autotune`.
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use storaged::ops;
use storaged::registry::BucketRegistry;
use tokio::sync::RwLock;
use crate::index_registry::IndexRegistry;
use crate::trial::HnswConfig;
const PROMOTION_PREFIX: &str = "_hnsw_promotions";
/// One promotion record. The `trial_id` is the origin of the config —
/// lets operators trace back "why was this config picked?" to the exact
/// trial in the trial journal.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PromotionEntry {
pub config: HnswConfig,
pub trial_id: String,
pub promoted_at: DateTime<Utc>,
#[serde(default)]
pub promoted_by: String,
#[serde(default)]
pub note: Option<String>,
}
/// Serialized form of an index's promotion history.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PromotionFile {
pub index_name: String,
pub current: Option<PromotionEntry>,
#[serde(default)]
pub history: Vec<PromotionEntry>,
}
#[derive(Clone)]
pub struct PromotionRegistry {
buckets: Arc<BucketRegistry>,
index_registry: IndexRegistry,
cache: Arc<RwLock<HashMap<String, PromotionFile>>>,
}
impl PromotionRegistry {
pub fn new(buckets: Arc<BucketRegistry>, index_registry: IndexRegistry) -> Self {
Self {
buckets,
index_registry,
cache: Arc::new(RwLock::new(HashMap::new())),
}
}
fn key(index_name: &str) -> String {
// Sanitize for object-store safety.
let safe: String = index_name
.chars()
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' })
.collect();
format!("{PROMOTION_PREFIX}/{safe}.json")
}
/// Resolve which bucket's store holds this index's promotion file.
/// Same rules as TrialJournal::bucket_for — follows IndexMeta.bucket,
/// defaults to primary when metadata is missing.
async fn store_for(&self, index_name: &str) -> Result<Arc<dyn object_store::ObjectStore>, String> {
let bucket = self.index_registry
.get(index_name)
.await
.map(|m| m.bucket)
.unwrap_or_else(|| "primary".to_string());
self.buckets.get(&bucket)
}
/// Load (and cache) the promotion file for an index.
pub async fn load(&self, index_name: &str) -> Result<PromotionFile, String> {
if let Some(cached) = self.cache.read().await.get(index_name) {
return Ok(cached.clone());
}
let store = self.store_for(index_name).await?;
let key = Self::key(index_name);
let file = match ops::get(&store, &key).await {
Ok(bytes) => serde_json::from_slice::<PromotionFile>(&bytes)
.map_err(|e| format!("parse promotion file: {e}"))?,
Err(_) => PromotionFile {
index_name: index_name.to_string(),
current: None,
history: Vec::new(),
},
};
self.cache.write().await.insert(index_name.to_string(), file.clone());
Ok(file)
}
/// Promote a config to the active slot. Pushes the current promotion
/// (if any) onto the history stack. Persists before returning — the
/// config is durable by the time this call completes.
pub async fn promote(
&self,
index_name: &str,
entry: PromotionEntry,
) -> Result<PromotionFile, String> {
let mut file = self.load(index_name).await?;
if let Some(prior) = file.current.take() {
file.history.push(prior);
// Cap history to something sensible so this file doesn't grow
// unbounded. 50 entries = 50 promotions — way more than any
// sane workflow needs.
const HISTORY_CAP: usize = 50;
if file.history.len() > HISTORY_CAP {
let drop = file.history.len() - HISTORY_CAP;
file.history.drain(0..drop);
}
}
// Bind `entry` ref-captured for the log line below so the log
// doesn't double-unwrap file.current — entry is Some-by-construction
// at the function boundary; past versions reached in via
// `.as_ref().unwrap()` twice, which compiled but would panic if
// the construction above ever changed.
file.current = Some(entry);
file.index_name = index_name.to_string();
let store = self.store_for(index_name).await?;
let key = Self::key(index_name);
let json = serde_json::to_vec_pretty(&file).map_err(|e| e.to_string())?;
ops::put(&store, &key, json.into()).await?;
self.cache.write().await.insert(index_name.to_string(), file.clone());
if let Some(cur) = &file.current {
tracing::info!(
"promoted '{}' to config {:?} (trial={})",
index_name, cur.config, cur.trial_id,
);
}
Ok(file)
}
/// Pop the latest promotion back onto the current slot (if any
/// history exists). If current is set but history is empty, the
/// current promotion is cleared — the index falls back to defaults.
pub async fn rollback(&self, index_name: &str) -> Result<PromotionFile, String> {
let mut file = self.load(index_name).await?;
match file.history.pop() {
Some(prev) => {
file.current = Some(prev);
}
None => {
if file.current.is_none() {
return Err(format!("no promotion to rollback for '{index_name}'"));
}
file.current = None;
}
}
let store = self.store_for(index_name).await?;
let key = Self::key(index_name);
let json = serde_json::to_vec_pretty(&file).map_err(|e| e.to_string())?;
ops::put(&store, &key, json.into()).await?;
self.cache.write().await.insert(index_name.to_string(), file.clone());
tracing::info!("rolled back promotion for '{}'", index_name);
Ok(file)
}
/// Get the currently-promoted config (if any). Callers use this to
/// pick the right HnswConfig at build time.
pub async fn get_current(&self, index_name: &str) -> Option<PromotionEntry> {
self.load(index_name).await.ok().and_then(|f| f.current)
}
/// Convenience: return the promoted config or the provided default.
pub async fn config_or(&self, index_name: &str, default: HnswConfig) -> HnswConfig {
match self.get_current(index_name).await {
Some(entry) => entry.config,
None => default,
}
}
/// List every index that has a promotion recorded (for operator UI).
///
/// Federation: scans EVERY registered bucket for promotion files.
/// Per-profile buckets each have their own `_hnsw_promotions/` so we
/// aggregate across them. Dedups by index_name — if the same index
/// somehow has promotion files in multiple buckets, the one from the
/// bucket recorded in IndexMeta wins.
pub async fn list_all(&self) -> Result<Vec<PromotionFile>, String> {
let bucket_infos = self.buckets.list().await;
let mut by_name: HashMap<String, PromotionFile> = HashMap::new();
for b in &bucket_infos {
let store = match self.buckets.get(&b.name) {
Ok(s) => s,
Err(_) => continue,
};
let keys = ops::list(&store, Some(&format!("{PROMOTION_PREFIX}/")))
.await.unwrap_or_default();
for key in keys {
if !key.ends_with(".json") { continue; }
let bytes = match ops::get(&store, &key).await {
Ok(b) => b,
Err(_) => continue,
};
if let Ok(f) = serde_json::from_slice::<PromotionFile>(&bytes) {
by_name.insert(f.index_name.clone(), f);
}
}
}
Ok(by_name.into_values().collect())
}
}
+183
View File
@@ -0,0 +1,183 @@
/// RAG pipeline: question → embed → search → rerank → generate answer.
///
/// The rerank step (added 2026-04-17) uses the LLM as a cross-encoder
/// between retrieval and generation. This catches cases where the
/// embedding model scores documents as similar but the content isn't
/// actually relevant to the question — a known weakness of small
/// general-purpose embed models on domain-specific text.
use object_store::ObjectStore;
use std::sync::Arc;
use aibridge::client::{AiClient, EmbedRequest, GenerateRequest};
use aibridge::continuation::{generate_continuable, ContinuableOpts, ResponseShape};
use crate::search::{self, SearchResult};
use crate::store;
/// Cross-encoder rerank: ask the LLM to re-sort retrieved chunks by
/// relevance. Falls back to the original order if the model returns
/// garbage (which happens ~5% of the time with 7B models).
async fn rerank(
question: &str,
mut results: Vec<SearchResult>,
ai_client: &AiClient,
) -> Vec<SearchResult> {
if results.len() <= 1 {
return results;
}
let chunk_list: String = results.iter().enumerate().map(|(i, r)| {
let text: String = r.chunk_text.chars().take(200).collect();
format!("[{i}] {text}")
}).collect::<Vec<_>>().join("\n");
let resp = ai_client.generate(GenerateRequest {
prompt: format!(
"Rank these text chunks by relevance to the question.\n\
Return ONLY a comma-separated list of indices, most relevant first.\n\n\
Question: {question}\n\n\
Chunks:\n{chunk_list}\n\n\
Ranking:"
),
model: None,
system: None,
temperature: Some(0.0),
max_tokens: Some(50),
// Reranker returns a comma-separated int list — pure structured
// output, zero benefit from hidden reasoning. Opt out to avoid
// the empty-response failure mode Phase 21 catalogued.
think: Some(false),
}).await;
match resp {
Ok(gen_resp) => {
let text = gen_resp.text.trim();
let indices: Vec<usize> = text
.split(|c: char| c == ',' || c.is_whitespace())
.filter_map(|s| s.trim().parse::<usize>().ok())
.filter(|&i| i < results.len())
.collect();
if indices.is_empty() {
tracing::debug!("reranker returned unparseable output: {text}");
return results;
}
let mut reranked: Vec<SearchResult> = Vec::with_capacity(results.len());
let mut used = vec![false; results.len()];
for &i in &indices {
if !used[i] {
reranked.push(results[i].clone());
used[i] = true;
}
}
// Append any the model didn't mention (preserves all results).
for (i, r) in results.drain(..).enumerate() {
if !used[i] {
reranked.push(r);
}
}
tracing::info!("reranker reordered {}/{} chunks", indices.len(), reranked.len());
reranked
}
Err(e) => {
tracing::debug!("reranker failed (using original order): {e}");
results
}
}
}
/// Full RAG answer with provenance.
#[derive(Debug, Clone, serde::Serialize)]
pub struct RagResponse {
pub answer: String,
pub model: String,
pub sources: Vec<SearchResult>,
pub tokens_generated: Option<u64>,
}
/// Execute full RAG: embed question → search index → retrieve context → generate answer.
pub async fn query(
question: &str,
index_name: &str,
top_k: usize,
object_store: &Arc<dyn ObjectStore>,
ai_client: &AiClient,
) -> Result<RagResponse, String> {
// 1. Embed the question
tracing::info!("RAG: embedding question");
let embed_resp = ai_client.embed(EmbedRequest {
texts: vec![question.to_string()],
model: None,
}).await?;
if embed_resp.embeddings.is_empty() {
return Err("no embedding returned for question".into());
}
let query_vec: Vec<f32> = embed_resp.embeddings[0].iter().map(|&x| x as f32).collect();
// 2. Load index and search
tracing::info!("RAG: searching index '{index_name}'");
let embeddings = store::load_embeddings(object_store, index_name).await?;
let results = search::search(&query_vec, &embeddings, top_k);
if results.is_empty() {
return Ok(RagResponse {
answer: "No relevant information found.".into(),
model: String::new(),
sources: vec![],
tokens_generated: None,
});
}
// 3. Rerank: ask the LLM to sort retrieved chunks by relevance.
// This cross-encoder step catches cases where embedding similarity
// is high but semantic relevance is low — common with small embed
// models on domain-specific text.
let results = rerank(question, results, ai_client).await;
// 4. Build context from (reranked) top chunks
let context: String = results.iter().enumerate().map(|(i, r)| {
format!("[{}] (source: {}, doc: {}) {}", i + 1, r.source, r.doc_id, r.chunk_text)
}).collect::<Vec<_>>().join("\n\n");
// 5. Generate answer
tracing::info!("RAG: generating answer from {} chunks", results.len());
let prompt = format!(
"You are a helpful assistant answering questions about a staffing database.\n\n\
Use ONLY the following retrieved records to answer. Be specific — cite names,\n\
numbers, skills, and cities from the records. If the context doesn't contain\n\
enough information to fully answer, say what you can and note what's missing.\n\
Cite sources by their number [1], [2], etc.\n\n\
Context:\n{context}\n\n\
Question: {question}\n\n\
Answer:"
);
// Route the answer call through Phase 21's generate_continuable so
// a thinking-model empty-response or a mid-JSON truncation self-
// repairs instead of silently returning half an answer. Shape is
// Text (the answer is prose, not JSON), think is Some(false) to
// opt out of hidden reasoning on the hot path. This is the first
// production caller of the Phase 21 primitives — see audit finding
// "Phase 21 Rust primitives are wired but not CALLED by any
// production surface" from 2026-04-21.
let mut cont_opts = ContinuableOpts::new("qwen2.5:latest");
cont_opts.max_tokens = Some(512);
cont_opts.temperature = Some(0.2);
cont_opts.shape = ResponseShape::Text;
cont_opts.think = Some(false);
let outcome = generate_continuable(ai_client, &prompt, &cont_opts).await?;
Ok(RagResponse {
answer: outcome.text.trim().to_string(),
// generate_continuable doesn't surface the model name (sidecar
// echoes whatever Ollama loaded). Use the configured tier model
// for now; if RAG needs to report the actual resolved model,
// the runner can add a post-call ps probe later.
model: "qwen2.5:latest".to_string(),
sources: results,
tokens_generated: None,
})
}
+315
View File
@@ -0,0 +1,315 @@
//! Phase C: Decoupled embedding refresh.
//!
//! When a dataset's row-level data changes, the vector index it feeds is
//! stale. Historically we coupled ingest and embedding — writing new rows
//! also re-embedded every row, which is fine at 500 rows but blows up at
//! 100K+. The llms3.com architecture calls out "asynchronous vector
//! refresh cycles independent from transactional mutations" as the right
//! pattern.
//!
//! This module implements the refresh side. The ingest path marks
//! embeddings stale (see catalogd::registry::mark_embeddings_stale); this
//! code clears that staleness by embedding only rows whose `doc_id` isn't
//! already in the existing vector index.
//!
//! Scope — MVP:
//! - Reads the dataset's Parquet, extracts (doc_id, text) pairs from named
//! columns
//! - Loads existing embeddings via EmbeddingCache
//! - Filters to rows whose doc_id is NOT in the existing set
//! - Chunks, embeds via Ollama, appends to the index parquet
//! - Clears stale flag on success
//!
//! Not in MVP:
//! - UPDATE semantics (same doc_id, new content) — would need content-hash
//! comparison per row
//! - Large-scale resilience (batching with checkpoints like the
//! supervisor) — MVP does it inline
//! - Lance backend — ADR-019 makes this straightforward later; MVP stays
//! on Parquet sidecar indexes
use std::collections::HashSet;
use std::sync::Arc;
use aibridge::client::{AiClient, EmbedRequest};
use arrow::array::{Array, Int32Array, Int64Array, StringArray};
use catalogd::registry::Registry;
use object_store::ObjectStore;
use crate::chunker::{self, TextChunk};
use crate::embedding_cache::EmbeddingCache;
use crate::index_registry::IndexRegistry;
use crate::store::{self, StoredEmbedding};
#[derive(Debug, Clone, serde::Deserialize)]
pub struct RefreshRequest {
pub index_name: String,
/// Column with the document id (row identity).
pub id_column: String,
/// Column with the text to embed.
pub text_column: String,
/// Chunk size (chars). Defaults to 500.
#[serde(default)]
pub chunk_size: Option<usize>,
/// Overlap (chars). Defaults to 50.
#[serde(default)]
pub overlap: Option<usize>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct RefreshResult {
pub index_name: String,
pub dataset_name: String,
pub pre_existing_docs: usize,
pub dataset_docs: usize,
pub new_docs_embedded: usize,
pub new_chunks_embedded: usize,
pub total_embeddings_after: usize,
pub duration_secs: f32,
pub stale_cleared: bool,
}
/// Full refresh pipeline. Takes dataset_name as URL param, body has the
/// column selectors.
pub async fn refresh_index(
dataset_name: &str,
req: &RefreshRequest,
store_: &Arc<dyn ObjectStore>,
registry: &Registry,
ai_client: &AiClient,
embedding_cache: &EmbeddingCache,
index_registry: &IndexRegistry,
) -> Result<RefreshResult, String> {
let t0 = std::time::Instant::now();
// 1. Find the dataset manifest, pull its object storage key
let manifest = registry
.get_by_name(dataset_name)
.await
.ok_or_else(|| format!("dataset not found: {dataset_name}"))?;
if manifest.objects.is_empty() {
return Err(format!("dataset '{dataset_name}' has no object references"));
}
// 2. Read dataset rows — extract (doc_id, text) pairs from the
// specified columns. Multi-object datasets: concat across all.
let mut doc_id_to_text: Vec<(String, String)> = Vec::new();
for obj in &manifest.objects {
let data = storaged::ops::get(store_, &obj.key).await
.map_err(|e| format!("read {}: {e}", obj.key))?;
let (schema, batches) = shared::arrow_helpers::parquet_to_record_batches(&data)
.map_err(|e| format!("parse {}: {e}", obj.key))?;
let id_idx = schema.index_of(&req.id_column)
.map_err(|_| format!("id column '{}' not in dataset schema", req.id_column))?;
let text_idx = schema.index_of(&req.text_column)
.map_err(|_| format!("text column '{}' not in dataset schema", req.text_column))?;
for batch in &batches {
let text_col = batch
.column(text_idx)
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| format!("text column '{}' is not Utf8", req.text_column))?;
// Accept Utf8, Int32, or Int64 as id — that covers CSV-derived
// and Postgres-imported schemas without forcing upstream casts.
let id_reader: Box<dyn Fn(usize) -> Option<String>> = {
let col = batch.column(id_idx);
if let Some(s) = col.as_any().downcast_ref::<StringArray>() {
let s = s.clone();
Box::new(move |row| if s.is_null(row) { None } else { Some(s.value(row).to_string()) })
} else if let Some(a) = col.as_any().downcast_ref::<Int32Array>() {
let a = a.clone();
Box::new(move |row| if a.is_null(row) { None } else { Some(a.value(row).to_string()) })
} else if let Some(a) = col.as_any().downcast_ref::<Int64Array>() {
let a = a.clone();
Box::new(move |row| if a.is_null(row) { None } else { Some(a.value(row).to_string()) })
} else {
return Err(format!(
"id column '{}' must be Utf8, Int32, or Int64 — got {}",
req.id_column,
col.data_type(),
));
}
};
for row in 0..batch.num_rows() {
if text_col.is_null(row) { continue; }
let Some(id) = id_reader(row) else { continue; };
let text = text_col.value(row).to_string();
if text.trim().is_empty() { continue; }
doc_id_to_text.push((id, text));
}
}
}
let dataset_docs = doc_id_to_text.len();
tracing::info!("refresh '{}': dataset has {dataset_docs} rows", dataset_name);
// 3. Load existing embeddings (empty if no index yet)
let existing: Vec<StoredEmbedding> = match embedding_cache.get_or_load(&req.index_name).await {
Ok(arc) => arc.as_ref().clone(),
Err(_) => Vec::new(), // first-time index build
};
let pre_existing_docs: HashSet<String> = existing
.iter()
.map(|e| e.doc_id.clone())
.collect();
let pre_existing_count = pre_existing_docs.len();
// 4. Delta — rows whose doc_id isn't already indexed
let new_rows: Vec<(String, String)> = doc_id_to_text
.into_iter()
.filter(|(id, _)| !pre_existing_docs.contains(id))
.collect();
let new_docs = new_rows.len();
if new_docs == 0 {
tracing::info!("refresh '{}': no new docs to embed", dataset_name);
// Even on no-op, make sure index metadata is registered so
// downstream discovery (profile activation, A/B search) sees
// indexes built by earlier refresh calls that predated the
// auto-register behavior.
let _ = try_update_index_meta(index_registry, &req.index_name, existing.len()).await;
registry.clear_embeddings_stale(dataset_name).await?;
return Ok(RefreshResult {
index_name: req.index_name.clone(),
dataset_name: dataset_name.to_string(),
pre_existing_docs: pre_existing_count,
dataset_docs,
new_docs_embedded: 0,
new_chunks_embedded: 0,
total_embeddings_after: existing.len(),
duration_secs: t0.elapsed().as_secs_f32(),
stale_cleared: true,
});
}
// 5. Chunk the new rows
let chunk_size = req.chunk_size.unwrap_or(500);
let overlap = req.overlap.unwrap_or(50);
let doc_ids: Vec<String> = new_rows.iter().map(|(id, _)| id.clone()).collect();
let texts: Vec<String> = new_rows.iter().map(|(_, t)| t.clone()).collect();
let chunks: Vec<TextChunk> = chunker::chunk_column(
dataset_name, &doc_ids, &texts, chunk_size, overlap,
);
let new_chunks = chunks.len();
tracing::info!("refresh '{}': {} new docs -> {} chunks", dataset_name, new_docs, new_chunks);
// 6. Embed via Ollama (batched)
let batch_size = 32;
let mut all_vectors: Vec<Vec<f64>> = Vec::with_capacity(new_chunks);
for batch in chunks.chunks(batch_size) {
let batch_texts: Vec<String> = batch.iter().map(|c| c.text.clone()).collect();
let resp = ai_client
.embed(EmbedRequest { texts: batch_texts, model: None })
.await
.map_err(|e| format!("embed: {e}"))?;
all_vectors.extend(resp.embeddings);
}
// 7. Combine existing + new and write back as a single parquet
// (MVP — append-as-rewrite. ADR-019 points to Lance for true native
// append; for the Parquet sidecar we rewrite at refresh time.)
let mut new_stored: Vec<StoredEmbedding> = existing.clone();
for (chunk, vector) in chunks.into_iter().zip(all_vectors.iter()) {
new_stored.push(StoredEmbedding {
source: chunk.source,
doc_id: chunk.doc_id,
chunk_idx: chunk.chunk_idx,
chunk_text: chunk.text,
vector: vector.iter().map(|&x| x as f32).collect(),
});
}
// We have to reconstruct chunks + vectors for store_embeddings — but
// store_embeddings takes &[TextChunk] and vectors separately. Convert
// the combined StoredEmbedding back to that shape.
let combined_chunks: Vec<TextChunk> = new_stored
.iter()
.map(|e| TextChunk {
source: e.source.clone(),
doc_id: e.doc_id.clone(),
chunk_idx: e.chunk_idx,
text: e.chunk_text.clone(),
})
.collect();
let combined_vectors: Vec<Vec<f64>> = new_stored
.iter()
.map(|e| e.vector.iter().map(|&f| f as f64).collect())
.collect();
let _key = store::store_embeddings(
store_, &req.index_name, &combined_chunks, &combined_vectors,
).await?;
// 8. Evict embedding cache — next read will pick up the new file
let _ = embedding_cache.evict(&req.index_name).await;
// 9. Update index registry metadata (row/chunk counts; others unchanged
// so we don't disturb the existing entry if present)
let _ = try_update_index_meta(index_registry, &req.index_name, new_stored.len()).await;
// 10. Clear stale flag
registry.clear_embeddings_stale(dataset_name).await?;
let total = new_stored.len();
Ok(RefreshResult {
index_name: req.index_name.clone(),
dataset_name: dataset_name.to_string(),
pre_existing_docs: pre_existing_count,
dataset_docs,
new_docs_embedded: new_docs,
new_chunks_embedded: new_chunks,
total_embeddings_after: total,
duration_secs: t0.elapsed().as_secs_f32(),
stale_cleared: true,
})
}
/// Best-effort refresh of index registry metadata.
/// - If the index already exists, bump the chunk_count.
/// - If it's brand new (first-time build via refresh), REGISTER it so
/// downstream discovery (profile activation, A/B search) can find it.
async fn try_update_index_meta(
index_registry: &IndexRegistry,
index_name: &str,
chunk_count: usize,
) -> Result<(), String> {
if let Some(mut meta) = index_registry.get(index_name).await {
meta.chunk_count = chunk_count;
return index_registry.register(meta).await;
}
// First-time registration — infer reasonable defaults. Convention:
// index name `{source}_v1` or `{source}_vN` implies a source dataset
// named by stripping the `_vN` suffix. Otherwise use the index name
// as the source and let the caller patch later.
let source = match index_name.rsplit_once('_') {
Some((base, suffix)) if suffix.starts_with('v') && suffix[1..].chars().all(|c| c.is_ascii_digit()) => {
base.to_string()
}
_ => index_name.to_string(),
};
let meta = crate::index_registry::IndexMeta {
index_name: index_name.to_string(),
source,
model_name: "nomic-embed-text".to_string(),
model_version: "latest".to_string(),
dimensions: 768,
chunk_count,
doc_count: chunk_count,
chunk_size: 500,
overlap: 50,
storage_key: format!("vectors/{index_name}.parquet"),
created_at: chrono::Utc::now(),
build_time_secs: 0.0,
chunks_per_sec: 0.0,
bucket: "primary".to_string(),
vector_backend: shared::types::VectorBackend::Parquet,
id_prefix: None,
last_used: None,
build_signature: None,
};
index_registry.register(meta).await
}
+107
View File
@@ -0,0 +1,107 @@
/// Brute-force vector search with cosine similarity.
/// Works well up to ~100K vectors. HNSW index would go here for larger scale.
use crate::store::StoredEmbedding;
/// A search result with score.
#[derive(Debug, Clone, serde::Serialize)]
pub struct SearchResult {
pub source: String,
pub doc_id: String,
pub chunk_idx: u32,
pub chunk_text: String,
pub score: f32,
}
/// Search embeddings by cosine similarity. Returns top_k results.
pub fn search(
query_vector: &[f32],
embeddings: &[StoredEmbedding],
top_k: usize,
) -> Vec<SearchResult> {
let query_norm = norm(query_vector);
if query_norm == 0.0 {
return vec![];
}
let mut scored: Vec<SearchResult> = embeddings.iter().map(|emb| {
let score = cosine_similarity(query_vector, &emb.vector, query_norm);
SearchResult {
source: emb.source.clone(),
doc_id: emb.doc_id.clone(),
chunk_idx: emb.chunk_idx,
chunk_text: emb.chunk_text.clone(),
score,
}
}).collect();
// Sort descending by score
scored.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
scored.truncate(top_k);
scored
}
fn cosine_similarity(a: &[f32], b: &[f32], a_norm: f32) -> f32 {
if a.len() != b.len() {
return 0.0;
}
let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
let b_norm = norm(b);
if b_norm == 0.0 {
return 0.0;
}
dot / (a_norm * b_norm)
}
fn norm(v: &[f32]) -> f32 {
v.iter().map(|x| x * x).sum::<f32>().sqrt()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn identical_vectors_score_1() {
let v = vec![1.0, 2.0, 3.0];
let emb = StoredEmbedding {
source: "test".into(),
doc_id: "1".into(),
chunk_idx: 0,
chunk_text: "hello".into(),
vector: v.clone(),
};
let results = search(&v, &[emb], 1);
assert!((results[0].score - 1.0).abs() < 0.001);
}
#[test]
fn orthogonal_vectors_score_0() {
let q = vec![1.0, 0.0];
let emb = StoredEmbedding {
source: "test".into(),
doc_id: "1".into(),
chunk_idx: 0,
chunk_text: "hello".into(),
vector: vec![0.0, 1.0],
};
let results = search(&q, &[emb], 1);
assert!(results[0].score.abs() < 0.001);
}
#[test]
fn returns_top_k() {
let q = vec![1.0, 0.0, 0.0];
let embs: Vec<StoredEmbedding> = (0..10).map(|i| StoredEmbedding {
source: "test".into(),
doc_id: format!("{i}"),
chunk_idx: 0,
chunk_text: format!("doc {i}"),
vector: vec![1.0 - i as f32 * 0.1, i as f32 * 0.1, 0.0],
}).collect();
let results = search(&q, &embs, 3);
assert_eq!(results.len(), 3);
assert!(results[0].score >= results[1].score);
assert!(results[1].score >= results[2].score);
}
}
File diff suppressed because it is too large Load Diff
+115
View File
@@ -0,0 +1,115 @@
/// Vector storage as Parquet files.
/// Each embedding index is stored as: source, doc_id, chunk_idx, chunk_text, vector (binary blob).
/// Vectors are stored as raw f32 bytes for compact storage and fast loading.
use arrow::array::{ArrayRef, BinaryArray, Int32Array, RecordBatch, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use object_store::ObjectStore;
use std::sync::Arc;
use storaged::ops;
use crate::chunker::TextChunk;
/// A stored embedding — chunk text + its vector.
#[derive(Debug, Clone)]
pub struct StoredEmbedding {
pub source: String,
pub doc_id: String,
pub chunk_idx: u32,
pub chunk_text: String,
pub vector: Vec<f32>,
}
/// Store embeddings as a Parquet file in object storage.
pub async fn store_embeddings(
store: &Arc<dyn ObjectStore>,
index_name: &str,
chunks: &[TextChunk],
vectors: &[Vec<f64>], // from embedding API (f64), we store as f32
) -> Result<String, String> {
if chunks.len() != vectors.len() {
return Err(format!("chunk count ({}) != vector count ({})", chunks.len(), vectors.len()));
}
let n = chunks.len();
let sources: Vec<&str> = chunks.iter().map(|c| c.source.as_str()).collect();
let doc_ids: Vec<&str> = chunks.iter().map(|c| c.doc_id.as_str()).collect();
let chunk_idxs: Vec<i32> = chunks.iter().map(|c| c.chunk_idx as i32).collect();
let texts: Vec<&str> = chunks.iter().map(|c| c.text.as_str()).collect();
// Store vectors as raw f32 bytes (compact binary blob)
let vector_bytes: Vec<Vec<u8>> = vectors.iter().map(|v| {
v.iter().map(|&x| x as f32).flat_map(|f| f.to_le_bytes()).collect()
}).collect();
let vector_refs: Vec<&[u8]> = vector_bytes.iter().map(|v| v.as_slice()).collect();
let schema = Arc::new(Schema::new(vec![
Field::new("source", DataType::Utf8, false),
Field::new("doc_id", DataType::Utf8, false),
Field::new("chunk_idx", DataType::Int32, false),
Field::new("chunk_text", DataType::Utf8, false),
Field::new("vector", DataType::Binary, false),
]));
let arrays: Vec<ArrayRef> = vec![
Arc::new(StringArray::from(sources)),
Arc::new(StringArray::from(doc_ids)),
Arc::new(Int32Array::from(chunk_idxs)),
Arc::new(StringArray::from(texts)),
Arc::new(BinaryArray::from(vector_refs)),
];
let batch = RecordBatch::try_new(schema, arrays)
.map_err(|e| format!("RecordBatch error: {e}"))?;
let parquet = shared::arrow_helpers::record_batch_to_parquet(&batch)?;
let key = format!("vectors/{index_name}.parquet");
ops::put(store, &key, parquet).await?;
tracing::info!("stored {n} embeddings in {key}");
Ok(key)
}
/// Load all embeddings from a vector index file.
pub async fn load_embeddings(
store: &Arc<dyn ObjectStore>,
index_name: &str,
) -> Result<Vec<StoredEmbedding>, String> {
let key = format!("vectors/{index_name}.parquet");
let data = ops::get(store, &key).await?;
let (_, batches) = shared::arrow_helpers::parquet_to_record_batches(&data)?;
let mut embeddings = Vec::new();
for batch in &batches {
let sources = batch.column(0).as_any().downcast_ref::<StringArray>()
.ok_or("source column not string")?;
let doc_ids = batch.column(1).as_any().downcast_ref::<StringArray>()
.ok_or("doc_id column not string")?;
let chunk_idxs = batch.column(2).as_any().downcast_ref::<Int32Array>()
.ok_or("chunk_idx column not int")?;
let texts = batch.column(3).as_any().downcast_ref::<StringArray>()
.ok_or("chunk_text column not string")?;
let vectors = batch.column(4).as_any().downcast_ref::<BinaryArray>()
.ok_or("vector column not binary")?;
for i in 0..batch.num_rows() {
let vec_bytes = vectors.value(i);
let vector: Vec<f32> = vec_bytes.chunks_exact(4)
.map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
.collect();
embeddings.push(StoredEmbedding {
source: sources.value(i).to_string(),
doc_id: doc_ids.value(i).to_string(),
chunk_idx: chunk_idxs.value(i) as u32,
chunk_text: texts.value(i).to_string(),
vector,
});
}
}
tracing::info!("loaded {} embeddings from {key}", embeddings.len());
Ok(embeddings)
}
+286
View File
@@ -0,0 +1,286 @@
/// Dual-pipeline supervisor for embedding ingestion.
/// Splits work into ranges, runs parallel pipelines, handles failures
/// with round-robin retry and checkpointing.
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use aibridge::client::{AiClient, EmbedRequest};
use crate::chunker::TextChunk;
use crate::jobs::JobTracker;
use crate::store;
use object_store::ObjectStore;
/// A range of chunks to embed.
#[derive(Debug, Clone)]
struct ChunkRange {
start: usize,
end: usize,
attempts: u32,
}
/// Checkpoint: tracks which ranges are done, persisted for crash recovery.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Checkpoint {
pub job_id: String,
pub index_name: String,
pub total_chunks: usize,
pub completed_ranges: Vec<(usize, usize)>, // (start, end) pairs
pub failed_ranges: Vec<(usize, usize, String)>, // (start, end, error)
pub embedded_count: usize,
}
/// Pipeline status for monitoring.
#[derive(Debug, Clone, serde::Serialize)]
pub struct PipelineStatus {
pub pipeline_id: String,
pub current_range: Option<(usize, usize)>,
pub chunks_done: usize,
pub rate: f32,
pub alive: bool,
}
/// Supervisor configuration.
pub struct SupervisorConfig {
pub num_pipelines: usize, // default 2
pub batch_size: usize, // embeddings per API call (default 32)
pub range_size: usize, // chunks per range (default 2000)
pub max_retries: u32, // per range (default 3)
pub checkpoint_interval: usize, // checkpoint every N chunks (default 1000)
}
impl Default for SupervisorConfig {
fn default() -> Self {
Self {
num_pipelines: 4, // i9 + A4000 can handle parallel embedding
batch_size: 64, // larger batches for GPU throughput
range_size: 2500, // bigger ranges, fewer coordination overhead
max_retries: 3,
checkpoint_interval: 1000,
}
}
}
/// Run supervised dual-pipeline embedding.
pub async fn run_supervised(
job_id: &str,
index_name: &str,
chunks: Vec<TextChunk>,
ai_client: &AiClient,
obj_store: &Arc<dyn ObjectStore>,
tracker: &JobTracker,
config: SupervisorConfig,
) -> Result<String, String> {
let total = chunks.len();
let chunks = Arc::new(chunks);
tracing::info!("supervisor: starting {} pipelines for {} chunks (range_size={})",
config.num_pipelines, total, config.range_size);
// Split into ranges
let mut ranges: Vec<ChunkRange> = Vec::new();
let mut start = 0;
while start < total {
let end = (start + config.range_size).min(total);
ranges.push(ChunkRange { start, end, attempts: 0 });
start = end;
}
tracing::info!("supervisor: {} ranges created", ranges.len());
// Shared state
let work_queue: Arc<RwLock<Vec<ChunkRange>>> = Arc::new(RwLock::new(ranges));
let all_vectors: Arc<RwLock<HashMap<usize, Vec<Vec<f64>>>>> = Arc::new(RwLock::new(HashMap::new()));
let dead_letter: Arc<RwLock<Vec<(usize, usize, String)>>> = Arc::new(RwLock::new(Vec::new()));
let global_count: Arc<std::sync::atomic::AtomicUsize> = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let start_time = std::time::Instant::now();
// Checkpoint state
let checkpoint: Arc<RwLock<Checkpoint>> = Arc::new(RwLock::new(Checkpoint {
job_id: job_id.to_string(),
index_name: index_name.to_string(),
total_chunks: total,
completed_ranges: vec![],
failed_ranges: vec![],
embedded_count: 0,
}));
// Try to load existing checkpoint (crash recovery)
if let Ok(data) = storaged::ops::get(obj_store, &format!("checkpoints/{job_id}.json")).await {
if let Ok(saved) = serde_json::from_slice::<Checkpoint>(&data) {
tracing::info!("supervisor: recovered checkpoint — {}/{} already done",
saved.embedded_count, saved.total_chunks);
// Remove already-completed ranges from work queue
let mut queue = work_queue.write().await;
queue.retain(|r| !saved.completed_ranges.iter().any(|(s, e)| *s == r.start && *e == r.end));
global_count.store(saved.embedded_count, std::sync::atomic::Ordering::Relaxed);
*checkpoint.write().await = saved;
}
}
// Spawn pipeline workers
let mut handles = Vec::new();
for pipeline_id in 0..config.num_pipelines {
let queue = work_queue.clone();
let vectors = all_vectors.clone();
let dead = dead_letter.clone();
let count = global_count.clone();
let ckpt = checkpoint.clone();
let client = ai_client.clone();
let store = obj_store.clone();
let chunks = chunks.clone();
let tracker = tracker.clone();
let jid = job_id.to_string();
let batch_size = config.batch_size;
let max_retries = config.max_retries;
let ckpt_interval = config.checkpoint_interval;
let handle = tokio::spawn(async move {
let pid = format!("pipeline-{pipeline_id}");
tracing::info!("{pid}: started");
loop {
// Grab next range from queue
let range = {
let mut q = queue.write().await;
q.pop()
};
let range = match range {
Some(r) => r,
None => {
tracing::info!("{pid}: no more work, shutting down");
break;
}
};
tracing::debug!("{pid}: processing range [{}, {})", range.start, range.end);
// Embed the range
match embed_range(&chunks[range.start..range.end], &client, batch_size).await {
Ok(range_vectors) => {
let n = range_vectors.len();
vectors.write().await.insert(range.start, range_vectors);
let prev = count.fetch_add(n, std::sync::atomic::Ordering::Relaxed);
let done = prev + n;
// Update job tracker
let elapsed = start_time.elapsed().as_secs_f32();
let rate = if elapsed > 0.0 { done as f32 / elapsed } else { 0.0 };
tracker.update_embed_progress(&jid, done, rate).await;
// Update checkpoint
let mut ckpt = ckpt.write().await;
ckpt.completed_ranges.push((range.start, range.end));
ckpt.embedded_count = done;
// Persist checkpoint periodically
if done % ckpt_interval < n {
let json = serde_json::to_vec(&*ckpt).unwrap_or_default();
let _ = storaged::ops::put(&store, &format!("checkpoints/{jid}.json"), json.into()).await;
}
tracing::info!("{pid}: range [{},{}) done — {done}/{} total",
range.start, range.end, chunks.len());
}
Err(e) => {
let attempt = range.attempts + 1;
tracing::warn!("{pid}: range [{},{}) failed (attempt {attempt}/{max_retries}): {e}",
range.start, range.end);
if attempt < max_retries {
// Push back to queue for retry (round-robin to other pipeline)
queue.write().await.push(ChunkRange {
start: range.start,
end: range.end,
attempts: attempt,
});
} else {
// Dead letter
tracing::error!("{pid}: range [{},{}) dead-lettered after {max_retries} attempts",
range.start, range.end);
dead.write().await.push((range.start, range.end, e));
let mut ckpt = ckpt.write().await;
ckpt.failed_ranges.push((range.start, range.end, format!("max retries exceeded")));
}
}
}
}
});
handles.push(handle);
}
// Wait for all pipelines
for handle in handles {
let _ = handle.await;
}
// Check for dead letters
let dead = dead_letter.read().await;
if !dead.is_empty() {
tracing::warn!("supervisor: {} ranges failed permanently", dead.len());
}
// Assemble vectors in order
let vectors_map = all_vectors.read().await;
let mut sorted_starts: Vec<usize> = vectors_map.keys().cloned().collect();
sorted_starts.sort();
let mut final_vectors: Vec<Vec<f64>> = Vec::with_capacity(total);
for start in sorted_starts {
final_vectors.extend(vectors_map[&start].clone());
}
let embedded_count = final_vectors.len();
tracing::info!("supervisor: {embedded_count}/{total} chunks embedded, storing index");
// Store — only the successfully embedded chunks
// We need to match chunks to vectors
let successful_chunks: Vec<TextChunk> = {
let ckpt = checkpoint.read().await;
let mut result = Vec::new();
for (s, e) in &ckpt.completed_ranges {
result.extend(chunks[*s..*e].iter().cloned());
}
result
};
if successful_chunks.is_empty() {
return Err("no chunks were successfully embedded".into());
}
let key = store::store_embeddings(obj_store, index_name, &successful_chunks, &final_vectors).await?;
// Clean up checkpoint
let _ = storaged::ops::delete(obj_store, &format!("checkpoints/{job_id}.json")).await;
let elapsed = start_time.elapsed().as_secs_f32();
let rate = embedded_count as f32 / elapsed;
tracing::info!("supervisor: completed — {embedded_count} vectors in {elapsed:.0}s ({rate:.0}/sec)");
Ok(key)
}
/// Embed a range of chunks. Returns vectors or error.
async fn embed_range(
chunks: &[TextChunk],
ai_client: &AiClient,
batch_size: usize,
) -> Result<Vec<Vec<f64>>, String> {
let mut vectors = Vec::with_capacity(chunks.len());
for batch in chunks.chunks(batch_size) {
let texts: Vec<String> = batch.iter().map(|c| c.text.clone()).collect();
let resp = ai_client.embed(EmbedRequest {
texts,
model: None,
}).await.map_err(|e| format!("embed error: {e}"))?;
vectors.extend(resp.embeddings);
}
Ok(vectors)
}
+239
View File
@@ -0,0 +1,239 @@
/// Trial journal for HNSW parameter tuning.
///
/// Every HNSW build+eval is recorded as a Trial. The journal is append-only
/// and stored under `_hnsw_trials/{index_name}/` as batched JSONL files —
/// an AI agent iterating on configs reads prior trials to decide what to
/// try next, and writes a new trial on each attempt.
///
/// Storage uses the shared `storaged::append_log::AppendLog` so appends are
/// write-once (new file per batch) rather than rewriting a single growing
/// JSONL on every event. See `append_log.rs` for the full rationale.
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use storaged::append_log::AppendLog;
use storaged::registry::BucketRegistry;
use crate::index_registry::IndexRegistry;
/// HNSW build/search parameters the agent can tune.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HnswConfig {
pub ef_construction: usize,
pub ef_search: usize,
#[serde(default)]
pub seed: Option<u64>,
}
impl Default for HnswConfig {
/// Production default, locked in 2026-04-16 based on trial grid against
/// resumes_100k_v2 (100K vectors, 20 queries, recall@10):
/// ec=20 es=30 → recall 0.960, p50 509us, build 8s
/// ec=80 es=30 → recall 1.000, p50 873us, build 230s ← sweet spot
/// ec=200 es=30 → recall 1.000, p50 874us, build 106s (no recall gain)
///
/// `ec=80` is the smallest value that reaches 100% recall. Higher values
/// waste build time. `es=30` gives faster search than `es=100` with no
/// recall penalty at this scale.
fn default() -> Self {
Self {
ef_construction: 80,
ef_search: 30,
seed: None,
}
}
}
/// Metrics collected on every trial. All latencies in microseconds.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrialMetrics {
pub build_time_secs: f32,
pub search_latency_p50_us: f32,
pub search_latency_p95_us: f32,
pub search_latency_p99_us: f32,
pub recall_at_k: f32,
pub memory_bytes: u64,
pub vectors: usize,
pub eval_queries: usize,
/// Brute-force latency for comparison — how much speedup did HNSW buy us?
pub brute_force_latency_us: f32,
}
/// A single tuning attempt.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Trial {
pub id: String,
pub index_name: String,
pub eval_set: String,
pub config: HnswConfig,
pub metrics: TrialMetrics,
pub created_at: DateTime<Utc>,
/// Free-form note — the agent can record why it tried this config.
#[serde(default)]
pub note: Option<String>,
}
impl Trial {
pub fn new_id() -> String {
format!(
"trial-{}-{}",
Utc::now().timestamp_millis(),
&uuid::Uuid::new_v4().to_string()[..8]
)
}
}
/// Per-index append log, lazy-created on first write.
///
/// Federation layer 2: the journal resolves each index's bucket from the
/// index registry and writes its JSONL batches to THAT bucket, not
/// primary. Back-compat is preserved by `IndexMeta::bucket` defaulting
/// to "primary" for pre-federation indexes. Indexes the registry has
/// never heard of (edge case — trials run before first register) fall
/// through to primary as well.
#[derive(Clone)]
pub struct TrialJournal {
buckets: Arc<BucketRegistry>,
index_registry: IndexRegistry,
/// Cache per (bucket, index) AppendLog so the in-memory buffer persists
/// across calls. Keyed by `(bucket, index_name)` so moving an index
/// between buckets is clean — the old journal stays intact.
logs: Arc<RwLock<HashMap<(String, String), Arc<AppendLog>>>>,
}
impl TrialJournal {
pub fn new(buckets: Arc<BucketRegistry>, index_registry: IndexRegistry) -> Self {
Self {
buckets,
index_registry,
logs: Arc::new(RwLock::new(HashMap::new())),
}
}
fn prefix(index_name: &str) -> String {
format!("_hnsw_trials/{}", index_name)
}
/// Resolve which bucket holds this index's trial artifacts.
/// Falls back to primary for indexes without recorded metadata.
async fn bucket_for(&self, index_name: &str) -> String {
self.index_registry
.get(index_name)
.await
.map(|m| m.bucket)
.unwrap_or_else(|| "primary".to_string())
}
async fn log_for(&self, index_name: &str) -> Result<Arc<AppendLog>, String> {
let bucket = self.bucket_for(index_name).await;
let key = (bucket.clone(), index_name.to_string());
if let Some(log) = self.logs.read().await.get(&key) {
return Ok(log.clone());
}
let mut guard = self.logs.write().await;
if let Some(log) = guard.get(&key) {
return Ok(log.clone());
}
let store = self.buckets.get(&bucket)?;
// Trials arrive one at a time during human/agent iteration — a low
// threshold gives "hit /trials and see my latest attempt" immediacy
// without creating one file per event.
let log = Arc::new(
AppendLog::new(store, Self::prefix(index_name))
.with_flush_threshold(4),
);
guard.insert(key, log.clone());
Ok(log)
}
/// Append a trial record. In-memory buffered; persisted in batches.
pub async fn append(&self, trial: &Trial) -> Result<(), String> {
let line = serde_json::to_vec(trial).map_err(|e| e.to_string())?;
let log = self.log_for(&trial.index_name).await?;
log.append(line).await
}
/// Read all trials for an index (flushed batches + unflushed buffer).
pub async fn list(&self, index_name: &str) -> Result<Vec<Trial>, String> {
let log = self.log_for(index_name).await?;
let lines = log.read_all().await?;
let mut trials = Vec::with_capacity(lines.len());
for line in lines {
match serde_json::from_slice::<Trial>(&line) {
Ok(t) => trials.push(t),
Err(e) => tracing::warn!("trial journal: skip malformed line: {e}"),
}
}
Ok(trials)
}
/// Explicit flush for callers that want write-through semantics
/// (e.g. an agent that wants to commit a trial before querying stats).
pub async fn flush(&self, index_name: &str) -> Result<(), String> {
let log = self.log_for(index_name).await?;
log.flush().await
}
/// Compact all batch files for an index into one.
pub async fn compact(&self, index_name: &str) -> Result<storaged::append_log::CompactStats, String> {
let log = self.log_for(index_name).await?;
log.compact().await
}
/// Current champion for an index by the named metric.
/// Valid metrics: `recall`, `latency`, `pareto`.
///
/// The `pareto` strategy is a placeholder — J should tune the scoring
/// function to match what matters in production. Right now it's a simple
/// weighted sum.
pub async fn best(
&self,
index_name: &str,
metric: &str,
) -> Result<Option<Trial>, String> {
let trials = self.list(index_name).await?;
if trials.is_empty() {
return Ok(None);
}
let best = match metric {
"recall" => trials
.into_iter()
.max_by(|a, b| {
a.metrics
.recall_at_k
.partial_cmp(&b.metrics.recall_at_k)
.unwrap_or(std::cmp::Ordering::Equal)
})
.unwrap(),
"latency" => trials
.into_iter()
.min_by(|a, b| {
a.metrics
.search_latency_p95_us
.partial_cmp(&b.metrics.search_latency_p95_us)
.unwrap_or(std::cmp::Ordering::Equal)
})
.unwrap(),
"pareto" | _ => trials
.into_iter()
.max_by(|a, b| pareto_score(a).partial_cmp(&pareto_score(b)).unwrap())
.unwrap(),
};
Ok(Some(best))
}
}
/// Simple Pareto-style score: reward recall, penalize p95 latency.
/// Tunable — J should swap this in production to match what matters.
fn pareto_score(t: &Trial) -> f32 {
// Recall is [0, 1]. Latency is us — assume 100us baseline.
let recall = t.metrics.recall_at_k;
let latency_penalty = (t.metrics.search_latency_p95_us / 1000.0).min(1.0); // cap at 1ms
recall - 0.2 * latency_penalty
}