AI/ontology_platform/tests/unit/test_phase7_hybrid_extraction.py

from crawler_platform.app.config.loader import ProjectConfig
from crawler_platform.app.api.routes import extraction_log_summary
from crawler_platform.app.core.extractor.base import ExtractedClaim, ExtractionBundle, ExtractionPageContext
from crawler_platform.app.core.extractor.hybrid import HybridExtractor, fallback_bundle, mark_bundle, merge_bundles
from crawler_platform.app.core.extractor.validation import validate_extraction_bundle


def project_config() -> ProjectConfig:
    return ProjectConfig(
        project_name="phase7-test",
        domain="generic",
        target_entities=["Thing"],
        fields=[],
        sources=[],
        ontology={
            "entity_types": ["Thing"],
            "predicates": ["hasPart", "color", "size"],
            "relation_types": {
                "hasPart": {
                    "allowed_subject_types": ["Thing"],
                    "allowed_object_types": ["Thing"],
                    "confidence_rules": {"min_confidence": 0.1},
                },
                "color": {
                    "allowed_subject_types": ["Thing"],
                    "literal_value": True,
                    "confidence_rules": {"min_confidence": 0.1},
                },
                "size": {
                    "allowed_subject_types": ["Thing"],
                    "literal_value": True,
                    "confidence_rules": {"min_confidence": 0.1},
                },
            },
        },
    )


def test_hybrid_merge_marks_agreement_conflict_and_counts() -> None:
    rule = ExtractionBundle(
        claims=[
            ExtractedClaim("A", "Thing", "hasPart", "B", "Thing", confidence=0.7),
            ExtractedClaim("A", "Thing", "color", object_value="red", confidence=0.6),
        ]
    )
    llm = ExtractionBundle(
        claims=[
            ExtractedClaim("A", "Thing", "hasPart", "B", "Thing", confidence=0.8),
            ExtractedClaim("A", "Thing", "color", object_value="blue", confidence=0.65),
            ExtractedClaim("A", "Thing", "size", object_value="large", confidence=0.5),
        ]
    )
    mark_bundle(rule, source="rule", mode="compare")
    mark_bundle(llm, source="llm", mode="compare")

    merged = merge_bundles(rule, llm, HybridExtractor("generic", mode="compare"))

    assert merged.extractor_name == "compare_rule_llm_extractor"
    assert merged.raw_output["comparison"] == {
        "both_agree": 1,
        "rule_only": 0,
        "llm_only": 1,
        "conflict": 1,
        "rejected_by_validation": 0,
    }
    by_predicate = {(claim.predicate, claim.object_name or claim.object_value): claim for claim in merged.claims}
    assert by_predicate[("hasPart", "B")].metadata["agreement"] == "rule_and_llm"
    assert by_predicate[("color", "red")].metadata["review_required"] is True
    assert by_predicate[("color", "blue")].metadata["conflict_status"] == "rule_llm_conflict"
    assert by_predicate[("size", "large")].metadata["agreement"] == "llm_only"


def test_validation_keeps_evidence_missing_llm_only_claim_as_candidate() -> None:
    bundle = ExtractionBundle(
        claims=[
            ExtractedClaim(
                "A",
                "Thing",
                "color",
                object_value="blue",
                confidence=0.8,
                metadata={
                    "agreement": "llm_only",
                    "extraction_source": "llm",
                    "llm_confidence": 0.8,
                },
            )
        ],
        extractor_name="hybrid_rule_llm_extractor",
        provider="lm_studio",
        raw_output={"extraction_mode": "hybrid"},
    )

    result = validate_extraction_bundle(bundle, project_config())

    assert len(result.bundle.claims) == 1
    claim = result.bundle.claims[0]
    assert claim.metadata["validation_status"] == "candidate_claim"
    assert claim.metadata["review_required"] is True
    assert claim.metadata["review_reason"] == "claim has no evidence"


def test_validation_marks_rule_only_claim_as_rule_candidate() -> None:
    bundle = ExtractionBundle(
        claims=[
            ExtractedClaim(
                "Alpha",
                "Thing",
                "color",
                object_value="red",
                evidence_text="Alpha color red",
                confidence=0.7,
                metadata={
                    "agreement": "rule_only",
                    "extraction_source": "rule",
                    "rule_confidence": 0.7,
                },
            )
        ],
        extractor_name="hybrid_rule_llm_extractor",
        provider="lm_studio",
        raw_output={"extraction_mode": "hybrid"},
    )

    result = validate_extraction_bundle(bundle, project_config())

    assert result.bundle.claims[0].metadata["validation_status"] == "rule_candidate"
    assert result.bundle.claims[0].metadata["claim_kind"] == "rule_candidate"


def test_hybrid_smart_routing_skips_llm_for_category_page() -> None:
    context = ExtractionPageContext(
        url="https://example.test/category",
        final_url=None,
        title="Category",
        page_type="CategoryPage",
        clean_text="Category listing " * 40,
    )

    bundle = HybridExtractor("generic", mode="hybrid").extract_from_context(context, project_config())

    assert bundle.extractor_name == "hybrid_rule_only_routed"
    assert bundle.raw_output["llm_skipped"] is True
    assert bundle.raw_output["effective_extraction_mode"] == "rule_only"
    assert "CategoryPage" in bundle.raw_output["llm_skip_reason"]


def test_rule_only_mode_records_rule_llm_summary_counts() -> None:
    rule_bundle = ExtractionBundle(
        claims=[
            ExtractedClaim("A", "Thing", "color", object_value="red"),
            ExtractedClaim("A", "Thing", "size", object_value="large"),
        ]
    )

    class FixedRuleExtractor:
        def extract(self, page_text, project_config):  # noqa: ANN001
            return rule_bundle

        def extract_from_context(self, context, project_config):  # noqa: ANN001
            return rule_bundle

    extractor = HybridExtractor("generic", mode="rule_only")
    extractor.rule_extractor = FixedRuleExtractor()

    bundle = extractor.extract("Alpha color red size large", project_config())

    assert bundle.raw_output["rule_claim_count"] == 2
    assert bundle.raw_output["llm_claim_count"] == 0
    assert bundle.raw_output["comparison"]["rule_only"] == 2


def test_fallback_records_rule_counts_and_comparison() -> None:
    rule_bundle = ExtractionBundle(
        claims=[
            ExtractedClaim("A", "Thing", "color", object_value="red"),
            ExtractedClaim("A", "Thing", "size", object_value="large"),
        ]
    )
    mark_bundle(rule_bundle, source="rule", mode="hybrid")

    bundle = fallback_bundle(rule_bundle, HybridExtractor("generic", mode="hybrid"), RuntimeError("boom"))

    assert bundle.raw_output["rule_claim_count"] == 2
    assert bundle.raw_output["llm_claim_count"] == 0
    assert bundle.raw_output["comparison"]["rule_only"] == 2


def test_extraction_log_summary_backfills_legacy_rule_only_counts() -> None:
    summary = extraction_log_summary(
        {
            "validation": {"rejected_claim_count": 4},
            "candidate_claims": [
                {"metadata": {"agreement": "rule_only", "extraction_source": "rule"}},
                {"metadata": {"agreement": "rule_only", "extraction_source": "rule"}},
            ]
        }
    )

    assert summary["rule_claim_count"] == 2
    assert summary["llm_claim_count"] == 0
    assert summary["comparison"]["rule_only"] == 2
    assert summary["comparison"]["rejected_by_validation"] == 4