77 lines
2.1 KiB
Rust
77 lines
2.1 KiB
Rust
use anyhow::Result;
|
|
use reqwest::Client;
|
|
use serde_json::json;
|
|
|
|
/// Rerank response item (bare array, not OpenAI envelope).
|
|
#[derive(serde::Deserialize, Debug)]
|
|
pub struct RerankScore {
|
|
pub index: usize,
|
|
pub score: f32,
|
|
}
|
|
|
|
/// Rerank client (BAAI/bge-reranker-base via TEI).
|
|
pub struct RerankClient {
|
|
base_url: String,
|
|
api_key: String,
|
|
model: String,
|
|
timeout_secs: u64,
|
|
}
|
|
|
|
impl RerankClient {
|
|
/// Create rerank client.
|
|
pub fn new(base_url: &str, api_key: &str, model: &str) -> Result<Self> {
|
|
Ok(Self {
|
|
base_url: base_url.to_string(),
|
|
api_key: api_key.to_string(),
|
|
model: model.to_string(),
|
|
timeout_secs: 300,
|
|
})
|
|
}
|
|
|
|
/// Rerank query against texts, return scored items in score order.
|
|
/// Returns Vec<(index, score)> mapping back to input positions.
|
|
pub async fn rerank(&self, query: &str, texts: &[&str]) -> Result<Vec<(usize, f32)>> {
|
|
// Empty input returns empty without request
|
|
if texts.is_empty() {
|
|
return Ok(vec![]);
|
|
}
|
|
|
|
let url = format!("{}/rerank", self.base_url);
|
|
|
|
let client = Client::builder()
|
|
.timeout(std::time::Duration::from_secs(self.timeout_secs))
|
|
.build()?;
|
|
|
|
let payload = json!({
|
|
"query": query,
|
|
"texts": texts,
|
|
});
|
|
|
|
let response = client
|
|
.post(&url)
|
|
.header("apikey", &self.api_key)
|
|
.header("Content-Type", "application/json")
|
|
.json(&payload)
|
|
.send()
|
|
.await?;
|
|
|
|
if !response.status().is_success() {
|
|
return Err(anyhow::anyhow!("Rerank failed: {}", response.status()));
|
|
}
|
|
|
|
// Parse bare array (not OpenAI envelope)
|
|
let scores: Vec<RerankScore> = response.json().await?;
|
|
|
|
// Map back to input positions and scores
|
|
let mut results: Vec<(usize, f32)> = scores
|
|
.into_iter()
|
|
.map(|s| (s.index, s.score))
|
|
.collect();
|
|
|
|
// Sort by score descending (highest first)
|
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
|
|
|
Ok(results)
|
|
}
|
|
}
|