Add data agent record tools

src/agent_tools.py

@@ -15,6 +15,16 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Iterator, Union
 
 from .agent_types import AgentPermissions, AgentRuntimeConfig, ToolExecutionResult
+from .data_agent_records import (
+    DataRecordError,
+    confirm_generation_plan,
+    get_generation_plan,
+    normalize_dataset_draft,
+    prepare_generation_plan,
+    records_from_tool_argument,
+    update_generation_plan,
+    validate_dataset_records,
+)
 from .session_env_vars import get_session_env_vars
 
 if TYPE_CHECKING:
@@ -1241,6 +1251,172 @@ def default_tool_registry() -> dict[str, AgentTool]:
             },
             handler=_execute_skill,
         ),
+        AgentTool(
+            name='data_agent_prepare_generation_plan',
+            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.'
+            ),
+            parameters={
+                'type': 'object',
+                'properties': {
+                    'dataset_label': {'type': 'string'},
+                    'target': {
+                        'type': 'string',
+                        'description': 'Single target label. For multi-target boundary tasks, use a short summary and fill target_definitions.',
+                    },
+                    'target_definitions': {
+                        'type': 'array',
+                        'description': 'Optional target labels for multi-class boundary data.',
+                        'items': {
+                            'type': 'object',
+                            'properties': {
+                                'name': {'type': 'string'},
+                                'target': {'type': 'string'},
+                                'rule': {'type': 'string'},
+                            },
+                            'required': ['target'],
+                        },
+                    },
+                    'total_count': {'type': 'integer', 'minimum': 1},
+                    'turn_mix': {'type': 'string', 'description': 'Single-turn and multi-turn count or ratio.'},
+                    'coverage': {'type': 'string', 'description': 'Query types, intent boundaries, or error types to cover.'},
+                    'exclusions': {'type': 'string', 'description': 'Negative examples or boundaries to avoid.'},
+                    'output_path': {'type': 'string', 'description': 'Where draft, records, and validation files should be written.'},
+                    'notes': {'type': 'string'},
+                },
+                'required': [
+                    'dataset_label',
+                    'total_count',
+                    'turn_mix',
+                    'coverage',
+                    'exclusions',
+                    'output_path',
+                ],
+            },
+            handler=_data_agent_prepare_generation_plan_tool,
+        ),
+        AgentTool(
+            name='data_agent_show_generation_plan',
+            description='Show a data-agent generation plan for human review.',
+            parameters={
+                'type': 'object',
+                'properties': {
+                    'plan_id': {'type': 'string'},
+                },
+                'required': ['plan_id'],
+            },
+            handler=_data_agent_show_generation_plan_tool,
+        ),
+        AgentTool(
+            name='data_agent_update_generation_plan',
+            description=(
+                'Update a pending data-agent generation plan from human review feedback, then show it again for review.'
+            ),
+            parameters={
+                'type': 'object',
+                'properties': {
+                    'plan_id': {'type': 'string'},
+                    'review_feedback': {
+                        'type': 'string',
+                        'description': 'Human review feedback that explains why the plan is changing.',
+                    },
+                    'updates': {
+                        'type': 'object',
+                        'description': 'Fields to update on the plan.',
+                    },
+                },
+                'required': ['plan_id', 'review_feedback', 'updates'],
+            },
+            handler=_data_agent_update_generation_plan_tool,
+        ),
+        AgentTool(
+            name='data_agent_confirm_generation_plan',
+            description=(
+                'Mark a pending data-agent generation plan as confirmed after the user explicitly approves it.'
+            ),
+            parameters={
+                'type': 'object',
+                'properties': {
+                    'plan_id': {'type': 'string'},
+                    'confirmation': {
+                        'type': 'string',
+                        'description': 'The user approval text, such as 确认, 开始生成, or approve.',
+                    },
+                    'reviewed_revision': {
+                        'type': 'integer',
+                        'description': 'The plan revision shown to the user before confirmation.',
+                    },
+                },
+                'required': ['plan_id', 'confirmation'],
+            },
+            handler=_data_agent_confirm_generation_plan_tool,
+        ),
+        AgentTool(
+            name='data_agent_normalize_dataset_draft',
+            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.'
+            ),
+            parameters={
+                'type': 'object',
+                'properties': {
+                    'draft_text': {
+                        'type': 'string',
+                        'description': 'Dataset draft text in dataset draft text v1 format.',
+                    },
+                    'batch_id': {
+                        'type': 'string',
+                        'description': 'Batch identifier used in generated record_id values.',
+                    },
+                    'source_type': {
+                        'type': 'string',
+                        'enum': ['generated', 'online', 'manual', 'mixed'],
+                    },
+                    'base_timestamp': {
+                        'type': 'integer',
+                        'description': 'Optional millisecond timestamp for deterministic generated records.',
+                    },
+                    'timestamp_step_ms': {
+                        'type': 'integer',
+                        'minimum': 1,
+                        'description': 'Millisecond gap between adjacent turns. Defaults to 60000.',
+                    },
+                    'default_request_id': {
+                        'type': 'string',
+                        'description': 'Request ID to use when a case does not provide a real request_id.',
+                    },
+                    'confirmed_plan_id': {
+                        'type': 'string',
+                        'description': 'Required for generated data; returned by data_agent_confirm_generation_plan.',
+                    },
+                },
+                'required': ['draft_text'],
+            },
+            handler=_normalize_dataset_draft_tool,
+        ),
+        AgentTool(
+            name='data_agent_validate_dataset_records',
+            description=(
+                'Validate canonical data-agent records for required fields, labels, source metadata, '
+                'prev_session structure, timestamp order, and 5-minute prompt stitching constraints.'
+            ),
+            parameters={
+                'type': 'object',
+                'properties': {
+                    'records': {
+                        'oneOf': [
+                            {'type': 'array'},
+                            {'type': 'string'},
+                        ],
+                        'description': 'Canonical records as an array, or a JSON string containing the array.',
+                    },
+                },
+                'required': ['records'],
+            },
+            handler=_validate_dataset_records_tool,
+        ),
     ]
     return {tool.name: tool for tool in tools}
 
