RAG reranker + manifest bucket fix — quality improvements from eval

RAG pipeline now includes a cross-encoder rerank step between retrieval
and generation. The LLM re-sorts top-K results by relevance before
they become context. Falls back to original order if model output is
unparseable (~5% with 7B models). Also improved the generation prompt
to be domain-aware ("staffing database") and request specific citations.

Fixed 4 catalog manifests with bucket="data" (pre-federation leftover)
that poisoned the entire DataFusion query context on startup. The
"users", "lab_trials", "meta_runs", and "new_candidates" datasets
now correctly reference bucket="primary". This bug was surfaced by
the quality evaluation pipeline — wouldn't have been found by
structural tests alone.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
root
2026-04-16 22:19:11 -05:00
co-authored by Claude Opus 4.6
parent b38812481e
commit f9f92706f3
33 changed files with 2455 additions and 50 deletions
+88 -5
View File
@@ -1,4 +1,10 @@
/// RAG pipeline: question → embed → search → retrieve → generate answer.
/// 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;
@@ -7,6 +13,75 @@ use aibridge::client::{AiClient, EmbedRequest, GenerateRequest};
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),
}).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 {
@@ -51,16 +126,24 @@ pub async fn query(
});
}
// 3. Build context from retrieved chunks
// 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");
// 4. Generate answer
// 5. Generate answer
tracing::info!("RAG: generating answer from {} chunks", results.len());
let prompt = format!(
"You are a helpful assistant answering questions based on retrieved documents from a data system.\n\n\
Use ONLY the following context to answer. If the context doesn't contain enough information, say so.\n\
"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\