From 12f6ced2529d8bbd18826a65efea3ddc4f467872 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=AD=A6=E9=98=B3?= Date: Wed, 29 Apr 2026 19:31:16 +0800 Subject: [PATCH] Add data agent record tools --- src/agent_tools.py | 317 +++++++++++++++ src/data_agent_records.py | 655 +++++++++++++++++++++++++++++++ tests/test_data_agent_records.py | 359 +++++++++++++++++ 3 files changed, 1331 insertions(+) create mode 100644 src/data_agent_records.py create mode 100644 tests/test_data_agent_records.py diff --git a/src/agent_tools.py b/src/agent_tools.py index dc7c25b..3265c27 100644 --- a/src/agent_tools.py +++ b/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 diff --git a/src/data_agent_records.py b/src/data_agent_records.py new file mode 100644 index 0000000..06edb0b --- /dev/null +++ b/src/data_agent_records.py @@ -0,0 +1,655 @@ +from __future__ import annotations + +"""数据智能体记录辅助逻辑。 + +这个模块把数据智能体专属的记录处理逻辑从通用工具注册表里拆出来。 +Skill 应该先让模型生成紧凑、便于人工 review 的 dataset draft text， +再由这里的逻辑解析并补齐为 canonical records，供后续校验和导出使用。 +""" + +import json +import re +import time +from typing import Any + + +DEFAULT_REQUEST_ID = 'aabbccdd' +DEFAULT_TIMESTAMP_STEP_MS = 60_000 +DEFAULT_PLAN_STATE_FILE = '.port_sessions/data_agent_generation_plans.json' + + +class DataRecordError(ValueError): + """Raised when dataset draft text or records are structurally invalid.""" + + +def normalize_dataset_draft( + draft_text: str, + *, + batch_id: str = DEFAULT_REQUEST_ID, + source_type: str = 'generated', + base_timestamp: int | None = None, + timestamp_step_ms: int = DEFAULT_TIMESTAMP_STEP_MS, + default_request_id: str = DEFAULT_REQUEST_ID, + confirmed_plan_id: str | None = None, + plan_state_root: str | None = None, +) -> dict[str, Any]: + if not draft_text.strip(): + raise DataRecordError('draft_text must be non-empty') + if timestamp_step_ms <= 0: + raise DataRecordError('timestamp_step_ms must be greater than 0') + confirmed_plan: dict[str, Any] | None = None + if plan_state_root is not None: + if not confirmed_plan_id: + raise DataRecordError('confirmed_plan_id is required before normalizing generated data') + confirmed_plan = _require_confirmed_generation_plan(plan_state_root, confirmed_plan_id) + + base = int(time.time() * 1000) if base_timestamp is None else base_timestamp + parsed = _parse_draft_text(draft_text) + records: list[dict[str, Any]] = [] + warnings: list[str] = [] + + for index, case in enumerate(parsed['cases'], start=1): + record, case_warnings = _case_to_record( + case, + global_dataset_label=parsed.get('dataset_label') or '', + index=index, + batch_id=batch_id, + source_type=source_type, + base_timestamp=base, + timestamp_step_ms=timestamp_step_ms, + default_request_id=default_request_id, + ) + records.append(record) + warnings.extend(case_warnings) + + if confirmed_plan is not None: + _validate_records_against_plan(records, confirmed_plan) + + return {'records': records, 'warnings': warnings} + + +def prepare_generation_plan( + *, + root: str, + dataset_label: str, + target: str, + total_count: int, + turn_mix: str, + coverage: str, + exclusions: str, + output_path: str, + target_definitions: list[dict[str, Any]] | None = None, + notes: str = '', +) -> dict[str, Any]: + missing = _missing_plan_fields( + { + 'dataset_label': dataset_label, + 'turn_mix': turn_mix, + 'coverage': coverage, + 'exclusions': exclusions, + 'output_path': output_path, + } + ) + normalized_targets = _normalize_target_definitions(target_definitions) + if not target.strip() and not normalized_targets: + missing.append('target') + if missing: + raise DataRecordError('missing required plan fields: ' + ', '.join(missing)) + if total_count <= 0: + raise DataRecordError('total_count must be greater than 0') + + state = _load_plan_state(root) + sequence = int(state.get('next_sequence') or 1) + plan_id = f'data_plan_{sequence:06d}' + resolved_output_path = _dedupe_output_path(root, output_path.strip(), plan_id) + plan = { + 'plan_id': plan_id, + 'status': 'pending_confirmation', + 'dataset_label': dataset_label.strip(), + 'target': target.strip(), + 'target_definitions': normalized_targets, + 'total_count': total_count, + 'turn_mix': turn_mix.strip(), + 'coverage': coverage.strip(), + 'exclusions': exclusions.strip(), + 'output_path': resolved_output_path, + 'requested_output_path': output_path.strip(), + 'notes': notes.strip(), + 'revision': 1, + 'review_history': [], + 'created_at_ms': int(time.time() * 1000), + } + plans = state.setdefault('plans', {}) + if not isinstance(plans, dict): + plans = {} + state['plans'] = plans + plans[plan_id] = plan + state['next_sequence'] = sequence + 1 + _save_plan_state(root, state) + return { + 'plan': plan, + 'message': '计划已创建，状态为 pending_confirmation。请向用户展示计划，并等待用户确认后再调用确认工具。', + } + + +def get_generation_plan(*, root: str, plan_id: str) -> dict[str, Any]: + state = _load_plan_state(root) + plan = _get_plan(state, plan_id) + return { + 'plan': plan, + 'message': '请把计划展示给用户 review。用户有意见时调用 update；用户明确确认后再调用 confirm。', + } + + +def update_generation_plan( + *, + root: str, + plan_id: str, + review_feedback: str, + updates: dict[str, Any], +) -> dict[str, Any]: + if not review_feedback.strip(): + raise DataRecordError('review_feedback is required') + if not isinstance(updates, dict) or not updates: + raise DataRecordError('updates must be a non-empty object') + state = _load_plan_state(root) + plan = _get_plan(state, plan_id) + if plan.get('status') == 'confirmed': + raise DataRecordError(f'generation plan is already confirmed: {plan_id}') + + allowed_keys = { + 'dataset_label', + 'target', + 'target_definitions', + 'total_count', + 'turn_mix', + 'coverage', + 'exclusions', + 'output_path', + 'notes', + } + unknown = sorted(key for key in updates if key not in allowed_keys) + if unknown: + raise DataRecordError('unsupported update fields: ' + ', '.join(unknown)) + + if 'target_definitions' in updates: + updates = dict(updates) + updates['target_definitions'] = _normalize_target_definitions(updates['target_definitions']) + if 'total_count' in updates: + value = updates['total_count'] + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise DataRecordError('total_count must be a positive integer') + for key, value in updates.items(): + if key == 'target_definitions' or key == 'total_count': + continue + if not isinstance(value, str): + raise DataRecordError(f'{key} must be a string') + updates[key] = value.strip() + + plan.update(updates) + plan['status'] = 'pending_confirmation' + plan['revision'] = int(plan.get('revision') or 1) + 1 + history = plan.get('review_history') + if not isinstance(history, list): + history = [] + plan['review_history'] = history + history.append( + { + 'revision': plan['revision'], + 'feedback': review_feedback.strip(), + 'updates': updates, + 'updated_at_ms': int(time.time() * 1000), + } + ) + _save_plan_state(root, state) + return { + 'plan': plan, + 'message': '计划已根据用户反馈更新。请重新展示给用户 review，确认前不要生成数据。', + } + + +def confirm_generation_plan( + *, + root: str, + plan_id: str, + confirmation: str, + reviewed_revision: int | None = None, +) -> dict[str, Any]: + if not plan_id.strip(): + raise DataRecordError('plan_id is required') + if not _is_confirmation_text(confirmation): + raise DataRecordError('confirmation must clearly approve the plan') + state = _load_plan_state(root) + plan = _get_plan(state, plan_id) + current_revision = int(plan.get('revision') or 1) + if reviewed_revision is not None and reviewed_revision != current_revision: + raise DataRecordError( + f'reviewed_revision must match current plan revision {current_revision}' + ) + plan['status'] = 'confirmed' + plan['confirmed_at_ms'] = int(time.time() * 1000) + plan['confirmation'] = confirmation.strip() + plan['confirmed_revision'] = current_revision + _save_plan_state(root, state) + return { + 'plan': plan, + 'confirmed_plan_id': plan_id, + 'message': '计划已确认，可以把 confirmed_plan_id 传给 data_agent_normalize_dataset_draft。', + } + + +def validate_dataset_records(records: list[dict[str, Any]]) -> dict[str, Any]: + errors: list[dict[str, Any]] = [] + warnings: list[dict[str, Any]] = [] + + if not isinstance(records, list): + raise DataRecordError('records must be a list') + + seen_ids: set[str] = set() + for index, record in enumerate(records): + prefix = f'records[{index}]' + if not isinstance(record, dict): + errors.append(_issue(prefix, 'record must be an object')) + continue + + record_id = record.get('record_id') + if not _non_empty_str(record_id): + errors.append(_issue(f'{prefix}.record_id', 'record_id is required')) + elif record_id in seen_ids: + errors.append(_issue(f'{prefix}.record_id', 'record_id must be unique')) + else: + seen_ids.add(record_id) + + source = record.get('source') + if not isinstance(source, dict): + errors.append(_issue(f'{prefix}.source', 'source is required')) + else: + if source.get('type') not in {'generated', 'online', 'manual', 'mixed'}: + errors.append(_issue(f'{prefix}.source.type', 'source.type is invalid')) + if not _non_empty_str(source.get('request_id')): + errors.append(_issue(f'{prefix}.source.request_id', 'source.request_id is required')) + if not _int_like(source.get('timestamp')): + errors.append(_issue(f'{prefix}.source.timestamp', 'source.timestamp must be an integer')) + + turn = record.get('turn') + if not isinstance(turn, dict): + errors.append(_issue(f'{prefix}.turn', 'turn is required')) + turn_timestamp = None + else: + if not _non_empty_str(turn.get('query')): + errors.append(_issue(f'{prefix}.turn.query', 'turn.query is required')) + if 'tts' in turn: + errors.append(_issue(f'{prefix}.turn.tts', 'current turn must not contain tts')) + turn_timestamp = turn.get('timestamp') + if not _int_like(turn_timestamp): + errors.append(_issue(f'{prefix}.turn.timestamp', 'turn.timestamp must be an integer')) + + prev_session = record.get('prev_session') + if not isinstance(prev_session, list): + errors.append(_issue(f'{prefix}.prev_session', 'prev_session must be a list')) + prev_timestamps: list[int] = [] + else: + if len(prev_session) > 10: + errors.append(_issue(f'{prefix}.prev_session', 'prev_session must contain at most 10 turns')) + prev_timestamps = [] + for prev_index, item in enumerate(prev_session): + item_prefix = f'{prefix}.prev_session[{prev_index}]' + if not isinstance(item, dict): + errors.append(_issue(item_prefix, 'prev_session item must be an object')) + continue + if not _non_empty_str(item.get('query')): + errors.append(_issue(f'{item_prefix}.query', 'query is required')) + if not _non_empty_str(item.get('tts')): + errors.append(_issue(f'{item_prefix}.tts', 'tts is required')) + timestamp = item.get('timestamp') + if not _int_like(timestamp): + errors.append(_issue(f'{item_prefix}.timestamp', 'timestamp must be an integer')) + else: + prev_timestamps.append(int(timestamp)) + + if prev_timestamps != sorted(prev_timestamps): + errors.append(_issue(f'{prefix}.prev_session', 'timestamps must be sorted from early to late')) + if _int_like(turn_timestamp): + all_timestamps = [*prev_timestamps, int(turn_timestamp)] + for left, right in zip(all_timestamps, all_timestamps[1:]): + if right <= left: + errors.append(_issue(f'{prefix}.timestamp', 'turn timestamps must be strictly increasing')) + break + if right - left > 300_000: + warnings.append(_issue(f'{prefix}.timestamp', 'adjacent turns are more than 5 minutes apart')) + + context = record.get('context') + if not isinstance(context, dict): + errors.append(_issue(f'{prefix}.context', 'context must be an object')) + + label = record.get('label') + if not isinstance(label, dict): + errors.append(_issue(f'{prefix}.label', 'label is required')) + else: + if not _non_empty_str(label.get('dataset_label')): + warnings.append(_issue(f'{prefix}.label.dataset_label', 'dataset_label is empty')) + if not _non_empty_str(label.get('target')): + errors.append(_issue(f'{prefix}.label.target', 'label.target is required')) + if label.get('target_type') not in {'agent', 'function', 'unknown'}: + errors.append(_issue(f'{prefix}.label.target_type', 'label.target_type is invalid')) + + return { + 'ok': not errors, + 'error_count': len(errors), + 'warning_count': len(warnings), + 'errors': errors, + 'warnings': warnings, + } + + +def _parse_draft_text(draft_text: str) -> dict[str, Any]: + dataset_label = '' + cases: list[dict[str, Any]] = [] + current: dict[str, Any] | None = None + + for raw_line in draft_text.splitlines(): + line = raw_line.strip() + if not line: + continue + dataset_match = re.match(r'^#\s*dataset_label\s*[:：]\s*(.+)$', line, re.IGNORECASE) + if dataset_match: + dataset_label = dataset_match.group(1).strip() + continue + case_match = re.match(r'^#{3,}\s*case\s*[:：]\s*(.+)$', line, re.IGNORECASE) + if case_match: + current = {'case_name': case_match.group(1).strip(), 'turns': []} + cases.append(current) + continue + if current is None: + continue + + field_match = re.match(r'^(用户|小爱|target|notes|dataset_label|request_id|timestamp)\s*[:：]\s*(.*)$', line, re.IGNORECASE) + if not field_match: + continue + key = field_match.group(1).lower() + value = field_match.group(2).strip() + if key == '用户': + current['turns'].append({'role': 'user', 'text': value}) + elif key == '小爱': + current['turns'].append({'role': 'assistant', 'text': value}) + else: + current[key] = value + + if not cases: + raise DataRecordError('draft_text must contain at least one "### case:" block') + return {'dataset_label': dataset_label, 'cases': cases} + + +def _missing_plan_fields(fields: dict[str, str]) -> list[str]: + return [name for name, value in fields.items() if not isinstance(value, str) or not value.strip()] + + +def _dedupe_output_path(root: str, output_path: str, plan_id: str) -> str: + from pathlib import Path + + path = Path(output_path) + resolved = path if path.is_absolute() else Path(root).resolve() / path + if not resolved.exists(): + return output_path + parts = path.parts + if len(parts) >= 2 and parts[-1] == 'artifacts': + task_dir = Path(*parts[:-1]) + deduped = Path(f'{task_dir}-{plan_id}') / 'artifacts' + else: + deduped = Path(f'{output_path.rstrip("/")}-{plan_id}') + return deduped.as_posix() + + +def _normalize_target_definitions(value: list[dict[str, Any]] | None) -> list[dict[str, str]]: + if value is None: + return [] + if not isinstance(value, list): + raise DataRecordError('target_definitions must be a list') + normalized: list[dict[str, str]] = [] + for index, item in enumerate(value): + if not isinstance(item, dict): + raise DataRecordError(f'target_definitions[{index}] must be an object') + name = str(item.get('name') or '').strip() + target = str(item.get('target') or '').strip() + rule = str(item.get('rule') or '').strip() + if not target: + raise DataRecordError(f'target_definitions[{index}].target is required') + normalized.append({'name': name, 'target': target, 'rule': rule}) + return normalized + + +def _allowed_targets_for_plan(plan: dict[str, Any]) -> set[str]: + target_definitions = plan.get('target_definitions') + allowed: set[str] = set() + if isinstance(target_definitions, list): + for item in target_definitions: + if isinstance(item, dict) and isinstance(item.get('target'), str) and item['target'].strip(): + allowed.add(item['target'].strip()) + target = plan.get('target') + if isinstance(target, str) and target.strip(): + allowed.add(target.strip()) + return allowed + + +def _validate_records_against_plan(records: list[dict[str, Any]], plan: dict[str, Any]) -> None: + allowed_targets = _allowed_targets_for_plan(plan) + if not allowed_targets: + return + unexpected: list[str] = [] + for index, record in enumerate(records): + label = record.get('label') + target = label.get('target') if isinstance(label, dict) else None + if isinstance(target, str) and target.strip() in allowed_targets: + continue + unexpected.append(f'records[{index}].label.target={target!r}') + if unexpected: + raise DataRecordError( + 'record targets must match the confirmed generation plan: ' + + '; '.join(unexpected) + ) + + +def _is_confirmation_text(text: str) -> bool: + if not isinstance(text, str): + return False + normalized = text.strip().lower() + if not normalized: + return False + approval_words = ( + '确认', + '同意', + '开始生成', + '按这个计划', + '批准', + 'approve', + 'approved', + 'confirmed', + 'yes', + ) + return any(word in normalized for word in approval_words) + + +def _plan_state_path(root: str): + from pathlib import Path + + return Path(root).resolve() / DEFAULT_PLAN_STATE_FILE + + +def _load_plan_state(root: str) -> dict[str, Any]: + path = _plan_state_path(root) + if not path.exists(): + return {'next_sequence': 1, 'plans': {}} + try: + payload = json.loads(path.read_text(encoding='utf-8')) + except json.JSONDecodeError as exc: + raise DataRecordError(f'invalid generation plan state file: {path}') from exc + if not isinstance(payload, dict): + raise DataRecordError(f'invalid generation plan state file: {path}') + if not isinstance(payload.get('plans'), dict): + payload['plans'] = {} + if not isinstance(payload.get('next_sequence'), int): + payload['next_sequence'] = 1 + return payload + + +def _save_plan_state(root: str, state: dict[str, Any]) -> None: + path = _plan_state_path(root) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding='utf-8') + + +def _get_plan(state: dict[str, Any], plan_id: str) -> dict[str, Any]: + plans = state.get('plans') + if not isinstance(plans, dict): + raise DataRecordError('generation plan state is invalid') + plan = plans.get(plan_id) + if not isinstance(plan, dict): + raise DataRecordError(f'generation plan not found: {plan_id}') + return plan + + +def _require_confirmed_generation_plan(root: str, plan_id: str) -> dict[str, Any]: + state = _load_plan_state(root) + plan = _get_plan(state, plan_id) + if plan.get('status') != 'confirmed': + raise DataRecordError(f'generation plan is not confirmed: {plan_id}') + return plan + + +def _case_to_record( + case: dict[str, Any], + *, + global_dataset_label: str, + index: int, + batch_id: str, + source_type: str, + base_timestamp: int, + timestamp_step_ms: int, + default_request_id: str, +) -> tuple[dict[str, Any], list[str]]: + warnings: list[str] = [] + turns = case.get('turns') + if not isinstance(turns, list): + raise DataRecordError(f'case {index} has invalid turns') + + user_indexes = [turn_index for turn_index, turn in enumerate(turns) if turn.get('role') == 'user'] + if not user_indexes: + raise DataRecordError(f'case {index} must contain at least one 用户 line') + final_user_index = user_indexes[-1] + current_query = str(turns[final_user_index].get('text') or '').strip() + if not current_query: + raise DataRecordError(f'case {index} current 用户 line must be non-empty') + + target = str(case.get('target') or '').strip() + if not target: + raise DataRecordError(f'case {index} target is required') + + prev_session = _build_prev_session(turns[:final_user_index], case_index=index) + if len(prev_session) > 10: + warnings.append(f'case {index} prev_session has more than 10 turns; keeping the latest 10') + prev_session = prev_session[-10:] + + request_id = str(case.get('request_id') or default_request_id).strip() + case_timestamp = _optional_int(case.get('timestamp')) + turn_timestamp = case_timestamp if case_timestamp is not None else base_timestamp + (index - 1) * timestamp_step_ms * 20 + first_prev_timestamp = turn_timestamp - len(prev_session) * timestamp_step_ms + for prev_index, item in enumerate(prev_session): + item['timestamp'] = first_prev_timestamp + prev_index * timestamp_step_ms + + dataset_label = str(case.get('dataset_label') or global_dataset_label or '').strip() + record = { + 'record_id': _record_id(source_type=source_type, batch_id=batch_id, request_id=request_id, index=index), + 'source': { + 'type': source_type, + 'request_id': request_id, + 'timestamp': turn_timestamp, + }, + 'turn': { + 'query': current_query, + 'timestamp': turn_timestamp, + }, + 'prev_session': prev_session, + 'context': {}, + 'label': { + 'dataset_label': dataset_label, + 'target': target, + 'target_type': _target_type(target), + }, + 'meta': { + 'case_name': str(case.get('case_name') or '').strip(), + 'notes': str(case.get('notes') or '').strip(), + }, + } + return record, warnings + + +def _build_prev_session(turns: list[dict[str, Any]], *, case_index: int) -> list[dict[str, Any]]: + prev_session: list[dict[str, Any]] = [] + index = 0 + while index < len(turns): + turn = turns[index] + if turn.get('role') != 'user': + raise DataRecordError(f'case {case_index} prev_session must start with 用户 before 小爱') + if index + 1 >= len(turns) or turns[index + 1].get('role') != 'assistant': + raise DataRecordError(f'case {case_index} each previous 用户 line must be followed by 小爱') + query = str(turn.get('text') or '').strip() + tts = str(turns[index + 1].get('text') or '').strip() + if not query or not tts: + raise DataRecordError(f'case {case_index} previous 用户/小爱 lines must be non-empty') + prev_session.append({'query': query, 'tts': tts}) + index += 2 + return prev_session + + +def _target_type(target: str) -> str: + if re.match(r'^Agent\s*\(\s*tag\s*=', target): + return 'agent' + if target: + return 'function' + return 'unknown' + + +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}' + if source_type == 'online': + return f'online_{_safe_token(request_id)}_{index:06d}' + return f'{_safe_token(source_type)}_{_safe_token(batch_id)}_{index:06d}' + + +def _safe_token(value: str) -> str: + token = re.sub(r'[^A-Za-z0-9_.-]+', '_', value.strip()) + return token.strip('_') or 'unknown' + + +def _optional_int(value: Any) -> int | None: + if value is None or value == '': + return None + if isinstance(value, bool): + raise DataRecordError('timestamp must be an integer') + try: + return int(value) + except (TypeError, ValueError) as exc: + raise DataRecordError('timestamp must be an integer') from exc + + +def _non_empty_str(value: Any) -> bool: + return isinstance(value, str) and bool(value.strip()) + + +def _int_like(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _issue(path: str, message: str) -> dict[str, str]: + return {'path': path, 'message': message} + + +def records_from_tool_argument(value: Any) -> list[dict[str, Any]]: + if isinstance(value, str): + decoded = json.loads(value) + else: + decoded = value + if not isinstance(decoded, list): + raise DataRecordError('records must be a JSON array') + return decoded diff --git a/tests/test_data_agent_records.py b/tests/test_data_agent_records.py new file mode 100644 index 0000000..bfcdd73 --- /dev/null +++ b/tests/test_data_agent_records.py @@ -0,0 +1,359 @@ +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + +from src.agent_tools import build_tool_context, default_tool_registry, execute_tool +from src.agent_types import AgentRuntimeConfig +from src.data_agent_records import ( + confirm_generation_plan, + get_generation_plan, + normalize_dataset_draft, + prepare_generation_plan, + update_generation_plan, + validate_dataset_records, +) + + +class DataAgentRecordTests(unittest.TestCase): + def test_normalize_dataset_draft_builds_canonical_records(self) -> None: + draft = ''' +# dataset_label: 地图和生活边界数据 + +### case: 附近餐饮查询 +用户: 帮我看看附近有什么好吃的 +target: Agent(tag="life_service") +notes: 附近生活服务查询 + +### case: 多轮承接附近餐饮 +用户: 我想出门逛逛 +小爱: 好的 +用户: 看看附近有什么好吃的 +target: Agent(tag="life_service") +notes: 多轮承接附近生活服务查询 +'''.strip() + + payload = normalize_dataset_draft( + draft, + batch_id='demo', + base_timestamp=1_755_567_930_500, + timestamp_step_ms=60_000, + ) + + self.assertEqual(payload['warnings'], []) + records = payload['records'] + self.assertEqual(len(records), 2) + self.assertEqual(records[0]['record_id'], 'gen_demo_000001') + self.assertEqual(records[0]['turn']['query'], '帮我看看附近有什么好吃的') + self.assertEqual(records[0]['prev_session'], []) + self.assertEqual(records[0]['label']['target_type'], 'agent') + self.assertEqual(records[1]['turn']['query'], '看看附近有什么好吃的') + self.assertEqual( + records[1]['prev_session'], + [ + { + 'query': '我想出门逛逛', + 'tts': '好的', + 'timestamp': 1_755_569_070_500, + } + ], + ) + + def test_validate_dataset_records_reports_valid_payload(self) -> None: + records = normalize_dataset_draft( + ''' +# dataset_label: 地图和生活边界数据 +### case: 附近餐饮查询 +用户: 附近有什么好吃的 +target: Agent(tag="life_service") +'''.strip(), + base_timestamp=1_755_567_930_500, + )['records'] + + result = validate_dataset_records(records) + + self.assertTrue(result['ok']) + self.assertEqual(result['error_count'], 0) + + def test_validate_dataset_records_rejects_current_turn_tts(self) -> None: + records = normalize_dataset_draft( + ''' +# dataset_label: 地图和生活边界数据 +### case: 附近餐饮查询 +用户: 附近有什么好吃的 +target: Agent(tag="life_service") +'''.strip(), + base_timestamp=1_755_567_930_500, + )['records'] + records[0]['turn']['tts'] = '不应该出现' + + result = validate_dataset_records(records) + + self.assertFalse(result['ok']) + self.assertEqual(result['errors'][0]['path'], 'records[0].turn.tts') + + def test_normalize_online_draft_uses_real_request_metadata(self) -> None: + records = normalize_dataset_draft( + ''' +# dataset_label: 线上误召回badcase专项 +### case: 线上样例 +request_id: rid-123 +timestamp: 1755567930500 +用户: 附近有什么好吃的 +target: Agent(tag="life_service") +'''.strip(), + source_type='online', + )['records'] + + self.assertEqual(records[0]['record_id'], 'online_rid-123_000001') + self.assertEqual(records[0]['source']['request_id'], 'rid-123') + self.assertEqual(records[0]['source']['timestamp'], 1_755_567_930_500) + + def test_normalize_requires_confirmed_plan_when_state_root_is_provided(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + with self.assertRaisesRegex(Exception, 'confirmed_plan_id is required'): + normalize_dataset_draft( + ''' +# dataset_label: 地图和生活边界数据 +### case: 附近餐饮查询 +用户: 附近有什么好吃的 +target: Agent(tag="life_service") +'''.strip(), + plan_state_root=tmp_dir, + ) + + def test_generation_plan_can_be_confirmed_then_used(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + plan_payload = prepare_generation_plan( + root=tmp_dir, + dataset_label='地图和生活边界数据', + target='Agent(tag="life_service")', + total_count=2, + turn_mix='1 条单轮，1 条多轮', + coverage='附近吃喝玩乐', + exclusions='不要生成导航路线类 query', + output_path='tasks/demo/artifacts', + ) + plan_id = plan_payload['plan']['plan_id'] + confirmed = confirm_generation_plan(root=tmp_dir, plan_id=plan_id, confirmation='确认，开始生成') + records = normalize_dataset_draft( + ''' +# dataset_label: 地图和生活边界数据 +### case: 附近餐饮查询 +用户: 附近有什么好吃的 +target: Agent(tag="life_service") +'''.strip(), + confirmed_plan_id=confirmed['confirmed_plan_id'], + plan_state_root=tmp_dir, + )['records'] + + self.assertEqual(records[0]['label']['dataset_label'], '地图和生活边界数据') + + def test_generation_plan_dedupes_existing_task_output_path(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + existing = Path(tmp_dir) / 'tasks' / 'demo' / 'artifacts' + existing.mkdir(parents=True) + plan_payload = prepare_generation_plan( + root=tmp_dir, + dataset_label='地图和生活边界数据', + target='Agent(tag="life_service")', + total_count=2, + turn_mix='1 条单轮，1 条多轮', + coverage='附近吃喝玩乐', + exclusions='不要生成导航路线类 query', + output_path='tasks/demo/artifacts', + ) + + self.assertEqual( + plan_payload['plan']['output_path'], + 'tasks/demo-data_plan_000001/artifacts', + ) + self.assertEqual(plan_payload['plan']['requested_output_path'], 'tasks/demo/artifacts') + + def test_generation_plan_review_update_and_revision_confirmation(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + plan_payload = prepare_generation_plan( + root=tmp_dir, + dataset_label='地图和生活边界数据', + target='Agent(tag="life_service")', + total_count=2, + turn_mix='1 条单轮，1 条多轮', + coverage='附近吃喝玩乐', + exclusions='不要生成导航路线类 query', + output_path='tasks/demo/artifacts', + ) + plan_id = plan_payload['plan']['plan_id'] + updated = update_generation_plan( + root=tmp_dir, + plan_id=plan_id, + review_feedback='多轮数据多一点', + updates={'turn_mix': '1 条单轮，3 条多轮', 'total_count': 4}, + ) + + self.assertEqual(updated['plan']['revision'], 2) + self.assertEqual(updated['plan']['turn_mix'], '1 条单轮，3 条多轮') + self.assertEqual(len(updated['plan']['review_history']), 1) + shown = get_generation_plan(root=tmp_dir, plan_id=plan_id) + self.assertEqual(shown['plan']['revision'], 2) + with self.assertRaisesRegex(Exception, 'reviewed_revision must match'): + confirm_generation_plan( + root=tmp_dir, + plan_id=plan_id, + confirmation='确认，开始生成', + reviewed_revision=1, + ) + confirmed = confirm_generation_plan( + root=tmp_dir, + plan_id=plan_id, + confirmation='确认，开始生成', + reviewed_revision=2, + ) + + self.assertEqual(confirmed['plan']['status'], 'confirmed') + self.assertEqual(confirmed['plan']['confirmed_revision'], 2) + + def test_multi_target_plan_allows_only_declared_targets(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + plan_payload = prepare_generation_plan( + root=tmp_dir, + dataset_label='餐饮和导航边界数据', + target='', + target_definitions=[ + { + 'name': '餐饮服务', + 'target': 'Agent(tag="餐饮服务")', + 'rule': '找附近美食但不导航', + }, + { + 'name': '地图导航', + 'target': 'Agent(tag="地图导航")', + 'rule': '明确要求导航去某地', + }, + ], + total_count=2, + turn_mix='2 条单轮', + coverage='餐饮服务和地图导航边界', + exclusions='不要生成无关闲聊', + output_path='tasks/demo/artifacts', + ) + plan_id = plan_payload['plan']['plan_id'] + confirm_generation_plan(root=tmp_dir, plan_id=plan_id, confirmation='确认，开始生成') + records = normalize_dataset_draft( + ''' +# dataset_label: 餐饮和导航边界数据 +### case: 找奶茶 +用户: 附近有没有奶茶店 +target: Agent(tag="餐饮服务") +### case: 导航去奶茶店 +用户: 导航去最近的奶茶店 +target: Agent(tag="地图导航") +'''.strip(), + confirmed_plan_id=plan_id, + plan_state_root=tmp_dir, + )['records'] + + self.assertEqual(records[0]['label']['target'], 'Agent(tag="餐饮服务")') + self.assertEqual(records[1]['label']['target'], 'Agent(tag="地图导航")') + + def test_multi_target_plan_rejects_undeclared_targets(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + plan_payload = prepare_generation_plan( + root=tmp_dir, + dataset_label='餐饮和导航边界数据', + target='', + target_definitions=[ + {'target': 'Agent(tag="餐饮服务")'}, + {'target': 'Agent(tag="地图导航")'}, + ], + total_count=1, + turn_mix='1 条单轮', + coverage='餐饮服务和地图导航边界', + exclusions='不要生成无关闲聊', + output_path='tasks/demo/artifacts', + ) + plan_id = plan_payload['plan']['plan_id'] + confirm_generation_plan(root=tmp_dir, plan_id=plan_id, confirmation='确认，开始生成') + with self.assertRaisesRegex(Exception, 'record targets must match'): + normalize_dataset_draft( + ''' +# dataset_label: 餐饮和导航边界数据 +### case: 天气 +用户: 明天天气怎么样 +target: Agent(tag="天气") +'''.strip(), + confirmed_plan_id=plan_id, + plan_state_root=tmp_dir, + ) + + def test_tools_execute_against_registry(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + context = build_tool_context(AgentRuntimeConfig(cwd=Path(tmp_dir))) + plan_result = execute_tool( + default_tool_registry(), + 'data_agent_prepare_generation_plan', + { + 'dataset_label': '地图和生活边界数据', + 'target': 'Agent(tag="life_service")', + 'total_count': 1, + 'turn_mix': '1 条单轮', + 'coverage': '附近吃喝玩乐', + 'exclusions': '不要生成导航路线类 query', + 'output_path': 'tasks/demo/artifacts', + }, + context, + ) + self.assertTrue(plan_result.ok) + plan_id = json.loads(plan_result.content)['plan']['plan_id'] + update_result = execute_tool( + default_tool_registry(), + 'data_agent_update_generation_plan', + { + 'plan_id': plan_id, + 'review_feedback': '多轮不需要，先只测单轮', + 'updates': {'turn_mix': '1 条单轮'}, + }, + context, + ) + self.assertTrue(update_result.ok) + revision = json.loads(update_result.content)['plan']['revision'] + confirm_result = execute_tool( + default_tool_registry(), + 'data_agent_confirm_generation_plan', + {'plan_id': plan_id, 'confirmation': '确认，开始生成', 'reviewed_revision': revision}, + context, + ) + self.assertTrue(confirm_result.ok) + normalize_result = execute_tool( + default_tool_registry(), + 'data_agent_normalize_dataset_draft', + { + 'draft_text': ''' +# dataset_label: 地图和生活边界数据 +### case: 附近餐饮查询 +用户: 附近有什么好吃的 +target: Agent(tag="life_service") +'''.strip(), + 'batch_id': 'demo', + 'base_timestamp': 1_755_567_930_500, + 'confirmed_plan_id': plan_id, + }, + context, + ) + self.assertTrue(normalize_result.ok) + records = json.loads(normalize_result.content)['records'] + + validate_result = execute_tool( + default_tool_registry(), + 'data_agent_validate_dataset_records', + {'records': records}, + context, + ) + + self.assertTrue(validate_result.ok) + self.assertTrue(json.loads(validate_result.content)['ok']) + + +if __name__ == '__main__': + unittest.main()