Phase 7: Vector index + RAG pipeline
- vectord crate: chunk → embed → store → search → RAG - chunker: configurable chunk size + overlap, sentence-boundary aware splitting - store: embeddings as Parquet (binary blob f32 vectors), portable format - search: brute-force cosine similarity (works up to ~100K vectors) - rag: full pipeline — embed question → search index → retrieve context → LLM answer - Endpoints: POST /vectors/index, /vectors/search, /vectors/rag - Gateway wired with vectord service - Tested: 200 candidate resumes indexed in 5.4s, semantic search + RAG working - 20 unit tests passing (chunker, search, ingestd, shared) - AI gives honest "no match found" when context doesn't support an answer Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,175 @@
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::State,
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
use aibridge::client::{AiClient, EmbedRequest};
|
||||
use crate::{chunker, rag, search, store};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct VectorState {
|
||||
pub store: Arc<dyn ObjectStore>,
|
||||
pub ai_client: AiClient,
|
||||
}
|
||||
|
||||
pub fn router(state: VectorState) -> Router {
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/index", post(create_index))
|
||||
.route("/search", post(search_index))
|
||||
.route("/rag", post(rag_query))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
async fn health() -> &'static str {
|
||||
"vectord ok"
|
||||
}
|
||||
|
||||
// --- Index creation: chunk text → embed → store ---
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CreateIndexRequest {
|
||||
/// Name for this vector index
|
||||
index_name: String,
|
||||
/// Source identifier
|
||||
source: String,
|
||||
/// List of documents to index
|
||||
documents: Vec<DocInput>,
|
||||
/// Chunk size in characters (default 500)
|
||||
chunk_size: Option<usize>,
|
||||
/// Overlap in characters (default 50)
|
||||
overlap: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct DocInput {
|
||||
id: String,
|
||||
text: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct CreateIndexResponse {
|
||||
index_name: String,
|
||||
documents: usize,
|
||||
chunks: usize,
|
||||
storage_key: String,
|
||||
}
|
||||
|
||||
async fn create_index(
|
||||
State(state): State<VectorState>,
|
||||
Json(req): Json<CreateIndexRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let chunk_size = req.chunk_size.unwrap_or(500);
|
||||
let overlap = req.overlap.unwrap_or(50);
|
||||
|
||||
tracing::info!("creating vector index '{}' from {} documents", req.index_name, req.documents.len());
|
||||
|
||||
// 1. Chunk all documents
|
||||
let doc_ids: Vec<String> = req.documents.iter().map(|d| d.id.clone()).collect();
|
||||
let texts: Vec<String> = req.documents.iter().map(|d| d.text.clone()).collect();
|
||||
let chunks = chunker::chunk_column(&req.source, &doc_ids, &texts, chunk_size, overlap);
|
||||
|
||||
if chunks.is_empty() {
|
||||
return Err((StatusCode::BAD_REQUEST, "no text to index".to_string()));
|
||||
}
|
||||
|
||||
tracing::info!("{} documents → {} chunks", req.documents.len(), chunks.len());
|
||||
|
||||
// 2. Embed all chunks (batch to avoid timeout)
|
||||
let batch_size = 32;
|
||||
let mut all_vectors: Vec<Vec<f64>> = Vec::new();
|
||||
|
||||
for batch in chunks.chunks(batch_size) {
|
||||
let texts: Vec<String> = batch.iter().map(|c| c.text.clone()).collect();
|
||||
let embed_resp = state.ai_client.embed(EmbedRequest {
|
||||
texts,
|
||||
model: None,
|
||||
}).await.map_err(|e| (StatusCode::BAD_GATEWAY, format!("embed error: {e}")))?;
|
||||
all_vectors.extend(embed_resp.embeddings);
|
||||
}
|
||||
|
||||
// 3. Store
|
||||
let key = store::store_embeddings(&state.store, &req.index_name, &chunks, &all_vectors)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||
|
||||
Ok((StatusCode::CREATED, Json(CreateIndexResponse {
|
||||
index_name: req.index_name,
|
||||
documents: req.documents.len(),
|
||||
chunks: chunks.len(),
|
||||
storage_key: key,
|
||||
})))
|
||||
}
|
||||
|
||||
// --- Search ---
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SearchRequest {
|
||||
index_name: String,
|
||||
query: String,
|
||||
top_k: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SearchResponse {
|
||||
results: Vec<search::SearchResult>,
|
||||
query: String,
|
||||
}
|
||||
|
||||
async fn search_index(
|
||||
State(state): State<VectorState>,
|
||||
Json(req): Json<SearchRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let top_k = req.top_k.unwrap_or(5);
|
||||
|
||||
// Embed query
|
||||
let embed_resp = state.ai_client.embed(EmbedRequest {
|
||||
texts: vec![req.query.clone()],
|
||||
model: None,
|
||||
}).await.map_err(|e| (StatusCode::BAD_GATEWAY, format!("embed error: {e}")))?;
|
||||
|
||||
if embed_resp.embeddings.is_empty() {
|
||||
return Err((StatusCode::BAD_GATEWAY, "no embedding returned".to_string()));
|
||||
}
|
||||
|
||||
let query_vec: Vec<f32> = embed_resp.embeddings[0].iter().map(|&x| x as f32).collect();
|
||||
|
||||
// Load index and search
|
||||
let embeddings = store::load_embeddings(&state.store, &req.index_name)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("index not found: {e}")))?;
|
||||
|
||||
let results = search::search(&query_vec, &embeddings, top_k);
|
||||
|
||||
Ok(Json(SearchResponse {
|
||||
results,
|
||||
query: req.query,
|
||||
}))
|
||||
}
|
||||
|
||||
// --- RAG ---
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RagRequest {
|
||||
index_name: String,
|
||||
question: String,
|
||||
top_k: Option<usize>,
|
||||
}
|
||||
|
||||
async fn rag_query(
|
||||
State(state): State<VectorState>,
|
||||
Json(req): Json<RagRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let top_k = req.top_k.unwrap_or(5);
|
||||
|
||||
match rag::query(&req.question, &req.index_name, top_k, &state.store, &state.ai_client).await {
|
||||
Ok(resp) => Ok(Json(resp)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user