Initial Open-SWE-Traces cleanup pipeline
This commit is contained in:
@@ -0,0 +1,303 @@
|
||||
"""Command-line interface for static Open-SWE-Traces quality control."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .client import GLMClient
|
||||
from .config import Settings
|
||||
from .features import extract_static_signals
|
||||
from .io import append_jsonl, get_sample_id, iter_records, load_completed_ids, load_jsonl_index
|
||||
from .repair import apply_static_repair
|
||||
from .workflow import audit_trajectory, classify_record, plan_static_repair, review_repaired_record
|
||||
|
||||
|
||||
def _write_error(path: Path | None, sample_id: str, stage: str, error: Exception) -> None:
|
||||
"""Write a compact non-sensitive error record for later retry."""
|
||||
|
||||
value = {
|
||||
"sample_id": sample_id,
|
||||
"stage": stage,
|
||||
"error_type": type(error).__name__,
|
||||
"error": str(error)[:2000],
|
||||
}
|
||||
if path is None:
|
||||
print(json.dumps(value, ensure_ascii=False), file=sys.stderr)
|
||||
else:
|
||||
append_jsonl(path, value)
|
||||
|
||||
|
||||
def _run_streaming_stage(
|
||||
*,
|
||||
input_path: Path,
|
||||
output_path: Path,
|
||||
errors_path: Path | None,
|
||||
limit: int | None,
|
||||
resume: bool,
|
||||
stage_name: str,
|
||||
processor: Callable[[dict[str, Any]], dict[str, Any] | None],
|
||||
) -> int:
|
||||
"""Run one append-only stage with resume support and progress reporting."""
|
||||
|
||||
completed = load_completed_ids(output_path) if resume else set()
|
||||
processed = 0
|
||||
failures = 0
|
||||
for record in iter_records(input_path):
|
||||
sample_id = get_sample_id(record)
|
||||
if sample_id in completed:
|
||||
continue
|
||||
if limit is not None and processed >= limit:
|
||||
break
|
||||
try:
|
||||
result = processor(record)
|
||||
if result is not None:
|
||||
append_jsonl(output_path, result)
|
||||
except Exception as exc: # noqa: BLE001 - each sample must fail independently.
|
||||
failures += 1
|
||||
_write_error(errors_path, sample_id, stage_name, exc)
|
||||
processed += 1
|
||||
if processed % 10 == 0:
|
||||
print(f"{stage_name}: processed={processed} failures={failures}", file=sys.stderr)
|
||||
print(f"{stage_name}: completed={processed} failures={failures}", file=sys.stderr)
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
def command_features(args: argparse.Namespace) -> int:
|
||||
"""Extract deterministic features without requiring an API key."""
|
||||
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
errors_path=args.errors,
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
stage_name="features",
|
||||
processor=lambda record: {
|
||||
"sample_id": get_sample_id(record),
|
||||
"resolved": record.get("resolved"),
|
||||
"static_signals": extract_static_signals(record),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def command_classify(args: argparse.Namespace) -> int:
|
||||
"""Classify records through GLM-5.2 and local policy enforcement."""
|
||||
|
||||
settings = Settings.from_env()
|
||||
with GLMClient(settings) as client:
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
errors_path=args.errors,
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
stage_name="classification",
|
||||
processor=lambda record: classify_record(record, client),
|
||||
)
|
||||
|
||||
|
||||
def command_audit(args: argparse.Namespace) -> int:
|
||||
"""Score educational process quality and locate causal truncation points."""
|
||||
|
||||
settings = Settings.from_env()
|
||||
with GLMClient(settings) as client:
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
errors_path=args.errors,
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
stage_name="trajectory_audit",
|
||||
processor=lambda record: audit_trajectory(record, client),
|
||||
)
|
||||
|
||||
|
||||
def command_repair_plan(args: argparse.Namespace) -> int:
|
||||
"""Ask GLM-5.2 for allowlisted static repair plans."""
|
||||
|
||||
classifications = load_jsonl_index(args.classifications)
|
||||
settings = Settings.from_env()
|
||||
with GLMClient(settings) as client:
|
||||
def processor(record: dict[str, Any]) -> dict[str, Any] | None:
|
||||
sample_id = get_sample_id(record)
|
||||
classification = classifications[sample_id]
|
||||
if classification["qc_decision"] != "STATIC_REPAIR" and not args.include_nonrepair:
|
||||
return None
|
||||
return plan_static_repair(record, classification, client)
|
||||
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
errors_path=args.errors,
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
stage_name="repair_plan",
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
|
||||
def command_apply_repair(args: argparse.Namespace) -> int:
|
||||
"""Apply validated static repair plans deterministically without API calls."""
|
||||
|
||||
plans = load_jsonl_index(args.plans)
|
||||
completed = load_completed_ids(args.output) if args.resume else set()
|
||||
processed = 0
|
||||
failures = 0
|
||||
for record in iter_records(args.input):
|
||||
sample_id = get_sample_id(record)
|
||||
if sample_id in completed or sample_id not in plans:
|
||||
continue
|
||||
if args.limit is not None and processed >= args.limit:
|
||||
break
|
||||
try:
|
||||
repaired, diff = apply_static_repair(record, plans[sample_id])
|
||||
repaired["sample_id"] = sample_id
|
||||
repaired["static_repair_diff"] = diff
|
||||
append_jsonl(args.output, repaired)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
failures += 1
|
||||
_write_error(args.errors, sample_id, "apply_repair", exc)
|
||||
processed += 1
|
||||
print(f"apply_repair: completed={processed} failures={failures}", file=sys.stderr)
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
def command_review(args: argparse.Namespace) -> int:
|
||||
"""Run independent GLM review of applied static repairs."""
|
||||
|
||||
classifications = load_jsonl_index(args.classifications)
|
||||
plans = load_jsonl_index(args.plans)
|
||||
repaired_records = load_jsonl_index(args.repaired)
|
||||
settings = Settings.from_env()
|
||||
with GLMClient(settings) as client:
|
||||
def processor(record: dict[str, Any]) -> dict[str, Any] | None:
|
||||
sample_id = get_sample_id(record)
|
||||
if sample_id not in repaired_records:
|
||||
return None
|
||||
repaired = repaired_records[sample_id]
|
||||
return review_repaired_record(
|
||||
original_record=record,
|
||||
classification=classifications[sample_id],
|
||||
repair_plan=plans[sample_id],
|
||||
repaired_record=repaired,
|
||||
structured_diff=repaired.get("static_repair_diff") or {},
|
||||
client=client,
|
||||
)
|
||||
|
||||
return _run_streaming_stage(
|
||||
input_path=args.input,
|
||||
output_path=args.output,
|
||||
errors_path=args.errors,
|
||||
limit=args.limit,
|
||||
resume=args.resume,
|
||||
stage_name="repair_review",
|
||||
processor=processor,
|
||||
)
|
||||
|
||||
|
||||
def command_smoke_test(_: argparse.Namespace) -> int:
|
||||
"""Perform one minimal authenticated structured-output request."""
|
||||
|
||||
settings = Settings.from_env()
|
||||
schema = {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"required": ["status", "model"],
|
||||
"properties": {
|
||||
"status": {"const": "ok"},
|
||||
"model": {"type": "string", "minLength": 1},
|
||||
},
|
||||
}
|
||||
with GLMClient(settings) as client:
|
||||
response = client.invoke_json(
|
||||
system_prompt=(
|
||||
"Return a JSON object with status='ok' and model set to the model name you were asked to use."
|
||||
),
|
||||
payload={"requested_model": settings.model},
|
||||
schema=schema,
|
||||
)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"endpoint": settings.endpoint,
|
||||
"model": settings.model,
|
||||
"response": response.data,
|
||||
"request_id": response.request_id,
|
||||
"usage": response.usage,
|
||||
"compatibility_fallback_used": response.compatibility_fallback_used,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def _add_stream_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
"""Add arguments shared by streaming pipeline stages."""
|
||||
|
||||
parser.add_argument("--input", type=Path, required=True, help="JSONL file or dataset directory")
|
||||
parser.add_argument("--output", type=Path, required=True, help="Append-only output JSONL")
|
||||
parser.add_argument("--errors", type=Path, help="Optional error JSONL")
|
||||
parser.add_argument("--limit", type=int, help="Maximum number of new records")
|
||||
parser.add_argument("--resume", action="store_true", help="Skip IDs already present in output")
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
"""Construct the complete command-line parser."""
|
||||
|
||||
parser = argparse.ArgumentParser(prog="swe-qc", description=__doc__)
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
features = subparsers.add_parser("features", help="Extract deterministic static QC features")
|
||||
_add_stream_arguments(features)
|
||||
features.set_defaults(func=command_features)
|
||||
|
||||
classify = subparsers.add_parser("classify", help="Classify trajectories through GLM-5.2")
|
||||
_add_stream_arguments(classify)
|
||||
classify.set_defaults(func=command_classify)
|
||||
|
||||
audit = subparsers.add_parser(
|
||||
"audit", help="Score process quality and locate causal truncation points"
|
||||
)
|
||||
_add_stream_arguments(audit)
|
||||
audit.set_defaults(func=command_audit)
|
||||
|
||||
repair_plan = subparsers.add_parser("repair-plan", help="Create allowlisted static repair plans")
|
||||
_add_stream_arguments(repair_plan)
|
||||
repair_plan.add_argument("--classifications", type=Path, required=True)
|
||||
repair_plan.add_argument("--include-nonrepair", action="store_true")
|
||||
repair_plan.set_defaults(func=command_repair_plan)
|
||||
|
||||
apply_repair = subparsers.add_parser("apply-repair", help="Apply static plans deterministically")
|
||||
_add_stream_arguments(apply_repair)
|
||||
apply_repair.add_argument("--plans", type=Path, required=True)
|
||||
apply_repair.set_defaults(func=command_apply_repair)
|
||||
|
||||
review = subparsers.add_parser("review", help="Independently review repaired records")
|
||||
_add_stream_arguments(review)
|
||||
review.add_argument("--classifications", type=Path, required=True)
|
||||
review.add_argument("--plans", type=Path, required=True)
|
||||
review.add_argument("--repaired", type=Path, required=True)
|
||||
review.set_defaults(func=command_review)
|
||||
|
||||
smoke_test = subparsers.add_parser("smoke-test", help="Test API authentication and JSON output")
|
||||
smoke_test.set_defaults(func=command_smoke_test)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
"""CLI entry point used by the ``swe-qc`` console script."""
|
||||
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
return int(args.func(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user