228 lines
6.6 KiB
Python
228 lines
6.6 KiB
Python
"""Disposable Python probe scaffold.
|
|
|
|
Copy this file to a temporary location and adapt it for one narrow question.
|
|
Recommended usage from the repository root:
|
|
|
|
uv run python /tmp/probe.py
|
|
|
|
If you want structured artifacts for repeat-heavy or benchmark probes:
|
|
|
|
PROBE_OUTPUT_DIR=/tmp/probe-run uv run python /tmp/probe.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import statistics
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from collections import Counter, defaultdict
|
|
from importlib import metadata
|
|
from pathlib import Path
|
|
|
|
SCENARIO = "replace-me"
|
|
RUN_LABEL = "replace-me"
|
|
MODE = "single-shot"
|
|
APPROVED_ENV_VARS: list[str] = []
|
|
OUTPUT_DIR_ENV = "PROBE_OUTPUT_DIR"
|
|
|
|
RESULTS: list[dict[str, object]] = []
|
|
|
|
|
|
def _git_value(*args: str) -> str:
|
|
result = subprocess.run(
|
|
["git", *args],
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
if result.returncode != 0:
|
|
return "unknown"
|
|
return result.stdout.strip() or "unknown"
|
|
|
|
|
|
def _package_version(name: str) -> str | None:
|
|
try:
|
|
return metadata.version(name)
|
|
except metadata.PackageNotFoundError:
|
|
return None
|
|
|
|
|
|
def _output_dir() -> Path | None:
|
|
value = os.getenv(OUTPUT_DIR_ENV)
|
|
if not value:
|
|
return None
|
|
return Path(value)
|
|
|
|
|
|
def _write_json(path: Path, payload: object) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
|
|
|
|
|
def emit(kind: str, **payload: object) -> None:
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"ts": round(time.time(), 3),
|
|
"kind": kind,
|
|
**payload,
|
|
},
|
|
sort_keys=True,
|
|
)
|
|
)
|
|
|
|
|
|
def runtime_context() -> dict[str, object]:
|
|
approved = {name: ("set" if os.getenv(name) else "unset") for name in APPROVED_ENV_VARS}
|
|
package_versions = {
|
|
name: version
|
|
for name in ("openai", "agents")
|
|
if (version := _package_version(name)) is not None
|
|
}
|
|
return {
|
|
"scenario": SCENARIO,
|
|
"run_label": RUN_LABEL,
|
|
"mode": MODE,
|
|
"cwd": os.getcwd(),
|
|
"script_path": str(Path(__file__).resolve()),
|
|
"python_executable": sys.executable,
|
|
"python_version": sys.version.split()[0],
|
|
"platform": platform.platform(),
|
|
"git_commit": _git_value("rev-parse", "HEAD"),
|
|
"git_branch": _git_value("rev-parse", "--abbrev-ref", "HEAD"),
|
|
"uv_path": shutil.which("uv"),
|
|
"package_versions": package_versions,
|
|
"approved_env_vars": approved,
|
|
"output_dir": str(_output_dir()) if _output_dir() else None,
|
|
}
|
|
|
|
|
|
def start_case(case_id: str, *, mode: str = MODE, note: str | None = None) -> None:
|
|
emit("case_start", case_id=case_id, mode=mode, note=note)
|
|
|
|
|
|
def record_case_result(
|
|
case_id: str,
|
|
observation_summary: str,
|
|
result_flag: str,
|
|
*,
|
|
mode: str = MODE,
|
|
is_warmup: bool = False,
|
|
total_latency_s: float | None = None,
|
|
first_token_latency_s: float | None = None,
|
|
metrics: dict[str, object] | None = None,
|
|
error: str | None = None,
|
|
) -> None:
|
|
payload: dict[str, object] = {
|
|
"case_id": case_id,
|
|
"mode": mode,
|
|
"is_warmup": is_warmup,
|
|
"observation_summary": observation_summary,
|
|
"result_flag": result_flag,
|
|
"metrics": metrics or {},
|
|
"error": error,
|
|
}
|
|
if total_latency_s is not None:
|
|
payload["total_latency_s"] = total_latency_s
|
|
if first_token_latency_s is not None:
|
|
payload["first_token_latency_s"] = first_token_latency_s
|
|
RESULTS.append(payload)
|
|
emit("case_result", **payload)
|
|
|
|
|
|
def summarize_results() -> dict[str, object]:
|
|
by_case: defaultdict[str, list[dict[str, object]]] = defaultdict(list)
|
|
for result in RESULTS:
|
|
by_case[str(result["case_id"])].append(result)
|
|
|
|
summary_cases: dict[str, object] = {}
|
|
for case_id, items in by_case.items():
|
|
measured = [item for item in items if not bool(item.get("is_warmup"))]
|
|
latencies = [
|
|
float(item["total_latency_s"])
|
|
for item in measured
|
|
if item.get("total_latency_s") is not None
|
|
]
|
|
first_token_latencies = [
|
|
float(item["first_token_latency_s"])
|
|
for item in measured
|
|
if item.get("first_token_latency_s") is not None
|
|
]
|
|
result_flags = Counter(str(item["result_flag"]) for item in measured or items)
|
|
observations = [str(item["observation_summary"]) for item in (measured or items)[:3]]
|
|
summary_cases[case_id] = {
|
|
"mode": str(items[-1]["mode"]),
|
|
"runs": len(measured),
|
|
"warmups": len(items) - len(measured),
|
|
"result_flags": dict(result_flags),
|
|
"median_total_latency_s": (statistics.median(latencies) if latencies else None),
|
|
"mean_total_latency_s": statistics.mean(latencies) if latencies else None,
|
|
"median_first_token_latency_s": (
|
|
statistics.median(first_token_latencies) if first_token_latencies else None
|
|
),
|
|
"observations": observations,
|
|
}
|
|
|
|
return {
|
|
"scenario": SCENARIO,
|
|
"run_label": RUN_LABEL,
|
|
"mode": MODE,
|
|
"result_count": len(RESULTS),
|
|
"cases": summary_cases,
|
|
"result_flags": dict(Counter(str(item["result_flag"]) for item in RESULTS)),
|
|
}
|
|
|
|
|
|
def finalize(exit_code: int) -> None:
|
|
metadata_payload = {
|
|
"exit_code": exit_code,
|
|
"runtime_context": runtime_context(),
|
|
}
|
|
summary_payload = summarize_results()
|
|
emit("summary", metadata=metadata_payload, summary=summary_payload)
|
|
|
|
output_dir = _output_dir()
|
|
if not output_dir:
|
|
return
|
|
|
|
metadata_path = output_dir / "metadata.json"
|
|
results_path = output_dir / "results.json"
|
|
summary_path = output_dir / "summary.json"
|
|
_write_json(metadata_path, metadata_payload)
|
|
_write_json(results_path, RESULTS)
|
|
_write_json(summary_path, summary_payload)
|
|
emit(
|
|
"artifact_paths",
|
|
metadata_path=str(metadata_path),
|
|
results_path=str(results_path),
|
|
summary_path=str(summary_path),
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
case_id = os.getenv("PROBE_CASE_ID", f"case-{uuid.uuid4().hex[:8]}")
|
|
emit("banner", context=runtime_context())
|
|
start_case(case_id)
|
|
|
|
# Replace this block with the narrow runtime question you want to test.
|
|
observation = "replace-me"
|
|
result_flag = "expected"
|
|
|
|
record_case_result(
|
|
case_id=case_id,
|
|
observation_summary=observation,
|
|
result_flag=result_flag,
|
|
)
|
|
finalize(exit_code=0)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|