diff --git a/apps/openant-cli/cmd/enhance.go b/apps/openant-cli/cmd/enhance.go index 5381213f..0c437b48 100644 --- a/apps/openant-cli/cmd/enhance.go +++ b/apps/openant-cli/cmd/enhance.go @@ -32,6 +32,7 @@ var ( enhanceCheckpoint string enhanceWorkers int enhanceBackoff int + enhanceLimit int ) func init() { @@ -42,6 +43,7 @@ func init() { enhanceCmd.Flags().StringVar(&enhanceCheckpoint, "checkpoint", "", "Path to save/resume checkpoint (agentic mode)") enhanceCmd.Flags().IntVar(&enhanceWorkers, "workers", 8, "Number of parallel workers for LLM steps (default: 8)") enhanceCmd.Flags().IntVar(&enhanceBackoff, "backoff", 30, "Seconds to wait when rate-limited (default: 30)") + enhanceCmd.Flags().IntVar(&enhanceLimit, "limit", 0, "Max units to enhance (0 = no limit)") } func runEnhance(cmd *cobra.Command, args []string) { @@ -104,6 +106,9 @@ func runEnhance(cmd *cobra.Command, args []string) { if enhanceBackoff != 30 { pyArgs = append(pyArgs, "--backoff", fmt.Sprintf("%d", enhanceBackoff)) } + if enhanceLimit > 0 { + pyArgs = append(pyArgs, "--limit", fmt.Sprintf("%d", enhanceLimit)) + } result, err := python.Invoke(rt.Path, pyArgs, "", quiet, requireAPIKey()) if err != nil { diff --git a/apps/openant-cli/cmd/enhance_limit_test.go b/apps/openant-cli/cmd/enhance_limit_test.go new file mode 100644 index 00000000..91d69b40 --- /dev/null +++ b/apps/openant-cli/cmd/enhance_limit_test.go @@ -0,0 +1,15 @@ +package cmd + +import "testing" + +// The `enhance` command must expose a `--limit` flag for +// parity with `analyze`/`scan`, so a user can cost-limit / test-run enhancement on N units. +func TestEnhanceCommandHasLimitFlag(t *testing.T) { + f := enhanceCmd.Flags().Lookup("limit") + if f == nil { + t.Fatal("enhance command is missing the --limit flag (parity with analyze/scan)") + } + if f.DefValue != "0" { + t.Errorf("--limit default = %q, want \"0\" (0 = no limit)", f.DefValue) + } +} diff --git a/libs/openant-core/core/enhancer.py b/libs/openant-core/core/enhancer.py index 70879b81..bc766a96 100644 --- a/libs/openant-core/core/enhancer.py +++ b/libs/openant-core/core/enhancer.py @@ -30,6 +30,7 @@ def enhance_dataset( model: str = "sonnet", workers: int = 8, backoff_seconds: int = 30, + limit: int | None = None, ) -> EnhanceResult: """Enhance a parsed dataset with security context. @@ -44,6 +45,7 @@ def enhance_dataset( model: "sonnet" (default, cost-effective). workers: Number of parallel workers (default: 8). backoff_seconds: Seconds to wait on rate limit before retry (default: 30). + limit: Max number of units to enhance (None = all). Mirrors `analyze --limit`. Returns: EnhanceResult with output path, stats, and usage. @@ -72,6 +74,9 @@ def enhance_dataset( print(f"[Enhance] Loading dataset: {dataset_path}", file=sys.stderr) dataset = read_json(dataset_path) units = dataset.get("units", []) + if limit: + units = units[:limit] + dataset["units"] = units # so the agentic/single-shot paths enhance only these print(f"[Enhance] Units to enhance: {len(units)}", file=sys.stderr) # Set up progress reporter diff --git a/libs/openant-core/openant/cli.py b/libs/openant-core/openant/cli.py index c303c647..5a23d367 100644 --- a/libs/openant-core/openant/cli.py +++ b/libs/openant-core/openant/cli.py @@ -189,6 +189,7 @@ def cmd_enhance(args): checkpoint_path=args.checkpoint, workers=args.workers, backoff_seconds=args.backoff, + limit=args.limit, ) ctx.summary = { @@ -1049,6 +1050,7 @@ def main(): enhance_p.add_argument("--repo-path", help="Path to the repository (required for agentic mode)") enhance_p.add_argument("--output", "-o", help="Output path for enhanced dataset (default: {input}_enhanced.json)") enhance_p.add_argument("--checkpoint", help="Path to save/resume checkpoint (agentic mode)") + enhance_p.add_argument("--limit", type=int, help="Max units to enhance") enhance_p.add_argument( "--mode", choices=["agentic", "single-shot"], diff --git a/libs/openant-core/tests/test_enhance_limit.py b/libs/openant-core/tests/test_enhance_limit.py new file mode 100644 index 00000000..bf096da3 --- /dev/null +++ b/libs/openant-core/tests/test_enhance_limit.py @@ -0,0 +1,82 @@ +"""Regression tests: `enhance` lacked a `--limit` (max-units) capability at all +three layers (Go CLI, Python CLI, core), while `analyze`/`scan` have it. + +These tests cover the two Python layers behaviorally: +- the core `enhance_dataset(limit=N)` actually slices the dataset to N units (the LLM is stubbed); +- the CLI path `enhance ... --limit N` (parser `enhance_p` + `cmd_enhance`) forwards N to enhance_dataset. + +(The Go layer is covered by apps/openant-cli/cmd/enhance_limit_test.go.) +""" +import sys + +import core.enhancer as enh +import utilities.context_enhancer as _ce +import utilities.llm_client as _llm + + +_received = {} + + +class _DummyEnhancer: + def __init__(self, **kw): + pass + + def enhance_dataset(self, dataset, progress_callback=None, workers=8): + _received["units"] = len(dataset.get("units", [])) + return dataset + + def enhance_dataset_agentic(self, dataset, **kw): + _received["units"] = len(dataset.get("units", [])) + return dataset + + +def test_enhance_dataset_limit_slices_units(tmp_path, monkeypatch): + five = [{"id": f"u{i}", "llm_context": {}} for i in range(5)] + monkeypatch.setattr(enh, "read_json", lambda p: {"units": list(five)}) + monkeypatch.setattr(enh, "write_json", lambda p, d: None) + monkeypatch.setattr(enh, "configure_rate_limiter", lambda **k: None) + monkeypatch.setattr(_llm, "AnthropicClient", lambda **k: object()) + monkeypatch.setattr(_llm, "get_global_tracker", lambda: None) + monkeypatch.setattr(_ce, "ContextEnhancer", _DummyEnhancer) + + _received.clear() + res = enh.enhance_dataset( + dataset_path=str(tmp_path / "x.json"), + output_path=str(tmp_path / "out.json"), + mode="single-shot", + limit=2, + ) + # The enhancer must RECEIVE only the limited units -- this guards the load-bearing + # `dataset["units"] = units` slice (a local-only slice would leave 5 units reaching the enhancer). + assert _received["units"] == 2, f"enhancer received {_received.get('units')} units, expected 2" + assert res.units_enhanced == 2, f"limit=2 should enhance 2 of 5 units, got {res.units_enhanced}" + + +def test_enhance_cli_forwards_limit_end_to_end(tmp_path, monkeypatch): + from openant import cli + + captured = {} + + def _fake_enhance_dataset(**kwargs): + captured.update(kwargs) + + class _R: + enhanced_dataset_path = str(tmp_path / "out.json") + units_enhanced = 3 + error_count = 0 + classifications: dict = {} + error_summary: dict = {} + + def to_dict(self): + return {"units_enhanced": 3} + + return _R() + + # cmd_enhance does `from core.enhancer import enhance_dataset` at call time -> patch the source attr + monkeypatch.setattr(enh, "enhance_dataset", _fake_enhance_dataset) + monkeypatch.setattr(sys, "argv", ["openant", "enhance", str(tmp_path / "x.json"), "--limit", "3"]) + + rc = cli.main() + assert rc == 0, f"enhance command returned {rc}" + assert captured.get("limit") == 3, \ + f"`enhance --limit 3` must forward limit=3 to enhance_dataset, got {captured.get('limit')!r}"