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

Stabilize product data export flow

Browse Source
This commit is contained in:
wuyang6
2026-05-11 11:46:09 +08:00
parent 1752c81913
commit 0e2a00990f
19 changed files with 659 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_agent_runtime.py
+58
View File
@@ -769,6 +769,64 @@ class AgentRuntimeTests(unittest.TestCase):
self.assertGreaterEqual(mutation_totals.get('assistant_tool_call_delta', 0), 2)
self.assertEqual(mutation_totals.get('assistant_finalize', 0), 1)
def test_agent_reports_malformed_streaming_tool_arguments(self) -> None:
responses = [
[
{
'choices': [
{
'delta': {
'tool_calls': [
{
'index': 0,
'id': 'call_1',
'function': {
'name': 'data_agent_normalize_dataset_draft',
'arguments': '{"draft_text": "target: Agent(tag="地图导航")"}',
},
}
]
},
'finish_reason': None,
}
]
},
{
'choices': [{'delta': {}, 'finish_reason': 'tool_calls'}],
'usage': {'prompt_tokens': 9, 'completion_tokens': 4},
},
]
]
with tempfile.TemporaryDirectory() as tmp_dir:
workspace = Path(tmp_dir)
with patch(
'src.openai_compat.request.urlopen',
side_effect=make_streaming_urlopen_side_effect(responses),
):
agent = LocalCodingAgent(
model_config=ModelConfig(
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
base_url='http://127.0.0.1:8000/v1',
),
runtime_config=AgentRuntimeConfig(
cwd=workspace,
stream_model_responses=True,
),
)
result = agent.run('Generate records')
self.assertEqual(result.stop_reason, 'backend_error')
self.assertIn('模型返回的工具参数不是合法 JSON', result.final_output)
self.assertIn('data_agent_normalize_dataset_draft', result.final_output)
malformed = [
message
for message in result.transcript
if message.get('stop_reason') == 'malformed_tool_arguments'
]
self.assertEqual(len(malformed), 1)
self.assertNotIn('tool_calls', malformed[0])
self.assertEqual(result.transcript[-1]['content'], result.final_output)
def test_transcript_entries_include_structured_blocks(self) -> None:
responses = [
{
tests/test_data_agent_records.py
+117
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import csv
import json
import tempfile
import unittest
@@ -15,6 +16,7 @@ from src.data_agent_records import (
normalize_dataset_draft,
prepare_generation_goal,
prepare_generation_plan,
render_dataset_records_table_csv,
update_generation_plan,
validate_dataset_records,
)
@@ -64,6 +66,22 @@ notes: 多轮承接附近生活服务查询
],
)
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_validate_dataset_records_reports_valid_payload(self) -> None:
records = normalize_dataset_draft(
'''
@@ -116,12 +134,50 @@ target: Agent(tag="life_service")
)
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(table_path.read_text(encoding='utf-8-sig').splitlines()))
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'], 'Agent(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(render_dataset_records_table_csv(records).splitlines()))
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'], 'Agent(tag="地图导航")')
def test_normalize_online_draft_uses_real_request_metadata(self) -> None:
records = normalize_dataset_draft(
@@ -532,6 +588,65 @@ target: Agent(tag="life_service")
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()
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)
def test_data_agent_tools_route_relative_outputs_to_session_output(self) -> None:
with tempfile.TemporaryDirectory() as tmp_dir:
root = Path(tmp_dir)
@@ -578,8 +693,10 @@ target: Agent(tag="life_service")
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: