Federation foundation + HNSW trial system + Postgres streaming + PRD reframe
Four shipped features and a PRD realignment, all measured end-to-end:
HNSW trial system (Phase 15 horizon item → complete)
- vectord: EmbeddingCache, harness (eval sets + brute-force ground truth),
TrialJournal, parameterized HnswConfig on build_index_with_config
- /vectors/hnsw/trial, /hnsw/trials/{idx}, /hnsw/trials/{idx}/best,
/hnsw/evals/{name}/autogen, /hnsw/cache/stats
- Measured on resumes_100k_v2 (100K × 768d): brute-force 44ms -> HNSW 873us
at 100% recall@10. ec=80 es=30 locked as HnswConfig::default()
- Lower ec values trade recall for build time: 20/30 = 0.96 recall in 8s,
80/30 = 1.00 recall in 230s
Catalog manifest repair
- catalogd: resync_from_parquet reads parquet footers to restore row_count
and columns on drifted manifests
- POST /catalog/datasets/{name}/resync + POST /catalog/resync-missing
- All 7 staffing tables recovered to PRD-matching 2,469,278 rows
Federation foundation (ADR-017)
- shared::secrets: SecretsProvider trait + FileSecretsProvider (reads
/etc/lakehouse/secrets.toml, enforces 0600 perms)
- storaged::registry::BucketRegistry — multi-bucket resolution with
rescue_bucket read fallback and reachability probing
- storaged::error_journal — bucket op failures visible in one HTTP call
- storaged::append_log — write-once batched append pattern (fixes the RMW
anti-pattern llms3.com calls out; errors and trial journals both use it)
- /storage/buckets, /storage/errors, /storage/bucket-health,
/storage/errors/{flush,compact}
- Bucket-aware I/O at /storage/buckets/{bucket}/objects/{*key} with
X-Lakehouse-Rescue-Used observability headers on fallback
Postgres streaming ingest
- ingestd::pg_stream: DSN parser, batched ORDER BY + LIMIT/OFFSET pagination
into ArrowWriter, lineage redacts password
- POST /ingest/db — verified against live knowledge_base.team_runs
(586 rows × 13 cols, 6 batches, 196ms end-to-end)
PRD realignment (2026-04-16)
- Dual use case: staffing analytics + local LLM knowledge substrate
- Removed "multi-tenancy (single-owner system)" from non-goals
- Added invariants 8-11: indexes hot-swappable, per-reader profiles,
trials-as-data, operational failures findable in one HTTP call
- New phases 16 (hot-swap generations), 17 (model profiles + dataset
bindings), 18 (Lance vs Parquet+sidecar evaluation)
- Known ceilings table documents the 5M vector wall and escape hatches
- ADR-017 (federation), ADR-018 (append-log pattern) added
- EXECUTION_PLAN.md sequences phases B-E with success gates and
decision rules
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
root
@@ -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,216 @@
|
||||
/// 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::sync::Arc;
|
||||
|
||||
use aibridge::client::{AiClient, EmbedRequest};
|
||||
use storaged::ops;
|
||||
|
||||
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())
|
||||
}
|
||||
}
|
||||
|
||||
/// 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;
|
||||
}
|
||||
let gt_set: std::collections::HashSet<&String> =
|
||||
ground_truth.iter().take(k).collect();
|
||||
let hits = predicted
|
||||
.iter()
|
||||
.take(k)
|
||||
.filter(|d| gt_set.contains(d))
|
||||
.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,
|
||||
}
|
||||
}
|
||||
@@ -7,6 +7,7 @@ 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)]
|
||||
@@ -62,18 +63,32 @@ impl HnswStore {
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an HNSW index from stored embeddings.
|
||||
/// 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 index '{}' from {} vectors", index_name, n);
|
||||
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
|
||||
@@ -92,14 +107,17 @@ impl HnswStore {
|
||||
values.push(i);
|
||||
}
|
||||
|
||||
// Build HNSW — this is the expensive part
|
||||
let map = Builder::default()
|
||||
.ef_construction(40) // balanced for 100K scale
|
||||
.ef_search(50) // fast search with good recall
|
||||
.build(points, values);
|
||||
// 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 index '{}' built: {} vectors in {:.1}s", index_name, n, build_time);
|
||||
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);
|
||||
@@ -111,6 +129,43 @@ impl HnswStore {
|
||||
})
|
||||
}
|
||||
|
||||
/// 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,
|
||||
@@ -166,3 +221,9 @@ pub struct BuildStats {
|
||||
pub vectors: usize,
|
||||
pub build_time_secs: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BenchResult {
|
||||
pub latencies_us: Vec<f32>,
|
||||
pub retrieved: Vec<Vec<String>>,
|
||||
}
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
pub mod chunker;
|
||||
pub mod embedding_cache;
|
||||
pub mod harness;
|
||||
pub mod hnsw;
|
||||
pub mod index_registry;
|
||||
pub mod jobs;
|
||||
@@ -7,3 +9,4 @@ pub mod search;
|
||||
pub mod rag;
|
||||
pub mod supervisor;
|
||||
pub mod service;
|
||||
pub mod trial;
|
||||
|
||||
+304
-10
@@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
use aibridge::client::{AiClient, EmbedRequest};
|
||||
use crate::{chunker, hnsw, index_registry, jobs, rag, search, store, supervisor};
|
||||
use crate::{chunker, embedding_cache, harness, hnsw, index_registry, jobs, rag, search, store, supervisor, trial};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct VectorState {
|
||||
@@ -19,6 +19,8 @@ pub struct VectorState {
|
||||
pub job_tracker: jobs::JobTracker,
|
||||
pub index_registry: index_registry::IndexRegistry,
|
||||
pub hnsw_store: hnsw::HnswStore,
|
||||
pub embedding_cache: embedding_cache::EmbeddingCache,
|
||||
pub trial_journal: trial::TrialJournal,
|
||||
}
|
||||
|
||||
pub fn router(state: VectorState) -> Router {
|
||||
@@ -34,6 +36,17 @@ pub fn router(state: VectorState) -> Router {
|
||||
.route("/hnsw/build", post(build_hnsw))
|
||||
.route("/hnsw/search", post(search_hnsw))
|
||||
.route("/hnsw/list", get(list_hnsw))
|
||||
// Trial system — parameterized tuning loop
|
||||
.route("/hnsw/trial", post(run_trial))
|
||||
.route("/hnsw/trials/{index_name}", get(list_trials))
|
||||
.route("/hnsw/trials/{index_name}/best", get(best_trial))
|
||||
// Eval sets
|
||||
.route("/hnsw/evals", get(list_evals))
|
||||
.route("/hnsw/evals/{name}", get(get_eval).put(put_eval))
|
||||
.route("/hnsw/evals/{name}/autogen", post(autogen_eval))
|
||||
// Cache management
|
||||
.route("/hnsw/cache/stats", get(cache_stats))
|
||||
.route("/hnsw/cache/{index_name}", axum::routing::delete(cache_evict))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
@@ -307,26 +320,35 @@ async fn rag_query(
|
||||
struct BuildHnswRequest {
|
||||
/// Name of the stored vector index to build HNSW from
|
||||
index_name: String,
|
||||
/// Optional config override. Omit to use the production default
|
||||
/// (ec=80 es=30 — see HnswConfig::default docs for rationale).
|
||||
#[serde(default)]
|
||||
config: Option<trial::HnswConfig>,
|
||||
}
|
||||
|
||||
/// Build an HNSW index from an existing stored vector index.
|
||||
/// Loads embeddings from Parquet, builds HNSW in memory.
|
||||
/// Uses the embedding cache so repeated builds don't reload from Parquet.
|
||||
async fn build_hnsw(
|
||||
State(state): State<VectorState>,
|
||||
Json(req): Json<BuildHnswRequest>,
|
||||
) -> impl IntoResponse {
|
||||
tracing::info!("building HNSW for '{}'", req.index_name);
|
||||
let config = req.config.unwrap_or_default();
|
||||
tracing::info!(
|
||||
"building HNSW for '{}' ef_construction={} ef_search={}",
|
||||
req.index_name, config.ef_construction, config.ef_search,
|
||||
);
|
||||
|
||||
// Load embeddings from Parquet
|
||||
let embeddings = store::load_embeddings(&state.store, &req.index_name)
|
||||
let embeddings = state
|
||||
.embedding_cache
|
||||
.get_or_load(&req.index_name)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("index not found: {e}")))?;
|
||||
|
||||
let n = embeddings.len();
|
||||
tracing::info!("loaded {} embeddings, building HNSW...", n);
|
||||
|
||||
// Build HNSW
|
||||
match state.hnsw_store.build_index(&req.index_name, embeddings).await {
|
||||
match state
|
||||
.hnsw_store
|
||||
.build_index_with_config(&req.index_name, (*embeddings).clone(), &config)
|
||||
.await
|
||||
{
|
||||
Ok(stats) => Ok(Json(stats)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
@@ -372,3 +394,275 @@ async fn search_hnsw(
|
||||
async fn list_hnsw(State(state): State<VectorState>) -> impl IntoResponse {
|
||||
Json(state.hnsw_store.list().await)
|
||||
}
|
||||
|
||||
// --- Trial System: parameterized HNSW tuning loop ---
|
||||
//
|
||||
// Flow:
|
||||
// 1. Agent picks an HnswConfig
|
||||
// 2. POST /hnsw/trial builds HNSW with that config against cached embeddings,
|
||||
// runs every query in the harness, measures latency + recall vs the
|
||||
// harness's ground truth, appends a Trial record to _hnsw_trials/{idx}.jsonl
|
||||
// 3. Agent reads GET /hnsw/trials/{index}, sees history, decides next config
|
||||
// 4. Repeat until converged.
|
||||
//
|
||||
// The first trial triggers embedding load (slow). Every subsequent trial reuses
|
||||
// the cache — so the agent iterates in seconds, not minutes.
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TrialRequest {
|
||||
index_name: String,
|
||||
harness: String,
|
||||
#[serde(default)]
|
||||
config: trial::HnswConfig,
|
||||
#[serde(default)]
|
||||
note: Option<String>,
|
||||
}
|
||||
|
||||
async fn run_trial(
|
||||
State(state): State<VectorState>,
|
||||
Json(req): Json<TrialRequest>,
|
||||
) -> Result<Json<trial::Trial>, (StatusCode, String)> {
|
||||
let mut harness_set = harness::EvalSet::load(&state.store, &req.harness)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("harness not found: {e}")))?;
|
||||
|
||||
if harness_set.index_name != req.index_name {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!(
|
||||
"harness '{}' is for index '{}', not '{}'",
|
||||
req.harness, harness_set.index_name, req.index_name
|
||||
),
|
||||
));
|
||||
}
|
||||
if harness_set.queries.is_empty() {
|
||||
return Err((StatusCode::BAD_REQUEST, "harness has no queries".into()));
|
||||
}
|
||||
|
||||
let embeddings = state
|
||||
.embedding_cache
|
||||
.get_or_load(&req.index_name)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("load embeddings: {e}")))?;
|
||||
|
||||
if !harness_set.ground_truth_built {
|
||||
tracing::info!("trial: computing ground truth for harness '{}'", harness_set.name);
|
||||
let t0 = std::time::Instant::now();
|
||||
harness::compute_ground_truth(&mut harness_set, &embeddings, &state.ai_client)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("ground truth: {e}")))?;
|
||||
tracing::info!("trial: ground truth built in {:.1}s", t0.elapsed().as_secs_f32());
|
||||
harness_set
|
||||
.save(&state.store)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("save harness: {e}")))?;
|
||||
}
|
||||
|
||||
let trial_id = trial::Trial::new_id();
|
||||
let hnsw_slot = format!("{}__{}", req.index_name, trial_id);
|
||||
|
||||
let build_stats = state
|
||||
.hnsw_store
|
||||
.build_index_with_config(&hnsw_slot, (*embeddings).clone(), &req.config)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("build: {e}")))?;
|
||||
|
||||
let query_vectors: Vec<Vec<f32>> = harness_set
|
||||
.queries
|
||||
.iter()
|
||||
.filter_map(|q| q.query_embedding.clone())
|
||||
.collect();
|
||||
let bench = state
|
||||
.hnsw_store
|
||||
.bench_search(&hnsw_slot, &query_vectors, harness_set.k)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("search: {e}")))?;
|
||||
|
||||
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)]
|
||||
};
|
||||
|
||||
// One brute-force reference latency — keeps the cost proportional to
|
||||
// whatever the agent is willing to pay per trial.
|
||||
let brute_latency_us = if let Some(qv) = query_vectors.first() {
|
||||
let t0 = std::time::Instant::now();
|
||||
let _ = harness::brute_force_top_k(qv, &embeddings, harness_set.k);
|
||||
t0.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_record = trial::Trial {
|
||||
id: trial_id.clone(),
|
||||
index_name: req.index_name.clone(),
|
||||
eval_set: req.harness.clone(),
|
||||
config: req.config.clone(),
|
||||
metrics: trial::TrialMetrics {
|
||||
build_time_secs: build_stats.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_stats.vectors,
|
||||
eval_queries: harness_set.queries.len(),
|
||||
brute_force_latency_us: brute_latency_us,
|
||||
},
|
||||
created_at: chrono::Utc::now(),
|
||||
note: req.note,
|
||||
};
|
||||
|
||||
state
|
||||
.trial_journal
|
||||
.append(&trial_record)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("journal: {e}")))?;
|
||||
|
||||
state.hnsw_store.drop(&hnsw_slot).await;
|
||||
|
||||
Ok(Json(trial_record))
|
||||
}
|
||||
|
||||
async fn list_trials(
|
||||
State(state): State<VectorState>,
|
||||
Path(index_name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match state.trial_journal.list(&index_name).await {
|
||||
Ok(trials) => Ok(Json(trials)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct BestTrialQuery {
|
||||
#[serde(default = "default_metric")]
|
||||
metric: String,
|
||||
}
|
||||
|
||||
fn default_metric() -> String {
|
||||
"pareto".to_string()
|
||||
}
|
||||
|
||||
async fn best_trial(
|
||||
State(state): State<VectorState>,
|
||||
Path(index_name): Path<String>,
|
||||
Query(q): Query<BestTrialQuery>,
|
||||
) -> impl IntoResponse {
|
||||
match state.trial_journal.best(&index_name, &q.metric).await {
|
||||
Ok(Some(t)) => Ok(Json(t)),
|
||||
Ok(None) => Err((StatusCode::NOT_FOUND, "no trials yet".to_string())),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Harness management ---
|
||||
|
||||
async fn list_evals(State(state): State<VectorState>) -> impl IntoResponse {
|
||||
match harness::EvalSet::list(&state.store).await {
|
||||
Ok(names) => Ok(Json(names)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_eval(
|
||||
State(state): State<VectorState>,
|
||||
Path(name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match harness::EvalSet::load(&state.store, &name).await {
|
||||
Ok(e) => Ok(Json(e)),
|
||||
Err(err) => Err((StatusCode::NOT_FOUND, err)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn put_eval(
|
||||
State(state): State<VectorState>,
|
||||
Path(name): Path<String>,
|
||||
Json(mut harness_set): Json<harness::EvalSet>,
|
||||
) -> impl IntoResponse {
|
||||
harness_set.name = name;
|
||||
harness_set.ground_truth_built = harness_set
|
||||
.queries
|
||||
.iter()
|
||||
.all(|q| q.ground_truth.is_some());
|
||||
match harness_set.save(&state.store).await {
|
||||
Ok(()) => Ok(Json(harness_set)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct AutogenRequest {
|
||||
index_name: String,
|
||||
#[serde(default = "default_sample_count")]
|
||||
sample_count: usize,
|
||||
#[serde(default = "default_k")]
|
||||
k: usize,
|
||||
}
|
||||
|
||||
fn default_sample_count() -> usize { 100 }
|
||||
fn default_k() -> usize { 10 }
|
||||
|
||||
async fn autogen_eval(
|
||||
State(state): State<VectorState>,
|
||||
Path(name): Path<String>,
|
||||
Json(req): Json<AutogenRequest>,
|
||||
) -> Result<Json<harness::EvalSet>, (StatusCode, String)> {
|
||||
let embeddings = state
|
||||
.embedding_cache
|
||||
.get_or_load(&req.index_name)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("load embeddings: {e}")))?;
|
||||
|
||||
let mut harness_set = harness::synthetic_from_chunks(
|
||||
&name,
|
||||
&req.index_name,
|
||||
&embeddings,
|
||||
req.sample_count,
|
||||
req.k,
|
||||
);
|
||||
|
||||
harness::compute_ground_truth(&mut harness_set, &embeddings, &state.ai_client)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("ground truth: {e}")))?;
|
||||
|
||||
harness_set
|
||||
.save(&state.store)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("save: {e}")))?;
|
||||
|
||||
Ok(Json(harness_set))
|
||||
}
|
||||
|
||||
// --- Embedding cache management ---
|
||||
|
||||
async fn cache_stats(State(state): State<VectorState>) -> impl IntoResponse {
|
||||
Json(state.embedding_cache.stats().await)
|
||||
}
|
||||
|
||||
async fn cache_evict(
|
||||
State(state): State<VectorState>,
|
||||
Path(index_name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
let ok = state.embedding_cache.evict(&index_name).await;
|
||||
Json(serde_json::json!({ "evicted": ok, "index_name": index_name }))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
/// 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 object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use storaged::append_log::AppendLog;
|
||||
|
||||
/// 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.
|
||||
#[derive(Clone)]
|
||||
pub struct TrialJournal {
|
||||
store: Arc<dyn ObjectStore>,
|
||||
/// Cache per-index AppendLog instances so the in-memory buffer persists
|
||||
/// across calls.
|
||||
logs: Arc<RwLock<HashMap<String, Arc<AppendLog>>>>,
|
||||
}
|
||||
|
||||
impl TrialJournal {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
logs: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
fn prefix(index_name: &str) -> String {
|
||||
format!("_hnsw_trials/{}", index_name)
|
||||
}
|
||||
|
||||
async fn log_for(&self, index_name: &str) -> Arc<AppendLog> {
|
||||
if let Some(log) = self.logs.read().await.get(index_name) {
|
||||
return log.clone();
|
||||
}
|
||||
let mut guard = self.logs.write().await;
|
||||
if let Some(log) = guard.get(index_name) {
|
||||
return log.clone();
|
||||
}
|
||||
// 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(self.store.clone(), Self::prefix(index_name))
|
||||
.with_flush_threshold(4),
|
||||
);
|
||||
guard.insert(index_name.to_string(), log.clone());
|
||||
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