kg-agent/scripts/test_simple_plan.py

"""Offline correctness and model-call regression for the simple planning path."""
from copy import deepcopy
from types import SimpleNamespace
import os
import unittest
from unittest.mock import patch
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command
from step3_qa_agent.agent import nodes, production, llm
from step3_qa_agent.agent.graph import build_agent_graph
from step3_qa_agent.agent.simple_plan import build_simple_plan
from step3_qa_agent.agent.timing import measure_llm_calls, llm_calls
from step3_qa_agent.retrieval.production_query import compile_query

SCHEMA = {'meta': {'schema_version': 2}, 'nodes': [
    {'id': '人员信息', 'active': True, 'attributes': ['工号', '姓名', '岗位名称', '部门'],
     'identity_fields': ['工号'], 'multivalue_fields': {'岗位名称': 1}},
    {'id': '项目信息', 'attributes': ['项目编号', '项目名称']}],
    'relations': [{'id': 'works', 'source': '人员信息', 'target': '项目信息', 'type': '服务于'}]}
CONTEXT = {'schema': SCHEMA, 'build_id': 'v1', 'catalog': [], 'assessment_text': '',
           'selected_datasets': ['人员信息'], 'loaded_datasets': [], 'selection_source': 'understanding',
           'issues': []}

def state(question='共有多少名员工？'):
    return {'question': question, 'category': '图谱检索', 'qa_context': deepcopy(CONTEXT),
            'slots': {'anchors': [], 'filters': {}, 'aggregation': '汇总'},
            'simple_query_needs': {'datasets': ['人员信息'], 'relations': []},
            'simple_query': {'kind': 'count', 'dataset': '人员信息', 'filters': [],
                             'fields': [], 'limit': 50, 'complete': True}}

