wuyang6
1025 lines
38 KiB
Python
1025 lines
38 KiB
Python
from __future__ import annotations
|
||
|
||
"""product-data 可迁移脚本的共享逻辑。
|
||
|
||
这里不依赖 ZK Data Agent 的运行时，便于把整个 skill 复制到其他 Agent
|
||
或普通 Python 环境中使用。平台注册工具时，也应该只薄封装这些脚本。
|
||
"""
|
||
|
||
import argparse
|
||
import csv
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
import time
|
||
from io import StringIO
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
|
||
DEFAULT_REQUEST_ID = "aabbccdd"
|
||
DEFAULT_TIMESTAMP_STEP_MS = 60_000
|
||
DEFAULT_SYSTEM_PROMPT = "你是小爱同学，中文智能语音助手。"
|
||
DEFAULT_CONTEXT_FIELDS = ("location", "rag")
|
||
|
||
|
||
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 = resolve_portable_path(path)
|
||
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 isinstance(decoded, dict) and isinstance(decoded.get("records"), list):
|
||
decoded = decoded["records"]
|
||
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|complex|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")
|
||
|
||
raw_target, prefixed_complex = split_complex_prefixed_target(str(case.get("target") or "").strip())
|
||
target = canonical_target(raw_target)
|
||
if not target:
|
||
raise ProductDataError(f"case {index} target is required")
|
||
complex_value = parse_complex(
|
||
case.get("complex"),
|
||
default=False if prefixed_complex is None else prefixed_complex,
|
||
)
|
||
|
||
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),
|
||
},
|
||
"dimensions": {
|
||
"complex": complex_value,
|
||
},
|
||
"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"))
|
||
|
||
dimensions = record.get("dimensions")
|
||
if not isinstance(dimensions, dict):
|
||
warnings.append(issue(f"{prefix}.dimensions", "dimensions.complex is missing; false will be used as fallback"))
|
||
elif "complex" not in dimensions:
|
||
warnings.append(issue(f"{prefix}.dimensions.complex", "complex is missing; false will be used as fallback"))
|
||
elif not isinstance(dimensions.get("complex"), bool):
|
||
errors.append(issue(f"{prefix}.dimensions.complex", "complex must be a boolean"))
|
||
|
||
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,
|
||
export_table: 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()
|
||
table_path = path.with_name("records.csv")
|
||
if path.exists() and not overwrite:
|
||
raise ProductDataError(f"output_path already exists: {output_path}")
|
||
if export_table and table_path.exists() and not overwrite:
|
||
raise ProductDataError("table output_path already exists: records.csv")
|
||
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")
|
||
payload = {
|
||
"output_path": str(path),
|
||
"output_format": output_format,
|
||
"record_count": len(records),
|
||
"bytes_written": len(content.encode("utf-8")),
|
||
"validation": validation,
|
||
}
|
||
if export_table:
|
||
table_content = render_dataset_records_table_csv(records)
|
||
table_path.write_text(table_content, encoding="utf-8-sig")
|
||
payload.update(
|
||
{
|
||
"table_output_path": str(table_path),
|
||
"table_output_format": "csv",
|
||
"table_bytes_written": len(table_content.encode("utf-8-sig")),
|
||
}
|
||
)
|
||
return payload
|
||
|
||
|
||
def export_dataset_table(
|
||
records: list[dict[str, Any]],
|
||
*,
|
||
output_path: str,
|
||
require_validation_ok: bool = True,
|
||
overwrite: bool = True,
|
||
) -> dict[str, Any]:
|
||
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)
|
||
content = render_dataset_records_table_csv(records)
|
||
path.write_text(content, encoding="utf-8-sig")
|
||
return {
|
||
"output_path": str(path),
|
||
"output_format": "csv",
|
||
"record_count": len(records),
|
||
"bytes_written": len(content.encode("utf-8-sig")),
|
||
"validation": validation,
|
||
}
|
||
|
||
|
||
def export_training_jsonl(
|
||
records: list[dict[str, Any]],
|
||
*,
|
||
output_path: str,
|
||
session_num: int = 5,
|
||
session_time_minutes: int = 5,
|
||
context_fields: list[str] | None = None,
|
||
system_prompt: str = DEFAULT_SYSTEM_PROMPT,
|
||
require_validation_ok: bool = True,
|
||
overwrite: bool = True,
|
||
) -> dict[str, Any]:
|
||
if session_num <= 0:
|
||
raise ProductDataError("session_num must be greater than 0")
|
||
if session_time_minutes <= 0:
|
||
raise ProductDataError("session_time_minutes must be greater than 0")
|
||
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)
|
||
fields = normalize_context_fields(context_fields)
|
||
content = "".join(
|
||
training_jsonl_line(
|
||
record,
|
||
session_num=session_num,
|
||
session_time_minutes=session_time_minutes,
|
||
context_fields=fields,
|
||
system_prompt=system_prompt,
|
||
)
|
||
+ "\n"
|
||
for record in records
|
||
)
|
||
path.write_text(content, encoding="utf-8")
|
||
return {
|
||
"output_path": str(path),
|
||
"output_format": "jsonl",
|
||
"record_count": len(records),
|
||
"bytes_written": len(content.encode("utf-8")),
|
||
"validation": validation,
|
||
}
|
||
|
||
|
||
def export_planning_eval_csv(
|
||
records: list[dict[str, Any]],
|
||
*,
|
||
output_path: str,
|
||
session_num: int = 5,
|
||
session_time_minutes: int = 5,
|
||
context_fields: list[str] | None = None,
|
||
system_prompt: str = DEFAULT_SYSTEM_PROMPT,
|
||
complex_default: bool = False,
|
||
require_validation_ok: bool = True,
|
||
overwrite: bool = True,
|
||
) -> dict[str, Any]:
|
||
if session_num <= 0:
|
||
raise ProductDataError("session_num must be greater than 0")
|
||
if session_time_minutes <= 0:
|
||
raise ProductDataError("session_time_minutes must be greater than 0")
|
||
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)
|
||
content = render_planning_eval_csv(
|
||
records,
|
||
session_num=session_num,
|
||
session_time_minutes=session_time_minutes,
|
||
context_fields=context_fields,
|
||
system_prompt=system_prompt,
|
||
complex_default=complex_default,
|
||
)
|
||
path.write_text(content, encoding="utf-8-sig")
|
||
return {
|
||
"output_path": str(path),
|
||
"output_format": "csv",
|
||
"record_count": len(records),
|
||
"bytes_written": len(content.encode("utf-8-sig")),
|
||
"validation": validation,
|
||
}
|
||
|
||
|
||
def render_dataset_records_table_csv(records: list[dict[str, Any]]) -> str:
|
||
output = StringIO()
|
||
writer = csv.DictWriter(output, fieldnames=dataset_table_columns(), lineterminator="\n")
|
||
writer.writeheader()
|
||
for record in records:
|
||
writer.writerow(dataset_record_table_row(record))
|
||
return output.getvalue()
|
||
|
||
|
||
def render_planning_eval_csv(
|
||
records: list[dict[str, Any]],
|
||
*,
|
||
session_num: int = 5,
|
||
session_time_minutes: int = 5,
|
||
context_fields: list[str] | None = None,
|
||
system_prompt: str = DEFAULT_SYSTEM_PROMPT,
|
||
complex_default: bool = False,
|
||
) -> str:
|
||
fields = normalize_context_fields(context_fields)
|
||
output = StringIO()
|
||
writer = csv.DictWriter(
|
||
output,
|
||
fieldnames=["request_id", "newPrompt", "query", "类别真实标签", "code标签", "complex"],
|
||
lineterminator="\n",
|
||
)
|
||
writer.writeheader()
|
||
for record in records:
|
||
source = record.get("source") if isinstance(record.get("source"), dict) else {}
|
||
turn = record.get("turn") if isinstance(record.get("turn"), dict) else {}
|
||
target = record_target(record)
|
||
writer.writerow(
|
||
{
|
||
"request_id": str(source.get("request_id") or ""),
|
||
"newPrompt": build_planning_prompt(
|
||
record,
|
||
session_num=session_num,
|
||
session_time_minutes=session_time_minutes,
|
||
context_fields=fields,
|
||
system_prompt=system_prompt,
|
||
),
|
||
"query": str(turn.get("query") or ""),
|
||
"类别真实标签": category_label_from_target(target),
|
||
"code标签": target,
|
||
"complex": eval_complex_literal(record_complex(record, default=complex_default)),
|
||
}
|
||
)
|
||
return output.getvalue()
|
||
|
||
|
||
def build_training_instruction(
|
||
record: dict[str, Any],
|
||
*,
|
||
session_num: int = 5,
|
||
session_time_minutes: int = 5,
|
||
context_fields: list[str] | None = None,
|
||
) -> str:
|
||
turn = record.get("turn") if isinstance(record.get("turn"), dict) else {}
|
||
query = str(turn.get("query") or "")
|
||
context = record.get("context") if isinstance(record.get("context"), dict) else {}
|
||
prev_session = record.get("prev_session") if isinstance(record.get("prev_session"), list) else []
|
||
current_ts = optional_int_for_export(turn.get("timestamp"))
|
||
if current_ts is None:
|
||
source = record.get("source") if isinstance(record.get("source"), dict) else {}
|
||
current_ts = optional_int_for_export(source.get("timestamp"))
|
||
session, last_tts = training_history(
|
||
prev_session,
|
||
current_ts=current_ts,
|
||
session_num=session_num,
|
||
session_time_minutes=session_time_minutes,
|
||
)
|
||
fields = normalize_context_fields(context_fields)
|
||
|
||
instruction = "请参考用户的[当前query]、[对话历史]、[知识注入]、[系统状态]识别出[当前query]的[function]结果，[function]是python的code形式。\n"
|
||
instruction += "[知识注入]\n"
|
||
instruction += f"{context_prompt_block(context, fields)}\n"
|
||
instruction += "[系统状态]\n"
|
||
instruction += "{}\n"
|
||
instruction += "[对话历史]\n"
|
||
if session:
|
||
history_parts = [f"用户: {session_query}" for session_query in session]
|
||
if last_tts is not None:
|
||
history_parts.append(f"小爱: {last_tts}")
|
||
instruction += "\n".join(history_parts)
|
||
instruction += "\n"
|
||
instruction += "[当前query]\n"
|
||
instruction += f"用户: {query}\n"
|
||
instruction += "[function]\n"
|
||
return instruction
|
||
|
||
|
||
def build_planning_prompt(
|
||
record: dict[str, Any],
|
||
*,
|
||
session_num: int = 5,
|
||
session_time_minutes: int = 5,
|
||
context_fields: list[str] | None = None,
|
||
system_prompt: str = DEFAULT_SYSTEM_PROMPT,
|
||
) -> str:
|
||
"""把训练 instruction 包成评测侧使用的 chat template。"""
|
||
|
||
instruction = prompt_section_with_one_trailing_newline(
|
||
build_training_instruction(
|
||
record,
|
||
session_num=session_num,
|
||
session_time_minutes=session_time_minutes,
|
||
context_fields=context_fields,
|
||
)
|
||
)
|
||
system_prompt = prompt_section_without_trailing_newline(system_prompt)
|
||
return (
|
||
f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
|
||
f"<|im_start|>user\n{instruction}<|im_end|>\n"
|
||
"<|im_start|>assistant\n"
|
||
)
|
||
|
||
|
||
def prompt_section_without_trailing_newline(text: str) -> str:
|
||
return str(text or "").rstrip("\n")
|
||
|
||
|
||
def prompt_section_with_one_trailing_newline(text: str) -> str:
|
||
return prompt_section_without_trailing_newline(text) + "\n"
|
||
|
||
|
||
def training_jsonl_line(
|
||
record: dict[str, Any],
|
||
*,
|
||
session_num: int,
|
||
session_time_minutes: int,
|
||
context_fields: list[str],
|
||
system_prompt: str,
|
||
) -> str:
|
||
payload = {
|
||
"system": prompt_section_without_trailing_newline(system_prompt),
|
||
"instruction": build_training_instruction(
|
||
record,
|
||
session_num=session_num,
|
||
session_time_minutes=session_time_minutes,
|
||
context_fields=context_fields,
|
||
),
|
||
"output": combined_function_label(record),
|
||
}
|
||
return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
|
||
|
||
|
||
def training_history(
|
||
prev_session: list[Any],
|
||
*,
|
||
current_ts: int | None,
|
||
session_num: int,
|
||
session_time_minutes: int,
|
||
) -> tuple[list[str], str | None]:
|
||
if current_ts is None:
|
||
return [], None
|
||
session_time_ms = session_time_minutes * 60 * 1000
|
||
valid_items: list[tuple[int, str, str]] = []
|
||
for item in prev_session:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
ts = optional_int_for_export(item.get("timestamp"))
|
||
query = str(item.get("query") or "")
|
||
tts = str(item.get("tts") or "")
|
||
if ts is not None and ts > 0 and query:
|
||
valid_items.append((ts, query, tts))
|
||
valid_items.sort(key=lambda item: item[0])
|
||
|
||
session: list[str] = []
|
||
last_tts: str | None = None
|
||
last_ts = current_ts
|
||
count = 0
|
||
for ts, query, tts in reversed(valid_items):
|
||
if last_ts - ts <= session_time_ms and last_ts >= ts:
|
||
session.append(query)
|
||
if count == 0:
|
||
last_tts = tts
|
||
last_ts = ts
|
||
count += 1
|
||
if count >= session_num:
|
||
break
|
||
else:
|
||
break
|
||
session.reverse()
|
||
return session, last_tts
|
||
|
||
|
||
def context_prompt_block(context: dict[str, Any], fields: list[str]) -> str:
|
||
lines = ["{"]
|
||
for index, field in enumerate(fields):
|
||
comma = "," if index < len(fields) - 1 else ""
|
||
value = str(context.get(field, ""))
|
||
lines.append(f'"{field}": {json.dumps(value, ensure_ascii=False)}{comma}')
|
||
lines.append("}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def normalize_context_fields(context_fields: list[str] | None) -> list[str]:
|
||
if context_fields is not None and not isinstance(context_fields, list):
|
||
raise ProductDataError("context_fields must be an array of strings")
|
||
fields = context_fields or list(DEFAULT_CONTEXT_FIELDS)
|
||
normalized = [field.strip() for field in fields if isinstance(field, str) and field.strip()]
|
||
return normalized or list(DEFAULT_CONTEXT_FIELDS)
|
||
|
||
|
||
def optional_int_for_export(value: Any) -> int | None:
|
||
if isinstance(value, bool) or value is None or value == "":
|
||
return None
|
||
try:
|
||
return int(value)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
def category_label_from_target(target: str) -> str:
|
||
target = strip_complex_prefix(target)
|
||
agent_match = re.match(r'''^Agent\s*\(\s*tag\s*=\s*["']([^"']+)["']\s*\)$''', target.strip())
|
||
if agent_match:
|
||
return agent_match.group(1)
|
||
function_names = re.findall(r"\b([A-Za-z_][A-Za-z0-9_]*)\s*\(", target)
|
||
if function_names:
|
||
return function_names[-1]
|
||
return target
|
||
|
||
|
||
def dataset_table_columns() -> list[str]:
|
||
return [
|
||
"request_id",
|
||
"timestamp",
|
||
"query",
|
||
"prev_session",
|
||
"context",
|
||
"label",
|
||
"是否迁移Function",
|
||
"function",
|
||
]
|
||
|
||
|
||
def dataset_record_table_row(record: dict[str, Any]) -> dict[str, str]:
|
||
source = record.get("source") if isinstance(record.get("source"), dict) else {}
|
||
turn = record.get("turn") if isinstance(record.get("turn"), dict) else {}
|
||
label = record.get("label") if isinstance(record.get("label"), dict) else {}
|
||
context = record.get("context") if isinstance(record.get("context"), dict) else {}
|
||
prev_session = record.get("prev_session")
|
||
if not isinstance(prev_session, list):
|
||
prev_session = []
|
||
return {
|
||
"request_id": str(source.get("request_id") or ""),
|
||
"timestamp": str(source.get("timestamp") or turn.get("timestamp") or ""),
|
||
"query": str(turn.get("query") or ""),
|
||
"prev_session": json.dumps(table_prev_session(prev_session), ensure_ascii=False, separators=(",", ":")),
|
||
"context": json.dumps(context, ensure_ascii=False, separators=(",", ":")),
|
||
"label": str(label.get("dataset_label") or ""),
|
||
"是否迁移Function": "",
|
||
"function": combined_function_label(record),
|
||
}
|
||
|
||
|
||
def table_prev_session(prev_session: list[Any]) -> list[dict[str, str]]:
|
||
rows: list[dict[str, str]] = []
|
||
for item in prev_session[-10:]:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
rows.append(
|
||
{
|
||
"query": str(item.get("query") or ""),
|
||
"tts": str(item.get("tts") or ""),
|
||
"timestamp": str(item.get("timestamp") or ""),
|
||
}
|
||
)
|
||
return rows
|
||
|
||
|
||
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 str(resolve_portable_path(payload["output_path"].strip()))
|
||
if isinstance(payload.get("output_dir"), str) and payload["output_dir"].strip():
|
||
return str(resolve_portable_path(payload["output_dir"].strip()) / filename)
|
||
return str(default_output_root() / filename)
|
||
|
||
|
||
def default_table_output_path(payload: dict[str, Any]) -> str:
|
||
if isinstance(payload.get("output_path"), str) and payload["output_path"].strip():
|
||
return str(resolve_portable_path(payload["output_path"].strip()))
|
||
if isinstance(payload.get("output_dir"), str) and payload["output_dir"].strip():
|
||
return str(resolve_portable_path(payload["output_dir"].strip()) / "records.csv")
|
||
if isinstance(payload.get("records_path"), str) and payload["records_path"].strip():
|
||
return str(resolve_portable_path(payload["records_path"].strip()).with_name("records.csv"))
|
||
return str(default_output_root() / "records.csv")
|
||
|
||
|
||
def default_training_output_path(payload: dict[str, Any]) -> str:
|
||
if isinstance(payload.get("output_path"), str) and payload["output_path"].strip():
|
||
return str(resolve_portable_path(payload["output_path"].strip()))
|
||
if isinstance(payload.get("output_dir"), str) and payload["output_dir"].strip():
|
||
return str(resolve_portable_path(payload["output_dir"].strip()) / "training.jsonl")
|
||
if isinstance(payload.get("records_path"), str) and payload["records_path"].strip():
|
||
return str(resolve_portable_path(payload["records_path"].strip()).with_name("training.jsonl"))
|
||
return str(default_output_root() / "training.jsonl")
|
||
|
||
|
||
def default_planning_eval_output_path(payload: dict[str, Any]) -> str:
|
||
if isinstance(payload.get("output_path"), str) and payload["output_path"].strip():
|
||
return str(resolve_portable_path(payload["output_path"].strip()))
|
||
if isinstance(payload.get("output_dir"), str) and payload["output_dir"].strip():
|
||
return str(resolve_portable_path(payload["output_dir"].strip()) / "eval_planning.csv")
|
||
if isinstance(payload.get("records_path"), str) and payload["records_path"].strip():
|
||
return str(resolve_portable_path(payload["records_path"].strip()).with_name("eval_planning.csv"))
|
||
return str(default_output_root() / "eval_planning.csv")
|
||
|
||
|
||
def default_output_root() -> Path:
|
||
"""返回 portable 脚本默认输出目录。
|
||
|
||
在 ZK Data Agent 中，python_exec 的 cwd 是当前会话 scratchpad；
|
||
通过 PYTHON_EXEC_SCRATCHPAD 可以定位同级 output 目录。离开平台独立
|
||
运行时，保持原来的相对 output/ 行为。
|
||
"""
|
||
|
||
scratchpad = os.environ.get("PYTHON_EXEC_SCRATCHPAD")
|
||
if scratchpad:
|
||
return Path(scratchpad).expanduser().parent / "output"
|
||
return Path("output")
|
||
|
||
|
||
def resolve_portable_path(path: str) -> Path:
|
||
"""解析 portable 脚本路径，并兼容平台逻辑路径。
|
||
|
||
平台内运行时，把 output/、scratchpad/、input/ 映射到当前会话目录；
|
||
其他相对路径仍按当前 cwd 解析，便于普通 Python 环境复用。
|
||
"""
|
||
|
||
raw_path = Path(path).expanduser()
|
||
if raw_path.is_absolute():
|
||
return raw_path
|
||
scratchpad = os.environ.get("PYTHON_EXEC_SCRATCHPAD")
|
||
if not scratchpad:
|
||
return raw_path
|
||
session_root = Path(scratchpad).expanduser().parent
|
||
parts = raw_path.parts
|
||
if not parts:
|
||
return raw_path
|
||
head, *tail = parts
|
||
tail_path = Path(*tail) if tail else Path()
|
||
if head in {"output", "outputs"}:
|
||
return session_root / "output" / tail_path
|
||
if head in {"scratchpad", "scratch"}:
|
||
return Path(scratchpad).expanduser() / tail_path
|
||
if head in {"input", "inputs"}:
|
||
return session_root / "input" / tail_path
|
||
return raw_path
|
||
|
||
|
||
def normalize_target_expression(target: str) -> str:
|
||
"""把可能带 complex 前缀的标签表达式收敛为纯 target。"""
|
||
|
||
return canonical_target(strip_complex_prefix(target))
|
||
|
||
|
||
def strip_complex_prefix(target: str) -> str:
|
||
stripped_target, _complex_value = split_complex_prefixed_target(target)
|
||
return stripped_target
|
||
|
||
|
||
def split_complex_prefixed_target(target: str) -> tuple[str, bool | None]:
|
||
lines = target.strip().splitlines()
|
||
if not lines:
|
||
return "", None
|
||
first_line = lines[0].strip()
|
||
match = re.fullmatch(r"complex\s*=\s*(.+)", first_line, flags=re.I)
|
||
if not match:
|
||
return target.strip(), None
|
||
complex_value = parse_complex(match.group(1), default=False)
|
||
return "\n".join(lines[1:]).strip(), complex_value
|
||
|
||
|
||
def parse_complex(value: Any, *, default: bool) -> bool:
|
||
if value is None or value == "":
|
||
return default
|
||
if isinstance(value, bool):
|
||
return value
|
||
text = str(value).strip().lower()
|
||
if text in {"true", "1", "yes", "y", "是", "复杂", "complex"}:
|
||
return True
|
||
if text in {"false", "0", "no", "n", "否", "不复杂", "简单", "simple"}:
|
||
return False
|
||
raise ProductDataError("complex must be a boolean value such as true/false")
|
||
|
||
|
||
def complex_literal(value: bool) -> str:
|
||
return "true" if value else "false"
|
||
|
||
|
||
def eval_complex_literal(value: bool) -> str:
|
||
return "TRUE" if value else "FALSE"
|
||
|
||
|
||
def record_target(record: dict[str, Any]) -> str:
|
||
label = record.get("label") if isinstance(record.get("label"), dict) else {}
|
||
return normalize_target_expression(str(label.get("target") or ""))
|
||
|
||
|
||
def record_complex(record: dict[str, Any], *, default: bool = False) -> bool:
|
||
dimensions = record.get("dimensions")
|
||
if isinstance(dimensions, dict) and "complex" in dimensions:
|
||
return parse_complex(dimensions.get("complex"), default=default)
|
||
label = record.get("label") if isinstance(record.get("label"), dict) else {}
|
||
if "complex" in label:
|
||
return parse_complex(label.get("complex"), default=default)
|
||
_target, prefixed_complex = split_complex_prefixed_target(str(label.get("target") or ""))
|
||
if prefixed_complex is not None:
|
||
return prefixed_complex
|
||
return default
|
||
|
||
|
||
def combined_function_label(record: dict[str, Any], *, default_complex: bool = False) -> str:
|
||
target = record_target(record)
|
||
return f"complex={complex_literal(record_complex(record, default=default_complex))}\n{target}"
|
||
|
||
|
||
def target_type(target: str) -> str:
|
||
target = strip_complex_prefix(target)
|
||
if re.match(r"^Agent\s*\(\s*tag\s*=", target):
|
||
return "agent"
|
||
if target:
|
||
return "function"
|
||
return "unknown"
|
||
|
||
|
||
def canonical_target(target: str) -> str:
|
||
target = strip_complex_prefix(target)
|
||
agent_match = re.match(r'''^Agent\s*\(\s*tag\s*=\s*["']([^"']+)["']\s*\)$''', target)
|
||
if agent_match:
|
||
return f'Agent(tag="{agent_match.group(1)}")'
|
||
return target
|
||
|
||
|
||
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}
|