wuyang6
458 lines
17 KiB
Python
458 lines
17 KiB
Python
from __future__ import annotations
|
||
|
||
"""product-data 可迁移脚本的共享逻辑。
|
||
|
||
这里不依赖 ZK Data Agent 的运行时，便于把整个 skill 复制到其他 Agent
|
||
或普通 Python 环境中使用。平台注册工具时，也应该只薄封装这些脚本。
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
|
||
DEFAULT_REQUEST_ID = "aabbccdd"
|
||
DEFAULT_TIMESTAMP_STEP_MS = 60_000
|
||
|
||
|
||
class ProductDataError(ValueError):
|
||
"""数据格式或参数错误。"""
|
||
|
||
|
||
def load_json_payload(argv: list[str] | None = None) -> dict[str, Any]:
|
||
parser = argparse.ArgumentParser(add_help=True)
|
||
parser.add_argument("--input", "-i", help="JSON 参数文件；不传则从 stdin 读取")
|
||
args = parser.parse_args(argv)
|
||
|
||
if args.input:
|
||
text = Path(args.input).expanduser().read_text(encoding="utf-8")
|
||
else:
|
||
text = sys.stdin.read()
|
||
if not text.strip():
|
||
raise ProductDataError("input JSON is required")
|
||
payload = json.loads(text)
|
||
if not isinstance(payload, dict):
|
||
raise ProductDataError("input JSON must be an object")
|
||
return payload
|
||
|
||
|
||
def emit_success(payload: dict[str, Any]) -> None:
|
||
result = {"ok": True, **payload}
|
||
print(json.dumps(result, ensure_ascii=False, separators=(",", ":")))
|
||
|
||
|
||
def emit_error(exc: BaseException) -> None:
|
||
print(
|
||
json.dumps(
|
||
{"ok": False, "error": str(exc)},
|
||
ensure_ascii=False,
|
||
separators=(",", ":"),
|
||
)
|
||
)
|
||
|
||
|
||
def load_records(payload: dict[str, Any]) -> list[dict[str, Any]]:
|
||
if "records" in payload:
|
||
records = payload["records"]
|
||
elif "records_path" in payload:
|
||
records = read_records_file(str(payload["records_path"]))
|
||
else:
|
||
raise ProductDataError("records or records_path is required")
|
||
if isinstance(records, str):
|
||
records = json.loads(records)
|
||
if not isinstance(records, list):
|
||
raise ProductDataError("records must be an array")
|
||
normalized: list[dict[str, Any]] = []
|
||
for index, item in enumerate(records):
|
||
if not isinstance(item, dict):
|
||
raise ProductDataError(f"records[{index}] must be an object")
|
||
normalized.append(item)
|
||
return normalized
|
||
|
||
|
||
def read_records_file(path: str) -> list[dict[str, Any]]:
|
||
file_path = Path(path).expanduser()
|
||
text = file_path.read_text(encoding="utf-8")
|
||
if file_path.suffix.lower() == ".jsonl":
|
||
records: list[dict[str, Any]] = []
|
||
for line_number, line in enumerate(text.splitlines(), start=1):
|
||
if not line.strip():
|
||
continue
|
||
item = json.loads(line)
|
||
if not isinstance(item, dict):
|
||
raise ProductDataError(f"{path}:{line_number} must be a JSON object")
|
||
records.append(item)
|
||
return records
|
||
decoded = json.loads(text)
|
||
if not isinstance(decoded, list):
|
||
raise ProductDataError(f"{path} must contain a JSON array")
|
||
return decoded
|
||
|
||
|
||
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,
|
||
) -> dict[str, Any]:
|
||
if source_type not in {"generated", "online", "manual", "mixed"}:
|
||
raise ProductDataError("source_type is invalid")
|
||
if not draft_text.strip():
|
||
raise ProductDataError("draft_text must be non-empty")
|
||
if timestamp_step_ms <= 0:
|
||
raise ProductDataError("timestamp_step_ms must be greater than 0")
|
||
|
||
base = int(time.time() * 1000) if base_timestamp is None else int(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=str(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)
|
||
|
||
return {"records": records, "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.I)
|
||
if dataset_match:
|
||
dataset_label = dataset_match.group(1).strip()
|
||
continue
|
||
case_match = re.match(r"^#{3,}\s*case\s*[:：]\s*(.+)$", line, re.I)
|
||
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.I,
|
||
)
|
||
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 ProductDataError('draft_text must contain at least one "### case:" block')
|
||
return {"dataset_label": dataset_label, "cases": cases}
|
||
|
||
|
||
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 ProductDataError(f"case {index} has invalid turns")
|
||
|
||
user_indexes = [
|
||
turn_index
|
||
for turn_index, turn in enumerate(turns)
|
||
if isinstance(turn, dict) and turn.get("role") == "user"
|
||
]
|
||
if not user_indexes:
|
||
raise ProductDataError(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 ProductDataError(f"case {index} current 用户 line must be non-empty")
|
||
|
||
target = str(case.get("target") or "").strip()
|
||
if not target:
|
||
raise ProductDataError(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()
|
||
turn_timestamp = optional_int(case.get("timestamp"))
|
||
if turn_timestamp is None:
|
||
turn_timestamp = 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 not isinstance(turn, dict) or turn.get("role") != "user":
|
||
raise ProductDataError(f"case {case_index} prev_session must start with 用户 before 小爱")
|
||
if index + 1 >= len(turns) or turns[index + 1].get("role") != "assistant":
|
||
raise ProductDataError(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 ProductDataError(f"case {case_index} previous 用户/小爱 lines must be non-empty")
|
||
prev_session.append({"query": query, "tts": tts})
|
||
index += 2
|
||
return prev_session
|
||
|
||
|
||
def validate_dataset_records(records: list[dict[str, Any]]) -> dict[str, Any]:
|
||
errors: list[dict[str, str]] = []
|
||
warnings: list[dict[str, str]] = []
|
||
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_value = record.get("record_id")
|
||
if not non_empty_str(record_id_value):
|
||
errors.append(issue(f"{prefix}.record_id", "record_id is required"))
|
||
elif record_id_value in seen_ids:
|
||
errors.append(issue(f"{prefix}.record_id", "record_id must be unique"))
|
||
else:
|
||
seen_ids.add(str(record_id_value))
|
||
|
||
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")
|
||
prev_timestamps: list[int] = []
|
||
if not isinstance(prev_session, list):
|
||
errors.append(issue(f"{prefix}.prev_session", "prev_session must be a list"))
|
||
else:
|
||
if len(prev_session) > 10:
|
||
errors.append(issue(f"{prefix}.prev_session", "prev_session must contain at most 10 turns"))
|
||
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 "tts" not in item or not isinstance(item.get("tts"), str):
|
||
errors.append(issue(f"{item_prefix}.tts", "tts must be a string"))
|
||
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"))
|
||
|
||
if not isinstance(record.get("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 export_dataset_records(
|
||
records: list[dict[str, Any]],
|
||
*,
|
||
output_path: str,
|
||
output_format: str = "jsonl",
|
||
require_validation_ok: bool = True,
|
||
overwrite: bool = True,
|
||
) -> dict[str, Any]:
|
||
if output_format not in {"jsonl", "json"}:
|
||
raise ProductDataError("output_format must be jsonl or json")
|
||
validation = validate_dataset_records(records)
|
||
if require_validation_ok and not validation["ok"]:
|
||
raise ProductDataError(f"records failed validation with {validation['error_count']} errors")
|
||
|
||
path = Path(output_path).expanduser()
|
||
if path.exists() and not overwrite:
|
||
raise ProductDataError(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")
|
||
return {
|
||
"output_path": str(path),
|
||
"output_format": output_format,
|
||
"record_count": len(records),
|
||
"bytes_written": len(content.encode("utf-8")),
|
||
"validation": validation,
|
||
}
|
||
|
||
|
||
def default_output_path(payload: dict[str, Any]) -> str:
|
||
output_format = str(payload.get("output_format") or "jsonl")
|
||
filename = "records.json" if output_format == "json" else "records.jsonl"
|
||
if isinstance(payload.get("output_path"), str) and payload["output_path"].strip():
|
||
return payload["output_path"].strip()
|
||
if isinstance(payload.get("output_dir"), str) and payload["output_dir"].strip():
|
||
return str(Path(payload["output_dir"]).expanduser() / filename)
|
||
return str(Path("output") / filename)
|
||
|
||
|
||
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 ProductDataError("timestamp must be an integer")
|
||
try:
|
||
return int(value)
|
||
except (TypeError, ValueError) as exc:
|
||
raise ProductDataError("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}
|