use crate::executor::{LlmJudgeEffort, LlmJudgeParams};
use codespan_reporting::diagnostic::{Diagnostic, Label};
use globetrotter_llm_judge::{Effort, Judge, KeyInput, LanguageText, Options};
use globetrotter_model::{
Translations,
diagnostics::{FileId, Span},
lint::{LintCode, is_allowed},
};
use std::collections::HashMap;
pub use globetrotter_llm_judge::{Error, Stats};
pub struct BarProgress(
pub indicatif::ProgressBar,
);
impl globetrotter_llm_judge::Progress for BarProgress {
fn set_length(&self, total: u64) {
self.0.set_length(total);
}
fn inc(&self, delta: u64) {
self.0.inc(delta);
}
}
pub fn judge(params: &LlmJudgeParams) -> Result<Judge, Error> {
Judge::new(Options {
base_url: params.base_url.clone(),
model: params.model.clone(),
api_key_env: params.api_key_env.clone(),
concurrency: params.concurrency,
temperature: params.temperature,
effort: params.effort.map(|effort| match effort {
LlmJudgeEffort::Low => Effort::Low,
LlmJudgeEffort::Medium => Effort::Medium,
LlmJudgeEffort::High => Effort::High,
}),
template: params.template.clone(),
min_confidence: params.min_confidence,
cache_dir: params.cache_dir.clone(),
cache_capacity: params.cache_capacity,
})
}
struct KeySpans<'a> {
file_id: FileId,
key_span: Span,
language_spans: HashMap<&'a str, Span>,
}
pub async fn stream(
judge: &Judge,
translations: &Translations,
progress: &dyn globetrotter_llm_judge::Progress,
emit: &mut dyn FnMut(Diagnostic<FileId>),
) -> Result<Stats, Error> {
let mut inputs: Vec<KeyInput<'_>> = Vec::new();
let mut spans: HashMap<&str, KeySpans<'_>> = HashMap::new();
for (key, translation) in translations {
if is_allowed(&translation.allow, LintCode::LlmDrift) {
continue;
}
let mut languages: Vec<LanguageText<'_>> = Vec::new();
let mut language_spans: HashMap<&str, Span> = HashMap::new();
for (language, text) in &translation.language {
let code: &'static str = (*language).into();
languages.push(LanguageText {
language: code,
text: text.as_ref().as_str(),
});
language_spans.insert(code, text.span.clone());
}
if languages.len() < 2 {
continue;
}
let key_str = key.as_ref().as_str();
spans.insert(
key_str,
KeySpans {
file_id: translation.file_id,
key_span: key.span.clone(),
language_spans,
},
);
inputs.push(KeyInput {
key: key_str,
languages,
});
}
if inputs.is_empty() {
return Ok(Stats::default());
}
judge
.judge(&inputs, progress, &mut |finding| {
if let Some(diagnostic) = diagnostic_for(&spans, &finding) {
emit(diagnostic);
}
})
.await
}
fn confidence_badge(confidence: f64) -> String {
use colored::Colorize;
let percent = format!("{:.0}% confident", confidence * 100.0);
let badge = if confidence >= 0.8 {
percent.red()
} else if confidence >= 0.5 {
percent.yellow()
} else {
percent.dimmed()
};
badge.to_string()
}
impl crate::executor::Executor {
pub(crate) async fn stream_llm_judge(
&self,
judge: &Judge,
translations: &std::sync::Arc<Translations>,
) -> Result<(), crate::error::Error> {
let bar = judge_progress_bar();
let progress = BarProgress(bar.clone());
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Diagnostic<FileId>>();
let judge_future = async move {
let mut sink = |diagnostic| {
let _ = tx.send(diagnostic);
};
stream(judge, translations.as_ref(), &progress, &mut sink).await
};
let drain_future = async {
let interactive = std::io::IsTerminal::is_terminal(&std::io::stderr());
while let Some(diagnostic) = rx.recv().await {
if interactive {
let rendered = self.diagnostic_printer.render(&diagnostic).await?;
bar.println(rendered.trim_end());
} else {
self.diagnostic_printer.emit(&diagnostic).await?;
}
}
Ok::<_, crate::error::Error>(())
};
let (stats, drained) = tokio::join!(judge_future, drain_future);
let stats = stats?;
drained?;
bar.finish_and_clear();
tracing::info!(
judged = stats.judged,
cached = stats.cached,
failed = stats.failed,
flagged = stats.flagged,
suppressed = stats.suppressed,
"llm judge finished"
);
Ok(())
}
}
fn judge_progress_bar() -> indicatif::ProgressBar {
use std::io::IsTerminal;
let bar = indicatif::ProgressBar::new(0);
if std::io::stderr().is_terminal() {
let style = indicatif::ProgressStyle::with_template(
"{spinner:.cyan} llm-judge: judging {bar:30.cyan/blue} {pos}/{len} keys ({eta})",
)
.unwrap_or_else(|_| indicatif::ProgressStyle::default_bar());
bar.set_style(style);
} else {
bar.set_draw_target(indicatif::ProgressDrawTarget::hidden());
}
bar.enable_steady_tick(std::time::Duration::from_millis(120));
bar
}
fn diagnostic_for(
spans: &HashMap<&str, KeySpans<'_>>,
finding: &globetrotter_llm_judge::Finding,
) -> Option<Diagnostic<FileId>> {
let key_spans = spans.get(finding.key.as_str())?;
let label = match key_spans.language_spans.get(finding.language.as_str()) {
Some(span) => Label::primary(key_spans.file_id, span.clone())
.with_message(format!("`{}` translation", finding.language)),
None => {
Label::primary(key_spans.file_id, key_spans.key_span.clone()).with_message("this key")
}
};
Some(
Diagnostic::note()
.with_code(LintCode::LlmDrift)
.with_message(format!(
"`{}`: `{}` may tell users something different ({}): {}",
finding.key,
finding.language,
confidence_badge(finding.confidence),
finding.problem
))
.with_labels(vec![label]),
)
}