dp-zp-agent/test/test_graph_query_regressions.py

"""Regressions for actual Claude questions and adversarial same-person joins."""

import copy
import json
from pathlib import Path

import pytest

from tools.answers.checked_generation import generate_checked_answer
from tools.answers.context_policy import answer_context_issues
from tools.answers.graph_context import graph_context
from tools.answers.render_answer import render_graph_answer
from tools.answers.answer_delivery import deliver_answer
from tools.graphrag import graph_search as gs
from tools.graphrag.tests.test_graph_search import FakeSession
from tools.retrieval.question_constraints import unsupported_constraint

PEOPLE = ["Ján Pták", "Ján Malinovský", "Valerii Kutsenko", "Jakub Schwarc"]
PROMPT = Path("tools/answers/graph_answer_prompt.txt").read_text()


def person(index, name):
    return {
        "kind": "Person",
        "name": name,
        "element_id": f"p{index}",
        "source_document": f"pages/students/2022/p{index}/README.md",
    }


def source(index):
    return {
        "document_path": f"pages/students/2022/p{index}/README.md",
        "source_url": f"https://example.test/p{index}",
    }


@pytest.fixture
def tagged_projects():
    catalog = [person(i, name) for i, name in enumerate(PEOPLE, 1)]
    catalog.append({"kind": "Keyword", "name": "RAG", "id": "rag", "element_id": "krag"})
    projects = [
        {
            "person": name,
            "person_element_id": f"p{i}",
            "year": 2026,
            "work_type": "diploma_project",
            "work_id": f"work:p{i}",
            "heading": "Diplomový projekt 2026",
            "title": f"Téma projektu {name}",
            **source(i),
        }
        for i, name in enumerate(PEOPLE, 1)
    ]
    tags = [
        {
            "person": name,
            "person_element_id": f"p{i}",
            "keyword": "RAG",
            "keyword_id": "rag",
            "keyword_element_id": "krag",
            **source(i),
        }
        for i, name in enumerate(PEOPLE[:3], 1)
    ]
    # A fifth person has the tag, but no work in the requested year/type.
    tags.append(
        {
            "person": "Iná Osoba",
            "person_element_id": "p5",
            "keyword": "RAG",
            "keyword_id": "rag",
            "keyword_element_id": "krag",
            **source(5),
        }
    )
    return FakeSession(
        {
            gs.CATALOG_QUERY: catalog,
            gs.WORK_FACTS_QUERY: projects,
            gs.PERSON_KEYWORD_FACTS_QUERY: tags,
        }
    )


@pytest.mark.parametrize(
    "question",
    [
        "Ktorí študenti majú tag RAG a zároveň diplomový projekt z roku 2026?",
        "Ktorí študenti s tagom RAG majú diplomový projekt v roku 2026?",
        "Ktorí študenti majú tag RAG a diplomový projekt z roku 2026?",
        "Ktorí študenti majú tag „RAG“ a zároveň diplomový projekt z roku 2026?",
        "Ktorí študenti s tagom RAG a zároveň diplomovým projektom z roku 2026?",
        "Ktorí študenti s tagom RAG a diplomovým projektom z roku 2026?",
    ],
)
def test_tag_and_project_are_witnessed_by_the_same_person(tagged_projects, question):
    r = gs.search_graph_with_session(tagged_projects, question, limit=50)
    assert r["status"] == "ok" and r["constraints"]["keyword_phrase"] == "rag"
    assert {f["person"] for f in r["facts"]} == set(PEOPLE[:3])
    work = [params for q, params in tagged_projects.calls if q == gs.WORK_FACTS_QUERY][0]
    tags = [params for q, params in tagged_projects.calls if q == gs.PERSON_KEYWORD_FACTS_QUERY][0]
    assert work["years"] == [2026] and work["work_types"] == ["diploma_project"]
    assert tags["keyword_ids"] == ["krag"]
    context = graph_context(r)
    assert {c["person"] for c in context["answer_candidates"]} == set(PEOPLE[:3])
    assert all(c["work_type"] == "diploma_project" for c in context["answer_candidates"])


def test_diploma_thesis_cannot_be_replaced_by_a_project(tagged_projects):
    r = gs.search_graph_with_session(
        tagged_projects,
        "Ktorí študenti majú tag RAG a zároveň diplomovú prácu z roku 2026?",
        limit=50,
    )
    assert r["status"] == "no_answer" and r["reason"] == "no_matching_facts"
    assert not r["facts"] and not r["constraints"]["unresolved_keyword"]


