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
src/agent_tools.py
+366 -3
View File
@@ -17,14 +17,24 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, Union
from .agent_types import AgentPermissions, AgentRuntimeConfig, ToolExecutionResult
from .data_agent_records import (
DataRecordError,
confirm_generation_goal,
confirm_generation_plan,
export_dataset_records,
get_generation_plan,
normalize_dataset_draft,
prepare_generation_goal,
prepare_generation_plan,
records_from_tool_argument,
update_generation_plan,
validate_dataset_records,
)
from .data_agent_inputs import (
DataAgentInputError,
dumps_payload as data_agent_dumps_payload,
extract_case_evidence,
load_input_sources,
render_source_context,
)
from .session_env_vars import get_session_env_vars
if TYPE_CHECKING:
@@ -1252,14 +1262,92 @@ def default_tool_registry() -> dict[str, AgentTool]:
handler=_execute_skill,
),
AgentTool(
name='data_agent_prepare_generation_plan',
name='data_agent_prepare_generation_goal',
description=(
'Create a pending data-agent generation plan. Use this before generating dataset draft text; '
'show the returned plan to the user and wait for confirmation.'
'Create a pending data-agent generation goal from source evidence or user rules. '
'Use this before data_agent_prepare_generation_plan; show the returned goal to the user and wait for confirmation.'
),
parameters={
'type': 'object',
'properties': {
'dataset_label': {'type': 'string'},
'goal_summary': {'type': 'string'},
'target': {
'type': 'string',
'description': 'Single target label. For multi-target boundary tasks, use target_definitions.',
},
'target_definitions': {
'type': 'array',
'items': {
'type': 'object',
'properties': {
'name': {'type': 'string'},
'target': {'type': 'string'},
'rule': {'type': 'string'},
},
'required': ['target'],
},
},
'plan_hint': {
'type': 'string',
'description': (
'Optional plain-language hint for the later plan, such as '
'"建议先生成 50 条单轮，输出到 tasks/.../records.jsonl". '
'Structured total_count, turn_mix, and output_path belong to data_agent_prepare_generation_plan.'
),
},
'coverage': {'type': 'string'},
'exclusions': {'type': 'string'},
'open_questions': {
'type': 'array',
'items': {'type': 'string'},
},
'source_refs': {
'type': 'array',
'items': {'type': 'string'},
},
'notes': {'type': 'string'},
},
'required': ['dataset_label', 'goal_summary', 'coverage', 'exclusions'],
},
handler=_data_agent_prepare_generation_goal_tool,
),
AgentTool(
name='data_agent_confirm_generation_goal',
description=(
'Mark a pending data-agent generation goal as confirmed after the user explicitly approves it. '
'The returned confirmed_goal_id is required by data_agent_prepare_generation_plan.'
),
parameters={
'type': 'object',
'properties': {
'goal_id': {'type': 'string'},
'confirmation': {
'type': 'string',
'description': 'The user approval text, such as 确认, 开始生成, or approve.',
},
'reviewed_revision': {
'type': 'integer',
'description': 'The goal revision shown to the user before confirmation.',
},
},
'required': ['goal_id', 'confirmation'],
},
handler=_data_agent_confirm_generation_goal_tool,
),
AgentTool(
name='data_agent_prepare_generation_plan',
description=(
'Create a pending data-agent generation plan. Use this before generating dataset draft text; '
'requires confirmed_goal_id from data_agent_confirm_generation_goal; show the returned plan to the user and wait for confirmation.'
),
parameters={
'type': 'object',
'properties': {
'confirmed_goal_id': {
'type': 'string',
'description': 'Required; returned by data_agent_confirm_generation_goal after the user reviews the generation goal.',
},
'dataset_label': {'type': 'string'},
'target': {
'type': 'string',
@@ -1286,6 +1374,7 @@ def default_tool_registry() -> dict[str, AgentTool]:
'notes': {'type': 'string'},
},
'required': [
'confirmed_goal_id',
'dataset_label',
'total_count',
'turn_mix',
@@ -1296,6 +1385,87 @@ def default_tool_registry() -> dict[str, AgentTool]:
},
handler=_data_agent_prepare_generation_plan_tool,
),
AgentTool(
name='data_agent_load_input_sources',
description=(
'Load data-agent input files or directories and extract structured paragraphs and table previews '
'from xlsx, csv, docx, pdf, txt, md, json, or jsonl sources.'
),
parameters={
'type': 'object',
'properties': {
'paths': {
'type': 'array',
'items': {'type': 'string'},
'description': 'Workspace-relative files or directories to load.',
},
'max_files': {'type': 'integer', 'minimum': 1, 'maximum': 100},
'max_paragraphs_per_file': {'type': 'integer', 'minimum': 1, 'maximum': 300},
'max_tables_per_file': {'type': 'integer', 'minimum': 1, 'maximum': 100},
'max_rows_per_table': {'type': 'integer', 'minimum': 1, 'maximum': 200},
'max_cell_chars': {'type': 'integer', 'minimum': 20, 'maximum': 2000},
},
'required': ['paths'],
},
handler=_data_agent_load_input_sources_tool,
),
AgentTool(
name='data_agent_extract_case_evidence',
description=(
'Extract query/badcase evidence from loaded input sources or source paths. '
'Profiles table columns first and returns required questions when query or expected label columns are ambiguous.'
),
parameters={
'type': 'object',
'properties': {
'paths': {
'type': 'array',
'items': {'type': 'string'},
'description': 'Workspace-relative files or directories. Omit when loaded_sources is provided.',
},
'loaded_sources': {
'type': 'object',
'description': 'Output from data_agent_load_input_sources.',
},
'field_mapping': {
'type': 'object',
'description': 'Optional role-to-column mapping, e.g. {"query":"query","expected_label":"预期domain"}.',
},
'max_cases': {'type': 'integer', 'minimum': 1, 'maximum': 1000},
},
},
handler=_data_agent_extract_case_evidence_tool,
),
AgentTool(
name='data_agent_render_source_context',
description=(
'Render loaded data-agent sources into LLM-readable evidence text with source refs. '
'Use this after data_agent_load_input_sources before asking the model to structure product semantics.'
),
parameters={
'type': 'object',
'properties': {
'paths': {
'type': 'array',
'items': {'type': 'string'},
'description': 'Workspace-relative files or directories. Omit when loaded_sources is provided.',
},
'loaded_sources': {
'type': 'object',
'description': 'Output from data_agent_load_input_sources.',
},
'max_chars': {'type': 'integer', 'minimum': 1000, 'maximum': 200000},
'max_tables_per_source': {'type': 'integer', 'minimum': 1, 'maximum': 100},
'max_rows_per_table': {'type': 'integer', 'minimum': 1, 'maximum': 200},
'focus_keywords': {
'type': 'array',
'items': {'type': 'string'},
'description': 'Optional keywords used to keep only matching paragraphs/table rows.',
},
},
},
handler=_data_agent_render_source_context_tool,
),
AgentTool(
name='data_agent_show_generation_plan',
description='Show a data-agent generation plan for human review.',
@@ -1417,6 +1587,44 @@ def default_tool_registry() -> dict[str, AgentTool]:
},
handler=_validate_dataset_records_tool,
),
AgentTool(
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.'
),
parameters={
'type': 'object',
'properties': {
'records': {
'oneOf': [
{'type': 'array'},
{'type': 'string'},
],
'description': 'Canonical records as an array, or a JSON string containing the array.',
},
'output_path': {
'type': 'string',
'description': 'Workspace-relative path to write, usually ending in .jsonl.',
},
'output_format': {
'type': 'string',
'enum': ['jsonl', 'json'],
'description': 'Defaults to jsonl. json writes a compact JSON array.',
},
'require_validation_ok': {
'type': 'boolean',
'description': 'When true, refuse to write records with validation errors.',
},
'overwrite': {
'type': 'boolean',
'description': 'When false, refuse to overwrite an existing file.',
},
},
'required': ['records', 'output_path'],
},
handler=_export_dataset_records_tool,
),
]
return {tool.name: tool for tool in tools}
@@ -1488,6 +1696,64 @@ def _coerce_float(arguments: dict[str, Any], key: str, default: float) -> float:
return float(value)
def _optional_string_list(arguments: dict[str, Any], key: str) -> list[str] | None:
value = arguments.get(key)
if value is None:
return None
if not isinstance(value, list) or not all(isinstance(item, str) for item in value):
raise ToolExecutionError(f'{key} must be an array of strings')
return value
def _data_agent_prepare_generation_goal_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
notes = arguments.get('notes', '')
if not isinstance(notes, str):
raise ToolExecutionError('notes must be a string')
try:
payload = prepare_generation_goal(
root=str(context.root),
dataset_label=_require_string(arguments, 'dataset_label'),
goal_summary=_require_string(arguments, 'goal_summary'),
target=_optional_string(arguments, 'target'),
target_definitions=_optional_target_definitions(arguments.get('target_definitions')),
plan_hint=_optional_string(arguments, 'plan_hint'),
coverage=_require_string(arguments, 'coverage'),
exclusions=_require_string(arguments, 'exclusions'),
open_questions=_optional_string_list(arguments, 'open_questions'),
source_refs=_optional_string_list(arguments, 'source_refs'),
notes=notes,
)
except DataRecordError as exc:
raise ToolExecutionError(str(exc)) from exc
return (
json.dumps(payload, ensure_ascii=False, indent=2),
{
'action': 'data_agent_prepare_generation_goal',
'requires_user_review': True,
'goal_id': payload.get('goal', {}).get('goal_id'),
'goal_revision': payload.get('goal', {}).get('revision'),
},
)
def _data_agent_confirm_generation_goal_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
reviewed_revision = arguments.get('reviewed_revision')
if reviewed_revision is not None and (
isinstance(reviewed_revision, bool) or not isinstance(reviewed_revision, int)
):
raise ToolExecutionError('reviewed_revision must be an integer')
try:
payload = confirm_generation_goal(
root=str(context.root),
goal_id=_require_string(arguments, 'goal_id'),
confirmation=_require_string(arguments, 'confirmation'),
reviewed_revision=reviewed_revision,
)
except DataRecordError as exc:
raise ToolExecutionError(str(exc)) from exc
return json.dumps(payload, ensure_ascii=False, indent=2)
def _data_agent_prepare_generation_plan_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
total_count = _coerce_int(arguments, 'total_count', 0)
notes = arguments.get('notes', '')
@@ -1505,6 +1771,7 @@ def _data_agent_prepare_generation_plan_tool(arguments: dict[str, Any], context:
output_path=_require_string(arguments, 'output_path'),
target_definitions=_optional_target_definitions(arguments.get('target_definitions')),
notes=notes,
confirmed_goal_id=_require_string(arguments, 'confirmed_goal_id'),
)
except DataRecordError as exc:
raise ToolExecutionError(str(exc)) from exc
@@ -1519,6 +1786,75 @@ def _data_agent_prepare_generation_plan_tool(arguments: dict[str, Any], context:
)
def _data_agent_load_input_sources_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
paths = arguments.get('paths')
if not isinstance(paths, list) or not all(isinstance(item, str) for item in paths):
raise ToolExecutionError('paths must be an array of strings')
try:
payload = load_input_sources(
context.root,
paths,
max_files=_coerce_int(arguments, 'max_files', 20),
max_paragraphs_per_file=_coerce_int(arguments, 'max_paragraphs_per_file', 80),
max_tables_per_file=_coerce_int(arguments, 'max_tables_per_file', 20),
max_rows_per_table=_coerce_int(arguments, 'max_rows_per_table', 30),
max_cell_chars=_coerce_int(arguments, 'max_cell_chars', 240),
)
except DataAgentInputError as exc:
raise ToolExecutionError(str(exc)) from exc
return data_agent_dumps_payload(payload)
def _data_agent_extract_case_evidence_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
paths = arguments.get('paths')
if paths is not None and (not isinstance(paths, list) or not all(isinstance(item, str) for item in paths)):
raise ToolExecutionError('paths must be an array of strings')
loaded_sources = arguments.get('loaded_sources')
if loaded_sources is not None and not isinstance(loaded_sources, dict):
raise ToolExecutionError('loaded_sources must be an object')
field_mapping = arguments.get('field_mapping')
if field_mapping is not None and not isinstance(field_mapping, dict):
raise ToolExecutionError('field_mapping must be an object')
try:
payload = extract_case_evidence(
context.root,
paths=paths,
loaded_sources=loaded_sources,
field_mapping=field_mapping,
max_cases=_coerce_int(arguments, 'max_cases', 200),
)
except DataAgentInputError as exc:
raise ToolExecutionError(str(exc)) from exc
return data_agent_dumps_payload(payload)
def _data_agent_render_source_context_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
paths = arguments.get('paths')
if paths is not None and (not isinstance(paths, list) or not all(isinstance(item, str) for item in paths)):
raise ToolExecutionError('paths must be an array of strings')
loaded_sources = arguments.get('loaded_sources')
if loaded_sources is not None and not isinstance(loaded_sources, dict):
raise ToolExecutionError('loaded_sources must be an object')
focus_keywords = arguments.get('focus_keywords')
if focus_keywords is not None and (
not isinstance(focus_keywords, list) or not all(isinstance(item, str) for item in focus_keywords)
):
raise ToolExecutionError('focus_keywords must be an array of strings')
try:
payload = render_source_context(
context.root,
paths=paths,
loaded_sources=loaded_sources,
max_chars=_coerce_int(arguments, 'max_chars', 40_000),
max_tables_per_source=_coerce_int(arguments, 'max_tables_per_source', 12),
max_rows_per_table=_coerce_int(arguments, 'max_rows_per_table', 20),
focus_keywords=focus_keywords,
)
except DataAgentInputError as exc:
raise ToolExecutionError(str(exc)) from exc
return data_agent_dumps_payload(payload)
def _data_agent_show_generation_plan_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
try:
payload = get_generation_plan(
@@ -1609,6 +1945,33 @@ def _validate_dataset_records_tool(arguments: dict[str, Any], _context: ToolExec
return json.dumps(payload, ensure_ascii=False, indent=2)
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')
require_validation_ok = arguments.get('require_validation_ok', True)
if not isinstance(require_validation_ok, bool):
raise ToolExecutionError('require_validation_ok must be a boolean')
overwrite = arguments.get('overwrite', True)
if not isinstance(overwrite, bool):
raise ToolExecutionError('overwrite must be a boolean')
try:
records = records_from_tool_argument(arguments['records'])
payload = export_dataset_records(
records,
root=str(context.root),
output_path=_require_string(arguments, 'output_path'),
output_format=output_format,
require_validation_ok=require_validation_ok,
overwrite=overwrite,
)
except (DataRecordError, json.JSONDecodeError) as exc:
raise ToolExecutionError(str(exc)) from exc
return json.dumps(payload, ensure_ascii=False, indent=2)
def _resolve_path(raw_path: str, context: ToolExecutionContext, *, allow_missing: bool = True) -> Path:
expanded = Path(raw_path).expanduser()
candidate = expanded if expanded.is_absolute() else context.root / expanded