Add data agent input and export workflow

src/data_agent_records.py

@@ -10,18 +10,109 @@ Skill 应该先让模型生成紧凑、便于人工 review 的 dataset draft tex
 import json
 import re
 import time
+from pathlib import Path
 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'
+DEFAULT_GOAL_STATE_FILE = '.port_sessions/data_agent_generation_goals.json'
 
 
 class DataRecordError(ValueError):
     """Raised when dataset draft text or records are structurally invalid."""
 
 
+def prepare_generation_goal(
+    *,
+    root: str,
+    dataset_label: str,
+    goal_summary: str,
+    coverage: str,
+    exclusions: str,
+    target: str = '',
+    target_definitions: list[dict[str, Any]] | None = None,
+    plan_hint: str = '',
+    open_questions: list[str] | None = None,
+    source_refs: list[str] | None = None,
+    notes: str = '',
+) -> dict[str, Any]:
+    missing = _missing_plan_fields(
+        {
+            'dataset_label': dataset_label,
+            'goal_summary': goal_summary,
+            'coverage': coverage,
+            'exclusions': exclusions,
+        }
+    )
+    if missing:
+        raise DataRecordError('missing required goal fields: ' + ', '.join(missing))
+    normalized_targets = _normalize_target_definitions(target_definitions)
+    state = _load_goal_state(root)
+    sequence = int(state.get('next_sequence') or 1)
+    goal_id = f'data_goal_{sequence:06d}'
+    goal = {
+        'goal_id': goal_id,
+        'status': 'pending_confirmation',
+        'dataset_label': dataset_label.strip(),
+        'goal_summary': goal_summary.strip(),
+        'target': target.strip(),
+        'target_definitions': normalized_targets,
+        'plan_hint': plan_hint.strip(),
+        'coverage': coverage.strip(),
+        'exclusions': exclusions.strip(),
+        'open_questions': _normalize_string_list(open_questions, 'open_questions'),
+        'source_refs': _normalize_string_list(source_refs, 'source_refs'),
+        'notes': notes.strip(),
+        'revision': 1,
+        'created_at_ms': int(time.time() * 1000),
+    }
+    goals = state.setdefault('goals', {})
+    if not isinstance(goals, dict):
+        goals = {}
+        state['goals'] = goals
+    goals[goal_id] = goal
+    state['next_sequence'] = sequence + 1
+    _save_goal_state(root, state)
+    return {
+        'goal': goal,
+        'message': 'generation goal 已创建，状态为 pending_confirmation。请向用户展示 goal，并等待用户确认后再创建 generation plan。',
+    }
+
+
+def confirm_generation_goal(
+    *,
+    root: str,
+    goal_id: str,
+    confirmation: str,
+    reviewed_revision: int | None = None,
+) -> dict[str, Any]:
+    if not goal_id.strip():
+        raise DataRecordError('goal_id is required')
+    if not _is_confirmation_text(confirmation):
+        raise DataRecordError('confirmation must clearly approve the generation goal')
+    state = _load_goal_state(root)
+    goal = _get_goal(state, goal_id)
+    current_revision = int(goal.get('revision') or 1)
+    if reviewed_revision is not None and reviewed_revision != current_revision:
+        raise DataRecordError(
+            f'reviewed_revision must match current goal revision {current_revision}'
+        )
+    if not _allowed_targets_for_goal(goal):
+        raise DataRecordError('confirmed generation goal must include target or target_definitions')
+    goal['status'] = 'confirmed'
+    goal['confirmed_at_ms'] = int(time.time() * 1000)
+    goal['confirmation'] = confirmation.strip()
+    goal['confirmed_revision'] = current_revision
+    _save_goal_state(root, state)
+    return {
+        'goal': goal,
+        'confirmed_goal_id': goal_id,
+        'message': 'generation goal 已确认，可以把 confirmed_goal_id 传给 data_agent_prepare_generation_plan。',
+    }
+
+
 def normalize_dataset_draft(
     draft_text: str,
     *,
@@ -80,7 +171,11 @@ def prepare_generation_plan(
     output_path: str,
     target_definitions: list[dict[str, Any]] | None = None,
     notes: str = '',
+    confirmed_goal_id: str | None = None,
 ) -> dict[str, Any]:
+    confirmed_goal: dict[str, Any] | None = None
+    if confirmed_goal_id is not None:
+        confirmed_goal = _require_confirmed_generation_goal(root, confirmed_goal_id)
     missing = _missing_plan_fields(
         {
             'dataset_label': dataset_label,
@@ -97,6 +192,13 @@ def prepare_generation_plan(
         raise DataRecordError('missing required plan fields: ' + ', '.join(missing))
     if total_count <= 0:
         raise DataRecordError('total_count must be greater than 0')
+    if confirmed_goal is not None:
+        _validate_plan_against_goal(
+            dataset_label=dataset_label,
+            target=target,
+            target_definitions=normalized_targets,
+            goal=confirmed_goal,
+        )
 
     state = _load_plan_state(root)
     sequence = int(state.get('next_sequence') or 1)
@@ -115,6 +217,7 @@ def prepare_generation_plan(
         'output_path': resolved_output_path,
         'requested_output_path': output_path.strip(),
         'notes': notes.strip(),
+        'confirmed_goal_id': confirmed_goal_id or '',
         'revision': 1,
         'review_history': [],
         'created_at_ms': int(time.time() * 1000),
@@ -342,6 +445,56 @@ def validate_dataset_records(records: list[dict[str, Any]]) -> dict[str, Any]:
     }
 
 
+def export_dataset_records(
+    records: list[dict[str, Any]],
+    *,
+    root: str,
+    output_path: str,
+    output_format: str = 'jsonl',
+    require_validation_ok: bool = True,
+    overwrite: bool = True,
+) -> dict[str, Any]:
+    """校验并导出 canonical records。
+
+    这个函数用于收口数据智能体的最终落盘动作，避免模型手写 JSON/JSONL。
+    默认导出紧凑 JSONL：一行一个 canonical record，方便 diff、抽样和增量处理。
+    """
+
+    if output_format not in {'jsonl', 'json'}:
+        raise DataRecordError('output_format must be jsonl or json')
+    if not output_path.strip():
+        raise DataRecordError('output_path is required')
+
+    validation = validate_dataset_records(records)
+    if require_validation_ok and not validation['ok']:
+        raise DataRecordError(
+            f'records failed validation with {validation["error_count"]} errors'
+        )
+
+    path = _resolve_output_path(root, output_path)
+    if path.exists() and not overwrite:
+        raise DataRecordError(f'output_path already exists: {output_path}')
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    if output_format == 'jsonl':
+        content = ''.join(
+            json.dumps(record, ensure_ascii=False, separators=(',', ':')) + '\n'
+            for record in records
+        )
+    else:
+        content = json.dumps(records, ensure_ascii=False, separators=(',', ':')) + '\n'
+
+    path.write_text(content, encoding='utf-8')
+    rel_path = str(path.relative_to(Path(root).resolve()))
+    return {
+        'output_path': rel_path,
+        'output_format': output_format,
+        'record_count': len(records),
+        'bytes_written': len(content.encode('utf-8')),
+        'validation': validation,
+    }
+
+
 def _parse_draft_text(draft_text: str) -> dict[str, Any]:
     dataset_label = ''
     cases: list[dict[str, Any]] = []
@@ -418,6 +571,25 @@ def _normalize_target_definitions(value: list[dict[str, Any]] | None) -> list[di
     return normalized
 
 
+def _normalize_string_list(value: list[str] | None, field_name: str) -> list[str]:
+    if value is None:
+        return []
+    if not isinstance(value, list):
+        raise DataRecordError(f'{field_name} must be a list')
+    normalized: list[str] = []
+    for index, item in enumerate(value):
+        if not isinstance(item, str):
+            raise DataRecordError(f'{field_name}[{index}] must be a string')
+        text = item.strip()
+        if text:
+            normalized.append(text)
+    return normalized
+
+
+def _allowed_targets_for_goal(goal: dict[str, Any]) -> set[str]:
+    return _allowed_targets_for_plan(goal)
+
+
 def _allowed_targets_for_plan(plan: dict[str, Any]) -> set[str]:
     target_definitions = plan.get('target_definitions')
     allowed: set[str] = set()
@@ -431,6 +603,33 @@ def _allowed_targets_for_plan(plan: dict[str, Any]) -> set[str]:
     return allowed
 
 
+def _validate_plan_against_goal(
+    *,
+    dataset_label: str,
+    target: str,
+    target_definitions: list[dict[str, str]],
+    goal: dict[str, Any],
+) -> None:
+    goal_label = str(goal.get('dataset_label') or '').strip()
+    if goal_label and dataset_label.strip() != goal_label:
+        raise DataRecordError(
+            f'plan dataset_label must match confirmed generation goal: {goal_label}'
+        )
+    goal_targets = _allowed_targets_for_goal(goal)
+    plan_targets = set()
+    if target.strip():
+        plan_targets.add(target.strip())
+    for item in target_definitions:
+        item_target = item.get('target', '').strip()
+        if item_target:
+            plan_targets.add(item_target)
+    unexpected = sorted(plan_targets - goal_targets)
+    if unexpected:
+        raise DataRecordError(
+            'plan targets must match confirmed generation goal: ' + ', '.join(unexpected)
+        )
+
+
 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:
@@ -475,6 +674,12 @@ def _plan_state_path(root: str):
     return Path(root).resolve() / DEFAULT_PLAN_STATE_FILE
 
 
+def _goal_state_path(root: str):
+    from pathlib import Path
+
+    return Path(root).resolve() / DEFAULT_GOAL_STATE_FILE
+
+
 def _load_plan_state(root: str) -> dict[str, Any]:
     path = _plan_state_path(root)
     if not path.exists():
@@ -492,12 +697,35 @@ def _load_plan_state(root: str) -> dict[str, Any]:
     return payload
 
 
+def _load_goal_state(root: str) -> dict[str, Any]:
+    path = _goal_state_path(root)
+    if not path.exists():
+        return {'next_sequence': 1, 'goals': {}}
+    try:
+        payload = json.loads(path.read_text(encoding='utf-8'))
+    except json.JSONDecodeError as exc:
+        raise DataRecordError(f'invalid generation goal state file: {path}') from exc
+    if not isinstance(payload, dict):
+        raise DataRecordError(f'invalid generation goal state file: {path}')
+    if not isinstance(payload.get('goals'), dict):
+        payload['goals'] = {}
+    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 _save_goal_state(root: str, state: dict[str, Any]) -> None:
+    path = _goal_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):
@@ -508,6 +736,16 @@ def _get_plan(state: dict[str, Any], plan_id: str) -> dict[str, Any]:
     return plan
 
 
+def _get_goal(state: dict[str, Any], goal_id: str) -> dict[str, Any]:
+    goals = state.get('goals')
+    if not isinstance(goals, dict):
+        raise DataRecordError('generation goal state is invalid')
+    goal = goals.get(goal_id)
+    if not isinstance(goal, dict):
+        raise DataRecordError(f'generation goal not found: {goal_id}')
+    return goal
+
+
 def _require_confirmed_generation_plan(root: str, plan_id: str) -> dict[str, Any]:
     state = _load_plan_state(root)
     plan = _get_plan(state, plan_id)
@@ -516,6 +754,14 @@ def _require_confirmed_generation_plan(root: str, plan_id: str) -> dict[str, Any
     return plan
 
 
+def _require_confirmed_generation_goal(root: str, goal_id: str) -> dict[str, Any]:
+    state = _load_goal_state(root)
+    goal = _get_goal(state, goal_id)
+    if goal.get('status') != 'confirmed':
+        raise DataRecordError(f'generation goal is not confirmed: {goal_id}')
+    return goal
+
+
 def _case_to_record(
     case: dict[str, Any],
     *,
@@ -622,6 +868,18 @@ def _safe_token(value: str) -> str:
     return token.strip('_') or 'unknown'
 
 
+def _resolve_output_path(root: str, output_path: str) -> Path:
+    root_path = Path(root).resolve()
+    candidate = Path(output_path).expanduser()
+    path = candidate if candidate.is_absolute() else root_path / candidate
+    resolved = path.resolve(strict=False)
+    try:
+        resolved.relative_to(root_path)
+    except ValueError as exc:
+        raise DataRecordError(f'output_path escapes workspace root: {output_path}') from exc
+    return resolved
+
+
 def _optional_int(value: Any) -> int | None:
     if value is None or value == '':
         return None