@pytest.mark.parametrize("limit", [1, 2, 3, 4, 5, 6])
def test_tag_predicate_stays_bound_to_visible_records_after_truncation(tagged_projects, limit):
    q = "Ktorí študenti majú tag RAG a zároveň diplomový projekt z roku 2026?"
    raw = gs.search_graph_with_session(tagged_projects, q, limit=limit)
    assert len(raw["facts"]) <= limit and raw["truncated"] == (limit < 3)
    context = graph_context(raw)
    visible = {r["person"] for r in raw["facts"] if r["fact_type"] == "person_work"}
    assert {c["person"] for c in context["answer_candidates"]} == visible
    assert all(r["person"] in visible for r in raw["supporting_facts"])
    assert all(r["fact_type"] == "person_keyword" for r in raw["supporting_facts"])
    claims = [
        {k: c[k] for k in ["person", "year", "keyword", "work_type"]}
        for c in context["answer_candidates"]
    ]
    answer = {
        "status": "ok",
        "claims": claims,
        "source_urls": [s["source_url"] for s in context["sources"]],
    }
    rendered = render_graph_answer(answer, context)
    assert not answer_context_issues(rendered, context)
    assert "Jakub Schwarc" not in rendered["answer"] and "Iná Osoba" not in rendered["answer"]
    assert ("Výpis je neúplný." in rendered["answer"]) == (limit < 3)


def test_count_model_input_keeps_actual_total_and_only_claim_fields(tagged_projects):
    import json
    from tools.answers.answer_format import Claim

    q = "Koľko študentov má tag RAG a zároveň diplomový projekt z roku 2026?"
    context = graph_context(gs.search_graph_with_session(tagged_projects, q, limit=1))
    before = copy.deepcopy(context)

    def select(messages):
        sent = json.loads(messages[1]["content"])["graph_context"]
        assert sent["aggregation"]["count"] == 3 and "candidate_count" not in sent["aggregation"]
        assert "count=3" in messages[0]["content"]
        assert len(sent["answer_candidates"]) == 1
        assert set(sent["answer_candidates"][0]) <= set(Claim.model_fields)
        a = {
            "status": "ok",
            "answer": "",
            "count": sent["aggregation"]["count"],
            "claims": sent["answer_candidates"],
            "source_urls": context["answer_candidates"][0]["source_urls"],
        }
        return {"choices": [{"message": {"content": json.dumps(a)}}]}

    r = generate_checked_answer(q, context, PROMPT, select)
    assert r["accepted_for_delivery"] and r["answer"]["count"] == 3 and context == before


def test_unknown_tag_cannot_trigger_a_global_project_listing(tagged_projects):
    r = gs.search_graph_with_session(
        tagged_projects, "Ktorí študenti majú tag XNEZNAMY a zároveň diplomový projekt z roku 2026?"
    )
    assert r["status"] == "no_answer" and r["reason"] == "unknown_keyword" and not r["facts"]


@pytest.mark.parametrize(
    "question,expected",
    [
        ("Ľudia s tagom strojové učenie a zároveň diplomovou prácou.", "strojove ucenie"),
        (
            "Ľudia s tagom „výskum a aplikácie“ a zároveň diplomovým projektom.",
            "vyskum a aplikacie",
        ),
        ("Ľudia s tagom RAG z roku 2026.", "rag"),
    ],
)
def test_multiword_tag_and_condition_boundary(question, expected):
    assert gs.explicit_keyword_phrase(question) == expected


@pytest.mark.parametrize(
    "question,code",
    [
        ("Ktorí študenti majú aspoň tri práce?", "quantity_condition"),
        ("Ktorí študenti majú aspoň 3 práce?", "quantity_condition"),
        ("Kto má najmenej tri práce?", "quantity_condition"),
        ("Študenti s najviac dvoma projektmi?", "quantity_condition"),
        ("Študenti s minimálne piatimi prácami?", "quantity_condition"),
        (
            "ktorý študenti týmto tempom ukončia svoje diplomové práce v roku 2027?",
            "completion_prediction",
        ),
        ("Ktorí študenti ukončia diplomové práce v roku 2027?", "completion_prediction"),
        ("Odhadni termín dokončenia diplomovej práce Jána Holpa.", "completion_prediction"),
        ("Kedy Ján Holp ukončil diplomovú prácu?", "work_event_date"),
        ("Aké metódy predikcie počasia používa diplomová práca?", None),
        ("Aký je názov práce „Kto dokončí prácu v roku 2027“?", None),
    ],
)
def test_unsupported_condition_wording(question, code):
    assert unsupported_constraint(question) == code


@pytest.mark.parametrize("engine", ["rag", "graphrag"])
@pytest.mark.parametrize(
    "question,word",
    [
        ("Ktorí študenti majú aspoň tri práce?", "počet"),
        ("Ktorí študenti týmto tempom ukončia diplomové práce v roku 2027?", "predpovedať"),
    ],
)
def test_unsupported_condition_never_calls_a_model(engine, question, word):
    r = deliver_answer(
        question,
        engine=engine,
        complete=lambda _: pytest.fail("No model for unsupported condition"),
    )
    assert r["status"] == "ambiguous" and r["accepted_for_delivery"] and word in r["answer"]
    assert not r["claims"] and not r["source_urls"]