class SimplePlanTests(unittest.TestCase):
    def setUp(self):
        self.env = patch.dict(os.environ, {'QA_FAST_PLAN_ENABLED': 'true'})
        self.env.start()
        self.addCleanup(self.env.stop)

    def test_count_is_distinct_node_query_not_schema_count(self):
        s = state()
        plan, _ = build_simple_plan(s)
        spec = plan['steps'][0]['params']
        compiled = compile_query(spec, SCHEMA, 'v1')
        self.assertIn('count(DISTINCT n._kg_id)', compiled.query)
        self.assertEqual(compiled.params['build'], 'v1')
        with patch.object(production, 'chat_json') as model:
            result = production.plan(s)
        model.assert_not_called()
        self.assertEqual(result['trace'][-1]['mode'], 'simple')
        self.assertEqual(result['trace'][-1]['llm_calls'], [])

    def test_property_query_keeps_anchor_and_provenance(self):
        s = state('张三的岗位名称是什么？')
        s['slots'] = {'anchors': [{'raw': '张三'}]}
        s['simple_query'].update(kind='list', fields=['姓名', '岗位名称'],
             filters=[{'field': '姓名', 'op': 'eq', 'value': '张三'}])
        plan, _ = build_simple_plan(s)
        compiled = compile_query(plan['steps'][0]['params'], SCHEMA, 'v1')
        self.assertEqual(compiled.params['v0'], '张三')
        self.assertIn('_sources', compiled.query)
        self.assertIn('any(x IN CASE', compiled.query)

    def test_list_preserves_all_and_filters_and_limit(self):
        s = state('列出运营部岗位名称包含保安的员工姓名')
        s['slots'] = {'filters': {'部门': '运营部', '岗位': '保安'}}
        s['simple_query'].update(kind='list', fields=['工号', '姓名'], limit=20,
            filters=[{'field': '部门', 'op': 'eq', 'value': '运营部'},
                     {'field': '岗位名称', 'op': 'contains', 'value': '保安'}])
        plan, _ = build_simple_plan(s)
        compiled = compile_query(plan['steps'][0]['params'], SCHEMA, 'v1')
        self.assertEqual((compiled.params['v0'], compiled.params['v1']), ('运营部', '保安'))
        self.assertEqual(compiled.params['result_limit'], 21)
        self.assertIn(' AND ', compiled.query)

    def test_complex_questions_reject_even_when_model_claims_complete(self):
        for question in ['他们有多少人', '按部门统计人数', '每个项目多少人', '保安或电工有多少',
                         '没有证书的员工', '去年入职人数', '工资大于5000人数',
                         '前10名员工', '删除所有员工', '比较项目人数']:
            with self.subTest(question=question):
                self.assertIsNone(build_simple_plan(state(question))[0])

    def test_relation_and_invalid_selection_never_disappear(self):
        for needs in [{'datasets': ['人员信息'], 'relations': ['unknown']},
                      {'datasets': ['人员信息', '项目信息'], 'relations': ['works']},
                      {'datasets': ['人员信息', 'unknown'], 'relations': []}]:
            s = state()
            s['simple_query_needs'] = needs
            self.assertIsNone(build_simple_plan(s)[0])

    def test_dropped_anchor_or_filter_is_rejected(self):
        for slots in [{'anchors': [{'raw': '张三'}]}, {'filters': {'岗位': '保安'}},
                      {'time_range': {'start': '2026-01'}}, {'filters': {'岗位': ['保安', '电工']}}]:
            s = state('张三是保安吗')
            s['slots'] = slots
            self.assertIsNone(build_simple_plan(s)[0])

    def test_unknown_field_and_nonliteral_value_fall_back(self):
        for condition in [{'field': '不存在', 'op': 'eq', 'value': '员工'},
                          {'field': '岗位名称', 'op': 'eq', 'value': '保安'},
                          {'field': '岗位名称', 'op': 'ne', 'value': '员工'}]:
            s = state()
            s['simple_query']['filters'] = [condition]
            self.assertIsNone(build_simple_plan(s)[0])

    def test_malformed_candidates_fail_closed(self):
        for candidate in [None, [], {}, {'complete': True},
                          {**state()['simple_query'], 'fields': None},
                          {**state()['simple_query'], 'fields': [['姓名']]},
                          {**state()['simple_query'], 'limit': True},
                          {**state()['simple_query'], 'limit': 201},
                          {**state()['simple_query'], 'filters': [None]},
                          {**state()['simple_query'], 'extra': 'ignored?'},
                          {**state()['simple_query'], 'kind': 'sum'},
                          {**state()['simple_query'], 'complete': 'true'}]:
            s = state()
            s['simple_query'] = candidate
            self.assertIsNone(build_simple_plan(s)[0])

    def test_feedback_always_uses_full_planner(self):
        original, _ = build_simple_plan(state())
        for key in ['plan_feedback', 'run_feedback', 'slots_feedback', 'feedback']:
            s = state()
            s[key] = '改成只查保安'
            with patch.object(production, 'chat_json', return_value=original) as model:
                result = production.plan(s)
            model.assert_called_once()
            self.assertEqual(result['trace'][-1]['mode'], 'llm')
            self.assertEqual(result['trace'][-1]['reason'], 'feedback_requires_replanning')

    def test_disabled_and_legacy_fall_back(self):
        original, _ = build_simple_plan(state())
        with patch.dict(os.environ, {'QA_FAST_PLAN_ENABLED': 'false'}):
            with patch.object(production, 'chat_json', return_value=original) as model:
                self.assertEqual(production.plan(state())['trace'][-1]['reason'], 'disabled')
                model.assert_called_once()
        s = state()
        s['qa_context']['schema']['meta']['schema_version'] = 1
        self.assertIsNone(build_simple_plan(s)[0])

    def model_output(self):
        s = state()
        return {'category': '图谱检索', 'slots': s['slots'], 'data_needs': s['simple_query_needs'],
                'simple_query': s['simple_query']}

    def test_real_graph_skips_second_model_and_still_confirms(self):
        with patch.object(nodes, 'load_question_context', return_value=deepcopy(CONTEXT)), \
             patch.object(nodes, 'select_assessments', side_effect=lambda ctx, *a, **kw: ctx), \
             patch.object(nodes, 'chat_json', return_value=self.model_output()) as understanding, \
             patch.object(production, 'chat_json') as planner:
            graph = build_agent_graph(InMemorySaver())
            result = graph.invoke({'question': '共有多少名员工？'},
                                  {'configurable': {'thread_id': 'simple'}})
        understanding.assert_called_once()
        planner.assert_not_called()
        self.assertEqual(result['__interrupt__'][0].value['type'], 'confirm_plan')
        self.assertEqual(result['trace'][-1]['mode'], 'simple')

    def test_confirm_edit_replans_and_then_executes(self):
        original, _ = build_simple_plan(state())
        driver = SimpleNamespace(execute_query=lambda *a, **kw: SimpleNamespace(records=[{'记录数': 2}]))
        with patch.object(nodes, 'load_question_context', return_value=deepcopy(CONTEXT)), \
             patch.object(nodes, 'select_assessments', side_effect=lambda ctx, *a, **kw: ctx), \
             patch.object(nodes, 'chat_json', return_value=self.model_output()), \
             patch.object(production, 'chat_json', return_value=original) as planner, \
             patch.object(production, 'get_driver', return_value=driver), \
             patch.object(production, 'release_snapshot', return_value={'data_version': 'v1'}), \
             patch.object(production, 'chat_text', return_value='共2条记录'):
            graph = build_agent_graph(InMemorySaver())
            config = {'configurable': {'thread_id': 'edit'}}
            graph.invoke({'question': '共有多少名员工？'}, config)
            result = graph.invoke(Command(resume='请重新核对计数口径'), config)
            self.assertIn('__interrupt__', result)
            planner.assert_called_once()
            result = graph.invoke(Command(resume='确认'), config)
            self.assertEqual(result['answer'], '共2条记录')
            self.assertEqual(result['subgraph']['results']['s1']['rows'], [{'记录数': 2}])

    def test_understanding_resets_old_candidate_and_switch_restores_prompt(self):
        for enabled in ['true', 'false']:
            s = state()
            with patch.dict(os.environ, {'QA_FAST_PLAN_ENABLED': enabled}), \
                 patch.object(nodes, 'load_question_context', return_value=deepcopy(CONTEXT)), \
                 patch.object(nodes, 'select_assessments', side_effect=lambda ctx, *a, **kw: ctx), \
                 patch.object(nodes, 'chat_json', return_value={'category': '闲聊'}) as model:
                result = nodes.understand(s)
            self.assertIsNone(result['simple_query'])
            self.assertIsNone(result['simple_query_needs'])
            self.assertEqual('simple_query' in model.call_args.args[0], enabled == 'true')

    def test_full_planner_still_retries_invalid_plan(self):
        s = state()
        s['simple_query'] = None
        original, _ = build_simple_plan(state())
        with patch.object(production, 'chat_json', side_effect=[{'steps': []}, original]) as model:
            result = production.plan(s)
        self.assertEqual(model.call_count, 2)
        self.assertEqual(result['trace'][-1]['validation_attempts'], 2)

    def test_minimal_candidate_matches_explicit_defaults(self):
        full = state()
        minimal = state()
        minimal['simple_query'] = {'kind': 'count', 'dataset': '人员信息', 'complete': True}
        self.assertEqual(build_simple_plan(full), build_simple_plan(minimal))
        plan, _ = build_simple_plan(minimal)
        step = plan['steps'][0]
        self.assertNotIn('fields', step)
        self.assertNotIn('depends', step)
        self.assertNotIn('select', step['params'])
        self.assertNotIn('filters', step['params'])
        self.assertNotIn('limit', step['params'])

    def test_compact_full_plan_preserves_filters_and_does_not_add_empty_fields(self):
        s = state('按部门分别统计运营部员工')
        s['simple_query'] = None
        plan = {'steps': [{'step_id': 's1', 'tool': '图谱查询', 'params': {
            'nodes': [{'alias': 'n', 'type': '人员信息'}],
            'filters': [{'alias': 'n', 'field': '部门', 'op': 'eq', 'value': '运营部'}],
            'select': [{'alias': 'n', 'field': '部门', 'as': '部门'}],
            'aggregates': [{'alias': 'n', 'op': 'count', 'as': '人数'}]}}]}
        with patch.dict(os.environ, {'QA_COMPACT_INTERMEDIATE_ENABLED': 'true'}), \
             patch.object(production, 'chat_json', return_value=plan) as model:
            result = production.plan(s)
        self.assertEqual(result['plan'], plan)
        self.assertIn('省略所有空数组', model.call_args.args[0])
        self.assertNotIn('"history"', model.call_args.args[1])
        self.assertNotIn('"feedback"', model.call_args.args[1])
        compiled = compile_query(plan['steps'][0]['params'], SCHEMA, 'v1')
        self.assertEqual(compiled.params['v0'], '运营部')
        self.assertIn('n.`部门` AS `部门`', compiled.query)

    def test_both_switches_off_restore_original_prompt(self):
        s = state()
        original, _ = build_simple_plan(s)
        with patch.dict(os.environ, {'QA_FAST_PLAN_ENABLED': 'false',
                                     'QA_COMPACT_INTERMEDIATE_ENABLED': 'false'}), \
             patch.object(nodes, 'load_question_context', return_value=deepcopy(CONTEXT)), \
             patch.object(nodes, 'select_assessments', side_effect=lambda ctx, *a, **kw: ctx), \
             patch.object(nodes, 'chat_json', return_value=self.model_output()) as understand_model, \
             patch.object(production, 'chat_json', return_value=original) as plan_model:
            s.update(nodes.understand(s))
            production.plan(s)
        self.assertNotIn('simple_query', understand_model.call_args.args[0])
        self.assertNotIn('仅输出紧凑JSON', understand_model.call_args.args[0])
        self.assertIn('不需要的数组填[]', plan_model.call_args.args[0])

    def test_invalid_switch_is_not_silently_accepted(self):
        from step2_graph_building.config import get_qa_fast_plan_enabled
        with patch.dict(os.environ, {'QA_FAST_PLAN_ENABLED': 'maybe'}):
            with self.assertRaises(ValueError):
                get_qa_fast_plan_enabled()

    def test_graph_records_one_vs_two_real_wrapper_streams(self):
        import json
        original, _ = build_simple_plan(state())
        output = self.model_output()
        class Model:
            def stream(self, messages):
                body = output if '问题理解器' in messages[0]['content'] else original
                yield SimpleNamespace(content=json.dumps(body, ensure_ascii=False), usage_metadata=None)
        for enabled, expected in [('true', 1), ('false', 2)]:
            with patch.dict(os.environ, {'QA_FAST_PLAN_ENABLED': enabled,
                                         'QA_COMPACT_INTERMEDIATE_ENABLED': enabled}), \
                 patch.object(nodes, 'load_question_context', return_value=deepcopy(CONTEXT)), \
                 patch.object(nodes, 'select_assessments', side_effect=lambda ctx, *a, **kw: ctx), \
                 patch.object(llm, 'get_chat', return_value=Model()) as client:
                result = build_agent_graph(InMemorySaver()).invoke({'question': '共有多少名员工？'},
                    {'configurable': {'thread_id': enabled}})
            self.assertEqual(client.call_count, expected)
            self.assertEqual(sum(len(t['llm_calls']) for t in result['trace']), expected)
            self.assertEqual(result['__interrupt__'][0].value['type'], 'confirm_plan')

