"""Tests for the confidence feature pipeline. Covers feature extraction, calibrator fit/predict, conservative defaults, artifact save/load, and ECE/Brier computation. """ from __future__ import annotations import tempfile from pathlib import Path import numpy as np import pytest from services.intelligence_pipeline_v3.confidence.artifacts import ( list_versions, load_artifact, load_metadata, save_artifact, ) from services.intelligence_pipeline_v3.confidence.calibrator import ( ConfidenceCalibrator, _compute_brier, _compute_ece, compare_methods, ) from services.intelligence_pipeline_v3.confidence.defaults import ( get_default_confidence, is_underrepresented, ) from services.intelligence_pipeline_v3.confidence.features import ( AgreementStageResult, ConfidenceFeatureExtractor, EvidenceStageResult, ExtractionStageResult, ResolutionStageResult, SentimentStageResult, ) from services.intelligence_pipeline_v3.confidence.models import ( ConfidenceFeatures, ConfidenceResult, ) # --- Fixtures --- def _make_extraction_result( entity_scores: list[float] | None = None, relation_scores: list[float] | None = None, total_facts: int = 10, valid_numeric_facts: int = 8, populated_fields: int = 7, expected_fields: int = 10, ) -> ExtractionStageResult: return ExtractionStageResult( entity_scores=[0.9, 0.85, 0.7] if entity_scores is None else entity_scores, relation_scores=[0.8, 0.75] if relation_scores is None else relation_scores, total_facts=total_facts, valid_numeric_facts=valid_numeric_facts, populated_fields=populated_fields, expected_fields=expected_fields, ) def _make_resolution_result( margins: list[float] | None = None, ) -> ResolutionStageResult: return ResolutionStageResult( ambiguity_margins=[0.9, 0.6] if margins is None else margins, ) def _make_evidence_result( total: int = 10, supported: int = 8, ) -> EvidenceStageResult: return EvidenceStageResult( total_claims=total, supported_claims=supported, ) def _make_sentiment_result( probs: list[float] | None = None, ) -> SentimentStageResult: return SentimentStageResult( max_class_probabilities=[0.85, 0.9] if probs is None else probs, calibration_version="v1.0", ) def _make_agreement_result() -> AgreementStageResult: return AgreementStageResult( agreement_ratio=0.8, novelty_certainty=0.7, hard_case_score=0.2, ) def _make_features( entity_span_score: float = 0.85, document_type: str = "news", ) -> ConfidenceFeatures: return ConfidenceFeatures( entity_span_score=entity_span_score, alias_resolution_margin=0.75, numeric_parser_validity=0.8, evidence_coverage=0.8, relation_score=0.775, sentiment_calibration_confidence=0.875, cross_stage_agreement=0.8, duplicate_novelty_certainty=0.7, document_completeness=0.7, document_type=document_type, known_hard_case_patterns=0.2, ) def _generate_training_data( n_samples: int = 100, seed: int = 42, ) -> tuple[list[ConfidenceFeatures], list[bool]]: """Generate synthetic training data for calibrator tests.""" rng = np.random.default_rng(seed) features = [] labels = [] doc_types = ["news", "filing", "transcript", "press_release", "macro_event"] for _ in range(n_samples): # Generate features with some correlation to the label base_quality = rng.uniform(0.3, 0.95) noise = rng.normal(0, 0.1) f = ConfidenceFeatures( entity_span_score=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)), alias_resolution_margin=float(np.clip(base_quality + rng.normal(0, 0.15), 0, 1)), numeric_parser_validity=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)), evidence_coverage=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)), relation_score=float(np.clip(base_quality + rng.normal(0, 0.15), 0, 1)), sentiment_calibration_confidence=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)), cross_stage_agreement=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)), duplicate_novelty_certainty=float(np.clip(base_quality + rng.normal(0, 0.15), 0, 1)), document_completeness=float(np.clip(base_quality + rng.normal(0, 0.1), 0, 1)), document_type=rng.choice(doc_types), known_hard_case_patterns=float(np.clip(rng.uniform(0, 0.5), 0, 1)), ) features.append(f) # Label correlates with base quality label = bool(rng.random() < (base_quality + noise)) labels.append(label) return features, labels # --- Test Feature Extraction --- class TestFeatureExtraction: """Test that feature extraction produces valid feature vectors.""" def test_extract_features_produces_valid_vector(self): """Feature extraction from all stages produces a valid ConfidenceFeatures.""" extractor = ConfidenceFeatureExtractor() features = extractor.extract_features( extraction_result=_make_extraction_result(), resolution_result=_make_resolution_result(), evidence_result=_make_evidence_result(), sentiment_result=_make_sentiment_result(), agreement_result=_make_agreement_result(), document_type="news", ) assert isinstance(features, ConfidenceFeatures) assert 0.0 <= features.entity_span_score <= 1.0 assert 0.0 <= features.alias_resolution_margin <= 1.0 assert 0.0 <= features.numeric_parser_validity <= 1.0 assert 0.0 <= features.evidence_coverage <= 1.0 assert 0.0 <= features.relation_score <= 1.0 assert 0.0 <= features.sentiment_calibration_confidence <= 1.0 assert 0.0 <= features.cross_stage_agreement <= 1.0 assert 0.0 <= features.duplicate_novelty_certainty <= 1.0 assert 0.0 <= features.document_completeness <= 1.0 assert 0.0 <= features.known_hard_case_patterns <= 1.0 assert features.document_type == "news" def test_extract_features_without_agreement(self): """Feature extraction uses sensible defaults when agreement is not available.""" extractor = ConfidenceFeatureExtractor() features = extractor.extract_features( extraction_result=_make_extraction_result(), resolution_result=_make_resolution_result(), evidence_result=_make_evidence_result(), sentiment_result=_make_sentiment_result(), agreement_result=None, document_type="filing", ) assert features.cross_stage_agreement == 0.5 assert features.duplicate_novelty_certainty == 0.5 assert features.known_hard_case_patterns == 0.0 def test_extract_features_empty_entities(self): """Feature extraction handles empty entity scores gracefully.""" extractor = ConfidenceFeatureExtractor() features = extractor.extract_features( extraction_result=_make_extraction_result(entity_scores=[]), resolution_result=_make_resolution_result(), evidence_result=_make_evidence_result(), sentiment_result=_make_sentiment_result(), ) assert features.entity_span_score == 0.0 def test_extract_features_no_claims(self): """Feature extraction handles zero claims gracefully.""" extractor = ConfidenceFeatureExtractor() features = extractor.extract_features( extraction_result=_make_extraction_result(), resolution_result=_make_resolution_result(), evidence_result=_make_evidence_result(total=0, supported=0), sentiment_result=_make_sentiment_result(), ) assert features.evidence_coverage == 0.0 def test_to_vector_produces_correct_length(self): """Feature vector has expected dimensionality.""" features = _make_features() vector = features.to_vector() assert len(vector) == 11 assert all(isinstance(v, float) for v in vector) def test_unknown_document_type_defaults(self): """Unknown document types are normalized to 'unknown'.""" extractor = ConfidenceFeatureExtractor() features = extractor.extract_features( extraction_result=_make_extraction_result(), resolution_result=_make_resolution_result(), evidence_result=_make_evidence_result(), sentiment_result=_make_sentiment_result(), document_type="exotic_type", ) assert features.document_type == "unknown" # --- Test Calibrator --- class TestCalibrator: """Test calibrator fit/predict roundtrip and method comparison.""" def test_fit_predict_isotonic(self): """Isotonic calibrator can fit and produce predictions in [0, 1].""" features, labels = _generate_training_data(n_samples=50) cal = ConfidenceCalibrator(method="isotonic") cal.fit(features, labels, version="test-v1") assert cal.is_fitted assert cal.version == "test-v1" prediction = cal.predict(features[0]) assert 0.0 <= prediction <= 1.0 def test_fit_predict_platt(self): """Platt calibrator can fit and produce predictions in [0, 1].""" features, labels = _generate_training_data(n_samples=50) cal = ConfidenceCalibrator(method="platt") cal.fit(features, labels, version="test-v1") assert cal.is_fitted prediction = cal.predict(features[0]) assert 0.0 <= prediction <= 1.0 def test_unfitted_returns_neutral(self): """Unfitted calibrator returns 0.5 as neutral default.""" cal = ConfidenceCalibrator() features = _make_features() prediction = cal.predict(features) assert prediction == 0.5 def test_predict_batch(self): """Batch prediction returns correct number of results.""" features, labels = _generate_training_data(n_samples=50) cal = ConfidenceCalibrator(method="isotonic") cal.fit(features, labels) batch_predictions = cal.predict_batch(features[:10]) assert len(batch_predictions) == 10 assert all(0.0 <= p <= 1.0 for p in batch_predictions) def test_fit_empty_raises(self): """Fitting with empty data raises ValueError.""" cal = ConfidenceCalibrator() with pytest.raises(ValueError, match="must not be empty"): cal.fit([], []) def test_fit_mismatched_lengths_raises(self): """Fitting with mismatched lengths raises ValueError.""" features, labels = _generate_training_data(n_samples=10) cal = ConfidenceCalibrator() with pytest.raises(ValueError, match="must have the same length"): cal.fit(features, labels[:5]) def test_metadata_after_fit(self): """Metadata is populated after fitting.""" features, labels = _generate_training_data(n_samples=50) cal = ConfidenceCalibrator(method="isotonic") cal.fit(features, labels, version="v1.0.0", training_range="2024-01-01 to 2024-06-30") assert cal.metadata is not None assert cal.metadata.version == "v1.0.0" assert cal.metadata.method == "isotonic" assert cal.metadata.training_count == 50 assert cal.metadata.training_range == "2024-01-01 to 2024-06-30" assert 0.0 <= cal.metadata.ece <= 1.0 assert 0.0 <= cal.metadata.brier_score <= 1.0 def test_compare_methods(self): """Method comparison returns ECE and Brier for both methods.""" features, labels = _generate_training_data(n_samples=50) results = compare_methods(features, labels, n_folds=3) assert "isotonic" in results assert "platt" in results assert "ece" in results["isotonic"] assert "brier" in results["isotonic"] assert "ece" in results["platt"] assert "brier" in results["platt"] # --- Test ECE and Brier --- class TestMetrics: """Test ECE and Brier score computation.""" def test_ece_perfect_calibration(self): """ECE is 0 for perfectly calibrated predictions.""" # Perfect: predict 1.0 for positives, 0.0 for negatives predictions = np.array([1.0, 1.0, 0.0, 0.0, 1.0]) labels = np.array([1.0, 1.0, 0.0, 0.0, 1.0]) ece = _compute_ece(predictions, labels) assert ece == pytest.approx(0.0, abs=1e-10) def test_ece_worst_calibration(self): """ECE is high for badly calibrated predictions.""" # Predict 1.0 but all are actually 0 predictions = np.array([0.9, 0.9, 0.9, 0.9, 0.9]) labels = np.array([0.0, 0.0, 0.0, 0.0, 0.0]) ece = _compute_ece(predictions, labels) assert ece > 0.5 def test_brier_perfect_predictions(self): """Brier score is 0 for perfect predictions.""" predictions = np.array([1.0, 0.0, 1.0, 0.0]) labels = np.array([1.0, 0.0, 1.0, 0.0]) brier = _compute_brier(predictions, labels) assert brier == pytest.approx(0.0, abs=1e-10) def test_brier_worst_predictions(self): """Brier score is 1 for worst possible predictions.""" predictions = np.array([1.0, 1.0, 0.0, 0.0]) labels = np.array([0.0, 0.0, 1.0, 1.0]) brier = _compute_brier(predictions, labels) assert brier == pytest.approx(1.0, abs=1e-10) def test_brier_uniform_predictions(self): """Brier score for uniform 0.5 predictions against balanced labels is 0.25.""" predictions = np.array([0.5, 0.5, 0.5, 0.5]) labels = np.array([1.0, 0.0, 1.0, 0.0]) brier = _compute_brier(predictions, labels) assert brier == pytest.approx(0.25, abs=1e-10) def test_ece_empty_returns_zero(self): """ECE of empty arrays is 0.""" ece = _compute_ece(np.array([]), np.array([])) assert ece == 0.0 def test_brier_empty_returns_zero(self): """Brier of empty arrays is 0.""" brier = _compute_brier(np.array([]), np.array([])) assert brier == 0.0 # --- Test Conservative Defaults --- class TestDefaults: """Test conservative defaults for underrepresented classes.""" def test_known_document_type(self): """Known document types return conservative probabilities in [0.3, 0.5].""" result = get_default_confidence("news", "earnings_beat") assert isinstance(result, ConfidenceResult) assert 0.3 <= result.probability <= 0.5 assert result.under_calibrated is True assert result.is_calibrated is False assert "conservative-default" in result.calibration_version def test_unknown_document_type(self): """Unknown document types return the most conservative default (0.3).""" result = get_default_confidence("exotic_type", "unknown_event") assert result.probability == 0.3 assert result.under_calibrated is True def test_unknown_event_class(self): """Unknown event classes use the lowest default.""" result = get_default_confidence("news", "never_seen_before") assert result.probability == 0.30 assert result.under_calibrated is True def test_all_document_types_conservative(self): """All defined document types have defaults in [0.3, 0.5].""" doc_types = ["news", "filing", "transcript", "press_release", "macro_event", "unknown"] for dt in doc_types: result = get_default_confidence(dt, "earnings_beat") assert 0.3 <= result.probability <= 0.5, f"Failed for {dt}" def test_is_underrepresented_no_counts(self): """Without known counts, unknown types are underrepresented.""" assert is_underrepresented("exotic", "unknown_event") is True assert is_underrepresented("news", "earnings_beat") is False def test_is_underrepresented_with_counts(self): """With known counts, low-count classes are underrepresented.""" counts = {("news", "earnings_beat"): 100, ("filing", "merger"): 5} assert is_underrepresented("news", "earnings_beat", known_counts=counts) is False assert is_underrepresented("filing", "merger", known_counts=counts) is True assert is_underrepresented("news", "unknown", known_counts=counts) is True # --- Test Artifact Save/Load --- class TestArtifacts: """Test calibration artifact persistence.""" def test_save_load_roundtrip(self): """Save and load preserves calibrator state.""" features, labels = _generate_training_data(n_samples=50) cal = ConfidenceCalibrator(method="isotonic") cal.fit(features, labels, version="v1.0.0", training_range="test") with tempfile.TemporaryDirectory() as tmpdir: save_artifact(cal, "v1.0.0", tmpdir) loaded = load_artifact(Path(tmpdir) / "v1.0.0") assert loaded.is_fitted assert loaded.version == "v1.0.0" assert loaded.method == "isotonic" # Predictions should match test_features = _make_features() original_pred = cal.predict(test_features) loaded_pred = loaded.predict(test_features) assert original_pred == pytest.approx(loaded_pred, abs=1e-10) def test_save_unfitted_raises(self): """Saving an unfitted calibrator raises ValueError.""" cal = ConfidenceCalibrator() with tempfile.TemporaryDirectory() as tmpdir: with pytest.raises(ValueError, match="unfitted"): save_artifact(cal, "v1.0.0", tmpdir) def test_load_nonexistent_raises(self): """Loading from a missing path raises FileNotFoundError.""" with pytest.raises(FileNotFoundError): load_artifact("/nonexistent/path") def test_load_metadata(self): """Metadata can be loaded independently.""" features, labels = _generate_training_data(n_samples=50) cal = ConfidenceCalibrator(method="platt") cal.fit(features, labels, version="v2.0.0", training_range="2024-01-01 to 2024-12-31") with tempfile.TemporaryDirectory() as tmpdir: save_artifact(cal, "v2.0.0", tmpdir) metadata = load_metadata(Path(tmpdir) / "v2.0.0") assert metadata.version == "v2.0.0" assert metadata.method == "platt" assert metadata.training_count == 50 def test_list_versions(self): """list_versions finds all saved artifact versions.""" features, labels = _generate_training_data(n_samples=50) with tempfile.TemporaryDirectory() as tmpdir: for version in ["v1.0.0", "v1.1.0", "v2.0.0"]: cal = ConfidenceCalibrator(method="isotonic") cal.fit(features, labels, version=version) save_artifact(cal, version, tmpdir) versions = list_versions(tmpdir) assert versions == ["v1.0.0", "v1.1.0", "v2.0.0"] def test_list_versions_empty_dir(self): """list_versions returns empty list for empty or missing directory.""" with tempfile.TemporaryDirectory() as tmpdir: assert list_versions(tmpdir) == [] assert list_versions("/nonexistent") == []