from __future__ import annotations 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_records import ( confirm_generation_goal, confirm_generation_plan, export_dataset_records, get_generation_plan, normalize_dataset_draft, prepare_generation_goal, prepare_generation_plan, update_generation_plan, validate_dataset_records, ) class DataAgentRecordTests(unittest.TestCase): def test_normalize_dataset_draft_builds_canonical_records(self) -> None: draft = ''' # dataset_label: 地图和生活边界数据 ### case: 附近餐饮查询 用户: 帮我看看附近有什么好吃的 target: Agent(tag="life_service") notes: 附近生活服务查询 ### case: 多轮承接附近餐饮 用户: 我想出门逛逛 小爱: 好的 用户: 看看附近有什么好吃的 target: Agent(tag="life_service") notes: 多轮承接附近生活服务查询 '''.strip() payload = normalize_dataset_draft( draft, batch_id='demo', base_timestamp=1_755_567_930_500, timestamp_step_ms=60_000, ) self.assertEqual(payload['warnings'], []) records = payload['records'] self.assertEqual(len(records), 2) self.assertEqual(records[0]['record_id'], 'gen_demo_000001') self.assertEqual(records[0]['turn']['query'], '帮我看看附近有什么好吃的') self.assertEqual(records[0]['prev_session'], []) self.assertEqual(records[0]['label']['target_type'], 'agent') self.assertEqual(records[1]['turn']['query'], '看看附近有什么好吃的') self.assertEqual( records[1]['prev_session'], [ { 'query': '我想出门逛逛', 'tts': '好的', 'timestamp': 1_755_569_070_500, } ], ) def test_validate_dataset_records_reports_valid_payload(self) -> None: records = normalize_dataset_draft( ''' # dataset_label: 地图和生活边界数据 ### case: 附近餐饮查询 用户: 附近有什么好吃的 target: Agent(tag="life_service") '''.strip(), base_timestamp=1_755_567_930_500, )['records'] result = validate_dataset_records(records) self.assertTrue(result['ok']) self.assertEqual(result['error_count'], 0) def test_validate_dataset_records_rejects_current_turn_tts(self) -> None: records = normalize_dataset_draft( ''' # dataset_label: 地图和生活边界数据 ### case: 附近餐饮查询 用户: 附近有什么好吃的 target: Agent(tag="life_service") '''.strip(), base_timestamp=1_755_567_930_500, )['records'] records[0]['turn']['tts'] = '不应该出现' result = validate_dataset_records(records) self.assertFalse(result['ok']) self.assertEqual(result['errors'][0]['path'], 'records[0].turn.tts') def test_export_dataset_records_writes_compact_jsonl(self) -> None: records = normalize_dataset_draft( ''' # dataset_label: 地图和生活边界数据 ### case: 附近餐饮查询 用户: 附近有什么好吃的 target: Agent(tag="life_service") '''.strip(), batch_id='demo', base_timestamp=1_755_567_930_500, )['records'] with tempfile.TemporaryDirectory() as tmp_dir: payload = export_dataset_records( records, root=tmp_dir, output_path='tasks/demo/records.jsonl', ) output_path = Path(tmp_dir) / payload['output_path'] lines = output_path.read_text(encoding='utf-8').splitlines() self.assertEqual(payload['output_format'], 'jsonl') self.assertEqual(payload['record_count'], 1) self.assertEqual(len(lines), 1) self.assertNotIn(': ', lines[0]) self.assertEqual(json.loads(lines[0])['turn']['query'], '附近有什么好吃的') def test_normalize_online_draft_uses_real_request_metadata(self) -> None: records = normalize_dataset_draft( ''' # dataset_label: 线上误召回badcase专项 ### case: 线上样例 request_id: rid-123 timestamp: 1755567930500 用户: 附近有什么好吃的 target: Agent(tag="life_service") '''.strip(), source_type='online', )['records'] self.assertEqual(records[0]['record_id'], 'online_rid-123_000001') self.assertEqual(records[0]['source']['request_id'], 'rid-123') self.assertEqual(records[0]['source']['timestamp'], 1_755_567_930_500) def test_normalize_requires_confirmed_plan_when_state_root_is_provided(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: with self.assertRaisesRegex(Exception, 'confirmed_plan_id is required'): normalize_dataset_draft( ''' # dataset_label: 地图和生活边界数据 ### case: 附近餐饮查询 用户: 附近有什么好吃的 target: Agent(tag="life_service") '''.strip(), plan_state_root=tmp_dir, ) def test_generation_plan_can_be_confirmed_then_used(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: plan_payload = prepare_generation_plan( root=tmp_dir, dataset_label='地图和生活边界数据', target='Agent(tag="life_service")', total_count=2, turn_mix='1 条单轮，1 条多轮', coverage='附近吃喝玩乐', exclusions='不要生成导航路线类 query', output_path='tasks/demo/artifacts', ) plan_id = plan_payload['plan']['plan_id'] confirmed = confirm_generation_plan(root=tmp_dir, plan_id=plan_id, confirmation='确认，开始生成') records = normalize_dataset_draft( ''' # dataset_label: 地图和生活边界数据 ### case: 附近餐饮查询 用户: 附近有什么好吃的 target: Agent(tag="life_service") '''.strip(), confirmed_plan_id=confirmed['confirmed_plan_id'], plan_state_root=tmp_dir, )['records'] self.assertEqual(records[0]['label']['dataset_label'], '地图和生活边界数据') def test_generation_goal_can_gate_plan_creation(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: goal_payload = prepare_generation_goal( root=tmp_dir, dataset_label='地图和生活边界数据', goal_summary='生成附近生活服务边界数据', target='Agent(tag="life_service")', plan_hint='建议先生成 2 条单轮，输出到 tasks/demo/records.jsonl', coverage='附近吃喝玩乐', exclusions='不要生成导航路线类 query', source_refs=['manual:user'], ) goal_id = goal_payload['goal']['goal_id'] with self.assertRaisesRegex(Exception, 'generation goal is not confirmed'): prepare_generation_plan( root=tmp_dir, confirmed_goal_id=goal_id, dataset_label='地图和生活边界数据', target='Agent(tag="life_service")', total_count=2, turn_mix='2 条单轮', coverage='附近吃喝玩乐', exclusions='不要生成导航路线类 query', output_path='tasks/demo/artifacts', ) confirmed_goal = confirm_generation_goal( root=tmp_dir, goal_id=goal_id, confirmation='确认 goal', reviewed_revision=1, ) plan_payload = prepare_generation_plan( root=tmp_dir, confirmed_goal_id=confirmed_goal['confirmed_goal_id'], dataset_label='地图和生活边界数据', target='Agent(tag="life_service")', total_count=2, turn_mix='2 条单轮', coverage='附近吃喝玩乐', exclusions='不要生成导航路线类 query', output_path='tasks/demo/artifacts', ) self.assertEqual(plan_payload['plan']['confirmed_goal_id'], goal_id) self.assertEqual( goal_payload['goal']['plan_hint'], '建议先生成 2 条单轮，输出到 tasks/demo/records.jsonl', ) def test_generation_plan_dedupes_existing_task_output_path(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: existing = Path(tmp_dir) / 'tasks' / 'demo' / 'artifacts' existing.mkdir(parents=True) plan_payload = prepare_generation_plan( root=tmp_dir, dataset_label='地图和生活边界数据', target='Agent(tag="life_service")', total_count=2, turn_mix='1 条单轮，1 条多轮', coverage='附近吃喝玩乐', exclusions='不要生成导航路线类 query', output_path='tasks/demo/artifacts', ) self.assertEqual( plan_payload['plan']['output_path'], 'tasks/demo-data_plan_000001/artifacts', ) self.assertEqual(plan_payload['plan']['requested_output_path'], 'tasks/demo/artifacts') def test_generation_plan_review_update_and_revision_confirmation(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: plan_payload = prepare_generation_plan( root=tmp_dir, dataset_label='地图和生活边界数据', target='Agent(tag="life_service")', total_count=2, turn_mix='1 条单轮，1 条多轮', coverage='附近吃喝玩乐', exclusions='不要生成导航路线类 query', output_path='tasks/demo/artifacts', ) plan_id = plan_payload['plan']['plan_id'] updated = update_generation_plan( root=tmp_dir, plan_id=plan_id, review_feedback='多轮数据多一点', updates={'turn_mix': '1 条单轮，3 条多轮', 'total_count': 4}, ) self.assertEqual(updated['plan']['revision'], 2) self.assertEqual(updated['plan']['turn_mix'], '1 条单轮，3 条多轮') self.assertEqual(len(updated['plan']['review_history']), 1) shown = get_generation_plan(root=tmp_dir, plan_id=plan_id) self.assertEqual(shown['plan']['revision'], 2) with self.assertRaisesRegex(Exception, 'reviewed_revision must match'): confirm_generation_plan( root=tmp_dir, plan_id=plan_id, confirmation='确认，开始生成', reviewed_revision=1, ) confirmed = confirm_generation_plan( root=tmp_dir, plan_id=plan_id, confirmation='确认，开始生成', reviewed_revision=2, ) self.assertEqual(confirmed['plan']['status'], 'confirmed') self.assertEqual(confirmed['plan']['confirmed_revision'], 2) def test_multi_target_plan_allows_only_declared_targets(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: plan_payload = prepare_generation_plan( root=tmp_dir, dataset_label='餐饮和导航边界数据', target='', target_definitions=[ { 'name': '餐饮服务', 'target': 'Agent(tag="餐饮服务")', 'rule': '找附近美食但不导航', }, { 'name': '地图导航', 'target': 'Agent(tag="地图导航")', 'rule': '明确要求导航去某地', }, ], total_count=2, turn_mix='2 条单轮', coverage='餐饮服务和地图导航边界', exclusions='不要生成无关闲聊', output_path='tasks/demo/artifacts', ) plan_id = plan_payload['plan']['plan_id'] confirm_generation_plan(root=tmp_dir, plan_id=plan_id, confirmation='确认，开始生成') records = normalize_dataset_draft( ''' # dataset_label: 餐饮和导航边界数据 ### case: 找奶茶 用户: 附近有没有奶茶店 target: Agent(tag="餐饮服务") ### case: 导航去奶茶店 用户: 导航去最近的奶茶店 target: Agent(tag="地图导航") '''.strip(), confirmed_plan_id=plan_id, plan_state_root=tmp_dir, )['records'] self.assertEqual(records[0]['label']['target'], 'Agent(tag="餐饮服务")') self.assertEqual(records[1]['label']['target'], 'Agent(tag="地图导航")') def test_multi_target_plan_rejects_undeclared_targets(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: plan_payload = prepare_generation_plan( root=tmp_dir, dataset_label='餐饮和导航边界数据', target='', target_definitions=[ {'target': 'Agent(tag="餐饮服务")'}, {'target': 'Agent(tag="地图导航")'}, ], total_count=1, turn_mix='1 条单轮', coverage='餐饮服务和地图导航边界', exclusions='不要生成无关闲聊', output_path='tasks/demo/artifacts', ) plan_id = plan_payload['plan']['plan_id'] confirm_generation_plan(root=tmp_dir, plan_id=plan_id, confirmation='确认，开始生成') with self.assertRaisesRegex(Exception, 'record targets must match'): normalize_dataset_draft( ''' # dataset_label: 餐饮和导航边界数据 ### case: 天气 用户: 明天天气怎么样 target: Agent(tag="天气") '''.strip(), confirmed_plan_id=plan_id, plan_state_root=tmp_dir, ) def test_tools_execute_against_registry(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: context = build_tool_context(AgentRuntimeConfig(cwd=Path(tmp_dir))) goal_result = execute_tool( default_tool_registry(), 'data_agent_prepare_generation_goal', { 'dataset_label': '地图和生活边界数据', 'goal_summary': '生成附近生活服务边界数据', 'target': 'Agent(tag="life_service")', 'coverage': '附近吃喝玩乐', 'exclusions': '不要生成导航路线类 query', 'source_refs': ['manual:user'], }, context, ) self.assertTrue(goal_result.ok, goal_result.content) goal_id = json.loads(goal_result.content)['goal']['goal_id'] blocked_plan_result = execute_tool( default_tool_registry(), 'data_agent_prepare_generation_plan', { 'dataset_label': '地图和生活边界数据', 'target': 'Agent(tag="life_service")', 'total_count': 1, 'turn_mix': '1 条单轮', 'coverage': '附近吃喝玩乐', 'exclusions': '不要生成导航路线类 query', 'output_path': 'tasks/demo/artifacts', }, context, ) self.assertFalse(blocked_plan_result.ok) confirmed_goal_result = execute_tool( default_tool_registry(), 'data_agent_confirm_generation_goal', {'goal_id': goal_id, 'confirmation': '确认 goal', 'reviewed_revision': 1}, context, ) self.assertTrue(confirmed_goal_result.ok, confirmed_goal_result.content) plan_result = execute_tool( default_tool_registry(), 'data_agent_prepare_generation_plan', { 'confirmed_goal_id': goal_id, 'dataset_label': '地图和生活边界数据', 'target': 'Agent(tag="life_service")', 'total_count': 1, 'turn_mix': '1 条单轮', 'coverage': '附近吃喝玩乐', 'exclusions': '不要生成导航路线类 query', 'output_path': 'tasks/demo/artifacts', }, context, ) self.assertTrue(plan_result.ok) plan_id = json.loads(plan_result.content)['plan']['plan_id'] update_result = execute_tool( default_tool_registry(), 'data_agent_update_generation_plan', { 'plan_id': plan_id, 'review_feedback': '多轮不需要，先只测单轮', 'updates': {'turn_mix': '1 条单轮'}, }, context, ) self.assertTrue(update_result.ok) revision = json.loads(update_result.content)['plan']['revision'] confirm_result = execute_tool( default_tool_registry(), 'data_agent_confirm_generation_plan', {'plan_id': plan_id, 'confirmation': '确认，开始生成', 'reviewed_revision': revision}, context, ) self.assertTrue(confirm_result.ok) normalize_result = execute_tool( default_tool_registry(), 'data_agent_normalize_dataset_draft', { 'draft_text': ''' # dataset_label: 地图和生活边界数据 ### case: 附近餐饮查询 用户: 附近有什么好吃的 target: Agent(tag="life_service") '''.strip(), 'batch_id': 'demo', 'base_timestamp': 1_755_567_930_500, 'confirmed_plan_id': plan_id, }, context, ) self.assertTrue(normalize_result.ok) records = json.loads(normalize_result.content)['records'] validate_result = execute_tool( default_tool_registry(), 'data_agent_validate_dataset_records', {'records': records}, context, ) export_result = execute_tool( default_tool_registry(), 'data_agent_export_dataset_records', { 'records': records, 'output_path': 'tasks/demo/records.jsonl', }, context, ) export_payload = json.loads(export_result.content) if export_result.ok else {} exported_lines = ( (Path(tmp_dir) / export_payload['output_path']).read_text(encoding='utf-8').splitlines() if export_result.ok else [] ) self.assertTrue(validate_result.ok) self.assertTrue(json.loads(validate_result.content)['ok']) self.assertTrue(export_result.ok, export_result.content) self.assertEqual(export_payload['record_count'], 1) self.assertEqual(len(exported_lines), 1) if __name__ == '__main__': unittest.main()