def selected(context):
    fields = [
        "person",
        "record_id",
        "title",
        "year",
        "work_type",
        "work_label",
        "keyword",
        "category",
        "document_author",
    ]
    rows = context.get("required_academic_records") or context["facts"]
    claims = [{k: r[k] for k in fields if r.get(k) is not None} for r in rows]
    return {
        "status": "ok",
        "answer": "",
        "claims": claims,
        "source_urls": [s["source_url"] for s in context["sources"]],
    }


def generated(answer, context):
    return generate_checked_answer(
        context.get("question", ""),
        context,
        PROMPT,
        lambda _: {"choices": [{"message": {"content": json.dumps(answer)}}]},
    )


def test_titleless_records_survive_and_are_delivered_without_invented_titles():
    catalog = [person(1, "Ján Holp")]
    records = [
        {
            "person": "Ján Holp",
            "person_element_id": "p1",
            "work_id": f"project-{i}",
            "title": None,
            "heading": label,
            "work_type": "diploma_project",
            "year": 2020,
            **source(1),
        }
        for i, label in enumerate(["Diplomový projekt 2020", "Diplomový projekt 2 2020"])
    ]
    s = FakeSession({gs.CATALOG_QUERY: catalog, gs.WORK_FACTS_QUERY: records})
    raw = gs.search_graph_with_session(s, "Aký názov má diplomový projekt Jána Holpa z roku 2020?")
    assert raw["status"] == "ok" and len(raw["facts"]) == 2
    context = graph_context(raw)
    a = selected(context)
    result = generated(a, context)
    assert result["accepted_for_delivery"]
    assert result["answer"]["answer"].count("názov v zdroji nie je uvedený") == 2
    assert len(result["answer"]["claims"]) == 2 and not any(
        c.get("title") for c in result["answer"]["claims"]
    )
    bad = copy.deepcopy(a)
    bad["claims"][0]["title"] = "Vymyslený názov"
    assert "unsupported_academic_binding" in answer_context_issues(bad, context, check_text=False)
    bad = copy.deepcopy(a)
    bad["claims"].pop()
    assert "missing_academic_record_details" in answer_context_issues(
        bad, context, check_text=False
    )


def shared_context(three=False):
    people = PEOPLE[:3] if three else PEOPLE[:2]
    raw = {
        "status": "ok",
        "reason": None,
        "complete": True,
        "truncated": False,
        "aggregation": None,
        "query": "Aké spoločné kľúčové slová majú " + " a ".join(people) + "?",
        "constraints": gs.extract_query_constraints("spoločné kľúčové slová"),
        "resolved_entities": [person(i, p) for i, p in enumerate(people, 1)],
        "sources": [source(i) for i in range(1, len(people) + 1)],
        "facts": [],
    }
    for keyword, left, right in [("NLP", 1, 2), ("RAG", 1, 2)] + ([("NLP", 1, 3)] if three else []):
        raw["facts"].append(
            {
                "fact_type": "shared_keyword",
                "keyword": keyword,
                "left_person": people[left - 1],
                "right_person": people[right - 1],
                "left_document_path": source(left)["document_path"],
                "right_document_path": source(right)["document_path"],
                "left_source_url": source(left)["source_url"],
                "right_source_url": source(right)["source_url"],
            }
        )
    return graph_context(raw)


def test_shared_tags_are_summarized_once_and_require_all_peoples_evidence():
    context = shared_context()
    a = selected(context)
    result = generated(a, context)
    assert result["accepted_for_delivery"]
    text = result["answer"]["answer"]
    assert "Spoločné kľúčové slová: **NLP**, **RAG**." in text
    assert text.count("NLP") == 1 and text.count("RAG") == 1 and "Pri osobe" not in text
    bad = copy.deepcopy(a)
    bad["claims"].pop()
    assert "missing_shared_property_evidence" in answer_context_issues(
        bad, context, check_text=False
    )


def test_pairwise_shared_tag_does_not_mean_shared_by_everyone():
    context = shared_context(three=True)
    assert {r["value"] for r in context["shared_properties"]} == {"NLP"}
    text = render_graph_answer(selected(context), context)["answer"]
    assert "Spoločné kľúčové slová: **NLP**." in text and "**RAG**" not in text


def test_incomplete_shared_tags_cannot_claim_exhaustive_results():
    context = shared_context()
    context.update(complete=False, truncated=True, truncation_reasons=["fact_limit"])
    result = generated(selected(context), context)
    assert result["accepted_for_delivery"]
    assert "Doložené spoločné kľúčové slová" in result["answer"]["answer"]
    assert "Výpis je neúplný." in result["answer"]["answer"]