dp-zp-agent/test/test_question_plan.py

"""Question semantics, conservative LLM assistance and actual shared execution."""
import copy
import json

import pytest
from pydantic import ValidationError

from tools.answers.question_plan import (QuestionPlan,build_question_plan,validate_model_plan,PlanRejected,PlanModelUnavailable)
from tools.knowledge.sources import profile_units
from tools.knowledge.planned_answers import answer_facets

CATALOG=[{'kind':'Person','name':name,'source_document':path} for name,path in [
 ('Ondrej Megela','pages/students/2018/ondrej_megela/README.md'),
 ('Ján Holp','pages/students/2016/jan_holp/README.md')]]


def plan(q,**kwargs):return build_question_plan(q,CATALOG,**kwargs)


def test_general_question_is_an_overview_but_specific_duplicate_project_needs_selection():
    broad=plan('Aké technické problémy riešil Ondrej Megela?')
    assert broad.spec.intent=='facets' and broad.spec.scope=='across_works' and broad.spec.facets==['problems']
    specific=plan('Aké metódy používal Ján Holp pri diplomovom projekte v roku 2020?')
    assert specific.spec.scope=='single_work' and specific.spec.works[0].years==[2020]
    assert specific.spec.works[0].work_type=='diploma_project'


def test_two_people_type_and_year_remain_bound_to_their_own_work():
    p=plan('Porovnaj zadanie bakalárskej práce Ondreja Megelu z roku 2020 a diplomovej práce Jána Holpa z roku 2021.').spec
    assert p.matches_work('Ondrej Megela','Bakalárska práca 2020','bachelor_thesis',2020)
    assert not p.matches_work('Ondrej Megela','Bakalárska práca 2021','bachelor_thesis',2021)
    assert p.matches_work('Ján Holp','Diplomová práca 2021','diploma_thesis',2021)
    assert not p.matches_work('Ján Holp','Diplomová práca 2020','diploma_thesis',2020)


@pytest.mark.parametrize('q,reason',[
 ('Aké práce má Xaver Neexistujúci?','unknown_person'),
 ('Aké práce má Ján?','incomplete_person_name'),
 ('Kto týmto tempom dokončí prácu v roku 2027?','unsupported_constraint'),
 ('Daj obe.','missing_followup_context'),
])
def test_ambiguity_unknown_identity_and_unsupported_conditions_never_call_model(q,reason):
    result=plan(q,complete=lambda _:pytest.fail('No model may invent missing identity/context'))
    assert result.reason==reason and result.clarification


def test_work_interval_covers_intermediate_years():
    p=plan('Aké práce má Ján Holp od roku 2019 do roku 2021?').spec
    assert p.period.field=='work_year'
    assert p.matches_work('Ján Holp','Diplomový projekt 2020','diploma_project',2020)
    assert not p.matches_work('Ján Holp','Diplomový projekt 2022','diploma_project',2022)


def test_record_period_is_not_work_year():
    p=plan('Ako pokročil Ondrej Megela v roku 2021?').spec
    assert p.period.field=='record_date' and p.works[0].years==[]


@pytest.mark.parametrize('field,value',[
 ('question','Aké práce má niekto iný?'),('people',[]),('works',[]),('years',[2099]),
 ('conditions',[{'kind':'tag','quote':'tag RAG'}]),('scope','across_works'),('facets',[]),
])
def test_a_structurally_valid_model_plan_cannot_change_known_meaning(field,value):
    baseline=plan('Aké problémy riešil Ondrej Megela pri bakalárskej práci v roku 2020?').spec
    altered=baseline.model_dump();altered[field]=value
    with pytest.raises(PlanRejected):validate_model_plan(QuestionPlan.model_validate(altered),baseline)


