diff --git a/README.md b/README.md index a7ce72e1..b2866908 100644 --- a/README.md +++ b/README.md @@ -37,11 +37,12 @@ The course follows a four-week learning path: the scheduler does not rebuild dense history on every step. - **Week 4: Build a Coding Agent.** Start with a bounded, validated agent loop, then connect it to a small workspace. The course is publishing one reviewed - checkpoint at a time; Days 1 through 7 now cover inspection, approved edits, + checkpoint at a time; Days 1 through 8 now cover inspection, approved edits, one validation command, simple effect receipts, and one visible checkpoint-and-resume boundary, receipt-backed context compaction, and one visible inspect-and-steer pause, and deterministic evaluation of observable - outcomes. + outcomes, then tokenizer/KV-prefix reuse for two isolated steered branches + and one explicit evidence-backed selection. ## Why MLX and Qwen3? @@ -74,7 +75,7 @@ implementation, test, and publication readiness is tracked below. ## Roadmap -The table tracks implementation (`Code`), tests (`Test`), rendered chapters (`Doc`), and Chi's review of learner-facing material (`Audit`). Week 4 is publishing one reviewed day at a time; Days 1 through 7 are currently available to learners. The Audit column reflects Chi's personal editorial pass on the published course content and is independent of code/test/doc readiness. +The table tracks implementation (`Code`), tests (`Test`), rendered chapters (`Doc`), and Chi's review of learner-facing material (`Audit`). Week 4 is publishing one reviewed day at a time; Days 1 through 8 are currently available to learners. The Audit column reflects Chi's personal editorial pass on the published course content and is independent of code/test/doc readiness. Day 3 can send file contents to the model, modify files after approval, and run one exact configured command. Use a disposable workspace without secrets and @@ -88,6 +89,10 @@ Day 6 inspects one complete-observation checkpoint, appends one visible operator instruction, and resumes a fresh model without replaying the completed effect. Day 7 evaluates one completed run from declared final, file, result, and receipt facts without grading hidden reasoning or exact transcript shape. +Day 8 reuses one real tokenizer/KV checkpoint for two differently steered, +effect-isolated continuations, evaluates both with Day 7's harness, and makes +one explicit passing selection without pretending completed effects were +rewound. | Week + Chapter | Topic | Code | Test | Doc | Audit | |---|---|---|---|---|---| @@ -119,6 +124,7 @@ facts without grading hidden reasoning or exact transcript shape. | 4.5 | Compact Completed Work | ✅ | ✅ | ✅ | 🚧 | | 4.6 | Inspect and Steer a Paused Agent | ✅ | ✅ | ✅ | 🚧 | | 4.7 | Evaluate Observable Outcomes | ✅ | ✅ | ✅ | 🚧 | +| 4.8 | Fork, Steer, and Select | ✅ | ✅ | ✅ | 🚧 | Other topics not covered include quantized or compressed KV caches, cross-request prefix caching, fine-tuning, and long-context techniques. diff --git a/book/src/SUMMARY.md b/book/src/SUMMARY.md index 944cf3fa..7fd07ac7 100644 --- a/book/src/SUMMARY.md +++ b/book/src/SUMMARY.md @@ -38,6 +38,7 @@ - [🚧 Day 5: Compact Completed Work](./week4-05-compaction.md) - [🚧 Day 6: Inspect and Steer a Paused Agent](./week4-06-steering.md) - [🚧 Day 7: Evaluate Observable Outcomes](./week4-07-evaluation.md) + - [🚧 Day 8: Fork, Steer, and Select](./week4-08-fork-steer-select.md) - [🚧 Appendix: Performance Evidence Ledger](./appendix-performance.md) - [Sponsored by Raft.build](./sponsor.md) diff --git a/book/src/week4-07-evaluation.md b/book/src/week4-07-evaluation.md index 11fa407a..ebc8db0d 100644 --- a/book/src/week4-07-evaluation.md +++ b/book/src/week4-07-evaluation.md @@ -162,9 +162,9 @@ observable outcomes. This harness samples the facts a particular case names. It does not prove general task correctness, model quality, security, or production safety, and it is not a hidden grader, benchmark suite, or LLM-as-judge system. -That closes the seven-day Week 4 path: build a bounded loop, inspect and change -a workspace with approvals and receipts, pause and resume, compact completed -evidence, steer at a safe boundary, and finally evaluate one run by what it -actually left behind. +You now have the evidence needed to compare continuations. Continue with [Day +8: Fork, Steer, and Select](week4-08-fork-steer-select.md) to reuse one real +token/KV prefix, steer two isolated branches, and explicitly choose a passing +outcome without rewinding completed effects. {{#include copyright.md}} diff --git a/book/src/week4-08-fork-steer-select.md b/book/src/week4-08-fork-steer-select.md new file mode 100644 index 00000000..803c5fcd --- /dev/null +++ b/book/src/week4-08-fork-steer-select.md @@ -0,0 +1,202 @@ +# Day 8: Fork, Steer, and Select + +> 🚧 **Early-review WIP:** Use only pre-created disposable workspaces. A +> control-state fork does not undo a file edit or any other completed effect. + +Days 4 and 6 paused one agent and resumed one continuation. Day 8 asks a new +question: after the model has inspected or changed the workspace, can we reuse +the same inference prefix, try two explicit directions, and select the branch +whose observable result is better? + +The answer reconnects Week 4 to the inference system from Weeks 1–3. The course +tokenizer renders the checkpoint conversation once. `TinyKvFullCache` stores +the prefix keys and values for every model layer. Each branch gets a fresh +control object and cache handles that share only those immutable prefix arrays, +then decodes its own suffix. The branch report exposes the reused token count, +the layer offsets, and the full-prefix prefill that was avoided. + +This is control-state reuse, not effect rollback. Copy the already-modified +disposable workspace and its completed receipt log before running either +branch. Both copies begin with the same files and evidence; later receipts stay +inside their branch. + +## The Starter Surface + +Day 8 adds one module and extends the approval result: + +| File | Public names | Purpose | +| --- | --- | --- | +| `src/tiny_llm/agent/workspace.py` | `ApprovalDecision` | Carry an operator's denial reason back as one ordinary model-visible observation. | +| `src/tiny_llm/agent/branching.py` | `PrefixReuse`, `KvPrefixGenerator`, `BranchOutcome`, `run_branch`, `select_branch` | Reuse a dense KV prefix, run isolated steered continuations, evaluate them, and make one explicit choice. | +| `src/tiny_llm/agent/__init__.py` | the names above | Export the cumulative Day 8 API. | + +Copy and run the five learner tasks: + +```bash +pdm run copy-test --week 4 --day 8 +pdm run test --week 4 --day 8 +``` + +Use this command for the supplied implementation: + +```bash +pdm run test-refsol --week 4 --day 8 +``` + +Before you implement the TODOs, all five Day 8 tasks are expected to fail. + +## Task 1: Return a Reason with a Denial + +Add the immutable decision: + +```python +ApprovalDecision(approved=False, reason="keep the requested answer at 2") +``` + +A structured denial requires a nonblank reason. `Workspace.execute` returns +that reason in its normal `error:` result so the next model turn can react to +the operator's instruction. It does not execute the effect or append a receipt. +Existing callbacks that return plain `True` or `False` remain compatible. + +The reason is steering, not a secret channel. Keep it short and suitable for +the model-visible transcript. + +## Task 2: Save One Real Token and KV Prefix + +`KvPrefixGenerator.save_checkpoint(messages)` renders the checkpoint messages +without a generation prompt, tokenizes them with the course tokenizer, and +prefills one `TinyKvFullCache` per layer. It records the exact token IDs and +layer offsets in the existing Day 4 `ModelCheckpoint`. + +The saved prompt must be an exact token prefix of every later steered prompt. +Reject a continuation if even a same-length token differs. This binds cache +reuse to content, not merely to a position. + +`fork()` creates a fresh generator whose cache handles point at the frozen +prefix arrays. When one branch grows, `TinyKvFullCache` assigns newly +concatenated arrays to that branch. The frozen prefix and its sibling remain +unchanged. This lesson deliberately uses the dense compatibility path; paged +copy-on-write and radix serving are separate scaling topics. + +## Task 3: Expose What Was Reused + +Each continuation reports: + +```python +PrefixReuse( + reused_tokens=prefix_length, + layer_offsets=(prefix_length, ...), + avoided_prefill_tokens=prefix_length, +) +``` + +The first suffix model call starts at `prefix_length`; it must not call the +model again at offset zero. These numbers make the inference boundary visible: +the branch is not cloning only a Python transcript and silently recomputing the +whole prompt. + +## Task 4: Fork Effects and Evidence Explicitly + +Suppose the completed prefix changed `app.py` from `answer = 1` to +`answer = 2` and wrote `call-1`, the edit receipt. Copy both the post-effect +workspace and `receipts.jsonl` into two roots: + +```text +base after checkpoint +├── app.py answer = 2 +└── receipts.jsonl call-1: edit_file + +validate-only/ try-extra-edit/ +├── app.py ├── app.py +└── receipts.jsonl └── receipts.jsonl +``` + +Both receipt files begin byte-identical and contain `call-1`. The +`validate-only` branch appends `call-2` after its exact allowed validation +command. The other branch asks to change the answer again; the operator denies +it with a reason, so its file and receipt bytes remain unchanged. + +Construct each branch with its own `Workspace` and `ReceiptStore`, then call: + +```python +outcome = run_branch( + "validate-only", + "validate without another edit", + checkpoint, + prefix_generator.fork(), + workspace, + receipts, + evaluation_case, +) +``` + +`run_branch` composes the Day 6 steered resume with the Day 7 observable-outcome +evaluator. It does not copy a directory, infer an evaluation case, or merge +effects for you. + +## Task 5: Select One Passing Branch + +Make the choice explicit: + +```python +selected = select_branch(outcomes, "validate-only") +``` + +The name must identify exactly one outcome, and that outcome must pass its Day +7 report. Reject an absent selected name, a selected name that matches multiple +outcomes, or a failing branch. Day 8 does not invent a hidden score or ask +another model to judge the traces. + +## Manual Qwen/MLX Walkthrough + +Complete Weeks 1–3 and the Day 8 TODOs first. Use a cached local Qwen model and +the same dense compatibility path: + +```python +from mlx_lm import load +from tiny_llm import Qwen3ModelWeek3 +from tiny_llm.agent import KvPrefixGenerator, create_checkpoint + +mlx_model, tokenizer = load("Qwen/Qwen3-0.6B-MLX-4bit") +model = Qwen3ModelWeek3(mlx_model, enable_paged_attention=False) +prefix_generator = KvPrefixGenerator(model, tokenizer, max_tokens=128) +``` + +First create a Day 4 checkpoint named `paused` after a complete tool +observation. Its messages are the control boundary you want to share, while its +model field belongs to the generator that created it. Rebind those exact +messages to the real tokenizer and dense KV cache before resuming: + +```python +messages = [ + {"role": role, "content": content} + for role, content in paused.messages +] +model_checkpoint = prefix_generator.save_checkpoint(messages) +checkpoint = create_checkpoint(paused.task, messages, model_checkpoint) +``` + +Now fork two fresh generators with `prefix_generator.fork()`. Give them +different visible steering messages and the two workspace/receipt copies +described above, passing the rebound `checkpoint` to each `run_branch` call. +Print each `outcome.reuse`, render both evaluation reports, and select the +passing name. + +Model responses are nondeterministic, so this walkthrough is manual. Inspect +the actual proposed actions, approval reason, final file bytes, receipt logs, +and evaluation reports. The deterministic learner test covers the same public +boundary with a tiny tokenizer/model and no download. + +## Checkpoint + +You can now connect a Day 4 control checkpoint to the actual tokenizer and KV +cache path, reuse one immutable prefix for two isolated continuations, expose a +denial reason to the model without recording an effect, evaluate both branches +from declared evidence, and choose one passing result. + +Completed effects were copied, not rewound. The two branches do not run +concurrently, share a mutable workspace, merge receipts, or provide a session +server/tree. Day 8 teaches the boundary visibly before adding any serving-scale +machinery. + +{{#include copyright.md}} diff --git a/book/src/week4-overview.md b/book/src/week4-overview.md index a7d3d04f..4c744ad2 100644 --- a/book/src/week4-overview.md +++ b/book/src/week4-overview.md @@ -1,7 +1,7 @@ # 🚧 Week 4: Build a Coding Agent > **Course status:** Week 4 is being published one checkpoint at a time. Days 1 -> through 7 are ready to learn and review. Additional capabilities will appear +> through 8 are ready to learn and review. Additional capabilities will appear > only after their implementation, starter, and reviews are ready. Weeks 1 through 3 turn tokens into text and make serving that text efficient. @@ -18,6 +18,9 @@ adds one visible operator steering message, and resumes a fresh model without replaying the completed effect. Day 7 evaluates one completed run from declared final, file, result, and receipt facts without grading hidden reasoning or exact transcript shape. +Day 8 reconnects the agent to the inference system from Weeks 1–3: it saves one +real tokenizer/KV prefix, forks two isolated steered continuations without +rewinding completed effects, evaluates both, and makes one explicit selection. ## What Day 1 Builds @@ -115,8 +118,15 @@ Outcomes](week4-07-evaluation.md). Its cumulative command is: pdm run test --week 4 --day 7 ``` -Only the Day 1 through Day 7 starter modules are published. Do not add session -trees, rewind, reconciliation, an LLM judge, or other later public APIs to your -solution. +After Day 7 passes, continue with [Day 8: Fork, Steer, and +Select](week4-08-fork-steer-select.md). Its cumulative command is: + +```bash +pdm run test --week 4 --day 8 +``` + +Only the Day 1 through Day 8 starter modules are published. Do not add session +trees, effect rewind, reconciliation, an LLM judge, radix serving, or other +later public APIs to your solution. {{#include copyright.md}} diff --git a/docs/week4-day-split.md b/docs/week4-day-split.md index 69db84c9..37add743 100644 --- a/docs/week4-day-split.md +++ b/docs/week4-day-split.md @@ -1,10 +1,9 @@ # Week 4 Day Split (reference for reviewers) -Status: Days 1--4 are published checkpoints. Days 2--4 follow Chi's approved -simplified agent-loop cycle, with each day shipping as one cumulative learner -PR so reviewers can see exactly what belongs to that checkpoint. +Status: Days 1--8 are published checkpoints. Each day ships as one cumulative +learner PR so reviewers can see exactly what belongs to that checkpoint. -## 7-day structure +## 8-day structure | Day | Theme | Features (PRs) | Modules | |---|---|---|---| @@ -12,9 +11,10 @@ PR so reviewers can see exactly what belongs to that checkpoint. | 2 | Inspect a workspace | read-only list/read tools | `workspace.py` | | 3 | Edit, validate, and record | approved edits, one command, simple receipts | `workspace.py`, `receipts.py` | | 4 | Checkpoint and resume | one conversation + fake-model cache snapshot | `checkpoint.py`, `loop.py` | -| 5 | Reserved learner checkpoint | unpublished | — | -| 6 | Reserved learner checkpoint | unpublished | — | -| 7 | Reserved learner checkpoint | unpublished | — | +| 5 | Compact completed work | bounded receipt-backed transcript view | `compaction.py` | +| 6 | Inspect and steer | safe-boundary status and one visible steering message | `steering.py` | +| 7 | Evaluate outcomes | declared final/file/result/receipt facts | `evaluation.py` | +| 8 | Fork, steer, and select | dense tokenizer/KV prefix reuse, isolated branches, explicit selection | `branching.py`, `workspace.py` | Extension (not a day): COW/radix cache — `docs/week4-cow-radix-extension-plan.md`. @@ -26,8 +26,10 @@ Extension (not a day): COW/radix cache — `docs/week4-cow-radix-extension-plan. - Days are implemented sequentially. A later day does not leak API or prose into the current starter. -## Why 7 days +## Why 8 days -The existing `week4-01..07` chapter numbering stays stable. Each remaining day -adds one visible agent-loop concept; scaling and production-hardening machinery -stay outside the core course unless a later checkpoint explicitly teaches it. +The first seven days establish the agent loop and its observable evidence. Day +8 reconnects that control path to the tokenizer and KV cache built in Weeks +1--3. Each day adds one visible concept; scaling and production-hardening +machinery stay outside the core course unless a later checkpoint explicitly +teaches it. diff --git a/src/tiny_llm/agent/__init__.py b/src/tiny_llm/agent/__init__.py index 24768f52..9d9a6f28 100644 --- a/src/tiny_llm/agent/__init__.py +++ b/src/tiny_llm/agent/__init__.py @@ -1,6 +1,13 @@ # WARNING: Under review - generated by LLM. from .checkpoint import AgentCheckpoint, ModelCheckpoint, create_checkpoint +from .branching import ( + BranchOutcome, + KvPrefixGenerator, + PrefixReuse, + run_branch, + select_branch, +) from .compaction import CompactionResult, compact_completed_interactions from .evaluation import ( EvaluationCase, @@ -29,7 +36,7 @@ ) from .receipts import EffectReceipt, ReceiptStore from .steering import AgentStatus, inspect_checkpoint, resume_with_steering -from .workspace import ToolPolicy, Workspace +from .workspace import ApprovalDecision, ToolPolicy, Workspace __all__ = [ @@ -39,6 +46,8 @@ "AgentLimits", "AgentRun", "AgentStatus", + "ApprovalDecision", + "BranchOutcome", "CompactionResult", "EvaluationCase", "EvaluationCheck", @@ -46,10 +55,12 @@ "EffectReceipt", "FinalAction", "FileExpectation", + "KvPrefixGenerator", "ModelCheckpoint", "ReceiptStore", "ReceiptExpectation", "ResultExpectation", + "PrefixReuse", "ToolAction", "ToolPolicy", "Workspace", @@ -64,5 +75,7 @@ "resume_agent", "resume_with_steering", "run_agent", + "run_branch", "run_to_checkpoint", + "select_branch", ] diff --git a/src/tiny_llm/agent/branching.py b/src/tiny_llm/agent/branching.py new file mode 100644 index 00000000..b5944e26 --- /dev/null +++ b/src/tiny_llm/agent/branching.py @@ -0,0 +1,98 @@ +# WARNING: Under review - generated by LLM. + +"""Week 4, Day 8 learner surface for cached branch continuations.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from .checkpoint import AgentCheckpoint, ModelCheckpoint +from .evaluation import EvaluationCase, EvaluationReport +from .generation import Generate, Message +from .loop import AgentLimits, AgentRun +from .receipts import ReceiptStore +from .workspace import Workspace + + +@dataclass(frozen=True) +class PrefixReuse: + """Observable token and cache positions reused by one continuation.""" + + reused_tokens: int + layer_offsets: tuple[int, ...] + avoided_prefill_tokens: int + + +@dataclass(frozen=True) +class BranchOutcome: + """One steered branch, its observable evaluation, and prefix reuse.""" + + name: str + steering: str + run: AgentRun + report: EvaluationReport + reuse: PrefixReuse + + +class KvPrefixGenerator: + """Greedy course-model generation from one frozen dense KV prefix.""" + + def __init__( + self, + model: Any, + tokenizer: Any, + max_tokens: int, + enable_thinking: bool = False, + ) -> None: + pass + + @property + def reuse(self) -> PrefixReuse: + """Return the prefix positions used by the latest continuation.""" + + pass + + def save_checkpoint(self, messages: list[Message]) -> ModelCheckpoint: + """Render and prefill one checkpoint prefix exactly once.""" + + pass + + def restore_checkpoint(self, checkpoint: ModelCheckpoint) -> None: + """Bind a fresh continuation to this generator's frozen prefix.""" + + pass + + def fork(self) -> KvPrefixGenerator: + """Create a fresh generator that shares only immutable prefix arrays.""" + + pass + + def __call__(self, messages: list[Message]) -> str: + """Generate one continuation from the restored prefix.""" + + pass + + +def run_branch( + name: str, + steering: str, + checkpoint: AgentCheckpoint, + generate: Generate, + workspace: Workspace, + receipts: ReceiptStore, + case: EvaluationCase, + limits: AgentLimits | None = None, +) -> BranchOutcome: + """Resume, evaluate, and retain the reuse facts for one isolated branch.""" + + pass + + +def select_branch( + outcomes: Sequence[BranchOutcome], selected_name: str +) -> BranchOutcome: + """Return one explicitly named passing branch.""" + + pass diff --git a/src/tiny_llm/agent/workspace.py b/src/tiny_llm/agent/workspace.py index 51a65ceb..9d711c63 100644 --- a/src/tiny_llm/agent/workspace.py +++ b/src/tiny_llm/agent/workspace.py @@ -12,6 +12,18 @@ from .receipts import ReceiptStore +@dataclass(frozen=True) +class ApprovalDecision: + """One operator approval or a model-visible denial reason.""" + + approved: bool + reason: str = "" + + def __post_init__(self) -> None: + # TODO: validate the decision and require a reason for denials. + pass + + @dataclass(frozen=True) class ToolPolicy: """Tools and limits authorized for one explicit workspace root.""" @@ -34,7 +46,7 @@ class Workspace: """Run the Day 3 tools under one policy and approval callback.""" policy: ToolPolicy - confirm_tool: Callable[[ToolAction], bool] | None = None + confirm_tool: Callable[[ToolAction], bool | ApprovalDecision] | None = None receipt_store: ReceiptStore = field(default_factory=ReceiptStore) _observed: dict[str, str] = field(default_factory=dict, init=False) _modified: set[str] = field(default_factory=set, init=False) diff --git a/src/tiny_llm_ref/agent/__init__.py b/src/tiny_llm_ref/agent/__init__.py index 24768f52..9d9a6f28 100644 --- a/src/tiny_llm_ref/agent/__init__.py +++ b/src/tiny_llm_ref/agent/__init__.py @@ -1,6 +1,13 @@ # WARNING: Under review - generated by LLM. from .checkpoint import AgentCheckpoint, ModelCheckpoint, create_checkpoint +from .branching import ( + BranchOutcome, + KvPrefixGenerator, + PrefixReuse, + run_branch, + select_branch, +) from .compaction import CompactionResult, compact_completed_interactions from .evaluation import ( EvaluationCase, @@ -29,7 +36,7 @@ ) from .receipts import EffectReceipt, ReceiptStore from .steering import AgentStatus, inspect_checkpoint, resume_with_steering -from .workspace import ToolPolicy, Workspace +from .workspace import ApprovalDecision, ToolPolicy, Workspace __all__ = [ @@ -39,6 +46,8 @@ "AgentLimits", "AgentRun", "AgentStatus", + "ApprovalDecision", + "BranchOutcome", "CompactionResult", "EvaluationCase", "EvaluationCheck", @@ -46,10 +55,12 @@ "EffectReceipt", "FinalAction", "FileExpectation", + "KvPrefixGenerator", "ModelCheckpoint", "ReceiptStore", "ReceiptExpectation", "ResultExpectation", + "PrefixReuse", "ToolAction", "ToolPolicy", "Workspace", @@ -64,5 +75,7 @@ "resume_agent", "resume_with_steering", "run_agent", + "run_branch", "run_to_checkpoint", + "select_branch", ] diff --git a/src/tiny_llm_ref/agent/branching.py b/src/tiny_llm_ref/agent/branching.py new file mode 100644 index 00000000..bff4de56 --- /dev/null +++ b/src/tiny_llm_ref/agent/branching.py @@ -0,0 +1,245 @@ +# WARNING: Under review - generated by LLM. + +"""Week 4, Day 8: fork one cached control prefix and select an outcome.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from ..kv_cache import TinyKvFullCache +from .checkpoint import AgentCheckpoint, ModelCheckpoint +from .evaluation import EvaluationCase, EvaluationReport, evaluate_run +from .generation import Generate, Message +from .loop import AgentLimits, AgentRun +from .protocol import AgentError +from .receipts import ReceiptStore +from .steering import resume_with_steering +from .workspace import Workspace + + +@dataclass(frozen=True) +class PrefixReuse: + """Observable token and cache positions reused by one continuation.""" + + reused_tokens: int + layer_offsets: tuple[int, ...] + avoided_prefill_tokens: int + + +@dataclass(frozen=True) +class BranchOutcome: + """One steered branch, its observable evaluation, and prefix reuse.""" + + name: str + steering: str + run: AgentRun + report: EvaluationReport + reuse: PrefixReuse + + +class KvPrefixGenerator: + """Greedy course-model generation from one frozen dense KV prefix.""" + + def __init__( + self, + model: Any, + tokenizer: Any, + max_tokens: int, + enable_thinking: bool = False, + ) -> None: + if type(max_tokens) is not int or max_tokens <= 0: + raise ValueError("max_tokens must be a positive integer") + if type(enable_thinking) is not bool: + raise ValueError("enable_thinking must be a boolean") + layers = getattr(model, "num_hidden_layers", None) + if type(layers) is not int or layers <= 0: + raise ValueError("model must expose a positive num_hidden_layers") + self._model = model + self._tokenizer = tokenizer + self._max_tokens = max_tokens + self._enable_thinking = enable_thinking + self._layer_count = layers + self._response_index = 0 + self._checkpoint: ModelCheckpoint | None = None + self._prefix_tokens: tuple[int, ...] = () + self._prefix_key_values: tuple[tuple[Any, Any], ...] = () + self._restored = False + self._reuse = PrefixReuse(0, (), 0) + + @property + def reuse(self) -> PrefixReuse: + """Return the prefix positions used by the latest continuation.""" + + return self._reuse + + def save_checkpoint(self, messages: list[Message]) -> ModelCheckpoint: + """Render and prefill one checkpoint prefix exactly once.""" + + import mlx.core as mx + + if self._checkpoint is not None: + raise AgentError("prefix checkpoint was already saved") + prompt = self._render(messages, add_generation_prompt=False) + token_ids = self._encode(prompt) + if not token_ids: + raise AgentError("checkpoint prompt must contain at least one token") + caches = [TinyKvFullCache() for _ in range(self._layer_count)] + tokens = mx.array(token_ids) + self._model(tokens[None], 0, caches, logits_to_keep=1) + for cache in caches: + cache.materialize() + if any(cache.key_values is None for cache in caches): + raise AgentError("model did not populate every dense cache layer") + self._prefix_tokens = token_ids + self._prefix_key_values = tuple(cache.key_values for cache in caches) + offsets = tuple(cache.offset for cache in caches) + checkpoint = ModelCheckpoint( + len(messages), self._response_index, token_ids, offsets + ) + self._checkpoint = checkpoint + self._reuse = PrefixReuse(len(token_ids), offsets, len(token_ids)) + return checkpoint + + def restore_checkpoint(self, checkpoint: ModelCheckpoint) -> None: + """Bind a fresh continuation to this generator's frozen prefix.""" + + if not isinstance(checkpoint, ModelCheckpoint): + raise AgentError("model checkpoint is invalid") + if self._checkpoint is None or checkpoint != self._checkpoint: + raise AgentError("model checkpoint does not match the saved KV prefix") + self._response_index = checkpoint.response_index + self._restored = True + + def fork(self) -> KvPrefixGenerator: + """Create a fresh generator that shares only immutable prefix arrays.""" + + if self._checkpoint is None: + raise AgentError("save a prefix checkpoint before forking") + branch = KvPrefixGenerator( + self._model, + self._tokenizer, + self._max_tokens, + self._enable_thinking, + ) + branch._checkpoint = self._checkpoint + branch._prefix_tokens = self._prefix_tokens + branch._prefix_key_values = self._prefix_key_values + branch._response_index = self._checkpoint.response_index + branch._reuse = self._reuse + return branch + + def __call__(self, messages: list[Message]) -> str: + import mlx.core as mx + + if not self._restored or self._checkpoint is None: + raise AgentError("restore the checkpoint before generating") + prompt = self._render(messages, add_generation_prompt=True) + token_ids = self._encode(prompt) + prefix_size = len(self._prefix_tokens) + if token_ids[:prefix_size] != self._prefix_tokens: + raise AgentError("steered prompt does not extend the saved token prefix") + suffix = token_ids[prefix_size:] + if not suffix: + raise AgentError("steered prompt must add tokens after the saved prefix") + + caches = self._fork_caches() + tokens = mx.array(suffix) + offset = prefix_size + logits = self._model(tokens[None], offset, caches, logits_to_keep=1)[:, -1, :] + offset += len(suffix) + output: list[int] = [] + for _ in range(self._max_tokens): + token = int(mx.argmax(logits, axis=-1).item()) + if token == self._tokenizer.eos_token_id: + break + output.append(token) + tokens = mx.array([token]) + logits = self._model(tokens[None], offset, caches, logits_to_keep=1)[ + :, -1, : + ] + offset += 1 + self._response_index += 1 + self._reuse = PrefixReuse( + prefix_size, + self._checkpoint.layer_offsets, + prefix_size, + ) + return self._tokenizer.decode(output) + + def _render(self, messages: list[Message], *, add_generation_prompt: bool) -> str: + try: + prompt = self._tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=add_generation_prompt, + enable_thinking=self._enable_thinking, + ) + except (KeyError, TypeError, ValueError) as error: + raise AgentError("could not render checkpoint messages") from error + if not isinstance(prompt, str): + raise AgentError("chat template must render text") + return prompt + + def _encode(self, prompt: str) -> tuple[int, ...]: + try: + encoded = self._tokenizer.encode(prompt, add_special_tokens=False) + tokens = tuple(int(token) for token in encoded) + except (TypeError, ValueError) as error: + raise AgentError("tokenizer returned invalid token ids") from error + if any(token < 0 for token in tokens): + raise AgentError("tokenizer returned invalid token ids") + return tokens + + def _fork_caches(self) -> list[TinyKvFullCache]: + caches: list[TinyKvFullCache] = [] + assert self._checkpoint is not None + for key_values, offset in zip( + self._prefix_key_values, + self._checkpoint.layer_offsets, + strict=True, + ): + cache = TinyKvFullCache() + cache.key_values = key_values + cache.offset = offset + caches.append(cache) + return caches + + +def run_branch( + name: str, + steering: str, + checkpoint: AgentCheckpoint, + generate: Generate, + workspace: Workspace, + receipts: ReceiptStore, + case: EvaluationCase, + limits: AgentLimits | None = None, +) -> BranchOutcome: + """Resume, evaluate, and retain the reuse facts for one isolated branch.""" + + if not isinstance(name, str) or not name.strip(): + raise ValueError("branch name must not be blank") + run = resume_with_steering(checkpoint, steering, generate, workspace, limits) + report = evaluate_run(run, workspace, receipts, case) + reuse = getattr(generate, "reuse", None) + if not isinstance(reuse, PrefixReuse): + raise AgentError("branch generator did not report prefix reuse") + return BranchOutcome(name, steering, run, report, reuse) + + +def select_branch( + outcomes: Sequence[BranchOutcome], selected_name: str +) -> BranchOutcome: + """Return one explicitly named passing branch.""" + + if not isinstance(selected_name, str) or not selected_name.strip(): + raise ValueError("selected branch name must not be blank") + matches = [outcome for outcome in outcomes if outcome.name == selected_name] + if len(matches) != 1: + raise AgentError("selected branch name must match exactly one outcome") + selected = matches[0] + if not selected.report.passed: + raise AgentError("selected branch must pass evaluation") + return selected diff --git a/src/tiny_llm_ref/agent/workspace.py b/src/tiny_llm_ref/agent/workspace.py index d27345db..765245e4 100644 --- a/src/tiny_llm_ref/agent/workspace.py +++ b/src/tiny_llm_ref/agent/workspace.py @@ -35,6 +35,22 @@ def _is_protected(path: Path) -> bool: ) +@dataclass(frozen=True) +class ApprovalDecision: + """One operator approval or a model-visible denial reason.""" + + approved: bool + reason: str = "" + + def __post_init__(self) -> None: + if type(self.approved) is not bool: + raise ValueError("approved must be a boolean") + if not isinstance(self.reason, str): + raise ValueError("approval reason must be a string") + if not self.approved and not self.reason.strip(): + raise ValueError("a denial requires a nonblank operator reason") + + @dataclass(frozen=True) class ToolPolicy: """Tools and limits authorized for one explicit workspace root.""" @@ -96,7 +112,7 @@ class Workspace: """Run the Day 3 tools under one policy and approval callback.""" policy: ToolPolicy - confirm_tool: Callable[[ToolAction], bool] | None = None + confirm_tool: Callable[[ToolAction], bool | ApprovalDecision] | None = None receipt_store: ReceiptStore = field(default_factory=ReceiptStore) _observed: dict[str, str] = field(default_factory=dict, init=False) _modified: set[str] = field(default_factory=set, init=False) @@ -260,7 +276,13 @@ def execute(self, action: ToolAction, tool_call_id: str | None = None) -> str: ) else: self._allowed_command(action.arguments["argv"]) - if self.confirm_tool is None or self.confirm_tool(action) is not True: + decision = self.confirm_tool(action) if self.confirm_tool else False + if isinstance(decision, ApprovalDecision): + if not decision.approved: + raise AgentError( + f"operator denied the tool action: {decision.reason}" + ) + elif decision is not True: raise AgentError("operator denied the tool action") if action.tool == "write_file": result = self.write_file( @@ -365,6 +387,8 @@ def _allowed_command(self, argv: list[str]) -> tuple[str, ...]: return command def _new_call_id(self) -> str: - call_id = f"call-{self._next_call_number}" - self._next_call_number += 1 - return call_id + while True: + call_id = f"call-{self._next_call_number}" + self._next_call_number += 1 + if self.receipt_store.get(call_id) is None: + return call_id diff --git a/tests_refsol/test_week_4_day_8.py b/tests_refsol/test_week_4_day_8.py new file mode 100644 index 00000000..bea513b4 --- /dev/null +++ b/tests_refsol/test_week_4_day_8.py @@ -0,0 +1,399 @@ +# WARNING: Under review - generated by LLM. + +"""Week 4 Day 8 cached fork, steering, and selection course-code tests.""" + +import json +import shutil +import sys +from dataclasses import replace + +import mlx.core as mx +import pytest + +from .tiny_llm_base import ( + AgentError, + ApprovalDecision, + EvaluationCase, + FileExpectation, + KvPrefixGenerator, + ModelCheckpoint, + PrefixReuse, + ReceiptExpectation, + ReceiptStore, + ResultExpectation, + ToolAction, + ToolPolicy, + Workspace, + run_branch, + run_to_checkpoint, + select_branch, +) + + +class CharacterTokenizer: + eos_token_id = 0 + + def apply_chat_template( + self, + messages, + *, + tokenize, + add_generation_prompt, + enable_thinking, + ): + assert tokenize is False + assert enable_thinking is False + rendered = "".join( + f"<{message['role']}>{message['content']}\n" for message in messages + ) + if add_generation_prompt: + rendered += "" + return rendered + + @staticmethod + def encode(text, *, add_special_tokens): + assert add_special_tokens is False + return [ord(character) for character in text] + + @staticmethod + def decode(tokens): + return "".join(chr(token) for token in tokens) + + +class DenseEosModel: + """Populate real dense caches, then emit EOS immediately.""" + + num_hidden_layers = 2 + + def __init__(self): + self.calls = [] + + def __call__(self, tokens, offset, caches, logits_to_keep=None): + assert logits_to_keep == 1 + size = int(tokens.shape[1]) + before = tuple( + None if cache.key_values is None else id(cache.key_values[0]) + for cache in caches + ) + keys = mx.reshape(tokens.astype(mx.float32), (1, 1, size, 1)) + for layer, cache in enumerate(caches): + cache.update_and_fetch(keys + layer, keys + layer) + after = tuple(id(cache.key_values[0]) for cache in caches) + self.calls.append((offset, size, tuple(caches), before, after)) + return mx.concatenate( + [mx.ones((1, size, 1)), mx.zeros((1, size, 127))], axis=-1 + ) + + +class ScriptedCheckpointModel: + """Resume deterministic responses while exposing the recorded reuse facts.""" + + def __init__(self, responses): + self.responses = tuple(responses) + self.response_index = 0 + self.calls = [] + self.reuse = PrefixReuse(0, (), 0) + + @staticmethod + def _tokens(messages): + return tuple(len(message["content"]) for message in messages) + + def __call__(self, messages): + self.calls.append([dict(message) for message in messages]) + response = self.responses[self.response_index] + self.response_index += 1 + return response + + def save_checkpoint(self, messages): + tokens = self._tokens(messages) + return ModelCheckpoint( + len(messages), self.response_index, tokens, (len(tokens), len(tokens)) + ) + + def restore_checkpoint(self, checkpoint): + self.response_index = checkpoint.response_index + self.reuse = PrefixReuse( + len(checkpoint.cached_token_ids), + checkpoint.layer_offsets, + len(checkpoint.cached_token_ids), + ) + + +def test_task_1_structured_denial_requires_and_exposes_one_operator_reason(tmp_path): + (tmp_path / "app.py").write_text("answer = 2\n", encoding="utf-8") + receipts = ReceiptStore(tmp_path / "receipts.jsonl") + workspace = Workspace( + ToolPolicy(tmp_path, allow_writes=True), + lambda _action: ApprovalDecision(False, "keep the requested answer at 2"), + receipts, + ) + workspace.read_file("app.py") + action = ToolAction("edit_file", {"path": "app.py", "old": "2", "new": "3"}) + + result = workspace.execute(action) + + assert result == ( + "error: operator denied the tool action: keep the requested answer at 2" + ) + assert result.count("keep the requested answer at 2") == 1 + assert (tmp_path / "app.py").read_text(encoding="utf-8") == "answer = 2\n" + assert receipts.get("call-1") is None + with pytest.raises(ValueError, match="nonblank operator reason"): + ApprovalDecision(False, " ") + + +def test_task_1_legacy_boolean_approval_remains_compatible(tmp_path): + (tmp_path / "app.py").write_text("answer = 1\n", encoding="utf-8") + allowed = Workspace(ToolPolicy(tmp_path, allow_writes=True), lambda _action: True) + allowed.read_file("app.py") + assert ( + allowed.execute( + ToolAction("edit_file", {"path": "app.py", "old": "1", "new": "2"}) + ) + == "edited app.py" + ) + denied = Workspace(ToolPolicy(tmp_path, allow_writes=True), lambda _action: False) + denied.read_file("app.py") + assert ( + denied.execute( + ToolAction("edit_file", {"path": "app.py", "old": "2", "new": "3"}) + ) + == "error: operator denied the tool action" + ) + + +def test_task_2_cached_prefix_is_prefilled_once_and_reused_by_both_forks(): + model = DenseEosModel() + generator = KvPrefixGenerator(model, CharacterTokenizer(), max_tokens=4) + messages = [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "task"}, + {"role": "assistant", "content": '{"tool":"read_file","path":"app.py"}'}, + {"role": "user", "content": "Tool result:\nanswer = 2\n"}, + ] + checkpoint = generator.save_checkpoint(messages) + first = generator.fork() + second = generator.fork() + first.restore_checkpoint(checkpoint) + second.restore_checkpoint(checkpoint) + + assert ( + first([*messages, {"role": "user", "content": "Operator steering:\nvalidate"}]) + == "" + ) + assert ( + second([*messages, {"role": "user", "content": "Operator steering:\ninspect"}]) + == "" + ) + + prefix_call, first_call, second_call = model.calls + assert prefix_call[0] == 0 + assert prefix_call[1] == len(checkpoint.cached_token_ids) + assert first_call[0] == second_call[0] == len(checkpoint.cached_token_ids) + assert all( + first_cache is not second_cache + for first_cache, second_cache in zip(first_call[2], second_call[2], strict=True) + ) + assert first_call[3] == second_call[3] == prefix_call[4] + assert first_call[4] != first_call[3] + assert second_call[4] != second_call[3] + assert ( + first.reuse + == second.reuse + == PrefixReuse( + len(checkpoint.cached_token_ids), + checkpoint.layer_offsets, + len(checkpoint.cached_token_ids), + ) + ) + assert checkpoint.layer_offsets == ( + len(checkpoint.cached_token_ids), + len(checkpoint.cached_token_ids), + ) + + +def test_task_3_steered_prompt_must_extend_the_exact_saved_token_prefix(): + generator = KvPrefixGenerator(DenseEosModel(), CharacterTokenizer(), 2) + messages = [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "task"}, + ] + checkpoint = generator.save_checkpoint(messages) + branch = generator.fork() + foreign_generator = KvPrefixGenerator(DenseEosModel(), CharacterTokenizer(), 2) + foreign_checkpoint = foreign_generator.save_checkpoint( + [ + {"role": "system", "content": "foreign system"}, + {"role": "user", "content": "foreign task"}, + ] + ) + + with pytest.raises(AgentError, match="does not match the saved KV prefix"): + branch.restore_checkpoint(foreign_checkpoint) + branch.restore_checkpoint(checkpoint) + changed = [ + {"role": "system", "content": "changed"}, + {"role": "user", "content": "task"}, + {"role": "user", "content": "Operator steering:\ncontinue"}, + ] + + with pytest.raises(AgentError, match="does not extend the saved token prefix"): + branch(changed) + + +def _forkable_effect(tmp_path): + base = tmp_path / "base" + base.mkdir() + (base / "app.py").write_text("answer = 1\n", encoding="utf-8") + receipts = ReceiptStore(base / "receipts.jsonl") + workspace = Workspace( + ToolPolicy(base, allow_writes=True), lambda _action: True, receipts + ) + responses = ( + '{"tool":"read_file","path":"app.py"}', + '{"tool":"edit_file","path":"app.py","old":"1","new":"2"}', + ) + checkpoint = run_to_checkpoint( + "set answer = 2 and validate", + ScriptedCheckpointModel(responses), + workspace, + after_tool_calls=2, + ) + return base, checkpoint + + +def _copy_branch(base, destination, command, approval): + shutil.copytree(base, destination) + receipts = ReceiptStore(destination / "receipts.jsonl") + workspace = Workspace( + ToolPolicy( + destination, + allow_writes=True, + allowed_commands=(command,), + ), + approval, + receipts, + ) + return workspace, receipts + + +def test_task_4_forks_effects_and_receipts_then_isolates_later_branch_evidence( + tmp_path, +): + base, checkpoint = _forkable_effect(tmp_path) + initial_receipts = (base / "receipts.jsonl").read_bytes() + command = ( + sys.executable, + "-c", + "from pathlib import Path; assert Path('app.py').read_text() == 'answer = 2\\n'; print('validation passed')", + ) + validate_workspace, validate_receipts = _copy_branch( + base, tmp_path / "validate-only", command, lambda _action: True + ) + denied_workspace, denied_receipts = _copy_branch( + base, + tmp_path / "try-extra-edit", + command, + lambda _action: ApprovalDecision(False, "keep the requested answer at 2"), + ) + assert (validate_receipts.path.read_bytes(), denied_receipts.path.read_bytes()) == ( + initial_receipts, + initial_receipts, + ) + + case = EvaluationCase( + final_contains="validated", + files=(FileExpectation("app.py", "answer = 2\n"),), + results=(ResultExpectation("run_command", "validation passed"),), + receipts=( + ReceiptExpectation( + "call-1", "edit_file", "ok", "edited app.py", ("app.py",) + ), + ReceiptExpectation("call-2", "run_command", "ok", "validation passed"), + ), + ) + validate_responses = ( + "unused read", + "unused edit", + json.dumps({"tool": "run_command", "argv": list(command)}), + '{"final":"validated answer = 2"}', + ) + denied_responses = ( + "unused read", + "unused edit", + '{"tool":"read_file","path":"app.py"}', + '{"tool":"edit_file","path":"app.py","old":"2","new":"3"}', + '{"final":"the extra edit was denied"}', + ) + validate_model = ScriptedCheckpointModel(validate_responses) + denied_model = ScriptedCheckpointModel(denied_responses) + + passing = run_branch( + "validate-only", + "validate without another edit", + checkpoint, + validate_model, + validate_workspace, + validate_receipts, + case, + ) + failing = run_branch( + "try-extra-edit", + "try changing the answer again", + checkpoint, + denied_model, + denied_workspace, + denied_receipts, + case, + ) + + validate_steering = "validate without another edit" + denied_steering = "try changing the answer again" + assert passing.steering == validate_steering + assert failing.steering == denied_steering + assert validate_model.calls[0][-1] == { + "role": "user", + "content": f"Operator steering:\n{validate_steering}", + } + assert denied_model.calls[0][-1] == { + "role": "user", + "content": f"Operator steering:\n{denied_steering}", + } + assert passing.reuse == validate_model.reuse + assert failing.reuse == denied_model.reuse + assert passing.reuse.reused_tokens > 0 + assert failing.reuse.reused_tokens > 0 + assert passing.report.passed + assert not failing.report.passed + assert select_branch((passing, failing), "validate-only") is passing + assert (tmp_path / "validate-only" / "app.py").read_text() == "answer = 2\n" + assert (tmp_path / "try-extra-edit" / "app.py").read_text() == "answer = 2\n" + assert validate_receipts.get("call-1").tool == "edit_file" + assert validate_receipts.get("call-2").tool == "run_command" + assert denied_receipts.get("call-1").tool == "edit_file" + assert denied_receipts.get("call-2") is None + assert denied_receipts.path.read_bytes() == initial_receipts + assert validate_receipts.path.read_bytes().startswith(initial_receipts) + visible = "\n".join( + message["content"] for call in denied_model.calls for message in call + ) + assert visible.count("keep the requested answer at 2") == 1 + + +def test_task_5_selection_requires_one_named_passing_outcome(tmp_path): + base, checkpoint = _forkable_effect(tmp_path) + workspace, receipts = _copy_branch( + base, tmp_path / "branch", (sys.executable, "-V"), lambda _action: True + ) + model = ScriptedCheckpointModel(("unused", "unused", '{"final":"done"}')) + case = EvaluationCase(final_contains="missing") + failing = run_branch( + "candidate", "finish", checkpoint, model, workspace, receipts, case + ) + + with pytest.raises(AgentError, match="exactly one"): + select_branch((failing,), "absent") + with pytest.raises(AgentError, match="exactly one"): + select_branch((failing, replace(failing)), "candidate") + with pytest.raises(AgentError, match="must pass"): + select_branch((failing,), "candidate") diff --git a/tests_refsol/test_week_4_starter_sync.py b/tests_refsol/test_week_4_starter_sync.py index c7ce0ad3..9e9b7744 100644 --- a/tests_refsol/test_week_4_starter_sync.py +++ b/tests_refsol/test_week_4_starter_sync.py @@ -1,6 +1,6 @@ # WARNING: Under review - generated by LLM. -"""Course-code guards for the cumulative Week 4 Day 1--7 starter.""" +"""Course-code guards for the cumulative Week 4 Day 1--8 starter.""" import ast import importlib.util @@ -19,6 +19,7 @@ REFSOL = ROOT / "src" / "tiny_llm_ref" / "agent" REAL_MODEL_CLI = ROOT / "agent.py" MODULES = ( + "branching", "checkpoint", "compaction", "evaluation", @@ -126,7 +127,7 @@ def test_starter_is_solution_free_and_reference_is_implemented(module): ) -def test_package_exports_match_the_published_day_7_surface(): +def test_package_exports_match_the_published_day_8_surface(): import tiny_llm.agent as starter import tiny_llm_ref.agent as refsol @@ -137,6 +138,8 @@ def test_package_exports_match_the_published_day_7_surface(): "AgentLimits", "AgentRun", "AgentStatus", + "ApprovalDecision", + "BranchOutcome", "CompactionResult", "EvaluationCase", "EvaluationCheck", @@ -144,10 +147,12 @@ def test_package_exports_match_the_published_day_7_surface(): "EffectReceipt", "FinalAction", "FileExpectation", + "KvPrefixGenerator", "ModelCheckpoint", "ReceiptStore", "ReceiptExpectation", "ResultExpectation", + "PrefixReuse", "ToolAction", "ToolPolicy", "Workspace", @@ -162,14 +167,17 @@ def test_package_exports_match_the_published_day_7_surface(): "resume_agent", "resume_with_steering", "run_agent", + "run_branch", "run_to_checkpoint", + "select_branch", } assert set(starter.__all__) == set(refsol.__all__) == expected -def test_only_day_1_through_day_7_modules_exist_in_the_starter(): +def test_only_day_1_through_day_8_modules_exist_in_the_starter(): allowed = { "__init__.py", + "branching.py", "checkpoint.py", "compaction.py", "evaluation.py", @@ -185,7 +193,7 @@ def test_only_day_1_through_day_7_modules_exist_in_the_starter(): def test_day_3_workspace_surface_remains_complete(): surface = _public_surface(STARTER, "workspace") - assert set(surface) == {"ToolPolicy", "Workspace"} + assert set(surface) == {"ApprovalDecision", "ToolPolicy", "Workspace"} assert {name for name, *_ in surface["Workspace"]} == { "available_tools", "modified_files", @@ -254,6 +262,53 @@ def test_day_7_evaluation_surface_is_complete_and_has_no_future_api(): assert future_name not in source +def test_day_8_branching_surface_is_complete_and_has_no_future_api(): + branching = _public_surface(STARTER, "branching") + assert set(branching) == { + "BranchOutcome", + "KvPrefixGenerator", + "PrefixReuse", + "run_branch", + "select_branch", + } + assert {name for name, *_ in branching["KvPrefixGenerator"]} == { + "fork", + "restore_checkpoint", + "reuse", + "save_checkpoint", + } + starter_class = next( + node + for node in _tree(STARTER, "branching").body + if isinstance(node, ast.ClassDef) and node.name == "KvPrefixGenerator" + ) + reference_class = next( + node + for node in _tree(REFSOL, "branching").body + if isinstance(node, ast.ClassDef) and node.name == "KvPrefixGenerator" + ) + starter_call = next( + node + for node in starter_class.body + if isinstance(node, ast.FunctionDef) and node.name == "__call__" + ) + reference_call = next( + node + for node in reference_class.body + if isinstance(node, ast.FunctionDef) and node.name == "__call__" + ) + assert (ast.unparse(starter_call.args), ast.unparse(starter_call.returns)) == ( + ast.unparse(reference_call.args), + ast.unparse(reference_call.returns), + ) + assert _is_stub(starter_call.body) and not _is_stub(reference_call.body) + source = "\n".join( + path.read_text(encoding="utf-8") for path in STARTER.glob("*.py") + ) + for future_name in ("Session", "reconcile", "radix", "ArtifactStore"): + assert future_name not in source + + def test_removed_catalog_hash_is_not_exported_or_declared(): for root in (STARTER, REFSOL): source = (root / "protocol.py").read_text(encoding="utf-8")