From 52ee5227737e18ca1937919c215c3e7d3e230c7e Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 31 Mar 2026 00:53:32 +0000 Subject: [PATCH 1/2] Re-enable mypy typecheck in CI pipeline Fix all 16 mypy errors: add type parameters to generic dicts, fix variable reuse type conflict in cli.py, add float() casts for no-any-return, remove stale type: ignore comment. Add types-PyYAML to dev deps, configure mypy overrides for optional backend stubs. Use python -m mypy in Makefile for correct interpreter resolution. https://claude.ai/code/session_01E5v76N5vK2Et2qHkqgKXcc --- .github/workflows/ci.yml | 3 +++ Makefile | 2 +- pyproject.toml | 5 +++++ src/broadside_ai/benchmark.py | 2 +- src/broadside_ai/cli.py | 4 ++-- src/broadside_ai/parsing.py | 9 +++++---- src/broadside_ai/quality.py | 2 +- src/broadside_ai/strategies/weighted_merge.py | 10 +++++----- 8 files changed, 23 insertions(+), 14 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2493c28..1b45ac1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -27,5 +27,8 @@ jobs: - name: Lint run: make lint + - name: Typecheck + run: make typecheck + - name: Test run: make test diff --git a/Makefile b/Makefile index 40f21a8..4f864eb 100644 --- a/Makefile +++ b/Makefile @@ -11,7 +11,7 @@ lint: ruff format --check src/ tests/ typecheck: - mypy src/broadside_ai/ + python -m mypy src/broadside_ai/ clean: rm -rf dist/ build/ *.egg-info src/*.egg-info diff --git a/pyproject.toml b/pyproject.toml index 2a90403..d9b1ce5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ dev = [ "pytest-asyncio>=0.24.0", "ruff>=0.8.0", "mypy>=1.13.0", + "types-PyYAML>=6.0", ] [project.scripts] @@ -72,3 +73,7 @@ testpaths = ["tests"] [tool.mypy] python_version = "3.10" strict = true + +[[tool.mypy.overrides]] +module = ["anthropic", "openai"] +ignore_missing_imports = true diff --git a/src/broadside_ai/benchmark.py b/src/broadside_ai/benchmark.py index 3ee7aeb..4c43c33 100644 --- a/src/broadside_ai/benchmark.py +++ b/src/broadside_ai/benchmark.py @@ -296,7 +296,7 @@ class MEMORYSTATUSEX(ctypes.Structure): kernel32.GlobalMemoryStatusEx(ctypes.byref(stat)) info["ram_gb"] = round(stat.ullTotalPhys / (1024**3), 1) else: - mem_bytes = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") # type: ignore[attr-defined] + mem_bytes = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") info["ram_gb"] = round(mem_bytes / (1024**3), 1) except Exception: pass diff --git a/src/broadside_ai/cli.py b/src/broadside_ai/cli.py index 4d4e8ad..ae2b3c3 100644 --- a/src/broadside_ai/cli.py +++ b/src/broadside_ai/cli.py @@ -233,8 +233,8 @@ async def _run_pipeline( for i in range(n): progress.update(ptask, description=f"Agent {i + 1}/{n} thinking") try: - result = await llm.complete(prompt) - results.append(result) + agent_result = await llm.complete(prompt) + results.append(agent_result) except Exception as exc: console.print(f" [yellow]Agent {i + 1} failed: {exc}[/yellow]") progress.update(ptask, advance=1) diff --git a/src/broadside_ai/parsing.py b/src/broadside_ai/parsing.py index 8a7b03a..4a8e9b6 100644 --- a/src/broadside_ai/parsing.py +++ b/src/broadside_ai/parsing.py @@ -11,9 +11,10 @@ import json import re +from typing import Any -def try_parse_json(text: str) -> dict | None: +def try_parse_json(text: str) -> dict[str, Any] | None: """Attempt to parse JSON from agent output text. Returns the parsed dict on success, None on failure. Never raises. @@ -41,7 +42,7 @@ def try_parse_json(text: str) -> dict | None: return None -def _try_loads(text: str) -> dict | None: +def _try_loads(text: str) -> dict[str, Any] | None: try: parsed = json.loads(text) return parsed if isinstance(parsed, dict) else None @@ -52,14 +53,14 @@ def _try_loads(text: str) -> dict | None: _FENCE_RE = re.compile(r"```(?:json)?\s*\n?(.*?)\n?\s*```", re.DOTALL) -def _try_fenced(text: str) -> dict | None: +def _try_fenced(text: str) -> dict[str, Any] | None: match = _FENCE_RE.search(text) if match: return _try_loads(match.group(1).strip()) return None -def _try_extract_object(text: str) -> dict | None: +def _try_extract_object(text: str) -> dict[str, Any] | None: """Find the first { ... } block and try to parse it.""" start = text.find("{") if start == -1: diff --git a/src/broadside_ai/quality.py b/src/broadside_ai/quality.py index 4a4ccf1..cea818c 100644 --- a/src/broadside_ai/quality.py +++ b/src/broadside_ai/quality.py @@ -93,7 +93,7 @@ def _check_agreement(results: list[AgentResult], threshold: float) -> bool: return most_common_count / n >= threshold -def _dict_signature(d: dict) -> str: +def _dict_signature(d: dict[str, object]) -> str: """Create a canonical string from a dict for comparison. Ignores 'confidence' fields since those are expected to vary. diff --git a/src/broadside_ai/strategies/weighted_merge.py b/src/broadside_ai/strategies/weighted_merge.py index aaf61c1..353f391 100644 --- a/src/broadside_ai/strategies/weighted_merge.py +++ b/src/broadside_ai/strategies/weighted_merge.py @@ -95,7 +95,7 @@ def _merge_fields(outputs: list[dict[str, Any]], weights: list[float]) -> dict[s merged: dict[str, Any] = {} for key in sorted(all_keys): - values = [(out.get(key), w) for out, w in zip(outputs, weights) if key in out] + values = [(out[key], w) for out, w in zip(outputs, weights) if key in out] if not values: continue @@ -105,9 +105,9 @@ def _merge_fields(outputs: list[dict[str, Any]], weights: list[float]) -> dict[s if all(isinstance(v, (int, float)) for v in raw_vals): merged[key] = _merge_numeric(raw_vals, val_weights) elif all(isinstance(v, str) for v in raw_vals): - merged[key] = _merge_categorical(raw_vals) + merged[key] = _merge_categorical([str(v) for v in raw_vals]) elif all(isinstance(v, list) for v in raw_vals): - merged[key] = _merge_lists(raw_vals) + merged[key] = _merge_lists([list(v) for v in raw_vals]) else: # Mixed types or unsupported — take the most common merged[key] = _merge_categorical([str(v) for v in raw_vals]) @@ -119,8 +119,8 @@ def _merge_numeric(values: list[Any], weights: list[float]) -> float: """Weighted average of numeric values.""" total_weight = sum(weights) if total_weight == 0: - return sum(values) / len(values) - return round(sum(v * w for v, w in zip(values, weights)) / total_weight, 4) + return float(sum(values) / len(values)) + return round(float(sum(v * w for v, w in zip(values, weights)) / total_weight), 4) def _merge_categorical(values: list[str]) -> str: From d56c59251c5a4a691444587a1b740cdf84c35080 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 31 Mar 2026 01:13:23 +0000 Subject: [PATCH 2/2] Update docs for Phase 1 completion, implement HITL checkpoints Docs: Add weighted_merge strategy, structured outputs, and early termination to README, llms.txt, and CONTRIBUTING. Remove already- implemented items from contribution ideas. Checkpoints: Implement Checkpoint protocol with three async hooks (pre_scatter, post_gather, post_synthesis). Add TerminalCheckpoint for interactive terminal use. Wire into run.py pipeline and CLI (--checkpoint flag). 9 new tests, 88 total. https://claude.ai/code/session_01E5v76N5vK2Et2qHkqgKXcc --- CONTRIBUTING.md | 5 +- README.md | 37 +++++-- llms.txt | 5 +- src/broadside_ai/checkpoints/__init__.py | 118 +++++++++++++++++++- src/broadside_ai/cli.py | 30 +++++ src/broadside_ai/run.py | 25 ++++- tests/test_checkpoints.py | 135 +++++++++++++++++++++++ 7 files changed, 341 insertions(+), 14 deletions(-) create mode 100644 tests/test_checkpoints.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c848930..51efc43 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -13,7 +13,7 @@ pip install -e ".[dev]" make test ``` -All 12 tests should pass. If they don't, open an issue. +All tests should pass. If they don't, open an issue. ## Ways to Contribute @@ -36,9 +36,8 @@ python -m broadside_ai.task_validator tasks/my_task.yaml New ways to collapse N outputs into one. Look at `src/broadside_ai/strategies/consensus.py` for the pattern. Ideas: -- Weighted merge (scored recommendations) -- Structured extraction (pull specific fields from each output) - Embedding-based clustering (group similar outputs, pick representative) +- Chain-of-thought synthesis (multi-step reasoning across outputs) Open an issue first to discuss the approach. diff --git a/README.md b/README.md index 3dc09f6..25cd496 100644 --- a/README.md +++ b/README.md @@ -77,20 +77,43 @@ Results are saved to `broadside_ai_output/` organized by model and topic for eas ## Synthesis Strategies -Broadside ships with three strategies for collapsing N outputs into one: +Broadside ships with four strategies for collapsing N outputs into one: -| Strategy | Best for | How it works | -| --------------- | --------------- | --------------------------------------------------------------------------------------------------- | -| `llm` (default) | General tasks | Sends all outputs to a model that identifies consensus, flags outliers, and writes a unified answer | -| `consensus` | Knowledge tasks | Extracts agreed-upon claims, disagreements, and unique insights | -| `voting` | Reasoning tasks | Each output votes on an answer; majority wins | +| Strategy | Best for | How it works | +| ---------------- | ------------------ | --------------------------------------------------------------------------------------------------- | +| `llm` (default) | General tasks | Sends all outputs to a model that identifies consensus, flags outliers, and writes a unified answer | +| `consensus` | Knowledge tasks | Extracts agreed-upon claims, disagreements, and unique insights | +| `voting` | Reasoning tasks | Each output votes on an answer; majority wins | +| `weighted_merge` | Scored/structured | Confidence-weighted field merge — zero LLM calls on happy path | -Note: the default `llm` strategy calls the model one additional time for synthesis. In practice, synthesis can use as many tokens as the scatter itself. For cost-sensitive workloads, `consensus` or `voting` are lighter alternatives. +Note: the default `llm` strategy calls the model one additional time for synthesis, which can use as many tokens as the scatter itself. `consensus` and `voting` are lighter alternatives. `weighted_merge` uses no LLM calls at all — it merges structured fields algorithmically. ```bash python -m broadside_ai run --prompt "Your task" --n 3 --synthesis voting ``` +### Structured Outputs + +Pass an `output_schema` to get structured JSON from each agent. Works with any strategy but pairs best with `weighted_merge`: + +```python +task = Task( + prompt="Classify this support ticket", + output_schema={"label": "string", "confidence": "float", "reasoning": "string"}, +) +result = run_sync(task, n=3, backend="ollama", strategy="weighted_merge") +``` + +### Early Termination + +Stop scattering early when enough agents agree — saves time and tokens: + +```bash +python -m broadside_ai run --prompt "Your task" --n 5 --early-stop 3 --agreement 0.66 +``` + +`--early-stop 3` requires at least 3 results before checking. `--agreement 0.66` stops when 66%+ of results match. + ## Backends ### Ollama (default — no API key) diff --git a/llms.txt b/llms.txt index e16babd..d37096d 100644 --- a/llms.txt +++ b/llms.txt @@ -9,10 +9,13 @@ Broadside implements the scatter/gather pattern as its core organizing principle - Task: a tightly scoped unit of work with a clear output - Scatter: fan a task to N parallel agents (performance plateaus beyond 4-5, DeepMind 2025) - Gather: collect and normalize outputs -- Synthesize: collapse N outputs into one — three strategies: +- Synthesize: collapse N outputs into one — four strategies: - `llm` (default): general-purpose, sends all outputs to a model for unified synthesis - `consensus`: best for knowledge tasks, extracts agreements and disagreements - `voting`: best for reasoning tasks, majority-wins + - `weighted_merge`: best for scored/structured outputs, confidence-weighted field merge (zero LLM calls) +- Structured outputs: pass `output_schema` dict to Task for JSON-structured agent responses +- Early termination: quality signals (min_complete, agreement_threshold) cancel remaining branches when results converge - Budget: per-scatter token limits (scatter/gather multiplies costs ~15x; synthesis adds more) ## Backends diff --git a/src/broadside_ai/checkpoints/__init__.py b/src/broadside_ai/checkpoints/__init__.py index bf466c6..50a4785 100644 --- a/src/broadside_ai/checkpoints/__init__.py +++ b/src/broadside_ai/checkpoints/__init__.py @@ -1,4 +1,118 @@ -"""Human-in-the-loop checkpoints. +"""Human-in-the-loop checkpoints for the scatter/gather/synthesize pipeline. -Phase 5 will add: pre-scatter review, gather-point review, per-tool approval. +Checkpoints let a human review and approve/reject at three points: +1. Pre-scatter — review the task before scattering to N agents +2. Post-gather — review raw results before synthesis +3. Post-synthesis — review the final output before accepting + +Usage:: + + from broadside_ai.checkpoints import TerminalCheckpoint + from broadside_ai import Task, run_sync + + result = run_sync(task, n=3, checkpoint=TerminalCheckpoint()) """ + +from __future__ import annotations + +import asyncio +from typing import Protocol, runtime_checkable + +from broadside_ai.gather import GatherResult +from broadside_ai.synthesize import Synthesis +from broadside_ai.task import Task + + +class CheckpointRejected(Exception): + """Raised when a human rejects at a checkpoint.""" + + def __init__(self, stage: str, reason: str = "") -> None: + self.stage = stage + self.reason = reason + msg = f"Rejected at {stage}" + if reason: + msg += f": {reason}" + super().__init__(msg) + + +@runtime_checkable +class Checkpoint(Protocol): + """Protocol for human-in-the-loop checkpoints. + + Implement this to create custom checkpoint handlers (Slack bots, + web UIs, approval queues, etc.). Each method returns True to proceed + or False to reject. + """ + + async def pre_scatter(self, task: Task, n: int) -> bool: + """Review the task before scattering. Return True to proceed.""" + ... + + async def post_gather(self, gathered: GatherResult) -> bool: + """Review gathered results before synthesis. Return True to synthesize.""" + ... + + async def post_synthesis(self, synthesis: Synthesis) -> bool: + """Review the final synthesis. Return True to accept.""" + ... + + +class TerminalCheckpoint: + """Interactive terminal checkpoint — prompts via stdin. + + Displays a summary at each stage and asks the user to approve or reject. + Uses asyncio.to_thread so it doesn't block the event loop. + """ + + async def pre_scatter(self, task: Task, n: int) -> bool: + """Show task details and ask for approval.""" + prompt_preview = task.prompt[:200] + if len(task.prompt) > 200: + prompt_preview += "..." + + print("\n--- Pre-Scatter Checkpoint ---") + print(f"Task: {prompt_preview}") + print(f"Agents: {n}") + if task.output_schema: + print(f"Output schema: {task.output_schema}") + print() + + return await self._ask("Proceed with scatter?") + + async def post_gather(self, gathered: GatherResult) -> bool: + """Show gather summary and ask for approval.""" + print("\n--- Post-Gather Checkpoint ---") + print(f"Completed: {gathered.n_completed}/{gathered.n_completed + gathered.n_failed}") + print(f"Tokens used: {gathered.total_tokens:,}") + print(f"Wall clock: {gathered.wall_clock_ms / 1000:.1f}s") + + for i, text in enumerate(gathered.texts, 1): + preview = text[:150].replace("\n", " ") + if len(text) > 150: + preview += "..." + print(f"\n Agent {i}: {preview}") + print() + + return await self._ask("Proceed with synthesis?") + + async def post_synthesis(self, synthesis: Synthesis) -> bool: + """Show synthesis result and ask for approval.""" + print("\n--- Post-Synthesis Checkpoint ---") + print(f"Strategy: {synthesis.strategy}") + print(f"Synthesis tokens: {synthesis.synthesis_tokens:,}") + + result_preview = synthesis.result[:300] + if len(synthesis.result) > 300: + result_preview += "..." + print(f"\nResult:\n{result_preview}") + print() + + return await self._ask("Accept this result?") + + async def _ask(self, question: str) -> bool: + """Prompt user and return True for yes, False for no.""" + response = await asyncio.to_thread(input, f"{question} [Y/n] ") + return response.strip().lower() in ("", "y", "yes") + + +__all__ = ["Checkpoint", "CheckpointRejected", "TerminalCheckpoint"] diff --git a/src/broadside_ai/cli.py b/src/broadside_ai/cli.py index ae2b3c3..f6b473c 100644 --- a/src/broadside_ai/cli.py +++ b/src/broadside_ai/cli.py @@ -77,6 +77,11 @@ def main() -> None: default=None, help="Stop when this fraction of results agree (0.0-1.0).", ) +@click.option( + "--checkpoint", + is_flag=True, + help="Enable interactive checkpoints (pause for human review at each stage).", +) def run( task_file: str | None, prompt: str | None, @@ -92,6 +97,7 @@ def run( json_out: bool, early_stop: int | None, agreement: float | None, + checkpoint: bool, ) -> None: """Run a scatter/gather cycle on a task.""" if not task_file and not prompt: @@ -120,6 +126,13 @@ def run( es = EarlyStop(min_complete=early_stop, agreement_threshold=agreement) + # Build checkpoint + cp = None + if checkpoint: + from broadside_ai.checkpoints import TerminalCheckpoint + + cp = TerminalCheckpoint() + # Run the scatter/gather/synthesize pipeline try: asyncio.run( @@ -136,6 +149,7 @@ def run( raw, json_out, es, + cp, ) ) except RuntimeError as exc: @@ -156,6 +170,7 @@ async def _run_pipeline( raw: bool, json_out: bool, early_stop: Any | None = None, + checkpoint: Any | None = None, ) -> None: """Execute the full pipeline with rich output.""" import time @@ -210,6 +225,11 @@ async def _run_pipeline( ) console.print() + # --- Pre-scatter checkpoint --- + if checkpoint and not await checkpoint.pre_scatter(task, n): + console.print("[yellow]Scatter rejected at checkpoint.[/yellow]") + return + # --- Scatter phase with live progress --- scatter_start = time.perf_counter() @@ -296,6 +316,11 @@ async def _run_pipeline( _save_raw(gathered.texts, task, output_dir, model_hint=model_hint) return + # --- Post-gather checkpoint --- + if checkpoint and not await checkpoint.post_gather(gathered): + console.print("[yellow]Synthesis rejected at checkpoint.[/yellow]") + return + # --- Synthesis phase --- synth_start = time.perf_counter() progress = Progress( @@ -322,6 +347,11 @@ async def _run_pipeline( ) console.print() + # --- Post-synthesis checkpoint --- + if checkpoint and not await checkpoint.post_synthesis(result): + console.print("[yellow]Result rejected at checkpoint.[/yellow]") + return + # --- Results --- _show_rich(result, total_wall) diff --git a/src/broadside_ai/run.py b/src/broadside_ai/run.py index 3f0daf5..4ef7174 100644 --- a/src/broadside_ai/run.py +++ b/src/broadside_ai/run.py @@ -10,6 +10,7 @@ from typing import Any from broadside_ai.budget import ScatterBudget +from broadside_ai.checkpoints import Checkpoint, CheckpointRejected from broadside_ai.gather import gather from broadside_ai.quality import EarlyStop from broadside_ai.scatter import scatter @@ -29,6 +30,7 @@ async def run( max_tokens: int | None = None, parallel: bool | None = None, early_stop: EarlyStop | None = None, + checkpoint: Checkpoint | None = None, ) -> Synthesis: """Run a full scatter/gather/synthesize cycle. @@ -47,12 +49,21 @@ async def run( max_tokens: Per-scatter token budget. None = unbounded. parallel: Run branches concurrently. Defaults to True for cloud backends, False for local (Ollama) to avoid overloading. + checkpoint: Optional HITL checkpoint handler. When provided, the + pipeline pauses at each stage for human review. Returns: Synthesis with the final result, raw outputs, and cost data. + + Raises: + CheckpointRejected: If a human rejects at any checkpoint. """ budget = ScatterBudget(max_tokens=max_tokens) if max_tokens else None + # Pre-scatter checkpoint + if checkpoint and not await checkpoint.pre_scatter(task, n): + raise CheckpointRejected("pre_scatter") + # Scatter t0 = time.perf_counter() results = await scatter( @@ -75,11 +86,15 @@ async def run( output_schema=task.output_schema, ) + # Post-gather checkpoint + if checkpoint and not await checkpoint.post_gather(gathered): + raise CheckpointRejected("post_gather") + # Synthesize syn_backend = synthesis_backend or backend syn_bk = backend_kwargs if syn_backend == backend else None - return await synthesize( + result = await synthesize( gathered=gathered, strategy=synthesis_strategy, backend=syn_backend, @@ -88,6 +103,12 @@ async def run( output_schema=task.output_schema, ) + # Post-synthesis checkpoint + if checkpoint and not await checkpoint.post_synthesis(result): + raise CheckpointRejected("post_synthesis") + + return result + def run_sync( task: Task, @@ -101,6 +122,7 @@ def run_sync( max_tokens: int | None = None, parallel: bool | None = None, early_stop: EarlyStop | None = None, + checkpoint: Checkpoint | None = None, ) -> Synthesis: """Synchronous version of run() — no asyncio knowledge needed. @@ -122,5 +144,6 @@ def run_sync( max_tokens=max_tokens, parallel=parallel, early_stop=early_stop, + checkpoint=checkpoint, ) ) diff --git a/tests/test_checkpoints.py b/tests/test_checkpoints.py new file mode 100644 index 0000000..30f2298 --- /dev/null +++ b/tests/test_checkpoints.py @@ -0,0 +1,135 @@ +"""Tests for human-in-the-loop checkpoints.""" + +from __future__ import annotations + +import pytest + +from broadside_ai.checkpoints import Checkpoint, CheckpointRejected, TerminalCheckpoint +from broadside_ai.gather import GatherResult +from broadside_ai.run import run +from broadside_ai.synthesize import Synthesis +from broadside_ai.task import Task + +# --------------------------------------------------------------------------- +# Mock checkpoint that records calls and returns configurable decisions +# --------------------------------------------------------------------------- + + +class MockCheckpoint: + """A checkpoint that approves or rejects based on constructor args.""" + + def __init__( + self, + pre_scatter: bool = True, + post_gather: bool = True, + post_synthesis: bool = True, + ) -> None: + self.approve_pre_scatter = pre_scatter + self.approve_post_gather = post_gather + self.approve_post_synthesis = post_synthesis + self.calls: list[str] = [] + + async def pre_scatter(self, task: Task, n: int) -> bool: + self.calls.append("pre_scatter") + return self.approve_pre_scatter + + async def post_gather(self, gathered: GatherResult) -> bool: + self.calls.append("post_gather") + return self.approve_post_gather + + async def post_synthesis(self, synthesis: Synthesis) -> bool: + self.calls.append("post_synthesis") + return self.approve_post_synthesis + + +# --------------------------------------------------------------------------- +# Protocol conformance +# --------------------------------------------------------------------------- + + +def test_terminal_checkpoint_implements_protocol() -> None: + assert isinstance(TerminalCheckpoint(), Checkpoint) + + +def test_mock_checkpoint_implements_protocol() -> None: + assert isinstance(MockCheckpoint(), Checkpoint) + + +# --------------------------------------------------------------------------- +# CheckpointRejected exception +# --------------------------------------------------------------------------- + + +def test_checkpoint_rejected_basic() -> None: + exc = CheckpointRejected("pre_scatter") + assert exc.stage == "pre_scatter" + assert exc.reason == "" + assert "pre_scatter" in str(exc) + + +def test_checkpoint_rejected_with_reason() -> None: + exc = CheckpointRejected("post_gather", "too many failures") + assert exc.stage == "post_gather" + assert exc.reason == "too many failures" + assert "too many failures" in str(exc) + + +# --------------------------------------------------------------------------- +# Pipeline integration via run() +# --------------------------------------------------------------------------- + + +async def test_all_checkpoints_approved() -> None: + """When all checkpoints approve, the full pipeline runs.""" + cp = MockCheckpoint() + task = Task(prompt="test prompt") + + result = await run(task, n=2, backend="mock", checkpoint=cp) + + assert result.result # got a synthesis + assert cp.calls == ["pre_scatter", "post_gather", "post_synthesis"] + + +async def test_pre_scatter_rejection() -> None: + """When pre_scatter rejects, scatter never runs.""" + cp = MockCheckpoint(pre_scatter=False) + task = Task(prompt="test prompt") + + with pytest.raises(CheckpointRejected) as exc_info: + await run(task, n=2, backend="mock", checkpoint=cp) + + assert exc_info.value.stage == "pre_scatter" + assert cp.calls == ["pre_scatter"] + + +async def test_post_gather_rejection() -> None: + """When post_gather rejects, synthesis never runs.""" + cp = MockCheckpoint(post_gather=False) + task = Task(prompt="test prompt") + + with pytest.raises(CheckpointRejected) as exc_info: + await run(task, n=2, backend="mock", checkpoint=cp) + + assert exc_info.value.stage == "post_gather" + assert cp.calls == ["pre_scatter", "post_gather"] + + +async def test_post_synthesis_rejection() -> None: + """When post_synthesis rejects, the result is not returned.""" + cp = MockCheckpoint(post_synthesis=False) + task = Task(prompt="test prompt") + + with pytest.raises(CheckpointRejected) as exc_info: + await run(task, n=2, backend="mock", checkpoint=cp) + + assert exc_info.value.stage == "post_synthesis" + assert cp.calls == ["pre_scatter", "post_gather", "post_synthesis"] + + +async def test_no_checkpoint_runs_normally() -> None: + """Without a checkpoint, the pipeline runs end-to-end.""" + task = Task(prompt="test prompt") + + result = await run(task, n=2, backend="mock", checkpoint=None) + + assert result.result # got a synthesis