@pytest.mark.parametrize('intent',['facets','progress'])
def test_complex_paraphrase_uses_typed_plan_not_a_rewritten_question(intent):
    calls=[]
    def model(messages):
        calls.append(messages)
        value=json.loads(messages[1]['content'])['baseline']
        value.update(intent=intent,facets=['problems'])
        return {'choices':[{'message':{'content':json.dumps(value)}}]}
    result=plan('Povedz mi, s čím zápasil Ondrej Megela počas svojej roboty.',complete=model,verify_complete=lambda _: {'choices':[{'message':{'content':'{"decision":"supported"}'}}]})
    assert result.origin=='model' and result.spec.intent=='facets'
    assert len(calls)==1 and result.spec.people[0].name=='Ondrej Megela'


def test_invalid_model_json_retries_but_never_executes_an_unvalidated_plan():
    calls=[]
    def model(messages):calls.append(1);return {'choices':[{'message':{'content':'not JSON'}}]}
    with pytest.raises(PlanRejected):plan('Povedz mi, s čím zápasil Ondrej Megela počas svojej roboty.',complete=model)
    assert len(calls)==2


def test_provider_failure_is_not_no_answer():
    def model(messages):raise RuntimeError('PRIVATE_ENDPOINT')
    with pytest.raises(PlanModelUnavailable):plan('Povedz mi, s čím zápasil Ondrej Megela počas svojej roboty.',complete=model)


def source_rows():
    return profile_units({'title':'Ján Holp','path':CATALOG[1]['source_document'],
        'content':'# Diplomový projekt 2020\nStav:\n- Vyriešený problém Cassandra.\n# Diplomový projekt 2 2020\nStav:\n- Vznikol problém BERT.\nÚlohy:\n- Otestujte model.'},include_records=True)


@pytest.mark.parametrize('engine',['rag','graphrag'])
def test_actual_facet_executor_groups_works_and_preserves_source_bindings(engine):
    p=plan('Aké problémy riešil Ján Holp?')
    rows=source_rows();index={'index_version':{'version':'v'},'version':'content','units':rows}
    result=answer_facets(p,index,rows,engine=engine,limit=5)
    assert result['status']=='ok' and result['accepted_for_delivery']
    assert 'Cassandra' in result['answer'] and 'BERT' in result['answer']
    assert 'Otestujte' not in result['answer']
    assert result['source_urls']==['https://zp.kemt.fei.tuke.sk/students/2016/jan_holp']
    assert {r['work_label'] for r in result['evidence']}=={'Diplomový projekt 2020','Diplomový projekt 2 2020'}


def test_single_project_asks_for_concrete_choice_but_both_projects_are_returned():
    rows=source_rows();index={'index_version':{'version':'v'},'version':'content','units':rows}
    a=answer_facets(plan('Aké problémy riešil Ján Holp pri diplomovom projekte v roku 2020?'),index,rows,engine='rag',limit=5)
    assert a['reason']=='ambiguous_work' and 'projekt 2' in a['answer']
    b=answer_facets(plan('Aké problémy riešil Ján Holp pri oboch diplomových projektoch v roku 2020?'),index,rows,engine='rag',limit=5)
    assert b['status']=='ok' and 'Cassandra' in b['answer'] and 'BERT' in b['answer']


@pytest.mark.parametrize('question',[
    'Aké problémy riešil Ján Holp?',
    'Aké problémy riešil Ján Holp pri oboch diplomových projektoch v roku 2020?',
    'Porovnaj problémy Jána Holpa pri oboch diplomových projektoch v roku 2020.',
])
def test_planned_facet_output_passes_the_existing_mcp_client_contract(monkeypatch,question):
    from packages.claude import client
    rows=source_rows();index={'index_version':{'version':'v'},'version':'content','units':rows}
    result=answer_facets(plan(question),index,rows,engine='rag',limit=5)
    assert result['status']=='ok'
    monkeypatch.setattr(client,'request_json',lambda *a,**kw:{**result,'query':question,'engine':'rag'})
    delivered=client.query_backend(question,'rag',config=object())
    assert delivered['status']=='ok' and delivered['accepted_for_delivery']
    assert delivered['answer']==result['answer'] and delivered['source_urls']==result['source_urls']


