from __future__ import annotations import importlib.util import json import tempfile import unittest from pathlib import Path from src.agent_tools import build_tool_context, default_tool_registry, execute_tool from src.agent_types import AgentRuntimeConfig from src.data_agent_router_sessions import ( convert_router_candidates_to_records, profile_router_sessions, sample_router_candidates, search_router_sessions, ) HAS_PYARROW = importlib.util.find_spec('pyarrow') is not None class DataAgentRouterSessionTests(unittest.TestCase): def test_sample_router_candidates_accepts_search_payload(self) -> None: payload = { 'candidates': [ _candidate('s1', 'phone', 'QA', 'unknown', '怎么查找设备'), _candidate('s2', 'car', 'mapCopilot', '导航', '导航去公司'), _candidate('s3', 'phone', 'QA', 'unknown', '手机丢了怎么办'), ] } result = sample_router_candidates(payload, sample_size=2, strategy='stride') self.assertEqual(result['total_candidate_count'], 3) self.assertEqual(result['sampled_count'], 2) self.assertEqual(result['summary']['device'][0], {'value': 'phone', 'count': 2}) self.assertEqual(result['candidates'][0]['semantic_session_id'], 's1') self.assertEqual(result['candidates'][1]['semantic_session_id'], 's3') def test_sample_router_candidates_accepts_json_string(self) -> None: result = sample_router_candidates( json.dumps([_candidate('s1', 'phone', 'QA', 'unknown', '怎么查找设备')], ensure_ascii=False), sample_size=1, ) self.assertEqual(result['sampled_count'], 1) self.assertEqual(result['candidates'][0]['matched_turn']['query'], '怎么查找设备') def test_convert_router_candidates_to_records_keeps_online_metadata(self) -> None: candidate = _candidate('s1', 'phone', 'QA', 'unknown', '总结一下') candidate['turn_count'] = 2 candidate['matched_turn_index'] = 1 candidate['matched_turn']['timestamp'] = 20 candidate['prev_turns'] = [{'timestamp': 10, 'query': '这篇文章讲了什么', 'tts': ''}] result = convert_router_candidates_to_records( [candidate], dataset_label='总结类边界评测集', default_target='Summarize', batch_id='summary', ) self.assertEqual(result['record_count'], 1) record = result['records'][0] self.assertEqual(record['source']['type'], 'online') self.assertEqual(record['source']['request_id'], 's1-rid') self.assertEqual(record['turn']['query'], '总结一下') self.assertEqual(record['prev_session'], [{'query': '这篇文章讲了什么', 'tts': '', 'timestamp': 10}]) self.assertEqual(record['context']['domain'], 'QA') self.assertEqual(record['label']['target'], 'Summarize') def test_convert_router_candidates_to_records_applies_review_decisions(self) -> None: candidates = [ _candidate('s1', 'phone', 'QA', 'unknown', '总结一下'), _candidate('s2', 'phone', 'QA', 'unknown', '总结一下红楼梦'), ] result = convert_router_candidates_to_records( candidates, dataset_label='总结类边界评测集', review_decisions=[ { 'semantic_session_id': 's1', 'matched_turn_index': 0, 'decision': 'include', 'target': 'Summarize', 'notes': '前文有可总结内容', }, { 'semantic_session_id': 's2', 'matched_turn_index': 0, 'decision': 'exclude', 'notes': '开放知识对象', }, ], ) self.assertEqual(result['record_count'], 1) self.assertEqual(result['skipped_count'], 1) self.assertEqual(result['records'][0]['meta']['notes'], '前文有可总结内容') @unittest.skipUnless(HAS_PYARROW, 'pyarrow is required for parquet tests') def test_profile_and_search_router_sessions_read_parquet(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: _write_router_parquet(Path(tmp_dir)) profile = profile_router_sessions(tmp_dir, dates=['20260428'], max_files=1) search = search_router_sessions( tmp_dir, dates=['20260428'], domains=['mapCopilot'], query_keywords=['导航'], max_files=1, max_candidates=5, ) self.assertEqual(profile['sampled_row_count'], 2) self.assertEqual(profile['missing_req_id_count'], 0) self.assertTrue(any(item['value'] == 'phone' for item in profile['distributions']['device'])) self.assertEqual(search['candidate_count'], 1) candidate = search['candidates'][0] self.assertEqual(candidate['req_id'], 'rid-1') self.assertEqual(candidate['matched_turn']['query'], '导航去公司') self.assertEqual(candidate['prev_turns'][0]['query'], '你好小爱') @unittest.skipUnless(HAS_PYARROW, 'pyarrow is required for parquet tests') def test_router_session_tools_allow_external_absolute_paths(self) -> None: with tempfile.TemporaryDirectory() as workspace_dir, tempfile.TemporaryDirectory() as data_dir: _write_router_parquet(Path(data_dir)) external_partition = Path(data_dir) / 'router_session_parquet' / 'date=20260428' profile = profile_router_sessions( workspace_dir, paths=[str(external_partition)], max_files=1, ) search = search_router_sessions( workspace_dir, paths=[str(external_partition)], query_keywords=['导航'], max_files=1, max_candidates=5, ) self.assertEqual(profile['sampled_row_count'], 2) self.assertEqual(search['candidate_count'], 1) self.assertEqual(search['candidates'][0]['matched_turn']['query'], '导航去公司') @unittest.skipUnless(HAS_PYARROW, 'pyarrow is required for parquet tests') def test_router_session_tools_execute_against_registry(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: _write_router_parquet(Path(tmp_dir)) context = build_tool_context(AgentRuntimeConfig(cwd=Path(tmp_dir))) registry = default_tool_registry() search_result = execute_tool( registry, 'data_agent_search_router_sessions', {'dates': ['20260428'], 'query_regex': '查找设备|手机丢了', 'max_files': 1}, context, ) self.assertTrue(search_result.ok, search_result.content) sample_result = execute_tool( registry, 'data_agent_sample_router_candidates', {'candidates': json.loads(search_result.content), 'sample_size': 1, 'strategy': 'first'}, context, ) convert_result = execute_tool( registry, 'data_agent_convert_router_candidates_to_records', { 'candidates': json.loads(sample_result.content), 'dataset_label': '查找设备评测集', 'default_target': 'QA', 'batch_id': 'find_device', }, context, ) self.assertTrue(sample_result.ok, sample_result.content) self.assertEqual(json.loads(sample_result.content)['sampled_count'], 1) self.assertTrue(convert_result.ok, convert_result.content) self.assertTrue(json.loads(convert_result.content)['validation']['ok']) def _candidate(session_id: str, device: str, domain: str, intent: str, query: str) -> dict[str, object]: return { 'semantic_session_id': session_id, 'req_id': f'{session_id}-rid', 'device': device, 'turn_count': 1, 'matched_turn_index': 0, 'matched_turn': { 'timestamp': 1, 'query': query, 'tts': '', 'action': {'domain': domain, 'intent': intent}, }, 'prev_turns': [], 'source_path': 'router_session_parquet/date=20260428/part-0.parquet', 'date': '20260428', } def _write_router_parquet(root: Path) -> None: import pyarrow as pa import pyarrow.parquet as pq partition = root / 'router_session_parquet' / 'date=20260428' partition.mkdir(parents=True) rows = [ { 'semantic_session_id': 'phone_20260428_1', 'device': 'phone', 'req_id': 'rid-1', 'turn_count': 2, 'turns_array': [ { 'timestamp': 1, 'query': '你好小爱', 'action_json': '{"domain":"dialogCopilot","intent":"问候"}', 'tts': '你好', }, { 'timestamp': 2, 'query': '导航去公司', 'action_json': '{"domain":"mapCopilot","intent":"导航","func":"start_navigation"}', 'tts': '开始导航', }, ], }, { 'semantic_session_id': 'phone_20260428_2', 'device': 'phone', 'req_id': 'rid-2', 'turn_count': 1, 'turns_array': [ { 'timestamp': 3, 'query': '手机丢了怎么查找设备', 'action_json': '{"domain":"QA","intent":"unknown"}', 'tts': '', } ], }, ] pq.write_table(pa.Table.from_pylist(rows), partition / 'part-0.parquet') if __name__ == '__main__': unittest.main()