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
src/agent_runtime.py
+48 -3
View File
@@ -832,8 +832,15 @@ class LocalCodingAgent:
return result
# fall through to the normal tool-call branch below
# normal error path if not recovered
final_output = str(exc)
session.append_assistant(
final_output,
message_id=f'assistant_{len(session.messages)}',
stop_reason='backend_error',
)
_append_final_text_stream_events(stream_events, final_output)
result = AgentRunResult(
final_output=str(exc),
final_output=final_output,
turns=max(turn_index - 1, 0),
tool_calls=tool_calls,
transcript=session.transcript(),
@@ -1535,15 +1542,42 @@ class LocalCodingAgent:
usage=usage,
)
assistant_message = session.messages[assistant_index]
try:
parsed_tool_calls = self._tool_calls_from_message(assistant_message.tool_calls)
except OpenAICompatError:
self._scrub_malformed_assistant_tool_calls(session, assistant_index)
raise
turn = AssistantTurn(
content=assistant_message.content,
tool_calls=self._tool_calls_from_message(assistant_message.tool_calls),
tool_calls=parsed_tool_calls,
finish_reason=finish_reason,
raw_message=assistant_message.to_openai_message(),
usage=usage,
)
return turn, tuple(events)
def _scrub_malformed_assistant_tool_calls(
self,
session: AgentSessionState,
assistant_index: int,
) -> None:
message = session.messages[assistant_index]
filtered_blocks = tuple(
block
for block in message.blocks
if block.get('type') not in {'tool_use', 'tool_call'}
)
session.messages[assistant_index] = replace(
message,
tool_calls=(),
blocks=filtered_blocks,
stop_reason='malformed_tool_arguments',
metadata={
**message.metadata,
'malformed_tool_arguments': True,
},
)
def _tool_calls_from_message(
self,
tool_calls: tuple[dict[str, object], ...],
@@ -1558,7 +1592,18 @@ class LocalCodingAgent:
continue
raw_arguments = function_block.get('arguments', '')
if isinstance(raw_arguments, str) and raw_arguments.strip():
arguments = json.loads(raw_arguments)
try:
arguments = json.loads(raw_arguments)
except json.JSONDecodeError as exc:
preview = raw_arguments[:240].replace('\n', '\\n')
raise OpenAICompatError(
'模型返回的工具参数不是合法 JSON，已停止本轮以避免执行错误工具。'
f'工具：{name}；位置：line {exc.lineno} column {exc.colno}；'
f'原因：{exc.msg}；参数片段：{preview!r}。'
'请回复“继续”，我会要求模型重新生成合法 JSON 参数；'
'对于大段 dataset draft，应分批或使用 draft_path/records_path，'
'不要把大量带引号的文本一次性塞进工具参数。'
) from exc
if not isinstance(arguments, dict):
raise OpenAICompatError(
f'Tool arguments must decode to an object, got {type(arguments).__name__}'
src/agent_tool_specs/data_agent.py
+33 -5
View File
@@ -429,7 +429,8 @@ def build_data_agent_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentToo
description=(
'Parse dataset draft text written with 用户/小爱/target blocks and build canonical '
'data-agent records with generated record IDs, timestamps, source metadata, context, and labels. '
'Generated data requires a confirmed_plan_id from data_agent_confirm_generation_plan.'
'Generated data requires a confirmed_plan_id from data_agent_confirm_generation_plan. '
'Use draft_path instead of draft_text when the draft is large or contains many quoted targets.'
),
parameters={
'type': 'object',
@@ -438,6 +439,10 @@ def build_data_agent_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentToo
'type': 'string',
'description': 'Dataset draft text in dataset draft text v1 format.',
},
'draft_path': {
'type': 'string',
'description': 'Path to a UTF-8 dataset draft text file. Prefer this for large drafts.',
},
'batch_id': {
'type': 'string',
'description': 'Batch identifier used in generated record_id values.',
@@ -464,7 +469,10 @@ def build_data_agent_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentToo
'description': 'Required for generated data; returned by data_agent_confirm_generation_plan.',
},
},
'required': ['draft_text'],
'anyOf': [
{'required': ['draft_text']},
{'required': ['draft_path']},
],
},
handler=resolve_handler(handlers, 'data_agent_normalize_dataset_draft', 'data-agent'),
),
@@ -484,8 +492,15 @@ def build_data_agent_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentToo
],
'description': 'Canonical records as an array, or a JSON string containing the array.',
},
'records_path': {
'type': 'string',
'description': 'Path to records JSON/JSONL. Prefer this when records are large.',
},
},
'required': ['records'],
'anyOf': [
{'required': ['records']},
{'required': ['records_path']},
],
},
handler=resolve_handler(handlers, 'data_agent_validate_dataset_records', 'data-agent'),
),
@@ -493,7 +508,8 @@ def build_data_agent_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentToo
name='data_agent_export_dataset_records',
description=(
'Validate canonical data-agent records and write them to a workspace file. '
'Defaults to compact JSONL, one canonical record per line, so the model does not hand-write JSON output.'
'Defaults to compact JSONL, one canonical record per line, and also writes records.csv '
'in the same output directory for human sharing.'
),
parameters={
'type': 'object',
@@ -505,6 +521,10 @@ def build_data_agent_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentToo
],
'description': 'Canonical records as an array, or a JSON string containing the array.',
},
'records_path': {
'type': 'string',
'description': 'Path to records JSON/JSONL. Prefer this when records are large.',
},
'output_path': {
'type': 'string',
'description': (
@@ -526,8 +546,16 @@ def build_data_agent_tools(handlers: Mapping[str, ToolHandler]) -> list[AgentToo
'type': 'boolean',
'description': 'When false, refuse to overwrite an existing file.',
},
'export_table': {
'type': 'boolean',
'description': 'Defaults to true. Also write records.csv with the shared spreadsheet columns.',
},
},
'required': ['records', 'output_path'],
'required': ['output_path'],
'anyOf': [
{'required': ['records']},
{'required': ['records_path']},
],
},
handler=resolve_handler(handlers, 'data_agent_export_dataset_records', 'data-agent'),
),
src/agent_tools.py
+61 -7
View File
@@ -1474,7 +1474,7 @@ def _data_agent_confirm_generation_plan_tool(arguments: dict[str, Any], context:
def _normalize_dataset_draft_tool(arguments: dict[str, Any], _context: ToolExecutionContext) -> str:
draft_text = _require_string(arguments, 'draft_text')
draft_text = _draft_text_from_arguments(arguments, _context)
batch_id = arguments.get('batch_id', 'aabbccdd')
if not isinstance(batch_id, str) or not batch_id.strip():
raise ToolExecutionError('batch_id must be a non-empty string')
@@ -1508,10 +1508,8 @@ def _normalize_dataset_draft_tool(arguments: dict[str, Any], _context: ToolExecu
def _validate_dataset_records_tool(arguments: dict[str, Any], _context: ToolExecutionContext) -> str:
if 'records' not in arguments:
raise ToolExecutionError('records is required')
try:
records = records_from_tool_argument(arguments['records'])
records = _records_from_arguments(arguments, _context)
payload = validate_dataset_records(records)
except (DataRecordError, json.JSONDecodeError) as exc:
raise ToolExecutionError(str(exc)) from exc
@@ -1519,8 +1517,6 @@ def _validate_dataset_records_tool(arguments: dict[str, Any], _context: ToolExec
def _export_dataset_records_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
if 'records' not in arguments:
raise ToolExecutionError('records is required')
output_format = arguments.get('output_format', 'jsonl')
if not isinstance(output_format, str):
raise ToolExecutionError('output_format must be a string')
@@ -1530,8 +1526,11 @@ def _export_dataset_records_tool(arguments: dict[str, Any], context: ToolExecuti
overwrite = arguments.get('overwrite', True)
if not isinstance(overwrite, bool):
raise ToolExecutionError('overwrite must be a boolean')
export_table = arguments.get('export_table', True)
if not isinstance(export_table, bool):
raise ToolExecutionError('export_table must be a boolean')
try:
records = records_from_tool_argument(arguments['records'])
records = _records_from_arguments(arguments, context)
canonical_filename = 'records.json' if output_format == 'json' else 'records.jsonl'
payload = export_dataset_records(
records,
@@ -1544,12 +1543,67 @@ def _export_dataset_records_tool(arguments: dict[str, Any], context: ToolExecuti
output_format=output_format,
require_validation_ok=require_validation_ok,
overwrite=overwrite,
export_table=export_table,
)
except (DataRecordError, json.JSONDecodeError) as exc:
raise ToolExecutionError(str(exc)) from exc
return json.dumps(payload, ensure_ascii=False, indent=2)
def _draft_text_from_arguments(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
draft_text = arguments.get('draft_text')
draft_path = arguments.get('draft_path')
has_text = isinstance(draft_text, str) and bool(draft_text.strip())
has_path = isinstance(draft_path, str) and bool(draft_path.strip())
if has_text and has_path:
raise ToolExecutionError('provide only one of draft_text or draft_path')
if has_text:
return draft_text
if not has_path:
raise ToolExecutionError('draft_text or draft_path is required')
path = _resolve_path(draft_path, context, allow_missing=False, allow_outside_root=True)
if not path.is_file():
raise ToolExecutionError(f'draft_path is not a file: {path}')
return path.read_text(encoding='utf-8', errors='replace')
def _records_from_arguments(arguments: dict[str, Any], context: ToolExecutionContext) -> list[dict[str, Any]]:
has_records = 'records' in arguments and arguments.get('records') is not None
records_path = arguments.get('records_path')
has_path = isinstance(records_path, str) and bool(records_path.strip())
if has_records and has_path:
raise ToolExecutionError('provide only one of records or records_path')
if has_records:
return records_from_tool_argument(arguments['records'])
if not has_path:
raise ToolExecutionError('records or records_path is required')
path = _resolve_path(records_path, context, allow_missing=False, allow_outside_root=True)
if not path.is_file():
raise ToolExecutionError(f'records_path is not a file: {path}')
return _records_from_file(path)
def _records_from_file(path: Path) -> list[dict[str, Any]]:
text = path.read_text(encoding='utf-8', errors='replace').strip()
if not text:
raise DataRecordError('records_path is empty')
if path.suffix == '.jsonl':
records: list[dict[str, Any]] = []
for line_number, line in enumerate(text.splitlines(), start=1):
stripped = line.strip()
if not stripped:
continue
item = json.loads(stripped)
if not isinstance(item, dict):
raise DataRecordError(f'records_path line {line_number} must decode to an object')
records.append(item)
return records_from_tool_argument(records)
payload = json.loads(text)
if isinstance(payload, dict) and isinstance(payload.get('records'), list):
payload = payload['records']
return records_from_tool_argument(payload)
def _data_agent_output_path(
output_path: str,
context: ToolExecutionContext,
src/data_agent_records.py
+89 -2
View File
@@ -7,9 +7,11 @@ Skill 应该先让模型生成紧凑、便于人工 review 的 dataset draft tex
再由这里的逻辑解析并补齐为 canonical records，供后续校验和导出使用。
"""
import csv
import json
import re
import time
from io import StringIO
from pathlib import Path
from typing import Any
@@ -461,6 +463,7 @@ def export_dataset_records(
output_format: str = 'jsonl',
require_validation_ok: bool = True,
overwrite: bool = True,
export_table: bool = True,
) -> dict[str, Any]:
"""校验并导出 canonical records。
@@ -480,8 +483,11 @@ def export_dataset_records(
)
path = _resolve_output_path(root, output_path)
table_path = path.with_name('records.csv')
if path.exists() and not overwrite:
raise DataRecordError(f'output_path already exists: {output_path}')
if export_table and table_path.exists() and not overwrite:
raise DataRecordError('table output_path already exists: records.csv')
path.parent.mkdir(parents=True, exist_ok=True)
if output_format == 'jsonl':
@@ -494,13 +500,83 @@ def export_dataset_records(
path.write_text(content, encoding='utf-8')
rel_path = str(path.relative_to(Path(root).resolve()))
return {
payload = {
'output_path': rel_path,
'output_format': output_format,
'record_count': len(records),
'bytes_written': len(content.encode('utf-8')),
'validation': validation,
}
if export_table:
table_content = render_dataset_records_table_csv(records)
table_path.write_text(table_content, encoding='utf-8-sig')
payload.update(
{
'table_output_path': str(table_path.relative_to(Path(root).resolve())),
'table_output_format': 'csv',
'table_bytes_written': len(table_content.encode('utf-8-sig')),
}
)
return payload
def render_dataset_records_table_csv(records: list[dict[str, Any]]) -> str:
"""把 canonical records 转成同事流转用的表格 CSV。"""
output = StringIO()
writer = csv.DictWriter(output, fieldnames=_dataset_table_columns(), lineterminator='\n')
writer.writeheader()
for record in records:
writer.writerow(_dataset_record_table_row(record))
return output.getvalue()
def _dataset_table_columns() -> list[str]:
return [
'request_id',
'timestamp',
'query',
'prev_session',
'context',
'label',
'是否迁移Function',
'function',
]
def _dataset_record_table_row(record: dict[str, Any]) -> dict[str, str]:
source = record.get('source') if isinstance(record.get('source'), dict) else {}
turn = record.get('turn') if isinstance(record.get('turn'), dict) else {}
label = record.get('label') if isinstance(record.get('label'), dict) else {}
context = record.get('context') if isinstance(record.get('context'), dict) else {}
prev_session = record.get('prev_session')
if not isinstance(prev_session, list):
prev_session = []
return {
'request_id': str(source.get('request_id') or ''),
'timestamp': str(source.get('timestamp') or turn.get('timestamp') or ''),
'query': str(turn.get('query') or ''),
'prev_session': json.dumps(_table_prev_session(prev_session), ensure_ascii=False, separators=(',', ':')),
'context': json.dumps(context, ensure_ascii=False, separators=(',', ':')),
'label': str(label.get('dataset_label') or ''),
'是否迁移Function': '',
'function': str(label.get('target') or ''),
}
def _table_prev_session(prev_session: list[Any]) -> list[dict[str, str]]:
rows: list[dict[str, str]] = []
for item in prev_session[-10:]:
if not isinstance(item, dict):
continue
rows.append(
{
'query': str(item.get('query') or ''),
'tts': str(item.get('tts') or ''),
'timestamp': str(item.get('timestamp') or ''),
}
)
return rows
def _parse_draft_text(draft_text: str) -> dict[str, Any]:
@@ -797,7 +873,7 @@ def _case_to_record(
if not current_query:
raise DataRecordError(f'case {index} current 用户 line must be non-empty')
target = str(case.get('target') or '').strip()
target = _canonical_target(str(case.get('target') or '').strip())
if not target:
raise DataRecordError(f'case {index} target is required')
@@ -866,6 +942,17 @@ def _target_type(target: str) -> str:
return 'unknown'
def _canonical_target(target: str) -> str:
"""规范常见 Agent 标签写法，允许草稿里用单引号降低 JSON 转义风险。"""
match = re.fullmatch(r'Agent\s*\(\s*tag\s*=\s*([\'"])(.+?)\1\s*\)', target)
if not match:
return target
tag = match.group(2).strip()
escaped = tag.replace('\\', '\\\\').replace('"', '\\"')
return f'Agent(tag="{escaped}")'
def _record_id(*, source_type: str, batch_id: str, request_id: str, index: int) -> str:
if source_type == 'generated':
return f'gen_{_safe_token(batch_id)}_{index:06d}'
src/openai_compat.py
+4 -1
View File
@@ -75,8 +75,11 @@ def _parse_tool_arguments(raw_arguments: Any) -> dict[str, Any]:
try:
parsed = json.loads(raw_arguments)
except json.JSONDecodeError as exc:
preview = raw_arguments[:240].replace('\n', '\\n')
raise OpenAICompatError(
f'Invalid tool arguments returned by model: {raw_arguments!r}'
'Invalid tool arguments returned by model: '
f'line {exc.lineno} column {exc.colno}: {exc.msg}; '
f'preview={preview!r}'
) from exc
if not isinstance(parsed, dict):
raise OpenAICompatError(