wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add data agent input and export workflow

Browse Source
This commit is contained in:
武阳
2026-04-30 15:11:37 +08:00
parent e3664f8843
commit 5bdc0ffb01
9 changed files with 1787 additions and 237 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_data_agent_records.py
+134
View File
@@ -8,9 +8,12 @@ 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,
@@ -94,6 +97,32 @@ target: Agent(tag="life_service")
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(
'''
@@ -151,6 +180,55 @@ target: Agent(tag="life_service")
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'
@@ -290,10 +368,48 @@ target: Agent(tag="天气")
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,
@@ -350,9 +466,27 @@ target: Agent(tag="life_service")
{'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__':