Initial AITuner study orchestrator

src/aituner/cli.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from .job import append_job, build_trial_job
+from .llm import build_prompt, call_llm_for_proposal, load_capability_profile, parse_proposal_text
+from .spec import Proposal, SpecError, load_study_spec
+from .store import StudyStore
+from .trace import load_trace_requests, summarize_window
+from .worker import run_trial
+
+
+def _study_source_path(study_root: Path) -> Path:
+    return Path((study_root / "study_spec.source").read_text(encoding="utf-8").strip())
+
+
+def cmd_study_init(args: argparse.Namespace) -> int:
+    spec_path = Path(args.spec).resolve()
+    study = load_study_spec(spec_path)
+    store = StudyStore(Path(args.store_root) if args.store_root else None)
+    root = store.init_study(spec_path=spec_path, study=study)
+    print(root)
+    return 0
+
+
+def cmd_study_prompt(args: argparse.Namespace) -> int:
+    store = StudyStore(Path(args.store_root) if args.store_root else None)
+    study_root = Path(args.study_root).resolve()
+    source_path = _study_source_path(study_root)
+    study = load_study_spec(source_path)
+    state = store.load_state(study.study_id)
+    capability_profile = load_capability_profile(study, study_spec_path=source_path)
+    window, requests = load_trace_requests(study, study_spec_path=source_path)
+    prompt = build_prompt(
+        study=study,
+        window_summary=summarize_window(requests, window),
+        state=state,
+        capability_profile=capability_profile,
+    )
+    prompt_name = args.prompt_name or f"prompt-{state.next_trial_index:04d}"
+    path = store.write_prompt(study.study_id, prompt_name, prompt)
+    print(path)
+    return 0
+
+
+def cmd_study_llm_propose(args: argparse.Namespace) -> int:
+    store = StudyStore(Path(args.store_root) if args.store_root else None)
+    study_root = Path(args.study_root).resolve()
+    source_path = _study_source_path(study_root)
+    study = load_study_spec(source_path)
+    state = store.load_state(study.study_id)
+    capability_profile = load_capability_profile(study, study_spec_path=source_path)
+    window, requests = load_trace_requests(study, study_spec_path=source_path)
+    prompt = build_prompt(
+        study=study,
+        window_summary=summarize_window(requests, window),
+        state=state,
+        capability_profile=capability_profile,
+    )
+    proposal_text = call_llm_for_proposal(policy=study.llm, prompt=prompt)
+    proposal = parse_proposal_text(proposal_text, study)
+    name = args.proposal_name or f"proposal-{state.next_trial_index:04d}"
+    path = store.write_proposal(study.study_id, name, proposal)
+    print(path)
+    return 0
+
+
+def cmd_study_register_proposal(args: argparse.Namespace) -> int:
+    store = StudyStore(Path(args.store_root) if args.store_root else None)
+    study_root = Path(args.study_root).resolve()
+    source_path = _study_source_path(study_root)
+    study = load_study_spec(source_path)
+    proposal = parse_proposal_text(Path(args.proposal_file).read_text(encoding="utf-8"), study)
+    name = args.proposal_name or Path(args.proposal_file).stem
+    path = store.write_proposal(study.study_id, name, proposal)
+    print(path)
+    return 0
+
+
+def cmd_study_emit_job(args: argparse.Namespace) -> int:
+    store = StudyStore(Path(args.store_root) if args.store_root else None)
+    study_root = Path(args.study_root).resolve()
+    source_path = _study_source_path(study_root)
+    study = load_study_spec(source_path)
+    state = store.load_state(study.study_id)
+    proposal_text = Path(args.proposal_file).read_text(encoding="utf-8")
+    proposal = parse_proposal_text(proposal_text, study)
+    trial, _ = store.materialize_trial(study=study, state=state, proposal=proposal)
+    repo_root = Path(__file__).resolve().parents[2]
+    job = build_trial_job(study=study, trial=trial, repo_root=repo_root)
+    append_job(Path(args.jobs_file).resolve(), job)
+    print(trial.trial_id)
+    return 0
+
+
+def cmd_study_ingest(args: argparse.Namespace) -> int:
+    store = StudyStore(Path(args.store_root) if args.store_root else None)
+    study_root = Path(args.study_root).resolve()
+    study_id = study_root.name
+    state = store.ingest_trial_results(study_id)
+    print(json.dumps({"best_trial_id": state.best_trial_id, "best_request_rate": state.best_request_rate}))
+    return 0
+
+
+def cmd_worker_run_trial(args: argparse.Namespace) -> int:
+    result = run_trial(Path(args.trial_spec).resolve())
+    print(json.dumps(result))
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="AITuner CLI")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    study = subparsers.add_parser("study")
+    study_sub = study.add_subparsers(dest="study_command", required=True)
+
+    init = study_sub.add_parser("init")
+    init.add_argument("--spec", required=True)
+    init.add_argument("--store-root")
+    init.set_defaults(func=cmd_study_init)
+
+    prompt = study_sub.add_parser("prompt")
+    prompt.add_argument("--study-root", required=True)
+    prompt.add_argument("--store-root")
+    prompt.add_argument("--prompt-name")
+    prompt.set_defaults(func=cmd_study_prompt)
+
+    llm_propose = study_sub.add_parser("llm-propose")
+    llm_propose.add_argument("--study-root", required=True)
+    llm_propose.add_argument("--store-root")
+    llm_propose.add_argument("--proposal-name")
+    llm_propose.set_defaults(func=cmd_study_llm_propose)
+
+    register = study_sub.add_parser("register-proposal")
+    register.add_argument("--study-root", required=True)
+    register.add_argument("--store-root")
+    register.add_argument("--proposal-file", required=True)
+    register.add_argument("--proposal-name")
+    register.set_defaults(func=cmd_study_register_proposal)
+
+    emit = study_sub.add_parser("emit-job")
+    emit.add_argument("--study-root", required=True)
+    emit.add_argument("--store-root")
+    emit.add_argument("--proposal-file", required=True)
+    emit.add_argument("--jobs-file", required=True)
+    emit.set_defaults(func=cmd_study_emit_job)
+
+    ingest = study_sub.add_parser("ingest")
+    ingest.add_argument("--study-root", required=True)
+    ingest.add_argument("--store-root")
+    ingest.set_defaults(func=cmd_study_ingest)
+
+    worker = subparsers.add_parser("worker")
+    worker_sub = worker.add_subparsers(dest="worker_command", required=True)
+    run = worker_sub.add_parser("run-trial")
+    run.add_argument("--trial-spec", required=True)
+    run.set_defaults(func=cmd_worker_run_trial)
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    try:
+        return int(args.func(args))
+    except SpecError as exc:
+        print(str(exc), file=sys.stderr)
+        return 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())