def test_no_unsupported_fields_in_provider_schema_and_full_local_validation():
    from tools.answers.model_client import strict_output_schema
    schema=strict_output_schema(QuestionPlan.model_json_schema())
    assert schema['additionalProperties'] is False and set(schema['required'])==set(schema['properties'])
    with pytest.raises(ValidationError):QuestionPlan(question='q',intent='execute_cypher')


@pytest.mark.parametrize('verdict',['unsupported','uncertain'])
def test_plausible_but_semantically_wrong_route_is_vetoed(verdict):
    def proposer(messages):
        value=json.loads(messages[1]['content'])['baseline']
        value.update(intent='facets',facets=['results'])
        return {'choices':[{'message':{'content':json.dumps(value)}}]}
    def judge(messages):
        return {'choices':[{'message':{'content':json.dumps({'decision':verdict})}}]}
    with pytest.raises(PlanRejected):
        plan('Povedz mi, s čím zápasil Ondrej Megela počas svojej roboty.',complete=proposer,verify_complete=judge)


def test_malformed_provider_message_is_rejected_not_a_retrieval_error():
    with pytest.raises(PlanRejected):
        plan('Povedz mi, s čím zápasil Ondrej Megela počas svojej roboty.',complete=lambda _: {'choices':[{'message':{'content':None}}]})


def test_explicit_staff_role_does_not_disappear_when_staff_has_a_student_profile():
    p=plan('Ktoré práce vedie Ján Holp?')
    assert p.spec.people[0].role=='supervisor'


@pytest.mark.parametrize('provider,key',[('anthropic','output_config'),('openai_compatible','response_format')])
def test_schema_is_sent_using_provider_specific_transport(monkeypatch,provider,key):
    from tools.answers import model_client
    sent=[]
    class Reply:
        def __enter__(self):return self
        def __exit__(self,*args):pass
        def read(self,*args):
            return json.dumps({'stop_reason':'end_turn','content':[{'type':'text','text':'{}'}]} if provider=='anthropic' else {'choices':[{'finish_reason':'stop','message':{'content':'{}'}}]}).encode()
    class Opener:
        def open(self,request,**kwargs):sent.append(json.loads(request.data));return Reply()
    monkeypatch.setattr(model_client.urllib.request,'build_opener',lambda *a:Opener())
    config={'provider':provider,'model':'test','key':'private','endpoint':'https://example.test/api'}
    model_client._complete_request([{'role':'user','content':'q'}],config,output_schema={'type':'object','properties':{},'additionalProperties':False})
    assert key in sent[0]


def test_requested_table_and_real_limit_are_respected():
    rows=source_rows();index={'index_version':{'version':'v'},'version':'content','units':rows}
    p=plan('Aké problémy riešil Ján Holp? Uveď ich v tabuľke.')
    value=answer_facets(p,index,rows,engine='rag',limit=5)
    assert '| Oblasť |' in value['answer']
    limited=answer_facets(p,index,rows,engine='rag',limit=1)
    assert limited['truncated'] and 'Zúž otázku' in limited['answer']
    assert len({r['work_label'] for r in limited['evidence']})==1


def test_current_status_is_preserved_without_inference_from_old_notes():
    p=plan('Aké problémy teraz riešil Ján Holp?')
    assert any(c.kind=='current_status' for c in p.spec.conditions)
    rows=source_rows();value=answer_facets(p,{'index_version':{},'version':'v'},rows,engine='rag',limit=5)
    assert 'nepotvrdzujú aktuálny stav' in value['answer']


def test_explicit_record_dates_are_not_silently_replaced_by_work_years():
    p=plan('Aké problémy riešil Ondrej Megela od 1.9.2020 do 31.1.2021?').spec
    assert p.period.field=='record_date' and p.period.start=='2020-09-01' and p.period.end=='2021-01-31'
    assert p.works[0].years==[]


