use std::sync::Arc;
use smooth_operator::rerank::{LexicalReranker, NoopReranker, Reranker};
use smooth_operator_adapter_postgres::{GatewayReranker, DEFAULT_RERANK_MODEL};
#[derive(Debug, Clone)]
pub struct RerankerConfig {
pub gateway_url: String,
pub gateway_key: Option<String>,
pub model: String,
pub mode: RerankMode,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RerankMode {
#[default]
Off,
Gateway,
Lexical,
}
impl RerankMode {
#[must_use]
pub fn parse(s: &str) -> Self {
match s.trim().to_ascii_lowercase().as_str() {
"gateway" | "on" | "1" | "true" => Self::Gateway,
"lexical" => Self::Lexical,
_ => Self::Off,
}
}
}
impl RerankerConfig {
#[must_use]
pub fn mode_from_env() -> RerankMode {
std::env::var("SMOOTH_AGENT_RERANK")
.ok()
.map(|s| RerankMode::parse(&s))
.unwrap_or_default()
}
#[must_use]
pub fn from_gateway(gateway_url: impl Into<String>, gateway_key: Option<String>) -> Self {
Self {
gateway_url: gateway_url.into(),
gateway_key,
model: DEFAULT_RERANK_MODEL.to_string(),
mode: Self::mode_from_env(),
}
}
#[must_use]
pub fn from_server_config(config: &crate::config::ServerConfig) -> Self {
Self::from_gateway(config.gateway_url.clone(), config.gateway_key.clone())
}
}
#[must_use]
pub fn build_reranker(config: &RerankerConfig) -> Option<Arc<dyn Reranker>> {
match config.mode {
RerankMode::Off => {
tracing::info!("rerank stage disabled (default) — retrieval order unchanged");
None
}
RerankMode::Gateway => match &config.gateway_key {
Some(key) if !key.trim().is_empty() => {
tracing::info!(
model = %config.model,
"using GatewayReranker (cross-encoder /v1/rerank) for retrieval reorder"
);
Some(Arc::new(GatewayReranker::new(
config.gateway_url.clone(),
key.clone(),
config.model.clone(),
)))
}
_ => {
tracing::warn!(
"SMOOTH_AGENT_RERANK=gateway but no gateway key — \
falling back to the offline LexicalReranker"
);
Some(Arc::new(LexicalReranker::new()))
}
},
RerankMode::Lexical => {
tracing::info!(
"using offline LexicalReranker (BM25-ish, no network) for retrieval reorder"
);
Some(Arc::new(LexicalReranker::new()))
}
}
}
#[must_use]
pub fn noop_reranker() -> Arc<dyn Reranker> {
Arc::new(NoopReranker)
}
#[cfg(test)]
mod tests {
use super::*;
fn cfg(mode: RerankMode, key: Option<&str>) -> RerankerConfig {
RerankerConfig {
gateway_url: "https://example.test/v1".into(),
gateway_key: key.map(str::to_string),
model: DEFAULT_RERANK_MODEL.to_string(),
mode,
}
}
#[test]
fn default_mode_is_off_yielding_no_reranker() {
assert!(build_reranker(&cfg(RerankMode::Off, Some("sk-test"))).is_none());
assert!(build_reranker(&cfg(RerankMode::default(), None)).is_none());
}
#[test]
fn gateway_mode_with_key_selects_a_reranker() {
assert!(build_reranker(&cfg(RerankMode::Gateway, Some("sk-test"))).is_some());
}
#[test]
fn gateway_mode_without_key_falls_back_to_lexical() {
assert!(build_reranker(&cfg(RerankMode::Gateway, None)).is_some());
assert!(build_reranker(&cfg(RerankMode::Gateway, Some(" "))).is_some());
}
#[test]
fn lexical_mode_selects_a_reranker_without_a_key() {
assert!(build_reranker(&cfg(RerankMode::Lexical, None)).is_some());
}
#[test]
fn rerank_mode_parse() {
assert_eq!(RerankMode::parse("gateway"), RerankMode::Gateway);
assert_eq!(RerankMode::parse("ON"), RerankMode::Gateway);
assert_eq!(RerankMode::parse("lexical"), RerankMode::Lexical);
assert_eq!(RerankMode::parse("off"), RerankMode::Off);
assert_eq!(RerankMode::parse(""), RerankMode::Off);
assert_eq!(RerankMode::parse("nonsense"), RerankMode::Off);
}
}