tenant.openagents/omega

1use crate::{
2    PredictArgs, PredictionProvider,
3    example::Example,
4    format_prompt::TeacherPrompt,
5    headless::EpAppState,
6    parse_output::parse_prediction_output,
7    predict::run_prediction,
8    progress::{ExampleProgress, Step},
9};
10use anyhow::Context as _;
11use edit_prediction_context::limit_retrieved_context_to_bytes;
12use edit_prediction_metrics::{
13    ActualPredictionCursor, Excerpt, PredictionReversalContext, PredictionScoringInput,
14};
15use gpui::{AppContext as _, AsyncApp};
16use std::fs::File;
17use std::io::BufWriter;
18use std::path::Path;
19use std::sync::Arc;
20use zeta_prompt::{ContextSource, RelatedFile};
21
22pub const EVAL_RELATED_CONTEXT_TOKENS_LIMIT: usize = 4000;
23
24pub async fn run_scoring(
25    example: &mut Example,
26    args: &PredictArgs,
27    app_state: Arc<EpAppState>,
28    example_progress: &ExampleProgress,
29    cx: AsyncApp,
30    allow_missing_predictions: bool,
31    retrieved_context_byte_limit: Option<usize>,
32    context_source_filter: Option<Vec<ContextSource>>,
33) -> anyhow::Result<()> {
34    if !(allow_missing_predictions && args.provider.is_none() && example.predictions.is_empty()) {
35        run_prediction(example, args, app_state, example_progress, cx.clone()).await?;
36    }
37
38    let progress = example_progress.start(Step::Score);
39
40    progress.set_substatus("computing metrics");
41    let example_for_scoring = example.clone();
42    example.score = cx
43        .background_spawn(async move {
44            let prompt_inputs = example_for_scoring
45                .prompt_inputs
46                .as_ref()
47                .context("prompt_inputs is required for scoring - run prediction first or ensure JSON includes prompt_inputs")?;
48            let original_text: &str = prompt_inputs.cursor_excerpt.as_ref();
49            let expected_patches_with_cursors = example_for_scoring
50                .spec
51                .expected_patches_with_cursor_positions();
52
53            let old_editable_region = if let Some(p) = example_for_scoring.prompt.as_ref() {
54                if matches!(
55                    p.provider,
56                    PredictionProvider::Teacher(_, _) | PredictionProvider::TeacherNonBatching(_, _)
57                ) {
58                    Some(
59                        TeacherPrompt::extract_editable_region(&p.input)?
60                            .replace(TeacherPrompt::USER_CURSOR_MARKER, ""),
61                    )
62                } else {
63                    None
64                }
65            } else {
66                None
67            };
68
69            let cursor_path = example_for_scoring.spec.cursor_path.as_ref();
70            let context = context_excerpts(
71                &example_for_scoring,
72                prompt_inputs,
73                retrieved_context_byte_limit,
74                context_source_filter.as_deref(),
75            );
76
77            let prepared_expected_patches = match edit_prediction_metrics::prepare_expected_patches(
78                &expected_patches_with_cursors,
79                original_text,
80                old_editable_region.as_deref(),
81            ) {
82                Ok(prepared_expected_patches) => prepared_expected_patches,
83                Err(_) if !context.is_empty() => expected_patches_with_cursors
84                    .iter()
85                    .map(|(patch, cursor_offset)| edit_prediction_metrics::PreparedExpectedPatch {
86                        patch: patch.clone(),
87                        text: original_text.to_string(),
88                        cursor_editable_region_offset: *cursor_offset,
89                    })
90                    .collect(),
91                Err(error) => {
92                    return Err(error).with_context(|| {
93                        format!(
94                            "Expected patch did not apply for {}",
95                            example_for_scoring.spec.name
96                        )
97                    });
98                }
99            };
100
101            let mut scores = vec![];
102            if allow_missing_predictions && example_for_scoring.predictions.is_empty() {
103                scores.push(edit_prediction_metrics::score_prediction(
104                    PredictionScoringInput {
105                        original_text,
106                        expected_patches: &prepared_expected_patches,
107                        actual_patch: None,
108                        actual_cursor: None,
109                        reversal_context: Some(PredictionReversalContext {
110                            edit_history: &prompt_inputs.events,
111                            excerpt_start_row: prompt_inputs.excerpt_start_row,
112                            cursor_path,
113                        }),
114                        cumulative_logprob: None,
115                        avg_logprob: None,
116                        context: Some(&context),
117                    },
118                ));
119            }
120
121            for prediction in &example_for_scoring.predictions {
122                let actual_patch = prediction.actual_patch.clone().or_else(|| {
123                    parse_prediction_output(
124                        &example_for_scoring,
125                        &prediction.actual_output,
126                        prediction.provider,
127                    )
128                    .ok()
129                    .map(|(patch, _)| patch)
130                });
131
132                let actual_cursor = prediction.actual_cursor.as_ref().map(|cursor| {
133                    ActualPredictionCursor {
134                        row: cursor.row,
135                        editable_region_offset: cursor.editable_region_offset,
136                    }
137                });
138
139                scores.push(edit_prediction_metrics::score_prediction(
140                    PredictionScoringInput {
141                        original_text,
142                        expected_patches: &prepared_expected_patches,
143                        actual_patch: actual_patch.as_deref(),
144                        actual_cursor,
145                        reversal_context: Some(PredictionReversalContext {
146                            edit_history: &prompt_inputs.events,
147                            excerpt_start_row: prompt_inputs.excerpt_start_row,
148                            cursor_path,
149                        }),
150                        cumulative_logprob: prediction.cumulative_logprob,
151                        avg_logprob: prediction.avg_logprob,
152                        context: Some(&context),
153                    },
154                ));
155            }
156
157            anyhow::Ok(scores)
158        })
159        .await?;
160    Ok(())
161}
162
163pub fn run_context_coverage_scoring(
164    example: &mut Example,
165    example_progress: &ExampleProgress,
166    retrieved_context_byte_limit: Option<usize>,
167    context_source_filter: Option<&[ContextSource]>,
168) -> anyhow::Result<()> {
169    let progress = example_progress.start(Step::Score);
170
171    progress.set_substatus("computing context coverage");
172    let prompt_inputs = example
173        .prompt_inputs
174        .as_ref()
175        .context("prompt_inputs is required for context coverage scoring")?;
176    let context = context_excerpts(
177        example,
178        prompt_inputs,
179        retrieved_context_byte_limit,
180        context_source_filter,
181    );
182
183    let editable_context_coverage = example
184        .spec
185        .expected_patches_with_cursor_positions()
186        .iter()
187        .map(|(expected_patch, _)| {
188            edit_prediction_metrics::editable_context_coverage(expected_patch, &context)
189        })
190        .max_by(|left, right| {
191            left.lines_f1
192                .total_cmp(&right.lines_f1)
193                .then_with(|| left.files_f1.total_cmp(&right.files_f1))
194        });
195
196    let mut score = edit_prediction_metrics::PredictionScore::zero();
197    score.editable_context_coverage = editable_context_coverage;
198    example.score = vec![score];
199
200    Ok(())
201}
202
203fn context_excerpts(
204    _example: &Example,
205    prompt_inputs: &zeta_prompt::Zeta2PromptInput,
206    retrieved_context_byte_limit: Option<usize>,
207    context_source_filter: Option<&[ContextSource]>,
208) -> Vec<Excerpt> {
209    let mut context = Vec::new();
210
211    if let Some(excerpt_start_row) = prompt_inputs.excerpt_start_row {
212        let row_count = prompt_inputs.cursor_excerpt.lines().count() as u32;
213
214        context.push(Excerpt {
215            path: prompt_inputs.cursor_path.to_string_lossy().to_string(),
216            row_range: excerpt_start_row..excerpt_start_row.saturating_add(row_count),
217            content: prompt_inputs.cursor_excerpt.to_string(),
218        });
219    }
220
221    if let Some(related_files) = &prompt_inputs.related_files {
222        let related_files = filtered_related_files(related_files, context_source_filter);
223        let related_files = if let Some(max_bytes) = retrieved_context_byte_limit {
224            limit_retrieved_context_to_bytes(&related_files, max_bytes)
225        } else {
226            related_files
227        };
228        for related_file in &related_files {
229            for excerpt in &related_file.excerpts {
230                // First component is a project name which is not present in expected patch, skip it
231                let path = related_file
232                    .path
233                    .iter()
234                    .skip(1)
235                    .collect::<std::path::PathBuf>()
236                    .to_string_lossy()
237                    .to_string();
238                context.push(Excerpt {
239                    path,
240                    row_range: excerpt.row_range.clone(),
241                    content: excerpt.text.to_string(),
242                });
243            }
244        }
245    }
246
247    context
248}
249
250fn filtered_related_files(
251    related_files: &[RelatedFile],
252    context_source_filter: Option<&[ContextSource]>,
253) -> Vec<RelatedFile> {
254    let Some(context_source_filter) = context_source_filter else {
255        return related_files.to_vec();
256    };
257
258    related_files
259        .iter()
260        .filter_map(|related_file| {
261            let excerpts = related_file
262                .excerpts
263                .iter()
264                .filter(|excerpt| context_source_filter.contains(&excerpt.context_source))
265                .cloned()
266                .collect::<Vec<_>>();
267            if excerpts.is_empty() {
268                None
269            } else {
270                Some(RelatedFile {
271                    path: related_file.path.clone(),
272                    max_row: related_file.max_row,
273                    excerpts,
274                    in_open_source_repo: related_file.in_open_source_repo,
275                })
276            }
277        })
278        .collect()
279}
280
281fn retrieved_context_bytes(
282    example: &Example,
283    retrieved_context_byte_limit: Option<usize>,
284    context_source_filter: Option<&[ContextSource]>,
285) -> Option<usize> {
286    let related_files = example.prompt_inputs.as_ref()?.related_files.as_ref()?;
287    let related_files = filtered_related_files(related_files, context_source_filter);
288    let related_files = if let Some(max_bytes) = retrieved_context_byte_limit {
289        limit_retrieved_context_to_bytes(&related_files, max_bytes)
290    } else {
291        related_files
292    };
293    Some(
294        related_files
295            .iter()
296            .flat_map(|file| file.excerpts.iter())
297            .map(|excerpt| excerpt.text.len())
298            .sum::<usize>(),
299    )
300}
301
302pub fn print_report(
303    examples: &[Example],
304    verbose: bool,
305    context_only: bool,
306    retrieved_context_byte_limit: Option<usize>,
307    context_source_filter: Option<&[ContextSource]>,
308) {
309    const MAX_EXAMPLES_DEFAULT: usize = 20;
310    const LINE_WIDTH: usize = 101;
311
312    if context_only {
313        print_context_coverage_report(
314            examples,
315            verbose,
316            retrieved_context_byte_limit,
317            context_source_filter,
318        );
319        return;
320    }
321
322    let separator = "─".repeat(LINE_WIDTH);
323
324    println!("{}", separator);
325    println!(
326        "{:<40} {:>8} {:>5} {:>7} {:>7} {:>7} {:>7} {:>6} {:>5}",
327        "Example", "DeltaChrF", "Brace", "F1", "Revert", "QaRev", "QaConf", "Cursor", "WrgER"
328    );
329    println!("{}", separator);
330
331    let mut patch_inserted_tokens: Vec<usize> = Vec::new();
332    let mut patch_deleted_tokens: Vec<usize> = Vec::new();
333    let mut predictions_with_patch: usize = 0;
334
335    let mut printed_lines: usize = 0;
336    let mut skipped_lines: usize = 0;
337
338    for example in examples {
339        for (score_idx, score) in example.score.iter().enumerate() {
340            let exact_lines = score.exact_lines_counts();
341
342            // Get QA results for this prediction if available
343            let qa_result = example.qa.get(score_idx).and_then(|q| q.as_ref());
344            let qa_reverts_str = qa_result
345                .and_then(|q| q.reverts_edits)
346                .map(|v| if v { "yes" } else { "no" })
347                .unwrap_or("-");
348            let qa_conf_str = qa_result
349                .and_then(|q| q.confidence)
350                .map(|v| format!("{}", v))
351                .unwrap_or("-".to_string());
352
353            // Format wrong editable region metric
354            let wrong_er_str = match score.wrong_editable_region {
355                Some(true) => "✗",
356                Some(false) => "",
357                None => "",
358            };
359
360            // Format cursor metric
361            let cursor_str = match (score.cursor_exact_match, score.cursor_distance) {
362                (Some(true), _) => "✓".to_string(),
363                (Some(false), Some(dist)) => format!("±{}", dist),
364                (Some(false), None) => "✗".to_string(),
365                (None, _) => "-".to_string(),
366            };
367
368            if verbose || printed_lines < MAX_EXAMPLES_DEFAULT {
369                println!(
370                    "{:<40} {:>8.2} {:>5} {:>6.1}% {:>6.1}% {:>7} {:>7} {:>6} {:>5}",
371                    truncate_name(&example.spec.name, 40),
372                    score.delta_chr_f,
373                    score.braces_disbalance,
374                    exact_lines.f1() * 100.0,
375                    score.reversal_ratio * 100.0,
376                    qa_reverts_str,
377                    qa_conf_str,
378                    cursor_str,
379                    wrong_er_str
380                );
381                printed_lines += 1;
382            } else {
383                skipped_lines += 1;
384            }
385
386            // Token change percentiles need the raw per-prediction values, so
387            // they are collected here rather than in `compute_summary`.
388            let has_patch = example
389                .predictions
390                .get(score_idx)
391                .and_then(|p| p.actual_patch.as_ref())
392                .is_some_and(|p| !p.is_empty());
393            if has_patch {
394                predictions_with_patch += 1;
395                patch_inserted_tokens.push(score.inserted_tokens);
396                patch_deleted_tokens.push(score.deleted_tokens);
397            }
398        }
399    }
400
401    if skipped_lines > 0 {
402        println!(
403            "{:<40} (use --verbose to see all {} examples)",
404            format!("... and {} more", skipped_lines),
405            printed_lines + skipped_lines
406        );
407    }
408    println!("{}", separator);
409
410    let summary = compute_summary(
411        examples,
412        retrieved_context_byte_limit,
413        context_source_filter,
414    );
415
416    if summary.total_examples > 0 {
417        let total_scores = summary.total_examples;
418        let format_rate = |rate: Option<f32>, precision: usize| {
419            rate.map(|rate| format!("{:.*}%", precision, rate * 100.0))
420                .unwrap_or_else(|| "-".to_string())
421        };
422        let qa_reverts_str = format_rate(summary.qa_avg_reverts_edits, 1);
423        let qa_conf_str = summary
424            .qa_avg_confidence
425            .map(|confidence| format!("{:.1}", confidence))
426            .unwrap_or_else(|| "-".to_string());
427        let cursor_str = format_rate(summary.cursor_exact_match_rate, 0);
428        let wrong_er_str = format_rate(summary.wrong_editable_region_rate, 2);
429
430        println!(
431            "{:<40} {:>8.2} {:>5.1} {:>6.1}% {:>6.1}% {:>7} {:>7} {:>6} {:>5}",
432            "TOTAL / AVERAGE",
433            summary.avg_delta_chr_f,
434            summary.avg_braces_disbalance,
435            summary.exact_lines_f1 * 100.0,
436            summary.avg_reversal_ratio * 100.0,
437            qa_reverts_str,
438            qa_conf_str,
439            cursor_str,
440            wrong_er_str
441        );
442        println!("{}", separator);
443        println!(
444            "Delta chrF (β={:.1}): TP={}, FP={}, FN={}, P={:.1}%, R={:.1}%",
445            summary.delta_chr_f_beta,
446            summary.delta_chr_f_true_positives,
447            summary.delta_chr_f_false_positives,
448            summary.delta_chr_f_false_negatives,
449            summary.delta_chr_f_precision * 100.0,
450            summary.delta_chr_f_recall * 100.0
451        );
452
453        if let Some(avg_distance) = summary.cursor_avg_distance {
454            println!(
455                "Cursor: {}/{} exact matches ({:.0}%), avg distance: {:.1} bytes",
456                summary.cursor_exact_matches.unwrap_or(0),
457                summary.cursor_total_evaluated.unwrap_or(0),
458                summary.cursor_exact_match_rate.unwrap_or(0.0) * 100.0,
459                avg_distance
460            );
461        }
462
463        if let (Some(count), Some(rate)) = (
464            summary.isolated_whitespace_count,
465            summary.isolated_whitespace_rate,
466        ) {
467            println!(
468                "Isolated whitespace changes: {}/{} ({:.1}%)",
469                count,
470                total_scores,
471                rate * 100.0
472            );
473        }
474
475        if let (Some(avg_kept_rate), Some(evaluated)) =
476            (summary.avg_kept_rate, summary.kept_rate_examples)
477        {
478            println!(
479                "Kept rate: {:.1}% avg ({} evaluated, kept chars: {}, correctly deleted chars: {}, discarded chars: {})",
480                avg_kept_rate * 100.0,
481                evaluated,
482                summary.total_kept_chars.unwrap_or(0),
483                summary.total_correctly_deleted_chars.unwrap_or(0),
484                summary.total_discarded_chars.unwrap_or(0)
485            );
486        }
487        if let (Some(avg_recall_rate), Some(evaluated)) =
488            (summary.avg_recall_rate, summary.recall_rate_examples)
489        {
490            println!(
491                "Recall rate: {:.1}% avg ({} evaluated)",
492                avg_recall_rate * 100.0,
493                evaluated
494            );
495        }
496        if let (Some(avg_bytes), Some(example_count)) = (
497            summary.avg_retrieved_context_bytes,
498            summary.retrieved_context_examples,
499        ) {
500            println!(
501                "Retrieved context size: {:.0} bytes avg ({} examples)",
502                avg_bytes, example_count
503            );
504        }
505
506        print_prf_line(
507            "Editable context lines",
508            summary.editable_context_examples,
509            summary.avg_editable_context_lines_precision,
510            summary.avg_editable_context_lines_recall,
511            summary.avg_editable_context_lines_f1,
512            summary.editable_context_lines_tp,
513            summary.editable_context_lines_fp,
514            summary.editable_context_lines_fn,
515        );
516        print_prf_line(
517            "Editable context files",
518            summary.editable_context_examples,
519            summary.avg_editable_context_files_precision,
520            summary.avg_editable_context_files_recall,
521            summary.avg_editable_context_files_f1,
522            summary.editable_context_files_tp,
523            summary.editable_context_files_fp,
524            summary.editable_context_files_fn,
525        );
526        print_prf_line(
527            "Jump location lines",
528            summary.jump_location_examples,
529            summary.avg_jump_location_lines_precision,
530            summary.avg_jump_location_lines_recall,
531            summary.avg_jump_location_lines_f1,
532            summary.jump_location_lines_tp,
533            summary.jump_location_lines_fp,
534            summary.jump_location_lines_fn,
535        );
536        print_prf_line(
537            "Jump location files",
538            summary.jump_location_examples,
539            summary.avg_jump_location_files_precision,
540            summary.avg_jump_location_files_recall,
541            summary.avg_jump_location_files_f1,
542            summary.jump_location_files_tp,
543            summary.jump_location_files_fp,
544            summary.jump_location_files_fn,
545        );
546
547        // Print token change percentile summary (only for predictions with a patch)
548        if !patch_inserted_tokens.is_empty() {
549            patch_inserted_tokens.sort_unstable();
550            patch_deleted_tokens.sort_unstable();
551            let mut patch_total_tokens: Vec<usize> = patch_inserted_tokens
552                .iter()
553                .zip(patch_deleted_tokens.iter())
554                .map(|(i, d)| i + d)
555                .collect();
556            patch_total_tokens.sort_unstable();
557
558            let patch_rate = predictions_with_patch as f32 / total_scores as f32 * 100.0;
559            println!();
560            println!(
561                "Token changes ({}/{} predictions produced a patch, {:.1}% — table includes only those)",
562                predictions_with_patch, total_scores, patch_rate
563            );
564            println!(
565                "{:<20} {:>8} {:>8} {:>8} {:>8} {:>8}",
566                "", "p25", "p50", "p75", "p90", "p99"
567            );
568            println!("{}", "─".repeat(LINE_WIDTH));
569            println!(
570                "{:<20} {:>8} {:>8} {:>8} {:>8} {:>8}",
571                "Inserted tokens",
572                percentile(&patch_inserted_tokens, 25),
573                percentile(&patch_inserted_tokens, 50),
574                percentile(&patch_inserted_tokens, 75),
575                percentile(&patch_inserted_tokens, 90),
576                percentile(&patch_inserted_tokens, 99),
577            );
578            println!(
579                "{:<20} {:>8} {:>8} {:>8} {:>8} {:>8}",
580                "Deleted tokens",
581                percentile(&patch_deleted_tokens, 25),
582                percentile(&patch_deleted_tokens, 50),
583                percentile(&patch_deleted_tokens, 75),
584                percentile(&patch_deleted_tokens, 90),
585                percentile(&patch_deleted_tokens, 99),
586            );
587            println!(
588                "{:<20} {:>8} {:>8} {:>8} {:>8} {:>8}",
589                "Total tokens",
590                percentile(&patch_total_tokens, 25),
591                percentile(&patch_total_tokens, 50),
592                percentile(&patch_total_tokens, 75),
593                percentile(&patch_total_tokens, 90),
594                percentile(&patch_total_tokens, 99),
595            );
596        }
597    }
598
599    println!("\n");
600}
601
602fn print_context_coverage_report(
603    examples: &[Example],
604    verbose: bool,
605    retrieved_context_byte_limit: Option<usize>,
606    context_source_filter: Option<&[ContextSource]>,
607) {
608    const MAX_EXAMPLES_DEFAULT: usize = 20;
609    const LINE_WIDTH: usize = 120;
610
611    let separator = "─".repeat(LINE_WIDTH);
612    println!("{}", separator);
613    println!(
614        "{:<40} {:>6} {:>6} {:>6} {:>5} {:>5} {:>5} {:>6} {:>6} {:>6} {:>5} {:>5} {:>5}",
615        "Example",
616        "LineP",
617        "LineR",
618        "LineF1",
619        "LTP",
620        "LFP",
621        "LFN",
622        "FileP",
623        "FileR",
624        "FileF1",
625        "FTP",
626        "FFP",
627        "FFN"
628    );
629    println!("{}", separator);
630
631    let mut printed_lines = 0;
632    let mut skipped_lines = 0;
633
634    for example in examples {
635        for score in &example.score {
636            let Some(coverage) = &score.editable_context_coverage else {
637                continue;
638            };
639
640            if verbose || printed_lines < MAX_EXAMPLES_DEFAULT {
641                println!(
642                    "{:<40} {:>5.1}% {:>5.1}% {:>5.1}% {:>5} {:>5} {:>5} {:>5.1}% {:>5.1}% {:>5.1}% {:>5} {:>5} {:>5}",
643                    truncate_name(&example.spec.name, 40),
644                    coverage.lines_precision * 100.0,
645                    coverage.lines_recall * 100.0,
646                    coverage.lines_f1 * 100.0,
647                    coverage.lines_tp,
648                    coverage.lines_fp,
649                    coverage.lines_fn,
650                    coverage.files_precision * 100.0,
651                    coverage.files_recall * 100.0,
652                    coverage.files_f1 * 100.0,
653                    coverage.files_tp,
654                    coverage.files_fp,
655                    coverage.files_fn
656                );
657                printed_lines += 1;
658            } else {
659                skipped_lines += 1;
660            }
661        }
662    }
663
664    if skipped_lines > 0 {
665        println!(
666            "{:<40} (use --verbose to see all {} examples)",
667            format!("... and {} more", skipped_lines),
668            printed_lines + skipped_lines
669        );
670    }
671
672    println!("{}", separator);
673
674    let summary = compute_summary(
675        examples,
676        retrieved_context_byte_limit,
677        context_source_filter,
678    );
679
680    if let Some(total_scores) = summary.editable_context_examples {
681        println!(
682            "{:<40} {:>5.1}% {:>5.1}% {:>5.1}% {:>5} {:>5} {:>5} {:>5.1}% {:>5.1}% {:>5.1}% {:>5} {:>5} {:>5}",
683            "TOTAL / AVERAGE",
684            summary.avg_editable_context_lines_precision.unwrap_or(0.0) * 100.0,
685            summary.avg_editable_context_lines_recall.unwrap_or(0.0) * 100.0,
686            summary.avg_editable_context_lines_f1.unwrap_or(0.0) * 100.0,
687            summary.editable_context_lines_tp.unwrap_or(0),
688            summary.editable_context_lines_fp.unwrap_or(0),
689            summary.editable_context_lines_fn.unwrap_or(0),
690            summary.avg_editable_context_files_precision.unwrap_or(0.0) * 100.0,
691            summary.avg_editable_context_files_recall.unwrap_or(0.0) * 100.0,
692            summary.avg_editable_context_files_f1.unwrap_or(0.0) * 100.0,
693            summary.editable_context_files_tp.unwrap_or(0),
694            summary.editable_context_files_fp.unwrap_or(0),
695            summary.editable_context_files_fn.unwrap_or(0)
696        );
697        println!("{}", separator);
698        println!(
699            "Evaluated editable context coverage for {} examples",
700            total_scores
701        );
702        if let (Some(avg_bytes), Some(example_count)) = (
703            summary.avg_retrieved_context_bytes,
704            summary.retrieved_context_examples,
705        ) {
706            println!(
707                "Retrieved context size: {:.0} bytes avg ({} examples)",
708                avg_bytes, example_count
709            );
710        }
711    }
712
713    println!("\n");
714}
715
716/// Print one "P/R/F1 avg + pooled TP/FP/FN" summary line, or nothing when
717/// the metric was never evaluated.
718fn print_prf_line(
719    label: &str,
720    examples: Option<usize>,
721    precision: Option<f64>,
722    recall: Option<f64>,
723    f1: Option<f64>,
724    true_positives: Option<usize>,
725    false_positives: Option<usize>,
726    false_negatives: Option<usize>,
727) {
728    let (
729        Some(examples),
730        Some(precision),
731        Some(recall),
732        Some(f1),
733        Some(true_positives),
734        Some(false_positives),
735        Some(false_negatives),
736    ) = (
737        examples,
738        precision,
739        recall,
740        f1,
741        true_positives,
742        false_positives,
743        false_negatives,
744    )
745    else {
746        return;
747    };
748    println!(
749        "{}: P={:.1}%, R={:.1}%, F1={:.1}% avg ({} evaluated, TP={}, FP={}, FN={})",
750        label,
751        precision * 100.0,
752        recall * 100.0,
753        f1 * 100.0,
754        examples,
755        true_positives,
756        false_positives,
757        false_negatives
758    );
759}
760
761fn percentile(sorted_values: &[usize], p: usize) -> usize {
762    if sorted_values.is_empty() {
763        return 0;
764    }
765    let idx = (p as f64 / 100.0 * (sorted_values.len() as f64 - 1.0)).round() as usize;
766    sorted_values[idx.min(sorted_values.len() - 1)]
767}
768
769fn truncate_name(name: &str, max_len: usize) -> String {
770    if name.len() <= max_len {
771        name.to_string()
772    } else {
773        format!("{}...", &name[..max_len - 3])
774    }
775}
776
777pub type SummaryJson = edit_prediction_metrics::SummaryJson;
778
779pub fn compute_summary(
780    examples: &[Example],
781    retrieved_context_byte_limit: Option<usize>,
782    context_source_filter: Option<&[ContextSource]>,
783) -> SummaryJson {
784    edit_prediction_metrics::compute_summary(examples.iter().flat_map(|example| {
785        let retrieved_context_bytes =
786            retrieved_context_bytes(example, retrieved_context_byte_limit, context_source_filter);
787        example
788            .score
789            .iter()
790            .enumerate()
791            .map(move |(score_idx, score)| {
792                let qa = example
793                    .qa
794                    .get(score_idx)
795                    .and_then(|qa| qa.as_ref())
796                    .map(|qa| edit_prediction_metrics::QaSummaryData {
797                        reverts_edits: qa.reverts_edits,
798                        confidence: qa.confidence,
799                    });
800                let retrieved_context_bytes = (score_idx == 0)
801                    .then_some(retrieved_context_bytes)
802                    .flatten();
803
804                edit_prediction_metrics::PredictionSummaryInput {
805                    score,
806                    qa,
807                    retrieved_context_bytes,
808                }
809            })
810    }))
811}
812
813pub fn write_summary_json(
814    examples: &[Example],
815    path: &Path,
816    retrieved_context_byte_limit: Option<usize>,
817    context_source_filter: Option<&[ContextSource]>,
818) -> anyhow::Result<()> {
819    let summary = compute_summary(
820        examples,
821        retrieved_context_byte_limit,
822        context_source_filter,
823    );
824    let file = File::create(path)
825        .with_context(|| format!("Failed to create summary JSON file: {}", path.display()))?;
826    let writer = BufWriter::new(file);
827    serde_json::to_writer_pretty(writer, &summary)
828        .with_context(|| format!("Failed to write summary JSON to: {}", path.display()))?;
829    eprintln!("Wrote summary JSON to: {}", path.display());
830    Ok(())
831}
832
833#[cfg(test)]
834mod tests {
835    use super::*;
836    use edit_prediction::example_spec::ExampleSpec;
837    use edit_prediction_metrics::PredictionScore;
838    use std::path::Path;
839    use zeta_prompt::{ExcerptRanges, RelatedExcerpt, Zeta2PromptInput};
840
841    #[test]
842    fn summary_includes_limited_filtered_retrieved_context_bytes_once_per_example() {
843        let examples = vec![
844            example_with_related_files(
845                Some(vec![RelatedFile {
846                    path: Path::new("project/src/lib.rs").into(),
847                    max_row: 10,
848                    excerpts: vec![
849                        related_excerpt("abcd", 0..1, 0, ContextSource::CurrentFile),
850                        related_excerpt("ignored by source filter", 1..2, 1, ContextSource::Lsp),
851                        related_excerpt("efghij", 2..3, 2, ContextSource::CurrentFile),
852                    ],
853                    in_open_source_repo: false,
854                }]),
855                2,
856            ),
857            example_with_related_files(None, 1),
858        ];
859
860        let summary = compute_summary(&examples, Some(10), Some(&[ContextSource::CurrentFile]));
861
862        assert_eq!(summary.total_examples, 3);
863        assert_eq!(summary.avg_retrieved_context_bytes, Some(10.0));
864        assert_eq!(summary.total_retrieved_context_bytes, Some(10));
865        assert_eq!(summary.retrieved_context_examples, Some(1));
866    }
867
868    fn example_with_related_files(
869        related_files: Option<Vec<RelatedFile>>,
870        score_count: usize,
871    ) -> Example {
872        Example {
873            spec: ExampleSpec {
874                name: "example".to_string(),
875                repository_url: "https://github.com/zed-industries/zed.git".to_string(),
876                revision: "revision".to_string(),
877                tags: Vec::new(),
878                reasoning: None,
879                uncommitted_diff: String::new(),
880                recently_opened_files: Vec::new(),
881                recently_viewed_files: Vec::new(),
882                uncommitted_diff_contains_edit_history: false,
883                cursor_path: Path::new("project/src/main.rs").into(),
884                cursor_position: String::new(),
885                edit_history: String::new(),
886                expected_patches: Vec::new(),
887                rejected_patch: None,
888                telemetry: None,
889                human_feedback: Vec::new(),
890                rating: None,
891            },
892            prompt_inputs: Some(Zeta2PromptInput {
893                cursor_path: Path::new("project/src/main.rs").into(),
894                cursor_excerpt: "".into(),
895                cursor_offset_in_excerpt: 0,
896                excerpt_start_row: None,
897                events: Vec::new(),
898                related_files,
899                active_buffer_diagnostics: Vec::new(),
900                excerpt_ranges: ExcerptRanges::default(),
901                syntax_ranges: None,
902                in_open_source_repo: false,
903                can_collect_data: false,
904                repo_url: None,
905            }),
906            prompt: None,
907            predictions: Vec::new(),
908            score: vec![PredictionScore::zero(); score_count],
909            qa: Vec::new(),
910            zed_version: None,
911            state: None,
912        }
913    }
914
915    fn related_excerpt(
916        text: &str,
917        row_range: std::ops::Range<u32>,
918        order: usize,
919        context_source: ContextSource,
920    ) -> RelatedExcerpt {
921        RelatedExcerpt {
922            row_range,
923            text: text.into(),
924            order,
925            context_source,
926        }
927    }
928}
929