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:
@@ -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"] }
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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 }
|
||||
}
|
||||
}
|
||||
@@ -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(>_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, >, 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, >, 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, >, 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, >, 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, >, 2);
|
||||
assert!((r - 1.0).abs() < 1e-6);
|
||||
}
|
||||
}
|
||||
@@ -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>>,
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user