From 9e448401b5429e199175aa89aacb1402ea32c07e Mon Sep 17 00:00:00 2001 From: Harvey Tseng Date: Fri, 4 Sep 2026 17:23:12 -0400 Subject: [PATCH 1/4] feat: rerank stage behind an optional [rerank] section - Fused top_n candidates go to a Jina-style POST /rerank and come back in the reranker's order with its relevance as the score; absent config means no call, and --no-rerank skips it per request - gate stays a config knob (none, natural-language) until measured --- crates/scry-core/src/config.rs | 47 +++++++++ crates/scry-core/src/error.rs | 2 + crates/scry-core/src/lib.rs | 1 + crates/scry-core/src/rerank.rs | 126 ++++++++++++++++++++++++ crates/scry-server/src/api.rs | 7 ++ crates/scry-server/src/lib.rs | 3 + crates/scry-server/src/routes/search.rs | 45 ++++++++- crates/scry-server/tests/http.rs | 91 +++++++++++++++++ crates/scry/src/cli.rs | 11 ++- crates/scry/src/commands/eval.rs | 11 ++- crates/scry/src/commands/search.rs | 1 + deploy/config.example.toml | 11 +++ docs/search.md | 8 ++ 13 files changed, 357 insertions(+), 7 deletions(-) create mode 100644 crates/scry-core/src/rerank.rs diff --git a/crates/scry-core/src/config.rs b/crates/scry-core/src/config.rs index f6ded5a..5cd53a6 100644 --- a/crates/scry-core/src/config.rs +++ b/crates/scry-core/src/config.rs @@ -20,6 +20,8 @@ pub struct Config { #[serde(default)] pub tavily: Option, #[serde(default)] + pub rerank: Option, + #[serde(default)] pub client: ClientConfig, #[serde(default)] pub index: IndexConfig, @@ -101,6 +103,35 @@ pub struct ChatConfig { pub thinking: bool, } +/// Cross-encoder rerank over the fused candidates through an +/// OpenAI/Jina-style `POST {base_url}/rerank`. Absent means no rerank. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RerankConfig { + pub base_url: String, + /// Empty sends no Authorization header. + #[serde(default)] + pub api_key: String, + pub model: String, + #[serde(default = "default_rerank_top_n")] + pub top_n: usize, + #[serde(default)] + pub gate: RerankGate, +} + +/// Which queries pay for a rerank call. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Default)] +#[serde(rename_all = "kebab-case")] +pub enum RerankGate { + #[default] + None, + NaturalLanguage, +} + +fn default_rerank_top_n() -> usize { + 20 +} + #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields)] pub struct TavilyConfig { @@ -307,6 +338,10 @@ impl Config { tavily.api_key = resolve_secret(Some(std::mem::take(&mut tavily.api_key))).unwrap_or_default(); } + if let Some(rerank) = self.rerank.as_mut() { + rerank.api_key = + resolve_secret(Some(std::mem::take(&mut rerank.api_key))).unwrap_or_default(); + } } } @@ -364,6 +399,18 @@ mod tests { assert_eq!(config.index.max_file_count, 500); } + #[test] + fn parses_rerank_section() { + let config: Config = toml::from_str( + "[rerank]\nbase_url = \"http://r:8080/v1\"\nmodel = \"bge\"\ngate = \"natural-language\"\n", + ) + .unwrap(); + let rerank = config.rerank.unwrap(); + assert_eq!(rerank.top_n, 20); + assert_eq!(rerank.gate, RerankGate::NaturalLanguage); + assert!(Config::default().rerank.is_none()); + } + #[test] fn rejects_unknown_keys() { assert!(toml::from_str::("[server]\nlissten = \"x\"\n").is_err()); diff --git a/crates/scry-core/src/error.rs b/crates/scry-core/src/error.rs index 1339d72..fa8743b 100644 --- a/crates/scry-core/src/error.rs +++ b/crates/scry-core/src/error.rs @@ -14,6 +14,8 @@ pub enum Error { Embedding(String), #[error("chat: {0}")] Chat(String), + #[error("rerank: {0}")] + Rerank(String), } pub type Result = std::result::Result; diff --git a/crates/scry-core/src/lib.rs b/crates/scry-core/src/lib.rs index 1e9075b..8968224 100644 --- a/crates/scry-core/src/lib.rs +++ b/crates/scry-core/src/lib.rs @@ -9,6 +9,7 @@ pub mod hashing; pub mod index; pub mod memory; pub mod repo; +pub mod rerank; pub mod search; pub mod store; pub mod walk; diff --git a/crates/scry-core/src/rerank.rs b/crates/scry-core/src/rerank.rs new file mode 100644 index 0000000..9290cb0 --- /dev/null +++ b/crates/scry-core/src/rerank.rs @@ -0,0 +1,126 @@ +//! Cross-encoder rerank of the fused candidates. The endpoint is the +//! Jina/Cohere shape llama-server exposes at `/v1/rerank`. + +use serde::Deserialize; + +use crate::config::{RerankConfig, RerankGate}; +use crate::search::SearchHit; +use crate::{Error, Result}; + +pub struct RerankClient { + client: reqwest::Client, + config: RerankConfig, +} + +#[derive(Deserialize)] +struct RerankResponse { + results: Vec, +} + +#[derive(Debug, Clone, Copy, Deserialize)] +pub struct RerankResult { + pub index: usize, + pub relevance_score: f64, +} + +impl RerankClient { + pub fn new(config: RerankConfig) -> Self { + Self { + client: reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .expect("client build"), + config, + } + } + + pub fn top_n(&self) -> usize { + self.config.top_n + } + + pub fn gate(&self) -> RerankGate { + self.config.gate + } + + /// Scores every document against the query; results come back best + /// first with the index into `documents`. + pub async fn rerank(&self, query: &str, documents: &[String]) -> Result> { + let url = format!("{}/rerank", self.config.base_url.trim_end_matches('/')); + let body = serde_json::json!({ + "model": self.config.model, + "query": query, + "documents": documents, + "top_n": documents.len(), + }); + let mut request = self.client.post(&url).json(&body); + if !self.config.api_key.is_empty() { + request = request.bearer_auth(&self.config.api_key); + } + let response = request.send().await?.error_for_status()?; + let parsed: RerankResponse = response.json().await?; + let mut results = parsed.results; + if results.iter().any(|r| r.index >= documents.len()) { + return Err(Error::Rerank("index out of range in response".to_string())); + } + results.sort_by(|a, b| b.relevance_score.total_cmp(&a.relevance_score)); + Ok(results) + } +} + +/// Reorders `hits` by the reranker's ranking, reports its relevance clamped +/// to [0, 1] as the score, and keeps `limit`. Hits the reranker did not +/// score are dropped. +pub fn reorder(hits: Vec, ranked: &[RerankResult], limit: usize) -> Vec { + let mut slots: Vec> = hits.into_iter().map(Some).collect(); + ranked + .iter() + .filter_map(|r| { + slots[r.index].take().map(|hit| SearchHit { + score: r.relevance_score.clamp(0.0, 1.0), + ..hit + }) + }) + .take(limit) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn hit(relpath: &str) -> SearchHit { + SearchHit { + repo_key: "r".to_string(), + relpath: relpath.to_string(), + start_line: 1, + end_line: 1, + symbol: None, + score: 0.5, + content: String::new(), + } + } + + #[test] + fn reorder_follows_the_reranker_and_truncates() { + let hits = vec![hit("a"), hit("b"), hit("c")]; + let ranked = [ + RerankResult { + index: 2, + relevance_score: 3.5, + }, + RerankResult { + index: 0, + relevance_score: 0.4, + }, + RerankResult { + index: 1, + relevance_score: -1.0, + }, + ]; + let out = reorder(hits, &ranked, 2); + let paths: Vec<&str> = out.iter().map(|h| h.relpath.as_str()).collect(); + assert_eq!(paths, vec!["c", "a"]); + assert_eq!(out[0].score, 1.0); + assert_eq!(out[1].score, 0.4); + } +} diff --git a/crates/scry-server/src/api.rs b/crates/scry-server/src/api.rs index e792ef3..6c4693b 100644 --- a/crates/scry-server/src/api.rs +++ b/crates/scry-server/src/api.rs @@ -12,12 +12,19 @@ pub struct SearchRequest { pub limit: usize, #[serde(default)] pub path_prefix: Option, + /// Applies the server's `[rerank]` stage when one is configured. + #[serde(default = "default_true")] + pub rerank: bool, } fn default_limit() -> usize { 10 } +fn default_true() -> bool { + true +} + #[derive(Debug, Serialize, Deserialize)] pub struct SearchResponse { pub hits: Vec, diff --git a/crates/scry-server/src/lib.rs b/crates/scry-server/src/lib.rs index ac123fd..85888bd 100644 --- a/crates/scry-server/src/lib.rs +++ b/crates/scry-server/src/lib.rs @@ -19,6 +19,7 @@ use scry_core::Result; use scry_core::chat::ChatClient; use scry_core::config::{Config, HydeMode, IndexConfig, MemoryConfig}; use scry_core::embed::{Embedder, HttpEmbedder}; +use scry_core::rerank::RerankClient; use scry_core::store::Store; const MAX_SYNC_BODY_BYTES: usize = 64 * 1024 * 1024; @@ -32,6 +33,7 @@ pub struct AppState { pub index_config: IndexConfig, pub memory_config: MemoryConfig, pub tavily: Option, + pub rerank: Option, } impl AppState { @@ -54,6 +56,7 @@ impl AppState { index_config: config.index.clone(), memory_config: config.memory.clone(), tavily: config.tavily.clone().map(tavily::TavilyClient::new), + rerank: config.rerank.clone().map(RerankClient::new), } } } diff --git a/crates/scry-server/src/routes/search.rs b/crates/scry-server/src/routes/search.rs index b94e2a6..c332905 100644 --- a/crates/scry-server/src/routes/search.rs +++ b/crates/scry-server/src/routes/search.rs @@ -2,12 +2,24 @@ use std::sync::Arc; use axum::Json; use axum::extract::State; -use scry_core::search::{SearchOptions, query_vector, search_with_vector}; +use scry_core::config::RerankGate; +use scry_core::index::embed_input; +use scry_core::rerank::reorder; +use scry_core::search::{SearchOptions, query_vector, route_query, search_with_vector}; use crate::AppState; use crate::api::{Hit, SearchRequest, SearchResponse}; use crate::error::ApiError; +/// A reranker is an improvement layer: past this budget or on any error +/// the fused order is returned, so search is never worse than without it. +const RERANK_BUDGET: std::time::Duration = std::time::Duration::from_secs(6); + +fn truncated(mut hits: Vec, limit: usize) -> Vec { + hits.truncate(limit); + hits +} + pub async fn search( State(state): State>, Json(request): Json, @@ -19,6 +31,15 @@ pub async fn search( &request.query, ) .await?; + let rerank = state.rerank.as_ref().filter(|client| { + request.rerank + && match client.gate() { + RerankGate::None => true, + RerankGate::NaturalLanguage => route_query(&request.query).natural_language, + } + }); + let pool = rerank.map_or(request.limit, |client| request.limit.max(client.top_n())); + let (query, limit) = (request.query.clone(), request.limit); let hits = state .store .call(move |store| { @@ -30,13 +51,33 @@ pub async fn search( None => None, }; let options = SearchOptions { - limit: request.limit, + limit: pool, path_prefix: request.path_prefix.clone(), }; search_with_vector(store, repo_id, &request.query, &vector, &options).map(Some) }) .await? .ok_or_else(|| ApiError::NotFound("repo not indexed".to_string()))?; + let hits = match rerank { + Some(client) => { + let documents: Vec = hits + .iter() + .map(|h| embed_input(&h.repo_key, &h.relpath, h.symbol.as_deref(), &h.content)) + .collect(); + match tokio::time::timeout(RERANK_BUDGET, client.rerank(&query, &documents)).await { + Ok(Ok(ranked)) => reorder(hits, &ranked, limit), + Ok(Err(error)) => { + tracing::warn!("rerank failed, returning fused order: {error}"); + truncated(hits, limit) + } + Err(_) => { + tracing::warn!("rerank exceeded {RERANK_BUDGET:?}, returning fused order"); + truncated(hits, limit) + } + } + } + None => hits, + }; Ok(Json(SearchResponse { hits: hits .into_iter() diff --git a/crates/scry-server/tests/http.rs b/crates/scry-server/tests/http.rs index 560cde3..f0116c7 100644 --- a/crates/scry-server/tests/http.rs +++ b/crates/scry-server/tests/http.rs @@ -1,7 +1,9 @@ use std::sync::Arc; use scry_core::config::Config; +use scry_core::config::RerankConfig; use scry_core::embed::HashEmbedder; +use scry_core::rerank::RerankClient; use scry_core::store::Store; use scry_server::api::{ FileUpload, ManifestRequest, SearchRequest, SearchResponse, StatusResponse, SyncRequest, @@ -13,6 +15,10 @@ const TOKEN: &str = "test-token"; const REPO: &str = "github.com/test/http"; async fn spawn_server() -> String { + spawn_server_with(None).await +} + +async fn spawn_server_with(rerank: Option) -> String { let store = Store::open_in_memory("hash-test", 128).unwrap(); let config: Config = toml::from_str(&format!( "[server]\nauth_token = \"{TOKEN}\"\n[embedding]\ndim = 128\nmodel = \"hash-test\"\n" @@ -27,6 +33,7 @@ async fn spawn_server() -> String { index_config: config.index.clone(), memory_config: config.memory.clone(), tavily: None, + rerank, }; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let base = format!("http://{}", listener.local_addr().unwrap()); @@ -121,6 +128,7 @@ async fn auth_sync_search_roundtrip() { query: "load configuration".to_string(), limit: 5, path_prefix: None, + rerank: true, }) .send() .await @@ -140,6 +148,7 @@ async fn auth_sync_search_roundtrip() { query: "load configuration".to_string(), limit: 5, path_prefix: None, + rerank: true, }) .send() .await @@ -159,9 +168,91 @@ async fn auth_sync_search_roundtrip() { query: "x".to_string(), limit: 5, path_prefix: None, + rerank: true, }) .send() .await .unwrap(); assert_eq!(missing_repo.status(), 404); } + +/// Scores each document by its position, so the last candidate wins. +async fn spawn_mock_reranker() -> String { + let app = axum::Router::new().route( + "/v1/rerank", + axum::routing::post( + |axum::Json(body): axum::Json| async move { + let n = body["documents"].as_array().map_or(0, |d| d.len()); + let results: Vec = (0..n) + .map(|i| serde_json::json!({ "index": i, "relevance_score": i as f64 })) + .collect(); + axum::Json(serde_json::json!({ "results": results })) + }, + ), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}/v1", listener.local_addr().unwrap()); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + base +} + +async fn search(http: &reqwest::Client, base: &str, rerank: bool) -> Vec { + let response: SearchResponse = http + .post(format!("{base}/v1/search")) + .bearer_auth(TOKEN) + .json(&SearchRequest { + repo_key: Some(REPO.to_string()), + query: "load configuration".to_string(), + limit: 2, + path_prefix: None, + rerank, + }) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap(); + response.hits.into_iter().map(|h| h.relpath).collect() +} + +#[tokio::test] +async fn rerank_stage_reorders_and_no_rerank_bypasses_it() { + let reranker = spawn_mock_reranker().await; + let config: RerankConfig = toml::from_str(&format!( + "base_url = \"{reranker}\"\nmodel = \"mock\"\ntop_n = 2\n" + )) + .unwrap(); + let base = spawn_server_with(Some(RerankClient::new(config))).await; + let http = reqwest::Client::new(); + http.post(format!("{base}/v1/sync")) + .bearer_auth(TOKEN) + .json(&SyncRequest { + repo_key: REPO.to_string(), + upserts: vec![ + upload( + "src/config.rs", + "pub fn load_configuration(path: &str) -> u16 {\n path.len() as u16\n}\n", + ), + upload( + "src/net.rs", + "pub fn send_request(url: &str) -> usize {\n url.len()\n}\n", + ), + ], + deletes: vec![], + }) + .send() + .await + .unwrap() + .error_for_status() + .unwrap(); + + let plain = search(&http, &base, false).await; + let reranked = search(&http, &base, true).await; + assert_eq!(plain[0], "src/config.rs"); + assert_eq!(reranked.len(), 2); + assert_eq!(reranked[0], plain[1]); + assert_eq!(reranked[1], plain[0]); +} diff --git a/crates/scry/src/cli.rs b/crates/scry/src/cli.rs index 7460ed3..96c33ec 100644 --- a/crates/scry/src/cli.rs +++ b/crates/scry/src/cli.rs @@ -29,6 +29,7 @@ SEARCH OPTIONS: -w, --web include web results --repo search a specific indexed repo from anywhere; outside a repo, all indexed repos are searched + --no-rerank skip the server's rerank stage, if configured "; #[derive(Debug, PartialEq, Eq)] @@ -40,6 +41,7 @@ pub struct SearchArgs { pub answer: bool, pub web: bool, pub repo: Option, + pub rerank: bool, } const VALUE_FLAGS: &[&str] = &["-m", "--max-count", "--max-file-size", "--max-file-count"]; @@ -47,7 +49,7 @@ const VALUE_FLAGS: &[&str] = &["-m", "--max-count", "--max-file-size", "--max-fi pub fn parse_search_args(args: &[String]) -> Result { let mut positionals: Vec<&str> = Vec::new(); let mut max_count = 10; - let (mut content, mut answer, mut web) = (false, false, false); + let (mut content, mut answer, mut web, mut rerank) = (false, false, false, true); let mut repo = None; let mut it = args.iter().peekable(); while let Some(arg) = it.next() { @@ -64,7 +66,8 @@ pub fn parse_search_args(args: &[String]) -> Result { flag if VALUE_FLAGS.contains(&flag) => { it.next(); } - "-i" | "-r" | "-s" | "-d" | "--sync" | "--dry-run" | "--no-rerank" => {} + "--no-rerank" => rerank = false, + "-i" | "-r" | "-s" | "-d" | "--sync" | "--dry-run" => {} flag if flag.starts_with('-') => {} positional => positionals.push(positional), } @@ -80,6 +83,7 @@ pub fn parse_search_args(args: &[String]) -> Result { answer, web, repo, + rerank, }) } @@ -149,6 +153,9 @@ mod tests { assert_eq!(parsed.max_count, 5); assert!(parsed.content && parsed.web && parsed.answer); assert_eq!(parsed.query, "query"); + assert!(parsed.rerank); + let parsed = parse_search_args(&strings(&["--no-rerank", "query"])).unwrap(); + assert!(!parsed.rerank); } #[test] diff --git a/crates/scry/src/commands/eval.rs b/crates/scry/src/commands/eval.rs index 42fba3b..95d7573 100644 --- a/crates/scry/src/commands/eval.rs +++ b/crates/scry/src/commands/eval.rs @@ -7,7 +7,7 @@ use serde::Deserialize; use super::repo_context; use crate::client::ApiClient; -const USAGE: &str = "usage: scry eval [--runs N] [--limit N]"; +const USAGE: &str = "usage: scry eval [--runs N] [--limit N] [--no-rerank]"; #[derive(Deserialize)] struct EvalFile { @@ -54,16 +54,18 @@ struct Args { file: String, runs: usize, limit: usize, + rerank: bool, } fn parse_args(args: &[String]) -> Result { let mut file = None; - let (mut runs, mut limit) = (1, 10); + let (mut runs, mut limit, mut rerank) = (1, 10, true); let mut it = args.iter(); while let Some(arg) = it.next() { match arg.as_str() { "--runs" => runs = it.next().and_then(|v| v.parse().ok()).unwrap_or(runs), "--limit" => limit = it.next().and_then(|v| v.parse().ok()).unwrap_or(limit), + "--no-rerank" => rerank = false, _ => file = Some(arg.clone()), } } @@ -74,6 +76,7 @@ fn parse_args(args: &[String]) -> Result { file, runs: runs.max(1), limit: limit.max(1), + rerank, }) } @@ -82,6 +85,7 @@ async fn search_case( repo_key: &str, case: &EvalCase, limit: usize, + rerank: bool, ) -> Result { let started = Instant::now(); let response = client @@ -90,6 +94,7 @@ async fn search_case( query: case.query.clone(), limit, path_prefix: case.path_prefix.clone(), + rerank, }) .await?; let rank = response.hits.iter().position(|hit| { @@ -175,7 +180,7 @@ pub async fn run(args: &[String]) -> Result<()> { for _ in 0..args.runs { let mut results = Vec::with_capacity(cases.case.len()); for case in &cases.case { - results.push(search_case(&ctx.client, &repo_key, case, args.limit).await?); + results.push(search_case(&ctx.client, &repo_key, case, args.limit, args.rerank).await?); } runs.push(results); } diff --git a/crates/scry/src/commands/search.rs b/crates/scry/src/commands/search.rs index 4752d36..c7a6054 100644 --- a/crates/scry/src/commands/search.rs +++ b/crates/scry/src/commands/search.rs @@ -75,6 +75,7 @@ pub async fn run(args: SearchArgs) -> Result<()> { query: args.query.clone(), limit: args.max_count, path_prefix: scope.path_prefix, + rerank: args.rerank, }) .await?; match &scope.local_root { diff --git a/deploy/config.example.toml b/deploy/config.example.toml index 8babb88..b75aa67 100644 --- a/deploy/config.example.toml +++ b/deploy/config.example.toml @@ -26,6 +26,17 @@ batch_size = 32 # [tavily] # api_key = "env:TAVILY_API_KEY" +# Cross-encoder rerank of the fused top candidates through a Jina-style +# POST {base_url}/rerank (llama-server with --reranking serves one). +# gate = "none" reranks every query; "natural-language" skips identifier +# queries. --no-rerank on the CLI bypasses it per query. Errors and calls +# over six seconds fall back to the fused order. api_key is optional. +# [rerank] +# base_url = "http://localhost:12434/v1" +# model = "bge-reranker-v2-m3" +# top_n = 20 +# gate = "none" + [client] server_url = "http://127.0.0.1:7345" # token = "env:SCRY_TOKEN" diff --git a/docs/search.md b/docs/search.md index ed41be2..ae12e5d 100644 --- a/docs/search.md +++ b/docs/search.md @@ -15,6 +15,14 @@ queries lean on BM25, natural-language questions lean on the dense leg. Fused candidates get a small recency boost for recently edited files and a greedy near-duplicate filter before the final ranking. +With a `[rerank]` endpoint configured, the fused top `top_n` candidates +are rescored by a cross-encoder and the reranker's relevance, clamped to +[0, 1], becomes the displayed score; `gate` decides which queries pay for +the call and `--no-rerank` skips it per query. A reranker that errors or +takes longer than six seconds is skipped for that query and the fused +order is returned, so search is never worse than without it. The score +transform for logit-shaped rerankers is decided once one is measured. + Chunks are function-level where a tree-sitter grammar exists (16 languages), blank-line-snapped windows elsewhere, and each chunk is embedded with a `repo > path > symbol` header for context. From c0ce9d4dfd952303d8544d9ca2500d3510bc7059 Mon Sep 17 00:00:00 2001 From: Harvey Tseng Date: Fri, 4 Sep 2026 18:26:25 -0400 Subject: [PATCH 2/4] test: offline rerank fusion simulator from a dump run --- eval/simulate_rerank.py | 78 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 eval/simulate_rerank.py diff --git a/eval/simulate_rerank.py b/eval/simulate_rerank.py new file mode 100644 index 0000000..ae6aede --- /dev/null +++ b/eval/simulate_rerank.py @@ -0,0 +1,78 @@ +"""Offline rerank simulation from a dump run. + +Usage: simulate_rerank.py + +Dump format, one JSON line per query, written by a measurement build of +the search route with SCRY_RERANK_DUMP set: + {"query": str, + "candidates": [{"relpath", "start", "end", "symbol", "score"}, ...], + "rerank": [{"index": int, "score": float}, ...]} +`candidates` is the pool search_with_vector returns to the route (after +weighted RRF, recency boost, and Jaccard dedup), in rank order, top_n +long; `rerank` holds the reranker's score per candidate index. + +Fusion is computed exactly where the shipped leg sits, at the route level +after the pool exists: score[i] = 1/(k + pool_rank_i + 1) ++ w / (k + rerank_rank_i + 1), k = 60, then sort descending. The NL gate +applies the leg only to queries route_query marks natural language. +""" +import json, re, sys, tomllib +from pathlib import Path + +dump_path, eval_dir = sys.argv[1], Path(sys.argv[2]) +K = 60.0 + +def tokens(q): return [t for t in re.split(r'[^0-9A-Za-z_:]+', q) if len(t) >= 2] +def inner_upper(t): return t[:1].islower() and any(c.isupper() for c in t[1:]) +def natural_language(q): + ts = tokens(q); ident = any('_' in t or '::' in t or inner_upper(t) for t in ts) + return not (len(ts) <= 3 or (ident and len(ts) <= 6)) + +def matches(cand, expectation): + path, _, line = expectation.rpartition(':') + if path and line.isdigit(): + return cand['relpath'] == path and cand['start'] <= int(line) <= cand['end'] + return cand['relpath'] == expectation + +dumps = {} +for line in open(dump_path): + d = json.loads(line); dumps.setdefault(d['query'], d) + +def order_for(d, mode, weight): + cands = d['candidates']; n = len(cands) + if mode == 'fused': + return list(range(n)) + rerank = sorted(d['rerank'], key=lambda r: -r['score']) + if mode == 'replace': + return [r['index'] for r in rerank] + score = [1.0 / (K + i + 1) for i in range(n)] + for rank, r in enumerate(rerank): + score[r['index']] += weight / (K + rank + 1) + return sorted(range(n), key=lambda i: -score[i]) + +def evaluate(cases, mode, weight=1.0, gate=False): + hits, rr = 0, 0.0 + for case in cases: + d = dumps.get(case['query']) + if d is None: + continue + m = 'fused' if (gate and not natural_language(case['query'])) else mode + order = order_for(d, m, weight)[:10] + rank = next((i for i, idx in enumerate(order) + if any(matches(d['candidates'][idx], e) for e in case['expect'])), None) + if rank is not None: + hits += 1; rr += 1 / (rank + 1) + n = len(cases) + return f"{hits/n:.3f}/{rr/n:.3f}" + +for set_name in ['scry', 'finance-query', 'soothfast']: + cases = tomllib.load(open(eval_dir / f'{set_name}.toml', 'rb'))['case'] + covered = sum(c['query'] in dumps for c in cases) + row = [f"{set_name:14} ({covered}/{len(cases)} dumped)", + f"fused {evaluate(cases, 'fused')}", + f"replace {evaluate(cases, 'replace')}"] + for w in (1.0, 2.0, 3.0, 5.0): + row.append(f"rrf w{w:g} {evaluate(cases, 'rrf', w)}") + row.append(f"rrf w2 nl-gate {evaluate(cases, 'rrf', 2.0, gate=True)}") + row.append(f"replace nl-gate {evaluate(cases, 'replace', gate=True)}") + print(' '.join(row)) From 6f19c6c87f02a24cc80d234826020b59ca948e8d Mon Sep 17 00:00:00 2001 From: Harvey Tseng Date: Fri, 4 Sep 2026 19:38:14 -0400 Subject: [PATCH 3/4] feat: fuse the reranker as a weighted RRF leg - Offline simulation from one dump run (bge-reranker-v2-m3, 20-50 candidates, index-e6c.db copy): replacing the order lost 8 scry hits while gaining 4 on finance-query; fusing the reranker's ranking as a third RRF leg at weight 1 over the top 20 gains on every set - Summed over the three sets, recall@10 1.934 -> 2.064 and MRR 1.290 -> 1.383; the natural-language gate only lowered MRR, so it is gone and top_n defaults to 20, which also halves the CPU cost --- crates/scry-core/src/config.rs | 26 +++++------ crates/scry-core/src/rerank.rs | 57 ++++++++++++++++--------- crates/scry-server/src/routes/search.rs | 15 ++----- crates/scry-server/tests/http.rs | 2 +- deploy/config.example.toml | 12 +++--- docs/search.md | 11 +++-- 6 files changed, 65 insertions(+), 58 deletions(-) diff --git a/crates/scry-core/src/config.rs b/crates/scry-core/src/config.rs index 5cd53a6..d0829ed 100644 --- a/crates/scry-core/src/config.rs +++ b/crates/scry-core/src/config.rs @@ -103,8 +103,9 @@ pub struct ChatConfig { pub thinking: bool, } -/// Cross-encoder rerank over the fused candidates through an -/// OpenAI/Jina-style `POST {base_url}/rerank`. Absent means no rerank. +/// Cross-encoder over the fused top `top_n`, reached through an +/// OpenAI/Jina-style `POST {base_url}/rerank`, fused back into the +/// ranking as a weighted reciprocal-rank leg. Absent means no rerank. #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields)] pub struct RerankConfig { @@ -115,23 +116,18 @@ pub struct RerankConfig { pub model: String, #[serde(default = "default_rerank_top_n")] pub top_n: usize, - #[serde(default)] - pub gate: RerankGate, -} - -/// Which queries pay for a rerank call. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Default)] -#[serde(rename_all = "kebab-case")] -pub enum RerankGate { - #[default] - None, - NaturalLanguage, + #[serde(default = "default_rerank_weight")] + pub weight: f64, } fn default_rerank_top_n() -> usize { 20 } +fn default_rerank_weight() -> f64 { + 1.0 +} + #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields)] pub struct TavilyConfig { @@ -402,12 +398,12 @@ mod tests { #[test] fn parses_rerank_section() { let config: Config = toml::from_str( - "[rerank]\nbase_url = \"http://r:8080/v1\"\nmodel = \"bge\"\ngate = \"natural-language\"\n", + "[rerank]\nbase_url = \"http://r:8080/v1\"\nmodel = \"bge\"\nweight = 2.0\n", ) .unwrap(); let rerank = config.rerank.unwrap(); assert_eq!(rerank.top_n, 20); - assert_eq!(rerank.gate, RerankGate::NaturalLanguage); + assert_eq!(rerank.weight, 2.0); assert!(Config::default().rerank.is_none()); } diff --git a/crates/scry-core/src/rerank.rs b/crates/scry-core/src/rerank.rs index 9290cb0..f1e8fb9 100644 --- a/crates/scry-core/src/rerank.rs +++ b/crates/scry-core/src/rerank.rs @@ -3,7 +3,7 @@ use serde::Deserialize; -use crate::config::{RerankConfig, RerankGate}; +use crate::config::RerankConfig; use crate::search::SearchHit; use crate::{Error, Result}; @@ -38,8 +38,8 @@ impl RerankClient { self.config.top_n } - pub fn gate(&self) -> RerankGate { - self.config.gate + pub fn weight(&self) -> f64 { + self.config.weight } /// Scores every document against the query; results come back best @@ -67,19 +67,28 @@ impl RerankClient { } } -/// Reorders `hits` by the reranker's ranking, reports its relevance clamped -/// to [0, 1] as the score, and keeps `limit`. Hits the reranker did not -/// score are dropped. -pub fn reorder(hits: Vec, ranked: &[RerankResult], limit: usize) -> Vec { +/// Fuses the reranker's ranking into the pool as a reciprocal-rank leg +/// weighted by `weight` against the pool's own order, then keeps +/// `limit`. Scores are left as they were; the reranker changes order only. +pub fn fuse( + hits: Vec, + ranked: &[RerankResult], + weight: f64, + limit: usize, +) -> Vec { + const K: f64 = 60.0; + let mut score: Vec = (0..hits.len()) + .map(|rank| 1.0 / (K + rank as f64 + 1.0)) + .collect(); + for (rank, r) in ranked.iter().enumerate() { + score[r.index] += weight / (K + rank as f64 + 1.0); + } + let mut order: Vec = (0..hits.len()).collect(); + order.sort_by(|a, b| score[*b].total_cmp(&score[*a])); let mut slots: Vec> = hits.into_iter().map(Some).collect(); - ranked - .iter() - .filter_map(|r| { - slots[r.index].take().map(|hit| SearchHit { - score: r.relevance_score.clamp(0.0, 1.0), - ..hit - }) - }) + order + .into_iter() + .filter_map(|i| slots[i].take()) .take(limit) .collect() } @@ -101,7 +110,7 @@ mod tests { } #[test] - fn reorder_follows_the_reranker_and_truncates() { + fn fuse_promotes_by_reranker_rank_and_keeps_scores() { let hits = vec![hit("a"), hit("b"), hit("c")]; let ranked = [ RerankResult { @@ -117,10 +126,20 @@ mod tests { relevance_score: -1.0, }, ]; - let out = reorder(hits, &ranked, 2); + let out = fuse(hits, &ranked, 2.0, 2); let paths: Vec<&str> = out.iter().map(|h| h.relpath.as_str()).collect(); assert_eq!(paths, vec!["c", "a"]); - assert_eq!(out[0].score, 1.0); - assert_eq!(out[1].score, 0.4); + assert_eq!(out[0].score, 0.5); + } + + #[test] + fn fuse_with_zero_weight_keeps_the_pool_order() { + let hits = vec![hit("a"), hit("b")]; + let ranked = [RerankResult { + index: 1, + relevance_score: 9.0, + }]; + let out = fuse(hits, &ranked, 0.0, 2); + assert_eq!(out[0].relpath, "a"); } } diff --git a/crates/scry-server/src/routes/search.rs b/crates/scry-server/src/routes/search.rs index c332905..fc72061 100644 --- a/crates/scry-server/src/routes/search.rs +++ b/crates/scry-server/src/routes/search.rs @@ -2,10 +2,9 @@ use std::sync::Arc; use axum::Json; use axum::extract::State; -use scry_core::config::RerankGate; use scry_core::index::embed_input; -use scry_core::rerank::reorder; -use scry_core::search::{SearchOptions, query_vector, route_query, search_with_vector}; +use scry_core::rerank::fuse; +use scry_core::search::{SearchOptions, query_vector, search_with_vector}; use crate::AppState; use crate::api::{Hit, SearchRequest, SearchResponse}; @@ -31,13 +30,7 @@ pub async fn search( &request.query, ) .await?; - let rerank = state.rerank.as_ref().filter(|client| { - request.rerank - && match client.gate() { - RerankGate::None => true, - RerankGate::NaturalLanguage => route_query(&request.query).natural_language, - } - }); + let rerank = state.rerank.as_ref().filter(|_| request.rerank); let pool = rerank.map_or(request.limit, |client| request.limit.max(client.top_n())); let (query, limit) = (request.query.clone(), request.limit); let hits = state @@ -65,7 +58,7 @@ pub async fn search( .map(|h| embed_input(&h.repo_key, &h.relpath, h.symbol.as_deref(), &h.content)) .collect(); match tokio::time::timeout(RERANK_BUDGET, client.rerank(&query, &documents)).await { - Ok(Ok(ranked)) => reorder(hits, &ranked, limit), + Ok(Ok(ranked)) => fuse(hits, &ranked, client.weight(), limit), Ok(Err(error)) => { tracing::warn!("rerank failed, returning fused order: {error}"); truncated(hits, limit) diff --git a/crates/scry-server/tests/http.rs b/crates/scry-server/tests/http.rs index f0116c7..33694e1 100644 --- a/crates/scry-server/tests/http.rs +++ b/crates/scry-server/tests/http.rs @@ -222,7 +222,7 @@ async fn search(http: &reqwest::Client, base: &str, rerank: bool) -> Vec async fn rerank_stage_reorders_and_no_rerank_bypasses_it() { let reranker = spawn_mock_reranker().await; let config: RerankConfig = toml::from_str(&format!( - "base_url = \"{reranker}\"\nmodel = \"mock\"\ntop_n = 2\n" + "base_url = \"{reranker}\"\nmodel = \"mock\"\ntop_n = 2\nweight = 2.0\n" )) .unwrap(); let base = spawn_server_with(Some(RerankClient::new(config))).await; diff --git a/deploy/config.example.toml b/deploy/config.example.toml index b75aa67..0f69be7 100644 --- a/deploy/config.example.toml +++ b/deploy/config.example.toml @@ -26,16 +26,16 @@ batch_size = 32 # [tavily] # api_key = "env:TAVILY_API_KEY" -# Cross-encoder rerank of the fused top candidates through a Jina-style -# POST {base_url}/rerank (llama-server with --reranking serves one). -# gate = "none" reranks every query; "natural-language" skips identifier -# queries. --no-rerank on the CLI bypasses it per query. Errors and calls -# over six seconds fall back to the fused order. api_key is optional. +# Cross-encoder over the fused top candidates through a Jina-style +# POST {base_url}/rerank (llama-server with --reranking serves one), fused +# back in as a weighted reciprocal-rank leg. Errors and calls over six +# seconds fall back to the fused order; --no-rerank skips it per query. +# api_key is optional. # [rerank] # base_url = "http://localhost:12434/v1" # model = "bge-reranker-v2-m3" # top_n = 20 -# gate = "none" +# weight = 1.0 [client] server_url = "http://127.0.0.1:7345" diff --git a/docs/search.md b/docs/search.md index ae12e5d..5b4f401 100644 --- a/docs/search.md +++ b/docs/search.md @@ -16,12 +16,11 @@ Fused candidates get a small recency boost for recently edited files and a greedy near-duplicate filter before the final ranking. With a `[rerank]` endpoint configured, the fused top `top_n` candidates -are rescored by a cross-encoder and the reranker's relevance, clamped to -[0, 1], becomes the displayed score; `gate` decides which queries pay for -the call and `--no-rerank` skips it per query. A reranker that errors or -takes longer than six seconds is skipped for that query and the fused -order is returned, so search is never worse than without it. The score -transform for logit-shaped rerankers is decided once one is measured. +are scored by a cross-encoder and its ranking joins the fusion as a +third reciprocal-rank leg with `weight` (k=60, default 1.0). The +reranker changes order only; the displayed score stays the dense cosine. +A reranker that errors or takes longer than six seconds is skipped for +that query, and `--no-rerank` skips it per query. Chunks are function-level where a tree-sitter grammar exists (16 languages), blank-line-snapped windows elsewhere, and each chunk is From c05273e27d896cfe54306437cee67f696d71790a Mon Sep 17 00:00:00 2001 From: Harvey Tseng Date: Fri, 4 Sep 2026 19:50:04 -0400 Subject: [PATCH 4/4] refactor: share RRF_K with the rerank leg, reject repeated indices - fuse() reads the same k as the dense and lexical legs so the three cannot drift; a response naming an index twice is an error instead of a double-counted vote - The simulator drops the rejected natural-language gate and its copy of the routing rule --- crates/scry-core/src/rerank.rs | 17 ++++++++----- crates/scry-core/src/search.rs | 2 +- eval/simulate_rerank.py | 46 ++++++++++++++++------------------ 3 files changed, 33 insertions(+), 32 deletions(-) diff --git a/crates/scry-core/src/rerank.rs b/crates/scry-core/src/rerank.rs index f1e8fb9..f25e06b 100644 --- a/crates/scry-core/src/rerank.rs +++ b/crates/scry-core/src/rerank.rs @@ -4,7 +4,7 @@ use serde::Deserialize; use crate::config::RerankConfig; -use crate::search::SearchHit; +use crate::search::{RRF_K, SearchHit}; use crate::{Error, Result}; pub struct RerankClient { @@ -59,8 +59,14 @@ impl RerankClient { let response = request.send().await?.error_for_status()?; let parsed: RerankResponse = response.json().await?; let mut results = parsed.results; - if results.iter().any(|r| r.index >= documents.len()) { - return Err(Error::Rerank("index out of range in response".to_string())); + let mut seen = std::collections::HashSet::with_capacity(results.len()); + if results + .iter() + .any(|r| r.index >= documents.len() || !seen.insert(r.index)) + { + return Err(Error::Rerank( + "bad or repeated index in response".to_string(), + )); } results.sort_by(|a, b| b.relevance_score.total_cmp(&a.relevance_score)); Ok(results) @@ -76,12 +82,11 @@ pub fn fuse( weight: f64, limit: usize, ) -> Vec { - const K: f64 = 60.0; let mut score: Vec = (0..hits.len()) - .map(|rank| 1.0 / (K + rank as f64 + 1.0)) + .map(|rank| 1.0 / (RRF_K + rank as f64 + 1.0)) .collect(); for (rank, r) in ranked.iter().enumerate() { - score[r.index] += weight / (K + rank as f64 + 1.0); + score[r.index] += weight / (RRF_K + rank as f64 + 1.0); } let mut order: Vec = (0..hits.len()).collect(); order.sort_by(|a, b| score[*b].total_cmp(&score[*a])); diff --git a/crates/scry-core/src/search.rs b/crates/scry-core/src/search.rs index baaca43..49efbbb 100644 --- a/crates/scry-core/src/search.rs +++ b/crates/scry-core/src/search.rs @@ -11,7 +11,7 @@ use crate::embed::Embedder; use crate::store::Store; use crate::{Error, Result}; -const RRF_K: f64 = 60.0; +pub const RRF_K: f64 = 60.0; const CANDIDATES: usize = 50; const JACCARD_DEDUP: f64 = 0.8; const RECENCY_BOOST: f64 = 0.1; diff --git a/eval/simulate_rerank.py b/eval/simulate_rerank.py index ae6aede..f342f70 100644 --- a/eval/simulate_rerank.py +++ b/eval/simulate_rerank.py @@ -13,21 +13,15 @@ Fusion is computed exactly where the shipped leg sits, at the route level after the pool exists: score[i] = 1/(k + pool_rank_i + 1) -+ w / (k + rerank_rank_i + 1), k = 60, then sort descending. The NL gate -applies the leg only to queries route_query marks natural language. ++ w / (k + rerank_rank_i + 1), k = 60, then sort descending; `top_n` +limits which pool positions the reranker scored. """ -import json, re, sys, tomllib +import json, sys, tomllib from pathlib import Path dump_path, eval_dir = sys.argv[1], Path(sys.argv[2]) K = 60.0 -def tokens(q): return [t for t in re.split(r'[^0-9A-Za-z_:]+', q) if len(t) >= 2] -def inner_upper(t): return t[:1].islower() and any(c.isupper() for c in t[1:]) -def natural_language(q): - ts = tokens(q); ident = any('_' in t or '::' in t or inner_upper(t) for t in ts) - return not (len(ts) <= 3 or (ident and len(ts) <= 6)) - def matches(cand, expectation): path, _, line = expectation.rpartition(':') if path and line.isdigit(): @@ -38,26 +32,27 @@ def matches(cand, expectation): for line in open(dump_path): d = json.loads(line); dumps.setdefault(d['query'], d) -def order_for(d, mode, weight): +def order_for(d, mode, weight, top_n=None): cands = d['candidates']; n = len(cands) if mode == 'fused': return list(range(n)) rerank = sorted(d['rerank'], key=lambda r: -r['score']) + if top_n is not None: + rerank = [r for r in rerank if r['index'] < top_n] if mode == 'replace': - return [r['index'] for r in rerank] + return [r['index'] for r in rerank] + [i for i in range(n) if i >= (top_n or n)] score = [1.0 / (K + i + 1) for i in range(n)] for rank, r in enumerate(rerank): score[r['index']] += weight / (K + rank + 1) return sorted(range(n), key=lambda i: -score[i]) -def evaluate(cases, mode, weight=1.0, gate=False): +def evaluate(cases, mode, weight=1.0, top_n=None): hits, rr = 0, 0.0 for case in cases: d = dumps.get(case['query']) if d is None: continue - m = 'fused' if (gate and not natural_language(case['query'])) else mode - order = order_for(d, m, weight)[:10] + order = order_for(d, mode, weight, top_n)[:10] rank = next((i for i, idx in enumerate(order) if any(matches(d['candidates'][idx], e) for e in case['expect'])), None) if rank is not None: @@ -65,14 +60,15 @@ def evaluate(cases, mode, weight=1.0, gate=False): n = len(cases) return f"{hits/n:.3f}/{rr/n:.3f}" -for set_name in ['scry', 'finance-query', 'soothfast']: - cases = tomllib.load(open(eval_dir / f'{set_name}.toml', 'rb'))['case'] - covered = sum(c['query'] in dumps for c in cases) - row = [f"{set_name:14} ({covered}/{len(cases)} dumped)", - f"fused {evaluate(cases, 'fused')}", - f"replace {evaluate(cases, 'replace')}"] - for w in (1.0, 2.0, 3.0, 5.0): - row.append(f"rrf w{w:g} {evaluate(cases, 'rrf', w)}") - row.append(f"rrf w2 nl-gate {evaluate(cases, 'rrf', 2.0, gate=True)}") - row.append(f"replace nl-gate {evaluate(cases, 'replace', gate=True)}") - print(' '.join(row)) +sets = ['scry', 'finance-query', 'soothfast'] +cases_by_set = {n: tomllib.load(open(eval_dir / f'{n}.toml', 'rb'))['case'] for n in sets} +print(f"{'arm':22}" + ''.join(f"{n:>16}" for n in sets) + f"{'sum':>16}") +def row(label, **kw): + cells = [evaluate(cases_by_set[n], **kw) for n in sets] + tot_r = sum(float(c.split('/')[0]) for c in cells); tot_m = sum(float(c.split('/')[1]) for c in cells) + print(f"{label:22}" + ''.join(f"{c:>16}" for c in cells) + f"{tot_r:>9.3f}/{tot_m:.3f}") +row('fused (no rerank)', mode='fused') +for top_n in (20, 50): + row(f'replace top{top_n}', mode='replace', top_n=top_n) + for w in (0.5, 1.0, 1.5, 2.0, 3.0): + row(f'rrf w{w:g} top{top_n}', mode='rrf', weight=w, top_n=top_n)