tenant.openagents/omega

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    ops::Range,
4};
5
6use crate::patch::{Patch, PatchLine};
7use crate::patch_metrics::ClassificationMetrics;
8
9const LINE_RELEVANCE_WINDOW: u32 = 20;
10
11#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
12pub struct Excerpt {
13    pub path: String,
14    pub row_range: Range<u32>,
15    pub content: String,
16}
17
18/// Line- and file-level precision/recall/F1 over TP/FP/FN counts.
19///
20/// Shared shape for two metrics with the same structure: how much expected
21/// edit context was retrieved (`EditableContextCoverage`) and how well
22/// predicted edit locations match expected ones (`PatchLocationMatch`).
23#[derive(Clone, Debug, PartialEq, serde::Deserialize, serde::Serialize)]
24pub struct LineFileClassification {
25    pub lines_tp: usize,
26    pub lines_fp: usize,
27    pub lines_fn: usize,
28    pub lines_precision: f64,
29    pub lines_recall: f64,
30    pub lines_f1: f64,
31
32    pub files_tp: usize,
33    pub files_fp: usize,
34    pub files_fn: usize,
35    pub files_precision: f64,
36    pub files_recall: f64,
37    pub files_f1: f64,
38}
39
40pub type EditableContextCoverage = LineFileClassification;
41pub type PatchLocationMatch = LineFileClassification;
42
43impl LineFileClassification {
44    pub fn new(
45        lines_tp: usize,
46        lines_fp: usize,
47        lines_fn: usize,
48        files_tp: usize,
49        files_fp: usize,
50        files_fn: usize,
51    ) -> Self {
52        Self {
53            lines_tp,
54            lines_fp,
55            lines_fn,
56            lines_precision: precision(lines_tp, lines_fp, lines_fn),
57            lines_recall: recall(lines_tp, lines_fp, lines_fn),
58            lines_f1: f1(lines_tp, lines_fp, lines_fn),
59            files_tp,
60            files_fp,
61            files_fn,
62            files_precision: precision(files_tp, files_fp, files_fn),
63            files_recall: recall(files_tp, files_fp, files_fn),
64            files_f1: f1(files_tp, files_fp, files_fn),
65        }
66    }
67
68    pub fn lines_counts(&self) -> ClassificationMetrics {
69        ClassificationMetrics {
70            true_positives: self.lines_tp,
71            false_positives: self.lines_fp,
72            false_negatives: self.lines_fn,
73        }
74    }
75
76    pub fn files_counts(&self) -> ClassificationMetrics {
77        ClassificationMetrics {
78            true_positives: self.files_tp,
79            false_positives: self.files_fp,
80            false_negatives: self.files_fn,
81        }
82    }
83}
84
85/// Measures how well an actual patch's edited file/line locations match an expected patch.
86pub fn patch_location_match(expected_patch: &str, actual_patch: &str) -> PatchLocationMatch {
87    let expected_patch = Patch::parse_unified_diff(expected_patch);
88    let actual_patch = Patch::parse_unified_diff(actual_patch);
89    let (expected_files, expected_anchor_lines, _) = expected_context(&expected_patch);
90    let (actual_files, mut actual_anchor_lines, _) = expected_context(&actual_patch);
91
92    normalize_actual_paths(
93        &expected_files,
94        &mut actual_anchor_lines,
95        actual_files.iter().cloned().collect(),
96    );
97    let actual_files = actual_anchor_lines.keys().cloned().collect();
98
99    let (lines_tp, lines_fp, lines_fn) =
100        match_anchor_rows(&expected_anchor_lines, &actual_anchor_lines);
101    let (files_tp, files_fp, files_fn) = classify_values(&expected_files, &actual_files);
102
103    PatchLocationMatch::new(lines_tp, lines_fp, lines_fn, files_tp, files_fp, files_fn)
104}
105
106/// Measures how much expected edit context was retrieved and how much unrelated context was retrieved.
107pub fn editable_context_coverage(
108    expected_patch: &str,
109    context: &[Excerpt],
110) -> EditableContextCoverage {
111    let patch = Patch::parse_unified_diff(expected_patch);
112    let (expected_files, expected_anchor_lines, relevant_lines) = expected_context(&patch);
113    let (retrieved_files, retrieved_lines) = retrieved_context(context);
114
115    let (lines_tp, lines_fp) = classify_retrieved_rows(&relevant_lines, &retrieved_lines);
116    let (lines_fn, expected_anchor_line_count) =
117        count_missing_rows(&expected_anchor_lines, &retrieved_lines);
118    let (files_tp, files_fp, files_fn) = classify_values(&expected_files, &retrieved_files);
119
120    let lines_precision = precision(lines_tp, lines_fp, lines_fn);
121    let lines_recall = line_recall(expected_anchor_line_count, lines_fn);
122    let lines_f1 = f1_from_precision_and_recall(lines_precision, lines_recall);
123    let files_precision = precision(files_tp, files_fp, files_fn);
124    let files_recall = recall(files_tp, files_fp, files_fn);
125    let files_f1 = f1_from_precision_and_recall(files_precision, files_recall);
126
127    EditableContextCoverage {
128        lines_tp,
129        lines_fp,
130        lines_fn,
131        lines_precision,
132        lines_recall,
133        lines_f1,
134        files_tp,
135        files_fp,
136        files_fn,
137        files_precision,
138        files_recall,
139        files_f1,
140    }
141}
142
143fn expected_context(
144    patch: &Patch,
145) -> (
146    BTreeSet<String>,
147    BTreeMap<String, BTreeSet<u32>>,
148    BTreeMap<String, BTreeSet<u32>>,
149) {
150    let mut expected_files = BTreeSet::new();
151    let mut expected_anchor_lines = BTreeMap::new();
152    let mut relevant_lines = BTreeMap::new();
153
154    for hunk in &patch.hunks {
155        if hunk
156            .lines
157            .iter()
158            .any(|line| matches!(line, PatchLine::Addition(_) | PatchLine::Deletion(_)))
159        {
160            expected_files.insert(hunk.filename.clone());
161        }
162
163        let mut old_row = hunk.old_start.saturating_sub(1).max(0) as u32;
164        let mut previous_context_row = None;
165        let mut index = 0;
166
167        while index < hunk.lines.len() {
168            match &hunk.lines[index] {
169                PatchLine::Context(_) => {
170                    previous_context_row = Some(old_row);
171                    old_row = old_row.saturating_add(1);
172                    index += 1;
173                }
174                PatchLine::Addition(_) | PatchLine::Deletion(_) => {
175                    let mut deletion_rows = Vec::new();
176                    let mut has_addition = false;
177
178                    while index < hunk.lines.len() {
179                        match &hunk.lines[index] {
180                            PatchLine::Addition(_) => {
181                                has_addition = true;
182                                index += 1;
183                            }
184                            PatchLine::Deletion(_) => {
185                                deletion_rows.push(old_row);
186                                old_row = old_row.saturating_add(1);
187                                index += 1;
188                            }
189                            _ => break,
190                        }
191                    }
192
193                    if deletion_rows.is_empty() {
194                        if has_addition {
195                            if let Some(row) = previous_context_row {
196                                insert_anchor_row(
197                                    &mut expected_anchor_lines,
198                                    &mut relevant_lines,
199                                    &hunk.filename,
200                                    row,
201                                );
202                            }
203                            if matches!(hunk.lines.get(index), Some(PatchLine::Context(_))) {
204                                insert_anchor_row(
205                                    &mut expected_anchor_lines,
206                                    &mut relevant_lines,
207                                    &hunk.filename,
208                                    old_row,
209                                );
210                            }
211                        }
212                    } else {
213                        for row in deletion_rows {
214                            insert_anchor_row(
215                                &mut expected_anchor_lines,
216                                &mut relevant_lines,
217                                &hunk.filename,
218                                row,
219                            );
220                        }
221                    }
222
223                    previous_context_row = None;
224                }
225                PatchLine::Garbage(_) => {
226                    index += 1;
227                }
228            }
229        }
230    }
231
232    (expected_files, expected_anchor_lines, relevant_lines)
233}
234
235fn retrieved_context(context: &[Excerpt]) -> (BTreeSet<String>, BTreeMap<String, BTreeSet<u32>>) {
236    let mut retrieved_files = BTreeSet::new();
237    let mut retrieved_lines = BTreeMap::new();
238
239    for excerpt in context {
240        retrieved_files.insert(excerpt.path.clone());
241        let rows = retrieved_lines
242            .entry(excerpt.path.clone())
243            .or_insert_with(BTreeSet::new);
244        rows.extend(excerpt.row_range.clone());
245    }
246
247    (retrieved_files, retrieved_lines)
248}
249
250fn normalize_actual_paths(
251    expected_files: &BTreeSet<String>,
252    actual_anchor_lines: &mut BTreeMap<String, BTreeSet<u32>>,
253    actual_files: BTreeSet<String>,
254) {
255    let mut normalized = BTreeMap::new();
256
257    for actual_path in actual_files {
258        let normalized_path = normalize_actual_path(expected_files, &actual_path);
259        if let Some(rows) = actual_anchor_lines.remove(&actual_path) {
260            normalized
261                .entry(normalized_path)
262                .or_insert_with(BTreeSet::new)
263                .extend(rows);
264        }
265    }
266
267    *actual_anchor_lines = normalized;
268}
269
270fn normalize_actual_path(expected_files: &BTreeSet<String>, actual_path: &str) -> String {
271    if expected_files.contains(actual_path) {
272        return actual_path.to_string();
273    }
274
275    if let Some(stripped_path) = strip_first_path_component(actual_path) {
276        if expected_files.contains(stripped_path) {
277            return stripped_path.to_string();
278        }
279    }
280
281    actual_path.to_string()
282}
283
284fn strip_first_path_component(path: &str) -> Option<&str> {
285    path.split_once('/')
286        .map(|(_, rest)| rest)
287        .filter(|rest| !rest.is_empty())
288}
289
290fn match_anchor_rows(
291    expected: &BTreeMap<String, BTreeSet<u32>>,
292    actual: &BTreeMap<String, BTreeSet<u32>>,
293) -> (usize, usize, usize) {
294    let mut true_positives = 0;
295    let mut false_positives = 0;
296    let mut false_negatives = 0;
297    let paths = expected
298        .keys()
299        .chain(actual.keys())
300        .cloned()
301        .collect::<BTreeSet<_>>();
302
303    for path in paths {
304        let expected_rows = expected.get(&path).cloned().unwrap_or_default();
305        let actual_rows = actual.get(&path).cloned().unwrap_or_default();
306        let matched = match_rows_with_window(&expected_rows, &actual_rows);
307        true_positives += matched;
308        false_positives += actual_rows.len().saturating_sub(matched);
309        false_negatives += expected_rows.len().saturating_sub(matched);
310    }
311
312    (true_positives, false_positives, false_negatives)
313}
314
315fn match_rows_with_window(expected: &BTreeSet<u32>, actual: &BTreeSet<u32>) -> usize {
316    let mut unmatched_expected = expected.clone();
317    let mut matched = 0;
318
319    for actual_row in actual {
320        let start = actual_row.saturating_sub(LINE_RELEVANCE_WINDOW);
321        let end = actual_row.saturating_add(LINE_RELEVANCE_WINDOW);
322        if let Some(expected_row) = unmatched_expected.range(start..=end).next().copied() {
323            unmatched_expected.remove(&expected_row);
324            matched += 1;
325        }
326    }
327
328    matched
329}
330
331fn insert_anchor_row(
332    anchor_lines_by_file: &mut BTreeMap<String, BTreeSet<u32>>,
333    relevant_lines_by_file: &mut BTreeMap<String, BTreeSet<u32>>,
334    path: &str,
335    row: u32,
336) {
337    insert_row(anchor_lines_by_file, path, row);
338
339    let start = row.saturating_sub(LINE_RELEVANCE_WINDOW);
340    let end = row.saturating_add(LINE_RELEVANCE_WINDOW);
341    for relevant_row in start..=end {
342        insert_row(relevant_lines_by_file, path, relevant_row);
343    }
344}
345
346fn insert_row(lines_by_file: &mut BTreeMap<String, BTreeSet<u32>>, path: &str, row: u32) {
347    lines_by_file
348        .entry(path.to_string())
349        .or_insert_with(BTreeSet::new)
350        .insert(row);
351}
352
353fn classify_retrieved_rows(
354    relevant: &BTreeMap<String, BTreeSet<u32>>,
355    retrieved: &BTreeMap<String, BTreeSet<u32>>,
356) -> (usize, usize) {
357    let mut true_positives = 0;
358    let mut false_positives = 0;
359
360    for (path, rows) in retrieved {
361        for row in rows {
362            if relevant
363                .get(path)
364                .is_some_and(|relevant_rows| relevant_rows.contains(row))
365            {
366                true_positives += 1;
367            } else {
368                false_positives += 1;
369            }
370        }
371    }
372
373    (true_positives, false_positives)
374}
375
376fn count_missing_rows(
377    expected: &BTreeMap<String, BTreeSet<u32>>,
378    retrieved: &BTreeMap<String, BTreeSet<u32>>,
379) -> (usize, usize) {
380    let mut false_negatives = 0;
381    let mut expected_count = 0;
382
383    for (path, rows) in expected {
384        for row in rows {
385            expected_count += 1;
386            if !retrieved
387                .get(path)
388                .is_some_and(|retrieved_rows| retrieved_rows.contains(row))
389            {
390                false_negatives += 1;
391            }
392        }
393    }
394
395    (false_negatives, expected_count)
396}
397
398fn classify_values<T: Ord>(
399    expected: &BTreeSet<T>,
400    retrieved: &BTreeSet<T>,
401) -> (usize, usize, usize) {
402    let true_positives = expected.intersection(retrieved).count();
403    let false_positives = retrieved.difference(expected).count();
404    let false_negatives = expected.difference(retrieved).count();
405    (true_positives, false_positives, false_negatives)
406}
407
408fn precision(true_positives: usize, false_positives: usize, false_negatives: usize) -> f64 {
409    if true_positives + false_positives + false_negatives == 0 {
410        return 1.0;
411    }
412
413    let denominator = true_positives + false_positives;
414    if denominator == 0 {
415        1.0
416    } else {
417        true_positives as f64 / denominator as f64
418    }
419}
420
421fn recall(true_positives: usize, false_positives: usize, false_negatives: usize) -> f64 {
422    if true_positives + false_positives + false_negatives == 0 {
423        return 1.0;
424    }
425
426    let denominator = true_positives + false_negatives;
427    if denominator == 0 {
428        1.0
429    } else {
430        true_positives as f64 / denominator as f64
431    }
432}
433
434fn line_recall(expected_anchor_line_count: usize, false_negatives: usize) -> f64 {
435    if expected_anchor_line_count == 0 {
436        1.0
437    } else {
438        (expected_anchor_line_count - false_negatives) as f64 / expected_anchor_line_count as f64
439    }
440}
441
442fn f1(true_positives: usize, false_positives: usize, false_negatives: usize) -> f64 {
443    let precision = precision(true_positives, false_positives, false_negatives);
444    let recall = recall(true_positives, false_positives, false_negatives);
445
446    f1_from_precision_and_recall(precision, recall)
447}
448
449fn f1_from_precision_and_recall(precision: f64, recall: f64) -> f64 {
450    if precision + recall == 0.0 {
451        0.0
452    } else {
453        2.0 * precision * recall / (precision + recall)
454    }
455}
456
457#[cfg(test)]
458mod tests {
459    use super::*;
460    use indoc::indoc;
461
462    fn excerpt(path: &str, row_range: Range<u32>) -> Excerpt {
463        Excerpt {
464            path: path.to_string(),
465            row_range,
466            content: String::new(),
467        }
468    }
469
470    #[test]
471    fn deletion_counts_deleted_old_line_as_true_positive() {
472        let patch = indoc! {"
473            --- a/src/main.rs
474            +++ b/src/main.rs
475            @@ -2,1 +2,0 @@
476            -let value = 1;
477        "};
478
479        let score = editable_context_coverage(patch, &[excerpt("src/main.rs", 1..2)]);
480
481        assert_eq!(score, EditableContextCoverage::new(1, 0, 0, 1, 0, 0));
482    }
483
484    #[test]
485    fn retrieved_lines_inside_relevance_window_are_true_positives() {
486        let patch = indoc! {"
487            --- a/src/main.rs
488            +++ b/src/main.rs
489            @@ -4,1 +4,0 @@
490            -let value = 4;
491        "};
492
493        let score = editable_context_coverage(
494            patch,
495            &[excerpt("src/main.rs", 0..1), excerpt("src/main.rs", 3..4)],
496        );
497
498        assert_eq!(score, EditableContextCoverage::new(2, 0, 0, 1, 0, 0));
499    }
500
501    #[test]
502    fn replacement_counts_deleted_old_line_without_addition_anchor() {
503        let patch = indoc! {"
504            --- a/src/main.rs
505            +++ b/src/main.rs
506            @@ -1,3 +1,3 @@
507             fn main() {
508            -    let value = 1;
509            +    let value = 2;
510             }
511        "};
512
513        let score = editable_context_coverage(patch, &[excerpt("src/main.rs", 1..2)]);
514
515        assert_eq!(score, EditableContextCoverage::new(1, 0, 0, 1, 0, 0));
516    }
517
518    #[test]
519    fn pure_insertion_counts_previous_and_next_old_lines_as_expected_context() {
520        let patch = indoc! {"
521            --- a/src/main.rs
522            +++ b/src/main.rs
523            @@ -1,2 +1,3 @@
524             line 1
525            +inserted
526             line 2
527        "};
528
529        let score = editable_context_coverage(patch, &[excerpt("src/main.rs", 0..1)]);
530
531        assert_eq!(score, EditableContextCoverage::new(1, 0, 1, 1, 0, 0));
532    }
533
534    #[test]
535    fn pure_insertion_at_file_boundary_uses_available_neighboring_context() {
536        let patch = indoc! {"
537            --- a/src/main.rs
538            +++ b/src/main.rs
539            @@ -1,1 +1,2 @@
540            +inserted
541             line 1
542        "};
543
544        let score = editable_context_coverage(patch, &[excerpt("src/main.rs", 0..1)]);
545
546        assert_eq!(score, EditableContextCoverage::new(1, 0, 0, 1, 0, 0));
547    }
548
549    #[test]
550    fn counts_false_negatives_and_file_false_positives() {
551        let patch = indoc! {"
552            --- a/src/main.rs
553            +++ b/src/main.rs
554            @@ -1,3 +1,3 @@
555            -let first = 1;
556            +let first = 2;
557             let middle = 3;
558            -let last = 4;
559            +let last = 5;
560        "};
561
562        let score = editable_context_coverage(
563            patch,
564            &[excerpt("src/main.rs", 0..1), excerpt("src/lib.rs", 0..1)],
565        );
566
567        assert_eq!(score, EditableContextCoverage::new(1, 1, 1, 1, 1, 0));
568    }
569
570    #[test]
571    fn overlapping_excerpts_are_counted_once() {
572        let patch = indoc! {"
573            --- a/src/main.rs
574            +++ b/src/main.rs
575            @@ -2,1 +2,0 @@
576            -let value = 1;
577        "};
578
579        let score = editable_context_coverage(
580            patch,
581            &[excerpt("src/main.rs", 0..2), excerpt("src/main.rs", 1..3)],
582        );
583
584        assert_eq!(score, EditableContextCoverage::new(3, 0, 0, 1, 0, 0));
585    }
586
587    #[test]
588    fn nearby_lines_do_not_satisfy_line_recall_without_exact_anchor_lines() {
589        let patch = indoc! {"
590            --- a/src/main.rs
591            +++ b/src/main.rs
592            @@ -1,2 +1,3 @@
593             line 1
594            +inserted
595             line 2
596        "};
597
598        let score = editable_context_coverage(patch, &[excerpt("src/main.rs", 2..3)]);
599
600        assert_eq!(score.lines_tp, 1);
601        assert_eq!(score.lines_fp, 0);
602        assert_eq!(score.lines_fn, 2);
603        assert_eq!(score.lines_precision, 1.0);
604        assert_eq!(score.lines_recall, 0.0);
605        assert_eq!(score.lines_f1, 0.0);
606    }
607
608    #[test]
609    fn retrieved_lines_outside_relevance_window_are_false_positives() {
610        let patch = indoc! {"
611            --- a/src/main.rs
612            +++ b/src/main.rs
613            @@ -1,1 +1,0 @@
614            -line 1
615        "};
616
617        let score = editable_context_coverage(patch, &[excerpt("src/main.rs", 21..22)]);
618
619        assert_eq!(score, EditableContextCoverage::new(0, 1, 1, 1, 0, 0));
620    }
621
622    #[test]
623    fn empty_patch_with_no_context_has_perfect_f1() {
624        let score = editable_context_coverage(
625            indoc! {"
626            "},
627            &[],
628        );
629
630        assert_eq!(score, EditableContextCoverage::new(0, 0, 0, 0, 0, 0));
631    }
632
633    #[test]
634    fn patch_location_match_counts_file_and_nearby_line_matches() {
635        let expected = indoc! {"
636            --- a/src/main.rs
637            +++ b/src/main.rs
638            @@ -50,1 +50,1 @@
639            -let value = 1;
640            +let value = 2;
641        "};
642        let actual = indoc! {"
643            --- a/src/main.rs
644            +++ b/src/main.rs
645            @@ -55,1 +55,1 @@
646            -let value = 1;
647            +let value = 2;
648        "};
649
650        let score = patch_location_match(expected, actual);
651
652        assert_eq!(score, PatchLocationMatch::new(1, 0, 0, 1, 0, 0));
653    }
654
655    #[test]
656    fn patch_location_match_counts_missing_and_extra_files() {
657        let expected = indoc! {"
658            --- a/src/main.rs
659            +++ b/src/main.rs
660            @@ -1,1 +1,1 @@
661            -let value = 1;
662            +let value = 2;
663        "};
664        let actual = indoc! {"
665            --- a/src/lib.rs
666            +++ b/src/lib.rs
667            @@ -1,1 +1,1 @@
668            -let value = 1;
669            +let value = 2;
670        "};
671
672        let score = patch_location_match(expected, actual);
673
674        assert_eq!(score, PatchLocationMatch::new(0, 1, 1, 0, 1, 1));
675    }
676
677    #[test]
678    fn patch_location_match_normalizes_project_prefixed_actual_path() {
679        let expected = indoc! {"
680            --- a/src/main.rs
681            +++ b/src/main.rs
682            @@ -1,1 +1,1 @@
683            -let value = 1;
684            +let value = 2;
685        "};
686        let actual = indoc! {"
687            --- a/project/src/main.rs
688            +++ b/project/src/main.rs
689            @@ -1,1 +1,1 @@
690            -let value = 1;
691            +let value = 2;
692        "};
693
694        let score = patch_location_match(expected, actual);
695
696        assert_eq!(score, PatchLocationMatch::new(1, 0, 0, 1, 0, 0));
697    }
698}
699