from __future__ import annotations import csv import io 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 ( build_planning_prompt, confirm_generation_goal, confirm_generation_plan, export_dataset_records, export_planning_eval_csv, export_training_jsonl, get_generation_plan, normalize_dataset_draft, prepare_generation_goal, prepare_generation_plan, render_dataset_records_table_csv, render_planning_eval_csv, 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[0]['dimensions'], {'complex': False}) self.assertEqual(records[1]['turn']['query'], '看看附近有什么好吃的') self.assertEqual( records[1]['prev_session'], [ { 'query': '我想出门逛逛', 'tts': '好的', 'timestamp': 1_755_569_070_500, } ], ) def test_normalize_dataset_draft_canonicalizes_single_quote_agent_target(self) -> None: payload = normalize_dataset_draft( ''' # dataset_label: 地图导航边界数据 ### case: 顺路停车场 用户: 帮我找个顺路的停车场 target: Agent(tag='地图导航') '''.strip(), batch_id='demo', base_timestamp=1_755_567_930_500, ) record = payload['records'][0] self.assertEqual(record['label']['target'], 'Agent(tag="地图导航")') self.assertEqual(record['label']['target_type'], 'agent') def test_normalize_dataset_draft_keeps_complex_dimension_independent(self) -> None: payload = normalize_dataset_draft( ''' # dataset_label: 地图导航复杂样例 ### case: 复杂导航 用户: 帮我规划一条先去加油站再去公司的路线 complex: true target: Agent(tag='地图导航') '''.strip(), batch_id='demo', base_timestamp=1_755_567_930_500, ) record = payload['records'][0] self.assertEqual(record['label']['target'], 'Agent(tag="地图导航")') self.assertEqual(record['dimensions']['complex'], True) 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() table_path = Path(tmp_dir) / payload['table_output_path'] table_rows = list(csv.DictReader(io.StringIO(table_path.read_text(encoding='utf-8-sig')))) self.assertEqual(payload['output_format'], 'jsonl') self.assertEqual(payload['table_output_format'], 'csv') self.assertEqual(payload['record_count'], 1) self.assertEqual(len(lines), 1) self.assertNotIn(': ', lines[0]) self.assertEqual(json.loads(lines[0])['turn']['query'], '附近有什么好吃的') self.assertEqual(len(table_rows), 1) self.assertEqual( list(table_rows[0].keys()), ['request_id', 'timestamp', 'query', 'prev_session', 'context', 'label', '是否迁移Function', 'function'], ) self.assertEqual(table_rows[0]['request_id'], 'aabbccdd') self.assertEqual(table_rows[0]['query'], '附近有什么好吃的') self.assertEqual(table_rows[0]['prev_session'], '[]') self.assertEqual(table_rows[0]['context'], '{}') self.assertEqual(table_rows[0]['label'], '地图和生活边界数据') self.assertEqual(table_rows[0]['是否迁移Function'], '') self.assertEqual(table_rows[0]['function'], 'complex=false\nAgent(tag="life_service")') def test_render_dataset_records_table_csv_keeps_prev_session_json(self) -> None: records = normalize_dataset_draft( ''' # dataset_label: 地图导航边界数据 ### case: 多轮顺路停车场 用户: 查一下附近停车场 小爱: 找到了附近停车场 用户: 帮我找个最顺路的 target: Agent(tag='地图导航') '''.strip(), batch_id='demo', base_timestamp=1_755_567_930_500, timestamp_step_ms=60_000, )['records'] rows = list(csv.DictReader(io.StringIO(render_dataset_records_table_csv(records)))) prev_session = json.loads(rows[0]['prev_session']) self.assertEqual(prev_session[0]['query'], '查一下附近停车场') self.assertEqual(prev_session[0]['tts'], '找到了附近停车场') self.assertEqual(prev_session[0]['timestamp'], '1755567870500') self.assertEqual(rows[0]['function'], 'complex=false\nAgent(tag="地图导航")') def test_export_training_jsonl_uses_prompt_template_and_history(self) -> None: records = normalize_dataset_draft( ''' # dataset_label: 地图导航边界数据 ### case: 多轮顺路停车场 用户: 查一下附近停车场 小爱: 找到了附近停车场 用户: 帮我找个顺路的 target: Agent(tag='地图导航') '''.strip(), batch_id='demo', base_timestamp=1_755_567_930_500, timestamp_step_ms=60_000, )['records'] records[0]['context'] = {'location': '北京', 'rag': '地图服务可用', 'other': '忽略'} with tempfile.TemporaryDirectory() as tmp_dir: payload = export_training_jsonl( records, root=tmp_dir, output_path='output/training.jsonl', ) output_path = Path(tmp_dir) / payload['output_path'] line = json.loads(output_path.read_text(encoding='utf-8').splitlines()[0]) self.assertEqual(payload['output_format'], 'jsonl') self.assertEqual(line['system'], '你是小爱同学，中文智能语音助手。') self.assertEqual(line['output'], 'complex=false\nAgent(tag="地图导航")') self.assertIn('[知识注入]\n{\n"location": "北京",\n"rag": "地图服务可用"\n}', line['instruction']) self.assertIn('用户: 查一下附近停车场\n小爱: 找到了附近停车场', line['instruction']) self.assertIn('[当前query]\n用户: 帮我找个顺路的', line['instruction']) def test_render_planning_eval_csv_uses_expected_columns(self) -> None: records = normalize_dataset_draft( ''' # dataset_label: 地图导航边界数据 ### case: 多轮顺路停车场 用户: 查一下附近停车场 小爱: 找到了附近停车场 用户: 帮我找个顺路的 target: Agent(tag='地图导航') '''.strip(), batch_id='demo', base_timestamp=1_755_567_930_500, )['records'] rows = list(csv.DictReader(io.StringIO(render_planning_eval_csv(records)))) self.assertEqual( list(rows[0].keys()), ['request_id', 'newPrompt', 'query', '类别真实标签', 'code标签', 'complex'], ) self.assertEqual(rows[0]['request_id'], 'aabbccdd') self.assertEqual(rows[0]['query'], '帮我找个顺路的') self.assertEqual(rows[0]['类别真实标签'], '地图导航') self.assertEqual(rows[0]['code标签'], 'Agent(tag="地图导航")') self.assertEqual(rows[0]['complex'], 'FALSE') self.assertTrue(rows[0]['newPrompt'].startswith('<|im_start|>system\n你是小爱同学，中文智能语音助手。<|im_end|>')) self.assertIn('<|im_start|>user\n请参考用户的[当前query]', rows[0]['newPrompt']) self.assertIn('[function]', rows[0]['newPrompt']) self.assertTrue(rows[0]['newPrompt'].endswith('<|im_start|>assistant\n')) def test_build_planning_prompt_matches_standard_chat_template(self) -> None: records = normalize_dataset_draft( ''' # dataset_label: 地图导航边界数据 ### case: 多轮顺路停车场 用户: 帮我查一下最近的停车场 小爱: 最近的停车场离你27米 用户: 帮我找一个最顺路的停车场 target: Agent(tag='地图导航') '''.strip(), batch_id='demo', base_timestamp=1_755_567_930_500, timestamp_step_ms=60_000, )['records'] records[0]['context'] = {'location': '', 'rag': ''} prompt = build_planning_prompt(records[0], system_prompt='你是小爱同学，中文智能语音助手。\n') self.assertEqual( prompt, '<|im_start|>system\n' '你是小爱同学，中文智能语音助手。<|im_end|>\n' '<|im_start|>user\n' '请参考用户的[当前query]、[对话历史]、[知识注入]、[系统状态]识别出[当前query]的[function]结果，[function]是python的code形式。\n' '[知识注入]\n' '{\n' '"location": "",\n' '"rag": ""\n' '}\n' '[系统状态]\n' '{}\n' '[对话历史]\n' '用户: 帮我查一下最近的停车场\n' '小爱: 最近的停车场离你27米\n' '[当前query]\n' '用户: 帮我找一个最顺路的停车场\n' '[function]\n' '<|im_end|>\n' '<|im_start|>assistant\n', ) def test_export_planning_eval_csv_writes_file(self) -> None: records = normalize_dataset_draft( ''' # dataset_label: 时间工具数据 ### case: 几点 用户: 现在几点 complex: true target: CalendarQA(type="TIME") '''.strip(), batch_id='demo', base_timestamp=1_755_567_930_500, )['records'] with tempfile.TemporaryDirectory() as tmp_dir: payload = export_planning_eval_csv( records, root=tmp_dir, output_path='output/eval_planning.csv', ) output_path = Path(tmp_dir) / payload['output_path'] rows = list(csv.DictReader(io.StringIO(output_path.read_text(encoding='utf-8-sig')))) self.assertEqual(payload['output_format'], 'csv') self.assertEqual(rows[0]['类别真实标签'], 'CalendarQA') self.assertEqual(rows[0]['code标签'], 'CalendarQA(type="TIME")') self.assertEqual(rows[0]['complex'], 'TRUE') 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_goal_requires_declared_targets_before_review(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: with self.assertRaisesRegex(Exception, 'must include target or target_definitions'): prepare_generation_goal( root=tmp_dir, dataset_label='地图和生活边界数据', goal_summary='生成附近生活服务边界数据', coverage='附近吃喝玩乐', exclusions='不要生成导航路线类 query', ) def test_generation_plan_inherits_targets_from_confirmed_goal(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: goal_payload = prepare_generation_goal( root=tmp_dir, dataset_label='餐饮和导航边界数据', goal_summary='生成餐饮服务和地图导航边界数据', target_definitions=[ {'name': '餐饮服务', 'target': 'Agent(tag="餐饮服务")', 'rule': '找附近餐饮'}, {'name': '地图导航', 'target': 'Agent(tag="地图导航")', 'rule': '明确导航'}, ], coverage='餐饮服务和地图导航边界', exclusions='不要生成无关闲聊', ) goal_id = goal_payload['goal']['goal_id'] confirm_generation_goal(root=tmp_dir, goal_id=goal_id, confirmation='确认目标') plan_payload = prepare_generation_plan( root=tmp_dir, confirmed_goal_id=goal_id, dataset_label='餐饮和导航边界数据', target='', total_count=2, turn_mix='2 条单轮', coverage='餐饮服务和地图导航边界', exclusions='不要生成无关闲聊', output_path='tasks/demo/artifacts', ) self.assertEqual( [item['target'] for item in plan_payload['plan']['target_definitions']], ['Agent(tag="餐饮服务")', 'Agent(tag="地图导航")'], ) 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) def test_data_agent_tools_accept_draft_and_records_paths(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) context = build_tool_context(AgentRuntimeConfig(cwd=root)) draft_path = root / 'draft.txt' draft_path.write_text( ''' # dataset_label: 地图导航边界数据 ### case: 顺路停车场 用户: 帮我找个顺路的停车场 target: Agent(tag='地图导航') '''.strip(), encoding='utf-8', ) normalize_result = execute_tool( default_tool_registry(), 'data_agent_normalize_dataset_draft', { 'draft_path': str(draft_path), 'batch_id': 'demo', 'source_type': 'manual', 'base_timestamp': 1_755_567_930_500, }, context, ) self.assertTrue(normalize_result.ok, normalize_result.content) records = json.loads(normalize_result.content)['records'] self.assertEqual(records[0]['label']['target'], 'Agent(tag="地图导航")') records_path = root / 'records.jsonl' records_path.write_text( '\n'.join(json.dumps(record, ensure_ascii=False) for record in records), encoding='utf-8', ) validate_result = execute_tool( default_tool_registry(), 'data_agent_validate_dataset_records', {'records_path': str(records_path)}, context, ) export_result = execute_tool( default_tool_registry(), 'data_agent_export_dataset_records', { 'records_path': str(records_path), 'output_path': 'output/records.jsonl', }, context, ) self.assertTrue(export_result.ok, export_result.content) export_payload = json.loads(export_result.content) table_exists = (root / export_payload['table_output_path']).exists() training_result = execute_tool( default_tool_registry(), 'data_agent_export_training_jsonl', {'records_path': str(records_path)}, context, ) self.assertTrue(training_result.ok, training_result.content) training_payload = json.loads(training_result.content) training_exists = (root / training_payload['output_path']).exists() eval_result = execute_tool( default_tool_registry(), 'data_agent_export_planning_eval_csv', {'records_path': str(records_path)}, context, ) self.assertTrue(eval_result.ok, eval_result.content) eval_payload = json.loads(eval_result.content) eval_exists = (root / eval_payload['output_path']).exists() self.assertTrue(validate_result.ok, validate_result.content) self.assertTrue(json.loads(validate_result.content)['ok']) self.assertEqual(export_payload['record_count'], 1) self.assertTrue(table_exists) self.assertEqual(training_payload['output_path'], '.port_sessions/data_agent_output/training.jsonl') self.assertTrue(training_exists) self.assertEqual(eval_payload['output_path'], '.port_sessions/data_agent_output/eval_planning.csv') self.assertTrue(eval_exists) def test_tool_schemas_avoid_top_level_composition_keywords(self) -> None: # 部分 Bedrock/Anthropic 兼容后端不接受顶层 anyOf/oneOf/allOf。 blocked_keywords = {'anyOf', 'oneOf', 'allOf'} violations = [ name for name, tool in default_tool_registry().items() if blocked_keywords.intersection(tool.parameters) ] self.assertEqual(violations, []) def test_data_agent_tools_route_relative_outputs_to_session_output(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) scratchpad = root / '.port_sessions' / 'accounts' / 'alice' / 'sessions' / 's1' / 'scratchpad' scratchpad.mkdir(parents=True) context = build_tool_context(AgentRuntimeConfig(cwd=root), scratchpad_directory=scratchpad) registry = default_tool_registry() plan_result = execute_tool( registry, 'data_agent_prepare_generation_plan', { 'direct_review': True, 'dataset_label': '地图和生活边界数据', 'target': 'Agent(tag="life_service")', 'total_count': 1, 'turn_mix': '1 条单轮', 'coverage': '附近吃喝玩乐', 'exclusions': '不要生成导航路线类 query', 'output_path': 'output/demo/records.jsonl', }, context, ) self.assertTrue(plan_result.ok, plan_result.content) plan_payload = json.loads(plan_result.content) records = normalize_dataset_draft( ''' # dataset_label: 地图和生活边界数据 ### case: 附近餐饮查询 用户: 附近有什么好吃的 target: Agent(tag="life_service") '''.strip(), batch_id='demo', base_timestamp=1_755_567_930_500, )['records'] export_result = execute_tool( registry, 'data_agent_export_dataset_records', {'records': records, 'output_path': 'output/地图和生活边界数据/自定义名字.jsonl'}, context, ) self.assertTrue(export_result.ok, export_result.content) export_payload = json.loads(export_result.content) expected = '.port_sessions/accounts/alice/sessions/s1/output/records.jsonl' expected_table = '.port_sessions/accounts/alice/sessions/s1/output/records.csv' self.assertEqual(plan_payload['plan']['output_path'], expected) self.assertEqual(export_payload['output_path'], expected) self.assertEqual(export_payload['table_output_path'], expected_table) def test_data_agent_session_records_path_stays_stable_when_file_exists(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) scratchpad = root / '.port_sessions' / 'accounts' / 'alice' / 'sessions' / 's1' / 'scratchpad' output_root = scratchpad.parent / 'output' output_root.mkdir(parents=True) (output_root / 'records.jsonl').write_text('{}\n', encoding='utf-8') context = build_tool_context(AgentRuntimeConfig(cwd=root), scratchpad_directory=scratchpad) plan_result = execute_tool( default_tool_registry(), 'data_agent_prepare_generation_plan', { 'direct_review': True, 'dataset_label': '地图和生活边界数据', 'target': 'Agent(tag="life_service")', 'total_count': 1, 'turn_mix': '1 条单轮', 'coverage': '附近吃喝玩乐', 'exclusions': '不要生成导航路线类 query', 'output_path': 'output/任意子目录/任意名字.jsonl', }, context, ) self.assertTrue(plan_result.ok, plan_result.content) plan_payload = json.loads(plan_result.content) self.assertEqual( plan_payload['plan']['output_path'], '.port_sessions/accounts/alice/sessions/s1/output/records.jsonl', ) if __name__ == '__main__': unittest.main()