Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into staging

This commit is contained in:
dal 2025-06-24 08:25:55 -07:00 committed by GitHub
parent 9452bb548d b8b83bc31c
commit f7df68880e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
api/libs/rerank/src/lib.rs
View File

@ -9,24 +9,15 @@ pub struct Reranker {
base_url: String,
model: String,
client: Client,
environment: String,
}
impl Reranker {
pub fn new() -> Result<Self, Box<dyn Error>> {
dotenv().ok();
let environment = env::var("ENVIRONMENT").unwrap_or_else(|_| "production".to_string());
// If local environment, we don't need these values
let (api_key, model, base_url) = if environment == "local" {
(String::new(), String::new(), String::new())
} else {
(
env::var("RERANK_API_KEY")?,
env::var("RERANK_MODEL")?,
env::var("RERANK_BASE_URL")?,
)
};
let api_key = env::var("RERANK_API_KEY")?;
let model = env::var("RERANK_MODEL")?;
let base_url = env::var("RERANK_BASE_URL")?;
let client = Client::new();
Ok(Self {
@ -34,7 +25,6 @@ impl Reranker {
base_url,
model,
client,
environment,
})
}
@ -44,13 +34,13 @@ impl Reranker {
documents: &[&str],
top_n: usize,
) -> Result<Vec<RerankResult>, Box<dyn Error>> {
// Otherwise use the remote API
let request_body = RerankRequest {
query: query.to_string(),
documents: documents.iter().map(|s| s.to_string()).collect(),
top_n,
model: self.model.clone(),
};
let response = self
.client
.post(&self.base_url)
@ -58,6 +48,7 @@ impl Reranker {
.json(&request_body)
.send()
.await?;
let response_body: RerankResponse = response.json().await?;
Ok(response_body.results)
}