kg-agent/scripts/test_value_clarification.py

"""Offline regression for empty-result suggestions and explicit selection."""
from copy import deepcopy
from types import SimpleNamespace
import unittest
from unittest.mock import patch

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import StateGraph, START, END
from langgraph.types import Command
from step3_qa_agent.agent import production as p
from step3_qa_agent.agent.state import AgentState
from step3_qa_agent.agent.value_clarification import empty_result, find_candidates, selected_option, selected_options

SCHEMA = {'nodes': [{'id': '人员证书', 'attributes': ['证书名称', '级别']}], 'relations': []}
SPEC = {'nodes': [{'alias': 'n', 'type': '人员证书'}], 'relations': [],
        'filters': [{'alias': 'n', 'field': '证书名称', 'op': 'eq', 'value': '电工'},
                    {'alias': 'n', 'field': '级别', 'op': 'eq', 'value': '中级'}],
        'select': [], 'aggregates': [{'alias': 'n', 'op': 'count', 'as': '人数'}]}
CONTEXT = {'schema': SCHEMA, 'build_id': 'v1'}
PLAN = {'steps': [{'step_id': 's1', 'tool': '图谱查询', 'params': SPEC}]}

class Driver:
    def __init__(self):
        self.calls = []
    def execute_query(self, query, **params):
        self.calls.append((query, params))
        if 'AS `候选值`' in query:
            rows = [{'候选值': ['水电工', '维修电工', '保安'], '记录数': 2}] if 'n.`证书名称` AS' in query else [{'候选值': '中级', '记录数': 1}]
        else:
            value = params.get('v0')
            matched = bool(set(value if isinstance(value, list) else [value]) & {'水电工', '维修电工'})
            rows = [{'人数': 2 if matched else 0}]
        return SimpleNamespace(records=rows)

