|
|
@@ -0,0 +1,307 @@
|
|
|
+"""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)
|