use crate::jev::{client::JevClient, questions, schemas::*};
use serde_json::json;
pub async fn rank_kept(
client: &JevClient,
model: &str,
concept: &str,
kept: &[(String, String)],
any_doc: bool,
) -> anyhow::Result<Vec<(String, f64)>> {
let request = JevRequest {
state: json!({
"concept": concept,
"files": kept.iter().enumerate().map(|(i, (path, head))| json!({
"id": format!("f{i}"), "path": path, "head": head,
})).collect::<Vec<_>>(),
}),
model: model.into(),
questions: [(
"rank".to_string(),
questions::listwise_rank((0..kept.len()).map(|i| format!("f{i}")), any_doc),
)]
.into(),
};
let resp = client.system_one(&request).await?;
let Some(Answer::Choice { probabilities, .. }) = resp.answers.get("rank") else {
anyhow::bail!("expected choice answer for rank");
};
let mut ranked: Vec<_> = kept
.iter()
.enumerate()
.map(|(i, (path, _))| {
let p = probabilities.get(&format!("f{i}")).copied().unwrap_or(0.0);
(path.clone(), p)
})
.collect();
ranked.sort_by(|a, b| b.1.total_cmp(&a.1));
Ok(ranked)
}