dp-zp-agent/tools/graphrag/tests/test_graphrag_evaluator.py

from __future__ import annotations

from tools.graphrag.evaluation.graphrag_evaluator import (
    aggregate_rows,
    collect_actual_relationships,
    evaluate_result,
    normalize_text,
    property_present,
)


def test_normalize_text() -> None:
    assert normalize_text("Dárius Lindvai") == "darius lindvai"


def test_relationships_can_be_derived_from_fact() -> None:
    result = {
        "paths": [],
        "facts": [
            {
                "fact_type": ("person_keyword"),
                "person": ("Jakub Maruniak"),
                "keyword": "NLP",
            }
        ],
    }

    relationships = collect_actual_relationships(result)

    assert relationships == {
        "HAS_DOCUMENT",
        "HAS_KEYWORD",
    }


def test_thesis_properties_from_fact() -> None:
    result = {
        "resolved_entities": [],
        "paths": [],
        "facts": [
            {
                "fact_type": ("person_thesis"),
                "person": ("Dárius Lindvai"),
                "title": ("Obnovenie interpunkcie"),
                "year": 2021,
                "thesis_type": ("diploma_thesis"),
            }
        ],
    }

    assert property_present(
        "Thesis.title",
        result,
    )

    assert property_present(
        "Thesis.year",
        result,
    )

    assert property_present(
        "Thesis.year|Work.year",
        result,
    )


def test_strict_thesis_result() -> None:
    question = {
        "id": "g-test",
        "split": "dev",
        "category": ("person_title_year"),
        "difficulty": "medium",
        "question": ("Aký názov a rok práce?"),
        "expected_answer_contains": [
            "Dárius Lindvai",
            ("Obnovenie interpunkcie " "pomocou hlbokých neurónových sietí"),
            "2021",
        ],
        "expected_documents": ["pages/students/2016/" "darius_lindvai/README.md"],
        "expected_source_urls": ["https://zp.kemt.fei.tuke.sk/" "students/2016/darius_lindvai"],
        "graph": {
            "schema_version": 2,
            "expected_entities": [
                {
                    "type": ("thesis_title"),
                    "value": ("Obnovenie interpunkcie " "pomocou hlbokých neurónových sietí"),
                },
                {
                    "type": "year",
                    "value": "2021",
                },
            ],
            "expected_relations": [
                "AUTHORED",
            ],
            "expected_properties": [
                "Thesis.title",
                "Thesis.year",
            ],
        },
    }

    result = {
        "strategies": [
            "entity_resolution",
            "thesis_constraints",
        ],
        "resolved_entities": [],
        "paths": [],
        "facts": [
            {
                "fact_type": ("person_thesis"),
                "person": ("Dárius Lindvai"),
                "title": ("Obnovenie interpunkcie " "pomocou hlbokých neurónových sietí"),
                "year": 2021,
                "thesis_type": ("diploma_thesis"),
                "document_path": ("pages/students/2016/" "darius_lindvai/README.md"),
                "source_url": ("https://zp.kemt.fei.tuke.sk/" "students/2016/darius_lindvai"),
            }
        ],
        "documents": [
            {
                "document_path": ("pages/students/2016/" "darius_lindvai/README.md"),
                "source_url": ("https://zp.kemt.fei.tuke.sk/" "students/2016/darius_lindvai"),
            }
        ],
    }

    row = evaluate_result(
        question,
        result,
        latency_seconds=0.01,
    )

    assert row["strict"] is True

    assert row["value_score"] == 1.0

    assert row["document_recall"] == 1.0


def test_failed_result_is_not_strict() -> None:
    question = {
        "id": "g-test",
        "split": "dev",
        "category": ("shared_keyword_pair"),
        "difficulty": "hard",
        "question": "Otázka",
        "expected_answer_contains": [
            "NLP",
        ],
        "expected_documents": [
            "a.md",
        ],
        "expected_source_urls": [],
        "graph": {
            "schema_version": 2,
            "expected_entities": [
                {
                    "type": "keyword",
                    "value": "NLP",
                }
            ],
            "expected_relations": [
                "HAS_KEYWORD",
            ],
            "expected_properties": [],
        },
    }

    result = {
        "strategies": [],
        "resolved_entities": [],
        "paths": [],
        "facts": [],
        "documents": [],
    }

    row = evaluate_result(
        question,
        result,
        latency_seconds=0.01,
    )

    assert row["strict"] is False


def test_aggregate_rows() -> None:
    rows = [
        {
            "strict": True,
            "error": None,
            "value_score": 1.0,
            "entity_score": 1.0,
            "relation_score": 1.0,
            "property_score": 1.0,
            "relation_choice_score": 1.0,
            "document_recall": 1.0,
            "source_recall": 1.0,
            "latency_seconds": 0.1,
        },
        {
            "strict": False,
            "error": None,
            "value_score": 0.5,
            "entity_score": 0.5,
            "relation_score": 0.5,
            "property_score": 0.5,
            "relation_choice_score": 0.5,
            "document_recall": 0.5,
            "source_recall": 0.5,
            "latency_seconds": 0.2,
        },
    ]

    summary = aggregate_rows(rows)

    assert summary["questions"] == 2

    assert summary["strict_count"] == 1

    assert summary["strict_accuracy"] == 0.5

    assert summary["latency_mean_seconds"] == 0.15