class ClarificationTests(unittest.TestCase):
    def state(self):
        return {'question': '有多少人是中级电工', 'plan': deepcopy(PLAN),
                'qa_context': deepcopy(CONTEXT)}

    def test_zero_count_not_merely_empty_rows(self):
        self.assertTrue(empty_result(SPEC, [{'人数': 0}]))
        self.assertTrue(empty_result(SPEC, []))
        self.assertFalse(empty_result(SPEC, [{'人数': 2}]))
        self.assertTrue(empty_result({'aggregates': [{'op': 'count_distinct', 'as': '人数'}]}, [{'人数': 0}]))
        self.assertFalse(empty_result({'aggregates': [{'op': 'sum', 'as': '金额'}]}, [{'金额': 0}]))

    def test_real_multivalue_candidates_preserve_other_filters(self):
        driver = Driver()
        result = find_candidates(PLAN, {'s1': {'rows': [{'人数': 0}]}}, CONTEXT, driver)
        self.assertEqual({c['value'] for c in result['options']}, {'水电工', '维修电工'})
        query, args = driver.calls[0]
        self.assertIn('中级', args.values())
        self.assertNotIn('电工', args.values())
        self.assertEqual(args['build'], 'v1')
        self.assertEqual(SPEC['filters'][0]['value'], '电工')

    def test_positive_results_do_not_probe(self):
        driver = Driver()
        self.assertFalse(find_candidates(PLAN, {'s1': {'rows': [{'人数': 2}]}}, CONTEXT, driver)['options'])
        self.assertFalse(driver.calls)

    def test_confirmation_is_not_selection(self):
        options = [{'label': 'name', 'value': '水电工'}]
        self.assertIsNone(selected_option('确认', options))
        for reply in ('1', '选择1', '水电工', 'name'):
            self.assertEqual(selected_option(reply, options), options[0])
        self.assertIsNone(selected_option('2', options))

    def test_multiple_numbers_values_and_json_are_supported(self):
        options = [{'label': '1. 水电工', 'value': '水电工'},
                   {'label': '2. 维修电工', 'value': '维修电工'}]
        for reply in ('1,2', '选择1，选2', '水电工、维修电工',
                      '["1. 水电工", "2. 维修电工"]'):
            self.assertEqual([item['value'] for item in selected_options(reply, options)],
                             ['水电工', '维修电工'])
        self.assertEqual([item['value'] for item in selected_options('1,1', options)], ['水电工'])
        self.assertEqual(selected_options('确认', options), [])

    def graph(self):
        graph = StateGraph(AgentState)
        graph.add_node('run', p.execute)
        graph.add_node('clarify', p.clarify)
        graph.add_edge(START, 'run')
        graph.add_edge('run', 'clarify')
        graph.add_conditional_edges('clarify', lambda s: 'run' if s['plan_confirm'] else END)
        return graph.compile(checkpointer=InMemorySaver())

    def test_interrupt_resume_requeries_without_changing_level(self):
        driver = Driver()
        with patch.object(p, 'get_driver', return_value=driver), patch.object(p, 'release_snapshot', return_value={'data_version': 'v1'}):
            graph = self.graph(); config = {'configurable': {'thread_id': 'choice'}}
            result = graph.invoke(self.state(), config)
            self.assertEqual(result['__interrupt__'][0].value['type'], 'clarify_value')
            self.assertEqual(result['subgraph']['results']['s1']['rows'], [{'人数': 0}])
            result = graph.invoke(Command(resume='确认'), config)
            self.assertIn('__interrupt__', result)
            result = graph.invoke(Command(resume='1'), config)
            self.assertNotIn('__interrupt__', result)
            self.assertEqual(result['subgraph']['results']['s1']['rows'], [{'人数': 2}])
            self.assertEqual(result['plan']['steps'][0]['params']['filters'][1]['value'], '中级')

    def test_multiple_choices_become_in_filter(self):
        driver = Driver()
        with patch.object(p, 'get_driver', return_value=driver), patch.object(p, 'release_snapshot', return_value={'data_version': 'v1'}):
            graph = self.graph(); config = {'configurable': {'thread_id': 'multi-choice'}}
            first = graph.invoke(self.state(), config)
            self.assertEqual(len(first['subgraph']['suggestions']['options']), 2)
            result = graph.invoke(Command(resume='1,2'), config)
            condition = result['plan']['steps'][0]['params']['filters'][0]
            self.assertEqual(condition['op'], 'in')
            self.assertEqual(set(condition['value']), {'水电工', '维修电工'})
            self.assertEqual(result['subgraph']['results']['s1']['rows'], [{'人数': 2}])

    def test_decline_keeps_original_query(self):
        with patch.object(p, 'get_driver', return_value=Driver()), patch.object(p, 'release_snapshot', return_value={'data_version': 'v1'}):
            graph = self.graph(); config = {'configurable': {'thread_id': 'decline'}}
            graph.invoke(self.state(), config)
            result = graph.invoke(Command(resume='都不是'), config)
            self.assertEqual(result['plan'], PLAN)
            self.assertFalse(result['plan_confirm'])

    def test_version_change_on_resume_rejects_stale_choice(self):
        with patch.object(p, 'get_driver', return_value=Driver()), patch.object(p, 'release_snapshot', return_value={'data_version': 'v1'}) as release:
            graph = self.graph(); config = {'configurable': {'thread_id': 'version'}}
            graph.invoke(self.state(), config)
            release.return_value = {'data_version': 'v2'}
            result = graph.invoke(Command(resume='1'), config)
            self.assertIn('数据已更新', result['subgraph']['error'])

    def test_probe_failure_is_not_zero_data_evidence(self):
        with patch.object(p, 'get_driver', return_value=Driver()), patch.object(p, 'release_snapshot', return_value={'data_version': 'v1'}), patch.object(p, 'find_candidates', side_effect=RuntimeError('secret')):
            result = p.execute(self.state())
            self.assertIn('暂不可用', result['subgraph']['suggestions']['error'])
            self.assertNotIn('secret', str(result))

    def test_no_candidates_and_bounded_results(self):
        driver = SimpleNamespace(execute_query=lambda *a, **k: SimpleNamespace(records=[{'候选值': 'xyz', '记录数': 1}] * 201))
        result = find_candidates(PLAN, {'s1': {'rows': []}}, CONTEXT, driver)
        self.assertFalse(result['options'])
        self.assertTrue(result['limited'])


class ApiClarificationTests(unittest.IsolatedAsyncioTestCase):
    async def test_auto_confirm_stops_for_value_choice(self):
        from step4_web import api
        class FakeGraph:
            calls = 0
            async def astream(self, *args, **kwargs):
                self.calls += 1
                if self.calls > 1:
                    raise AssertionError('自动确认不能替用户选择')
                yield {'__interrupt__': [SimpleNamespace(value={
                    'type': 'clarify_value', 'message': '请选择实际值',
                    'options': ['水电工', '都不是']})]}
        fake = FakeGraph()
        with patch.object(api, '_get_graph', return_value=fake):
            events = [event async for event in api._astream_run('test', '中级电工', True)]
        self.assertEqual(fake.calls, 1)
        self.assertEqual(events[-1]['type'], 'confirm')
        self.assertEqual(events[-1]['options'], ['水电工', '都不是'])
        self.assertEqual(events[-1]['confirm_type'], 'clarify_value')

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