def test_work_year_is_distinct_from_a_record_interval():
    p=plan('Aké problémy riešil Ondrej Megela pri bakalárskej práci 2020 od 1.1.2021 do 31.12.2021?').spec
    assert p.period.field=='record_date' and p.works[0].years==[2020]


def test_structural_format_never_licenses_added_sql_or_cypher():
    value=plan('Aké práce má Ján Holp?').spec.model_dump();value['cypher']='MATCH (n) DELETE n'
    with pytest.raises(ValidationError):QuestionPlan.model_validate(value)


def test_plain_work_question_with_year_is_direct_not_an_llm_paraphrase():
    p=plan('Akú prácu mal Ján Holp v roku 1999?',complete=lambda _:pytest.fail('Clear year/work query must be direct'))
    assert p.origin=='direct' and p.spec.intent=='metadata' and p.spec.years==[1999]


@pytest.mark.parametrize('q',[
 'Aké problémy riešil Ján Holp pri diplomovej práci v roku 99?',
 'Aké problémy riešil Ján Holp bez Cassandry?',
 'Aké problémy riešil Ján Holp v marci?',
])
def test_recognized_but_unimplemented_filter_never_expands_to_general_results(q):
    p=plan(q,complete=lambda _:pytest.fail('Unsupported filters need clarification, not a model'))
    assert p.clarification and p.spec.unsupported


def test_explicit_consultant_is_a_known_identity_not_inferred_from_authorship():
    catalog=CATALOG+[{'kind':'DocumentAuthor','name':'Daniel Hladek'}]
    p=build_question_plan('Ktoré práce konzultoval Daniel Hladek?',catalog)
    assert p.resolution['status']=='ok' and p.spec.people[0].role=='consultant'


def test_plural_work_request_is_not_forced_to_pick_one_work():
    p=plan('Aké diplomové práce má Ján Holp?').spec
    assert p.scope=='selected_works'


def test_graph_adapter_filters_each_persons_work_year_before_counting():
    from tools.graphrag import graph_search as gs
    from tools.graphrag.tests.test_graph_search import FakeSession
    q='Aké názvy majú bakalárska práca Ondrej Megela 2020 a diplomová práca Ján Holp 2021?'
    spec=plan(q).spec
    rows=[{'person':name,'person_element_id':name,'thesis_element_id':str(i),'thesis_id':str(i),'thesis_type':kind,'heading':label,'year':year,'title':'Téma '+str(i),'document_path':path,'source_url':'https://example.test/'+str(i)}
          for i,(name,kind,label,year,path) in enumerate([
          ('Ondrej Megela','bachelor_thesis','Bakalárska práca 2020',2020,CATALOG[0]['source_document']),
          ('Ondrej Megela','bachelor_thesis','Bakalárska práca 2021',2021,CATALOG[0]['source_document']),
          ('Ján Holp','diploma_thesis','Diplomová práca 2021',2021,CATALOG[1]['source_document'])])]
    catalog=[{**r,'element_id':r['name']} for r in CATALOG]
    result=gs.search_graph_with_session(FakeSession({gs.CATALOG_QUERY:catalog,gs.THESIS_FACTS_QUERY:rows}),q,question_plan=spec)
    assert not any(f.get('person')=='Ondrej Megela' and f.get('year')==2021 for f in result['facts'])
    assert any(f.get('person')=='Ondrej Megela' and f.get('year')==2020 for f in result['facts'])


