Add data agent record tools

2026-04-29 19:31:16 +08:00
parent a040d80096
commit 12f6ced252
3 changed files with 1331 additions and 0 deletions
@@ -0,0 +1,359 @@
+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_plan,
+    get_generation_plan,
+    normalize_dataset_draft,
+    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_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_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)))
+            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.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,
+            )
+
+        self.assertTrue(validate_result.ok)
+        self.assertTrue(json.loads(validate_result.content)['ok'])
+
+
+if __name__ == '__main__':
+    unittest.main()