use crate::config::RerankerConfig;
use crate::store::{CodeBlock, CommitBlock, DocumentBlock, TextBlock};
use anyhow::Result;
pub async fn rerank_code_blocks_with_octolib(
query: &str,
mut blocks: Vec<CodeBlock>,
config: &RerankerConfig,
) -> Result<Vec<CodeBlock>> {
if !config.enabled || blocks.is_empty() {
return Ok(blocks);
}
blocks.truncate(config.top_k_candidates);
let (provider, model) = parse_reranker_model(&config.model)?;
let documents: Vec<String> = blocks.iter().map(|b| b.content.clone()).collect();
let response = octolib::reranker::rerank(
query,
documents,
&provider,
&model,
Some(config.final_top_k),
)
.await?;
let mut reranked: Vec<CodeBlock> = response
.results
.into_iter()
.filter_map(|r| {
blocks.get(r.index).map(|b| {
let mut block = b.clone();
block.distance = Some(1.0 - r.relevance_score as f32);
block
})
})
.collect();
reranked.sort_by(|a, b| {
a.distance
.unwrap_or(1.0)
.partial_cmp(&b.distance.unwrap_or(1.0))
.unwrap_or(std::cmp::Ordering::Equal)
});
Ok(reranked)
}
pub async fn rerank_doc_blocks_with_octolib(
query: &str,
mut blocks: Vec<DocumentBlock>,
config: &RerankerConfig,
) -> Result<Vec<DocumentBlock>> {
if !config.enabled || blocks.is_empty() {
return Ok(blocks);
}
blocks.truncate(config.top_k_candidates);
let (provider, model) = parse_reranker_model(&config.model)?;
let documents: Vec<String> = blocks
.iter()
.map(|b| format!("{}\n{}", b.title, b.content))
.collect();
let response = octolib::reranker::rerank(
query,
documents,
&provider,
&model,
Some(config.final_top_k),
)
.await?;
let mut reranked: Vec<DocumentBlock> = response
.results
.into_iter()
.filter_map(|r| {
blocks.get(r.index).map(|b| {
let mut block = b.clone();
block.distance = Some(1.0 - r.relevance_score as f32);
block
})
})
.collect();
reranked.sort_by(|a, b| {
a.distance
.unwrap_or(1.0)
.partial_cmp(&b.distance.unwrap_or(1.0))
.unwrap_or(std::cmp::Ordering::Equal)
});
Ok(reranked)
}
pub async fn rerank_text_blocks_with_octolib(
query: &str,
mut blocks: Vec<TextBlock>,
config: &RerankerConfig,
) -> Result<Vec<TextBlock>> {
if !config.enabled || blocks.is_empty() {
return Ok(blocks);
}
blocks.truncate(config.top_k_candidates);
let (provider, model) = parse_reranker_model(&config.model)?;
let documents: Vec<String> = blocks.iter().map(|b| b.content.clone()).collect();
let response = octolib::reranker::rerank(
query,
documents,
&provider,
&model,
Some(config.final_top_k),
)
.await?;
let mut reranked: Vec<TextBlock> = response
.results
.into_iter()
.filter_map(|r| {
blocks.get(r.index).map(|b| {
let mut block = b.clone();
block.distance = Some(1.0 - r.relevance_score as f32);
block
})
})
.collect();
reranked.sort_by(|a, b| {
a.distance
.unwrap_or(1.0)
.partial_cmp(&b.distance.unwrap_or(1.0))
.unwrap_or(std::cmp::Ordering::Equal)
});
Ok(reranked)
}
pub async fn rerank_commit_blocks_with_octolib(
query: &str,
mut blocks: Vec<CommitBlock>,
config: &RerankerConfig,
) -> Result<Vec<CommitBlock>> {
if !config.enabled || blocks.is_empty() {
return Ok(blocks);
}
blocks.truncate(config.top_k_candidates);
let (provider, model) = parse_reranker_model(&config.model)?;
let documents: Vec<String> = blocks.iter().map(|b| b.content.clone()).collect();
let response = octolib::reranker::rerank(
query,
documents,
&provider,
&model,
Some(config.final_top_k),
)
.await?;
let mut reranked: Vec<CommitBlock> = response
.results
.into_iter()
.filter_map(|r| {
blocks.get(r.index).map(|b| {
let mut block = b.clone();
block.distance = Some(1.0 - r.relevance_score as f32);
block
})
})
.collect();
reranked.sort_by(|a, b| {
a.distance
.unwrap_or(1.0)
.partial_cmp(&b.distance.unwrap_or(1.0))
.unwrap_or(std::cmp::Ordering::Equal)
});
Ok(reranked)
}
fn parse_reranker_model(model_str: &str) -> Result<(String, String)> {
match model_str.split_once(':') {
Some((provider, model)) => Ok((provider.to_string(), model.to_string())),
None => Err(anyhow::anyhow!(
"Invalid reranker model format: '{}'. Expected 'provider:model' (e.g., 'voyage:rerank-2.5')",
model_str
)),
}
}