Story Crater Bot
70e3f7c9a5
- EmbeddingsClient: use /v1/embeddings from gateway (nomic-ai/nomic-embed-text-v2-moe) - RerankClient: use /v1/rerank from gateway (BAAI/bge-reranker-base) - ChatClient: support gateway without auth headers (future: Bearer token) - All clients use LLM_API_BASE env var (default: https://api.riotpiao.com) - Ready for Kubernetes deployment with proper API routing
97 lines
2.7 KiB
Rust
97 lines
2.7 KiB
Rust
use anyhow::Result;
|
|
use reqwest::Client;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::time::Duration;
|
|
|
|
/// Rerank score result
|
|
#[derive(Serialize, Deserialize, Debug, Clone)]
|
|
pub struct RerankScore {
|
|
pub index: usize,
|
|
pub score: f32,
|
|
}
|
|
|
|
/// Rerank response from gateway
|
|
#[derive(Deserialize)]
|
|
struct RerankResponse {
|
|
results: Vec<RerankScore>,
|
|
}
|
|
|
|
/// Rerank client using api.riotpiao.com gateway (BAAI/bge-reranker-base model)
|
|
pub struct RerankClient {
|
|
base_url: String,
|
|
model: String,
|
|
timeout_secs: u64,
|
|
}
|
|
|
|
impl RerankClient {
|
|
/// Create rerank client pointing to gateway
|
|
pub fn new(base_url: &str, _api_key: &str, model: &str) -> Result<Self> {
|
|
Ok(Self {
|
|
base_url: base_url.to_string(),
|
|
model: model.to_string(),
|
|
timeout_secs: 300,
|
|
})
|
|
}
|
|
|
|
/// Create from environment (uses api.riotpiao.com)
|
|
pub fn from_env() -> Result<Self> {
|
|
let base_url = std::env::var("LLM_API_BASE")
|
|
.unwrap_or_else(|_| "https://api.riotpiao.com".to_string());
|
|
let model = "BAAI/bge-reranker-base".to_string();
|
|
|
|
Ok(Self {
|
|
base_url,
|
|
model,
|
|
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!("{}/v1/rerank", self.base_url);
|
|
|
|
let client = Client::builder()
|
|
.timeout(Duration::from_secs(self.timeout_secs))
|
|
.build()?;
|
|
|
|
let payload = serde_json::json!({
|
|
"model": self.model,
|
|
"query": query,
|
|
"texts": texts,
|
|
"top_k": texts.len(),
|
|
});
|
|
|
|
let response = client
|
|
.post(&url)
|
|
.header("Content-Type", "application/json")
|
|
.json(&payload)
|
|
.send()
|
|
.await?;
|
|
|
|
if !response.status().is_success() {
|
|
let error_text = response.text().await.unwrap_or_default();
|
|
return Err(anyhow::anyhow!("Rerank failed: {}", error_text));
|
|
}
|
|
|
|
// Parse gateway response (OpenAI format with results field)
|
|
let resp: RerankResponse = response.json().await?;
|
|
|
|
// Map to (index, score) and sort by score descending
|
|
let mut results: Vec<(usize, f32)> = resp
|
|
.results
|
|
.into_iter()
|
|
.map(|s| (s.index, s.score))
|
|
.collect();
|
|
|
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
|
|
|
Ok(results)
|
|
}
|
|
}
|