Files
2026-07-13 13:12:00 +08:00

243 lines
8.4 KiB
Python

"""Command-line entry point for the harness bench."""
from __future__ import annotations
import argparse
import asyncio
import sys
from tests.harness_bench.bench import run_bench
from tests.harness_bench.events import LineSink
from tests.harness_bench.manifest import OFFICIAL_PROFILES
from tests.harness_bench.profile import BenchProfile, resolve_profile
from tests.harness_bench.report import render_json, render_markdown, render_table
from tests.harness_bench.transport import driver_registry, resolve_transport_name
def _parse_args(argv: list[str]) -> argparse.Namespace:
parser = argparse.ArgumentParser(
prog="python -m tests.harness_bench",
description="Probe a harness and report a verdict per capability dimension.",
)
parser.add_argument(
"--harness",
action="append",
metavar="NAME",
help="Harness to probe (repeatable). An official name, or a "
"'module:attr' / 'module.ATTR' reference to a community "
"BenchProfile. Defaults to every official harness.",
)
parser.add_argument(
"--profile",
metavar="NAME",
default=None,
help="Databricks gateway profile override. Optional: without it the "
"bench derives creds the way `omni run` does (a configured "
"~/.omnigent profile or ambient OPENAI_*). The live layer turns on "
"whenever creds are resolvable; use --no-live for the offline matrix.",
)
parser.add_argument(
"--live",
dest="live",
action="store_true",
default=None,
help="Force the live layer (needs resolvable creds: --profile, a "
"configured ~/.omnigent profile, or ambient OPENAI_*).",
)
parser.add_argument(
"--no-live",
dest="live",
action="store_false",
help="Force the offline (declared-only) render.",
)
transport_grp = parser.add_mutually_exclusive_group()
transport_grp.add_argument(
"--transport",
metavar="NAME",
default=None,
help="Transport driver override (e.g. 'sdk-inproc', 'full-server'). "
"Wins over the family default. By default SDK harnesses run on "
"full-server (fullest coverage: Tool calling + Policy DENY); natives "
"run on native-tui.",
)
transport_grp.add_argument(
"--fast",
action="store_true",
help="Run SDK harnesses on sdk-inproc instead of the full-server "
"default: skips the server boot for a quicker run, at the cost of the "
"Tool calling + Policy DENY dimensions (reported SKIPPED). No effect on "
"native harnesses. Mutually exclusive with --transport.",
)
fmt = parser.add_mutually_exclusive_group()
fmt.add_argument(
"--markdown",
action="store_true",
help="Emit the GitHub-flavored Markdown table (for docs / PRs).",
)
fmt.add_argument("--json", action="store_true", help="Emit JSON.")
parser.add_argument(
"--no-color", action="store_true", help="Disable ANSI color in the terminal table."
)
parser.add_argument(
"--jobs",
"-j",
type=int,
default=1,
metavar="N",
help="Run up to N harnesses concurrently (default 1 = sequential). "
"Probes within a harness stay sequential. Higher N cuts wall-clock but "
"raises process / gateway load; 3-4 is a reasonable ceiling on one host.",
)
rich_grp = parser.add_mutually_exclusive_group()
rich_grp.add_argument(
"--rich",
dest="rich",
action="store_true",
default=None,
help="Force the live rich progress table (needs a TTY + rich).",
)
rich_grp.add_argument(
"--no-rich",
dest="rich",
action="store_false",
help="Force plain per-line progress (no live table).",
)
parser.add_argument(
"--report",
metavar="PATH",
default=None,
help="Also write the final matrix to PATH. Format follows --json / "
"--markdown, else inferred from the extension (.json / .md), else plain text.",
)
parser.add_argument("--list", action="store_true", help="List official harnesses and exit.")
return parser.parse_args(argv)
def _resolve_profiles(names: list[str] | None) -> list[BenchProfile]:
if not names:
return list(OFFICIAL_PROFILES.values())
return [resolve_profile(name) for name in names]
def main(argv: list[str] | None = None) -> int:
args = _parse_args(argv if argv is not None else sys.argv[1:])
if args.list:
for name, profile in sorted(OFFICIAL_PROFILES.items()):
transport = resolve_transport_name(profile, override=None, fast=False)
print(f"{name}\t{transport}\t{profile.model}")
return 0
try:
profiles = _resolve_profiles(args.harness)
except KeyError as exc:
print(str(exc), file=sys.stderr)
return 2
if args.transport is not None and args.transport not in driver_registry():
known = ", ".join(sorted(driver_registry()))
print(
f"unknown --transport {args.transport!r}; known transports: {known}", file=sys.stderr
)
return 2
if args.jobs < 1:
print("--jobs must be >= 1", file=sys.stderr)
return 2
from tests.harness_bench.runtime_env import bench_creds_skip_reason
creds_skip = bench_creds_skip_reason(args.profile)
live = args.live if args.live is not None else creds_skip is None
if live and creds_skip is not None:
print(f"--live needs resolvable gateway creds: {creds_skip}", file=sys.stderr)
return 2
sink = None
if live:
sink = _select_progress_sink(args.rich)
matrix = asyncio.run(
run_bench(
profiles,
databricks_profile=args.profile,
live=live,
transport=args.transport,
fast=args.fast,
progress=sink,
jobs=args.jobs,
)
)
if sink is not None:
sink.close()
declared = not live
if args.json:
output = render_json(matrix)
elif args.markdown:
output = render_markdown(matrix, declared=declared)
else:
color = sys.stdout.isatty() and not args.no_color
# Avoid duplicating a live grid, but preserve it in redirected output.
grid = not (_grid_already_shown(sink) and sys.stdout.isatty())
output = render_table(matrix, color=color, declared=declared, grid=grid)
print(output, end="")
if args.report:
_write_report(args.report, matrix, json_flag=args.json, markdown_flag=args.markdown)
return 1 if matrix.has_drift else 0
def _grid_already_shown(sink) -> bool:
"""Whether the progress sink already painted the glyph grid to the terminal.
True only for the rich live table (which sets ``drew_grid = True``); the
plain :class:`LineSink` and a silent run do not, so their report prints the
grid in full.
"""
return bool(getattr(sink, "drew_grid", False))
def _select_progress_sink(rich_flag: bool | None):
"""Pick the progress sink for a live run.
``rich_flag``: ``True`` forces rich, ``False`` forces plain, ``None`` =
auto (rich on a TTY, plain otherwise). Falls back to the plain
:class:`LineSink` whenever rich is unavailable or not a terminal.
"""
def _line(msg: str) -> None:
print(msg, file=sys.stderr, flush=True)
if rich_flag is not False:
from tests.harness_bench.richreport import rich_sink_or_none
rich_sink = rich_sink_or_none(force=bool(rich_flag))
if rich_sink is not None:
return rich_sink
if rich_flag is True:
print("--rich requested but rich/TTY unavailable; using plain output", file=sys.stderr)
return LineSink(_line)
def _write_report(path: str, matrix, *, json_flag: bool, markdown_flag: bool) -> None:
"""Write the matrix to *path*; format from flags, else the extension."""
if json_flag:
content = render_json(matrix)
elif markdown_flag:
content = render_markdown(matrix, declared=False)
elif path.endswith(".json"):
content = render_json(matrix)
elif path.endswith((".md", ".markdown")):
content = render_markdown(matrix, declared=False)
else:
content = render_table(matrix, color=False, declared=False)
with open(path, "w", encoding="utf-8") as fh:
fh.write(content if content.endswith("\n") else content + "\n")
print(f"report written to {path}", file=sys.stderr)
if __name__ == "__main__":
raise SystemExit(main())