@@ -1278,6 +1454,26 @@ def _require_string(arguments: dict[str, Any], key: str) -> str:
     return value
 
 
+def _optional_string(arguments: dict[str, Any], key: str) -> str:
+    value = arguments.get(key, '')
+    if value is None:
+        return ''
+    if not isinstance(value, str):
+        raise ToolExecutionError(f'{key} must be a string')
+    return value
+
+
+def _optional_target_definitions(value: Any) -> list[dict[str, Any]] | None:
+    if value is None:
+        return None
+    if not isinstance(value, list):
+        raise ToolExecutionError('target_definitions must be an array')
+    for index, item in enumerate(value):
+        if not isinstance(item, dict):
+            raise ToolExecutionError(f'target_definitions[{index}] must be an object')
+    return value
+
+
 def _coerce_int(arguments: dict[str, Any], key: str, default: int) -> int:
     value = arguments.get(key, default)
     if isinstance(value, bool) or not isinstance(value, int):
@@ -1292,6 +1488,127 @@ def _coerce_float(arguments: dict[str, Any], key: str, default: float) -> float:
     return float(value)
 
 
+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', '')
+    if not isinstance(notes, str):
+        raise ToolExecutionError('notes must be a string')
+    try:
+        payload = prepare_generation_plan(
+            root=str(context.root),
+            dataset_label=_require_string(arguments, 'dataset_label'),
+            target=_optional_string(arguments, 'target'),
+            total_count=total_count,
+            turn_mix=_require_string(arguments, 'turn_mix'),
+            coverage=_require_string(arguments, 'coverage'),
+            exclusions=_require_string(arguments, 'exclusions'),
+            output_path=_require_string(arguments, 'output_path'),
+            target_definitions=_optional_target_definitions(arguments.get('target_definitions')),
+            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_plan',
+            'requires_user_review': True,
+            'plan_id': payload.get('plan', {}).get('plan_id'),
+            'plan_revision': payload.get('plan', {}).get('revision'),
+        },
+    )
+
+
+def _data_agent_show_generation_plan_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
+    try:
+        payload = get_generation_plan(
+            root=str(context.root),
+            plan_id=_require_string(arguments, 'plan_id'),
+        )
+    except DataRecordError as exc:
+        raise ToolExecutionError(str(exc)) from exc
+    return json.dumps(payload, ensure_ascii=False, indent=2)
+
+
+def _data_agent_update_generation_plan_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
+    updates = arguments.get('updates')
+    if not isinstance(updates, dict):
+        raise ToolExecutionError('updates must be an object')
+    try:
+        payload = update_generation_plan(
+            root=str(context.root),
+            plan_id=_require_string(arguments, 'plan_id'),
+            review_feedback=_require_string(arguments, 'review_feedback'),
+            updates=updates,
+        )
+    except DataRecordError as exc:
+        raise ToolExecutionError(str(exc)) from exc
+    return json.dumps(payload, ensure_ascii=False, indent=2)
+
+
+def _data_agent_confirm_generation_plan_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_plan(
+            root=str(context.root),
+            plan_id=_require_string(arguments, 'plan_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 _normalize_dataset_draft_tool(arguments: dict[str, Any], _context: ToolExecutionContext) -> str:
+    draft_text = _require_string(arguments, 'draft_text')
+    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')
+    source_type = arguments.get('source_type', 'generated')
+    if source_type not in {'generated', 'online', 'manual', 'mixed'}:
+        raise ToolExecutionError('source_type must be one of generated, online, manual, mixed')
+    base_timestamp = arguments.get('base_timestamp')
+    if base_timestamp is not None and (isinstance(base_timestamp, bool) or not isinstance(base_timestamp, int)):
+        raise ToolExecutionError('base_timestamp must be an integer')
+    timestamp_step_ms = _coerce_int(arguments, 'timestamp_step_ms', 60_000)
+    default_request_id = arguments.get('default_request_id', 'aabbccdd')
+    if not isinstance(default_request_id, str) or not default_request_id.strip():
+        raise ToolExecutionError('default_request_id must be a non-empty string')
+    confirmed_plan_id = arguments.get('confirmed_plan_id')
+    if confirmed_plan_id is not None and not isinstance(confirmed_plan_id, str):
+        raise ToolExecutionError('confirmed_plan_id must be a string')
+    try:
+        payload = normalize_dataset_draft(
+            draft_text,
+            batch_id=batch_id.strip(),
+            source_type=source_type,
+            base_timestamp=base_timestamp,
+            timestamp_step_ms=timestamp_step_ms,
+            default_request_id=default_request_id.strip(),
+            confirmed_plan_id=confirmed_plan_id,
+            plan_state_root=str(_context.root) if source_type == 'generated' else None,
+        )
+    except DataRecordError as exc:
+        raise ToolExecutionError(str(exc)) from exc
+    return json.dumps(payload, ensure_ascii=False, indent=2)
+
+
+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'])
+        payload = validate_dataset_records(records)
+    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