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:
@@ -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\
|
||||
|
||||
Reference in New Issue
Block a user