"""Unit tests for entity/ticker precision, recall, F1, and ambiguity accuracy. Validates: Requirements 16.3, 16.4 """ from __future__ import annotations from services.intelligence_pipeline_v3.evaluation.entity_metrics import ( PRF1, AmbiguityResult, EntityMetricsResult, EntitySpan, MatchMode, TickerMention, TickerMetricsResult, compute_ambiguity_accuracy, compute_entity_metrics, compute_ticker_metrics, evaluate_entities, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _entity( text: str, entity_type: str, start: int, end: int, is_ambiguous: bool = False, ) -> EntitySpan: return EntitySpan( text=text, entity_type=entity_type, start_char=start, end_char=end, is_ambiguous=is_ambiguous, ) def _ticker( text: str, ticker: str, start: int, end: int, is_ambiguous: bool = False, ) -> TickerMention: return TickerMention( text=text, ticker=ticker, start_char=start, end_char=end, is_ambiguous=is_ambiguous, ) # --------------------------------------------------------------------------- # Entity Metrics - Strict Mode # --------------------------------------------------------------------------- class TestEntityMetricsStrict: def test_perfect_match(self) -> None: gold = [_entity("Apple", "company", 0, 5)] pred = [_entity("Apple", "company", 0, 5)] result = compute_entity_metrics(pred, gold, MatchMode.strict) assert result.overall.precision == 1.0 assert result.overall.recall == 1.0 assert result.overall.f1 == 1.0 def test_no_predictions(self) -> None: gold = [_entity("Apple", "company", 0, 5)] result = compute_entity_metrics([], gold, MatchMode.strict) assert result.overall.precision == 1.0 # no false positives assert result.overall.recall == 0.0 assert result.overall.f1 == 0.0 def test_no_gold(self) -> None: pred = [_entity("Apple", "company", 0, 5)] result = compute_entity_metrics(pred, [], MatchMode.strict) assert result.overall.precision == 0.0 assert result.overall.recall == 1.0 # no false negatives assert result.overall.f1 == 0.0 def test_both_empty(self) -> None: result = compute_entity_metrics([], [], MatchMode.strict) assert result.overall.precision == 1.0 assert result.overall.recall == 1.0 assert result.overall.f1 == 1.0 def test_partial_match(self) -> None: gold = [ _entity("Apple", "company", 0, 5), _entity("Tim Cook", "person", 10, 18), ] pred = [ _entity("Apple", "company", 0, 5), _entity("iPhone", "product", 20, 26), ] result = compute_entity_metrics(pred, gold, MatchMode.strict) # 1 TP out of 2 predicted -> precision = 0.5 assert result.overall.precision == 0.5 # 1 TP out of 2 gold -> recall = 0.5 assert result.overall.recall == 0.5 assert result.overall.f1 == 0.5 def test_wrong_type_no_match(self) -> None: gold = [_entity("Apple", "company", 0, 5)] pred = [_entity("Apple", "product", 0, 5)] result = compute_entity_metrics(pred, gold, MatchMode.strict) assert result.overall.precision == 0.0 assert result.overall.recall == 0.0 def test_off_by_one_no_strict_match(self) -> None: gold = [_entity("Apple", "company", 0, 5)] pred = [_entity("Apple", "company", 0, 6)] # end_char differs result = compute_entity_metrics(pred, gold, MatchMode.strict) assert result.overall.precision == 0.0 assert result.overall.recall == 0.0 def test_per_type_breakdown(self) -> None: gold = [ _entity("Apple", "company", 0, 5), _entity("Google", "company", 10, 16), _entity("Tim Cook", "person", 20, 28), ] pred = [ _entity("Apple", "company", 0, 5), _entity("Tim Cook", "person", 20, 28), ] result = compute_entity_metrics(pred, gold, MatchMode.strict) assert result.per_type["company"].precision == 1.0 assert result.per_type["company"].recall == 0.5 assert result.per_type["person"].precision == 1.0 assert result.per_type["person"].recall == 1.0 assert result.per_type["person"].f1 == 1.0 def test_match_mode_in_result(self) -> None: result = compute_entity_metrics([], [], MatchMode.strict) assert result.match_mode == "strict" # --------------------------------------------------------------------------- # Entity Metrics - Relaxed Mode # --------------------------------------------------------------------------- class TestEntityMetricsRelaxed: def test_overlapping_span_matches(self) -> None: gold = [_entity("Apple Inc.", "company", 0, 10)] pred = [_entity("Apple", "company", 0, 5)] # subset overlap result = compute_entity_metrics(pred, gold, MatchMode.relaxed) assert result.overall.precision == 1.0 assert result.overall.recall == 1.0 assert result.overall.f1 == 1.0 def test_non_overlapping_no_match(self) -> None: gold = [_entity("Apple", "company", 0, 5)] pred = [_entity("Google", "company", 10, 16)] result = compute_entity_metrics(pred, gold, MatchMode.relaxed) assert result.overall.precision == 0.0 assert result.overall.recall == 0.0 def test_adjacent_spans_no_overlap(self) -> None: gold = [_entity("Apple", "company", 0, 5)] pred = [_entity("Inc", "company", 5, 8)] # adjacent, not overlapping result = compute_entity_metrics(pred, gold, MatchMode.relaxed) assert result.overall.precision == 0.0 assert result.overall.recall == 0.0 def test_partial_overlap_different_type(self) -> None: gold = [_entity("Apple", "company", 0, 5)] pred = [_entity("Apple", "product", 0, 5)] result = compute_entity_metrics(pred, gold, MatchMode.relaxed) assert result.overall.precision == 0.0 def test_match_mode_in_result(self) -> None: result = compute_entity_metrics([], [], MatchMode.relaxed) assert result.match_mode == "relaxed" # --------------------------------------------------------------------------- # Ticker Metrics # --------------------------------------------------------------------------- class TestTickerMetrics: def test_perfect_match_strict(self) -> None: gold = [_ticker("Apple Inc.", "AAPL", 0, 10)] pred = [_ticker("Apple Inc.", "AAPL", 0, 10)] result = compute_ticker_metrics(pred, gold, MatchMode.strict) assert result.overall.precision == 1.0 assert result.overall.recall == 1.0 assert result.overall.f1 == 1.0 def test_wrong_ticker_no_match(self) -> None: gold = [_ticker("Apple", "AAPL", 0, 5)] pred = [_ticker("Apple", "APLE", 0, 5)] result = compute_ticker_metrics(pred, gold, MatchMode.strict) assert result.overall.precision == 0.0 assert result.overall.recall == 0.0 def test_relaxed_overlapping_ticker(self) -> None: gold = [_ticker("Apple Inc.", "AAPL", 0, 10)] pred = [_ticker("Apple", "AAPL", 0, 5)] result = compute_ticker_metrics(pred, gold, MatchMode.relaxed) assert result.overall.precision == 1.0 assert result.overall.recall == 1.0 def test_multiple_tickers(self) -> None: gold = [ _ticker("Apple", "AAPL", 0, 5), _ticker("Google", "GOOGL", 10, 16), _ticker("Microsoft", "MSFT", 20, 29), ] pred = [ _ticker("Apple", "AAPL", 0, 5), _ticker("Microsoft", "MSFT", 20, 29), ] result = compute_ticker_metrics(pred, gold, MatchMode.strict) assert result.overall.precision == 1.0 assert result.overall.recall == 2 / 3 def test_per_ticker_breakdown(self) -> None: gold = [ _ticker("Apple", "AAPL", 0, 5), _ticker("Google", "GOOGL", 10, 16), ] pred = [ _ticker("Apple", "AAPL", 0, 5), ] result = compute_ticker_metrics(pred, gold, MatchMode.strict) assert "AAPL" in result.per_type assert "GOOGL" in result.per_type assert result.per_type["AAPL"].f1 == 1.0 assert result.per_type["GOOGL"].recall == 0.0 def test_empty_inputs(self) -> None: result = compute_ticker_metrics([], [], MatchMode.strict) assert result.overall.f1 == 1.0 # --------------------------------------------------------------------------- # Ambiguity Accuracy # --------------------------------------------------------------------------- class TestAmbiguityAccuracy: def test_perfect_ambiguity_detection(self) -> None: gold = [ _entity("Apple", "company", 0, 5, is_ambiguous=True), _entity("Google", "company", 10, 16, is_ambiguous=False), ] pred = [ _entity("Apple", "company", 0, 5, is_ambiguous=True), _entity("Google", "company", 10, 16, is_ambiguous=False), ] result = compute_ambiguity_accuracy(pred, gold) assert result.accuracy == 1.0 assert result.true_positives == 1 assert result.true_negatives == 1 def test_all_wrong(self) -> None: gold = [ _entity("Apple", "company", 0, 5, is_ambiguous=True), _entity("Google", "company", 10, 16, is_ambiguous=False), ] pred = [ _entity("Apple", "company", 0, 5, is_ambiguous=False), _entity("Google", "company", 10, 16, is_ambiguous=True), ] result = compute_ambiguity_accuracy(pred, gold) assert result.accuracy == 0.0 assert result.false_negatives == 1 assert result.false_positives == 1 def test_no_aligned_spans(self) -> None: gold = [_entity("Apple", "company", 0, 5, is_ambiguous=True)] pred = [_entity("Apple", "company", 10, 15, is_ambiguous=True)] result = compute_ambiguity_accuracy(pred, gold) assert result.support == 0 assert result.accuracy == 1.0 # vacuously true def test_mixed_results(self) -> None: gold = [ _entity("Apple", "company", 0, 5, is_ambiguous=True), _entity("Google", "company", 10, 16, is_ambiguous=False), _entity("Tesla", "company", 20, 25, is_ambiguous=True), ] pred = [ _entity("Apple", "company", 0, 5, is_ambiguous=True), # TP _entity("Google", "company", 10, 16, is_ambiguous=True), # FP _entity("Tesla", "company", 20, 25, is_ambiguous=False), # FN ] result = compute_ambiguity_accuracy(pred, gold) assert result.true_positives == 1 assert result.false_positives == 1 assert result.false_negatives == 1 assert result.true_negatives == 0 assert result.support == 3 assert abs(result.accuracy - 1 / 3) < 1e-9 def test_ticker_mentions_supported(self) -> None: gold = [_ticker("Apple", "AAPL", 0, 5, is_ambiguous=True)] pred = [_ticker("Apple", "AAPL", 0, 5, is_ambiguous=True)] result = compute_ambiguity_accuracy(pred, gold) assert result.accuracy == 1.0 # --------------------------------------------------------------------------- # Full Evaluation Report # --------------------------------------------------------------------------- class TestEvaluateEntities: def test_full_evaluation(self) -> None: gold_entities = [ _entity("Apple", "company", 0, 5), _entity("Tim Cook", "person", 10, 18), ] pred_entities = [ _entity("Apple", "company", 0, 5), _entity("Tim Cook", "person", 10, 18), ] gold_tickers = [_ticker("Apple Inc.", "AAPL", 0, 10)] pred_tickers = [_ticker("Apple Inc.", "AAPL", 0, 10)] report = evaluate_entities( pred_entities, gold_entities, pred_tickers, gold_tickers, mode=MatchMode.strict, document_count=1, ) assert isinstance(report.entity_metrics, EntityMetricsResult) assert isinstance(report.ticker_metrics, TickerMetricsResult) assert isinstance(report.ambiguity_accuracy, AmbiguityResult) assert report.document_count == 1 assert report.entity_metrics.overall.f1 == 1.0 assert report.ticker_metrics.overall.f1 == 1.0 def test_multiple_documents(self) -> None: # Simulating aggregated results from multiple documents gold_entities = [ _entity("Apple", "company", 0, 5, is_ambiguous=True), _entity("Google", "company", 100, 106), ] pred_entities = [ _entity("Apple", "company", 0, 5, is_ambiguous=True), ] gold_tickers = [ _ticker("Apple", "AAPL", 0, 5), _ticker("Google", "GOOGL", 100, 106), ] pred_tickers = [ _ticker("Apple", "AAPL", 0, 5), ] report = evaluate_entities( pred_entities, gold_entities, pred_tickers, gold_tickers, mode=MatchMode.strict, document_count=2, ) assert report.document_count == 2 assert report.entity_metrics.overall.recall == 0.5 assert report.ticker_metrics.overall.recall == 0.5 # --------------------------------------------------------------------------- # PRF1 Model validation # --------------------------------------------------------------------------- class TestPRF1Model: def test_valid_prf1(self) -> None: prf1 = PRF1(precision=0.8, recall=0.6, f1=0.686, support_predicted=10, support_gold=12) assert prf1.precision == 0.8 assert prf1.recall == 0.6 def test_f1_harmonic_mean(self) -> None: """F1 should be the harmonic mean when computed by the metric functions.""" gold = [ _entity("Apple", "company", 0, 5), _entity("Google", "company", 10, 16), _entity("Tesla", "company", 20, 25), ] pred = [ _entity("Apple", "company", 0, 5), _entity("Microsoft", "company", 30, 39), ] result = compute_entity_metrics(pred, gold, MatchMode.strict) p = result.overall.precision r = result.overall.recall expected_f1 = 2 * p * r / (p + r) if (p + r) > 0 else 0.0 assert abs(result.overall.f1 - expected_f1) < 1e-9