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

from __future__ import annotations

import pytest

from tools.graphrag import graph_search as gs
from tools.graphrag.evaluation.audit_benchmark import audit_questions
from tools.graphrag.evaluation.graphrag_evaluator import (
    evaluate_result,
    property_present,
    value_is_present,
)
from .test_graph_search import FakeSession, sample_catalog


@pytest.mark.parametrize("word", ["názov", "názvom", "názvu"])
def test_title_inflections(word):
    assert gs.extract_query_constraints(f"Prepoj osobu s {word} jej práce.")["mentions_title"]


@pytest.mark.parametrize(
    "question",
    [
        "V ktorom rovnakom roku je uvedená práca pri osobách?",
        "Aký rok práce spája tieto osoby?",
        "Prepoj osobu s názvom jej práce a rokom.",
    ],
)
def test_academic_year_intent_does_not_require_a_numeric_year(question):
    assert gs.extract_query_constraints(question)["mentions_academic_year"]


def test_tag_is_extracted_without_following_clause():
    assert gs.explicit_keyword_phrase("Kľúčové slovo RAG a jej práca v roku 2024?") == "rag"


def test_incidental_zp_tag_does_not_filter_title_lookup():
    catalog = sample_catalog() + [{"element_id": "zp", "kind": "Keyword", "id": "zp", "name": "zp"}]
    fact = {
        "person_element_id": "p2",
        "person": "Jakub Maruniak",
        "thesis_element_id": "th1",
        "title": "Ukážková práca",
        "year": 2021,
        "thesis_type": "diploma_thesis",
    }
    session = FakeSession({gs.CATALOG_QUERY: catalog, gs.THESIS_FACTS_QUERY: [fact]})
    result = gs.search_graph_with_session(
        session, "Aký názov práce je v ZP Wiki pri osobe Jakub Maruniak?"
    )
    assert any(row["fact_type"] == "person_thesis" for row in result["facts"])
    assert gs.PERSON_KEYWORD_FACTS_QUERY not in {query for query, _ in session.calls}


def test_keyword_and_year_cannot_be_satisfied_by_different_people():
    session = FakeSession(
        {
            gs.CATALOG_QUERY: sample_catalog(),
            gs.THESIS_FACTS_QUERY: [
                {
                    "person_element_id": "p2",
                    "person": "Jakub Maruniak",
                    "year": 2024,
                    "thesis_element_id": "th1",
                }
            ],
            gs.PERSON_KEYWORD_FACTS_QUERY: [
                {
                    "person_element_id": "p1",
                    "person": "Maroš Harahus",
                    "keyword_element_id": "k2",
                    "keyword": "RAG",
                }
            ],
            gs.THESIS_PATHS_QUERY: [
                {
                    "nodes": [{"kind": "Thesis", "element_id": "th1", "year": 2024}],
                    "relationships": ["AUTHORED"],
                }
            ],
        }
    )
    result = gs.search_graph_with_session(
        session,
        "Ktorá z osôb Maroš Harahus a Jakub Maruniak má kľúčové slovo RAG a jej práca je z roku 2024?",
    )
    assert result["facts"] == []
    assert result["paths"] == []


def test_unknown_explicit_tag_does_not_return_unfiltered_titles():
    session = FakeSession(
        {gs.CATALOG_QUERY: sample_catalog(), gs.THESIS_FACTS_QUERY: [{"title": "Nepodložené"}]}
    )
    result = gs.search_graph_with_session(
        session, "Aký názov práce má Jakub Maruniak, ak má kľúčové slovo neexistujucitag?"
    )
    assert result["facts"] == []
    assert result["constraints"]["unresolved_keyword"]


def test_unknown_quoted_title_does_not_list_every_thesis():
    session = FakeSession(
        {gs.CATALOG_QUERY: sample_catalog(), gs.THESIS_FACTS_QUERY: [{"title": "Iná práca"}]}
    )
    result = gs.search_graph_with_session(
        session, "Kto riešil diplomovú prácu „Úplne vymyslený názov“?"
    )
    assert result["facts"] == []


def test_recognized_query_entities_are_not_answer_evidence():
    result = {
        "resolved_entities": [{"kind": "Thesis", "title": "Vymyslený názov", "year": 2024}],
        "facts": [],
        "paths": [],
        "documents": [{"document_path": "Vymyslený názov"}],
    }
    assert not value_is_present("Vymyslený názov", result)
    assert not property_present("Thesis.title", result)


def test_value_matching_has_word_boundaries():
    assert not value_is_present("202", {"facts": [{"year": 2024}]})


def test_wrong_entity_type_cannot_satisfy_expected_keyword():
    question = {"graph": {"expected_entities": [{"type": "keyword", "value": "RAG"}]}}
    row = evaluate_result(
        question, {"facts": [{"title": "RAG", "thesis_type": "diploma_thesis"}]}, latency_seconds=0
    )
    assert row["entity_score"] == 0


def test_correct_set_with_wrong_count_fails_strict():
    question = {
        "category": "candidate_set_aggregation",
        "graph": {
            "expected_count": 1,
            "expected_entity_set": ["Eva Nová"],
            "entity_set_mode": "exact",
        },
    }
    result = {
        "facts": [{"fact_type": "person_thesis", "person": "Eva Nová"}],
        "aggregation": {"count": 2, "complete": True},
    }
    row = evaluate_result(question, result, latency_seconds=0)
    assert not row["strict"] and not row["count_ok"]


def test_truncated_result_cannot_pass_strict():
    assert not evaluate_result({"graph": {}}, {"truncated": True}, latency_seconds=0)["strict"]


def test_source_audit_detects_nonexistent_tag_and_omitted_child_category():
    payload = {
        "people": [{"id": "p", "name": "Eva Nová", "source_document": "p.md"}],
        "documents": [
            {"path": "p.md", "categories": ["bp2024"]},
            {"path": "child.md", "categories": ["tp2023"]},
        ],
        "person_documents": [
            {"person_id": "p", "document_path": path} for path in ["p.md", "child.md"]
        ],
        "keywords": [{"id": "nlp", "name": "NLP"}],
        "document_keywords": [],
        "theses": [],
        "works": [],
    }
    question = {
        "id": "example",
        "split": "dev",
        "category": "person_multiple_categories",
        "question": "",
        "graph": {
            "start_entities": [{"type": "person", "value": "Eva Nová"}],
            "expected_entities": [{"type": "keyword", "value": "Chýbajúci tag"}],
            "expected_entity_set": ["bp2024"],
            "entity_set_mode": "exact",
        },
    }
    issues = audit_questions([question], payload)["questions"][0]["issues"]
    assert "keyword_not_in_frontmatter_catalog" in issues
    assert "exact_categories_omit_or_add_source_categories" in issues