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

Improve data skill routing and tool organization

Browse Source
This commit is contained in:
武阳
2026-05-06 21:46:09 +08:00
parent 4d3981cccd
commit 1aafc0fce7
16 changed files with 1342 additions and 955 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_agent_prompting.py
+9 -8
View File
@@ -32,12 +32,13 @@ class AgentPromptingTests(unittest.TestCase):
)
prompt = render_system_prompt(parts)
self.assertIn('# System', prompt)
self.assertIn('# Doing tasks', prompt)
self.assertIn('# Using your tools', prompt)
self.assertIn('# Environment', prompt)
self.assertIn('# 系统规则', prompt)
self.assertIn('# 处理任务', prompt)
self.assertIn('# 使用工具', prompt)
self.assertIn('# Skills', prompt)
self.assertIn('product-data', prompt)
self.assertIn('__SYSTEM_PROMPT_DYNAMIC_BOUNDARY__', prompt)
self.assertIn('Primary working directory:', prompt)
self.assertIn('主工作目录:', prompt)
def test_session_state_exports_messages_in_order(self) -> None:
state = AgentSessionState.create(['sys one', 'sys two'], 'hello')
@@ -57,9 +58,9 @@ class AgentPromptingTests(unittest.TestCase):
runtime_config=AgentRuntimeConfig(cwd=Path(tmp_dir)),
)
prompt = agent.render_system_prompt()
self.assertIn('Claw Code Python', prompt)
self.assertIn('# System', prompt)
self.assertIn('# Environment', prompt)
self.assertIn('Zhongkong Agent', prompt)
self.assertIn('# 系统规则', prompt)
self.assertIn('# 环境信息', prompt)
def test_prompt_builder_mentions_plugins_when_cache_is_loaded(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
tests/test_agent_runtime.py
+17 -5
View File
@@ -144,6 +144,17 @@ class AgentRuntimeTests(unittest.TestCase):
self.assertEqual(child_model.model, 'child-model')
self.assertEqual(sorted(child_tools), ['grep_search', 'read_file'])
def test_builtin_child_model_alias_inherits_non_claude_parent_model(self) -> None:
agent = LocalCodingAgent(
model_config=ModelConfig(model='xiaomi/mimo-v2.5-pro'),
runtime_config=AgentRuntimeConfig(cwd=Path.cwd()),
)
agent_def = agent._resolve_agent_definition({'subagent_type': 'Explore'})
child_model = agent._resolve_child_model_config({}, agent_def)
self.assertEqual(agent_def.model, 'haiku')
self.assertEqual(child_model.model, 'xiaomi/mimo-v2.5-pro')
def test_openai_client_parses_tool_calls(self) -> None:
responses = [
{
@@ -322,9 +333,9 @@ class AgentRuntimeTests(unittest.TestCase):
self.assertEqual(result.stop_reason, 'user_review_required')
self.assertEqual(result.tool_calls, 1)
self.assertIn('本轮已暂停', result.final_output)
self.assertIn('先确认边界', result.final_output)
self.assertIn('goal_id', result.final_output)
self.assertIn('计划提示', result.final_output)
self.assertIn('下一步', result.final_output)
self.assertIn('建议先生成 10 条单轮', result.final_output)
review_messages = [
entry
@@ -333,7 +344,7 @@ class AgentRuntimeTests(unittest.TestCase):
and entry.get('stop_reason') == 'user_review_required'
]
self.assertEqual(len(review_messages), 1)
self.assertIn('数据生成目标', str(review_messages[0].get('content', '')))
self.assertIn('生成目标', str(review_messages[0].get('content', '')))
self.assertTrue(
any(event.get('type') == 'user_review_required' for event in result.events)
)
@@ -416,8 +427,9 @@ class AgentRuntimeTests(unittest.TestCase):
result = agent.run('Prepare data generation plan')
self.assertEqual(result.stop_reason, 'user_review_required')
self.assertIn('数据生成计划', result.final_output)
self.assertIn('只补充生成参数', result.final_output)
self.assertIn('plan_id', result.final_output)
self.assertNotIn('target_definitions', result.final_output)
review_messages = [
entry
for entry in result.transcript
@@ -425,7 +437,7 @@ class AgentRuntimeTests(unittest.TestCase):
and entry.get('stop_reason') == 'user_review_required'
]
self.assertEqual(len(review_messages), 1)
self.assertIn('数据生成计划', str(review_messages[0].get('content', '')))
self.assertIn('生成参数', str(review_messages[0].get('content', '')))
def test_write_tool_is_blocked_without_permission(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
tests/test_data_agent_records.py
+93
View File
@@ -229,6 +229,50 @@ target: Agent(tag="life_service")
'建议先生成 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'
@@ -488,6 +532,55 @@ target: Agent(tag="life_service")
self.assertEqual(export_payload['record_count'], 1)
self.assertEqual(len(exported_lines), 1)
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/demo/records.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/demo/records.jsonl'
self.assertEqual(plan_payload['plan']['output_path'], expected)
self.assertEqual(export_payload['output_path'], expected)
if __name__ == '__main__':
unittest.main()