class TimingTests(unittest.TestCase):
    def test_json_retries_and_usage_are_measured(self):
        class Model:
            def __init__(self):
                self.count = 0
            def stream(self, messages):
                self.count += 1
                yield SimpleNamespace(content='', usage_metadata=None)
                yield SimpleNamespace(content='invalid' if self.count == 1 else '{"ok":true}',
                                      usage_metadata={'input_tokens': 4, 'output_tokens': 2})
        with measure_llm_calls() as calls, patch.object(llm, 'get_chat', return_value=Model()):
            self.assertEqual(llm.chat_json('system', 'user'), {'ok': True})
        self.assertEqual(len(calls), 2)
        self.assertEqual(calls[-1]['usage']['input_tokens'], 4)
        self.assertGreaterEqual(calls[-1]['duration_ms'], calls[-1]['first_text_ms'])
        self.assertNotIn('system', str(calls))
        self.assertIsNone(llm_calls.get())

    def test_nested_collectors_do_not_leak(self):
        with measure_llm_calls() as outer:
            with measure_llm_calls() as inner:
                self.assertIs(llm_calls.get(), inner)
            self.assertIs(llm_calls.get(), outer)
        self.assertIsNone(llm_calls.get())

    def test_stream_failure_is_measured_and_propagated(self):
        def fail(messages):
            raise RuntimeError('transport failed')
            yield
        with measure_llm_calls() as calls, patch.object(llm, 'get_chat',
                return_value=SimpleNamespace(stream=fail)):
            with self.assertRaises(RuntimeError):
                llm.chat_json('system', 'user')
        self.assertEqual(calls[0]['status'], 'error')
        self.assertIsNone(calls[0]['first_text_ms'])

if __name__ == '__main__':
    unittest.main(verbosity=2)