File size: 6,100 Bytes
0ffbb74 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """Run one real Muse Spark agent run on a simulated scenario.
Prints each model call with its tool calls and reasons, then the report, the
revealed truth, total tokens and the estimated cost, and saves the run as JSON
in ``runs/`` (git-ignored). Only simulated data is sent to the API.
Usage (from the repository root, PowerShell)::
.venv\\Scripts\\python.exe scripts\\run_agent.py --scenario clean --seed 0
"""
from __future__ import annotations
import argparse
import json
import sys
from dataclasses import replace
from datetime import datetime
from pathlib import Path
from typing import Any
from beampilot.agent import (
DEFAULT_EFFORT,
DEFAULT_MODEL,
PRICES_PER_MTOK,
AgentConfig,
make_client,
run_agent,
save_run,
)
from beampilot.session import (
SCENARIOS,
LensKit,
MeasurementSession,
random_scenario,
reveal_truth,
session_record,
)
RUNS_DIR = Path(__file__).resolve().parents[1] / "runs"
def short(result: dict[str, Any]) -> str:
"""One-line summary of a tool result."""
if "error" in result or "internal_error" in result:
return json.dumps(result, ensure_ascii=False)
if "flags" in result:
flags = [name for name, value in result["flags"].items() if value]
return (
f"{result['measurement_id']} at {result['z_mm']} mm: d_x "
f"{result['d_x_um']['mean']} um, d_y {result['d_y_um']['mean']} um, peak "
f"{result['details']['peak_counts']}, flags {flags or 'none'}, "
f"suggested {result['suggested_verdict']}, "
f"{result['positions_left']} left"
)
if result.get("status") == "ok":
axes = []
for axis in ("x", "y"):
a = result[axis]
coverage = a["coverage"]["summary"] if a["coverage"] else "no coverage"
axes.append(
f"M²{axis} {a['m2']} z0 {a['z0_mm']} zR {a['z_r_mm']} [{coverage}]"
)
return "; ".join(axes) + f"; ISO coverage met {result['iso_coverage_met']}"
if "report" in result:
return f"report: verdict {result['report']['verdict']}"
text = json.dumps(result, ensure_ascii=False)
return text if len(text) < 200 else text[:200] + " ..."
def print_step(step: dict[str, Any]) -> None:
usage = step["usage"]
print(
f"\n-- model call {step['call']} ({step['latency_s']} s, in "
f"{usage['input_tokens']} [cached {usage['cached_tokens']}], out "
f"{usage['output_tokens']} [reasoning {usage['reasoning_tokens']}], "
f"${step['cost_usd']:.4f})"
)
if step["text"]:
print(f" says: {step['text'].strip()}")
for call in step["tool_calls"]:
arguments = {k: v for k, v in call["arguments"].items() if k != "reason"}
print(f" {call['name']}({arguments})")
print(f" reason: {call.get('reason')}")
print(f" -> {short(call['result'])}")
def main() -> int:
# Model text may contain any Unicode; never crash a paid run on printing,
# and flush each line so a redirected log shows progress.
sys.stdout.reconfigure(encoding="utf-8", errors="replace", line_buffering=True)
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument(
"--scenario", choices=[*sorted(SCENARIOS), "mystery"], default="clean"
)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--note", default=None, help="operator note for the agent")
parser.add_argument("--model", default=DEFAULT_MODEL)
parser.add_argument("--effort", default=DEFAULT_EFFORT)
parser.add_argument("--max-calls", type=int, default=60)
parser.add_argument("--max-cost", type=float, default=0.50, help="USD per run")
parser.add_argument(
"--lens", action="store_true", help="add the lens kit and offer change_lens"
)
parser.add_argument("--max-positions", type=int, default=None)
args = parser.parse_args()
factory = (
random_scenario if args.scenario == "mystery" else SCENARIOS[args.scenario]
)
scenario = factory(seed=args.seed)
if args.lens:
scenario = replace(scenario, lens_kit=LensKit())
if args.max_positions is not None:
scenario = replace(scenario, max_positions=args.max_positions)
session = MeasurementSession(scenario)
config = AgentConfig(
model=args.model,
reasoning_effort=args.effort,
max_model_calls=args.max_calls,
max_cost_usd=args.max_cost,
change_lens=args.lens,
)
print(
f"== Agent run: {args.scenario}, seed {args.seed}, {args.model}, "
f"effort {args.effort}"
)
if args.note:
print(f" operator note: {args.note}")
run = run_agent(
session, make_client(), config, on_step=print_step, operator_note=args.note
)
print(f"\n== Stop reason: {run.stop_reason}")
print("\n== Report")
print(json.dumps(run.report, indent=2, ensure_ascii=False))
truth = reveal_truth(session)
print("\n== Revealed truth")
print(json.dumps(truth, indent=2, ensure_ascii=False))
usage = run.usage
print(
f"\n== Totals: {len(run.steps)} model calls, {session.positions_used} "
f"positions, input {usage.input_tokens} (cached {usage.cached_tokens}), "
f"output {usage.output_tokens} (reasoning {usage.reasoning_tokens})"
)
print(
f" estimated cost ${run.cost_usd:.4f} (prices {PRICES_PER_MTOK} per "
"million tokens -- check pricing page)"
)
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
path = RUNS_DIR / f"{stamp}_{args.scenario}_s{args.seed}_{args.model}.json"
save_run(
run,
path,
{"scenario": args.scenario, "seed": args.seed, **session_record(session)},
)
print(f" saved {path}")
return 0 if run.stop_reason == "finished" else 2
if __name__ == "__main__":
sys.exit(main())
|