def test_llm_can_bind_supplied_years_but_never_invent_or_drop_them():
    q='Porovnaj zadania: Ondrej Megela bakalárska práca 2020, Ján Holp diplomová práca 2021.'
    def proposer(messages):
        value=json.loads(messages[1]['content'])['baseline']
        value['works'][0]['years']=[2020];value['works'][1]['years']=[2021]
        return {'choices':[{'message':{'content':json.dumps(value)}}]}
    def judge(messages):return {'choices':[{'message':{'content':'{"decision":"supported"}'}}]}
    p=plan(q,complete=proposer,verify_complete=judge)
    assert p.origin=='model' and p.spec.works[0].years==[2020] and p.spec.works[1].years==[2021]
    base=plan(q).spec;invented=p.spec.model_copy(deep=True);invented.works[0].years=[2099]
    with pytest.raises(PlanRejected):validate_model_plan(invented,base,allow_work_binding=True)


def test_slovak_summary_imperative_is_not_misclassified_as_gibberish():
    p=plan('Zhrň zadanie a výsledky diplomovej práce Jána Holpa.')
    assert not p.clarification and p.spec.intent=='facets'


@pytest.mark.parametrize('alias,kind',[('BP','bachelor_thesis'),('DP','diploma_thesis'),('PhD','doctoral_thesis'),('TP','team_project')])
def test_existing_graph_work_type_aliases_remain_explicit_in_the_plan(alias,kind):
    p=plan('Aký názov má '+alias+' Jána Holpa z roku 2021?').spec
    assert p.works[0].work_type==kind and p.works[0].years==[2021]


def test_model_repairs_changed_per_person_scope_without_relaxing_validation():
    q='Porovnaj zadania: Ondrej Megela bakalárska práca 2020, Ján Holp diplomová práca 2021.'
    calls=[]
    def propose(messages):
        calls.append(messages[-1]['content'])
        value=json.loads(messages[1]['content'])['baseline']
        value['works'][0]['years']=[2020];value['works'][1]['years']=[2021]
        if len(calls)==1:value['scope']='selected_works'
        return {'choices':[{'message':{'content':json.dumps(value)}}]}
    judge=lambda _: {'choices':[{'message':{'content':'{"decision":"supported"}'}}]}
    result=plan(q,complete=propose,verify_complete=judge)
    assert result.spec.scope=='single_work' and len(calls)==2
    assert 'Changed work selection scope' in calls[1] and 'scope=single_work' in calls[1]


def test_work_limit_does_not_starve_the_second_requested_person():
    first=profile_units({'title':'Ondrej Megela','path':CATALOG[0]['source_document'],
        'content':'# Bakalárska práca 2020\nStav:\n- Vznikol problém s Pytorch.\n# Diplomová práca 2023\nStav:\n- Vznikol problém s vyhodnotením.'},include_records=True)
    rows=first+source_rows()
    p=plan('Aké problémy riešili Ondrej Megela a Ján Holp?')
    result=answer_facets(p,{'index_version':{},'version':'v'},rows,engine='rag',limit=2)
    assert {r['person'] for r in result['evidence']}=={'Ondrej Megela','Ján Holp'}
    assert result['truncated']


def test_budget_exhaustion_is_not_misreported_as_absent_evidence():
    text='\n'.join('# Diplomová práca '+str(y)+'\nStav:\n- Vznikol problém. '+('Podrobný záznam. '*350) for y in range(2020,2024))
    rows=profile_units({'title':'Ján Holp','path':CATALOG[1]['source_document'],'content':text},include_records=True)
    result=answer_facets(plan('Aké problémy riešil Ján Holp?'),{'index_version':{},'version':'v'},rows,engine='rag',limit=4)
    assert result['status']=='ambiguous' and result['reason']=='content_budget_exceeded'
    assert 'údaje chýbajú' in result['answer']


@pytest.mark.parametrize('suffix',['Uveď to v tabuľke.','Odpovedz stručne.'])
def test_model_cannot_drop_an_explicit_presentation_requirement(suffix):
    baseline=plan('Aké problémy riešil Ján Holp? '+suffix).spec
    changed=baseline.model_copy(update={'presentation':'list'})
    with pytest.raises(PlanRejected):validate_model_plan(changed,baseline)